Neuron DSpark (#15): greedy verification and MTP ops, CP followers run the drafts

- fused_sample reject mode for greedy rows (accept iff the draft is the
  argmax), get_accept_result / finalize_verification (static final rows),
  store_mtp_info (static pool writes, slot 0 kept), confidence_head_fused,
  dflash_input_store, get_fused_dflash_get_input: the frontend refs with the
  kernels' argument lists.
- CP followers run the draft forward (its MoE has collectives); the driver's
  accepted offsets reach them through one all-reduce on Neuron (no eager
  broadcast), the request first rows through the graph metadata.
- run_deepinfer_tiny.py --nextn: free-running greedy with drafts. Host tiny,
  nextn=5, eager: short 9/9 golden tokens; long matches up to its known near
  tie at step 5.
- bench17: the two-level MoE all-reduce does not move TPOT (433 ms).

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

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

@@ -1,18 +1,23 @@
-"""Mirrors ``tilelang_ops/asc/dflash.py`` (confidence_head_fused_tl, dflash_input_store_tl).
-
-DFlash (MTP drafting) only: the draft module and the decoder-layer DFlash input store exist only with
-DEEPINFER_USE_NEXTN_DECODE>0. Out of scope on Neuron basic mode (NEURON_PORT.md "Scope"): the factories
-keep the asc signatures and raise NotImplementedError.
-"""
+"""Mirrors ``tilelang_ops/asc/dflash.py`` (confidence_head_fused_tl, dflash_input_store_tl) for DSpark drafting:
+the frontend pure-torch references with the kernels' argument lists."""
 
 from typing import Any
 
-from ._unsupported import unsupported
+import torch
 
 
 def confidence_head_fused_tl(n_embed: int, markov_rank: int, has_markov_emb: bool) -> Any:
-    unsupported("confidence_head_fused_tl")
+    from ..dflash import confidence_head_fused_ref
+
+    def _confidence(hidden, markov_emb, w, out, temperature):  # noqa: ANN001
+        confidence_head_fused_ref(hidden, w.view(-1, 1), temperature, markov_emb if has_markov_emb else None, out=out)
+
+    return _confidence
 
 
 def dflash_input_store_tl(n_embed: int, hc_mult: int) -> Any:
-    unsupported("dflash_input_store_tl")
+    def _store(hidden_states, out):  # noqa: ANN001
+        h = hidden_states.float().mean(dim=-2) if hc_mult > 1 else hidden_states.reshape(-1, n_embed)
+        out.copy_(h.to(out.dtype).reshape(out.shape))
+
+    return _store

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

@@ -166,6 +166,25 @@ def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_si
     return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size)
 
 
+def _greedy_verify(
+    temperatures, tokens, target_logprobs, is_verify, draft_logprobs, read_rows, draft_ids, accept, draft_selected, accept_probs
+):  # noqa: ANN001
+    """Speculative verification for greedy rows (DSpark on Neuron, issue #15): a one-hot target accepts a draft iff it
+    is the argmax, and the residual's argmax is the target argmax, so ``tokens`` is already the reject/fresh token.
+    Sampled (non-greedy) verification is not ported (its int64 RNG reject draw)."""
+    from deepinfer.sampling_params import SAMPLING_EPS
+
+    if not bool((temperatures.cpu() <= SAMPLING_EPS).all()):
+        unsupported("fused_sample_tl(do_reject_sampling=True) for non-greedy rows")
+    rows = read_rows.long()
+    draft = draft_ids.index_select(0, rows)
+    ok = is_verify & (draft >= 0) & (draft == tokens)
+    accept.copy_(ok)
+    lp = draft_logprobs.index_select(0, rows).gather(1, draft.clamp(min=0).long().unsqueeze(1)).squeeze(1)
+    draft_selected.copy_(torch.where(draft >= 0, lp, torch.full_like(lp, float("-inf"))))
+    accept_probs.copy_(torch.where(is_verify, ok.float(), torch.ones_like(lp)))
+
+
 def fused_sample_tl(
     vocab_size: int,
     do_sample: bool,
@@ -173,8 +192,6 @@ def fused_sample_tl(
 ) -> Any:
     """Factory mirroring asc ``fused_sample_tl``; the kernel takes the asc argument list."""
     assert not do_reject_sampling or do_sample, "reject sampling requires sampling"
-    if do_reject_sampling:
-        unsupported("fused_sample_tl(do_reject_sampling=True)")
 
     def _fused_sample(
         logits: torch.Tensor,
@@ -204,5 +221,10 @@ def fused_sample_tl(
         out_logprobs, tokens = _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size)
         logprobs.copy_(out_logprobs)
         candidate_token_ids.copy_(tokens)
+        if do_reject_sampling:
+            _greedy_verify(
+                temperatures, tokens, out_logprobs, is_verify, last_nextn_token_logprobs, mtp_info_read_rows,
+                last_next_n_token_ids, accept_result_verify, draft_selected_logprobs, accept_probs,
+            )
 
     return _fused_sample

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

@@ -53,5 +53,23 @@ def get_fused_dflash_get_input_tl(
     has_valid_token_mask: bool,
     has_gate_force_random: bool,
 ) -> Any:
-    # DSpark/DFlash drafting only (DEEPINFER_USE_NEXTN_DECODE=0 on Neuron).
-    unsupported("get_fused_dflash_get_input_tl")
+    """The frontend ``get_fused_dflash_get_input_ref`` with the kernel's argument list (DSpark drafting)."""
+    from ..get_kv_index import get_fused_dflash_get_input_ref
+
+    def _fill(slot_mapping_swa, positions, ftp, last_valid_offset, newx_out, positions_out, input_ids_out,
+              valid_out, force_out, kv_index_out, kv_len, window_size):  # noqa: ANN001
+        block = newx_out.shape[1]
+        newx, pos, ids, valid, force, kv_index = get_fused_dflash_get_input_ref(
+            slot_mapping_swa, positions, ftp, last_valid_offset, swa_block_size, block, window_size,
+            has_valid_token_mask, has_gate_force_random,
+        )
+        newx_out.copy_(newx.view_as(newx_out))
+        positions_out.copy_(pos.view_as(positions_out))
+        input_ids_out.copy_(ids.to(input_ids_out.dtype).view_as(input_ids_out))
+        if has_valid_token_mask:
+            valid_out.copy_(valid.view_as(valid_out))
+        if has_gate_force_random:
+            force_out.copy_(force.view_as(force_out))
+        kv_index_out.copy_(kv_index)
+
+    return _fill

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

@@ -110,7 +110,19 @@ def remap_draft_topk_tl(force_topk: int, vocab_size: int) -> Any:
 
 
 def get_accept_result_tl(has_golden: bool) -> Any:
-    unsupported("get_accept_result_tl")
+    """The frontend ``get_accept_result_ref`` (static shapes: a cumsum prefix), with the kernel's argument list."""
+    from ..sampler import get_accept_result_ref
+
+    def _accept(raw_accept, golden, draft_pool, read_rows, candidates, is_verify, ftp, row_to_req, token_ids, accept, is_final):  # noqa: ANN001
+        draft = draft_pool.index_select(0, read_rows.long())
+        tokens, eff, final = get_accept_result_ref(
+            raw_accept, golden if has_golden else None, draft, candidates, is_verify, ftp, row_to_req
+        )
+        token_ids.copy_(tokens)
+        accept.copy_(eff)
+        is_final.copy_(final)
+
+    return _accept
 
 
 def finalize_verification_tl(
@@ -119,4 +131,53 @@ def finalize_verification_tl(
     num_steps: int,
     num_stop_tokens: int,
 ) -> Any:
-    unsupported("finalize_verification_tl")
+    """The frontend ``finalize_verification_ref`` with the kernel's argument list, its boolean-mask final-row lookup
+    replaced by a static one (exactly one final row per request: ``ftp + last_valid_offset``)."""
+    from ..sampler import compute_confidence_error_ref
+
+    def _finalize(
+        effective_accept, is_final, effective_token_ids, token_logprobs, ssp, top_lp, top_ids, ftp, seqlen_q,
+        row_to_req, is_verify, accept_probs, confidence_pool, read_rows, stop_token_ids,
+        last_valid_offset, num_generated_tokens, num_verified_tokens, batch_confidence_error, emitted_stop_token,
+        sampled_token_ids, sampled_token_logprobs, sampled_ssp, sampled_top_lp, sampled_top_ids,
+        verified_token_ids, verified_token_logprobs, verified_ssp, verified_top_lp, verified_top_ids,
+    ):  # noqa: ANN001
+        q = effective_accept.shape[0]
+        req = row_to_req.long()
+        rows = torch.arange(q, dtype=torch.int64, device=ftp.device)
+        lvo = torch.zeros(ftp.shape[0], dtype=torch.int64, device=ftp.device)
+        lvo.scatter_add_(0, req, effective_accept.to(torch.int64))
+        last_valid_offset.copy_(lvo.to(last_valid_offset.dtype))
+        final_index = ftp.long() + lvo
+        local_step = rows - ftp.long()[req]
+        step_index = local_step.clamp(max=num_steps - 1)
+        gen = torch.zeros(num_steps, dtype=torch.int32, device=ftp.device)
+        ver = torch.zeros(num_steps, dtype=torch.int32, device=ftp.device)
+        gen.scatter_add_(0, step_index, effective_accept.to(torch.int32))
+        ver.scatter_add_(0, step_index, is_verify.to(torch.int32))
+        num_generated_tokens.copy_(gen)
+        num_verified_tokens.copy_(ver)
+        if has_confidence:
+            rates = confidence_pool.index_select(0, read_rows.long())
+            batch_confidence_error.copy_(compute_confidence_error_ref(accept_probs, rates, is_verify).reshape(1))
+            row_stops = stop_token_ids.index_select(0, req)
+            hits = ((effective_token_ids.unsqueeze(1) == row_stops).any(dim=1)) & (effective_accept | is_final)
+            counts = torch.zeros(ftp.shape[0], dtype=torch.int32, device=ftp.device)
+            counts.scatter_add_(0, req, hits.to(torch.int32))
+            emitted_stop_token.copy_(counts > 0)
+        else:
+            batch_confidence_error.zero_()
+        sampled_token_ids.copy_(effective_token_ids.index_select(0, final_index))
+        sampled_token_logprobs.copy_(token_logprobs.index_select(0, final_index))
+        if has_sum_smaller_probs:
+            sampled_ssp.copy_(ssp.index_select(0, final_index))
+            verified_ssp.copy_(torch.where(effective_accept, ssp, 0.0))
+        if top_lp.shape[-1]:
+            sampled_top_lp.copy_(top_lp.index_select(0, final_index))
+            sampled_top_ids.copy_(top_ids.index_select(0, final_index))
+            verified_top_lp.copy_(torch.where(effective_accept.unsqueeze(-1), top_lp, 0.0))
+            verified_top_ids.copy_(torch.where(effective_accept.unsqueeze(-1), top_ids, 0).to(torch.int32))
+        verified_token_ids.copy_(torch.where(effective_accept, effective_token_ids, -1).to(torch.int32))
+        verified_token_logprobs.copy_(torch.where(effective_accept, token_logprobs, 0.0))
+
+    return _finalize

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

@@ -1,13 +1,32 @@
-"""Mirrors ``tilelang_ops/asc/store_mtp_info.py`` (store_mtp_info_tl).
+"""Mirrors ``tilelang_ops/asc/store_mtp_info.py`` (store_mtp_info_tl): MTP info pool writes for DSpark drafting.
 
-MTP info pool writes; MTPInfoPool is only built by the DFlash draft module (MTP drafting). Out of scope on
-Neuron basic mode (NEURON_PORT.md "Scope"): the factory keeps the asc signature and raises.
+The frontend ``store_mtp_info_ref`` with static shapes: every request writes its block, and a slot-0 request writes
+back the pool's own values (slot 0 holds the padding drafts and must not change).
 """
 
 from typing import Any
 
-from ._unsupported import unsupported
+import torch
 
 
 def store_mtp_info_tl(nextn: int, vocab_size: int, has_index: bool, has_confidence: bool) -> Any:
-    unsupported("store_mtp_info_tl")
+    def _store(write_slots, token_ids, logprobs, logprobs_index, confidence, token_ids_pool, logprobs_pool,
+               logprobs_index_pool, confidence_pool):  # noqa: ANN001
+        req = write_slots.shape[0]
+        rows = (write_slots.long().unsqueeze(1) * nextn + torch.arange(nextn, device=write_slots.device)).reshape(-1)
+        keep = (write_slots == 0).unsqueeze(1).expand(req, nextn).reshape(-1)
+
+        def put(pool, values):  # noqa: ANN001, ANN202
+            values = values.reshape(req * nextn, *pool.shape[1:])
+            old = pool.index_select(0, rows)
+            mask = keep.view(-1, *([1] * (pool.dim() - 1)))
+            pool.index_copy_(0, rows, torch.where(mask, old, values))
+
+        put(token_ids_pool, token_ids)
+        put(logprobs_pool, logprobs)
+        if has_index:
+            put(logprobs_index_pool, logprobs_index)
+        if has_confidence:
+            put(confidence_pool, confidence)
+
+    return _store

deepinfer/model_executor/models/mtp/dflash.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING
 import torch
 from torch import nn
 
-from deepinfer import AttrDict, envs
+from deepinfer import PLATFORM, AttrDict, envs
 from deepinfer.engine.engine_type import EngineType
 from deepinfer.logger import init_logger
 from deepinfer.model_executor.input_metadata import InputMetadataV4
@@ -362,7 +362,18 @@ class DFlash(nn.Module):
             assert PG.context_parallel_group().size() > 1, "DFlash MTP follower ranks require context parallelism"
             last_valid_offset = torch.empty_like(decode_input.first_token_positions)
             golden_tokens = torch.empty_like(decode_input.first_token_positions)
-        cp_broadcast_multi_tensor([last_valid_offset, golden_tokens])
+        if PLATFORM.is_neuron and PG.context_parallel_group().size() > 1:
+            # Neuron: eager dist.broadcast is not available over the CP group; an all-reduce of the driver's values
+            # (zeros from the followers) is the same broadcast, in one collective.
+            import torch.distributed._functional_collectives as funcol
+
+            packed = torch.stack([last_valid_offset, golden_tokens]).to(torch.int32)
+            if driver_input is None:
+                packed = torch.zeros_like(packed)
+            packed = funcol.wait_tensor(funcol.all_reduce(packed, "sum", PG.context_parallel_group()))
+            last_valid_offset, golden_tokens = packed[0].to(last_valid_offset.dtype), packed[1].to(golden_tokens.dtype)
+        else:
+            cp_broadcast_multi_tensor([last_valid_offset, golden_tokens])
 
         dflash_input_metadata = main_input_metadata.build_dflash_input(
             decode_input.first_token_positions,

deepinfer/worker/model_runner.py

@@ -253,6 +253,8 @@ class ModelRunner:
                     # No in-graph broadcast on Neuron (the GPU graphs broadcast their device metadata themselves):
                     # send the graph's host metadata through shm; followers copy it into their own source buffers.
                     input_metadata_dict["neuron_graph"] = {
+                        # DSpark: each request's first packed row, for the followers' draft forward.
+                        "mtp_ftp": [sum(seq.seqlen_q for seq in seq_metadata_list[:i]) for i in range(num_reqs)],
                         "graph_size": graph_size,
                         "prompt_as_decode": prompt_as_decode,
                         "ctx_limit": getattr(input_metadata, "neuron_ctx_limit", None),
@@ -284,6 +286,7 @@ class ModelRunner:
                 (target[sub] if sub else target).copy_(value)
             input_metadata = self._graph_input_metadata(cuda_graph_used, num_reqs)
             input_metadata.neuron_ctx_limit = neuron_graph["ctx_limit"]
+            input_metadata.neuron_mtp_ftp = neuron_graph["mtp_ftp"]
             if neuron_graph["prompt_as_decode"]:
                 input_metadata = copy.copy(input_metadata)
                 input_metadata.is_prompt = True
@@ -418,6 +421,18 @@ class ModelRunner:
             # load-statistics all-gather spans the EP group.
             from deepinfer.model_executor import eplb
 
+            mtp = self.model.MTP_module
+            if self.nextn and mtp is not None and not input_metadata.is_prompt:
+                # DSpark: the draft layers have collectives (MoE), so followers run them too, as the GPU sampler's
+                # follower branch does; the driver's accepted offsets reach them inside run_drafts.
+                from deepinfer.model_executor.sampling_inputs import MTPDecodeInput
+
+                ftp = torch.tensor(input_metadata.neuron_mtp_ftp, dtype=torch.int32).cuda()
+                mtp.decode_forward(
+                    main_input_metadata=input_metadata,
+                    nextn_kv_caches=kv_caches,
+                    decode_input=MTPDecodeInput(first_token_positions=ftp, driver=None),
+                )
             if eplb.current() is not None:
                 eplb.current().step(input_metadata.is_prompt)
             return SamplerOutput(post_funcs=[])

examples/neuron/run_deepinfer_tiny.py

@@ -46,6 +46,9 @@ def parse_args():
     ap.add_argument("--hidden", action="store_true", help="also record per-layer block outputs (hc residual stream)")
     ap.add_argument("--compare", action="store_true", help="run compare_golden on every prompt afterwards")
     ap.add_argument("--attn-p-bf16", action="store_true", help="experiment: patch the flash_mla shim with the proposed bf16-P fix")
+    ap.add_argument("--nextn", type=int, default=0,
+                    help="DSpark speculative decoding with this many drafts (free-running greedy, no teacher forcing): "
+                         "the engine's tokens are checked against the golden greedy tokens")
     ap.add_argument("--cp", type=int, default=1,
                     help="context-parallel group of one engine (decode replicated on every member, EP over the group): "
                          "CUDA_VISIBLE_DEVICES lists its cores (default 0..cp-1)")
@@ -57,11 +60,11 @@ def parse_args():
     return ap.parse_args()
 
 
-def setup_env(device: str, cp: int = 1) -> None:
+def setup_env(device: str, cp: int = 1, nextn: int = 0) -> None:
     # Must happen before deepinfer is imported (see examples/_config.py).
     os.environ["DEEPINFER_ACCELERATOR"] = "neuron"
     os.environ["DEEPINFER_NEURON_DEVICE"] = device
-    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
+    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = str(nextn)
     os.environ.setdefault("DEEPINFER_ENFORCE_EAGER", "1")
     os.environ.setdefault("DEEPINFER_ZERO_OVERHEAD_SCHEDULING", "0")
     os.environ.setdefault("CUDA_VISIBLE_DEVICES", ",".join(str(i) for i in range(cp)))  # mapped to the target
@@ -251,7 +254,7 @@ def _ep_rank_main(args, rank: int, core: str) -> None:
 
 
 def run(args, rank: int = 0):
-    setup_env(args.device, args.cp)
+    setup_env(args.device, args.cp, args.nextn)
     os.makedirs(args.out, exist_ok=True)
 
     import torch
@@ -285,6 +288,7 @@ def run(args, rank: int = 0):
         max_num_batched_tokens=args.max_num_batched_tokens,
         max_num_seqs=1,
         gpu_memory_utilization=0.5,
+        **({"lookahead_tokens": 1} if args.nextn else {}),
     )
     print(f"[run] engine up in {time.time() - t0:.1f}s (ep={args.ep}, rank={rank})", flush=True)
     status = install_hidden_hooks(rec, llm) if args.hidden else "hidden not requested"
@@ -295,11 +299,21 @@ def run(args, rank: int = 0):
         ids, gtok = list(gold["prompt_ids"]), list(gold["tokens"])
         steps = len(gtok)
         rec.reset()
-        sp = SamplingParams(temperature=0.0, max_tokens=steps, ignore_eos=True, golden_token_ids=gtok)
+        if args.nextn:
+            sp = SamplingParams(temperature=0.0, max_tokens=steps, ignore_eos=True)
+        else:
+            sp = SamplingParams(temperature=0.0, max_tokens=steps, ignore_eos=True, golden_token_ids=gtok)
         t1 = time.time()
         outputs = llm.generate(prompt_token_ids=[ids], sampling_params=sp, use_tqdm=False)
         dt = time.time() - t1
         out_tokens = [o.token_id for o in outputs[0]]
+        if args.nextn:
+            match = sum(a == b for a, b in zip(out_tokens, gtok))
+            steps_run = len(rec.logits)
+            print(f"[run] {name}: nextn={args.nextn} free-running greedy {out_tokens} vs golden {gtok}: {match}/{steps} "
+                  f"match, {steps_run} sampler calls ({steps / max(steps_run - 1, 1):.2f} tokens/decode step), {dt:.1f}s",
+                  flush=True)
+            continue
         rows = [r for r in rec.logits]
         assert all(r.size(0) == 1 for r in rows), [tuple(r.shape) for r in rows]
         logits = torch.cat(rows)[-steps:] if len(rows) >= steps else torch.cat(rows)