From fdbb81f14f4071ae420178319309aa43e3f4af57 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 13:41:17 +0800 Subject: [PATCH] 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 --- deepinfer/model_executor/models/mtp/dflash.py | 22 ++++++++++++++---- deepinfer/neuron_ops/nki_attention.py | 15 ++++++++++-- examples/neuron/bench_neuron.py | 23 +++++++++++-------- 3 files changed, 44 insertions(+), 16 deletions(-) diff --git a/deepinfer/model_executor/models/mtp/dflash.py b/deepinfer/model_executor/models/mtp/dflash.py index c23b455..32465ca 100644 --- a/deepinfer/model_executor/models/mtp/dflash.py +++ b/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, diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py index 56231e6..98221af 100644 --- a/deepinfer/neuron_ops/nki_attention.py +++ b/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). diff --git a/examples/neuron/bench_neuron.py b/examples/neuron/bench_neuron.py index 4de6244..1dff11a 100644 --- a/examples/neuron/bench_neuron.py +++ b/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)")