[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(