[None][fix] Keep CuTe-DSL MLA decode for Kimi K3 H=96 speculative-verify batches

[moraxu]

Dev Engineer Review

• Updated Kimi K3 MLA backend selection.
• Keeps cute-dsl for all configurations with 96 query heads per rank.
• Preserves the trtllm-gen fallback for smaller per-rank head counts.
• Updated the developer guide to document the H=96 restriction and TRTLLM-Gen incompatibility.
• No public API changes.
• Reported validation passed 12 unit tests.
• B200 speculative-decoding logits parity passed with zero drift.
• GB300 4-GPU speculative-decoding parity coverage passed.

QA Engineer Review

• Modified tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py.
• Updated backend-policy coverage for generation-only multi-token batches:
  • H=96 selects cute-dsl.
  • H=12 selects trtllm-gen.
• No tests/integration/test_lists/ coverage entry was reported.
• Verdict: needs follow-up because CI or manual QA test-list coverage is not confirmed.

Description

Stack trace: L0_MergeRequest_PR/56537 → GB300-4_GPUs-PyTorch-Post-Merge-1 → test_kimi_k3_specdec → test_kimi_k3_sa_specdec_logits_parity

Blocks: #17921 — the 4-GPU SA spec-dec logits-parity test wired into the GB300 post-merge stage there fails deterministically with this error (builds 56259 and 56537).

Since the recent MLA decode backend changes on main, Kimi K3 with suffix-automaton (SA) speculative decoding + attention-DP fails deterministically at engine initialization on every rank with:

ValueError: trtllm-gen MLA decode does not support 64 < num_heads_q < 128; got num_heads_q=96.
Use backend='cute-dsl' instead when the remaining configuration is CuTeDSL-compatible.

The non-speculative (baseline) engine with the identical configuration initializes and generates correctly. Only the SA spec-dec engine fails.

Root cause

Interaction of two recent main commits:

1. f51e323 — [TRTLLM-15033][feat] Upstream Kimi K3 MLA decode backend selection to main #17800 introduced the per-batch backend policy _kimi_k3_mla_decode_backend_policy() in tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py.
2. b057f77 — [https://nvbugs/6617948][fix] Restore trtllm-gen MLA decode perf gate dropped by #15300 #18054 restored the trtllm-gen MLA decode gate that raises the ValueError above for 64 < num_heads_q < 128.

The policy kept the requested cute-dsl backend for (a) generation-only single-token batches and (b) mixed batches with num_heads == 96 (its explicit "H=96 correctness exception"). However, for generation-only multi-token batches — which is exactly what SA speculative verification produces — it fell back to trtllm-gen even at num_heads == 96:

is_single_token_generation = num_gen_tokens == metadata.num_generations
requires_cute_dsl_for_mixed_batch = metadata.num_contexts > 0 and num_heads == 96
if (requested_backend == "cute-dsl"
        and not requires_cute_dsl_for_mixed_batch
        and (metadata.num_contexts > 0 or not is_single_token_generation)):
    return "trtllm-gen"

With attention-DP each rank keeps the full 96 query heads, so the trtllm-gen path raises the restored gate error at init. The H=96 correctness exception was simply never extended to the spec-verify case (the CuTe-DSL kernel itself accepts multi-token queries; the fallback is a decode-tuning-coverage preference, not a correctness requirement).

The TLLM_K3_MLA_GEN_BACKEND environment variable cannot work around this: the policy downgrades cute-dsl to trtllm-gen after the env value is read.

Fix

Extend the existing H=96 exception to every fallback candidate: any batch shape with num_heads == 96 stays on cute-dsl — the gate's own error message recommends exactly that — while smaller per-rank head counts keep the existing decode-tuning fallback to trtllm-gen. Also updates the policy-contract paragraph in ATTENTION_DEVELOPER_GUIDE.md and the unit-test matrix (one pre-existing matrix row enshrined the buggy expectation).

Impact if unfixed

• Kimi K3 SA speculative decoding is unusable with attention-DP on main (engine init failure); attention-DP (DEP16) is the primary qualified K3 deployment shape.
• Blocks [TRTLLM-15035][test] Wire Kimi K3 spec-dec and suffix-automaton tests into L0 CI #17921's GB300 post-merge stage and will break the weekly QA GSM8K sa leg being added in [TRTLLM-15036][test] Add Kimi K3 GSM8K/MMMU accuracy tests and register them in QA's weekly multinode list #17922 as soon as it runs on a main tree containing the two commits above.

Test Coverage

• tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py — policy matrix updated: generation-only multi-token at H=96 now expects cute-dsl; an H=12 row was added to prove smaller per-rank head counts keep the trtllm-gen fallback. 12/12 pass.
• A/B validation on a single B200 (SM100) with tests/integration/defs/kimi_k3_sa_harness.py at KIMI_K3_TP=1 with attention-DP (96 heads per rank), a 4-layer truncated Kimi-K3, KIMI_K3_SPEC_MODE=sa, logits-parity mode:
  • without this change: the baseline engine generates correctly, then the SA engine fails initialization with the exact ValueError above (harness exit 1) — same signature as the CI failures;
  • with this change: [sanity] PASS — spec-dec logits parity verified (4 prompts, 0 drift).
• End-to-end guard: tests/integration/defs/test_kimi_k3_specdec.py::test_kimi_k3_sa_specdec_logits_parity (the failing GB300 post-merge test from [TRTLLM-15035][test] Wire Kimi K3 spec-dec and suffix-automaton tests into L0 CI #17921) exercises this path on 4 GPUs.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

🤖 Generated with Claude Code

[moraxu]

/bot run --disable-fail-fast

[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: 7c3779d7-73b8-4269-92dd-750f890cc8a5

📥 Commits

Reviewing files that changed from the base of the PR and between b5875ec and fa1620d.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
• tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py
• tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py

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

---

Walkthrough

The Kimi K3 MLA backend now keeps H=96 configurations on CuTe-DSL across mixed, attention-DP, and multi-token speculative-verification batches. Tests and developer guidance document this policy.

Changes

Kimi K3 MLA backend policy

Layer / File(s)	Summary
Enforce CuTe-DSL for H=96 tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py	H=96 fallback candidates now use CuTe-DSL unconditionally. Other configurations retain TRTLLM-Gen fallback behavior where applicable.
Validate and document backend selection tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py, tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md	Tests cover multi-token generation for H=96 and H=12. The guide documents the incompatible TRTLLM-Gen tile and decode gate for H=96.

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

Merge Risk: ⚪ Minimal · up to fa162

The change is localized and has no actionable merge-blocking risk remaining 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 clearly and concisely describes the main fix: keeping CuTe-DSL MLA decode for Kimi K3 H=96 speculative-verification batches.
Description check	✅ Passed	The description explains the failure, root cause, fix, impact, test coverage, and checklist status. It matches the required template and provides specific validation results.
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.}

[tensorrt-cicd]

PR_Github #69336 [ run ] triggered by Bot. Commit: fa1620d Link to invocation

[tensorrt-cicd]

PR_Github #69336 [ run ] completed with state FAILURE. Commit: fa1620d
/LLM/main/L0_MergeRequest_PR pipeline #56682 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[moraxu]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #69499 [ run ] triggered by Bot. Commit: fa1620d Link to invocation

[tensorrt-cicd]

PR_Github #69499 [ run ] completed with state SUCCESS. Commit: fa1620d
/LLM/main/L0_MergeRequest_PR pipeline #56825 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[moraxu]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #69529 [ run ] triggered by Bot. Commit: fa1620d Link to invocation

[tensorrt-cicd]

PR_Github #69529 [ run ] completed with state SUCCESS. Commit: fa1620d
/LLM/main/L0_MergeRequest_PR pipeline #56852 completed with status: 'SUCCESS'

CI Report

Link to invocation

[pengbowang-nv]

The changes LGTM but I doubt there are some back-and-forth changes related to this affair. Please confirm with @SimengLiu-nv and @brnguyen2 . Thanks!