From ed1f54c4a175a7f858369f43bdb44e358f31d54d Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Tue, 27 Oct 2026 17:48:20 -0700 Subject: [PATCH 1/6] fix(kv-cache): publish reset events without waiting for inference Signed-off-by: tanmayv25 --- .../kv_cache_manager_v2/blockRadixTree.cpp | 4 + .../kv_cache_manager_v2/eventSink.h | 3 + .../streamingEventSink.cpp | 9 ++ .../kv_cache_manager_v2/streamingEventSink.h | 9 +- .../batch_manager/kvCacheManagerV2.cpp | 2 + .../unit_tests/batch_manager/CMakeLists.txt | 12 ++ .../kvCacheManagerV2StreamingEventTest.cpp | 40 +++++++ docs/source/features/kvcache.md | 7 +- .../kv_cache/kv_cache_manager_v2.py | 2 + .../_torch/pyexecutor/kv_cache_events.py | 2 + tensorrt_llm/grpc/openengine/README.md | 2 + .../runtime/kv_cache_manager_v2/__init__.py | 2 + .../runtime/kv_cache_manager_v2/__init__.pyi | 6 +- .../test_kv_cache_event_manager.py | 104 ++++++++++++++++++ 14 files changed, 198 insertions(+), 6 deletions(-) create mode 100644 cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StreamingEventTest.cpp diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index 18ed536c27f0..b37f09d253bd 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/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 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h index eace2073f4b0..efae4f9de227 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h +++ b/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 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.cpp index 81278c7f3d54..197fe12e4aaf 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.cpp +++ b/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 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. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.h index 89e519ce7bd5..1491857f99d2 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.h @@ -47,7 +47,13 @@ struct StreamingBlockRemovedData std::vector blockHashes; }; -using StreamingEventData = std::variant; +//! Semantic data for one wire-level AllBlocksCleared event. +struct StreamingAllBlocksClearedData +{ +}; + +using StreamingEventData + = std::variant; //! 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; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index c503e0250648..a8d9194228ff 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1090,6 +1090,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "StreamingBlockRemovedData") .def_ro("block_hashes", &kv::StreamingBlockRemovedData::blockHashes); + nb::class_(m, "StreamingAllBlocksClearedData"); + nb::class_(m, "StreamingEventStats") .def_ro("stored_blocks", &kv::StreamingEventStats::storedBlocks) .def_ro("removed_blocks", &kv::StreamingEventStats::removedBlocks) diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 0f2879b1dd77..96994ff3bbdd 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/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 + $ + LABELS + "unit;batch-manager;host") + add_tllm_gtest( kvCacheManagerV2ConfigTest OWNER diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StreamingEventTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StreamingEventTest.cpp new file mode 100644 index 000000000000..a1df317c0e6c --- /dev/null +++ b/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 + +#include +#include + +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(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(event)); + } + EXPECT_TRUE(sink->drainIterationEvents().empty()); + EXPECT_EQ(sink->getStats().droppedEvents, 0); +} + +} // namespace diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 3f46d098c863..a7c421efb6e6 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -426,10 +426,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 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a963738f9ff8..fc841f7dcf3c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/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() diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 29533afa3c37..1095ff84f58c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/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 diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 5ce68cf169c1..615d2455cd26 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/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. +A successful prefix-cache reset publishes ordered `AllBlocksCleared` after preceding cache events, without waiting for another inference iteration. This also applies while inference is idle or paused. Reset markers participate in bounded replay and do not change the engine lifetime ID. + 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. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 3eb903af9973..be9ef13dcb65 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/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", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index e418a261853e..8604e198e5ab 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/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: ... diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index 8198529e4069..e22c155f33b9 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -19,12 +19,14 @@ import sys import threading import time +import uuid from importlib.util import find_spec from types import SimpleNamespace from typing import TYPE_CHECKING, cast import msgspec import pytest +import zmq from tensorrt_llm._torch.pyexecutor.kv_cache_events import ( KVEventBatch, @@ -347,6 +349,108 @@ def test_native_streaming_sink_to_python_wire_structs(real_block_factory, lora_i manager.shutdown() +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("populated", [False, True]) +def test_streaming_cache_reset_is_published_while_idle_and_replayed(populated): + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + + init_cuda_once() + endpoint = f"inproc://kv-reset-{uuid.uuid4().hex}" + events = StreamingKVCacheEventManager( + KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint=endpoint, + replay_endpoint=f"{endpoint}-replay", + ), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + ) + events.set_layer_group_window_sizes({0: 128}) + manager = _create_test_manager(events.event_sink) + wrapper = KVCacheManagerV2.__new__(KVCacheManagerV2) + wrapper.impl = manager + wrapper.event_manager = events + wrapper.conversation_manager = None + wrapper._sharing = None + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + replay = zmq.Context.instance().socket(zmq.DEALER) + replay.setsockopt(zmq.RCVTIMEO, 5000) + replay.setsockopt(zmq.SNDTIMEO, 5000) + subscriber = zmq.Context.instance().socket(zmq.SUB) + subscriber.setsockopt(zmq.SUBSCRIBE, b"") + subscriber.setsockopt(zmq.RCVTIMEO, 5000) + events.start() + replay.connect(f"{endpoint}-replay") + subscriber.connect(endpoint) + + def replay_through(expected_batches, start_seq): + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + replay.send_multipart((b"", start_seq.to_bytes(8, "big"))) + batches = [] + while True: + _, _, seq, payload = replay.recv_multipart() + if seq == events._publisher.END_SEQ: + break + batches.append((int.from_bytes(seq, "big"), msgspec.msgpack.decode(payload))) + if len(batches) >= expected_batches: + return batches + pytest.fail("Reset was not published without another inference iteration") + + try: + # Establish the SUB subscription using real resets, without relying on a + # sleep or on future inference traffic to make the connection usable. + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + wrapper.reset_reuse_state() + if subscriber.poll(100): + _, seq, payload = subscriber.recv_multipart() + # A subscription may become ready after an earlier reset was + # queued. Consume through this attempt before fixing the cursor. + while int.from_bytes(seq, "big") < events.enqueued_batches - 1: + _, seq, payload = subscriber.recv_multipart() + assert msgspec.msgpack.decode(payload)[1][-1] == {"type": "AllBlocksCleared"} + start_seq = int.from_bytes(seq, "big") + 1 + break + else: + pytest.fail("Idle reset did not reach the live subscriber") + + if populated: + _commit_and_close(manager, stream, _token_ids(1, 5)) + # Reset must flush both pending cache changes and the reset itself. No + # executor iteration or explicit flush is allowed between reset and replay. + wrapper.reset_reuse_state() + _, seq, payload = subscriber.recv_multipart() + assert int.from_bytes(seq, "big") == start_seq + assert msgspec.msgpack.decode(payload)[1][-1] == {"type": "AllBlocksCleared"} + batches = replay_through(1, start_seq) + first_events = batches[0][1][1] + assert first_events[-1] == {"type": "AllBlocksCleared"} + if populated: + assert first_events[0]["type"] == "BlockStored" + assert first_events[0]["token_ids"] == [1, 2, 3, 4] + + # Reusing the same tokens after reset must emit fresh stored data. + _commit_and_close(manager, stream, _token_ids(1, 5)) + events.flush_iteration_events() + wrapper.reset_reuse_state() + batches = replay_through(3, start_seq) + assert [seq for seq, _ in batches] == list(range(start_seq, start_seq + 3)) + assert batches[1][1][1][0]["type"] == "BlockStored" + assert batches[1][1][1][0]["token_ids"] == [1, 2, 3, 4] + assert batches[2][1][1][-1] == {"type": "AllBlocksCleared"} + assert events.dropped_events == 0 + finally: + replay.close(linger=0) + subscriber.close(linger=0) + manager.shutdown() + events.shutdown() + + def test_native_streaming_sink_separates_lora_scopes(real_block_factory): manager = StreamingKVCacheEventManager( KVEventsConfig(enable_kv_cache_events=True, publisher="null"), From d0a3a0d7f7e1985d1f9404ed6946dec71cfa8bf1 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Tue, 27 Oct 2026 18:01:33 -0700 Subject: [PATCH 2/6] docs(kv-cache): centralize streaming event documentation Signed-off-by: tanmayv25 --- docs/source/features/kvcache.md | 15 ++------------- tensorrt_llm/grpc/openengine/README.md | 2 +- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index a7c421efb6e6..9c5c1978e7f6 100644 --- a/docs/source/features/kvcache.md +++ b/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 @@ -450,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` 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 diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 615d2455cd26..17d9cb830ecc 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -130,7 +130,7 @@ 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. -A successful prefix-cache reset publishes ordered `AllBlocksCleared` after preceding cache events, without waiting for another inference iteration. This also applies while inference is idle or paused. Reset markers participate in bounded replay and do not change the engine lifetime ID. +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. From 0eae22c6518645206f5c1ba172d57ad3dc31347a Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Tue, 27 Oct 2026 22:22:17 -0700 Subject: [PATCH 3/6] test(kv-cache): cover paused control resets and attention DP Signed-off-by: tanmayv25 --- .../test_lists/test-db/l0_dgx_h100.yml | 1 + .../test_lists/test-db/l0_h100.yml | 1 + .../test_kv_cache_reset_events.py | 179 ++++++++++++++++++ 3 files changed, 181 insertions(+) create mode 100644 tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 8354a7290642..7a8503c63dad 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/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" diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 209c866aa838..ab1cb5692df7 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/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" diff --git a/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py b/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py new file mode 100644 index 000000000000..ac473c688c94 --- /dev/null +++ b/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import time + +import msgspec +import pytest +import ray +import zmq +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM +from tensorrt_llm.executor.ray.utils import control_action_decorator +from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams +from tensorrt_llm.llmapi.llm_args import KVEventsConfig +from tensorrt_llm.llmapi.rlhf_utils import WorkerExtension +from tensorrt_llm.scheduling_params import SchedulingParams + + +class ResetEventsWorkerExtension(WorkerExtension): + @control_action_decorator + def reset_and_wait_for_observer(self, endpoint): + # Keep the real executor parked until the driver observes the events; + # resuming it early could hide a missing reset-time flush. + assert self.engine.control_request_barrier.is_set() + assert not self.engine.control_action_done.is_set() + WorkerExtension.reset_prefix_cache.__wrapped__(self) + rank = self.engine.dist.rank + with zmq.Context.instance().socket(zmq.DEALER) as observer: + observer.setsockopt(zmq.LINGER, 0) + observer.setsockopt(zmq.IDENTITY, str(rank).encode()) + observer.setsockopt(zmq.RCVTIMEO, 30_000) + observer.setsockopt(zmq.SNDTIMEO, 30_000) + observer.connect(endpoint) + observer.send(b"reset") + assert observer.recv() == b"observed" + return rank + + +def _replay_until(replay, start_seq, event_type): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + replay.send_multipart([b"", start_seq.to_bytes(8, "big")]) + batches = [] + while True: + _, _, sequence, payload = replay.recv_multipart() + if not payload: + break + batches.append((int.from_bytes(sequence, "big"), msgspec.msgpack.decode(payload))) + if any(event["type"] == event_type for _, batch in batches for event in batch[1]): + return batches + pytest.fail(f"No {event_type} event published while the executor was parked") + + +def _buffered_events(llm): + events = llm._executor.rpc_client.fetch_kv_cache_events_wait_async(timeout=1).remote() + return [json.loads(event) if isinstance(event, str) else event for event in events] + + +def _check_cache_reset_events_through_control_rpc(streaming, tmp_path): + ranks = 1 if streaming else 2 + endpoint = f"ipc://{tmp_path}/events" + replay_endpoint = f"ipc://{tmp_path}/replay" + observer_endpoint = f"ipc://{tmp_path}/observer" + config = KvCacheConfig( + enable_block_reuse=True, + use_kv_cache_manager_v2=True, + max_tokens=1024, + tokens_per_block=32, + event_buffer_max_size=0 if streaming else 1024, + kv_events_config=( + KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint=endpoint, + replay_endpoint=replay_endpoint, + ) + if streaming + else None + ), + ) + with ( + zmq.Context.instance().socket(zmq.ROUTER) as observer, + zmq.Context.instance().socket(zmq.DEALER) as replay, + ): + for socket in (observer, replay): + socket.setsockopt(zmq.LINGER, 0) + socket.setsockopt(zmq.RCVTIMEO, 30_000) + socket.setsockopt(zmq.SNDTIMEO, 30_000) + observer.bind(observer_endpoint) + if streaming: + replay.connect(replay_endpoint) + with LLM( + model=llm_models_root() / "Qwen3/Qwen3-0.6B", + orchestrator_type="ray", + ray_worker_extension_cls=( + "_torch.ray_orchestrator.test_kv_cache_reset_events.ResetEventsWorkerExtension" + ), + tensor_parallel_size=ranks, + enable_attention_dp=not streaming, + enable_lm_head_tp_in_adp=not streaming, + kv_cache_config=config, + load_format="dummy", + model_kwargs={"num_hidden_layers": 1}, + max_seq_len=256, + max_batch_size=2, + max_num_tokens=256, + enable_autotuner=False, + ) as llm: + # Leave rank 1 empty: it must still participate in the buffered + # reset gather while rank 0 has removal events to publish. + output = llm.generate( + [1] * 65, + SamplingParams(max_tokens=1, ignore_eos=True), + scheduling_params=( + SchedulingParams(attention_dp_rank=0, attention_dp_relax=False) + if not streaming + else None + ), + ) + assert output.outputs[0].token_ids + + stored = {rank: set() for rank in range(ranks)} + if streaming: + before = _replay_until(replay, 0, "BlockStored") + start_seq = before[-1][0] + 1 + else: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline and not stored[0]: + for event in _buffered_events(llm): + if event and event["data"]["type"] == "stored": + stored[event["attention_dp_rank"]].update( + block["block_hash"] for block in event["data"]["blocks"] + ) + assert stored[0], "Rank 0 must have cached blocks" + assert not stored[1], "Rank 1 must exercise the empty-cache gather path" + + refs = llm._collective_rpc( + "reset_and_wait_for_observer", args=(observer_endpoint,), non_block=True + ) + parked = [] + try: + for _ in range(ranks): + identity, message = observer.recv_multipart() + parked.append(identity) + assert message == b"reset" + assert set(parked) == {str(rank).encode() for rank in range(ranks)} + if streaming: + after = _replay_until(replay, start_seq, "AllBlocksCleared") + assert after[0][0] == start_seq + assert after[-1][1][1][-1] == {"type": "AllBlocksCleared"} + else: + removed = {rank: set() for rank in range(ranks)} + deadline = time.monotonic() + 10 + while time.monotonic() < deadline and removed != stored: + for event in _buffered_events(llm): + if event and event["data"]["type"] == "removed": + removed[event["attention_dp_rank"]].update( + event["data"]["block_hashes"] + ) + assert removed == stored + finally: + for identity in parked: + observer.send_multipart([identity, b"observed"]) + assert sorted(ray.get(refs, timeout=40)) == list(range(ranks)) + + +@pytest.mark.ray +@pytest.mark.timeout(240) +def test_streaming_reset_is_published_before_control_rpc_resumes(tmp_path): + _check_cache_reset_events_through_control_rpc(True, tmp_path) + + +@pytest.mark.ray +@pytest.mark.gpu2 +@pytest.mark.timeout(240) +def test_buffered_attention_dp_reset_with_an_empty_rank(tmp_path): + _check_cache_reset_events_through_control_rpc(False, tmp_path) From 4682397a6019beea9a9e4af4fbd7200b7b8d30ee Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Tue, 27 Oct 2026 23:40:42 -0700 Subject: [PATCH 4/6] test(kv-cache): use Gloo for buffered reset coverage Signed-off-by: tanmayv25 --- .../_torch/ray_orchestrator/test_kv_cache_reset_events.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py b/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py index ac473c688c94..6348d6431131 100644 --- a/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py +++ b/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py @@ -175,5 +175,7 @@ def test_streaming_reset_is_published_before_control_rpc_resumes(tmp_path): @pytest.mark.ray @pytest.mark.gpu2 @pytest.mark.timeout(240) -def test_buffered_attention_dp_reset_with_an_empty_rank(tmp_path): +def test_buffered_attention_dp_reset_with_an_empty_rank(tmp_path, monkeypatch): + """Exercise the event gather through Gloo's CPU object transport.""" + monkeypatch.setenv("TLLM_CUDA_OBJECT_COLLECTIVES", "0") _check_cache_reset_events_through_control_rpc(False, tmp_path) From 9bdd9f04a7bd89afe70c77f5ed003ac8ed6fc418 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 28 Oct 2026 00:51:48 -0700 Subject: [PATCH 5/6] test(kv-cache): stabilize reset baseline and repair CI fixtures Signed-off-by: tanmayv25 --- .../_torch/executor/test_pytorch_model_engine.py | 4 ++-- .../ray_orchestrator/test_kv_cache_reset_events.py | 12 ++++++++++-- tests/unittest/_torch/test_connector_adp.py | 7 ++++++- 3 files changed, 18 insertions(+), 5 deletions(-) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 6bae538cca81..d3310d812328 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/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 diff --git a/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py b/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py index 6348d6431131..280eb0c16930 100644 --- a/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py +++ b/tests/unittest/_torch/ray_orchestrator/test_kv_cache_reset_events.py @@ -21,6 +21,7 @@ class ResetEventsWorkerExtension(WorkerExtension): @control_action_decorator def reset_and_wait_for_observer(self, endpoint): + """Reset reuse state and hold every executor until its events are observed.""" # Keep the real executor parked until the driver observes the events; # resuming it early could hide a missing reset-time flush. assert self.engine.control_request_barrier.is_set() @@ -39,6 +40,7 @@ def reset_and_wait_for_observer(self, endpoint): def _replay_until(replay, start_seq, event_type): + """Read retained batches until the requested event appears, without resuming inference.""" deadline = time.monotonic() + 10 while time.monotonic() < deadline: replay.send_multipart([b"", start_seq.to_bytes(8, "big")]) @@ -59,6 +61,7 @@ def _buffered_events(llm): def _check_cache_reset_events_through_control_rpc(streaming, tmp_path): + """Observe reset events through public transports while the control RPC is parked.""" ranks = 1 if streaming else 2 endpoint = f"ipc://{tmp_path}/events" replay_endpoint = f"ipc://{tmp_path}/replay" @@ -126,14 +129,19 @@ def _check_cache_reset_events_through_control_rpc(streaming, tmp_path): before = _replay_until(replay, 0, "BlockStored") start_seq = before[-1][0] + 1 else: + # The 65-token prompt stores two full blocks and one partial + # tail, which can arrive in separate publisher batches. + expected_blocks = 3 deadline = time.monotonic() + 10 - while time.monotonic() < deadline and not stored[0]: + while time.monotonic() < deadline and len(stored[0]) < expected_blocks: for event in _buffered_events(llm): if event and event["data"]["type"] == "stored": stored[event["attention_dp_rank"]].update( block["block_hash"] for block in event["data"]["blocks"] ) - assert stored[0], "Rank 0 must have cached blocks" + assert len(stored[0]) == expected_blocks, ( + "Rank 0 must publish all three cached blocks" + ) assert not stored[1], "Rank 1 must exercise the empty-cache gather path" refs = llm._collective_rpc( diff --git a/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index ad1e6829630e..a046fa6fe884 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -78,9 +78,13 @@ def _manager(rank: int = 1) -> KvCacheConnectorManager: def _configure_recovery_controls(executor: "PyExecutor") -> None: + from tensorrt_llm._torch.distributed import Distributed from tensorrt_llm._torch.pyexecutor.executor_request_queue import ExecutorRequestQueue - executor.dist = MagicMock(tp_size=1, cp_size=1, pp_size=1) + executor.dist = MagicMock(tp_size=1, cp_size=1, pp_size=1, tp_rank=0, has_cp_helix=False) + executor.dist.tp_cp_allgather_int64.side_effect = partial( + Distributed.tp_cp_allgather_int64, executor.dist + ) executor.enable_attention_dp = True executor.enable_early_first_token_response = False executor.kv_connector_manager = None @@ -414,6 +418,7 @@ def test_batch_start_combines_readiness_with_status(batch_sizes: tuple[int, int] _configure_recovery_controls(executor) executor.kv_connector_manager = _manager(rank) executor.dist.tp_size = 2 + executor.dist.tp_rank = rank executor.dist.tp_allgather_int64.side_effect = partial(collective.gather_int64, rank) executors.append(executor) From b1e30ece8c5899b1d0466ca41072a23f03ef3184 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 28 Oct 2026 01:01:45 -0700 Subject: [PATCH 6/6] test(executor): initialize cache transfer state in ADP fixture Signed-off-by: tanmayv25 --- tests/unittest/_torch/test_connector_adp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index a046fa6fe884..b97ba7539111 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -88,6 +88,7 @@ def _configure_recovery_controls(executor: "PyExecutor") -> None: executor.enable_attention_dp = True executor.enable_early_first_token_response = False executor.kv_connector_manager = None + executor.kv_cache_transceiver = None executor._is_kv_manager_v2 = False executor._kv_load_tracker = KvLoadTracker(None, rank=0, enabled=False) executor._kv_cache_capacity = {}