[Spec] Add trtllm_mha support for Gemma 4 MTP draft attention backend (#25545)

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/srt/arg_groups/overrides.py

@@ -1841,6 +1841,7 @@ def _mla_backend_page_constraints(view: Any) -> dict:
         view.attention_backend == "trtllm_mha"
         or view.decode_attention_backend == "trtllm_mha"
         or view.prefill_attention_backend == "trtllm_mha"
+        or view.speculative_draft_attention_backend == "trtllm_mha"
     ):
         # 128 runs on trtllm-gen's dynamic tokens-per-page kernels (flashinfer
         # >= 0.6.12), which require GQA and equal QK/V head dims — validated at

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -125,7 +125,9 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
         self.max_bs = max(self.capture_bs)
         self.max_num_token = self.max_bs * self.captured_req_width
 
-        self.draft_attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token)
+        self.draft_attn_backend.init_cuda_graph_state(
+            self.max_bs * self.topk, self.max_num_token
+        )
         self.seq_len_fill_value = (
             self.draft_attn_backend.get_cuda_graph_seq_len_fill_value()
         )

python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py

@@ -233,12 +233,14 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
             return self.draft_model_runner.attn_backend
 
         backend_type = self._resolve_draft_backend_type()
-        if backend_type != "triton":
-            raise ValueError(
-                "Frozen-KV MTP topk > 1 currently supports only the triton "
-                f"attention backend, got {backend_type}."
-            )
-        return self._init_triton_draft_attn_backend()
+        if backend_type == "triton":
+            return self._init_triton_draft_attn_backend()
+        if backend_type == "trtllm_mha":
+            return self._init_trtllm_mha_draft_attn_backend()
+        raise ValueError(
+            "Frozen-KV MTP topk > 1 currently supports triton and trtllm_mha "
+            f"attention backends, got {backend_type}."
+        )
 
     def _init_triton_draft_attn_backend(self):
         from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
@@ -253,6 +255,11 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
             kv_indptr_buf=kv_indptr_buf,
         )
 
+    def _init_trtllm_mha_draft_attn_backend(self):
+        from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend
+
+        return TRTLLMHAAttnBackend(self.draft_model_runner, skip_prefill=True)
+
     def _bind_kv_context(self) -> None:
         draft_model = self.draft_model_runner.model
         if not hasattr(draft_model, "build_frozen_kv_mtp_context") or not hasattr(

test/registered/attention/unittests/dense/test_trtllm_mha.py

@@ -36,7 +36,7 @@ register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
 @unittest.skipIf(
     not torch.cuda.is_available()
     or not is_flashinfer_available()
-    or not (is_sm90_supported() or is_sm120_supported()),
+    or not (is_sm90_supported() or is_sm100_supported() or is_sm120_supported()),
     "CUDA + FlashInfer TRT-LLM MHA decode support are required",
 )
 class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase):
@@ -149,8 +149,7 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase):
         ),
     )
 
-    # Frozen-KV MTP draft CG runner (chain, topk=1) — records the fused
-    # in-graph metadata rebuild inside FrozenKVMTPCudaGraphRunner's capture.
+    # Frozen-KV MTP draft CG runner (chain, topk=1).
     FROZEN_KV_MTP_RUNNER_CASES = (
         DenseAttentionCase(
             name="runner_frozen_kv_mtp_decode_trtllm_mha_cuda_graph",