From 1454f28a18fa6539bee8bc1b7c110b3c0b7b6753 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Wed, 16 Sep 2026 01:59:52 -0700 Subject: [PATCH 1/3] [None][feat] support C++ streaming KV event capture Signed-off-by: Guan Luo --- .../kv_cache_manager_v2/CMakeLists.txt | 1 + .../streamingEventSink.cpp | 245 ++++++++++++++++++ .../kv_cache_manager_v2/streamingEventSink.h | 93 +++++++ .../batch_manager/kvCacheManagerV2.cpp | 54 +++- docs/source/features/kvcache.md | 12 +- .../kv_cache/kv_cache_manager_v2.py | 24 +- .../_torch/pyexecutor/kv_cache_events.py | 87 +++++-- .../runtime/kv_cache_manager_v2/__init__.py | 15 +- .../kv_cache_manager_v2/_introspection.py | 36 +++ .../test_kv_cache_event_manager.py | 76 ++++++ .../test_streaming_kv_events.py | 9 +- 11 files changed, 612 insertions(+), 40 deletions(-) create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.cpp create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.h diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt index a2a737db9821..c52bbd86734c 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt @@ -57,6 +57,7 @@ set(KV_CACHE_MANAGER_V2_SRCS kv_cache_manager_v2/batchedPageCopy.cu kv_cache_manager_v2/coldPageCodec.cpp kv_cache_manager_v2/eventManager.cpp + kv_cache_manager_v2/streamingEventSink.cpp kv_cache_manager_v2/page.cpp kv_cache_manager_v2/storageManager.cpp kv_cache_manager_v2/introspection.cpp 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 new file mode 100644 index 000000000000..c98cdd8b5b8b --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.cpp @@ -0,0 +1,245 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/streamingEventSink.h" + +#include "kv_cache_manager_v2/blockRadixTree.h" +#include "kv_cache_manager_v2/page.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +StreamingEventSink::StreamingEventSink(int tokensPerBlock, int maxEntries) + : mTokensPerBlock(tokensPerBlock) + , mMaxEntries(maxEntries) +{ + if (mTokensPerBlock <= 0) + { + throw std::invalid_argument("tokensPerBlock must be positive"); + } + if (mMaxEntries <= 0) + { + throw std::invalid_argument("maxEntries must be positive"); + } +} + +void StreamingEventSink::setTargetLifeCycle(LifeCycleId lifeCycle) +{ + std::lock_guard lock(mMutex); + mTargetLifeCycle = lifeCycle; +} + +std::vector StreamingEventSink::drainIterationEvents() +{ + std::lock_guard lock(mMutex); + auto events = std::move(mPendingEvents); + mPendingEvents.clear(); + mPendingEntries = 0; + return events; +} + +StreamingEventStats StreamingEventSink::getStats() const +{ + std::lock_guard lock(mMutex); + return mStats; +} + +void StreamingEventSink::addStoredBlock(Block const& block) +{ + std::lock_guard lock(mMutex); + addStoredBlockUnlocked(block); +} + +void StreamingEventSink::addStoredLifeCycle(Block const& block, LifeCycleId lifeCycle) +{ + std::lock_guard lock(mMutex); + if (!mTargetLifeCycle.has_value()) + { + return; + } + if (lifeCycle != *mTargetLifeCycle) + { + ++mStats.nonTargetLifeCyclesIgnored; + return; + } + addStoredBlockUnlocked(block); +} + +void StreamingEventSink::addRemovedBlock(Digest const& blockKey) +{ + std::lock_guard lock(mMutex); + addRemovedBlockUnlocked(blockKey); +} + +void StreamingEventSink::addRemovedLifeCycle(Digest const& blockKey, LifeCycleId lifeCycle) +{ + std::lock_guard lock(mMutex); + if (!mTargetLifeCycle.has_value()) + { + return; + } + if (lifeCycle != *mTargetLifeCycle) + { + ++mStats.nonTargetLifeCyclesIgnored; + return; + } + addRemovedBlockUnlocked(blockKey); +} + +void StreamingEventSink::addCacheLevelUpdated(Digest const&, CacheLevel, CacheLevel, LifeCycleId) +{ + // The streaming protocol currently tracks radix-tree residency, not cache-tier movement. +} + +void StreamingEventSink::addStoredBlockUnlocked(Block const& block) +{ + if (!mTargetLifeCycle.has_value() || *mTargetLifeCycle >= block.storage.size()) + { + return; + } + auto const* page = block.getPage(*mTargetLifeCycle); + if (page == nullptr) + { + return; + } + if (!block.isFull() || page->numTokensInBlock < static_cast(block.tokens.size())) + { + ++mStats.partialBlocksSuppressed; + return; + } + if (mStoredBlocks.count(block.key) != 0) + { + return; + } + for (auto const& token : block.tokens) + { + if (token.isDigest()) + { + ++mStats.multimodalBlocksSuppressed; + return; + } + } + if (!reserveEntryUnlocked()) + { + return; + } + if (block.prev == nullptr) + { + throw std::logic_error("Cannot publish an orphan KV cache block"); + } + + int64_t const blockHash = wireHash(block.key); + std::optional parentHash; + if (block.prev->type() == NodeBase::Type::kBLOCK) + { + parentHash = wireHash(static_cast(block.prev)->key); + } + + std::vector tokenIds; + tokenIds.reserve(block.tokens.size()); + for (auto const& token : block.tokens) + { + tokenIds.push_back(token.tokenId()); + } + + mStoredBlocks.emplace(block.key, blockHash); + if (!mPendingEvents.empty()) + { + auto* stored = std::get_if(&mPendingEvents.back()); + if (stored != nullptr && !stored->blockHashes.empty() && parentHash.has_value() + && stored->blockHashes.back() == *parentHash) + { + stored->blockHashes.push_back(blockHash); + stored->tokenIds.insert(stored->tokenIds.end(), tokenIds.begin(), tokenIds.end()); + ++mStats.storedBlocks; + return; + } + } + mPendingEvents.emplace_back(StreamingBlockStoredData{{blockHash}, parentHash, std::move(tokenIds)}); + ++mStats.storedBlocks; +} + +void StreamingEventSink::addRemovedBlockUnlocked(Digest const& blockKey) +{ + auto const stored = mStoredBlocks.find(blockKey); + if (stored == mStoredBlocks.end()) + { + return; + } + int64_t const blockHash = stored->second; + mStoredBlocks.erase(stored); + addRemovedHashUnlocked(blockHash); +} + +void StreamingEventSink::addRemovedHashUnlocked(int64_t blockHash) +{ + if (!mPendingEvents.empty()) + { + auto* removed = std::get_if(&mPendingEvents.back()); + if (removed != nullptr) + { + removed->blockHashes.push_back(blockHash); + ++mStats.removedBlocks; + return; + } + } + mPendingEvents.emplace_back(StreamingBlockRemovedData{{blockHash}}); + ++mStats.removedBlocks; +} + +bool StreamingEventSink::reserveEntryUnlocked() +{ + if (mPendingEntries < mMaxEntries) + { + ++mPendingEntries; + return true; + } + ++mStats.droppedEvents; + int64_t const dropped = mStats.droppedEvents; + if (dropped == 1 || (dropped & (dropped - 1)) == 0) + { + TLLM_LOG_WARNING( + "Dropping streaming KV events because the per-iteration safety cap was exceeded; " + "dropped_events=%" PRId64, + dropped); + } + return false; +} + +int64_t StreamingEventSink::wireHash(Digest const& digest) +{ + uint64_t value = 0; + for (size_t i = 0; i < sizeof(value); ++i) + { + value = (value << 8U) | std::to_integer(digest[i]); + } + uint64_t constexpr kSignedMax = static_cast(std::numeric_limits::max()); + if (value <= kSignedMax) + { + return static_cast(value); + } + return -static_cast(~value) - 1; +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 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 new file mode 100644 index 000000000000..22e314314d0a --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/streamingEventSink.h @@ -0,0 +1,93 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/eventSink.h" + +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +//! Semantic data for one wire-level BlockStored event. +struct StreamingBlockStoredData +{ + std::vector blockHashes; + std::optional parentBlockHash; + std::vector tokenIds; +}; + +//! Semantic data for one wire-level BlockRemoved event. +struct StreamingBlockRemovedData +{ + std::vector blockHashes; +}; + +using StreamingEventData = std::variant; + +//! Counters accumulated over the lifetime of a streaming event sink. +struct StreamingEventStats +{ + int64_t storedBlocks = 0; + int64_t removedBlocks = 0; + int64_t partialBlocksSuppressed = 0; + int64_t multimodalBlocksSuppressed = 0; + int64_t nonTargetLifeCyclesIgnored = 0; + int64_t droppedEvents = 0; +}; + +//! Captures streaming KV-cache lifecycle events without depending on Python or a transport. +class StreamingEventSink final : public EventSink +{ +public: + StreamingEventSink(int tokensPerBlock, int maxEntries); + + void setTargetLifeCycle(LifeCycleId lifeCycle); + [[nodiscard]] std::vector drainIterationEvents(); + [[nodiscard]] StreamingEventStats getStats() const; + + void addStoredBlock(Block const& block) override; + void addStoredLifeCycle(Block const& block, LifeCycleId lifeCycle) override; + void addRemovedBlock(Digest const& blockKey) override; + void addRemovedLifeCycle(Digest const& blockKey, LifeCycleId lifeCycle) override; + void addCacheLevelUpdated( + Digest const& blockKey, CacheLevel oldLevel, CacheLevel newLevel, LifeCycleId lifeCycle) override; + +private: + void addStoredBlockUnlocked(Block const& block); + void addRemovedBlockUnlocked(Digest const& blockKey); + void addRemovedHashUnlocked(int64_t blockHash); + [[nodiscard]] bool reserveEntryUnlocked(); + [[nodiscard]] static int64_t wireHash(Digest const& digest); + + int mTokensPerBlock; + int mMaxEntries; + std::optional mTargetLifeCycle; + int mPendingEntries = 0; + std::unordered_map mStoredBlocks; + std::vector mPendingEvents; + StreamingEventStats mStats; + mutable std::mutex mMutex; +}; + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index f45743b79290..6811cfae4946 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -31,6 +31,7 @@ #include "kv_cache_manager_v2/stats.h" #include "kv_cache_manager_v2/storage/config.h" #include "kv_cache_manager_v2/storage/core.h" +#include "kv_cache_manager_v2/streamingEventSink.h" #include "kv_cache_manager_v2/utils/optionalGilRelease.h" #include @@ -1049,7 +1050,36 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) self.attentionDpRank, self.layerGroupId)); }); - nb::class_(m, "KVCacheEventManager") + nb::class_(m, "KVCacheEventSink"); + + nb::class_(m, "StreamingBlockStoredData") + .def_ro("block_hashes", &kv::StreamingBlockStoredData::blockHashes) + .def_ro("parent_block_hash", &kv::StreamingBlockStoredData::parentBlockHash) + .def_ro("token_ids", &kv::StreamingBlockStoredData::tokenIds); + + nb::class_(m, "StreamingBlockRemovedData") + .def_ro("block_hashes", &kv::StreamingBlockRemovedData::blockHashes); + + nb::class_(m, "StreamingEventStats") + .def_ro("stored_blocks", &kv::StreamingEventStats::storedBlocks) + .def_ro("removed_blocks", &kv::StreamingEventStats::removedBlocks) + .def_ro("partial_blocks_suppressed", &kv::StreamingEventStats::partialBlocksSuppressed) + .def_ro("multimodal_blocks_suppressed", &kv::StreamingEventStats::multimodalBlocksSuppressed) + .def_ro("non_target_life_cycles_ignored", &kv::StreamingEventStats::nonTargetLifeCyclesIgnored) + .def_ro("dropped_events", &kv::StreamingEventStats::droppedEvents); + + nb::class_(m, "StreamingEventSink") + .def(nb::init(), nb::arg("tokens_per_block"), nb::arg("max_entries") = 50'000) + .def( + "set_target_life_cycle", + [](kv::StreamingEventSink& self, int lifeCycleId) + { self.setTargetLifeCycle(kv::LifeCycleId{lifeCycleId}); }, + nb::arg("life_cycle_id"), nb::call_guard()) + .def("drain_iteration_events", &kv::StreamingEventSink::drainIterationEvents, + nb::call_guard()) + .def_prop_ro("stats", &kv::StreamingEventSink::getStats, nb::call_guard()); + + nb::class_(m, "KVCacheEventManager") .def( "__init__", [](kv::EventManager* self, int maxKvEventEntries, int windowSize, std::optional attentionDpRank, @@ -2032,6 +2062,26 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) [](kv::EventManager& eventManager, EventManagerTestBlock const& block, int lifeCycleId) { eventManager.addStoredLifeCycle(*block.block, kv::LifeCycleId{lifeCycleId}); }, nb::arg("event_manager"), nb::arg("block"), nb::arg("life_cycle_id"), nb::call_guard()); + mIntrospection.def( + "streaming_event_sink_add_stored_block", + [](kv::StreamingEventSink& eventSink, EventManagerTestBlock const& block) + { eventSink.addStoredBlock(*block.block); }, + nb::arg("event_sink"), nb::arg("block"), nb::call_guard()); + mIntrospection.def( + "streaming_event_sink_add_stored_life_cycle", + [](kv::StreamingEventSink& eventSink, EventManagerTestBlock const& block, int lifeCycleId) + { eventSink.addStoredLifeCycle(*block.block, kv::LifeCycleId{lifeCycleId}); }, + nb::arg("event_sink"), nb::arg("block"), nb::arg("life_cycle_id"), nb::call_guard()); + mIntrospection.def( + "streaming_event_sink_add_removed_block", + [](kv::StreamingEventSink& eventSink, EventManagerTestBlock const& block) + { eventSink.addRemovedBlock(block.block->key); }, + nb::arg("event_sink"), nb::arg("block"), nb::call_guard()); + mIntrospection.def( + "streaming_event_sink_add_removed_life_cycle", + [](kv::StreamingEventSink& eventSink, EventManagerTestBlock const& block, int lifeCycleId) + { eventSink.addRemovedLifeCycle(block.block->key, kv::LifeCycleId{lifeCycleId}); }, + nb::arg("event_sink"), nb::arg("block"), nb::arg("life_cycle_id"), nb::call_guard()); mIntrospection.def( "active_page_stats", [](kv::KvCache const& kvCache) @@ -2245,7 +2295,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) std::shared_ptr eventSink; if (!eventManager.is_none()) { - eventSink = nb::cast>(eventManager); + eventSink = nb::cast>(eventManager); } std::unique_ptr codec; diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 7ccea7387c3a..b66dd3909ba2 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -297,12 +297,12 @@ kv_cache_config = KvCacheConfig( ) ``` -**Constraints.** The streaming path requires KV cache manager V2 running on its Python -backend (`TLLM_KV_CACHE_MANAGER_V2_BACKEND=python`); the default `cpp` backend cannot -consume the Python event sink and raises an error naming this variable. Pipeline -parallelism and context parallelism are rejected. Events are not published for draft -models or during KV-cache-size estimation. When streaming is enabled the buffered pull API -returns an empty list rather than raising. +**Constraints.** The streaming path supports both KV cache manager V2 backends. With the +default `cpp` backend, a native event sink captures compact semantic event data and Python +converts it to the wire structs at the once-per-iteration flush boundary; no Python callback +runs from the native cache hot path. Pipeline parallelism and context parallelism are +rejected. Events are not published for draft models or during KV-cache-size estimation. +When streaming is enabled the buffered pull API returns an empty list rather than raising. **Endpoint convention.** Every attention-DP rank binds `base_port + rank` using its **global** rank, so `N` ranks occupy `[base_port, base_port + N - 1]` cluster-wide and 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 0774d0affc9f..74300e22202c 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 @@ -47,6 +47,7 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KVEventsConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime import kv_cache_manager_v2 as kv_cache_manager_v2_runtime from tensorrt_llm.runtime.kv_cache_hash import get_effective_kv_cache_event_hash_algo from tensorrt_llm.runtime.kv_cache_manager_v2 import ( _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, @@ -1291,8 +1292,8 @@ def __init__( "will return no events." ) assert kv_events_config is not None - # Rejects unsupported parallelism, a non-Python V2 backend and colliding - # publish/replay port ranges, all before any socket is bound. + # Reject unsupported parallelism and colliding publish/replay port ranges + # before any socket is bound. validate_streaming_support( kv_events_config, pp_size=mapping.pp_size, @@ -1306,11 +1307,19 @@ def __init__( # Constructing it is side-effect free; start() below binds the socket # and starts the publisher thread once every other check has passed. event_rank = mapping.rank if mapping.enable_attention_dp else 0 + native_event_sink = ( + kv_cache_manager_v2_runtime.StreamingEventSink( + tokens_per_block=self.tokens_per_block, + ) + if KV_CACHE_MANAGER_V2_BACKEND == "cpp" + else None + ) self.event_manager = StreamingKVCacheEventManager( kv_events_config, data_parallel_rank=event_rank, block_size=self.tokens_per_block, max_window_size=event_window_size, + native_event_sink=native_event_sink, ) elif self.event_buffer_max_size > 0: if mapping.enable_attention_dp: @@ -1525,10 +1534,15 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) candidate: Optional[KVCacheManagerPy] = None + event_sink = ( + self.event_manager.event_sink + if isinstance(self.event_manager, StreamingKVCacheEventManager) + else self.event_manager + ) if not has_host_cache_tier: candidate = KVCacheManagerPy( config, - event_manager=self.event_manager, + event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), ) else: @@ -1537,7 +1551,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: try: candidate = KVCacheManagerPy( config, - event_manager=self.event_manager, + event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), ) except Exception as error: @@ -1577,7 +1591,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) candidate = KVCacheManagerPy( config, - event_manager=self.event_manager, + event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), ) except Exception as error: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index f2d4dab1f7b7..31b2da1e4035 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -39,6 +39,7 @@ from tensorrt_llm.llmapi.llm_args import KVEventsConfig from tensorrt_llm.logger import logger +from tensorrt_llm.runtime import kv_cache_manager_v2 as kv_cache_manager_v2_runtime from tensorrt_llm.runtime.kv_cache_hash import truncate_sha256_hash_to_int64 from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import KVCacheEvent, KVCacheEventDiff @@ -412,17 +413,8 @@ def validate_streaming_support( raise ValueError("Streaming KV events do not support pipeline parallelism") if cp_size > 1: raise ValueError("Streaming KV events do not support context parallelism") - if backend != "python": - # StreamingKVCacheEventManager is a duck-typed Python event sink, which cannot - # satisfy the nanobind constructor's nb::cast> - # (and the C++ radix tree calls the sink natively, not through Python). Fail - # with an actionable message instead of an opaque TypeError from the cast. - raise ValueError( - "Streaming KV events (kv_cache_config.kv_events_config) are only supported " - f"by the Python KV cache manager V2 backend, but '{backend}' is active. Set " - "TLLM_KV_CACHE_MANAGER_V2_BACKEND=python to enable streaming KV events, or " - "use the buffered path via kv_cache_config.event_buffer_max_size." - ) + if backend not in ("cpp", "python"): + raise ValueError(f"Unsupported KV cache manager V2 backend: {backend!r}") validate_endpoint_ranges(config, ranks_per_host, data_parallel_size) @@ -515,13 +507,13 @@ class _MultimodalBlockError(ValueError): class StreamingKVCacheEventManager: - """Scheduler-local fast path that produces KV cache event wire messages directly. + """Scheduler-local facade for streaming KV event capture and publishing. - Implements the V2 KV-cache-manager event-sink hook interface by duck - typing rather than inheriting ``KVCacheEventManager``: it fully replaces - event production (reusing the radix block hashes) and shares none of the - base manager's state, so subclassing would only risk partially initialised - base attributes. + With the Python KV-cache backend this object is the duck-typed event sink. With + the C++ backend, ``native_event_sink`` captures the same semantics without a + Python callback on the cache hot path, and this facade drains its DTOs at the + once-per-iteration flush boundary. Both modes share publisher lifecycle, + batching, counters, and wire structs here. """ def __init__( @@ -532,12 +524,14 @@ def __init__( block_size: int, max_window_size: int, max_entries: int = 50_000, + native_event_sink: object | None = None, ) -> None: self._rank = data_parallel_rank self._publisher = create_event_publisher(config, data_parallel_rank) self._block_size = block_size self._max_window_size = max_window_size self._max_entries = max_entries + self._native_event_sink = native_event_sink self._target_life_cycle_id: int | None = None self._stored_blocks: dict[bytes, int] = {} self._pending_events: list[BlockStored | BlockRemoved | AllBlocksCleared] = [] @@ -582,6 +576,8 @@ def set_layer_group_window_sizes(self, window_sizes: dict[int, int]) -> None: if not target_ids: raise ValueError("Streaming KV events require an attention KV cache life cycle") self._target_life_cycle_id = min(target_ids) + if self._native_event_sink is not None: + self._native_event_sink.set_target_life_cycle(self._target_life_cycle_id) logger.info( "Streaming KV event fast path selected " f"lifecycle_id={self._target_life_cycle_id} " @@ -761,11 +757,18 @@ def _reserve_entries(self, num_entries: int) -> bool: return False def flush_iteration_events(self) -> None: - if self._closed or not self._pending_events: + if self._closed: + return + if self._native_event_sink is not None: + events = self._drain_native_events() + else: + if not self._pending_events: + return + events = self._pending_events + self._pending_events = [] + self._pending_entries = 0 + if not events: return - events = self._pending_events - self._pending_events = [] - self._pending_entries = 0 batch = KVEventBatch( ts=time.time(), events=events, @@ -784,6 +787,47 @@ def flush_iteration_events(self) -> None: f"{traceback.format_exc()}" ) + @property + def event_sink(self) -> object: + """Return the backend-specific sink installed in KVCacheManager.""" + return self if self._native_event_sink is None else self._native_event_sink + + def _drain_native_events( + self, + ) -> list[BlockStored | BlockRemoved | AllBlocksCleared]: + assert self._native_event_sink is not None + result: list[BlockStored | BlockRemoved | AllBlocksCleared] = [] + for event in self._native_event_sink.drain_iteration_events(): + if isinstance(event, kv_cache_manager_v2_runtime.StreamingBlockStoredData): + result.append( + BlockStored( + block_hashes=list(event.block_hashes), + parent_block_hash=event.parent_block_hash, + token_ids=list(event.token_ids), + block_size=self._block_size, + lora_id=None, + medium="GPU", + lora_name=None, + ) + ) + elif isinstance(event, kv_cache_manager_v2_runtime.StreamingBlockRemovedData): + result.append(BlockRemoved(block_hashes=list(event.block_hashes), medium="GPU")) + else: + raise TypeError(f"Unsupported native streaming KV event: {type(event)!r}") + self._sync_native_stats() + return result + + def _sync_native_stats(self) -> None: + if self._native_event_sink is None: + return + stats = self._native_event_sink.stats + self.stored_blocks = stats.stored_blocks + self.removed_blocks = stats.removed_blocks + self.partial_blocks_suppressed = stats.partial_blocks_suppressed + self.multimodal_blocks_suppressed = stats.multimodal_blocks_suppressed + self.non_target_life_cycles_ignored = stats.non_target_life_cycles_ignored + self.dropped_events = stats.dropped_events + def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: # Streaming publishing pushes events out-of-band, so the pull API has # nothing to return. Return empty instead of raising so callers of the @@ -796,6 +840,7 @@ def shutdown(self) -> None: self.flush_iteration_events() self._closed = True self._publisher.shutdown() + self._sync_native_stats() logger.info( "Streaming KV event fast path " f"rank={self._rank} " diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 5a5eee4ecda0..60137c94caba 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -196,6 +196,10 @@ class _BatchDescFieldSpec: KVCacheIterationStatsDelta = _cpp.KVCacheIterationStatsDelta KVCacheManager = _cpp.KVCacheManager KVCacheManagerConfig = _cpp.KVCacheManagerConfig + StreamingBlockRemovedData = _cpp.StreamingBlockRemovedData + StreamingBlockStoredData = _cpp.StreamingBlockStoredData + StreamingEventSink = _cpp.StreamingEventSink + StreamingEventStats = _cpp.StreamingEventStats IKvCacheColdPageCodec = _cpp.IKvCacheColdPageCodec create_default_kv_cache_cold_page_codec = _cpp.create_default_kv_cache_cold_page_codec # The C++ KVCacheManagerConfig binding replaces the Python @dataclass, but @@ -404,4 +408,13 @@ def typed_range(*args: int) -> range: ] if _BACKEND != "python": - __all__.extend(["IKvCacheColdPageCodec", "create_default_kv_cache_cold_page_codec"]) + __all__.extend( + [ + "IKvCacheColdPageCodec", + "StreamingBlockRemovedData", + "StreamingBlockStoredData", + "StreamingEventSink", + "StreamingEventStats", + "create_default_kv_cache_cold_page_codec", + ] + ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index e18d4f4d6bc7..4ed6075dadf4 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -172,6 +172,42 @@ def event_manager_add_stored_life_cycle(event_manager: Any, block: Any, life_cyc event_manager.add_stored_life_cycle_event_from_block(block.block, life_cycle_id) +def streaming_event_sink_add_stored_block(event_sink: Any, block: Any) -> None: + """Feed a real test block to the native streaming sink.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is None: + raise RuntimeError("the streaming event sink requires the C++ backend") + cpp_introspection.streaming_event_sink_add_stored_block(event_sink, block) + + +def streaming_event_sink_add_stored_life_cycle( + event_sink: Any, block: Any, life_cycle_id: int +) -> None: + """Feed one lifecycle of a real test block to the native streaming sink.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is None: + raise RuntimeError("the streaming event sink requires the C++ backend") + cpp_introspection.streaming_event_sink_add_stored_life_cycle(event_sink, block, life_cycle_id) + + +def streaming_event_sink_add_removed_block(event_sink: Any, block: Any) -> None: + """Remove a real test block from the native streaming sink.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is None: + raise RuntimeError("the streaming event sink requires the C++ backend") + cpp_introspection.streaming_event_sink_add_removed_block(event_sink, block) + + +def streaming_event_sink_add_removed_life_cycle( + event_sink: Any, block: Any, life_cycle_id: int +) -> None: + """Remove one lifecycle of a real test block from the native streaming sink.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is None: + raise RuntimeError("the streaming event sink requires the C++ backend") + cpp_introspection.streaming_event_sink_add_removed_life_cycle(event_sink, block, life_cycle_id) + + def active_page_stats(kv_cache: Any) -> tuple[list[int], list[int]]: """Return active pages and unscheduled evictable active pages by cache level.""" cpp_introspection = _cpp_introspection_module() 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 16854c737b9e..7157669e0809 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 @@ -23,7 +23,10 @@ import pytest +from tensorrt_llm._torch.pyexecutor.kv_cache_events import StreamingKVCacheEventManager from tensorrt_llm._utils import KVCacheEventSerializer +from tensorrt_llm.llmapi.llm_args import KVEventsConfig +from tensorrt_llm.runtime import kv_cache_manager_v2 as kv_cache_manager_v2_runtime from tensorrt_llm.runtime.kv_cache_hash import ( KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2_SHA256_64, @@ -172,6 +175,22 @@ def _add_stored_life_cycle(event_manager, block, life_cycle_id): _introspection.event_manager_add_stored_life_cycle(event_manager, block, life_cycle_id) +def _add_streaming_stored_block(event_sink, block): + _introspection.streaming_event_sink_add_stored_block(event_sink, block) + + +def _add_streaming_stored_life_cycle(event_sink, block, life_cycle_id): + _introspection.streaming_event_sink_add_stored_life_cycle(event_sink, block, life_cycle_id) + + +def _add_streaming_removed_block(event_sink, block): + _introspection.streaming_event_sink_add_removed_block(event_sink, block) + + +def _add_streaming_removed_life_cycle(event_sink, block, life_cycle_id): + _introspection.streaming_event_sink_add_removed_life_cycle(event_sink, block, life_cycle_id) + + def _token_ids(start, end): return [TokenId(token_id) for token_id in range(start, end)] @@ -261,6 +280,63 @@ def test_native_event_manager_queue_and_stored_coalescing(): ] +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="requires the native C++ streaming sink") +def test_native_streaming_sink_to_python_wire_structs(real_block_factory): + event_sink = kv_cache_manager_v2_runtime.StreamingEventSink( + tokens_per_block=2, + max_entries=8, + ) + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=2, + max_window_size=128, + native_event_sink=event_sink, + ) + manager.start() + try: + manager.set_layer_group_window_sizes({0: 128, 1: 64}) + published = [] + manager._publisher.publish = lambda batch: published.append(batch) or True + make_block = real_block_factory(event_sink, num_life_cycles=2, tokens_per_block=2) + + first = make_block(_token_ids(1, 3), [2, 2]) + partial = make_block(_token_ids(5, 7), [1, 2], parent=first) + second = make_block(_token_ids(3, 5), [2, 2], parent=first) + + _add_streaming_stored_block(event_sink, first) + _add_streaming_stored_block(event_sink, partial) + _add_streaming_stored_life_cycle(event_sink, second, 1) + _add_streaming_stored_life_cycle(event_sink, second, 0) + manager.flush_iteration_events() + + first_hash = int.from_bytes(_block_key(first)[:8], byteorder="big", signed=True) + second_hash = int.from_bytes(_block_key(second)[:8], byteorder="big", signed=True) + assert len(published) == 1 + stored = published[0].events + assert len(stored) == 1 + assert stored[0].block_hashes == [first_hash, second_hash] + assert stored[0].parent_block_hash is None + assert stored[0].token_ids == [1, 2, 3, 4] + + _add_streaming_removed_life_cycle(event_sink, second, 1) + _add_streaming_removed_block(event_sink, first) + _add_streaming_removed_life_cycle(event_sink, second, 0) + manager.flush_iteration_events() + + assert len(published) == 2 + removed = published[1].events + assert len(removed) == 1 + assert removed[0].block_hashes == [first_hash, second_hash] + assert manager.stored_blocks == 2 + assert manager.removed_blocks == 2 + assert manager.partial_blocks_suppressed == 1 + assert manager.non_target_life_cycles_ignored == 2 + assert manager.dropped_events == 0 + finally: + manager.shutdown() + + def test_native_event_manager_v1_hash_matches_legacy_cpp_hasher(): _tb = pytest.importorskip("tensorrt_llm.bindings") block_key = _tb.internal.batch_manager.BlockKey diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py index c9c25731e3b6..a69531328212 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py @@ -380,17 +380,16 @@ def test_validate_streaming_support_rejects_unsupported_setups() -> None: config = KVEventsConfig(enable_kv_cache_events=True, endpoint="tcp://*:5557") supported = dict(pp_size=1, cp_size=1, ranks_per_host=1, data_parallel_size=1, backend="python") - # The supported baseline must not raise, or the negative cases prove nothing. + # Both implementations use a backend-specific sink behind the same facade. validate_streaming_support(config, **supported) + validate_streaming_support(config, **{**supported, "backend": "cpp"}) with pytest.raises(ValueError, match="pipeline parallelism"): validate_streaming_support(config, **{**supported, "pp_size": 2}) with pytest.raises(ValueError, match="context parallelism"): validate_streaming_support(config, **{**supported, "cp_size": 2}) - # The default backend is "cpp", whose nanobind KVCacheManager cannot accept a - # duck-typed Python event sink; the error must name the env var that fixes it. - with pytest.raises(ValueError, match="TLLM_KV_CACHE_MANAGER_V2_BACKEND=python"): - validate_streaming_support(config, **{**supported, "backend": "cpp"}) + with pytest.raises(ValueError, match="Unsupported KV cache manager V2 backend"): + validate_streaming_support(config, **{**supported, "backend": "invalid"}) @pytest.mark.parametrize( From 63db7b7a7f412abca38a0341072f14c775e28269 Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Wed, 16 Sep 2026 02:39:53 -0700 Subject: [PATCH 2/3] [None][fix] select streaming KV lifecycle with C++ backend Signed-off-by: Guan Luo --- .../kv_cache/kv_cache_manager_v2.py | 9 ++++--- .../kv_cache/test_kv_cache_manager_v2.py | 24 +++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) 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 74300e22202c..da747701bafc 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 @@ -71,7 +71,6 @@ KVCacheEventManager, KVCacheIterationStatsDelta, LayerId, - LifeCycleId, PageIndexMode, PlannedDropHandle, PoolGroupPeakBlockStats, @@ -2447,9 +2446,13 @@ def get_event_window_size(layer_id: int) -> int: window_sizes: Dict[int, int] = {} for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping): + layer_config = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] if attention_only: - life_cycle = self.impl._life_cycles.get_life_cycle(LifeCycleId(layer_group_id)) - if not isinstance(life_cycle, AttnLifeCycle): + # The Python implementation exposes its internal life-cycle + # registry, but the C++ backend deliberately does not. The + # public layer configuration carries the same distinction and + # keeps this selection backend-independent. + if not isinstance(layer_config, AttentionLayerConfig): continue window_sizes[int(layer_group_id)] = get_event_window_size(int(layer_ids[0])) return window_sizes diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index 55c917fd6f9e..0e26fe86f291 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py @@ -533,6 +533,30 @@ def test_zero_size_filter_rejects_empty_local_cache() -> None: manager._remove_zero_size_buffers(config) +def test_event_window_sizes_filter_attention_without_backend_internals() -> None: + manager = object.__new__(KVCacheManagerV2) + manager.max_seq_len = MAX_SEQ_LEN + manager.kv_cache_manager_py_config = SimpleNamespace( + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=Role.KEY, size=128)], + sliding_window_size=8, + ), + SsmLayerConfig( + layer_id=LayerId(1), + buffers=[BufferConfig(role=DataRole("ssm_state"), size=128)], + ), + ] + ) + # Match the C++ binding surface: layer_grouping is public, while the + # Python implementation's private _life_cycles registry is absent. + manager.impl = SimpleNamespace(layer_grouping=[[0], [1]]) + + assert manager._get_event_window_sizes_by_layer_group(attention_only=True) == {0: 8} + assert manager._get_event_window_sizes_by_layer_group() == {0: 8, 1: MAX_SEQ_LEN} + + def test_draft_token_relocation_uses_local_cache_layout(monkeypatch: pytest.MonkeyPatch) -> None: request = SimpleNamespace( state=LlmRequestState.GENERATION_IN_PROGRESS, From 3e0a0df65be84c8b7b6dc86e3bf2704ab0375dba Mon Sep 17 00:00:00 2001 From: Guan Luo Date: Mon, 21 Sep 2026 01:43:58 -0700 Subject: [PATCH 3/3] [None][fix] preserve multimodal streaming KV events Signed-off-by: Guan Luo --- .../kv_cache_manager_v2/CMakeLists.txt | 1 + .../kv_cache_manager_v2/eventData.cpp | 90 +++++++++++++++ .../kv_cache_manager_v2/eventData.h | 61 ++++++++++ .../kv_cache_manager_v2/eventManager.cpp | 63 +---------- .../kv_cache_manager_v2/eventManager.h | 17 +-- .../streamingEventSink.cpp | 32 +++--- .../kv_cache_manager_v2/streamingEventSink.h | 14 ++- .../batch_manager/kvCacheManagerV2.cpp | 26 ++++- docs/source/features/kvcache.md | 6 + .../kv_cache/kv_cache_manager_v2.py | 2 + .../_torch/pyexecutor/kv_cache_events.py | 105 ++++++++++++------ .../runtime/kv_cache_manager_v2/__init__.pyi | 40 +++++++ .../test_kv_cache_event_manager.py | 66 +++++++++++ .../test_streaming_kv_events.py | 74 ++++++++++++ 14 files changed, 467 insertions(+), 130 deletions(-) create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventData.cpp create mode 100644 cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventData.h diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt index c52bbd86734c..eb0bcb8fa8c1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt @@ -56,6 +56,7 @@ set(KV_CACHE_MANAGER_V2_SRCS kv_cache_manager_v2/blockRadixTree.cpp kv_cache_manager_v2/batchedPageCopy.cu kv_cache_manager_v2/coldPageCodec.cpp + kv_cache_manager_v2/eventData.cpp kv_cache_manager_v2/eventManager.cpp kv_cache_manager_v2/streamingEventSink.cpp kv_cache_manager_v2/page.cpp diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventData.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventData.cpp new file mode 100644 index 000000000000..d3667b64236f --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventData.cpp @@ -0,0 +1,90 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_cache_manager_v2/eventData.h" + +#include "kv_cache_manager_v2/blockRadixTree.h" + +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +std::string digestToHex(Digest const& digest) +{ + constexpr char kHex[] = "0123456789abcdef"; + std::string result; + result.resize(digest.size() * 2); + for (size_t i = 0; i < digest.size(); ++i) + { + auto const value = std::to_integer(digest[i]); + result[2 * i] = kHex[value >> 4U]; + result[2 * i + 1] = kHex[value & 0x0FU]; + } + return result; +} + +DecodedEventBlock decodeEventBlock(Block const& block, std::optional mmTokenIdOffset) +{ + DecodedEventBlock result; + result.tokenIds.reserve(block.tokens.size()); + + Digest const* itemDigest = nullptr; + if (mmTokenIdOffset.has_value() && block.prev != nullptr && block.prev->type() == NodeBase::Type::kBLOCK) + { + itemDigest = static_cast(block.prev)->getLastTokenDigest().get(); + } + bool inMmRun = false; + for (auto const& token : block.tokens) + { + if (token.isDigest()) + { + result.tokenIds.emplace_back(std::in_place_index<1>, digestToHex(token.digest())); + if (mmTokenIdOffset.has_value()) + { + itemDigest = &token.digest(); + result.mmKeys.push_back( + {std::string(reinterpret_cast(itemDigest->data()), itemDigest->size()), 0, + std::nullopt, false}); + inMmRun = true; + } + continue; + } + + auto const tokenId = token.tokenId(); + result.tokenIds.emplace_back(std::in_place_index<0>, tokenId); + if (itemDigest != nullptr && tokenId > *mmTokenIdOffset) + { + if (!inMmRun) + { + result.mmKeys.push_back( + {std::string(reinterpret_cast(itemDigest->data()), itemDigest->size()), + tokenId - *mmTokenIdOffset, std::nullopt, false}); + } + inMmRun = true; + } + else + { + // Text separates runs of the same item, so retain its digest for later continuations. + inMmRun = false; + } + } + return result; +} + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventData.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventData.h new file mode 100644 index 000000000000..f4f1f13ba0c1 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventData.h @@ -0,0 +1,61 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "kv_cache_manager_v2/common.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 +{ + +class Block; + +using EventTokenId = std::variant; + +struct MmKey +{ + std::string hash; + int startOffset = 0; + std::optional uuid; + bool hasUuidField = false; + + bool operator==(MmKey const& other) const + { + return hash == other.hash && startOffset == other.startOffset && uuid == other.uuid + && hasUuidField == other.hasUuidField; + } +}; + +struct DecodedEventBlock +{ + std::vector tokenIds; + std::vector mmKeys; +}; + +[[nodiscard]] std::string digestToHex(Digest const& digest); + +//! Decode the V2 digest-first multimodal representation used by KV event consumers. +//! When mmTokenIdOffset is absent, tokens are still preserved but no MM segments are derived. +[[nodiscard]] DecodedEventBlock decodeEventBlock(Block const& block, std::optional mmTokenIdOffset); + +} // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp index 8ca96b4edde8..bd7067080b1d 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp @@ -485,20 +485,6 @@ int EventManager::getWindowSize(EventLayerGroupId layerGroupId) const return windowSize == mWindowSizeByLayerGroup.end() ? mWindowSize : windowSize->second; } -std::string EventManager::digestToHex(Digest const& digest) -{ - constexpr char kHex[] = "0123456789abcdef"; - std::string result; - result.resize(digest.size() * 2); - for (size_t i = 0; i < digest.size(); ++i) - { - auto const value = std::to_integer(digest[i]); - result[2 * i] = kHex[value >> 4U]; - result[2 * i + 1] = kHex[value & 0x0FU]; - } - return result; -} - uint64_t EventManager::truncateDigestToInt64(Digest const& digest) { uint64_t result = 0; @@ -566,54 +552,15 @@ std::optional EventManager::storedBlockFromBlock( return std::nullopt; } - std::vector mmKeys; - Digest const* itemDigest = nullptr; - if (mMmTokenIdOffset.has_value() && block.prev != nullptr && block.prev->type() == NodeBase::Type::kBLOCK) - { - itemDigest = static_cast(block.prev)->getLastTokenDigest().get(); - } - bool inMmRun = false; + auto decoded = decodeEventBlock(block, mMmTokenIdOffset); std::vector tokens; - tokens.reserve(block.tokens.size()); - for (auto const& token : block.tokens) + tokens.reserve(decoded.tokenIds.size()); + for (auto& tokenId : decoded.tokenIds) { - if (!token.isDigest()) - { - UniqueToken uniqueToken; - uniqueToken.tokenId = EventTokenId{std::in_place_index<0>, token.tokenId()}; - tokens.push_back(std::move(uniqueToken)); - if (itemDigest != nullptr && token.tokenId() > *mMmTokenIdOffset) - { - if (!inMmRun) - { - mmKeys.push_back( - {std::string(reinterpret_cast(itemDigest->data()), itemDigest->size()), - token.tokenId() - *mMmTokenIdOffset, std::nullopt, false}); - } - inMmRun = true; - } - else - { - // Text separates runs of the same item, so retain its digest for later continuations. - inMmRun = false; - } - } - else - { - UniqueToken uniqueToken; - uniqueToken.tokenId = EventTokenId{std::in_place_index<1>, digestToHex(token.digest())}; - tokens.push_back(std::move(uniqueToken)); - if (mMmTokenIdOffset.has_value()) - { - itemDigest = &token.digest(); - mmKeys.push_back({std::string(reinterpret_cast(itemDigest->data()), itemDigest->size()), 0, - std::nullopt, false}); - inMmRun = true; - } - } + tokens.push_back(UniqueToken{std::move(tokenId)}); } return KVCacheStoredBlockData{ - hashFromBlock(block), std::move(tokens), cacheLevel.value(), priority, std::move(mmKeys), std::nullopt}; + hashFromBlock(block), std::move(tokens), cacheLevel.value(), priority, std::move(decoded.mmKeys), std::nullopt}; } uint64_t EventManager::hashV1BlockKey(std::vector const& tokens, uint64_t parentHash, diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h index c1ca4ae83450..351120c5c47b 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h @@ -18,6 +18,7 @@ #pragma once #include "kv_cache_manager_v2/common.h" +#include "kv_cache_manager_v2/eventData.h" #include "kv_cache_manager_v2/eventSink.h" #include @@ -40,7 +41,6 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { using EventBlockHash = std::variant; -using EventTokenId = std::variant; using EventLayerGroupId = std::optional; struct UniqueToken @@ -64,20 +64,6 @@ struct KVCacheCreatedData } }; -struct MmKey -{ - std::string hash; - int startOffset = 0; - std::optional uuid; - bool hasUuidField = false; - - bool operator==(MmKey const& other) const - { - return hash == other.hash && startOffset == other.startOffset && uuid == other.uuid - && hasUuidField == other.hasUuidField; - } -}; - struct KVCacheStoredBlockData { EventBlockHash blockHash; @@ -220,7 +206,6 @@ class EventManager final : public EventSink using V1RootAttrs = std::pair, std::optional>; static std::pair parseHashAlgorithm(std::string const& hashAlgo); - static std::string digestToHex(Digest const& digest); static uint64_t truncateDigestToInt64(Digest const& digest); static std::vector trimEvents(std::vector events, int maxKvEventEntries); 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 c98cdd8b5b8b..90fd70ba659e 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 @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -30,9 +31,10 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { -StreamingEventSink::StreamingEventSink(int tokensPerBlock, int maxEntries) +StreamingEventSink::StreamingEventSink(int tokensPerBlock, int maxEntries, std::optional mmTokenIdOffset) : mTokensPerBlock(tokensPerBlock) , mMaxEntries(maxEntries) + , mMmTokenIdOffset(mmTokenIdOffset) { if (mTokensPerBlock <= 0) { @@ -42,6 +44,10 @@ StreamingEventSink::StreamingEventSink(int tokensPerBlock, int maxEntries) { throw std::invalid_argument("maxEntries must be positive"); } + if (mMmTokenIdOffset.has_value() && *mMmTokenIdOffset < 0) + { + throw std::invalid_argument("mmTokenIdOffset must be non-negative"); + } } void StreamingEventSink::setTargetLifeCycle(LifeCycleId lifeCycle) @@ -132,14 +138,6 @@ void StreamingEventSink::addStoredBlockUnlocked(Block const& block) { return; } - for (auto const& token : block.tokens) - { - if (token.isDigest()) - { - ++mStats.multimodalBlocksSuppressed; - return; - } - } if (!reserveEntryUnlocked()) { return; @@ -156,12 +154,7 @@ void StreamingEventSink::addStoredBlockUnlocked(Block const& block) parentHash = wireHash(static_cast(block.prev)->key); } - std::vector tokenIds; - tokenIds.reserve(block.tokens.size()); - for (auto const& token : block.tokens) - { - tokenIds.push_back(token.tokenId()); - } + auto decoded = decodeEventBlock(block, mMmTokenIdOffset); mStoredBlocks.emplace(block.key, blockHash); if (!mPendingEvents.empty()) @@ -171,12 +164,17 @@ void StreamingEventSink::addStoredBlockUnlocked(Block const& block) && stored->blockHashes.back() == *parentHash) { stored->blockHashes.push_back(blockHash); - stored->tokenIds.insert(stored->tokenIds.end(), tokenIds.begin(), tokenIds.end()); + stored->tokenIds.insert(stored->tokenIds.end(), std::make_move_iterator(decoded.tokenIds.begin()), + std::make_move_iterator(decoded.tokenIds.end())); + stored->mmKeys.push_back(std::move(decoded.mmKeys)); ++mStats.storedBlocks; return; } } - mPendingEvents.emplace_back(StreamingBlockStoredData{{blockHash}, parentHash, std::move(tokenIds)}); + std::vector> mmKeys; + mmKeys.push_back(std::move(decoded.mmKeys)); + mPendingEvents.emplace_back( + StreamingBlockStoredData{{blockHash}, parentHash, std::move(decoded.tokenIds), std::move(mmKeys)}); ++mStats.storedBlocks; } 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 22e314314d0a..784aa10a4e1f 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 @@ -17,6 +17,7 @@ #pragma once +#include "kv_cache_manager_v2/eventData.h" #include "kv_cache_manager_v2/eventSink.h" #include @@ -34,7 +35,9 @@ struct StreamingBlockStoredData { std::vector blockHashes; std::optional parentBlockHash; - std::vector tokenIds; + std::vector tokenIds; + //! One entry per blockHash. Empty block entries represent text-only blocks. + std::vector> mmKeys; }; //! Semantic data for one wire-level BlockRemoved event. @@ -51,7 +54,6 @@ struct StreamingEventStats int64_t storedBlocks = 0; int64_t removedBlocks = 0; int64_t partialBlocksSuppressed = 0; - int64_t multimodalBlocksSuppressed = 0; int64_t nonTargetLifeCyclesIgnored = 0; int64_t droppedEvents = 0; }; @@ -60,7 +62,12 @@ struct StreamingEventStats class StreamingEventSink final : public EventSink { public: - StreamingEventSink(int tokensPerBlock, int maxEntries); + StreamingEventSink(int tokensPerBlock, int maxEntries, std::optional mmTokenIdOffset = std::nullopt); + + bool needsTokenDigestContext() const override + { + return mMmTokenIdOffset.has_value(); + } void setTargetLifeCycle(LifeCycleId lifeCycle); [[nodiscard]] std::vector drainIterationEvents(); @@ -82,6 +89,7 @@ class StreamingEventSink final : public EventSink int mTokensPerBlock; int mMaxEntries; + std::optional mMmTokenIdOffset; std::optional mTargetLifeCycle; int mPendingEntries = 0; std::unordered_map mStoredBlocks; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 6811cfae4946..3f0a135638ed 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -461,10 +461,10 @@ static std::vector castMmKeys(nb::handle values) return result; } -static nb::list castMmKeys(kv::KVCacheStoredBlockData const& data) +static nb::list castMmKeys(std::vector const& mmKeys) { nb::list result; - for (auto const& mmKey : data.mmKeys) + for (auto const& mmKey : mmKeys) { auto hash = nb::bytes(mmKey.hash.data(), mmKey.hash.size()); if (mmKey.hasUuidField) @@ -479,6 +479,21 @@ static nb::list castMmKeys(kv::KVCacheStoredBlockData const& data) return result; } +static nb::list castMmKeys(kv::KVCacheStoredBlockData const& data) +{ + return castMmKeys(data.mmKeys); +} + +static nb::list castStreamingMmKeys(kv::StreamingBlockStoredData const& data) +{ + nb::list result; + for (auto const& mmKeys : data.mmKeys) + { + result.append(castMmKeys(mmKeys)); + } + return result; +} + static nb::object castEventData(kv::KVCacheEventData const& data) { return std::visit([](auto const& concreteData) { return nb::cast(concreteData); }, data); @@ -1055,7 +1070,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "StreamingBlockStoredData") .def_ro("block_hashes", &kv::StreamingBlockStoredData::blockHashes) .def_ro("parent_block_hash", &kv::StreamingBlockStoredData::parentBlockHash) - .def_ro("token_ids", &kv::StreamingBlockStoredData::tokenIds); + .def_ro("token_ids", &kv::StreamingBlockStoredData::tokenIds) + .def_prop_ro("mm_keys", [](kv::StreamingBlockStoredData const& self) { return castStreamingMmKeys(self); }); nb::class_(m, "StreamingBlockRemovedData") .def_ro("block_hashes", &kv::StreamingBlockRemovedData::blockHashes); @@ -1064,12 +1080,12 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_ro("stored_blocks", &kv::StreamingEventStats::storedBlocks) .def_ro("removed_blocks", &kv::StreamingEventStats::removedBlocks) .def_ro("partial_blocks_suppressed", &kv::StreamingEventStats::partialBlocksSuppressed) - .def_ro("multimodal_blocks_suppressed", &kv::StreamingEventStats::multimodalBlocksSuppressed) .def_ro("non_target_life_cycles_ignored", &kv::StreamingEventStats::nonTargetLifeCyclesIgnored) .def_ro("dropped_events", &kv::StreamingEventStats::droppedEvents); nb::class_(m, "StreamingEventSink") - .def(nb::init(), nb::arg("tokens_per_block"), nb::arg("max_entries") = 50'000) + .def(nb::init>(), nb::arg("tokens_per_block"), nb::arg("max_entries") = 50'000, + nb::arg("mm_token_id_offset") = std::nullopt) .def( "set_target_life_cycle", [](kv::StreamingEventSink& self, int lifeCycleId) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index b66dd3909ba2..e0dc2bd7c904 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -304,6 +304,12 @@ runs from the native cache hot path. Pipeline parallelism and context parallelis rejected. Events are not published for draft models or during KV-cache-size estimation. When streaming is enabled the buffered pull API returns an empty list rather than raising. +For V2 multimodal prefixes, streaming uses the same digest-first representation as the +buffered path. A digest token is emitted as a hexadecimal string in `token_ids`, and +`mm_keys` is aligned one-for-one with `block_hashes`; each nested list contains that block's +multimodal segments using the `hash` and `start_offset` semantics described above. Consumers +must normalize these typed tokens before applying their ordinary token hashing logic. + **Endpoint convention.** Every attention-DP rank binds `base_port + rank` using its **global** rank, so `N` ranks occupy `[base_port, base_port + N - 1]` cluster-wide and each rank's port is distinct — on a multi-node deployment, rank 8 binds `base_port + 8` 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 da747701bafc..2ca3d29331fa 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 @@ -1309,6 +1309,7 @@ def __init__( native_event_sink = ( kv_cache_manager_v2_runtime.StreamingEventSink( tokens_per_block=self.tokens_per_block, + mm_token_id_offset=vocab_size, ) if KV_CACHE_MANAGER_V2_BACKEND == "cpp" else None @@ -1318,6 +1319,7 @@ def __init__( data_parallel_rank=event_rank, block_size=self.tokens_per_block, max_window_size=event_window_size, + mm_token_id_offset=vocab_size, native_event_sink=native_event_sink, ) elif self.event_buffer_max_size > 0: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 31b2da1e4035..ee9a6aedb76d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -46,6 +46,7 @@ # Subscribers decode block hashes as 64-bit ints, so a bytes value would fail the # decode for the entire batch. ExternalBlockHash = int +EventTokenId = int | str class EventBatch( @@ -70,12 +71,24 @@ class KVCacheWireEvent( """Base class for KV cache event wire messages.""" +class MultimodalKey( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] + tag="mm_key", +): + """One continuous multimodal segment within a stored block.""" + + hash: str + start_offset: int + + class BlockStored(KVCacheWireEvent): """A sequence of full KV cache blocks was stored.""" block_hashes: list[ExternalBlockHash] parent_block_hash: ExternalBlockHash | None - token_ids: list[int] + token_ids: list[EventTokenId] block_size: int lora_id: int | None medium: str | None @@ -85,6 +98,8 @@ class BlockStored(KVCacheWireEvent): kv_cache_spec_kind: str | None = None kv_cache_spec_sliding_window: int | None = None locality: str | None = None + # Aligned one-for-one with block_hashes when multimodal decoding is enabled. + mm_keys: list[list[MultimodalKey]] | None = None class BlockRemoved(KVCacheWireEvent): @@ -497,15 +512,6 @@ def _kv_event_wire_hash_from_radix_key(block_key: bytes) -> int: return unsigned_hash - 2**64 if unsigned_hash >= 2**63 else unsigned_hash -class _MultimodalBlockError(ValueError): - """A block token is a multimodal cache-key digest (bytes), not a wire int. - - ``gen_multimodal_cache_key_tokens`` stores the per-item digest as ``bytes``, - which has no integer wire representation. Such blocks are skipped - quietly rather than routed through the malformed-data traceback path. - """ - - class StreamingKVCacheEventManager: """Scheduler-local facade for streaming KV event capture and publishing. @@ -524,6 +530,7 @@ def __init__( block_size: int, max_window_size: int, max_entries: int = 50_000, + mm_token_id_offset: int | None = None, native_event_sink: object | None = None, ) -> None: self._rank = data_parallel_rank @@ -531,6 +538,9 @@ def __init__( self._block_size = block_size self._max_window_size = max_window_size self._max_entries = max_entries + if mm_token_id_offset is not None and mm_token_id_offset < 0: + raise ValueError("mm_token_id_offset must be non-negative") + self._mm_token_id_offset = mm_token_id_offset self._native_event_sink = native_event_sink self._target_life_cycle_id: int | None = None self._stored_blocks: dict[bytes, int] = {} @@ -540,7 +550,6 @@ def __init__( self.stored_blocks = 0 self.removed_blocks = 0 self.partial_blocks_suppressed = 0 - self.multimodal_blocks_suppressed = 0 self.non_target_life_cycles_ignored = 0 self.dropped_events = 0 self.enqueued_batches = 0 @@ -548,8 +557,7 @@ def __init__( self.dropped_batches = 0 def needs_token_digest_context(self) -> bool: - # Streaming events do not emit multimodal keys. - return False + return self._mm_token_id_offset is not None def start(self) -> None: """Bind the publisher's sockets and start its background thread. @@ -633,14 +641,8 @@ def _add_full_block(self, block: Any) -> None: if not self._reserve_entries(1): return try: - token_ids = self._token_ids(block.tokens) + token_ids, mm_keys = self._decode_block(block) block_hash, parent_hash = self._block_hashes(block) - except _MultimodalBlockError: - # Expected for multimodal cache-key blocks; skip without the - # malformed-data traceback that would otherwise flood the log. - self.multimodal_blocks_suppressed += 1 - self._pending_entries -= 1 - return except ValueError: self.dropped_events += 1 self._pending_entries -= 1 @@ -655,6 +657,8 @@ def _add_full_block(self, block: Any) -> None: if previous.block_hashes and previous.block_hashes[-1] == parent_hash: previous.block_hashes.append(block_hash) previous.token_ids.extend(token_ids) + if previous.mm_keys is not None: + previous.mm_keys.append(mm_keys) self.stored_blocks += 1 return self._pending_events.append( @@ -666,21 +670,50 @@ def _add_full_block(self, block: Any) -> None: lora_id=None, medium="GPU", lora_name=None, + mm_keys=[mm_keys] if self.needs_token_digest_context() else None, ) ) self.stored_blocks += 1 - @staticmethod - def _token_ids(tokens: Any) -> list[int]: - token_ids: list[int] = [] - for token in tokens: + def _decode_block(self, block: Any) -> tuple[list[EventTokenId], list[MultimodalKey]]: + """Decode the same digest-first V2 event representation as the buffered path.""" + parent = block.prev + digest = ( + getattr(parent, "last_token_digest", None) + if getattr(parent, "ordinal", -1) >= 0 + else None + ) + in_mm_run = False + token_ids: list[EventTokenId] = [] + mm_keys: list[MultimodalKey] = [] + for token in block.tokens: if type(token) is bytes: - # Multimodal cache-key digest; not representable as a wire int. - raise _MultimodalBlockError - if type(token) is not int: - raise ValueError("KV cache event wire format requires integer token IDs") - token_ids.append(token) - return token_ids + digest = token + token_ids.append(token.hex()) + if self.needs_token_digest_context(): + mm_keys.append(MultimodalKey(hash=token.hex(), start_offset=0)) + in_mm_run = True + elif type(token) is int: + token_ids.append(token) + if ( + self._mm_token_id_offset is not None + and digest is not None + and token > self._mm_token_id_offset + ): + if not in_mm_run: + mm_keys.append( + MultimodalKey( + hash=bytes(digest).hex(), + start_offset=token - self._mm_token_id_offset, + ) + ) + in_mm_run = True + else: + # Text separates runs of the same item, but not its inherited context. + in_mm_run = False + else: + raise ValueError("KV cache event tokens must be int or digest bytes") + return token_ids, mm_keys def _block_hashes( self, @@ -808,6 +841,17 @@ def _drain_native_events( lora_id=None, medium="GPU", lora_name=None, + mm_keys=( + [ + [ + MultimodalKey(hash=bytes(key[0]).hex(), start_offset=key[1]) + for key in block_keys + ] + for block_keys in event.mm_keys + ] + if self.needs_token_digest_context() + else None + ), ) ) elif isinstance(event, kv_cache_manager_v2_runtime.StreamingBlockRemovedData): @@ -824,7 +868,6 @@ def _sync_native_stats(self) -> None: self.stored_blocks = stats.stored_blocks self.removed_blocks = stats.removed_blocks self.partial_blocks_suppressed = stats.partial_blocks_suppressed - self.multimodal_blocks_suppressed = stats.multimodal_blocks_suppressed self.non_target_life_cycles_ignored = stats.non_target_life_cycles_ignored self.dropped_events = stats.dropped_events diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 91ccd9142528..96d54e053748 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -295,6 +295,46 @@ class KVCacheEvent: attention_dp_rank: int | None = None layer_group_id: int | None = None +class StreamingBlockStoredData: + @property + def block_hashes(self) -> list[int]: ... + @property + def parent_block_hash(self) -> int | None: ... + @property + def token_ids(self) -> list[EventTokenId]: ... + @property + def mm_keys(self) -> list[list[MmKey]]: ... + +class StreamingBlockRemovedData: + @property + def block_hashes(self) -> list[int]: ... + +class StreamingEventStats: + @property + def stored_blocks(self) -> int: ... + @property + def removed_blocks(self) -> int: ... + @property + def partial_blocks_suppressed(self) -> int: ... + @property + def non_target_life_cycles_ignored(self) -> int: ... + @property + def dropped_events(self) -> int: ... + +class StreamingEventSink: + def __init__( + self, + tokens_per_block: int, + max_entries: int = ..., + mm_token_id_offset: int | None = None, + ) -> None: ... + def set_target_life_cycle(self, life_cycle_id: int) -> None: ... + def drain_iteration_events( + self, + ) -> list[StreamingBlockStoredData | StreamingBlockRemovedData]: ... + @property + def stats(self) -> StreamingEventStats: ... + class KVCacheEventManager: def __init__( self, 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 7157669e0809..2a0ea1e971e3 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 @@ -337,6 +337,72 @@ def test_native_streaming_sink_to_python_wire_structs(real_block_factory): manager.shutdown() +@pytest.mark.skipif(not _USING_CPP_BACKEND, reason="requires the native C++ streaming sink") +def test_native_streaming_sink_preserves_multimodal_event_data(real_block_factory): + event_sink = kv_cache_manager_v2_runtime.StreamingEventSink( + tokens_per_block=4, + max_entries=8, + mm_token_id_offset=1000, + ) + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + mm_token_id_offset=1000, + native_event_sink=event_sink, + ) + manager.start() + try: + manager.set_layer_group_window_sizes({0: 128}) + published = [] + manager._publisher.publish = lambda batch: published.append(batch) or True + make_block = real_block_factory(event_sink, tokens_per_block=4) + digest_a = bytes(range(32)) + digest_b = bytes(reversed(range(32))) + first = make_block([1, digest_a, 1001, 1002], [4]) + gap = make_block([2, 3, 4, 5], [4], parent=first) + continued = make_block([1003, 7, 1004, 1005], [4], parent=gap) + last = make_block([1006, digest_b, 1001, 9], [4], parent=continued) + + for block in (first, gap, continued, last): + _add_streaming_stored_block(event_sink, block) + manager.flush_iteration_events() + + assert len(published) == 1 + assert len(published[0].events) == 1 + stored = published[0].events[0] + assert stored.token_ids == [ + 1, + digest_a.hex(), + 1001, + 1002, + 2, + 3, + 4, + 5, + 1003, + 7, + 1004, + 1005, + 1006, + digest_b.hex(), + 1001, + 9, + ] + assert [ + [(key.hash, key.start_offset) for key in block_keys] for block_keys in stored.mm_keys + ] == [ + [(digest_a.hex(), 0)], + [], + [(digest_a.hex(), 3), (digest_a.hex(), 4)], + [(digest_a.hex(), 6), (digest_b.hex(), 0)], + ] + assert manager.stored_blocks == 4 + finally: + manager.shutdown() + + def test_native_event_manager_v1_hash_matches_legacy_cpp_hasher(): _tb = pytest.importorskip("tensorrt_llm.bindings") block_key = _tb.internal.batch_manager.BlockKey diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py index a69531328212..f1e5258bc6ce 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py @@ -139,6 +139,80 @@ def test_streaming_sink_supports_real_radix_blocks(monkeypatch: pytest.MonkeyPat manager.shutdown() +def test_streaming_sink_emits_multimodal_keys_across_blocks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Streaming uses the buffered V2 digest and continuation-token contract.""" + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + mm_token_id_offset=1000, + ) + life_cycles = LifeCycleRegistry( + KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[GpuCacheTierConfig(quota=4096)], + layers=[], + ) + ) + tree = BlockRadixTree(life_cycles, tokens_per_block=4, event_manager=manager) + published: list[KVEventBatch] = [] + monkeypatch.setattr( + manager._publisher, "publish", lambda batch: published.append(batch) or True + ) + try: + digest_a = bytes(range(32)) + digest_b = bytes(reversed(range(32))) + root = tree.add_or_get_existing(ReuseScope()) + first = Block([1, digest_a, 1001, 1002], root) + gap = Block([2, 3, 4, 5], first) + continued = Block([1003, 7, 1004, 1005], gap) + last = Block([1006, digest_b, 1001, 9], continued) + + for block in (first, gap, continued, last): + manager._add_full_block(block) + manager.flush_iteration_events() + + decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(published[0])) + stored = decoded[1][0] + assert stored["token_ids"] == [ + 1, + digest_a.hex(), + 1001, + 1002, + 2, + 3, + 4, + 5, + 1003, + 7, + 1004, + 1005, + 1006, + digest_b.hex(), + 1001, + 9, + ] + assert stored["mm_keys"] == [ + [{"type": "mm_key", "hash": digest_a.hex(), "start_offset": 0}], + [], + [ + {"type": "mm_key", "hash": digest_a.hex(), "start_offset": 3}, + {"type": "mm_key", "hash": digest_a.hex(), "start_offset": 4}, + ], + [ + {"type": "mm_key", "hash": digest_a.hex(), "start_offset": 6}, + {"type": "mm_key", "hash": digest_b.hex(), "start_offset": 0}, + ], + ] + assert manager.stored_blocks == 4 + finally: + tree.clear() + manager.shutdown() + + def test_streaming_fast_path_publishes_only_full_max_window_blocks() -> None: """Protect radix hash reuse, filtering, wire format, and shutdown.""" topic = "kv-events"