minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Khoa Pham
2026-06-01 22:27:25 -07:00
committed by GitHub
co-authored by Cursor
parent f27fa0da93
commit 08526c7fca
3 changed files with 97 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
+15 -3
View File
@@ -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:
python/sglang/srt/speculative/frozen_kv_mtp_worker.py
+69 -57
View File
@@ -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(
python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py
+13 -4
View File
@@ -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,