From 131bd51b01d96de247fbd1966fb45bdbecbb80b5 Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Sat, 1 Aug 2026 16:16:04 -0700 Subject: [PATCH] [Spec] Add `trtllm_mha` support for Gemma 4 MTP draft attention backend (#25545) Co-authored-by: Cursor Co-authored-by: Claude Opus 5 (1M context) --- python/sglang/srt/arg_groups/overrides.py | 1 + .../frozen_kv_mtp_cuda_graph_runner.py | 4 +++- .../speculative/frozen_kv_mtp_worker_v2.py | 19 +++++++++++++------ .../unittests/dense/test_trtllm_mha.py | 5 ++--- 4 files changed, 19 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 69bc7a7cc..f1290b9c6 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/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 diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 3fdf45e77..df7fab33c 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/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() ) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 98c210b4e..2e251bfc1 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/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( diff --git a/test/registered/attention/unittests/dense/test_trtllm_mha.py b/test/registered/attention/unittests/dense/test_trtllm_mha.py index 7c2a8413e..ac65a048a 100644 --- a/test/registered/attention/unittests/dense/test_trtllm_mha.py +++ b/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",