diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index cdbc908477c7..c947895f3102 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -214,6 +214,76 @@ The property ```copy_on_partial_reuse``` specifies whether a block should be cop Property ```max_attention_window``` specifies the maximum attention window size for each layer in the model as a list of integer values. If the length of this list is less than number of layers, the list is repeated as many times as necessary. For instance, if the model has only full attention layers and maximum sequence length is 4096, you can specify this as ```max_attention_window = [4096]```. If the first layer is full attention, the second layer is limited attention with window size 256 and then this repeats for the remaining layers, you specify this as ```max_attention_window = [4096,256]```. This means first layer is full attention, second layer is limited attention, third layer is full attention, fourth layer is limited attention and so on. +### KV Cache Events + +KV cache events report block **stored**, **removed**, **created** and **updated** operations +so an external KV-cache-aware router (for example NVIDIA Dynamo) can route a request to the +engine that already holds its prefix. Two delivery paths are available. + +#### Buffered path (default) + +Set ```event_buffer_max_size``` to a positive integer and ```enable_block_reuse``` to True. +Events are buffered per rank, gathered onto rank 0 under attention data parallelism, and +pulled per iteration through `LLM.get_kv_cache_events()` / `LLM.get_kv_cache_events_async()`, +or over the `/kv_cache_events` endpoint of `trtllm-serve`. + +#### Streaming path (prototype) + +Configured with ```kv_cache_config.kv_events_config```. Each rank encodes its own events and +publishes them directly over a ZeroMQ `PUB` socket from a background thread, so there is no +rank-0 gather and no per-iteration pull. + +```python +from tensorrt_llm.llmapi import KvCacheConfig, KVEventsConfig + +kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + kv_events_config=KVEventsConfig( + enable_kv_cache_events=True, + endpoint="tcp://*:5557", + replay_endpoint="tcp://*:5657", + ), +) +``` + +**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. + +**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` +whichever node it runs on. Co-located engines — for example disaggregated prefill and +decode on one host — must use base ports at least `N` apart. + +```replay_endpoint``` follows the same convention. Because only ranks co-located on one +host actually contend for a port, and a host holds a contiguous run of ranks, its base +port must be at least *ranks-per-host* away from ```endpoint```'s rather than `N` away. +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`, `BlockRemoved` or `AllBlocksCleared` — 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. + +**Delivery guarantees.** Delivery is best effort, but loss is observable. Every accepted +batch reserves a sequence number up front, so a batch dropped by a full publisher queue +(```max_queue_size```) or by a failed send leaves a hole in the sequence. Subscribers must +treat any gap as lost KV-cache state and resynchronize rather than assuming continuity. + +**Replay.** If ```replay_endpoint``` is set, the publisher also binds a `ROUTER` socket. A +subscriber sends an empty delimiter frame plus an 8-byte big-endian start sequence, and +receives each retained batch as `[delimiter, topic, seq, payload]`, terminated by a sentinel +with an empty payload. Only the last ```buffer_steps``` batches are retained, so a replay +can legitimately start above the requested sequence — that too is a gap. + ### Deprecated Properties Property ```use_uvm``` has been deprecated and will be removed in a future release. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 084da74711ed..0114c575e310 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -29,7 +29,7 @@ # isort: off from tensorrt_llm.llmapi.llm_args import ( CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, - KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, + KVEventsConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, MultimodalEncoderSchedulingPolicy, PeftCacheConfig, SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy) # isort: on @@ -1398,6 +1398,9 @@ def _create_kv_cache_manager( execution_stream=self._execution_stream, layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, + kv_events_config=None + if estimating_kv_cache or model_engine.is_draft_model else + self._llm_args.kv_cache_config.kv_events_config, cold_page_codec_provider=cold_page_codec_provider, ) @@ -2242,7 +2245,8 @@ def _create_kv_cache_manager( head_dim: Optional[int] = None, kv_cache_type=None, is_disagg: bool = False, - cold_page_codec_provider: Optional[object] = None) -> KVCacheManager: + cold_page_codec_provider: Optional[object] = None, + kv_events_config: Optional[KVEventsConfig] = None) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config @@ -2381,6 +2385,12 @@ def _create_kv_cache_manager( manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats manager_extra_kwargs[ "cold_page_codec_provider"] = cold_page_codec_provider + manager_extra_kwargs["kv_events_config"] = kv_events_config + elif kv_events_config is not None and kv_events_config.enable_kv_cache_events: + logger.warning( + "kv_cache_config.kv_events_config is set but streaming KV event " + "publishing requires KV cache manager V2; events will not be " + f"published for {kv_cache_manager_cls.__name__}.") if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py new file mode 100644 index 000000000000..2352df07cc06 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -0,0 +1,785 @@ +# 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. +# +# This module defines TensorRT-LLM's KV cache event wire format: msgpack event batches +# published over ZeroMQ in the three-frame (topic, seq, payload) framing that external +# KV-cache-aware routers expect. Each event encodes as a map tagged with a "type" key, +# the form documented for custom router backends, so routers consume these batches +# without translation. This differs from vLLM's vllm/distributed/kv_events.py, whose +# structs set array_like=True and encode as tagged positional arrays; keeping the map +# form leaves field order out of the wire contract. The batch envelope is positional +# in both. + +from __future__ import annotations + +import queue +import threading +import time +import traceback +from abc import ABC, abstractmethod +from collections import deque +from itertools import count +from queue import Queue +from typing import Any, Optional + +import msgspec +import zmq + +from tensorrt_llm.llmapi.llm_args import KVEventsConfig +from tensorrt_llm.logger import logger +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 + +# Subscribers decode block hashes as 64-bit ints, so a bytes value would fail the +# decode for the entire batch. +ExternalBlockHash = int + + +class EventBatch( + msgspec.Struct, + array_like=True, # type: ignore[call-arg] + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] +): + """KV cache event wire batch envelope.""" + + ts: float + events: list[Any] + data_parallel_rank: int | None = None + + +class KVCacheWireEvent( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] + tag=True, +): + """Base class for KV cache event wire messages.""" + + +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] + block_size: int + lora_id: int | None + medium: str | None + lora_name: str | None + extra_keys: list[tuple[Any, ...] | None] | None = None + group_idx: int | None = None + kv_cache_spec_kind: str | None = None + kv_cache_spec_sliding_window: int | None = None + locality: str | None = None + + +class BlockRemoved(KVCacheWireEvent): + """A sequence of KV cache blocks was removed.""" + + block_hashes: list[ExternalBlockHash] + medium: str | None + group_idx: int | None = None + locality: str | None = None + + +class AllBlocksCleared(KVCacheWireEvent): + """All KV cache blocks were cleared.""" + + +class KVEventBatch(EventBatch): + """A batch containing only KV cache lifecycle events.""" + + events: list[BlockStored | BlockRemoved | AllBlocksCleared] + + +class EventPublisher(ABC): + """Publishes KV cache event wire batches for one cache rank.""" + + def __init__(self, data_parallel_rank: int = 0) -> None: + self._data_parallel_rank = data_parallel_rank + + def start(self) -> None: + """Acquire external resources. + + Split from ``__init__`` so constructing a publisher has no side effects: the + owner can build it early, finish its own validation, and only then commit to + binding sockets and running threads. + """ + + @abstractmethod + def publish(self, events: EventBatch) -> bool: + """Enqueue an event batch without blocking the scheduler.""" + + @abstractmethod + def shutdown(self) -> None: + """Flush pending batches and stop the publisher.""" + + +class NullEventPublisher(EventPublisher): + """Drains event batches locally without external I/O.""" + + def publish(self, events: EventBatch) -> bool: + return True + + def shutdown(self) -> None: + return + + +class ZmqEventPublisher(EventPublisher): + """Publishes event batches over the three-frame ZeroMQ wire protocol. + + Delivery is best effort, but loss is observable: :meth:`publish` reserves a sequence + number per accepted batch, so a dropped batch leaves a gap. Subscribers must treat a + gap -- including a replay that starts above the requested ``start_seq`` because + ``buffer_steps`` evicted older batches -- as lost KV-cache state and resynchronize. + """ + + SHUTDOWN_TIMEOUT = 1.0 + END_SEQ = (-1).to_bytes(8, "big", signed=True) + + def __init__( + self, + data_parallel_rank: int, + endpoint: str = "tcp://*:5557", + replay_endpoint: str | None = None, + buffer_steps: int = 10_000, + hwm: int = 100_000, + max_queue_size: int = 100_000, + topic: str = "", + ) -> None: + super().__init__(data_parallel_rank) + self._event_queue = Queue[Optional[tuple[int, EventBatch]]](maxsize=max_queue_size) + self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps) + self._ctx = zmq.Context.instance() + self._pub: Optional[zmq.Socket] = None + self._replay: Optional[zmq.Socket] = None + self._rank = data_parallel_rank + self._endpoint = self.offset_endpoint_port(endpoint, self._rank) + self._replay_endpoint = self.offset_endpoint_port(replay_endpoint, self._rank) + self._hwm = hwm + self._seq_gen = count() + self._topic_bytes = topic.encode("utf-8") + self._running = True + self._shutdown_lock = threading.Lock() + self.enqueued_batches = 0 + self.published_batches = 0 + self._queue_full_drops = 0 + self._send_error_drops = 0 + self._topic = topic + # Nothing is bound and no thread runs until start(); see EventPublisher.start(). + self._thread: Optional[threading.Thread] = None + + def start(self) -> None: + if self._thread is not None: + return + try: + self._socket_setup() + except Exception: + # start() never returns on failure, so close whatever was opened rather + # than leaking it on the shared context. + if self._pub is not None: + self._pub.close(linger=0) + self._pub = None + if self._replay is not None: + self._replay.close(linger=0) + self._replay = None + raise + self._thread = threading.Thread( + target=self._publisher_thread, + daemon=True, + name=f"trtllm-kv-events-rank-{self._rank}", + ) + self._thread.start() + logger.info( + f"Started streaming KV event publisher rank={self._rank} " + f"endpoint={self._endpoint} topic={self._topic!r}" + ) + + @property + def dropped_batches(self) -> int: + # Two independent writers: the scheduler thread bumps _queue_full_drops + # (queue full) and the publisher thread bumps _send_error_drops (send + # failure). Each counter has a single writer, so the sum needs no lock. + return self._queue_full_drops + self._send_error_drops + + def publish(self, events: EventBatch) -> bool: + if not self._running: + return False + if events.data_parallel_rank is None: + events.data_parallel_rank = self._data_parallel_rank + # Reserve the sequence number here rather than in the publisher thread, so a + # batch lost to a full queue or a failed send leaves a detectable gap instead of + # a contiguous stream that hides the loss. publish() is the only allocator. + seq = next(self._seq_gen) + try: + self._event_queue.put_nowait((seq, events)) + self.enqueued_batches += 1 + return True + except queue.Full: + self._queue_full_drops += 1 + drops = self._queue_full_drops + if drops == 1 or (drops & (drops - 1) == 0): + logger.warning( + f"Dropping streaming KV event batch on rank={self._rank} because " + f"the publisher queue is full; seq={seq} will be missing from the " + f"stream; dropped_batches={self.dropped_batches}" + ) + return False + + def shutdown(self) -> None: + with self._shutdown_lock: + if not self._running: + return + self._running = False + try: + self._event_queue.put_nowait(None) + except queue.Full: + # The thread exits after draining the full queue. + pass + if self._thread is not None: + self._thread.join(timeout=self.SHUTDOWN_TIMEOUT) + if self._thread.is_alive(): + logger.warning( + f"Streaming KV event publisher rank={self._rank} did not stop " + f"within {self.SHUTDOWN_TIMEOUT:.1f}s" + ) + logger.info( + f"Stopped streaming KV event publisher rank={self._rank} " + f"enqueued_batches={self.enqueued_batches} " + f"published_batches={self.published_batches} " + f"dropped_batches={self.dropped_batches}" + ) + + def _socket_setup(self) -> None: + self._pub = self._ctx.socket(zmq.PUB) + self._pub.set_hwm(self._hwm) + if not self._endpoint: + raise ValueError("KV event publisher endpoint must not be empty") + if not self._endpoint.startswith(("tcp://", "ipc://", "inproc://")): + raise ValueError(f"Unsupported KV event endpoint scheme: {self._endpoint!r}") + # The publisher owns its endpoint and subscribers connect to it, so the + # PUB socket always binds -- including explicit-host TCP binds like + # tcp://0.0.0.0:5557 that the previous '*'-only heuristic wrongly + # treated as connect targets (silently dropping every event). + self._pub.bind(self._endpoint) + + if self._replay_endpoint is not None: + self._replay = self._ctx.socket(zmq.ROUTER) + self._replay.bind(self._replay_endpoint) + + def _publisher_thread(self) -> None: + encoder = msgspec.msgpack.Encoder() + assert self._pub is not None + try: + while self._running or not self._event_queue.empty(): + if self._replay is not None and self._replay.poll(0): + try: + self._service_replay() + except Exception: + logger.error( + "Failed to service streaming KV event replay request\n" + f"{traceback.format_exc()}" + ) + try: + item = self._event_queue.get(timeout=0.1) + except queue.Empty: + continue + if item is None: + self._event_queue.task_done() + break + seq, event = item + try: + payload = encoder.encode(event) + self._pub.send_multipart( + ( + self._topic_bytes, + seq.to_bytes(8, "big"), + payload, + ) + ) + self._buffer.append((seq, payload)) + self.published_batches += 1 + except Exception: + self._send_error_drops += 1 + logger.error( + f"Failed to publish streaming KV event batch rank={self._rank}; " + f"seq={seq} will be missing from the stream\n" + f"{traceback.format_exc()}" + ) + time.sleep(0.1) + finally: + self._event_queue.task_done() + finally: + self._pub.close(linger=0) + if self._replay is not None: + self._replay.close(linger=0) + + def _service_replay(self) -> None: + assert self._replay is not None + frame = self._replay.recv_multipart() + if len(frame) != 3: + logger.warning(f"Invalid streaming KV event replay request: {frame}") + return + client_id, _, start_seq_bytes = frame + start_seq = int.from_bytes(start_seq_bytes, "big") + for seq, payload in self._buffer: + if seq >= start_seq: + self._replay.send_multipart( + ( + client_id, + b"", + self._topic_bytes, + seq.to_bytes(8, "big"), + payload, + ) + ) + self._replay.send_multipart((client_id, b"", b"", self.END_SEQ, b"")) + + @staticmethod + def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | None: + """Apply the base-port-plus-rank endpoint convention (each rank binds base_port + rank).""" + if not endpoint: + return endpoint + # Match the scheme with startswith so detection agrees with + # _socket_setup (substring tests misclassify hosts like "ipc-host"). + # ipc/inproc have no port; give each rank a distinct suffix instead. + if endpoint.startswith(("inproc://", "ipc://")): + return endpoint if data_parallel_rank == 0 else f"{endpoint}_dp{data_parallel_rank}" + if endpoint.startswith("tcp://"): + host_port = endpoint[len("tcp://") :] + if ":" not in host_port: + raise ValueError(f"TCP KV event endpoint must include a port: {endpoint!r}") + last_colon_idx = endpoint.rfind(":") + base_addr = endpoint[:last_colon_idx] + port_text = endpoint[last_colon_idx + 1 :] + # Validate the port value up front so a bad port names the endpoint + # instead of surfacing as an opaque int()/ZeroMQ bind error on ranks > 0. + if not (port_text.isdigit() and 1 <= int(port_text) <= 65_535): + raise ValueError( + f"TCP KV event endpoint must have a port in [1, 65535]: {endpoint!r}" + ) + if data_parallel_rank == 0: + return endpoint + base_port = int(port_text) + new_port = base_port + data_parallel_rank + if new_port > 65_535: + raise ValueError( + f"KV event endpoint port exceeds 65535 for rank {data_parallel_rank}" + ) + return f"{base_addr}:{new_port}" + raise ValueError("Invalid endpoint: must start with 'inproc://', 'ipc://', or 'tcp://'") + + +def _tcp_base_port(endpoint: str | None) -> int | None: + """Return the base port of a TCP endpoint, or None if it is not TCP.""" + if not endpoint or not endpoint.startswith("tcp://"): + return None + last_colon_idx = endpoint.rfind(":") + port_text = endpoint[last_colon_idx + 1 :] + if not port_text.isdigit(): + return None + return int(port_text) + + +def validate_streaming_support( + config: KVEventsConfig, + *, + pp_size: int, + cp_size: int, + ranks_per_host: int, + data_parallel_size: int, + backend: str, +) -> None: + """Reject streaming-KV-event configurations the engine cannot honour. + + Split out of ``KVCacheManagerV2.__init__`` so the preconditions are testable + without building a manager, which needs a GPU. + """ + if pp_size > 1: + 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." + ) + validate_endpoint_ranges(config, ranks_per_host, data_parallel_size) + + +def validate_endpoint_ranges( + config: KVEventsConfig, ranks_per_host: int, data_parallel_size: int +) -> None: + """Reject configurations whose publish and replay port ranges overlap. + + Ranks bind ``base_port + rank`` using their **global** rank, so each rank's port is + distinct cluster-wide and the sockets span ``[base, base + world - 1]``. Only ranks + co-located on one host actually contend for a port, and a host holds a contiguous + run of ranks, so the required spacing between the two base ports is the per-host + rank count rather than the total. Catch it before any socket is created rather than + as an opaque ``EADDRINUSE``. + """ + pub_base = _tcp_base_port(config.endpoint) + replay_base = _tcp_base_port(config.replay_endpoint) + if pub_base is None or replay_base is None: + return + span = max(1, ranks_per_host) + distance = abs(pub_base - replay_base) + if distance < span: + world = max(1, data_parallel_size) + raise ValueError( + f"KV event endpoint {config.endpoint!r} and replay_endpoint " + f"{config.replay_endpoint!r} overlap: ranks bind base_port+rank by global " + f"rank, so with {world} rank(s) the publish sockets span " + f"[{pub_base}, {pub_base + world - 1}] and the replay sockets span " + f"[{replay_base}, {replay_base + world - 1}]. Ranks co-located on a host " + f"contend for ports, so the base ports must be at least {span} apart (the " + f"per-host rank count) but are {distance} apart." + ) + + +def create_event_publisher(config: KVEventsConfig, data_parallel_rank: int) -> EventPublisher: + """Create the configured publisher for one cache rank.""" + if config.publisher == "null": + return NullEventPublisher(data_parallel_rank) + if config.publisher == "zmq": + return ZmqEventPublisher( + data_parallel_rank=data_parallel_rank, + endpoint=config.endpoint, + replay_endpoint=config.replay_endpoint, + buffer_steps=config.buffer_steps, + hwm=config.hwm, + max_queue_size=config.max_queue_size, + topic=config.topic, + ) + raise ValueError(f"Unsupported KV event publisher: {config.publisher!r}") + + +def _kv_event_wire_hash_from_radix_key(block_key: bytes) -> int: + """Reuse an existing SHA-256 radix key as the KV cache event's signed int64 wire hash.""" + if len(block_key) < 8: + raise ValueError("V2 radix block keys must contain at least 8 bytes") + # Reuse the canonical SHA-256 -> int64 truncation (first 8 bytes) shared with + # the rest of the KV-cache-event machinery instead of a second, divergent + # truncation, then reinterpret the low 64 bits as the signed int64 wire hash. + unsigned_hash = truncate_sha256_hash_to_int64(block_key) + 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 fast path that produces KV cache event wire messages directly. + + 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. + """ + + def __init__( + self, + config: KVEventsConfig, + *, + data_parallel_rank: int, + block_size: int, + max_window_size: int, + max_entries: int = 50_000, + ) -> 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._target_life_cycle_id: int | None = None + self._stored_blocks: dict[bytes, int] = {} + self._pending_events: list[BlockStored | BlockRemoved | AllBlocksCleared] = [] + self._pending_entries = 0 + self._closed = False + 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 + self.enqueued_events = 0 + self.dropped_batches = 0 + + def start(self) -> None: + """Bind the publisher's sockets and start its background thread. + + Construction is side-effect free, so the owner calls this only once every + other initialization check has passed. A failure before this point therefore + leaves no socket bound and no thread running. + """ + self._publisher.start() + + def set_layer_group_window_sizes(self, window_sizes: dict[int, int]) -> None: + target_ids = [ + int(life_cycle_id) + for life_cycle_id, window_size in window_sizes.items() + if int(window_size) == self._max_window_size + ] + if not target_ids and window_sizes: + largest_window = max(window_sizes.values()) + target_ids = [ + int(life_cycle_id) + for life_cycle_id, window_size in window_sizes.items() + if window_size == largest_window + ] + if not target_ids: + raise ValueError("Streaming KV events require an attention KV cache life cycle") + self._target_life_cycle_id = min(target_ids) + logger.info( + "Streaming KV event fast path selected " + f"lifecycle_id={self._target_life_cycle_id} " + f"window_size={self._max_window_size}" + ) + + def add_created_event( + self, + num_blocks_per_cache_level: Any, + layer_group_ids: Any = None, + ) -> None: + return + + def add_stored_event(self, *args: Any, **kwargs: Any) -> None: + # Streaming publishing derives stored events from the per-block hooks + # below; the aggregate stored-event hook is intentionally unused. + return + + def add_stored_block_event_from_block(self, block: Any) -> None: + if self._closed or self._target_life_cycle_id is None: + return + life_cycle_id = self._target_life_cycle_id + if life_cycle_id >= len(block.storage): + return + page_ref = block.storage[life_cycle_id] + page = None if page_ref is None else page_ref() + if page is None: + return + # A non-null page does not imply it covers the whole radix block: V2 can attach + # a page adopted from a shorter sibling. Publishing that as a BlockStored would + # tell the router the engine holds a prefix it cannot fully reuse. The buffered + # manager applies the same rule in _life_cycle_ids_from_radix_block(). + if page.num_tokens_in_block < len(block.tokens): + self.partial_blocks_suppressed += 1 + return + self._add_full_block(block) + + def add_stored_life_cycle_event_from_block(self, block: Any, life_cycle_id: int) -> None: + if life_cycle_id is None or self._target_life_cycle_id is None: + return + if int(life_cycle_id) != self._target_life_cycle_id: + self.non_target_life_cycles_ignored += 1 + return + self.add_stored_block_event_from_block(block) + + def _add_full_block(self, block: Any) -> None: + key = bytes(block.key) + if key in self._stored_blocks: + return + if len(block.tokens) != self._block_size: + self.partial_blocks_suppressed += 1 + return + if not self._reserve_entries(1): + return + try: + token_ids = self._token_ids(block.tokens) + 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 + logger.error( + "Dropping streaming KV store event with unsupported token data\n" + f"{traceback.format_exc()}" + ) + return + self._stored_blocks[key] = block_hash + if self._pending_events and isinstance(self._pending_events[-1], BlockStored): + previous = self._pending_events[-1] + if previous.block_hashes and previous.block_hashes[-1] == parent_hash: + previous.block_hashes.append(block_hash) + previous.token_ids.extend(token_ids) + self.stored_blocks += 1 + return + self._pending_events.append( + BlockStored( + block_hashes=[block_hash], + parent_block_hash=parent_hash, + token_ids=token_ids, + block_size=self._block_size, + lora_id=None, + medium="GPU", + lora_name=None, + ) + ) + self.stored_blocks += 1 + + @staticmethod + def _token_ids(tokens: Any) -> list[int]: + token_ids: list[int] = [] + for token in 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 + + def _block_hashes( + self, + block: Any, + ) -> tuple[int, int | None]: + parent = block.prev + is_root_child = getattr(parent, "ordinal", -1) == -1 + block_hash = _kv_event_wire_hash_from_radix_key(bytes(block.key)) + parent_hash = ( + None if is_root_child else _kv_event_wire_hash_from_radix_key(bytes(parent.key)) + ) + return block_hash, parent_hash + + def add_removed_event(self, block_hashes: Any) -> None: + if self._closed: + return + if isinstance(block_hashes, (bytes, str, int)): + block_hashes = (block_hashes,) + removed_hashes: list[ExternalBlockHash] = [] + for block_key in block_hashes: + if not isinstance(block_key, bytes): + continue + stored_hash = self._stored_blocks.pop(block_key, None) + if stored_hash is not None: + removed_hashes.append(stored_hash) + self._add_removed_hashes(removed_hashes) + + def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: + if self._closed or life_cycle_id is None or self._target_life_cycle_id is None: + return + if int(life_cycle_id) != self._target_life_cycle_id: + self.non_target_life_cycles_ignored += 1 + return + stored_hash = self._stored_blocks.pop(block_hash, None) + if stored_hash is not None: + self._add_removed_hashes([stored_hash]) + + def _add_removed_hashes(self, block_hashes: list[ExternalBlockHash]) -> None: + if not block_hashes: + return + # Removals are never dropped by the per-iteration cap and, unlike stores, + # do not consume the _pending_entries budget: each hash was already + # reported as stored (so removals are bounded by the stored set), and + # counting them against the store budget would starve legitimate + # BlockStored events in a removal-heavy iteration. + if self._pending_events and isinstance(self._pending_events[-1], BlockRemoved): + self._pending_events[-1].block_hashes.extend(block_hashes) + else: + self._pending_events.append(BlockRemoved(block_hashes=block_hashes, medium="GPU")) + self.removed_blocks += len(block_hashes) + + def add_updated_event( + self, + block_hash: Any, + *, + cache_level: KVCacheEventDiff | None = None, + priority: KVCacheEventDiff | None = None, + layer_group_id: int | None = None, + ) -> None: + return + + def _reserve_entries(self, num_entries: int) -> bool: + if self._pending_entries + num_entries <= self._max_entries: + self._pending_entries += num_entries + return True + self.dropped_events += num_entries + if self.dropped_events == num_entries or ( + self.dropped_events & (self.dropped_events - 1) == 0 + ): + logger.warning( + "Dropping streaming KV events because the per-iteration safety " + f"cap was exceeded; dropped_events={self.dropped_events}" + ) + return False + + def flush_iteration_events(self) -> None: + if self._closed or not self._pending_events: + return + events = self._pending_events + self._pending_events = [] + self._pending_entries = 0 + batch = KVEventBatch( + ts=time.time(), + events=events, + data_parallel_rank=self._rank, + ) + try: + if self._publisher.publish(batch): + self.enqueued_batches += 1 + self.enqueued_events += len(events) + else: + self.dropped_batches += 1 + except Exception: + self.dropped_batches += 1 + logger.error( + f"Dropping streaming KV event iteration batch on rank={self._rank}\n" + f"{traceback.format_exc()}" + ) + + 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 + # buffered polling path degrade cleanly rather than erroring. + return [] + + def shutdown(self) -> None: + if self._closed: + return + self.flush_iteration_events() + self._closed = True + self._publisher.shutdown() + logger.info( + "Streaming KV event fast path " + f"rank={self._rank} " + f"stored_blocks={self.stored_blocks} " + f"removed_blocks={self.removed_blocks} " + f"partial_blocks_suppressed={self.partial_blocks_suppressed} " + f"non_target_life_cycles_ignored={self.non_target_life_cycles_ignored} " + f"dropped_events={self.dropped_events} " + f"enqueued_batches={self.enqueued_batches} " + f"dropped_batches={self.dropped_batches}" + ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index e3bfd3c2db4c..cc1391d309cd 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -38,7 +38,7 @@ IndexMapper, copy_batch_block_offsets_to_device, ) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KVEventsConfig 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, @@ -61,6 +61,7 @@ KVCacheEventManager, KVCacheIterationStatsDelta, LayerId, + LifeCycleId, PageIndexMode, PlannedDropHandle, PoolGroupPeakBlockStats, @@ -74,6 +75,7 @@ gen_multimodal_cache_key_tokens, typed_range, ) +from tensorrt_llm.runtime.kv_cache_manager_v2 import BACKEND as KV_CACHE_MANAGER_V2_BACKEND from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as KVCacheManagerPy from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfMemoryError as KVCacheOutOfMemoryError @@ -85,6 +87,7 @@ from ..utils import maybe_compile from .config_utils import uses_vswa_kv_cache_layout from .connectors.kv_cache_connector import KvCacheConnectorManager +from .kv_cache_events import StreamingKVCacheEventManager, validate_streaming_support from .kv_cache_stats import ( KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, @@ -818,6 +821,7 @@ def __init__( is_disagg: bool = False, enable_stats: bool = False, num_reserved_index_slots: int = 1, + kv_events_config: Optional[KVEventsConfig] = None, is_estimating_kv_cache: bool = False, cold_page_codec_provider: Optional[object] = None, **kwargs, @@ -952,8 +956,41 @@ def __init__( self.max_seq_len if window_size is None else int(window_size) for window_size in self.max_attention_window_vec ) - self.event_manager: Optional[KVCacheEventManager] = None - if self.event_buffer_max_size > 0: + self.event_manager: Optional[KVCacheEventManager | StreamingKVCacheEventManager] = None + streaming_events_enabled = ( + kv_events_config is not None and kv_events_config.enable_kv_cache_events + ) + if streaming_events_enabled: + if self.event_buffer_max_size > 0: + logger.warning( + "Both kv_cache_config.event_buffer_max_size and streaming " + "kv_events_config are enabled; streaming publishing takes " + "precedence and the buffered get_kv_cache_events() poll path " + "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. + validate_streaming_support( + kv_events_config, + pp_size=mapping.pp_size, + cp_size=mapping.cp_size, + # Ranks bind by global rank; only those sharing a host can collide. + ranks_per_host=min(mapping.dp_size, mapping.gpus_per_node), + data_parallel_size=mapping.dp_size, + backend=KV_CACHE_MANAGER_V2_BACKEND, + ) + if mapping.enable_attention_dp or mpi_rank() == 0: + # 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 + self.event_manager = StreamingKVCacheEventManager( + kv_events_config, + data_parallel_rank=event_rank, + block_size=self.tokens_per_block, + max_window_size=event_window_size, + ) + elif self.event_buffer_max_size > 0: if mapping.enable_attention_dp: self.event_manager = KVCacheEventManager( self.event_buffer_max_size, @@ -1229,7 +1266,9 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self.can_evict = len(config.cache_tiers) > 1 if self.event_manager is not None: self.event_manager.set_layer_group_window_sizes( - self._get_event_window_sizes_by_layer_group() + self._get_event_window_sizes_by_layer_group( + attention_only=isinstance(self.event_manager, StreamingKVCacheEventManager) + ) ) self.event_manager.add_created_event( self._get_event_num_blocks_per_cache_level(config.cache_tiers, tokens_per_block), @@ -1349,6 +1388,14 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._log_kv_cache_pool_lifecycle_mapping() + # Last: bind the publisher socket and start its thread only once every check + # above has passed. Constructing the manager is side-effect free, so a failure + # anywhere earlier -- including the rank-coordinated aborts, where this rank + # raises because a peer failed -- leaves nothing bound to clean up. + if isinstance(self.event_manager, StreamingKVCacheEventManager): + self.event_manager.start() + logger.info("Streaming KV event fast path reuses V2 radix block hashes") + def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: """Return the roles represented by the two page-table index lanes. @@ -1675,21 +1722,33 @@ def _get_event_num_blocks_per_cache_level( def _get_event_layer_group_ids(self) -> List[int]: return [int(layer_group_id) for layer_group_id in range(len(self.impl.layer_grouping))] - def _get_event_window_sizes_by_layer_group(self) -> Dict[int, int]: + def _get_event_window_sizes_by_layer_group( + self, attention_only: bool = False + ) -> Dict[int, int]: # Assumes every layer in a group shares the same sliding_window_size, # which is how `impl.layer_grouping` partitions layers today. Only the # first layer's window is read; if the grouping policy ever permits # mixed windows in one group, this needs to fan out per-layer. + # + # `attention_only` is set for the streaming event manager, which tracks + # attention prefix reuse only: excluding SSM and other non-attention life cycles + # prevents a state life cycle (which reports max_seq_len as its window) from + # tying with the attention life cycle and being selected as the event target. + # The buffered manager keeps every layer group, so its windows are unchanged. def get_event_window_size(layer_id: int) -> int: layer_config = self.kv_cache_manager_py_config.layers[layer_id] window_size = getattr(layer_config, "sliding_window_size", None) return self.max_seq_len if window_size is None else int(window_size) - return { - int(layer_group_id): get_event_window_size(int(layer_ids[0])) - for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping) - } + window_sizes: Dict[int, int] = {} + for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping): + if attention_only: + life_cycle = self.impl._life_cycles.get_life_cycle(LifeCycleId(layer_group_id)) + if not isinstance(life_cycle, AttnLifeCycle): + continue + window_sizes[int(layer_group_id)] = get_event_window_size(int(layer_ids[0])) + return window_sizes def _format_kv_cache_pool_lifecycle_entry(self, layer_id: LayerId, role: DataRole) -> str: for pool_group in self.impl.pool_group_descs: @@ -3309,13 +3368,23 @@ def get_kv_cache_stats(self): return kv_cache_stats def flush_iteration_events(self): - if self.event_manager is not None: - self.event_manager.flush_iteration_events() + event_manager = self.event_manager + if event_manager is not None: + event_manager.flush_iteration_events() def get_latest_events(self, timeout_ms: Optional[float] = None): - if self.event_manager is None: + # Streaming publishing pushes events out-of-band; in that mode the event + # manager's get_latest_events returns [], so the buffered pull path + # degrades cleanly instead of raising. Snapshot event_manager once so a + # concurrent shutdown cannot turn it into None between the check and use. + event_manager = self.event_manager + if event_manager is None: return [] - return self.event_manager.get_latest_events(timeout_ms) + return event_manager.get_latest_events(timeout_ms) + + @property + def streaming_kv_events_enabled(self) -> bool: + return isinstance(self.event_manager, StreamingKVCacheEventManager) def get_iteration_stats(self): if not self.enable_stats: @@ -3894,6 +3963,13 @@ def shutdown(self): self.kv_cache_map.clear() self._request_stats_enabled_ids.clear() self.impl.shutdown() + # Shut the streaming event manager down last so removals emitted during + # cache / impl teardown (via the radix tree's own event-manager + # reference) are still flushed before the publisher stops. Do not null + # event_manager: get_latest_events/flush snapshot it and operate safely + # on a closed manager, so there is no teardown-time None race. + if isinstance(self.event_manager, StreamingKVCacheEventManager): + self.event_manager.shutdown() if self.conversation_manager is not None: self.conversation_manager.clear() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 836bcce336b7..1ebf687ef79c 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -598,7 +598,9 @@ def __init__( self._is_kv_manager_v2 = isinstance(self.kv_cache_manager, KVCacheManagerV2) self._prefetched_request_ids: set[int] = set() - self.enable_kv_cache_events = self.kv_cache_manager is not None and self.kv_cache_manager.event_buffer_max_size > 0 + self.enable_kv_cache_events = self.kv_cache_manager is not None and ( + self.kv_cache_manager.event_buffer_max_size > 0 or getattr( + self.kv_cache_manager, "streaming_kv_events_enabled", False)) self.enable_kv_cache_reuse = self.kv_cache_manager is not None and self.kv_cache_manager.enable_block_reuse # AsyncTransferManager pin/unpin path is V1-only; V2 holds blocks via _KVCache refcount. self.enable_partial_reuse_for_disagg = ( diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index e4520f2aadfa..25d4c856470a 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -16,17 +16,18 @@ DraftTargetDecodingConfig, DSparkDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, - LookaheadDecodingConfig, MambaStateConfig, - MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, - MoeConfig, MTPDecodingConfig, MultimodalConfig, - NGramDecodingConfig, PARDDecodingConfig, - PrefillCudaGraphBackend, PrometheusMetricsConfig, - ReorderRequestPolicyConfig, RocketSparseAttentionConfig, - SADecodingConfig, SAEnhancerConfig, - SaveHiddenStatesDecodingConfig, SchedulerConfig, - SkipSoftmaxAttentionConfig, TorchCompileConfig, - TorchLlmArgs, TriAttentionKvCacheCompressionConfig, + ExtendedRuntimePerfKnobConfig, KvCacheConfig, + KVEventsConfig, LlmArgs, LookaheadDecodingConfig, + MambaStateConfig, MedusaDecodingConfig, + MiniMaxM3SparseAttentionConfig, MoeConfig, + MTPDecodingConfig, MultimodalConfig, NGramDecodingConfig, + PARDDecodingConfig, PrefillCudaGraphBackend, + PrometheusMetricsConfig, ReorderRequestPolicyConfig, + RocketSparseAttentionConfig, SADecodingConfig, + SAEnhancerConfig, SaveHiddenStatesDecodingConfig, + SchedulerConfig, SkipSoftmaxAttentionConfig, + TorchCompileConfig, TorchLlmArgs, + TriAttentionKvCacheCompressionConfig, UserProvidedDecodingConfig) from .llm_utils import KvCacheRetentionConfig, QuantAlgo, QuantConfig from .mm_encoder import MultimodalEncoder @@ -47,6 +48,7 @@ 'DisaggScheduleStyle', 'BlockReuseConfig', 'KvCacheConfig', + 'KVEventsConfig', 'MambaStateConfig', 'KvCacheRetentionConfig', 'CudaGraphConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 79e8daffb1ea..0c6989506a34 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3973,6 +3973,59 @@ class MambaStateConfig(StrictBaseModel): "snapshots require KV cache manager V2.") +class KVEventsConfig(StrictBaseModel): + """Configuration for streaming (push-based) KV cache event publishing.""" + + enable_kv_cache_events: bool = Field( + default=False, + description= + "Whether to produce and publish KV cache events over the streaming (push) path." + ) + publisher: Optional[Literal["null", "zmq"]] = Field( + default=None, + description= + "Publisher implementation. Defaults to 'zmq' when events are enabled and 'null' otherwise." + ) + endpoint: str = Field( + default="tcp://*:5557", + min_length=1, + description= + "Base ZeroMQ endpoint the publisher binds. Each attention-DP rank binds " + "base_port+rank, so co-located engines (e.g. disaggregated prefill and " + "decode on one host) must use distinct base ports.") + replay_endpoint: Optional[str] = Field( + default=None, + min_length=1, + description= + "Optional base ZeroMQ endpoint used to replay KV cache events. Ranks apply " + "the same global base_port+rank convention as `endpoint`. Only ranks sharing a " + "host contend for a port, so the two base ports must be at least " + "ranks-per-host apart or a rank's replay bind collides with another rank's " + "publish bind on that host.") + buffer_steps: int = Field( + default=10_000, + gt=0, + description="Number of previously published batches retained for replay." + ) + hwm: int = Field(default=100_000, + gt=0, + description="ZeroMQ publisher socket high-water mark. " + "0 means unlimited in ZeroMQ, so it is disallowed here.") + max_queue_size: int = Field( + default=100_000, + gt=0, + description="Maximum number of batches queued for background publishing. " + "Must be positive; 0 would make the queue unbounded.") + topic: str = Field( + default="", + description="ZeroMQ subscription topic used for KV cache event batches." + ) + + def model_post_init(self, __context) -> None: + if self.publisher is None: + self.publisher = "zmq" if self.enable_kv_cache_events else "null" + + class BlockReuseConfig(StrictBaseModel): """Configuration for KV cache block reuse policies.""" @@ -4077,6 +4130,14 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description= "The period in milliseconds to gather attention DP events across ranks." ) + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. + kv_events_config: Optional[KVEventsConfig] = Field( + default=None, + status="prototype", + description= + "Streaming (push-based) KV cache event publishing (KV cache manager V2 only). When set, " + "each rank publishes its own events directly (e.g. over ZeroMQ) instead " + "of the buffered event_buffer_max_size gather/poll path.") enable_partial_reuse: bool = Field( default=True, description= diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index b5c50e86a9dc..b2ac11496829 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -28,10 +28,10 @@ from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - KvCacheConfig, LlmArgs, LookaheadDecodingConfig, - MedusaDecodingConfig, MTPDecodingConfig, - NGramDecodingConfig, SchedulerConfig, TorchLlmArgs, - UserProvidedDecodingConfig, _ModelWrapper, + KvCacheConfig, KVEventsConfig, LlmArgs, + LookaheadDecodingConfig, MedusaDecodingConfig, + MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig, + TorchLlmArgs, UserProvidedDecodingConfig, _ModelWrapper, _ParallelConfig, update_llm_args_with_extra_dict, update_llm_args_with_extra_options) # yapf: enable @@ -502,6 +502,7 @@ class LlmBuildStats: 'DecodeCudaGraphConfig', 'EncodeCudaGraphConfig', 'KvCacheConfig', + 'KVEventsConfig', 'CachedModelLoader', 'EagleDecodingConfig', 'Eagle3DecodingConfig', diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index ac0c604d02c6..657eed29d22f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -22,6 +22,10 @@ _BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() +#: Name of the active backend ("cpp" or "python"). Exposed so callers can gate +#: Python-only extension points, such as duck-typed event sinks, on the selection. +BACKEND = _BACKEND + if _BACKEND == "python": from . import rawref # noqa: F401 from ._block_radix_tree import ( # noqa: F401 @@ -292,6 +296,7 @@ def typed_range(*args: int) -> range: __all__ = [ "AggregatedPageDesc", "AttentionLayerConfig", + "BACKEND", "BAD_PAGE_INDEX", "CACHE_LEVEL1", "BatchDesc", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 2a467edef33c..f447a911c664 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -814,6 +814,44 @@ "kind": "categorical", "path": "kv_cache_config.kv_cache_event_hash_algo" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_config.kv_events_config.buffer_steps" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_config.kv_events_config.enable_kv_cache_events" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_config.kv_events_config.hwm" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_config.kv_events_config.max_queue_size" + }, + { + "allowed_values": [ + "null", + "zmq" + ], + "annotation": "Optional[Literal['null', 'zmq']]", + "converter": "", + "kind": "categorical", + "path": "kv_cache_config.kv_events_config.publisher" + }, { "allowed_values": [ "auto", 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 new file mode 100644 index 000000000000..c0c6775d800a --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py @@ -0,0 +1,481 @@ +# 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. + +import socket +from types import SimpleNamespace +from typing import Callable + +import msgspec +import pytest +import zmq + +from tensorrt_llm._torch.pyexecutor.kv_cache_events import ( + KVEventBatch, + StreamingKVCacheEventManager, + ZmqEventPublisher, + validate_endpoint_ranges, + validate_streaming_support, +) +from tensorrt_llm.llmapi.llm_args import KVEventsConfig + +_ZMQ_SETUP_ATTEMPTS = 4 +_RECEIVE_TIMEOUT_MS = 2_000 +_SUBSCRIBE_ATTEMPTS = 50 +_PROBE_TIMEOUT_MS = 100 + + +class _NotReceived(Exception): + """No batch arrived: the subscription had not propagated yet.""" + + +def _unused_tcp_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _await_subscription(publisher: ZmqEventPublisher, subscriber: zmq.Socket) -> int: + """Publish probe batches until the subscriber's subscription is live. + + A PUB socket silently drops everything published before a subscriber's + subscription has propagated, and that window is not bounded by any delay the test + can pick -- so synchronise on an actual received message instead of sleeping. + Returns the number of probes published, which is the sequence number the next real + batch will carry. + """ + for probes in range(1, _SUBSCRIBE_ATTEMPTS + 1): + publisher.publish(KVEventBatch(ts=0.0, events=[])) + if subscriber.poll(_PROBE_TIMEOUT_MS): + while subscriber.poll(0): + subscriber.recv_multipart() + return probes + raise _NotReceived("subscription never propagated") + + +def _run_on_fresh_port(scenario: Callable[[int], None]) -> None: + """Retry `scenario(port)` on a fresh port if its sockets could not come up. + + `_unused_tcp_port()` releases its port before the publisher binds it, so another + process can take it in between. Assertion failures inside `scenario` are not + retried. + """ + for _ in range(_ZMQ_SETUP_ATTEMPTS): + try: + scenario(_unused_tcp_port()) + return + except _NotReceived: + pass + except zmq.ZMQError as exc: + if exc.errno != zmq.EADDRINUSE: + raise + pytest.fail(f"ZeroMQ setup failed after {_ZMQ_SETUP_ATTEMPTS} attempts") + + +def test_streaming_fast_path_publishes_only_full_max_window_blocks() -> None: + """Protect radix hash reuse, filtering, wire format, and shutdown.""" + topic = "kv-events" + context = zmq.Context.instance() + + # Wire hashes come from truncate_sha256_hash_to_int64 = the FIRST 8 bytes + # of the radix key, so put the distinguishing bytes -- including the high + # bit that exercises the signed-wraparound branch -- at the front. + first_hash = b"\x80\x00\x00\x00\x00\x00\x00\x01" + b"\x11" * 24 + partial_hash = b"\x22" * 32 + second_hash = b"\x00\x00\x00\x00\x00\x00\x00\x02" + b"\x33" * 24 + first_wire_hash = int.from_bytes(first_hash[:8], "big") - 2**64 + second_wire_hash = int.from_bytes(second_hash[:8], "big") + + # A fresh manager per attempt restarts sequence numbers at 0 and clears the + # stored-block dedup state, so a retry replays the scenario exactly. + def scenario(port: int) -> None: + bind_endpoint = f"tcp://*:{port}" + subscriber = context.socket(zmq.SUB) + subscriber.setsockopt_string(zmq.SUBSCRIBE, topic) + subscriber.connect(f"tcp://127.0.0.1:{port}") + manager = None + try: + manager = StreamingKVCacheEventManager( + KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint=bind_endpoint, + topic=topic, + max_queue_size=8, + ), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + ) + manager.start() + base_seq = _await_subscription(manager._publisher, subscriber) + manager.set_layer_group_window_sizes({0: 128, 1: 64}) + + root = SimpleNamespace(ordinal=-1) + + def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: + max_window_page = SimpleNamespace(num_tokens_in_block=len(tokens)) + smaller_window_page = SimpleNamespace(num_tokens_in_block=len(tokens)) + return SimpleNamespace( + key=key, + tokens=tokens, + prev=prev, + ordinal=getattr(prev, "ordinal", -1) + 1, + storage=[lambda: max_window_page, lambda: smaller_window_page], + ) + + first = block(first_hash, [1, 2, 3, 4], root) + partial = block(partial_hash, [5, 6], first) + second = block(second_hash, [5, 6, 7, 8], first) + + manager.add_stored_block_event_from_block(first) + manager.add_stored_block_event_from_block(partial) + manager.add_stored_life_cycle_event_from_block(second, 1) + manager.add_stored_life_cycle_event_from_block(second, 0) + manager.flush_iteration_events() + manager.add_removed_event([first_hash, partial_hash, second_hash]) + manager.flush_iteration_events() + + frames = [] + for _ in range(2): + if not subscriber.poll(_RECEIVE_TIMEOUT_MS): + raise _NotReceived(port) + frames.append(subscriber.recv_multipart()) + + assert [frame[0] for frame in frames] == [topic.encode(), topic.encode()] + # Sequence numbers stay dense across the probes and the real batches. + assert [int.from_bytes(frame[1], "big") for frame in frames] == [ + base_seq, + base_seq + 1, + ] + stored_batch = msgspec.msgpack.decode(frames[0][2]) + removed_batch = msgspec.msgpack.decode(frames[1][2]) + assert stored_batch[2] == 0 + assert stored_batch[1] == [ + { + "type": "BlockStored", + "block_hashes": [first_wire_hash, second_wire_hash], + "parent_block_hash": None, + "token_ids": [1, 2, 3, 4, 5, 6, 7, 8], + "block_size": 4, + "lora_id": None, + "medium": "GPU", + "lora_name": None, + } + ] + assert removed_batch[1] == [ + { + "type": "BlockRemoved", + "block_hashes": [first_wire_hash, second_wire_hash], + "medium": "GPU", + } + ] + assert manager.stored_blocks == 2 + assert manager.removed_blocks == 2 + assert manager.partial_blocks_suppressed == 1 + assert manager.non_target_life_cycles_ignored == 1 + assert manager.dropped_events == 0 + + # Streaming publishing pushes events out-of-band, so the buffered pull + # API must degrade to an empty result rather than raising. + assert manager.get_latest_events() == [] + + # shutdown() must be idempotent and must release the bound port. + manager.shutdown() + manager.shutdown() + replacement = context.socket(zmq.PUB) + replacement.bind(bind_endpoint) + replacement.close(linger=0) + finally: + if manager is not None: + manager.shutdown() + subscriber.close(linger=0) + + _run_on_fresh_port(scenario) + + +def test_streaming_removals_are_never_dropped_by_the_entry_cap() -> None: + """Removals must survive the per-iteration cap or the consumer desyncs.""" + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=2, + max_window_size=128, + max_entries=2, + ) + manager.start() + try: + manager.set_layer_group_window_sizes({0: 128}) + + # Capture what actually reaches the publisher so the test proves the + # removals are emitted on flush, not merely queued in _pending_events. + published: list[object] = [] + manager._publisher.publish = lambda batch: published.append(batch) or True + + root = SimpleNamespace(ordinal=-1) + + def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: + page = SimpleNamespace(num_tokens_in_block=len(tokens)) + return SimpleNamespace( + key=key, + tokens=tokens, + prev=prev, + ordinal=getattr(prev, "ordinal", -1) + 1, + storage=[lambda: page], + ) + + first = block(b"\x01" * 32, [1, 2], root) + second = block(b"\x02" * 32, [3, 4], first) + manager.add_stored_block_event_from_block(first) + manager.add_stored_block_event_from_block(second) + + # Both stores fill the entry cap (max_entries=2); the removals must still + # be emitted rather than dropped, or the consumer treats the blocks as + # resident forever. + manager.add_removed_event([b"\x01" * 32, b"\x02" * 32]) + manager.flush_iteration_events() + + assert manager.removed_blocks == 2 + assert len(published) == 1 + # Round-trip through msgpack to prove the removals reach the wire as a + # BlockRemoved batch carrying both hashes. + decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(published[0])) + removed = [event for event in decoded[1] if event.get("type") == "BlockRemoved"] + assert sum(len(event["block_hashes"]) for event in removed) == 2 + finally: + manager.shutdown() + + +def test_dropped_batches_leave_a_sequence_gap() -> None: + """A batch lost to a full queue must be observable as a missing sequence number.""" + # Left unstarted on purpose: publish() only touches the queue, so the drop path is + # exercised without binding a socket or draining the queue from a live thread. + publisher = ZmqEventPublisher( + data_parallel_rank=0, + endpoint="inproc://kv-events-drop-test", + max_queue_size=1, + ) + try: + assert publisher.publish(KVEventBatch(ts=0.0, events=[])) is True + assert publisher.publish(KVEventBatch(ts=1.0, events=[])) is False + assert publisher.dropped_batches == 1 + + # The accepted batch kept seq 0 and the dropped batch consumed seq 1, so the + # next batch is seq 2: subscribers see a hole rather than a contiguous stream + # that hides the loss. + seq, _ = publisher._event_queue.get_nowait() + assert seq == 0 + assert publisher.publish(KVEventBatch(ts=2.0, events=[])) is True + next_seq, _ = publisher._event_queue.get_nowait() + assert next_seq == 2 + finally: + publisher.shutdown() + + +def test_construction_binds_nothing_until_start() -> None: + """A constructed-but-unstarted publisher must hold no socket and no thread.""" + port = _unused_tcp_port() + endpoint = f"tcp://127.0.0.1:{port}" + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="zmq", endpoint=endpoint), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + ) + try: + publisher = manager._publisher + assert publisher._pub is None + assert publisher._thread is None + + manager.start() + assert publisher._pub is not None + assert publisher._thread is not None and publisher._thread.is_alive() + # start() is idempotent. + manager.start() + finally: + manager.shutdown() + + +def test_shutdown_without_start_is_safe() -> None: + """Tearing down a manager that never started must not raise.""" + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://127.0.0.1:1"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + ) + # Never started, so nothing was bound -- shutdown must still be a clean no-op. + manager.shutdown() + manager.shutdown() + + +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. + validate_streaming_support(config, **supported) + + 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"}) + + +@pytest.mark.parametrize( + "endpoint,replay_endpoint,ranks_per_host,overlaps", + [ + # 2 ranks bind 5557-5558 and 5558-5559: rank 1's publish hits rank 0's replay. + ("tcp://*:5557", "tcp://*:5558", 2, True), + ("tcp://*:5557", "tcp://*:5558", 1, False), + ("tcp://*:5557", "tcp://*:5657", 2, False), + # Replay below the publish base overlaps just the same. + ("tcp://*:5558", "tcp://*:5557", 2, True), + ( + "tcp://*:5557", + "tcp://*:5559", + 2, + False, + ), + # No replay endpoint means no second range to collide with. + ("tcp://*:5557", None, 8, False), + # 16 attention-DP ranks over 2 nodes collide only within a node, so spacing + # equal to the per-host rank count is legal even though it is under dp_size. + ("tcp://*:5557", "tcp://*:5565", 8, False), + # ipc/inproc endpoints have no ports, so the check does not apply. + ("ipc:///tmp/kv-events", "ipc:///tmp/kv-replay", 8, False), + ], +) +def test_validate_endpoint_ranges(endpoint, replay_endpoint, ranks_per_host, overlaps) -> None: + kwargs = {"replay_endpoint": replay_endpoint} if replay_endpoint else {} + config = KVEventsConfig(enable_kv_cache_events=True, endpoint=endpoint, **kwargs) + if overlaps: + with pytest.raises(ValueError, match="overlap"): + validate_endpoint_ranges(config, ranks_per_host, ranks_per_host) + else: + validate_endpoint_ranges(config, ranks_per_host, ranks_per_host) + + +def test_partial_target_page_coverage_is_suppressed_until_fully_covered() -> None: + """A page adopted from a shorter sibling must not be published as a full block.""" + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + ) + manager.start() + try: + manager.set_layer_group_window_sizes({0: 128}) + published: list[object] = [] + manager._publisher.publish = lambda batch: published.append(batch) or True + + root = SimpleNamespace(ordinal=-1) + # The block holds 4 tokens but its target page only covers 2 of them. + page = SimpleNamespace(num_tokens_in_block=2) + block = SimpleNamespace( + key=b"\x01" * 32, + tokens=[1, 2, 3, 4], + prev=root, + ordinal=0, + storage=[lambda: page], + ) + + manager.add_stored_block_event_from_block(block) + manager.flush_iteration_events() + assert manager.stored_blocks == 0 + assert manager.partial_blocks_suppressed == 1 + assert published == [] + + # Once the page covers the whole block, the same block is published. + page.num_tokens_in_block = 4 + manager.add_stored_life_cycle_event_from_block(block, 0) + manager.flush_iteration_events() + assert manager.stored_blocks == 1 + assert len(published) == 1 + decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(published[0])) + stored = [event for event in decoded[1] if event["type"] == "BlockStored"] + assert sum(len(event["block_hashes"]) for event in stored) == 1 + finally: + manager.shutdown() + + +def test_life_cycle_hooks_ignore_none_ids() -> None: + """A None life-cycle id must not reach int() before the target is configured.""" + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + ) + manager.start() + try: + # Before set_layer_group_window_sizes(), and with a None id, both hooks are + # no-ops rather than raising TypeError. + manager.add_stored_life_cycle_event_from_block(object(), None) + manager.add_removed_life_cycle_event(b"\x01" * 32, None) + manager.set_layer_group_window_sizes({0: 128}) + manager.add_stored_life_cycle_event_from_block(object(), None) + manager.add_removed_life_cycle_event(b"\x01" * 32, None) + assert manager.stored_blocks == 0 + assert manager.removed_blocks == 0 + finally: + manager.shutdown() + + +def test_kv_events_config_publisher_default() -> None: + """model_post_init resolves the publisher default (the common user path).""" + assert KVEventsConfig(enable_kv_cache_events=True).publisher == "zmq" + assert KVEventsConfig().publisher == "null" + assert KVEventsConfig(enable_kv_cache_events=False).publisher == "null" + # An explicitly set publisher is always respected. + assert KVEventsConfig(enable_kv_cache_events=True, publisher="null").publisher == "null" + assert KVEventsConfig(enable_kv_cache_events=False, publisher="zmq").publisher == "zmq" + + +@pytest.mark.parametrize( + "endpoint,rank,expected", + [ + ("tcp://*:5557", 0, "tcp://*:5557"), # rank 0 is identity + ("tcp://*:5557", 3, "tcp://*:5560"), # tcp base_port + rank + ("tcp://127.0.0.1:5557", 1, "tcp://127.0.0.1:5558"), + ("ipc:///tmp/kv-events", 2, "ipc:///tmp/kv-events_dp2"), # no port -> suffix + ("inproc://kv-events", 2, "inproc://kv-events_dp2"), + (None, 5, None), + ], +) +def test_offset_endpoint_port(endpoint, rank, expected) -> None: + assert ZmqEventPublisher.offset_endpoint_port(endpoint, rank) == expected + + +def test_offset_endpoint_port_rejects_bad_input() -> None: + # base_port + rank must stay within the u16 range. + with pytest.raises(ValueError): + ZmqEventPublisher.offset_endpoint_port("tcp://*:65535", 1) + # Unknown scheme is rejected for a non-zero rank. + with pytest.raises(ValueError): + ZmqEventPublisher.offset_endpoint_port("http://host:5557", 1) + # A TCP endpoint without a port is rejected instead of raising an opaque + # int() error on the scheme colon. + with pytest.raises(ValueError): + ZmqEventPublisher.offset_endpoint_port("tcp://host", 1) + # Non-numeric or out-of-range ports are rejected with an endpoint-naming + # error instead of an opaque int()/bind failure. + for bad in ("tcp://host:abc", "tcp://host:0", "tcp://host:-5"): + with pytest.raises(ValueError): + ZmqEventPublisher.offset_endpoint_port(bad, 1)