From a45e0e5df4e6e2d765afcaa00b16a67694f5170e Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Sat, 17 Jan 2026 23:56:21 -0800 Subject: [PATCH] [SPEC_V2] Enable cudagraph draft_extend for trtllm_mla_backend and Acclen Fix for DP under cudagraph mode (#16974) --- .../eagle_draft_extend_cuda_graph_runner.py | 18 +++++++++++++++--- .../sglang/srt/speculative/eagle_worker_v2.py | 7 +++++++ 2 files changed, 22 insertions(+), 3 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 f5a262ad9..e1afdd84b 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 @@ -253,6 +253,12 @@ class EAGLEDraftExtendCudaGraphRunner: : bs if self.forward_mode == ForwardMode.DRAFT_EXTEND else num_tokens ] + # V1 (DRAFT_EXTEND): pruned_states = bs (last token per seq) + # V2 (DRAFT_EXTEND_V2): pruned_states = num_tokens (all tokens) + num_tokens_for_logprob = ( + num_tokens if self.forward_mode.is_draft_extend_v2() else bs + ) + if self.require_mlp_tp_gather: self.global_num_tokens_gpu.copy_( torch.tensor( @@ -263,7 +269,7 @@ class EAGLEDraftExtendCudaGraphRunner: ) self.global_num_tokens_for_logprob_gpu.copy_( torch.tensor( - [bs] * self.dp_size, + [num_tokens_for_logprob] * self.dp_size, dtype=torch.int32, device=self.input_ids.device, ) @@ -279,7 +285,7 @@ class EAGLEDraftExtendCudaGraphRunner: ) self.global_num_tokens_for_logprob_gpu.copy_( torch.tensor( - [bs], + [num_tokens_for_logprob], dtype=torch.int32, device=self.input_ids.device, ) @@ -419,7 +425,13 @@ class EAGLEDraftExtendCudaGraphRunner: # TODO(ch-wan): support num_token_non_padded if self.require_gathered_buffer: self.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs) - self.global_num_tokens_for_logprob_gpu.fill_(bs) + # V1: pruned_states = bs; V2: pruned_states = num_tokens + if self.forward_mode.is_draft_extend_v2(): + self.global_num_tokens_for_logprob_gpu.fill_( + bs * self.num_tokens_per_bs + ) + else: + self.global_num_tokens_for_logprob_gpu.fill_(bs) if forward_batch.seq_lens_cpu is not None: if bs != raw_bs: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index d3dcc2afd..92f7886da 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -13,6 +13,9 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner i EAGLEDraftNpuGraphRunner, ) from sglang.srt.layers.attention.triton_backend import TritonMultiStepDraftBackend +from sglang.srt.layers.attention.trtllm_mla_backend import ( + TRTLLMMLAMultiStepDraftBackend, +) from sglang.srt.layers.dp_attention import get_attention_tp_group from sglang.srt.layers.moe.utils import ( speculative_moe_a2a_backend_context, @@ -277,6 +280,10 @@ class EagleDraftWorker(BaseDraftWorker): _is_cuda and isinstance(self.draft_attn_backend, TritonMultiStepDraftBackend) ) + or ( + _is_cuda + and isinstance(self.draft_attn_backend, TRTLLMMLAMultiStepDraftBackend) + ) ): tic = time.perf_counter() before_mem = get_available_gpu_memory(self.device, self.gpu_id)