[https://nvbugs/6430674][fix] Surgical revert of PR #15838's dispatch guard removal: re-add…#16167
Open
chenfeiz0326 wants to merge 1 commit into
Open
[https://nvbugs/6430674][fix] Surgical revert of PR #15838's dispatch guard removal: re-add…#16167chenfeiz0326 wants to merge 1 commit into
chenfeiz0326 wants to merge 1 commit into
Conversation
… 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>
Contributor
📝 WalkthroughWalkthroughAdds 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. ChangesMLA Generation Support Validation
Estimated code review effort: 1 (Trivial) | ~5 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py (1)
411-413: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider 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_DIMSright 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
📒 Files selected for processing (1)
tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
Collaborator
Author
|
/bot run --disable-fail-fast --stage-list "*PerfSanity*" |
Collaborator
|
PR_Github #58412 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58412 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Test plan
Links
Summary by CodeRabbit