From 55c1963df4cf9932fdb2c6454410ebfed7fd77ea Mon Sep 17 00:00:00 2001 From: weireweire Date: Fri, 31 Jul 2026 08:45:01 +0800 Subject: [PATCH] Remove unused draft-extend CUDA graph top-k (#31430) Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com> --- .../eagle_draft_extend_cuda_graph_runner.py | 20 +-------- .../speculative_draft_extend_runner.py | 45 ++----------------- 2 files changed, 5 insertions(+), 60 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index c4b7f92cc..667064c69 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -38,12 +38,8 @@ from sglang.srt.model_executor.runner_backend_utils import ( from sglang.srt.runtime_context import get_flags from sglang.srt.speculative.eagle_info import EagleDraftExtendInput from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim -from sglang.srt.speculative.spec_utils import ( - fast_topk, - resolve_num_tokens_per_req, -) +from sglang.srt.speculative.spec_utils import resolve_num_tokens_per_req from sglang.srt.utils import ( - is_hip, require_attn_tp_gather, require_gathered_buffer, require_mlp_sync, @@ -51,8 +47,6 @@ from sglang.srt.utils import ( ) from sglang.srt.utils.device_timer import device_timer_ctx -_is_hip = is_hip() - if TYPE_CHECKING: from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker @@ -116,7 +110,6 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): if speculative_num_steps is None else speculative_num_steps ) - self.topk = model_runner.server_args.speculative_eagle_topk self.draft_extend_attn_backend = ( draft_extend_attn_backend or eagle_worker.draft_extend_attn_backend ) @@ -435,17 +428,6 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): forward_batch.positions, forward_batch, ) - # ROCm's argmax tie-breaks differently from CUDA's softmax+max - # path on FP8 logits, which corrupts MTP draft selection on AMD. - # Keep the fastpath CUDA-only. - if self.topk == 1 and not _is_hip: - ret.topk_index = torch.argmax( - ret.next_token_logits, dim=-1, keepdim=True - ) - ret.topk_p = torch.ones_like(ret.topk_index, dtype=torch.float32) - else: - probs = torch.softmax(ret.next_token_logits, dim=-1) - ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1) forward_batch.out_cache_loc = output_cache_loc_backup forward_batch.spec_info.hidden_states = hidden_states_backup diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index 81ff540d6..b3e7dd235 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -22,7 +22,6 @@ from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( 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 from ..attention_methods.dense_attention import DEFAULT_DEVICE from ..attention_methods.dense_attention import DEFAULT_DEVICE as DENSE_DEFAULT_DEVICE @@ -388,36 +387,13 @@ def run_mla_draft_extend_v2_cuda_graph_case( # etc.) is imported from there. -def _assert_draft_extend_outputs_close(actual, expected, settings) -> None: - torch.testing.assert_close( - actual.next_token_logits, - expected.next_token_logits, - atol=settings.atol, - rtol=settings.rtol, - ) - torch.testing.assert_close( - actual.hidden_states, - expected.hidden_states, - atol=settings.atol, - rtol=settings.rtol, - ) - torch.testing.assert_close( - actual.topk_p, - expected.topk_p, - atol=settings.atol, - rtol=settings.rtol, - ) - torch.testing.assert_close(actual.topk_index, expected.topk_index) - - def _assert_draft_extend_v2_outputs_close(actual, expected, settings) -> None: # DRAFT_EXTEND_V2 graph runner only anchors the full-row # `next_token_logits` / `hidden_states`; the selected-row `topk_p` / # `topk_index` are owned by EAGLEWorkerV2 and computed *after* replay (see - # `eagle_worker_v2._draft_extend_for_decode` and the early-return in - # `EAGLEDraftExtendCudaGraphRunner.replay` for DRAFT_EXTEND_V2). The V2 - # production runner output therefore carries no topk fields, so the - # runner-mode reference must only compare what the graph actually anchors. + # `eagle_worker_v2._draft_extend_for_decode`). The graph computes no topk + # and `EAGLEDraftExtendCudaGraphRunner.replay` returns no topk fields, so + # the runner-mode reference must only compare what the graph anchors. torch.testing.assert_close( actual.next_token_logits, expected.next_token_logits, @@ -445,7 +421,7 @@ class EagleDraftExtendCudaGraphRunnerAdapter: lambda _case, _settings: None ) assert_outputs_close: Callable[[Any, Any, EagleDraftRunnerSettings], None] = ( - _assert_draft_extend_outputs_close + _assert_draft_extend_v2_outputs_close ) @@ -604,19 +580,6 @@ def _run_eagle_draft_extend_eager( batch.positions, batch, ) - # Mirror the production fast path from - # EAGLEDraftExtendCudaGraphRunner.replay (#26397): when topk == 1 - # production skips the full-vocab softmax and returns - # `topk_p = ones_like(topk_index)` (the value is unused downstream). - # The eager reference must match this for assert_outputs_close. - from sglang.srt.utils import is_hip - - if settings.topk == 1 and not is_hip(): - ret.topk_index = torch.argmax(ret.next_token_logits, dim=-1, keepdim=True) - ret.topk_p = torch.ones_like(ret.topk_index, dtype=torch.float32) - else: - probs = torch.softmax(ret.next_token_logits, dim=-1) - ret.topk_p, ret.topk_index = fast_topk(probs, settings.topk, dim=-1) return ret