[None][fix] Unwaive supported attention backend cases by yuxianq · Pull Request #15838 · NVIDIA/TensorRT-LLM

yuxianq · 2026-07-01T13:13:03Z

@coderabbitai summary

Description

This is a follow-up to #15536 that carries only the attention-backend changes
not included in that PR's squash merge.

Exercise TRTLLM MLA context through the standalone backend harness using the
production input contract, including fused RoPE and latent-cache validation.
Enable the available Blackwell TRTLLM-Gen MLA decode shape and adapt to the
current KV-cache metadata return contract.
Remove the obsolete SM90 paged-context FMHA workaround and its dependent
standalone-harness exclusion now that the underlying issue is fixed.
Remove capability exclusions for TRTLLM cross-attention, FP8 generation, and
Blackwell sliding-window decode cases that now pass the full suite.

Test Coverage

B200: tests/unittest/_torch/attention/test_attention_backends.py
- 935 passed in 376.14s (395s wall time)
H200 NVL, clean SM90 build: tests/unittest/_torch/attention/test_attention_backends.py
- 935 passed in 78.43s (101s wall time)

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.

coderabbitai · 2026-07-01T13:21:20Z

📝 Walkthrough

Walkthrough

Removes a missing-kernel gating mechanism from trtllm-gen MLA generation support, adjusts KV cache metadata unpacking, narrows an SM90 workaround to context-request cases in TrtllmAttention, and updates test harness capability gating and MLA context test input/RoPE generation logic accordingly.

Changes

MLA generation support and attention forward changes

Layer / File(s)	Summary
Remove missing-kernel gating in MLA generation support `tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py`	Deletes `MISSING_MLA_GENERATION_KERNELS` and its check in `_check_mla_generation_support`; `run_mla_generation` now unpacks an extra unused value from the KV cache metadata builder.
Narrow SM90 paged-context workaround `tensorrt_llm/_torch/attention_backend/trtllm.py`	SM90 workaround forcing `use_paged_context_fmha=True` now only applies when `metadata.num_contexts > 0`.
Update backend capability skip rules `tests/unittest/_torch/attention/backend_capability.py`	Removes TRTLLM Blackwell sliding-window pure-decode skip list/branch; adds a new SM90 FP8-KV mixed-phase skip for equal head counts.
Rework MLA context test input generation and forward invocation `tests/unittest/_torch/attention/backend_case.py`	Reworks `generate_mla_context_inputs` to build q/k/v, compressed_kv, k_pe, latent_cache; adds `_prepare_mla_context_inputs` for RoPE handling; updates `create_attention`, `_run_mla_context_backend`, forward invocation, cache tolerance checks, and cross-attention `prompt_lens`.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Suggested reviewers: QiJune, PerkzZheng, dongfengy

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows the required bracketed format and matches the main change of unwaiving supported attention backend cases.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Description check	✅ Passed	The PR description follows the template well and includes the issue, solution, test coverage, and checklist acknowledgment.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/attention/backend_capability.py`:
- Around line 187-200: The SM90/TRTLLM skip in backend_capability is too broad
because it catches no-cache cases as well as paged-cache cases. Tighten the
unsupported_reason() predicate by adding a cache check (only skip when the case
uses paged cache) while keeping the existing SM90, fp8 KV, head-shape, and
mixed-context conditions in place. Use the existing symbols unsupported_reason,
case.kv_dtype, case.num_heads, case.num_kv_heads, and case.num_contexts to
locate and adjust the gate.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ac2f1836-0245-44ae-abd9-8d996b7f1a4b

📥 Commits

Reviewing files that changed from the base of the PR and between d567f92 and e97de04.

📒 Files selected for processing (4)

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
tensorrt_llm/_torch/attention_backend/trtllm.py
tests/unittest/_torch/attention/backend_capability.py
tests/unittest/_torch/attention/backend_case.py

yuxianq · 2026-07-01T13:41:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-01T13:47:44Z

PR_Github #56940 [ run ] triggered by Bot. Commit: 9a32abb Link to invocation

yuxianq · 2026-07-01T14:59:46Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-01T15:08:09Z

PR_Github #56952 [ run ] triggered by Bot. Commit: 0afb7e9 Link to invocation

tensorrt-cicd · 2026-07-01T15:12:34Z

PR_Github #56940 [ run ] completed with state ABORTED. Commit: 9a32abb

Link to invocation

tensorrt-cicd · 2026-07-01T23:34:40Z

PR_Github #56952 [ run ] completed with state SUCCESS. Commit: 0afb7e9
/LLM/main/L0_MergeRequest_PR pipeline #45756 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

yuxianq · 2026-07-02T01:31:23Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-02T01:39:10Z

PR_Github #57056 [ run ] triggered by Bot. Commit: 73e9402 Link to invocation

yihwang-nv

LGTM

tensorrt-cicd · 2026-07-02T06:03:26Z

PR_Github #57056 [ run ] completed with state SUCCESS. Commit: 73e9402
/LLM/main/L0_MergeRequest_PR pipeline #45849 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

yuxianq · 2026-07-02T10:01:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-03T02:46:51Z

PR_Github #57322 [ run ] triggered by Bot. Commit: 4f25a1d Link to invocation

tensorrt-cicd · 2026-07-03T06:21:32Z

PR_Github #57322 [ run ] completed with state FAILURE. Commit: 4f25a1d
/LLM/main/L0_MergeRequest_PR pipeline #46079 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

yuxianq · 2026-07-03T06:28:24Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-03T06:34:54Z

PR_Github #57384 [ run ] triggered by Bot. Commit: ab64fff Link to invocation

tensorrt-cicd · 2026-07-03T13:07:32Z

PR_Github #57384 [ run ] completed with state SUCCESS. Commit: ab64fff
/LLM/main/L0_MergeRequest_PR pipeline #46134 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

yuxianq · 2026-07-06T07:09:32Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T07:15:18Z

PR_Github #57719 [ run ] triggered by Bot. Commit: e9e0179 Link to invocation

tensorrt-cicd · 2026-07-06T10:46:12Z

PR_Github #57719 [ run ] completed with state SUCCESS. Commit: e9e0179
/LLM/main/L0_MergeRequest_PR pipeline #46429 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

yuxianq · 2026-07-07T02:49:14Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T02:55:43Z

PR_Github #57905 [ run ] triggered by Bot. Commit: 969a6cc Link to invocation

tensorrt-cicd · 2026-07-07T13:40:21Z

PR_Github #57972 [ run ] completed with state SUCCESS. Commit: d6fec9b
/LLM/main/L0_MergeRequest_PR pipeline #46646 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

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq · 2026-07-07T14:41:51Z

/bot run --disable-fail-fast

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

tensorrt-cicd · 2026-07-07T14:48:16Z

PR_Github #58019 [ run ] triggered by Bot. Commit: 3a40842 Link to invocation

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq · 2026-07-07T15:17:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T15:24:15Z

PR_Github #58025 [ run ] triggered by Bot. Commit: 7483843 Link to invocation

tensorrt-cicd · 2026-07-07T15:28:42Z

PR_Github #58019 [ run ] completed with state ABORTED. Commit: 3a40842

Link to invocation

tensorrt-cicd · 2026-07-07T18:57:40Z

PR_Github #58025 [ run ] completed with state SUCCESS. Commit: 7483843
/LLM/main/L0_MergeRequest_PR pipeline #46695 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

yuxianq · 2026-07-08T02:08:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-08T02:14:14Z

PR_Github #58112 [ run ] triggered by Bot. Commit: 7483843 Link to invocation

yuxianq · 2026-07-08T02:23:11Z

/bot skip --comment "All failed tests are pre-existed and unrelated to the attention backend tests in this PR, skip CI"

tensorrt-cicd · 2026-07-08T02:31:23Z

PR_Github #58119 [ skip ] triggered by Bot. Commit: 7483843 Link to invocation

tensorrt-cicd · 2026-07-08T02:34:58Z

PR_Github #58112 [ run ] completed with state ABORTED. Commit: 7483843

Link to invocation

tensorrt-cicd · 2026-07-08T02:38:44Z

PR_Github #58119 [ skip ] completed with state SUCCESS. Commit: 7483843
Skipping testing for commit 7483843

Link to invocation

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq requested a review from a team as a code owner July 1, 2026 13:13

yuxianq requested a review from brb-nv July 1, 2026 13:13

github-actions Bot assigned yuxianq Jul 1, 2026

yuxianq requested a review from yihwang-nv July 1, 2026 13:14

coderabbitai Bot reviewed Jul 1, 2026

View reviewed changes

Comment thread tests/unittest/_torch/attention/backend_capability.py Outdated

yihwang-nv approved these changes Jul 2, 2026

View reviewed changes

yuxianq force-pushed the unwaive-attn-backend-test-cases branch from ab64fff to e9e0179 Compare July 6, 2026 06:43

yuxianq force-pushed the unwaive-attn-backend-test-cases branch 2 times, most recently from 60a9660 to d6fec9b Compare July 7, 2026 06:59

yuxianq added 11 commits July 7, 2026 14:28

[None][fix] enable TRTLLM MLA context backend tests

bb93032

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] enable TRTLLM Blackwell MLA generation

ccf8516

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] reduce TRTLLM attention backend exclusions

3e2a315

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] unwaive stable attention backend cases

ca66d0c

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] remove obsolete Hopper FMHA workaround

4f259ca

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] support large FlashInfer KV appends

52d1b18

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] restore missing MLA generation kernel guard

96f9813

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] support FP8 KV MLA in TRTLLM-gen

aadd5a4

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] fall back for FP8 multi-token MLA

ec49210

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] gate FP8 multi-token MLA fallback by FlashInfer version

f497c9c

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] remove obsolete FlashInfer MLA fallback

a385754

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq force-pushed the unwaive-attn-backend-test-cases branch from d6fec9b to 3a40842 Compare July 7, 2026 14:41

[None][fix] fix FP8 KV-only attention fallback

3a40842

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

[None][fix] avoid forcing paged context FMHA for FP8 KV

7483843

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq merged commit 4860595 into NVIDIA:main Jul 8, 2026
7 checks passed

BrianLi23 pushed a commit to BrianLi23/TensorRT-LLM that referenced this pull request Jul 9, 2026

[None][fix] Unwaive supported attention backend cases (NVIDIA#15838)

608eb72

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

chenfeiz0326 mentioned this pull request Jul 9, 2026

[https://nvbugs/6430674][fix] Surgical revert of PR #15838's dispatch guard removal: re-add… #16167

Open

2 tasks

Uh oh!

Conversation

yuxianq commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yuxianq commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

yuxianq commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

yuxianq commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

yihwang-nv left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

yuxianq commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

yuxianq commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

yuxianq commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

yuxianq commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

yuxianq commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

yuxianq commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

yuxianq commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

yuxianq commented Jul 8, 2026

yuxianq commented Jul 1, 2026 •

edited

Loading

coderabbitai Bot commented Jul 1, 2026 •

edited

Loading