[None][test] Cover KVCacheManagerV2 C++ pool rebalance path by thorjohnsen · Pull Request #17387 · NVIDIA/TensorRT-LLM

thorjohnsen · 2026-08-07T00:45:31Z

Description

The C++ KVCacheManagerV2 backend is now the default (TLLM_KV_CACHE_MANAGER_V2_BACKEND, default cpp), but the pool-rebalance path had no test that actually reached it:

TestSlotAllocatorShrink — the NVBug 6225866 regression test — imports SlotAllocator from kv_cache_manager_v2._storage._core, i.e. the Python class, so it never exercised the C++ implementation.
The C++ V2 gtests covered only HostMem / Stats / TypedIndex.
tests/unittest/_torch/executor/test_kv_pool_rebalance.py mocks the PyExecutor hook and constructs no real manager.

Net effect: the NVBug 6225866 fix was shipping unguarded in the default code path. This PR adds test coverage only — no production code is modified.

1. `kvCacheManagerV2SlotAllocatorTest` (new gtest)

Ports both cases of the Python TestSlotAllocatorShrink to the C++ SlotAllocator.

ShrinkUnderusedPool is the regression case. Without the max(0, numActiveSlots - targetCapacity) guard in SlotAllocator::finishShrink, the expected-overflow count goes negative, never matches mOverflowSlots.size(), and finishShrink() throws "cannot finish shrink yet".
ShrinkTouchedPool covers the ordinary migration path.

Neither needs a CUDA context, so this runs in the normal C++ unit-test stage.

2. `TestPoolRebalance` (new Python test class)

Backend-agnostic coverage of need_adjustment → adjust() → adjust_cache_level → shrink/expand_pool_group, driven through the manager so it runs on whichever backend is selected, and requiring no model weights.

The assertions are what keep it non-vacuous: it checks that pool group 0's slot total grows while group 1's shrinks, proving both the expand and the shrink path ran, and that KV committed before the resize still verifies through block reuse afterwards.

Note for reviewers: pool groups are formed per distinct slot layout, not per sliding-window size, so the test config uses two attention layers with different buffer sizes to obtain two pool groups. The even/odd SWA split used elsewhere in this suite yields only one pool group and would make the test a silent no-op.

Test Coverage

New tests:

cpp/tests/unit_tests/batch_manager/kvCacheManagerV2SlotAllocatorTest.cpp — ShrinkUnderusedPool, ShrinkTouchedPool
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py::TestPoolRebalance — test_adjust_resizes_pool_groups, test_kv_survives_adjust

Validated on 8×H100 (sm90):

Test	Backend	Result
`kvCacheManagerV2SlotAllocatorTest` (2)	C++	PASS
`TestPoolRebalance` (2)	C++ (default)	PASS
`TestPoolRebalance` (2)	Python	PASS (reference)
Full `kv_cache_manager_v2_tests` (127)	C++	PASS — `OK (skipped=13)`
Full `kv_cache_manager_v2_tests` (127)	Python	PASS — `OK (skipped=12)`
`tests/unittest/_torch/executor/test_kv_pool_rebalance.py` (15)	mock	PASS
`accuracy/test_kv_pool_rebalance_accuracy.py` (2)	C++	PASS

The 13-vs-12 skip differential is accounted for: test_planned_drop_handle_rejects_partial_coverage carries @requires_python_backend because forcing the state under test needs a direct write into a pure-Python page object. It is unrelated to rebalancing.

Mutation-tested. A passing test proves nothing unless it fails when the bug returns, so the fix was temporarily reverted in storage/core.cpp:

// fixed
SlotCount const expectedOverflow = std::max(SlotCount{0}, mNumActiveSlots - mTargetCapacity);
// reverted to reproduce NVBug 6225866
SlotCount const expectedOverflow = (mNumActiveSlots - mTargetCapacity);

Result — exactly the predicted failure, and only the regression case failed:

[  FAILED  ] KvCacheManagerV2SlotAllocatorTest.ShrinkUnderusedPool
  C++ exception with description "SlotAllocator::finishShrink: cannot finish shrink yet"
[       OK ] KvCacheManagerV2SlotAllocatorTest.ShrinkTouchedPool

The revert was undone and rebuilt; both gtests pass again.

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.

🤖 Generated with Claude Code

Dev Engineer Review

Added the kvCacheManagerV2SlotAllocatorTest CMake target.
Added C++ regression tests for underused-pool shrinking and touched-pool migration.
Added Python coverage for pool adjustment, expansion, shrinking, and KV block reuse.
No production code or public APIs changed.
No configuration or test-list changes were identified.
Reported H100 validation covered the new tests and related suites.

QA Engineer Review

Added tests:

ShrinkUnderusedPool
ShrinkTouchedPool
TestPoolRebalance.test_adjust_resizes_pool_groups
TestPoolRebalance.test_kv_survives_adjust

Supporting methods added:

prepare_two_pool_groups
_run_sequence

No matching entries were found in tests/integration/test_lists/, test-db/, or qa/. The added tests are not explicitly listed for CI or manual QA coverage.

Verdict: needs follow-up.

The C++ KVCacheManagerV2 backend is now the default (TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp), but the rebalance path had no test that reached it: - TestSlotAllocatorShrink, the NVBug 6225866 regression test, imports SlotAllocator from kv_cache_manager_v2._storage._core, i.e. the Python class, so it never exercised the C++ implementation. - The C++ V2 gtests covered only HostMem / Stats / TypedIndex. - tests/unittest/_torch/executor/test_kv_pool_rebalance.py mocks the PyExecutor hook and constructs no real manager. So the NVBug 6225866 fix shipped unguarded in the default code path. Add two pieces of coverage: 1. kvCacheManagerV2SlotAllocatorTest - ports both cases of the Python TestSlotAllocatorShrink to the C++ SlotAllocator. ShrinkUnderusedPool is the regression case: without the max(0, numActiveSlots - targetCapacity) guard in SlotAllocator::finishShrink the expected-overflow count goes negative and finishShrink throws. ShrinkTouchedPool covers the ordinary migration path. Neither needs a CUDA context. 2. TestPoolRebalance - backend-agnostic coverage of need_adjustment -> adjust() -> adjust_cache_level -> shrink/expand_pool_group, driven through the manager so it runs on whichever backend is selected, and needing no model weights. The tests assert that pool group 0's slot total grows while group 1's shrinks, which is what proves both the expand and the shrink path ran, and that KV committed before the resize still verifies through block reuse afterwards. Note that pool groups are formed per distinct slot layout, not per sliding window size, so the test config uses two attention layers with different buffer sizes to obtain two pool groups. The new gtest was mutation-tested: reverting the max(0, ...) guard makes ShrinkUnderusedPool fail with "cannot finish shrink yet" while ShrinkTouchedPool stays green. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

coderabbitai · 2026-08-07T00:48:52Z

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: c84c5d30-5359-4dc6-b4cf-e16c877719c4

📥 Commits

Reviewing files that changed from the base of the PR and between 925ed55 and 0ee57f8.

📒 Files selected for processing (2)

cpp/tests/unit_tests/batch_manager/CMakeLists.txt
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

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

cpp/tests/unit_tests/batch_manager/CMakeLists.txt

Walkthrough

Added C++ unit tests for SlotAllocator pool shrinking and Python end-to-end tests for GPU pool rebalancing. The tests verify capacity changes, overflow-slot handling, active allocations, and KV data reuse after adjustment.

Changes

KV cache pool rebalancing tests

Layer / File(s)	Summary
Slot allocator shrink coverage `cpp/tests/unit_tests/batch_manager/CMakeLists.txt`, `cpp/tests/unit_tests/batch_manager/kvCacheManagerV2SlotAllocatorTest.cpp`	Added the `kvCacheManagerV2SlotAllocatorTest` target, allocation helpers, and tests for shrinking underused and fully touched pools.
Pool rebalance test setup `tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`	Added two GPU pool group configuration, manager initialization, workload execution with reference checks, and pool slot inspection.
Rebalance and KV persistence scenarios `tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`	Added tests for pool ratio and capacity changes and for KV reuse after pool adjustment and page migration.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#16060: Implements KV cache manager V2 pool behavior exercised by these slot allocator and pool rebalance tests.
NVIDIA/TensorRT-LLM#17308: Modifies KVCacheManagerV2 slot and pool lifecycle behavior in the same batch-manager test infrastructure.

Suggested labels: ci: full pre-merge approved

Suggested reviewers: lowsfer, nvpohanh, allisonlim-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows the required format and clearly identifies the added C++ pool rebalance test coverage.
Description check	✅ Passed	The description explains the motivation, implementation, test coverage, validation results, and checklist status.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.}

thorjohnsen · 2026-08-07T00:49:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-07T00:55:19Z

PR_Github #64430 [ run ] triggered by Bot. Commit: 65af434 Link to invocation

tensorrt-cicd · 2026-08-07T05:03:43Z

PR_Github #64430 [ run ] completed with state FAILURE. Commit: 65af434
/LLM/main/L0_MergeRequest_PR pipeline #52310 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

thorjohnsen · 2026-08-07T13:46:13Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-07T13:51:57Z

PR_Github #64608 [ run ] triggered by Bot. Commit: 65af434 Link to invocation

tensorrt-cicd · 2026-08-07T16:08:40Z

PR_Github #64608 [ run ] completed with state SUCCESS. Commit: 65af434
/LLM/main/L0_MergeRequest_PR pipeline #52472 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

thorjohnsen · 2026-08-07T16:24:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-07T16:30:42Z

PR_Github #64648 [ run ] triggered by Bot. Commit: 65af434 Link to invocation

tensorrt-cicd · 2026-08-07T18:06:33Z

PR_Github #64648 [ run ] completed with state SUCCESS. Commit: 65af434
/LLM/main/L0_MergeRequest_PR pipeline #52507 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

thorjohnsen · 2026-08-10T14:21:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-10T14:28:19Z

PR_Github #65042 [ run ] triggered by Bot. Commit: 65af434 Link to invocation

thorjohnsen · 2026-08-10T15:14:31Z

/bot kill

coderabbitai · 2026-08-10T15:14:52Z

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.

tensorrt-cicd · 2026-08-10T15:21:12Z

PR_Github #65059 [ kill ] triggered by Bot. Commit: 925ed55 Link to invocation

tensorrt-cicd · 2026-08-10T15:23:25Z

PR_Github #65042 [ run ] completed with state ABORTED. Commit: 65af434

Link to invocation

tensorrt-cicd · 2026-08-10T15:23:35Z

PR_Github #65059 [ kill ] completed with state SUCCESS. Commit: 925ed55
Successfully killed previous jobs for commit 925ed55

Link to invocation

thorjohnsen · 2026-08-10T15:24:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-10T15:29:30Z

PR_Github #65061 [ run ] triggered by Bot. Commit: bee6ee9 Link to invocation

tensorrt-cicd · 2026-08-10T19:16:05Z

PR_Github #65061 [ run ] completed with state SUCCESS. Commit: bee6ee9
/LLM/main/L0_MergeRequest_PR pipeline #52868 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

thorjohnsen · 2026-08-10T21:24:23Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-10T21:30:58Z

PR_Github #65126 [ run ] triggered by Bot. Commit: bee6ee9 Link to invocation

tensorrt-cicd · 2026-08-10T23:42:39Z

PR_Github #65126 [ run ] completed with state SUCCESS. Commit: bee6ee9
/LLM/main/L0_MergeRequest_PR pipeline #52925 completed with status: 'SUCCESS'

CI Report

Link to invocation

thorjohnsen · 2026-08-11T01:42:03Z

@lowsfer @nvpohanh @lori-ren @allisonlim-nv — could you take a look when you get a chance? Need one approval from trt-llm-kv-cache-manager-devs (lowsfer/nvpohanh) and one from trt-llm-runtime-devs (lori-ren/allisonlim-nv).

Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

thorjohnsen · 2026-08-11T16:48:32Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-11T16:54:55Z

PR_Github #65342 [ ] completed with state FAILURE. Commit: ``

Link to invocation

SimengLiu-nv

LGTM.

thorjohnsen · 2026-08-11T17:12:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-11T17:17:45Z

PR_Github #65347 [ run ] triggered by Bot. Commit: 0ee57f8 Link to invocation

tensorrt-cicd · 2026-08-11T21:58:16Z

PR_Github #65347 [ run ] completed with state SUCCESS. Commit: 0ee57f8
/LLM/main/L0_MergeRequest_PR pipeline #53116 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

thorjohnsen · 2026-08-11T23:54:22Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-11T23:59:42Z

PR_Github #65425 [ run ] triggered by Bot. Commit: 0ee57f8 Link to invocation

tensorrt-cicd · 2026-08-12T01:37:09Z

PR_Github #65425 [ run ] completed with state FAILURE. Commit: 0ee57f8
/LLM/main/L0_MergeRequest_PR pipeline #53179 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

thorjohnsen · 2026-08-12T02:02:27Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-12T02:08:51Z

PR_Github #65458 [ run ] triggered by Bot. Commit: 0ee57f8 Link to invocation

tensorrt-cicd · 2026-08-12T03:11:47Z

PR_Github #65458 [ run ] completed with state FAILURE. Commit: 0ee57f8
/LLM/main/L0_MergeRequest_PR pipeline #53207 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

thorjohnsen · 2026-08-12T03:38:18Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-12T03:43:51Z

PR_Github #65481 [ run ] triggered by Bot. Commit: 0ee57f8 Link to invocation

tensorrt-cicd · 2026-08-12T04:24:30Z

PR_Github #65481 [ run ] completed with state SUCCESS. Commit: 0ee57f8
/LLM/main/L0_MergeRequest_PR pipeline #53230 completed with status: 'SUCCESS'

CI Report

Link to invocation

thorjohnsen requested review from a team as code owners August 7, 2026 00:45

thorjohnsen requested review from allisonlim-nv, lori-ren, lowsfer and nvpohanh August 7, 2026 00:45

github-actions Bot assigned thorjohnsen Aug 7, 2026

Merge branch 'main' into thor/kvcmv2-cpp-rebalance

65af434

coderabbitai Bot mentioned this pull request Aug 7, 2026

[None][fix] Agree the KVCacheManagerV2 rebalance trigger across TP ranks #17391

Open

1 task

coderabbitai Bot mentioned this pull request Aug 7, 2026

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

Merged

1 task

Merge branch 'main' into thor/kvcmv2-cpp-rebalance

925ed55

Merge branch 'main' into thor/kvcmv2-cpp-rebalance

bee6ee9

lori-ren approved these changes Aug 11, 2026

View reviewed changes

Merge branch 'main' into thor/kvcmv2-cpp-rebalance

0ee57f8

Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

thorjohnsen requested a review from SimengLiu-nv August 11, 2026 17:09

SimengLiu-nv approved these changes Aug 11, 2026

View reviewed changes

thorjohnsen enabled auto-merge (squash) August 12, 2026 00:15

thorjohnsen merged commit d5b1818 into NVIDIA:main Aug 12, 2026
11 checks passed

Uh oh!

Conversation

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

Uh oh!

Description

1. kvCacheManagerV2SlotAllocatorTest (new gtest)

2. TestPoolRebalance (new Python test class)

Test Coverage

PR Checklist

GitHub Bot Help

Dev Engineer Review

QA Engineer Review

Uh oh!

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

Uh oh!

Walkthrough

Changes

Uh oh!

thorjohnsen commented Aug 7, 2026

Uh oh!

tensorrt-cicd commented Aug 7, 2026

Uh oh!

tensorrt-cicd commented Aug 7, 2026

Uh oh!

thorjohnsen commented Aug 7, 2026

Uh oh!

tensorrt-cicd commented Aug 7, 2026

Uh oh!

tensorrt-cicd commented Aug 7, 2026

Uh oh!

thorjohnsen commented Aug 7, 2026

Uh oh!

tensorrt-cicd commented Aug 7, 2026

Uh oh!

tensorrt-cicd commented Aug 7, 2026

Uh oh!

thorjohnsen commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

thorjohnsen commented Aug 10, 2026

Uh oh!

coderabbitai Bot commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

thorjohnsen commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

thorjohnsen commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

thorjohnsen commented Aug 11, 2026

Uh oh!

thorjohnsen commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

SimengLiu-nv left a comment

Choose a reason for hiding this comment

Uh oh!

thorjohnsen commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

thorjohnsen commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

thorjohnsen commented Aug 7, 2026 •

edited by coderabbitai Bot

Loading

1. `kvCacheManagerV2SlotAllocatorTest` (new gtest)

2. `TestPoolRebalance` (new Python test class)

coderabbitai Bot commented Aug 7, 2026 •

edited

Loading