[https://nvbugs/6669206][fix] Gate the escape on the head count alone via trtllm-gen's real constraint (`not…

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:
 

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"

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)

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,