[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>
This commit is contained in:
Khoa Pham
2026-08-01 16:16:04 -07:00
committed by GitHub
co-authored by Cursor Claude Opus 5
parent e2cf21b9e5
commit 131bd51b01
4 changed files with 19 additions and 10 deletions
@@ -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",