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 <noreply@anthropic.com>

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()

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:

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))

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

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)
 

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("<I", struct.pack("<f", 8.0 / (8 + m)))[0] for m in range(8))
 

deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py

@@ -20,6 +20,7 @@ from typing import Any
 import torch
 
 from ..common import FLOAT8E4M3_MAX, FP8_AMAX_MARGIN
+from deepinfer.neuron_ops.fp8 import FP8_MAX as FLOAT8E4M3_MAX  # trn2 hardware e4m3 range (#24)
 from .kv_quant import fp8_e4m3_bytes
 from .rstd_1ulp import _fma, make_rstd_impl
 

deepinfer/model_executor/layers/tilelang_ops/neuron/store_kv_cache.py

@@ -88,12 +88,19 @@ def store_kv_cache_tl(
             _masked_row_store(kv_cache, slot_mapping, torch.cat([values, scales], dim=-1), block_size)
         else:
             num_blocks = kv_cache.shape[0]
-            page_stride = kv_cache.stride(0)
-            base = kv_cache.storage_offset()
-            cache_values = kv_cache.as_strided((num_blocks, block_size, data_bytes), (page_stride, data_bytes, 1), base)
-            cache_sf = kv_cache.as_strided(
-                (num_blocks, block_size, sf_bytes), (page_stride, sf_bytes, 1), base + block_size * data_bytes
-            )
+            if kv_cache.stride(1) == payload_bytes:
+                # Page-wide sections by slicing a page view (no as_strided: torch.compile cannot functionalize
+                # mutations through it).
+                pages = kv_cache.flatten(1)
+                cache_values = pages[:, : block_size * data_bytes].view(num_blocks, block_size, data_bytes)
+                cache_sf = pages[:, block_size * data_bytes : block_size * payload_bytes].view(num_blocks, block_size, sf_bytes)
+            else:
+                page_stride = kv_cache.stride(0)
+                base = kv_cache.storage_offset()
+                cache_values = kv_cache.as_strided((num_blocks, block_size, data_bytes), (page_stride, data_bytes, 1), base)
+                cache_sf = kv_cache.as_strided(
+                    (num_blocks, block_size, sf_bytes), (page_stride, sf_bytes, 1), base + block_size * data_bytes
+                )
             _masked_row_store(cache_values, slot_mapping, values, block_size)
             _masked_row_store(cache_sf, slot_mapping, scales, block_size)
 

deepinfer/model_executor/layers/tilelang_ops/norm.py

@@ -52,10 +52,14 @@ def _alloc_fp8_sf(
     if use_packed_ue8m0:
         packs = (d_scale + pack_factor - 1) // pack_factor
         storage = torch.empty(batch, packs, m_pad, dtype=pack_dtype, device=device)
-        kernel_arg = storage.view(torch.uint8).as_strided(
-            (batch, seqlen, packs, pack_factor),
-            (packs * m_pad * pack_factor, pack_factor, m_pad * pack_factor, 1),
-        )
+        if PLATFORM.is_neuron:
+            # Same layout without as_strided, which torch.compile cannot functionalize when the kernel writes through it.
+            kernel_arg = storage.view(torch.uint8).view(batch, packs, m_pad, pack_factor).permute(0, 2, 1, 3)[:, :seqlen]
+        else:
+            kernel_arg = storage.view(torch.uint8).as_strided(
+                (batch, seqlen, packs, pack_factor),
+                (packs * m_pad * pack_factor, pack_factor, m_pad * pack_factor, 1),
+            )
         return storage.mT[:, :seqlen], kernel_arg
     storage = torch.empty(batch, d_scale, m_pad, dtype=torch.float32, device=device)
     sf = storage.mT[:, :seqlen]

deepinfer/model_executor/models/deepseek.py

@@ -8,7 +8,7 @@ import tile_kernels
 import torch
 from torch import nn
 
-from deepinfer import AttrDict, envs
+from deepinfer import PLATFORM, AttrDict, envs
 from deepinfer.engine.engine_type import EngineType
 from deepinfer.logger import init_logger
 from deepinfer.model_executor import eplb
@@ -509,6 +509,12 @@ class DeepseekForCausalLM(nn.Module):
         if not skip_load:
             load_hf_checkpoint(gpt, hf_dir=hf_dir, hf_coverage=hf_coverage)
         started = log_load_phase("checkpoint", started)
+        if PLATFORM.is_neuron:
+            # trn2 reads fp8 as IEEE-style e4m3 (max 240): bring the checkpoint's e4m3fn weights into range.
+            from deepinfer.neuron_ops.fp8 import rescale_fp8_weights
+
+            logger.info(f"Neuron: rescaled {rescale_fp8_weights(gpt)} fp8 weights into the e4m3 range")
+            started = log_load_phase("fp8_rescale", started)
 
         allocated_before = torch.cuda.memory_allocated()
         model = cls.from_model_cfg(gpt, model_cfg)

deepinfer/model_executor/models/v4_modules/mlp.py

@@ -33,6 +33,7 @@ from deepinfer.model_executor.models.v4_modules.moe import FusedExperts, Replica
 from deepinfer.model_executor.parallel_utils.parallel_state import PG
 from deepinfer.model_executor.parallel_utils.userbuffer import limit_num_sms, userbuffer
 from deepinfer.model_executor.utils import reconstruct_param
+from deepinfer.neuron_ops.static_moe import static_moe_forward
 from deepinfer.timer import cuda_graph_timer
 
 logger = init_logger(__name__)
@@ -574,6 +575,9 @@ class DeepseekMoE(nn.Module):
             return self.mega_moe_forward(hidden_states, hidden_states_bf16, input_metadata)
 
         topk_idx, topk_weights = self.gate_forward(hidden_states_bf16, input_metadata)
+        if PLATFORM.is_neuron and not input_metadata.is_prompt:
+            # Fixed shapes for the decode graphs: all-gather + dense local experts + reduce-scatter.
+            return static_moe_forward(self, hidden_states, topk_idx, topk_weights)
 
         try:
             # Dispatch

deepinfer/model_executor/utils.py

@@ -81,9 +81,12 @@ class Movable[T]:
     """
 
     __slots__ = ("_value",)
+    # Declared here, not on the assignment in __init__: an annotated attribute assignment evaluates `T | None` on
+    # every call, and torch.compile cannot trace a TypeVar.
+    _value: T | None
 
     def __init__(self, value: T) -> None:
-        self._value: T | None = value
+        self._value = value
 
     def move(self) -> T:
         if self._value is None:

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', ''))}"

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

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))

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

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)."""

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:

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,
 )

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

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,

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,

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

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

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)

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]
 

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.<dtype>` 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,
 }

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

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")
 

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,
             )

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

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.
+
+<details><summary>中文</summary>
+
+# 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 缩小。
+
+</details>

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.
+
+<details>
+<summary>中文</summary>
+
+# 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，合并后再核对一次。
+
+</details>

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

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

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 <stage> <command...>
   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/^/    | /'