[None][feat] Add DSA sparse attention reference to VanillaAttention and unify sparse backend tests - #16309
[None][feat] Add DSA sparse attention reference to VanillaAttention and unify sparse backend tests#16309yihwang-nv wants to merge 19 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 |
|
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 |
| 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.
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
|
/bot run |
|
PR_Github #63935 [ run ] triggered by Bot. Commit: |
|
PR_Github #63935 [ run ] completed with state
|
Signed-off-by: Yihan Wang <yihwang@nvidia.com>
The origin merge refactored AttentionForwardArgs and replaced the DSA backend, breaking the vanilla-vs-TRTLLM sparse MLA comparison: - Route caller-provided top-k through sparse_backend_args.topk_indices (VanillaAttention.sparse_attn_predict) and rename the dropped forward_args.sparse_prediction to sparse_runtime_params in vanilla.py and backend_case.py. - Restore the is_vanilla branching in the test (direct selection for vanilla, indexer mock + sparse_backend_args for TRTLLM) and drop the duplicated generation forward kwargs left by the merge. - VanillaAttention applies no RoPE itself while the kernel does via rope_append, so pre-apply the matching yarn RoPE to the vanilla q_pe/k_pe at absolute positions before the golden forward. All 12 parametrized cases pass. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
yihwang-nv
changed the title
VanillaAttention.sparse_kv_predict only ever returned (None, None) and its sole caller wrote those Nones back as the already-default sparse_runtime_params values. Vanilla handles sparse selection inline via sparse_attn_predict (not through the sparse hooks), so the method is dead. Drop it and the corresponding kv-index assignments. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (8)
tests/unittest/_torch/attention/backend_case.py (2)
316-320: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the new
sparse_configparameter.Every other parameter of
_build_mla_kv_cache_manageris annotated. Addsparse_config: Optional[SparseAttentionConfig] = None, as required by the Python guideline to annotate every function.As per coding guidelines: "Annotate every function, use
Nonefor procedures, avoid unnecessaryAnyandtype: ignore".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/backend_case.py` around lines 316 - 320, Annotate the sparse_config parameter in _build_mla_kv_cache_manager as Optional[SparseAttentionConfig] with a default of None, adding or reusing the necessary imports while preserving the existing function behavior.Source: Coding guidelines
1218-1222: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe non-MLA sparse guard is duplicated.
run_backendandrun_caseeach raiseValueError(f"Unsupported sparse contract: is_mla={case.is_mla}")._validate_sparse_casealready exists as the single validation point for sparse cases. Move theis_mlacheck there and let both callers rely on it.Also applies to: 1412-1417
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/backend_case.py` around lines 1218 - 1222, Move the non-MLA sparse validation into _validate_sparse_case, making it raise the existing “Unsupported sparse contract” ValueError for sparse cases where is_mla is false. Remove the duplicated guard and raise from both run_backend and run_case, while preserving their MLA sparse execution flow and relying on _validate_sparse_case as the single validation point.tensorrt_llm/_torch/attention_backend/vanilla.py (3)
679-709: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTwo MLA latent gather helpers now coexist.
_gather_paged_mla_latent(lines 189-208) and_load_mla_latent_cacheboth materialize a request's latent cache from paged blocks. The differences are HND support,BAD_PAGE_INDEXhandling, and the read-offset assumption. Consider folding the HND branch into_gather_paged_mla_latentand deleting one helper, so the dense and sparse MLA paths cannot drift.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/vanilla.py` around lines 679 - 709, Consolidate the duplicate MLA page-materialization logic by extending _gather_paged_mla_latent with the HND layout handling, BAD_PAGE_INDEX filtering, and required read-offset behavior currently implemented by _load_mla_latent_cache. Update all callers to use the unified helper, then remove _load_mla_latent_cache so dense and sparse MLA paths share one implementation.
818-845: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueThe per-token Python loop dominates the golden runtime.
Each query token runs one
index_selectand two matmuls in Python. For the context phase this is one iteration per packed token, per layer. When every row of a request has the same valid count, you can gather once withlatent.index_select(0, rows.flatten()), reshape to[q_len, k, d_latent], and run a single batchedbmm+ softmax. Keep the current loop only for ragged rows. This is optional; the current form is correct.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/vanilla.py` around lines 818 - 845, Optimize the per-token processing in the attention path around _load_mla_latent_cache by detecting when each topk_indices row has the same valid count. For uniform rows, flatten the selected indices, perform one latent.index_select, reshape to [q_len, k, d_latent], and replace the per-token matmuls with batched bmm and softmax; retain the existing loop for ragged rows.
926-927: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReplace the constant-name
getattrwith attribute access.
getattr(self.sparse_params, "indices_block_size")behaves exactly likeself.sparse_params.indices_block_size, including theAttributeErroron a missing field. Use direct attribute access.♻️ Proposed change
- sparse_attn_indices_block_size=getattr( - self.sparse_params, "indices_block_size"), + sparse_attn_indices_block_size=self.sparse_params. + indices_block_size,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/vanilla.py` around lines 926 - 927, In the attention configuration call, replace the getattr expression for self.sparse_params.indices_block_size with direct attribute access. Keep the existing sparse_attn_indices_block_size argument and behavior unchanged.Source: Linters/SAST tools
tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py (2)
194-198: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winSelection determinism depends on identical RNG call ordering between the two runs.
_test_sparse_attention_mlaruns once per backend and reseedstopk_generatorfrom the sametopk_seed. The two runs produce matchingtopk_indicesonly while both take exactly the same number ofrandpermdraws in the same order. The Vanilla and TRTLLM branches diverge inside the forward loop, so a future extra draw on one branch would silently desynchronize the selections and turn the comparison into a false pass.Consider generating the selections once outside the backend run and passing them in, or asserting that both runs produced identical
topk_indices.Also applies to: 215-219
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py` around lines 194 - 198, Make top-k index generation in the shared test setup independent of backend-specific forward-loop execution: generate the complete topk_indices once and reuse it for both Vanilla and TRTLLM runs, or explicitly compare the selections after each run and fail if they differ. Update the _test_sparse_attention_mla flow and its selection-generation helper so backend branches cannot silently consume different randperm sequences.
424-427: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winLower
sparse_configbefore backend dispatch.get_attention_backendacceptsSparseParams, whileDeepSeekSparseAttentionConfigis a user-facing config. Usesparse_config.to_sparse_params()for the lookup. The"dsa"registry entry already resolves"VANILLA"toVanillaAttention; keepis_vanillafor the later backend-specific setup.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py` around lines 424 - 427, Update the backend lookup in the AttentionCls initialization to pass sparse_config.to_sparse_params() to get_attention_backend, while preserving the existing is_vanilla calculation for later backend-specific setup and retaining the direct VanillaAttention selection behavior.tests/unittest/_torch/attention/model_attn_config.py (1)
93-99: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCentralize
sparse_topkon the sparse config base.SparseAttentionConfigincludes variants withoutindex_topk, so both accessors can raiseAttributeError. Define a safe base-class contract, override it for DeepSeek configs, and delegate from both callers instead of readingcfg.index_topkdirectly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/model_attn_config.py` around lines 93 - 99, Centralize the safe sparse_topk contract on the SparseAttentionConfig base class, returning None for configurations without index_topk, and override it in the DeepSeek config to expose the appropriate value. Update both model_attn_config.py lines 93-99 and backend_case.py lines 151-157 so their callers delegate to the config’s sparse_topk accessor instead of reading cfg.index_topk directly; both sites require this delegation change.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/vanilla.py`:
- Around line 132-143: Update the AttentionForwardArgs construction in
backend_case.py to pass the selection through
SparseBackendForwardArgs(topk_indices=...) via sparse_backend_args, rather than
supplying topk_indices as a direct AttentionForwardArgs field. Preserve
sparse_attn_predict()’s existing read from forward_args.sparse_backend_args.
In `@tests/unittest/_torch/attention/backend_case.py`:
- Around line 1456-1458: Add a TODO marker to the existing sparse-case exclusion
comment near the case.is_gen_only condition, explicitly tracking the follow-up
work to add CUDA-graph coverage for sparse cases while preserving the current
skip behavior.
- Around line 776-789: Update the _assert_cache_contains_new_tokens call in the
MLA cache validation path to use exact tolerances (atol=0.0 and rtol=0.0),
matching _run_mla_gen_backend. Keep the cache assertion exact while retaining
_tolerances(case, case.compute_dtype) only for the output comparison.
- Around line 100-104: Update BackendCase serialization to explicitly model_dump
sparse_attention_config instead of relying on asdict, and use
TypeAdapter(SparseAttentionConfig).validate_python(...) in from_dict to
reconstruct the discriminated union before sparse_topk is accessed. Add a
round-trip test covering serialization and deserialization of a BackendCase with
SparseAttentionConfig.
In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py`:
- Around line 877-887: Update the per-layer and final backend print messages in
the sparse MLA test to state that outputs were recorded rather than that the
test passed. Adjust the messages near the output assignment and final backend
log, while preserving the existing step, layer, and backend context.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/vanilla.py`:
- Around line 679-709: Consolidate the duplicate MLA page-materialization logic
by extending _gather_paged_mla_latent with the HND layout handling,
BAD_PAGE_INDEX filtering, and required read-offset behavior currently
implemented by _load_mla_latent_cache. Update all callers to use the unified
helper, then remove _load_mla_latent_cache so dense and sparse MLA paths share
one implementation.
- Around line 818-845: Optimize the per-token processing in the attention path
around _load_mla_latent_cache by detecting when each topk_indices row has the
same valid count. For uniform rows, flatten the selected indices, perform one
latent.index_select, reshape to [q_len, k, d_latent], and replace the per-token
matmuls with batched bmm and softmax; retain the existing loop for ragged rows.
- Around line 926-927: In the attention configuration call, replace the getattr
expression for self.sparse_params.indices_block_size with direct attribute
access. Keep the existing sparse_attn_indices_block_size argument and behavior
unchanged.
In `@tests/unittest/_torch/attention/backend_case.py`:
- Around line 316-320: Annotate the sparse_config parameter in
_build_mla_kv_cache_manager as Optional[SparseAttentionConfig] with a default of
None, adding or reusing the necessary imports while preserving the existing
function behavior.
- Around line 1218-1222: Move the non-MLA sparse validation into
_validate_sparse_case, making it raise the existing “Unsupported sparse
contract” ValueError for sparse cases where is_mla is false. Remove the
duplicated guard and raise from both run_backend and run_case, while preserving
their MLA sparse execution flow and relying on _validate_sparse_case as the
single validation point.
In `@tests/unittest/_torch/attention/model_attn_config.py`:
- Around line 93-99: Centralize the safe sparse_topk contract on the
SparseAttentionConfig base class, returning None for configurations without
index_topk, and override it in the DeepSeek config to expose the appropriate
value. Update both model_attn_config.py lines 93-99 and backend_case.py lines
151-157 so their callers delegate to the config’s sparse_topk accessor instead
of reading cfg.index_topk directly; both sites require this delegation change.
In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py`:
- Around line 194-198: Make top-k index generation in the shared test setup
independent of backend-specific forward-loop execution: generate the complete
topk_indices once and reuse it for both Vanilla and TRTLLM runs, or explicitly
compare the selections after each run and fail if they differ. Update the
_test_sparse_attention_mla flow and its selection-generation helper so backend
branches cannot silently consume different randperm sequences.
- Around line 424-427: Update the backend lookup in the AttentionCls
initialization to pass sparse_config.to_sparse_params() to
get_attention_backend, while preserving the existing is_vanilla calculation for
later backend-specific setup and retaining the direct VanillaAttention selection
behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b96c7533-b57f-4100-bd09-17ecbab8925c
📒 Files selected for processing (8)
tensorrt_llm/_torch/attention_backend/sparse/registry.pytensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytests/unittest/_torch/attention/backend_capability.pytests/unittest/_torch/attention/backend_case.pytests/unittest/_torch/attention/model_attn_config.pytests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.pytests/unittest/_torch/attention/test_attention_backends.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/unittest/_torch/attention/backend_capability.py
| # User-facing sparse config, lowered into backend params, metadata params, and | ||
| # the sparse KV-cache manager exactly as in production. The selection unit and | ||
| # top-k are derived from it (properties below); the attention family is | ||
| # ``is_mla``. | ||
| sparse_attention_config: Optional[SparseAttentionConfig] = None |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find round-trip users of BackendCase.to_dict / from_dict.
set -euo pipefail
rg -n -C 5 '\b(to_dict|from_dict)\s*\(' tests/unittest/_torch/attention/
rg -n -C 3 'class SparseAttentionConfig|SparseAttentionConfig =' tensorrt_llm/llmapi/llm_args.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1462
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- BackendCase serialization and sparse properties ---'
sed -n '1,240p' tests/unittest/_torch/attention/backend_case.py
rg -n -C 8 'class SparseAttentionConfig|SparseAttentionConfig' --glob '*.py' .
printf '%s\n' '--- All BackendCase serialization call sites ---'
rg -n -C 5 'BackendCase\.(to_dict|from_dict)|\.(to_dict|from_dict)\s*\(' --glob '*.py' tests tensorrt_llm | head -n 400
printf '%s\n' '--- Sparse-case construction and minimizer/persistence paths ---'
rg -n -C 6 'sparse_attention_config|sparse_topk|BackendCase' tests/unittest/_torch/attention --glob '*.py' | head -n 600Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- BackendCase implementation ---'
sed -n '60,225p' tests/unittest/_torch/attention/backend_case.py
printf '%s\n' '--- SparseAttentionConfig declarations ---'
rg -n -C 6 '^class .*SparseAttentionConfig|^SparseAttentionConfig' tensorrt_llm/llmapi/llm_args.py
printf '%s\n' '--- Explicit BackendCase serialization references ---'
rg -n -C 4 'BackendCase\.(to_dict|from_dict)|BackendCase\.__|from_dict\(.*BackendCase|to_dict\(.*BackendCase' --glob '*.py' tests tensorrt_llm
printf '%s\n' '--- Serialization-related references in attention tests ---'
rg -n -C 4 '\b(to_dict|from_dict)\b|json\.(dump|dumps|load|loads)|pickle|yaml' tests/unittest/_torch/attention --glob '*.py'Repository: NVIDIA/TensorRT-LLM
Length of output: 14440
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import json
from dataclasses import asdict, dataclass
from typing import Optional
try:
from pydantic import BaseModel, TypeAdapter
except ImportError as exc:
print(f"pydantic unavailable: {exc}")
raise SystemExit(0)
class SparseConfig(BaseModel):
algorithm: str
index_topk: int
`@dataclass`
class Case:
sparse_attention_config: Optional[SparseConfig] = None
case = Case(sparse_attention_config=SparseConfig(algorithm="dsa", index_topk=8))
payload = asdict(case)
print("asdict field type:", type(payload["sparse_attention_config"]).__name__)
try:
json.dumps(payload)
except TypeError as exc:
print("json.dumps(asdict(case)):", type(exc).__name__, str(exc))
else:
print("json.dumps(asdict(case)): succeeded")
json_payload = {"sparse_attention_config": case.sparse_attention_config.model_dump()}
restored = Case(**json_payload)
print("dict-loaded field type:", type(restored.sparse_attention_config).__name__)
try:
print("dict-loaded index_topk:", restored.sparse_attention_config.index_topk)
except AttributeError as exc:
print("dict-loaded sparse_topk access:", type(exc).__name__, str(exc))
validated = TypeAdapter(SparseConfig).validate_python(
json_payload["sparse_attention_config"])
print("TypeAdapter validation type:", type(validated).__name__)
PY
printf '%s\n' '--- Complete sparse config type alias ---'
sed -n '3540,3585p' tensorrt_llm/llmapi/llm_args.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 1835
Preserve SparseAttentionConfig in BackendCase serialization.
SparseAttentionConfig is a discriminated Pydantic union, not a dataclass. Therefore, asdict(self) leaves it as a model, and json.dumps(case.to_dict()) fails. A dictionary loaded by from_dict() also makes sparse_topk fail. Use model_dump() and TypeAdapter(SparseAttentionConfig).validate_python(...); the type alias has no model_validate() method. Add a round-trip test.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/attention/backend_case.py` around lines 100 - 104,
Update BackendCase serialization to explicitly model_dump
sparse_attention_config instead of relying on asdict, and use
TypeAdapter(SparseAttentionConfig).validate_python(...) in from_dict to
reconstruct the discriminated union before sparse_topk is accessed. Add a
round-trip test covering serialization and deserialization of a BackendCase with
SparseAttentionConfig.
| expected_latents = _split_packed_tokens(inputs["expected_new_latent"], case.seq_lens) | ||
| cache_atol, cache_rtol = _tolerances(case, case.compute_dtype) | ||
| _assert_cache_contains_new_tokens( | ||
| mgr, | ||
| 0, | ||
| request_ids, | ||
| case.seq_lens, | ||
| case.num_cached_tokens, | ||
| expected_latents, | ||
| kv_layout=metadata.kv_layout, | ||
| cache_kind="mla", | ||
| atol=cache_atol, | ||
| rtol=cache_rtol, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The cache-content assertion uses the loose sparse tolerance.
_tolerances(case, ...) returns SPARSE_ATOL=1e-1 / SPARSE_RTOL=1e-2 for every sparse case, because of the new early return at line 1095. This assertion compares cache bytes that the backend copied verbatim from expected_new_latent; it should be exact. _run_mla_gen_backend performs the same check with the default (0.0, 0.0).
An atol of 1e-1 on a bf16 latent lets a wrong-block or wrong-offset append pass whenever the values happen to be close. Use exact tolerances here and keep the loose tolerance only for the output comparison.
🐛 Proposed fix
expected_latents = _split_packed_tokens(inputs["expected_new_latent"], case.seq_lens)
- cache_atol, cache_rtol = _tolerances(case, case.compute_dtype)
_assert_cache_contains_new_tokens(
mgr,
0,
request_ids,
case.seq_lens,
case.num_cached_tokens,
expected_latents,
kv_layout=metadata.kv_layout,
cache_kind="mla",
- atol=cache_atol,
- rtol=cache_rtol,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| expected_latents = _split_packed_tokens(inputs["expected_new_latent"], case.seq_lens) | |
| cache_atol, cache_rtol = _tolerances(case, case.compute_dtype) | |
| _assert_cache_contains_new_tokens( | |
| mgr, | |
| 0, | |
| request_ids, | |
| case.seq_lens, | |
| case.num_cached_tokens, | |
| expected_latents, | |
| kv_layout=metadata.kv_layout, | |
| cache_kind="mla", | |
| atol=cache_atol, | |
| rtol=cache_rtol, | |
| ) | |
| expected_latents = _split_packed_tokens(inputs["expected_new_latent"], case.seq_lens) | |
| _assert_cache_contains_new_tokens( | |
| mgr, | |
| 0, | |
| request_ids, | |
| case.seq_lens, | |
| case.num_cached_tokens, | |
| expected_latents, | |
| kv_layout=metadata.kv_layout, | |
| cache_kind="mla", | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/attention/backend_case.py` around lines 776 - 789,
Update the _assert_cache_contains_new_tokens call in the MLA cache validation
path to use exact tolerances (atol=0.0 and rtol=0.0), matching
_run_mla_gen_backend. Keep the cache assertion exact while retaining
_tolerances(case, case.compute_dtype) only for the output comparison.
| # eager golden. The sparse runner rebuilds each request's logical cache on | ||
| # the host, which is not graph-capturable, so sparse cases are skipped. | ||
| if case.is_gen_only and not case.is_sparse: |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the requested TODO for sparse CUDA-graph coverage.
A previous review asked for a TODO here so the follow-up PR that adds CUDA-graph coverage for sparse cases can be tracked. The current comment explains the exclusion but carries no marker.
Do you want me to open an issue to track sparse CUDA-graph coverage?
📝 Proposed change
# A gen-only batch also exercises the captured-CUDA-graph path
# (production replays a captured decode graph); it must still match the
# eager golden. The sparse runner rebuilds each request's logical cache on
# the host, which is not graph-capturable, so sparse cases are skipped.
+ # TODO: add CUDA-graph coverage for sparse cases in a follow-up PR.
if case.is_gen_only and not case.is_sparse:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/attention/backend_case.py` around lines 1456 - 1458,
Add a TODO marker to the existing sparse-case exclusion comment near the
case.is_gen_only condition, explicitly tracking the follow-up work to add
CUDA-graph coverage for sparse cases while preserving the current skip behavior.
| print( | ||
| f"{backend_name} output mean: {result.abs().mean().item()}, max: {result.abs().max().item()}" | ||
| ) | ||
| print( | ||
| f"Reference output mean: {ref_result.abs().mean().item()}, max: {ref_result.abs().max().item()}" | ||
| ) | ||
| print( | ||
| f"Difference mean: {(result - ref_result).abs().mean().item()}, \ | ||
| max: {(result - ref_result).abs().max().item()}" | ||
| ) | ||
|
|
||
| # Assert results are close | ||
| atol, rtol = accuracy_dict[kv_cache_dtype] | ||
| assert torch.allclose(result, ref_result, atol=atol, rtol=rtol), ( | ||
| f"Results for sparse MLA in {backend_name} backend don't match reference implementation \ | ||
| at layer {layer_idx} in step {step}" | ||
| ) | ||
|
|
||
| print( | ||
| f"Test for sparse MLA in {backend_name} backend passed at layer {layer_idx} in step {step}" | ||
| ) | ||
| print(f"---- step {step} layer {layer_idx} end ----") | ||
| phase = "context" if step == 0 else f"generation_{step - 1}" | ||
| outputs[f"{phase}_layer_{layer_idx}"] = result.detach().clone() | ||
|
|
||
| print(f"Test for sparse MLA in {backend_name} backend passed") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
The "passed" messages no longer describe what happened.
The per-layer and per-backend "passed" prints were accurate when this function asserted internally. The function now only records outputs; _assert_matches_vanilla performs the comparison afterwards. A reader of the log sees "passed" for a run that has not been checked, and still sees it on a run that later fails.
Reword these to state that the output was recorded.
📝 Proposed change
- print(
- f"Test for sparse MLA in {backend_name} backend passed at layer {layer_idx} in step {step}"
- )
+ print(
+ f"Recorded sparse MLA output for {backend_name} at layer {layer_idx} in step {step}"
+ )
print(f"---- step {step} layer {layer_idx} end ----")
phase = "context" if step == 0 else f"generation_{step - 1}"
outputs[f"{phase}_layer_{layer_idx}"] = result.detach().clone()
- print(f"Test for sparse MLA in {backend_name} backend passed")
+ print(f"Collected all sparse MLA outputs for the {backend_name} backend")
return outputs📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| print( | |
| f"{backend_name} output mean: {result.abs().mean().item()}, max: {result.abs().max().item()}" | |
| ) | |
| print( | |
| f"Reference output mean: {ref_result.abs().mean().item()}, max: {ref_result.abs().max().item()}" | |
| ) | |
| print( | |
| f"Difference mean: {(result - ref_result).abs().mean().item()}, \ | |
| max: {(result - ref_result).abs().max().item()}" | |
| ) | |
| # Assert results are close | |
| atol, rtol = accuracy_dict[kv_cache_dtype] | |
| assert torch.allclose(result, ref_result, atol=atol, rtol=rtol), ( | |
| f"Results for sparse MLA in {backend_name} backend don't match reference implementation \ | |
| at layer {layer_idx} in step {step}" | |
| ) | |
| print( | |
| f"Test for sparse MLA in {backend_name} backend passed at layer {layer_idx} in step {step}" | |
| ) | |
| print(f"---- step {step} layer {layer_idx} end ----") | |
| phase = "context" if step == 0 else f"generation_{step - 1}" | |
| outputs[f"{phase}_layer_{layer_idx}"] = result.detach().clone() | |
| print(f"Test for sparse MLA in {backend_name} backend passed") | |
| print( | |
| f"{backend_name} output mean: {result.abs().mean().item()}, max: {result.abs().max().item()}" | |
| ) | |
| print( | |
| f"Recorded sparse MLA output for {backend_name} at layer {layer_idx} in step {step}" | |
| ) | |
| print(f"---- step {step} layer {layer_idx} end ----") | |
| phase = "context" if step == 0 else f"generation_{step - 1}" | |
| outputs[f"{phase}_layer_{layer_idx}"] = result.detach().clone() | |
| print(f"Collected all sparse MLA outputs for the {backend_name} backend") | |
| return outputs |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py` around
lines 877 - 887, Update the per-layer and final backend print messages in the
sparse MLA test to state that outputs were recorded rather than that the test
passed. Adjust the messages near the output assignment and final backend log,
while preserving the existing step, layer, and backend context.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py (2)
377-399: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd complete annotations to the changed helper functions.
tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py#L377-L399: Add precise parameter annotations to_run_test_for_backend.tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py#L569-L585: Add a precise return annotation tocreate_layer.As per coding guidelines, “Annotate every function” and avoid unnecessary
Any.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py` around lines 377 - 399, Add precise type annotations for every parameter of _run_test_for_backend in tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py lines 377-399, avoiding unnecessary Any, and add the precise return annotation to create_layer at lines 569-585. No direct changes are needed elsewhere.Source: Coding guidelines
407-409: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUse snake_case for the local backend class variable.
AttentionClsis a local variable. Rename it toattention_clsand update its uses at Lines 575, 642, and 669.As per coding guidelines, “Use snake_case for files, functions, methods, locals, and mutable globals.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py` around lines 407 - 409, Rename the local variable AttentionCls to attention_cls in the backend selection code, and update all references at the indicated uses so the same class remains instantiated or accessed consistently.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py`:
- Around line 377-399: Add precise type annotations for every parameter of
_run_test_for_backend in
tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py lines 377-399,
avoiding unnecessary Any, and add the precise return annotation to create_layer
at lines 569-585. No direct changes are needed elsewhere.
- Around line 407-409: Rename the local variable AttentionCls to attention_cls
in the backend selection code, and update all references at the indicated uses
so the same class remains instantiated or accessed consistently.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0a84bd2b-e66b-4fcd-9630-e33b64f5b2f8
📒 Files selected for processing (1)
tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py
Merge the q_pe/k_pe rotate helpers into one _rope_qk_for_vanilla over a shared _apply_yarn_rope core, unify the context vanilla/non-vanilla forward() into a single call, and trim redundant comments. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
ca9744a to
4f7ba69
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py (1)
116-117: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse built-in generic annotations.
Replace
List[int]withlist[int]for both parameters.Proposed change
- seq_lens: List[int], - past_lens: List[int], + seq_lens: list[int], + past_lens: list[int],As per coding guidelines, “prefer built-in generic types.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py` around lines 116 - 117, Update the affected function signature by replacing both seq_lens and past_lens annotations from List[int] to the built-in list[int] syntax, preserving their parameter types and behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py`:
- Around line 116-117: Update the affected function signature by replacing both
seq_lens and past_lens annotations from List[int] to the built-in list[int]
syntax, preserving their parameter types and behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6e4e414a-7e8e-4cb2-acae-111510089533
📒 Files selected for processing (1)
tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py
3 participants
Dev Engineer Review
VanillaAttention.VanillaAttentionfor thedsasparse algorithm.SparseAttentionConfig.deepseekv3_2_dsa_mla.QA Engineer Review
Modified test code
tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py.VanillaAttentiongolden-output comparison.tests/unittest/_torch/attention/test_attention_backends.py.backend_case.py,backend_capability.py, andmodel_attn_config.py.Test-list coverage
tests/integration/test_lists/,test-db/, orqa/are reported.Verdict: needs follow-up.
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.