Neuron DSpark (#15): draft layers and main-cache writes as region graphs, decode rows padded to a power of two

- DFlash forward_drafts per layer and update_main_cache (projection, then each
  layer's cache write) run as region graphs: eagerly the fp8 cast of the
  projection norm does not compile.
- neuronx-cc rejects the sparse decode attention at 6 rows (a DSpark step with
  5 drafts; ISA validation 'same_element_count_t2d'), 4 and 8 rows compile:
  multi-row decode steps pad to a power of two with empty rows.
- bench_neuron.py --nextn: DSpark drafts, reports tokens per decode step.
- Tiny checkpoint on device, cp=2, free-running greedy with drafts: nextn=3,
  5, 7 all 9/9 golden tokens (short), nextn=5 long 9/9.

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

deepinfer/model_executor/models/mtp/dflash.py

@@ -23,6 +23,7 @@ from deepinfer.model_executor.models.v4_modules.utils import convert_rmsnorm, me
 from deepinfer.model_executor.parallel_utils.communication_op import cp_broadcast_multi_tensor
 from deepinfer.model_executor.parallel_utils.parallel_state import PG
 from deepinfer.model_executor.sampling_inputs import MTPDecodeInput, MTPDriverInput, SamplingKernelInputs
+from deepinfer.neuron_ops import regions as neuron_regions
 from deepinfer.sampling_params import SAMPLING_EPS
 from deepinfer.sequence import MTPPredictInfo
 from deepinfer.timer import cpu_timer
@@ -216,17 +217,25 @@ class DFlash(nn.Module):
             main_positions = input_metadata.positions_cp
             assert dflash_input.shape[0] == input_metadata.cp_mirror_len
 
-        main_hiddens = self.projection_norm(self.projection(dflash_input))
-        main_k = self.main_kv_proj(main_hiddens).unflatten(-1, (self.MTP_num_layers, self.qk_head_dim))
+        use_regions = neuron_regions.enabled(input_metadata)
+        # Neuron: the projection and each layer's cache write are region graphs (the eager fp8 cast does not compile).
+        main_k = (neuron_regions.entry(self, "_main_k") if use_regions else self._main_k)(dflash_input)
         for i, layer in enumerate(self.layers):
-            layer.update_main_cache(
-                main_k=main_k[:, i],
+            update = neuron_regions.entry(layer, "update_main_cache") if use_regions else layer.update_main_cache
+            update(
+                main_k=main_k[i],
                 kv_cache=nextn_kv_caches[layer.layer_idx],
                 main_positions=main_positions,
                 slot_mapping_store_swa=input_metadata.slot_mapping_store_swa,
                 mirror_cp_reverse_index=mirror_cp_reverse_index,
             )
 
+    def _main_k(self, dflash_input: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        """Each DFlash layer's main-token keys (contiguous, one tensor a layer)."""
+        main_hiddens = self.projection_norm(self.projection(dflash_input))
+        main_k = self.main_kv_proj(main_hiddens).unflatten(-1, (self.MTP_num_layers, self.qk_head_dim))
+        return tuple(main_k[:, i].contiguous() for i in range(self.MTP_num_layers))
+
     def _embed_draft_init_tokens(self, golden_tokens: torch.Tensor) -> torch.Tensor:
         req_cnt = golden_tokens.shape[0]
         draft_init_tokens = torch.full(
@@ -247,8 +256,11 @@ class DFlash(nn.Module):
         """Run DFlash draft layers on the fixed-stride metadata built by build_dflash_input."""
         draft_hiddens = self._embed_draft_init_tokens(golden_tokens)
         residual = None
+        use_regions = neuron_regions.enabled(input_metadata)
         for layer in self.layers:
-            draft_hiddens, residual = layer.forward_drafts(
+            # Neuron: each draft layer is one region graph, shared by the draft layers (same kind and shapes).
+            forward = neuron_regions.entry(layer, "forward_drafts") if use_regions else layer.forward_drafts
+            draft_hiddens, residual = forward(
                 draft_hiddens=draft_hiddens,
                 kv_cache=nextn_kv_caches[layer.layer_idx],
                 input_metadata=input_metadata,

deepinfer/neuron_ops/nki_attention.py

@@ -262,8 +262,19 @@ def flash_mla_decode(
         # so the kernel takes them as in decode (the torch path was 41.6 ms a layer for a 128-token chunk, bench4).
         out = _chunk_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
     else:
-        out = sparse_decode(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
-    return out.view(b, s_q, h_q, d)
+        q_rows = q.reshape(rows, h_q, d)
+        pad = (1 << (rows - 1).bit_length()) - rows if rows > 1 else 0
+        if pad:
+            # neuronx-cc rejects the record decode at some row counts (6 rows, a DSpark step with 5 drafts: ISA
+            # validation 'same_element_count_t2d'); power-of-two rows compile. Pad rows have no valid slot.
+            q_rows = torch.cat([q_rows, q_rows.new_zeros((pad, h_q, d))])
+            scopes = [
+                (cache, torch.cat([idx, idx.new_zeros((pad, idx.shape[1]))]),
+                 torch.cat([invalid, invalid.new_ones((pad, invalid.shape[1]))]))
+                for cache, idx, invalid in scopes
+            ]
+        out = sparse_decode(q_rows, scopes, softmax_scale, attn_sink)[:rows]
+    return out.reshape(b, s_q, h_q, d)
 
 
 #: Largest KV cache (slots, all scopes) the chunk path decodes whole (DEEPINFER_NEURON_ATTN_TABLE_SLOTS).

examples/neuron/bench_neuron.py

@@ -45,6 +45,7 @@ parser.add_argument("--num_kv_blocks", default=400, type=int)
 parser.add_argument("--max_model_len", default=16384, type=int)
 parser.add_argument("--max_num_batched_tokens", default=128, type=int)
 parser.add_argument("--max_num_seqs", default=1, type=int)
+parser.add_argument("--nextn", default=0, type=int, help="DSpark drafts per step (0 = plain decode)")
 parser.add_argument("--out", default="", help="write the results as JSON here (rank 0)")
 args = parser.parse_args()
 dp_size = args.ep // args.cp
@@ -70,17 +71,18 @@ def _random_prompt(length: int, seed: int, vocab: int) -> list[int]:
 
 
 def _run_one(llm, prompt: list[int] | None, params) -> tuple[float, float, int] | None:  # noqa: ANN001
-    """Run the engine until it drains. With a prompt, returns (ttft_s, latency_s, output_tokens) for it."""
+    """Run the engine until it drains. With a prompt, returns (ttft_s, latency_s, output_tokens, steps_with_output) for it."""
     engine = llm.llm_engine
     start = time.perf_counter()
     if prompt is not None:
         llm._add_request(None, params, prompt)
     first = last = None
-    tokens = 0
+    tokens = steps = 0
     while True:
         unfinished = engine.has_unfinished_requests()
         outputs = engine.step() if unfinished else engine.flush_inflight_reqs()
         now = time.perf_counter()
+        steps += bool(outputs)
         for out in outputs:
             tokens += 1
             first = first if first is not None else now
@@ -89,11 +91,11 @@ def _run_one(llm, prompt: list[int] | None, params) -> tuple[float, float, int]
             break
     if prompt is None:
         return None
-    return first - start, last - start, tokens
+    return first - start, last - start, tokens, steps
 
 
 def main(dp_size: int, dp_rank: int) -> None:
-    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
+    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = str(args.nextn)
     os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
     setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=False, enable_eplb=False)
     from deep_cluster.hardware import get_actual_num_gpus
@@ -111,7 +113,7 @@ def main(dp_size: int, dp_rank: int) -> None:
         max_model_len=args.max_model_len,
         max_num_seqs=args.max_num_seqs,
         gpu_memory_utilization=0.85,
-        lookahead_tokens=0,
+        lookahead_tokens=1 if args.nextn else 0,
     )
     from deepinfer.sampling_params import SamplingParams
 
@@ -130,16 +132,19 @@ def main(dp_size: int, dp_rank: int) -> None:
         for r in range(args.repeats):
             got = _run_one(llm, _random_prompt(length, 1000 * length + r + 7919 * dp_rank, vocab) if busy else None, params)
             if got is not None:
-                ttft, latency, n = got
+                ttft, latency, n, steps = got
                 tpot = (latency - ttft) / max(n - 1, 1)
-                cell.append((ttft, tpot, n))
+                accept = (n - 1) / max(steps - 1, 1)  # tokens per decode step (the acceptance length with DSpark)
+                cell.append((ttft, tpot, n, accept))
                 if dp_rank == 0:
-                    print(f"BENCH input={length} rep={r} ttft_ms={ttft * 1e3:.1f} tpot_ms={tpot * 1e3:.1f} tokens={n}", flush=True)
+                    print(f"BENCH input={length} rep={r} ttft_ms={ttft * 1e3:.1f} tpot_ms={tpot * 1e3:.1f} tokens={n} "
+                          f"tokens_per_step={accept:.2f}", flush=True)
         if dp_rank == 0 and cell:
             ttft = statistics.median(c[0] for c in cell)
             tpot = statistics.median(c[1] for c in cell)
             results.append({"input": length, "output": args.output_len, "ttft_p50_ms": ttft * 1e3, "tpot_p50_ms": tpot * 1e3,
-                            "decode_tok_s_per_request": 1.0 / tpot})
+                            "decode_tok_s_per_request": 1.0 / tpot,
+                            "tokens_per_step": statistics.median(c[3] for c in cell)})
     if dp_rank == 0:
         print("BENCH_RESULTS " + json.dumps(results), flush=True)
         print(f"| input | output | TTFT p50 (ms) | TPOT p50 (ms) | decode tok/s per request |   (batch 1 per DP rank)")