Draft: [None][feat] Integrate M3 sparse attention kernels form MSA - #15998
Draft: [None][feat] Integrate M3 sparse attention kernels form MSA#15998brb-nv wants to merge 1 commit into
Conversation
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughAdds an optional fmha_sm100 (MSA) dependency and new FMHA backend abstractions (BlockSparseFmha, IndexerProxyFmha) to route MiniMax-M3 sparse attention through external kernels, including a CUDA-graph-safe decode driver, plan cache, backend/config wiring, and model integration. Separately adds piecewise CUDA graph runner tracking/clearing and a compile-boundary custom op for MiniMax-M3 attention. ChangesMSA-backed MiniMax-M3 sparse attention
Piecewise CUDA graph runner lifecycle and compile boundary op
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant Model as MiniMaxM3Attention
participant Backend as MsaBackend
participant ProxyFmha as MsaProxyMqaFmha
participant Driver as M3DecodeKernelDriver
participant SparseFmha as MsaSparseGqaFmha
Model->>Backend: forward_sparse(q, k, v, idx_q, idx_k, metadata)
Backend->>Driver: proxy_max_score(idx_q, idx_k_paged)
Driver->>ProxyFmha: forward_proxy(...)
ProxyFmha-->>Driver: max_score
Driver->>Driver: select_blocks(max_score, seq_lens)
Driver->>SparseFmha: forward_block_sparse(q, k, v, kv_block_indexes)
SparseFmha-->>Driver: attention output
Driver-->>Backend: output buffer
Backend-->>Model: o_proj
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: 3
🧹 Nitpick comments (3)
tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py (1)
58-86: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
is_availableboilerplate withMsaSparseGqaFmha.The
find_specprobe + CUDA-availability check + device-capability try/except is duplicated almost verbatim inmsa_sparse_gqa.py'sis_available. Consider extracting a shared helper (e.g. in a small internal module) that both classes call, passing only the class name for log messages, to avoid the two copies drifting out of sync.🤖 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/msa_proxy_mqa.py` around lines 58 - 86, The is_available logic in MsaProxyMqaFmha duplicates the same fmha_sm100/CUDA/device-capability probing used by MsaSparseGqaFmha, so extract that shared availability check into a small internal helper and have both classes call it. Keep the class-specific log text configurable (for example via the caller’s name) so MsaProxyMqaFmha and MsaSparseGqaFmha can share the same implementation without drifting out of sync.tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/dispatch.py (1)
29-29: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer builtin generics over
typing.Dict/typing.Tuple.This file already uses
from __future__ import annotations; the class/module annotations can usedict[...]/tuple[...].As per coding guidelines: "Prefer built-in types
list,dict,tupleovertyping.List,typing.Dict,typing.Tuple".🤖 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/decode_wrapper/dispatch.py` at line 29, Update the type imports in dispatch.py to stop using typing.Dict and typing.Tuple, since the module already has future annotations enabled. Replace any annotations in the relevant symbols in this file with builtin generics like dict[...] and tuple[...] so the type hints follow the codebase guideline.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py (1)
999-1019: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsolidate the MiniMax M3 backend dispatch
get_minimax_m3_attention_backend_cls_with_msa()isn’t referenced anywhere, whiletensorrt_llm/_torch/attention_backend/sparse/utils.pyalready reimplements the sameuse_msabranch. Route callers through one resolver or remove this helper and its__all__export.🤖 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 999 - 1019, `get_minimax_m3_attention_backend_cls_with_msa()` is duplicated and currently unused, while the same `use_msa` dispatch is already handled in `tensorrt_llm/_torch/attention_backend/sparse/utils.py`. Either update callers to use this resolver consistently, or remove `get_minimax_m3_attention_backend_cls_with_msa` and its `__all__` export from `msa_backend.py` if it is not needed anywhere else.
🤖 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 `@requirements.txt`:
- Around line 58-61: Make the fmha_sm100 git dependency optional instead of
always included in install_requires, since setup.py consumes requirements.txt
directly. Move the MSA-specific dependency behind an extra or environment marker
so only SM100-enabled/runtime MSA installs request it, and keep the default
requirements usable for offline, restricted, or non-SM100 environments. Update
the requirements entry and the packaging path that reads it so the optional
dependency is only activated when needed.
In `@tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py`:
- Around line 76-79: The exception handling in the capability check inside the
MSA proxy backend is too broad; narrow the `except` around
`torch.cuda.get_device_capability()` in `msa_proxy_mqa.py` so it only catches
the CUDA-related failures expected there instead of all `Exception`s. Update the
`get_device_capability` try-except in the backend probe to use the smallest
specific exception type(s) that indicate CUDA/device capability lookup is
unavailable, while keeping the existing False fallback behavior.
In `@tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py`:
- Around line 75-78: The exception handling in MsaSparseGqaFmha.is_available is
too broad and should be narrowed to match MsaProxyMqaFmha.is_available. Update
the try/except around torch.cuda.get_device_capability() to catch RuntimeError
instead of Exception, keeping the existing False fallback behavior unchanged.
Use the MsaSparseGqaFmha.is_available method as the location to make the same
BLE001-safe adjustment seen in msa_proxy_mqa.py.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py`:
- Around line 58-86: The is_available logic in MsaProxyMqaFmha duplicates the
same fmha_sm100/CUDA/device-capability probing used by MsaSparseGqaFmha, so
extract that shared availability check into a small internal helper and have
both classes call it. Keep the class-specific log text configurable (for example
via the caller’s name) so MsaProxyMqaFmha and MsaSparseGqaFmha can share the
same implementation without drifting out of sync.
In
`@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/dispatch.py`:
- Line 29: Update the type imports in dispatch.py to stop using typing.Dict and
typing.Tuple, since the module already has future annotations enabled. Replace
any annotations in the relevant symbols in this file with builtin generics like
dict[...] and tuple[...] so the type hints follow the codebase guideline.
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py`:
- Around line 999-1019: `get_minimax_m3_attention_backend_cls_with_msa()` is
duplicated and currently unused, while the same `use_msa` dispatch is already
handled in `tensorrt_llm/_torch/attention_backend/sparse/utils.py`. Either
update callers to use this resolver consistently, or remove
`get_minimax_m3_attention_backend_cls_with_msa` and its `__all__` export from
`msa_backend.py` if it is not needed anywhere else.
🪄 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: 62da48ad-d31f-4a68-906f-54b06c48c345
📒 Files selected for processing (30)
requirements.txttensorrt_llm/_torch/attention_backend/fmha/__init__.pytensorrt_llm/_torch/attention_backend/fmha/block_sparse.pytensorrt_llm/_torch/attention_backend/fmha/indexer_proxy.pytensorrt_llm/_torch/attention_backend/fmha/interface.pytensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/fmha/registry.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/__init__.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/__init__.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/dispatch.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/topk.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/worklist.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_plan_cache.pytensorrt_llm/_torch/attention_backend/sparse/utils.pytensorrt_llm/_torch/compilation/backend.pytensorrt_llm/_torch/compilation/piecewise_optimizer.pytensorrt_llm/_torch/compilation/utils.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/pyexecutor/model_engine.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/unittest/_torch/attention/sparse/test_minimax_m3_decode_driver_vs_msa.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
|
PR_Github #57816 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #57844 [ run ] triggered by Bot. Commit: |
|
PR_Github #57816 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #57867 [ run ] triggered by Bot. Commit: |
|
PR_Github #57844 [ run ] completed with state |
9513ad7 to
87b4178
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58196 [ run ] completed with state
|
5741d30 to
818a73e
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58294 [ run ] triggered by Bot. Commit: |
…5809) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
818a73e to
3314c00
Compare
|
Rebase needed because tests that are failing in ongoing run seem to be fixed and unwaived here: |
|
/bot run --disable-fail-fast |
|
PR_Github #58319 [ run ] triggered by Bot. Commit: |
|
PR_Github #58294 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #58385 [ run ] triggered by Bot. Commit: |
|
PR_Github #58319 [ run ] completed with state
|
|
PR_Github #58385 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58415 [ run ] triggered by Bot. Commit: |
|
PR_Github #58415 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
| - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.9-fp8kv=True] | ||
| - disaggregated/test_disaggregated.py::test_disaggregated_mamba_conc_greater_than_mbs[NVIDIA-Nemotron-3-Super-120B-A12B-FP8] | ||
| - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_attn_multi_gpus TIMEOUT (60) | ||
| - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] TIMEOUT (180) |
There was a problem hiding this comment.
We only add disable-MSA cases to post-merge CI, can we add the enable-MSA cases to per-merge CI and ensure that it can pass?
There was a problem hiding this comment.
To speedup CI, we can only run the stages that contain M3 MSA test
There was a problem hiding this comment.
Hi Yuxian, this is intentional. Currently, installing MSA in TRTLLM needs to go through some approval process. To unblock the team, we decided to do that in a follow-up MR.
|
PR_Github #58441 [ run ] triggered by Bot. Commit: |
|
PR_Github #58441 [ run ] completed with state
|
|
/bot skip --comment "Failing flaky test has been waived already" |
|
PR_Github #58469 [ skip ] triggered by Bot. Commit: |
|
PR_Github #58469 [ skip ] completed with state |
pcastonguay
changed the title
15 participants
Description
This MR integrates Minimax's MSA kernels to support Minimax M3. Original MR by @WeiHaocheng :
#15809
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.Summary by CodeRabbit
New Features
Bug Fixes
Tests