From b9bc48d14580078b66adae20790909b42cd6f73c Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Sat, 1 Aug 2026 09:25:49 +0000 Subject: [PATCH 1/3] [None][fix] drain in-flight requests before clearing the KV cache reuse state WorkerExtension.reset_prefix_cache() was missing @control_action_decorator, unlike both of its neighbours: update_weights(), whose docstring calls out that it "uses the control_action_decorator to ensure all active requests are finished", and wait_for_engine_idle(). It is a public Ray worker-extension method, so _collective_rpc("reset_prefix_cache") reached KvCacheManager::clearReusableBlocks() with requests still in flight. Clearing detaches the whole radix tree, but a live request keeps owning the tree blocks it matched via SeqBlock::treeBlock, and goes on using them as the parent for the block its close() commits. resume() deliberately un-commits a partial trailing block so it is re-committed at close, so this is the normal path, not a corner case. The resulting block is grafted onto a detached subtree: unreachable from any root, so the work is silently discarded. On C++ it is worse than silent, because addOrGetExistingBlock() derives the block size via prev->tokensPerBlock(), which reads the now-null grandparent link and segfaults in release builds. Add the decorator so the reuse state is only cleared once the engine has drained. Also guard both entry points -- clearReusableBlocks() and shutdown(), which frees the storage those pages live in -- so misuse reports itself at the offending call rather than surfacing later as discarded work or a null dereference. Both backends name the API in the message, and the check counts sequences that are still open rather than objects that are still referenced, since KvCache::close() / _KVCache.close() drop the entry. shutdown() runs from ~KvCacheManager and __del__, which must not propagate, so both report and skip teardown instead -- leaving the storage alive is safer than freeing it under live pages. Signed-off-by: Yao Yao --- .../kv_cache_manager_v2/kvCacheManager.cpp | 18 ++++++++++++++- .../kv_cache_manager_v2/kvCacheManager.h | 4 ++++ tensorrt_llm/llmapi/rlhf_utils.py | 8 ++++++- .../_core/_kv_cache_manager.py | 22 ++++++++++++++++++- 4 files changed, 49 insertions(+), 3 deletions(-) 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..e97bce916d35 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 KvCache(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( From bbcfa5c6db97b40f274c5c2a7b6b07aebf0159fc Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Sun, 2 Aug 2026 04:33:17 +0000 Subject: [PATCH 2/3] [None][test] cover the living-KvCache guards and align the message across backends The guards added with the drain fix had no regression coverage: every existing clear_reusable_blocks() call site closes its caches first, so nothing exercised the rejection path, and nothing pinned the promise that a rejected call leaves the storage intact. Add TestLivingKvCacheGuard: both entry points reject an open cache, the operation succeeds once it is closed, and the count tracks close() rather than object liveness -- the last one matters because the guard counts sequences that are still open, not objects that are still referenced. Writing it surfaced a divergence the fix had claimed to remove: C++ said "KvCache(s)" where Python said "KV cache(s)", so the two backends reported the same mistake differently. Use the Python wording in C++ so the message is identical either way, and assert the shared text so the two cannot drift apart again unnoticed. The exception types still differ -- Python raises LogicError while TLLM_CHECK_WITH_INFO surfaces as RuntimeError, since only RequestSpecificException is translated -- so the test accepts either rather than pretending they agree. Signed-off-by: Yao Yao --- .../kv_cache_manager_v2/kvCacheManager.cpp | 2 +- .../test_kv_cache_manager_v2.py | 80 ++++++++++++++++++- 2 files changed, 79 insertions(+), 3 deletions(-) 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 e97bce916d35..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 @@ -141,7 +141,7 @@ KvCacheManager::~KvCacheManager() void KvCacheManager::_checkNoLivingKvCaches(char const* api) const { TLLM_CHECK_WITH_INFO(mLivingKvCaches.empty(), - "%s with %zu KvCache(s) still open; close them (or drain the engine) first", api, mLivingKvCaches.size()); + "%s with %zu KV cache(s) still open; close them (or drain the engine) first", api, mLivingKvCaches.size()); } void KvCacheManager::shutdown() 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..4844aff1d1f2 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,82 @@ 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 _prepare_with_open_cache(self) -> _KVCache: + self.prepare(32 << 20, 32 << 20, 1 << 30, 4, 128, 1) + prompt = [self.next_token() for _ in range(64)] + return self.manager.create_kv_cache(ReuseScope(lora_id=None), prompt) + + def test_clear_reusable_blocks_rejects_open_kv_cache(self) -> None: + kv_cache = self._prepare_with_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)) + finally: + kv_cache.close() + + # Storage must survive the rejected call: the operation succeeds once closed. + self.manager.clear_reusable_blocks() + + def test_shutdown_rejects_open_kv_cache(self) -> None: + kv_cache = self._prepare_with_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)) + finally: + kv_cache.close() + + # The rejected shutdown() must not have torn down storage, so this one works. + 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: + for kv_cache in caches[1:]: + 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 From 1b56f684b7b58eadd97f0c6a9e8efd45e8e5a665 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Sun, 2 Aug 2026 13:47:24 +0000 Subject: [PATCH 3/3] [None][test] make the living-KvCache guard tests non-vacuous The guards were exercised against an empty radix tree: the caches were created but never resumed or committed, so "the rejected call left the storage intact" had nothing to observe. An implementation that cleared the tree and *then* raised passed just as well as the correct one. Seed committed, reusable content first, and assert it is still reusable after the rejection -- then assert it is gone after the call that is allowed through, so the first check cannot pass vacuously. Verified by mutation: reordering clear_reusable_blocks() to clear before checking now fails with 0 != 64, and neutering the clear fails the follow-up assertion. For shutdown(), which frees the storage rather than just the tree, also allocate and resume a fresh sequence after the rejection to show the pool itself survived. Close every cache in the counting test's cleanup rather than all but the first. The first was closed inside the try block, so an assertion failing ahead of it left it open, and tearDown's shutdown() then raised the guard error over the top of the real failure. close() is idempotent, so closing all of them is enough. Signed-off-by: Yao Yao --- .../test_kv_cache_manager_v2.py | 55 ++++++++++++++++--- 1 file changed, 47 insertions(+), 8 deletions(-) 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 4844aff1d1f2..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 @@ -984,13 +984,35 @@ class TestLivingKvCacheGuard(TestKVCacheManagerV2): # this test is meaningful under both. GuardError = (LogicError, RuntimeError) - def _prepare_with_open_cache(self) -> _KVCache: - self.prepare(32 << 20, 32 << 20, 1 << 30, 4, 128, 1) + 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)] - return self.manager.create_kv_cache(ReuseScope(lora_id=None), prompt) + 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: - kv_cache = self._prepare_with_open_cache() + 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() @@ -998,23 +1020,37 @@ def test_clear_reusable_blocks_rejects_open_kv_cache(self) -> None: # 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() - # Storage must survive the rejected call: the operation succeeds once closed. + # ...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: - kv_cache = self._prepare_with_open_cache() + 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() - # The rejected shutdown() must not have torn down storage, so this one works. self.manager.shutdown() del self.manager @@ -1037,7 +1073,10 @@ def test_guard_counts_only_open_caches(self) -> None: self.manager.clear_reusable_blocks() self.assertIn("2 KV cache(s) still open", str(ctx.exception)) finally: - for kv_cache in caches[1:]: + # 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