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

[trtllm-agent]

Summary

• Root cause: The H=96 correctness escape in K3's MLA decode backend policy was gated on metadata.num_contexts > 0, so generation-only multi-token (spec-verification) batches took the perf fallback to trtllm-gen, whose MLA decode rejects 64 < num_heads_q < 128.
• Fix: Gate the escape on the head count alone via trtllm-gen's real constraint (not 64 < num_heads < 128), so neither perf fallback can select an unsupported head count, while both still apply where trtllm-gen is legal (H=128 control); updated the docstring, developer guide, unit-test truth table, and removed the bug's waiver.
• Original test: pytest tests/unittest/tools/test_layer_wise_benchmarks.py::test_kimi_k3_gen_dep -k "1" -v
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6669206

Dev Engineer Review

• The decode fallback now rejects trtllm-gen for 65 <= num_heads <= 127.
• This fixes the unsupported 96-head generation-only speculative-verification path.
• The 128-head fallback remains on trtllm-gen.
• The developer guide, docstring, and unit-test truth table match the new policy.
• The waiver removal has the correct format and targets the affected test.
• No API or unrelated configuration changes were identified.

QA Engineer Review

• Modified the Kimi K3 MLA backend-selection unit test.
• Added coverage for generation-only speculative verification with 96 and 128 heads.
• The test is no longer waived in tests/integration/test_lists/waives.txt, so it is covered by CI.
• Verdict: sufficient.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 826ececc-1c2e-4917-bec5-20f510248a56

📥 Commits

Reviewing files that changed from the base of the PR and between 528a530 and 8240f36.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
• tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py
• tests/integration/test_lists/waives.txt
• tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

Kimi K3 MLA decode backend selection now avoids TRTLLM-Gen for head counts from 65 through 127. Tests verify H=96 uses CuTe DSL and H=128 uses TRTLLM-Gen. Related guidance and test waivers are updated.

Changes

Kimi K3 MLA backend selection

Layer / File(s)	Summary
Update decode fallback policy tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py, tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md	The decode policy and guidance exclude head counts from 65 through 127 from TRTLLM-Gen selection.
Validate supported head counts tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py, tests/integration/test_lists/waives.txt	Tests verify H=96 uses cute-dsl and H=128 uses trtllm-gen. The related skip waiver is removed.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 8240f

This localized change adjusts the K3 MLA backend selection gate and updates its tests and documentation; no actionable merge-blocking risk remains after normal checks and review.

Suggested reviewers: juney-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title identifies the NVBugs fix, change type, and main correction to the K3 MLA backend escape. It is specific and related to the primary change.
Description check	✅ Passed	The description explains the root cause and fix, identifies affected behavior, lists relevant testing, and links the bug. It omits the template's full checklist but remains sufficiently complete.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 2 files. (1 skipped: 1 …
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 2 files. (1 skipped: 1 unsupported.)

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[brnguyen2]

Superseded. The underlying fix for this test landed in #18226 (keep CuTe-DSL MLA decode for Kimi K3 H=96 speculative-verify batches). The only remaining piece, removing the test waiver, is handled by #18371, which has passed the single-GPU B200 CI stage. Closing this as superseded.