Neuron: full DeepSeek-V4.1-Flash load + decode region graphs at EP=64 (issues #21, #25)

- Layer-major KV / state layout so per-layer views are contiguous graph inputs
  (the executor copied non-contiguous inputs every call).
- One compiled entry per (class, method), rank as a registered buffer, no
  layer_idx guard: every layer reuses the same graphs and loaded NEFFs.
- fp8 weight rescale on the host (numpy LUT, ~2 min instead of 35+ min of
  eager per-shape compiles); engram sf packing on the host.
- int16 ue8m0 packing in fp32 (int32 interleaved slices narrowed to int16
  crash neuronx-cc, NCC_IBCG902).
- Token embedding and the text model's vision tower stay in host memory on
  Neuron (2.5 GiB of HBM a rank).
- offline example: --max_num_seqs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/environ.py

@@ -114,6 +114,8 @@ class envs:
     DEEPINFER_NEURON_GRAPH_MODE = EnvVar.str(default="regions")
     # Neuron only: cut the region graphs further at the stage boundaries marked by neuron_ops.regions.split().
     DEEPINFER_NEURON_REGION_SPLIT = EnvVar.bool(default=True)
+    # Keep the token embedding (and the text model's unused vision tower) in host memory on Neuron.
+    DEEPINFER_NEURON_HOST_EMBED = EnvVar.bool(default=True)
     # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break).
     DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool()
 

deepinfer/model_executor/layers/tilelang_ops/fused_transform.py

@@ -275,7 +275,8 @@ def fused_norm_rope_hadamard_cast_store_kvcache(
         block_size, cache = split_kv_cache(cache_out.kv_cache, cache_layout)
         cache_quant, cache_sf = cache.quantized_data, cache.sf_data
         assert cache_quant.stride(1) % cache_layout.stride_alignment == 0
-        assert cache_quant.stride(0) % 32 == 0
+        # Page alignment is a GPU/Ascend kernel requirement; Neuron's layer-major caches pack pages tightly.
+        assert PLATFORM.is_neuron or cache_quant.stride(0) % 32 == 0
         out_bf16 = _restore_single_head_axis(cache_out.out, x.shape, original_ndim)
         bf16_stage = cache_out.out_stage
         if out_bf16 is not None:

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

@@ -373,11 +373,13 @@ def fused_norm_rope_hadamard_cast_store_kvcache_tl(num_heads: int, head_dim: int
                 # Neuron graphs: join the scale byte pairs into the int16 storage arithmetically (low byte first), as
                 # neuron/norm.py does. Writing them through a uint8 view that the same graph reads back as int16
                 # crashes neuronx-cc (NCC_IINL902, issue #25).
-                e = scale.to(torch.int32)
+                # Pair in fp32 and narrow once (int32 interleaved slices -> int16 crash codegen, NCC_IBCG902).
+                f = scale.to(torch.float32)
                 if num_scales % 2:
-                    e = torch.nn.functional.pad(e, (0, 1))
-                v = e[:, 0::2] | (e[:, 1::2] << 8)
-                osf[:, : v.shape[-1]].copy_(torch.where(v >= 32768, v - 65536, v).to(torch.int16))
+                    f = torch.nn.functional.pad(f, (0, 1))
+                lo, hi = f[:, 0::2], f[:, 1::2]
+                v = lo + hi * 256.0 - torch.where(hi >= 128.0, 65536.0, 0.0)
+                osf[:, : v.shape[-1]].copy_(v.to(torch.int16))
             else:
                 osf.view(torch.uint8)[:, :num_scales].copy_(scale)
 

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

@@ -161,11 +161,13 @@ def norm_tl(
 
         if use_packed_ue8m0 and sf.dtype == torch.int16:
             # int16 storage (compiled regions): join exponent byte pairs arithmetically, low byte first.
-            e = sf_values.to(torch.int32)
-            if e.shape[-1] % 2:
-                e = torch.nn.functional.pad(e, (0, 1))
-            v = e[..., 0::2] | (e[..., 1::2] << 8)
-            sf.copy_(torch.where(v >= 32768, v - 65536, v).to(torch.int16).view(sf.shape))
+            # Pair the bytes in fp32 (exact below 2^17) and narrow once: int32 interleaved slices narrowed to int16
+            # crash neuronx-cc codegen (NCC_IBCG902, see nki_norm.pack_ue8m0_int16).
+            f = sf_values.to(torch.float32)
+            if f.shape[-1] % 2:
+                f = torch.nn.functional.pad(f, (0, 1))
+            lo, hi = f[..., 0::2], f[..., 1::2]
+            sf.copy_((lo + hi * 256.0 - torch.where(hi >= 128.0, 65536.0, 0.0)).to(torch.int16).view(sf.shape))
         elif use_packed_ue8m0 and sf.ndim == 4:
             # col-major packed view: (B, M, num_groups // pack, pack) bytes
             sf.copy_(sf_values.view(*lead, sf.shape[-2], sf.shape[-1]))

deepinfer/model_executor/layers/vocab_parallel_embedding.py

@@ -20,6 +20,9 @@ class VocabParallelEmbedding(torch.nn.Module):
         self.weight = reconstruct_param(weight)
 
     def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        if self.weight.device != input_.device:
+            # Host-resident table (Neuron, see DeepseekForCausalLM.load_model): move the looked-up rows only.
+            return gather(self.weight, input_.cpu()).to(input_.device)
         output = gather(self.weight, input_)
         return output
 

deepinfer/model_executor/model_loader.py

@@ -76,6 +76,8 @@ def get_model(
     # The model device is "cuda", or the mapped Neuron target (neuron:0, or the host in Neuron CPU mode).
     device_type = PLATFORM.device.type if PLATFORM.is_neuron else "cuda"
     for n, p in model.named_parameters():
+        if PLATFORM.is_neuron and p.device.type == "cpu" and n.startswith("model.embed_tokens."):
+            continue  # host-resident on Neuron (DeepseekForCausalLM._neuron_host_resident)
         assert p.device.type == device_type, f"{n} not on {device_type}"
         if hasattr(p, "sf"):
             assert p.sf.device.type == device_type, f"{n}.sf not on {device_type}"

deepinfer/model_executor/models/deepseek.py

@@ -148,7 +148,7 @@ class DeepseekModel(nn.Module):
         )
 
     def allocate_output(self, num_tokens: int) -> CausalLMForwardOutput:
-        hidden_states = self.embed_tokens.weight.new_zeros((num_tokens, self.hidden_size))
+        hidden_states = torch.zeros((num_tokens, self.hidden_size), dtype=self.embed_tokens.weight.dtype, device="cuda")
         return CausalLMForwardOutput(hidden_states=hidden_states)
 
     def forward(
@@ -475,6 +475,8 @@ class DeepseekForCausalLM(nn.Module):
         for nm, m in gpt.named_modules():
             for pn, p in m.named_parameters(recurse=False):
                 device = engram_device if getattr(p, "is_engram", False) else "cuda"
+                if PLATFORM.is_neuron and cls._neuron_host_resident(nm):
+                    device = "cpu"
                 packed_fp4 = ("experts.linear1" in nm or "experts.linear2" in nm) and model_cfg.gpt.mxfp4
                 alias_source = hf_coverage.aliases.get(f"{nm}.{pn}")
                 if alias_source is not None:
@@ -526,6 +528,15 @@ class DeepseekForCausalLM(nn.Module):
 
         return model
 
+    @classmethod
+    def _neuron_host_resident(cls, module_name: str) -> bool:
+        """Parameters kept in host memory on Neuron, where a NeuronCore has 24 GiB of HBM: the token embedding (1.2 GiB;
+        a lookup moves only its rows, see VocabParallelEmbedding) and the vision tower the text model drops after
+        loading (1.3 GiB that would otherwise stay in the device caching allocator)."""
+        if not envs.DEEPINFER_NEURON_HOST_EMBED:
+            return False
+        return module_name.endswith("emb.tok_emb") or module_name.startswith("vision_pipeline")
+
     @classmethod
     def from_model_cfg(cls, gpt: nn.Module, model_cfg: AttrDict) -> nn.Module:
         return cls(gpt, model_cfg.gpt)

deepinfer/model_executor/models/deepseek_vl.py

@@ -38,6 +38,11 @@ class DeepseekVLForCausalLM(DeepseekForCausalLM):
             vision_source=model.vision_pipeline.aligner,
         )
 
+    @classmethod
+    def _neuron_host_resident(cls, module_name: str) -> bool:
+        # The vision tower runs here: only the token embedding stays on the host.
+        return super()._neuron_host_resident(module_name) and not module_name.startswith("vision_pipeline")
+
     @staticmethod
     def _split_cached_prefix(
         segment_input: "SegmentInput",

deepinfer/model_executor/models/v4_modules/engram.py

@@ -27,6 +27,9 @@ from deepinfer.model_executor.parallel_utils.userbuffer import userbuffer
 def _pack_ue8m0_to_int(x: torch.Tensor) -> torch.Tensor:
     """Pack the ue8m0 scales represented in fp32 into int32 (CUDA) / int16 (ASC), 4/2 per group."""
     assert x.dtype == torch.float32 and x.size(-1) % 4 == 0
+    if PLATFORM.is_neuron and x.device.type != "cpu":
+        # One-time load step on the host: eager device ops on this large tensor compile a big NEFF per shape.
+        return _pack_ue8m0_to_int(x.cpu()).to(x.device)
     x_int32 = x.view(torch.int32)
     assert (x_int32 >= 0).all() and (x_int32 & 0x7FFFFF == 0).all()
     return torch.floor_divide(x_int32, 1 << 23).to(torch.uint8).view(PLATFORM.fp8_sf_type)

deepinfer/model_executor/models/v4_modules/mlp.py

@@ -190,6 +190,12 @@ class DeepseekMoE(nn.Module):
             and PG.is_within_nvl_domain(PG.expert_parallel_group())
         )
 
+        if PLATFORM.is_neuron:
+            # A registered buffer (see neuron_rank_tensor): Dynamo lifts buffers as graph inputs, while a plain tensor
+            # attribute was constant-folded into the graph (one NEFF per rank again).
+            self.register_buffer(
+                "_neuron_rank_t", torch.tensor([PG.expert_parallel_group().rank()], dtype=torch.int32), persistent=False
+            )
         self.ep_buffer: EPBuffer | None = None
         self.sym_buffer: deep_gemm.SymmBuffer | None = None
         self.ep_group = PG.expert_parallel_group()
@@ -326,6 +332,11 @@ class DeepseekMoE(nn.Module):
                 use_fp8_dispatch=self.use_fp8_act,
             )
 
+    def neuron_rank_tensor(self) -> torch.Tensor:
+        """This rank's EP index as a device tensor: compiled graphs read it as an input, so every rank traces the
+        same graph and neuronx-cc compiles it once (a Python int would be baked in, one NEFF per rank)."""
+        return self._neuron_rank_t
+
     def _eplb_layer(self) -> eplb.ExpertLoadBalancerLayerState:
         return eplb.current().layer(self.layer_idx)
 
@@ -379,21 +390,23 @@ class DeepseekMoE(nn.Module):
         else:
             gate_force_random = input_metadata.gate_force_random
 
-        eplb_layer = self._eplb_layer()
+        # Neuron without redundant experts: the logical->physical map is the identity, so skip it. Reading it indexes
+        # a Python list by layer_idx, which made every layer's MoE region graph its own (Dynamo guard on layer_idx).
+        eplb_layer = None if PLATFORM.is_neuron and self.num_extra_experts == 0 else self._eplb_layer()
 
         gate_kwargs = dict(
             num_topk=self.num_topk,
             use_shared_as_routed=False,
             num_shared_experts=0,
             routed_scaling_factor=self.routed_scaling_factor,
-            ep_rank=self.ep_rank,
+            ep_rank=self.neuron_rank_tensor() if PLATFORM.is_neuron else self.ep_rank,
             mask=input_metadata.valid_token_mask,
             bias=self.e_score_correction_bias,
             image_bias=self.e_score_correction_bias_vl if images_seq_mask is not None else None,
             image_token_mask=images_seq_mask,
             force_random=gate_force_random,
-            to_physical_map=eplb_layer.logical_to_physical_map,
-            logical_count=eplb_layer.logical_count,
+            to_physical_map=None if eplb_layer is None else eplb_layer.logical_to_physical_map,
+            logical_count=None if eplb_layer is None else eplb_layer.logical_count,
             scoring_func=self.scoring_func,
             out=out,
         )

deepinfer/neuron_ops/fp8.py

@@ -15,9 +15,11 @@ identically whenever the exponent field is below 15, i.e. for |x| < 256. So on N
 
 from __future__ import annotations
 
+import functools
 import os
 import shlex
 
+import numpy as np
 import torch
 
 FP8_MAX = float(os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224"))
@@ -65,13 +67,37 @@ _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)
+    """In place, on the host in fixed-size chunks (bounded int16 temporaries: the Engram tables are ~200 GB).
+    Device weights make one host round trip: eager device ops recompile per shape / view offset, which made this
+    step take over half an hour for the full model."""
+    if data.device.type != "cpu":
+        host = data.cpu()
+        _halve_inplace(host)
+        data.copy_(host.to(data.device))
+        return
+    flat = data.view(-1).view(torch.uint8)
+    table = _halve_table().numpy()
+    for start in range(0, flat.numel(), _CHUNK_BYTES):
         chunk = flat[start : start + _CHUNK_BYTES]
-        chunk.copy_(halve_e4m3fn_bytes(chunk))
+        # A clone per chunk: .numpy() on the parameter itself would pin its storage as non-resizable, and the
+        # runtime adaptation later frees merged weights with untyped_storage().resize_(0).
+        arr = chunk.clone().numpy()
+        np.take(table, arr, out=arr)  # one 256-entry lookup per byte
+        chunk.copy_(torch.from_numpy(arr))
+
+
+@functools.cache
+def _halve_table() -> torch.Tensor:
+    # Explicitly on the host: model loading runs under a default-device context that maps to the NeuronCore.
+    return halve_e4m3fn_bytes(torch.arange(256, dtype=torch.int32, device="cpu").to(torch.uint8))
 
 
 def _double_scale(sf: torch.Tensor) -> None:
+    if sf.device.type != "cpu":  # small: on the host, no per-shape eager compiles on the device
+        host = sf.cpu()
+        _double_scale(host)
+        sf.copy_(host.to(sf.device))
+        return
     if sf.dtype in (torch.float32, torch.bfloat16, torch.float16):
         sf.mul_(2)
         return

deepinfer/neuron_ops/nki_gemm.py

@@ -47,8 +47,11 @@ def prepare_model(model: torch.nn.Module) -> int:
         sf = getattr(p, "sf", None)
         if sf is None or p.dtype not in (torch.float8_e4m3fn, torch.int8) or p.dim() not in (2, 3):
             continue
-        scale = _kernel_scale(p.data, sf.data if isinstance(sf, torch.nn.Parameter) else sf)
+        # On the host: the scale tensors are small, and eager device ops would compile once per weight shape.
+        sf_host = (sf.data if isinstance(sf, torch.nn.Parameter) else sf).cpu()
+        scale = _kernel_scale(p.data, sf_host)
         if scale is not None:
+            scale = scale.to(p.device)
             p._nki_scale = scale
             # The weight bytes as their own uint8 tensor (same storage): a dtype bitcast of a graph input feeding the
             # kernel crashes neuronx-cc (InferNonlocalTensors), a uint8 input does not.

deepinfer/neuron_ops/nki_norm.py

@@ -77,11 +77,17 @@ def ue8m0_exponent(v: torch.Tensor) -> torch.Tensor:
 
 
 def pack_ue8m0_int16(e: torch.Tensor) -> torch.Tensor:
-    """int32 exponent bytes ``[..., g]`` -> int16 words ``[..., ceil(g / 2)]``, low byte first (arithmetic only)."""
+    """int32 exponent bytes ``[..., g]`` -> int16 words ``[..., ceil(g / 2)]``, low byte first (arithmetic only).
+
+    The word is assembled in fp32 (exact: |v| < 2**17) and converted fp32 -> int16 once. An int32 -> int16 narrowing
+    of the strided (``0::2`` / ``1::2``) int32 result -- with or without a select -- makes neuronx-cc emit a
+    same-shape reinterpret between 4- and 2-byte dtypes and die in codegen (NCC_IBCG902 "Cannot reinterpret to
+    dtype with different size", issue #13) when the exponents come from the NKI RMSNorm output."""
     if e.shape[-1] % 2:
         e = torch.nn.functional.pad(e, (0, 1))
-    v = e[..., 0::2] + e[..., 1::2] * 256
-    return torch.where(v >= 32768, v - 65536, v).to(torch.int16)
+    f = e.to(torch.float32)
+    lo, hi = f[..., 0::2], f[..., 1::2]
+    return (lo + hi * 256.0 - torch.where(hi >= 128.0, 65536.0, 0.0)).to(torch.int16)
 
 
 def fp8_cast(

deepinfer/neuron_ops/regions.py

@@ -7,6 +7,7 @@ fullgraph. Layers of the same kind trace to the same graph, so neuronx-cc's NEFF
 
 from __future__ import annotations
 
+import functools
 from collections.abc import Callable
 from typing import Any
 
@@ -33,25 +34,33 @@ def raise_dynamo_limits() -> None:
             setattr(torch._dynamo.config, name, 1 << 30)
 
 
+_ENTRIES: dict[tuple[type, str], Callable[..., Any]] = {}
+
+
 def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]:
-    """``module.<name>`` as its own fullgraph region, compiled once per module (cached on it)."""
-    cache = module.__dict__.setdefault("_neuron_region_entries", {})
-    fn = cache.get(name)
+    """``type(module).<name>`` as a region, compiled once per (class, method) and called with the module.
+
+    One compiled callable for all layers of a class: Dynamo then reuses the same graph (and the same loaded NEFF)
+    for every layer whose guards match (same layer kind and shapes). A per-module compile gave every layer its own
+    compiled artifact and NEFF load -- the same code loaded 43 times filled ~6 GB of HBM with model code.
+    """
+    key = (type(module), name)
+    fn = _ENTRIES.get(key)
     if fn is None:
         from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda
 
         raise_dynamo_limits()
         # fullgraph=False: split() marks sub-region boundaries inside the entry (smaller graphs for neuronx-cc).
         compiled = torch.compile(
-            getattr(module, name), backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, fullgraph=False, dynamic=False
+            getattr(type(module), name), backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, fullgraph=False, dynamic=False
         )
 
-        def fn(*args: Any, **kwargs: Any) -> Any:
+        def fn(mod: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
             with hide_cuda():
-                return compiled(*args, **kwargs)
+                return compiled(mod, *args, **kwargs)
 
-        cache[name] = fn
-    return fn
+        _ENTRIES[key] = fn
+    return functools.partial(fn, module)
 
 
 @torch._dynamo.disable

deepinfer/neuron_ops/static_moe.py

@@ -148,7 +148,10 @@ def _nki_moe_forward(
     num_local = moe.num_local_physical_experts
     num_tokens = x.shape[0]
     cap = max(moe.max_num_decoding_tokens, num_tokens)
-    idx, w = topk_idx.to(torch.int32), topk_weights.to(torch.float32)
+    x_local = x
+    # Local expert ids relative to this rank, from a rank *tensor*: the graph stays identical across ranks.
+    idx = topk_idx.to(torch.int32)
+    w = topk_weights.to(torch.float32)
     if size > 1:
         pad = cap - num_tokens
         x = torch.nn.functional.pad(x, (0, 0, 0, pad))
@@ -156,8 +159,9 @@ def _nki_moe_forward(
         w = torch.nn.functional.pad(w, (0, 0, 0, pad))
         x, idx, w = (_all_gather(t, group, size) for t in (x, idx, w))
     acc = torch.zeros((x.shape[0], moe.hidden_size), dtype=torch.float32, device=x.device)
+    local_idx = idx - moe.neuron_rank_tensor() * num_local
     for e in range(num_local):
-        w_e = torch.where(idx == rank * num_local + e, w, torch.zeros_like(w)).sum(-1)
+        w_e = torch.where(local_idx == e, w, torch.zeros_like(w)).sum(-1)
         gate_up = nki_gemm.linear_qdq(x, moe.gate_up_proj_weight, group=e)
         act = swiglu_forward_and_per_token_cast(
             gate_up, fmt="bf16", num_per_channels=None, topk_weights=w_e, clamp_value=moe.swiglu_clamp_value
@@ -165,8 +169,7 @@ def _nki_moe_forward(
         acc = acc + nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float()
     out = _reduce_scatter(acc, group, size, rank)[:num_tokens]
     if moe.has_shared:
-        shared = nki_gemm.linear_qdq(x[:num_tokens] if size == 1 else x.narrow(0, rank * cap, num_tokens),
-                                     moe.shared_gate_up_proj.weight)
+        shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)
         shared = swiglu_forward_and_per_token_cast(
             shared, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value
         )

deepinfer/neuron_ops/vendor/tile_kernels/moe/moe_topk_gate_forward_kernel.py

@@ -22,7 +22,7 @@ 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:
+def _hash_uniform(a: torch.Tensor, b: torch.Tensor, seed) -> 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)
@@ -31,14 +31,19 @@ def _hash_uniform(a: torch.Tensor, b: torch.Tensor, seed: int) -> torch.Tensor:
     return (h & 0xFFFFFF).to(torch.float32) * (1.0 / (1 << 24))
 
 
+def _seed(ep_rank):
+    # A tensor rank keeps compiled graphs identical across ranks (one NEFF for all instead of one per rank).
+    return ep_rank.to(torch.int32) if isinstance(ep_rank, torch.Tensor) else int(ep_rank)
+
+
 def _force_random(num_tokens, num_physical_topk, num_physical_experts, max_physical_experts, ep_rank, 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))
+    keys = _hash_uniform(token, cols[None, :], _seed(ep_rank))
     weights = _hash_uniform(token, torch.arange(num_physical_topk, dtype=torch.int32, device=device)[None, :] + max_physical_experts,
-                            int(ep_rank))
+                            _seed(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/worker/base.py

@@ -18,7 +18,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 clone_and_zero_states_and_blocks, set_engram_lookup_table
 from deepinfer.config import CacheConfig, ModelConfig, ParallelConfig, SchedulerConfig
 from deepinfer.engine.engine_type import EngineType
@@ -163,6 +163,8 @@ class WorkerBase:
         clone_blocks: np.ndarray[np.int32],
         zero_states: np.ndarray[np.int32],
     ) -> tuple[np.ndarray[np.int32], np.ndarray[np.int32]]:
+        if PLATFORM.is_neuron:
+            return _neuron_clone_and_zero(self.cache_engine.gpu_cache, clone_states, clone_blocks, zero_states)
         return clone_and_zero_states_and_blocks(
             self.cache_engine.gpu_states, self.cache_engine.gpu_blocks, clone_states, clone_blocks, zero_states
         )
@@ -355,3 +357,38 @@ def _bind_numa(local_rank: int) -> None:
         bind_numa(which_numa(local_rank))
     except Exception:
         logger.warning("NUMA 绑定失败，内存落位退化为 first-touch，可能导致 H2D 性能下降", exc_info=True)
+
+
+def _neuron_clone_and_zero(
+    gpu_cache: object, clone_states: np.ndarray, clone_blocks: np.ndarray, zero_states: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+    """``clone_and_zero_states_and_blocks`` for Neuron's layer-major cache layout (cache_engine._LayerMajorColumns):
+    a block or state slot is spread over every layer's view, so copy / zero row ``i`` of each view instead of one
+    contiguous block. Same contract as the C++ version: returns the cloned source ids (states, blocks)."""
+    block_views, state_views, seen = [], [], set()
+    for cache in gpu_cache.cache_list:
+        for name in ("ica_cache", "ca_cache"):
+            t = getattr(cache, name)
+            if t is not None and id(t) not in seen:
+                seen.add(id(t))
+                block_views.append(t)
+        for name in ("swa_cache", "ica_state", "ca_state"):
+            t = getattr(cache, name)
+            if t is not None and id(t) not in seen:
+                seen.add(id(t))
+                state_views.append(t)
+    released_states, released_blocks = [], []
+    for src, dst in clone_states.reshape(-1, 2) if clone_states.size else []:
+        if dst:
+            for t in state_views:
+                t[int(dst)].copy_(t[int(src)])
+        released_states.append(int(src))
+    for src, dst in clone_blocks.reshape(-1, 2) if clone_blocks.size else []:
+        if dst:
+            for t in block_views:
+                t[int(dst)].copy_(t[int(src)])
+        released_blocks.append(int(src))
+    for i in zero_states.reshape(-1) if zero_states.size else []:
+        for t in state_views:
+            t[int(i)].zero_()
+    return np.asarray(released_states, dtype=np.int32), np.asarray(released_blocks, dtype=np.int32)

deepinfer/worker/cache_engine.py

@@ -306,6 +306,50 @@ class StateBlock:
         self.state_size_to_disk = self.swa_store_total_bytes_inference_pad
 
 
+class _LayerMajorColumns:
+    """Neuron: a ``[rows, cols]`` byte buffer laid out column-region-major (vllm-neuron's per-layer layout).
+
+    The engine slices each layer's cache as a column range ``[:, a:b]`` of a block-major (or state-slot-major) buffer.
+    Here the same bytes hold the range ``[a, b)`` as one contiguous ``[rows, b - a]`` region at ``flat[a*rows:b*rows]``:
+    every layer view is then contiguous and has no per-layer offset inside it, so region graphs take the views as
+    inputs and same-kind layers share one compiled NEFF (strided column views were made contiguous -- copied -- by
+    the executor on every call, and slicing them in-graph baked each layer's offsets into its own NEFF).
+    """
+
+    def __init__(self, flat: torch.Tensor, rows: int, col0: int, cols: int) -> None:
+        self.flat, self.rows, self.col0, self.cols = flat, rows, col0, cols
+
+    @classmethod
+    def of(cls, t: torch.Tensor) -> _LayerMajorColumns:
+        assert t.dim() == 2 and t.dtype == torch.uint8 and t.is_contiguous()
+        return cls(t.view(-1), t.shape[0], 0, t.shape[1])
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return (self.rows, self.cols)
+
+    def stride(self, dim: int | None = None) -> int | tuple[int, int]:
+        return (self.cols, 1) if dim is None else (self.cols, 1)[dim]
+
+    def __getitem__(self, idx: tuple[slice, slice]) -> _LayerMajorColumns:
+        rows, cols = idx
+        assert rows == slice(None), idx
+        a, b, step = cols.indices(self.cols)
+        assert step == 1
+        return _LayerMajorColumns(self.flat, self.rows, self.col0 + a, b - a)
+
+    def tensor(self) -> torch.Tensor:
+        start = self.col0 * self.rows
+        return self.flat[start : start + self.cols * self.rows].view(self.rows, self.cols)
+
+    def view(self, *shape: int) -> torch.Tensor:
+        return self.tensor().view(*shape)
+
+    def as_strided(self, size: tuple[int, ...], stride: tuple[int, ...]) -> torch.Tensor:
+        # The block-major strides asked for describe this same logical layout; the region is contiguous.
+        return self.tensor().view(*size)
+
+
 class KVCacheManager:
     """管理 KV cache 的分配和布局。
 
@@ -315,6 +359,10 @@ class KVCacheManager:
 
     def __init__(self, cfg: AttrDict, kv_tensor: torch.Tensor | None, state_tensor: torch.Tensor | None) -> None:
         self.cfg = cfg
+        if PLATFORM.is_neuron and isinstance(kv_tensor, torch.Tensor):
+            kv_tensor = _LayerMajorColumns.of(kv_tensor)
+        if PLATFORM.is_neuron and isinstance(state_tensor, torch.Tensor):
+            state_tensor = _LayerMajorColumns.of(state_tensor)
         self.kv_tensor = kv_tensor
 
         model_metadata = ModelMetadata(cfg)

deepinfer/worker/neuron_graph_runner.py

@@ -41,7 +41,12 @@ class _GraphKVCaches:
     """
 
     def __init__(self, kv_caches: GPUCacheInterface) -> None:
+        self.cache_list = kv_caches.cache_list
         self.roots: list[torch.Tensor] = []
+        if envs.DEEPINFER_NEURON_GRAPH_MODE == "regions":
+            # The cache engine lays every layer view out contiguously on Neuron (_LayerMajorColumns); the region
+            # graphs take them as inputs as they are.
+            return
         self.recipes: list[dict[str, tuple | None]] = []
         for cache in kv_caches.cache_list:
             recipe = {}
@@ -82,15 +87,16 @@ class _GraphKVCaches:
 
     def __getitem__(self, layer_idx: int) -> KVCache | _LazyKVCache:
         if envs.DEEPINFER_NEURON_GRAPH_MODE == "regions":
-            # Views built on attribute access, i.e. inside the layer's region graph (roots are its inputs).
-            return _LazyKVCache(self, layer_idx)
+            # The layer's own contiguous, non-overlapping views as region graph inputs (like vllm-neuron): slicing
+            # them from the roots inside the graph baked each layer's byte offsets in, one NEFF per layer.
+            return self.cache_list[layer_idx]
         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)
+        return layer_idx < len(self.cache_list)
 
 
 class _LazyKVCache:

examples/neuron/offline_inference_neuron.py

@@ -40,6 +40,8 @@ parser.add_argument("--max_tokens", default=16, type=int)
 parser.add_argument("--num_kv_blocks", default=400, type=int, help="fixed KV pool per rank; the scheduler admits a request only if max_model_len tokens fit, per running sequence")
 parser.add_argument("--max_model_len", default=16384, type=int, help="the candidate-block filter needs >= candidate_topk_blocks * candidate_block_size (2048 * 8)")
 parser.add_argument("--max_num_batched_tokens", default=1024, type=int)
+parser.add_argument("--max_num_seqs", default=1, type=int,
+                    help="decode batch; every captured batch size adds one set of region NEFFs per layer (HBM: model code)")
 args = parser.parse_args()
 dp_size = args.ep
 
@@ -77,7 +79,7 @@ def main(dp_size: int, dp_rank: int) -> None:
         data_parallel_size=dp_size,
         max_num_batched_tokens=args.max_num_batched_tokens,
         max_model_len=args.max_model_len,
-        max_num_seqs=len(PROMPTS),
+        max_num_seqs=args.max_num_seqs,
         gpu_memory_utilization=0.85,
         lookahead_tokens=0,
     )

tests/neuron/compile/test_nki_hc.py

@@ -345,3 +345,11 @@ def test_shifted_post_pre_chain(cfg: str, n: int, mode: str) -> None:
     close(got[4], ref[4])
     npc = get_fp8_cast_params()["num_per_channels"]
     torch.testing.assert_close(_dequant(got[5], got[6], npc), _dequant(ref[5], ref[6], npc), rtol=0.13, atol=1e-2)
+
+
+def test_pack_ue8m0_int16_exhaustive() -> None:
+    """Every exponent byte pair packs to the same int16 word as the byte view (host only)."""
+    e = torch.stack(torch.meshgrid(torch.arange(256), torch.arange(256), indexing="ij"), -1).reshape(-1, 2).to(torch.int32)
+    words = nki_norm.pack_ue8m0_int16(e)
+    assert words.dtype == torch.int16 and words.shape == (65536, 1)
+    assert torch.equal(words.contiguous().view(torch.uint8).view(-1, 2).to(torch.int32), e)