[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 standalone DSAVanillaAttention and DSAVanillaIndexer PyTorch golden implementations with the same module-facing contracts as the TRTLLM DSA backend.
Implement the fake-fused RoPE path entirely with PyTorch APIs. Context and generation paths apply RoPE exactly once and preserve the same cache-ownership contract as DSATrtllmAttention.
Cover paged latent-cache reads and writes, selected sparse MLA, FP8-DS-MLA inline scales, FP8/FP4 indexer quantization, ragged and paged MQA logits, TopK selection, compressed index coordinates, cached-prefix generation, and the TRTLLM_DSA_INDEXER_BF16 projection mode without dispatching through the production fused kernels.
Register Vanilla DSA and update the sparse-MLA/backend test harnesses to use Vanilla Attention plus the Vanilla Indexer as independent golden references.
Fix the paged-MQA PyTorch reference for a max_model_len whose final page is partial, and validate compressed generation coordinates when a cached prefix is present.

The implementation is intentionally standalone: projection, quantization, RoPE, cache reconstruction, scoring, and selection are expressed in PyTorch so a defect in a TRTLLM/DeepGEMM/CuTe kernel cannot be hidden by sharing its implementation with the golden.

Test Coverage

pytest -q tests/unittest/_torch/attention/sparse/dsa — passed on B200 before the upstream rebase.
pytest -q tests/unittest/_torch/attention/test_attention_backends.py -k deepseekv3_2_dsa_mla — passed on B200 before the upstream rebase.
Repository pre-commit hooks on the rebased commits — passed.
ruff check, ruff format --check, and git diff --check for the rebased DSA changes — passed.

The current host cannot collect TensorRT-LLM pytest tests because its Python environment cannot load libmpi; CI should revalidate the rebased head.

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

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.

Last reviewed commit: ca9744a
Latest commit, not reviewed: 513e7a3

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 | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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.

✅ Passed checks (4 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 clearly identifies the feature: DSA sparse attention in the Vanilla backend. It follows the required ticket and type format.
Description check	✅ Passed	The description explains the implementation, scope, dependency, and relevant test coverage. It does not include explicit PR Checklist confirmations, but the required technical content is mostly comple…

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

tensorrt-cicd · 2026-09-10T21:00:09Z

PR_Github #72765 [ run ] completed with state SUCCESS. Commit: 33054de
/LLM/main/L0_MergeRequest_PR pipeline #59752 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: Yihan Wang <yihwang@nvidia.com>

yihwang-nv · 2026-09-11T02:51:05Z

/bot run

tensorrt-cicd · 2026-09-11T02:57:59Z

PR_Github #72837 [ run ] triggered by Bot. Commit: 55edaaa Link to invocation

tensorrt-cicd · 2026-09-11T05:53:23Z

PR_Github #72837 [ run ] completed with state SUCCESS. Commit: 55edaaa
/LLM/main/L0_MergeRequest_PR pipeline #59820 completed with status: 'SUCCESS'

CI Report

Link to invocation

yuxianq · 2026-10-10T03:08:56Z

@yuxianq, this PR is already assigned to you as trt-llm-torch-attention-devs based on your earlier review. Could you help with the renewed review request? Thanks~

Assignment explanation

Components	backends/sparse (60%); backends/vanilla (5%); integration/attention-tests (35%)
Tags	algorithm:sparse (100%); specific-algorithm:dsa (100%); representation:mla (70%); representation:mqa (30%); phase:prefill (50%); phase:decode (50%); features:rope (35%); features:quantization (35%); features:paged-kv (30%); context:llm (100%)
PR effort	8 points — The new reference must preserve RoPE, cache ownership, quantization and sparse-coordinate contracts while remaining sufficiently independent of production implementations. This requires deep algorithm and oracle review, not a new-model classification.
Importance	Important — Introduces a specialized DSA reference backend and changes the golden comparisons used to validate sparse selection and attention. Errors in this reference could conceal production correctness defects.
Affinity (0–1)	1.000 component/tag match; 1.000 after PIC preference
Outstanding points at selection	185 → 185 (+0)
Points already reserved on this PR	8
Capacity	30 points/week
Projected load	6.17 capacity weeks
Selection rule	existing review owner; renewed attention team request

Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once.

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>

yihwang-nv · 2026-10-15T07:50:14Z

/bot run

tensorrt-cicd · 2026-10-15T07:58:01Z

PR_Github #78057 [ run ] triggered by Bot. Commit: 72ceed3 Link to invocation

tensorrt-cicd · 2026-10-15T11:28:18Z

PR_Github #78057 [ run ] completed with state FAILURE. Commit: 72ceed3
/LLM/main/L0_MergeRequest_PR pipeline #64339 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

One issue with the FP8 block-scale projection reference.

… indexer end to end - Pick UE8M0 vs float32 1x128 activation scales in the Vanilla indexer wq_b reference by mirroring FP8BlockScalesLinearMethod.apply dispatch (SM100 family, SM107 CuTe DSL MXFP8, SM120, int32 weight scales). - Dequantize FP8 QDQ in float32 before casting back to the input dtype. - Tile weighted_relu_scores over queries to bound [H, T, S] intermediates on long sequences. - Run the indexer end to end in all sparse MLA cases, including the long ones, instead of injecting random top-k indices. - Cover DSA sparse backend cases on both v1 and v2 KV cache managers. Signed-off-by: yihwang-nv <yihwang@nvidia.com>

yihwang-nv · 2026-10-28T06:42:55Z

/bot run

- Follow fp8_quantize_1x128's arithmetic in _qdq_fp8_blocks (quant = 448 / amax, dequant = 1 / quant, quantize x * quant). - Add a GPU test comparing _qdq_fp8_blocks with the native quantizer: exact for UE8M0 scales, and at most one adjacent E4M3 code on rare elements for float32 scales, since the kernel is built with --use_fast_math. - Use an asymmetric input in the CPU dispatch test so quantization errors cannot cancel in the projection. Signed-off-by: yihwang-nv <yihwang@nvidia.com>

yihwang-nv · 2026-10-28T09:45:07Z

/bot run

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

yihwang-nv · 2026-10-29T03:25:01Z

/bot run

tensorrt-cicd · 2026-10-29T03:33:29Z

PR_Github #81324 [ run ] triggered by Bot. Commit: bc9e4ef Link to invocation

tensorrt-cicd · 2026-10-29T07:30:34Z

PR_Github #81324 [ run ] completed with state FAILURE. Commit: bc9e4ef
/LLM/main/L0_MergeRequest_PR pipeline #67009 completed with status: 'FAILURE'

CI Report

⚠️ Resource Maintenance — Test Coverage Reduced:

Resource maintenance in LLM/main/L0_Test-x86_64-Single-GPU #9635: L40S-PyTorch-2 is temporarily paused due to Kubernetes agent provisioning and node issues.
Affected stages:
L40S-PyTorch-2

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

The CuTe DSL paged-MQA atom split reads its expanded schedule from on_update_kv_lens(), which the executor runs after prepare() on every step. Call it in the generation steps so the TRTLLM indexer matches the Vanilla golden. Signed-off-by: yihwang-nv <yihwang@nvidia.com>

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

yihwang-nv · 2026-10-30T05:45:00Z

/bot run

tensorrt-cicd · 2026-10-30T05:51:59Z

PR_Github #81604 [ run ] triggered by Bot. Commit: 513e7a3 Link to invocation

tensorrt-cicd · 2026-10-30T09:24:04Z

PR_Github #81604 [ run ] completed with state SUCCESS. Commit: 513e7a3
/LLM/main/L0_MergeRequest_PR pipeline #67273 completed with status: 'SUCCESS'

CI Report

Link to invocation

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 Resolved

yuxianq reviewed Jul 15, 2026

View reviewed changes

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

yuxianq reviewed Jul 15, 2026

View reviewed changes

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

yuxianq reviewed Jul 15, 2026

View reviewed changes

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

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