From 8240f3603ebd4453f7c69bc2cbb9b53b5bd0ce03 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:10:13 -0700 Subject: [PATCH] [https://nvbugs/6669206][fix] Keep Kimi K3 H=96 MLA decode on cute-dsl for every batch shape trtllm-gen MLA decode rejects 64 < num_heads_q < 128 outright: its Q tile is 64 or 128 heads and neither divides such a head count once K3's padding to 128 heads was removed. K3's per-batch decode policy gated that correctness escape on num_contexts > 0, so it only covered mixed batches. A generation-only speculative-verification batch (num_contexts=0, num_gen_tokens != num_ generations) took the multi-token perf fallback and was downgraded to trtllm-gen, which then raised ValueError for num_heads_q=96. Gate the escape on the head count alone, expressed as trtllm-gen's real constraint, so neither perf-motivated fallback can target an illegal head count. Both fallbacks still apply where trtllm-gen is legal (e.g. H=128). Removes the waiver the failure was parked under. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../modules/ATTENTION_DEVELOPER_GUIDE.md | 8 +++++--- .../kimi_k3_mla/kimi_k3_mla_attention.py | 20 ++++++++++--------- tests/integration/test_lists/waives.txt | 1 - .../modules/test_kimi_k3_mla_backend.py | 10 ++++++++-- 4 files changed, 24 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md index 618ac27fd0e6..5074f7034696 100644 --- a/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md @@ -337,9 +337,11 @@ cache (override with `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen`; other values are rejected at model build). FP8 KV cache forces `trtllm-gen`. K3 also installs a per-batch policy that falls back to `trtllm-gen` for mixed context/generation batches and multi-token generation (speculative -verification), keeping `cute-dsl` for plain one-token-per-request decode. A -mixed H=96 batch remains on `cute-dsl`: TRTLLM-Gen may select a 64-head Q tile, -which does not divide 96 after K3's head padding removal. +verification), keeping `cute-dsl` for plain one-token-per-request decode. Both +fallbacks are perf tuning, so neither applies when `64 < num_heads < 128` +(e.g. K3's H=96 after its padding to 128 heads was removed): TRTLLM-Gen's MLA +decode rejects those head counts outright, since its Q tile is 64 or 128 heads +and neither divides them. The FMHA package is split by role: diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index 889ac18c7efa..411550471740 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -77,19 +77,21 @@ def _kimi_k3_mla_decode_backend_policy( CuTe-DSL reuses one staged page table across MLA layers for a generation-only, one-token-per-request batch. Other mixed batches repeat the staging copies in every MLA layer and regress time to first token, so - they fall back to TRTLLM-Gen. The H=96 path is the correctness exception: - TRTLLM-Gen may select a 64-head Q tile, which does not divide 96 and - produces an invalid configuration after K3's head padding was removed. - - The CuTe-DSL kernel itself accepts multi-token queries, but K3's decode - tuning covers only the one-token-per-request regime, so generation-only - speculative verification also falls back. + they fall back to TRTLLM-Gen. The CuTe-DSL kernel itself accepts + multi-token queries, but K3's decode tuning covers only the + one-token-per-request regime, so generation-only speculative verification + also falls back. + + Both fallbacks are perf tuning and neither may override correctness: + TRTLLM-Gen's MLA decode rejects ``64 < num_heads < 128`` outright, because + its Q tile is 64 or 128 heads and neither divides such a head count once + K3's padding to 128 heads was removed. """ is_single_token_generation = num_gen_tokens == metadata.num_generations - requires_cute_dsl_for_mixed_batch = metadata.num_contexts > 0 and num_heads == 96 + trtllm_gen_supports_num_heads = not (64 < num_heads < 128) if ( requested_backend == "cute-dsl" - and not requires_cute_dsl_for_mixed_batch + and trtllm_gen_supports_num_heads and (metadata.num_contexts > 0 or not is_single_token_generation) ): return "trtllm-gen" diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 03282837778b..e582cdb30cbf 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -356,5 +356,4 @@ unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_pp2[False-Fals unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504) unittest/llmapi/test_memory_profiling.py::test_profile_kvcache SKIP (https://nvbugs/5580781) -unittest/tools/test_layer_wise_benchmarks.py::test_kimi_k3_gen_dep[1] SKIP (https://nvbugs/6669206) verl/test_verl_cases.py::test_trtllm_abort SKIP (https://nvbugs/6272653) diff --git a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py index 937daeabb2f9..2c73e4bab228 100644 --- a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py +++ b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py @@ -66,7 +66,13 @@ def test_select_kimi_k3_mla_generation_backend_uses_trtllm_gen_for_fp8_kv_cache( ("cute-dsl", 0, 4, 4, 96, "cute-dsl"), ("cute-dsl", 1, 3, 3, 12, "trtllm-gen"), ("cute-dsl", 1, 3, 3, 96, "cute-dsl"), - ("cute-dsl", 0, 4, 8, 96, "trtllm-gen"), + # Generation-only speculative verification at H=96, the shape the + # layer-wise benchmark hits in https://nvbugs/6669206: trtllm-gen + # rejects 64 < num_heads < 128, so the multi-token perf fallback + # must not fire. + ("cute-dsl", 0, 4, 8, 96, "cute-dsl"), + # H=128 divides trtllm-gen's Q tile, so the fallback still applies. + ("cute-dsl", 0, 4, 8, 128, "trtllm-gen"), ("trtllm-gen", 1, 3, 3, 96, "trtllm-gen"), ], ) @@ -78,7 +84,7 @@ def test_kimi_k3_mla_decode_backend_policy_by_batch_shape( num_heads: int, expected_backend: str, ) -> None: - """K3 falls back outside plain decode except for unsafe H=96 mixed batches.""" + """K3 falls back outside plain decode, except where trtllm-gen cannot run.""" assert ( _kimi_k3_mla_decode_backend_policy( requested_backend,