Neuron: full DeepSeek-V4.1-Flash loads and prefills on 64 NeuronCores (EP=64)

Device-memory fixes found bringing up the released checkpoint (16.45 GiB of
weights per 24 GB core):
- Engram: the Neuron EngramBuffer keeps only this rank's SF rows, so skip the
  all-gather of the replicated SF table (~3 GB per layer, allocated on the
  device -> "Failed to allocate resource").
- gather_tl: index_select / advanced indexing of an int8 table allocates
  several times the table on the device (+4.9 GiB for the 1.3 GB embedding);
  gather through F.embedding on an int32 / fp16 view instead (bit-exact, no
  extra memory).
- deep_gemm FP8/FP4 GEMMs: dequantize the weight in N chunks of 32M elements
  and one expert at a time instead of whole projections / all local experts
  in fp32 (OOM on a fragmented pool); bit-identical results.
- transfer_to_neuron: import torch.utils._device explicitly (missing when
  torch_neuronx is not auto-loaded).

examples/neuron/offline_inference_neuron.py: offline inference in basic mode
(text prompts, no MTP / EPLB, fixed KV pool, max_model_len >= 2048 * 8 for
the candidate-block filter), one process per NeuronCore.

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

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

@@ -5,6 +5,21 @@ from collections.abc import Callable
 import torch
 
 
+def _gather_rows(x: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
+    """``x[index]`` for a byte table ``x [n, D]``.
+
+    On the neuron device, index_select / advanced indexing of an int8 table allocates several times the table
+    (+4.9 GiB for the 1.3 GB V4.1 embedding, which ran the full model out of HBM), while F.embedding through a
+    4- or 2-byte view copies only the selected rows and is bit-exact (verified on arbitrary bytes, NaN patterns
+    included).
+    """
+    if x.is_contiguous():
+        for elem, dtype in ((4, torch.int32), (2, torch.float16)):
+            if x.shape[1] % elem == 0 and x.storage_offset() % elem == 0:
+                return torch.nn.functional.embedding(index, x.view(dtype)).view(torch.int8)
+    return torch.index_select(x, 0, index)
+
+
 def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: int) -> Callable[..., None]:
     """Kernel args: ``x [n, D]`` (row stride ``x_stride``), ``index [m]``, ``y [m, D]`` contiguous.
 
@@ -21,7 +36,7 @@ def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: in
             return
         valid = (index >= 0) & (index < n)
         safe = torch.where(valid, index, torch.zeros_like(index))
-        rows = torch.index_select(x[:, :D], 0, safe)
+        rows = _gather_rows(x[:, :D], safe)
         y.copy_(torch.where(valid.view(m, 1), rows, torch.zeros((), dtype=rows.dtype, device=rows.device)))
 
     return _gather

deepinfer/model_executor/models/v4_modules/engram.py

@@ -96,12 +96,12 @@ class NgramEmbedding(nn.Module):
     def gather_sf(self) -> torch.Tensor:
         assert self.sf is not None
         local_sf = _pack_ue8m0_to_int(self.sf)
-        sf = local_sf.new_empty(self.engram_group.size() * self.num_local_embeddings, local_sf.size(1))
         if PLATFORM.is_neuron:
-            # Gather raw bytes: gloo (host mode) has no int16/int32 collectives for this dtype.
-            dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.contiguous().view(torch.uint8), group=self.engram_group)
-        else:
-            dist.all_gather_into_tensor(sf, local_sf, group=self.engram_group)
+            # The Neuron EngramBuffer keeps only this rank's SF rows (it serves its own table shard), so skip the
+            # all-gather: the replicated table is ~3 GB per layer for the released checkpoint, allocated on the device.
+            return local_sf.contiguous()
+        sf = local_sf.new_empty(self.engram_group.size() * self.num_local_embeddings, local_sf.size(1))
+        dist.all_gather_into_tensor(sf, local_sf, group=self.engram_group)
         return sf
 
     def release_storage(self) -> None:

deepinfer/neuron_ops/transfer_to_neuron.py

@@ -24,6 +24,7 @@ import types
 from typing import Any
 
 import torch
+import torch.utils._device  # not imported by `import torch` alone (e.g. TORCH_DEVICE_BACKEND_AUTOLOAD=0)
 
 _REAL_DEVICE = torch.device
 _TARGET: torch.device | None = None

deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py

@@ -176,8 +176,12 @@ class EngramBuffer(BufferBase):
         if sfs is not None:
             for layer_idx, sf in enumerate(sfs):
                 num_entries = self.num_entries_per_layer[layer_idx]
-                assert sf.dim() == 2 and sf.is_contiguous()
-                assert sf.size(0) == self.num_ranks * num_entries and sf.size(1) == self.num_sf_packs
+                assert sf.dim() == 2 and sf.is_contiguous() and sf.size(1) == self.num_sf_packs
+                if sf.size(0) == num_entries:
+                    # Already this rank's slice (DeepInfer on Neuron skips the replicated-table all-gather).
+                    self._sf_shards.append(sf.clone())
+                    continue
+                assert sf.size(0) == self.num_ranks * num_entries
                 begin = self.rank_idx * num_entries
                 self._sf_shards.append(sf[begin:begin + num_entries].clone())
 

deepinfer/neuron_ops/vendor/deep_gemm/gemm.py

@@ -56,6 +56,23 @@ def _dequant_operand(
     return dequantize(data, sf, gran_k=recipe[2])
 
 
+# Dequantized weight rows per chunk (fp32 elements). The device is a 24 GB NeuronCore that already holds the
+# packed weights: dequantizing a whole projection (or all local experts) to fp32 at once ran the full model out
+# of HBM. Chunking along N keeps every output element's K-reduction intact, so results are unchanged.
+_CHUNK_ELEMS = 1 << 25
+
+
+def _row_chunks(n: int, k: int) -> list[tuple[int, int]]:
+    step = max(1, _CHUNK_ELEMS // max(1, k))
+    return [(n0, min(n, n0 + step)) for n0 in range(0, n, step)]
+
+
+def _required_sf(sf, mn, k, recipe, is_sfa, num_groups=None, disable_ue8m0_cast=False):
+    return transform_sf_into_required_layout(
+        sf, mn, k, recipe, num_groups=num_groups, is_sfa=is_sfa, disable_ue8m0_cast=disable_ue8m0_cast
+    )
+
+
 # --------------------------------------------------------------------------- #
 # BF16
 # --------------------------------------------------------------------------- #
@@ -177,8 +194,12 @@ def _fp8_fp4_gemm(
         return
     full_recipe = get_recipe(recipe, recipe_a, recipe_b)
     a_f32 = _dequant_operand(a, sfa, m, k, full_recipe, True, disable_ue8m0_cast=disable_ue8m0_cast)
-    b_f32 = _dequant_operand(b, sfb, n, k, full_recipe, False, disable_ue8m0_cast=disable_ue8m0_cast)
-    _write_output(a_f32 @ b_f32.transpose(0, 1), d, c, alpha)
+    sfb_req = _required_sf(sfb, n, k, full_recipe, False, disable_ue8m0_cast=disable_ue8m0_cast)
+    c_is_d = c is not None and _is_alias(c, d)
+    for n0, n1 in _row_chunks(n, k):
+        b_f32 = dequantize(b[n0:n1], sfb_req[n0:n1], gran_k=full_recipe[2])
+        c_chunk = None if c is None else (d[:, n0:n1] if c_is_d else c[..., n0:n1])
+        _write_output(a_f32 @ b_f32.transpose(0, 1), d[:, n0:n1], c_chunk, alpha)
 
 
 def fp8_fp4_gemm_nt(
@@ -324,9 +345,7 @@ def m_grouped_fp8_fp4_gemm_nt_contiguous(
         return
     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)
-    b_f32 = _dequant_operand(
-        b, sfb, n, k, full_recipe, False, num_groups=num_groups, 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)
 
     spans = _psum_group_spans(grouped_layout, get_mk_alignment_for_contiguous_layout(), total_m)
     assert len(spans) == num_groups, f"{len(spans)=}, {num_groups=}"
@@ -334,8 +353,11 @@ def m_grouped_fp8_fp4_gemm_nt_contiguous(
     for group_idx, (start, end) in enumerate(spans):
         if end <= start:
             continue
-        out = a_f32[start:end] @ b_f32[group_idx].transpose(0, 1)
-        d[start:end].copy_(out.to(d.dtype))
+        # One expert (and one N chunk) at a time: never all local experts dequantized at once.
+        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])
+            out = a_f32[start:end] @ b_f32.transpose(0, 1)
+            d[start:end, n0:n1].copy_(out.to(d.dtype))
 
 
 def m_grouped_fp8_fp4_gemm_nn_contiguous(

examples/neuron/offline_inference_neuron.py

@@ -0,0 +1,98 @@
+"""Offline inference with the full DeepSeek-V4.1-Flash on one trn2.48xlarge (Neuron basic mode).
+
+Same flow as examples/offline_inference.py (launch_dp -> one engine process per DP rank, each owning one
+NeuronCore), with the settings basic mode needs:
+  * text-only prompts (VL kernels are stubbed in basic mode);
+  * no DSpark/MTP drafting (DEEPINFER_USE_NEXTN_DECODE=0) and no EPLB;
+  * a fixed KV-block pool instead of the engine's memory-profiling prefill (a max_num_batched_tokens dummy
+    prefill does not fit next to the eager pure-torch kernels);
+  * small batches and few output tokens: eager execution compiles every new op/shape on the first run.
+
+    export DEEPINFER_ACCELERATOR=neuron WORLD_SIZE=1 RANK=0
+    python examples/neuron/offline_inference_neuron.py --model_name /mnt/nvme/models/DeepSeek-V4.1-Flash-1017 --ep 64
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import sys
+from pathlib import Path
+
+EXAMPLES = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(EXAMPLES))
+
+from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env, setup_cuda_visible_devices  # noqa: E402
+from _launch import launch_dp  # noqa: E402
+
+PROMPTS = [
+    "问题：法国的首都是哪里？\n答案：巴黎\n\n问题：一年有几个月？\n答案：12 个月\n\n问题：中国的首都是哪里？\n答案：",
+    "Question: What is the largest planet in the solar system?\nAnswer: Jupiter\n\n"
+    "Question: How many legs does a spider have?\nAnswer: Eight\n\n"
+    "Question: Who wrote Romeo and Juliet?\nAnswer:",
+]
+
+parser = argparse.ArgumentParser(description="V4.1 offline inference on Neuron (basic mode)")
+parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME, help="HF model id or local HF checkpoint directory")
+parser.add_argument("--ep", default=64, type=int)
+parser.add_argument("--max_tokens", default=16, type=int)
+parser.add_argument("--num_kv_blocks", default=176, type=int, help="fixed KV pool per rank; must hold max_model_len tokens")
+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)
+args = parser.parse_args()
+dp_size = args.ep
+
+
+def _fixed_kv_pool(num_blocks: int) -> None:
+    """Skip WorkerBase.profile_num_available_blocks (a max_num_batched_tokens dummy prefill)."""
+    from deepinfer.worker.base import WorkerBase
+
+    def profile_num_available_blocks(self):  # noqa: ANN001, ANN202
+        self.model_runner.set_block_size(self.cache_config.block_size)
+        self.cache_config.num_gpu_blocks = num_blocks
+        return num_blocks
+
+    WorkerBase.profile_num_available_blocks = profile_num_available_blocks
+
+
+def main(dp_size: int, dp_rank: int) -> None:
+    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
+    os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
+    setup_common_env(model_name=args.model_name, sampler_force_topk=4096, set_override_experts=False, enable_eplb=False)
+    # Pick this rank's NeuronCore before anything initializes deepinfer.PLATFORM (importing deepinfer.worker does):
+    # PlatformNeuron claims the core from CUDA_VISIBLE_DEVICES once, at init.
+    from deep_cluster.hardware import get_actual_num_gpus
+
+    setup_cuda_visible_devices(dp_rank, 1, get_actual_num_gpus())
+    _fixed_kv_pool(args.num_kv_blocks)
+
+    llm = create_llm(
+        args.model_name,
+        dp_rank,
+        ep_size=args.ep,
+        cp_size=1,
+        env={"DEEPINFER_LOGGING_INTERVAL_STEP": "1"},
+        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),
+        gpu_memory_utilization=0.85,
+        lookahead_tokens=0,
+    )
+
+    from deepinfer.sampling_params import SamplingParams
+
+    params = [SamplingParams(temperature=0.0, max_tokens=args.max_tokens) for _ in PROMPTS]
+    outputs = llm.generate(prompts=PROMPTS, sampling_params=params, use_tqdm=False)
+    for i, output in enumerate(outputs):
+        text = llm.llm_engine.tokenizer.decode([x.token_id for x in output])
+        logps = [x.token_logprobs for x in output]
+        ppl = math.exp(-sum(logps) / len(logps)) if logps else float("nan")
+        if dp_rank == 0:
+            print(f"DP {dp_rank}; INPUT: {PROMPTS[i]!r}; OUTPUT: {text!r}; output_tokens={len(output)}; decode_ppl={ppl}", flush=True)
+    llm.close()
+
+
+if __name__ == "__main__":
+    launch_dp(main, dp_size)