[None][feat] Add DSA Vanilla sparse attention - #16309
Conversation
|
/bot run --disable-fail-fast |
yihwang-nv
changed the title
|
PR_Github #59376 [ run ] triggered by Bot. Commit: |
| # (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: |
There was a problem hiding this comment.
Why we cannot enable cuda graph for sparse case?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Please add a TODO here, let's try to add it in a follow-up PR
|
PR_Github #59376 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59435 [ run ] triggered by Bot. Commit: |
|
PR_Github #59435 [ run ] completed with state |
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughVanillaAttention 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. ChangesDSA sparse MLA support
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
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #67030 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #67323 [ run ] triggered by Bot. Commit: |
|
PR_Github #67323 [ run ] completed with state |
|
/bot run |
|
PR_Github #68108 [ run ] triggered by Bot. Commit: |
|
PR_Github #68108 [ run ] completed with state
|
mikeiovine
left a comment
mikeiovine
left a comment
There was a problem hiding this comment.
Stamp on behalf of runtime devs; delegating review to @NVIDIA/trt-llm-torch-attention-devs
|
/bot run --disable-fail-fast |
|
PR_Github #69658 [ run ] triggered by Bot. Commit: |
|
PR_Github #69658 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #69800 [ run ] triggered by Bot. Commit: |
|
PR_Github #69800 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #71968 [ run ] triggered by Bot. Commit: |
|
PR_Github #71968 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #71981 [ run ] triggered by Bot. Commit: |
|
PR_Github #71981 [ run ] completed with state
|
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
|
/bot run |
|
PR_Github #72100 [ run ] triggered by Bot. Commit: |
|
PR_Github #72100 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #72146 [ run ] triggered by Bot. Commit: |
|
PR_Github #72146 [ run ] completed with state
|
4 participants
Depends on #18044.
Description
DSAVanillaAttentionandDSAVanillaIndexerPyTorch golden implementations with the same module-facing contracts as the TRTLLM DSA backend.DSATrtllmAttention.TRTLLM_DSA_INDEXER_BF16projection mode without dispatching through the production fused kernels.max_model_lenwhose 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.ruff check,ruff format --check, andgit diff --checkfor 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.