[None][feat] Add DSA sparse attention + indexer reference to VanillaAttention; unify sparse backend tests - #16309
[None][feat] Add DSA sparse attention + indexer reference to VanillaAttention; unify sparse backend tests#16309yihwang-nv wants to merge 12 commits into
Conversation
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
|
/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 |
- Register DSA in the vanilla sparse-backend factory so VanillaAttention is built through create_attention like every other backend (no special-casing in the sparse runner); update_quant_config is now called unconditionally. - Derive sparse family / selection unit / top-k from is_mla and sparse_attention_config, and move the model-agnostic sweep dimensions to shared constants, so ModelAttnConfig no longer carries a SparseHarnessConfig or hf_config_overrides/atol/rtol. Sparse tolerance is derived in _tolerances. - Clarify why the backend-only sparse oracle skips the captured-CUDA-graph path. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #59435 [ run ] triggered by Bot. Commit: |
|
PR_Github #59435 [ run ] completed with state |
📝 WalkthroughWalkthroughDSA sparse MLA support is added to ChangesDSA sparse MLA support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant VanillaAttention
participant SparsePredictors
participant PagedKVCache
participant DSAAttention
VanillaAttention->>SparsePredictors: compute sparse KV and attention selections
VanillaAttention->>DSAAttention: pass selected indices and MLA inputs
DSAAttention->>PagedKVCache: load latent-cache prefix
DSAAttention->>VanillaAttention: return selected-token attention output
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #59621 [ run ] completed with state |
|
PR_Github #59643 [ kill ] completed with state |
| if case.is_mla and case.sparse_selection_unit == "token": | ||
| inputs = generate_sparse_mla_inputs(case, seed) | ||
| elif case.sparse_selection_unit == "block": | ||
| inputs = generate_sparse_block_inputs(case, seed) |
There was a problem hiding this comment.
Since sparse mla uses absorption, can we view it as a special case of generate_mla_gen_inputs? Thus we can use a single generate_mla_gen_inputs for all three cases generate_mla_gen_inputs/generate_sparse_block_inputs/generate_sparse_mla_inputs, just need to generate an extra topk_indices for sparse case.
| block_size = getattr(sparse_params, "indices_block_size") | ||
| selection_gen = torch.Generator(device="cuda").manual_seed(seed + 1) | ||
| _ = _build_sparse_block_indices(case, selection_gen, block_size) | ||
| raise NotImplementedError( |
There was a problem hiding this comment.
If we don't support block sparse attention now, I think we should avoid to introduce abstraction for it, e.g., sparse_selection_unit. Whether an abstraction is proper should be validated by an real implementation, otherwise agents will inject some seems-good garbage in some corner.
| return indices, singleton_oracle_mask | ||
|
|
||
|
|
||
| def _build_sparse_block_indices( |
There was a problem hiding this comment.
Same comments for all block sparse related code.
| if sparse_config is not None and backend != "VANILLA": | ||
| cls = get_sparse_attn_kv_cache_manager(sparse_config) | ||
| if model_config is None or pretrained_config is None: | ||
| raise ValueError("Sparse cache manager requires model and pretrained configs") |
There was a problem hiding this comment.
model_config/pretrained_config are optional for sparse cache manager, so that we don't need to create and pass them.
| dtype=torch_dtype_to_binding(case.compute_dtype), | ||
| ) | ||
|
|
||
| if sparse_config is not None and backend != "VANILLA": |
There was a problem hiding this comment.
Can vanilla use sparse cache manager? If it can, we don't need to exclude vanilla.
| cu_q_seqlens = torch.empty(case.num_seqs + 1, dtype=torch.int32, device="cuda") | ||
| cu_kv_seqlens = torch.empty(case.num_seqs + 1, dtype=torch.int32, device="cuda") | ||
| fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device="cuda") | ||
| attn.mla_rope_generation( |
There was a problem hiding this comment.
Don't need to call mla_rope_generation for trtllm backend. trtllm backend does not call rope inside backend forward in this case, we should remove rope from vanilla to align with it (vanilla does not fuse rope, MLA will run rope for it, we should keep this behavior). But mla_rope_generation also update KV cache, to ensure KV cache is updated for trtllm backend, we should pass skip_mla_rope_generation=True in forward args like _run_mla_gen_backend, which hints the trtllm backend that mla_rope_generation is skipped, and it will add an update KV cache step before running fmha.
| torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) | ||
|
|
||
|
|
||
| def _run_sparse_mla_backend( |
There was a problem hiding this comment.
Can we merge _run_sparse_mla_backend back to _run_mla_gen_backend?
| num_kv_heads: Optional[int] = None, | ||
| quant_config: Optional[QuantConfig] = None, | ||
| q_scaling: Optional[float] = None, | ||
| pos_embd_params: Optional[PositionalEmbeddingParams] = None, |
There was a problem hiding this comment.
vanilla does not fuse rope, should we remove rope from vanilla?
| return attn_output | ||
|
|
||
|
|
||
| class VanillaIndexer: |
There was a problem hiding this comment.
It is only used by tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py, not our test suite. It has different API with DSA's Indexer (
). Only DSA and its derived ds v4 attention use Indexer class, not other sparse attention uses it, so it is no general. How about removing VanillaIndexer and revert change for tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py?
| "Vanilla selected MLA currently supports only DSA") | ||
| kv_idx, kv_off = self.sparse_kv_predict(q, k, metadata, | ||
| forward_args) | ||
| at_idx, at_off = self.sparse_attn_predict( |
There was a problem hiding this comment.
Use full name for kv_idx/kv_off/at_idx/at_off to make them readable. E.g., you can use the same field name with sparse_prediction sparse_kv_indices/sparse_kv_offsets/sparse_attn_indices/sparse_attn_offsets so that we don't need to convert names in our mind.
A token selection is a block selection with block_size == 1, so the harness no longer distinguishes token vs block selection units. Selection granularity is derived from the config (get_indices_block_size(); 1 for DSA) and the single selected-MLA path is exercised by DSA. Removes the unverified block-selection scaffolding (builder + runner/generator stubs) and the selection-unit plumbing; VanillaAttention._mla_forward_sparse guards block_size == 1. block_size > 1 is deferred to a follow-up PR (RocketKV / MiniMax-M3). Signed-off-by: Yihan Wang <yihwang@nvidia.com>
…la forward Signed-off-by: Yihan Wang <yihwang@nvidia.com>
| sparse_attn_indices, | ||
| forward_args.attention_input_type, | ||
| ) | ||
| if self.sparse_params is not None: |
There was a problem hiding this comment.
This block is unreachable, we can remove it.
| rotated = torch.stack((out1, out2), dim=-1).flatten(-2) | ||
| return torch.cat((rotated, x_pass), dim=-1) | ||
|
|
||
| def _prepare_sparse_mla_inputs( |
There was a problem hiding this comment.
Without rope, both _apply_rotary_embedding and _prepare_sparse_mla_inputs are unnecessary.
…o its test VanillaIndexer is a test-only fp32 reference used only by test_sparse_mla_forward.py and has a different API from the production DSA Indexer, so it does not belong in the production attention backend. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
| raise ValueError( | ||
| "Vanilla sparse MLA expects absorbed queries and latent cache, " | ||
| "not explicit K/V tensors") | ||
| return self._mla_forward_sparse( |
There was a problem hiding this comment.
The main diff between _mla_forward_sparse and _mla_forward_generation is that _mla_forward_sparse use paged cache but _mla_forward_generation uses linear cache. Can we make vanilla always use linear cache, so that we can rename _mla_forward_generation to _mla_forward_absorption and add topk_indices handle to it to support sparse attention?
…_config The sparse KV-cache manager takes sparse_attention_config directly, so model_config/pretrained_config are not needed, and Vanilla can use the same sparse cache manager as the other backends (no VANILLA special-case). Signed-off-by: Yihan Wang <yihwang@nvidia.com>
| mla_bmm2_scale = None | ||
| quant_q_buffer = None | ||
|
|
||
| gen_layers[layer_idx].mla_rope_generation( |
There was a problem hiding this comment.
Both of the rope and page cache design come from the fact that tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py uses this design. We can also update the tests in this file to remove rope for both vanilla and trtllm backend, and use linear cache for vanilla, so it aligns with the original design of vanilla.
…tion - Remove VanillaAttention's internal sparse-MLA RoPE. The harness now applies RoPE on the host and feeds pre-formed inputs, using skip_mla_rope_generation for the absorbed generation path (matches production: the MLA module runs RoPE, the attention backend does not). - Context phase feeds raw inputs so the TRTLLM MLA context kernel ropes once (avoids double RoPE); Vanilla and the generation path use the pre-RoPE'd inputs. - Gate TRTLLM DSA to sm>=100: the trtllm-gen DynamicTokenSparse FMHA kernels are Blackwell-only; on Hopper MLA generation falls back to dense FlashMLA, which has no per-token sparse path. - Drop the block-unit (sparse_block_size) abstraction; the suite is token-selection only (DSA / DeepSeek-V4). RocketKV block sparse is out of scope for this suite. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #60331 [ run ] triggered by Bot. Commit: |
|
PR_Github #60331 [ run ] completed with state
|
| num_kv_heads: Optional[int] = None, | ||
| quant_config: Optional[QuantConfig] = None, | ||
| q_scaling: Optional[float] = None, | ||
| pos_embd_params: Optional[PositionalEmbeddingParams] = None, |
There was a problem hiding this comment.
pos_embd_params is unused now
There was a problem hiding this comment.
Removed — dropped pos_embd_params from VanillaAttention.__init__ and its import; it only fed the sparse-MLA RoPE that was removed. create_attention passes it as a kwarg, which the base backend swallows, so it's a no-op.
| HAS_FLASH_MLA = False | ||
|
|
||
|
|
||
| class VanillaIndexer: |
There was a problem hiding this comment.
If we don't use VanillaIndexer in our test suite, should we fully revert changes in this file?
There was a problem hiding this comment.
Done — fully reverted this file to main. VanillaIndexer is not used by the model-driven suite and has a different API than the production DSA Indexer, so the original free-function version is restored.
…indexer test - Remove the now-unused pos_embd_params from VanillaAttention.__init__ (and its import); it only fed the sparse-MLA RoPE that was removed. create_attention passes it as a kwarg, which the base backend swallows, so this is a no-op. - Fully revert tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py: VanillaIndexer is not used by the model-driven suite and has a different API than the production DSA Indexer, so restore the original file. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
|
|
||
| per_token_outputs = [] | ||
| for token_idx in range(q_len): | ||
| row = topk_indices[token_offset + token_idx] |
There was a problem hiding this comment.
Can we use topk_indices to select kv and concat them in advance, so that we can directly use torch.nn.functional.scaled_dot_product_attention like dense mla case? If so, we can easily merge _mla_forward_sparse to _mla_forward_generation.
3 participants
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Description
Integrate DSA (DeepSeek Sparse Attention) into the
VanillaAttentionreference backend and express sparse-attention tests through the unified model-driven backend harness, so a sparse algorithm plugs into Vanilla the same way it does into TRTLLM. Vanilla is intended as the reference platform for the planned TRTLLM sparse refactor.Vanilla DSA attention-over-selection (
vanilla.py). Backend-neutralsparse_kv_predict/sparse_attn_predicthooks matching the production contract (results written intoforward_args.sparse_prediction), plus the_mla_forward_sparsegolden: applies MLA RoPE, appends and reconstructs paged latent KV across context and generation, validates top-k shape/dtype/padding/bounds/causality before mutating the cache, gathers selected latent K/V, and evaluates absorbed MLA with an FP32 softmax.Unified sparse test harness (
backend_case.py,model_attn_config.py,backend_capability.py,test_attention_backends.py). A sparse workload is an ordinaryModelAttnConfig(deepseekv3_2_dsa_mla) sitting next to the dense MLA configs. It carries the real user-facingsparse_attention_config— lowered via the productionto_sparse_params/to_sparse_metadata_params/ sparse KV-cache manager — plus a backend-neutralSparseHarnessConfig. The runner injects deterministic causal selections and compares every supported backend against the Vanilla golden. This replaces the standalone DSA test files.VanillaIndexer(vanilla.py). An FP32 reference for the production DSA / DeepSeek-V4Indexer, wrapping a production indexer instance so it runs against that indexer's own weights. It is a standalone reference (not wired intoforward), andtest_sparse_mla_forward.pynow builds its DeepSeek-V4 reference top-k on it, removing the duplicated projection / scoring / top-k helpers.Block-selection scaffold (
backend_case.py). A reusable causal block-index builder plus validation and dispatch stubs for the futureselection_unit="block"family (e.g. RocketKV); inert until a block config lands.Scope is limited to the internal Vanilla backend and sparse tests. The DSA indexer stays external, the production backend remains
DSATrtllmAttention, and no other sparse algorithm or public API is changed. The explicit Vanilla paths are correctness oracles, not performance paths.Test Coverage
Verified on a Blackwell GPU (sm 100), where the TRTLLM DSA path runs and is compared against the Vanilla golden:
pytest tests/unittest/_torch/attention/test_attention_backends.py -k dsa→ 6 passed (deepseekv3_2_dsa_mla× {ctx, gen, mix} × {random, singleton}).pytest tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py -k "deepseek_v4 and auto and (small_prefill or small_decode or small_mixed)"→ passed (DeepSeek-V4 indexer reference now driven byVanillaIndexer).py_compileand import checks pass.On
sm < 100, TRTLLM DSA is skipped and only the Vanilla golden (plus its analytic singleton oracle) runs.PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, comment
/bot help.