Remove unused draft-extend CUDA graph top-k (#31430)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>

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 (
+from sglang.srt.speculative.spec_utils import resolve_num_tokens_per_req
-    fast_topk,
-    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

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
+    # `eagle_worker_v2._draft_extend_for_decode`). The graph computes no topk
-    # `EAGLEDraftExtendCudaGraphRunner.replay` for DRAFT_EXTEND_V2). The V2
+    # and `EAGLEDraftExtendCudaGraphRunner.replay` returns no topk fields, so
-    # production runner output therefore carries no topk fields, so the
+    # the runner-mode reference must only compare what the graph anchors.
-    # runner-mode reference must only compare what the graph actually 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