[None][refactor] Unify sparse attention framework with clean backend interfaces by lfr-0531 · Pull Request #12733 · NVIDIA/TensorRT-LLM

lfr-0531 · 2026-04-03T09:40:20Z

Dev Engineer Review

Major sparse-attention refactor around explicit contracts:
- Introduced SparseBackendForwardArgs and SparseRuntimeParams, removing SparsePrediction and replacing SkipSoftmaxKernelParams usage with runtime fields.
- Added sparse hook facade/dispatcher in tensorrt_llm/_torch/attention_backend/sparse/hooks.py, and rewired MLA/attention execution to route sparse execution through hooks and consume predicted indices from forward_args.sparse_runtime_params.*.
Split and reorganized sparse backends under tensorrt_llm/_torch/attention_backend/sparse/<algorithm>/:
- DSA: new sparse/dsa/ package (backend, indexer, kernels, metadata, cache manager, custom CUDA-graph-friendly ops, params).
- DeepSeek-V4: new sparse/deepseek_v4/ package (runtime-only metadata, backend/indexer/kernels/compressor, params).
- RocketKV: migrated from sparse/rocket.py to sparse/rocket/ package; the monolithic rocket.py was deleted and replaced by package modules (backend/cache_manager/kernels/metadata/params).
Updated call sites to read/write sparse indices/offsets from SparseRuntimeParams across multiple backends and FMHA fallbacks, including SkipSoftmax threshold scaling fields.
Updated lint/config allowlists to match the new sparse package layout:
- .pre-commit-config.yaml, legacy-files.txt, pyproject.toml, ruff-legacy.toml.
Notable correctness/performance risks to watch:
- Triton/DSP boundary behavior: DSA kernels (e.g., global index conversion) only explicitly invalidate on some conditions (e.g., tok < 0, page-range check), not necessarily on all negative/out-of-range intermediates—verify assumptions about block_table/pool-page validity.
- Cache/index lifecycle: DSA shared-topk lifecycle across on_update_kv_lens() and CUDA-graph capture paths via custom ops; ensure buffer object identity/stability and cache revalidation are consistent.
- RocketKV migration/deletion: ensure all prior imports and expected symbols are correctly re-exported from the new sparse/rocket/ package and that any stale references to the removed rocket.py are gone.
- Interface migration completeness: ensure every sparse consumer now correctly uses 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.
CI/comments (from provided context): mixed outcomes with multiple pipelines failing and several helper jobs aborted/succeeded; failures likely require inspecting the specific sparse-attention tests/pipelines impacted by the refactor.

QA Engineer Review

No test changes (no tests/, tests/integration/test_lists/, test-db/, qa/, or waives.txt modifications 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:

Attention and MLA remain the shared modules. They resolve a validated set of
optional sparse-attention hooks for module initialization, weights, forward,
custom-op execution, and output projection.
Algorithm-specific module hooks and implementations live under
_torch/attention_backend/sparse/<algorithm>/.
SparseBackendForwardArgs carries sparse inputs from Attention/MLA to the
selected backend, while SparseRuntimeParams carries prediction and
SkipSoftmax runtime inputs from the backend to the attention op.
Backend prediction dispatch is shared, while each algorithm implements only
the hooks it needs.
DSA, DeepSeek-V4, and RocketKV are organized by responsibility into
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

No C/C++ changes relative to main; a native rebuild is not required for
this refactor.
pre-commit on all files changed by this PR.
Sparse Attention/MLA forward coverage, including DSA and DeepSeek-V4:
70 passed.
Sparse LLM argument coverage: 7 passed, 245 deselected.
DSA pinned block-table staging smoke test on B200.
Verified that tests imported Python sources and native bindings from this
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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

lfr-0531 · 2026-04-03T11:20:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-03T11:26:17Z

PR_Github #41657 [ run ] triggered by Bot. Commit: eaab4c3 Link to invocation

tensorrt-cicd · 2026-04-03T22:15:14Z

PR_Github #41657 [ run ] completed with state SUCCESS. Commit: eaab4c3
/LLM/main/L0_MergeRequest_PR pipeline #32562 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

lfr-0531 · 2026-04-04T07:46:02Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-04T07:51:31Z

PR_Github #41813 [ run ] triggered by Bot. Commit: 4004b1a Link to invocation

tensorrt-cicd · 2026-04-04T07:51:32Z

PR_Github #41813 [ run ] completed with state DISABLED
CI server is currently disabled for scheduled maintenance. Estimated completion time: 9 PM PST on 4/4.

Link to invocation

lfr-0531 · 2026-04-06T08:55:06Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-06T09:00:55Z

PR_Github #41921 [ run ] triggered by Bot. Commit: 4004b1a Link to invocation

tensorrt-cicd · 2026-04-06T14:43:43Z

PR_Github #41921 [ run ] completed with state SUCCESS. Commit: 4004b1a
/LLM/main/L0_MergeRequest_PR pipeline #32780 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

lfr-0531 · 2026-04-06T17:20:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-06T17:26:54Z

PR_Github #41962 [ run ] triggered by Bot. Commit: 0a79557 Link to invocation

tensorrt-cicd · 2026-04-06T23:19:16Z

PR_Github #41962 [ run ] completed with state SUCCESS. Commit: 0a79557
/LLM/main/L0_MergeRequest_PR pipeline #32816 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

heyuhhh · 2026-04-07T02:16:28Z


        TLLM_CHECK(host_kv_cache_pool_mapping.has_value());
        int32_t const layer_num = host_kv_cache_pool_mapping.value().size(0);



We can add an assert sentence here to make sure sparse_mla_topk_value is a valid value

heyuhhh · 2026-04-07T02:23:49Z

-        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 |


Why we just remove there files but not add new files?

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.

tensorrt-cicd · 2026-07-27T06:39:37Z

PR_Github #61837 [ run ] triggered by Bot. Commit: 7a75644 Link to invocation

lfr-0531 · 2026-07-27T06:53:12Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-27T06:59:03Z

PR_Github #61843 [ run ] triggered by Bot. Commit: 749fd2b Link to invocation

tensorrt-cicd · 2026-07-27T07:02:55Z

PR_Github #61837 [ run ] completed with state ABORTED. Commit: 7a75644

Link to invocation

tensorrt-cicd · 2026-07-27T10:29:35Z

PR_Github #61843 [ run ] completed with state FAILURE. Commit: 749fd2b
/LLM/main/L0_MergeRequest_PR pipeline #50037 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

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>

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>

Group all accepted module hook signatures under one contract table and document the lifecycle of each sparse attention hook. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

Document that the model-specific integration is a compatibility exception and should not be used as the pattern for new sparse attention algorithms. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

lfr-0531 · 2026-07-27T12:52:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-27T12:59:38Z

PR_Github #61912 [ run ] triggered by Bot. Commit: 87e815e Link to invocation

tensorrt-cicd · 2026-07-27T16:52:10Z

PR_Github #61912 [ run ] completed with state FAILURE. Commit: 87e815e
/LLM/main/L0_MergeRequest_PR pipeline #50103 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

lfr-0531 · 2026-07-28T02:00:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-28T02:07:07Z

PR_Github #62059 [ run ] triggered by Bot. Commit: a56b81c Link to invocation

bobboli · 2026-07-28T04:06:32Z

Changes related to Skip Softmax Attention look nice.
Maybe also worth mentioning AttentionForwardArgs in sparse attention developer guide?

mingyangHao · 2026-07-28T05:56:56Z

+        self.aux_stream = aux_stream
+        self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]
+        self.kv_a_layernorm = RMSNorm(hidden_size=self.kv_lora_rank, dtype=dtype, eps=rms_norm_eps)
+        self.kv_b_proj = Linear(


Blocking: this unconditionally constructs kv_b_proj (and later v_b_proj, o_proj, and the
dense MHA backend) before the sparse hook is initialized. The DeepSeek-V4 hook then deletes
these modules.

This is not only an abstraction issue: skip_create_weights_in_init defaults to false, so the
unused weights are actually allocated. With the dimensions used by the DeepSeek-V4 test, the
three deleted BF16 parameters alone account for roughly 84.5 MiB of transient allocation per
layer; larger production configurations can be substantially worse. Depending on the default
device used during model construction, this can increase host memory or GPU memory peaks and
potentially cause initialization OOMs.

The pre-refactor implementation explicitly avoided constructing these modules for DeepSeek-V4.
Please resolve the algorithm-specific build requirements before constructing the common MLA
modules, e.g. through a declarative build spec/capability object or an algorithm-specific
component factory, rather than constructing and deleting unused modules.

mingyangHao · 2026-07-28T05:57:24Z

+    _ATTENTION_HOOKS[algorithm] = hooks
+
+
+@lru_cache(maxsize=None)


Major: this makes the public registration API ineffective for a new algorithm. Even after
register_mla_sparse_hooks("new_algo", NewHooks) succeeds, this lookup returns None unless
"new_algo" is also added to the private _MLA_HOOK_MODULE_PATHS map.

The same issue exists in the Attention hook getter. The getter should first consult the hook
registry and use the module-path map only to lazily import built-in algorithms. Please also
add a test that registers an arbitrary algorithm name and retrieves its hook without modifying
a central built-in map.

If external registration is intentionally unsupported, the register_* functions should not
be exported as public extension points.

mingyangHao · 2026-07-28T05:57:47Z

+
+
+@lru_cache(maxsize=None)
+def _get_sparse_mla_hooks_for_algorithm(algorithm: str) -> Optional[MLASparseHooks]:


Major: this cache stores the instantiated hook object, so every layer and every model using
the same algorithm shares one global adapter instance.

The current adapters happen to be stateless, but neither the interface nor the type system
enforces that constraint. A future hook that stores initialization or layer-specific state on
self would silently leak state across layers or model instances.

Please cache only the lazy-import or hook class resolution and instantiate the adapter per
owning MLA/Attention module. A state-isolation test with two module instances would help
protect this contract.

mingyangHao · 2026-07-28T05:57:58Z

+    @abstractmethod
+    def initialize(
+        self,
+        mla: "MLA",


Major architectural concern: although this adapter is typed, passing the concrete MLA object
gives each sparse algorithm unrestricted access to the entire module implementation. The
current implementations directly inspect, replace, and delete MLA attributes, so any internal
MLA refactor can break an algorithm at runtime.

Moreover, most of the actual algorithm logic is delegated to free functions whose self
parameter is untyped, which means the typed adapter does not provide meaningful static
checking for those internal dependencies.

Please consider modeling the sparse implementation as a composed SparseMLAComponent that
owns its algorithm-specific weights and state, and pass it a narrow, stable context/services
interface instead of the full MLA object. At minimum, the set of MLA fields available to hooks
should be explicitly defined and contract-tested.

mingyangHao · 2026-07-28T05:59:07Z

+    sparse_attn_offsets: Optional[torch.Tensor] = None
+    sparse_attn_indices_block_size: int = 0
+    # DeepSeek-V4 compressed-cache inputs.
+    sparse_mla_topk_lens: Optional[torch.Tensor] = None


Major extensibility concern: this supposedly shared runtime contract already mixes generic
sparse-index fields with DeepSeek-V4-specific cache inputs and SkipSoftmax-specific
thresholds.

With this design, every new algorithm-specific runtime value requires modifying the common
dataclass and likely the shared THOP/C++ attention ABI, causing unrelated algorithms and
backends to depend on fields they never use.

Please keep only genuinely algorithm-independent indexing data in SparseRuntimeParams.
Algorithm-specific values should remain in a typed/discriminated payload and be lowered by the
selected backend at the THOP boundary. This would prevent the common carrier from growing into
a sparse-attention parameter bag.

github-actions Bot assigned lfr-0531 Apr 3, 2026

lfr-0531 requested a review from heyuhhh April 3, 2026 09:45

lfr-0531 force-pushed the user/fanrongl/sparse-attention-refactor branch 8 times, most recently from a1e4402 to eaab4c3 Compare April 3, 2026 11:18

lfr-0531 mentioned this pull request Apr 3, 2026

[None][doc] add attention developer guide #12693

Merged

1 task

lfr-0531 force-pushed the user/fanrongl/sparse-attention-refactor branch 2 times, most recently from 6d89705 to 4004b1a Compare April 4, 2026 06:23

heyuhhh reviewed Apr 7, 2026

View reviewed changes

lfr-0531 force-pushed the user/fanrongl/sparse-attention-refactor branch 3 times, most recently from 79c8913 to 53a6ee5 Compare April 7, 2026 03:07

heyuhhh approved these changes Jul 27, 2026

View reviewed changes

yihwang-nv reviewed Jul 27, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/attention_backend/sparse/hooks.py Outdated

lfr-0531 added 14 commits July 27, 2026 11:35

[None][refactor] isolate sparse attention implementations

3a635f7

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][refactor] unify sparse MLA backend dispatch

a0f2849

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][refactor] route sparse MLA through module facade

9dbf096

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][test] preserve inference mode in sparse MLA tests

db8cfae

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>

[None][fix] fix sparse attention fallback paths

3f9d28c

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][fix] fix sparse imports and PP validation

0bbdccd

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][refactor] adopt standard linting for sparse attention

6e49e73

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>

[None][fix] fix sparse attention compile interfaces

17b703d

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][refactor] clarify sparse attention hook contracts

c924efc

Group all accepted module hook signatures under one contract table and document the lifecycle of each sparse attention hook. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][docs] clarify MiniMax M3 sparse integration

efb0b47

Document that the model-specific integration is a compatibility exception and should not be used as the pattern for new sparse attention algorithms. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][refactor] add typed sparse attention hook adapters

87e815e

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

[None][test] update DeepSeek V4 hook assertion

a56b81c

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

bobboli approved these changes Jul 28, 2026

View reviewed changes

mingyangHao reviewed Jul 28, 2026

View reviewed changes


		TLLM_CHECK(host_kv_cache_pool_mapping.has_value());
		int32_t const layer_num = host_kv_cache_pool_mapping.value().size(0);



		@lru_cache(maxsize=None)
		def _get_sparse_mla_hooks_for_algorithm(algorithm: str) -> Optional[MLASparseHooks]:

Uh oh!

Conversation

lfr-0531 commented Apr 3, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

lfr-0531 commented Apr 3, 2026

Uh oh!

tensorrt-cicd commented Apr 3, 2026

Uh oh!

tensorrt-cicd commented Apr 3, 2026

Uh oh!

lfr-0531 commented Apr 4, 2026

Uh oh!

tensorrt-cicd commented Apr 4, 2026

Uh oh!

tensorrt-cicd commented Apr 4, 2026

Uh oh!

lfr-0531 commented Apr 6, 2026

Uh oh!

tensorrt-cicd commented Apr 6, 2026

Uh oh!

tensorrt-cicd commented Apr 6, 2026

Uh oh!

lfr-0531 commented Apr 6, 2026

Uh oh!

tensorrt-cicd commented Apr 6, 2026

Uh oh!

tensorrt-cicd commented Apr 6, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

lfr-0531 commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

lfr-0531 commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

tensorrt-cicd commented Jul 27, 2026

Uh oh!

lfr-0531 commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

bobboli commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

lfr-0531 commented Apr 3, 2026 •

edited by coderabbitai Bot

Loading

bobboli commented Jul 28, 2026 •

edited

Loading