[None][test] Cover KVCacheManagerV2 C++ pool rebalance path

cpp/tests/unit_tests/batch_manager/CMakeLists.txt

@@ -37,6 +37,11 @@ target_include_directories(
   kvCacheManagerV2StatsTest
   PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager
           ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256)
+add_gtest(kvCacheManagerV2SlotAllocatorTest
+          kvCacheManagerV2SlotAllocatorTest.cpp)
+target_include_directories(
+  kvCacheManagerV2SlotAllocatorTest
+  PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager)
 add_gtest(kvCacheManagerV2DigestPoolTest kvCacheManagerV2DigestPoolTest.cpp)
 target_include_directories(
   kvCacheManagerV2DigestPoolTest

cpp/tests/unit_tests/batch_manager/kvCacheManagerV2SlotAllocatorTest.cpp

@@ -0,0 +1,103 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h"
+
+#include <gtest/gtest.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+namespace
+{
+
+using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2;
+
+std::vector<Slot> allocateSlots(SlotAllocator& allocator, std::size_t count)
+{
+    std::vector<Slot> slots;
+    slots.reserve(count);
+    for (std::size_t i = 0; i < count; ++i)
+    {
+        slots.push_back(allocator.allocate());
+    }
+    return slots;
+}
+
+// Release slots[first, last). Slot's implicit move leaves the source's slot id
+// intact (std::optional<SlotId> is trivially copyable), so hand ownership over
+// with setSlot(), which resets the source and makes double release impossible.
+void releaseSlots(
+    SlotAllocator& allocator, std::vector<Slot>& slots, std::size_t first = 0, std::size_t last = SIZE_MAX)
+{
+    for (std::size_t i = first; i < std::min(last, slots.size()); ++i)
+    {
+        Slot slot;
+        slot.setSlot(slots[i]);
+        allocator.release(std::move(slot));
+    }
+}
+
+// Pool rebalance shrinks a pool group whose new size is still above the slot-ID
+// high-water mark. Regression for NVBug 6225866: finishShrink() used to compute
+// the expected overflow count as numActiveSlots - targetCapacity, which goes
+// negative here, so the count never matched and the shrink threw instead of
+// completing. Mirrors TestSlotAllocatorShrink.test_shrink_underused_pool in
+// tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py, which
+// only covers the Python SlotAllocator.
+TEST(KvCacheManagerV2SlotAllocatorTest, ShrinkUnderusedPool)
+{
+    SlotAllocator allocator{SlotCount{184064}};
+    auto slots = allocateSlots(allocator, 2048);
+    releaseSlots(allocator, slots);
+    EXPECT_EQ(allocator.numActiveSlots(), 2048);
+
+    allocator.prepareForShrink(SlotCount{122624});
+    EXPECT_EQ(allocator.numOverflowSlots(), 0);
+
+    EXPECT_TRUE(allocator.finishShrink());
+    EXPECT_EQ(allocator.numSlots(), 122624);
+    EXPECT_EQ(allocator.numActiveSlots(), 2048);
+    EXPECT_FALSE(allocator.shrinkInProgress());
+}
+
+// The non-trivial migration path: every ID is issued, half are released, and the
+// pool shrinks to half. The released overflow-range slots must be reclaimed by
+// finishShrink(). Mirrors test_shrink_touched_pool.
+TEST(KvCacheManagerV2SlotAllocatorTest, ShrinkTouchedPool)
+{
+    SlotAllocator allocator{SlotCount{16}};
+    auto slots = allocateSlots(allocator, 16);
+    releaseSlots(allocator, slots, 8);
+    EXPECT_EQ(allocator.numActiveSlots(), 16);
+
+    allocator.prepareForShrink(SlotCount{8});
+    EXPECT_EQ(allocator.numOverflowSlots(), 8);
+
+    EXPECT_TRUE(allocator.finishShrink());
+    EXPECT_EQ(allocator.numSlots(), 8);
+    EXPECT_EQ(allocator.numActiveSlots(), 8);
+    EXPECT_FALSE(allocator.shrinkInProgress());
+
+    // Leave the allocator quiescent so the debug-build destructor checks pass.
+    releaseSlots(allocator, slots, 0, 8);
+}
+
+} // namespace

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

@@ -4322,6 +4322,146 @@ def test_page_coverage_only_grows(self) -> None:
         )
 
 
+class TestPoolRebalance(TestKVCacheManagerV2):
+    """Drive the auto-tuner's pool rebalance end to end.
+
+    TestSlotAllocatorShrink below pokes the Python SlotAllocator directly, so it
+    never reaches the backend selected by TLLM_KV_CACHE_MANAGER_V2_BACKEND. This
+    class goes through the manager instead, covering
+    need_adjustment -> adjust() -> adjust_cache_level -> shrink/expand_pool_group
+    on whichever backend is active (C++ by default).
+    """
+
+    _TOKENS_PER_BLOCK = 32
+    _PROMPT_LEN = 64
+    _DECODE_LEN = 96
+
+    def prepare_two_pool_groups(self, gpu_quota: int = 256 << 20) -> None:
+        """Two attention life cycles with different slot sizes -> two pool groups.
+
+        Pool groups are formed per distinct slot layout, so differing window
+        sizes alone are not enough -- the buffer sizes have to differ too.
+        """
+        self.cfg = KVCacheManagerConfig(
+            tokens_per_block=self._TOKENS_PER_BLOCK,
+            cache_tiers=[GpuCacheTierConfig(quota=gpu_quota)],
+            layers=[
+                AttentionLayerConfig(
+                    layer_id=LayerId(0),
+                    buffers=[
+                        BufferConfig(role=Role.KEY, size=8192),
+                        BufferConfig(role=Role.VALUE, size=8192),
+                    ],
+                    sliding_window_size=128,
+                    num_sink_tokens=0,
+                ),
+                AttentionLayerConfig(
+                    layer_id=LayerId(1),
+                    buffers=[
+                        BufferConfig(role=Role.KEY, size=2048),
+                        BufferConfig(role=Role.VALUE, size=2048),
+                    ],
+                    sliding_window_size=None,
+                ),
+            ],
+            typical_step=BatchDesc(kv_caches=[KVCacheDesc(capacity=160, history_length=0)]),
+            constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=160, history_length=0)])],
+        )
+        self.engine = FakeEngine(self.cfg)
+        self.manager = KVCacheManager(self.cfg)
+
+    def _gpu_ratios(self) -> list[float]:
+        return list(_introspection.current_gpu_ratio(self.manager))
+
+    def _run_sequence(
+        self, prompt: list[TokenIdExt] | None = None, expect_reuse: bool = False
+    ) -> list[TokenIdExt]:
+        """Prefill + decode one sequence with reference checking; return its prompt.
+
+        Passing a previously used prompt exercises block reuse, so the reference
+        check reads back KV from blocks committed before the pool resize.
+        """
+        if prompt is None:
+            prompt = [self.next_token() for _ in range(self._PROMPT_LEN)]
+        kv_cache = self.manager.create_kv_cache(ReuseScope(), prompt)
+        with TemporaryCudaStream([]) as s:
+            stream = cast(CudaStream, s.handle)
+            self.assertTrue(kv_cache.resume(stream))
+            num_reused = kv_cache.num_committed_tokens
+            if expect_reuse:
+                self.assertGreater(
+                    num_reused, 0, "replay reused no blocks; the KV check would be vacuous"
+                )
+            self.assertTrue(kv_cache.resize(round_up(len(prompt), self._TOKENS_PER_BLOCK)))
+            capacity = kv_cache.capacity
+            history = prompt[:num_reused]
+            new_tokens = prompt[num_reused:]
+            self.engine.execute([Step(kv_cache, new_tokens, history)], stream)
+            if new_tokens:
+                kv_cache.commit(new_tokens)
+                history.extend(new_tokens)
+            for _ in range(self._DECODE_LEN):
+                if len(history) + 1 > capacity:
+                    kv_cache.commit(history[kv_cache.history_length :])
+                    self.assertTrue(
+                        kv_cache.resize(round_up(len(history) + 1, self._TOKENS_PER_BLOCK))
+                    )
+                    capacity = kv_cache.capacity
+                token = self.next_token()
+                self.engine.execute([Step(kv_cache, [token], history)], stream)
+                history.append(token)
+            kv_cache.commit(history[kv_cache.history_length :])
+            self.engine.execute([Step(kv_cache, [], history)], stream)
+        s.take_finish_event().synchronize()
+        kv_cache.close()
+        return prompt
+
+    def _slot_totals(self) -> list[int]:
+        return [s.total for s in _introspection.storage_statistics(self.manager, GPU_LEVEL)]
+
+    def test_adjust_resizes_pool_groups(self) -> None:
+        self.prepare_two_pool_groups()
+        before = self._gpu_ratios()
+        self.assertEqual(len(before), 2, f"expected two pool groups, got {before}")
+
+        self._run_sequence()
+        slots_before = self._slot_totals()
+
+        # Bypass the sample-count / cooldown gates and skew the target ratio so
+        # pool group 0 must grow and pool group 1 must shrink.
+        _introspection.force_rebalance_precondition(self.manager, skew=2.0)
+        self.assertTrue(self.manager.need_adjustment)
+
+        self.manager.adjust()
+
+        after = self._gpu_ratios()
+        self.assertEqual(len(after), len(before))
+        self.assertGreater(
+            after[0] / after[1],
+            before[0] / before[1],
+            f"adjust() did not skew the pool ratios: {before} -> {after}",
+        )
+        # Guard against a no-op adjust(): the pools must really have been
+        # resized, which means both the expand and the shrink path ran.
+        slots_after = self._slot_totals()
+        self.assertGreater(slots_after[0], slots_before[0], f"{slots_before} -> {slots_after}")
+        self.assertLess(slots_after[1], slots_before[1], f"{slots_before} -> {slots_after}")
+
+    def test_kv_survives_adjust(self) -> None:
+        """Committed blocks must still verify after pages migrate between slots."""
+        self.prepare_two_pool_groups()
+        self.assertEqual(len(self._gpu_ratios()), 2)
+
+        prompt = self._run_sequence()
+
+        _introspection.force_rebalance_precondition(self.manager, skew=2.0)
+        self.manager.adjust()
+
+        # Replay the same prompt: it reuses the committed blocks, and the fake
+        # engine's reference check reads back the KV those blocks point at.
+        self._run_sequence(prompt=prompt, expect_reuse=True)
+
+
 class TestSlotAllocatorShrink(unittest.TestCase):
     def test_shrink_underused_pool(self) -> None:
         # Regression for NVBug 6225866: shrinking a pool whose new size is