[None][feat] Add KV cache manager V2 support for DSA

[yizhang-nv]

Dev Engineer Review

• Added DSA KV cache manager V2 support while preserving V1 selection.
• Registered indexer K cache storage as a native per-layer Role.INDEX_KEY buffer.
• Added V2 page-index conversion for gather, scatter, and paged-logits paths.
• Added FP8 and FP4 indexer storage accounting for cache sizing, quotas, and warm-up estimates.
• Added stable representative-layer selection for coalesced KV cache pools.
• Added draft-cache buffer handling for FlashMLA and DSA CUDA-graph replay.
• Added model defaults that enable KV cache manager V2 for DeepSeek V3.
• Updated configuration validation and documentation for unsupported V2 combinations.
• Updated copyright and license metadata in affected files.
• No test-list files were changed.
• Review should confirm V2 page scaling, FP4 byte accounting, draft-buffer restoration, and V1 fallback behavior across all supported context-parallel configurations.

QA Engineer Review

Test code changes are present.

Added or updated coverage includes:

• DSA V2 manager selection, page mapping, layer masks, cache sizing, and metadata behavior.
• DSA draft replay and cached full-KV mapping restoration.
• KV cache V2 stable-layer page-table mapping.
• DSA with STAR context parallelism validation.
• Nested model-default preservation.
• DSA accuracy coverage through TestGLM52.test_nvfp4.
• C++ custom-op documentation coverage for indexer_k_cache_scatter_op.

No corresponding tests/integration/test_lists/ changes were reported for these tests. Test-list coverage cannot be confirmed from the available changes.

Verdict: needs follow-up.

Description

Port the DSA KV cache manager V2 work from liji-nv@492422b onto the latest main.

This change:

• Adds DSACacheManagerV2 selection through use_kv_cache_manager_v2 while preserving the V1 path.
• Registers the indexer K cache as a native per-layer Role.INDEX_KEY buffer managed by V2.
• Makes indexer gather/scatter and paged-logits paths aware of V2 page-index scaling for coalesced pools.
• Accounts for FP8/FP4 indexer storage in Role.ALL, runtime quota calculation, and static KV cache warm-up estimation.
• Adds DSA V2 accuracy coverage and targeted unit tests for cache layout and scaled page indices.

Test Coverage

• Clean C++ source builds passed for SM80 (A100), SM90 (H100/H200), SM100 (B200), and SM103 (B300).
• DSA V2 manager selection, FP8/FP4 layout, role sizing, and warm-up quota: 5 passed.
• Indexer gather/scatter page-index scaling with FP8/FP4: 4 passed.
• DSA V1 scatter regression: 1 passed.
• YAPF, Ruff lint/format, and git diff --check passed.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58215 [ run ] triggered by Bot. Commit: 69404fe Link to invocation

[tensorrt-cicd]

PR_Github #58215 [ run ] completed with state FAILURE. Commit: 69404fe
/LLM/main/L0_MergeRequest_PR pipeline #46857 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60273 [ run ] triggered by Bot. Commit: 767b3be Link to invocation

[tensorrt-cicd]

PR_Github #60273 [ run ] completed with state FAILURE. Commit: 767b3be
/LLM/main/L0_MergeRequest_PR pipeline #48631 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60657 [ run ] triggered by Bot. Commit: eeadb01 Link to invocation

[coderabbitai]

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

The PR adds DSA support for V1 and V2 KV cache managers. It updates indexer cache sizing and mappings, draft replay metadata handling, stable pool mapping, manager selection, model defaults, tests, and documentation.

Changes

DSA cache manager contracts and V2 implementation

Layer / File(s)	Summary
DSA cache manager contracts and V2 implementation tensorrt_llm/_torch/attention_backend/sparse/dsa.py, tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py	DSA adds V2 cache management, indexer cache sizing, page conversion, stable-slot lookup, per-layer accounting, and manager-backed metadata buffers. Tests cover V1 and V2 behavior, page mapping, shared-indexer layers, and cache sizing.

Indexer offsets and full-KV mappings

Layer / File(s)	Summary
Indexer offsets and full-KV mappings tensorrt_llm/_torch/attention_backend/sparse/dsa.py	Indexer preparation uses manager-provided page parameters. Full context-KV gather mappings are recomputed through a shared utility. Packed FP4 gathering uses half the logical head dimension.

KV cache manager selection and compatibility

Layer / File(s)	Summary
KV cache manager selection and compatibility tensorrt_llm/_torch/attention_backend/sparse/utils.py, tensorrt_llm/_torch/pyexecutor/_util.py, tensorrt_llm/_torch/models/modeling_deepseekv3.py, tests/unittest/_torch/executor/test_mamba_cache_manager.py, tests/unittest/llmapi/test_config_database.py, docs/source/features/kvcache.md, tests/integration/defs/accuracy/test_llm_api_pytorch.py	Runtime routing forwards the V2 selection flag and rejects DSA with STAR context parallelism. DeepSeek V3 defaults to V2. Configuration tests preserve nested defaults, and documentation describes V1, V2, and automatic selection.

Stable pool representatives and mapping validation

Layer / File(s)	Summary
Stable pool representatives and mapping validation tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py, tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py	Pool mapping and page-table initialization use configuration-order representative layers. Regression tests verify behavior when internal layer grouping is reordered.

Indexer mappings and draft replay restoration

Layer / File(s)	Summary
Indexer mappings and draft replay restoration tensorrt_llm/_torch/attention_backend/trtllm.py, tensorrt_llm/_torch/speculative/interface.py, tensorrt_llm/_torch/modules/mla.py, tensorrt_llm/_torch/pyexecutor/model_engine.py, tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py	FlashMLA receives dedicated draft buffers. Draft replay swaps and restores FlashMLA and DSA metadata, including cached-KV full mappings. Sparse FlashMLA derives the active layer index from attention metadata.

Kernel documentation and supporting metadata

Layer / File(s)	Summary
Kernel documentation and supporting metadata cpp/tensorrt_llm/kernels/IndexerKCacheGather.h, cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h, cpp/tensorrt_llm/kernels/indexerKCacheGather.cu, cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu, tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py, tests/unittest/_torch/speculative/test_eagle3.py	Copyright headers and kernel parameter comments are updated. The custom-operator documentation lists the scatter operator.

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

Sequence Diagram(s)

sequenceDiagram
  participant LLMConfig
  participant get_sparse_attn_kv_cache_manager
  participant DSACacheManagerV2
  participant DSAtrtllmAttentionMetadata
  participant prepare_attn_metadata_for_draft_replay
  LLMConfig->>get_sparse_attn_kv_cache_manager: select KV cache manager V2
  get_sparse_attn_kv_cache_manager->>DSACacheManagerV2: construct DSA V2 manager
  DSACacheManagerV2->>DSAtrtllmAttentionMetadata: provide page parameters and cache buffers
  prepare_attn_metadata_for_draft_replay->>DSAtrtllmAttentionMetadata: bind draft buffers and recompute mappings
  DSAtrtllmAttentionMetadata-->>prepare_attn_metadata_for_draft_replay: restore target metadata after replay

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16598: Extends shared KVCacheManagerV2 and hybrid cache-manager infrastructure.
• NVIDIA/TensorRT-LLM#16887: Shares DSA draft-replay metadata handling and separate draft KV-cache buffers.
• NVIDIA/TensorRT-LLM#17334: Overlaps with hybrid KDA/KV-cache handling and related DSA/KV-cache-manager changes.

Suggested reviewers: bowenfu, qijune, tburt-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 59.55% 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
Title check	✅ Passed	The title clearly identifies the feature and matches the primary change: DSA KV cache manager V2 support.
Description check	✅ Passed	The description explains the change, lists relevant test coverage, and includes the required checklist with the review confirmation checked.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[tensorrt-cicd]

PR_Github #60657 [ run ] completed with state FAILURE. Commit: eeadb01
/LLM/main/L0_MergeRequest_PR pipeline #48955 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65049 [ run ] triggered by Bot. Commit: f3146d7 Link to invocation

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65112 [ run ] triggered by Bot. Commit: f3146d7 Link to invocation

[tensorrt-cicd]

PR_Github #65049 [ run ] completed with state ABORTED. Commit: f3146d7

Link to invocation

[tensorrt-cicd]

PR_Github #65112 [ run ] completed with state FAILURE. Commit: f3146d7
/LLM/main/L0_MergeRequest_PR pipeline #52910 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

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65191 [ run ] triggered by Bot. Commit: f3146d7 Link to invocation

[tensorrt-cicd]

PR_Github #65191 [ run ] completed with state SUCCESS. Commit: f3146d7
/LLM/main/L0_MergeRequest_PR pipeline #52980 completed with status: 'SUCCESS'

CI Report

Link to invocation

[QiJune]

LGTM