[None][feat] Integrate M3 sparse attention kernels form MSA - #15809
Closed
WeiHaocheng wants to merge 4 commits into
Closed
[None][feat] Integrate M3 sparse attention kernels form MSA#15809WeiHaocheng wants to merge 4 commits into
WeiHaocheng wants to merge 4 commits into
Conversation
Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com>
WeiHaocheng
force-pushed
the
feat/main_msa_integration
branch
from
July 6, 2026 12:44
7b2721c to
077ccb1
Compare
pcicotti
marked this pull request as ready for review
July 6, 2026 12:47
Collaborator
Author
|
/bot run --disable-fail-fast |
Contributor
📝 WalkthroughWalkthroughThis PR adds an external MSA (fmha_sm100) sparse attention path for MiniMax-M3. It introduces new FMHA backend base classes (BlockSparseFmha, IndexerProxyFmha), MSA kernel wrappers, a CUDA-graph-safe decode driver, plan caching, config/metadata plumbing (sparse_use_msa), model-level index-head TP geometry changes, and tests. ChangesMSA fmha_sm100 sparse attention integration
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant MiniMaxM3MSARuntimeBackend
participant minimax_m3_msa_sparse_prefill
participant IndexerProxyFmha
participant BlockSparseFmha
MiniMaxM3MSARuntimeBackend->>minimax_m3_msa_sparse_prefill: forward(q, k_cache, v_cache, idx_q, idx_k_cache)
minimax_m3_msa_sparse_prefill->>IndexerProxyFmha: forward_proxy(idx_q, idx_k_paged)
IndexerProxyFmha-->>minimax_m3_msa_sparse_prefill: max_score
minimax_m3_msa_sparse_prefill->>minimax_m3_msa_sparse_prefill: select top-k kv_block_indexes
minimax_m3_msa_sparse_prefill->>BlockSparseFmha: forward_block_sparse(q, k_paged, v_paged, kv_block_indexes)
BlockSparseFmha-->>minimax_m3_msa_sparse_prefill: output tensor
sequenceDiagram
participant proxy_mqa_decode
participant M3DecodeKernelDriver
participant sparse_gqa_decode
proxy_mqa_decode->>M3DecodeKernelDriver: proxy_max_score(idx_q, idx_k_paged)
M3DecodeKernelDriver-->>proxy_mqa_decode: max_score view
M3DecodeKernelDriver->>M3DecodeKernelDriver: select_blocks(max_score, seq_lens)
sparse_gqa_decode->>M3DecodeKernelDriver: sparse_attention(q, k_paged, v_paged, kv_block_indexes)
M3DecodeKernelDriver-->>sparse_gqa_decode: output view
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
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: 6
🧹 Nitpick comments (5)
tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py (1)
18-21: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse built-in
listinstead oftyping.List.
from typing import ListandList[int] | None(line 148) should use the built-inlist[int] | None, especially sincefrom __future__ import annotations(line 18) already makes this safe pre-3.9. As per coding guidelines, "Prefer using the built-in typeslist,dict, andtupleto the legacytyping.List,typing.Dict, andtyping.Tuple."♻️ Proposed fix
import importlib -from typing import List import pytest import torchextend_seq_lens: List[int] | None = None, + extend_seq_lens: list[int] | None = None,Also applies to: 144-150
🤖 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 18 - 21, The test module is still importing and using the legacy typing alias `List` instead of the built-in generic `list`, so update the type annotation in `test_minimax_m3_msa_backend` to use `list[int] | None` and remove the `typing.List` import. Keep the change consistent anywhere the same annotation pattern appears in the affected test block, using the existing `from __future__ import annotations` support and the relevant test helper/function names in that section to locate the type hints.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py (1)
67-86: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to the new helper/factory functions.
Several new functions omit return types, including
_require_msa_module, the FMHA selector helpers, the backend factory, nested__init__, andget_minimax_m3_attention_backend_cls_with_msa.As per coding guidelines, “Always annotate functions with return types (use
Noneif no return).”Also applies to: 220-241, 457-475, 837-861, 999-1011
🤖 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 67 - 86, Several newly added helper and factory functions are missing explicit return annotations, including _require_msa_module, the FMHA selector helpers, the backend factory, the nested __init__ methods, and get_minimax_m3_attention_backend_cls_with_msa. Update each of these definitions to include the correct return type hint, using None for constructors/initializers and the appropriate concrete return types for module selectors and factory helpers, so the new MSA backend APIs follow the project’s annotation guidelines.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/fmha/indexer_proxy.py (1)
92-125: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring uses NumPy style instead of Google style.
Same issue as
block_sparse.py:forward_proxy's docstring usesParameters/Returnssections (NumPy convention) rather than Google style. As per coding guidelines, "Use Google style docstrings for classes and functions in Python, parseable by Sphinx" applies here.📝 Convert to Google style
- Parameters - ---------- - idx_q : torch.Tensor - Shape ``[total_q, num_qo_heads, head_dim]`` (bf16/fp16). - idx_k_paged : torch.Tensor - Paged KV in HND layout - ``[num_pages, num_kv_heads, page_size, head_dim]``. For the - canonical MQA proxy, ``num_kv_heads == 1`` and ``idx_k`` is - broadcast across every QO head during scoring. - qo_lens_cpu, kv_lens_cpu : torch.Tensor - Shape ``[batch]``, dtype int32, on CPU. Per-request Q/O - and KV lengths. - qo_offset_cpu : torch.Tensor, optional - Shape ``[batch]``, dtype int32, on CPU. Per-request causal - offset (i.e. prefix length). Ignored when ``causal=False``. - kv_indices : torch.Tensor - Shape ``[sum_pages_across_batch]``, dtype int32, on the - cache device. Flattened paged-KV page table. - sm_scale : float - Softmax scale applied to the QK scores prior to the - per-block max reduction. - causal : bool - Whether to apply a causal mask. Prefill batches typically - use ``True``; pure-decode batches use ``False``. - - Returns - ------- - torch.Tensor - Shape ``[num_qo_heads, max_k_tiles, total_q]``, dtype - float32. Out-of-range tile slots are padded with ``-inf`` - so a subsequent top-k selector can ignore them. - """ + Args: + idx_q: Shape ``[total_q, num_qo_heads, head_dim]`` (bf16/fp16). + idx_k_paged: Paged KV in HND layout + ``[num_pages, num_kv_heads, page_size, head_dim]``. For the + canonical MQA proxy, ``num_kv_heads == 1`` and ``idx_k`` is + broadcast across every QO head during scoring. + qo_lens_cpu: Shape ``[batch]``, dtype int32, on CPU. Per-request Q/O lengths. + kv_lens_cpu: Shape ``[batch]``, dtype int32, on CPU. Per-request KV lengths. + qo_offset_cpu: Shape ``[batch]``, dtype int32, on CPU, optional. Per-request + causal offset (i.e. prefix length). Ignored when ``causal=False``. + kv_indices: Shape ``[sum_pages_across_batch]``, dtype int32, on the + cache device. Flattened paged-KV page table. + sm_scale: Softmax scale applied to the QK scores prior to the + per-block max reduction. + causal: Whether to apply a causal mask. Prefill batches typically + use ``True``; pure-decode batches use ``False``. + + Returns: + Shape ``[num_qo_heads, max_k_tiles, total_q]``, dtype float32. + Out-of-range tile slots are padded with ``-inf`` so a subsequent + top-k selector can ignore them. + """🤖 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/indexer_proxy.py` around lines 92 - 125, The docstring on the max-score helper in indexer_proxy.py is using NumPy-style sections instead of the required Google style. Update the docstring for the proxy method that computes the per-(qo_head, kv_tile) max-score tensor to use Google-style Args and Returns formatting, matching the style used elsewhere such as in block_sparse.py, while preserving the same parameter descriptions and output shape details.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/fmha/block_sparse.py (1)
96-131: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring uses NumPy style instead of Google style.
forward_block_sparse's docstring usesParameters/Returnssections with dashed underlines (NumPy convention), not Google style (Args:/Returns:). As per coding guidelines, "Use Google style docstrings for classes and functions in Python, parseable by Sphinx" applies to**/*.py.📝 Convert to Google style
- Parameters - ---------- - q : torch.Tensor - Shape ``[total_q, num_qo_heads, head_dim]`` (bf16/fp16). - k_paged : torch.Tensor - Paged K cache in HND layout - ``[num_pages, num_kv_heads, page_size, head_dim]``. - v_paged : torch.Tensor - Paged V cache, same shape as ``k_paged``. - kv_block_indexes : torch.Tensor - Shape ``[total_q, num_kv_heads, topk]``, dtype int32, - ascending per row with ``-1`` padding at the tail. Encodes - the per-query subset of KV blocks selected by the - preceding sparse predictor. - qo_lens_cpu, kv_lens_cpu : torch.Tensor - Shape ``[batch]``, dtype int32, on CPU. Per-request Q/O - and KV lengths. - qo_offset_cpu : torch.Tensor, optional - Shape ``[batch]``, dtype int32, on CPU. Per-request causal - offset (i.e. prefix length). Ignored when ``causal=False``. - kv_indices : torch.Tensor - Shape ``[sum_pages_across_batch]``, dtype int32, on the - cache device. Flattened paged-KV page table. - sm_scale : float - Softmax scale. - causal : bool - Whether to apply a causal mask. - - Returns - ------- - torch.Tensor - Shape ``[total_q, num_qo_heads, head_dim]``, dtype - bfloat16. The attention output over the selected KV blocks. - """ + Args: + q: Shape ``[total_q, num_qo_heads, head_dim]`` (bf16/fp16). + k_paged: Paged K cache in HND layout + ``[num_pages, num_kv_heads, page_size, head_dim]``. + v_paged: Paged V cache, same shape as ``k_paged``. + kv_block_indexes: Shape ``[total_q, num_kv_heads, topk]``, dtype + int32, ascending per row with ``-1`` padding at the tail. + Encodes the per-query subset of KV blocks selected by the + preceding sparse predictor. + qo_lens_cpu: Shape ``[batch]``, dtype int32, on CPU. Per-request Q/O lengths. + kv_lens_cpu: Shape ``[batch]``, dtype int32, on CPU. Per-request KV lengths. + qo_offset_cpu: Shape ``[batch]``, dtype int32, on CPU, optional. Per-request + causal offset (i.e. prefix length). Ignored when ``causal=False``. + kv_indices: Shape ``[sum_pages_across_batch]``, dtype int32, on the + cache device. Flattened paged-KV page table. + sm_scale: Softmax scale. + causal: Whether to apply a causal mask. + + Returns: + Shape ``[total_q, num_qo_heads, head_dim]``, dtype bfloat16. The + attention output over the selected KV blocks. + """🤖 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/block_sparse.py` around lines 96 - 131, The docstring on forward_block_sparse is using NumPy-style sections instead of the required Google style. Update the function documentation to use Google-style headings like Args and Returns, keeping the same parameter descriptions for q, k_paged, v_paged, kv_block_indexes, qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, kv_indices, sm_scale, and causal so the docstring remains Sphinx-parseable.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py (1)
58-86: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
is_available()logic withMsaSparseGqaFmha.This method is nearly byte-for-byte identical to
MsaSparseGqaFmha.is_available(). Extracting a shared helper (e.g._is_fmha_sm100_available()in a small shared module) would avoid future drift between the two copies.🤖 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 MsaProxyMqaFmha.is_available() method duplicates the same availability checks used by MsaSparseGqaFmha.is_available(), so refactor the shared CUDA/SM100 probing logic into a common helper such as _is_fmha_sm100_available() in a shared module and have both classes call it. Keep the existing fmha_sm100 import probe, CUDA availability check, and compute-capability gate in one place, and preserve the class-specific debug messages only if needed around that shared result to prevent the two implementations from drifting.
🤖 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/fmha/msa_proxy_mqa.py`:
- Around line 76-79: The exception handling in the capability check is too
broad; narrow the catch in the `get_supported_msa_mqa_backend` path around
`torch.cuda.get_device_capability()` so it only handles the expected
CUDA-related failure. Update the `try`/`except` in `msa_proxy_mqa.py` to catch
`RuntimeError` (or the specific CUDA error type used here) instead of
`Exception`, while keeping the fallback `return False` behavior unchanged.
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 around torch.cuda.get_device_capability(); narrow the catch from
Exception to RuntimeError to match the existing MsaProxyMqaFmha.is_available
behavior. Update the try/except block in that method so only RuntimeError is
swallowed and the function still returns False in that case.
In
`@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/topk.py`:
- Line 65: The unpack in decode_wrapper/topk.py is triggering a Ruff
unused-unpack warning because the first value from max_score_kv.shape is never
used. Update the assignment in the topk.py logic that unpacks max_score_kv.shape
so the unused num_kv_heads binding is marked intentionally unused, while keeping
the remaining max_k_tiles and total_q values unchanged.
In
`@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/decode_wrapper/worklist.py`:
- Around line 62-68: In the worklist validation logic, add an explicit check for
num_ctas before it is used as a divisor so invalid values are rejected early
with a clear ValueError instead of triggering a divide-by-zero later. Update the
validation block in the worklist builder/initializer near the existing
batch_size and num_packed_heads checks, and ensure the error message clearly
states that num_ctas must be positive and references the invalid value.
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_plan_cache.py`:
- Around line 185-186: The global MSA geometry initialization in MSA plan cache
currently treats any later assignment as a no-op, which can silently accept
conflicting geometry across models in the same process. Update the logic around
_GLOBAL_MSA_GEOMETRY in msa_plan_cache.py so the first write still sets the
value, but any subsequent non-identical geometry is explicitly rejected with an
error or assertion instead of being ignored, and keep identical repeats as
no-ops.
In `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 1217-1230: The MSA geometry on the metadata class is only
initialized once, so a later MiniMax-M3 model can reuse stale sparse dimensions
from a previous config. Update the logic around the attn_metadata._msa_geometry
check in modeling_minimaxm3.py so it refreshes the class-level
MsaPlanCacheGeometry when the current self.attn.m3_config differs from the
existing geometry, rather than only checking for None. Use the existing symbols
attn_metadata, _msa_geometry, and MsaPlanCacheGeometry to locate and update the
assignment path.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/block_sparse.py`:
- Around line 96-131: The docstring on forward_block_sparse is using NumPy-style
sections instead of the required Google style. Update the function documentation
to use Google-style headings like Args and Returns, keeping the same parameter
descriptions for q, k_paged, v_paged, kv_block_indexes, qo_lens_cpu,
kv_lens_cpu, qo_offset_cpu, kv_indices, sm_scale, and causal so the docstring
remains Sphinx-parseable.
In `@tensorrt_llm/_torch/attention_backend/fmha/indexer_proxy.py`:
- Around line 92-125: The docstring on the max-score helper in indexer_proxy.py
is using NumPy-style sections instead of the required Google style. Update the
docstring for the proxy method that computes the per-(qo_head, kv_tile)
max-score tensor to use Google-style Args and Returns formatting, matching the
style used elsewhere such as in block_sparse.py, while preserving the same
parameter descriptions and output shape details.
In `@tensorrt_llm/_torch/attention_backend/fmha/msa_proxy_mqa.py`:
- Around line 58-86: The MsaProxyMqaFmha.is_available() method duplicates the
same availability checks used by MsaSparseGqaFmha.is_available(), so refactor
the shared CUDA/SM100 probing logic into a common helper such as
_is_fmha_sm100_available() in a shared module and have both classes call it.
Keep the existing fmha_sm100 import probe, CUDA availability check, and
compute-capability gate in one place, and preserve the class-specific debug
messages only if needed around that shared result to prevent the two
implementations from drifting.
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py`:
- Around line 67-86: Several newly added helper and factory functions are
missing explicit return annotations, including _require_msa_module, the FMHA
selector helpers, the backend factory, the nested __init__ methods, and
get_minimax_m3_attention_backend_cls_with_msa. Update each of these definitions
to include the correct return type hint, using None for
constructors/initializers and the appropriate concrete return types for module
selectors and factory helpers, so the new MSA backend APIs follow the project’s
annotation guidelines.
In `@tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py`:
- Around line 18-21: The test module is still importing and using the legacy
typing alias `List` instead of the built-in generic `list`, so update the type
annotation in `test_minimax_m3_msa_backend` to use `list[int] | None` and remove
the `typing.List` import. Keep the change consistent anywhere the same
annotation pattern appears in the affected test block, using the existing `from
__future__ import annotations` support and the relevant test helper/function
names in that section to locate the type hints.
🪄 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: 2df7663a-c7b9-4479-8b8c-ca69602a7dc9
📒 Files selected for processing (24)
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/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/models/modeling_minimaxm3.pytensorrt_llm/llmapi/llm_args.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/qa/llm_function_core.txttests/unittest/_torch/attention/sparse/test_minimax_m3_decode_driver_vs_msa.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
Comment on lines
+76
to
+79
| try: | ||
| major, _ = torch.cuda.get_device_capability() | ||
| except Exception: | ||
| return False |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Narrow the exception type.
Catching blind Exception around get_device_capability() masks unrelated bugs. torch.cuda.get_device_capability() only raises RuntimeError/CUDA errors in practice; catch that specifically.
🔧 Proposed fix
- try:
- major, _ = torch.cuda.get_device_capability()
- except Exception:
- return False
+ try:
+ major, _ = torch.cuda.get_device_capability()
+ except RuntimeError:
+ return False📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| try: | |
| major, _ = torch.cuda.get_device_capability() | |
| except Exception: | |
| return False | |
| try: | |
| major, _ = torch.cuda.get_device_capability() | |
| except RuntimeError: | |
| return False |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 78-78: Do not catch blind exception: Exception
(BLE001)
🤖 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 76
- 79, The exception handling in the capability check is too broad; narrow the
catch in the `get_supported_msa_mqa_backend` path around
`torch.cuda.get_device_capability()` so it only handles the expected
CUDA-related failure. Update the `try`/`except` in `msa_proxy_mqa.py` to catch
`RuntimeError` (or the specific CUDA error type used here) instead of
`Exception`, while keeping the fallback `return False` behavior unchanged.
Sources: Coding guidelines, Linters/SAST tools
Comment on lines
+75
to
+78
| try: | ||
| major, _ = torch.cuda.get_device_capability() | ||
| except Exception: | ||
| return False |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Narrow the exception type.
Same issue as MsaProxyMqaFmha.is_available — catching blind Exception around get_device_capability() should be narrowed to RuntimeError.
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 77-77: Do not catch blind exception: Exception
(BLE001)
🤖 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_sparse_gqa.py` around lines 75
- 78, The exception handling in MsaSparseGqaFmha.is_available is too broad
around torch.cuda.get_device_capability(); narrow the catch from Exception to
RuntimeError to match the existing MsaProxyMqaFmha.is_available behavior. Update
the try/except block in that method so only RuntimeError is swallowed and the
function still returns False in that case.
Sources: Coding guidelines, Linters/SAST tools
| ``-1`` padded at the tail (the sparse FMHA kernel's | ||
| ``kv_block_indexes`` contract). | ||
| """ | ||
| num_kv_heads, max_k_tiles, total_q = max_score_kv.shape |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Address the Ruff unused-unpack warning.
num_kv_heads is unpacked but never used; prefix it as intentionally unused.
Proposed fix
- num_kv_heads, max_k_tiles, total_q = max_score_kv.shape
+ _num_kv_heads, max_k_tiles, total_q = max_score_kv.shape📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| num_kv_heads, max_k_tiles, total_q = max_score_kv.shape | |
| _num_kv_heads, max_k_tiles, total_q = max_score_kv.shape |
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 65-65: Unpacked variable num_kv_heads is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🤖 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/topk.py`
at line 65, The unpack in decode_wrapper/topk.py is triggering a Ruff
unused-unpack warning because the first value from max_score_kv.shape is never
used. Update the assignment in the topk.py logic that unpacks max_score_kv.shape
so the unused num_kv_heads binding is marked intentionally unused, while keeping
the remaining max_k_tiles and total_q values unchanged.
Source: Linters/SAST tools
Comment on lines
+62
to
+68
| if batch_size <= 0: | ||
| raise ValueError(f"batch_size must be positive, got {batch_size}") | ||
| if batch_size > 0xFFFF: | ||
| raise ValueError(f"batch_size {batch_size} exceeds the 16-bit work-item field") | ||
| if num_packed_heads <= 0 or num_packed_heads > 0xFFFF: | ||
| raise ValueError(f"num_packed_heads out of range: {num_packed_heads}") | ||
|
|
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Validate num_ctas before using it as a divisor.
num_ctas <= 0 currently reaches Line 78 and fails with a less actionable divide-by-zero error.
Proposed fix
if num_packed_heads <= 0 or num_packed_heads > 0xFFFF:
raise ValueError(f"num_packed_heads out of range: {num_packed_heads}")
+ if num_ctas <= 0:
+ raise ValueError(f"num_ctas must be positive, got {num_ctas}")
n_items = batch_size * num_packed_headsAlso applies to: 78-78
🤖 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/worklist.py`
around lines 62 - 68, In the worklist validation logic, add an explicit check
for num_ctas before it is used as a divisor so invalid values are rejected early
with a clear ValueError instead of triggering a divide-by-zero later. Update the
validation block in the worklist builder/initializer near the existing
batch_size and num_packed_heads checks, and ensure the error message clearly
states that num_ctas must be positive and references the invalid value.
Comment on lines
+185
to
+186
| if _GLOBAL_MSA_GEOMETRY is None: | ||
| _GLOBAL_MSA_GEOMETRY = geometry |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject conflicting process-wide MSA geometry.
The docstring says later writes are identical no-ops, but the code silently ignores non-identical geometry. A second model/config in the same process could then stage CUDA-graph metadata with stale geometry.
Proposed fix
global _GLOBAL_MSA_GEOMETRY
if _GLOBAL_MSA_GEOMETRY is None:
_GLOBAL_MSA_GEOMETRY = geometry
+ elif _GLOBAL_MSA_GEOMETRY != geometry:
+ raise ValueError(
+ f"Conflicting MSA geometry registration: existing={_GLOBAL_MSA_GEOMETRY}, "
+ f"new={geometry}"
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if _GLOBAL_MSA_GEOMETRY is None: | |
| _GLOBAL_MSA_GEOMETRY = geometry | |
| global _GLOBAL_MSA_GEOMETRY | |
| if _GLOBAL_MSA_GEOMETRY is None: | |
| _GLOBAL_MSA_GEOMETRY = geometry | |
| elif _GLOBAL_MSA_GEOMETRY != geometry: | |
| raise ValueError( | |
| f"Conflicting MSA geometry registration: existing={_GLOBAL_MSA_GEOMETRY}, " | |
| f"new={geometry}" | |
| ) |
🤖 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_plan_cache.py`
around lines 185 - 186, The global MSA geometry initialization in MSA plan cache
currently treats any later assignment as a no-op, which can silently accept
conflicting geometry across models in the same process. Update the logic around
_GLOBAL_MSA_GEOMETRY in msa_plan_cache.py so the first write still sets the
value, but any subsequent non-identical geometry is explicitly rejected with an
error or assertion instead of being ignored, and keep identical repeats as
no-ops.
Comment on lines
+1217
to
+1230
| if getattr(attn_metadata, "_msa_geometry", None) is None: | ||
| from ..attention_backend.sparse.minimax_m3.msa_plan_cache import MsaPlanCacheGeometry | ||
|
|
||
| m3_config = self.attn.m3_config | ||
| type(attn_metadata)._msa_geometry = MsaPlanCacheGeometry( | ||
| num_q_heads=int(m3_config.num_q_heads), | ||
| num_kv_heads=int(m3_config.num_kv_heads), | ||
| num_index_heads=int(m3_config.num_index_heads), | ||
| head_dim=int(m3_config.head_dim), | ||
| block_size=int(m3_config.block_size), | ||
| topk=int(m3_config.topk), | ||
| init_blocks=int(m3_config.init_blocks), | ||
| local_blocks=int(m3_config.local_blocks), | ||
| ) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Refresh stale metadata-class MSA geometry.
The one-time None guard can keep an old class-level _msa_geometry when a later MiniMax-M3 model/config in the same process has different sparse geometry. prepare() may then build MSA plans with stale dimensions before this forward runs.
Proposed fix
- if getattr(attn_metadata, "_msa_geometry", None) is None:
- from ..attention_backend.sparse.minimax_m3.msa_plan_cache import MsaPlanCacheGeometry
-
- m3_config = self.attn.m3_config
- type(attn_metadata)._msa_geometry = MsaPlanCacheGeometry(
- num_q_heads=int(m3_config.num_q_heads),
- num_kv_heads=int(m3_config.num_kv_heads),
- num_index_heads=int(m3_config.num_index_heads),
- head_dim=int(m3_config.head_dim),
- block_size=int(m3_config.block_size),
- topk=int(m3_config.topk),
- init_blocks=int(m3_config.init_blocks),
- local_blocks=int(m3_config.local_blocks),
- )
+ from ..attention_backend.sparse.minimax_m3.msa_plan_cache import MsaPlanCacheGeometry
+
+ m3_config = self.attn.m3_config
+ msa_geometry = MsaPlanCacheGeometry(
+ num_q_heads=int(m3_config.num_q_heads),
+ num_kv_heads=int(m3_config.num_kv_heads),
+ num_index_heads=int(m3_config.num_index_heads),
+ head_dim=int(m3_config.head_dim),
+ block_size=int(m3_config.block_size),
+ topk=int(m3_config.topk),
+ init_blocks=int(m3_config.init_blocks),
+ local_blocks=int(m3_config.local_blocks),
+ )
+ if getattr(type(attn_metadata), "_msa_geometry", None) != msa_geometry:
+ type(attn_metadata)._msa_geometry = msa_geometry
+ if getattr(m3_meta, "msa_plans", None) is not None:
+ m3_meta.msa_plans = None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if getattr(attn_metadata, "_msa_geometry", None) is None: | |
| from ..attention_backend.sparse.minimax_m3.msa_plan_cache import MsaPlanCacheGeometry | |
| m3_config = self.attn.m3_config | |
| type(attn_metadata)._msa_geometry = MsaPlanCacheGeometry( | |
| num_q_heads=int(m3_config.num_q_heads), | |
| num_kv_heads=int(m3_config.num_kv_heads), | |
| num_index_heads=int(m3_config.num_index_heads), | |
| head_dim=int(m3_config.head_dim), | |
| block_size=int(m3_config.block_size), | |
| topk=int(m3_config.topk), | |
| init_blocks=int(m3_config.init_blocks), | |
| local_blocks=int(m3_config.local_blocks), | |
| ) | |
| from ..attention_backend.sparse.minimax_m3.msa_plan_cache import MsaPlanCacheGeometry | |
| m3_config = self.attn.m3_config | |
| msa_geometry = MsaPlanCacheGeometry( | |
| num_q_heads=int(m3_config.num_q_heads), | |
| num_kv_heads=int(m3_config.num_kv_heads), | |
| num_index_heads=int(m3_config.num_index_heads), | |
| head_dim=int(m3_config.head_dim), | |
| block_size=int(m3_config.block_size), | |
| topk=int(m3_config.topk), | |
| init_blocks=int(m3_config.init_blocks), | |
| local_blocks=int(m3_config.local_blocks), | |
| ) | |
| if getattr(type(attn_metadata), "_msa_geometry", None) != msa_geometry: | |
| type(attn_metadata)._msa_geometry = msa_geometry | |
| if getattr(m3_meta, "msa_plans", None) is not None: | |
| m3_meta.msa_plans = None |
🤖 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 1217 - 1230,
The MSA geometry on the metadata class is only initialized once, so a later
MiniMax-M3 model can reuse stale sparse dimensions from a previous config.
Update the logic around the attn_metadata._msa_geometry check in
modeling_minimaxm3.py so it refreshes the class-level MsaPlanCacheGeometry when
the current self.attn.m3_config differs from the existing geometry, rather than
only checking for None. Use the existing symbols attn_metadata, _msa_geometry,
and MsaPlanCacheGeometry to locate and update the assignment path.
1 task
Collaborator
|
Closing in favor of #15998. |
brb-nv
added a commit
to brb-nv/TensorRT-LLM
that referenced
this pull request
Jul 7, 2026
…5809) Squashed cherry-pick of NVIDIA/TensorRT-LLM PR NVIDIA#15809, combining: - [None][feat] Integrate M3 sparse attention kernels form MSA - [None][fix] Fix parallel config of Minimax M3 MXFP8 test - [TRTLLM-14019][feat] Support MSA kernel with enable cuda graph - [None][fix] Fix M3 attention semantic issue Conflicts with the already-integrated PR NVIDIA#15923 were resolved as follows: - requirements.txt: keep flashinfer-python==0.6.14 (main) and add the fmha_sm100 (MSA) dependency. - modeling_minimaxm3.py: graft the _msa_geometry publishing block onto the current _sparse_attention_core (which returns self.attn.forward(..., output=output, ...)). - llm_function_core.txt: keep the MiniMax M3 piecewise/nvfp4 entries. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
added a commit
to brb-nv/TensorRT-LLM
that referenced
this pull request
Jul 7, 2026
…5809) Squashed cherry-pick of NVIDIA/TensorRT-LLM PR NVIDIA#15809, combining: - [None][feat] Integrate M3 sparse attention kernels form MSA - [None][fix] Fix parallel config of Minimax M3 MXFP8 test - [TRTLLM-14019][feat] Support MSA kernel with enable cuda graph - [None][fix] Fix M3 attention semantic issue Conflicts with the already-integrated PR NVIDIA#15923 were resolved as follows: - requirements.txt: keep flashinfer-python==0.6.14 (main) and add the fmha_sm100 (MSA) dependency. - modeling_minimaxm3.py: graft the _msa_geometry publishing block onto the current _sparse_attention_core (which returns self.attn.forward(..., output=output, ...)). - llm_function_core.txt: keep the MiniMax M3 piecewise/nvfp4 entries. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
added a commit
to brb-nv/TensorRT-LLM
that referenced
this pull request
Jul 8, 2026
…5809) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
added a commit
to brb-nv/TensorRT-LLM
that referenced
this pull request
Jul 8, 2026
…5809) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
dsingal0
pushed a commit
to dsingal0/TensorRT-LLM
that referenced
this pull request
Jul 9, 2026
…5809) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.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
2 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.
Summary by CodeRabbit
New Features
Bug Fixes
Description
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.