From 02e099fff8d7484b59770309eebf7bf91fbe0f7c Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 22 Oct 2026 11:59:37 +0800 Subject: [PATCH] Neuron: static decode graphs, static-shape MoE, trn2 e4m3 fp8 (issues #9, #12, #24) - 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 --- deepinfer/environ.py | 5 + deepinfer/model_executor/layers/linear.py | 2 + .../layers/tilelang_ops/gather.py | 8 + .../tilelang_ops/neuron/fused_transform.py | 26 +- .../layers/tilelang_ops/neuron/gather.py | 2 +- .../layers/tilelang_ops/neuron/kv_quant.py | 1 + .../layers/tilelang_ops/neuron/norm.py | 1 + .../tilelang_ops/neuron/store_kv_cache.py | 19 +- .../layers/tilelang_ops/norm.py | 12 +- deepinfer/model_executor/models/deepseek.py | 8 +- .../model_executor/models/v4_modules/mlp.py | 4 + deepinfer/model_executor/utils.py | 5 +- deepinfer/neuron_ops/fp8.py | 111 +++++++ deepinfer/neuron_ops/neuron_compile.py | 289 +++++++++++++++++ deepinfer/neuron_ops/static_moe.py | 127 ++++++++ deepinfer/neuron_ops/transfer_to_neuron.py | 28 +- .../vendor/deep_ep/buffers/engram.py | 3 + .../neuron_ops/vendor/deep_ep/utils/math.py | 8 +- .../neuron_ops/vendor/deep_gemm/__init__.py | 1 + .../neuron_ops/vendor/deep_gemm/attention.py | 4 +- deepinfer/neuron_ops/vendor/deep_gemm/gemm.py | 29 ++ .../neuron_ops/vendor/tile_kernels/config.py | 6 +- .../tile_kernels/engram/engram_hash_kernel.py | 8 +- .../moe/moe_topk_gate_forward_kernel.py | 21 +- .../vendor/tile_kernels/quant/_core.py | 6 +- .../vendor/tile_kernels/torch/cast.py | 6 +- .../neuron_ops/vendor/tilelang/language.py | 6 +- deepinfer/platform.py | 3 + deepinfer/worker/cache_engine.py | 12 +- deepinfer/worker/model_runner.py | 6 +- deepinfer/worker/neuron_graph_runner.py | 158 ++++++++++ docs/neuron/vllm_neuron_compile_notes.md | 281 +++++++++++++++++ docs/neuron/vllm_neuron_pr71_analysis.md | 292 ++++++++++++++++++ tests/neuron/compile/test_fp8_e4m3.py | 41 +++ tests/neuron/compile/test_fp8_legalize.py | 141 +++++++++ tests/neuron/run_all.sh | 7 +- 36 files changed, 1651 insertions(+), 36 deletions(-) create mode 100644 deepinfer/neuron_ops/fp8.py create mode 100644 deepinfer/neuron_ops/neuron_compile.py create mode 100644 deepinfer/neuron_ops/static_moe.py create mode 100644 deepinfer/worker/neuron_graph_runner.py create mode 100644 docs/neuron/vllm_neuron_compile_notes.md create mode 100644 docs/neuron/vllm_neuron_pr71_analysis.md create mode 100644 tests/neuron/compile/test_fp8_e4m3.py create mode 100644 tests/neuron/compile/test_fp8_legalize.py diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 84d8fa6..d5ebc02 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -103,6 +103,11 @@ class envs: DEEPINFER_NEURON_DEVICE = EnvVar.str(default="neuron") # Neuron only, debug: "1" records float8 cat / advanced-indexing call sites (summary at exit), "raise" raises. DEEPINFER_NEURON_FP8_GUARD = EnvVar.str(default="") + # Neuron only: torch.compile backend for the decode graphs (DEEPINFER_ENFORCE_EAGER=0): "neuron"; "deepinfer_neuron" + # carries fp8 as uint8 (fallback, deepinfer/neuron_ops/neuron_compile.py); "eager"/"aot_eager" to debug on the host. + DEEPINFER_NEURON_COMPILE_BACKEND = EnvVar.str(default="neuron") + # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break). + DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool() # -- Runtime -- DEEPINFER_COMPUTE_SUM_SMALLER_PROBS = EnvVar.bool() diff --git a/deepinfer/model_executor/layers/linear.py b/deepinfer/model_executor/layers/linear.py index 87ce96f..8c12592 100644 --- a/deepinfer/model_executor/layers/linear.py +++ b/deepinfer/model_executor/layers/linear.py @@ -104,6 +104,8 @@ class ReplicatedLinear(torch.nn.Module): f"{output.shape=}, {output.stride()=}" ) m = output.shape[:-1].numel() + if PLATFORM.is_neuron: + return output.view(m, self.output_size) # a view, not as_strided: the GEMM writes through it under torch.compile return output.as_strided((m, self.output_size), (output.stride(-2), 1)) def _align_output(self, output: torch.Tensor) -> torch.Tensor: diff --git a/deepinfer/model_executor/layers/tilelang_ops/gather.py b/deepinfer/model_executor/layers/tilelang_ops/gather.py index 6a74841..e6699a9 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/gather.py +++ b/deepinfer/model_executor/layers/tilelang_ops/gather.py @@ -4,6 +4,7 @@ from typing import Any import pytest import torch +from deepinfer import PLATFORM from deepinfer.model_executor.layers.tilelang_ops import ops @@ -53,6 +54,13 @@ def gather( if out is None: out = torch.empty(out_shape, dtype=x.dtype, device=x.device) + if PLATFORM.is_neuron and torch.compiler.is_compiling(): + # neuronx-cc crashes on a dtype-changing view feeding a gather; gather the native dtype instead. + valid = (index >= 0) & (index < x.shape[0]) + rows = torch.index_select(x, 0, torch.where(valid, index, torch.zeros_like(index))) + mask = valid.view(-1, *([1] * (x.ndim - 1))) + out.copy_(torch.where(mask, rows, torch.zeros((), dtype=rows.dtype, device=rows.device))) + return out x = _flat_dense_int8(x) y = _flat_dense_int8(out) kernel = ops.gather_tl(x.shape[1], x.dtype, index.dtype, x.stride(0)) diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py index 335aaac..001bdff 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py @@ -17,18 +17,20 @@ order: index caches; slot ``-1`` is skipped. """ +import struct from collections.abc import Callable from typing import Any import torch from ..common import FLOAT4_E2M1_MAX, FLOAT8E4M3_MAX, FLOAT8E4M3_MIN_POSITIVE, FP4_AMAX_MARGIN, FP8_AMAX_MARGIN +from deepinfer.neuron_ops.fp8 import FP8_MAX as FLOAT8E4M3_MAX # trn2 hardware e4m3 range (#24) from .rstd_1ulp import _fma, make_rstd_impl def _f32(value: float) -> float: """Round a Python float to the nearest FP32 value.""" - return torch.tensor(value, dtype=torch.float32).item() + return struct.unpack("f", struct.pack("f", value))[0] def _lane_tree_sum(v: torch.Tensor) -> torch.Tensor: @@ -101,6 +103,24 @@ def _from_bits(b: torch.Tensor) -> torch.Tensor: return b.to(torch.int32).view(torch.float32) +def _store_static(kv_cache: torch.Tensor, slots: torch.Tensor, data: torch.Tensor, sf_bytes: torch.Tensor, + block_size: int, merged_store: bool, payload_bytes: int, data_bytes: int, sf_row_bytes: int) -> None: + """Graph-friendly cache store: every row is written, padding rows (slot -1) to the scratch block the Neuron cache + engine keeps as the last page of each cache view; indexes a page view, no as_strided.""" + nb, bs, row_bytes = kv_cache.shape + assert kv_cache.stride(1) == row_bytes, "page bytes must be contiguous" + pages = kv_cache.flatten(1) # [nb, bs * row_bytes] + valid = slots >= 0 + blk = torch.where(valid, torch.div(slots, block_size, rounding_mode="floor"), nb - 1).unsqueeze(-1) + off = torch.where(valid, torch.remainder(slots, block_size), 0).unsqueeze(-1) + if merged_store: + cols = off * row_bytes + torch.arange(payload_bytes, device=slots.device) + pages[blk, cols] = torch.cat([data, sf_bytes], dim=-1) + return + pages[blk, off * data_bytes + torch.arange(data_bytes, device=slots.device)] = data + pages[blk, bs * data_bytes + off * sf_row_bytes + torch.arange(sf_row_bytes, device=slots.device)] = sf_bytes + + def fused_norm_rope_hadamard_cast_store_kvcache_tl(num_heads: int, head_dim: int, **kwargs: Any) -> Callable[..., None]: """Translate the shared public ABI into static descriptors (same checks as the Ascend factory).""" do_norm = kwargs["do_norm"] @@ -334,6 +354,10 @@ def fused_norm_rope_hadamard_cast_store_kvcache_tl(num_heads: int, head_dim: int nb, bs, _ = kv_cache.shape page_stride, token_stride = kv_cache.stride(0), kv_cache.stride(1) slots = slot_mapping.reshape(-1)[:rows] + if torch.compiler.is_compiling(): + _store_static(kv_cache, slots, data, sf_bytes, block_size, merged_store, payload_bytes, data_bytes, + sf_storage * sf_itemsize) + return valid = (slots >= 0).nonzero().squeeze(1) # NEURON: host sync (skip slot -1) if valid.numel() == 0: return diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py index 02fa4d0..5024555 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/gather.py @@ -15,7 +15,7 @@ def _gather_rows(x: torch.Tensor, index: torch.Tensor) -> torch.Tensor: """ if x.is_contiguous(): for elem, dtype in ((4, torch.int32), (2, torch.float16)): - if x.shape[1] % elem == 0 and x.storage_offset() % elem == 0: + 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) diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/kv_quant.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/kv_quant.py index 5f35cf2..92ffb34 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/kv_quant.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/kv_quant.py @@ -34,6 +34,7 @@ from ..common import ( FP4_AMAX_MARGIN, FP8_AMAX_MARGIN, ) +from deepinfer.neuron_ops.fp8 import FP8_MAX as FLOAT8E4M3_MAX # trn2 hardware e4m3 range (#24) _RECIPROCAL_BITS = tuple(struct.unpack(" None: - self._value: T | None = value + self._value = value def move(self) -> T: if self._value is None: diff --git a/deepinfer/neuron_ops/fp8.py b/deepinfer/neuron_ops/fp8.py new file mode 100644 index 0000000..4fd1bd1 --- /dev/null +++ b/deepinfer/neuron_ops/fp8.py @@ -0,0 +1,111 @@ +"""fp8 on trn2: keep every ``float8_e4m3fn`` value inside the range of the hardware's e4m3 (issue #24). + +trn2 has no e4m3fn. With ``--experimental-unsafe-fp8e4m3fn-as-fp8e4m3`` the compiler accepts ``float8_e4m3fn`` +tensors and reads their bits as IEEE-style e4m3 (bias 7, max 240, exponent field 15 = inf/NaN). Both formats decode +identically whenever the exponent field is below 15, i.e. for |x| < 256. So on Neuron: + +* activations / KV records quantize to ``FP8_MAX`` (224, as vllm-neuron#71) instead of 448. With power-of-two + (ue8m0) or ``amax / max`` scales the scale just doubles, so dequantized values are unchanged except in the lowest + subnormal binade; +* checkpoint fp8 weights (full ±448 range) are rescaled once at load: values halved, scales doubled + (``rescale_fp8_weights``), exact except the lowest subnormal bit. + +``DEEPINFER_NEURON_FP8_MAX=448`` restores the e4m3fn range (host-mode debugging only: the device cannot decode it). +""" + +from __future__ import annotations + +import os +import shlex + +import torch + +FP8_MAX = float(os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224")) +_HW_E4M3_FLAG = "--experimental-unsafe-fp8e4m3fn-as-fp8e4m3" + + +def configure_compiler() -> None: + """Let neuronx-cc (eager and torch.compile) take float8_e4m3fn tensors as the hardware e4m3.""" + if FP8_MAX > 240.0: + return + os.environ.setdefault("UNSAFE_FP8FNCAST", "1") + flags = os.environ.get("NEURON_CC_FLAGS", "") + if _HW_E4M3_FLAG not in flags: + # The option belongs to hlo2tensorizer; neuronx-cc rejects it as a top-level flag (NCC_EARG002). All + # hlo2tensorizer options must share one --internal-hlo2tensorizer-options (a second one replaces the first), + # so append to an existing one if the user set it. + key = "--internal-hlo2tensorizer-options=" + if key in flags: + head, _, rest = flags.partition(key) + value, sep, tail = rest.partition(" --") + value = value.strip("'\"") + flags = f"{head}{key}'{value} {_HW_E4M3_FLAG}'" + (f" --{tail}" if sep else "") + else: + flags = f"{flags} {key}{_HW_E4M3_FLAG}".strip() + os.environ["NEURON_CC_FLAGS"] = flags + + +def halve_e4m3fn_bytes(b: torch.Tensor) -> torch.Tensor: + """``e4m3fn(x / 2)`` from the bits of ``x`` (uint8), round-to-nearest-even; integer ops only, any device.""" + v = b.to(torch.int16) + sign = v & 0x80 + mag = v & 0x7F + exp = mag >> 3 + man = mag & 7 + # exponent >= 2: decrement the exponent. exponent 1 -> subnormal (8 + m) / 2; subnormal m -> m / 2 (RNE). + sub_in = torch.where(exp == 1, man + 8, man) + half = sub_in >> 1 + rne = half + ((sub_in & 1) & (half & 1)) + out = torch.where(exp >= 2, mag - 8, torch.where(exp <= 1, rne, mag)) + out = torch.where(mag == 0x7F, mag, out) # NaN stays NaN + return (out | sign).to(torch.uint8) + + +_CHUNK_BYTES = 1 << 26 + + +def _halve_inplace(data: torch.Tensor) -> None: + flat = data.view(-1) + for start in range(0, flat.numel(), _CHUNK_BYTES): # bounded int16 temporaries (Engram tables are ~200 GB) + chunk = flat[start : start + _CHUNK_BYTES] + chunk.copy_(halve_e4m3fn_bytes(chunk)) + + +def _double_scale(sf: torch.Tensor) -> None: + if sf.dtype in (torch.float32, torch.bfloat16, torch.float16): + sf.mul_(2) + return + # ue8m0 exponents (float8_e8m0fnu, or bytes packed into int words): exponent + 1. + sf_bytes = sf.view(torch.uint8) + assert not bool((sf_bytes == 255).any()), "ue8m0 scale overflow" + sf_bytes.add_(1) + + +def rescale_fp8_weights(model: torch.nn.Module) -> int: + """Halve every fp8 weight (and double its scale) so it fits the hardware e4m3. Returns the count rescaled. + + Run once, on the checkpoint parameters right after loading (before runtime adaptation merges or repacks them). + Parameters sharing storage (tied weights) are rescaled once. + """ + if FP8_MAX > 240.0: + return 0 + by_storage: dict[int, torch.nn.Parameter] = {} + for p in model.parameters(): + if p.dtype != torch.float8_e4m3fn or p.numel() == 0: + continue + ptr = p.untyped_storage().data_ptr() + if getattr(p, "sf", None) is not None or ptr not in by_storage: + by_storage[ptr] = p # prefer the copy that carries the scale (tied aliases may not) + count = 0 + for p in by_storage.values(): + sf = getattr(p, "sf", None) + assert sf is not None, "fp8 weight without a scale cannot be rescaled" + assert p.is_contiguous() + _halve_inplace(p.data.view(torch.uint8)) + _double_scale(sf.data if isinstance(sf, torch.nn.Parameter) else sf) + count += 1 + return count + + +def describe() -> str: + return f"fp8 max {FP8_MAX:g}, NEURON_CC_FLAGS={shlex.quote(os.environ.get('NEURON_CC_FLAGS', ''))}" diff --git a/deepinfer/neuron_ops/neuron_compile.py b/deepinfer/neuron_ops/neuron_compile.py new file mode 100644 index 0000000..765bb49 --- /dev/null +++ b/deepinfer/neuron_ops/neuron_compile.py @@ -0,0 +1,289 @@ +"""``torch.compile`` backend for the Neuron decode graphs: torch-neuronx's ``neuron`` backend plus fp8 legalization. + +trn2 has no ``F8E4M3FN`` (neuronx-cc rejects any graph holding one; issue #24), but the model carries fp8 weights, +fp8 KV records and fp8-quantized activations. This backend makes every fp8 value a ``uint8`` before lowering: + +* fp8 graph inputs are passed as zero-copy ``uint8`` views; +* a pass over the Dynamo graph, ahead of AOT, carries every fp8 value as ``uint8``. + Data movement (views, slices, cat, index, copy, where, ...) is byte-exact as is; ``fp8 -> float`` becomes an + exact 256-entry table lookup and ``float -> fp8`` an arithmetic round-to-nearest-even encode with torch's + saturating semantics (|x| > 448 and inf -> 448, NaN -> 0x7F). Negative zero encodes as +0 (same value). + +Usage: ``torch.compile(fn, backend="deepinfer_neuron")`` (registered on import). +""" + +from __future__ import annotations + +import os +import operator +from collections.abc import Callable, Sequence +from typing import Any + +import torch +import torch.fx +from torch._dynamo import register_backend +from torch._subclasses.fake_tensor import unset_fake_temporarily + +FP8 = torch.float8_e4m3fn +U8 = torch.uint8 + + +def _is_fp8(t: Any) -> bool: + return isinstance(t, torch.Tensor) and t.dtype == FP8 + + +def _encode_scalar(value: float) -> int: + with unset_fake_temporarily(): + return int(torch.tensor(float(value)).to(FP8).view(U8).item()) + + +# --------------------------------------------------------------------------- # +# Reference codec (traced into the graph) +# --------------------------------------------------------------------------- # + + +def _decode(x_u8: torch.Tensor) -> torch.Tensor: + """Exact e4m3fn bits -> fp32, elementwise (no table, no bitcast): (8 + m) * 2^(e - 10), subnormals m * 2^-9.""" + b = x_u8.to(torch.int32) + e = (b >> 3) & 15 + m = (b & 7).to(torch.float32) + p = torch.full_like(m, 2.0**-10) + for bit, factor in ((1, 2.0), (2, 4.0), (4, 16.0), (8, 256.0)): + p = p * torch.where((e & bit) != 0, factor, 1.0) + v = torch.where(e == 0, m * 2.0**-9, (m + 8.0) * p) + v = torch.where(b >= 128, -v, v) + return torch.where((b & 127) == 127, torch.full_like(v, float("nan")), v) + + +def _decode_as(dtype: torch.dtype) -> Callable[[torch.Tensor], torch.Tensor]: + def decode(x_u8: torch.Tensor) -> torch.Tensor: + return _decode(x_u8).to(dtype) + + return decode + + +def _rne(m: torch.Tensor) -> torch.Tensor: + f = torch.floor(m) + d = m - f + odd = torch.remainder(f, 2.0) == 1.0 + return f + ((d > 0.5) | ((d == 0.5) & odd)).to(m.dtype) + + +def _encode(x: torch.Tensor) -> torch.Tensor: + """fp32 -> e4m3fn bits, round-to-nearest-even, saturating like torch's cast (no table, no bitcast).""" + xf = x.float() + nan = xf != xf + neg = xf < 0 + a = torch.clamp(torch.abs(torch.where(nan, torch.zeros_like(xf), xf)), max=448.0) + # e = floor(log2(a)) and 2^(3 - e) for a in [2^-6, 448], by comparisons: exact, no log/exp on the device. + e = torch.full_like(a, -6.0) + scale = torch.full_like(a, 2.0**9) + for k in range(-5, 9): + ge = a >= 2.0**k + e = e + ge.to(a.dtype) + scale = scale * torch.where(ge, 0.5, 1.0) + r = _rne(a * scale) # [8, 16]; 16 carries into the next exponent below + code_n = torch.clamp((e + 7.0) * 8.0 + (r - 8.0), max=126.0) + code_s = _rne(a * 512.0) # subnormals: a / 2^-9, 8 = smallest normal + code = torch.where(a >= 2.0**-6, code_n, code_s) + neg.to(a.dtype) * 128.0 + code = torch.where(nan, torch.full_like(code, 127.0), code) + return code.to(torch.int32).to(U8) + + +# --------------------------------------------------------------------------- # +# Graph pass (Dynamo / torch level, ahead of AOT: the codecs inserted here are traced with the graph's devices) +# --------------------------------------------------------------------------- # + +_CREATE = {torch.empty, torch.zeros, torch.full, torch.empty_like, torch.zeros_like, torch.full_like, torch.tensor, + torch.empty_strided, torch.ones, torch.ones_like} +_CREATE_METHODS = {"new_empty", "new_zeros", "new_full", "new_ones", "new_tensor", "new_empty_strided"} +_TO_METHODS = {"float": torch.float32, "bfloat16": torch.bfloat16, "half": torch.float16, "double": torch.float64} +_FILL_ARG = {torch.full: 1, torch.full_like: 1, "new_full": 1, "fill_": 0, "masked_fill": 1, "masked_fill_": 1} + + +def _ev(node: Any) -> Any: + return node.meta.get("example_value") if isinstance(node, torch.fx.Node) else node + + +def _dtype_arg(node: torch.fx.Node, positional: int | None) -> Any: + if "dtype" in node.kwargs: + return node.kwargs["dtype"] + if positional is not None and len(node.args) > positional and isinstance(node.args[positional], torch.dtype): + return node.args[positional] + return None + + +def legalize_fp8(gm: torch.fx.GraphModule) -> bool: + """Carry every fp8 value of a Dynamo graph as uint8 bytes. Returns True if the graph changed. + + Only dtype-sensitive operations are rewritten (bitcasts, conversions, creation with dtype=fp8, implicit + conversions in copy_/setitem, fp8 fill values); everything else on fp8 (views, slicing, cat, indexing, where + between fp8 tensors, ...) moves bytes and runs unchanged on the uint8 carrier. + """ + graph = gm.graph + carriers: set[torch.fx.Node] = set() # nodes whose fp8 value is now a uint8 tensor + changed = False + + def is_fp8(a: Any) -> bool: + return a in carriers or _is_fp8(_ev(a)) + + for node in list(graph.nodes): + if node.op not in ("call_function", "call_method"): + continue + tgt, args = node.target, node.args + out_fp8 = _is_fp8(_ev(node)) or ( + isinstance(_ev(node), (tuple, list)) and any(_is_fp8(v) for v in _ev(node)) + ) + with graph.inserting_before(node): + # Bitcasts: x.view(dtype) + if node.op == "call_method" and tgt == "view" and len(args) == 2 and isinstance(args[1], torch.dtype): + src, dst = args + # Always a fresh node (a same-dtype view is a no-op): a node is a carrier or a plain uint8 tensor, + # never both, even when the bitcast itself becomes an identity. + if dst == FP8: + repl = graph.call_method("view", (src, U8)) + _swap(node, repl, carriers, carrier=True) + elif is_fp8(src): + repl = graph.call_method("view", (src, dst)) + _swap(node, repl, carriers, carrier=False) + changed = True + continue + # Conversions: x.to(dtype) / x.float() / x.type(dtype) ... + dst = None + if node.op == "call_method" and tgt in _TO_METHODS: + dst = _TO_METHODS[tgt] + elif node.op == "call_method" and tgt in ("to", "type"): + dst = _dtype_arg(node, 1) + if dst is None and len(args) > 2 and isinstance(args[2], torch.dtype): + dst = args[2] + if dst is not None: + src = args[0] + if is_fp8(src) and dst != FP8: + repl = graph.call_function(_decode_as(dst), (src,)) + _swap(node, repl, carriers, carrier=False) + elif dst == FP8 and not is_fp8(src): + repl = graph.call_function(_encode, (src,)) + _swap(node, repl, carriers, carrier=True) + else: # fp8 -> fp8 (copy / device move): bytes + _set_dtype(node, U8) + carriers.add(node) + changed = True + continue + # Creation with dtype=fp8 + if (node.op == "call_function" and tgt in _CREATE) or (node.op == "call_method" and tgt in _CREATE_METHODS): + if _dtype_arg(node, None) == FP8 or (out_fp8 and "dtype" not in node.kwargs and is_fp8(args[0] if args else None)): + if "dtype" in node.kwargs: + node.kwargs = {**node.kwargs, "dtype": U8} + _encode_fill(node) + carriers.add(node) + changed = True + continue + # Implicit conversions on copy_ / setitem into an fp8 tensor (or out of one). + if (node.op == "call_method" and tgt == "copy_") or (node.op == "call_function" and tgt is operator.setitem): + dst_t, src = args[0], args[-1] + if isinstance(src, torch.fx.Node) and is_fp8(dst_t) != is_fp8(src): + conv = graph.call_function(_encode, (src,)) if is_fp8(dst_t) else graph.call_function(_decode_as(_ev(dst_t).dtype), (src,)) + node.args = (*args[:-1], conv) + changed = True + elif is_fp8(dst_t) and isinstance(src, (int, float)): + node.args = (*args[:-1], _encode_scalar(src)) + changed = True + if is_fp8(dst_t): + carriers.add(node) + continue + if node.op == "call_method" and tgt in ("fill_", "masked_fill", "masked_fill_") and is_fp8(args[0]): + _encode_fill(node) + carriers.add(node) + changed = True + continue + if out_fp8 or any(is_fp8(a) for a in node.all_input_nodes): + if out_fp8: + carriers.add(node) + changed = True + + if changed: + graph.lint() + gm.recompile() + return changed + + +def _swap(node: torch.fx.Node, repl: torch.fx.Node, carriers: set, carrier: bool) -> None: + node.replace_all_uses_with(repl) + if carrier: + carriers.add(repl) + node.graph.erase_node(node) + + +def _set_dtype(node: torch.fx.Node, dtype: torch.dtype) -> None: + if "dtype" in node.kwargs: + node.kwargs = {**node.kwargs, "dtype": dtype} + return + node.args = tuple(dtype if a == FP8 else a for a in node.args) + + +def _encode_fill(node: torch.fx.Node) -> None: + key = node.target + if key not in _FILL_ARG: + if key in (torch.ones, torch.ones_like) or key == "new_ones": + raise NotImplementedError(f"fp8 legalization: {key} with fp8") + return + i = _FILL_ARG[key] + (1 if node.op == "call_method" else 0) + if i < len(node.args) and isinstance(node.args[i], (int, float)): + a = list(node.args) + a[i] = _encode_scalar(a[i]) + node.args = tuple(a) + elif "fill_value" in node.kwargs and isinstance(node.kwargs["fill_value"], (int, float)): + node.kwargs = {**node.kwargs, "fill_value": _encode_scalar(node.kwargs["fill_value"])} + elif "value" in node.kwargs and isinstance(node.kwargs["value"], (int, float)): + node.kwargs = {**node.kwargs, "value": _encode_scalar(node.kwargs["value"])} + + +# --------------------------------------------------------------------------- # +# Backend +# --------------------------------------------------------------------------- # + + +def _fp8_inputs_as_bytes(gm: torch.fx.GraphModule, example_inputs: Sequence[Any]) -> tuple[list[Any], list[int]]: + placeholders = [n for n in gm.graph.nodes if n.op == "placeholder"] + positions = [i for i, x in enumerate(example_inputs) if _is_fp8(x)] + if not positions: + return list(example_inputs), [] + first_body = next(n for n in gm.graph.nodes if n.op != "placeholder") + for i in positions: + ph = placeholders[i] + with gm.graph.inserting_before(first_body): + as_fp8 = gm.graph.call_method("view", (ph, FP8)) + ph.replace_all_uses_with(as_fp8, delete_user_cb=lambda user: user is not as_fp8) + if "example_value" in ph.meta: + fake = ph.meta["example_value"] + as_fp8.meta["example_value"] = fake + with fake.fake_mode: # same fake mode, so the static shape (and its source) carry over + ph.meta["example_value"] = fake.view(U8) + gm.recompile() + # Fresh (non-view) uint8 fakes: a view of the fp8 input reaches the lowering with symbolic sizes. + fake_mode = next(ph.meta["example_value"].fake_mode for ph in placeholders if "example_value" in ph.meta) + + def as_bytes(x: torch.Tensor) -> torch.Tensor: + with fake_mode: + return torch.empty_strided(x.shape, x.stride(), dtype=U8, device=x.device) + + inputs = [as_bytes(x) if i in positions else x for i, x in enumerate(example_inputs)] + return inputs, positions + + +@register_backend +def deepinfer_neuron(gm: torch.fx.GraphModule, example_inputs: Sequence[Any], **kwargs: Any) -> Callable: + from torch_neuronx.neuron_dynamo_backend.backend import neuron_backend + + inputs, positions = _fp8_inputs_as_bytes(gm, example_inputs) + if legalize_fp8(gm) and os.environ.get("DEEPINFER_FP8_LEGALIZE_DEBUG"): + print(gm.graph, flush=True) + compiled = neuron_backend(gm, inputs, **kwargs) + if not positions: + return compiled + pos = set(positions) + + def run(*args: Any) -> Any: + return compiled(*[a.view(U8) if i in pos else a for i, a in enumerate(args)]) + + return run diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py new file mode 100644 index 0000000..045072d --- /dev/null +++ b/deepinfer/neuron_ops/static_moe.py @@ -0,0 +1,127 @@ +"""Static-shape expert parallelism for the Neuron decode graphs (issue #12). + +The dispatch -> m_grouped GEMM -> combine path sizes everything by data (how many tokens each rank receives, the +psum layout of each expert), which costs host syncs and recompiles. Here every shape is fixed by the batch size: + + 1. pad the local tokens to ``max_num_decoding_tokens`` rows and all-gather tokens, routing and gate weights + over the EP group (one packed uint8 collective); + 2. run each local expert densely over all gathered rows, the gate weight of non-routed rows being 0; + 3. reduce-scatter the fp32 sum back to the owning rank and add the shared expert. + +Decode batches are small, so computing every local expert over the gathered rows costs little next to the +weight traffic, and the step traces into one graph with no data-dependent shapes. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import deep_gemm +import torch +import torch.distributed._functional_collectives as funcol +from tile_kernels.quant import swiglu_forward_and_per_token_cast + +from deepinfer.model_executor.layers.linear import get_fp8_cast_params, get_fp8_casting_granularity + +if TYPE_CHECKING: + from deepinfer.model_executor.models.v4_modules.mlp import DeepseekMoE + +# gloo (host mode) has no reduce-scatter. +_HOST = os.environ.get("DEEPINFER_NEURON_DEVICE", "neuron") == "cpu" +_ALIGN = 16 + + +def _as_rows(t: torch.Tensor, rows: int) -> torch.Tensor: + return t.contiguous().view(rows, -1).view(torch.uint8) + + +def _pack(parts: list[torch.Tensor], cap: int) -> tuple[torch.Tensor, list[tuple[int, int]]]: + """Byte rows ``[cap, W]``; each part starts at a 16-byte aligned column. Rows past the tokens are 0xFF so the + padded routing reads as -1 (int32) / NaN (fp32), i.e. selects no expert.""" + num_tokens = parts[0].shape[0] + spans, cols, width = [], [], 0 + for part in parts: + b = _as_rows(part, num_tokens) + spans.append((width, b.shape[1])) + pad = -b.shape[1] % _ALIGN + cols.append(torch.nn.functional.pad(b, (0, pad)) if pad else b) + width += b.shape[1] + pad + packed = torch.cat(cols, dim=1) + if cap > num_tokens: + packed = torch.cat([packed, packed.new_full((cap - num_tokens, width), 0xFF)], dim=0) + return packed, spans + + +def _unpack(packed: torch.Tensor, span: tuple[int, int], like: torch.Tensor) -> torch.Tensor: + start, nbytes = span + rows = packed.shape[0] + return packed[:, start : start + nbytes].contiguous().view(like.dtype).view(rows, *like.shape[1:]) + + +def _all_gather(t: torch.Tensor, group: torch.distributed.ProcessGroup, size: int) -> torch.Tensor: + if size == 1: + return t + return funcol.wait_tensor(funcol.all_gather_tensor(t, 0, group)) + + +def _reduce_scatter(t: torch.Tensor, group: torch.distributed.ProcessGroup, size: int, rank: int) -> torch.Tensor: + if size == 1: + return t + if _HOST: + n = t.shape[0] // size + return funcol.wait_tensor(funcol.all_reduce(t, "sum", group)).narrow(0, rank * n, n) + return funcol.wait_tensor(funcol.reduce_scatter_tensor(t, "sum", 0, group)) + + +def static_moe_forward( + moe: DeepseekMoE, + hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, +) -> torch.Tensor: + group, size, rank = moe.ep_group, moe.ep_size, moe.ep_rank + num_local = moe.num_local_physical_experts + x_parts = list(hidden_states) if moe.use_fp8_act else [hidden_states] + num_tokens = x_parts[0].shape[0] + cap = max(moe.max_num_decoding_tokens, num_tokens) + + parts = x_parts + [topk_idx.to(torch.int32), topk_weights.to(torch.float32)] + packed, spans = _pack(parts, cap) + gathered = _all_gather(packed, group, size) # [size * cap, W] + g_parts = [_unpack(gathered, span, part) for span, part in zip(spans, parts)] + g_x = tuple(g_parts[:-2]) if moe.use_fp8_act else g_parts[0] + g_idx, g_w = g_parts[-2], g_parts[-1] + num_rows = g_idx.shape[0] + + recipe_a = (1, get_fp8_casting_granularity()) + recipe_b = (1, 32 if moe.use_fp4_weight else get_fp8_casting_granularity()) + swiglu_kwargs = get_fp8_cast_params() if moe.use_fp8_act else {"fmt": "bf16", "num_per_channels": None} + gate_up_w = moe._weight_and_sf(moe.gate_up_proj_weight) + down_w = moe._weight_and_sf(moe.down_proj_weight) + + acc = torch.zeros((num_rows, moe.hidden_size), dtype=torch.float32, device=g_idx.device) + for e in range(num_local): + hit = g_idx == rank * num_local + e + w_e = torch.where(hit, g_w, torch.zeros_like(g_w)).sum(-1) + gate_up = torch.empty((num_rows, moe.gate_up_proj_weight.shape[1]), dtype=torch.bfloat16, device=g_idx.device) + _expert_gemm(g_x, gate_up_w, gate_up, e, moe.use_fp8_act, recipe_a, recipe_b) + down_in = swiglu_forward_and_per_token_cast( + gate_up, **swiglu_kwargs, topk_weights=w_e, clamp_value=moe.swiglu_clamp_value + ) + down = torch.empty((num_rows, moe.down_proj_weight.shape[1]), dtype=torch.bfloat16, device=g_idx.device) + _expert_gemm(down_in, down_w, down, e, moe.use_fp8_act, recipe_a, recipe_b) + acc = acc + down.float() + + out = _reduce_scatter(acc, group, size, rank)[:num_tokens] + shared = moe.shared_forward(hidden_states) + if shared is not None: + out = out + shared.float() + return out.to(torch.bfloat16) + + +def _expert_gemm(a, b, d, e: int, fp8: bool, recipe_a, recipe_b) -> None: # noqa: ANN001 + if fp8: + deep_gemm.m_grouped_fp8_fp4_gemm_nt_shared_a(a, b, d, e, recipe_a=recipe_a, recipe_b=recipe_b) + else: + d.copy_((a.float() @ b[e].float().transpose(0, 1)).to(d.dtype)) diff --git a/deepinfer/neuron_ops/transfer_to_neuron.py b/deepinfer/neuron_ops/transfer_to_neuron.py index fd6a22e..8ec91e2 100644 --- a/deepinfer/neuron_ops/transfer_to_neuron.py +++ b/deepinfer/neuron_ops/transfer_to_neuron.py @@ -21,6 +21,7 @@ from __future__ import annotations import contextlib import functools import types +from collections.abc import Iterator from typing import Any import torch @@ -56,6 +57,17 @@ class _DeviceMeta(type): def __subclasscheck__(cls, subclass: type) -> bool: return issubclass(subclass, _REAL_DEVICE) + # Compare equal to the real class: torch.fx codegen registers the ``device`` global only when the builtin it + # recorded at import (the real torch.device) ``== torch.device``. + def __eq__(cls, other: object) -> bool: + return other is cls or other is _REAL_DEVICE + + def __ne__(cls, other: object) -> bool: + return not _DeviceMeta.__eq__(cls, other) + + def __hash__(cls) -> int: + return hash(_REAL_DEVICE) + class _Device(metaclass=_DeviceMeta): """Drop-in for ``torch.device`` that maps ``cuda`` to the target device.""" @@ -255,6 +267,7 @@ def _patch_cuda_module(host: bool) -> None: # in host mode the CPU generator is the only one and torch.manual_seed already seeds it. "manual_seed": lambda seed: None if backend is None else backend.manual_seed(seed), "manual_seed_all": lambda seed: None if backend is None else backend.manual_seed_all(seed), + "is_current_stream_capturing": lambda: False, "set_sync_debug_mode": lambda mode: sync_debug_mode.__setitem__(0, mode), "get_sync_debug_mode": lambda: sync_debug_mode[0], "CUDAGraph": CUDAGraph, @@ -270,7 +283,7 @@ def _patch_cuda_module(host: bool) -> None: funcs[name] = (lambda *a, **k: 0) if backend is None else getattr(backend, name) if backend is None: funcs["get_device_properties"] = lambda device=None: types.SimpleNamespace( - name="neuron-host", total_memory=mem_get_info()[1], multi_processor_count=1 + name="neuron-host", total_memory=mem_get_info()[1], multi_processor_count=1, major=0, minor=0 ) else: @@ -296,6 +309,19 @@ def _patch_cuda_module(host: bool) -> None: memory._dump_snapshot = getattr(backend, "_dump_snapshot", lambda *a, **k: None) +@contextlib.contextmanager +def hide_cuda() -> Iterator[None]: + """Report ``torch.cuda`` as unavailable, e.g. around torch.compile: dynamo/inductor otherwise take their CUDA + paths (RNG state save/restore, triton probing, stream capture checks) on the patched module.""" + cuda = torch.cuda + saved = cuda.is_available + cuda.is_available = lambda: False + try: + yield + finally: + cuda.is_available = saved + + def install(device: str) -> None: """Install the mapping; ``device`` is ``"neuron"`` or ``"cpu"``.""" global _TARGET diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py index 1444d3e..12ac71a 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py @@ -257,6 +257,9 @@ class EngramBuffer(BufferBase): fetched_sf = sf_flat return fetched, fetched_sf + # The tables live in host memory: the gather (and its host-sized all-to-all) runs eagerly, outside the decode + # graph; the graph resumes with the fetched rows. + @torch.compiler.disable def fetch(self, indices: torch.Tensor, num_qps: int = 0, use_tma_aligned_col_major_sf: bool = True) -> List[Callable]: """Fetch all layers; returns one result hook per layer (work happens eagerly here).""" diff --git a/deepinfer/neuron_ops/vendor/deep_ep/utils/math.py b/deepinfer/neuron_ops/vendor/deep_ep/utils/math.py index bca245d..87cd44d 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/utils/math.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/utils/math.py @@ -5,6 +5,10 @@ from __future__ import annotations from typing import Tuple import torch +import os as _os + +# trn2 reads fp8 as IEEE-style e4m3 (max 240); see deepinfer/neuron_ops/fp8.py (kept import-free of deepinfer). +FP8_MAX = float(_os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224")) def calc_diff(x: torch.Tensor, y: torch.Tensor) -> float: @@ -38,8 +42,8 @@ def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: x_padded = torch.nn.functional.pad(x, (0, aligned_n - n), mode='constant', value=0) x_padded_view = x_padded.view(m, -1, 128) x_amax = x_padded_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) - return (x_padded_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( - m, aligned_n)[:, :n].contiguous(), (x_amax / 448.0).view(m, -1) + return (x_padded_view * (FP8_MAX / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( + m, aligned_n)[:, :n].contiguous(), (x_amax / FP8_MAX).view(m, -1) def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor) -> torch.Tensor: diff --git a/deepinfer/neuron_ops/vendor/deep_gemm/__init__.py b/deepinfer/neuron_ops/vendor/deep_gemm/__init__.py index bf465b9..0523fad 100644 --- a/deepinfer/neuron_ops/vendor/deep_gemm/__init__.py +++ b/deepinfer/neuron_ops/vendor/deep_gemm/__init__.py @@ -63,6 +63,7 @@ from .gemm import ( m_grouped_bf16_gemm_nt_contiguous, m_grouped_fp8_fp4_gemm_nn_contiguous, m_grouped_fp8_fp4_gemm_nt_contiguous, + m_grouped_fp8_fp4_gemm_nt_shared_a, m_grouped_fp8_gemm_nn_contiguous, m_grouped_fp8_gemm_nt_contiguous, ) diff --git a/deepinfer/neuron_ops/vendor/deep_gemm/attention.py b/deepinfer/neuron_ops/vendor/deep_gemm/attention.py index ece906b..7c1eb22 100644 --- a/deepinfer/neuron_ops/vendor/deep_gemm/attention.py +++ b/deepinfer/neuron_ops/vendor/deep_gemm/attention.py @@ -62,7 +62,7 @@ def _get_logits_stride(num_logits: int) -> int: def _empty_logits(num_rows: int, num_logits: int, device: torch.device) -> torch.Tensor: stride = _get_logits_stride(num_logits) storage = torch.empty((num_rows, stride), dtype=torch.bfloat16, device=device) - return storage.as_strided((num_rows, num_logits), (stride, 1)) + return storage[:, :num_logits] def _chunk_columns(num_q: int, num_heads: int, gather_keys: bool, multiple: int = 1) -> int: @@ -181,7 +181,7 @@ def _paged_page_views(kv_cache: torch.Tensor, is_fp4: bool) -> tuple[torch.Tenso data_bytes = HEAD_DIM // 2 if is_fp4 else HEAD_DIM assert kv_cache.shape[2] == 1 and token_bytes == data_bytes + SPARSE_SF_BYTES, f"{kv_cache.shape=}" assert kv_cache.stride(3) == 1 and kv_cache.stride(1) == token_bytes, f"{kv_cache.stride()=}" - pages = kv_cache.as_strided((num_pages, page_kv * token_bytes), (kv_cache.stride(0), 1)) + pages = kv_cache.flatten(1) # [num_pages, page_kv * token_bytes], a view data = pages[:, : page_kv * data_bytes].reshape(num_pages, page_kv, data_bytes) sf = pages[:, page_kv * data_bytes :].reshape(num_pages, page_kv, SPARSE_SF_BYTES) return data, sf diff --git a/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py b/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py index 812c68d..f488f24 100644 --- a/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py +++ b/deepinfer/neuron_ops/vendor/deep_gemm/gemm.py @@ -360,6 +360,35 @@ def m_grouped_fp8_fp4_gemm_nt_contiguous( d[start:end, n0:n1].copy_(out.to(d.dtype)) +def m_grouped_fp8_fp4_gemm_nt_shared_a( + a_: tuple[torch.Tensor, torch.Tensor], + b_: tuple[torch.Tensor, torch.Tensor], + d: torch.Tensor, + group_idx: int, + recipe=None, + recipe_a=None, + recipe_b=None, + disable_ue8m0_cast: bool = False, +) -> None: + """NEURON extension, static shapes: ``D = A @ B[group_idx]^T`` with every row of A (no psum layout, no host + sync). The static decode MoE runs each local expert over all gathered rows.""" + from .layout import get_recipe + + a, sfa = a_ + b, sfb = b_ + total_m = a.shape[0] + num_groups, n = b.shape[0], b.shape[1] + k = _logical_k(a) + assert _logical_k(b) == k, f"{a.shape=}, {b.shape=}" + assert d.shape == (total_m, n), f"{d.shape=}, {total_m=}, {n=}" + full_recipe = get_recipe(recipe, recipe_a, recipe_b) + a_f32 = _dequant_operand(a, sfa, total_m, k, full_recipe, True, disable_ue8m0_cast=disable_ue8m0_cast) + sfb_req = _required_sf(sfb, n, k, full_recipe, False, num_groups=num_groups, disable_ue8m0_cast=disable_ue8m0_cast) + for n0, n1 in _row_chunks(n, k): + b_f32 = dequantize(b[group_idx, n0:n1], sfb_req[group_idx, n0:n1], gran_k=full_recipe[2]) + d[:, n0:n1].copy_((a_f32 @ b_f32.transpose(0, 1)).to(d.dtype)) + + def m_grouped_fp8_fp4_gemm_nn_contiguous( a, b, d, grouped_layout, recipe=None, recipe_a=None, recipe_b=None, compiled_dims="nk", disable_ue8m0_cast=False, use_psum_layout=True, diff --git a/deepinfer/neuron_ops/vendor/tile_kernels/config.py b/deepinfer/neuron_ops/vendor/tile_kernels/config.py index ce690bc..deb458a 100644 --- a/deepinfer/neuron_ops/vendor/tile_kernels/config.py +++ b/deepinfer/neuron_ops/vendor/tile_kernels/config.py @@ -25,12 +25,16 @@ import os import sys import torch +import os as _os + +# trn2 reads fp8 as IEEE-style e4m3 (max 240); see deepinfer/neuron_ops/fp8.py (kept import-free of deepinfer). +FP8_MAX = float(_os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224")) _DEFAULT_NUM_SMS = 0 _DEFAULT_USE_PDL = False _DEFAULT_DETERMINISTIC = False _DEFAULT_TOKEN_ALIGNMENT = 128 -_E4M3_MAX = 448.0 +_E4M3_MAX = FP8_MAX # trn2: hardware e4m3 range, see deepinfer/neuron_ops/fp8.py _E2M1_MAX = 6.0 _DEFAULT_CLAMP_MIN_VALUES = { ('e4m3', False): 1e-4, diff --git a/deepinfer/neuron_ops/vendor/tile_kernels/engram/engram_hash_kernel.py b/deepinfer/neuron_ops/vendor/tile_kernels/engram/engram_hash_kernel.py index 78b93fa..fe85f9c 100644 --- a/deepinfer/neuron_ops/vendor/tile_kernels/engram/engram_hash_kernel.py +++ b/deepinfer/neuron_ops/vendor/tile_kernels/engram/engram_hash_kernel.py @@ -2,7 +2,8 @@ ``h_i = xor_{j <= i} (token_j * mult_j)`` in int64; output column ``(i - 1, t)`` = ``h_i % vocab_sizes[layer, i - 1, t] + offsets[layer, col]`` (int32), ``-1`` for image tokens. -NOTE(neuron): needs int64 multiply / xor / remainder on device. +NOTE(neuron): the hash needs exact int64 multiply / xor / remainder, which the Neuron compiler narrows to int32, +so a device input is hashed on the host (and kept out of torch.compile graphs); the result goes back to its device. """ # Portions derived from DeepSeek TileKernels 2.1.0, Copyright (c) 2026 DeepSeek, MIT License (see LICENSE). @@ -10,6 +11,7 @@ NOTE(neuron): needs int64 multiply / xor / remainder on device. import torch +@torch.compiler.disable def engram_hash( ngram_token_ids: torch.Tensor, multipliers: torch.Tensor, @@ -29,6 +31,10 @@ def engram_hash( Returns: (num_ngram_layers, num_tokens, (max_ngram_size - 1) * num_embed_table_per_ngram), int32. """ + device = ngram_token_ids.device + if device.type != "cpu": + args = [t if t is None else t.cpu() for t in (ngram_token_ids, multipliers, vocab_sizes, offsets, image_token_mask)] + return engram_hash(*args).to(device) num_tokens, max_ngram_size = ngram_token_ids.shape num_ngram_layers, _, num_tables = vocab_sizes.shape num_out_cols = (max_ngram_size - 1) * num_tables diff --git a/deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py b/deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py index 571a885..93c1dc9 100644 --- a/deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py +++ b/deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py @@ -22,12 +22,23 @@ from tile_kernels.moe.sqrt_softplus import sqrt_softplus _LARGE_PRIME = 23333 +def _hash_uniform(a: torch.Tensor, b: torch.Tensor, seed: int) -> torch.Tensor: + """Deterministic uniform [0, 1) from int32 coordinates (an integer hash; no generator, traceable).""" + h = a * 73856093 + b * 19349663 + seed * 83492791 + h = h ^ (h >> 13) + h = h * 1274126177 + h = h ^ (h >> 16) + return (h & 0xFFFFFF).to(torch.float32) * (1.0 / (1 << 24)) + + def _force_random(num_tokens, num_physical_topk, num_physical_experts, max_physical_experts, ep_rank, device): - gen = torch.Generator(device='cpu') - gen.manual_seed(int(ep_rank)) - keys = torch.rand((num_tokens, max_physical_experts), generator=gen).to(device) - weights = torch.rand((num_tokens, num_physical_topk), generator=gen).to(device) - cols = torch.arange(max_physical_experts, device=device) + # NEURON: random routing for the forced tokens from an integer hash instead of a seeded CPU generator, so it + # stays on the device and inside the decode graph. + token = torch.arange(num_tokens, dtype=torch.int32, device=device)[:, None] + cols = torch.arange(max_physical_experts, dtype=torch.int32, device=device) + keys = _hash_uniform(token, cols[None, :], int(ep_rank)) + weights = _hash_uniform(token, torch.arange(num_physical_topk, dtype=torch.int32, device=device)[None, :] + max_physical_experts, + int(ep_rank)) keys = torch.where(cols[None, :] < num_physical_experts, keys, torch.full_like(keys, 2.0)) idx = torch.sort(keys, dim=1, stable=True).indices[:, :num_physical_topk] return idx.to(torch.int64), weights diff --git a/deepinfer/neuron_ops/vendor/tile_kernels/quant/_core.py b/deepinfer/neuron_ops/vendor/tile_kernels/quant/_core.py index b6c1aef..ecc6c27 100644 --- a/deepinfer/neuron_ops/vendor/tile_kernels/quant/_core.py +++ b/deepinfer/neuron_ops/vendor/tile_kernels/quant/_core.py @@ -21,8 +21,12 @@ import torch.nn.functional as F from tile_kernels.quant.common import BaseCastConfig, CastOutputConfig, get_packed_ue8m0_pack_factor, get_packed_ue8m0_torch_dtype from tile_kernels.utils import align, ceil_div +import os as _os -E4M3_MAX = 448.0 +# trn2 reads fp8 as IEEE-style e4m3 (max 240); see deepinfer/neuron_ops/fp8.py (kept import-free of deepinfer). +FP8_MAX = float(_os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224")) + +E4M3_MAX = FP8_MAX # trn2: hardware e4m3 range, see deepinfer/neuron_ops/fp8.py E2M1_MAX = 6.0 _E2M1_TABLE = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) diff --git a/deepinfer/neuron_ops/vendor/tile_kernels/torch/cast.py b/deepinfer/neuron_ops/vendor/tile_kernels/torch/cast.py index b226e14..e6375a6 100644 --- a/deepinfer/neuron_ops/vendor/tile_kernels/torch/cast.py +++ b/deepinfer/neuron_ops/vendor/tile_kernels/torch/cast.py @@ -16,10 +16,14 @@ from tile_kernels.quant.common import ( unpack_from_e2m1fn_x2, ) from tile_kernels.quant.types import QuantTensor +import os as _os + +# trn2 reads fp8 as IEEE-style e4m3 (max 240); see deepinfer/neuron_ops/fp8.py (kept import-free of deepinfer). +FP8_MAX = float(_os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224")) def get_max_quant_val(dtype: torch.dtype): - max_quant_by_dtype = {torch.int8: 6.0, torch.float8_e4m3fn: 448.0} + max_quant_by_dtype = {torch.int8: 6.0, torch.float8_e4m3fn: FP8_MAX} assert dtype in max_quant_by_dtype return max_quant_by_dtype[dtype] diff --git a/deepinfer/neuron_ops/vendor/tilelang/language.py b/deepinfer/neuron_ops/vendor/tilelang/language.py index 460cb1f..983b508 100644 --- a/deepinfer/neuron_ops/vendor/tilelang/language.py +++ b/deepinfer/neuron_ops/vendor/tilelang/language.py @@ -11,6 +11,10 @@ from __future__ import annotations from typing import Any import torch +import os as _os + +# trn2 reads fp8 as IEEE-style e4m3 (max 240); see deepinfer/neuron_ops/fp8.py (kept import-free of deepinfer). +FP8_MAX = float(_os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224")) # TileLang dtype strings (the real `T.` objects stringify to these). float8_e4m3fn = "float8_e4m3fn" @@ -44,7 +48,7 @@ TL_DTYPE_TO_TORCH: dict[str, torch.dtype] = { } _FINITE_MAX = { - float8_e4m3fn: 448.0, + float8_e4m3fn: FP8_MAX, # trn2: hardware e4m3 range, see deepinfer/neuron_ops/fp8.py float8_e5m2: 57344.0, float4_e2m1fn: 6.0, } diff --git a/deepinfer/platform.py b/deepinfer/platform.py index b977b80..6bfa363 100644 --- a/deepinfer/platform.py +++ b/deepinfer/platform.py @@ -267,6 +267,9 @@ class PlatformNeuron(Platform): transfer_to_neuron.install(device) self.device = transfer_to_neuron.target_device() + from deepinfer.neuron_ops import fp8 + + fp8.configure_compiler() # float8_e4m3fn tensors compile as the hardware e4m3 (issue #24) self.fp8_sf_type = torch.int16 import deepinfer.neuron_ops # noqa: F401 diff --git a/deepinfer/worker/cache_engine.py b/deepinfer/worker/cache_engine.py index c21a535..83ee80f 100644 --- a/deepinfer/worker/cache_engine.py +++ b/deepinfer/worker/cache_engine.py @@ -14,7 +14,7 @@ from functools import lru_cache import torch import torch.distributed as dist -from deepinfer import AttrDict, envs +from deepinfer import PLATFORM, AttrDict, envs from deepinfer._C import cuda_utils # pyright: ignore[reportMissingModuleSource] from deepinfer.config import CacheConfig from deepinfer.logger import init_logger @@ -934,8 +934,12 @@ class CacheEngine: padding_info = f"padding_bytes={padding_bytes}({padding_percent:.2f}%)" logger.info(f"{self.num_gpu_blocks=}, {block_shape=}, {padding_info}") # [num_gpu_blocks, num_layers, block_size, kv, num_heads, head_size] - blocks = self._alloc_gpu_uint8((self.num_gpu_blocks, *block_shape), "kvcache") - gpu_blocks = blocks.view(self.num_gpu_blocks, -1) + # Neuron: one trailing scratch block (and state row), part of the cache views but never handed out by the block + # manager (num_gpu_blocks excludes it). The decode graphs store padding tokens (slot -1) there instead of + # filtering them out with a data-dependent nonzero. + num_scratch_blocks = 1 if PLATFORM.is_neuron else 0 + blocks = self._alloc_gpu_uint8((self.num_gpu_blocks + num_scratch_blocks, *block_shape), "kvcache") + gpu_blocks = blocks.view(blocks.shape[0], -1) kb_per_block = blocks[0].numel() * blocks[0].element_size() // 1024 logger.info(f"kvcache block shape: {blocks.shape[1:]}, type: {blocks.dtype}, size: {kb_per_block} KiB") @@ -953,7 +957,7 @@ class CacheEngine: state_nbytes = state_shape[0] * state_shape[1] logger.info(f"State size: {state_nbytes / (1 << 30)} GiB, shape: {state_shape}") if state_nbytes > 0: - gpu_states = self._alloc_gpu_uint8(state_shape, "state") + gpu_states = self._alloc_gpu_uint8((state_shape[0] + num_scratch_blocks, state_shape[1]), "state") else: gpu_states = torch.empty(size=state_shape, dtype=torch.uint8, device="cuda") diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index 9a03065..c0cf54d 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -11,7 +11,7 @@ import numpy as np import torch import torch.distributed as dist -from deepinfer import envs +from deepinfer import PLATFORM, envs from deepinfer._C.engine import prepare_sampler_inputs, update_prompt_logprobs # pyright: ignore[reportMissingModuleSource] from deepinfer.config import CacheConfig, ModelConfig, ParallelConfig, SchedulerConfig from deepinfer.engine.engine_type import EngineType @@ -34,6 +34,7 @@ from deepinfer.worker.cache_engine import CacheEngine, GPUCacheInterface, get_kv from deepinfer.worker.cuda_graph_runner import CudaGraphData, CudaGraphRunner from deepinfer.worker.decode_perf_cache import DecodePerfCache from deepinfer.worker.dflash_cuda_graph_runner import DFlashCudaGraphRunner +from deepinfer.worker.neuron_graph_runner import NeuronGraphRunner from deepinfer.worker.zos_status import MetadataZOS if TYPE_CHECKING: @@ -73,7 +74,8 @@ class ModelRunner: self.use_cuda_graph = not envs.DEEPINFER_ENFORCE_EAGER and self.parallel_config.engine_type != EngineType.PROMPT self.cuda_graph_runner: CudaGraphRunner | None = None if self.use_cuda_graph: - self.cuda_graph_runner = CudaGraphRunner( + runner_cls = NeuronGraphRunner if PLATFORM.is_neuron else CudaGraphRunner + self.cuda_graph_runner = runner_cls( max_batch_size=self.scheduler_config.max_num_seqs * (1 + self.nextn), input_on_device=envs.DEEPINFER_ZERO_OVERHEAD_SCHEDULING, ) diff --git a/deepinfer/worker/neuron_graph_runner.py b/deepinfer/worker/neuron_graph_runner.py new file mode 100644 index 0000000..cfe0fae --- /dev/null +++ b/deepinfer/worker/neuron_graph_runner.py @@ -0,0 +1,158 @@ +"""Static decode graphs on Neuron: the CudaGraphRunner contract, backed by torch.compile(backend="neuron"). + +One compiled graph per captured batch size, over persistent device metadata buffers (the same buffers the CUDA +runner captures over). Dynamo specializes on their shapes (dynamic=False), so each batch size is one NEFF set. +Replay copies the host metadata into the device buffers and calls the compiled forward; the host->device copy stays +outside the graph. +""" + +from __future__ import annotations + +import dataclasses +import math +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch + +from deepinfer import envs +from deepinfer.environ import _descriptors as _env_descriptors +from deepinfer.logger import init_logger +from deepinfer.model_executor.input_metadata import InputMetadata +from deepinfer.model_executor.model_output import CausalLMForwardOutput +from deepinfer.neuron_ops import neuron_compile # noqa: F401 (registers the "deepinfer_neuron" backend) +from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda +from deepinfer.sequence import SequenceMetadata +from deepinfer.worker.cache_engine import GPUCacheInterface, KVCache +from deepinfer.worker.cuda_graph_runner import CudaGraphData, CudaGraphRunner + +if TYPE_CHECKING: + from deepinfer.model_executor.models.deepseek import DeepseekForCausalLM + +logger = init_logger(__name__) + + +class _GraphKVCaches: + """The per-layer cache views, rebuilt from their root allocations inside the graph. + + The layer caches are views of one KV buffer and one state buffer. Passed as graph inputs, they alias each other, + and AOT regenerates aliased mutated inputs with as_strided, which cannot be functionalized. Here only the two + roots are graph inputs; each layer's views are re-sliced from them (slice + view, no as_strided) on access. + """ + + def __init__(self, kv_caches: GPUCacheInterface) -> None: + self.roots: list[torch.Tensor] = [] + self.recipes: list[dict[str, tuple | None]] = [] + for cache in kv_caches.cache_list: + recipe = {} + for f in dataclasses.fields(KVCache): + t = getattr(cache, f.name) + recipe[f.name] = None if t is None else self._recipe(t) + self.recipes.append(recipe) + + def _recipe(self, t: torch.Tensor) -> tuple: + storage = t.untyped_storage() + for i, r in enumerate(self.roots): + if r.untyped_storage().data_ptr() == storage.data_ptr(): + root = r + break + else: + # One flat uint8 tensor over the whole allocation (dtype views of it have no usable _base). + root = torch.empty(0, dtype=torch.uint8, device=t.device).set_(storage) + self.roots.append(root) + i = len(self.roots) - 1 + rows, row_bytes = t.shape[0], t.stride(0) * t.element_size() + offset = t.storage_offset() * t.element_size() + assert offset + math.prod(t.shape[1:]) * t.element_size() <= row_bytes, "cache view must be a column slice of its allocation" + nbytes = math.prod(t.shape[1:]) * t.element_size() + recipe = (i, rows, row_bytes, offset, nbytes, tuple(t.shape[1:]), t.dtype) + rebuilt = self._build(root, recipe) + assert rebuilt.data_ptr() == t.data_ptr() and rebuilt.shape == t.shape and rebuilt.stride() == t.stride(), ( + f"cache view not expressible as a row slice: {t.shape=} {t.stride()=}" + ) + return recipe + + @staticmethod + def _build(root: torch.Tensor, recipe: tuple) -> torch.Tensor: + _, rows, row_bytes, offset, nbytes, tail, dtype = recipe + esize = torch.empty((), dtype=dtype).element_size() + flat = root[: rows * row_bytes] + view = flat.view(rows, row_bytes)[:, offset : offset + nbytes].view(rows, *tail[:-1], tail[-1] * esize) + return view if dtype == torch.uint8 else view.view(dtype) + + def __getitem__(self, layer_idx: int) -> KVCache: + tensors = {} + for name, recipe in self.recipes[layer_idx].items(): + tensors[name] = None if recipe is None else self._build(self.roots[recipe[0]], recipe) + return KVCache(**tensors) + + def is_valid_layeridx(self, layer_idx: int) -> bool: + return layer_idx < len(self.recipes) + + +class NeuronGraphRunner(CudaGraphRunner): + def __init__(self, max_batch_size: int, input_on_device: bool = False) -> None: + super().__init__(max_batch_size, input_on_device) + self.compiled: dict[int, Callable[..., CausalLMForwardOutput]] = {} + self.model: DeepseekForCausalLM | None = None + self.kv_caches: GPUCacheInterface | None = None + + def capture( + self, + model: DeepseekForCausalLM, + kv_caches: GPUCacheInterface, + prepare_input_fn: Callable[[list[SequenceMetadata]], InputMetadata], + padding_seq_metadata: SequenceMetadata, + ) -> None: + self.destroy() + self.model, self.kv_caches = model, kv_caches + # Resolve every env var now: the lazy descriptor (lock + os.getenv) does not trace. + for name in _env_descriptors: + getattr(envs, name) + backend, fullgraph = envs.DEEPINFER_NEURON_COMPILE_BACKEND, envs.DEEPINFER_NEURON_FULLGRAPH + logger.info(f"Neuron decode graphs: backend={backend} fullgraph={fullgraph} batch_sizes={self.batch_sizes}") + + self.padding_seq_metadata_list = [padding_seq_metadata] * self.max_batch_size + input_metadata_source_full = prepare_input_fn(self.padding_seq_metadata_list) + input_metadata_device_full = input_metadata_source_full.recursive_fn(lambda x: x.cuda(non_blocking=True)) + + graph_caches = _GraphKVCaches(kv_caches) + + def forward(input_metadata: InputMetadata) -> CausalLMForwardOutput: + return model(kv_caches=graph_caches, input_metadata=input_metadata) + + self.graphs = {} + for batch_size in reversed(self.batch_sizes): + + def _cut(x: torch.Tensor, _bs: int = batch_size) -> torch.Tensor: + assert x.shape[0] == self.max_batch_size + return x[:_bs] + + input_metadata_source = input_metadata_source_full.recursive_fn(_cut) + input_metadata_device = input_metadata_device_full.recursive_fn(_cut) + compiled = torch.compile(forward, backend=backend, fullgraph=fullgraph, dynamic=False) + # First call compiles (and runs) this batch size's graph. + input_metadata_device.recursive_copy_(input_metadata_source) + with hide_cuda(): + compiled(input_metadata_device) + for name in input_metadata_source.get_clean_list(): + setattr(input_metadata_source, name, None) + setattr(input_metadata_device, name, None) + self.compiled[batch_size] = compiled + self.graphs[batch_size] = CudaGraphData(None, None, None, input_metadata_source, input_metadata_device) + + def replay(self, cuda_graph_data: CudaGraphData) -> CausalLMForwardOutput: + input_metadata_device = cuda_graph_data.input_metadata_device + # Forward recomputes these (and stores them on the metadata); reset them to the state the graph was traced in. + for name in input_metadata_device.get_clean_list(): + setattr(input_metadata_device, name, None) + if not self.input_on_device: + input_metadata_device.recursive_copy_(cuda_graph_data.input_metadata_source) + batch_size = next(bs for bs, data in self.graphs.items() if data is cuda_graph_data) + with hide_cuda(): + return self.compiled[batch_size](input_metadata_device) + + def destroy(self) -> None: + super().destroy() + self.compiled = {} + self.model = self.kv_caches = None diff --git a/docs/neuron/vllm_neuron_compile_notes.md b/docs/neuron/vllm_neuron_compile_notes.md new file mode 100644 index 0000000..ad2f0d7 --- /dev/null +++ b/docs/neuron/vllm_neuron_compile_notes.md @@ -0,0 +1,281 @@ +# How vllm-neuron compiles a decode step (notes for the DeepInfer Neuron port) + +Source: local checkout `Trn2/vllm-neuron` (HEAD `f32079d`, "Merge pull request #63 add-contrib-folder"). Paths are relative to the repo root. Read-only survey. Nothing in vllm-neuron was changed. + +**Caveat.** vllm-neuron does **not** use `torch-neuronx` or its `"neuron"` backend. It compiles through `libtorch-neuronx-lite` (`requirements/core.txt:7`), which carries its own copy of the native lowering: Dynamo, then AOTAutograd with Neuron decomps, then torch-mlir, StableHLO and neuronx-cc (`docs/design/compilation/how_vllm_is_integrated_with_torchneuron.md`). The lite source is not in this checkout. Statements about what lite does internally come from comments and docs, not from code. + +## 1. Compile backend and compiler flags + +- `envs.get_compile_backend_name()` (`vllm_neuron/envs.py:789-803`) has two routes: + - **Native route** (the default, `VLLM_NEURON_BACKEND` unset, `neuron_native` or `neuron`): it calls `libtorch_neuronx_lite.compile.native_backend.register()` and returns the name that call registers. That name is deliberately **not** `"neuron"`, which "stays reserved for the torch_neuronx package". + - **XLA route** (deprecated) and CPU mode: it returns `"neuron_libtorch"`, registered at import (`vllm_neuron/__init__.py:~128`). + - `torch-neuronx` and lite cannot coexist unless `TORCH_DEVICE_BACKEND_AUTOLOAD=0`. Otherwise you get "multiple backend fallbacks for PrivateUse1". +- The torch.compile call is in `vllm_neuron/vllm/worker/neuron_model_runner_v2.py:1903-1908`: + `torch.compile(NeuronForwardAdapter(model), backend=envs.get_compile_backend_name(), fullgraph=not DEBUG, options=self.compile_options)`. + - `torch._dynamo.config.cache_size_limit = 2**62` (`:1860`). Buckets are meant to specialize; recompiles are caught by the stance instead. + - The serve-time forward runs under `torch.compiler.set_stance("fail_on_recompile")` (`utils/neuron_utils.py:26-32`, used at `_serve_forward_guard`, `neuron_model_runner_v2.py:1223`). + - Optional `guard_filter_fn` (`envs.py:814`) for warm start: `keep_portable_guards_unsafe` or `skip_all_guards_unsafe`. +- `compile_options` (`neuron_model_runner_v2.py:387`, `:672-701`): + ``` + {"alias_meta_to_neuron": True, # lite: thread aliased (mutated) outputs back onto inputs + "compiler_args": ["--auto-cast=none", "--verbose=35", f"-O{vllm optimization_level}", + "--internal-hlo2tensorizer-options=--modular-flow-mac-threshold=10" + + (" --experimental-unsafe-fp8e4m3fn-as-fp8e4m3" if fp8 and target not trn3), + "--internal-backend-options=--enable-verifier=false --enable-nested-dynamic-loop", + (+ explicit --lnc=N only when it differs from the compiler default)]} + ``` + - `--enable-nested-dynamic-loop` is needed because NKI kernels inlined into the parent BIR have nested `scf.for`, and birverifier rejects the extra back-edges (`neuron_model_runner.py:2071-2076`). + - All hlo2tensorizer options must go into **one** string. A second `--internal-hlo2tensorizer-options` replaces the first (`neuron_model_runner.py:2037-2047`). + - The V1 runner adds `--enable-annotate-fsdp-cc=false` and `--model-type=generic` for the GPT-OSS giga-kernel graph (`:2047`, `:2079`). torch-neuronx injects `--model-type transformer`, and neuronx-cc keeps the last flag it sees. + - **Per-graph args.** A model may define `compiler_args_for(graph_type, example_inputs, base)` (`neuron_model_runner_v2.py:703-723`; `docs/design/framework/per_bucket_compiler_args.md`). Qwen3-VL and ASR compile the decode and encoder graphs at **-O3**, because "-O3 disables neuronx-cc modular flow; the -O1 default adds a PSEUDO_FUNCTION_RETURN/EVENT_SEMAPHORE wrapper on each collective" (`model/qwen3_vl/model_common.py:290-301`). Prefill keeps -O2. + - The native route defaults to O2 and the XLA route to O1 (`docs/guides/how-to-choose-compilation-route.md`). + +## 2. KV cache as graph input (the synthetic-base problem) + +- **The KV cache is not a forward argument.** `bind_kv_cache` stores plain tensor attributes on each attention module: `layer.self_attn.k_cache = kv[name][0]`, `.v_cache = kv[name][1]` (`model/qwen3/model.py:1160-1166`). Dynamo lifts these attributes as graph inputs. The adapter signature (`worker/neuron_forward_adapter.py:209-227`) carries only metadata and state tensors. +- **Allocation** (V2, `neuron_model_runner_v2.py:917-1215`): + - One `int8` backing tensor is created (`:1140`). + - For each layer it takes `raw = backing[start:start+span]`, then `raw.view(dtype).view(2, num_blocks, kv_heads, block_size, head_size)`. **K = typed[0] and V = typed[1]**: two non-overlapping views of one per-layer region. +- **The trick that avoids the synthetic base.** View objects are memoized: `shared_dtype_view` and `shared_kv_slices` are keyed by `(id(raw), dtype, kv_shape)` (`:1148-1186`). Layers that share bytes then receive the **exact same tensor objects**, and Dynamo dedups them into one placeholder. The comment at `:1160-1166` says: "Re-slicing shared storage per layer gives Dynamo overlapping inputs that AOTAutograd merges into a whole-buffer synthetic base, widening a slot-sized KV update into a full-buffer write." + - In short, the views AOTAutograd sees never partially overlap: K and V are disjoint, and layers are disjoint. AOTAutograd's overlap check then keeps them as independent mutated inputs and creates no synthetic base and no `as_strided` replay. + - A host preflight rejects pooled layers whose `(dtype, heads, block, head_size)` differ, because they would produce distinct, overlapping view objects (`:1086-1110`). +- **2**32 offset hazard.** The V1 runner deliberately allocates **one `torch.zeros` per layer** rather than slicing one giant buffer. Its reason: "views into a multi-GiB shared storage carry element offsets past 2**32, which the lowering path can truncate (silent cross-layer KV corruption)" (`neuron_model_runner.py:10988-10995`). V2 slices a single backing. **For DeepInfer, per-layer allocations are the safe choice.** +- **KV writes are plain `index_put_` inside the graph**, not a custom op and not a kernel-side cache update: + - `paged_slot_indices(slot_mapping, block_size, num_blocks)` (`model/kv_cache.py:80`) clamps invalid or `-1` slots to null block 0 **before** `//` and `%`. A `-1` slot becomes an out-of-bounds DGE scatter, which wedges the NEFF execution queue. + - Then `k_cache.index_put_((block_idx, head_idx, pos_idx), k_flat)` (`model/qwen3/model.py:336-343` for prefill, `:582-589` for decode). + - The decode NKI kernel is called with `update_cache=False` and returns `K_new` and `V_new`. The model then scatters them itself (`:525-589`, `qwen3_vl/utils/decode_kv.py:60` `scatter_new_kv`). +- **Mutation becomes aliasing.** + - Lite's `alias_meta_to_neuron=True` turns every mutated input into an extra output plus an `input_output_alias` (donation). The XLA-route equivalents are documented in `docs/design/compilation/aliasing_output_rewrite_pass.md` (mutated inputs are appended to the outputs, giving `io_map{out_idx: in_idx}`) and `inplace_to_outofplace_pass.md`. NKI kernels declare `operand_output_aliases` for the same purpose. + - vllm-neuron uses **no** `mark_static_address`, **no** `keep_inference_input_mutations` and no custom `torch._functorch` config. +- **Runner state follows the same pattern.** `token_ids.index_put_(...)` runs in-graph (`neuron_forward_adapter.py:190`), and the aliased bundle is returned. + - The next decode reads a **clone** of the sampled store, a non-aliased output (`:345-348`). That keeps TP sampling collectives ordered across ranks. + - The output future is kept alive until the next step (`neuron_model_runner_v2.py:1453-1460`). Otherwise in-flight NRT input slices get freed and TP hangs. + +## 3. Graph granularity + +- **One graph = the whole forward**, for each (phase, bucket). It covers the input-id build, all decoder layers (a Python loop unrolled by Dynamo, `model/qwen3/model.py:896`), the LM head, on-device sampling and the state commit (`neuron_forward_adapter.py:272-356`). + - Prefill and decode are separate graphs, selected by whether `linear_idx` is None. + - NEFF count is `|num_seqs_buckets| × (|decode_ctx_buckets|+1)` (`docs/design/vllm/decode-context-length-bucketing.md:111-116`). +- **There is no per-layer compiled region**, and no mechanism to "compile one layer once and reuse it". Reuse of repeated layers is left to **neuronx-cc modular flow** (`--modular-flow-mac-threshold=10` at O1/O2). It is turned **off** with `-O3` for the latency-critical decode graph. +- The giga-kernel path goes further: one NKI kernel for all 36 GPT-OSS layers in a single decode NEFF, with weights stacked on a `[num_layers, ...]` axis (`docs/design/vllm/gpt-oss-gk-decode.md`). +- The ~120-NEFF design in the DeepSeek-V4.1 PR #71 (per-expert graphs, `@torch._dynamo.disable` MoE, host sync per layer) is the opposite of what the rest of the repo does. See `vllm_neuron_pr71_analysis.md`. +- Compile time is handled by parallel trace and compile across the TP group (`parallel_compile`, `neuron_model_runner_v2.py:2035-2056`), the compile cache, and the SPMD single-NEFF guard. The graph is not split. + +## 4. Buckets and static shapes + +- Defaults (`utils/bucket_utils.py:36-126`): + - Prefill token buckets are powers of 2 from 128 up to `max_num_batched_tokens`. + - Decode batch buckets are powers of 2 from 1 up to `max_num_seqs`. + - Optional KV-segment buckets (prefix caching) and decode context-length buckets. +- `_select_bucket` picks the smallest bucket that fits (`neuron_model_runner_v2.py:903`). +- Padding rules: + - Padded tokens use `PAD_TOKEN=0`. Padded slots use `slot=-1`, which `paged_slot_indices` redirects to null block 0. + - The block table uses the `-1` sentinel past each row's valid blocks, and the attention kernel skips OOB. + - Padded decode rows have `row_to_slot=-1` (`worker/neuron_input_build.py:1-13`; `docs/design/vllm/padding-batching-block-kv.md`). +- Metadata is built on the host with NumPy only and never reads device state. It is copied with `torch.from_numpy(...).to(device)` **outside** the graph each step (`neuron_model_runner_v2.py:2687-2712`). + - Python ints (`max_query_len`, `block_size`, `kv_segment_size`) are part of the metadata dict on purpose, so they specialize the graph. + - Decode `input_ids` are built **in-graph** from the previous step's on-device sampled tokens, so steady-state decode does no D2H. +- **No `.item()` in the graph.** Host reads happen only at output time (`.cpu().tolist()` of the sampled future). +- Scale floats are kept as Python floats next to the tensors, because `tensor.item()` causes a graph break (`model/llama3/model.py:443`). + +## 5. Collectives inside graphs + +- The world and model-parallel process groups are **Gloo**. Collectives are traced into the NEFF and never dispatched at runtime (`envs.py:858-866`). +- APIs used: + - `torch.distributed._functional_collectives.all_reduce(x, reduceOp="sum", group=...)`. The reduceOp must be a **string**, because a `ReduceOp` object breaks precompile pickling (`nn/embedding.py:256-260`, `nn/rpl.py:147-150`). + - `all_gather_tensor`. + - `torch.ops._c10d_functional.all_gather_into_tensor`, `all_to_all_single` and `wait_tensor` (`functional/collectives/all_to_all.py:127-193`). + - Plain in-place `dist.all_reduce` (through vLLM's `GroupCoordinator`) is also used. Dynamo remaps it to the functional form. +- Restrictions: + - **trn2: all_to_all is emulated** as all_gather plus a local slice, because "ENC lacks Mesh" (`all_to_all.py:10-33`). + - Every rank must trace byte-identical HLO (SPMD). Use a runtime `rank` tensor, not Python ints, and replace rank-dependent slices with `index_select(arange+offset)`. Each group registers its full replica-group mesh (`docs/design/compilation/spmd_compilation.md`, `patches/replica_mesh_patch.py`, `parallel/neuron_group_coordinator.py:122-130`). + - **A collective cannot read a NEFF I/O tensor** (`NCC_INLA001`). Copy the input to an intermediate first (`functional/collectives/reduce_scatter_v.py:80-95, 318-330`). + +## 6. Known neuronx-cc workarounds found + +| Symptom | Workaround | Where | +|---|---|---| +| birverifier rejects nested loops from inlined NKI | `--enable-verifier=false --enable-nested-dynamic-loop` | `neuron_model_runner.py:2071-2076` | +| `NCC_INLA001` collective reads an I/O tensor | copy into an intermediate buffer before the collective | `reduce_scatter_v.py:89,177,324` | +| Broadcast boolean AND miscompiles (`[4096,5120] & [1,5120]`) | fold validity into positions and use one `where` (CHRS-985) | `attention/attention_segmented_cte.py:546` | +| `NCC_ITEN406` "too many partition dims" (strided AP), often after graph breaks | keep fullgraph; avoid strided views | `docs/model-dev/debugging.md:150-160` | +| `NCC_ILSA901` transpose AP in fused MLP at T=2/4 | in-tree MLP path | `functional/qwen3_intree_mlp.py:174-202` | +| `NCC_INKI016` query seqlen must be a multiple of 128 | bucket validation | `utils/bucket_utils.py:455` | +| `NCC_IBTN243` SendRecv does not lower at lnc1 | prefill at lnc2, separate engine | `gpt-oss-gk-decode.md:19` | +| `-1` slot → OOB DGE scatter wedges the queue | `paged_slot_indices` clamps to null block | `model/kv_cache.py:80` | +| Hidden element offsets > 2**32 into giant buffers | allocate per-layer KV | `neuron_model_runner.py:10988` | + +No `InferNonlocalTensors` hit was found in this repo. + +## 7. FP8 + +- KV fp8 uses `torch.float8_e4m3fn` (`utils/dtype_utils.py:7-14`). Values are clamped to `FP8_CLAMP_MAX`: **240 on trn2** (the hardware/NKI `float8_e4m3`) and 448 on trn3 (`:16-41`). The constant is resolved at import so it is a trace-time constant. Write path: `(k*k_scale).clamp(±240).to(e4m3fn)` (`model/qwen3/model.py:311-321`). +- On trn2 the compiler flag `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3` (inside the hlo2tensorizer string) makes the compiler treat e4m3fn as e4m3 (`neuron_model_runner_v2.py:673-686`). +- Static-fp8 weights (ModelOpt e4m3fn, range 448) are **downscaled by 240/448 at load** and the dequant scale is multiplied by 448/240 (`model/llama3/weight_loaders_static_fp8.py:45-80`, `model_static_fp8.py:126-160`). The prefix match `trn2` / `trn2-*` matters (`:150`). +- Some kernels and the swa path use 448 or 240 differently (`attention/swa_fused.py:45`). Check each kernel's range. + +## 8. Patches relevant to torch/dynamo + +- `vllm/patches/*` patches vLLM, not Dynamo or AOT (DP balancer, shm spin, pin_memory off, `accelerator.empty_cache` guard, port hold, replica mesh). +- Settings worth copying: + - `cache_size_limit=2**62` plus a `fail_on_recompile` stance at serve time. + - Warmup and serve must agree on grad mode, because Dynamo guards GLOBAL_STATE. Lite traces under `no_grad`, not `inference_mode` (`utils/neuron_utils.py:35-55`). + - Avoid logging inside traced code. Even `is_native_backend()` can log and break the graph, so it is read at construction (`nn/embedding.py:98-100`). + - Use precompile and AOTAutograd caching via `TORCH_CACHING_PRECOMPILE=1` and `TORCHINDUCTOR_AUTOGRAD_CACHE=1` (`vllm/inductor_cache.py:1-22`). + +## 9. DeepSeek-V4 / V4.1 in this checkout + +None. `contrib/models/` contains only `template/`, and `vllm_neuron/model/` has gpt_oss, llama3, qwen3, qwen3_vl, qwen3_asr and synthetic. "deepseek" appears only in comments (the V3.2 indexer key-only spec in `model/kv_cache.py:54`). PR #71 is unmerged here. + +## What DeepInfer should copy (ranked by impact) + +1. **KV as per-layer, non-overlapping, memoized views.** Allocate one tensor per layer (`(2,N,H,B,D)` or separate K and V). Bind them as module attributes. Never re-slice the same region twice. This removes the synthetic base and the `as_strided` path, and avoids the 2**32 offset bug. +2. **KV writes as plain `index_put_` with clamped indices** (`paged_slot_indices`, invalid slots go to null block 0). Rely on the backend's mutation-to-output aliasing (donation), not a custom op. +3. **One full-graph forward per (phase, bucket)**, including sampling and the in-graph state commit. Let neuronx-cc modular flow dedup layers (O1/O2), and try `-O3` for decode. Do not split per layer or per expert. +4. **Use the compiler flag set verbatim**: `--auto-cast=none`, a single `--internal-hlo2tensorizer-options` string (`--modular-flow-mac-threshold=10`, plus the fp8 flag on trn2), `--internal-backend-options=--enable-verifier=false --enable-nested-dynamic-loop`, and `--model-type=generic` if `transformer` hurts. +5. **Host-built, fixed-shape metadata, H2D outside the graph.** Decode input_ids come from the on-device previous sample. No `.item()` in the graph. Use `fail_on_recompile` at serve time. +6. **Functional collectives with string reduceOp, SPMD rank tensor, no collective on graph I/O.** On trn2, verify all_to_all; vllm-neuron emulates it with all_gather. +7. **FP8: e4m3fn storage, clamp ±240, unsafe-e4m3fn-as-e4m3 flag**, and a 240/448 weight downscale for OCP checkpoints. + +
中文 + +# vllm-neuron 如何编译 decode 一步（给 DeepInfer Neuron 移植的笔记） + +来源：本地 `Trn2/vllm-neuron`（HEAD `f32079d`，"Merge pull request #63 add-contrib-folder"）。路径相对仓库根目录。本文只读调研，没有改动 vllm-neuron 的任何内容。 + +**前提说明。** vllm-neuron **不用** `torch-neuronx`，也不用它的 `"neuron"` 后端。它通过 `libtorch-neuronx-lite`（`requirements/core.txt:7`）编译。lite 自带一份原生 lowering：Dynamo → AOTAutograd（带 Neuron decomps）→ torch-mlir → StableHLO → neuronx-cc（`docs/design/compilation/how_vllm_is_integrated_with_torchneuron.md`）。lite 源码不在这个仓库里。下文凡是讲 lite 内部行为的，依据都是注释和文档，不是代码。 + +## 1. 编译后端与编译参数 + +- `envs.get_compile_backend_name()`（`vllm_neuron/envs.py:789-803`）有两条路线： + - **原生路线**（默认；`VLLM_NEURON_BACKEND` 不设，或设为 `neuron_native`/`neuron`）：调用 `libtorch_neuronx_lite.compile.native_backend.register()`，返回它注册的名字。这个名字刻意**不叫** `"neuron"`，`"neuron"` 留给 torch_neuronx 包。 + - **XLA 路线**（已弃用）和 CPU 模式：返回 `"neuron_libtorch"`，在 import 时注册（`vllm_neuron/__init__.py:~128`）。 + - `torch-neuronx` 和 lite 要共存，必须设 `TORCH_DEVICE_BACKEND_AUTOLOAD=0`，否则报 "multiple backend fallbacks for PrivateUse1"。 +- torch.compile 调用在 `vllm_neuron/vllm/worker/neuron_model_runner_v2.py:1903-1908`： + `torch.compile(NeuronForwardAdapter(model), backend=..., fullgraph=not DEBUG, options=self.compile_options)`。 + - `torch._dynamo.config.cache_size_limit = 2**62`（`:1860`）。bucket 本来就要特化，重编译由 stance 来拦。 + - serve 时的 forward 包在 `torch.compiler.set_stance("fail_on_recompile")` 里（`utils/neuron_utils.py:26-32`，在 `_serve_forward_guard` 用到，`neuron_model_runner_v2.py:1223`）。 + - warm start 可选 `guard_filter_fn`（`envs.py:814`）：`keep_portable_guards_unsafe` 或 `skip_all_guards_unsafe`。 +- `compile_options`（`neuron_model_runner_v2.py:387`、`:672-701`）： + - `alias_meta_to_neuron=True`：lite 把被改写的输入接回成别名输出。 + - `compiler_args` 依次是： + - `--auto-cast=none` + - `--verbose=35` + - `-O{vLLM optimization_level}` + - `--internal-hlo2tensorizer-options=--modular-flow-mac-threshold=10`，trn2 上有 fp8 时在同一字符串里追加 ` --experimental-unsafe-fp8e4m3fn-as-fp8e4m3` + - `--internal-backend-options=--enable-verifier=false --enable-nested-dynamic-loop` + - LNC 与编译器默认值不同时，再加显式 `--lnc=N` + - 需要 `--enable-nested-dynamic-loop`，是因为 NKI kernel 内联进父 BIR 后有嵌套 `scf.for`，birverifier 会拒多出来的回边（`neuron_model_runner.py:2071-2076`）。 + - hlo2tensorizer 选项必须写在**同一个**字符串里。第二个 `--internal-hlo2tensorizer-options` 会覆盖第一个（`neuron_model_runner.py:2037-2047`）。 + - V1 runner 给 GPT-OSS giga-kernel 图加了 `--enable-annotate-fsdp-cc=false` 和 `--model-type=generic`（`:2047`、`:2079`）。torch-neuronx 会注入 `--model-type transformer`，而 neuronx-cc 以最后出现的那个为准。 + - **按图设参数。** 模型可以实现 `compiler_args_for(graph_type, example_inputs, base)`（`neuron_model_runner_v2.py:703-723`；`docs/design/framework/per_bucket_compiler_args.md`）。Qwen3-VL/ASR 的 decode 和 encoder 图用 **-O3**，理由是："-O3 关掉 neuronx-cc modular flow；-O1 默认开着它，会给每个 collective 包一层 PSEUDO_FUNCTION_RETURN/EVENT_SEMAPHORE"（`model/qwen3_vl/model_common.py:290-301`）。prefill 保持 -O2。 + - 原生路线默认 O2，XLA 路线默认 O1（`docs/guides/how-to-choose-compilation-route.md`）。 + +## 2. KV cache 作为图输入（synthetic base 问题） + +- **KV cache 不是 forward 参数。** `bind_kv_cache` 把普通 tensor 挂成每个 attention 模块的属性：`layer.self_attn.k_cache = kv[name][0]`、`.v_cache = kv[name][1]`（`model/qwen3/model.py:1160-1166`）。Dynamo 把这些属性提升为图输入。adapter 签名（`worker/neuron_forward_adapter.py:209-227`）里只有元数据和状态张量。 +- **分配方式**（V2，`neuron_model_runner_v2.py:917-1215`）： + - 先建一个 `int8` backing（`:1140`）。 + - 每层取 `raw = backing[start:start+span]`，再 `raw.view(dtype).view(2, num_blocks, kv_heads, block_size, head_size)`。**K = typed[0]，V = typed[1]**，是同一块每层区域上两个不重叠的 view。 +- **避开 synthetic base 的关键。** view 对象做了记忆化：`shared_dtype_view` 和 `shared_kv_slices` 以 `(id(raw), dtype, kv_shape)` 为 key（`:1148-1186`）。共享字节的层因此拿到**完全相同的 tensor 对象**，Dynamo 会把它们去重成一个 placeholder。`:1160-1166` 的注释原话："对共享 storage 每层重新切片，会给 Dynamo 一组重叠输入，AOTAutograd 会把它们合并成整块 synthetic base，把 slot 大小的 KV 更新放大成整块写。" + - 也就是说，AOTAutograd 看到的 view 之间从不部分重叠：K 和 V 不相交，层与层不相交。AOTAutograd 的重叠检查会把它们当成各自独立的被改写输入，不会造 synthetic base，也不会回放 `as_strided`。 + - host 端预检查会拒绝 `(dtype, heads, block, head_size)` 不一致的共享层，因为那样会产生不同的、相互重叠的 view 对象（`:1086-1110`）。 +- **2**32 偏移隐患。** V1 runner 故意**每层单独 `torch.zeros`**，不从一个大 buffer 里切。原因是："指向多 GiB 共享 storage 的 view，元素偏移会超过 2**32，lowering 可能截断它（静默的跨层 KV 损坏）"（`neuron_model_runner.py:10988-10995`）。V2 是从单个 backing 里切的。**DeepInfer 用每层独立分配最稳。** +- **KV 写入就是图内普通的 `index_put_`**，不是自定义 op，也不是 kernel 内部更新 cache： + - `paged_slot_indices(slot_mapping, block_size, num_blocks)`（`model/kv_cache.py:80`）在做 `//`、`%` **之前**，把非法或 `-1` 的 slot 钳到 null block 0。`-1` 会变成越界的 DGE scatter，卡死 NEFF 执行队列。 + - 然后 `k_cache.index_put_((block_idx, head_idx, pos_idx), k_flat)`（prefill 在 `model/qwen3/model.py:336-343`，decode 在 `:582-589`）。 + - decode 的 NKI kernel 以 `update_cache=False` 调用，返回 `K_new`、`V_new`，再由模型自己 scatter 回去（`:525-589`，`qwen3_vl/utils/decode_kv.py:60` 的 `scatter_new_kv`）。 +- **改写如何变成别名。** + - lite 的 `alias_meta_to_neuron=True` 把每个被改写的输入变成一个额外输出加一条 `input_output_alias`（donation）。XLA 路线上的对应做法见 `docs/design/compilation/aliasing_output_rewrite_pass.md`（被改写的输入追加进输出，得到 `io_map{out_idx: in_idx}`）和 `inplace_to_outofplace_pass.md`。NKI kernel 用 `operand_output_aliases` 声明同样的别名关系。 + - vllm-neuron **没有用** `mark_static_address`、`keep_inference_input_mutations`，也没有自定义 `torch._functorch` 配置。 +- **runner 状态也是同一套。** `token_ids.index_put_(...)` 在图内执行（`neuron_forward_adapter.py:190`），并返回别名 bundle。 + - 下一步 decode 读的是采样结果的 **clone**，一个非别名输出（`:345-348`）。这样各 TP rank 上的采样 collective 保持有序。 + - 输出 future 要一直持有到下一步（`neuron_model_runner_v2.py:1453-1460`），否则在飞的 NRT 输入切片会被释放，TP 会卡死。 + +## 3. 图粒度 + +- **一个图就是整个 forward**，按 (phase, bucket) 各一个。图里包括 input_ids 构造、所有 decoder 层（Python 循环被 Dynamo 展开，`model/qwen3/model.py:896`）、LM head、设备端采样和状态提交（`neuron_forward_adapter.py:272-356`）。 + - prefill 和 decode 是两个图，以 `linear_idx` 是否为 None 区分。 + - NEFF 数 = `|num_seqs_buckets| × (|decode_ctx_buckets|+1)`（`docs/design/vllm/decode-context-length-bucketing.md:111-116`）。 +- **没有按层编译的区域**，也没有"一层编译一次、所有层复用"的机制。重复层的复用交给 **neuronx-cc modular flow**（O1/O2 下的 `--modular-flow-mac-threshold=10`）。对延迟敏感的 decode 图用 `-O3` 把它**关掉**。 +- giga-kernel 路线走得更远：GPT-OSS 全部 36 层用一个 NKI kernel 跑完，单个 decode NEFF，权重在 `[num_layers, ...]` 轴上堆叠（`docs/design/vllm/gpt-oss-gk-decode.md`）。 +- DeepSeek-V4.1 PR #71 约 120 个 NEFF 的做法（按 expert 分图、MoE 用 `@torch._dynamo.disable`、每层一次 host 同步）与仓库其它部分正好相反。见 `vllm_neuron_pr71_analysis.md`。 +- 编译耗时的应对办法是：TP 组内并行 trace 和编译（`parallel_compile`，`neuron_model_runner_v2.py:2035-2056`）、编译缓存、SPMD 单 NEFF 校验。不靠拆图。 + +## 4. bucket 与静态形状 + +- 默认值（`utils/bucket_utils.py:36-126`）： + - prefill token bucket：从 128 起取 2 的幂，直到 `max_num_batched_tokens`。 + - decode batch bucket：从 1 起取 2 的幂，直到 `max_num_seqs`。 + - 可选：KV segment bucket（前缀缓存用）和 decode 上下文长度 bucket。 +- `_select_bucket` 选能装下的最小 bucket（`neuron_model_runner_v2.py:903`）。 +- padding 规则： + - 填充 token 用 `PAD_TOKEN=0`。填充 slot 用 `-1`，由 `paged_slot_indices` 重定向到 null block 0。 + - block table 在每行有效块之后填 `-1` 哨兵，attention kernel 跳过越界部分。 + - 填充的 decode 行 `row_to_slot=-1`（`worker/neuron_input_build.py:1-13`；`docs/design/vllm/padding-batching-block-kv.md`）。 +- 元数据只在 host 上用 NumPy 构造，从不读设备状态。每一步在**图外**用 `torch.from_numpy(...).to(device)` 拷到设备（`neuron_model_runner_v2.py:2687-2712`）。 + - Python int（`max_query_len`、`block_size`、`kv_segment_size`）有意放进元数据 dict，让它们参与图特化。 + - decode 的 `input_ids` 在**图内**由上一步设备端采样结果构造，稳态 decode 没有 D2H。 +- **图内没有 `.item()`。** host 只在取输出时读一次（对采样 future 做 `.cpu().tolist()`）。 +- 标度另存一份 Python float，跟张量放在一起，因为 `tensor.item()` 会触发 graph break（`model/llama3/model.py:443`）。 + +## 5. 图内 collective + +- 世界组和模型并行组都是 **Gloo**。collective 被 trace 进 NEFF，运行时从不单独派发（`envs.py:858-866`）。 +- 用到的 API： + - `torch.distributed._functional_collectives.all_reduce(x, reduceOp="sum", group=...)`。reduceOp 必须是**字符串**，`ReduceOp` 对象会让 precompile 的 pickle 失败（`nn/embedding.py:256-260`、`nn/rpl.py:147-150`）。 + - `all_gather_tensor`。 + - `torch.ops._c10d_functional.all_gather_into_tensor`、`all_to_all_single`、`wait_tensor`（`functional/collectives/all_to_all.py:127-193`）。 + - 也用了普通的就地 `dist.all_reduce`（经 vLLM `GroupCoordinator`），Dynamo 会把它改写成函数式版本。 +- 限制： + - **trn2 上 all_to_all 是模拟的**：all_gather 加本地切片，理由是 "ENC 缺 Mesh 算法"（`all_to_all.py:10-33`）。 + - 所有 rank 必须 trace 出逐字节相同的 HLO（SPMD）。用运行时 `rank` 张量，不用 Python int；依赖 rank 的切片改成 `index_select(arange+offset)`。每个组注册完整的 replica-group mesh（`docs/design/compilation/spmd_compilation.md`、`patches/replica_mesh_patch.py`、`parallel/neuron_group_coordinator.py:122-130`）。 + - **collective 不能直接读 NEFF I/O 张量**（`NCC_INLA001`）。要先拷到中间 buffer（`functional/collectives/reduce_scatter_v.py:80-95, 318-330`）。 + +## 6. 已知 neuronx-cc 绕法 + +| 现象 | 绕法 | 位置 | +|---|---|---| +| 内联 NKI 后 birverifier 拒嵌套循环 | `--enable-verifier=false --enable-nested-dynamic-loop` | `neuron_model_runner.py:2071-2076` | +| `NCC_INLA001` collective 读了 I/O 张量 | collective 前先拷到中间 buffer | `reduce_scatter_v.py:89,177,324` | +| 广播布尔 AND 编错（`[4096,5120] & [1,5120]`） | 把有效性折进 position，只用一次 `where`（CHRS-985） | `attention/attention_segmented_cte.py:546` | +| `NCC_ITEN406` "too many partition dims"（跨步访问模式），常见于 graph break 之后 | 保持 fullgraph，避免跨步 view | `docs/model-dev/debugging.md:150-160` | +| `NCC_ILSA901` 融合 MLP 在 T=2/4 时 transpose 访问模式出错 | 走 in-tree MLP 路径 | `functional/qwen3_intree_mlp.py:174-202` | +| `NCC_INKI016` query seqlen 必须是 128 的倍数 | bucket 校验 | `utils/bucket_utils.py:455` | +| `NCC_IBTN243` lnc1 下 SendRecv 降不下去 | prefill 用 lnc2，单独起引擎 | `gpt-oss-gk-decode.md:19` | +| `-1` slot 变成越界 DGE scatter，卡死队列 | `paged_slot_indices` 钳到 null block | `model/kv_cache.py:80` | +| 大 buffer 里元素偏移超过 2**32 | KV 每层独立分配 | `neuron_model_runner.py:10988` | + +本仓库里没有搜到 `InferNonlocalTensors`。 + +## 7. FP8 + +- KV fp8 用 `torch.float8_e4m3fn` 存储（`utils/dtype_utils.py:7-14`）。数值钳到 `FP8_CLAMP_MAX`：**trn2 上是 240**（硬件/NKI 的 `float8_e4m3`），trn3 上是 448（`:16-41`）。这个常量在 import 时就解析好，所以在 trace 时是常量。写入路径：`(k*k_scale).clamp(±240).to(e4m3fn)`（`model/qwen3/model.py:311-321`）。 +- trn2 上的编译参数 `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3`（放在 hlo2tensorizer 字符串里）让编译器把 e4m3fn 当 e4m3 处理（`neuron_model_runner_v2.py:673-686`）。 +- 静态 fp8 权重（ModelOpt e4m3fn，范围 448）在**加载时乘 240/448 缩小**，反量化标度相应乘 448/240（`model/llama3/weight_loaders_static_fp8.py:45-80`、`model_static_fp8.py:126-160`）。平台判断要前缀匹配 `trn2` / `trn2-*`（`:150`）。 +- 个别 kernel 和 swa 路径对 448/240 的用法不一样（`attention/swa_fused.py:45`），要逐个核对 kernel 的数值范围。 + +## 8. 与 torch/dynamo 相关的 patch + +- `vllm/patches/*` 改的是 vLLM，不是 Dynamo/AOT（DP 负载均衡、shm 自旋、关掉 pin_memory、`accelerator.empty_cache` 保护、端口占位、replica mesh）。 +- 值得照搬的设置： + - `cache_size_limit=2**62`，serve 时加 `fail_on_recompile` stance。 + - warmup 和 serve 的 grad 模式必须一致，因为 Dynamo 会对 GLOBAL_STATE 加 guard。lite 在 `no_grad` 下 trace，不是 `inference_mode`（`utils/neuron_utils.py:35-55`）。 + - traced 代码里不要打日志。连 `is_native_backend()` 都可能打日志并导致 graph break，所以它在构造时就读好（`nn/embedding.py:98-100`）。 + - 用 `TORCH_CACHING_PRECOMPILE=1` 和 `TORCHINDUCTOR_AUTOGRAD_CACHE=1` 开 precompile 和 AOTAutograd 缓存（`vllm/inductor_cache.py:1-22`）。 + +## 9. 本仓库里的 DeepSeek-V4/V4.1 + +没有。`contrib/models/` 下只有 `template/`；`vllm_neuron/model/` 下是 gpt_oss、llama3、qwen3、qwen3_vl、qwen3_asr、synthetic。"deepseek" 只出现在注释里（V3.2 indexer 的 key-only spec，`model/kv_cache.py:54`）。PR #71 在这里没有合并。 + +## DeepInfer 该照搬什么（按影响排序） + +1. **KV 按层独立、互不重叠、view 记忆化。** 每层一个张量（`(2,N,H,B,D)`，或 K、V 分开），挂成模块属性。同一块区域绝不切第二次。这样就没有 synthetic base 和 `as_strided` 路径，也躲开了 2**32 偏移的 bug。 +2. **KV 写入用普通 `index_put_`，索引先钳位**（`paged_slot_indices`，非法 slot 指向 null block 0）。依赖后端把改写转成输出别名（donation），不写自定义 op。 +3. **每个 (phase, bucket) 一个 full-graph forward**，采样和状态提交都放在图内。重复层交给 neuronx-cc modular flow 去重（O1/O2），decode 可以试 `-O3`。不要按层或按 expert 拆图。 +4. **编译参数原样照抄**：`--auto-cast=none`；hlo2tensorizer 选项只写一个字符串（`--modular-flow-mac-threshold=10`，trn2 上再加 fp8 参数）；`--internal-backend-options=--enable-verifier=false --enable-nested-dynamic-loop`；如果 `transformer` 有负面影响，加 `--model-type=generic`。 +5. **元数据在 host 构造、形状固定、在图外做 H2D。** decode 的 input_ids 来自设备端上一步采样结果。图内不用 `.item()`。serve 时开 `fail_on_recompile`。 +6. **函数式 collective + 字符串 reduceOp + SPMD rank 张量；collective 不直接读图 I/O。** trn2 上要实测 all_to_all，vllm-neuron 是用 all_gather 模拟的。 +7. **FP8：e4m3fn 存储，钳 ±240，加 unsafe-e4m3fn-as-e4m3 参数**；OCP checkpoint 的权重按 240/448 缩小。 + +
diff --git a/docs/neuron/vllm_neuron_pr71_analysis.md b/docs/neuron/vllm_neuron_pr71_analysis.md new file mode 100644 index 0000000..9cd0e1c --- /dev/null +++ b/docs/neuron/vllm_neuron_pr71_analysis.md @@ -0,0 +1,292 @@ +# vllm-neuron PR #71 (DeepSeek-V4.1-Flash + DSpark on trn2): analysis for the DeepInfer Neuron port + +Source: [vllm-project/vllm-neuron#71](https://github.com/vllm-project/vllm-neuron/pull/71), head `30d54d1` (OswaldHe/vllm-neuron), open, not merged, read 2026-10-22. +Paths below are relative to `contrib/models/deepseek_v41/src/` unless prefixed. `model.py:N` = line N at that SHA. +Note: the call-site line numbers in `kernels/__init__.py` are stale; the ones here come from `model.py` itself. + +## 0. Summary + +- **No XLA in the model code.** It never imports `torch_xla` or NxDI. It uses `libtorch_neuronx_lite` (`wrap_nki` for NKI, `model.py:401`), `torch.compile(backend=vllm_neuron.envs.get_compile_backend_name())` (`model.py:445`, `:2234`), `torch.ops._c10d_functional` collectives (`model.py:225-240`) and `vllm_neuron.*` helpers. Caveat: the numbers were measured on vllm-neuron 0.24.0.1.1.0, and the README lists `torch`/`torch-xla` 2.11 in that stack. In 0.24 the default compile route is `neuron_libtorch` (lite on top of XLA); the native route needs `VLLM_NEURON_BACKEND=neuron_native`. The PR never says which route it ran. Its error strings ("Detected non-contiguous slicing for requested Device Tensor") look like the native device. Either way, the kernels and the model logic do not depend on XLA. +- **Not a whole-model static graph.** About 120 NEFFs (`__init__.py:83`), one graph execution costs about 1.12 ms (`model.py:2353`), and every MoE layer syncs with the host (`model.py:2775`). The decode step (ITL) is about 320-370 ms. TPOT is about 117 ms only because DSpark accepts about 2.75 tokens per step (`TPOT × (1 + 5 × accept)` = step time, `model.py:283`). +- **Parallel layout.** Attention runs TP=8 and is replicated 8x. Experts are partitioned by **expert id** (6 of 384 per rank), so this is EP-style, not TP-sharded experts. Tokens are replicated, and there is one world all-reduce per MoE layer. Engram is row-partitioned across 64 ranks and resident in HBM. Everything lives in HBM: about 13.9 GiB per rank. +- **Kernels.** There are 12 NKI kernels, all written against `nki` 0.6 (`import nki`, `nki.isa` keyword `dst=` form, `@nki.jit`), the same `nki>=0.6` that our `requirements-neuron.txt` pins. There is no MoE kernel. Expert GEMMs are `blockquant.fp4_linear_2core`, run once per expert from a host loop. + +## 1. Parallelism and placement + +| Part | Layout | Ref | +|---|---|---| +| MLA attention | `ATTENTION_TENSOR_PARALLEL_SIZE=8` (capped by `o_groups=8`), so the attention is replicated `64/8 = 8` times. Each rank holds 8 of 64 query heads. `wq_b` is column-split by head, `wo_a` is split by o-group (1 local group), `wo_b` is row-split. `wq_a`, `wkv`, norms and `attn_sink` (sharded with the heads) are replicated. One fp32 all-reduce of `[T,5120]` over the 8-rank attention group, which is 8 consecutive ranks on 2 neighbouring devices. | `config.py:30-43`, `model.py:4099-4124`, `:1917-1946` | +| Indexer | `index_n_heads=32` split by 8 (4 per rank). `wq_b` and `weights_proj` are split by head, `wk`/`k_norm` are replicated. The score is all-reduced over the attention group, then top-k with `index_topk=512`. Only `index_source_layer_ids` [2,8,14,20,24,28,32,36] run an indexer; the other layers reuse `shared.topk_idxs`. | `config.py:117-124`, `model.py:1471-1605`, `:1782-1861` | +| Compressor / compressed KV | Replicated weights (bf16). Only `kv_source_layer_ids` [2,8,14,20] own a compressed cache. Later layers with the same ratio reuse the rows read by their source layer; this saves memory only, with ±1.3% latency. | `model.py:1412-1468`, `:1640-1690` | +| Dense, shared expert, router | Replicated on all 64 ranks: shared expert in fp8, gate in bf16, mHC `fn/base/scale`, norms. | `model.py:4127-4140` | +| Embedding / LM head | Vocab-parallel over 64 ranks (aligned to 128, padded). Embedding uses mask + world all-reduce; logits use a world all-gather. | `model.py:4468-4512` | +| Routed experts | Partitioned by id: rank r owns experts `[6r, 6r+6)` of 384 in each of the 40 backbone layers. The 3 MTP stages have 128 experts each (2 per rank). All are resident in HBM in fp4 (4.3 GiB per rank), byte-for-byte as in the checkpoint. Weight `[n, k/2]` uint8 holds two e2m1 values per byte; scale `[n, k/32]` uint8 holds ue8m0. Fp8 weights are `[n,k]` uint8 e4m3 with a `[ceil(n/32), ceil(k/32)]` ue8m0 scale. Gate and tokens are replicated, so every rank computes the same routing. It runs its own experts on its rows and scatter-adds into an fp32 carry, then **one world all-reduce** of `[T,5120]` fp32 runs. There is no dispatch collective. | `weight_loader.py:13-33`, `model.py:2616-2928`, `:657-706` | +| Engram | 2 tables (layers 1 and 14) of 384M rows × 256. The fp8 checkpoint is dequantized to **bf16 at load** and row-partitioned: about 6.0M rows per rank, 2.86 GiB per table, **in HBM**. Hashing runs on device (`NgramHashState`). Each rank looks up its own rows (`F.embedding` + mask), then a world all-reduce of `[1,T,24,256]` bf16 runs. A host streaming fallback exists but is not used. | `weight_loader.py:29-33`, `model.py:3112-3215`, `:4680-4706` | +| mHC | 4 streams. All coefficients and the 20-step Sinkhorn run locally (replicated). There are no collectives. | `model.py:3217-3283` | +| DSpark (`mtp.0-2`) | Same `Attention`/`MoE`/`BlockQuantLinear` classes, so the same layout. Markov head embed/head are vocab-parallel. | `model.py:3549-3740`, `:4141-4157` | + +**Collectives per forward (batch 1):** each of the 40 layers has 1 attention-group all-reduce and 1 world all-reduce (MoE). The 8 indexer layers add 1 attention-group all-reduce each. The 2 engram layers add 1 world all-reduce each. Then there is 1 world all-reduce for the embedding and 1 world all-gather for the logits, and the DSpark stages add their own attention/MoE/embedding/markov collectives. This totals the "~43 world-wide all-reduces" that `model.py:248` cites. Measured cost: a 64-rank all-reduce takes about 1.9 ms against about 0.1 ms for 8 consecutive ranks. The cost is per-call latency, not bytes (`model.py:248-262`), so world collectives take about 80 ms per step. Factoring the world into 8×8 was 2.8% *slower*. TP×PP was probed (`DEEPSEEK_V41_PP_EMULATE`, `kernels/pp_permute.py`) but does not ship. + +## 2. Execution model + +- **Compile unit.** The runner compiles the model forward with graph breaks allowed (`VLLM_NEURON_DEBUG_MODE=1` gives `fullgraph=False`, `__init__.py:34`). The model also creates its own compiled entry points with `torch.compile(..., fullgraph=True, options=_COMPILE_OPTIONS)`: + - one per NKI kernel (`_nki_entry`, `model.py:386-450`); + - one per block for `attn_half`, which covers mHC, norm, attention, indexer, compressor and KV writes (`_block_entry`, `model.py:3432`, `:5025-5051`); + - one per small step: router (`_gate_step`), each routed-expert step, shared expert, `_post_into`, and the collectives (`_compiled`, `model.py:2257`). + + `MoE.forward` is `@torch._dynamo.disable` (`model.py:2851`). Dynamo limits are raised to 2^30 because more than 256 variants otherwise make dynamo silently fall back to eager (`model.py:192-223`). Request-time compilation is allowed because the graph set depends on what the router picks (`model.py:5215-5249`). +- **Compiler flags** (`model.py:180-190`): `--auto-cast=none -O1`, `--modular-flow-mac-threshold=10`, `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3` (required by the fp8 kernels), `--enable-verifier=false --enable-nested-dynamic-loop`. +- **Buckets.** There is no separate context-encoding / token-generation graph family; prefill and decode go through the same per-layer graphs, keyed by row count. + - Prefill is segmented: `max_num_batched_tokens` is 4096 at 64 ranks (`__init__.py:451-561`). TTFT fits `ceil(tokens/segment) × F + tokens × ~1.06 ms`, with F ≈ 325-783 ms depending on segment and context. + - Decode has `(1+5) × num_seqs` rows. `MAX_NUM_SEQS=1` because runs with more stalled (`config.py:45-49`). + - Expert row lists are padded to powers of two with a floor of 128 (`model.py:2503`, `:2790`), because the GEMM tiles M at 128. + - The compressed-KV read length is pinned to `max_model_len`, so prefill chunks and decode share graphs (`model.py:1178-1185`, `:4425-4440`). The cost is that a short request on a long-context engine reads 13x more. +- **KV layout.** + - The window KV is vLLM's paged allocation, one per layer, `[blocks, 1, block_size, 512]` bf16. It is declared as full attention because the paged slots also index the other per-slot caches (`model.py:4268-4312`). + - The model allocates flat per-slot buffers for kv-source layers: `compress_kv_cache [slots,512]` bf16, `index_k_cache [slots,128]` bf16, and compressor `pool_kv`/`pool_score [slots,512]` bf16 (`model.py:4314-4377`, `:1431-1439`). + - Quantization is **emulated**: values pass through fp8 QDQ (window KV) or fp4 QDQ (compressed KV, index K/Q) and are **stored as bf16** (`model.py:1770`, `:1543-1557`, `:1823`). + - Addressing goes through `PagedSpan` (`model.py:1302-1392`). +- **DSpark inside the target forward.** `forward` (`model.py:4515-4677`) runs 40 blocks and taps the attention inputs of layers 37/38/39 (`_dspark_tap`, mean over the 4 streams). It then runs the LM head, the on-device sampler and the rejection sampler. `_draft` → `forward_spec` (`model.py:5053-5145`) then runs the 3 MTP stages over a 5-token draft block (seed + noise tokens). Their attention reads the main window and the draft block's own KV, with no compressed KV. The last stage samples 5 tokens one at a time through a rank-256 Markov logit bias and adds a confidence head. Drafts are stored in `self.last_draft`, and `DSparkProposer.propose` (`drafter.py`) hands them to vLLM. There is no separate draft NEFF or warm-up. This needs the plugin's proposer registry (the `vllm_neuron/vllm/spec_decode` and runner hunks in this PR). +- **Host syncs are not avoided.** + - Each MoE layer copies `indices.cpu()` to the host and builds row lists there (`model.py:2775-2800`). + - Span and slot planning calls `.tolist()` on `slot_mapping` and `block_table` (`model.py:4766`, `:4831`, `:4943`). + - `int(start_pos)` is read in `Block.forward` (`model.py:3332`). + - The fixed cost per graph execution (about 1.12 ms) is the bottleneck the comments identify (`model.py:2351-2370`). Experts can be fused 2-3 per graph through `DEEPSEEK_V41_EXPERT_FUSE`, which is off by default. + +## 3. Kernel inventory + +Common to all kernels: +- `import nki; import nki.isa as nisa; import nki.language as nl`; entry points are `@nki.jit`. +- Outputs are allocated as `nl.shared_hbm` and **returned** (functional style). Scratch space is `nl.private_hbm`. +- 128-row tiles with a tail; reductions accumulate in fp32. +- Called through `wrap_nki(kernel)` (`libtorch_neuronx_lite.nki.nki_hop`) inside `torch.compile(fullgraph=True)`. `[2]` sets the grid for the 2-core variants. + +"Ours" means the shim or op it replaces in DeepInfer (`vendor/` = `deepinfer/neuron_ops/vendor/`, `tl/` = `deepinfer/model_executor/layers/tilelang_ops/neuron/`, `v4/` = `deepinfer/model_executor/models/v4_modules/`). + +| Kernel (entry, line) | Computes; I/O | Ours | Reuse on torch-neuronx 2.14 | +|---|---|---|---| +| `blockquant.fp8_linear` :399, `_raw` :417, `_2core` :477, `_raw_2core` :516 | Computes `x[m,k] @ W[n,k]^T`. The activation is first rounded through 1x32 fp8 (`_raw` skips this). The weight (uint8 e4m3 bytes plus a uint8 ue8m0 `[n/32,k/32]` scale) is decoded with bit arithmetic into **bf16 private HBM** on every call, then multiplied on the PE in bf16 (`nisa.nc_matmul`, fp32 PSUM). Output is `[ceil(m/128)*128, n]` bf16, and the caller slices it. `_2core` splits n across the 2 physical cores (`nl.program_id`, grid 2). | `vendor/deep_gemm/gemm.py:205` `fp8_fp4_gemm_nt` (dense, `v4/linear.py`), `einsum.py:77` `fp8_einsum` (`wo_a`, `v4/attn.py:215`), and the `tile_kernels.quant.per_token_cast` that comes before them | The kernel source can be reused as-is. Needs: the `nki_op` adapter; weights kept as checkpoint bytes (no `transform_sf_into_required_layout` repacking); the fp8 compiler flag; and grid=2 launch support on our stack for `_2core`. **Numerics: activation QDQ divides by 224, not 448** (Trn2 e4m3 saturates at 240), one mantissa bit off from DeepInfer's reference. | +| `blockquant.fp4_linear` :662, `fp4_linear_2core` :681 | Same pattern for fp4 experts: weight `[n,k/2]` holds two e2m1 values per byte, scale `[n,k/32]` ue8m0, output scaled by 2^64 after the matmul. 2core cut TPOT from 142.8 to 121.4 ms (`model.py:474-480`). | `gemm.py:316` `m_grouped_fp8_fp4_gemm_nt_contiguous` (`v4/mlp.py:119`) | Reusable. Called once per expert; our grouped API needs a loop or wrapper over local experts. | +| `blockquant.bf16_linear` :1221, `fp32w_linear` :1204 | Plain bf16 GEMM for the router, LM head, compressor, indexer `wk`/`weights_proj` and Markov head. `fp32w` narrows an **fp32 weight to bf16** in-kernel (mHC `fn`). | `gemm.py:107` `bf16_gemm_nt`; `vendor/deep_gemm/hc.py:10` `tf32_hc_prenorm_gemm` | Reusable. Precision note: bf16 instead of tf32 for the mHC prenorm GEMM. | +| `blockquant.*_fused`, `_tfused`, `_kmajor` | Variants that decode the weight in SBUF. `fp4_linear_fused` was measured **+50% TTFT / +24% TPOT** (`model.py:536-552`); the others are unreachable. | — | Skip. | +| `sparse_attn.gathered_attn_idx` :180 (live), `gathered_attn` :51 | For each token: indirect-DMA gather of K slots from `kv[n,512]` bf16 by `idx[T,K]` int32 (`kv.ap(vector_offset=…, indirect_dim=0)`); QK on bf16 operands into fp32; bias (column 0 = softmax scale, then -1e30 for empty slots); softmax with a learned sink; bf16 P·V. Input q is `[T*H,512]` bf16 with H=8 local heads; output `[T*H,512]` bf16. Requires K%128==0, d%128==0, H≤128. K = 128 window + 512 top-k = 640. | `vendor/flash_mla/flash_mla_interface.py:282` `flash_mla_with_kvcache`, `:155` `flash_mla_sparse_fwd` | The kernel can be reused, but **its KV layout differs**: it reads bf16 rows, while our cache stores fp8+ue8m0 records (528 B/token). Either switch the Neuron KV cache to bf16 (1024 B/token; fine at bs1 / 8K) or gather+dequantize in torch and call `gathered_attn`. Indirect DMA needs a device check under `nki_op`. | +| `index_score.weighted_scores` :41 | Computes Σ_h w[h,t]·relu(q_h·k). Inputs: q `[T*Hi,128]` bf16 (Hi=4 local), k `[n,128]` bf16 (n%128==0), w `[Hi,T]` fp32. Output `[T,n]` fp32 (PE matmul, then a second matmul for the head sum). | `vendor/deep_gemm/attention.py:123/225` `fp8_fp4_(paged_)mqa_logits` | Reusable once q/k are dequantized to bf16 (the PR runs fp4 QDQ in torch first). Top-k is **not** in the PR: it calls `vllm_neuron.functional.topk` (nkilib `gpsimd_topk` / vendored `rotational_topk`), which is also relevant to our #22 and `vendor/deep_select`. | +| `compress_pool.softmax_pool` :42 | Computes `(kv * softmax(score, dim=ratio)).sum(ratio)`; kv and score are `[G,ratio,512]` bf16, output `[G,512]` fp32. | `tl/compress.py:36/132` `compress_and_update_state_{prefill,decode}_tl` (pooling part) | Reusable. Small. | +| `hc_mixes.kernel` :66 | From raw `[R,24]` fp32 (the `fn` GEMM) and the stream `[R,4*5120]`: rms over the stream, then pre = σ+eps, post = 2σ, comb = softmax + **20 Sinkhorn iterations**. Outputs pre/post `[R,4]` and comb `[R,16]`, all fp32. | `tile_kernels.mhc` `norm_fn` / `pre_split_mixes` / `sinkhorn` (`v4/mhc_ops.py:10-16`) | Reusable. Check the iteration count and eps against our config (`sinkhorn_normalize_ref` defaults to 10 iterations). | +| `hc_expand.kernel` :28 | `[B,S,D]` → 4 copies `[B,S,4,D]` plus a one-hot pre_mix `[B,S,4]` fp32. DMA only. | `tile_kernels.mhc.expand_kernel.expand_to_mhc_fwd` | Reusable. | +| `hc_collapse.kernel` :24 | Σ_c pre[c]·x[c]: `[B,N,4,D]`, `[B,N,4]` fp32 → `[B,N,D]`. | `mhc_pre_apply_mix_fwd` | Reusable. | +| `hc_attn_out.kernel` :24 | post[j]·x + Σ_i comb[i,j]·res[i]: x `[B,S,D]`, res `[B,S,4,D]`, post `[B,S,4]`, comb `[B,S,4,4]` → `[B,S,4,D]`. | `mhc_post_fwd` | Reusable. | +| `rmsnorm.rms_norm` :30 | x `[R,D]`, weight `(1,D)`, eps as a `(1,1)` tensor → `[R,D]` in the input dtype. fp32 reduction plus one Newton step on rsqrt. | `tl/norm.py:75` `norm_tl`, `v4/norm.py` | Reusable. Check parity against `rstd_1ulp`. | +| `swiglu.swiglu` :40 | Clamped SwiGLU (up clamped to ±10, gate to ≤10) × routing weight: gate/up `[R,I]`, weights `(R,1)` or `(1,1)`, limit `(1,1)` → `[R,I]`. | `tile_kernels.quant.swiglu_forward_and_per_token_cast` (`v4/mlp.py:19`) | Reusable. It does not cast; the fp8 rounding is done by the next GEMM's QDQ. | +| `engram_gate.gated_add` :41 | For each of the 4 streams: rms(h)·rms(key) normalized dot → sign·sqrt → sigmoid gate; h + gate·value. h/key `[R,4*5120]`, value `[R,5120]`, weight `(1,4*5120)` = q_weight·k_weight. | `tile_kernels.engram.engram_gate_fwd` (`v4/engram.py:155`) | Reusable after checking it against our fused-weight semantics. | +| `pp_permute.kernel` :395 | `nki.collectives.collective_permute` plus `nisa.core_barrier` stage hand-off. **Not called**: the CC runtime refuses its replica group. | — | Skip. | + +**MoE GEMM and dequant without a MoE kernel** (`model.py:2287-2380`, `:2851-2928`): +1. Router: `bf16_linear` + sqrt-softplus + nkilib top-k. +2. The host copies the indices back and builds a per-local-expert row list padded to a bucket of at least 128. +3. One compiled graph per expert: `index_select` rows → `fp4_linear_2core` (w1, w3) → `swiglu` → `fp4_linear_2core` (w2) → `index_add` into the fp32 carry. +4. World all-reduce. +5. Shared expert: the same chain with `fp8_linear_2core`, applied to all rows. + +Dequantization happens inside every GEMM (weight bytes → bf16 in private HBM); there is no persistent bf16 copy. + +## 4. Where the time goes and what makes TPOT ≈ 120 ms possible + +The PR has no profile breakdown. The comments give these pieces: +- Graph execution: about 1.12 ms fixed each, about 120 per forward. A 128-row expert GEMM's roofline is about 154 µs, so most of the routed-expert time is overhead, not arithmetic (`model.py:2353-2363`). +- World all-reduces: about 43 × 1.9 ms ≈ 80 ms per step. +- Decode step ≈ 320-370 ms for 6 rows. TPOT = step / (1 + accepted), and the acceptance length is 2.75 on sonnet (36% per-position average). +- Prefill: F ≈ 0.3-0.8 s per forward plus about 1.06 ms/token. + +The decisions that matter: +1. **Everything resident in HBM** (13.9 GiB per rank): no host traffic during the forward. +2. **Attention TP=8, replicated**, so the attention reduction stays within 2 devices. +3. **Experts by id with replicated tokens**: no dispatch, one all-reduce. +4. **Using both physical cores** of each LNC=2 core in the GEMMs: fp4 then fp8 2core brought TPOT from 142.8 to 116.9 ms. +5. **NKI for every op family**, to cut the number of graphs and ops. +6. **Prefill segment 4096.** +7. **DSpark** (≈2.75x). + +Things they measured as no win: SBUF-fused dequant, the shared compressed read, `NEURON_RT_IO_RING_CACHE_SIZE`, 8×8 collective factoring, and `_EXPERT_ROW_FLOOR` below 128. + +## 5. Integration plan for our port (ordered by expected payoff) + +1. **#9 static graphs + #12 placement (structural, from about 150 s to under a second).** + - Compile one decode graph per bucket with the NKI ops inside it. + - Remove every host sync. To avoid PR #71's per-layer `.cpu()`, run **all 6 local experts over the bucketed rows**, with routing weight 0 for rows that are not routed to them. Their GEMM pads M to 128 anyway, so at decode this costs the same instructions per expert as their 128 bucket, and the graph stays static. + - At bs1, consider **attention TP=8 groups instead of attention DP**: today one rank does all 64 heads, and the other 63 wait. + - Once tokens are replicated, our all-gather dispatch is a no-op and the combine becomes one world all-reduce, like theirs. Measure all-reduce against reduce-scatter. + - Copy their `_COMPILE_OPTIONS` flags. Budget about 1.9 ms per world collective. +2. **#10 GEMM: reuse `blockquant`.** + - Put `fp8_linear(_raw)_2core` behind `fp8_fp4_gemm_nt`/`fp8_einsum`, `fp4_linear_2core` behind `m_grouped_fp8_fp4_gemm_nt_contiguous` (loop over local experts), and `bf16_linear`/`fp32w_linear` behind bf16/hc GEMMs. + - Keep weights as checkpoint uint8 bytes (this also follows our fp8-as-bytes rule). + - Decide whether to accept the 224 QDQ or patch the kernel to 448 and re-test the saturation flag. + - Highest kernel payoff. +3. **#15 DSpark.** Their TPOT includes a 2.75x multiplier. Without it we need a step of 117 ms or less, not about 320 ms. The design to follow is in-forward drafting (taps from layers 37-39, a 5-token block, a Markov head). +4. **#11 attention/indexer.** + - Choose the KV layout: bf16 cache plus `gathered_attn_idx`, or keep fp8 records with gather+dequant plus `gathered_attn`. + - Put `index_score` behind `fp8_fp4_(paged_)mqa_logits`, and `compress_pool` behind `tl/compress.py`. + - Borrow `vllm_neuron.functional.topk` (nkilib) for `deep_select.topk` and #22. +5. **#14 Engram.** Copy their placement: row-partitioned bf16 (or fp8, 2.9 GiB) tables in HBM, on-device hash and gather, and a world all-reduce. This removes the host gather and H2D entirely. It needs HBM headroom (#21): their dense part is 3.9 GiB per rank because attention is TP=8 and the vocab is split 64 ways. Add `engram_gate`. +6. **#13 small ops.** `hc_mixes`, `hc_expand`, `hc_collapse`, `hc_attn_out`, `rmsnorm` and `swiglu` are each small, but together they remove most of the per-layer op count. Do them after #9 so the profile is real. + +**Things to check before reusing any kernel on our stack:** +- Whether `torch_neuronx.nki_op` (2.14.3) accepts `@nki.jit` kernels that *return* `shared_hbm` outputs, use `private_hbm` scratch, run under both eager and `torch.compile(backend="neuron")`, and launch with grid=2 at LNC=2. +- Whether indirect DMA (`.ap(vector_offset=…)`) works through `nki_op`. +- The shape asserts (k%128, K%128, heads≤128). +- Numeric parity against our CPU shims within quantization noise. + +**License.** +- vllm-neuron is Apache-2.0 (`LICENSE`; `NOTICE`: "Copyright Amazon.com, Inc. or its affiliates"). DeepInfer is Apache-2.0 too, so the licenses are compatible. +- The kernel files have no per-file header. Only `kernels/__init__.py`, `model.py`, `weight_loader.py` and similar files carry `SPDX-License-Identifier: Apache-2.0`. +- When copying, add an SPDX header and an origin line to each copied file (repo, PR #71, SHA `30d54d1`, author Zifan He / OswaldHe), mark modified files as changed (Apache-2.0 §4b), and carry the relevant NOTICE text into our NOTICE or third-party notices. +- The PR is unmerged, so pin the SHA and re-check after merge. + +
+中文 + +# vllm-neuron PR #71（trn2 上的 DeepSeek-V4.1-Flash + DSpark）：面向 DeepInfer Neuron 移植的分析 + +来源：[vllm-project/vllm-neuron#71](https://github.com/vllm-project/vllm-neuron/pull/71)，head `30d54d1`（OswaldHe/vllm-neuron），状态 open、未合并，阅读日期 2026-10-22。 +下文路径若无前缀，都相对于 `contrib/models/deepseek_v41/src/`。`model.py:N` 指该 SHA 下的第 N 行。 +注意：`kernels/__init__.py` 里写的调用点行号已经过时，本文行号直接取自 `model.py`。 + +## 0. 摘要 + +- **模型代码里没有 XLA。** 不 import `torch_xla`，也不用 NxDI。它用的是 `libtorch_neuronx_lite`（NKI 走 `wrap_nki`，`model.py:401`）、`torch.compile(backend=vllm_neuron.envs.get_compile_backend_name())`（`model.py:445`、`:2234`）、`torch.ops._c10d_functional` collective（`model.py:225-240`），以及 `vllm_neuron.*` 工具函数。需要说明的是：性能数字测于 vllm-neuron 0.24.0.1.1.0，README 列出的这套环境里有 `torch`/`torch-xla` 2.11。0.24 的默认编译路线是 `neuron_libtorch`（lite 跑在 XLA 之上），要走 native 路线得设 `VLLM_NEURON_BACKEND=neuron_native`。PR 没写明用的是哪条路线；从报错字符串（"Detected non-contiguous slicing for requested Device Tensor"）看像是 native device。无论哪条路线，kernel 和模型逻辑都不依赖 XLA。 +- **不是整模型静态图。** 约 120 个 NEFF（`__init__.py:83`），每次图执行约 1.12 ms（`model.py:2353`），每个 MoE 层都要和 host 同步一次（`model.py:2775`）。decode 一步（ITL）约 320-370 ms。TPOT 能到约 117 ms，只是因为 DSpark 每步平均接受约 2.75 个 token（`TPOT × (1 + 5 × accept)` = 一步时间，`model.py:283`）。 +- **并行布局。** attention 用 TP=8，并复制 8 份。专家按**专家 id** 划分（每 rank 拿 384 个中的 6 个），也就是 EP 式，不是对专家做 TP 切分。token 在所有 rank 上复制，每个 MoE 层做一次全局 all-reduce。Engram 按行切到 64 个 rank，常驻 HBM。全部权重都在 HBM 里，每 rank 约 13.9 GiB。 +- **kernel。** 共 12 个 NKI kernel，全部基于 `nki` 0.6 编写（`import nki`、`nki.isa` 的 `dst=` 关键字写法、`@nki.jit`），与我们 `requirements-neuron.txt` 锁定的 `nki>=0.6` 是同一个包。没有 MoE kernel。专家 GEMM 用的是 `blockquant.fp4_linear_2core`，由 host 循环逐个专家调用。 + +## 1. 并行与放置 + +| 部分 | 布局 | 引用 | +|---|---|---| +| MLA attention | `ATTENTION_TENSOR_PARALLEL_SIZE=8`（上限由 `o_groups=8` 决定），所以 attention 复制 `64/8 = 8` 份。每 rank 持有 64 个 query head 中的 8 个。`wq_b` 按 head 列切，`wo_a` 按 o-group 切（本地 1 组），`wo_b` 行切。`wq_a`、`wkv`、各 norm 和 `attn_sink`（随 head 一起切）都复制。每层在 8-rank attention 组上做一次 `[T,5120]` 的 fp32 all-reduce，这个组是相邻 2 个 device 上连续的 8 个 rank。 | `config.py:30-43`、`model.py:4099-4124`、`:1917-1946` | +| Indexer | `index_n_heads=32` 切 8 份（每 rank 4 个）。`wq_b` 和 `weights_proj` 按 head 切，`wk`/`k_norm` 复制。score 在 attention 组上 all-reduce，然后做 `index_topk=512` 的 top-k。只有 `index_source_layer_ids` [2,8,14,20,24,28,32,36] 这几层跑 indexer，其他层复用 `shared.topk_idxs`。 | `config.py:117-124`、`model.py:1471-1605`、`:1782-1861` | +| Compressor / 压缩 KV | 权重复制（bf16）。只有 `kv_source_layer_ids` [2,8,14,20] 拥有压缩 cache。后面同 ratio 的层复用源层读出的行；这只省内存，延迟变化在 ±1.3% 以内。 | `model.py:1412-1468`、`:1640-1690` | +| 稠密层、共享专家、router | 在全部 64 个 rank 上复制：共享专家为 fp8，gate 为 bf16，还有 mHC 的 `fn/base/scale` 和各 norm。 | `model.py:4127-4140` | +| Embedding / LM head | 词表在 64 个 rank 上切分（按 128 对齐并补齐）。embedding 用掩码加全局 all-reduce；logits 用全局 all-gather。 | `model.py:4468-4512` | +| 路由专家 | 按 id 划分：40 个 backbone 层中，每层 384 个专家里 rank r 拥有 `[6r, 6r+6)`。3 个 MTP stage 每个有 128 个专家（每 rank 2 个）。全部以 fp4 常驻 HBM（每 rank 4.3 GiB），字节与 checkpoint 完全一致。权重 `[n, k/2]` uint8，每字节存两个 e2m1；标度 `[n, k/32]` uint8，存 ue8m0。fp8 权重为 `[n,k]` uint8 e4m3，配 `[ceil(n/32), ceil(k/32)]` 的 ue8m0 标度。gate 和 token 都是复制的，所以每个 rank 算出的路由完全相同。各 rank 只对路由到自己专家的行做计算，scatter-add 到 fp32 累加器，然后做**一次全局 all-reduce**（`[T,5120]` fp32）。没有 dispatch collective。 | `weight_loader.py:13-33`、`model.py:2616-2928`、`:657-706` | +| Engram | 2 张表（第 1、14 层），每张 384M 行 × 256。fp8 的 checkpoint 在**加载时反量化为 bf16**并按行切分：每 rank 约 6.0M 行，每张表 2.86 GiB，**放在 HBM**。hash 在 device 上计算（`NgramHashState`）。各 rank 查自己那部分行（`F.embedding` 加掩码），然后对 `[1,T,24,256]` bf16 做全局 all-reduce。代码里有 host 流式回退路径，但实际没用。 | `weight_loader.py:29-33`、`model.py:3112-3215`、`:4680-4706` | +| mHC | 4 条流。所有系数和 20 步 Sinkhorn 都在本地计算（复制）。没有 collective。 | `model.py:3217-3283` | +| DSpark（`mtp.0-2`） | 用的是同样的 `Attention`/`MoE`/`BlockQuantLinear` 类，所以布局相同。Markov head 的 embed/head 按词表切分。 | `model.py:3549-3740`、`:4141-4157` | + +**每次 forward 的 collective（batch 1）：** 40 层中每层有 1 次 attention 组 all-reduce 和 1 次全局 all-reduce（MoE）。8 个 indexer 层各多 1 次 attention 组 all-reduce。2 个 engram 层各多 1 次全局 all-reduce。另外还有 embedding 的 1 次全局 all-reduce、logits 的 1 次全局 all-gather，DSpark stage 也有自己的 attention/MoE/embedding/markov collective。合起来就是 `model.py:248` 说的"约 43 次全局 all-reduce"。实测成本：64-rank all-reduce 约 1.9 ms，连续 8 rank 约 0.1 ms。开销来自每次调用的固定延迟，与字节数无关（`model.py:248-262`），所以全局 collective 每步约占 80 ms。把全局拆成 8×8 反而**慢** 2.8%。TP×PP 做过探测（`DEEPSEEK_V41_PP_EMULATE`、`kernels/pp_permute.py`），但没有上线。 + +## 2. 执行模型 + +- **编译单元。** runner 编译模型 forward 时允许 graph break（`VLLM_NEURON_DEBUG_MODE=1` 对应 `fullgraph=False`，`__init__.py:34`）。模型自己还用 `torch.compile(..., fullgraph=True, options=_COMPILE_OPTIONS)` 建了一批编译入口： + - 每个 NKI kernel 一个（`_nki_entry`，`model.py:386-450`）； + - 每个 block 的 `attn_half` 一个，覆盖 mHC、norm、attention、indexer、compressor 和 KV 写入（`_block_entry`，`model.py:3432`、`:5025-5051`）； + - 每个小步骤一个：router（`_gate_step`）、每次路由专家计算、共享专家、`_post_into`，以及各 collective（`_compiled`，`model.py:2257`）。 + + `MoE.forward` 标了 `@torch._dynamo.disable`（`model.py:2851`）。dynamo 的上限调到了 2^30，因为图变体超过 256 个后，dynamo 会静默退回 eager（`model.py:192-223`）。由于图的集合取决于 router 的选择，所以允许在请求时编译（`model.py:5215-5249`）。 +- **编译参数**（`model.py:180-190`）：`--auto-cast=none -O1`、`--modular-flow-mac-threshold=10`、`--experimental-unsafe-fp8e4m3fn-as-fp8e4m3`（fp8 kernel 必需）、`--enable-verifier=false --enable-nested-dynamic-loop`。 +- **bucket。** 没有单独的 context-encoding / token-generation 两套图；prefill 和 decode 走同一组按行数做 key 的逐层图。 + - prefill 分段：64 rank 时 `max_num_batched_tokens` 为 4096（`__init__.py:451-561`）。TTFT 符合 `ceil(tokens/segment) × F + tokens × ~1.06 ms`，F 随分段大小和上下文长度在 325-783 ms 之间。 + - decode 每步 `(1+5) × num_seqs` 行。`MAX_NUM_SEQS=1`，因为大于 1 时会卡住（`config.py:45-49`）。 + - 专家的行列表补齐到 2 的幂、下限 128（`model.py:2503`、`:2790`），因为 GEMM 在 M 方向按 128 分块。 + - 压缩 KV 的读取长度固定为 `max_model_len`，这样 prefill 分块和 decode 能共用图（`model.py:1178-1185`、`:4425-4440`）。代价是短请求跑在长上下文引擎上时要多读 13 倍。 +- **KV 布局。** + - 窗口 KV 用的是 vLLM 的分页分配，每层一份，`[blocks, 1, block_size, 512]` bf16。之所以声明成 full attention，是因为分页的 slot 同时用来索引其他逐 slot cache（`model.py:4268-4312`）。 + - 模型为 kv-source 层自己分配平铺的逐 slot buffer：`compress_kv_cache [slots,512]` bf16、`index_k_cache [slots,128]` bf16，以及 compressor 的 `pool_kv`/`pool_score [slots,512]` bf16（`model.py:4314-4377`、`:1431-1439`）。 + - 量化是**模拟**的：数值经过 fp8 QDQ（窗口 KV）或 fp4 QDQ（压缩 KV、index K/Q）后**以 bf16 存储**（`model.py:1770`、`:1543-1557`、`:1823`）。 + - 寻址通过 `PagedSpan` 完成（`model.py:1302-1392`）。 +- **DSpark 在 target forward 内部。** `forward`（`model.py:4515-4677`）跑完 40 个 block，并抽取第 37/38/39 层的 attention 输入（`_dspark_tap`，对 4 条流取平均）。然后跑 LM head、device 上的 sampler 和拒绝采样。接着 `_draft` → `forward_spec`（`model.py:5053-5145`）在一个 5 token 的草稿块（种子 token 加噪声 token）上跑 3 个 MTP stage。这些 stage 的 attention 读主窗口和草稿块自己的 KV，不读压缩 KV。最后一个 stage 通过 rank-256 的 Markov logit 偏置逐个采样 5 个 token，并加一个 confidence head。草稿存在 `self.last_draft`，由 `DSparkProposer.propose`（`drafter.py`）交给 vLLM。没有单独的草稿 NEFF，也不用单独 warm-up。这需要插件的 proposer 注册机制（即本 PR 中 `vllm_neuron/vllm/spec_decode` 和 runner 的改动）。 +- **没有消除 host 同步。** + - 每个 MoE 层都把 `indices.cpu()` 拷回 host，在 host 上建行列表（`model.py:2775-2800`）。 + - span/slot 规划对 `slot_mapping` 和 `block_table` 调用 `.tolist()`（`model.py:4766`、`:4831`、`:4943`）。 + - `Block.forward` 里读了 `int(start_pos)`（`model.py:3332`）。 + - 注释指出的瓶颈是每次图执行的固定开销（约 1.12 ms）（`model.py:2351-2370`）。可以用 `DEEPSEEK_V41_EXPERT_FUSE` 把 2-3 个专家合进一张图，默认关闭。 + +## 3. Kernel 清单 + +所有 kernel 的共同点： +- `import nki; import nki.isa as nisa; import nki.language as nl`；入口为 `@nki.jit`。 +- 输出用 `nl.shared_hbm` 分配并**直接返回**（函数式写法）。临时空间用 `nl.private_hbm`。 +- 按 128 行分块并处理尾块；归约在 fp32 中累加。 +- 通过 `wrap_nki(kernel)`（`libtorch_neuronx_lite.nki.nki_hop`）在 `torch.compile(fullgraph=True)` 内调用。2core 变体用 `[2]` 指定 grid。 + +"我们"一栏指它在 DeepInfer 里替换的 shim 或算子（`vendor/` = `deepinfer/neuron_ops/vendor/`，`tl/` = `deepinfer/model_executor/layers/tilelang_ops/neuron/`，`v4/` = `deepinfer/model_executor/models/v4_modules/`）。 + +| Kernel（入口、行号） | 计算内容；输入/输出 | 我们 | 在 torch-neuronx 2.14 上复用 | +|---|---|---|---| +| `blockquant.fp8_linear` :399、`_raw` :417、`_2core` :477、`_raw_2core` :516 | 计算 `x[m,k] @ W[n,k]^T`。激活先做 1x32 的 fp8 舍入（`_raw` 跳过这一步）。权重（uint8 e4m3 字节加 uint8 ue8m0 的 `[n/32,k/32]` 标度）每次调用都用位运算解码到 **bf16 private HBM**，再在 PE 上做 bf16 matmul（`nisa.nc_matmul`，fp32 PSUM）。输出 `[ceil(m/128)*128, n]` bf16，由调用方切片。`_2core` 把 n 切到 2 个物理核（`nl.program_id`，grid 2）。 | `vendor/deep_gemm/gemm.py:205` `fp8_fp4_gemm_nt`（稠密层，`v4/linear.py`）、`einsum.py:77` `fp8_einsum`（`wo_a`，`v4/attn.py:215`），以及它们之前的 `tile_kernels.quant.per_token_cast` | kernel 源码可原样复用。需要：`nki_op` 适配层；权重保持 checkpoint 字节布局（不做 `transform_sf_into_required_layout` 重排）；fp8 编译参数；`_2core` 还需要我们的栈支持 grid=2 发射。**数值：激活 QDQ 除以 224 而不是 448**（Trn2 的 e4m3 在 240 饱和），比 DeepInfer 参考实现少一位尾数。 | +| `blockquant.fp4_linear` :662、`fp4_linear_2core` :681 | fp4 专家用同样的模式：权重 `[n,k/2]` 每字节存两个 e2m1，标度 `[n,k/32]` ue8m0，matmul 之后乘 2^64。2core 把 TPOT 从 142.8 ms 降到 121.4 ms（`model.py:474-480`）。 | `gemm.py:316` `m_grouped_fp8_fp4_gemm_nt_contiguous`（`v4/mlp.py:119`） | 可复用。它每个专家调用一次；我们的 grouped API 需要在本地专家上循环或包一层。 | +| `blockquant.bf16_linear` :1221、`fp32w_linear` :1204 | 普通 bf16 GEMM，用于 router、LM head、compressor、indexer 的 `wk`/`weights_proj` 和 Markov head。`fp32w` 在 kernel 内把 **fp32 权重收窄为 bf16**（mHC 的 `fn`）。 | `gemm.py:107` `bf16_gemm_nt`；`vendor/deep_gemm/hc.py:10` `tf32_hc_prenorm_gemm` | 可复用。精度注意：mHC 的 prenorm GEMM 从 tf32 变成了 bf16。 | +| `blockquant.*_fused`、`_tfused`、`_kmajor` | 在 SBUF 里解码权重的变体。`fp4_linear_fused` 实测 **TTFT +50%、TPOT +24%**（`model.py:536-552`）；其余变体代码路径根本走不到。 | — | 跳过。 | +| `sparse_attn.gathered_attn_idx` :180（实际使用）、`gathered_attn` :51 | 逐 token：按 `idx[T,K]` int32 从 `kv[n,512]` bf16 中间接 DMA gather K 个 slot（`kv.ap(vector_offset=…, indirect_dim=0)`）；QK 用 bf16 操作数、fp32 累加；bias（第 0 列为 softmax scale，空 slot 为 -1e30）；带可学习 sink 的 softmax；P·V 用 bf16。输入 q 为 `[T*H,512]` bf16，H=8 个本地 head；输出 `[T*H,512]` bf16。要求 K%128==0、d%128==0、H≤128。K = 128 窗口 + 512 top-k = 640。 | `vendor/flash_mla/flash_mla_interface.py:282` `flash_mla_with_kvcache`、`:155` `flash_mla_sparse_fwd` | kernel 可复用，但**KV 布局不同**：它读 bf16 行，而我们的 cache 存的是 fp8+ue8m0 记录（每 token 528 B）。两种办法：把 Neuron 的 KV cache 改成 bf16（每 token 1024 B，bs1 / 8K 下没问题）；或者先在 torch 里 gather+反量化，再调 `gathered_attn`。间接 DMA 在 `nki_op` 下能否工作需要上机验证。 | +| `index_score.weighted_scores` :41 | 计算 Σ_h w[h,t]·relu(q_h·k)。输入：q `[T*Hi,128]` bf16（Hi=4 个本地 head）、k `[n,128]` bf16（n%128==0）、w `[Hi,T]` fp32。输出 `[T,n]` fp32（PE matmul，再用第二次 matmul 对 head 求和）。 | `vendor/deep_gemm/attention.py:123/225` `fp8_fp4_(paged_)mqa_logits` | 把 q/k 反量化成 bf16 后即可复用（PR 是先在 torch 里做 fp4 QDQ）。top-k **不在** PR 里：它调用 `vllm_neuron.functional.topk`（nkilib `gpsimd_topk` / vendored 的 `rotational_topk`），这对我们的 #22 和 `vendor/deep_select` 也有参考价值。 | +| `compress_pool.softmax_pool` :42 | 计算 `(kv * softmax(score, dim=ratio)).sum(ratio)`；kv 和 score 为 `[G,ratio,512]` bf16，输出 `[G,512]` fp32。 | `tl/compress.py:36/132` `compress_and_update_state_{prefill,decode}_tl`（其中的 pooling 部分） | 可复用。体量小。 | +| `hc_mixes.kernel` :66 | 输入 raw `[R,24]` fp32（`fn` GEMM 的结果）和整条流 `[R,4*5120]`：先对整条流求 rms，然后 pre = σ+eps、post = 2σ、comb = softmax 加 **20 次 Sinkhorn**。输出 pre/post `[R,4]` 和 comb `[R,16]`，全为 fp32。 | `tile_kernels.mhc` 的 `norm_fn` / `pre_split_mixes` / `sinkhorn`（`v4/mhc_ops.py:10-16`） | 可复用。需要对照我们的配置核对迭代次数和 eps（`sinkhorn_normalize_ref` 默认 10 次）。 | +| `hc_expand.kernel` :28 | `[B,S,D]` → 4 份拷贝 `[B,S,4,D]`，外加 one-hot 的 pre_mix `[B,S,4]` fp32。只做 DMA。 | `tile_kernels.mhc.expand_kernel.expand_to_mhc_fwd` | 可复用。 | +| `hc_collapse.kernel` :24 | Σ_c pre[c]·x[c]：`[B,N,4,D]`、`[B,N,4]` fp32 → `[B,N,D]`。 | `mhc_pre_apply_mix_fwd` | 可复用。 | +| `hc_attn_out.kernel` :24 | post[j]·x + Σ_i comb[i,j]·res[i]：x `[B,S,D]`、res `[B,S,4,D]`、post `[B,S,4]`、comb `[B,S,4,4]` → `[B,S,4,D]`。 | `mhc_post_fwd` | 可复用。 | +| `rmsnorm.rms_norm` :30 | x `[R,D]`、weight `(1,D)`、eps 为 `(1,1)` 张量 → `[R,D]`，dtype 与输入相同。fp32 归约，rsqrt 后再做一步 Newton 迭代。 | `tl/norm.py:75` `norm_tl`、`v4/norm.py` | 可复用。需要对照 `rstd_1ulp` 核对一致性。 | +| `swiglu.swiglu` :40 | 带 clamp 的 SwiGLU（up 限制在 ±10，gate 限制 ≤10）× 路由权重：gate/up `[R,I]`，weights 为 `(R,1)` 或 `(1,1)`，limit `(1,1)` → `[R,I]`。 | `tile_kernels.quant.swiglu_forward_and_per_token_cast`（`v4/mlp.py:19`） | 可复用。它不做 cast；fp8 舍入交给下一个 GEMM 的 QDQ。 | +| `engram_gate.gated_add` :41 | 对 4 条流中的每一条：rms(h)·rms(key) 归一化点积 → sign·sqrt → sigmoid 门控；输出 h + gate·value。h/key `[R,4*5120]`，value `[R,5120]`，weight `(1,4*5120)` = q_weight·k_weight。 | `tile_kernels.engram.engram_gate_fwd`（`v4/engram.py:155`） | 对照我们 fused-weight 的语义核对后可复用。 | +| `pp_permute.kernel` :395 | 用 `nki.collectives.collective_permute` 加 `nisa.core_barrier` 做 stage 之间的交接。**没有被调用**：CC 运行时拒绝了它的 replica group。 | — | 跳过。 | + +**没有 MoE kernel 时，MoE GEMM 和反量化是怎么做的**（`model.py:2287-2380`、`:2851-2928`）： +1. router：`bf16_linear` + sqrt-softplus + nkilib top-k。 +2. host 把 indices 拷回来，为每个本地专家建一个行列表，补齐到不小于 128 的 bucket。 +3. 每个专家一张编译好的图：`index_select` 取行 → `fp4_linear_2core`（w1、w3）→ `swiglu` → `fp4_linear_2core`（w2）→ `index_add` 加回 fp32 累加器。 +4. 全局 all-reduce。 +5. 共享专家：同样的链路，换成 `fp8_linear_2core`，对所有行计算。 + +反量化在每次 GEMM 内部完成（权重字节 → private HBM 中的 bf16），不保留常驻的 bf16 副本。 + +## 4. 时间花在哪里，TPOT ≈ 120 ms 靠的是什么 + +PR 没有给出 profile 拆解。注释里能拼出这些数字： +- 图执行：每次固定约 1.12 ms，每次 forward 约 120 次。128 行专家 GEMM 的 roofline 约 154 µs，所以路由专家这部分时间大头是固定开销，而不是计算（`model.py:2353-2363`）。 +- 全局 all-reduce：约 43 × 1.9 ms ≈ 每步 80 ms。 +- decode 一步 ≈ 320-370 ms，处理 6 行。TPOT = 一步时间 / (1 + 接受数)，在 sonnet 上接受长度为 2.75（各位置平均接受率 36%）。 +- prefill：每次 forward 的固定开销 F ≈ 0.3-0.8 s，外加每 token 约 1.06 ms。 + +起决定作用的设计： +1. **全部常驻 HBM**（每 rank 13.9 GiB）：forward 期间没有任何 host 流量。 +2. **attention TP=8 并复制**，attention 的归约限制在 2 个 device 内。 +3. **专家按 id 划分、token 复制**：不需要 dispatch，只要一次 all-reduce。 +4. **GEMM 用满每个 LNC=2 核里的两个物理核**：先 fp4、再 fp8 的 2core 化把 TPOT 从 142.8 ms 降到 116.9 ms。 +5. **每类算子都用 NKI**，减少图和算子的数量。 +6. **prefill 分段大小 4096。** +7. **DSpark**（约 2.75 倍）。 + +他们实测过、但没有收益的尝试：在 SBUF 里融合反量化、共享压缩 KV 的读取、`NEURON_RT_IO_RING_CACHE_SIZE`、collective 拆成 8×8、`_EXPERT_ROW_FLOOR` 设到 128 以下。 + +## 5. 我们的集成计划（按预期收益排序） + +1. **#9 静态图 + #12 放置（结构性改动，从约 150 s 降到 1 秒以内）。** + - 每个 bucket 编译一张 decode 图，NKI 算子都放在图里。 + - 去掉所有 host 同步。为了避免 PR #71 那种每层一次的 `.cpu()`，**让 6 个本地专家都跑一遍 bucket 内的所有行**，没路由到该专家的行把路由权重设为 0。他们的 GEMM 本来就把 M 补到 128，所以 decode 时每个专家的指令量和他们 128 的 bucket 一样，而图保持静态。 + - bs1 下考虑**把 attention DP 换成 attention TP=8 组**：现在是一个 rank 算全部 64 个 head，其余 63 个 rank 在等。 + - token 复制以后，我们的 all-gather dispatch 就是空操作，combine 也变成一次全局 all-reduce，和他们的做法一样。需要实测 all-reduce 和 reduce-scatter 哪个更好。 + - 照搬他们 `_COMPILE_OPTIONS` 的编译参数。每次全局 collective 按约 1.9 ms 估算。 +2. **#10 GEMM：复用 `blockquant`。** + - `fp8_linear(_raw)_2core` 接到 `fp8_fp4_gemm_nt`/`fp8_einsum` 后面，`fp4_linear_2core` 接到 `m_grouped_fp8_fp4_gemm_nt_contiguous` 后面（在本地专家上循环），`bf16_linear`/`fp32w_linear` 接到 bf16 和 hc 的 GEMM 后面。 + - 权重保持 checkpoint 的 uint8 字节布局（这也符合我们"fp8 按字节搬运"的规则）。 + - 需要决定：接受 224 的 QDQ，还是把 kernel 改回 448 并重新验证饱和相关的编译参数。 + - 这是 kernel 层面收益最大的一项。 +3. **#15 DSpark。** 他们的 TPOT 里含有 2.75 倍的放大。没有 DSpark 的话，我们一步要做到 117 ms 以内，而不是约 320 ms。要照着做的设计是在 forward 内部起草（从第 37-39 层抽取、5 token 的草稿块、Markov head）。 +4. **#11 attention / indexer。** + - 选定 KV 布局：bf16 cache 配 `gathered_attn_idx`，或者保留 fp8 记录、先 gather+反量化再调 `gathered_attn`。 + - `index_score` 接到 `fp8_fp4_(paged_)mqa_logits` 后面，`compress_pool` 接到 `tl/compress.py` 后面。 + - `deep_select.topk` 和 #22 借用 `vllm_neuron.functional.topk`（nkilib）。 +5. **#14 Engram。** 照搬他们的放置：按行切分的 bf16（或 fp8，2.9 GiB）表放在 HBM，hash 和 gather 都在 device 上做，再做一次全局 all-reduce。这样 host gather 和 H2D 就完全去掉了。前提是 HBM 有余量（#21）：他们的稠密部分每 rank 只有 3.9 GiB，是因为 attention 用了 TP=8、词表切了 64 份。再加上 `engram_gate`。 +6. **#13 小算子。** `hc_mixes`、`hc_expand`、`hc_collapse`、`hc_attn_out`、`rmsnorm`、`swiglu` 单个都不大，但合起来能消掉每层大部分算子数。放在 #9 之后做，这样 profile 才是真实的。 + +**在我们的栈上复用任何 kernel 之前要先确认：** +- `torch_neuronx.nki_op`（2.14.3）是否接受*返回* `shared_hbm` 输出、使用 `private_hbm` 临时空间的 `@nki.jit` kernel，能否同时在 eager 和 `torch.compile(backend="neuron")` 下运行，能否在 LNC=2 下以 grid=2 发射。 +- 间接 DMA（`.ap(vector_offset=…)`）经过 `nki_op` 能否工作。 +- 形状断言（k%128、K%128、heads≤128）。 +- 与我们 CPU shim 的数值一致性，误差要在量化噪声以内。 + +**许可证。** +- vllm-neuron 是 Apache-2.0（`LICENSE`；`NOTICE`："Copyright Amazon.com, Inc. or its affiliates"）。DeepInfer 也是 Apache-2.0，两者兼容。 +- kernel 文件本身没有逐文件的许可证头。只有 `kernels/__init__.py`、`model.py`、`weight_loader.py` 等文件带 `SPDX-License-Identifier: Apache-2.0`。 +- 拷贝时：给每个拷过来的文件加 SPDX 头和来源说明（仓库、PR #71、SHA `30d54d1`、作者 Zifan He / OswaldHe）；改过的文件要注明已修改（Apache-2.0 §4b）；把相关的 NOTICE 文本带进我们的 NOTICE 或第三方声明。 +- PR 还没合并，所以要钉住 SHA，合并后再核对一次。 + +
diff --git a/tests/neuron/compile/test_fp8_e4m3.py b/tests/neuron/compile/test_fp8_e4m3.py new file mode 100644 index 0000000..08faaaf --- /dev/null +++ b/tests/neuron/compile/test_fp8_e4m3.py @@ -0,0 +1,41 @@ +"""trn2 e4m3 range handling (deepinfer/neuron_ops/fp8.py, issue #24), on the host.""" + +import torch + +from deepinfer.neuron_ops import fp8 + +FP8 = torch.float8_e4m3fn + + +def test_halve_matches_torch_cast(): + b = torch.arange(256, dtype=torch.int32).to(torch.uint8) + ref = (b.view(FP8).float() / 2).to(FP8).view(torch.uint8) + got = fp8.halve_e4m3fn_bytes(b) + nan = (b & 0x7F) == 0x7F + assert torch.equal(got[~nan], ref[~nan]) + assert bool(((got[nan] & 0x7F) == 0x7F).all()) + + +def test_halved_values_fit_hardware_e4m3(): + b = torch.arange(256, dtype=torch.int32).to(torch.uint8) + got = fp8.halve_e4m3fn_bytes(b) + finite = (b & 0x7F) != 0x7F + assert bool((((got[finite] >> 3) & 15) < 15).all()) # exponent field 15 is inf/NaN on trn2 + + +def test_rescale_keeps_dequantized_values(): + torch.manual_seed(0) + w = torch.nn.Parameter((torch.randn(64, 128) * 100).clamp(-448, 448).to(FP8), requires_grad=False) + w.sf = torch.randint(110, 130, (64, 4), dtype=torch.uint8) # ue8m0 per 32 columns + before = w.float() * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1) + lossy = ((w.data.view(torch.uint8) >> 3) & 15) <= 1 # halving rounds only below the smallest normal + model = torch.nn.Module() + model.w = w + if fp8.FP8_MAX > 240: + return + assert fp8.rescale_fp8_weights(model) == 1 + after = w.float() * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1) + assert torch.equal(after[~lossy], before[~lossy]) + ulp = 2.0**-9 * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1) # one smallest-subnormal step, after rescale + assert bool(((after - before).abs()[lossy] <= ulp[lossy]).all()) + assert float(w.float().abs().max()) <= 224.0 diff --git a/tests/neuron/compile/test_fp8_legalize.py b/tests/neuron/compile/test_fp8_legalize.py new file mode 100644 index 0000000..03efdb6 --- /dev/null +++ b/tests/neuron/compile/test_fp8_legalize.py @@ -0,0 +1,141 @@ +"""fp8 legalization for Neuron graphs (issue #24), checked on the host: the pass output runs eagerly on CPU.""" + +import pytest +import torch +from torch._dynamo.backends.common import aot_autograd + +from deepinfer.neuron_ops import neuron_compile as nc + +FP8 = torch.float8_e4m3fn + + +def test_encode_matches_torch_cast(): + torch.manual_seed(0) + x = torch.cat([ + torch.randn(100000) * s for s in (1e-4, 1e-2, 1.0, 30.0, 300.0) + ] + [torch.tensor([0.0, 448.0, 449.0, 464.0, 465.0, 480.0, 1e6, float("inf"), -float("inf"), float("nan"), + 2.0**-10, 1.5 * 2.0**-9, 2.5 * 2.0**-9, 2.0**-6, 2.0**-6 * 0.999, 2.0**-7 * 1.0625])]) + ref = x.to(FP8).view(torch.uint8) + got = nc._encode(x) + neg_zero = (ref == 128) & (got == 0) + assert torch.equal(torch.where(neg_zero, ref, got), ref) + + +def test_decode_matches_torch_cast(): + b = torch.arange(256, dtype=torch.int32).to(torch.uint8) + ref = b.view(FP8).float() + got = nc._decode(b) + assert torch.equal(torch.nan_to_num(ref, nan=1234.0), torch.nan_to_num(got, nan=1234.0)) + + +def _host_backend(): + def fw(gm, example_inputs): + bad = [n for n in gm.graph.nodes if nc._is_fp8(n.meta.get("val"))] + assert not bad, bad + return gm + + aot = aot_autograd(fw_compiler=fw) + + def backend(gm, example_inputs): + inputs, positions = nc._fp8_inputs_as_bytes(gm, example_inputs) + nc.legalize_fp8(gm) + compiled = aot(gm, inputs) + pos = set(positions) + return lambda *args: compiled(*[a.view(torch.uint8) if i in pos else a for i, a in enumerate(args)]) + + return backend + + +def _quant(x): + amax = x.abs().amax(-1, keepdim=True).clamp(min=1e-4) + scale = amax / 448.0 + return (x / scale).to(FP8), scale + + +@pytest.mark.parametrize("case", ["dequant_weight", "quant_act", "kv_bytes", "cat_index"]) +def test_graphs(case): + torch.manual_seed(0) + w, ws = _quant(torch.randn(64, 128)) + kv = torch.randint(0, 255, (4, 16, 32), dtype=torch.uint8) + + def f(x, w, ws, kv, idx): + if case == "dequant_weight": + return x @ (w.float() * ws).t() + if case == "quant_act": + q, s = _quant(x) + return q.float() * s, q.view(torch.uint8) + if case == "kv_bytes": + rows = kv.view(-1, 32)[idx] + vals = rows.view(FP8).float() + kv.view(-1, 32)[idx[:1]] = (vals[:1] * 0.5).to(FP8).view(torch.uint8) + return vals + if case == "cat_index": + q, _ = _quant(x) + both = torch.cat([q, w[: q.shape[0]]], dim=-1) + return both[idx].float(), torch.where(idx[:, None] > 2, both[idx], torch.zeros_like(both[idx])) + + x = torch.randn(8, 128) + idx = torch.tensor([1, 3, 5, 7]) + ref = f(x, w, ws, kv.clone(), idx) + kv2 = kv.clone() + got = torch.compile(f, backend=_host_backend(), fullgraph=True, dynamic=False)(x, w, ws, kv2, idx) + ref, got = (ref,) if isinstance(ref, torch.Tensor) else ref, (got,) if isinstance(got, torch.Tensor) else got + for r, g in zip(ref, got): + r = r.view(torch.uint8) if r.dtype == FP8 else r + g = g.view(torch.uint8) if g.dtype == FP8 else g + assert torch.equal(r, g), case + if case == "kv_bytes": + kv_ref = kv.clone() + f(x, w, ws, kv_ref, idx) + assert torch.equal(kv_ref, kv2) + + +@pytest.mark.parametrize("case", ["encode", "dequant_weight", "quant_act", "kv_bytes", "cat_index"]) +def test_graphs_on_device(case): + """Same graphs through ``backend="deepinfer_neuron"`` on a NeuronCore (skipped on the host).""" + torch_neuronx = pytest.importorskip("torch_neuronx") # noqa: F841 + if not hasattr(torch, "neuron") or not torch.neuron.is_available(): + pytest.skip("no NeuronCore") + import os + if os.environ.get("DEEPINFER_NEURON_DEVICE", "neuron") == "cpu": + pytest.skip("host mode") + dev = torch.device("neuron:0") + torch.manual_seed(0) + w, ws = _quant(torch.randn(64, 128)) + kv = torch.randint(0, 255, (4, 16, 32), dtype=torch.uint8) + x = torch.randn(8, 128) + idx = torch.tensor([1, 3, 5, 7], dtype=torch.int32) + + if case == "encode": + x = torch.cat([torch.randn(4096) * s for s in (1e-3, 1.0, 100.0)] + [torch.tensor([448.0, 500.0, 2.0**-9 * 1.5])]) + + def f(x, w, ws, kv, idx): + if case == "encode": + return x.to(FP8).view(torch.uint8), x.to(FP8).float() + if case == "dequant_weight": + return x @ (w.float() * ws).t() + if case == "quant_act": + q, s = _quant(x) + return q.float() * s, q.view(torch.uint8) + if case == "kv_bytes": + rows = torch.index_select(kv.view(-1, 32), 0, idx) + vals = rows.view(FP8).float() + return vals, (vals * 0.5).to(FP8).view(torch.uint8) + q, _ = _quant(x) + both = torch.cat([q, w[: q.shape[0]]], dim=-1) + return (torch.index_select(both, 0, idx).float(),) + + ref = f(x, w, ws, kv, idx) + got = torch.compile(f, backend="deepinfer_neuron", fullgraph=True, dynamic=False)( + x.to(dev), w.to(dev), ws.to(dev), kv.to(dev), idx.to(dev)) + ref = (ref,) if isinstance(ref, torch.Tensor) else ref + got = (got,) if isinstance(got, torch.Tensor) else got + for r, g in zip(ref, got): + g = g.cpu() + if r.dtype.is_floating_point and case in ("dequant_weight", "quant_act"): + # the device's matmul / division are not bit-identical to the host's + torch.testing.assert_close(g, r, rtol=7e-2, atol=1e-2) + elif r.dtype == torch.uint8 and case == "quant_act": + assert (r.int() - g.int()).abs().max() <= 1, case + else: + assert torch.equal(r, g), case diff --git a/tests/neuron/run_all.sh b/tests/neuron/run_all.sh index 77b917e..318acbc 100755 --- a/tests/neuron/run_all.sh +++ b/tests/neuron/run_all.sh @@ -24,6 +24,9 @@ cd "$REPO" || exit 1 VENDOR="$REPO/deepinfer/neuron_ops/vendor" T="$REPO/tests/neuron" export DEEPINFER_ACCELERATOR=neuron +# Kernel unit tests compare bit-for-bit against references quantizing to the e4m3fn range (448); the end-to-end +# stages run the default trn2 e4m3 range (224, deepinfer/neuron_ops/fp8.py) and compare against the golden. +UNIT_FP8_MAX="${UNIT_FP8_MAX:-448}" declare -a SUMMARY=() first_core() { echo "${CORES%%[-,]*}"; } @@ -43,7 +46,9 @@ run() { # run selected "$stage" || return 0 echo "=== $stage" local t0=$SECONDS - if "$@" > "$LOGDIR/$stage.log" 2>&1; then r=PASS; else r=FAIL; fi + local fp8_max="$UNIT_FP8_MAX" + case "$stage" in e2e_*|device_diff) fp8_max="${DEEPINFER_NEURON_FP8_MAX:-224}" ;; esac + if DEEPINFER_NEURON_FP8_MAX="$fp8_max" "$@" > "$LOGDIR/$stage.log" 2>&1; then r=PASS; else r=FAIL; fi SUMMARY+=("$(printf '%-6s %-28s %5ss %s' "$r" "$stage" $((SECONDS - t0)) "$LOGDIR/$stage.log")") echo " $r ($((SECONDS - t0))s)" [ "$r" = FAIL ] && tail -5 "$LOGDIR/$stage.log" | sed 's/^/ | /'