[Spec] FrozenKVMTP fold assistant seed into captured draft graph (#25539)

Co-authored-by: Cursor <cursoragent@cursor.com>

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)
+        # `topk_p`/`topk_index` are produced by the captured seed iter.
-        buffers.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index)
+        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:

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:
+        stashed = FrozenKVMTPDraftInput()
-            draft_input = FrozenKVMTPDraftInput()
+        stashed.bonus_tokens = last_token_ids.to(torch.int64)
-
+        stashed.hidden_states = last_hidden_states
-        draft_input.bonus_tokens = last_token_ids.to(torch.int64)
+        # Real-shaped zeros so inherited `filter_batch`/`merge_batch` can slice
-        draft_input.hidden_states = last_hidden_states
+        # them between iters; overwritten by the captured seed iter.
-        draft_input.capture_hidden_mode = CaptureHiddenMode.LAST
+        bs = last_token_ids.shape[0]
-        draft_input.num_tokens_per_req = 1
+        device = last_token_ids.device
-        draft_input.num_tokens_for_logprob_per_req = 1
+        stashed.topk_p = torch.zeros(
-        draft_input.positions = self._position_for_batch(batch)
+            (bs, self.topk), device=device, dtype=torch.float32
-
+        )
-        forward_mode_backup = batch.forward_mode
+        stashed.topk_index = torch.zeros(
-        input_ids_backup = batch.input_ids
+            (bs, self.topk), device=device, dtype=torch.int64
-        return_hidden_states_backup = batch.return_hidden_states
+        )
-        return_logprob_backup = batch.return_logprob
+        stashed.capture_hidden_mode = CaptureHiddenMode.LAST
-        spec_info_backup = batch.spec_info
+        stashed.num_tokens_per_req = 1
-
+        stashed.num_tokens_for_logprob_per_req = 1
-        batch.forward_mode = ForwardMode.DECODE
+        batch.spec_info = stashed
-        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
 
     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(

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"],
+        "bonus_tokens": bonus_tokens,
-        "topk_index": draft_inputs["topk_index"],
     }
 
 
@@ -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,