Repository navigation
[None][feat] Add CSA2 Vanilla sparse attention - #20431
Open
yihwang-nv wants to merge 6 commits into
Open
yihwang-nv wants to merge 6 commits into
yihwang-nv wants to merge 6 commits into
Conversation
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
Preserve the independent DSA Vanilla backend and its regression coverage. Reconcile the indexer test helper with upstream and limit signature exclusions to optional non-DSA modes. Signed-off-by: yihwang-nv <yihwang@nvidia.com>
Share phase, score, TopK mapping, and latent-cache helpers while preserving the independent PyTorch reference. Document quantization and cache contracts, improve input validation, and move kernel-only reference adapters into test utilities. Trim redundant tests while retaining quantized Q projection and generation RoPE/cache ownership coverage. Validation on umbriel-b200-043 (B200): 1305 passed, 4 existing skips, 0 failed. Ruff lint, formatting, and diff checks passed. Signed-off-by: yihwang-nv <yihwang@nvidia.com>
Add CSA2VanillaAttention and CSA2VanillaIndexer, a pure-PyTorch golden of CSA2 sparse attention: sliding-window and compressed-group cache writes, candidate publication, exact top-k selection inside the published candidates, and sink attention over the window plus selected rows. Geometry comes from request lengths and the cache manager's host page lists, never from the kernel-produced slot or page tables. Tests cover the golden on CPU tensors with hand-derived values and compare it with the TRTLLM backend on SM100. Move the RoPE and phase helpers shared with the DSA golden into sparse/_vanilla_utils.py. Signed-off-by: yihwang-nv <yihwang@nvidia.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
WalkthroughAdds pure-PyTorch reference implementations for DSA and CSA2 sparse attention. The changes connect DSA to the Vanilla backend registry and extend sparse-attention testing across indexer selection, cache handling, and attention outputs. ChangesDSA sparse MLA
CSA2 sparse attention
Priority: ➖ Normal Estimated code review effort: 5 (Critical) | ~120 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant Indexer as DSAVanillaIndexer
participant Metadata as DSA metadata
participant Cache as DSA cache manager
participant Attention as DSAVanillaAttention
Indexer->>Cache: Append projected indexer keys
Indexer->>Metadata: Publish TopK selections
Attention->>Metadata: Read selections and phase data
Attention->>Cache: Load selected latent rows
Attention->>Attention: Compute sparse MLA output
sequenceDiagram
participant State as CSA2ForwardState
participant Indexer as CSA2VanillaIndexer
participant Cache as CSA2 cache pools
participant Attention as CSA2VanillaAttention
State->>Indexer: Supply query and visible cache metadata
Indexer->>State: Publish selected indices and candidate blocks
State->>Attention: Supply layer state and query tile
Attention->>Cache: Write cache rows and resolve selected slots
Attention->>Attention: Compute sink attention over valid rows
Suggested reviewers: Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 66.83% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 205 functions across 16 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Title check The title is concise, specific, and accurately identifies the primary change: adding CSA2 Vanilla sparse attention. Description check The description explains the purpose, implementation scope, limitations, and test coverage. It includes the required Description, Test Coverage, and PR Checklist sections and reports concrete test res… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
- Fix all pre-merge checks with AI
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Collaborator
|
@farazkh80, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~ Assignment explanation
Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once. |
yuxianq
requested review from
farazkh80
and removed request for
PerkzZheng,
yunruis and
yuxianq
October 15, 2026 17:15
farazkh80
approved these changes
Oct 16, 2026
farazkh80
left a comment
farazkh80
left a comment
Collaborator
There was a problem hiding this comment.
Approved from attention perspective, with two questions (non-blocking).
| swa_slots, | ||
| state.swa_kv, | ||
| "swa", | ||
| state.swa_transform, |
Collaborator
There was a problem hiding this comment.
main_transform and index_transform RoPE positions are checked against the golden's geometry, but swa_transform isn't. A wrong SWA position would match in both golden and TRTLLM.
Same check apply to swa_transform?
| if query_start < batch.num_ctx_tokens < end: | ||
| raise ValueError("CSA2 query tiles must not cross the context/generation boundary") | ||
| if query_start == 0: | ||
| self._enter_layer(source, layer) |
Collaborator
There was a problem hiding this comment.
_enter_layer runs before the write-time checks. After a rejection, a retry on the same metadata fails with "cannot be reused" instead of the real error.
Should we move it after those checks?
|
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements. |
This branch has not been deployed
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
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
vanilla.py. It does not change runtime behavior.QA Engineer Review
No test changes.
Per-File QA Perspective
tensorrt_llm/_torch/attention/backends/vanilla.py: Import formatting only. No QA behavior change to verify.Depends on #16309. Only the last commit belongs to this PR.
Description
CSA2VanillaAttentionand a weightlessCSA2VanillaIndexer(tensorrt_llm/_torch/attention/backends/sparse/csa2/vanilla_backend.py), a pure-PyTorch golden reference for CSA2 sparse attention:kv_lens_cuda - seq_lensand pages from the cache manager's host page lists. The golden never reads the kernel-produced slot tables, page tables or visible lengths.sparse/_vanilla_utils.py(pure move; thephase_is_generationerror now says "Sparse MLA requires…" instead of "DSA requires…", andDSAVanillaEmbeddingis no longer listed in the DSA golden's__all__).Scope: the sparse registry and the CSA2 module are unchanged, so
attn_backend="VANILLA"still raises forcsa2(the module buildscreate_attention("TRTLLM", ...)). Tests construct the golden directly; plumbing the configured backend into the module is a separate change. Like the other Vanilla goldens, it rejects CUDA graph capture.Test Coverage
tests/unittest/_torch/attention/sparse/csa2/test_vanilla_backend.py:Results on B200:
pytest tests/unittest/_torch/attention/sparse/csa2/test_vanilla_backend.py— 18 passed (3 consecutive runs).pytest tests/unittest/_torch/attention/sparse/csa2(existing tests) — 592 passed, 12 skipped.pytest tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py(shared helper move) — 16 passed.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.