[attn backend] Make spec_v2 seq_lens_cpu optional in trtllm_mla backend (#26824)

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",