[None][fix] Drain in-flight requests before clearing the KV cache reuse state

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 <algorithm>
 #include <chrono>
 #include <cmath>
@@ -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();
 }

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h

@@ -298,6 +298,10 @@ class KvCacheManager : public std::enable_shared_from_this<KvCacheManager>
     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<PoolGroupIndex, float> const& _getTargetRatioList(CacheLevel level) const;

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

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(

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