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,