[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:
co-authored by
Cursor
Claude Opus 5
parent
e2cf21b9e5
commit
131bd51b01
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user