diff --git a/python/sglang/srt/kv_canary/plan_input.py b/python/sglang/srt/kv_canary/plan_input.py index 9b15a7f63..d4e6bbdf8 100644 --- a/python/sglang/srt/kv_canary/plan_input.py +++ b/python/sglang/srt/kv_canary/plan_input.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 21be4c096..a3391bb16 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -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`. diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 1849480ba..e54abb5a4 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index d3e4a3c61..8532df6e0 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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 diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index b794242d7..1f0e6d925 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -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,