Neuron: prompt chunks through the decode region graphs (issue #26)

- PrepareArgs.prompt_as_decode: decode-layout metadata for a prompt chunk;
  input_ids are the prompt's own tokens instead of last token + MTP drafts.
- DEEPINFER_NEURON_PREFILL_BUCKET adds one graph size; the model runner pads
  a prompt chunk to it with the decode placeholders (every EP rank takes the
  same bucket, so the static collectives match). Sampling still sees a prompt
  step.
- nki_attention.paged_mqa_logits: one gather + one kernel call for a chunk
  (one sequence) instead of a T-way loop.
- bf16 linear for the LM head eagerly on a NeuronCore (b.float() was a
  2.6 GiB copy a call); DEEPINFER_NEURON_REGION_TIMING; --num_prompts.

Host mode, tiny V4.1, chunk 128: same tokens as eager prefill (short 9/9,
long 715 tokens in 6 chunks 8/9 against golden).

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

csrc/engine/input_metadata/prepare_tensor_input.hpp

@@ -25,6 +25,9 @@ struct PrepareArgs {
     int32_t image_sentinel = -1;
     bool is_prompt = true;
     bool force_load_balance = false;
+    // decode-shaped metadata for a prompt chunk (Neuron prefill through the decode graphs): input_ids are the prompt's
+    // own tokens rather than the last token + MTP drafts; everything else follows the decode layout.
+    bool prompt_as_decode = false;
 
     // prompt-only
     int cp_size = 1;

csrc/engine/input_metadata/prepare_tensor_input_base.cpp

@@ -97,12 +97,13 @@ static py::dict prepare_common_base(const PrepareArgs &args) {
 
     auto input_ids = get_or_create_tensor(result, has_out, "input_ids", {total_q}, torch::kInt32);
     auto p_input_ids = input_ids.mutable_data_ptr<int32_t>();
+    const bool prompt_ids = is_prompt || args.prompt_as_decode;
     i_q = 0;
     for (auto p_seq : seq_metadata_list) {
         auto &token_ids = p_seq->seq_data->token_ids;
         static_assert(std::is_same_v<decltype(*p_input_ids), decltype(*token_ids.data())>);
         int seqlen_q_golden;
-        if (is_prompt) {
+        if (prompt_ids && p_seq->is_prompt) {
             TORCH_CHECK(int(token_ids.size()) >= p_seq->seqlen_k);
             seqlen_q_golden = p_seq->seqlen_q;
         } else {
@@ -110,7 +111,7 @@ static py::dict prepare_common_base(const PrepareArgs &args) {
             seqlen_q_golden = 1;
         }
         memcpy(p_input_ids + i_q, token_ids.data() + p_seq->seqlen_k - p_seq->seqlen_q, seqlen_q_golden * sizeof(int32_t));
-        if (!is_prompt) {
+        if (seqlen_q_golden < p_seq->seqlen_q) {
             int seqlen_q_draft = p_seq->seqlen_q - 1;
             TORCH_CHECK(int(p_seq->mtp.draft_tokens.size()) >= seqlen_q_draft);
             for (int j = 0; j < seqlen_q_draft; ++j) {

csrc/engine/pybind_output.cpp

@@ -67,6 +67,7 @@ void pybind_engine_output(pybind11::module &engine) {
         .def_readwrite("yoco_prefill_keep_tokens", &PrepareArgs::yoco_prefill_keep_tokens)
         .def_readwrite("max_backtracking_tokens", &PrepareArgs::max_backtracking_tokens)
         .def_readwrite("force_load_balance", &PrepareArgs::force_load_balance)
+        .def_readwrite("prompt_as_decode", &PrepareArgs::prompt_as_decode)
         .def_readwrite("dense_compress_ratios", &PrepareArgs::dense_compress_ratios)
         .def_readwrite("sparse_compress_ratios", &PrepareArgs::sparse_compress_ratios);
 

deepinfer/environ.py

@@ -114,8 +114,14 @@ 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)
+    # Neuron only: synchronize around every region graph call and log per-kind times each decode step (profiling).
+    DEEPINFER_NEURON_REGION_TIMING = EnvVar.bool()
     # 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, regions mode: run prompt chunks through the decode region graphs padded to this many tokens
+    # (issue #26, the vllm-neuron#71 layout). 0 = eager prefill. Needs max_num_batched_tokens <= bucket and
+    # DEEPINFER_MAX_BACKTRACKING_TOKENS >= bucket (the SWA / compressor rings must hold a whole chunk).
+    DEEPINFER_NEURON_PREFILL_BUCKET = EnvVar.int(default=0)
     # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break).
     DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool()
 

deepinfer/model_executor/model_metadata_utils.py

@@ -152,8 +152,10 @@ class ModelMetadata:
         window_size: int,
         swa_block_size: int,
         metadata_zos: MetadataZOS | None = None,
+        prompt_as_decode: bool = False,
     ) -> dict:
-        is_prompt = seq_metadata_list[0].is_prompt
+        # prompt_as_decode: decode-layout metadata for a prompt chunk (Neuron prefill through the decode graphs).
+        is_prompt = seq_metadata_list[0].is_prompt and not prompt_as_decode
 
         prepare_func = self.prepare_prompt if is_prompt else self.prepare_decode
 
@@ -168,6 +170,7 @@ class ModelMetadata:
             args.cp_size = PG.context_parallel_group().size()
         else:
             args.is_prompt = False
+            args.prompt_as_decode = prompt_as_decode
             args.max_num_blocks = max_num_blocks
             args.out = out_dict
         args.window_size = window_size

deepinfer/neuron_ops/nki_attention.py

@@ -29,6 +29,8 @@ import os
 import torch
 
 ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_ATTN", "1") == "1"
+#: DEEPINFER_NEURON_PREFILL_BUCKET (issue #26): a step of exactly this many query tokens is one prompt chunk.
+PREFILL_BUCKET = int(os.environ.get("DEEPINFER_NEURON_PREFILL_BUCKET", "0"))
 
 D_QK = 512
 INDEX_DIM = 128
@@ -320,6 +322,23 @@ def paged_mqa_logits(
     slot = torch.clamp(pos // page_kv, max=width - 1)
     lens = context_lens.reshape(num_q, 1).to(torch.int32)
     valid = (pos.view(1, n) < lens) & (pos.view(1, n) < width * page_kv)
+    w = weights.to(torch.float32)
+    kernel = _kernel("index_score", "weighted_scores")
+
+    if num_q == PREFILL_BUCKET:
+        # A prompt chunk run as a decode step (issue #26): its tokens are one sequence's, so they share token 0's block
+        # table and differ only in their context length. Gather the keys once and score every token in one kernel call
+        # (padding rows past the chunk score garbage on the same keys; their results are never read).
+        page = block_table[0].to(torch.int32).index_select(0, slot)
+        page = torch.where(pos < lens[0], page, torch.zeros_like(page)).view(1, n)
+        offset = (pos % page_kv).view(1, n)
+        data = gather_rows(kv_cache, page, offset, data_bytes, data_bytes).to(torch.int32)
+        sf = gather_rows(kv_cache, page, offset, SF_BYTES, SF_BYTES, base=page_kv * data_bytes).to(torch.int32)
+        values = _unpack_e2m1(data) if is_fp4 else _e4m3(data)  # [1, n, 128] fp32
+        keys = (values.view(n, SF_BYTES, 32) * _pow2(sf.view(n, SF_BYTES) - 127).unsqueeze(-1)).view(n, INDEX_DIM)
+        scores = kernel(q.reshape(num_q * heads, -1).contiguous(), keys.to(torch.bfloat16), w.transpose(0, 1).contiguous())
+        return _finish_logits(scores, valid, max_context_len, num_q, device)
+
     rows = torch.arange(num_q, device=device, dtype=torch.int32).view(num_q, 1) * width + slot.view(1, n)
     page = block_table.to(torch.int32).reshape(-1).index_select(0, rows.reshape(-1)).view(num_q, n)
     page = torch.where(valid, page, torch.zeros_like(page))  # stale table slots past the context: read page 0
@@ -331,12 +350,13 @@ def paged_mqa_logits(
     keys = (values.view(num_q, n, SF_BYTES, 32) * _pow2(sf - 127).unsqueeze(-1)).view(num_q, n, INDEX_DIM)
     keys = keys.to(torch.bfloat16)
 
-    w = weights.to(torch.float32)
-    kernel = _kernel("index_score", "weighted_scores")
     # Each query token has its own context: one kernel call per token (the kernel shares k across its tokens).
     scores = torch.cat([kernel(q[t].contiguous(), keys[t], w[t].reshape(heads, 1)) for t in range(num_q)], dim=0)
-    scores = torch.where(valid, scores, torch.full_like(scores, float("-inf")))[:, :max_context_len]
+    return _finish_logits(scores, valid, max_context_len, num_q, device)
 
+
+def _finish_logits(scores: torch.Tensor, valid: torch.Tensor, max_context_len: int, num_q: int, device) -> torch.Tensor:  # noqa: ANN001
+    scores = torch.where(valid, scores, torch.full_like(scores, float("-inf")))[:, :max_context_len]
     stride = _ceil(max_context_len, 256)  # deep_gemm's 512 B row alignment
     logits = torch.empty((num_q, stride), dtype=torch.bfloat16, device=device)[:, :max_context_len]
     logits.copy_(scores.to(torch.bfloat16))

deepinfer/neuron_ops/regions.py

@@ -8,6 +8,8 @@ fullgraph. Layers of the same kind trace to the same graph, so neuronx-cc's NEFF
 from __future__ import annotations
 
 import functools
+import time
+from collections import defaultdict
 from collections.abc import Callable
 from typing import Any
 
@@ -35,6 +37,15 @@ def raise_dynamo_limits() -> None:
 
 
 _ENTRIES: dict[tuple[type, str], Callable[..., Any]] = {}
+#: DEEPINFER_NEURON_REGION_TIMING: seconds per region kind in the current step (each call synchronized).
+_TIMES: dict[str, list[float]] = defaultdict(list)
+
+
+def take_timings() -> dict[str, list[float]]:
+    """The region timings recorded since the last call, then reset."""
+    out = dict(_TIMES)
+    _TIMES.clear()
+    return out
 
 
 def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]:
@@ -55,9 +66,18 @@ def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]:
             getattr(type(module), name), backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, fullgraph=False, dynamic=False
         )
 
+        label = f"{type(module).__name__}.{name}"
+
         def fn(mod: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
             with hide_cuda():
-                return compiled(mod, *args, **kwargs)
+                if not envs.DEEPINFER_NEURON_REGION_TIMING:
+                    return compiled(mod, *args, **kwargs)
+                torch.neuron.synchronize()
+                started = time.perf_counter()
+                out = compiled(mod, *args, **kwargs)
+                torch.neuron.synchronize()
+                _TIMES[label].append(time.perf_counter() - started)
+                return out
 
         _ENTRIES[key] = fn
     return functools.partial(fn, module)

deepinfer/neuron_ops/vendor/deep_gemm/gemm.py

@@ -106,9 +106,10 @@ def _bf16_gemm(
     if k == 0:
         _write_output(torch.zeros((m, n), dtype=torch.float32, device=d.device), d, c, alpha)
         return
-    if torch.compiler.is_compiling():
+    if torch.compiler.is_compiling() or b.device.type == "neuron":
         # NEURON graphs: one bf16 dot with fp32 accumulation and no materialized fp32 transpose (vllm-neuron#71 runs
         # these as bf16 linears; an fp32 transpose feeding a dot is a neuronx-cc InferNonlocalTensors trigger).
+        # Eagerly on a NeuronCore too: ``b.float()`` of the LM head is a 2.6 GiB copy a call.
         _write_output(torch.nn.functional.linear(a, b).float(), d, c, alpha)
         return
     _write_output(a.float() @ b.float().transpose(0, 1), d, c, alpha)

deepinfer/worker/model_runner.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import copy
 import time
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
@@ -79,6 +80,13 @@ class ModelRunner:
                 max_batch_size=self.scheduler_config.max_num_seqs * (1 + self.nextn),
                 input_on_device=envs.DEEPINFER_ZERO_OVERHEAD_SCHEDULING,
             )
+            bucket = getattr(self.cuda_graph_runner, "prefill_bucket", 0)
+            if bucket:
+                assert self.scheduler_config.max_num_batched_tokens <= bucket, (
+                    f"max_num_batched_tokens ({self.scheduler_config.max_num_batched_tokens}) > DEEPINFER_NEURON_PREFILL_BUCKET"
+                )
+                # The chunk's tokens are all stored before any of them attends: the rings must hold a whole chunk.
+                assert envs.DEEPINFER_MAX_BACKTRACKING_TOKENS >= bucket, "set DEEPINFER_MAX_BACKTRACKING_TOKENS >= the prefill bucket"
         self.dflash_use_cuda_graph = (
             self.use_cuda_graph
             and not envs.DEEPINFER_DFLASH_ENFORCE_EAGER
@@ -124,6 +132,7 @@ class ModelRunner:
         seq_metadata_list: list[SequenceMetadata],
         out_dict: dict | None = None,
         metadata_zos: MetadataZOS | None = None,
+        prompt_as_decode: bool = False,
     ) -> InputMetadata | None:
         """Build input metadata on CPU, or on device for ZOS decode.
 
@@ -138,6 +147,7 @@ class ModelRunner:
             window_size=self.window_size,
             swa_block_size=self.swa_block_size,
             metadata_zos=metadata_zos,
+            prompt_as_decode=prompt_as_decode,
         )
         if out_dict is not None:
             return None
@@ -181,16 +191,30 @@ class ModelRunner:
         if total_q is None:
             total_q = sum(seq.seqlen_q for seq in seq_metadata_list)
         is_prompt = seq_metadata_list[0].is_prompt
-        run_cuda_graph = not is_prompt and self.use_cuda_graph and bool(self.cuda_graph_runner.graphs)
+        # Neuron (issue #26): a prompt chunk runs through the decode graphs as a decode-shaped step padded to the
+        # prefill bucket. Every EP rank takes the same bucket (idle ranks run a placeholder prompt), so the static
+        # collectives inside the graphs match.
+        prefill_bucket = getattr(self.cuda_graph_runner, "prefill_bucket", 0) if self.use_cuda_graph else 0
+        prompt_as_decode = is_prompt and prefill_bucket > 0
+        if prompt_as_decode:
+            assert total_q <= prefill_bucket, f"prompt chunk of {total_q} tokens exceeds DEEPINFER_NEURON_PREFILL_BUCKET"
+            # The indexer scores the chunk against one block table (nki_attention.paged_mqa_logits).
+            assert num_reqs == 1, "a Neuron prefill step through the decode graphs takes one sequence"
+        run_cuda_graph = (not is_prompt or prompt_as_decode) and self.use_cuda_graph and bool(self.cuda_graph_runner.graphs)
 
         if run_cuda_graph:
-            graph_size, cuda_graph_used = self.cuda_graph_runner.select(total_q)
+            graph_size, cuda_graph_used = self.cuda_graph_runner.select(prefill_bucket if prompt_as_decode else total_q)
             self._prepare_input_tensors(
                 seq_metadata_list + self.cuda_graph_runner.padding_seq_metadata_list[: graph_size - total_q],
                 out_dict=cuda_graph_used.input_metadata_source.to_dict(),
                 metadata_zos=metadata_zos,
+                prompt_as_decode=prompt_as_decode,
             )
             input_metadata = self._graph_input_metadata(cuda_graph_used, num_reqs)
+            if prompt_as_decode:
+                # The graph runs on its decode-mode metadata; sampling and post-processing see a prompt step.
+                input_metadata = copy.copy(input_metadata)
+                input_metadata.is_prompt = True
         else:
             graph_size, cuda_graph_used = 0, None
             input_metadata = self._prepare_input_tensors(

deepinfer/worker/neuron_graph_runner.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import dataclasses
 import math
+import time
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
@@ -21,6 +22,7 @@ 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 import regions
 from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda
 from deepinfer.sequence import SequenceMetadata
 from deepinfer.worker.cache_engine import GPUCacheInterface, KVCache
@@ -120,6 +122,11 @@ class _LazyKVCache:
 class NeuronGraphRunner(CudaGraphRunner):
     def __init__(self, max_batch_size: int, input_on_device: bool = False) -> None:
         super().__init__(max_batch_size, input_on_device)
+        # Prompt chunks run as decode-shaped steps of this many tokens (issue #26): one more graph size.
+        self.prefill_bucket = envs.DEEPINFER_NEURON_PREFILL_BUCKET
+        if self.prefill_bucket > self.max_batch_size:
+            self.max_batch_size = self.prefill_bucket
+            self.batch_sizes = [*self.batch_sizes, self.prefill_bucket]
         self.compiled: dict[int, Callable[..., CausalLMForwardOutput]] = {}
         self.model: DeepseekForCausalLM | None = None
         self.kv_caches: GPUCacheInterface | None = None
@@ -179,8 +186,23 @@ class NeuronGraphRunner(CudaGraphRunner):
         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)
+        if not envs.DEEPINFER_NEURON_REGION_TIMING:
+            with hide_cuda():
+                return self.compiled[batch_size](input_metadata_device)
+        torch.neuron.synchronize()
+        started = time.perf_counter()
         with hide_cuda():
-            return self.compiled[batch_size](input_metadata_device)
+            out = self.compiled[batch_size](input_metadata_device)
+        torch.neuron.synchronize()
+        total = time.perf_counter() - started
+        timings = regions.take_timings()
+        in_regions = sum(sum(v) for v in timings.values())
+        if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
+            return out
+        parts = ", ".join(f"{k} {len(v)}x {1e3 * sum(v) / len(v):.2f} ms" for k, v in sorted(timings.items()))
+        logger.info(f"Neuron step bs={batch_size}: {1e3 * total:.1f} ms, regions {1e3 * in_regions:.1f} ms ({parts}), "
+                    f"outside regions {1e3 * (total - in_regions):.1f} ms")
+        return out
 
     def destroy(self) -> None:
         super().destroy()

examples/neuron/offline_inference_neuron.py

@@ -37,6 +37,7 @@ parser = argparse.ArgumentParser(description="V4.1 offline inference on Neuron (
 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_prompts", default=len(PROMPTS), type=int, help="run only the first N prompts (eager prefill compiles per prompt length)")
 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)
@@ -86,8 +87,9 @@ def main(dp_size: int, dp_rank: int) -> None:
 
     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)
+    prompts = PROMPTS[: args.num_prompts]
+    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]