[None][refactor] Unify sparse attention framework with clean backend interfaces#12733
Open
lfr-0531 wants to merge 11 commits into
Open
[None][refactor] Unify sparse attention framework with clean backend interfaces#12733lfr-0531 wants to merge 11 commits into
lfr-0531 wants to merge 11 commits into
Conversation
lfr-0531
force-pushed
the
user/fanrongl/sparse-attention-refactor
branch
8 times, most recently
from
April 3, 2026 11:18
a1e4402 to
eaab4c3
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #41657 [ run ] triggered by Bot. Commit: |
1 task
Collaborator
|
PR_Github #41657 [ run ] completed with state
|
lfr-0531
force-pushed
the
user/fanrongl/sparse-attention-refactor
branch
2 times, most recently
from
April 4, 2026 06:23
6d89705 to
4004b1a
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #41813 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #41813 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #41921 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #41921 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #41962 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #41962 [ run ] completed with state
|
heyuhhh
reviewed
Apr 7, 2026
|
|
||
| TLLM_CHECK(host_kv_cache_pool_mapping.has_value()); | ||
| int32_t const layer_num = host_kv_cache_pool_mapping.value().size(0); | ||
|
|
Collaborator
There was a problem hiding this comment.
We can add an assert sentence here to make sure sparse_mla_topk_value is a valid value
heyuhhh
reviewed
Apr 7, 2026
| tensorrt_llm/_torch/attention_backend/sparse/dsa.py | | ||
| tensorrt_llm/_torch/attention_backend/sparse/kernel.py | | ||
| tensorrt_llm/_torch/attention_backend/sparse/rocket.py | | ||
| tensorrt_llm/_torch/attention_backend/sparse/utils.py | |
Collaborator
There was a problem hiding this comment.
Why we just remove there files but not add new files?
Collaborator
Author
There was a problem hiding this comment.
The removed files (dsa.py, kernel.py, rocket.py, utils.py) were in the legacy exclusion list because they had historical lint violations that were grandfathered in.
The new replacement files (dsa/backend.py, dsa/custom_ops.py, dsa/indexer.py, dsa/metadata.py, dsa/cache_manager.py, rocket/backend.py, skip_softmax/backend.py, params.py, etc.) are written from scratch and fully comply with current lint standards — ruff check passes on all of them without any exclusions. So they don't need to be added to the legacy list.
The same applies to legacy-files.txt, pyproject.toml, and ruff-legacy.toml — all four config files are auto-generated from legacy-files.txt via scripts/legacy_utils.py gen-configs.
lfr-0531
force-pushed
the
user/fanrongl/sparse-attention-refactor
branch
3 times, most recently
from
April 7, 2026 03:07
79c8913 to
53a6ee5
Compare
karljang
approved these changes
Jul 23, 2026
karljang
left a comment
karljang
left a comment
Collaborator
There was a problem hiding this comment.
LGTM from the VisualGen perspective. The existing TRTLLM SkipSoftmax configuration and runtime behavior appear to be preserved by this lower-level refactor.
Collaborator
|
PR_Github #61321 [ run ] completed with state
|
![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: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py (1)
2331-2372: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPatch the alias this test actually uses
tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py:2331-2372
mla_forward_impl_with_deepseek_v4_wo_linearcalls the importedforward_context_deepseek_v4_mlabinding, so patchingdeepseek_v4_module.forward_context_sparse_attnnever intercepts the call. Patchforward_context_deepseek_v4_mlain this module, or route the helper through the module attribute, so the fused FP8/NaN path is exercised.🤖 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_sparse_mla_forward.py` around lines 2331 - 2372, Update test_forward_sparse_mla_unified_fused_q_fp8 to patch the forward_context_deepseek_v4_mla binding used by mla_forward_impl_with_deepseek_v4_wo_linear, rather than deepseek_v4_module.forward_context_sparse_attn. Preserve the existing patched behavior that populates the fused FP8 buffers and replaces q with NaN, ensuring the test actually exercises the fused FP8/NaN path.
♻️ Duplicate comments (1)
tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py (1)
145-157: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winUnresolved: generation-phase block offsets still sliced along the wrong dimension.
metadata.kv_cache_block_offsets[:, metadata.num_contexts:]slices the page/block dimension instead of selecting the generation-phase sequences (rows).kv_cache_block_offsetsis[num_requests, max_pages_per_req](confirmed by the block-table layout documented indsa/kernels.py), so for mixed ctx/gen batches this passes the wrong block table intomla_rope_append_paged_kv_assign_q, corrupting cache addressing for generation requests. This was already flagged as critical on a prior commit and does not appear to have been fixed.🐛 Proposed fix
if is_generation: cached_token_indptr = metadata.gen_cached_token_indptr kv_indptr = metadata.gen_kv_indptr num_seqs = metadata.num_generations max_seq_len = metadata.max_gen_seq_len - block_offsets = metadata.kv_cache_block_offsets[:, metadata. - num_contexts:] + block_offsets = metadata.kv_cache_block_offsets[metadata.num_contexts:, :]🤖 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/dsa/backend.py` around lines 145 - 157, Update the generation branch in the metadata-selection logic to select generation requests by rows from metadata.kv_cache_block_offsets, preserving all page columns; use metadata.num_contexts as the row offset. Keep the context branch unchanged and ensure mixed context/generation batches pass only generation block tables to mla_rope_append_paged_kv_assign_q.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py (1)
728-728: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead expression.
num_kv_heads * num_heads_per_kvcomputes a value that is never assigned or used. Remove it (or assign if it was intended for the grid/shape calculation).♻️ Proposed cleanup
num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim = q.shape - num_kv_heads * num_heads_per_kv🤖 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/rocket/kernels.py` at line 728, Remove the unused standalone expression `num_kv_heads * num_heads_per_kv` from the surrounding kernel configuration or shape-calculation code. If that multiplication is required for the intended grid or shape, assign its result to the appropriate existing value instead of leaving it unevaluated.tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py (1)
98-123: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd missing return annotations on the new methods to satisfy the annotation guideline.
post_load_weights(-> None),_qk_projection_and_rope(-> torch.Tensor), andforward(returns the topk-indices tensor) lack return types;__init__is missing-> None.As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions... prefer built-in generic types."🤖 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/deepseek_v4/indexer.py` around lines 98 - 123, Add return annotations to the new methods: annotate __init__ and post_load_weights with -> None, and _qk_projection_and_rope with -> torch.Tensor. Also annotate forward with the appropriate tensor return type for its top-k indices result, preserving its existing behavior.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.py (1)
97-110: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFix implicit Optional flagged by Ruff (RUF013).
stride_factor: int = Noneshould be typed as optional sinceNoneis a valid default.🔧 Proposed fix
+from typing import Optional + def triton_convert_req_index_to_global_index( req_id: torch.Tensor, block_table: torch.Tensor, token_indices: torch.Tensor, BLOCK_SIZE: int, NUM_TOPK_TOKENS: int = 2048, BLOCK_N: int = 128, - stride_factor: int = None, # elements per block in pool + stride_factor: Optional[int] = None, # elements per block in pool layer_id: int = 0, num_kv_heads: int = 1, kv_factor: int = 1, ):🤖 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/dsa/kernels.py` around lines 97 - 110, Update the stride_factor parameter in triton_convert_req_index_to_global_index to use an explicit optional integer type while preserving its existing None default and behavior.Source: Linters/SAST tools
🤖 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/rocket/kernels.py`:
- Around line 919-921: Update the output allocation in the surrounding kernel
function to use input.device instead of the hard-coded 'cuda' device, ensuring
output remains on the same device as input for device-specific GPU calls.
In `@tensorrt_llm/_torch/modules/attention.py`:
- Line 862: Validate or filter the arbitrary kwargs in Attention.forward before
passing them to AttentionForwardArgs, retaining only keys matching its dataclass
fields. Ensure unexpected kwargs are ignored rather than causing constructor
failure, while preserving all supported argument values.
In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py`:
- Around line 109-209: Register
tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py in the
appropriate test-list configuration under tests/integration/test_lists/test-db
or qa, alongside the existing DSA indexer test entry such as
test_dsa_fp4_indexer.py. Ensure the new test file is included in the relevant
test job.
---
Outside diff comments:
In `@tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py`:
- Around line 2331-2372: Update test_forward_sparse_mla_unified_fused_q_fp8 to
patch the forward_context_deepseek_v4_mla binding used by
mla_forward_impl_with_deepseek_v4_wo_linear, rather than
deepseek_v4_module.forward_context_sparse_attn. Preserve the existing patched
behavior that populates the fused FP8 buffers and replaces q with NaN, ensuring
the test actually exercises the fused FP8/NaN path.
---
Duplicate comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py`:
- Around line 145-157: Update the generation branch in the metadata-selection
logic to select generation requests by rows from
metadata.kv_cache_block_offsets, preserving all page columns; use
metadata.num_contexts as the row offset. Keep the context branch unchanged and
ensure mixed context/generation batches pass only generation block tables to
mla_rope_append_paged_kv_assign_q.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py`:
- Around line 98-123: Add return annotations to the new methods: annotate
__init__ and post_load_weights with -> None, and _qk_projection_and_rope with ->
torch.Tensor. Also annotate forward with the appropriate tensor return type for
its top-k indices result, preserving its existing behavior.
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.py`:
- Around line 97-110: Update the stride_factor parameter in
triton_convert_req_index_to_global_index to use an explicit optional integer
type while preserving its existing None default and behavior.
In `@tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py`:
- Line 728: Remove the unused standalone expression `num_kv_heads *
num_heads_per_kv` from the surrounding kernel configuration or shape-calculation
code. If that multiplication is required for the intended grid or shape, assign
its result to the appropriate existing value instead of leaving it unevaluated.
🪄 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: f6a84d92-3882-4ef3-a37c-a68972fc0cc8
📒 Files selected for processing (75)
.pre-commit-config.yamldocs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.mddocs/source/developer-guide/sparse-attention-development-guide.mddocs/source/torch/adding_custom_kernels.mdlegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.pytensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.pytensorrt_llm/_torch/attention_backend/sparse/dsa/backend.pytensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.pytensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.pytensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.pytensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.pytensorrt_llm/_torch/attention_backend/sparse/dsa/module.pytensorrt_llm/_torch/attention_backend/sparse/dsa/params.pytensorrt_llm/_torch/attention_backend/sparse/hooks.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/params.pytensorrt_llm/_torch/attention_backend/sparse/registry.pytensorrt_llm/_torch/attention_backend/sparse/rocket.pytensorrt_llm/_torch/attention_backend/sparse/rocket/__init__.pytensorrt_llm/_torch/attention_backend/sparse/rocket/backend.pytensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.pytensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.pytensorrt_llm/_torch/attention_backend/sparse/rocket/module.pytensorrt_llm/_torch/attention_backend/sparse/rocket/params.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/__init__.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_kernels.pytests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.pytests/unittest/_torch/attention/sparse/rocketkv/test_kernels.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/attention/sparse/test_sparse_attention.pytests/unittest/_torch/attention/sparse/test_sparse_mla_forward.pytests/unittest/_torch/attention/sparse/test_triton_topk.pytests/unittest/_torch/attention/test_attention_op_sync.pytests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.pytests/unittest/_torch/modeling/test_modeling_deepseekv4.pytests/unittest/_torch/modules/test_mla_helix.pytests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (2)
- tests/unittest/_torch/attention/sparse/test_triton_topk.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket.py
🚧 Files skipped from review as they are similar to previous changes (3)
- .pre-commit-config.yaml
- tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.py
Comment on lines
+919
to
+921
| output = torch.empty((row_size, num_tokens, head_num, dim_size), | ||
| device='cuda', | ||
| dtype=input.dtype) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Allocate the output on the input's device, not the default CUDA device. device='cuda' resolves to the current default device; if input is on a non-default GPU (e.g. cuda:1), the kernel launch operates on tensors residing on different devices. Use input.device to stay consistent with the callers that pass device-specific tensors.
🛡️ Proposed fix
output = torch.empty((row_size, num_tokens, head_num, dim_size),
- device='cuda',
+ device=input.device,
dtype=input.dtype)📝 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
| output = torch.empty((row_size, num_tokens, head_num, dim_size), | |
| device='cuda', | |
| dtype=input.dtype) | |
| output = torch.empty((row_size, num_tokens, head_num, dim_size), | |
| device=input.device, | |
| dtype=input.dtype) |
🤖 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/rocket/kernels.py` around lines
919 - 921, Update the output allocation in the surrounding kernel function to
use input.device instead of the hard-coded 'cuda' device, ensuring output
remains on the same device as input for device-specific GPU calls.
| relative_attention_bias=relative_attention_bias, | ||
| relative_attention_max_distance= | ||
| relative_attention_max_distance, | ||
| **kwargs, |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find callers passing extra kwargs to Attention.forward and inspect keys.
rg -nP -C2 '\.forward\(' --type=py tensorrt_llm/_torch | rg -n 'sparse_backend_args|attention_input_type|latent_cache|q_pe'
# List declared AttentionForwardArgs fields for cross-check.
ast-grep run --pattern 'class AttentionForwardArgs {$$$}' --lang python tensorrt_llm/_torch/attention_backend/interface.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 857
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant files and symbols first.
git ls-files 'tensorrt_llm/_torch/modules/attention.py' 'tensorrt_llm/_torch/attention_backend/interface.py' 'tensorrt_llm/_torch/modules/mla.py'
echo '--- outline: attention.py ---'
ast-grep outline tensorrt_llm/_torch/modules/attention.py --view expanded || true
echo '--- outline: interface.py ---'
ast-grep outline tensorrt_llm/_torch/attention_backend/interface.py --view expanded || true
echo '--- outline: mla.py ---'
ast-grep outline tensorrt_llm/_torch/modules/mla.py --view expanded || true
echo '--- forwarded kwargs sites in attention.py ---'
rg -n -C 3 'AttentionForwardArgs\(' tensorrt_llm/_torch/modules/attention.py
echo '--- caller/search for forward kwargs keys in _torch ---'
rg -n -C 2 '\.forward\(' tensorrt_llm/_torch | rg -n 'sparse_backend_args|attention_input_type|latent_cache|q_pe|softmax_stats_tensor|attention_mask|output'
echo '--- declared AttentionForwardArgs fields ---'
python3 - <<'PY'
from pathlib import Path
import re
path = Path('tensorrt_llm/_torch/attention_backend/interface.py')
text = path.read_text()
m = re.search(r'`@dataclass`\s+class\s+AttentionForwardArgs:\s*(.*?)(?:\n@dataclass|\Z)', text, re.S)
print('found_dataclass=', bool(m))
if m:
body = m.group(1)
fields = re.findall(r'^\s+([A-Za-z_][A-Za-z0-9_]*)\s*:\s*', body, re.M)
print('fields=', fields)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 19338
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read only the relevant sections around the constructor call and the dataclass.
echo '--- attention.py around line 862 ---'
sed -n '830,875p' tensorrt_llm/_torch/modules/attention.py | cat -n
echo '--- attention.py around line 912 ---'
sed -n '900,925p' tensorrt_llm/_torch/modules/attention.py | cat -n
echo '--- interface.py around AttentionForwardArgs ---'
rg -n -A 80 -B 10 'class AttentionForwardArgs' tensorrt_llm/_torch/attention_backend/interface.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 8258
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- merge_attention_forward_args ---'
sed -n '967,1025p' tensorrt_llm/_torch/attention_backend/interface.py | cat -n
echo '--- AttentionBackend.forward signature ---'
sed -n '1030,1085p' tensorrt_llm/_torch/attention_backend/interface.py | cat -n
echo '--- Attention.forward / forward_impl signatures and kwargs handling ---'
sed -n '780,930p' tensorrt_llm/_torch/modules/attention.py | cat -n
echo '--- call sites passing **kwargs into attention backend ---'
rg -n -C 3 'self\.(attn|mha|inner_backend)\.forward\(' tensorrt_llm/_torch | rg -n '\*\*kwargs|\*\*attn_kwargs'Repository: NVIDIA/TensorRT-LLM
Length of output: 14179
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
text = Path('tensorrt_llm/_torch/attention_backend/interface.py').read_text()
# Extract merge_attention_forward_args and AttentionBackend.forward bodies/signatures.
for name in ['merge_attention_forward_args', 'class AttentionBackend']:
idx = text.find(name)
print(f'--- {name} at {idx} ---')
if idx != -1:
print(text[idx:idx+1800])
text2 = Path('tensorrt_llm/_torch/modules/attention.py').read_text()
for name in ['def forward_impl(', 'def forward(']:
idx = text2.find(name)
print(f'--- attention.py {name} at {idx} ---')
if idx != -1:
print(text2[idx:idx+2200])
# Collect self.*.forward call sites with kwargs spread.
print('--- forward call sites with kwargs spread ---')
for m in re.finditer(r'self\.(?:attn|mha|inner_backend)\.forward\([^\n]*\n(?:.*\n){0,12}?\s*\*\*(kwargs|attn_kwargs)', Path('tensorrt_llm/_torch').read_text() if False else ''):
print(m.group(0))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 8384
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- where merge_attention_forward_args is used ---'
rg -n -C 3 'merge_attention_forward_args\(' tensorrt_llm/_torch
echo '--- concrete AttentionBackend.forward implementations ---'
rg -n -C 6 'def forward\(' tensorrt_llm/_torch/attention_backend tensorrt_llm/_torch/visual_gen/attention_backend tensorrt_llm/_torch/modules | rg -n 'merge_attention_forward_args|forward_args|kwargs'Repository: NVIDIA/TensorRT-LLM
Length of output: 8408
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
paths = [
Path('tensorrt_llm/_torch/attention_backend/interface.py'),
Path('tensorrt_llm/_torch/modules/attention.py'),
Path('tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py'),
Path('tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py'),
]
for path in paths:
text = path.read_text()
if 'merge_attention_forward_args' in text:
print(f'=== {path} contains merge_attention_forward_args ===')
for m in re.finditer(r'merge_attention_forward_args\(', text):
start = max(0, m.start() - 200)
end = min(len(text), m.start() + 500)
print(text[start:end])
print('---')
print('=== forward signatures with kwargs in attention_backend subtree ===')
for path in Path('tensorrt_llm/_torch').rglob('*.py'):
text = path.read_text()
if 'def forward(' in text and 'AttentionForwardArgs' in text and 'kwargs' in text:
if 'attention_backend' in str(path) or 'modules/attention.py' in str(path):
print(f'### {path}')
for m in re.finditer(r'def forward\([^\n]*\n(?:.*\n){0,20}?\s*return', text):
pass
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 1558
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 4 'merge_attention_forward_args\(' tensorrt_llm/_torchRepository: NVIDIA/TensorRT-LLM
Length of output: 6831
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,220p' tensorrt_llm/_torch/attention_backend/interface.py | cat -n | sed -n '1,220p'Repository: NVIDIA/TensorRT-LLM
Length of output: 11011
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 4 'merge_attention_forward_args\(' tensorrt_llm/_torch/attention_backend tensorrt_llm/_torch/visual_genRepository: NVIDIA/TensorRT-LLM
Length of output: 6831
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- attention.py forward_impl continuation ---'
sed -n '930,1045p' tensorrt_llm/_torch/modules/attention.py | cat -n
echo '--- attention.py calls to AttentionForwardArgs and kwargs handling ---'
rg -n -C 4 'AttentionForwardArgs\(|kwargs' tensorrt_llm/_torch/modules/attention.py
echo '--- public Attention.forward callers that pass nontrivial kwargs ---'
rg -n -C 2 'attn\.forward\(' tensorrt_llm/_torch | rg -n '\*\*kwargs|\*\*attn_kwargs|attention_mask_data|attention_sinks|relative_attention_bias|relative_attention_max_distance|mrope_config|out_scale|output_sf'Repository: NVIDIA/TensorRT-LLM
Length of output: 10448
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- visual_gen/modules/attention.py around the forward call ---'
sed -n '520,620p' tensorrt_llm/_torch/visual_gen/modules/attention.py | cat -n
echo '--- any explicit kwargs names forwarded into attention.forward ---'
rg -n -C 3 'self\.attn\.forward\(q=q, k=k, v=v, \*\*kwargs\)|self\.attn\.forward\(' tensorrt_llm/_torch/visual_gen/modules/attention.py tensorrt_llm/_torch/models tensorrt_llm/_torch/modules | head -n 80Repository: NVIDIA/TensorRT-LLM
Length of output: 8831
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- kwarg names passed into self.attn.forward in _torch models/modules ---'
rg -n -C 3 'self\.attn\.forward\(' tensorrt_llm/_torch/models tensorrt_llm/_torch/modules | rg -n 'batch_size=|seq_len=|seq_len_kv=|gate_compress=|gate_fine=|timestep=|output=|output_sf=|out_scale=|softmax_stats_tensor=|latent_cache=|q_pe=|attention_input_type=|sparse_backend_args='Repository: NVIDIA/TensorRT-LLM
Length of output: 356
Validate **kwargs before building AttentionForwardArgs.
Attention.forward() still accepts arbitrary kwargs, and this constructor only accepts dataclass fields. Unexpected keys will raise here instead of being ignored.
🤖 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/modules/attention.py` at line 862, Validate or filter the
arbitrary kwargs in Attention.forward before passing them to
AttentionForwardArgs, retaining only keys matching its dataclass fields. Ensure
unexpected kwargs are ignored rather than causing constructor failure, while
preserving all supported argument values.
Comment on lines
+109
to
+209
| def test_shared_topk_indices_lifecycle(): | ||
| sparse_config = DeepSeekSparseAttentionConfig( | ||
| index_n_heads=1, | ||
| index_head_dim=8, | ||
| index_topk=3, | ||
| skip_indexer_for_short_seqs=False, | ||
| ) | ||
| pretrained_config = SimpleNamespace( | ||
| num_hidden_layers=2, | ||
| index_topk_pattern=["F", "S"], | ||
| ) | ||
| sparse_metadata_params = sparse_config.to_sparse_metadata_params( | ||
| pretrained_config=pretrained_config | ||
| ) | ||
| assert sparse_metadata_params.has_shared_indexer_layers | ||
|
|
||
| metadata = object.__new__(DSAtrtllmAttentionMetadata) | ||
| metadata.sparse_metadata_params = sparse_metadata_params | ||
| metadata.max_num_sequences = 2 | ||
| metadata.max_num_tokens = 4 | ||
| metadata.num_sparse_topk = 3 | ||
| metadata.num_sms = 1 | ||
| metadata.cuda_graph_buffers = None | ||
| metadata.kv_cache_manager = SimpleNamespace(max_blocks_per_seq=2) | ||
| metadata.enable_context_mla_with_cached_kv = False | ||
| metadata.enable_indexer_skip = False | ||
| metadata.get_empty = Mock( | ||
| side_effect=lambda _, shape, **kwargs: torch.empty(tuple(shape), dtype=kwargs["dtype"]) | ||
| ) | ||
| metadata._create_kv_lens_2d_buffer = Mock() | ||
| metadata._create_radix_aux_buffers = Mock() | ||
| metadata.create_expanded_buffers = Mock() | ||
|
|
||
| with patch( | ||
| "tensorrt_llm._torch.attention_backend.sparse.dsa.metadata.prefer_pinned", | ||
| return_value=False, | ||
| ): | ||
| metadata.create_buffers_for_indexer() | ||
|
|
||
| buffer = metadata.shared_topk_indices | ||
| assert buffer.shape == (4, 3) | ||
|
|
||
| context_topk = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.int32) | ||
| generation_topk = torch.tensor([[2, 1, 0], [3, 2, 1]], dtype=torch.int32) | ||
| full_backend = SimpleNamespace( | ||
| indexer=SimpleNamespace( | ||
| forward_from_projected=Mock(side_effect=[context_topk, generation_topk]) | ||
| ), | ||
| get_local_layer_idx=Mock(return_value=0), | ||
| ) | ||
| shared_backend = SimpleNamespace( | ||
| indexer=None, | ||
| get_local_layer_idx=Mock(return_value=1), | ||
| ) | ||
| metadata._num_ctx_tokens = context_topk.shape[0] | ||
| metadata._num_tokens = context_topk.shape[0] + generation_topk.shape[0] | ||
| backend_args = DSABackendForwardArgs(indexer_intermediates=[]) | ||
|
|
||
| def _predict(backend, input_type): | ||
| return DSATrtllmAttention.sparse_attn_predict( | ||
| backend, | ||
| torch.empty((2, 1)), | ||
| None, | ||
| metadata, | ||
| AttentionForwardArgs( | ||
| attention_input_type=input_type, | ||
| sparse_backend_args=backend_args, | ||
| ), | ||
| )[0] | ||
|
|
||
| with patch( | ||
| "tensorrt_llm._torch.attention_backend.sparse.dsa.backend." | ||
| "transform_local_topk_and_prepare_pool_view", | ||
| side_effect=lambda topk, *_: (topk.clone(), None), | ||
| ): | ||
| _predict(full_backend, AttentionInputType.context_only) | ||
| _predict(full_backend, AttentionInputType.generation_only) | ||
| torch.testing.assert_close( | ||
| metadata.shared_topk_indices, | ||
| torch.cat([context_topk, generation_topk]), | ||
| ) | ||
| torch.testing.assert_close( | ||
| _predict(shared_backend, AttentionInputType.context_only), | ||
| context_topk, | ||
| ) | ||
| torch.testing.assert_close( | ||
| _predict(shared_backend, AttentionInputType.generation_only), | ||
| generation_topk, | ||
| ) | ||
|
|
||
| metadata._invalidate_pool_view_cache = Mock() | ||
| metadata.kv_cache_manager = None | ||
| metadata._num_tokens = 0 | ||
| metadata._num_generations = 0 | ||
| metadata.kv_lens_cuda = torch.empty(0, dtype=torch.int32) | ||
| metadata.prepare_dense_topk_indices = Mock() | ||
| metadata.on_update_kv_lens() | ||
|
|
||
| assert metadata.shared_topk_indices is buffer | ||
|
|
||
|
|
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n "test_shared_topk_indices_lifecycle|test_topk_indices_buffer_cuda_graph" tests/integration/test_listsRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant unittest file and inspect the changed region.
wc -l tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
sed -n '1,260p' tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
# Inspect the integration test-list layout.
find tests/integration/test_lists -maxdepth 2 -type f | sort | sed -n '1,200p'Repository: NVIDIA/TensorRT-LLM
Length of output: 14012
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- file line count ---'
wc -l tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
echo '--- changed test names in unittest file ---'
rg -n "def test_|test_shared_topk_indices_lifecycle|test_topk_indices_buffer_cuda_graph|test_metadata_cache_geometry_comes_from_sparse_metadata_params" tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
echo '--- integration test list files ---'
find tests/integration/test_lists -type f | sort
echo '--- matching names in integration test lists ---'
rg -n "test_shared_topk_indices_lifecycle|test_topk_indices_buffer_cuda_graph" tests/integration/test_lists || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 6888
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- search for test list policy references ---'
rg -n "test_lists|cbts_touchmap|test-db/|qa/" tests .github . 2>/dev/null | sed -n '1,200p'
echo '--- locate the specific unittest file symbols ---'
rg -n "def test_shared_topk_indices_lifecycle|def test_topk_indices_buffer_cuda_graph|def test_metadata_cache_geometry_comes_from_sparse_metadata_params" tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 25719
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- search for the unittest file path in test-db/qa lists ---'
rg -n "tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer\.py|_torch/attention/sparse/dsa/test_dsa_indexer|test_shared_topk_indices_lifecycle|test_topk_indices_buffer_cuda_graph" tests/integration/test_lists -S || true
echo '--- search for any unittest bridge entries in test-db ---'
rg -n "unittest/.*test_dsa_indexer\.py|unittest/.*_torch/attention/sparse/dsa" tests/integration/test_lists/test-db -S || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 398
Register tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py in the test list. test_shared_topk_indices_lifecycle is new here, but no unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py entry exists under tests/integration/test_lists/test-db/ or qa/ (the only nearby bridge is unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py in l0_dgx_b200.yml).
🤖 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/dsa/test_dsa_indexer.py` around lines
109 - 209, Register
tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py in the
appropriate test-list configuration under tests/integration/test_lists/test-db
or qa, alongside the existing DSA indexer test entry such as
test_dsa_fp4_indexer.py. Ensure the new test file is included in the relevant
test job.
Source: Path instructions
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61465 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61504 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61465 [ run ] completed with state |
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py (1)
9-13: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the new rocketkv tests to the integration lists
test_triton_bmm,test_triton_rocket_paged_kt_cache_bmm, andtest_topk_kernelare new intests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py.- No matching entries appear under
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/, so these tests won’t be picked up by the integration harness.Coverage verdict: insufficient.
🤖 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/rocketkv/test_kernels.py` around lines 9 - 13, Add integration-list entries for test_triton_bmm, test_triton_rocket_paged_kt_cache_bmm, and test_topk_kernel from test_kernels.py in the appropriate test-db and qa lists, using the existing list format so the integration harness discovers all three tests.Source: Path instructions
🧹 Nitpick comments (1)
tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py (1)
727-728: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead no-op expression. Line 728 computes
num_kv_heads * num_heads_per_kvbut discards the result — it has no effect and looks like a leftover (possibly intended to assign anum_heads/total_headslocal). Remove it or assign it if the value is needed.♻️ Suggested cleanup
num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim = q.shape - num_kv_heads * num_heads_per_kv🤖 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/rocket/kernels.py` around lines 727 - 728, Remove the discarded num_kv_heads * num_heads_per_kv expression after the q.shape unpacking, unless a later computation requires that product; if needed, assign it to the appropriate existing heads-count local and use that value.
🤖 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.
Outside diff comments:
In `@tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py`:
- Around line 9-13: Add integration-list entries for test_triton_bmm,
test_triton_rocket_paged_kt_cache_bmm, and test_topk_kernel from test_kernels.py
in the appropriate test-db and qa lists, using the existing list format so the
integration harness discovers all three tests.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py`:
- Around line 727-728: Remove the discarded num_kv_heads * num_heads_per_kv
expression after the q.shape unpacking, unless a later computation requires that
product; if needed, assign it to the appropriate existing heads-count local and
use that value.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 81865ea1-94a0-455d-88fc-d90fc7d04e41
📒 Files selected for processing (75)
.pre-commit-config.yamldocs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.mddocs/source/developer-guide/sparse-attention-development-guide.mddocs/source/torch/adding_custom_kernels.mdlegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.pytensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.pytensorrt_llm/_torch/attention_backend/sparse/dsa/backend.pytensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.pytensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.pytensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.pytensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.pytensorrt_llm/_torch/attention_backend/sparse/dsa/module.pytensorrt_llm/_torch/attention_backend/sparse/dsa/params.pytensorrt_llm/_torch/attention_backend/sparse/hooks.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/params.pytensorrt_llm/_torch/attention_backend/sparse/registry.pytensorrt_llm/_torch/attention_backend/sparse/rocket.pytensorrt_llm/_torch/attention_backend/sparse/rocket/__init__.pytensorrt_llm/_torch/attention_backend/sparse/rocket/backend.pytensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.pytensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.pytensorrt_llm/_torch/attention_backend/sparse/rocket/module.pytensorrt_llm/_torch/attention_backend/sparse/rocket/params.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/__init__.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_kernels.pytests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.pytests/unittest/_torch/attention/sparse/rocketkv/test_kernels.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/attention/sparse/test_sparse_attention.pytests/unittest/_torch/attention/sparse/test_sparse_mla_forward.pytests/unittest/_torch/attention/sparse/test_triton_topk.pytests/unittest/_torch/attention/test_attention_op_sync.pytests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.pytests/unittest/_torch/modeling/test_modeling_deepseekv4.pytests/unittest/_torch/modules/test_mla_helix.pytests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (2)
- tests/unittest/_torch/attention/sparse/test_triton_topk.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket.py
🚧 Files skipped from review as they are similar to previous changes (56)
- tensorrt_llm/_torch/attention_backend/sparse/registry.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/module.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/init.py
- tensorrt_llm/_torch/attention_backend/sparse/init.py
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/params.py
- tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/init.py
- tests/unittest/_torch/attention/sparse/dsa/test_kernels.py
- pyproject.toml
- tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/init.py
- legacy-files.txt
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py
- ruff-legacy.toml
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py
- tests/unittest/_torch/attention/test_attention_op_sync.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
- docs/source/torch/adding_custom_kernels.md
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py
- tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
- tests/unittest/_torch/modules/test_mla_helix.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py
- tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py
- tensorrt_llm/_torch/attention_backend/vanilla.py
- tensorrt_llm/_torch/models/modeling_deepseekv4.py
- docs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.md
- tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py
- tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
- tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.py
- tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
- docs/source/developer-guide/sparse-attention-development-guide.md
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/init.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py
- tensorrt_llm/_torch/attention_backend/sparse/hooks.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py
- tensorrt_llm/_torch/attention_backend/trtllm.py
- tensorrt_llm/_torch/attention_backend/fmha/fallback.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.py
- tensorrt_llm/_torch/attention_backend/sparse/params.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py
- tensorrt_llm/_torch/modules/attention.py
- tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
- tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py
- tensorrt_llm/_torch/modules/mla.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py (1)
727-728: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead expression statement.
num_kv_heads * num_heads_per_kvcomputes and discards a value with no effect. Remove it (or assign it if it was meant to be used).♻️ Proposed cleanup
num_gen_tokens, num_kv_heads, num_heads_per_kv, head_dim = q.shape - num_kv_heads * num_heads_per_kv🤖 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/rocket/kernels.py` around lines 727 - 728, Remove the unused expression statement `num_kv_heads * num_heads_per_kv` from the code immediately following the q.shape unpacking; leave the existing shape assignments unchanged unless the computed value is required by subsequent logic.tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py (1)
102-105: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing return-type annotations.
post_load_weights(Line 102),_qk_projection_and_rope(Line 107), andforward(Line 357) lack return-type annotations (-> None/-> torch.Tensor).As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions."✏️ Proposed fix
- def post_load_weights(self): + def post_load_weights(self) -> None: # V4 does not use the V3 fused fp32 wk+weights_proj GEMM, and the # base concat would now hit an fp32/bf16 dtype mismatch. return- def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor): + def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:def forward( self, qr: torch.Tensor, hidden_states: torch.Tensor, metadata: DeepseekV4TrtllmAttentionMetadata, position_ids: torch.Tensor, pre_aux: Optional[ Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] ] = None, - ): + ) -> torch.Tensor:Also applies to: 107-115, 357-366
🤖 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/deepseek_v4/indexer.py` around lines 102 - 105, Annotate post_load_weights with -> None, _qk_projection_and_rope with -> torch.Tensor, and forward with the appropriate torch.Tensor return annotation, preserving their existing behavior and signatures otherwise.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 `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py`:
- Around line 208-224: The FP8_BLOCKWISE path in _update_k_cache_if_needed
currently invokes _update_k_cache even though that override is a no-op, leaving
the indexer K cache unpopulated. Implement the required FP8 blockwise cache
update so this path populates the cache, or remove the unreachable FP8 handling
including the k_scale assertion if FP8 blockwise is unsupported; keep the MXFP4
early return unchanged.
In `@tests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py`:
- Around line 221-226: The test currently passes a flattened rank-2 view to
deepseek_v4_q_norm_fused_fp8 instead of exercising the documented 3D q_pe
contract. Update the invocation around q_pe and the custom op to pass the
original 3D tensor, or validate it through the sparse-MLA attention consumer,
while retaining assertions that verify the actual consumer path rather than only
the allocation.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.py`:
- Around line 102-105: Annotate post_load_weights with -> None,
_qk_projection_and_rope with -> torch.Tensor, and forward with the appropriate
torch.Tensor return annotation, preserving their existing behavior and
signatures otherwise.
In `@tensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.py`:
- Around line 727-728: Remove the unused expression statement `num_kv_heads *
num_heads_per_kv` from the code immediately following the q.shape unpacking;
leave the existing shape assignments unchanged unless the computed value is
required by subsequent logic.
🪄 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: d21421a4-1c83-4620-bfaa-d0675f36744a
📒 Files selected for processing (75)
.pre-commit-config.yamldocs/source/blogs/tech_blog/blog17_Sparse_Attention_in_TensorRT-LLM.mddocs/source/developer-guide/sparse-attention-development-guide.mddocs/source/torch/adding_custom_kernels.mdlegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/interface.pytensorrt_llm/_torch/attention_backend/sparse/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/indexer.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.pytensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.pytensorrt_llm/_torch/attention_backend/sparse/dsa/backend.pytensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/dsa/custom_ops.pytensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.pytensorrt_llm/_torch/attention_backend/sparse/dsa/kernels.pytensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.pytensorrt_llm/_torch/attention_backend/sparse/dsa/module.pytensorrt_llm/_torch/attention_backend/sparse/dsa/params.pytensorrt_llm/_torch/attention_backend/sparse/hooks.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.pytensorrt_llm/_torch/attention_backend/sparse/params.pytensorrt_llm/_torch/attention_backend/sparse/registry.pytensorrt_llm/_torch/attention_backend/sparse/rocket.pytensorrt_llm/_torch/attention_backend/sparse/rocket/__init__.pytensorrt_llm/_torch/attention_backend/sparse/rocket/backend.pytensorrt_llm/_torch/attention_backend/sparse/rocket/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/rocket/kernels.pytensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.pytensorrt_llm/_torch/attention_backend/sparse/rocket/module.pytensorrt_llm/_torch/attention_backend/sparse/rocket/params.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/__init__.pytensorrt_llm/_torch/attention_backend/sparse/skip_softmax/params.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_deepseekv4.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.pytests/unittest/_torch/attention/sparse/dsa/test_kernels.pytests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.pytests/unittest/_torch/attention/sparse/rocketkv/test_kernels.pytests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.pytests/unittest/_torch/attention/sparse/test_sparse_attention.pytests/unittest/_torch/attention/sparse/test_sparse_mla_forward.pytests/unittest/_torch/attention/sparse/test_triton_topk.pytests/unittest/_torch/attention/test_attention_op_sync.pytests/unittest/_torch/custom_ops/test_deepseek_v4_q_norm.pytests/unittest/_torch/modeling/test_modeling_deepseekv4.pytests/unittest/_torch/modules/test_mla_helix.pytests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (2)
- tests/unittest/_torch/attention/sparse/test_triton_topk.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket.py
🚧 Files skipped from review as they are similar to previous changes (54)
- tensorrt_llm/_torch/attention_backend/sparse/rocket/init.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/module.py
- tensorrt_llm/_torch/attention_backend/sparse/registry.py
- ruff-legacy.toml
- docs/source/torch/adding_custom_kernels.md
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/skip_softmax/init.py
- tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/params.py
- tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
- tensorrt_llm/_torch/attention_backend/sparse/init.py
- tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
- pyproject.toml
- tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py
- .pre-commit-config.yaml
- tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
- tests/unittest/_torch/attention/sparse/rocketkv/test_kernels.py
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py
- tests/unittest/_torch/attention/test_attention_op_sync.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py
- tensorrt_llm/_torch/models/modeling_deepseekv4.py
- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py
- tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
- tests/unittest/llmapi/test_llm_args.py
- tensorrt_llm/_torch/attention_backend/vanilla.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/params.py
- tensorrt_llm/_torch/attention_backend/sparse/hooks.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
- tensorrt_llm/llmapi/llm_args.py
- tests/unittest/_torch/modules/test_mla_helix.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
- tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py
- tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
- tensorrt_llm/_torch/attention_backend/trtllm.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/params.py
- tensorrt_llm/_torch/attention_backend/fmha/fallback.py
- tensorrt_llm/_torch/attention_backend/interface.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/backend.py
- tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/backend.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/metadata.py
- tensorrt_llm/_torch/modules/attention.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/init.py
- tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/kernels.py
- tensorrt_llm/_torch/attention_backend/sparse/rocket/metadata.py
- tests/unittest/_torch/attention/sparse/test_sparse_attention.py
- tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py
- tensorrt_llm/_torch/attention_backend/sparse/dsa/module.py
- tensorrt_llm/_torch/modules/mla.py
Collaborator
|
PR_Github #61533 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61504 [ run ] completed with state |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61541 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61533 [ run ] completed with state |
Shixiaowei02
approved these changes
Jul 24, 2026
Collaborator
|
PR_Github #61541 [ run ] completed with state
|
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Define validated sparse-attention hooks shared by Attention and MLA while keeping algorithm-specific module behavior under each sparse backend. Simplify DSA and DeepSeek-V4 module integration, move Rocket kernels into its backend directory, and align the related architecture tests. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Unify module-to-backend and backend-to-attention-op sparse runtime arguments. Split DeepSeek-V4 indexer, metadata, and parameter definitions, and route DSA prediction through the backend while preserving shared TopK buffer lifetime. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Allocate a single mixed-batch TopK buffer only when the model contains shared indexer layers. Reuse per-layer indexer routing to derive the metadata requirement and keep the buffer address stable across steps. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Keep split DSA indexer test paths under inference mode, matching the original integrated Indexer.forward contract. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Restore the shared lint configuration to match main and format the split sparse attention modules with Ruff. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Collaborator
Author
|
sparse_attn_hooks |
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
6 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
SparseBackendForwardArgsandSparseRuntimeParams, removingSparsePredictionand replacingSkipSoftmaxKernelParamsusage with runtime fields.tensorrt_llm/_torch/attention_backend/sparse/hooks.py, and rewired MLA/attention execution to route sparse execution through hooks and consume predicted indices fromforward_args.sparse_runtime_params.*.tensorrt_llm/_torch/attention_backend/sparse/<algorithm>/:sparse/dsa/package (backend, indexer, kernels, metadata, cache manager, custom CUDA-graph-friendly ops, params).sparse/deepseek_v4/package (runtime-only metadata, backend/indexer/kernels/compressor, params).sparse/rocket.pytosparse/rocket/package; the monolithicrocket.pywas deleted and replaced by package modules (backend/cache_manager/kernels/metadata/params).SparseRuntimeParamsacross multiple backends and FMHA fallbacks, including SkipSoftmax threshold scaling fields..pre-commit-config.yaml,legacy-files.txt,pyproject.toml,ruff-legacy.toml.tok < 0, page-range check), not necessarily on all negative/out-of-range intermediates—verify assumptions aboutblock_table/pool-page validity.on_update_kv_lens()and CUDA-graph capture paths via custom ops; ensure buffer object identity/stability and cache revalidation are consistent.sparse/rocket/package and that any stale references to the removedrocket.pyare gone.forward_args.sparse_runtime_params(not legacy prediction/kernel-param paths), especially where multiple FMHA/FlashInfer/MSA backends branch on “sparse vs dense plan” decisions.QA Engineer Review
No test changes (no
tests/,tests/integration/test_lists/,test-db/,qa/, orwaives.txtmodifications detected in the diff evidence available here).Description
TensorRT-LLM already provides a sparse-attention framework for DSA, DeepSeek-V4,
RocketKV, and SkipSoftmax, but algorithm-specific module, backend, prediction,
metadata, cache, and kernel logic had become mixed across the common Attention
and MLA paths. That made the extension boundary inconsistent and required
changes to shared modules when adding or maintaining an algorithm.
This PR refactors sparse attention around explicit module and backend contracts:
optional sparse-attention hooks for module initialization, weights, forward,
custom-op execution, and output projection.
_torch/attention_backend/sparse/<algorithm>/.SparseBackendForwardArgscarries sparse inputs from Attention/MLA to theselected backend, while
SparseRuntimeParamscarries prediction andSkipSoftmax runtime inputs from the backend to the attention op.
the hooks it needs.
module.py,backend.py,params.py,metadata.py, cache/indexer helpers,and algorithm-owned kernels. SkipSoftmax keeps its existing FMHA integration
through the shared runtime parameters.
The refactor preserves dense MLA fallback behavior, DSA cross-layer indexer
sharing and piecewise CUDA graph support, DeepSeek-V4 compressed-cache paths,
RocketKV, MiniMax-M3 integration, and SkipSoftmax.
Test Coverage
main; a native rebuild is not required forthis refactor.
pre-commiton all files changed by this PR.70 passed.
worktree.
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.