[None][fix] Publish cache reset events without waiting for inference

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp

@@ -1015,6 +1015,10 @@ void BlockRadixTree::clear()
     TLLM_CHECK_DEBUG(mRoots.size() == mPendingRootErases.size());
     mRoots.clear();
     mPendingRootErases.clear();
+    if (mEventSink)
+    {
+        mEventSink->addAllBlocksCleared();
+    }
 }
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h

@@ -47,6 +47,9 @@ class EventSink
 
     //! Protocols that expose retention metadata receive native page priority changes.
     virtual void addPriorityUpdated(Digest const&, Priority, Priority, LifeCycleId) {}
+
+    //! Notify streaming consumers after the complete reuse tree has been cleared.
+    virtual void addAllBlocksCleared() {}
 };
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.cpp

@@ -112,6 +112,15 @@ void StreamingEventSink::addRemovedLifeCycle(Digest const& blockKey, LifeCycleId
     addRemovedBlockUnlocked(blockKey);
 }
 
+void StreamingEventSink::addAllBlocksCleared()
+{
+    std::lock_guard<std::mutex> lock(mMutex);
+    // Reset markers must survive the entry cap, like removals. Keep queued
+    // events in order and discard all residency remembered before the reset.
+    mPendingEvents.emplace_back(StreamingAllBlocksClearedData{});
+    mStoredBlocks.clear();
+}
+
 void StreamingEventSink::addCacheLevelUpdated(Digest const&, CacheLevel, CacheLevel, LifeCycleId)
 {
     // The streaming protocol currently tracks radix-tree residency, not cache-tier movement.

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.h

@@ -47,7 +47,13 @@ struct StreamingBlockRemovedData
     std::vector<int64_t> blockHashes;
 };
 
-using StreamingEventData = std::variant<StreamingBlockStoredData, StreamingBlockRemovedData>;
+//! Semantic data for one wire-level AllBlocksCleared event.
+struct StreamingAllBlocksClearedData
+{
+};
+
+using StreamingEventData
+    = std::variant<StreamingBlockStoredData, StreamingBlockRemovedData, StreamingAllBlocksClearedData>;
 
 //! Counters accumulated over the lifetime of a streaming event sink.
 struct StreamingEventStats
@@ -78,6 +84,7 @@ class StreamingEventSink final : public EventSink
     void addStoredLifeCycle(Block const& block, LifeCycleId lifeCycle) override;
     void addRemovedBlock(Digest const& blockKey) override;
     void addRemovedLifeCycle(Digest const& blockKey, LifeCycleId lifeCycle) override;
+    void addAllBlocksCleared() override;
     void addCacheLevelUpdated(
         Digest const& blockKey, CacheLevel oldLevel, CacheLevel newLevel, LifeCycleId lifeCycle) override;
 

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp

@@ -1090,6 +1090,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
     nb::class_<kv::StreamingBlockRemovedData>(m, "StreamingBlockRemovedData")
         .def_ro("block_hashes", &kv::StreamingBlockRemovedData::blockHashes);
 
+    nb::class_<kv::StreamingAllBlocksClearedData>(m, "StreamingAllBlocksClearedData");
+
     nb::class_<kv::StreamingEventStats>(m, "StreamingEventStats")
         .def_ro("stored_blocks", &kv::StreamingEventStats::storedBlocks)
         .def_ro("removed_blocks", &kv::StreamingEventStats::removedBlocks)

cpp/tests/unit_tests/batch_manager/CMakeLists.txt

@@ -145,6 +145,18 @@ add_tllm_gtest(
   LABELS
   "unit;batch-manager;gpu")
 
+add_tllm_gtest(
+  kvCacheManagerV2StreamingEventTest
+  OWNER
+  batch-manager
+  SOURCES
+  kvCacheManagerV2StreamingEventTest.cpp
+  LINK_LIBRARIES
+  tllm::kv_cache
+  $<LINK_LIBRARY:WHOLE_ARCHIVE,tllm_kv_cache_test_support>
+  LABELS
+  "unit;batch-manager;host")
+
 add_tllm_gtest(
   kvCacheManagerV2ConfigTest
   OWNER

cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StreamingEventTest.cpp

@@ -0,0 +1,40 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "kvCacheManagerV2TestUtils.h"
+#include "kv_cache_manager_v2/blockRadixTree.h"
+#include "kv_cache_manager_v2/streamingEventSink.h"
+
+#include <gtest/gtest.h>
+
+#include <memory>
+#include <variant>
+
+namespace
+{
+
+using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2;
+
+TEST(KvCacheManagerV2StreamingEventTest, EmptyCacheResetsAreObservableWithoutInference)
+{
+    auto const config = test::makeConfig();
+    LifeCycleRegistry const lifeCycles(config);
+    auto const sink = std::make_shared<StreamingEventSink>(1);
+    BlockRadixTree tree(lifeCycles, config.tokensPerBlock, sink);
+
+    // Repeated empty resets must be visible even when capture has a one-entry cap.
+    tree.clear();
+    tree.clear();
+    auto const events = sink->drainIterationEvents();
+    ASSERT_EQ(events.size(), 2);
+    for (auto const& event : events)
+    {
+        EXPECT_TRUE(std::holds_alternative<StreamingAllBlocksClearedData>(event));
+    }
+    EXPECT_TRUE(sink->drainIterationEvents().empty());
+    EXPECT_EQ(sink->getStats().droppedEvents, 0);
+}
+
+} // namespace

docs/source/features/kvcache.md

@@ -402,12 +402,7 @@ exposed through `LLM.get_kv_cache_events()` / `LLM.get_kv_cache_events_async()`,
 
 #### Streaming path (prototype)
 
-Configured with ```kv_cache_config.kv_events_config```. Streaming is intended to reduce event
-publishing overhead under attention data parallelism: each emitting rank publishes its own
-events over a ZeroMQ `PUB` socket, removing the rank-0 gather and the consumer pull through the
-LLM API. Each emitting rank still drains its local native event source, converts the events to
-wire structs and enqueues one batch at the iteration boundary. A background thread performs
-msgpack encoding and socket I/O.
+Configured with ```kv_cache_config.kv_events_config```. Streaming is intended to reduce event publishing overhead under attention data parallelism: each emitting rank publishes its own events over a ZeroMQ `PUB` socket, removing the rank-0 gather and the consumer pull through the LLM API. Each emitting rank drains its local native event source, converts the events to wire structs and enqueues a batch at iteration boundaries and when a prefix-cache reset completes. A background thread performs msgpack encoding and socket I/O.
 
 ```python
 from tensorrt_llm.llmapi import KvCacheConfig, KVEventsConfig
@@ -426,10 +421,9 @@ kv_cache_config = KvCacheConfig(
 no events while streaming is enabled. Draft models and KV-cache-size estimation do not publish
 events. Use buffered mode for cache-tier, priority, and other lifecycle updates.
 
-Streaming exposes only the full-block residency information external routers need, using the
-attention lifecycle with the largest window. Conversion to `BlockStored`/`BlockRemoved` is
-centralized at the once-per-iteration publisher boundary, keeping internal lifecycle details
-contained and avoiding Python callbacks from native cache operations.
+A successful prefix-cache reset emits ordered `AllBlocksCleared` after preceding stored or removed events. Reset flushes events without waiting for another inference iteration, including while inference is idle or paused. The marker is retained in bounded replay history and does not change the engine lifetime ID.
+
+Streaming exposes only the full-block residency information external routers need, using the attention lifecycle with the largest window. Conversion to `BlockStored`, `BlockRemoved`, and `AllBlocksCleared` is centralized at iteration and reset publication boundaries, keeping internal lifecycle details contained and avoiding Python callbacks from native cache operations.
 
 **Multimodal payloads.** `token_ids` can contain integers and hexadecimal digest strings;
 integer-only consumers are incompatible. Each `mm_keys[i]` describes the multimodal segments
@@ -451,13 +445,7 @@ port must be at least *ranks-per-host* away from ```endpoint```'s rather than `N
 Overlapping ranges are rejected at startup. For `ipc://` and `inproc://` endpoints, which
 have no port, each rank appends a `_dp<rank>` suffix instead.
 
-**Wire format.** Each batch is sent as three ZeroMQ frames: the subscription ```topic```,
-an 8-byte big-endian sequence number, and a msgpack payload
-`[timestamp, [events], data_parallel_rank]`. Each event is a map tagged with a `type` key —
-`BlockStored` or `BlockRemoved` — carrying int64 block hashes derived from the V2 radix block
-keys. This is the format documented for custom router backends; it differs from vLLM's
-positional-array encoding of the individual events, though the batch envelope is positional in
-both.
+**Wire format.** Each batch is sent as three ZeroMQ frames: the subscription ```topic```, an 8-byte big-endian sequence number, and a msgpack payload `[timestamp, [events], data_parallel_rank]`. Each event is a map tagged with a `type` key: `BlockStored`, `BlockRemoved`, or `AllBlocksCleared`. Stored and removed events carry int64 block hashes derived from the V2 radix block keys. `AllBlocksCleared` has no additional fields. This is the format documented for custom router backends; it differs from vLLM's positional-array encoding of the individual events, though the batch envelope is positional in both.
 
 **Delivery guarantees.** Delivery is best effort, and batch loss is observable. Every accepted
 batch reserves a sequence number up front, so a batch dropped by a full publisher queue

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -8193,3 +8193,5 @@ def reset_reuse_state(self):
             self.conversation_manager.clear()
         if self._sharing is not None:
             self._sharing._on_reset()
+        # A reset can run while inference is paused or idle, with no next iteration.
+        self.flush_iteration_events()

tensorrt_llm/_torch/pyexecutor/kv_cache_events.py

@@ -591,6 +591,8 @@ def drain_events(self) -> list[BlockStored | BlockRemoved | AllBlocksCleared]:
                 )
             elif isinstance(event, kv_cache_manager_v2_runtime.StreamingBlockRemovedData):
                 result.append(BlockRemoved(block_hashes=list(event.block_hashes), medium="GPU"))
+            elif isinstance(event, kv_cache_manager_v2_runtime.StreamingAllBlocksClearedData):
+                result.append(AllBlocksCleared())
             else:
                 raise TypeError(f"Unsupported native streaming KV event: {type(event)!r}")
         return result

tensorrt_llm/grpc/openengine/README.md

@@ -130,6 +130,8 @@ Placement exchange occurs only during OpenEngine executor startup. Metadata RPCs
 
 Events report full-block radix-cache residency; the streaming sink does not report live GPU-versus-cold-tier movement. Load snapshots independently report GPU block usage. The upstream sink emits signed 64-bit hashes from the first eight bytes of each V2 radix key, regardless of `kv_cache_event_hash_algo`. `ServerInfo.extra.trtllm_kv_event_hash_algo` advertises `v2_sha256_64` when streaming is enabled. These hashes identify blocks in event parent chains and removals. A router can derive its own request-compatible token hashes from the stored block `token_ids`, as Dynamo does.
 
+See [KV-cache streaming events](../../../docs/source/features/kvcache.md#streaming-path-prototype) for cache-reset behavior and delivery guarantees.
+
 Consumers must track rank-local sequence numbers and resynchronize on gaps. Publisher queue drops and failed sends leave sequence gaps, but native capture overflow is logged without a wire gap. Replay retains only the configured number of batches; a slow replay client can lose its replay response and must retry. Replay is sent in bounded, nonblocking chunks so it cannot stall live publication. The upstream publisher emits no idle heartbeat, so source health cannot be inferred from silence alone. A reset or restart does not reconstruct already-resident cache contents, and recovery can temporarily under-report cache affinity.
 
 Direct ZMQ sockets are unauthenticated. Keep them on a trusted network and restrict access to intended subscribers.

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py

@@ -99,6 +99,7 @@ class _BatchDescFieldSpec:
 KVCacheIterationStatsDelta = _cpp.KVCacheIterationStatsDelta
 KVCacheManager = _cpp.KVCacheManager
 KVCacheManagerConfig = _cpp.KVCacheManagerConfig
+StreamingAllBlocksClearedData = _cpp.StreamingAllBlocksClearedData
 StreamingBlockRemovedData = _cpp.StreamingBlockRemovedData
 StreamingBlockStoredData = _cpp.StreamingBlockStoredData
 StreamingEventSink = _cpp.StreamingEventSink
@@ -284,6 +285,7 @@ def typed_range(*args: int) -> range:
     "SsmLayerConfig",
     "StreamingBlockRemovedData",
     "StreamingBlockStoredData",
+    "StreamingAllBlocksClearedData",
     "StreamingEventSink",
     "StreamingEventStats",
     "SwaScratchReuseConfig",

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi

@@ -407,6 +407,8 @@ class StreamingBlockRemovedData:
     @property
     def block_hashes(self) -> list[int]: ...
 
+class StreamingAllBlocksClearedData: ...
+
 class StreamingEventStats:
     @property
     def stored_blocks(self) -> int: ...
@@ -428,7 +430,9 @@ class StreamingEventSink:
     def set_target_life_cycle(self, life_cycle_id: int) -> None: ...
     def drain_iteration_events(
         self,
-    ) -> list[StreamingBlockStoredData | StreamingBlockRemovedData]: ...
+    ) -> list[
+        StreamingBlockStoredData | StreamingBlockRemovedData | StreamingAllBlocksClearedData
+    ]: ...
     @property
     def stats(self) -> StreamingEventStats: ...
 

tests/integration/test_lists/test-db/l0_dgx_h100.yml

@@ -288,6 +288,7 @@ l0_dgx_h100:
       orchestrator: ray
   tests:
     - unittest/_torch/ray_orchestrator/multi_gpu -m "gpu2"
+    - unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py::test_buffered_attention_dp_reset_with_an_empty_rank
     - unittest/_torch/ray_orchestrator/multi_gpu/test_inflight_weight_update.py
     - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part0"
     - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part1"

tests/integration/test_lists/test-db/l0_h100.yml

@@ -270,6 +270,7 @@ l0_h100:
     - unittest/executor/test_shim_ray.py
     - unittest/_torch/ray_orchestrator/single_gpu/test_llm_sleep.py
     - unittest/_torch/ray_orchestrator/single_gpu/test_postproc_workers.py
+    - unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py::test_streaming_reset_is_published_before_control_rpc_resumes
     - unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py -m "part0"
     - unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py -m "part1"
     - unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py -m "part2"

tests/unittest/_torch/executor/test_pytorch_model_engine.py

@@ -1904,12 +1904,12 @@ def test_mrope_multi_beam_generation_positions(self) -> None:
         """Every beam row gets its request's `past + delta` on all 3 axes."""
         beam_width = 2
         engine, kv_cache_manager = create_model_engine_and_kvcache(
+            self,
             llm_args=TorchLlmArgs(model="dummy",
                                   max_batch_size=4,
                                   max_num_tokens=64,
                                   max_beam_width=beam_width),
             model=DummyModel(torch.half, rope_scaling={"type": "mrope"}))
-        self.addCleanup(kv_cache_manager.shutdown)
         preparer = engine._runner._inputs
         # (num_tokens, delta): one delta arrives on the host and one is
         # already a CUDA int32 scalar.
@@ -1973,11 +1973,11 @@ def test_mrope_mixed_batched_and_fallback_positions_match_fallback(
             self) -> None:
         """Mixed scalar batching matches the old per-request writes exactly."""
         engine, kv_cache_manager = create_model_engine_and_kvcache(
+            self,
             llm_args=TorchLlmArgs(model="dummy",
                                   max_batch_size=2,
                                   max_num_tokens=64),
             model=DummyModel(torch.half, rope_scaling={"type": "mrope"}))
-        self.addCleanup(kv_cache_manager.shutdown)
         preparer = engine._runner._inputs
         sentinel = 123