From 3c6635ea1ccce5778dad091b39a191934a2a59b1 Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Fri, 7 Aug 2026 00:41:48 +0000 Subject: [PATCH] [None][test] Cover KVCacheManagerV2 C++ pool rebalance path 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> --- .../unit_tests/batch_manager/CMakeLists.txt | 5 + .../kvCacheManagerV2SlotAllocatorTest.cpp | 103 +++++++++++++ .../test_kv_cache_manager_v2.py | 140 ++++++++++++++++++ 3 files changed, 248 insertions(+) create mode 100644 cpp/tests/unit_tests/batch_manager/kvCacheManagerV2SlotAllocatorTest.cpp diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 02f214986712..a117a8ceee3a 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(kvCacheUtilsTest kvCacheUtilsTest.cpp) add_gtest(llmRequestTest llmRequestTest.cpp) add_gtest(microBatchSchedulerTest microBatchSchedulerTest.cpp) 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 48c10a006fc0..abd8875411f6 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 @@ -4265,6 +4265,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