[None][perf] AllReduce + ResidualAdd + RMSNorm

[brb-nv]

Description

This MR adds fusion for AllReduce, ResidualAdd and RMSNorm for MinimaxM3. The fusion doesn't seem to be supported for models with gemma-style RMSNorm. So, we fold +1 into norm's weight loading and make use_gemma_norm=False.

Test Coverage

$ pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True-eval_mode=default] -s -v

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.

Dev Engineer Review

• Added AllReduce fusion across MiniMax-M3 attention, decoder, MoE/MLP, and final RMSNorm paths.
• Added deferred reduction parameters to attention forwards and alias wiring via setup_aliases().
• Folded Gemma-style RMSNorm’s +1 into loaded weights and disabled Gemma norm behavior for fusion compatibility.
• Dense and sparse attention paths now propagate deferred reduction settings.
• Review focus: validate fusion configuration, environment-based disabling, API consistency, and residual/norm ordering.

QA Engineer Review

No test changes.

[brb-nv]

/bot run --disable-fail-fast

[coderabbitai]

Walkthrough

MiniMax-M3 now supports deferred AllReduce and residual/RMSNorm fusion across attention, decoder, MoE/MLP, and final normalization paths. Boundary norm aliases and Gemma weight folding are configured during model setup and checkpoint loading.

Changes

MiniMax-M3 AllReduce fusion

Layer / File(s)	Summary
Defer attention output reductions tensorrt_llm/_torch/models/modeling_minimaxm3.py	Dense and sparse attention paths accept deferred AllReduce parameters and forward them through output projections.
Fuse decoder residual and norm operations tensorrt_llm/_torch/models/modeling_minimaxm3.py	Decoder layers defer attention and feed-forward reductions, preserve residual state, and fuse residual updates with RMSNorm operations.
Update boundary norms and aliases tensorrt_llm/_torch/models/modeling_minimaxm3.py	Final normalization, decoder norm aliases, and Gemma boundary-weight folding support the fused execution path.

Estimated code review effort: 4 (Complex) | ~45 minutes

Suggested reviewers: bowenfu, weihaocheng

Sequence Diagram(s)

sequenceDiagram
  participant MiniMaxM3Attention
  participant MiniMaxM3DecoderLayer
  participant AllReduceFusionOp
  participant MoEOrMLP
  participant MiniMaxM3ForCausalLM
  MiniMaxM3DecoderLayer->>MiniMaxM3Attention: request deferred attention reduction
  MiniMaxM3Attention-->>MiniMaxM3DecoderLayer: return unreduced attention output
  MiniMaxM3DecoderLayer->>AllReduceFusionOp: fuse residual and boundary RMSNorm
  MiniMaxM3DecoderLayer->>MoEOrMLP: pass deferred feed-forward reduction
  MoEOrMLP-->>AllReduceFusionOp: return feed-forward output
  AllReduceFusionOp-->>MiniMaxM3DecoderLayer: return normalized layer state
  MiniMaxM3DecoderLayer->>MiniMaxM3ForCausalLM: apply aliased final norm on last layer

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the AllReduce, ResidualAdd, and RMSNorm fusion added for MiniMax-M3.
Description check	✅ Passed	The description explains the fusion, Gemma RMSNorm compatibility workaround, and provides relevant test coverage while following the required template.
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 💡 1

🛠️ Fix failing CI checks 💡

• Fix failing CI checks

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (3)

tensorrt_llm/_torch/models/modeling_minimaxm3.py (3)

1444-1468: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

TRTLLM_MINIMAX_M3_EAGER_FUSION_DISABLED only honors the literal "0".

Any other value ("false", "no", empty-ish typos) silently disables fusion. Consider a shared truthiness helper for consistency with the rest of the env-var handling in the repo.

🤖 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 1444 - 1468,
The enable_fusion configuration currently enables fusion only when
TRTLLM_MINIMAX_M3_EAGER_FUSION_DISABLED equals the literal "0". Update this
initialization to use the repository’s shared environment-variable truthiness
helper, preserving the intended disabled semantics for recognized truthy values
and consistency with other env-var handling; keep the attention-DP and TP-size
gating unchanged.

---

1470-1477: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Return annotation no longer matches the returned value.

forward returns (hidden_states, residual); the annotation still says torch.Tensor.

♻️ Proposed fix

-    ) -> torch.Tensor:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:

As per coding guidelines: "Annotate every function ... avoid Any and unnecessary type ignores".

🤖 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 1470 - 1477,
Update the return annotation of the MiniMaxM3 model’s forward method to describe
its actual two-value result, hidden_states and residual, using the appropriate
typed tuple annotation rather than torch.Tensor. Keep the existing parameter
annotations and return behavior unchanged.

Source: Coding guidelines

---

1730-1750: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Annotate and guard _fold_gemma_boundary_norm_weights against repeated application.

Use annotations matching the supported weight containers. Reapplying the helper changes weights from w + 1 to w + 2; ensure each mapping is folded only once. Restrict suffix matches to decoder layer keys plus the exact final norm key.

🤖 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 1730 - 1750,
Update _fold_gemma_boundary_norm_weights with type annotations matching the
supported weight mapping and tensor containers. Make the fold idempotent by
tracking whether each mapping has already been processed and returning it
unchanged on subsequent calls. Restrict suffix-based updates to decoder-layer
norm keys, while continuing to match the exact _M3_FINAL_NORM_KEY.

Source: Coding guidelines

🤖 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 1547-1571: Gate the deferred reduction in
_feed_forward_all_reduce_params() on both post_feed_forward_fusion and a wired
next_layer_layernorm. When next_layer_layernorm is None, return the default
parameters so the MoE or dense-MLP performs its internal AllReduce; preserve
deferred fusion only when the alias is present.
- Around line 1778-1783: Update the guard in the boundary-norm folding logic to
read use_gemma_norm from the pretrained ModelConfig stored in self.model_config,
rather than self.config. Preserve the existing default-to-false behavior when
the configuration field is absent, and continue invoking
_fold_gemma_boundary_norm_weights only when the flag is enabled.

---

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 1444-1468: The enable_fusion configuration currently enables
fusion only when TRTLLM_MINIMAX_M3_EAGER_FUSION_DISABLED equals the literal "0".
Update this initialization to use the repository’s shared environment-variable
truthiness helper, preserving the intended disabled semantics for recognized
truthy values and consistency with other env-var handling; keep the attention-DP
and TP-size gating unchanged.
- Around line 1470-1477: Update the return annotation of the MiniMaxM3 model’s
forward method to describe its actual two-value result, hidden_states and
residual, using the appropriate typed tuple annotation rather than torch.Tensor.
Keep the existing parameter annotations and return behavior unchanged.
- Around line 1730-1750: Update _fold_gemma_boundary_norm_weights with type
annotations matching the supported weight mapping and tensor containers. Make
the fold idempotent by tracking whether each mapping has already been processed
and returning it unchanged on subsequent calls. Restrict suffix-based updates to
decoder-layer norm keys, while continuing to match the exact _M3_FINAL_NORM_KEY.

🪄 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: 25d0bafc-232f-4c04-9d25-72c2f246f370

📥 Commits

Reviewing files that changed from the base of the PR and between 5c5ef98 and c59f666.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/models/modeling_minimaxm3.py

[tensorrt-cicd]

PR_Github #62897 [ run ] triggered by Bot. Commit: c59f666 Link to invocation

[tensorrt-cicd]

PR_Github #62897 [ run ] completed with state SUCCESS. Commit: c59f666
/LLM/main/L0_MergeRequest_PR pipeline #51019 completed with status: 'SUCCESS'

CI Report

Link to invocation