[None][feat] Add KV cache manager V2 support for DSA by yizhang-nv · Pull Request #16060 · NVIDIA/TensorRT-LLM

yizhang-nv · 2026-07-07T09:29:00Z

Summary by CodeRabbit

New Features
- Added DSA support for DSACacheManagerV2.
- Added native per-layer indexer K-cache management.
- Updated gather, scatter, and paged-logits paths for V2 page-index scaling.
- Added FP8 and FP4 indexer cache sizing and quota handling.
- Improved draft replay metadata recomputation and restoration.
Bug Fixes
- Improved cache-manager selection and fallback handling.
- Added stable-slot and physical-slot resolution for DSA cache mappings.
- Corrected FP4 indexer K-cache byte extraction.
Tests
- Added tests for V2 cache-manager selection, metadata acceptance, cache sizing, pool layout, page-index scaling, and draft replay metadata.
- Updated accuracy coverage to enable KV cache manager V2.
- No test-list files were modified.

Dev Engineer Review

The implementation preserves V1 behavior and adds V2 selection through use_kv_cache_manager_v2.
DSACacheManagerV2 provides role-based indexer K-cache sizing, pool layout, stable-slot resolution, and byte-per-token estimates.
FP8 and packed FP4 sizing paths are covered in runtime quota and warm-up calculations.
Draft replay now recomputes and restores DSA mappings for both host and device state.
The API changes are consistent with the new V2 manager flow.
Review should verify unsupported combinations, especially DSA with CpType.STAR, and confirm that fallback errors remain actionable.
Header and kernel documentation changes do not affect runtime behavior.

QA Engineer Review

Added test_dsa_cache_manager_selection_honors_v2_config.
Added test_dsa_metadata_accepts_kv_cache_manager_v2.
Added test_dsa_cache_manager_v2_cache_size_estimate_includes_indexer.
Added test_dsa_cache_manager_v2_indexer_pool_layout.
Added TestPrepareRestoreAttnMetadataForDraftReplay.test_draft_context_recomputes_and_restores_dsa_metadata.
Added TestPrepareRestoreAttnMetadataForDraftReplay.test_recompute_context_kv_gather_mappings_skips_generation.
Updated create_dsa_cache_manager to support V2 test setup.
Updated TestGLM52.test_nvfp4 to enable V2.
The changed test functions are not shown as covered by test-db/ or qa/ entries. 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 · 2026-07-08T09:37:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-08T09:42:28Z

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

tensorrt-cicd · 2026-07-08T15:38:22Z

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 · 2026-07-20T05:30:58Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-20T05:37:00Z

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

tensorrt-cicd · 2026-07-20T08:45:39Z

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 · 2026-07-21T08:57:04Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T09:03:08Z

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

coderabbitai · 2026-07-21T09:07:18Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ec7bcb2f-a9cf-47f1-8e4d-f933591167b5

📥 Commits

Reviewing files that changed from the base of the PR and between eeadb01 and ffcffbf.

📒 Files selected for processing (4)

cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h
cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu

🚧 Files skipped from review as they are similar to previous changes (4)

cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu
cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h
cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
cpp/tensorrt_llm/kernels/indexerKCacheGather.cu

Walkthrough

DSA now supports V1 and V2 KV cache managers. The change adds indexer cache sizing, pool-index resolution, full-KV mapping recomputation, FP4 byte indexing, and draft replay state restoration. Manager selection, compatibility checks, tests, integration configuration, and kernel documentation were updated.

Changes

DSA KV Cache Manager V2

Layer / File(s)	Summary
Cache manager contracts and V2 implementation `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`	Adds V2 cache-manager support, indexer byte accounting, stable-slot resolution, page-indexed pools, and expanded metadata validation.
KV cache manager selection and compatibility `tensorrt_llm/_torch/attention_backend/sparse/utils.py`, `tensorrt_llm/_torch/pyexecutor/_util.py`	Propagates `use_kv_cache_manager_v2` through sparse-attention manager creation and extends compatibility checks.
Indexer mappings and draft replay restoration `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`, `tensorrt_llm/_torch/speculative/interface.py`	Recomputes indexer offsets and full-KV gather mappings, and saves and restores expanded DSA state during draft replay.
V2 validation and kernel documentation `tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py`, `tests/integration/defs/accuracy/test_llm_api_pytorch.py`, `tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py`, `cpp/tensorrt_llm/kernels/*`	Adds V2 and replay tests, enables V2 in the NVFP4 integration test, and updates kernel comments, operator documentation, and copyright years.

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

Possibly related PRs

NVIDIA/TensorRT-LLM#13340: Related FP4 indexer cache and gather support.
NVIDIA/TensorRT-LLM#16558: Related DSA indexer K-cache pool mapping changes.
NVIDIA/TensorRT-LLM#17122: Related KVCacheManagerV2 page and lifecycle changes.

Suggested reviewers: qijune, emmaqiaoch, junyixu-nv

Sequence Diagram(s)

sequenceDiagram
  participant draft_kv_cache_context
  participant DSAtrtllmAttentionMetadata
  participant DSACacheManagerV2
  participant Indexer
  draft_kv_cache_context->>DSAtrtllmAttentionMetadata: swap draft cache state
  DSAtrtllmAttentionMetadata->>DSACacheManagerV2: resolve pool mappings
  DSAtrtllmAttentionMetadata->>Indexer: recompute slot and full-KV mappings
  Indexer-->>DSAtrtllmAttentionMetadata: return updated mappings
  draft_kv_cache_context->>DSAtrtllmAttentionMetadata: restore original state

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the addition of DSA KV cache manager V2 support and follows the required prefix format.
Description check	✅ Passed	The description explains the change, lists relevant test coverage, and includes the repository checklist with the final review item checked.
Docstring Coverage	✅ Passed	Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%.
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 · 2026-07-21T12:39:00Z

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 · 2026-07-21T13:53:04Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T13:59:33Z

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

tensorrt-cicd · 2026-07-21T17:57:58Z

PR_Github #60702 [ run ] completed with state FAILURE. Commit: eeadb01
/LLM/main/L0_MergeRequest_PR pipeline #48992 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 · 2026-07-21T18:20:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-21T18:26:43Z

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

tensorrt-cicd · 2026-07-21T20:29:56Z

PR_Github #60752 [ run ] completed with state FAILURE. Commit: eeadb01
/LLM/main/L0_MergeRequest_PR pipeline #49038 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 · 2026-07-22T05:36:46Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-22T05:43:27Z

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

tensorrt-cicd · 2026-07-22T07:15:24Z

PR_Github #60898 [ run ] completed with state SUCCESS. Commit: eeadb01
/LLM/main/L0_MergeRequest_PR pipeline #49166 completed with status: 'SUCCESS'

CI Report

Link to invocation

zhaoyangwang-nvidia

Approve for trt-llm-runtime-devs with few comments.

zhaoyangwang-nvidia · 2026-07-22T08:24:43Z

+        # contents in place can race those copies because CPU writes are not
+        # ordered with the CUDA stream.  The fresh host buffers stay attached
+        # to the metadata until after the draft forward has been enqueued.
+        m.host_indexer_k_cache_block_offsets = torch.empty_like(


This allocates three fresh pinned host buffers on every draft replay, i.e., every decode step. cudaHostAlloc is expensive and can implicitly synchronize, and PyTorch's caching host allocator doesn't guarantee a hit. Since the shapes are fixed, consider pre-allocating a dedicated set of draft-side replacement buffers once (double-buffering) at metadata init, and just swap references here.

zhaoyangwang-nvidia · 2026-07-22T08:26:53Z

-            kv_cache_manager_cls,
-            model_config: ModelConfig,
-            kv_cache_config: Optional[KvCacheConfig] = None):
+    def _fallback_if_unsupported_kv_cache_manager_v2(self, kv_cache_manager_cls,


consider renaming to_validate_or_fallback_kv_cache_manager_v2(or similar).

zhaoyangwang-nvidia · 2026-07-22T08:29:30Z

            if self._max_beam_width is not None and self._max_beam_width > 1:
                incompat.append("beam_width > 1")
+            sparse_attn_config = model_config.sparse_attention_config
+            if (sparse_attn_config is not None


This PR introduces two user-visible behavior changes: (a) DSA + STAR context parallelism now raises instead of running; (b) sparse-attention models no longer silently fall back from V2 to V1. Please update doc under docs/source/features/feature-combination-matrix.md or docs/source/features/kvcache.md

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

Signed-off-by: yizhang-nv <187001205+yizhang-nv@users.noreply.github.com> Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

Signed-off-by: Eric Tsai <ertsai@nvidia.com>

github-actions Bot assigned yizhang-nv Jul 7, 2026

yizhang-nv force-pushed the codex/dsa-kv-cache-v2-main branch from 8ff6b1c to 69404fe Compare July 8, 2026 09:03

yizhang-nv force-pushed the codex/dsa-kv-cache-v2-main branch from 69404fe to 767b3be Compare July 20, 2026 05:29

yizhang-nv marked this pull request as ready for review July 21, 2026 08:38

yizhang-nv requested review from a team as code owners July 21, 2026 08:38

yizhang-nv requested review from BowenFu, EmmaQiaoCh, JunyiXu-nv, QiJune, kris1025, tburt-nv, yingguo-trt, yufeiwu-nv, yunruis and zhaoyangwang-nvidia July 21, 2026 08:38

yizhang-nv force-pushed the codex/dsa-kv-cache-v2-main branch from 767b3be to eeadb01 Compare July 21, 2026 08:56

coderabbitai Bot mentioned this pull request Jul 21, 2026

[None][perf] Allocate DSA indexer k-cache only for layers that own an indexer #16558

Open

1 task

yufeiwu-nv approved these changes Jul 22, 2026

View reviewed changes

zhaoyangwang-nvidia approved these changes Jul 22, 2026

View reviewed changes

yingguo-trt approved these changes Jul 28, 2026

View reviewed changes

yizhang-nv and others added 5 commits August 2, 2026 21:13

[None][feat] add DSA KV cache manager V2 support

d40ecac

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

[None][chore] fix DSA pre-commit imports

328fbac

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

[None][fix] resolve DSA indexer cache page indices

50ef814

Signed-off-by: yizhang-nv <187001205+yizhang-nv@users.noreply.github.com> Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

[None][fix] stabilize DSA draft cache replay with V2

93b1d09

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

[None][fix] pre-allocate DSA draft indexer buffers for CUDA graph replay

ffcffbf

Signed-off-by: Eric Tsai <ertsai@nvidia.com>

yizhang-nv force-pushed the codex/dsa-kv-cache-v2-main branch from 2a417cc to ffcffbf Compare August 3, 2026 04:17

Uh oh!

Conversation

yizhang-nv commented Jul 7, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

yizhang-nv commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

yizhang-nv commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

tensorrt-cicd commented Jul 20, 2026

Uh oh!

yizhang-nv commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

coderabbitai Bot commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

yizhang-nv commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

yizhang-nv commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

tensorrt-cicd commented Jul 21, 2026

Uh oh!

yizhang-nv commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

tensorrt-cicd commented Jul 22, 2026

Uh oh!

zhaoyangwang-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

zhaoyangwang-nvidia Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

zhaoyangwang-nvidia Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

zhaoyangwang-nvidia Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

yizhang-nv commented Jul 7, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 21, 2026 •

edited

Loading