diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index ae136b0306f7..617dc86e233b 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -23,6 +23,7 @@ #include "kv_cache_manager_v2/utils/math.h" #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/logger.h" #include #include #include @@ -127,11 +128,25 @@ KvCacheManager::KvCacheManager(KVCacheManagerConfig const& config, std::shared_p KvCacheManager::~KvCacheManager() { - shutdown(); + try + { + shutdown(); + } + catch (std::exception const& e) + { + TLLM_LOG_ERROR("%s", e.what()); + } +} + +void KvCacheManager::_checkNoLivingKvCaches(char const* api) const +{ + TLLM_CHECK_WITH_INFO(mLivingKvCaches.empty(), + "%s with %zu KV cache(s) still open; close them (or drain the engine) first", api, mLivingKvCaches.size()); } void KvCacheManager::shutdown() { + _checkNoLivingKvCaches("shutdown()"); clearReusableBlocks(); TLLM_CHECK_DEBUG(mStorage); @@ -150,6 +165,7 @@ void KvCacheManager::shutdown() void KvCacheManager::clearReusableBlocks() { + _checkNoLivingKvCaches("clear_reusable_blocks()"); TLLM_CHECK_DEBUG(mRadixTree); mRadixTree->clear(); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h index a9737b5b282c..e51269f51284 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h @@ -298,6 +298,10 @@ class KvCacheManager : public std::enable_shared_from_this friend class KvCacheIntrospection; private: + // Throw unless every KvCache has been closed. `api` names the caller so the message + // points at the mistake rather than at whatever breaks later. + void _checkNoLivingKvCaches(char const* api) const; + void _adjustLevel(CacheLevel level, size_t quota); bool _needAdjustment(CacheLevel level) const; TypedVec const& _getTargetRatioList(CacheLevel level) const; diff --git a/tensorrt_llm/llmapi/rlhf_utils.py b/tensorrt_llm/llmapi/rlhf_utils.py index 313b1c181dbb..f6450e86181a 100644 --- a/tensorrt_llm/llmapi/rlhf_utils.py +++ b/tensorrt_llm/llmapi/rlhf_utils.py @@ -208,8 +208,14 @@ def update_weights(self, ipc_handles: Optional[dict] = None): logger.error("Encountered an error in update_weights") raise e + @control_action_decorator def reset_prefix_cache(self) -> None: - """Invalidate the KV cache prefix reuse state after weight updates.""" + """Invalidate the KV cache prefix reuse state after weight updates. + + Drains in-flight requests first, like update_weights(): clearing the reuse state + detaches the whole radix tree, and a request that is still holding blocks from it + would go on committing into the detached subtree. + """ self.engine.reset_prefix_cache() @control_action_decorator diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py index 74c73700d3e6..633c2f56c2e2 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py @@ -14,6 +14,7 @@ # limitations under the License. import time +import warnings from collections import defaultdict from collections.abc import Callable, Iterable, Sequence from copy import deepcopy @@ -37,6 +38,7 @@ TokenIdExt, ) from .._config import DataRole, KVCacheManagerConfig +from .._exceptions import LogicError from .._life_cycle_registry import LayerGroupId, LifeCycle, LifeCycleId, LifeCycleRegistry from .._page import Page, _PageHolder from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta, SsmSnapshotIterationStatsDelta @@ -292,13 +294,31 @@ def __init__( self._stats_excluded_kv_cache_ids = set() def __del__(self) -> None: - self.shutdown() + try: + self.shutdown() + except LogicError as e: + warnings.warn(str(e)) + + def _check_no_living_kv_caches(self, api: str) -> None: + """Raise unless every KV cache has been closed. + + `api` names the caller so the message points at the mistake rather than at + whatever breaks later. Entries are dropped by `_KVCache.close()`, so this counts + sequences that are still open, not merely un-collected objects. + """ + if self._living_kv_caches: + raise LogicError( + f"{api} with {len(self._living_kv_caches)} KV cache(s) still open; " + "close them (or drain the engine) first" + ) def shutdown(self) -> None: + self._check_no_living_kv_caches("shutdown()") self.clear_reusable_blocks() self._storage.destroy() def clear_reusable_blocks(self) -> None: + self._check_no_living_kv_caches("clear_reusable_blocks()") self._radix_tree.clear() def get_mem_pool_base_address( 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 e5b17c6bece3..30bd3eb2b591 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 @@ -65,7 +65,7 @@ SlidingWindowSize, ) from kv_cache_manager_v2._copy_engine import CopyTask, batched_copy - from kv_cache_manager_v2._exceptions import OutOfPagesError + from kv_cache_manager_v2._exceptions import LogicError, OutOfPagesError from kv_cache_manager_v2._storage._core import CacheLevelStorage, PoolGroupBase, SlotAllocator from kv_cache_manager_v2._storage_manager import StorageManager from kv_cache_manager_v2._utils import ( @@ -118,7 +118,7 @@ SlidingWindowSize, ) from tensorrt_llm.runtime.kv_cache_manager_v2._copy_engine import CopyTask, batched_copy - from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import OutOfPagesError + from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import LogicError, OutOfPagesError from tensorrt_llm.runtime.kv_cache_manager_v2._storage._core import ( CacheLevelStorage, PoolGroupBase, @@ -969,6 +969,121 @@ def test_naive_perf(self, interval, profile: bool) -> None: profiler.dump_stats("profiler.prof") +class TestLivingKvCacheGuard(TestKVCacheManagerV2): + """Guard against clearing/freeing the reuse state while KV caches are still open. + + `clear_reusable_blocks()` detaches the whole radix tree and `shutdown()` frees the + storage the pages live in. A request that is still open keeps committing into the + detached subtree, which silently discards work on the Python backend and segfaults on + the C++ one, so both entry points must reject the call instead. + """ + + # The two backends raise different types: the Python backend raises its own + # LogicError, while the C++ backend's TLLM_CHECK_WITH_INFO throws a TllmException + # (a std::runtime_error), which nanobind surfaces as RuntimeError. Accept either so + # this test is meaningful under both. + GuardError = (LogicError, RuntimeError) + + def _seed_reusable_prompt(self) -> list[TokenIdExt]: + """Commit and close a sequence so the radix tree actually holds reusable blocks. + + Without this the tree is empty, and a regression that cleared it *before* raising + would still pass — there would be nothing left to observe. + """ + prompt = [self.next_token() for _ in range(64)] + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + seed = self.manager.create_kv_cache() + seed.resume(stream) + seed.capacity = 32 + seed.commit(prompt[:32]) + seed.capacity = 64 + seed.commit(prompt[32:]) + seed.stop_committing() + seed.close() + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 64) + return prompt + + def _open_cache(self) -> _KVCache: + return self.manager.create_kv_cache( + ReuseScope(lora_id=None), [self.next_token() for _ in range(64)] + ) + + def test_clear_reusable_blocks_rejects_open_kv_cache(self) -> None: + self.prepare(32 << 20, 32 << 20, 1 << 30, 4, 128, 1) + prompt = self._seed_reusable_prompt() + kv_cache = self._open_cache() + try: + with self.assertRaises(self.GuardError) as ctx: + self.manager.clear_reusable_blocks() + # The message must name the API the caller actually invoked, and report how + # many sequences are still open. + self.assertIn("clear_reusable_blocks()", str(ctx.exception)) + self.assertIn("1 KV cache(s) still open", str(ctx.exception)) + # The rejected call must be a no-op: the check runs before the tree is + # touched, so every block is still reusable. + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 64) + finally: + kv_cache.close() + + # ...and once permitted it really does clear, which is what stops the assertion + # above from passing vacuously. + self.manager.clear_reusable_blocks() + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 0) + + def test_shutdown_rejects_open_kv_cache(self) -> None: + self.prepare(32 << 20, 32 << 20, 1 << 30, 4, 128, 1) + prompt = self._seed_reusable_prompt() + kv_cache = self._open_cache() + try: + with self.assertRaises(self.GuardError) as ctx: + self.manager.shutdown() + self.assertIn("shutdown()", str(ctx.exception)) + self.assertIn("1 KV cache(s) still open", str(ctx.exception)) + # shutdown() frees the storage the pages live in, so a rejected call must + # leave both the reuse state and the pool intact: the blocks are still + # reusable, and the manager can still hand out and resume a new sequence. + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 64) + stream_holder = CachedCudaStream() + probe = self.manager.create_kv_cache() + probe.resume(cast(CudaStream, stream_holder.handle)) + probe.close() + finally: + kv_cache.close() + + self.manager.shutdown() + del self.manager + + def test_guard_counts_only_open_caches(self) -> None: + """The guard counts sequences still open, not objects still referenced.""" + self.prepare(32 << 20, 32 << 20, 1 << 30, 4, 128, 1) + caches = [ + self.manager.create_kv_cache( + ReuseScope(lora_id=None), [self.next_token() for _ in range(64)] + ) + for _ in range(3) + ] + try: + with self.assertRaises(self.GuardError) as ctx: + self.manager.clear_reusable_blocks() + self.assertIn("3 KV cache(s) still open", str(ctx.exception)) + + caches[0].close() + with self.assertRaises(self.GuardError) as ctx: + self.manager.clear_reusable_blocks() + self.assertIn("2 KV cache(s) still open", str(ctx.exception)) + finally: + # Close every cache, not just caches[1:]: if an assertion above fails before + # caches[0] is closed, leaving it open makes tearDown's shutdown() raise and + # mask the real failure. close() is idempotent, so the double close is fine. + for kv_cache in caches: + kv_cache.close() + + # `caches` still holds all three references, so a passing call here proves the + # guard tracks close() rather than object liveness. + self.manager.clear_reusable_blocks() + + class TestBatching(TestKVCacheManagerV2): num_requests: int avg_length: int