[TRTLLM-14019][feat] Add MiniMax-M3 MSA sparse attention backend [revised]#16291
Merged
Conversation
brb-nv
force-pushed
the
user/brb/msa-with-cudagraph
branch
from
July 13, 2026 01:16
f60a78d to
a94a130
Compare
brb-nv
changed the title
brb-nv
force-pushed
the
user/brb/msa-with-cudagraph
branch
from
July 13, 2026 01:19
a94a130 to
84682b1
Compare
brb-nv
added
the
api-compatible
brb-nv
marked this pull request as ready for review
July 13, 2026 01:48
Contributor
WalkthroughAdds MiniMax-M3 MSA sparse attention through ChangesMiniMax-M3 MSA integration
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Config
participant Resolver
participant MsaBackend
participant Indexer
participant FMHA
Config->>Resolver: lower implementation
Resolver->>MsaBackend: select MSA backend
MsaBackend->>Indexer: select KV blocks
Indexer->>FMHA: compute proxy scores
MsaBackend->>FMHA: execute paged sparse GQA
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Checkov (3.3.8)requirements.txtTraceback (most recent call last): scripts/attribution/scan/metadata/msa.ymlTraceback (most recent call last): tensorrt_llm/usage/llm_args_golden_manifest.jsonTraceback (most recent call last):
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py (2)
47-51: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing positive-case coverage for
is_available.Only the negative cases (
None,object()) are tested. Consider adding an assertion thatMsaSparseGqaFmha.is_available(<actual MSA attention instance>)returnsTrue, to confirm the filter actually admits the intended layer type, not just rejects others.🤖 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/attention/sparse/test_minimax_m3_msa_backend.py` around lines 47 - 51, Add positive coverage to test_msa_fmha_is_available_filters_on_owning_attention by constructing the actual MSA attention instance used by this backend and asserting MsaSparseGqaFmha.is_available(...) returns True, while preserving the existing None and non-MSA rejection assertions.Source: Path instructions
135-140: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest doesn't verify the claimed rejection behavior.
The test name and comment claim "a mismatch must fail loudly," but the assertions only confirm
params.topkpreserves the passed-through value — no failure is actually exercised or expected here. As written, this test would pass even if the backend never gates badtopkvalues. Consider either renaming the test to reflect what it actually checks (config lowering is pass-through), or adding an assertion at the point where gating is supposed to occur (e.g. instantiating/using the MSA backend withtopk=8and expecting aValueError/RuntimeError).🤖 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/attention/sparse/test_minimax_m3_msa_backend.py` around lines 135 - 140, Update test_msa_params_reject_bad_topk_at_construction to verify the claimed rejection at the MSA backend construction or usage point, asserting the expected ValueError or RuntimeError for sparse_topk_blocks=8; alternatively, rename the test and comments to accurately describe that MiniMaxM3SparseAttentionConfig.to_sparse_params preserves the value without testing backend validation.Source: Path instructions
🤖 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/attention_backend/sparse/minimax_m3/msa_utils.py`:
- Around line 37-58: Update cache_view_to_msa_paged and msa_paged_kv so the hot
forward path does not call contiguous() or otherwise copy the full K/V pool per
step. Keep the KV cache stored in the fmha_sm100 HND paged layout, or perform
the permutation once during cache construction, while preserving the returned
per-layer K/V interface.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 677-682: Regenerate the LLM-args golden manifest after adding
sparse_use_msa to MiniMaxM3SparseAttentionConfig by running
scripts/generate_llm_args_golden_manifest.py, then commit the resulting
tensorrt_llm/usage/llm_args_golden_manifest.json and obtain the required
telemetry/privacy CODEOWNER approval.
---
Nitpick comments:
In `@tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py`:
- Around line 47-51: Add positive coverage to
test_msa_fmha_is_available_filters_on_owning_attention by constructing the
actual MSA attention instance used by this backend and asserting
MsaSparseGqaFmha.is_available(...) returns True, while preserving the existing
None and non-MSA rejection assertions.
- Around line 135-140: Update test_msa_params_reject_bad_topk_at_construction to
verify the claimed rejection at the MSA backend construction or usage point,
asserting the expected ValueError or RuntimeError for sparse_topk_blocks=8;
alternatively, rename the test and comments to accurately describe that
MiniMaxM3SparseAttentionConfig.to_sparse_params preserves the value without
testing backend validation.
🪄 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: 0f08efbb-b4ad-4e0e-871c-7093994206fe
📒 Files selected for processing (25)
tensorrt_llm/_torch/attention_backend/fmha/__init__.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/fmha/registry.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.pytensorrt_llm/_torch/attention_backend/sparse/utils.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/llmapi/llm_args.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/waives.txttests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_scaffolding.py
💤 Files with no reviewable changes (2)
- tensorrt_llm/_torch/attention_backend/sparse/init.py
- tests/integration/test_lists/waives.txt
brb-nv
commented
Jul 13, 2026
Collaborator
|
[by Codex] @yizhang-nv Friendly reminder to review this PR. Thanks! |
yuxianq
reviewed
Jul 13, 2026
yuxianq
reviewed
Jul 13, 2026
yuxianq
reviewed
Jul 13, 2026
yuxianq
reviewed
Jul 13, 2026
yuxianq
reviewed
Jul 13, 2026
yuxianq
reviewed
Jul 13, 2026
brb-nv
force-pushed
the
user/brb/msa-with-cudagraph
branch
2 times, most recently
from
July 13, 2026 14:39
1d4a737 to
dc708c7
Compare
Collaborator
|
PR_Github #60749 [ run ] completed with state |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (4)
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py (1)
419-651: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDead
static_buffersbranches after the early return.Since Lines 422-432 return via
_build_runtime_metadata_fresh(...)wheneverstatic_buffers is None, every subsequentif static_buffers is not None:guard (Lines 442, 504, 526, 557, 577, 631) is always true and each correspondingelsebranch (Lines 515-516, 538-540, 562-565, 606-610, 641-642) is unreachable. This dead code makes an already dense function harder to follow. Drop the redundant guards and their else branches so the remaining body reads as the static-buffers-only path it now is.🤖 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/sparse/minimax_m3/triton_metadata.py` around lines 419 - 651, Remove the unreachable non-static branches after the early return in the metadata-building flow. In the function containing _build_runtime_metadata_fresh, simplify each subsequent static_buffers guard to its static-buffer path, removing the corresponding else branches and redundant checks while preserving the existing allocation, copying, and metadata construction behavior.tensorrt_llm/llmapi/llm_args.py (1)
752-758: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd concrete method annotations.
The new validator and metadata-lowering method are unannotated, including the untyped
**kwargsboundary. Add explicit parameter and return types without introducingAny.As per coding guidelines, “Annotate every function” and “avoid
Any.”Also applies to: 782-803
🤖 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/llmapi/llm_args.py` around lines 752 - 758, Add explicit type annotations to _validate_msa_block_size and the metadata-lowering method in the same class, including a concrete type for the **kwargs boundary and each return type. Use existing project-specific types or precise mappings rather than Any, while preserving the validator’s current behavior.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py (1)
55-60: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNarrow the blind
except Exception: pass.Catching bare
Exceptionand swallowing it silently masks unexpected failures inget_buffers(0)and violates the repo's exception-handling guideline. Catch the specific error you expect (e.g.AttributeError/RuntimeError) so unrelated failures still surface.As per coding guidelines: "Avoid broad exception handling; catch specific exceptions instead of using bare
except:" and "Catch the narrowest possible exceptions".♻️ Proposed narrowing
if kv_cache_manager is not None: try: return kv_cache_manager.get_buffers(0).device - except Exception: + except (AttributeError, RuntimeError): pass return torch.device(f"cuda:{torch.cuda.current_device()}")🤖 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/sparse/minimax_m3/msa_backend.py` around lines 55 - 60, In the device-selection logic around kv_cache_manager.get_buffers(0), replace the broad silent Exception handler with only the specific expected exception types, such as AttributeError or RuntimeError, that indicate the buffer device cannot be retrieved. Preserve the fallback to torch.device using the current CUDA device, while allowing unrelated failures to propagate.Sources: Coding guidelines, Linters/SAST tools
tensorrt_llm/_torch/pyexecutor/config_utils.py (1)
99-101: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd type annotations.
Annotate the parameter and
-> boolreturn; the neighboring helpers in this module use forward-ref typed params.As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions".♻️ Proposed annotation
-def is_minimax_m3(sparse_attention_config): +def is_minimax_m3(sparse_attention_config: Optional["SparseAttentionConfig"]) -> bool: """True when the sparse attention config selects the MiniMax-M3 algorithm.""" return sparse_attention_config is not None and sparse_attention_config.algorithm == "minimax_m3"🤖 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/pyexecutor/config_utils.py` around lines 99 - 101, Update the is_minimax_m3 function signature to add a forward-reference type annotation for sparse_attention_config and an explicit bool return annotation, matching the neighboring helper functions’ typing style while preserving its existing behavior.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 `@scripts/attribution/scan/metadata/msa.yml`:
- Around line 1-5: Add the repository-standard NVIDIA copyright and SPDX header
at the beginning of the metadata file, before the existing YAML name,
description, source, and directory_matches entries. Preserve the current
attribution metadata unchanged.
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py`:
- Around line 776-780: In the cache-device initialization around
kv_cache_manager.get_buffers(0) and layer_buf.device, replace the broad except
Exception with only the expected lookup/access exceptions, such as
AttributeError and IndexError, so those cases retain the current CUDA fallback
while unexpected cache-manager failures propagate.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 731-742: Change the num_attention_heads and num_key_value_heads
fields in the LLM arguments model from Optional[int] to Optional[PositiveInt],
preserving their existing defaults and fallback behavior for unset values.
Update the corresponding generated golden manifest to reflect the schema change.
In `@tests/integration/test_lists/test-db/l0_dgx_b200.yml`:
- Line 239: Update the CUDA-graph integration test entry for
TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph to select use_msa=True instead of
use_msa=False, ensuring the graph-stable MSA execution path is covered while
preserving the existing timeout configuration.
In `@tests/integration/test_lists/waives.txt`:
- Around line 416-417: Remove the unrelated KV-transfer skip entries for
unittest/disaggregated/test_kv_transfer.py from the waivers list, including both
the general and parameterized test entries; leave MiniMax-M3 MSA-related waivers
unchanged.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py`:
- Around line 55-60: In the device-selection logic around
kv_cache_manager.get_buffers(0), replace the broad silent Exception handler with
only the specific expected exception types, such as AttributeError or
RuntimeError, that indicate the buffer device cannot be retrieved. Preserve the
fallback to torch.device using the current CUDA device, while allowing unrelated
failures to propagate.
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py`:
- Around line 419-651: Remove the unreachable non-static branches after the
early return in the metadata-building flow. In the function containing
_build_runtime_metadata_fresh, simplify each subsequent static_buffers guard to
its static-buffer path, removing the corresponding else branches and redundant
checks while preserving the existing allocation, copying, and metadata
construction behavior.
In `@tensorrt_llm/_torch/pyexecutor/config_utils.py`:
- Around line 99-101: Update the is_minimax_m3 function signature to add a
forward-reference type annotation for sparse_attention_config and an explicit
bool return annotation, matching the neighboring helper functions’ typing style
while preserving its existing behavior.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 752-758: Add explicit type annotations to _validate_msa_block_size
and the metadata-lowering method in the same class, including a concrete type
for the **kwargs boundary and each return type. Use existing project-specific
types or precise mappings rather than Any, while preserving the validator’s
current behavior.
🪄 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: dd141f4d-9c85-4729-bdec-4bcf65343c4f
📒 Files selected for processing (41)
.gitmodules3rdparty/MSALICENSEjenkins/Build.groovyjenkins/BuildDockerImage.groovyjenkins/L0_MergeRequest.groovyjenkins/L0_Test.groovyjenkins/TensorRT_LLM_PLC.groovyjenkins/runPerfSanityTriage.groovyrequirements.txtscripts/attribution/scan/metadata/msa.ymlsetup.pytensorrt_llm/_torch/attention_backend/fmha/__init__.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/fmha/registry.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.pytensorrt_llm/_torch/attention_backend/sparse/utils.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/pyexecutor/config_utils.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/accuracy/references/gsm8k.yamltests/integration/defs/accuracy/references/mmlu.yamltests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/waives.txttests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/attention_backend/sparse/init.py
🚧 Files skipped from review as they are similar to previous changes (12)
- .gitmodules
- tensorrt_llm/_torch/attention_backend/fmha/init.py
- 3rdparty/MSA
- LICENSE
- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
- tensorrt_llm/_torch/attention_backend/fmha/registry.py
- tests/integration/test_lists/qa/llm_function_core.txt
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
- tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60785 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60751 [ run ] completed with state |
Collaborator
|
PR_Github #60785 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60830 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60830 [ run ] completed with state
|
Collaborator
Author
|
Failing test waived here: #16698. Rerunning. |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60860 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
From what I see, the MSA test passed in DGX_B200-4_GPUs-PyTorch-1.accuracy.test_llm_api_pytorch.TestMiniMaxM3: |
Collaborator
|
PR_Github #60860 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60907 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60907 [ run ] completed with state
|
Add the MiniMax-M3 MSA sparse attention backend with CUDA-graph support and integrate the MSA fmha_sm100 kernels as a git submodule at 3rdparty/MSA. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
force-pushed
the
user/brb/msa-with-cudagraph
branch
from
July 22, 2026 14:10
fbf8252 to
43edd7d
Compare
Collaborator
Author
|
/bot skip --comment "rebase for waives.txt only - 2 failing multi-gpu stages are unrelated to changes here." |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py (1)
419-432: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winEarly return leaves the inline fresh-allocation branches dead.
Since this new guard returns via
_build_runtime_metadata_freshwheneverstatic_buffers is None, everyif static_buffers is not None:check further down (L442, L504, L526, L557, L577, L631) is now always-true, and eachelsefresh-allocation branch (e.g. L515-516, L538-540, L562-565, L606-610, L641-642) is unreachable. That dead path duplicates_build_runtime_metadata_freshand will silently diverge from it over time. Consider dropping the redundant guards/else branches so the static-buffer path reads unconditionally.🤖 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/sparse/minimax_m3/triton_metadata.py` around lines 419 - 432, Remove the redundant static_buffers conditionals and unreachable fresh-allocation else branches throughout the metadata-building flow after the early return in the surrounding function. Make the remaining logic unconditionally use static buffers, while retaining _build_runtime_metadata_fresh as the sole path for static_buffers=None and preserving the existing static-buffer behavior.
🤖 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/attention_backend/sparse/minimax_m3/msa_backend.py`:
- Around line 52-60: Update _cache_device to catch only the specific expected
lookup or device-access exceptions from kv_cache_manager.get_buffers(0),
preserving the current-device fallback for those cases. Remove the broad
Exception handler and allow errors such as NotImplementedError and ValueError to
propagate.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 1158-1163: Add the return annotation -> tuple[torch.Tensor,
torch.Tensor] to the _build_pool_mapping_tensors method, preserving its existing
behavior and overridable hook contract.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Line 7687: Annotate each modified test method, including test_mxfp8 and the
methods at the other referenced locations, with the use_msa: bool parameter type
and a -> None return annotation.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py`:
- Around line 419-432: Remove the redundant static_buffers conditionals and
unreachable fresh-allocation else branches throughout the metadata-building flow
after the early return in the surrounding function. Make the remaining logic
unconditionally use static buffers, while retaining
_build_runtime_metadata_fresh as the sole path for static_buffers=None and
preserving the existing static-buffer behavior.
🪄 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: 5f7d6559-da04-4d18-ae23-c03d7c3f9fff
📒 Files selected for processing (41)
.gitmodules3rdparty/MSALICENSEjenkins/Build.groovyjenkins/BuildDockerImage.groovyjenkins/L0_MergeRequest.groovyjenkins/L0_Test.groovyjenkins/TensorRT_LLM_PLC.groovyjenkins/runPerfSanityTriage.groovyrequirements.txtscripts/attribution/scan/metadata/msa.ymlsetup.pytensorrt_llm/_torch/attention_backend/fmha/__init__.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/fmha/registry.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.pytensorrt_llm/_torch/attention_backend/sparse/utils.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/pyexecutor/config_utils.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/accuracy/references/gsm8k.yamltests/integration/defs/accuracy/references/mmlu.yamltests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/waives.txttests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/attention_backend/sparse/init.py
🚧 Files skipped from review as they are similar to previous changes (32)
- 3rdparty/MSA
- jenkins/runPerfSanityTriage.groovy
- .gitmodules
- tensorrt_llm/_torch/attention_backend/fmha/init.py
- tensorrt_llm/_torch/pyexecutor/config_utils.py
- tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
- jenkins/BuildDockerImage.groovy
- LICENSE
- jenkins/Build.groovy
- jenkins/L0_Test.groovy
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py
- scripts/attribution/scan/metadata/msa.yml
- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
- tests/integration/test_lists/qa/llm_function_core.txt
- jenkins/L0_MergeRequest.groovy
- tests/integration/defs/accuracy/references/mmlu.yaml
- tensorrt_llm/llmapi/llm_args.py
- requirements.txt
- tests/integration/test_lists/waives.txt
- tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_indexer.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/init.py
- setup.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_kernels.py
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/utils.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py
- tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_backend.py
Collaborator
|
PR_Github #60998 [ skip ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60998 [ skip ] completed with state |
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 22, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 23, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 23, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
7 tasks
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 24, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
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
15 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.
Description
On main, MinimaxM3 is supported by a Triton+SDPA attn backend. This MR adds support for an MSA backend based on Minimax kernels from here: https://github.com/MiniMax-AI/MSA.
This MR:
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
MsaSparseGqaFmha(usesfmha_sm100), including CUDA-graph-stable plan/buffer handling and dispatch via the FMHA attention-layer registry (init_fmha_libs+ new"msa_sparse_gqa"entry).implementation: "triton" | "msa"with MSA validation (requiressparse_block_size == 128; additional MSA constraints enforced in code/tests).tokens_per_block=128whenimplementation=="msa").ensure_msa_available) and deferred loading of the MSA Python module with SM-version checks.sparse/minimax_m3/common.py.MiniMaxM3TritonSparseAttentionMetadata) and updates sparse runtime/orchestrator wiring and exports.cache_view_to_msa_paged/msa_paged_kv) andwrite_kv_slotsfor both flat and paged caches.sparse_prediction.sparse_attn_indicesand the top‑k indices routing contract.Verdict: needs follow-up. Reported MR CI pipelines repeatedly failed after rerun requests; although MSA accuracy reportedly passed on a local DGX B200 setup, CI stability still requires investigation.
QA Engineer Review
Test list / CI coverage files touched
tests/integration/test_lists/qa/llm_function_core.txtTestMiniMaxM3:test_mxfp8[use_msa=False]TIMEOUT now 180test_mxfp8_piecewise_cuda_graph[use_msa=False]TIMEOUT now 180test_nvfp4[use_msa=False]TIMEOUT now 180test_nvfp4[use_msa=True]TIMEOUT now 60test_auto_dtype[tp_size=8-ep_size=8]TIMEOUT 180)tests/integration/test_lists/test-db/l0_dgx_b200.ymlTestMiniMaxM3::test_nvfp4[use_msa=True]TIMEOUT 60TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False]TIMEOUT 180TestMiniMaxM2::test_4gpus[...]TIMEOUT 60tests/integration/test_lists/waives.txtB200/B300/GB200/GB300:TestMiniMaxM3::test_mxfp8[use_msa=False]andTestMiniMaxM3::test_nvfp4[use_msa=False](with corresponding NVbug links)B200andB300:TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8]SKIP entries addedunittest/disaggregated/test_kv_transfer.pytest_transfer_worker_v2[...]Verdict (test-list updates): needs follow-up (CI failures were reported; CBTS coverage not provided here).
Test code touched (non test-list)
tests/integration/defs/accuracy/test_llm_api_pytorch.pyTestMiniMaxM3:test_mxfp8(use_msa)test_mxfp8_piecewise_cuda_graph(use_msa)test_nvfp4(use_msa)MiniMaxM3SparseAttentionConfig(implementation="msa" if use_msa else "triton")and sets KV cache dtype conditionally.qa/llm_function_core.txtandtest-db/l0_dgx_b200.yml.tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pyensure_msa_availableis availablesparse_block_size == 128constraint validationtest-db/orqa/lists (unit-only).tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.pysparse_attention_config=None.test-db/orqa/lists (unit-only).Verdict (test code): insufficient pending CI failure investigation; unit coverage exists, but MR CI stability issues still need resolution.