diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 420d07dd543e..152feab05c0b 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/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 diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2SlotAllocatorTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2SlotAllocatorTest.cpp new file mode 100644 index 000000000000..d5be6b6bb6a4 --- /dev/null +++ b/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 + +#include +#include +#include +#include +#include + +namespace +{ + +using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; + +std::vector allocateSlots(SlotAllocator& allocator, std::size_t count) +{ + std::vector 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 is trivially copyable), so hand ownership over +// with setSlot(), which resets the source and makes double release impossible. +void releaseSlots( + SlotAllocator& allocator, std::vector& 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 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 41d1d9bcd81e..4d0900017593 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/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