Skip to content

[None][refactor] Unify sparse attention framework with clean backend interfaces - #12733

Open
lfr-0531 wants to merge 22 commits into
NVIDIA:mainfrom
lfr-0531:user/fanrongl/sparse-attention-refactor
Open

[None][refactor] Unify sparse attention framework with clean backend interfaces#12733
lfr-0531 wants to merge 22 commits into
NVIDIA:mainfrom
lfr-0531:user/fanrongl/sparse-attention-refactor

Conversation

@lfr-0531

@lfr-0531 lfr-0531 commented Apr 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

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
lfr-0531 requested a review from heyuhhh April 3, 2026 09:45
@lfr-0531
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

lfr-0531 commented Apr 3, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Apr 3, 2026

Copy link
Copy Markdown
Collaborator

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

@lfr-0531 lfr-0531 mentioned this pull request Apr 3, 2026
Merged
1 task
@tensorrt-cicd

tensorrt-cicd commented Apr 3, 2026

Copy link
Copy Markdown
Collaborator

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
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
@lfr-0531

lfr-0531 commented Apr 4, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Apr 4, 2026

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

tensorrt-cicd commented Apr 4, 2026

Copy link
Copy Markdown
Collaborator

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

lfr-0531 commented Apr 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Apr 6, 2026

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

tensorrt-cicd commented Apr 6, 2026

Copy link
Copy Markdown
Collaborator

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

lfr-0531 commented Apr 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Apr 6, 2026

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

tensorrt-cicd commented Apr 6, 2026

Copy link
Copy Markdown
Collaborator

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
heyuhhh reviewed Apr 7, 2026
View reviewed changes
Comment thread cpp/tensorrt_llm/thop/attentionOp.cpp
heyuhhh
heyuhhh reviewed Apr 7, 2026
View reviewed changes
Comment thread .pre-commit-config.yaml
@lfr-0531
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
@lfr-0531

lfr-0531 commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #64380 [ run ] triggered by Bot. Commit: b46b693 Link to invocation

lfr-0531 added 22 commits August 7, 2026 01:47
@lfr-0531
[None][refactor] isolate sparse attention implementations
4f2ad5d 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][refactor] unify sparse MLA backend dispatch
94fc988 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][refactor] route sparse MLA through module facade
25dff25 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][refactor] refine sparse attention module hooks
8675a55 
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>
@lfr-0531
[None][refactor] consolidate sparse attention runtime interfaces
fbd7f5b 
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>
@lfr-0531
[None][refactor] simplify shared DSA top-k storage
fc230c1 
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>
@lfr-0531
[None][test] preserve inference mode in sparse MLA tests
e78e616 
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>
@lfr-0531
[None][fix] fix sparse attention fallback paths
5e1f9a2 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][fix] fix sparse imports and PP validation
7638bff 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][refactor] adopt standard linting for sparse attention
6c91b8d 
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>
@lfr-0531
[None][fix] fix sparse attention compile interfaces
06d27b4 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][refactor] clarify sparse attention hook contracts
6f6eef7 
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>
@lfr-0531
[None][docs] clarify MiniMax M3 sparse integration
b5d54c8 
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>
@lfr-0531
[None][refactor] add typed sparse attention hook adapters
dba871c 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][test] update DeepSeek V4 hook assertion
6f0d72d 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][refactor] refine sparse attention hook contracts
3daa05b 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][fix] preserve DSA MTP top-k sharing after rebase
d1db820 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][fix] preserve sparse MLA auxiliary streams
a17e6f6 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][fix] Pass DeepSeek-V4 backend params to indexer
81c0d50 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][fix] handle phase-specific DSA MTP top-k stash
8bcdc66 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][fix] preserve masked DSA indexer cache allocation
41ef8a1 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531
[None][fix] preserve DSA runtime metadata after rebase
a181ebe 
Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
@lfr-0531

lfr-0531 commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #64447 [ run ] triggered by Bot. Commit: a181ebe Link to invocation

@tensorrt-cicd

tensorrt-cicd commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #64380 [ run ] completed with state ABORTED. Commit: b46b693

Link to invocation

@tensorrt-cicd

tensorrt-cicd commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #64447 [ run ] completed with state FAILURE. Commit: a181ebe
/LLM/main/L0_MergeRequest_PR pipeline #52322 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

@lfr-0531

lfr-0531 commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #64535 [ run ] triggered by Bot. Commit: a181ebe Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@coderabbitai coderabbitai[bot] coderabbitai[bot] left review comments
@yuxianq yuxianq yuxianq left review comments
@mingyangHao mingyangHao mingyangHao left review comments
@karljang karljang karljang approved these changes
@WeiHaocheng WeiHaocheng WeiHaocheng approved these changes
@bobboli bobboli bobboli approved these changes
@Shixiaowei02 Shixiaowei02 Shixiaowei02 approved these changes
@govind-ramnarayan govind-ramnarayan govind-ramnarayan approved these changes
@yihwang-nv yihwang-nv yihwang-nv approved these changes
@yizhang-nv yizhang-nv yizhang-nv approved these changes
@zhaoyangwang-nvidia zhaoyangwang-nvidia zhaoyangwang-nvidia approved these changes
@QiJune QiJune Awaiting requested review from QiJune
@brb-nv brb-nv Awaiting requested review from brb-nv brb-nv was automatically assigned from NVIDIA/trt-llm-torch-models-devs
@VALLIS-NERIA VALLIS-NERIA Awaiting requested review from VALLIS-NERIA VALLIS-NERIA is a code owner automatically assigned from NVIDIA/trt-llm-kv-cache-manager-devs
@pcastonguay pcastonguay Awaiting requested review from pcastonguay pcastonguay is a code owner automatically assigned from NVIDIA/trt-llm-disagg-devs
@chuangz0 chuangz0 Awaiting requested review from chuangz0 chuangz0 is a code owner automatically assigned from NVIDIA/trt-llm-disagg-devs
@kaiyux kaiyux Awaiting requested review from kaiyux kaiyux is a code owner automatically assigned from NVIDIA/trt-llm-doc-owners
@schetlur-nv schetlur-nv Awaiting requested review from schetlur-nv schetlur-nv was automatically assigned from NVIDIA/trt-llm-devs
+1 more reviewer
@heyuhhh heyuhhh heyuhhh approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

@lfr-0531 lfr-0531

Labels

api-compatible Accepted LLM API contract change that is backwards-compatible ci: full pre-merge approved

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.