[https://nvbugs/6430674][fix] Surgical revert of PR #15838's dispatch guard removal: re-add… by chenfeiz0326 · Pull Request #16167 · NVIDIA/TensorRT-LLM

chenfeiz0326 · 2026-07-09T03:40:28Z

Summary

Root cause: PR [None][fix] Unwaive supported attention backend cases #15838 removed the MISSING_MLA_GENERATION_KERNELS guard in flashinfer_trtllm_gen.py; R1 MLA at tokens_per_block=32 on Blackwell/FP8-KV now dispatches to a slower kernel path (~6.7% regression).
Fix: Surgical revert of PR [None][fix] Unwaive supported attention backend cases #15838's dispatch guard removal: re-add MISSING_MLA_GENERATION_KERNELS={(576,512,32)} and restore tokens_per_block parameter on _check_mla_generation_support and its caller. Other PR [None][fix] Unwaive supported attention backend cases #15838 runtime changes are intentionally left intact.
Perf metric: output_token_throughput (higher-is-better)
Bad perf: 5184
Good perf: 5642
Perf after fix: 5624
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/6430674

Summary by CodeRabbit

Bug Fixes
- Improved support checking for MLA generation so unsupported decode-kernel combinations are rejected earlier with a clearer error message.
- Added coverage for an additional blocked configuration, helping prevent runtime failures for certain model shape and block-size combinations.

… guard for R1 MLA decode on Blackwell PR NVIDIA#15838 removed the `MISSING_MLA_GENERATION_KERNELS={(576,512,32)}` guard and the `tokens_per_block` parameter from `_check_mla_generation_support` in `flashinfer_trtllm_gen.py`. Post-PR the check returns True for DeepSeek R1 MLA (head_dim_qk=576, head_dim_v=512, tokens_per_block=32) on SM100, so the dispatcher sends R1 MLA decode through `FlashInferTrtllmGenFmha.run_mla_generation` instead of the prior FMHA fallback. With `kv_cache_config.dtype: fp8` the newly-added FP8-KV branch runs, and this kernel + code path is ~6.7% slower on this specific R1 FP4 / FP8-KV shape. This is a surgical revert of just the dispatch guard removal from PR NVIDIA#15838. The other PR runtime changes (FP8-KV branch in `run_mla_generation`, flashinfer.py CTA-split helper, trtllm.py SM90-workaround removals, fmha/fallback.py timestep allow-list) are intentionally left intact -- they are either correctness-required or unrelated to this bug's dispatch path. Changes (Python-only, one file): * Re-add `MISSING_MLA_GENERATION_KERNELS = {(576, 512, 32)}` class attribute on `FlashInferTrtllmGenFmha`. * Restore the `tokens_per_block` parameter on `_check_mla_generation_support` and gate the (head_dim_qk, head_dim_v, tokens_per_block) triple to return `(False, ...)` for the missing kernel shape. * Pass `tokens_per_block=tokens_per_block` at the single caller inside `_is_supported_with_reason`. Verified on OCI GB200 in the auto-perf-fix-pro workflow, 3-rep median: test: r1_fp4_v2_tep8_mtp3-con32_iter12_1k1k good median (0457051): 5641.50 tok/s bad median (8c1b230): 5184.46 tok/s ToT median (5ce293f): 5299.95 tok/s Fix median (this diff): 5623.98 tok/s (+3.90% over threshold 5413, +6.11% over ToT) Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>

coderabbitai · 2026-07-09T03:43:05Z

📝 Walkthrough

Walkthrough

Adds a class-level mapping of missing TRTLLM-GEN decode kernel configurations keyed by head dims and tokens_per_block. Updates the MLA generation support check to accept tokens_per_block and reject configurations found in this mapping, with the caller updated accordingly.

Changes

MLA Generation Support Validation

Layer / File(s)	Summary
Missing kernel table and rejection logic `tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py`	Adds `MISSING_MLA_GENERATION_KERNELS` set, extends `_check_mla_generation_support` to accept `tokens_per_block` and reject listed `(head_dim_qk, head_dim_v, tokens_per_block)` combinations with a specific reason, and updates the caller in `_is_supported_with_reason` to pass `tokens_per_block`.

Estimated code review effort: 1 (Trivial) | ~5 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific and matches the main change: restoring the dispatch guard for the MLA decode path.
Description check	✅ Passed	The description covers the issue, fix, perf impact, tests, and bug link, though it does not use the exact template headings.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py (1)
411-413: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider a brief comment documenting why this shape is blocked.

The static analysis hint flags this as a mutable class-attribute default (RUF012), but it follows the exact same unannotated-mutable-set convention as SUPPORTED_MLA_GENERATION_HEAD_DIMS right above it, so it's consistent with existing local style. Given this entry exists specifically to prevent a re-occurrence of the throughput regression from PR #15838, a short inline comment (e.g. referencing the missing TRTLLM-GEN decode kernel / regression) would help future contributors avoid accidentally removing it again.
📝 Optional comment addition
+    # TRTLLM-GEN lacks a decode kernel for this (headDimQk, headDimV, tokens_per_block)
+    # combination; dispatching here regressed R1 MLA decode throughput on SM100 (see `#15838`).
     MISSING_MLA_GENERATION_KERNELS = {
         (576, 512, 32),
     }
🤖 Prompt for 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.

In `@tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py` around
lines 411 - 413, Add a brief inline comment near MISSING_MLA_GENERATION_KERNELS
in flashinfer_trtllm_gen.py explaining that this blocked shape is intentional
because it lacks a TRTLLM-GEN decode kernel and is kept to prevent the
throughput regression from PR `#15838`. Keep the existing set definition in place
and document the rationale so future edits to the MISSING_MLA_GENERATION_KERNELS
guard are less likely to remove it accidentally.

🤖 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.

Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py`:
- Around line 411-413: Add a brief inline comment near
MISSING_MLA_GENERATION_KERNELS in flashinfer_trtllm_gen.py explaining that this
blocked shape is intentional because it lacks a TRTLLM-GEN decode kernel and is
kept to prevent the throughput regression from PR `#15838`. Keep the existing set
definition in place and document the rationale so future edits to the
MISSING_MLA_GENERATION_KERNELS guard are less likely to remove it accidentally.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8861c620-f792-4aa1-9c90-17a533b79239

📥 Commits

Reviewing files that changed from the base of the PR and between 60583d2 and f833634.

📒 Files selected for processing (1)

tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py

chenfeiz0326 · 2026-07-09T07:12:21Z

/bot run --disable-fail-fast --stage-list "*PerfSanity*"

tensorrt-cicd · 2026-07-09T07:19:38Z

PR_Github #58412 [ run ] triggered by Bot. Commit: f833634 Link to invocation

tensorrt-cicd · 2026-07-09T12:25:58Z

PR_Github #58412 [ run ] completed with state FAILURE. Commit: f833634
/LLM/main/L0_MergeRequest_PR pipeline #47030 (Partly Tested) 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

chenfeiz0326 · 2026-07-10T08:31:22Z

/bot run --disable-fail-fast --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-PerfSanity-Node2-GPU8-Post-Merge-1,GB200-8_GPUs-2_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU1-GEN1-NODE1-GPU4-Post-Merge-2"

tensorrt-cicd · 2026-07-10T08:37:56Z

PR_Github #58646 [ run ] triggered by Bot. Commit: f833634 Link to invocation

tensorrt-cicd · 2026-07-10T11:23:12Z

PR_Github #58646 [ run ] completed with state SUCCESS. Commit: f833634
/LLM/main/L0_MergeRequest_PR pipeline #47233 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

chenfeiz0326 · 2026-07-11T09:52:16Z

From pre-merge test, the perf looks good.

…o SLOWER (review feedback) Address yuxianq's review on PR NVIDIA#16167: the (576, 512, 32) MLA generation kernel is not missing, just slower — reflect that in the identifier and the error string. - MISSING_MLA_GENERATION_KERNELS -> SLOWER_MLA_GENERATION_KERNELS - "Missing TRTLLM-GEN decode kernel" -> "slower TRTLLM-GEN decode kernel"

…o SLOWER (review feedback) Address yuxianq's review on PR NVIDIA#16167: the (576, 512, 32) MLA generation kernel is not missing, just slower — reflect that in the identifier and the error string. - MISSING_MLA_GENERATION_KERNELS -> SLOWER_MLA_GENERATION_KERNELS - "Missing TRTLLM-GEN decode kernel" -> "slower TRTLLM-GEN decode kernel" Signed-off-by: chenfeiz0326 <chenfeiz@nvidia.com>

chenfeiz0326 · 2026-07-11T15:02:08Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-11T15:08:23Z

PR_Github #58768 [ run ] triggered by Bot. Commit: a7c8641 Link to invocation

tensorrt-cicd · 2026-07-12T03:00:29Z

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

chenfeiz0326 · 2026-07-12T10:06:27Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-12T14:53:33Z

PR_Github #58810 [ run ] triggered by Bot. Commit: a7c8641 Link to invocation

tensorrt-cicd · 2026-07-12T14:53:33Z

PR_Github #58810 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance.

Link to invocation

chzblych · 2026-07-12T17:28:02Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-12T17:38:06Z

PR_Github #58829 [ run ] triggered by Bot. Commit: a7c8641 Link to invocation

tensorrt-cicd · 2026-07-12T18:24:32Z

PR_Github #58829 [ run ] completed with state SUCCESS. Commit: a7c8641
/LLM/main/L0_MergeRequest_PR pipeline #47380 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

chenfeiz0326 requested a review from a team as a code owner July 9, 2026 03:40

chenfeiz0326 requested a review from PerkzZheng July 9, 2026 03:40

github-actions Bot assigned chenfeiz0326 Jul 9, 2026

coderabbitai Bot reviewed Jul 9, 2026

View reviewed changes

yuxianq requested review from yihwang-nv and yuxianq July 9, 2026 06:09

yuxianq reviewed Jul 10, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py Outdated

yuxianq reviewed Jul 10, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py Outdated

chenfeiz0326 force-pushed the repair-bot-bug6430674 branch from bedb56f to a7c8641 Compare July 11, 2026 10:11

chenfeiz0326 requested a review from yuxianq July 12, 2026 10:06

yuxianq approved these changes Jul 13, 2026

View reviewed changes

chenfeiz0326 merged commit 9fe5d0f into NVIDIA:main Jul 13, 2026
8 checks passed

Uh oh!

Conversation

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

Uh oh!

Summary

Test plan

Links

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 9, 2026

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

chenfeiz0326 commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

chenfeiz0326 commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

chenfeiz0326 commented Jul 11, 2026

Uh oh!

chenfeiz0326 commented Jul 11, 2026

Uh oh!

tensorrt-cicd commented Jul 11, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

chenfeiz0326 commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

chzblych commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

tensorrt-cicd commented Jul 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

chenfeiz0326 commented Jul 9, 2026 •

edited by coderabbitai Bot

Loading