From 9e3bb9a3073554686ec48373032bbfd20e511377 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 18 May 2026 14:24:54 -0700 Subject: [PATCH] [Spec] fold can_run_cuda_graph into EagleVerifyOutput; drop dead extend-after-decode check (#25566) --- python/sglang/srt/speculative/eagle_info.py | 4 +++ python/sglang/srt/speculative/eagle_worker.py | 36 ++++++------------- .../sglang/srt/speculative/eagle_worker_v2.py | 2 +- .../srt/speculative/frozen_kv_mtp_worker.py | 17 +++++---- .../speculative/multi_layer_eagle_worker.py | 36 ++++++------------- .../multi_layer_eagle_worker_v2.py | 2 +- python/sglang/srt/speculative/spec_info.py | 6 ++-- .../srt/speculative/standalone_worker.py | 2 +- .../srt/speculative/standalone_worker_v2.py | 2 +- 9 files changed, 42 insertions(+), 65 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index d643f56db..767ef767a 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -1012,6 +1012,10 @@ class EagleVerifyOutput: num_correct_drafts_per_req_cpu: List[int] # Accepted indices from logits_output.next_token_logits accept_indices: torch.Tensor + # Whether the target verify forward ran a captured cuda graph. Set by + # the worker after `EagleVerifyInput.sample` returns; default kept so + # idle / direct constructions don't have to pass it. + can_run_cuda_graph: bool = False @classmethod def create_idle( diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index b1931214f..6ca7632bf 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -5,7 +5,6 @@ from typing import List, Optional, Tuple import torch -from sglang.srt.distributed import get_tp_group from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import ( EAGLEDraftNpuGraphRunner, ) @@ -162,7 +161,7 @@ class EAGLEWorker(TpModelWorker): server_args=server_args, gpu_id=gpu_id, tp_rank=tp_rank, - pp_rank=0, # FIXME + pp_rank=0, # spec workers don't support pipeline parallelism dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, attn_cp_rank=attn_cp_rank, @@ -492,8 +491,9 @@ class EAGLEWorker(TpModelWorker): set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True) set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True) + # Install verify_input as `batch.spec_info` for the verify forward. batch.spec_info = verify_input - logits_output, verify_output, can_run_cuda_graph = self.verify(batch) + verify_output = self.verify(batch) if get_global_tracing_enabled(): for idx, req in enumerate(batch.reqs): @@ -520,8 +520,9 @@ class EAGLEWorker(TpModelWorker): self.server_args.enable_dp_attention or draft_extend_input.input_ids.shape[0] > 0 ): - # decode is not finished; stash for extend, then restash - # the next-iter EagleDraftInput it returns. + # decode is not finished; install draft_extend_input for + # the extend forward, then install the next-iter + # EagleDraftInput it returns. batch.spec_info = draft_extend_input next_draft_input = self.forward_draft_extend_after_decode(batch) batch.spec_info = next_draft_input @@ -542,31 +543,13 @@ class EAGLEWorker(TpModelWorker): ) return GenerationBatchResult( - logits_output=logits_output, + logits_output=verify_output.logits_output, next_token_ids=verify_output.accept_tokens, num_correct_drafts=sum(verify_output.num_correct_drafts_per_req_cpu), num_correct_drafts_per_req_cpu=verify_output.num_correct_drafts_per_req_cpu, - can_run_cuda_graph=can_run_cuda_graph, + can_run_cuda_graph=verify_output.can_run_cuda_graph, ) - def check_forward_draft_extend_after_decode(self, verify_output: EagleVerifyOutput): - local_need_forward = verify_output.draft_extend_input.input_ids.shape[0] > 0 - if not self.server_args.enable_dp_attention: - return local_need_forward - - global_need_forward = torch.tensor( - [ - (local_need_forward), - ], - dtype=torch.int64, - ) - torch.distributed.all_reduce( - global_need_forward, group=get_tp_group().cpu_group - ) - global_need_forward_cnt = global_need_forward[0].item() - need_forward = global_need_forward_cnt > 0 - return need_forward - def forward_target_extend( self, batch: ScheduleBatch ) -> Tuple[LogitsProcessorOutput, torch.Tensor, Optional[torch.Tensor], bool]: @@ -1015,7 +998,8 @@ class EAGLEWorker(TpModelWorker): ForwardMode.DECODE if not batch.forward_mode.is_idle() else ForwardMode.IDLE ) - return logits_output, res, can_run_cuda_graph + res.can_run_cuda_graph = can_run_cuda_graph + return res def _mamba_verify_update( self, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index d5d64b844..172efd023 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -154,7 +154,7 @@ class EagleDraftWorker(BaseDraftWorker): server_args=server_args, gpu_id=gpu_id, tp_rank=tp_rank, - pp_rank=0, # FIXME + pp_rank=0, # spec workers don't support pipeline parallelism dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, attn_cp_rank=attn_cp_rank, diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index c52e01942..6e4ecdf03 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py @@ -449,8 +449,9 @@ class FrozenKVMTPWorker(TpModelWorker): set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True) set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True) + # Install verify_input as `batch.spec_info` for the verify forward. batch.spec_info = verify_input - logits_output, verify_output, can_run_cuda_graph = self.verify(batch) + verify_output = self.verify(batch) if get_global_tracing_enabled(): for idx, req in enumerate(batch.reqs): @@ -470,18 +471,19 @@ class FrozenKVMTPWorker(TpModelWorker): self.server_args.enable_dp_attention or draft_extend_input.input_ids.shape[0] > 0 ): - # Stash for the seed step; _run_assistant_seed_step swaps in - # a fresh FrozenKVMTPDraftInput for next iter. + # Install draft_extend_input as `batch.spec_info` for the seed + # step; `_run_assistant_seed_step` replaces it with a fresh + # `FrozenKVMTPDraftInput` for next iter. batch.spec_info = draft_extend_input self.forward_draft_extend_after_decode(batch) set_time_batch(batch.reqs, "set_spec_draft_extend_end_time", trace_only=True) return GenerationBatchResult( - logits_output=logits_output, + logits_output=verify_output.logits_output, next_token_ids=verify_output.accept_tokens, num_correct_drafts=sum(verify_output.num_correct_drafts_per_req_cpu), num_correct_drafts_per_req_cpu=verify_output.num_correct_drafts_per_req_cpu, - can_run_cuda_graph=can_run_cuda_graph, + can_run_cuda_graph=verify_output.can_run_cuda_graph, ) def forward_target_extend( @@ -518,7 +520,7 @@ class FrozenKVMTPWorker(TpModelWorker): input_is_idle = batch.forward_mode.is_idle() if not input_is_idle and draft_extend_input.input_ids.shape[0] == 0: - # All reqs finished; stash an idle FrozenKVMTPDraftInput so the + # All reqs finished. Install an idle FrozenKVMTPDraftInput so the # next-iter draft sees a valid spec_info. batch = batch.copy() batch.prepare_for_idle() @@ -775,4 +777,5 @@ class FrozenKVMTPWorker(TpModelWorker): ) del seq_lens_pre_verify - return logits_output, res, can_run_cuda_graph + res.can_run_cuda_graph = can_run_cuda_graph + return res diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index ebdeec9d8..c8f112360 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -18,7 +18,6 @@ from typing import TYPE_CHECKING, List, Optional, Tuple import torch -from sglang.srt.distributed import get_tp_group from sglang.srt.layers.dp_attention import get_attention_tp_group from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.moe.utils import speculative_moe_backend_context @@ -136,7 +135,7 @@ class MultiLayerEagleWorker(TpModelWorker): server_args=server_args, gpu_id=gpu_id, tp_rank=tp_rank, - pp_rank=0, # FIXME + pp_rank=0, # spec workers don't support pipeline parallelism dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, attn_cp_rank=attn_cp_rank, @@ -293,8 +292,9 @@ class MultiLayerEagleWorker(TpModelWorker): set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True) set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True) + # Install verify_input as `batch.spec_info` for the verify forward. batch.spec_info = verify_input - logits_output, verify_output, can_run_cuda_graph = self.verify(batch) + verify_output = self.verify(batch) if get_global_tracing_enabled(): for idx, req in enumerate(batch.reqs): @@ -320,8 +320,9 @@ class MultiLayerEagleWorker(TpModelWorker): self.server_args.enable_dp_attention or draft_extend_input.input_ids.shape[0] > 0 ): - # decode is not finished; stash for extend, then restash - # the next-iter EagleDraftInput it returns. + # decode is not finished; install draft_extend_input for + # the extend forward, then install the next-iter + # EagleDraftInput it returns. batch.spec_info = draft_extend_input next_draft_input = self.forward_draft_extend_after_decode(batch) batch.spec_info = next_draft_input @@ -337,31 +338,13 @@ class MultiLayerEagleWorker(TpModelWorker): ) return GenerationBatchResult( - logits_output=logits_output, + logits_output=verify_output.logits_output, next_token_ids=verify_output.accept_tokens, num_correct_drafts=sum(verify_output.num_correct_drafts_per_req_cpu), num_correct_drafts_per_req_cpu=verify_output.num_correct_drafts_per_req_cpu, - can_run_cuda_graph=can_run_cuda_graph, + can_run_cuda_graph=verify_output.can_run_cuda_graph, ) - def check_forward_draft_extend_after_decode(self, verify_output: EagleVerifyOutput): - local_need_forward = verify_output.draft_extend_input.input_ids.shape[0] > 0 - if not self.server_args.enable_dp_attention: - return local_need_forward - - global_need_forward = torch.tensor( - [ - (local_need_forward), - ], - dtype=torch.int64, - ) - torch.distributed.all_reduce( - global_need_forward, group=get_tp_group().cpu_group - ) - global_need_forward_cnt = global_need_forward[0].item() - need_forward = global_need_forward_cnt > 0 - return need_forward - def forward_target_extend( self, batch: ScheduleBatch ) -> Tuple[LogitsProcessorOutput, torch.Tensor, Optional[torch.Tensor], bool]: @@ -644,7 +627,8 @@ class MultiLayerEagleWorker(TpModelWorker): ForwardMode.DECODE if not batch.forward_mode.is_idle() else ForwardMode.IDLE ) - return logits_output, res, can_run_cuda_graph + res.can_run_cuda_graph = can_run_cuda_graph + return res def forward_draft_extend( self, 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 8cbf2dc57..abb95a2fa 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -128,7 +128,7 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): server_args=server_args, gpu_id=gpu_id, tp_rank=tp_rank, - pp_rank=0, # FIXME + pp_rank=0, # spec workers don't support pipeline parallelism dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, attn_cp_rank=attn_cp_rank, diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index ddfddf11a..6715fc122 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -232,9 +232,11 @@ class SpecInput(ABC): def __init__(self, spec_input_type: SpecInputType): self.spec_input_type = spec_input_type + # Cross-algorithm phase guards. Used by attention backends and + # ForwardBatch padding logic to dispatch on phase without hardcoding the + # specific algo class (EAGLE / FROZEN_KV_MTP / DFLASH / NGRAM each have + # their own draft / verify SpecInput subclasses). def is_draft_input(self) -> bool: - # FIXME: remove this function which is only used for assertion - # or use another variable name like `draft_input` to substitute `spec_info` return self.spec_input_type in { SpecInputType.EAGLE_DRAFT, SpecInputType.EAGLE_DRAFT_EXTEND, diff --git a/python/sglang/srt/speculative/standalone_worker.py b/python/sglang/srt/speculative/standalone_worker.py index 12d10533f..ea216b4ae 100644 --- a/python/sglang/srt/speculative/standalone_worker.py +++ b/python/sglang/srt/speculative/standalone_worker.py @@ -87,7 +87,7 @@ class StandaloneWorker(EAGLEWorker): server_args=server_args, gpu_id=gpu_id, tp_rank=tp_rank, - pp_rank=0, # FIXME + pp_rank=0, # spec workers don't support pipeline parallelism dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, attn_cp_rank=attn_cp_rank, diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index dacf2ae56..5d39f2d97 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -93,7 +93,7 @@ class StandaloneDraftWorker(EagleDraftWorker): server_args=server_args, gpu_id=gpu_id, tp_rank=tp_rank, - pp_rank=0, # FIXME + pp_rank=0, # spec workers don't support pipeline parallelism dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, attn_cp_rank=attn_cp_rank,