From 5af02c18ae706c1e421d65b4a253f46fc24b3406 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Thu, 4 Jun 2026 17:50:12 -0700 Subject: [PATCH] [spec_v2] Enable trtllm_mha draft-extend CUDA graph with v2 semantics (#25002) Co-authored-by: Lianmin Zheng Co-authored-by: Claude Opus 4.8 (1M context) --- .../layers/attention/trtllm_mha_backend.py | 38 ++++++++++++++----- .../eagle_draft_extend_cuda_graph_runner.py | 5 ++- .../sglang/srt/speculative/eagle_worker_v2.py | 15 ++++++-- .../speculative_draft_extend_runner.py | 24 ++++++++++++ 4 files changed, 67 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index de3ca31ad..93a028d20 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/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 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 422c2b26e..dc5f2505e 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 @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 2ba072dcd..a13cb890b 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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. 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 dfcaaa5c7..1efefeff5 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 @@ -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,