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

[Spec] Anchor GLM-5.2 MTP IndexShare topk on the draft-extend step (#29787)

Browse Source 
Co-authored-by: kpham-sgl <264503018+kpham-sgl@users.noreply.github.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
2026-07-06 20:36:48 -07:00
committed by GitHub
co-authored by kpham-sgl Xinyuan Tong Xinyuan Tong Claude Fable 5
parent 4145e595cf
commit 16372b4c5f
12 changed files with 185 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/batch_overlap/two_batch_overlap.py
+1 -1
View File
@@ -737,7 +737,7 @@ class TboForwardBatchPreparer:
"split_index", # for split prefill
"orig_seq_lens", # only used by qwen-1m, thus not care
"return_pooled_hidden_states",
"reuse_mtp_topk_indices", # forward-level flag, inherited by both child batches
"reuse_dsa_topk_indices", # forward-level flag, inherited by both child batches
]:
output_dict[key] = getattr(batch, key)
python/sglang/srt/managers/overlap_utils.py
+20
View File
@@ -111,6 +111,7 @@ class RelayPayload:
topk_index: Optional[torch.Tensor] = None
hidden_states: Optional[torch.Tensor] = None
draft_probs: Optional[torch.Tensor] = None
dsa_topk_indices: Optional[torch.Tensor] = None
@classmethod
def from_draft_input(cls, draft_input: EagleDraftInput) -> RelayPayload:
@@ -120,6 +121,7 @@ class RelayPayload:
topk_index=draft_input.topk_index,
hidden_states=draft_input.hidden_states,
draft_probs=getattr(draft_input, "draft_probs", None),
dsa_topk_indices=getattr(draft_input, "dsa_topk_indices", None),
)
@@ -216,6 +218,15 @@ class FutureMap:
device=self.device,
)
self.dsa_topk_indices_buf = None
if payload.dsa_topk_indices is not None:
seed0 = payload.dsa_topk_indices[0]
self.dsa_topk_indices_buf = torch.empty(
(self.req_pool_size, *seed0.shape),
dtype=payload.dsa_topk_indices.dtype,
device=self.device,
)
def _resolve_spec_extras(self, batch: ScheduleBatch) -> None:
if self.spec_algo.is_ngram():
# FIXME: remove once precomputed draft is supported.
@@ -255,6 +266,8 @@ class FutureMap:
draft_input.bonus_tokens = self.output_tokens_buf[indices]
if self.need_hidden_states and not self.need_topk:
draft_input.hidden_states = self.hidden_states_buf[indices]
if self.dsa_topk_indices_buf is not None:
draft_input.dsa_topk_indices = self.dsa_topk_indices_buf[indices]
if _DEBUG_ASSERT:
_assert_nonneg_and_invalidate(
draft_input.bonus_tokens, self.output_tokens_buf, indices
@@ -338,3 +351,10 @@ class FutureMap:
)
if self.draft_probs_buf is not None and payload.draft_probs is not None:
self.draft_probs_buf[indices] = payload.draft_probs
if (
self.dsa_topk_indices_buf is not None
and payload.dsa_topk_indices is not None
):
self.dsa_topk_indices_buf[indices] = payload.dsa_topk_indices.to(
self.dsa_topk_indices_buf.dtype
)
python/sglang/srt/model_executor/forward_batch_info.py
+2 -2
View File
@@ -431,8 +431,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
return_hidden_states_before_norm: bool = False
# Gate for reusing the first MTP draft step's indexer topk across steps;
# the carried topk lives on spec_info (see EagleDraftInput.mtp_topk_indices).
reuse_mtp_topk_indices: Optional[bool] = False
# the carried topk lives on spec_info (see EagleDraftInput.dsa_topk_indices).
reuse_dsa_topk_indices: Optional[bool] = False
# === Forward-derived (built in init_new on the forward stream; FB-owned) ===
# Position information
python/sglang/srt/models/deepseek_nextn.py
+13 -4
View File
@@ -237,13 +237,22 @@ class DeepseekModelNextN(nn.Module):
residual,
zero_allocator,
prev_topk_indices=(
forward_batch.spec_info.mtp_topk_indices
if forward_batch.reuse_mtp_topk_indices
forward_batch.spec_info.dsa_topk_indices
if forward_batch.reuse_dsa_topk_indices
else None
),
)
if forward_batch.reuse_mtp_topk_indices:
forward_batch.spec_info.mtp_topk_indices = topk_indices
if forward_batch.reuse_dsa_topk_indices:
forward_batch.spec_info.dsa_topk_indices = topk_indices
# MTP IndexShare: on draft-extend, publish the last-token DSA
# indexer top-k to seed (avoid recomputing in) the draft-decode loop.
if forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
seed_buf = forward_batch.spec_info.dsa_seed_topk_capture
if seed_buf is not None and topk_indices is not None:
sel = forward_batch.spec_info.dsa_seed_topk_select
src = topk_indices if sel is None else topk_indices[sel]
seed_buf[: src.shape[0]].copy_(src)
if not forward_batch.forward_mode.is_idle():
if residual is not None:
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+22
View File
@@ -69,6 +69,7 @@ class EagleDraftInputBuffers(ForwardInputBuffers):
hidden_states: Optional[torch.Tensor]
global_num_tokens_gpu: Optional[torch.Tensor]
global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
dsa_seed_topk: Optional[torch.Tensor] = None
class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
@@ -228,6 +229,16 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
(self.max_bs,), self.seq_len_fill_value, dtype=torch.int64, device="cpu"
)
dsa_seed_topk = (
torch.zeros(
(self.max_bs, self.eagle_worker.dsa_index_topk),
dtype=torch.int32,
device=model_runner.device,
)
if self.eagle_worker.seed_dsa_topk_from_draft_extend
else None
)
self.buffers = EagleDraftInputBuffers(
input_ids=input_ids,
req_pool_indices=req_pool_indices,
@@ -245,6 +256,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
hidden_states=hidden_states,
global_num_tokens_gpu=global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
dsa_seed_topk=dsa_seed_topk,
)
self.buffers.share_buffers()
@@ -372,6 +384,8 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
hidden_states=hidden_states,
capture_hidden_mode=capture_mode,
)
if self.buffers.dsa_seed_topk is not None:
spec_info.dsa_topk_indices = self.buffers.dsa_seed_topk[:num_seqs]
sampling_info = SamplingBatchInfo(
temperatures=self.temperatures[:num_seqs],
@@ -504,6 +518,8 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
buffers.draft_probs.zero_()
if buffers.hidden_states is not None:
buffers.hidden_states.zero_()
if buffers.dsa_seed_topk is not None:
buffers.dsa_seed_topk.zero_()
buffers.req_pool_indices.zero_()
num_tokens = bs * self.num_tokens_per_bs
@@ -563,6 +579,12 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
and forward_batch.spec_info.hidden_states is not None
):
buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
if buffers.dsa_seed_topk is not None:
seed = forward_batch.spec_info.dsa_topk_indices
if seed is not None:
buffers.dsa_seed_topk[:raw_bs].copy_(seed)
else:
buffers.dsa_seed_topk[:raw_bs].zero_()
# Only rejection sampling reads temperatures (renorm_draft_probs); skip
# the copy otherwise to keep the non-RS path free of extra work.
if (
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+18
View File
@@ -69,6 +69,7 @@ class EagleDraftExtendInputBuffers(ForwardInputBuffers):
next_token_logits_buffer: torch.Tensor
global_num_tokens_gpu: Optional[torch.Tensor]
global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
dsa_seed_topk_capture: Optional[torch.Tensor] = None
class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
@@ -234,6 +235,17 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
(self.max_bs,), self.seq_len_fill_value, dtype=torch.int64, device="cpu"
)
dsa_seed_topk_capture = (
torch.full(
(self.max_num_token, self.eagle_worker.dsa_index_topk),
-1,
dtype=torch.int32,
device=model_runner.device,
)
if self.eagle_worker.seed_dsa_topk_from_draft_extend
else None
)
self.buffers = EagleDraftExtendInputBuffers(
input_ids=input_ids,
req_pool_indices=req_pool_indices,
@@ -249,6 +261,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
next_token_logits_buffer=next_token_logits_buffer,
global_num_tokens_gpu=global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
dsa_seed_topk_capture=dsa_seed_topk_capture,
)
self.buffers.share_buffers()
@@ -385,6 +398,11 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
padded_static_len=self.padded_static_len,
)
if self.buffers.dsa_seed_topk_capture is not None:
spec_info.dsa_seed_topk_capture = self.buffers.dsa_seed_topk_capture[
:num_tokens
]
def run_once():
self.draft_extend_attn_backend.init_forward_metadata_in_graph(forward_batch)
python/sglang/srt/speculative/eagle_info.py
+15 -1
View File
@@ -161,7 +161,7 @@ class EagleDraftInput(SpecInput):
# Survives across draft steps: spec_info is shared by reference across the
# per-step forwards (each runs on a copied ForwardBatch, dropping writebacks).
mtp_topk_indices: Optional[torch.Tensor] = None
dsa_topk_indices: Optional[torch.Tensor] = None
# Per-req bonus token (the "+1" target prediction at end of each accept
# chain); the worker copies it here post-extend for next iter's draft.
@@ -233,6 +233,8 @@ class EagleDraftInput(SpecInput):
if self.hidden_states is not None:
self.hidden_states = self.hidden_states[: len(new_indices)]
self.bonus_tokens = self.bonus_tokens[: len(new_indices)]
if self.dsa_topk_indices is not None:
self.dsa_topk_indices = self.dsa_topk_indices[: len(new_indices)]
else:
# in some cases(e.g draft_extend), we have not filtered the batch by `unfinished_index`
self.topk_p = self.topk_p[new_indices]
@@ -242,6 +244,8 @@ class EagleDraftInput(SpecInput):
if self.hidden_states is not None:
self.hidden_states = self.hidden_states[new_indices]
self.bonus_tokens = self.bonus_tokens[new_indices]
if self.dsa_topk_indices is not None:
self.dsa_topk_indices = self.dsa_topk_indices[new_indices]
def merge_batch(self, spec_info: "EagleDraftInput"):
if self.future_indices is not None:
@@ -260,6 +264,7 @@ class EagleDraftInput(SpecInput):
self.topk_p = spec_info.topk_p
self.topk_index = spec_info.topk_index
self.draft_probs = spec_info.draft_probs
self.dsa_topk_indices = spec_info.dsa_topk_indices
return
if len(spec_info.topk_index) == 0:
return
@@ -272,6 +277,12 @@ class EagleDraftInput(SpecInput):
)
self.topk_p = torch.cat([self.topk_p, spec_info.topk_p])
self.topk_index = torch.cat([self.topk_index, spec_info.topk_index])
if self.dsa_topk_indices is not None and spec_info.dsa_topk_indices is not None:
self.dsa_topk_indices = torch.cat(
[self.dsa_topk_indices, spec_info.dsa_topk_indices]
)
else:
self.dsa_topk_indices = None
if self.draft_probs is not None and spec_info.draft_probs is not None:
self.draft_probs = torch.cat([self.draft_probs, spec_info.draft_probs])
@@ -317,6 +328,9 @@ class EagleDraftExtendInput(SpecInput):
num_tokens_per_req: int = -1
num_tokens_for_logprob_per_req: int = 1
dsa_seed_topk_capture: Optional[torch.Tensor] = None
dsa_seed_topk_select: Optional[torch.Tensor] = None
# None for draft-extend's idle batch; attention backends fall back to
# rebuilding plain metadata from seq_lens when this is None.
kv_indptr: torch.Tensor = None
python/sglang/srt/speculative/eagle_worker_v2.py
+87 -14
View File
@@ -14,6 +14,7 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner i
)
from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
from sglang.srt.kv_canary.runner.canary_manager import context_tuple
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend
from sglang.srt.layers.attention.tokenspeed_mla_backend import TokenspeedMLABackend
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
@@ -191,16 +192,9 @@ class EagleDraftWorker(EagleDraftWorkerBase):
# Alias for better readability
self.draft_runner = self.draft_worker.model_runner
# Reuse the first draft step's NSA/DSA indexer topk across the rest;
# topk == 1 only (select_top_k_tokens reorders rows, desyncing indices).
self.index_share_for_mtp_iteration = (
getattr(
self.draft_runner.model_config.hf_config,
"index_share_for_mtp_iteration",
False,
)
and self.topk == 1
)
self._init_dsa_index_share_state()
# Eager draft-extend seed buffer (graph paths use their own static ones).
self.dsa_extend_topk_buf: Optional[torch.Tensor] = None
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
@@ -264,6 +258,23 @@ class EagleDraftWorker(EagleDraftWorkerBase):
if (c := self.draft_runner.canary_manager) is not None:
c.mark_init_finished()
def _init_dsa_index_share_state(self) -> None:
# Populate DSA index-share fields from the draft runner's hf_config.
# Reused by the attention unit-test harnesses, which skip __init__.
hf_config = self.draft_runner.model_config.hf_config
# Reuse the first draft step's DSA indexer topk across the rest;
# topk == 1 only (select_top_k_tokens reorders rows, desyncing indices).
self.index_share_for_mtp_iteration = (
getattr(hf_config, "index_share_for_mtp_iteration", False)
and self.topk == 1
)
# GLM-5.2 MTP IndexShare: seed reused indexer top-k from draft-extend
# (last verified token), not draft-decode step 0.
self.dsa_index_topk = getattr(hf_config, "index_topk", None)
self.seed_dsa_topk_from_draft_extend = (
self.index_share_for_mtp_iteration and self.dsa_index_topk is not None
)
def _rebuild_topk1_chain_buffers(self) -> None:
# For topk=1 the draft tree degenerates to a chain, so parent_list and
# top_scores_index are runtime-invariant. Must be rebuilt after any
@@ -629,8 +640,13 @@ class EagleDraftWorker(EagleDraftWorkerBase):
# Forward multiple steps
scores = None
if self.index_share_for_mtp_iteration:
forward_batch.reuse_mtp_topk_indices = True
spec_info.mtp_topk_indices = None
forward_batch.reuse_dsa_topk_indices = True
# Keep the draft-extend seed so step 0 reuses it; else recompute it.
if not (
self.seed_dsa_topk_from_draft_extend
and spec_info.dsa_topk_indices is not None
):
spec_info.dsa_topk_indices = None
for i in range(self.speculative_num_steps):
input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
i, topk_p, topk_index, hidden_states, scores, self.topk
@@ -706,8 +722,8 @@ class EagleDraftWorker(EagleDraftWorkerBase):
forward_batch.positions.add_(1)
if self.index_share_for_mtp_iteration:
spec_info.mtp_topk_indices = None
forward_batch.reuse_mtp_topk_indices = False
spec_info.dsa_topk_indices = None
forward_batch.reuse_dsa_topk_indices = False
# Organize the results
if (
@@ -794,6 +810,23 @@ class EagleDraftWorker(EagleDraftWorkerBase):
if mm_input_embeds is not None:
forward_batch.mm_input_embeds = mm_input_embeds
# Seed the first draft-decode loop from each request's last prefill
# position. Gather last-per-req before the copy (prefill can be long).
# Skipped under context-parallel prefill (token layout wouldn't match).
seed_from_extend = (
self.seed_dsa_topk_from_draft_extend
and not forward_batch.forward_mode.is_idle()
and not dsa_use_prefill_cp(forward_batch)
)
if seed_from_extend:
bs = forward_batch.batch_size
forward_batch.spec_info.dsa_seed_topk_capture = (
self._get_dsa_extend_topk_buf(bs)
)
forward_batch.spec_info.dsa_seed_topk_select = (
torch.cumsum(forward_batch.extend_seq_lens, dim=0) - 1
).long()
canary_ctx = (
context_tuple(
c.with_ops_outside_graph(
@@ -810,6 +843,10 @@ class EagleDraftWorker(EagleDraftWorkerBase):
maybe_detect_nan(logits_output.next_token_logits, "draft_extend_for_prefill")
maybe_detect_inf(logits_output.next_token_logits, "draft_extend_for_prefill")
prefill_dsa_topk = None
if seed_from_extend:
prefill_dsa_topk = self.dsa_extend_topk_buf[:bs].clone()
# Assemble the next-iter draft spec_info from the extend output.
use_rejection_sampling = self.server_args.speculative_use_rejection_sampling
probs = renorm_draft_probs(
@@ -829,8 +866,22 @@ class EagleDraftWorker(EagleDraftWorkerBase):
bonus_tokens=next_token_ids,
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
dsa_topk_indices=prefill_dsa_topk,
)
def _get_dsa_extend_topk_buf(self, num_tokens: int) -> torch.Tensor:
"""Lazily-grown int32 [num_tokens, index_topk] eager draft-extend seed buffer."""
buf = self.dsa_extend_topk_buf
if buf is None or buf.shape[0] < num_tokens:
buf = torch.full(
(num_tokens, self.dsa_index_topk),
-1,
dtype=torch.int32,
device=self.device,
)
self.dsa_extend_topk_buf = buf
return buf[:num_tokens]
def _draft_extend_for_decode(
self, batch: ScheduleBatch, batch_result: GenerationBatchResult
):
@@ -882,6 +933,13 @@ class EagleDraftWorker(EagleDraftWorkerBase):
and self.cuda_graph_runner_for_draft_extend.can_run_graph(forward_batch)
)
# Eager path publishes the indexer top-k into a worker buffer (the graph
# path uses the runner's static buffer). Gathered at select_index below.
if self.seed_dsa_topk_from_draft_extend and not can_cuda_graph:
forward_batch.spec_info.dsa_seed_topk_capture = (
self._get_dsa_extend_topk_buf(forward_batch.input_ids.shape[0])
)
canary_ctx = (
context_tuple(
c.with_ops_outside_graph(
@@ -912,6 +970,19 @@ class EagleDraftWorker(EagleDraftWorkerBase):
f"draft_extend_for_decode (cuda_graph={can_cuda_graph})",
)
# Gather the per-request last-position indexer top-k as the next loop's
# seed (select_index already picks the last accepted position per req).
dsa_seed_topk_indices = None
if self.seed_dsa_topk_from_draft_extend:
if can_cuda_graph:
dsa_extend_topk_capture = (
self.cuda_graph_runner_for_draft_extend.buffers.dsa_seed_topk_capture
)
else:
dsa_extend_topk_capture = forward_batch.spec_info.dsa_seed_topk_capture
# Fancy indexing returns a fresh tensor (detached from the buffer).
dsa_seed_topk_indices = dsa_extend_topk_capture[select_index]
# Reorganize the spec info for the next batch
draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[
select_index
@@ -961,6 +1032,8 @@ class EagleDraftWorker(EagleDraftWorkerBase):
)
if self.server_args.speculative_use_rejection_sampling:
next_draft_input.draft_probs = ret_draft_probs
if self.seed_dsa_topk_from_draft_extend:
next_draft_input.dsa_topk_indices = dsa_seed_topk_indices
class EAGLEWorkerV2(BaseSpecWorker):
python/sglang/srt/speculative/standalone_worker_v2.py
+3
View File
@@ -113,6 +113,9 @@ class StandaloneDraftWorker(EagleDraftWorker):
)
and self.topk == 1
)
self.dsa_index_topk = None
self.seed_dsa_topk_from_draft_extend = False
self.dsa_extend_topk_buf = None
def alloc_memory_pool(
self,
python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py
+2
View File
@@ -20,6 +20,7 @@ from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
)
from sglang.srt.speculative.eagle_info import EagleDraftExtendInput
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import fast_topk
@@ -496,6 +497,7 @@ class _EagleDraftExtendV2WorkerHarness:
self.eagle_use_aux_hidden_state = False
self.hot_token_id = None
self.draft_runner.model = model_forward
EagleDraftWorker._init_dsa_index_share_state(self)
def _build_eagle_draft_extend_fixture(
python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py
+1 -9
View File
@@ -188,15 +188,7 @@ class _EagleDraftWorkerHarness:
self._topk1_parents_prealloc = None
self._topk1_score_indices_prealloc = None
EagleDraftWorker._rebuild_topk1_chain_buffers(self)
# draft_forward reads this (set in EagleDraftWorker.__init__, skipped here).
self.index_share_for_mtp_iteration = (
getattr(
self.model_config.hf_config,
"index_share_for_mtp_iteration",
False,
)
and self.topk == 1
)
EagleDraftWorker._init_dsa_index_share_state(self)
@property
def draft_model_runner(self):