[None][perf] Fuse index-q/index-k projections in MinimaxM3

[brb-nv]

Description

This MR fuses index-q / index-k projections in Minimax M3.

Test Coverage

$ pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] -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

• Fuses MiniMax-M3 index-Q and index-K projections into a single index_qk_proj, reducing projection execution from two calls to one.
• Adds checkpoint weight fusion through _load_index_qk_proj_weights and updates sparse attention to split the fused output into idx_q and idx_k.
• Updates documentation and preserves dense-attention behavior by ensuring only sparse attention exposes the fused projection.
• No configuration or test-list files were modified.
• Implementation is consistent with the stated performance optimization and checkpoint layout.

QA Engineer Review

• Updated MiniMax-M3 construction and checkpoint configuration tests to validate:
  • Fused index_qk_proj dimensions and replication.
  • Absence of separate index_q_proj and index_k_proj modules in sparse attention.
  • Absence of the fused projection in dense attention.
• No test functions were added or removed; existing assertions and documentation were updated.
• No tests/integration/test_lists/ coverage changes were included, so CI/manual QA coverage mapping is unavailable.
• Verdict: needs follow-up.

[brb-nv]

/bot run --disable-fail-fast

[coderabbitai]

Walkthrough

MiniMax-M3 sparse attention now uses a fused index_qk_proj for index Q/K computation. Checkpoint loading merges the original projection weights, and unit tests validate the fused module contract for dense and sparse configurations.

Changes

MiniMax-M3 sparse index projection

Layer / File(s)	Summary
Fused sparse index runtime tensorrt_llm/_torch/models/modeling_minimaxm3.py	Replaces separate index Q/K projections with one replicated projection, splits its concatenated output during sparse attention, and updates documentation.
Checkpoint weight fusion tensorrt_llm/_torch/models/modeling_minimaxm3.py	Merges checkpoint index_q_proj and index_k_proj tensors into index_qk_proj before generic weight loading.
Fused index contract tests tests/unittest/_torch/models/test_minimax_m3.py	Validates fused projection dimensions, replication, and removal of the former separate projection attributes in dense and sparse configurations.

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

Sequence Diagram(s)

sequenceDiagram
  participant Checkpoint
  participant MiniMaxM3ForCausalLM
  participant index_qk_proj
  participant SparseAttention
  Checkpoint->>MiniMaxM3ForCausalLM: load index_q_proj and index_k_proj tensors
  MiniMaxM3ForCausalLM->>index_qk_proj: fuse and load projection weights
  SparseAttention->>index_qk_proj: project hidden states once
  index_qk_proj->>SparseAttention: concatenated index Q/K output
  SparseAttention->>SparseAttention: split into idx_q and idx_k

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16468: Both changes modify MiniMax-M3 weight loading and the HF/loader-v2 integration path.

Suggested reviewers: qijune, bowenfu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title is concise and accurately describes the main change: fusing MiniMax-M3 index projections.
Description check	✅ Passed	The description follows the template with Description, Test Coverage, and PR Checklist sections filled in.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[tensorrt-cicd]

PR_Github #61963 [ run ] triggered by Bot. Commit: 3d88870 Link to invocation

[tensorrt-cicd]

PR_Github #61963 [ run ] completed with state SUCCESS. Commit: 3d88870
/LLM/main/L0_MergeRequest_PR pipeline #50154 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

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62567 [ run ] triggered by Bot. Commit: 3d88870 Link to invocation

[tensorrt-cicd]

PR_Github #62567 [ run ] completed with state SUCCESS. Commit: 3d88870
/LLM/main/L0_MergeRequest_PR pipeline #50711 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

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62606 [ run ] triggered by Bot. Commit: 3d88870 Link to invocation

[tensorrt-cicd]

PR_Github #62606 [ run ] completed with state SUCCESS. Commit: 3d88870
/LLM/main/L0_MergeRequest_PR pipeline #50747 completed with status: 'SUCCESS'

CI Report

Link to invocation