[None][fix] Honor MSA block-index strides for q_len <= 32 - #17284
Draft
peihu-nv wants to merge 8 commits into
Draft
[None][fix] Honor MSA block-index strides for q_len <= 32#17284peihu-nv wants to merge 8 commits into
peihu-nv wants to merge 8 commits into
Conversation
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
peihu-nv
requested review from
BowenFu,
brnguyen2,
kaiyux,
kris1025,
laikhtewari,
pengbowang-nv,
tburt-nv and
yuanjingx87
August 5, 2026 05:45
…-paged-kv-main Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> # Conflicts: # tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
peihu-nv
force-pushed
the
peihengh/m3-msa-selector-stride-trtllm
branch
from
August 5, 2026 05:49
d6e0b0b to
a92acdf
Compare
peihu-nv
marked this pull request as draft
August 5, 2026 05:51
Contributor
WalkthroughChangesMSA is fetched and packaged under MSA packaging and dependency acquisition
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant sparse_atten_func
participant FwdRunner
participant Sm100FmhaLoadTmaWarpspecialized
participant SparseMainloop
sparse_atten_func->>FwdRunner: pass prepared paged K/V tensors and block-index strides
FwdRunner->>Sm100FmhaLoadTmaWarpspecialized: construct arguments with 64-bit strides
Sm100FmhaLoadTmaWarpspecialized->>SparseMainloop: load sparse pages using strided offsets
SparseMainloop-->>sparse_atten_func: return sparse attention output
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
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: 4
🧹 Nitpick comments (1)
3rdparty/patches/msa_strided_paged_kv.patch (1)
273-282: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename
kv_block_indexes_numelto match its new meaning.This field no longer holds an element count. It now holds the maximum linear offset span of a possibly strided view. A future reader who divides by a stride or iterates
[0, numel)as dense elements will be wrong. Rename it tokv_block_indexes_extent(or_span) across the params struct, the load-mainloopArgs/Params, and this assignment. The bounds check itself stays correct, because the span is a valid conservative upper bound for every reachable offset.🤖 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 `@3rdparty/patches/msa_strided_paged_kv.patch` around lines 273 - 282, Rename kv_block_indexes_numel to kv_block_indexes_extent throughout the params struct, load-mainloop Args/Params, and the assignment shown here, preserving the existing strided offset calculation and bounds-check 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 `@3rdparty/patches/msa_strided_paged_kv.patch`:
- Around line 247-249: Initialize kv_block_indexes_stride_q,
kv_block_indexes_stride_h, and kv_block_indexes_stride_k to 0 in their
declarations within FMHACutlassSM100Params, matching the existing in-class
defaults and ensuring unset strides produce a zero offset.
In `@tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py`:
- Around line 165-171: Update
test_msa_sparse_attention_honors_noncontiguous_block_indexes to add a pytest
skip condition requiring get_sm_version() == 100 alongside the existing CUDA
availability guard. Leave the helper-only tests unchanged.
- Around line 83-112: Add an MSA availability guard to
test_msa_paged_kv_preserves_tma_compatible_outer_stride and the three similar
tests so environments lacking the optional fmha_sm100 package skip instead of
raising ImportError. Use pytest.importorskip("fmha_sm100") or the existing
module-level skip pattern before importing fmha_sm100.cute.interface, while
retaining the CUDA skip.
- Around line 83-217: Add package-availability and supported-architecture guards
to the CUDA-dependent tests, including
test_msa_paged_kv_preserves_tma_compatible_outer_stride,
test_msa_paged_hnd_input_materializes_unaligned_outer_stride, and
test_msa_sparse_attention_honors_noncontiguous_block_indexes, so they skip
unless fmha_sm100 is installed and the GPU is SM100/SM103. Add coverage for the
_prepare_paged_hnd_input tensor.ndim != 4 branch, preserving the existing
assertions and supported-hardware behavior.
---
Nitpick comments:
In `@3rdparty/patches/msa_strided_paged_kv.patch`:
- Around line 273-282: Rename kv_block_indexes_numel to kv_block_indexes_extent
throughout the params struct, load-mainloop Args/Params, and the assignment
shown here, preserving the existing strided offset calculation and bounds-check
behavior.
🪄 Autofix
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: dd38e8a0-411f-4cb2-824c-d98f97740817
📒 Files selected for processing (16)
.gitignore.gitmodules3rdparty/CMakeLists.txt3rdparty/MSA3rdparty/fetch_content.json3rdparty/patches/msa_strided_paged_kv.patchcpp/CMakeLists.txtdocs/source/installation/build-from-source.mdjenkins/UpdateTestDurations.groovyscripts/attribution/scan/metadata/msa.ymlscripts/build_wheel.pysetup.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
💤 Files with no reviewable changes (3)
- 3rdparty/MSA
- .gitmodules
- docs/source/installation/build-from-source.md
Comment on lines
+247
to
+249
| + int64_t kv_block_indexes_stride_q; | ||
| + int64_t kv_block_indexes_stride_h; | ||
| + int64_t kv_block_indexes_stride_k; |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Add = 0 default initializers to the three new stride fields.
The new kv_block_indexes_stride_* members have no in-class initializer, but later members in the same struct (for example pack_factor = 1) do. These strides are multiplied by batch_idx, kv_head_idx, and page_idx to form a global-memory offset into kv_block_indexes. If any producer of FMHACutlassSM100Params other than the generated FMHAVariantRun_* path leaves them unset, the kernel reads out-of-bounds memory. Default them to 0 so an unset value degrades to the single-slice offset instead of an arbitrary one.
🛡️ Proposed defaults
-+ int64_t kv_block_indexes_stride_q;
-+ int64_t kv_block_indexes_stride_h;
-+ int64_t kv_block_indexes_stride_k;
++ int64_t kv_block_indexes_stride_q = 0;
++ int64_t kv_block_indexes_stride_h = 0;
++ int64_t kv_block_indexes_stride_k = 0;📝 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
| + int64_t kv_block_indexes_stride_q; | |
| + int64_t kv_block_indexes_stride_h; | |
| + int64_t kv_block_indexes_stride_k; | |
| int64_t kv_block_indexes_stride_q = 0; | |
| int64_t kv_block_indexes_stride_h = 0; | |
| int64_t kv_block_indexes_stride_k = 0; |
🤖 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 `@3rdparty/patches/msa_strided_paged_kv.patch` around lines 247 - 249,
Initialize kv_block_indexes_stride_q, kv_block_indexes_stride_h, and
kv_block_indexes_stride_k to 0 in their declarations within
FMHACutlassSM100Params, matching the existing in-class defaults and ensuring
unset strides produce a zero offset.
Comment on lines
+83
to
+112
| pages, roles, heads, page_size, head_dim = 5, 2, 2, 128, 128 | ||
| pool = torch.empty( | ||
| pages, | ||
| roles, | ||
| heads, | ||
| page_size, | ||
| head_dim, | ||
| dtype=torch.bfloat16, | ||
| ) | ||
| hnd_cache = pool[:, 0] | ||
|
|
||
| class FakeCacheManager: | ||
| def __init__(self): | ||
| self.calls = [] | ||
|
|
||
| def get_index_k_buffer(self, layer_idx, kv_layout="NHD"): | ||
| self.calls.append((layer_idx, kv_layout)) | ||
| return hnd_cache | ||
|
|
||
| manager = FakeCacheManager() | ||
| metadata.kv_cache_manager = manager | ||
| metadata.msa_out_cache_loc = torch.tensor([2, page_size + 5], dtype=torch.int32) | ||
| values = torch.arange(2 * head_dim, dtype=torch.float32).reshape(2, 1, head_dim) | ||
|
|
||
| returned = metadata.msa_idx_k_cache(3) | ||
| metadata.msa_write_idx_k(3, values) | ||
|
|
||
| assert returned.data_ptr() == hnd_cache.data_ptr() | ||
| assert not returned.is_contiguous() | ||
| assert manager.calls == [(3, "HND"), (3, "HND")] | ||
| torch.testing.assert_close(hnd_cache[0, 0, 2], values[0, 0].to(torch.bfloat16)) | ||
| torch.testing.assert_close(hnd_cache[1, 0, 5], values[1, 0].to(torch.bfloat16)) | ||
|
|
||
|
|
||
| def test_msa_indexer_preserves_strided_hnd_index_k(monkeypatch): | ||
| import tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_indexer as indexer_module | ||
| from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import MiniMaxM3SparseConfig | ||
|
|
||
| config = MiniMaxM3SparseConfig( | ||
| num_q_heads=4, | ||
| num_kv_heads=1, | ||
| head_dim=128, | ||
| num_index_heads=4, | ||
| sparse_index_dim=128, | ||
| block_size=128, | ||
| topk=16, | ||
| ) | ||
| indexer = indexer_module.MsaIndexer(config) | ||
| pool = torch.randn(2, 7, 1, 128, 128, dtype=torch.bfloat16) | ||
| idx_k_paged = pool[:, 0] | ||
| captured = {} | ||
|
|
||
| def fake_proxy_max_score(idx_q, passed_idx_k, **kwargs): | ||
| del kwargs | ||
| captured["idx_k"] = passed_idx_k | ||
| return torch.zeros(4, 2, idx_q.shape[0]) | ||
|
|
||
| expected = torch.zeros(1, 1, 16, dtype=torch.int32) | ||
|
|
||
| def fake_select_blocks_from_maxscore(*args, **kwargs): | ||
| del args, kwargs | ||
| return expected | ||
|
|
||
| monkeypatch.setattr(indexer_module, "_proxy_max_score", fake_proxy_max_score) | ||
| monkeypatch.setattr( | ||
| indexer_module, | ||
| "select_blocks_from_maxscore", | ||
| fake_select_blocks_from_maxscore, | ||
| dtype=torch.float8_e4m3fn, | ||
| device="cuda", | ||
| ) | ||
| kv_cache_manager = Mock() | ||
| kv_cache_manager.get_buffers.return_value = pool | ||
|
|
||
| result = indexer.select_blocks( | ||
| torch.zeros(1, 4, 128, dtype=torch.bfloat16), | ||
| idx_k_paged, | ||
| idx_sm_scale=128**-0.5, | ||
| kv_indices=torch.arange(2, dtype=torch.int32), | ||
| qo_lens_cpu=torch.tensor([1], dtype=torch.int32), | ||
| kv_lens_cpu=torch.tensor([256], dtype=torch.int32), | ||
| qo_offset_cpu=torch.tensor([255], dtype=torch.int32), | ||
| ) | ||
| k_view, v_view = msa_paged_kv(kv_cache_manager, layer_idx=3) | ||
|
|
||
| assert captured["idx_k"] is idx_k_paged | ||
| assert captured["idx_k"].data_ptr() == idx_k_paged.data_ptr() | ||
| assert not captured["idx_k"].is_contiguous() | ||
| assert result is expected | ||
| kv_cache_manager.get_buffers.assert_called_once_with(3, kv_layout="HND") | ||
| for view in (k_view, v_view): | ||
| assert not view.is_contiguous() | ||
| prepared = sparse_interface._prepare_paged_hnd_input(view, page_size) | ||
| assert prepared.data_ptr() == view.data_ptr() | ||
| assert prepared.stride() == view.stride() | ||
|
|
||
| mismatched = sparse_interface._prepare_paged_hnd_input(k_view, page_size // 2) | ||
| assert mismatched.data_ptr() == k_view.data_ptr() | ||
| with pytest.raises(ValueError, match="page_size == blk_kv"): | ||
| sparse_interface._prepare_paged_kv_for_tma(mismatched, mismatched, page_size // 2) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Add an MSA-availability skip in addition to the CUDA skip.
This test is gated only on torch.cuda.is_available(). Line 85 then imports fmha_sm100.cute.interface. fmha_sm100 is an optional packaged wheel, so on a CUDA machine without it installed this test errors with ImportError instead of skipping. The same pattern repeats at lines 117, 145, and 172. Use pytest.importorskip("fmha_sm100") or a module-level skip helper.
🤖 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 83 - 112, Add an MSA availability guard to
test_msa_paged_kv_preserves_tma_compatible_outer_stride and the three similar
tests so environments lacking the optional fmha_sm100 package skip instead of
raising ImportError. Use pytest.importorskip("fmha_sm100") or the existing
module-level skip pattern before importing fmha_sm100.cute.interface, while
retaining the CUDA skip.
Comment on lines
+83
to
+217
|
|
||
| assert prepared.is_contiguous() | ||
| assert prepared.data_ptr() != view.data_ptr() | ||
|
|
||
|
|
||
| def _run_msa_sparse_attention( | ||
| q: torch.Tensor, | ||
| k: torch.Tensor, | ||
| v: torch.Tensor, | ||
| plan: tuple, | ||
| kv_indices: torch.Tensor, | ||
| selected: torch.Tensor, | ||
| ) -> torch.Tensor: | ||
| from fmha_sm100.api import fmha_sm100 | ||
|
|
||
| out = torch.empty_like(q) | ||
| returned, _ = fmha_sm100( | ||
| q, | ||
| k, | ||
| v, | ||
| plan, | ||
| kv_indices=kv_indices, | ||
| kv_block_indexes=selected, | ||
| out=out, | ||
| sm_scale=128.0**-0.5, | ||
| output_maxscore=False, | ||
| ) | ||
| kwargs = { | ||
| "qo_lens_cpu": torch.ones_like(kv_lens_cpu), | ||
| "kv_lens_cpu": kv_lens_cpu, | ||
| "qo_offset_cpu": kv_lens_cpu - 1, | ||
| "kv_indices": torch.arange(num_pages, device="cuda", dtype=torch.int32), | ||
| "sm_scale": head_dim**-0.5, | ||
| "causal": True, | ||
| } | ||
|
|
||
| strided_scores = _proxy_max_score(index_q, index_k_strided, **kwargs) | ||
| packed_scores = _proxy_max_score(index_q, index_k_packed, **kwargs) | ||
| torch.cuda.synchronize() | ||
| assert returned.data_ptr() == out.data_ptr() | ||
| assert torch.isfinite(out).all() | ||
| return out | ||
|
|
||
|
|
||
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") | ||
| @pytest.mark.parametrize("query_len", [1, 5]) | ||
| @pytest.mark.parametrize("num_kv_splits", [1, 2]) | ||
| def test_msa_sparse_attention_honors_noncontiguous_block_indexes( | ||
| query_len: int, | ||
| num_kv_splits: int, | ||
| ) -> None: | ||
| from fmha_sm100.api import fmha_sm100_plan | ||
|
|
||
| device = torch.device("cuda", 0) | ||
| torch.cuda.set_device(device) | ||
| total_selector_rows = query_len + 2 | ||
| q = torch.zeros((query_len, 32, 128), dtype=torch.bfloat16, device=device) | ||
| k = torch.zeros((8, 2, 128, 128), dtype=torch.bfloat16, device=device) | ||
| v = torch.empty_like(k) | ||
| for page in range(8): | ||
| v[page].fill_(page + 1) | ||
| kv_indices = torch.arange(8, dtype=torch.int32, device=device) | ||
|
|
||
| intended = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device) | ||
| poison = torch.tensor([4, 5, 6, 7], dtype=torch.int32, device=device) | ||
| logical = intended.expand(total_selector_rows, 2, 4).clone() | ||
| # These rows are outside the logical slice but occupy addresses reached by | ||
| # a pointer-only contiguous read for later query/head pairs. | ||
| logical[:2, 1, :] = poison | ||
| backing = logical.permute(1, 0, 2).contiguous() | ||
| selected_strided = backing.permute(1, 0, 2)[-query_len:] | ||
| selected_contiguous = logical[-query_len:].contiguous() | ||
|
|
||
| assert selected_strided.stride() == (4, total_selector_rows * 4, 1) | ||
| assert not selected_strided.is_contiguous() | ||
| assert torch.equal(selected_strided, selected_contiguous) | ||
|
|
||
| plan = fmha_sm100_plan( | ||
| torch.tensor([query_len], dtype=torch.int32), | ||
| torch.tensor([1024], dtype=torch.int32), | ||
| 32, | ||
| num_kv_heads=2, | ||
| qo_offset=torch.tensor([1024 - query_len], dtype=torch.int32), | ||
| num_kv_splits=num_kv_splits, | ||
| page_size=128, | ||
| output_maxscore=False, | ||
| kv_block_num=4, | ||
| causal=True, | ||
| device=device, | ||
| ) | ||
| short_plan = plan[3] | ||
| assert short_plan["MM-SA-Nv"] is False | ||
| assert short_plan["num_kv_splits"] == num_kv_splits | ||
|
|
||
| assert not index_k_strided.is_contiguous() | ||
| assert index_k_strided.stride(0) == coalescing_scale * page_size * head_dim | ||
| assert torch.equal(strided_scores, packed_scores) | ||
| expected = _run_msa_sparse_attention(q, k, v, plan, kv_indices, selected_contiguous) | ||
| actual = _run_msa_sparse_attention(q, k, v, plan, kv_indices, selected_strided) | ||
| torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether the MSA backend test module is registered in CI/QA test lists.
set -euo pipefail
fd -t d 'test_lists' tests | head
rg -rn 'test_minimax_m3_msa_backend|minimax_m3' tests/integration/test_lists 2>/dev/null || \
rg -n 'test_minimax_m3_msa_backend|minimax_m3' -g '**/test_lists/**' . || \
echo "No test-list entry found for test_minimax_m3_msa_backend"Repository: NVIDIA/TensorRT-LLM
Length of output: 1349
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- changed files ---'
git diff --name-status
printf '%s\n' '--- test module structure and imports ---'
ast-grep outline tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
sed -n '1,260p' tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
printf '%s\n' '--- test-list references ---'
rg -n -i 'minimax|msa|attention/sparse|_torch/attention|test_minimax_m3_msa_backend' tests/integration/test_lists || true
printf '%s\n' '--- optional dependency and architecture guards in nearby tests ---'
rg -n -i 'fmha_sm100|SM100|sm100|cuda.get_device_capability|importorskip|find_spec|skipif' tests/unittest/_torch/attention/sparse tests/unittest/_torch/attention 2>/dev/null | head -240
printf '%s\n' '--- preparation helper definitions and callers ---'
rg -n '_prepare_paged_hnd_input|_prepare_paged_kv_for_tma|def msa_paged_kv|msa_paged_kv' . -g '*.py' -g '*.cu' -g '*.cpp' -g '*.h' -g '*.hpp' | head -240Repository: NVIDIA/TensorRT-LLM
Length of output: 46467
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- MSA utility implementation ---'
sed -n '1,260p' tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py
printf '%s\n' '--- FMHA preparation implementation references ---'
rg -n -C 8 '_prepare_paged_hnd_input|_prepare_paged_kv_for_tma' . --glob '!tests/**' || true
printf '%s\n' '--- availability and import behavior ---'
sed -n '1,220p' tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.py
rg -n -C 5 'fmha_sm100' tensorrt_llm pyproject.toml setup.py requirements* docker .github 2>/dev/null | head -240
printf '%s\n' '--- exact relevant test-list entries ---'
sed -n '1,35p' tests/integration/test_lists/test-db/l0_b300.yml
sed -n '1,35p' tests/integration/test_lists/test-db/l0_dgx_b300.yml
rg -n -C 3 'unittest/_torch/attention' tests/integration/test_lists/test-db/l0_*.yml
printf '%s\n' '--- SM100 helper conventions ---'
sed -n '1,45p' tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py
sed -n '1,45p' tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
rg -n 'def getSMVersion|def get_sm_version|from .*getSMVersion|from .*get_sm_version' tests tensorrt_llm | head -100Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py")
tree = ast.parse(path.read_text())
for node in tree.body:
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
decorators = [
ast.unparse(dec) for dec in node.decorator_list
]
print(f"{node.name}: decorators={decorators}")
PY
printf '%s\n' '--- packaged MSA files tracked in the repository ---'
git ls-files | rg '(^|/)fmha_sm100(/|$)' | head -80 || true
printf '%s\n' '--- MSA FMHA availability gate ---'
sed -n '180,215p' tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py
printf '%s\n' '--- test-list path handling ---'
sed -n '70,145p' tests/integration/defs/test_list_parser.py
printf '%s\n' '--- all exact module references in test lists ---'
rg -n -F 'tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py' tests/integration/test_lists || true
rg -n -F 'unittest/_torch/attention' tests/integration/test_lists/test-db/l0_b300.yml tests/integration/test_lists/test-db/l0_dgx_b300.ymlRepository: NVIDIA/TensorRT-LLM
Length of output: 7343
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py")
tree = ast.parse(path.read_text())
for node in tree.body:
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
decorators = [ast.unparse(dec) for dec in node.decorator_list]
print(f"{node.name}: decorators={decorators}")
PY
printf '%s\n' '--- packaged MSA files tracked in the repository ---'
git ls-files | rg '(^|/)fmha_sm100(/|$)' | head -80 || true
printf '%s\n' '--- MSA FMHA availability gate ---'
sed -n '180,215p' tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py
printf '%s\n' '--- test-list path handling ---'
sed -n '70,145p' tests/integration/defs/test_list_parser.py
printf '%s\n' '--- all exact module references in test lists ---'
rg -n -F 'tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py' tests/integration/test_lists || true
rg -n -F 'unittest/_torch/attention' tests/integration/test_lists/test-db/l0_b300.yml tests/integration/test_lists/test-db/l0_dgx_b300.ymlRepository: NVIDIA/TensorRT-LLM
Length of output: 7343
🌐 Web query:
GitHub NVIDIA TensorRT-LLM "_prepare_paged_hnd_input"
💡 Result:
Based on a search of the NVIDIA TensorRT-LLM repository, there is no public function or method named _prepare_paged_hnd_input in the codebase [1][2][3][4][5]. It is highly likely that this term refers to a internal, project-specific, or recently introduced concept associated with "HND" (often referring to HND-format paged cache, used in advanced attention kernels like those in the MSA submodule) that has not been exposed in the main public API or documentation [1]. Relevant technical context includes: - HND Paged Cache: Recent developments in TensorRT-LLM involve optimizing the handling of HND (a format related to paged K/V cache layouts) to support zero-copy paths for attention kernels, particularly to satisfy TMA (Tensor Memory Accelerator) requirements [1]. - Input Preparation: TensorRT-LLM frequently uses internal methods like _prepare_inputs or _prepare_tp_inputs to manage runtime data, input pre-processing, and host-to-device transfers [2][5][6]. These methods are often targets for performance optimization [6]. - Paged Attention: The library relies heavily on paged KV cache management to enable features like chunked context, request scheduling, and KV cache reuse [7][8][9]. If you are encountering this term in a specific trace, error log, or experimental branch, it may be a private method within a specialized plugin or a very recent PR-specific implementation not yet merged into the main branch [1].
Citations:
- 1: [None][perf] Avoid paged MSA K/V materialization during prefill #16857
- 2: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/runtime/generation.py
- 3: https://github.com/NVIDIA/TensorRT-LLM/blob/7a8bd87f/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py
- 4: https://github.com/NVIDIA/TensorRT-LLM/blob/7a8bd87f/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
- 5: https://github.com/NVIDIA/TensorRT-LLM/blob/656091bf/tensorrt_llm/_torch/pyexecutor/model_engine.py
- 6: [TRTLLM-11148][perf] _prepare_inputs host time optimization #11704
- 7: https://nvidia.github.io/TensorRT-LLM/features/paged-attention-ifb-scheduler.html
- 8: https://nvidia.github.io/TensorRT-LLM/0.20.0rc3/performance/performance-tuning-guide/useful-build-time-flags.html
- 9: https://nvidia.github.io/TensorRT-LLM/advanced/kv-cache-reuse.html
🌐 Web query:
site:github.com/NVIDIA/TensorRT-LLM/pull/16857 "_prepare_paged_hnd_input" "ndim"
💡 Result:
PR #16857 defines _prepare_paged_hnd_input with this ndim guard:
if tensor.ndim != 4:
return tensor.contiguous()It preserves 4-D HND paged views only when page packing, strides, and 16-byte alignment satisfy TMA requirements; otherwise it materializes a contiguous copy. (github.com)
Citations:
Add MSA package and architecture guards to the CUDA tests.
- Added three tests and one helper. No tests were removed.
- The module is covered by CI directory entries in
l0_h100.yml,l0_b200.yml,l0_b300.yml,l0_dgx_b300.yml, andl0_gb300_multi_gpus.yml. No QA-list entry exists. - Coverage is insufficient. The tests skip only when CUDA is unavailable, but
fmha_sm100requires the package and SM100/SM103. Add both guards. Add a case for the_prepare_paged_hnd_inputtensor.ndim != 4branch. Runpytest tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pyon supported hardware.
🤖 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 83 - 217, Add package-availability and supported-architecture guards to
the CUDA-dependent tests, including
test_msa_paged_kv_preserves_tma_compatible_outer_stride,
test_msa_paged_hnd_input_materializes_unaligned_outer_stride, and
test_msa_sparse_attention_honors_noncontiguous_block_indexes, so they skip
unless fmha_sm100 is installed and the GPU is SM100/SM103. Add coverage for the
_prepare_paged_hnd_input tensor.ndim != 4 branch, preserving the existing
assertions and supported-hardware behavior.
Source: Path instructions
Comment on lines
+165
to
+171
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") | ||
| @pytest.mark.parametrize("query_len", [1, 5]) | ||
| @pytest.mark.parametrize("num_kv_splits", [1, 2]) | ||
| def test_msa_sparse_attention_honors_noncontiguous_block_indexes( | ||
| query_len: int, | ||
| num_kv_splits: int, | ||
| ) -> None: |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Gate this test on SM100 in addition to CUDA availability.
This test launches the packaged MSA kernel, which targets SM100. The only guard is torch.cuda.is_available(). On a CUDA machine that is not SM100, the kernel launch fails and the test reports an error rather than a skip. Add an SM-version guard, for example pytest.mark.skipif(get_sm_version() != 100, ...), alongside the existing CUDA guard. The helper-only tests at lines 83-134 do not need this guard, because they exercise the Python layout logic and never launch a kernel.
🤖 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 165 - 171, Update
test_msa_sparse_attention_honors_noncontiguous_block_indexes to add a pytest
skip condition requiring get_sm_version() == 100 alongside the existing CUDA
availability guard. Leave the helper-only tests unchanged.
1 task
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.
Dev Engineer Review
kv_block_indexesfor short queries.fmha_sm100package.QA Engineer Review
test_msa_paged_kv_preserves_tma_compatible_outer_stride.test_msa_paged_hnd_input_materializes_unaligned_outer_stride.test_msa_sparse_attention_honors_noncontiguous_block_indexes.tests/integration/test_lists/,test-db/, orqa/entries cover these tests.Prerequisite
Description
MiniMax-M3's MSA sparse-attention path for
q_len <= 32can receivekv_block_indexesas a non-contiguous, head-major tensor view. The legacySM100 kernel treated the underlying pointer as a contiguous
[query, KV head, top-k]allocation, which could select incorrect KV pagesand produce non-finite attention output.
This change passes the tensor's query, KV-head, and top-k strides through the
existing MSA native interface and uses them for every sparse block-index load.
It also updates the bounds-check extent for the strided view. The attention
algorithm is unchanged, and no materialization or additional buffer is added.
The fix is carried in the TensorRT-LLM-owned patch for the pinned MSA source;
it can be removed after the corresponding change lands in MSA.
Test Coverage
selectors are bit-identical and finite for
q_len1 and 5 with one and twoKV splits (4/4 combinations passed).
split-count matrix.
and current MSA
main.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.