From 118465f5b5e3b2283d9256103fc187689a702c39 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Sun, 31 May 2026 20:29:50 -0700 Subject: [PATCH] [attn backend] Make spec_v2 seq_lens_cpu optional in trtllm_mla backend (#26824) --- .../layers/attention/trtllm_mla_backend.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 932df4306..b9d55bfe9 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -262,6 +262,10 @@ class TRTLLMMLADecodeMetadata: class TRTLLMMLABackend(FlashInferMLAAttnBackend): """TRTLLM MLA attention kernel from flashinfer.""" + # trtllm-gen kernels rebuild metadata from preallocated buffers and never + # read seq_lens_cpu / seq_lens_sum; opt out of the D2H sync. + needs_cpu_seq_lens: bool = False + def __init__( self, model_runner: ModelRunner, @@ -335,6 +339,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ) self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens + self.cuda_graph_custom_mask = None def _calc_padded_blocks(self, max_seq_len: int) -> int: """ @@ -442,8 +447,20 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): device=self.device, ) + if self.num_draft_tokens and not self.skip_prefill: + # Worst-case FULL_MASK tree-mask scratch (bool); build_tree writes it + # in-place so the gpu_only path needs no seq_lens_sum. + self.cuda_graph_custom_mask = torch.zeros( + max_num_tokens * (self.max_context_len + self.num_draft_tokens), + dtype=torch.bool, + device=self.device, + ) + super().init_cuda_graph_state(max_bs, max_num_tokens, kv_indices_buf) + def get_verify_buffers_to_fill_after_draft(self): + return [self.cuda_graph_custom_mask, None] + def _init_cuda_graph_metadata( self, bs: int, @@ -1241,6 +1258,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): class TRTLLMMLAMultiStepDraftBackend(FlashInferMLAMultiStepDraftBackend): """Multi-step draft backend for TRT-LLM MLA used by EAGLE.""" + # Per-step draft decode never reads seq_lens_cpu / seq_lens_sum; opt out so + # decide_needs_cpu_seq_lens' OR over the backends stays False. + needs_cpu_seq_lens: bool = False + def __init__( self, model_runner: "ModelRunner",