[None][perf] Fuse QK-norm + RoPE and overlap qkv/idx_qk math - #16906
Conversation
WalkthroughMiniMax-M3 attention now uses shared auxiliary-stream parallelism and fused BF16 QK normalization with partial RoPE across dense and sparse paths. CUDA tests cover fused correctness, fallback behavior, and fusion predicates. ChangesMiniMax-M3 attention execution
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 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}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
tests/unittest/_torch/models/test_minimax_m3.py (3)
836-842: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring overstates coverage.
Only one bf16 config is constructed; MXFP8 and NVFP4 flavors are never instantiated here. Trim the claim to what the test actually asserts.
🤖 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 `@tests/unittest/_torch/models/test_minimax_m3.py` around lines 836 - 842, Update the docstring describing the fused-path expectation near the forward assertions to reflect only the single bf16 configuration constructed by this test. Remove claims about MXFP8 and NVFP4 coverage, while retaining the behavior described for non-bf16 activations or missing position_ids if those cases remain tested.
723-726: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTolerances are loose enough to mask small regressions.
atol=1e-1on bf16 activations of order 1 permits ~10% error. bf16 round-trip through two different kernels usually lands well underatol=2e-2; consider tightening (or asserting a relative-error norm) so an epsilon/scale mistake still fails.🤖 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 `@tests/unittest/_torch/models/test_minimax_m3.py` around lines 723 - 726, Tighten the q_f and k_f versus q_s and k_s assertions in the affected test by reducing the absolute tolerance from 1e-1 to a stricter value around 2e-2, while preserving the existing relative tolerance unless needed for consistency. Keep the exact v_f versus v_s assertion unchanged.
785-831: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFallback coverage stops at the helper.
Both cases exercise
_fused_qk_norm_ropein isolation. Nothing drives_main_norm_rope/_index_norm_ropedown the fallback branch, which is why the undefinedidx_q/idx_kdefect flagged intensorrt_llm/_torch/models/modeling_minimaxm3.pyslips through. A test that calls_sparse_forward's branch helpers with fp16 activations (expecting the separate path, with the assert relaxed) would close the gap.🤖 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 `@tests/unittest/_torch/models/test_minimax_m3.py` around lines 785 - 831, Extend test_minimax_m3_fused_qk_norm_rope_fallbacks to exercise _sparse_forward’s fallback branch through _main_norm_rope or _index_norm_rope, using fp16 activations and position_ids. Assert the separate path completes successfully rather than only checking _fused_qk_norm_rope returns None, and relax any branch assertion needed to reach the fallback implementation so undefined idx_q/idx_k usage is detected.tensorrt_llm/_torch/models/modeling_minimaxm3.py (1)
792-808: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNested
maybe_execute_in_parallelreuses the same two events.On the non-bf16 fallback path,
_main_norm_rope/_index_norm_ropealready run inside amaybe_execute_in_parallelregion that usesself.ln_events[0]/[1], and they then callapply_qk_norm/apply_index_qk_norm, which record/wait on those same two events. The re-records only over-synchronize today (no lost ordering), but it silently kills the inner overlap and is easy to break with future edits. Consider giving the inner norm helpers their own event pair (e.g.self.qk_norm_events).🤖 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/models/modeling_minimaxm3.py` around lines 792 - 808, Give the inner q/k normalization invoked by apply_qk_norm and apply_index_qk_norm a dedicated event pair, such as self.qk_norm_events, instead of reusing self.ln_events. Update the nested maybe_execute_in_parallel call and all corresponding record/wait references so the outer _main_norm_rope/_index_norm_rope synchronization remains independent and inner overlap is preserved.
🤖 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 `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 1353-1380: The sparse index helper _index_norm_rope uses an
undefined self.index_qk_proj and references idx_q/idx_k before assignment in the
fallback. Replace the projection with the existing index_q_proj/index_k_proj
flow (or define the combined module consistently), initialize both tensors
before apply_index_qk_norm, and preserve the fused path’s outputs and subsequent
RoPE handling.
- Around line 880-892: Update the fused RoPE path around the visible
rotary-dimension validation to return None unless head_dim is 64, 128, or 256,
matching the supported fused_qk_norm_rope kernel inputs. Add the same
supported-dimension guard to _expect_fused_qk_norm_rope() so its expectation
logic mirrors the runtime fallback behavior.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 792-808: Give the inner q/k normalization invoked by apply_qk_norm
and apply_index_qk_norm a dedicated event pair, such as self.qk_norm_events,
instead of reusing self.ln_events. Update the nested maybe_execute_in_parallel
call and all corresponding record/wait references so the outer
_main_norm_rope/_index_norm_rope synchronization remains independent and inner
overlap is preserved.
In `@tests/unittest/_torch/models/test_minimax_m3.py`:
- Around line 836-842: Update the docstring describing the fused-path
expectation near the forward assertions to reflect only the single bf16
configuration constructed by this test. Remove claims about MXFP8 and NVFP4
coverage, while retaining the behavior described for non-bf16 activations or
missing position_ids if those cases remain tested.
- Around line 723-726: Tighten the q_f and k_f versus q_s and k_s assertions in
the affected test by reducing the absolute tolerance from 1e-1 to a stricter
value around 2e-2, while preserving the existing relative tolerance unless
needed for consistency. Keep the exact v_f versus v_s assertion unchanged.
- Around line 785-831: Extend test_minimax_m3_fused_qk_norm_rope_fallbacks to
exercise _sparse_forward’s fallback branch through _main_norm_rope or
_index_norm_rope, using fp16 activations and position_ids. Assert the separate
path completes successfully rather than only checking _fused_qk_norm_rope
returns None, and relax any branch assertion needed to reach the fallback
implementation so undefined idx_q/idx_k usage is detected.
🪄 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: a92ab75c-8a5a-426c-a276-4004a0f430b8
📒 Files selected for processing (2)
tensorrt_llm/_torch/models/modeling_minimaxm3.pytests/unittest/_torch/models/test_minimax_m3.py
b8e8cd4 to
a95a6bb
Compare
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
a95a6bb to
3eca866
Compare
|
/bot run --disable-fail-fast |
There was a problem hiding this comment.
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 `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 795-811: Use a dedicated event pair for the outer sparse-forward
parallel scope instead of sharing self.ln_events and self.aux_stream with
apply_qk_norm and apply_index_qk_norm. Update _sparse_forward and its associated
maybe_execute_in_parallel calls to use the new attn_events synchronization
objects, while leaving the normalization helpers on ln_events so nested scopes
cannot overwrite each other’s events.
🪄 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: 41d84d93-1021-4465-97a0-b167a90740f6
📒 Files selected for processing (1)
tensorrt_llm/_torch/models/modeling_minimaxm3.py
|
PR_Github #62876 [ run ] triggered by Bot. Commit: |
|
PR_Github #62876 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63110 [ run ] triggered by Bot. Commit: |
|
PR_Github #63110 [ run ] completed with state |
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
3 participants
Description
This MR does the following:
Test Coverage
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.Dev Engineer Review
aux_streamparameter toMiniMaxM3Attentionand wires it through model initialization.QA Engineer Review
Added tests:
test_minimax_m3_fused_qk_norm_rope_main_matches_separatetest_minimax_m3_fused_qk_norm_rope_index_matches_separatetest_minimax_m3_fused_qk_norm_rope_fallbackstest_minimax_m3_expect_fused_qk_norm_rope_predicateThe tests cover fused main and index paths, fallback behavior, and fusion predicate guards. Coverage in
tests/integration/test_lists/was not confirmed.Verdict: needs follow-up