mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
02e099fff8
- NeuronGraphRunner: the CudaGraphRunner contract over torch.compile(dynamic=False), one graph per batch size (DEEPINFER_ENFORCE_EAGER=0, backend DEEPINFER_NEURON_COMPILE_BACKEND). Layer cache views are re-sliced from their root allocations inside the graph, so AOT never regenerates aliased mutated inputs with as_strided. - Decode MoE with fixed shapes: packed all-gather, dense local experts, fp32 reduce-scatter (neuron_ops/static_moe.py, deep_gemm.m_grouped_fp8_fp4_gemm_nt_shared_a). - Graph-break and compile fixes: hide torch.cuda during compile, torch.device equality for FX codegen, hash-based force_random routing, Engram hash/fetch on the host (exact int64), env vars resolved before tracing, Movable annotation, no .item()/nonzero/as_strided writes on the decode path, KV stores route padding slots to a scratch block kept past the pool, native-dtype gather under compile. - fp8: trn2 reads float8_e4m3fn as IEEE e4m3 (max 240). Activations/KV quantize to 224 (DEEPINFER_NEURON_FP8_MAX), checkpoint fp8 weights are rescaled at load (bytes / 2, ue8m0 + 1), and neuronx-cc gets --experimental-unsafe-fp8e4m3fn-as-fp8e4m3 via --internal-hlo2tensorizer-options. The uint8-carrier backend (neuron_ops/neuron_compile.py, "deepinfer_neuron") stays as a fallback. - Tests: tests/neuron/compile; run_all.sh runs kernel unit tests at the e4m3fn range (448) and the end-to-end stages at 224. Docs: vllm-neuron PR #71 analysis and compile notes. Host: tiny model 9/9 vs golden (all prompts) eager, compiled path (aot_eager) matches eager; CPU suite green. Device: the whole decode graph still crashes neuronx-cc (NCC_IINL902, being bisected). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
"""Mirrors ``tilelang_ops/asc/gather.py`` (``gather_tl``): bytewise row gather with zero-fill."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
import torch
|
|
|
|
|
|
def _gather_rows(x: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
|
|
"""``x[index]`` for a byte table ``x [n, D]``.
|
|
|
|
On the neuron device, index_select / advanced indexing of an int8 table allocates several times the table
|
|
(+4.9 GiB for the 1.3 GB V4.1 embedding, which ran the full model out of HBM), while F.embedding through a
|
|
4- or 2-byte view copies only the selected rows and is bit-exact (verified on arbitrary bytes, NaN patterns
|
|
included).
|
|
"""
|
|
if x.is_contiguous():
|
|
for elem, dtype in ((4, torch.int32), (2, torch.float16)):
|
|
if x.shape[1] % elem == 0 and (torch.compiler.is_compiling() or x.storage_offset() % elem == 0):
|
|
return torch.nn.functional.embedding(index, x.view(dtype)).view(torch.int8)
|
|
return torch.index_select(x, 0, index)
|
|
|
|
|
|
def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: int) -> Callable[..., None]:
|
|
"""Kernel args: ``x [n, D]`` (row stride ``x_stride``), ``index [m]``, ``y [m, D]`` contiguous.
|
|
|
|
``y[i] = x[index[i]]`` if ``0 <= index[i] < n`` else 0; written in place.
|
|
"""
|
|
|
|
def _gather(x: torch.Tensor, index: torch.Tensor, y: torch.Tensor) -> None:
|
|
n = x.shape[0]
|
|
m = index.shape[0]
|
|
if m == 0:
|
|
return
|
|
if n == 0:
|
|
y.zero_()
|
|
return
|
|
valid = (index >= 0) & (index < n)
|
|
safe = torch.where(valid, index, torch.zeros_like(index))
|
|
rows = _gather_rows(x[:, :D], safe)
|
|
y.copy_(torch.where(valid.view(m, 1), rows, torch.zeros((), dtype=rows.dtype, device=rows.device)))
|
|
|
|
return _gather
|