[spec_v2] Enable trtllm_mha draft-extend CUDA graph with v2 semantics (#25002)

Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -391,7 +391,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                 bs,
             )
             self.target_verify_metadata[bs] = metadata
-        elif forward_mode.is_draft_extend():
+        elif forward_mode.is_draft_extend(include_v2=True):
             num_tokens_per_bs = num_tokens // bs
             metadata.cache_seqlens_int32 = self.draft_extend_metadata["cache_seqlens"][
                 :bs
@@ -484,7 +484,8 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             ]
             metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size)
             self._copy_swa_page_table(metadata, page_indices, max_seq_pages)
-        elif forward_mode.is_draft_extend():
+            metadata.max_seq_len_q = self.speculative_num_draft_tokens
+        elif forward_mode.is_draft_extend(include_v2=True):
             metadata = self.draft_extend_metadata[bs]
             metadata.cache_seqlens_int32.copy_(seq_lens)
 
@@ -493,15 +494,34 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             metadata.cu_seqlens_k[1:].copy_(
                 torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
             )
-            extend_lens = spec_info.num_accept_tokens[:bs]
-            if spec_info.num_accept_tokens_cpu:
-                metadata.max_seq_len_q = max(spec_info.num_accept_tokens_cpu)
+            if forward_mode.is_draft_extend_v2():
+                num_tokens_per_bs = spec_info.num_tokens_per_req
+                if num_tokens_per_bs <= 0:
+                    # Capture uses a synthetic EagleDraftExtendInput; infer the
+                    # fixed V2 stride from the capture buffer when it is unset.
+                    num_tokens_per_bs = int(
+                        spec_info.num_accept_tokens[:bs].max().item()
+                    )
+                metadata.max_seq_len_q = num_tokens_per_bs
+                metadata.cu_seqlens_q[1:].copy_(
+                    torch.arange(
+                        num_tokens_per_bs,
+                        bs * num_tokens_per_bs + 1,
+                        num_tokens_per_bs,
+                        dtype=torch.int32,
+                        device=metadata.cu_seqlens_q.device,
+                    )
+                )
             else:
-                metadata.max_seq_len_q = 1
+                extend_lens = spec_info.num_accept_tokens[:bs]
+                if spec_info.num_accept_tokens_cpu:
+                    metadata.max_seq_len_q = max(spec_info.num_accept_tokens_cpu)
+                else:
+                    metadata.max_seq_len_q = 1
 
-            metadata.cu_seqlens_q[1:].copy_(
-                torch.cumsum(extend_lens, dim=0, dtype=torch.int32)
-            )
+                metadata.cu_seqlens_q[1:].copy_(
+                    torch.cumsum(extend_lens, dim=0, dtype=torch.int32)
+                )
 
             max_seq_pages = (
                 metadata.max_seq_len_k + self.page_size - 1

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -572,6 +572,7 @@ class EAGLEDraftExtendCudaGraphRunner:
                 next_token_logits=out.next_token_logits[:unpadding_bs],
                 hidden_states=out.hidden_states[:unpadding_bs],
             )
-            out.topk_p = out_copy.topk_p[:unpadding_bs]
-            out.topk_index = out_copy.topk_index[:unpadding_bs]
+            if self.forward_mode != ForwardMode.DRAFT_EXTEND_V2:
+                out.topk_p = out_copy.topk_p[:raw_bs]
+                out.topk_index = out_copy.topk_index[:raw_bs]
         return out

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -16,6 +16,7 @@ from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGra
 from sglang.srt.kv_canary.runner.canary_manager import context_tuple
 from sglang.srt.layers.attention.tokenspeed_mla_backend import TokenspeedMLABackend
 from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
+from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend
 from sglang.srt.layers.attention.trtllm_mla_backend import (
     TRTLLMMLABackend,
 )
@@ -365,10 +366,14 @@ class EagleDraftWorker(BaseDraftWorker):
                 self.draft_attn_backend, AiterMultiStepDraftBackend
             )
 
-        supports_cuda_draft_extend_graph = (_is_cuda or _is_musa) and (
-            isinstance(self.draft_extend_attn_backend, TritonAttnBackend)
-            or isinstance(self.draft_extend_attn_backend, TRTLLMMLABackend)
-            or isinstance(self.draft_extend_attn_backend, TokenspeedMLABackend)
+        supports_cuda_draft_extend_graph = (_is_cuda or _is_musa) and isinstance(
+            self.draft_extend_attn_backend,
+            (
+                TritonAttnBackend,
+                TRTLLMMLABackend,
+                TRTLLMHAAttnBackend,
+                TokenspeedMLABackend,
+            ),
         )
         # Capture extend
         # TODO: support draft extend cuda graph for more attention backends
@@ -769,6 +774,8 @@ class EagleDraftWorker(BaseDraftWorker):
             draft_logits_output.hidden_states = draft_logits_output.hidden_states[
                 select_index
             ]
+        # The draft-extend graph only anchors full logits; selected-row topk is
+        # owned by the worker for both graph and eager paths.
         if self.topk == 1 and not _is_hip:
             # Gated to CUDA: see #26358 — ROCm's argmax tie-break corrupts
             # MTP draft selection on FP8 logits.

python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py

@@ -1013,6 +1013,28 @@ def _assert_draft_extend_outputs_close(actual, expected, settings) -> None:
     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.
+    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,
+    )
+
+
 @dataclass(frozen=True)
 class EagleDraftExtendCudaGraphRunnerAdapter:
     build_fixture: Callable[..., Any]
@@ -1554,6 +1576,7 @@ def run_dense_eagle_draft_extend_v2_cuda_graph_runner_case(
         make_draft_inputs=_make_dense_draft_extend_inputs,
         prepare_replay_state=_prepare_dense_draft_extend_replay_state,
         make_forward_batch=_make_dense_eagle_draft_extend_v2_forward_batch,
+        assert_outputs_close=_assert_draft_extend_v2_outputs_close,
     )
     run_eagle_draft_extend_cuda_graph_runner_case(
         testcase,
@@ -1757,6 +1780,7 @@ def run_mla_eagle_draft_extend_v2_cuda_graph_runner_case(
         make_draft_inputs=_make_mla_draft_extend_inputs,
         prepare_replay_state=_prepare_mla_draft_extend_replay_state,
         make_forward_batch=_make_mla_eagle_draft_extend_v2_forward_batch,
+        assert_outputs_close=_assert_draft_extend_v2_outputs_close,
     )
     run_eagle_draft_extend_cuda_graph_runner_case(
         testcase,