[None][feat] Add DSA Vanilla sparse attention by yihwang-nv · Pull Request #16309 · NVIDIA/TensorRT-LLM

yihwang-nv · 2026-07-13T07:33:09Z

Depends on #18044.

Description

Add DSAVanillaAttention under sparse/dsa.
Register the Vanilla DSA backend.
Use Vanilla selected MLA as the golden reference in DSA sparse tests and the unified backend sweep.

yihwang-nv · 2026-07-15T06:14:24Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T06:21:08Z

PR_Github #59376 [ run ] triggered by Bot. Commit: b3238f5 Link to invocation

yuxianq · 2026-07-15T08:19:39Z

        # (production replays a captured decode graph); it must still match the
        # eager golden.
-        if case.is_gen_only:
+        if case.is_gen_only and not case.is_sparse:


Why we cannot enable cuda graph for sparse case?

This backend-only path is an eager correctness oracle: the standalone sparse runner validates the injected request-local selections on the host and rebuilds each request's logical cache with Python loops, neither of which is graph-capturable. Capturing would require rewriting the runner into a static-buffer form that no longer matches its role. Graph coverage of the DSA decode kernel is exercised at the model level, not by this backend oracle. I added a code comment making the exclusion explicit; happy to file a follow-up if you'd like graph coverage added here.

Please add a TODO here, let's try to add it in a follow-up PR

tensorrt-cicd · 2026-07-15T09:16:43Z

PR_Github #59376 [ run ] completed with state SUCCESS. Commit: b3238f5
/LLM/main/L0_MergeRequest_PR pipeline #47850 completed with status: 'SUCCESS'

CI Report

Link to invocation

yihwang-nv · 2026-07-15T10:26:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-15T10:33:32Z

PR_Github #59435 [ run ] triggered by Bot. Commit: 108df3f Link to invocation

tensorrt-cicd · 2026-07-15T13:47:47Z

PR_Github #59435 [ run ] completed with state SUCCESS. Commit: 108df3f
/LLM/main/L0_MergeRequest_PR pipeline #47904 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

coderabbitai · 2026-07-15T15:41:49Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

VanillaAttention now supports DSA sparse MLA with paged latent-cache handling, sparse selection routing, selected-token attention, and contiguous cache writes. The attention test harness adds sparse configuration, cache execution, deterministic inputs, capability checks, and Vanilla-versus-TRTLLM comparisons.

Changes

DSA sparse MLA support

Layer / File(s)	Summary
Sparse configuration and sweep generation `tests/unittest/_torch/attention/model_attn_config.py`, `tests/unittest/_torch/attention/backend_capability.py`, `tests/unittest/_torch/attention/test_attention_backends.py`	The harness defines DSA sparse configuration, enables Vanilla sparse capability, validates TRTLLM hardware requirements, and generates sparse context, generation, and mixed cases.
Vanilla sparse MLA execution `tensorrt_llm/_torch/attention_backend/vanilla.py`, `tensorrt_llm/_torch/attention_backend/sparse/registry.py`, `tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`	VanillaAttention materializes NHD or HND latent caches, validates sparse inputs, performs selected-token attention, routes DSA sparse requests, and uses contiguous dtype-reinterpreted cache copies. The registry dispatches DSA to VanillaAttention.
Sparse backend harness execution `tests/unittest/_torch/attention/backend_case.py`	BackendCase generates causal selections and sparse MLA inputs, builds sparse cache managers, runs context and generation phases, checks cache updates and singleton outputs, applies sparse tolerances, and skips sparse CUDA-graph replay.
Sparse golden comparison tests `tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py`	The DSA test runs VanillaAttention and TRTLLM with deterministic selections, applies backend-specific RoPE and metadata handling, and compares phase and layer outputs.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant BackendCase
  participant SparseKVCacheManager
  participant VanillaAttention
  participant TRTLLM
  BackendCase->>SparseKVCacheManager: create sparse paged cache
  BackendCase->>VanillaAttention: run sparse context and generation
  VanillaAttention->>SparseKVCacheManager: read and append latent cache
  BackendCase->>TRTLLM: run matching sparse phases
  BackendCase->>BackendCase: compare outputs and singleton oracle values

Possibly related PRs

NVIDIA/TensorRT-LLM#12733: Both changes use the shared sparse-attention backend interface and VanillaAttention.sparse_attn_predict integration.
NVIDIA/TensorRT-LLM#16060: Both changes modify DSA sparse MLA KV-cache handling and cache-manager integration.
NVIDIA/TensorRT-LLM#17273: Both changes modify the DeepSeek sparse MLA execution path.

Suggested reviewers: yunruis, cascade812, perkzzheng

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The description identifies the main changes but omits the required Test Coverage and PR Checklist sections and contains an incomplete description item.	Add the required Test Coverage and PR Checklist sections, and clarify the incomplete “Add under” item with the affected component and purpose.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title follows the required format and clearly identifies the addition of DSA sparse attention to the Vanilla backend.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

yihwang-nv · 2026-08-18T07:48:21Z

/bot run

tensorrt-cicd · 2026-08-18T07:54:38Z

PR_Github #67030 [ run ] triggered by Bot. Commit: 743f93d Link to invocation

tensorrt-cicd · 2026-08-18T10:06:39Z

PR_Github #67030 [ run ] completed with state FAILURE. Commit: 743f93d
/LLM/main/L0_MergeRequest_PR pipeline #54578 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

yihwang-nv · 2026-08-19T06:27:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-19T06:33:04Z

PR_Github #67323 [ run ] triggered by Bot. Commit: 743f93d Link to invocation

tensorrt-cicd · 2026-08-19T11:27:27Z

PR_Github #67323 [ run ] completed with state SUCCESS. Commit: 743f93d
/LLM/main/L0_MergeRequest_PR pipeline #54844 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

Signed-off-by: Yihan Wang <yihwang@nvidia.com>

yihwang-nv · 2026-08-21T03:02:13Z

/bot run

tensorrt-cicd · 2026-08-21T03:09:19Z

PR_Github #68108 [ run ] triggered by Bot. Commit: ba1eec2 Link to invocation

tensorrt-cicd · 2026-08-21T04:27:31Z

PR_Github #68108 [ run ] completed with state SUCCESS. Commit: ba1eec2
/LLM/main/L0_MergeRequest_PR pipeline #55555 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

yuxianq · 2026-08-21T09:41:32Z

+        sparse_params: Optional[DSAParams] = None,
+        **kwargs,
+    ) -> None:
+        if sparse_params is None:


I think we don't have to override the whole __init__ to just add an assertion, how about remove this override?

github-actions Bot assigned yihwang-nv Jul 13, 2026

yihwang-nv requested review from wenmingw and yuxianq July 13, 2026 08:01

yihwang-nv changed the title ~~[None][feat] Add DSA sparse attention to VanillaAttention~~ [None][feat] Add DSA sparse attention + indexer reference to VanillaAttention; unify sparse backend tests Jul 15, 2026

yuxianq reviewed Jul 15, 2026

View reviewed changes

Comment thread tests/unittest/_torch/attention/backend_case.py Outdated

yuxianq reviewed Jul 15, 2026

View reviewed changes

Comment thread tests/unittest/_torch/attention/backend_case.py Outdated

yuxianq reviewed Jul 15, 2026

View reviewed changes

Comment thread tests/unittest/_torch/attention/model_attn_config.py Outdated

yihwang-nv marked this pull request as ready for review July 15, 2026 15:05

yihwang-nv requested a review from a team as a code owner July 15, 2026 15:06

yihwang-nv requested review from PerkzZheng and yunruis July 15, 2026 15:06