[None][feat] add commit min snapshot for SSM reuse by lowsfer · Pull Request #15752 · NVIDIA/TensorRT-LLM

lowsfer · 2026-06-30T04:08:02Z

Summary:

Remove unused HelixConfig from KVCacheManagerV2 config and stubs.
Add commit_min_snapshot mode and require it for SSM layers.
Snapshot SSM state at each block-boundary commit and remove interval-based SSM reuse.
Allow SSM partial reuse.

Tests:

python3 -m py_compile init.pyi _config.py _core/_kv_cache_manager.py _core/_kv_cache.py _block_radix_tree.py _page.py test_kv_cache_manager_v2.py
TestSSMSupport
test_branch_reuse.py
TestScratchReuse
TestNoBatching.test_shrink_capacity_0 TestNoBatching.test_shrink_capacity_1 TestNoBatching.test_reuse_scope_isolates_reuse

Summary by CodeRabbit

New Features
- Added a new snapshot-commit option for KV cache management, with updated support for scratch reuse settings.
- Improved support for SSM workflows by allowing snapshots to be committed at block boundaries.
Bug Fixes
- Tightened commit validation so invalid snapshot commits are rejected earlier.
- Improved stability when clearing cached pages and handling reuse across multiple commits.
Tests
- Expanded coverage for SSM snapshot reuse, multiple commit scenarios, and configuration validation.

coderabbitai · 2026-06-30T04:16:20Z

📝 Walkthrough

Walkthrough

Removes HelixConfig and the ssm_reuse_interval-driven SSM snapshotting mechanism. Introduces commit_min_snapshot: bool on KVCacheManagerConfig and KVCacheManager, adds SwaScratchReuseConfig, updates _KVCache commit/snapshot logic to gate SSM snapshots on block-boundary commits, adds an expected_page guard to Block.unset_page, and updates tests accordingly.

Changes

SSM commit_min_snapshot refactor

Layer / File(s)	Summary
Config: remove HelixConfig, add commit_min_snapshot `tensorrt_llm/runtime/kv_cache_manager_v2/_config.py`, `tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi`	`HelixConfig` dataclass and `helix_config` field removed; `SwaScratchReuseConfig` and `swa_scratch_reuse` added; `commit_min_snapshot: bool = False` added to `KVCacheManagerConfig`; `__post_init__` asserts `commit_min_snapshot` when SSM layers are present.
KVCacheManager property swap `tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py`, `tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi`	`ssm_reuse_interval` property removed; `commit_min_snapshot` property added delegating to `init_config`.
CommittedPage destructor and Block.unset_page guard `tensorrt_llm/runtime/kv_cache_manager_v2/_page.py`, `tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py`	`CommittedPage.__del__` passes `self` to `block.unset_page`; `Block.unset_page` gains an optional `expected_page` parameter to guard against stale unsets.
_KVCache commit and snapshot logic `tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`	`commit()` enforces block-boundary alignment under `commit_min_snapshot` and forwards per-block `commit_ssm` flag; `_snapshot_ssm_to_tree_block` refactored for `num_tokens` and cache-level slot search; `_commit_block` uses `did_commit` flag to gate SSM snapshotting; `_on_stop_committing`, `_unlock_stale_blocks`, `_check_sanity` updated for `commit_min_snapshot` blocks.
Tests `tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`	Config helpers updated; `test_ssm_reuse_interval_boundary` replaced by `test_ssm_reuse_snapshots_each_commit`; new tests added for multi-commit snapshot persistence, boundary alignment requirements, and updated config validation.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#12644: Introduced ssm_reuse_interval-driven SSM snapshotting in _KVCache; this PR removes that mechanism and replaces it with the commit_min_snapshot flag.

Suggested reviewers

yizhang-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.84% 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 follows the required [None][feat] format and clearly names the new commit_min_snapshot SSM behavior.
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.
Description check	✅ Passed	The description covers the change summary and test coverage, with only the checklist section omitted and a minor heading mismatch.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py (1)
1279-1302: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

Do not silently commit without the required SSM snapshot.

If every cache level raises OutOfPagesError, this loop falls through and _commit_block() leaves the radix-tree block committed without tree_block.storage[ssm_lc_id]. Surface the failure or roll back/mark the block non-reusable for SSM instead; otherwise later reuse can hit the snapshot invariant asserted in _setup_for_reuse().
🤖 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/runtime/kv_cache_manager_v2/_core/_kv_cache.py` around lines
1279 - 1302, The commit path in `_commit_block()` is silently falling through
when every `storage.new_slots_for_pool_group(...)` call raises
`OutOfPagesError`, leaving the radix-tree block committed without setting
`tree_block.storage[ssm_lc_id]`. Update this logic to either propagate an
explicit failure or roll back/mark the block non-reusable for SSM before
returning, so a block is never committed unless the required snapshot slot is
allocated and recorded. Use the existing `_setup_for_reuse()` snapshot invariant
and the `CommittedPage`/`tree_block.storage[ssm_lc_id]` assignment as the points
to align the fix.

🤖 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 `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 1738-1749: The test cases around kv3.commit and kv4.commit are
only asserting a generic AssertionError, so they do not verify the intended
guard paths. Update the checks in test_kv_cache_manager_v2 to use
assertRaisesRegex on the commit calls, targeting the specific history-alignment
and block-boundary assertion messages from commit_min_snapshot or the related
validation path, so the test pins the exact failure reason instead of any
AssertionError.
- Around line 1693-1694: The multi-commit snapshot test closes kv1 without
explicitly finalizing the writer lifecycle first. Update the test around
kv1.commit and kv1.close to call kv1.stop_committing() before close so it
validates the finalized block-boundary commit path instead of depending on
close() to do that work implicitly.

---

Outside diff comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`:
- Around line 1279-1302: The commit path in `_commit_block()` is silently
falling through when every `storage.new_slots_for_pool_group(...)` call raises
`OutOfPagesError`, leaving the radix-tree block committed without setting
`tree_block.storage[ssm_lc_id]`. Update this logic to either propagate an
explicit failure or roll back/mark the block non-reusable for SSM before
returning, so a block is never committed unless the required snapshot slot is
allocated and recorded. Use the existing `_setup_for_reuse()` snapshot invariant
and the `CommittedPage`/`tree_block.storage[ssm_lc_id]` assignment as the points
to align the fix.

🪄 Autofix (Beta)

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: 8bb89a57-c336-40b2-8357-07179301b521

📥 Commits

Reviewing files that changed from the base of the PR and between 92147d6 and bf09f34.

📒 Files selected for processing (7)

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi
tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
tensorrt_llm/runtime/kv_cache_manager_v2/_config.py
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

lowsfer · 2026-06-30T13:06:55Z

/bot run

tensorrt-cicd · 2026-06-30T13:12:57Z

PR_Github #56609 [ run ] triggered by Bot. Commit: 266f0db Link to invocation

tensorrt-cicd · 2026-06-30T16:08:13Z

PR_Github #56609 [ run ] completed with state SUCCESS. Commit: 266f0db
/LLM/main/L0_MergeRequest_PR pipeline #45436 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

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

Add a two-prefill-turn test for TestScratchReuse: window=64, tokens_per_block=32, SWA scratch reuse enabled, commit_min_snapshot disabled. Turn 1 commits a 127-token prompt (= 2*window - 1); with scratch reuse the only KV data preserved after the turn is the last winSize-1 = 63 tokens (blocks 2,3), while the out-of-window blocks 0,1 keep no page. Turn 2's prompt shares those 127 tokens and must still reuse the full committed prefix, since the first input token is itself in the window and only winSize-1 history tokens are required. The reused KV is validated through FakeEngine, and removing the winSize-1 term from get_stale_range makes the test fail. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

lowsfer · 2026-07-07T10:19:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T10:25:39Z

PR_Github #57993 [ run ] triggered by Bot. Commit: 31567f9 Link to invocation

tensorrt-cicd · 2026-07-07T13:37:14Z

PR_Github #57993 [ run ] completed with state SUCCESS. Commit: 31567f9
/LLM/main/L0_MergeRequest_PR pipeline #46663 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

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

lowsfer · 2026-07-10T04:57:39Z

/bot run --disable-fail-fast

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

lowsfer · 2026-07-10T05:02:36Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-10T05:03:02Z

PR_Github #58594 [ run ] triggered by Bot. Commit: 67f10ad Link to invocation

tensorrt-cicd · 2026-07-10T05:07:55Z

PR_Github #58595 [ run ] triggered by Bot. Commit: 67f10ad Link to invocation

tensorrt-cicd · 2026-07-10T05:08:10Z

PR_Github #58594 [ run ] completed with state ABORTED. Commit: 67f10ad

Link to invocation

lowsfer · 2026-07-10T14:06:18Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-10T14:13:04Z

PR_Github #58666 [ run ] triggered by Bot. Commit: 67f10ad Link to invocation

tensorrt-cicd · 2026-07-10T14:13:13Z

PR_Github #58595 [ run ] completed with state ABORTED. Commit: 67f10ad
/LLM/main/L0_MergeRequest_PR pipeline #47187 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

tensorrt-cicd · 2026-07-10T14:59:06Z

PR_Github #58666 [ run ] completed with state SUCCESS. Commit: 67f10ad
/LLM/main/L0_MergeRequest_PR pipeline #47252 completed with status: 'SUCCESS'

CI Report

Link to invocation

jiaganc

LGTM

… KVCacheManagerV2 Translate two main-branch Python changes in runtime/kv_cache_manager_v2/ into the C++ backend under cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/: Translated from: - 17190cf [feat] add commit min snapshot for SSM reuse (NVIDIA#15752) config: drop HelixConfig/ssm_reuse_interval, add commit_min_snapshot; page: add SsmCommittedPage + UncommittedPage::convertToSsmCommitted; blockRadixTree: unlinkPage(expectedPage) guard + snapshot-length-aware SSM match truncation; kvCache: commit(is_end), _commitBlock(commitSsm,moveSsm), rewritten _snapshotSsmToTreeBlock, new _copyPageToTreeBlock and _snapshotPartialBlockToTree, plus commit_min_snapshot assert relaxations; nanobind: commit(is_end) and config/manager surface updates. - 52fc8f4 [fix] Reserve worst-case SWA slots to avoid single-request deadlock (NVIDIA#15588, nvbugs/6330273) storageManager::computeMinSlotsFromConstraints SWA worst-case floor; lifeCycleRegistry getStaleRange comment. Also in this change: - Mirror the Python Block._release_pages() eager reclaim into C++ Block::releasePages()/removeSubtree so page reclamation does not depend on ~Block() destruction timing. - Add kv::AssertionError + a nanobind translator to Python AssertionError so config validation and the commit history-alignment check match the Python backend's exception type. - Add white-box reuse-tree introspection helpers (attention/swa/ssm life-cycle ids, reuse_match_pages) to the C++ _introspection submodule and _introspection.py, and route the shared tests through them instead of Python-only internals. - Adapt the SWA event-manager tests (small window; deeper load for the migration test to force both pool groups to migrate under the new slot floor). - _core/_kv_cache.py: add the same VIRTUAL_STOP loop-break/partial-snapshot guard the C++ commit() uses, keeping the two backends behaviourally aligned. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

github-actions Bot assigned lowsfer Jun 30, 2026

coderabbitai Bot reviewed Jun 30, 2026

View reviewed changes

Comment thread tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

Comment thread tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py Outdated

jiaganc reviewed Jun 30, 2026

View reviewed changes

Comment thread tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py Outdated

Comment thread tensorrt_llm/runtime/kv_cache_manager_v2/_config.py

jiaganc reviewed Jun 30, 2026

View reviewed changes

Comment thread tensorrt_llm/runtime/kv_cache_manager_v2/_config.py Outdated

jiaganc reviewed Jun 30, 2026

View reviewed changes

Comment thread tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py Outdated

lowsfer force-pushed the commit_mode branch from bf09f34 to 266f0db Compare June 30, 2026 13:06

jiaganc reviewed Jul 1, 2026

View reviewed changes

Comment thread tensorrt_llm/runtime/kv_cache_manager_v2/_config.py

jiaganc mentioned this pull request Jul 3, 2026

[None][feat] add commit-min snapshots and per-conversation KV reuse #15916

Closed

1 task

lowsfer force-pushed the commit_mode branch from 266f0db to c21f439 Compare July 7, 2026 06:49

lowsfer and others added 7 commits July 7, 2026 10:18

[None][refactor] remove unused HelixConfig

2816be8

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

[None][feat] add commit min snapshot for SSM reuse

1bdc09e

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

[None][test] cover SWA commit min snapshot reuse

79789b6

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

[None][feat] Allow partial page commit for commit_min_snapshot/SSM

af97b6e

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

[None][fix] Fix dual-import in kv cache manager v2

9916ce2

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

[None][fix] Remove stale assertions.

31567f9

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

lowsfer force-pushed the commit_mode branch from 96cd834 to 31567f9 Compare July 7, 2026 10:19

jiaganc mentioned this pull request Jul 9, 2026

[None][feat] Integrate KVCMv2 commit_min_snapshot in runtime #16185

Merged

1 task

Merge branch 'main' into commit_mode

6ef3011

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

[None][fix] Remove trailing whitespace in _config.py

67f10ad

Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

jiaganc approved these changes Jul 13, 2026

View reviewed changes

lowsfer merged commit 17190cf into NVIDIA:main Jul 13, 2026
11 checks passed

Uh oh!

Conversation

lowsfer commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

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

Uh oh!

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lowsfer commented Jun 30, 2026

Uh oh!

tensorrt-cicd commented Jun 30, 2026

Uh oh!

tensorrt-cicd commented Jun 30, 2026

Uh oh!

Uh oh!

lowsfer commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

lowsfer commented Jul 10, 2026

Uh oh!

lowsfer commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

lowsfer commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

jiaganc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lowsfer commented Jun 30, 2026 •

edited

Loading

coderabbitai Bot commented Jun 30, 2026 •

edited

Loading