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

[Spec] Install EagleDraftExtendInput as the V2 draft-extend spec_info (#24860)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-12 00:46:16 -07:00
committed by GitHub
parent 1cd5cb1220
commit a52ccd2179
5 changed files with 62 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/kv_canary/plan_input.py
+1 -1
View File
@@ -118,7 +118,7 @@ def _extract_prefix_lens_and_extend_seq_lens(
out_prefix_lens.copy_(forward_batch.seq_lens[:bs].to(torch.int64))
out_extend_seq_lens.fill_(int(spec_info.draft_token_num))
elif forward_mode.is_draft_extend_v2():
# Evidence: EagleDraftInputV2Mixin.prepare_for_extend_to_fill_draft_kvcache bumps
# Evidence: EagleDraftExtendInputV2Mixin.prepare_for_extend_to_fill_draft_kvcache bumps
# seq_lens by num_draft_tokens. FlashAttentionBackend.init_forward_metadata reads the
# draft-extend-v2 query length from spec_info.extend_seq_lens_tensor when available.
# CUDA-graph replay passes extend_seq_lens but omits extend_prefix_lens, so derive the
python/sglang/srt/speculative/eagle_info.py
+10 -12
View File
@@ -9,6 +9,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
from sglang.srt.speculative.eagle_info_v2 import (
EagleDraftExtendInputV2Mixin,
EagleDraftInputV2Mixin,
EagleVerifyInputV2Mixin,
)
@@ -168,10 +169,6 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
# V2 overlap worker only: req_pool_indices used as buf slot keys.
future_indices: Optional[torch.Tensor] = None
# V2 reuses `EagleDraftInput` across phases (V1 has a separate
# `EagleDraftExtendInput` for these). Set during V2's draft-extend.
num_correct_drafts: Optional[torch.Tensor] = None
num_accept_tokens: Optional[torch.Tensor] = None
def __post_init__(self):
super().__init__(SpecInputType.EAGLE_DRAFT)
@@ -276,17 +273,18 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
@dataclass
class EagleDraftExtendInput(SpecInput):
"""Inputs to the draft-extend forward (the per-accepted-token pass after verify).
class EagleDraftExtendInput(SpecInput, EagleDraftExtendInputV2Mixin):
"""Inputs to the draft-extend forward (the fill-draft-kvcache pass after
target prefill / verify).
Produced by `EagleVerifyInput.verify`, installed on `batch.spec_info` for
the draft-extend forward, then replaced with a fresh `EagleDraftInput` for
the next iter's draft.
Installed on `batch.spec_info` by the worker's `_draft_extend_for_*`
(and synthetically by draft-extend cuda-graph capture), then replaced
with a fresh `EagleDraftInput` for the next iter's draft.
"""
# shape: (total_accepted, hidden_size). Sliced from verify-time hidden_states
# by accept_index; consumed by the draft-extend forward. None when the spec
# algorithm's draft doesn't read hidden_states (e.g., STANDALONE).
# Target-model hidden states for the draft-extend forward; None when the
# draft doesn't read hidden_states (e.g., STANDALONE). Shape: decode
# (bs * num_draft_tokens, hidden), prefill (extend_num_tokens, hidden).
hidden_states: Optional[torch.Tensor] = None
# Per-req accept counts. `num_accept_tokens = num_correct_drafts + 1`.
python/sglang/srt/speculative/eagle_info_v2.py
+8 -2
View File
@@ -60,7 +60,11 @@ if TYPE_CHECKING:
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
EAGLEDraftCudaGraphRunner,
)
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.eagle_info import (
EagleDraftExtendInput,
EagleDraftInput,
EagleVerifyInput,
)
if is_cuda() or is_musa():
from sgl_kernel import (
@@ -302,8 +306,10 @@ class EagleDraftInputV2Mixin:
can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
return forward_batch, can_cuda_graph
class EagleDraftExtendInputV2Mixin:
def prepare_for_extend_to_fill_draft_kvcache(
self,
self: EagleDraftExtendInput,
batch: ScheduleBatch,
predict: torch.Tensor,
num_draft_tokens: int,
python/sglang/srt/speculative/eagle_worker_v2.py
+22 -19
View File
@@ -55,7 +55,11 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
)
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.eagle_info import (
EagleDraftExtendInput,
EagleDraftInput,
EagleVerifyInput,
)
from sglang.srt.speculative.eagle_info_v2 import fill_bonus_tokens
from sglang.srt.speculative.eagle_utils import (
TreeMaskMode,
@@ -655,17 +659,15 @@ class EagleDraftWorker(BaseDraftWorker):
)
pt += extend_len
# Construct spec_info
next_draft_input = EagleDraftInput(
# Draft-extend spec_info for the extend forward; carries only
# hidden_states + shape info.
batch.spec_info = EagleDraftExtendInput(
hidden_states=target_hidden_states,
bonus_tokens=next_token_ids,
# draft mode is same with decode mode, only 1 token per req
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
)
batch.spec_info = next_draft_input
# Run forward (LAST mode: only the final hidden state per request,
# to feed the next draft step which expects [bs, hidden_dim]).
# STANDALONE skips hidden states end-to-end.
@@ -696,20 +698,27 @@ class EagleDraftWorker(BaseDraftWorker):
maybe_detect_nan(logits_output.next_token_logits, "draft_extend_for_prefill")
maybe_detect_inf(logits_output.next_token_logits, "draft_extend_for_prefill")
# Update spec_info for the next draft step
# Assemble the next-iter draft spec_info from the extend output.
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
next_draft_input.topk_p, next_draft_input.topk_index = fast_topk(
probs, self.topk, dim=-1
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
return EagleDraftInput(
topk_p=topk_p,
topk_index=topk_index,
hidden_states=logits_output.hidden_states,
bonus_tokens=next_token_ids,
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
)
next_draft_input.hidden_states = logits_output.hidden_states
return next_draft_input
def _draft_extend_for_decode(
self, batch: ScheduleBatch, batch_result: GenerationBatchResult
):
# Batch 2: Draft extend
draft_input = EagleDraftInput(
draft_extend_input = EagleDraftExtendInput(
hidden_states=batch_result.logits_output.hidden_states,
# accept_lens includes the bonus token; correct drafts exclude it.
num_correct_drafts=batch_result.accept_lens - 1,
num_accept_tokens=batch_result.accept_lens,
# Draft-extend fills the whole tree width (num_draft_tokens) per req,
# not num_steps + 1, so DP MLP-sync padding stays consistent for topk > 1.
num_tokens_per_req=self.speculative_num_draft_tokens,
@@ -724,7 +733,7 @@ class EagleDraftWorker(BaseDraftWorker):
# Prepare for draft extend in a separate stream
with self.plan_stream_ctx:
forward_batch = draft_input.prepare_for_extend_to_fill_draft_kvcache(
forward_batch = draft_extend_input.prepare_for_extend_to_fill_draft_kvcache(
batch,
batch_result.next_token_ids,
self.speculative_num_draft_tokens,
@@ -737,12 +746,6 @@ class EagleDraftWorker(BaseDraftWorker):
self.plan_stream
)
if forward_batch.spec_info.num_correct_drafts is None:
# `batch_result.accept_lens` already includes the bonus token, so use it
# directly for `num_accept_tokens` and subtract 1 for `num_correct_drafts`.
forward_batch.spec_info.num_correct_drafts = batch_result.accept_lens - 1
forward_batch.spec_info.num_accept_tokens = batch_result.accept_lens
# Run draft extend batch in the main compute stream
can_cuda_graph = (
self.cuda_graph_runner_for_draft_extend
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+21 -10
View File
@@ -43,7 +43,11 @@ from sglang.srt.model_executor.forward_batch_info import (
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import BaseDraftWorker, BaseSpecWorker
from sglang.srt.speculative.draft_utils import DraftBackendFactory
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.eagle_info import (
EagleDraftExtendInput,
EagleDraftInput,
EagleVerifyInput,
)
from sglang.srt.speculative.eagle_info_v2 import fill_bonus_tokens
from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient
from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
@@ -414,16 +418,15 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
"draft_extend_for_prefill: next_token_ids before draft embed",
)
# Construct spec_info
next_draft_input = EagleDraftInput(
# Draft-extend spec_info for the extend forward; carries only
# hidden_states + shape info.
extend_input = EagleDraftExtendInput(
hidden_states=target_hidden_states,
bonus_tokens=next_token_ids,
# draft mode is same with decode mode, only 1 token per req
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
)
batch.spec_info = next_draft_input
batch.spec_info = extend_input
# Chain-style MTP needs FULL to get all-token hidden states;
# non-chain only needs LAST (the target model's hidden states).
@@ -484,8 +487,16 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
forward_batch.extend_seq_lens,
topk_index,
)
next_draft_input.topk_p = torch.cat(topk_p_list, dim=1)
next_draft_input.topk_index = torch.cat(topk_index_list, dim=1)
next_draft_input = EagleDraftInput(
topk_p=torch.cat(topk_p_list, dim=1),
topk_index=torch.cat(topk_index_list, dim=1),
# Chain-style left the last step's hidden_states on the extend
# input; non-chain keeps the target hidden states.
hidden_states=extend_input.hidden_states,
bonus_tokens=next_token_ids,
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
)
# Update req_to_hidden_states_pool for KV Cache reversion
if forward_batch.extend_seq_lens is not None:
@@ -504,7 +515,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
self, batch: ScheduleBatch, batch_result: GenerationBatchResult
):
# Batch 2: Draft extend
draft_input = EagleDraftInput(
draft_extend_input = EagleDraftExtendInput(
hidden_states=batch_result.logits_output.hidden_states,
num_tokens_per_req=self.speculative_num_steps + 1,
num_tokens_for_logprob_per_req=1,
@@ -513,7 +524,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
# Prepare for draft extend in a separate stream
# Notice that here we use batch_result.next_token_ids as the input ids
with self.plan_stream_ctx:
forward_batch = draft_input.prepare_for_extend_to_fill_draft_kvcache(
forward_batch = draft_extend_input.prepare_for_extend_to_fill_draft_kvcache(
batch,
batch_result.next_token_ids,
self.speculative_num_draft_tokens,