From 08526c7fca6e0e5c8aea1e9d7071b7f57bb3af5e Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Mon, 1 Jun 2026 22:27:25 -0700 Subject: [PATCH] [Spec] `FrozenKVMTP` fold assistant seed into captured draft graph (#25539) Co-authored-by: Cursor --- .../frozen_kv_mtp_cuda_graph_runner.py | 18 ++- .../srt/speculative/frozen_kv_mtp_worker.py | 126 ++++++++++-------- .../runner_modes/speculative_draft_runner.py | 17 ++- 3 files changed, 97 insertions(+), 64 deletions(-) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index fd331dfde..a35084dbc 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -47,6 +47,8 @@ class FrozenKVMTPInputBuffers(ForwardInputBuffers): topk_p: torch.Tensor topk_index: torch.Tensor hidden_states: torch.Tensor + # Consumed by the captured seed iter; see `FrozenKVMTPWorker.draft_forward`. + bonus_tokens: torch.Tensor global_num_tokens_gpu: Optional[torch.Tensor] global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] @@ -107,6 +109,7 @@ class FrozenKVMTPCudaGraphRunner: (self.max_bs, frozen_kv_mtp_worker._recurrent_hidden_size), dtype=self.model_runner.dtype, ) + bonus_tokens = torch.zeros((self.max_bs,), dtype=torch.int64) if self.require_gathered_buffer: if self.require_mlp_tp_gather: @@ -135,6 +138,7 @@ class FrozenKVMTPCudaGraphRunner: topk_p=topk_p, topk_index=topk_index, hidden_states=hidden_states, + bonus_tokens=bonus_tokens, global_num_tokens_gpu=global_num_tokens_gpu, global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, ) @@ -208,6 +212,7 @@ class FrozenKVMTPCudaGraphRunner: topk_p = buffers.topk_p[:request_bs] topk_index = buffers.topk_index[:request_bs] hidden_states = buffers.hidden_states[:request_bs] + bonus_tokens = buffers.bonus_tokens[:request_bs] if self.require_mlp_tp_gather: buffers.global_num_tokens_gpu.copy_( @@ -254,6 +259,7 @@ class FrozenKVMTPCudaGraphRunner: topk_p=topk_p, topk_index=topk_index, hidden_states=hidden_states, + bonus_tokens=bonus_tokens, capture_hidden_mode=CaptureHiddenMode.LAST, ) spec_info.num_tokens_per_req = self.topk @@ -363,8 +369,8 @@ class FrozenKVMTPCudaGraphRunner: buffers.mrope_positions[:, :raw_num_token].copy_( forward_batch.mrope_positions ) - buffers.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p) - buffers.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index) + # `topk_p`/`topk_index` are produced by the captured seed iter. + buffers.bonus_tokens[:raw_bs].copy_(forward_batch.spec_info.bonus_tokens) buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states) buffers.req_pool_indices[:raw_expanded_bs].copy_(forward_batch.req_pool_indices) @@ -395,7 +401,13 @@ class FrozenKVMTPCudaGraphRunner: self.raw_bs = raw_bs self.bs = bs - self._replay() + # NVTX span: the graph bypasses `model_runner.forward`'s record_function. + span_name = f"step[DRAFT_LOOP raw_bs={raw_bs} bs={bs} topk={self.topk}]" + if torch.autograd._profiler_enabled(): + with torch.profiler.record_function(span_name): + self._replay() + else: + self._replay() out = self.output_buffers[bs] if bs != raw_bs: diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index ea8bcec72..267d9f797 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py @@ -360,7 +360,10 @@ class FrozenKVMTPWorker(TpModelWorker): mm_input_embeds: Optional[torch.Tensor] = None, draft_input: Optional[FrozenKVMTPDraftInput] = None, ) -> None: - """Run the one-token assistant seed step against frozen target KV.""" + """Stash seed inputs on ``batch.spec_info``; the forward runs inside + the captured draft graph (see ``draft_forward``'s seed iter).""" + del seq_lens_cpu, mm_input_embeds, draft_input + if batch.forward_mode.is_idle() or last_token_ids.numel() == 0: batch.spec_info = FrozenKVMTPDraftInput.create_idle_input( device=batch.device, @@ -371,55 +374,23 @@ class FrozenKVMTPWorker(TpModelWorker): ) return - if draft_input is None: - draft_input = FrozenKVMTPDraftInput() - - draft_input.bonus_tokens = last_token_ids.to(torch.int64) - draft_input.hidden_states = last_hidden_states - draft_input.capture_hidden_mode = CaptureHiddenMode.LAST - draft_input.num_tokens_per_req = 1 - draft_input.num_tokens_for_logprob_per_req = 1 - draft_input.positions = self._position_for_batch(batch) - - forward_mode_backup = batch.forward_mode - input_ids_backup = batch.input_ids - return_hidden_states_backup = batch.return_hidden_states - return_logprob_backup = batch.return_logprob - spec_info_backup = batch.spec_info - - batch.forward_mode = ForwardMode.DECODE - batch.input_ids = draft_input.bonus_tokens - batch.return_hidden_states = False - batch.return_logprob = False - batch.spec_info = draft_input - - try: - batch.seq_lens_cpu_cache = seq_lens_cpu - forward_batch = ForwardBatch.init_new(batch, self.draft_model_runner) - forward_batch.return_logprob = False - if mm_input_embeds is not None: - forward_batch.mm_input_embeds = mm_input_embeds - self._set_positions(forward_batch) - self._init_frozen_kv_metadata(forward_batch) - with ( - self._target_kv_pool_view(forward_batch), - forward_context(ForwardContext(attn_backend=self.draft_attn_backend)), - ): - logits_output = self.draft_model_runner.forward( - forward_batch, skip_attn_backend_init=True - ).logits_output - maybe_detect_nan(logits_output.next_token_logits, "frozen_kv_mtp_seed") - maybe_detect_inf(logits_output.next_token_logits, "frozen_kv_mtp_seed") - self._capture_for_decode(logits_output, draft_input) - finally: - batch.forward_mode = forward_mode_backup - batch.input_ids = input_ids_backup - batch.return_hidden_states = return_hidden_states_backup - batch.return_logprob = return_logprob_backup - # Keep the seeded draft state; only restore the old object on error paths - # before the assignment above could have happened. - if batch.spec_info is not draft_input: - batch.spec_info = spec_info_backup + stashed = FrozenKVMTPDraftInput() + stashed.bonus_tokens = last_token_ids.to(torch.int64) + stashed.hidden_states = last_hidden_states + # Real-shaped zeros so inherited `filter_batch`/`merge_batch` can slice + # them between iters; overwritten by the captured seed iter. + bs = last_token_ids.shape[0] + device = last_token_ids.device + stashed.topk_p = torch.zeros( + (bs, self.topk), device=device, dtype=torch.float32 + ) + stashed.topk_index = torch.zeros( + (bs, self.topk), device=device, dtype=torch.int64 + ) + stashed.capture_hidden_mode = CaptureHiddenMode.LAST + stashed.num_tokens_per_req = 1 + stashed.num_tokens_for_logprob_per_req = 1 + batch.spec_info = stashed def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult: if batch.forward_mode.is_extend() or batch.is_extend_in_batch: @@ -657,20 +628,61 @@ class FrozenKVMTPWorker(TpModelWorker): ): spec_info = forward_batch.spec_info assert isinstance(spec_info, FrozenKVMTPDraftInput) - topk_p, topk_index, hidden_states = ( - spec_info.topk_p, - spec_info.topk_index, - spec_info.hidden_states, - ) - maybe_detect_nan(topk_p, "frozen_kv_mtp_draft: initial topk_p") score_list: List[torch.Tensor] = [] token_list: List[torch.Tensor] = [] parents_list: List[torch.Tensor] = [] - if not skip_attn_backend_init and self.speculative_num_steps > 1: + # Seed + recurrent iters share the same `seq_lens - 1` rope position, + # so one init covers the loop. Must run even at num_steps == 1. + if not skip_attn_backend_init: self._init_frozen_kv_metadata(forward_batch) + # Seed iter: assistant forward on (bonus_token, target_h) to produce + # iter-0 `(topk_p, topk_index, hidden_states)`. For topk>1, replicate + # to `bs*topk` to match kernel shapes, then slice back per-req. + bonus_tokens = spec_info.bonus_tokens + target_hidden = spec_info.hidden_states + if self.topk > 1: + seed_input_ids = bonus_tokens.repeat_interleave(self.topk, dim=0) + seed_prev_hidden = target_hidden.repeat_interleave(self.topk, dim=0) + else: + seed_input_ids = bonus_tokens + seed_prev_hidden = target_hidden + + forward_batch.input_ids = seed_input_ids + forward_batch.spec_info.hidden_states = seed_prev_hidden + self._set_positions(forward_batch) + + with ( + self._target_kv_pool_view(forward_batch), + forward_context(ForwardContext(attn_backend=self.draft_attn_backend)), + ): + seed_output = self.draft_model_runner.forward( + forward_batch, skip_attn_backend_init=True + ).logits_output + + maybe_detect_nan( + seed_output.next_token_logits, "frozen_kv_mtp_draft: seed iter" + ) + + if self.topk > 1: + seed_next_logits = seed_output.next_token_logits[:: self.topk] + seed_hidden_per_req = seed_output.hidden_states[:: self.topk] + else: + seed_next_logits = seed_output.next_token_logits + seed_hidden_per_req = seed_output.hidden_states + + probs = torch.softmax(seed_next_logits, dim=-1) + topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + maybe_detect_oob( + topk_index, + 0, + seed_next_logits.shape[-1], + "frozen_kv_mtp_draft: seed topk_index OOB", + ) + hidden_states = seed_hidden_per_req + scores = None for i in range(self.speculative_num_steps): input_ids, hidden_states, scores, tree_info = select_top_k_tokens( diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index 11a52de30..c244afef9 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -713,10 +713,20 @@ def _make_dense_frozen_kv_mtp_draft_inputs( settings: EagleDraftRunnerSettings, ) -> dict[str, torch.Tensor]: draft_inputs = _make_dense_draft_inputs(case, settings) + # `draft_forward` now runs the assistant seed iter in-graph: it consumes the + # per-req bonus token + target hidden and derives iter-0 topk_p/topk_index + # itself, so the fixture supplies `bonus_tokens` rather than topk_p/index. + with _seeded_rng(4090 + len(case.name) + settings.topk, device=settings.device): + bonus_tokens = torch.randint( + 0, + settings.vocab_size, + (case.batch_size,), + dtype=torch.int64, + device=settings.device, + ) return { "hidden_states": draft_inputs["hidden_states"], - "topk_p": draft_inputs["topk_p"], - "topk_index": draft_inputs["topk_index"], + "bonus_tokens": bonus_tokens, } @@ -849,9 +859,8 @@ def _make_dense_frozen_kv_mtp_forward_batch( settings: EagleDraftRunnerSettings, ) -> ForwardBatch: spec_info = FrozenKVMTPDraftInput( - topk_p=draft_inputs["topk_p"].clone(), - topk_index=draft_inputs["topk_index"].clone(), hidden_states=draft_inputs["hidden_states"].clone(), + bonus_tokens=draft_inputs["bonus_tokens"].clone(), capture_hidden_mode=CaptureHiddenMode.LAST, num_tokens_per_req=settings.topk, num_tokens_for_logprob_per_req=settings.topk,