From b2d94e16b26f0fcfd433d2299a6de445261ec576 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sat, 24 Oct 2026 11:31:59 +0800 Subject: [PATCH] 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 --- .../layers/tilelang_ops/neuron/dflash.py | 23 ++++--- .../tilelang_ops/neuron/fused_sample.py | 26 +++++++- .../tilelang_ops/neuron/get_kv_index.py | 22 ++++++- .../layers/tilelang_ops/neuron/sampler.py | 65 ++++++++++++++++++- .../tilelang_ops/neuron/store_mtp_info.py | 29 +++++++-- deepinfer/model_executor/models/mtp/dflash.py | 15 ++++- deepinfer/worker/model_runner.py | 15 +++++ examples/neuron/run_deepinfer_tiny.py | 22 +++++-- 8 files changed, 191 insertions(+), 26 deletions(-) diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash.py index b413094..b3da998 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash.py +++ b/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 diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py index 741c401..8045ad0 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py +++ b/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 diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/get_kv_index.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/get_kv_index.py index bd2671e..0f11c2e 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/get_kv_index.py +++ b/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 diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/sampler.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/sampler.py index 6738aae..03d0bfa 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/sampler.py +++ b/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 diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/store_mtp_info.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/store_mtp_info.py index 142b62f..6611671 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/store_mtp_info.py +++ b/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 diff --git a/deepinfer/model_executor/models/mtp/dflash.py b/deepinfer/model_executor/models/mtp/dflash.py index c8f5902..c23b455 100644 --- a/deepinfer/model_executor/models/mtp/dflash.py +++ b/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, diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index 9ac9676..ced2df3 100644 --- a/deepinfer/worker/model_runner.py +++ b/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=[]) diff --git a/examples/neuron/run_deepinfer_tiny.py b/examples/neuron/run_deepinfer_tiny.py index 7cb0228..ff432f1 100644 --- a/examples/neuron/run_deepinfer_tiny.py +++ b/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)