From 92b8fd38a4fb162d490c124c9402af1fc6d7e540 Mon Sep 17 00:00:00 2001 From: Alec Flowers Date: Fri, 24 Jul 2026 23:26:56 -0700 Subject: [PATCH 01/25] feat: publish native v2 kv cache events Signed-off-by: Alec Flowers --- tensorrt_llm/_torch/pyexecutor/_util.py | 13 +- .../_torch/pyexecutor/kv_cache_events.py | 460 ++++++++++++++++++ .../_torch/pyexecutor/kv_cache_manager_v2.py | 43 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 5 +- tensorrt_llm/llmapi/__init__.py | 10 +- tensorrt_llm/llmapi/llm_args.py | 45 ++ tensorrt_llm/llmapi/llm_utils.py | 4 +- .../usage/llm_args_golden_manifest.json | 38 ++ .../test_native_kv_events.py | 159 ++++++ 9 files changed, 765 insertions(+), 12 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache_events.py create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 40ab2e3a64ed..4106f0260f06 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -30,9 +30,9 @@ # isort: off from tensorrt_llm.llmapi.llm_args import ( CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, - KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, PeftCacheConfig, - SamplerType, SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, - TorchLlmArgs, WaitingQueuePolicy) + KVEventsConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, + PeftCacheConfig, SamplerType, SchedulerConfig, SparseAttentionConfig, + SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy) # isort: on from tensorrt_llm.logger import logger from tensorrt_llm.lora_helper import (LoraConfig, @@ -1142,6 +1142,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_events_config, ) if not self._skip_est: @@ -1858,7 +1861,8 @@ def _create_kv_cache_manager( num_kv_heads: Optional[Union[int, List[int]]] = None, head_dim: Optional[int] = None, kv_cache_type=None, - is_disagg: bool = False) -> KVCacheManager: + is_disagg: bool = False, + kv_events_config: Optional[KVEventsConfig] = None) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config @@ -1989,6 +1993,7 @@ def _create_kv_cache_manager( manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats + manager_extra_kwargs["kv_events_config"] = kv_events_config 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..08ff38ed75e3 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -0,0 +1,460 @@ +# 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. +# +# The wire schema and ZeroMQ framing in this file are adapted from vLLM's +# vllm/distributed/kv_events.py. + +from __future__ import annotations + +import queue +import threading +import time +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_manager_v2._event_manager import ( + KVCacheCreatedData, + KVCacheEvent, + KVCacheRemovedData, + KVCacheStoredData, + KVCacheUpdatedData, +) + +ExternalBlockHash = bytes | 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] +): + """vLLM-compatible event 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 vLLM-compatible KV cache events.""" + + +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 vLLM-compatible event batches for one cache rank.""" + + def __init__(self, data_parallel_rank: int = 0) -> None: + self._data_parallel_rank = data_parallel_rank + + @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 with vLLM's three-frame ZeroMQ protocol.""" + + 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[EventBatch | None](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.dropped_batches = 0 + self._socket_setup() + 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 native KV event publisher rank={self._rank} " + f"endpoint={self._endpoint} topic={topic!r}") + + 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 + try: + self._event_queue.put_nowait(events) + self.enqueued_batches += 1 + return True + except queue.Full: + self.dropped_batches += 1 + if self.dropped_batches == 1 or (self.dropped_batches & + (self.dropped_batches - 1) == 0): + logger.warning( + f"Dropping native KV event batch on rank={self._rank} because " + "the publisher queue is full; " + f"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 + self._thread.join(timeout=self.SHUTDOWN_TIMEOUT) + if self._thread.is_alive(): + logger.warning( + f"Native KV event publisher rank={self._rank} did not stop " + f"within {self.SHUTDOWN_TIMEOUT:.1f}s") + logger.info(f"Stopped native 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 self._endpoint is None: + raise ValueError("KV event publisher endpoint must not be empty") + if ("*" in self._endpoint or "::" in self._endpoint + or self._endpoint.startswith(("ipc://", "inproc://"))): + self._pub.bind(self._endpoint) + else: + self._pub.connect(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.exception( + "Failed to service native KV event replay request") + try: + event = self._event_queue.get(timeout=0.1) + except queue.Empty: + continue + if event is None: + self._event_queue.task_done() + break + seq = next(self._seq_gen) + 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.dropped_batches += 1 + logger.exception(f"Failed to publish native KV event batch " + f"rank={self._rank} seq={seq}") + 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 native 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 vLLM's base-port-plus-rank endpoint convention.""" + if not endpoint or data_parallel_rank == 0: + return endpoint + if "inproc" in endpoint: + return f"{endpoint}_dp{data_parallel_rank}" + if "tcp" in endpoint and ":" in endpoint: + last_colon_idx = endpoint.rfind(":") + base_addr = endpoint[:last_colon_idx] + base_port = int(endpoint[last_colon_idx + 1:]) + 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 contain 'inproc' or 'tcp'") + + +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 _to_wire_hash(block_hash: int | str | None) -> ExternalBlockHash | None: + if block_hash is None: + return None + if isinstance(block_hash, int): + if block_hash >= 2**63: + return block_hash - 2**64 + if block_hash < -(2**63): + return ((block_hash + 2**63) % 2**64) - 2**63 + return block_hash + try: + return bytes.fromhex(block_hash) + except ValueError as error: + raise ValueError( + f"Invalid hexadecimal KV block hash: {block_hash!r}") from error + + +class KVEventAdapter: + """Converts local V2 events and publishes one wire batch per iteration.""" + + def __init__( + self, + config: KVEventsConfig, + *, + data_parallel_rank: int, + block_size: int, + max_window_size: int, + ) -> None: + self._rank = data_parallel_rank + self._block_size = block_size + self._max_window_size = max_window_size + self._publisher = create_event_publisher(config, data_parallel_rank) + self._partial_block_hashes: set[int | str] = set() + self._closed = False + self.local_batches = 0 + self.local_events = 0 + self.enqueued_batches = 0 + self.enqueued_events = 0 + self.dropped_batches = 0 + + def publish_local_events( + self, events: list[KVCacheEvent]) -> list[list[KVCacheEvent]]: + """Publish local events and return no gathered events to the manager.""" + if self._closed or not events: + return [] + self.local_batches += 1 + self.local_events += len(events) + try: + wire_events = [ + wire_event for event in events + if (wire_event := self._convert_event(event)) is not None + ] + if wire_events: + batch = KVEventBatch( + ts=time.time(), + events=wire_events, + data_parallel_rank=self._rank, + ) + if self._publisher.publish(batch): + self.enqueued_batches += 1 + self.enqueued_events += len(wire_events) + else: + self.dropped_batches += 1 + except Exception: + self.dropped_batches += 1 + logger.exception( + f"Dropping native KV event iteration batch on rank={self._rank}" + ) + return [] + + def _convert_event( + self, event: KVCacheEvent + ) -> BlockStored | BlockRemoved | AllBlocksCleared | None: + if event.window_size != self._max_window_size: + return None + data = event.data + if isinstance(data, (KVCacheCreatedData, KVCacheUpdatedData)): + return None + if isinstance(data, KVCacheStoredData): + block_hashes: list[ExternalBlockHash] = [] + token_ids: list[int] = [] + for block in data.blocks: + num_tokens = len(block.tokens) + if num_tokens > self._block_size: + raise ValueError( + f"KV block has {num_tokens} tokens, expected at most " + f"{self._block_size}") + if num_tokens < self._block_size: + self._partial_block_hashes.add(block.block_hash) + break + block_token_ids = [token.token_id for token in block.tokens] + if any(not isinstance(token_id, int) + for token_id in block_token_ids): + raise ValueError( + "vLLM-compatible KV events require integer token IDs") + wire_hash = _to_wire_hash(block.block_hash) + assert wire_hash is not None + block_hashes.append(wire_hash) + token_ids.extend(block_token_ids) + if not block_hashes: + return None + return BlockStored( + block_hashes=block_hashes, + parent_block_hash=_to_wire_hash(data.parent_hash), + token_ids=token_ids, + block_size=self._block_size, + lora_id=None, + medium="GPU", + lora_name=None, + ) + if isinstance(data, KVCacheRemovedData): + block_hashes: list[ExternalBlockHash] = [] + for block_hash in data.block_hashes: + if block_hash in self._partial_block_hashes: + self._partial_block_hashes.remove(block_hash) + continue + wire_hash = _to_wire_hash(block_hash) + assert wire_hash is not None + block_hashes.append(wire_hash) + if not block_hashes: + return None + return BlockRemoved(block_hashes=block_hashes, medium="GPU") + return None + + def shutdown(self) -> None: + """Close the publisher once and report direct-path counters.""" + if self._closed: + return + self._closed = True + self._publisher.shutdown() + logger.info( + f"Native KV events rank={self._rank} " + f"local_batches={self.local_batches} " + f"local_events={self.local_events} " + f"enqueued_batches={self.enqueued_batches} " + f"enqueued_events={self.enqueued_events} " + f"dropped_batches={self.dropped_batches} kv_event_allgathers=0") diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 518c7b711162..21ab6d715117 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -37,7 +37,7 @@ IndexMapper, copy_batch_block_offsets_to_device, ) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import KVEventsConfig, KvCacheConfig 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, @@ -82,6 +82,7 @@ from ...mapping import CpType, Mapping from ..utils import maybe_compile from .connectors.kv_cache_connector import KvCacheConnectorManager +from .kv_cache_events import KVEventAdapter from .kv_cache_stats import ( KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, @@ -767,6 +768,7 @@ def __init__( is_disagg: bool = False, enable_stats: bool = False, num_reserved_index_slots: int = 1, + kv_events_config: Optional[KVEventsConfig] = None, **kwargs, ) -> None: self.mapping = mapping @@ -867,7 +869,36 @@ def __init__( for window_size in self.max_attention_window_vec ) self.event_manager: Optional[KVCacheEventManager] = None - if self.event_buffer_max_size > 0: + self.kv_event_adapter: Optional[KVEventAdapter] = None + native_events_enabled = ( + kv_events_config is not None + and kv_events_config.enable_kv_cache_events + ) + if native_events_enabled: + if mapping.pp_size > 1: + raise ValueError( + "Native KV events do not support pipeline parallelism") + if mapping.cp_size > 1: + raise ValueError( + "Native KV events do not support context parallelism") + assert kv_events_config is not None + if mapping.enable_attention_dp or mpi_rank() == 0: + event_rank = mapping.rank if mapping.enable_attention_dp else 0 + self.kv_event_adapter = KVEventAdapter( + kv_events_config, + data_parallel_rank=event_rank, + block_size=self.tokens_per_block, + max_window_size=event_window_size, + ) + self.event_manager = KVCacheEventManager( + 50_000, + window_size=event_window_size, + attention_dp_rank=event_rank, + attention_dp_gather=self.kv_event_adapter. + publish_local_events, + hash_algo=kv_cache_event_hash_algo, + ) + elif self.event_buffer_max_size > 0: if mapping.enable_attention_dp: self.event_manager = KVCacheEventManager( self.event_buffer_max_size, @@ -2897,6 +2928,10 @@ def flush_iteration_events(self): self.event_manager.flush_iteration_events() def get_latest_events(self, timeout_ms: Optional[float] = None): + if self.kv_event_adapter is not None: + raise RuntimeError( + "KV cache event polling is unavailable while native publishing is enabled" + ) if self.event_manager is None: return [] return self.event_manager.get_latest_events(timeout_ms) @@ -3423,6 +3458,10 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool return bool(has_invalid_values) def shutdown(self): + if self.kv_event_adapter is not None: + self.flush_iteration_events() + self.kv_event_adapter.shutdown() + self.kv_event_adapter = None for kv_cache in self.kv_cache_map.values(): kv_cache.close() self.kv_cache_map.clear() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34ebd340a5ee..4253afc0cf4d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -645,7 +645,10 @@ 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, "kv_event_adapter", None) is not None + ) 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 1bd895dbd59b..6a430de2a17e 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -15,10 +15,11 @@ DraftTargetDecodingConfig, DSparkDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, - LookaheadDecodingConfig, MambaStateConfig, - MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, - MoeConfig, MTPDecodingConfig, NGramDecodingConfig, + ExtendedRuntimePerfKnobConfig, KVEventsConfig, + KvCacheConfig, LlmArgs, LookaheadDecodingConfig, + MambaStateConfig, MedusaDecodingConfig, + MiniMaxM3SparseAttentionConfig, MoeConfig, + MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, PrometheusMetricsConfig, ReorderRequestPolicyConfig, RocketSparseAttentionConfig, SADecodingConfig, SAEnhancerConfig, @@ -43,6 +44,7 @@ 'ConversationParams', 'DisaggScheduleStyle', 'KvCacheConfig', + 'KVEventsConfig', 'MambaStateConfig', 'KvCacheRetentionConfig', 'CudaGraphConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 68d9fdbeb207..884f7ae240ca 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3618,6 +3618,46 @@ class MambaStateConfig(StrictBaseModel): "snapshots require KV cache manager V2.") +class KVEventsConfig(StrictBaseModel): + """Configuration for native KV cache event publishing.""" + + enable_kv_cache_events: bool = Field( + default=False, + description="Whether to produce and publish native KV cache events.") + 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", + description="Base ZeroMQ endpoint used to publish KV cache events.") + replay_endpoint: Optional[str] = Field( + default=None, + description="Optional base ZeroMQ endpoint used to replay KV cache events." + ) + buffer_steps: int = Field( + default=10_000, + ge=0, + description="Number of previously published batches retained for replay." + ) + hwm: int = Field(default=100_000, + ge=0, + description="ZeroMQ publisher socket high-water mark.") + max_queue_size: int = Field( + default=100_000, + ge=0, + description="Maximum number of batches queued for background publishing." + ) + 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" + + @PybindMirror.mirror_pybind_fields(_KvCacheConfig) class KvCacheConfig(StrictBaseModel, PybindMirror): """Configuration for the KV cache.""" @@ -4336,6 +4376,10 @@ class BaseLlmArgs(StrictBaseModel): kv_cache_config: KvCacheConfig = Field(default_factory=KvCacheConfig, description="KV cache config.") + kv_events_config: Optional[KVEventsConfig] = Field( + default=None, + description="Native KV cache event publishing configuration.") + enable_chunked_prefill: bool = Field(default=False, description="Enable chunked prefill.") @@ -6072,6 +6116,7 @@ def update_llm_args_with_extra_dict( "attention_dp_config": AttentionDpConfig, "reorder_policy_config": ReorderRequestPolicyConfig, "kv_cache_config": KvCacheConfig, + "kv_events_config": KVEventsConfig, "dwdp_config": DwdpConfig, "multimodal_config": MultimodalConfig, "telemetry_config": TelemetryConfig, diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 63b360f5584b..c023c2259aae 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -28,7 +28,8 @@ from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - KvCacheConfig, LlmArgs, LookaheadDecodingConfig, + KVEventsConfig, KvCacheConfig, LlmArgs, + LookaheadDecodingConfig, MedusaDecodingConfig, MTPDecodingConfig, NGramDecodingConfig, SchedulerConfig, TorchLlmArgs, UserProvidedDecodingConfig, _ModelWrapper, @@ -478,6 +479,7 @@ class LlmBuildStats: 'DecodeCudaGraphConfig', 'EncodeCudaGraphConfig', 'KvCacheConfig', + 'KVEventsConfig', 'CachedModelLoader', 'EagleDecodingConfig', 'Eagle3DecodingConfig', diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f20e02169d62..aa9d63c68639 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -805,6 +805,44 @@ "kind": "categorical", "path": "kv_connector_config.connector" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_events_config.buffer_steps" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_events_config.enable_kv_cache_events" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_events_config.hwm" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_events_config.max_queue_size" + }, + { + "allowed_values": [ + "null", + "zmq" + ], + "annotation": "Optional[Literal['null', 'zmq']]", + "converter": "", + "kind": "categorical", + "path": "kv_events_config.publisher" + }, { "allowed_values": [], "annotation": "Optional[List[int]]", diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py new file mode 100644 index 000000000000..10897bc7e52c --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py @@ -0,0 +1,159 @@ +# 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 +import time + +import msgspec +import zmq + +from tensorrt_llm._torch.pyexecutor.kv_cache_events import KVEventAdapter +from tensorrt_llm.llmapi.llm_args import KVEventsConfig +from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import ( + KVCacheEvent, + KVCacheEventManager, + KVCacheRemovedData, + KVCacheStoredBlockData, + KVCacheStoredData, + UniqueToken, +) + + +def _stored_block(block_hash: int, tokens: list[int]) -> KVCacheStoredBlockData: + return KVCacheStoredBlockData( + block_hash=block_hash, + tokens=[UniqueToken(token) for token in tokens], + cache_level=0, + priority=0, + ) + + +def _unused_tcp_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_native_callback_drains_iteration_without_gathered_buffer(): + """The callback must leave the legacy pull buffer empty without a gather.""" + adapter = KVEventAdapter( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + ) + manager = KVCacheEventManager( + 50_000, + window_size=128, + attention_dp_rank=0, + attention_dp_gather=adapter.publish_local_events, + ) + manager.add_stored_event( + None, + [_stored_block(11, [1, 2, 3, 4])], + ) + + manager.flush_iteration_events() + + assert adapter.enqueued_batches == 1 + assert adapter.enqueued_events == 1 + assert manager.get_latest_events(timeout_ms=0) == [] + adapter.shutdown() + adapter.shutdown() + + +def test_native_zmq_wire_filters_partial_blocks_and_reuses_port(): + """Decode the vLLM wire format while filtering partial block lifecycle.""" + port = _unused_tcp_port() + bind_endpoint = f"tcp://*:{port}" + connect_endpoint = f"tcp://127.0.0.1:{port}" + topic = "kv-events" + context = zmq.Context.instance() + subscriber = context.socket(zmq.SUB) + subscriber.setsockopt_string(zmq.SUBSCRIBE, topic) + subscriber.connect(connect_endpoint) + + adapter = KVEventAdapter( + 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, + ) + time.sleep(0.2) + + full_hash = 2**63 + 5 + partial_hash = 29 + adapter.publish_local_events([ + KVCacheEvent( + event_id=0, + data=KVCacheStoredData( + parent_hash=7, + blocks=[ + _stored_block(full_hash, [1, 2, 3, 4]), + _stored_block(partial_hash, [5, 6]), + ], + ), + window_size=128, + attention_dp_rank=0, + ) + ]) + adapter.publish_local_events([ + KVCacheEvent( + event_id=1, + data=KVCacheRemovedData([full_hash, partial_hash]), + window_size=128, + attention_dp_rank=0, + ) + ]) + + frames = [] + for _ in range(2): + assert subscriber.poll(2_000) + frames.append(subscriber.recv_multipart()) + + assert [frame[0] for frame in frames] == [topic.encode(), topic.encode()] + assert [int.from_bytes(frame[1], "big") for frame in frames] == [0, 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": [-(2**63) + 5], + "parent_block_hash": 7, + "token_ids": [1, 2, 3, 4], + "block_size": 4, + "lora_id": None, + "medium": "GPU", + "lora_name": None, + }] + assert removed_batch[1] == [{ + "type": "BlockRemoved", + "block_hashes": [-(2**63) + 5], + "medium": "GPU", + }] + + adapter.shutdown() + adapter.shutdown() + subscriber.close(linger=0) + + replacement = context.socket(zmq.PUB) + replacement.bind(bind_endpoint) + replacement.close(linger=0) From 195fd1d7c8bf000b5300f9e76efc579223ef50a7 Mon Sep 17 00:00:00 2001 From: Alec Flowers Date: Sat, 25 Jul 2026 21:15:17 -0700 Subject: [PATCH 02/25] perf: optimize native v2 kv event production Signed-off-by: Alec Flowers --- .../_torch/pyexecutor/kv_cache_events.py | 268 ++++++++++++++++++ .../_torch/pyexecutor/kv_cache_manager_v2.py | 17 +- .../test_native_kv_events.py | 172 ++++++----- 3 files changed, 362 insertions(+), 95 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 08ff38ed75e3..a0a02ad4f941 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -35,6 +35,8 @@ from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import ( KVCacheCreatedData, KVCacheEvent, + KVCacheEventDiff, + KVCacheEventManager, KVCacheRemovedData, KVCacheStoredData, KVCacheUpdatedData, @@ -338,6 +340,16 @@ def _to_wire_hash(block_hash: int | str | None) -> ExternalBlockHash | None: f"Invalid hexadecimal KV block hash: {block_hash!r}") from error +def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: + """Convert an existing SHA-256 radix key like vLLM's integer event hashes.""" + if len(block_key) < 8: + raise ValueError("V2 radix block keys must contain at least 8 bytes") + unsigned_hash = int.from_bytes(block_key[-8:], "big", signed=False) + wire_hash = _to_wire_hash(unsigned_hash) + assert isinstance(wire_hash, int) + return wire_hash + + class KVEventAdapter: """Converts local V2 events and publishes one wire batch per iteration.""" @@ -391,6 +403,32 @@ def publish_local_events( ) return [] + def publish_wire_events( + self, + wire_events: list[BlockStored | BlockRemoved | AllBlocksCleared], + ) -> None: + """Enqueue an already-converted local iteration batch.""" + if self._closed or not wire_events: + return + self.local_batches += 1 + self.local_events += len(wire_events) + try: + batch = KVEventBatch( + ts=time.time(), + events=wire_events, + data_parallel_rank=self._rank, + ) + if self._publisher.publish(batch): + self.enqueued_batches += 1 + self.enqueued_events += len(wire_events) + else: + self.dropped_batches += 1 + except Exception: + self.dropped_batches += 1 + logger.exception( + f"Dropping native KV event iteration batch on rank={self._rank}" + ) + def _convert_event( self, event: KVCacheEvent ) -> BlockStored | BlockRemoved | AllBlocksCleared | None: @@ -458,3 +496,233 @@ def shutdown(self) -> None: f"enqueued_batches={self.enqueued_batches} " f"enqueued_events={self.enqueued_events} " f"dropped_batches={self.dropped_batches} kv_event_allgathers=0") + + +class _NativeStoredBlockState: + __slots__ = ("block_hash", ) + + def __init__(self, block_hash: int) -> None: + self.block_hash = block_hash + + +class NativeKVCacheEventManager(KVCacheEventManager): + """Scheduler-local fast path that produces vLLM wire events directly.""" + + def __init__( + self, + adapter: KVEventAdapter, + *, + block_size: int, + max_window_size: int, + max_entries: int = 50_000, + ) -> None: + self._adapter = adapter + 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, _NativeStoredBlockState] = {} + 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.non_target_life_cycles_ignored = 0 + self.dropped_events = 0 + + 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( + "Native KV events require an attention KV cache life cycle") + self._target_life_cycle_id = min(target_ids) + logger.info( + "Native 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_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] + if page_ref is None or page_ref() is None: + return + self._add_full_block(block) + + def add_stored_life_cycle_event_from_block(self, block: Any, + life_cycle_id: int) -> None: + 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, state = self._block_hashes(block) + except ValueError: + self.dropped_events += 1 + self._pending_entries -= 1 + logger.exception( + "Dropping native KV store event with unsupported token data") + return + self._stored_blocks[key] = state + 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 not int: + raise ValueError( + "vLLM-compatible KV events require integer token IDs") + token_ids.append(token) + return token_ids + + def _block_hashes( + self, + block: Any, + ) -> tuple[int, int | None, _NativeStoredBlockState]: + parent = block.prev + is_root_child = getattr(parent, "ordinal", -1) == -1 + block_hash = _vllm_wire_hash_from_radix_key(bytes(block.key)) + parent_hash = None if is_root_child else _vllm_wire_hash_from_radix_key( + bytes(parent.key)) + return block_hash, parent_hash, _NativeStoredBlockState(block_hash) + + def add_removed_event(self, block_hashes: Any) -> None: + 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 + state = self._stored_blocks.pop(block_key, None) + if state is not None: + removed_hashes.append(state.block_hash) + self._add_removed_hashes(removed_hashes) + + def add_removed_life_cycle_event(self, block_hash: bytes, + life_cycle_id: int) -> None: + if int(life_cycle_id) != self._target_life_cycle_id: + self.non_target_life_cycles_ignored += 1 + return + state = self._stored_blocks.pop(block_hash, None) + if state is not None: + self._add_removed_hashes([state.block_hash]) + + def _add_removed_hashes( + self, block_hashes: list[ExternalBlockHash]) -> None: + if not block_hashes: + return + if not self._reserve_entries(len(block_hashes)): + return + 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 native 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 + self._adapter.publish_wire_events(events) + + def get_latest_events( + self, timeout_ms: float | None = None) -> list[KVCacheEvent]: + raise RuntimeError( + "KV cache event polling is unavailable while native publishing " + "is enabled") + + def shutdown(self) -> None: + if self._closed: + return + self.flush_iteration_events() + self._closed = True + logger.info( + "Native KV event fast path " + 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=" + f"{self.non_target_life_cycles_ignored} " + f"dropped_events={self.dropped_events} " + f"kv_event_allgathers=0") diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 21ab6d715117..4cf2cd940fa9 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -82,7 +82,7 @@ from ...mapping import CpType, Mapping from ..utils import maybe_compile from .connectors.kv_cache_connector import KvCacheConnectorManager -from .kv_cache_events import KVEventAdapter +from .kv_cache_events import KVEventAdapter, NativeKVCacheEventManager from .kv_cache_stats import ( KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, @@ -890,14 +890,13 @@ def __init__( block_size=self.tokens_per_block, max_window_size=event_window_size, ) - self.event_manager = KVCacheEventManager( - 50_000, - window_size=event_window_size, - attention_dp_rank=event_rank, - attention_dp_gather=self.kv_event_adapter. - publish_local_events, - hash_algo=kv_cache_event_hash_algo, + self.event_manager = NativeKVCacheEventManager( + self.kv_event_adapter, + block_size=self.tokens_per_block, + max_window_size=event_window_size, ) + logger.info( + "Native KV event fast path reuses V2 radix block hashes") elif self.event_buffer_max_size > 0: if mapping.enable_attention_dp: self.event_manager = KVCacheEventManager( @@ -3460,6 +3459,8 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool def shutdown(self): if self.kv_event_adapter is not None: self.flush_iteration_events() + if isinstance(self.event_manager, NativeKVCacheEventManager): + self.event_manager.shutdown() self.kv_event_adapter.shutdown() self.kv_event_adapter = None for kv_cache in self.kv_cache_map.values(): diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py index 10897bc7e52c..b25b06d61bfe 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py @@ -15,29 +15,16 @@ import socket import time +from types import SimpleNamespace import msgspec import zmq -from tensorrt_llm._torch.pyexecutor.kv_cache_events import KVEventAdapter -from tensorrt_llm.llmapi.llm_args import KVEventsConfig -from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import ( - KVCacheEvent, - KVCacheEventManager, - KVCacheRemovedData, - KVCacheStoredBlockData, - KVCacheStoredData, - UniqueToken, +from tensorrt_llm._torch.pyexecutor.kv_cache_events import ( + KVEventAdapter, + NativeKVCacheEventManager, ) - - -def _stored_block(block_hash: int, tokens: list[int]) -> KVCacheStoredBlockData: - return KVCacheStoredBlockData( - block_hash=block_hash, - tokens=[UniqueToken(token) for token in tokens], - cache_level=0, - priority=0, - ) +from tensorrt_llm.llmapi.llm_args import KVEventsConfig def _unused_tcp_port() -> int: @@ -46,36 +33,8 @@ def _unused_tcp_port() -> int: return int(sock.getsockname()[1]) -def test_native_callback_drains_iteration_without_gathered_buffer(): - """The callback must leave the legacy pull buffer empty without a gather.""" - adapter = KVEventAdapter( - KVEventsConfig(enable_kv_cache_events=True, publisher="null"), - data_parallel_rank=0, - block_size=4, - max_window_size=128, - ) - manager = KVCacheEventManager( - 50_000, - window_size=128, - attention_dp_rank=0, - attention_dp_gather=adapter.publish_local_events, - ) - manager.add_stored_event( - None, - [_stored_block(11, [1, 2, 3, 4])], - ) - - manager.flush_iteration_events() - - assert adapter.enqueued_batches == 1 - assert adapter.enqueued_events == 1 - assert manager.get_latest_events(timeout_ms=0) == [] - adapter.shutdown() - adapter.shutdown() - - -def test_native_zmq_wire_filters_partial_blocks_and_reuses_port(): - """Decode the vLLM wire format while filtering partial block lifecycle.""" +def test_native_fast_path_publishes_only_full_max_window_blocks(): + """Protect radix hash reuse, filtering, wire format, and shutdown.""" port = _unused_tcp_port() bind_endpoint = f"tcp://*:{port}" connect_endpoint = f"tcp://127.0.0.1:{port}" @@ -97,32 +56,60 @@ def test_native_zmq_wire_filters_partial_blocks_and_reuses_port(): block_size=4, max_window_size=128, ) + manager = NativeKVCacheEventManager( + adapter, + block_size=4, + max_window_size=128, + ) + manager.set_layer_group_window_sizes({0: 128, 1: 64}) time.sleep(0.2) - full_hash = 2**63 + 5 - partial_hash = 29 - adapter.publish_local_events([ - KVCacheEvent( - event_id=0, - data=KVCacheStoredData( - parent_hash=7, - blocks=[ - _stored_block(full_hash, [1, 2, 3, 4]), - _stored_block(partial_hash, [5, 6]), - ], - ), - window_size=128, - attention_dp_rank=0, + root = SimpleNamespace(ordinal=-1) + + def block( + key: bytes, + tokens: list[int], + prev: object, + ) -> SimpleNamespace: + max_window_page = object() + smaller_window_page = object() + return SimpleNamespace( + key=key, + tokens=tokens, + prev=prev, + ordinal=getattr(prev, "ordinal", -1) + 1, + storage=[ + lambda: max_window_page, + lambda: smaller_window_page, + ], ) - ]) - adapter.publish_local_events([ - KVCacheEvent( - event_id=1, - data=KVCacheRemovedData([full_hash, partial_hash]), - window_size=128, - attention_dp_rank=0, - ) - ]) + + first_hash = b"\x11" * 24 + b"\x80\x00\x00\x00\x00\x00\x00\x01" + partial_hash = b"\x22" * 32 + second_hash = b"\x33" * 24 + b"\x00\x00\x00\x00\x00\x00\x00\x02" + first_wire_hash = int.from_bytes(first_hash[-8:], "big") + second_wire_hash = int.from_bytes(second_hash[-8:], "big") + first_wire_hash = ( + first_wire_hash - 2**64 + if first_wire_hash >= 2**63 + else first_wire_hash + ) + second_wire_hash = ( + second_wire_hash - 2**64 + if second_wire_hash >= 2**63 + else second_wire_hash + ) + 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): @@ -134,22 +121,33 @@ def test_native_zmq_wire_filters_partial_blocks_and_reuses_port(): 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": [-(2**63) + 5], - "parent_block_hash": 7, - "token_ids": [1, 2, 3, 4], - "block_size": 4, - "lora_id": None, - "medium": "GPU", - "lora_name": None, - }] - assert removed_batch[1] == [{ - "type": "BlockRemoved", - "block_hashes": [-(2**63) + 5], - "medium": "GPU", - }] - + 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 + + manager.shutdown() + manager.shutdown() adapter.shutdown() adapter.shutdown() subscriber.close(linger=0) From 7c58c0d0fe1a2232387f5a967e73ac2cbdcf5343 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 29 Jul 2026 13:25:51 -0700 Subject: [PATCH 03/25] refactor: streamline native KV event manager - pull API (get_latest_events) returns [] instead of raising, so LLM.get_kv_cache_events()/RPC fetch degrade cleanly in native mode instead of erroring and spamming tracebacks every poll - drop the dead generic conversion path (publish_local_events / _convert_event) superseded by the scheduler-local fast path - stop subclassing KVCacheEventManager; implement the event-sink hook interface by duck typing to avoid partially-initialised base state - remove the hardcoded kv_event_allgathers=0 log metric Signed-off-by: tanmayv25 --- .../_torch/pyexecutor/kv_cache_events.py | 291 +++++++----------- .../_torch/pyexecutor/kv_cache_manager_v2.py | 18 +- .../test_native_kv_events.py | 23 +- 3 files changed, 116 insertions(+), 216 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index a0a02ad4f941..7476fa7f8af0 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -32,24 +32,16 @@ from tensorrt_llm.llmapi.llm_args import KVEventsConfig from tensorrt_llm.logger import logger -from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import ( - KVCacheCreatedData, - KVCacheEvent, - KVCacheEventDiff, - KVCacheEventManager, - KVCacheRemovedData, - KVCacheStoredData, - KVCacheUpdatedData, -) +from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import KVCacheEvent, KVCacheEventDiff ExternalBlockHash = bytes | 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] + msgspec.Struct, + array_like=True, # type: ignore[call-arg] + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] ): """vLLM-compatible event batch envelope.""" @@ -59,10 +51,10 @@ class EventBatch( class KVCacheWireEvent( - msgspec.Struct, - omit_defaults=True, # type: ignore[call-arg] - gc=False, # type: ignore[call-arg] - tag=True, + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] + tag=True, ): """Base class for vLLM-compatible KV cache events.""" @@ -152,8 +144,7 @@ def __init__( 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._replay_endpoint = self.offset_endpoint_port(replay_endpoint, self._rank) self._hwm = hwm self._seq_gen = count() self._topic_bytes = topic.encode("utf-8") @@ -169,8 +160,10 @@ def __init__( name=f"trtllm-kv-events-rank-{self._rank}", ) self._thread.start() - logger.info(f"Started native KV event publisher rank={self._rank} " - f"endpoint={self._endpoint} topic={topic!r}") + logger.info( + f"Started native KV event publisher rank={self._rank} " + f"endpoint={self._endpoint} topic={topic!r}" + ) def publish(self, events: EventBatch) -> bool: if not self._running: @@ -183,12 +176,14 @@ def publish(self, events: EventBatch) -> bool: return True except queue.Full: self.dropped_batches += 1 - if self.dropped_batches == 1 or (self.dropped_batches & - (self.dropped_batches - 1) == 0): + if self.dropped_batches == 1 or ( + self.dropped_batches & (self.dropped_batches - 1) == 0 + ): logger.warning( f"Dropping native KV event batch on rank={self._rank} because " "the publisher queue is full; " - f"dropped_batches={self.dropped_batches}") + f"dropped_batches={self.dropped_batches}" + ) return False def shutdown(self) -> None: @@ -205,19 +200,25 @@ def shutdown(self) -> None: if self._thread.is_alive(): logger.warning( f"Native KV event publisher rank={self._rank} did not stop " - f"within {self.SHUTDOWN_TIMEOUT:.1f}s") - logger.info(f"Stopped native KV event publisher rank={self._rank} " - f"enqueued_batches={self.enqueued_batches} " - f"published_batches={self.published_batches} " - f"dropped_batches={self.dropped_batches}") + f"within {self.SHUTDOWN_TIMEOUT:.1f}s" + ) + logger.info( + f"Stopped native 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 self._endpoint is None: raise ValueError("KV event publisher endpoint must not be empty") - if ("*" in self._endpoint or "::" in self._endpoint - or self._endpoint.startswith(("ipc://", "inproc://"))): + if ( + "*" in self._endpoint + or "::" in self._endpoint + or self._endpoint.startswith(("ipc://", "inproc://")) + ): self._pub.bind(self._endpoint) else: self._pub.connect(self._endpoint) @@ -235,8 +236,7 @@ def _publisher_thread(self) -> None: try: self._service_replay() except Exception: - logger.exception( - "Failed to service native KV event replay request") + logger.exception("Failed to service native KV event replay request") try: event = self._event_queue.get(timeout=0.1) except queue.Empty: @@ -247,17 +247,20 @@ def _publisher_thread(self) -> None: seq = next(self._seq_gen) try: payload = encoder.encode(event) - self._pub.send_multipart(( - self._topic_bytes, - seq.to_bytes(8, "big"), - payload, - )) + 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.dropped_batches += 1 - logger.exception(f"Failed to publish native KV event batch " - f"rank={self._rank} seq={seq}") + logger.exception( + f"Failed to publish native KV event batch rank={self._rank} seq={seq}" + ) time.sleep(0.1) finally: self._event_queue.task_done() @@ -276,18 +279,19 @@ def _service_replay(self) -> None: 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"", + 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: + def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | None: """Apply vLLM's base-port-plus-rank endpoint convention.""" if not endpoint or data_parallel_rank == 0: return endpoint @@ -296,7 +300,7 @@ def offset_endpoint_port(endpoint: str | None, if "tcp" in endpoint and ":" in endpoint: last_colon_idx = endpoint.rfind(":") base_addr = endpoint[:last_colon_idx] - base_port = int(endpoint[last_colon_idx + 1:]) + base_port = int(endpoint[last_colon_idx + 1 :]) new_port = base_port + data_parallel_rank if new_port > 65_535: raise ValueError( @@ -306,8 +310,7 @@ def offset_endpoint_port(endpoint: str | None, raise ValueError("Invalid endpoint: must contain 'inproc' or 'tcp'") -def create_event_publisher(config: KVEventsConfig, - data_parallel_rank: int) -> EventPublisher: +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) @@ -336,8 +339,7 @@ def _to_wire_hash(block_hash: int | str | None) -> ExternalBlockHash | None: try: return bytes.fromhex(block_hash) except ValueError as error: - raise ValueError( - f"Invalid hexadecimal KV block hash: {block_hash!r}") from error + raise ValueError(f"Invalid hexadecimal KV block hash: {block_hash!r}") from error def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: @@ -358,14 +360,9 @@ def __init__( config: KVEventsConfig, *, data_parallel_rank: int, - block_size: int, - max_window_size: int, ) -> None: self._rank = data_parallel_rank - self._block_size = block_size - self._max_window_size = max_window_size self._publisher = create_event_publisher(config, data_parallel_rank) - self._partial_block_hashes: set[int | str] = set() self._closed = False self.local_batches = 0 self.local_events = 0 @@ -373,36 +370,6 @@ def __init__( self.enqueued_events = 0 self.dropped_batches = 0 - def publish_local_events( - self, events: list[KVCacheEvent]) -> list[list[KVCacheEvent]]: - """Publish local events and return no gathered events to the manager.""" - if self._closed or not events: - return [] - self.local_batches += 1 - self.local_events += len(events) - try: - wire_events = [ - wire_event for event in events - if (wire_event := self._convert_event(event)) is not None - ] - if wire_events: - batch = KVEventBatch( - ts=time.time(), - events=wire_events, - data_parallel_rank=self._rank, - ) - if self._publisher.publish(batch): - self.enqueued_batches += 1 - self.enqueued_events += len(wire_events) - else: - self.dropped_batches += 1 - except Exception: - self.dropped_batches += 1 - logger.exception( - f"Dropping native KV event iteration batch on rank={self._rank}" - ) - return [] - def publish_wire_events( self, wire_events: list[BlockStored | BlockRemoved | AllBlocksCleared], @@ -425,63 +392,7 @@ def publish_wire_events( self.dropped_batches += 1 except Exception: self.dropped_batches += 1 - logger.exception( - f"Dropping native KV event iteration batch on rank={self._rank}" - ) - - def _convert_event( - self, event: KVCacheEvent - ) -> BlockStored | BlockRemoved | AllBlocksCleared | None: - if event.window_size != self._max_window_size: - return None - data = event.data - if isinstance(data, (KVCacheCreatedData, KVCacheUpdatedData)): - return None - if isinstance(data, KVCacheStoredData): - block_hashes: list[ExternalBlockHash] = [] - token_ids: list[int] = [] - for block in data.blocks: - num_tokens = len(block.tokens) - if num_tokens > self._block_size: - raise ValueError( - f"KV block has {num_tokens} tokens, expected at most " - f"{self._block_size}") - if num_tokens < self._block_size: - self._partial_block_hashes.add(block.block_hash) - break - block_token_ids = [token.token_id for token in block.tokens] - if any(not isinstance(token_id, int) - for token_id in block_token_ids): - raise ValueError( - "vLLM-compatible KV events require integer token IDs") - wire_hash = _to_wire_hash(block.block_hash) - assert wire_hash is not None - block_hashes.append(wire_hash) - token_ids.extend(block_token_ids) - if not block_hashes: - return None - return BlockStored( - block_hashes=block_hashes, - parent_block_hash=_to_wire_hash(data.parent_hash), - token_ids=token_ids, - block_size=self._block_size, - lora_id=None, - medium="GPU", - lora_name=None, - ) - if isinstance(data, KVCacheRemovedData): - block_hashes: list[ExternalBlockHash] = [] - for block_hash in data.block_hashes: - if block_hash in self._partial_block_hashes: - self._partial_block_hashes.remove(block_hash) - continue - wire_hash = _to_wire_hash(block_hash) - assert wire_hash is not None - block_hashes.append(wire_hash) - if not block_hashes: - return None - return BlockRemoved(block_hashes=block_hashes, medium="GPU") - return None + logger.exception(f"Dropping native KV event iteration batch on rank={self._rank}") def shutdown(self) -> None: """Close the publisher once and report direct-path counters.""" @@ -495,18 +406,26 @@ def shutdown(self) -> None: f"local_events={self.local_events} " f"enqueued_batches={self.enqueued_batches} " f"enqueued_events={self.enqueued_events} " - f"dropped_batches={self.dropped_batches} kv_event_allgathers=0") + f"dropped_batches={self.dropped_batches}" + ) class _NativeStoredBlockState: - __slots__ = ("block_hash", ) + __slots__ = ("block_hash",) def __init__(self, block_hash: int) -> None: self.block_hash = block_hash -class NativeKVCacheEventManager(KVCacheEventManager): - """Scheduler-local fast path that produces vLLM wire events directly.""" +class NativeKVCacheEventManager: + """Scheduler-local fast path that produces vLLM wire events 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, @@ -522,8 +441,7 @@ def __init__( self._max_entries = max_entries self._target_life_cycle_id: int | None = None self._stored_blocks: dict[bytes, _NativeStoredBlockState] = {} - self._pending_events: list[ - BlockStored | BlockRemoved | AllBlocksCleared] = [] + self._pending_events: list[BlockStored | BlockRemoved | AllBlocksCleared] = [] self._pending_entries = 0 self._closed = False self.stored_blocks = 0 @@ -532,8 +450,7 @@ def __init__( self.non_target_life_cycles_ignored = 0 self.dropped_events = 0 - def set_layer_group_window_sizes(self, - window_sizes: dict[int, int]) -> None: + 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() @@ -547,13 +464,13 @@ def set_layer_group_window_sizes(self, if window_size == largest_window ] if not target_ids: - raise ValueError( - "Native KV events require an attention KV cache life cycle") + raise ValueError("Native KV events require an attention KV cache life cycle") self._target_life_cycle_id = min(target_ids) logger.info( "Native KV event fast path selected " f"lifecycle_id={self._target_life_cycle_id} " - f"window_size={self._max_window_size}") + f"window_size={self._max_window_size}" + ) def add_created_event( self, @@ -562,6 +479,11 @@ def add_created_event( ) -> None: return + def add_stored_event(self, *args: Any, **kwargs: Any) -> None: + # Native 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 @@ -573,8 +495,7 @@ def add_stored_block_event_from_block(self, block: Any) -> None: return self._add_full_block(block) - def add_stored_life_cycle_event_from_block(self, block: Any, - life_cycle_id: int) -> None: + def add_stored_life_cycle_event_from_block(self, block: Any, life_cycle_id: int) -> None: if int(life_cycle_id) != self._target_life_cycle_id: self.non_target_life_cycles_ignored += 1 return @@ -595,15 +516,12 @@ def _add_full_block(self, block: Any) -> None: except ValueError: self.dropped_events += 1 self._pending_entries -= 1 - logger.exception( - "Dropping native KV store event with unsupported token data") + logger.exception("Dropping native KV store event with unsupported token data") return self._stored_blocks[key] = state - if self._pending_events and isinstance(self._pending_events[-1], - BlockStored): + 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: + 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 @@ -617,7 +535,8 @@ def _add_full_block(self, block: Any) -> None: lora_id=None, medium="GPU", lora_name=None, - )) + ) + ) self.stored_blocks += 1 @staticmethod @@ -625,8 +544,7 @@ def _token_ids(tokens: Any) -> list[int]: token_ids: list[int] = [] for token in tokens: if type(token) is not int: - raise ValueError( - "vLLM-compatible KV events require integer token IDs") + raise ValueError("vLLM-compatible KV events require integer token IDs") token_ids.append(token) return token_ids @@ -637,13 +555,12 @@ def _block_hashes( parent = block.prev is_root_child = getattr(parent, "ordinal", -1) == -1 block_hash = _vllm_wire_hash_from_radix_key(bytes(block.key)) - parent_hash = None if is_root_child else _vllm_wire_hash_from_radix_key( - bytes(parent.key)) + parent_hash = None if is_root_child else _vllm_wire_hash_from_radix_key(bytes(parent.key)) return block_hash, parent_hash, _NativeStoredBlockState(block_hash) def add_removed_event(self, block_hashes: Any) -> None: if isinstance(block_hashes, (bytes, str, int)): - block_hashes = (block_hashes, ) + block_hashes = (block_hashes,) removed_hashes: list[ExternalBlockHash] = [] for block_key in block_hashes: if not isinstance(block_key, bytes): @@ -653,8 +570,7 @@ def add_removed_event(self, block_hashes: Any) -> None: removed_hashes.append(state.block_hash) self._add_removed_hashes(removed_hashes) - def add_removed_life_cycle_event(self, block_hash: bytes, - life_cycle_id: int) -> None: + def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: if int(life_cycle_id) != self._target_life_cycle_id: self.non_target_life_cycles_ignored += 1 return @@ -662,18 +578,15 @@ def add_removed_life_cycle_event(self, block_hash: bytes, if state is not None: self._add_removed_hashes([state.block_hash]) - def _add_removed_hashes( - self, block_hashes: list[ExternalBlockHash]) -> None: + def _add_removed_hashes(self, block_hashes: list[ExternalBlockHash]) -> None: if not block_hashes: return if not self._reserve_entries(len(block_hashes)): return - if self._pending_events and isinstance(self._pending_events[-1], - BlockRemoved): + 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._pending_events.append(BlockRemoved(block_hashes=block_hashes, medium="GPU")) self.removed_blocks += len(block_hashes) def add_updated_event( @@ -692,10 +605,12 @@ def _reserve_entries(self, num_entries: int) -> bool: return True self.dropped_events += num_entries if self.dropped_events == num_entries or ( - self.dropped_events & (self.dropped_events - 1) == 0): + self.dropped_events & (self.dropped_events - 1) == 0 + ): logger.warning( "Dropping native KV events because the per-iteration safety " - f"cap was exceeded; dropped_events={self.dropped_events}") + f"cap was exceeded; dropped_events={self.dropped_events}" + ) return False def flush_iteration_events(self) -> None: @@ -706,11 +621,11 @@ def flush_iteration_events(self) -> None: self._pending_entries = 0 self._adapter.publish_wire_events(events) - def get_latest_events( - self, timeout_ms: float | None = None) -> list[KVCacheEvent]: - raise RuntimeError( - "KV cache event polling is unavailable while native publishing " - "is enabled") + def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: + # Native publishing pushes events out-of-band, so the pull API has + # nothing to return. Return empty instead of raising so callers of the + # legacy polling path degrade cleanly rather than erroring. + return [] def shutdown(self) -> None: if self._closed: @@ -724,5 +639,5 @@ def shutdown(self) -> None: f"partial_blocks_suppressed={self.partial_blocks_suppressed} " f"non_target_life_cycles_ignored=" f"{self.non_target_life_cycles_ignored} " - f"dropped_events={self.dropped_events} " - f"kv_event_allgathers=0") + f"dropped_events={self.dropped_events}" + ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 4cf2cd940fa9..da03df136a9b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -37,7 +37,7 @@ IndexMapper, copy_batch_block_offsets_to_device, ) -from tensorrt_llm.llmapi.llm_args import KVEventsConfig, 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, @@ -868,35 +868,29 @@ 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 + self.event_manager: Optional[KVCacheEventManager | NativeKVCacheEventManager] = None self.kv_event_adapter: Optional[KVEventAdapter] = None native_events_enabled = ( - kv_events_config is not None - and kv_events_config.enable_kv_cache_events + kv_events_config is not None and kv_events_config.enable_kv_cache_events ) if native_events_enabled: if mapping.pp_size > 1: - raise ValueError( - "Native KV events do not support pipeline parallelism") + raise ValueError("Native KV events do not support pipeline parallelism") if mapping.cp_size > 1: - raise ValueError( - "Native KV events do not support context parallelism") + raise ValueError("Native KV events do not support context parallelism") assert kv_events_config is not None if mapping.enable_attention_dp or mpi_rank() == 0: event_rank = mapping.rank if mapping.enable_attention_dp else 0 self.kv_event_adapter = KVEventAdapter( kv_events_config, data_parallel_rank=event_rank, - block_size=self.tokens_per_block, - max_window_size=event_window_size, ) self.event_manager = NativeKVCacheEventManager( self.kv_event_adapter, block_size=self.tokens_per_block, max_window_size=event_window_size, ) - logger.info( - "Native KV event fast path reuses V2 radix block hashes") + logger.info("Native KV event fast path reuses V2 radix block hashes") elif self.event_buffer_max_size > 0: if mapping.enable_attention_dp: self.event_manager = KVCacheEventManager( diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py index b25b06d61bfe..80a307acab00 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py @@ -20,10 +20,7 @@ import msgspec import zmq -from tensorrt_llm._torch.pyexecutor.kv_cache_events import ( - KVEventAdapter, - NativeKVCacheEventManager, -) +from tensorrt_llm._torch.pyexecutor.kv_cache_events import KVEventAdapter, NativeKVCacheEventManager from tensorrt_llm.llmapi.llm_args import KVEventsConfig @@ -53,8 +50,6 @@ def test_native_fast_path_publishes_only_full_max_window_blocks(): max_queue_size=8, ), data_parallel_rank=0, - block_size=4, - max_window_size=128, ) manager = NativeKVCacheEventManager( adapter, @@ -89,16 +84,8 @@ def block( second_hash = b"\x33" * 24 + b"\x00\x00\x00\x00\x00\x00\x00\x02" first_wire_hash = int.from_bytes(first_hash[-8:], "big") second_wire_hash = int.from_bytes(second_hash[-8:], "big") - first_wire_hash = ( - first_wire_hash - 2**64 - if first_wire_hash >= 2**63 - else first_wire_hash - ) - second_wire_hash = ( - second_wire_hash - 2**64 - if second_wire_hash >= 2**63 - else second_wire_hash - ) + first_wire_hash = first_wire_hash - 2**64 if first_wire_hash >= 2**63 else first_wire_hash + second_wire_hash = second_wire_hash - 2**64 if second_wire_hash >= 2**63 else second_wire_hash 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) @@ -146,6 +133,10 @@ def block( assert manager.non_target_life_cycles_ignored == 1 assert manager.dropped_events == 0 + # Native publishing pushes events out-of-band, so the legacy pull API must + # degrade to an empty result rather than raising. + assert manager.get_latest_events() == [] + manager.shutdown() manager.shutdown() adapter.shutdown() From 5a66ee1a8c46b8a253387438c1f4ca6fb2b42881 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 29 Jul 2026 13:25:52 -0700 Subject: [PATCH 04/25] refactor: nest KV events config under KvCacheConfig Unify the KV-event configuration surface: move kv_events_config from a top-level TorchLlmArgs field into KvCacheConfig, alongside the existing event_buffer_max_size / attention_dp_events_gather_period_ms knobs, so there is a single place to configure KV-cache events. Mark the field prototype and warn when native events are requested on a non-V2 KV cache manager (where they are silently unsupported). Users now set kv_cache_config.kv_events_config instead of a top-level kv_events_config. Signed-off-by: tanmayv25 --- tensorrt_llm/_torch/pyexecutor/_util.py | 7 +- tensorrt_llm/llmapi/__init__.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 20 +++-- .../usage/llm_args_golden_manifest.json | 76 +++++++++---------- 4 files changed, 58 insertions(+), 49 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 4106f0260f06..1e125ad78eb6 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1144,7 +1144,7 @@ def _create_kv_cache_manager( is_disagg=self._is_disagg, kv_events_config=None if estimating_kv_cache or model_engine.is_draft_model else - self._llm_args.kv_events_config, + self._llm_args.kv_cache_config.kv_events_config, ) if not self._skip_est: @@ -1994,6 +1994,11 @@ def _create_kv_cache_manager( if issubclass(kv_cache_manager_cls, KVCacheManagerV2): manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats 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 native 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/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 6a430de2a17e..48fd5a0e048d 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -15,8 +15,8 @@ DraftTargetDecodingConfig, DSparkDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - ExtendedRuntimePerfKnobConfig, KVEventsConfig, - KvCacheConfig, LlmArgs, LookaheadDecodingConfig, + ExtendedRuntimePerfKnobConfig, KvCacheConfig, + KVEventsConfig, LlmArgs, LookaheadDecodingConfig, MambaStateConfig, MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, NGramDecodingConfig, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 884f7ae240ca..0835e158bfe9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3634,8 +3634,8 @@ class KVEventsConfig(StrictBaseModel): description="Base ZeroMQ endpoint used to publish KV cache events.") replay_endpoint: Optional[str] = Field( default=None, - description="Optional base ZeroMQ endpoint used to replay KV cache events." - ) + description= + "Optional base ZeroMQ endpoint used to replay KV cache events.") buffer_steps: int = Field( default=10_000, ge=0, @@ -3651,7 +3651,8 @@ class KVEventsConfig(StrictBaseModel): ) topic: str = Field( default="", - description="ZeroMQ subscription topic used for KV cache event batches.") + description="ZeroMQ subscription topic used for KV cache event batches." + ) def model_post_init(self, __context) -> None: if self.publisher is None: @@ -3725,6 +3726,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= + "Native 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 legacy event_buffer_max_size gather/poll path.") enable_partial_reuse: bool = Field( default=True, description= @@ -4376,10 +4385,6 @@ class BaseLlmArgs(StrictBaseModel): kv_cache_config: KvCacheConfig = Field(default_factory=KvCacheConfig, description="KV cache config.") - kv_events_config: Optional[KVEventsConfig] = Field( - default=None, - description="Native KV cache event publishing configuration.") - enable_chunked_prefill: bool = Field(default=False, description="Enable chunked prefill.") @@ -6116,7 +6121,6 @@ def update_llm_args_with_extra_dict( "attention_dp_config": AttentionDpConfig, "reorder_policy_config": ReorderRequestPolicyConfig, "kv_cache_config": KvCacheConfig, - "kv_events_config": KVEventsConfig, "dwdp_config": DwdpConfig, "multimodal_config": MultimodalConfig, "telemetry_config": TelemetryConfig, diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index aa9d63c68639..38cad09522d7 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -689,6 +689,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", @@ -805,44 +843,6 @@ "kind": "categorical", "path": "kv_connector_config.connector" }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_events_config.buffer_steps" - }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_events_config.enable_kv_cache_events" - }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_events_config.hwm" - }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_events_config.max_queue_size" - }, - { - "allowed_values": [ - "null", - "zmq" - ], - "annotation": "Optional[Literal['null', 'zmq']]", - "converter": "", - "kind": "categorical", - "path": "kv_events_config.publisher" - }, { "allowed_values": [], "annotation": "Optional[List[int]]", From 92f92dee292cbc0ef35399547dde62d8fcd3069d Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 29 Jul 2026 15:09:15 -0700 Subject: [PATCH 05/25] fix: address independent review of native KV events - get_latest_events: remove the raise in KVCacheManagerV2's wrapper so native mode returns [] on the pull path (the earlier fix only touched the inner manager, which the wrapper shadowed) - never drop block-removal events under the per-iteration entry cap; a dropped removal permanently desyncs the consumer (block reported stored but never removed). Add a socket-free regression test. - inline the single-use _to_wire_hash helper and drop its unreachable branches Signed-off-by: tanmayv25 --- .../_torch/pyexecutor/kv_cache_events.py | 29 +++-------- .../_torch/pyexecutor/kv_cache_manager_v2.py | 7 ++- .../test_native_kv_events.py | 50 ++++++++++++++++++- 3 files changed, 60 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 7476fa7f8af0..168bf749318b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -327,29 +327,13 @@ def create_event_publisher(config: KVEventsConfig, data_parallel_rank: int) -> E raise ValueError(f"Unsupported KV event publisher: {config.publisher!r}") -def _to_wire_hash(block_hash: int | str | None) -> ExternalBlockHash | None: - if block_hash is None: - return None - if isinstance(block_hash, int): - if block_hash >= 2**63: - return block_hash - 2**64 - if block_hash < -(2**63): - return ((block_hash + 2**63) % 2**64) - 2**63 - return block_hash - try: - return bytes.fromhex(block_hash) - except ValueError as error: - raise ValueError(f"Invalid hexadecimal KV block hash: {block_hash!r}") from error - - def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: - """Convert an existing SHA-256 radix key like vLLM's integer event hashes.""" + """Reuse an existing SHA-256 radix key as vLLM's signed integer event hash.""" if len(block_key) < 8: raise ValueError("V2 radix block keys must contain at least 8 bytes") unsigned_hash = int.from_bytes(block_key[-8:], "big", signed=False) - wire_hash = _to_wire_hash(unsigned_hash) - assert isinstance(wire_hash, int) - return wire_hash + # Reinterpret the low 64 bits as signed two's-complement for the wire format. + return unsigned_hash - 2**64 if unsigned_hash >= 2**63 else unsigned_hash class KVEventAdapter: @@ -581,8 +565,11 @@ def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> def _add_removed_hashes(self, block_hashes: list[ExternalBlockHash]) -> None: if not block_hashes: return - if not self._reserve_entries(len(block_hashes)): - return + # Removals are never dropped by the per-iteration cap: each hash was + # already reported as stored, so dropping its removal would leave the + # consumer believing the block is resident forever. They are bounded by + # the previously-stored set, so they cannot run away. + self._pending_entries += len(block_hashes) if self._pending_events and isinstance(self._pending_events[-1], BlockRemoved): self._pending_events[-1].block_hashes.extend(block_hashes) else: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index da03df136a9b..1c85239e99aa 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2921,10 +2921,9 @@ def flush_iteration_events(self): self.event_manager.flush_iteration_events() def get_latest_events(self, timeout_ms: Optional[float] = None): - if self.kv_event_adapter is not None: - raise RuntimeError( - "KV cache event polling is unavailable while native publishing is enabled" - ) + # Native publishing pushes events out-of-band; in that mode the event + # manager's get_latest_events returns [], so the legacy pull path + # degrades cleanly instead of raising. if self.event_manager is None: return [] return self.event_manager.get_latest_events(timeout_ms) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py index 80a307acab00..b15ccca90db6 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py @@ -20,7 +20,11 @@ import msgspec import zmq -from tensorrt_llm._torch.pyexecutor.kv_cache_events import KVEventAdapter, NativeKVCacheEventManager +from tensorrt_llm._torch.pyexecutor.kv_cache_events import ( + BlockRemoved, + KVEventAdapter, + NativeKVCacheEventManager, +) from tensorrt_llm.llmapi.llm_args import KVEventsConfig @@ -146,3 +150,47 @@ def block( replacement = context.socket(zmq.PUB) replacement.bind(bind_endpoint) replacement.close(linger=0) + + +def test_native_removals_are_never_dropped_by_the_entry_cap(): + """Removals must survive the per-iteration cap or the consumer desyncs.""" + adapter = KVEventAdapter( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + ) + manager = NativeKVCacheEventManager( + adapter, + block_size=2, + max_window_size=128, + max_entries=2, + ) + manager.set_layer_group_window_sizes({0: 128}) + + root = SimpleNamespace(ordinal=-1) + + def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: + page = object() + 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]) + + removed = [event for event in manager._pending_events if isinstance(event, BlockRemoved)] + assert manager.removed_blocks == 2 + assert sum(len(event.block_hashes) for event in removed) == 2 + + manager.shutdown() + adapter.shutdown() From 902c2ee41741d7b348914d1300cbe3d93e995573 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 29 Jul 2026 15:15:59 -0700 Subject: [PATCH 06/25] refactor: collapse KVEventAdapter into NativeKVCacheEventManager The adapter was a thin envelope: it wrapped wire events into a batch and owned the publisher lifecycle, duplicating the publisher's enqueued/dropped counters. Fold it into the manager, which now creates and owns the publisher directly and builds the batch in flush_iteration_events. Replace the kv_event_adapter presence flag with a native_kv_events_enabled property on KVCacheManagerV2. Signed-off-by: tanmayv25 --- .../_torch/pyexecutor/kv_cache_events.py | 91 ++++++------------- .../_torch/pyexecutor/kv_cache_manager_v2.py | 21 ++--- tensorrt_llm/_torch/pyexecutor/py_executor.py | 5 +- .../test_native_kv_events.py | 19 +--- 4 files changed, 41 insertions(+), 95 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 168bf749318b..534d8ef88a97 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -336,64 +336,6 @@ def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: return unsigned_hash - 2**64 if unsigned_hash >= 2**63 else unsigned_hash -class KVEventAdapter: - """Converts local V2 events and publishes one wire batch per iteration.""" - - def __init__( - self, - config: KVEventsConfig, - *, - data_parallel_rank: int, - ) -> None: - self._rank = data_parallel_rank - self._publisher = create_event_publisher(config, data_parallel_rank) - self._closed = False - self.local_batches = 0 - self.local_events = 0 - self.enqueued_batches = 0 - self.enqueued_events = 0 - self.dropped_batches = 0 - - def publish_wire_events( - self, - wire_events: list[BlockStored | BlockRemoved | AllBlocksCleared], - ) -> None: - """Enqueue an already-converted local iteration batch.""" - if self._closed or not wire_events: - return - self.local_batches += 1 - self.local_events += len(wire_events) - try: - batch = KVEventBatch( - ts=time.time(), - events=wire_events, - data_parallel_rank=self._rank, - ) - if self._publisher.publish(batch): - self.enqueued_batches += 1 - self.enqueued_events += len(wire_events) - else: - self.dropped_batches += 1 - except Exception: - self.dropped_batches += 1 - logger.exception(f"Dropping native KV event iteration batch on rank={self._rank}") - - def shutdown(self) -> None: - """Close the publisher once and report direct-path counters.""" - if self._closed: - return - self._closed = True - self._publisher.shutdown() - logger.info( - f"Native KV events rank={self._rank} " - f"local_batches={self.local_batches} " - f"local_events={self.local_events} " - f"enqueued_batches={self.enqueued_batches} " - f"enqueued_events={self.enqueued_events} " - f"dropped_batches={self.dropped_batches}" - ) - - class _NativeStoredBlockState: __slots__ = ("block_hash",) @@ -413,13 +355,15 @@ class NativeKVCacheEventManager: def __init__( self, - adapter: KVEventAdapter, + config: KVEventsConfig, *, + data_parallel_rank: int, block_size: int, max_window_size: int, max_entries: int = 50_000, ) -> None: - self._adapter = adapter + 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 @@ -433,6 +377,9 @@ def __init__( self.partial_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 set_layer_group_window_sizes(self, window_sizes: dict[int, int]) -> None: target_ids = [ @@ -606,7 +553,20 @@ def flush_iteration_events(self) -> None: events = self._pending_events self._pending_events = [] self._pending_entries = 0 - self._adapter.publish_wire_events(events) + 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.exception(f"Dropping native KV event iteration batch on rank={self._rank}") def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: # Native publishing pushes events out-of-band, so the pull API has @@ -619,12 +579,15 @@ def shutdown(self) -> None: return self.flush_iteration_events() self._closed = True + self._publisher.shutdown() logger.info( "Native 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=" - f"{self.non_target_life_cycles_ignored} " - f"dropped_events={self.dropped_events}" + 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 1c85239e99aa..a2c1c2ea80b5 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -82,7 +82,7 @@ from ...mapping import CpType, Mapping from ..utils import maybe_compile from .connectors.kv_cache_connector import KvCacheConnectorManager -from .kv_cache_events import KVEventAdapter, NativeKVCacheEventManager +from .kv_cache_events import NativeKVCacheEventManager from .kv_cache_stats import ( KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, @@ -869,7 +869,6 @@ def __init__( for window_size in self.max_attention_window_vec ) self.event_manager: Optional[KVCacheEventManager | NativeKVCacheEventManager] = None - self.kv_event_adapter: Optional[KVEventAdapter] = None native_events_enabled = ( kv_events_config is not None and kv_events_config.enable_kv_cache_events ) @@ -881,12 +880,9 @@ def __init__( assert kv_events_config is not None if mapping.enable_attention_dp or mpi_rank() == 0: event_rank = mapping.rank if mapping.enable_attention_dp else 0 - self.kv_event_adapter = KVEventAdapter( + self.event_manager = NativeKVCacheEventManager( kv_events_config, data_parallel_rank=event_rank, - ) - self.event_manager = NativeKVCacheEventManager( - self.kv_event_adapter, block_size=self.tokens_per_block, max_window_size=event_window_size, ) @@ -2928,6 +2924,10 @@ def get_latest_events(self, timeout_ms: Optional[float] = None): return [] return self.event_manager.get_latest_events(timeout_ms) + @property + def native_kv_events_enabled(self) -> bool: + return isinstance(self.event_manager, NativeKVCacheEventManager) + def get_iteration_stats(self): if not self.enable_stats: return None @@ -3450,12 +3450,9 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool return bool(has_invalid_values) def shutdown(self): - if self.kv_event_adapter is not None: - self.flush_iteration_events() - if isinstance(self.event_manager, NativeKVCacheEventManager): - self.event_manager.shutdown() - self.kv_event_adapter.shutdown() - self.kv_event_adapter = None + if isinstance(self.event_manager, NativeKVCacheEventManager): + self.event_manager.shutdown() + self.event_manager = None for kv_cache in self.kv_cache_map.values(): kv_cache.close() self.kv_cache_map.clear() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4253afc0cf4d..500f46aa4b22 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -646,9 +646,8 @@ def __init__( 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 - or getattr(self.kv_cache_manager, "kv_event_adapter", None) is not None - ) + self.kv_cache_manager.event_buffer_max_size > 0 or getattr( + self.kv_cache_manager, "native_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/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py index b15ccca90db6..f21f26473d4c 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py @@ -20,11 +20,7 @@ import msgspec import zmq -from tensorrt_llm._torch.pyexecutor.kv_cache_events import ( - BlockRemoved, - KVEventAdapter, - NativeKVCacheEventManager, -) +from tensorrt_llm._torch.pyexecutor.kv_cache_events import BlockRemoved, NativeKVCacheEventManager from tensorrt_llm.llmapi.llm_args import KVEventsConfig @@ -45,7 +41,7 @@ def test_native_fast_path_publishes_only_full_max_window_blocks(): subscriber.setsockopt_string(zmq.SUBSCRIBE, topic) subscriber.connect(connect_endpoint) - adapter = KVEventAdapter( + manager = NativeKVCacheEventManager( KVEventsConfig( enable_kv_cache_events=True, publisher="zmq", @@ -54,9 +50,6 @@ def test_native_fast_path_publishes_only_full_max_window_blocks(): max_queue_size=8, ), data_parallel_rank=0, - ) - manager = NativeKVCacheEventManager( - adapter, block_size=4, max_window_size=128, ) @@ -143,8 +136,6 @@ def block( manager.shutdown() manager.shutdown() - adapter.shutdown() - adapter.shutdown() subscriber.close(linger=0) replacement = context.socket(zmq.PUB) @@ -154,12 +145,9 @@ def block( def test_native_removals_are_never_dropped_by_the_entry_cap(): """Removals must survive the per-iteration cap or the consumer desyncs.""" - adapter = KVEventAdapter( + manager = NativeKVCacheEventManager( KVEventsConfig(enable_kv_cache_events=True, publisher="null"), data_parallel_rank=0, - ) - manager = NativeKVCacheEventManager( - adapter, block_size=2, max_window_size=128, max_entries=2, @@ -193,4 +181,3 @@ def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: assert sum(len(event.block_hashes) for event in removed) == 2 manager.shutdown() - adapter.shutdown() From bee841395584bdcba62f882e2be9479f3420f8ad Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 29 Jul 2026 16:49:14 -0700 Subject: [PATCH 07/25] fix: address xhigh code-review findings for native KV events Config validation: - require hwm/max_queue_size/buffer_steps > 0 (0 inverts ZMQ/Queue semantics into 'unlimited', defeating backpressure) - reject empty endpoint; document that co-located engines need distinct ports Endpoint handling: - PUB socket always binds (tcp/ipc/inproc) instead of connect()ing explicit hosts like tcp://0.0.0.0 (which silently dropped all events) - offset_endpoint_port handles ipc:// for DP rank>0 Correctness / teardown: - exclude non-attention (SSM) life cycles from native event target selection so hybrid Mamba models do not emit a corrupt/empty attention-reuse stream - warn when both legacy event_buffer_max_size and native events are enabled - removals no longer consume the store entry budget (was starving BlockStored) - guard removed-event hooks on _closed; split dropped_batches into two single-writer counters (lock-free); shut the event manager down last in teardown and stop nulling it (avoids a get/flush None race); tear the publisher down if manager construction fails after it bound Signed-off-by: tanmayv25 --- .../_torch/pyexecutor/kv_cache_events.py | 55 ++++++---- .../_torch/pyexecutor/kv_cache_manager_v2.py | 100 ++++++++++++------ tensorrt_llm/llmapi/llm_args.py | 19 ++-- 3 files changed, 111 insertions(+), 63 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 534d8ef88a97..77821f5f9c67 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -152,7 +152,8 @@ def __init__( self._shutdown_lock = threading.Lock() self.enqueued_batches = 0 self.published_batches = 0 - self.dropped_batches = 0 + self._queue_full_drops = 0 + self._send_error_drops = 0 self._socket_setup() self._thread = threading.Thread( target=self._publisher_thread, @@ -165,6 +166,13 @@ def __init__( f"endpoint={self._endpoint} topic={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 @@ -175,10 +183,9 @@ def publish(self, events: EventBatch) -> bool: self.enqueued_batches += 1 return True except queue.Full: - self.dropped_batches += 1 - if self.dropped_batches == 1 or ( - self.dropped_batches & (self.dropped_batches - 1) == 0 - ): + self._queue_full_drops += 1 + drops = self._queue_full_drops + if drops == 1 or (drops & (drops - 1) == 0): logger.warning( f"Dropping native KV event batch on rank={self._rank} because " "the publisher queue is full; " @@ -212,16 +219,15 @@ def shutdown(self) -> None: def _socket_setup(self) -> None: self._pub = self._ctx.socket(zmq.PUB) self._pub.set_hwm(self._hwm) - if self._endpoint is None: + if not self._endpoint: raise ValueError("KV event publisher endpoint must not be empty") - if ( - "*" in self._endpoint - or "::" in self._endpoint - or self._endpoint.startswith(("ipc://", "inproc://")) - ): - self._pub.bind(self._endpoint) - else: - self._pub.connect(self._endpoint) + 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) @@ -257,7 +263,7 @@ def _publisher_thread(self) -> None: self._buffer.append((seq, payload)) self.published_batches += 1 except Exception: - self.dropped_batches += 1 + self._send_error_drops += 1 logger.exception( f"Failed to publish native KV event batch rank={self._rank} seq={seq}" ) @@ -295,7 +301,8 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | """Apply vLLM's base-port-plus-rank endpoint convention.""" if not endpoint or data_parallel_rank == 0: return endpoint - if "inproc" in endpoint: + # ipc/inproc have no port; give each rank a distinct suffix instead. + if "inproc" in endpoint or "ipc" in endpoint: return f"{endpoint}_dp{data_parallel_rank}" if "tcp" in endpoint and ":" in endpoint: last_colon_idx = endpoint.rfind(":") @@ -307,7 +314,7 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | f"KV event endpoint port exceeds 65535 for rank {data_parallel_rank}" ) return f"{base_addr}:{new_port}" - raise ValueError("Invalid endpoint: must contain 'inproc' or 'tcp'") + raise ValueError("Invalid endpoint: must contain 'inproc', 'ipc', or 'tcp'") def create_event_publisher(config: KVEventsConfig, data_parallel_rank: int) -> EventPublisher: @@ -490,6 +497,8 @@ def _block_hashes( return block_hash, parent_hash, _NativeStoredBlockState(block_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] = [] @@ -502,6 +511,8 @@ def add_removed_event(self, block_hashes: Any) -> None: self._add_removed_hashes(removed_hashes) def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: + if self._closed: + return if int(life_cycle_id) != self._target_life_cycle_id: self.non_target_life_cycles_ignored += 1 return @@ -512,11 +523,11 @@ def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> def _add_removed_hashes(self, block_hashes: list[ExternalBlockHash]) -> None: if not block_hashes: return - # Removals are never dropped by the per-iteration cap: each hash was - # already reported as stored, so dropping its removal would leave the - # consumer believing the block is resident forever. They are bounded by - # the previously-stored set, so they cannot run away. - self._pending_entries += len(block_hashes) + # 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: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index a2c1c2ea80b5..963a8b880066 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -873,6 +873,13 @@ def __init__( kv_events_config is not None and kv_events_config.enable_kv_cache_events ) if native_events_enabled: + if self.event_buffer_max_size > 0: + logger.warning( + "Both kv_cache_config.event_buffer_max_size and native " + "kv_events_config are enabled; native publishing takes " + "precedence and the legacy get_kv_cache_events() poll path " + "will return no events." + ) if mapping.pp_size > 1: raise ValueError("Native KV events do not support pipeline parallelism") if mapping.cp_size > 1: @@ -1066,30 +1073,41 @@ def append_to_kv_heads_per_layer( self.kv_cache_manager_py_config = config + # The native event manager has already bound its ZMQ socket and started + # its background thread, so tear it down if impl construction or + # event-manager setup fails here -- otherwise the socket and daemon + # thread leak and an in-process retry cannot rebind the same endpoint. try: - self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) - except (CuError, KVCacheOutOfMemoryError): - if len(cache_tiers) > 1: - logger.warning( - "Failed to initialize KV cache manager with host cache " - "tier (cuMemHostRegister may have failed). " - "Retrying without host cache tier." - ) - cache_tiers_gpu_only = [t for t in cache_tiers if isinstance(t, GpuCacheTierConfig)] - config = replace(config, cache_tiers=cache_tiers_gpu_only) - cache_tiers = cache_tiers_gpu_only - self.kv_cache_manager_py_config = config + try: self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) - else: - raise - if self.event_manager is not None: - self.event_manager.set_layer_group_window_sizes( - self._get_event_window_sizes_by_layer_group() - ) - self.event_manager.add_created_event( - self._get_event_num_blocks_per_cache_level(cache_tiers, tokens_per_block), - self._get_event_layer_group_ids(), - ) + except (CuError, KVCacheOutOfMemoryError): + if len(cache_tiers) > 1: + logger.warning( + "Failed to initialize KV cache manager with host cache " + "tier (cuMemHostRegister may have failed). " + "Retrying without host cache tier." + ) + cache_tiers_gpu_only = [ + t for t in cache_tiers if isinstance(t, GpuCacheTierConfig) + ] + config = replace(config, cache_tiers=cache_tiers_gpu_only) + cache_tiers = cache_tiers_gpu_only + self.kv_cache_manager_py_config = config + self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) + else: + raise + if self.event_manager is not None: + self.event_manager.set_layer_group_window_sizes( + self._get_event_window_sizes_by_layer_group() + ) + self.event_manager.add_created_event( + self._get_event_num_blocks_per_cache_level(cache_tiers, tokens_per_block), + self._get_event_layer_group_ids(), + ) + except Exception: + if isinstance(self.event_manager, NativeKVCacheEventManager): + self.event_manager.shutdown() + raise self.num_pools = len(self.impl.layer_grouping) # num_pools is the physical pool count owned by the KV cache manager. @@ -1486,10 +1504,17 @@ def get_event_window_size(layer_id: int) -> int: 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): + life_cycle = self.impl._life_cycles.get_life_cycle(LifeCycleId(layer_group_id)) + # Native KV events track 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. + 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: attr = self.impl._storage.get_buffer_attr(layer_id, role) @@ -2913,16 +2938,19 @@ 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): # Native publishing pushes events out-of-band; in that mode the event # manager's get_latest_events returns [], so the legacy pull path - # degrades cleanly instead of raising. - if self.event_manager is None: + # 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 native_kv_events_enabled(self) -> bool: @@ -3450,13 +3478,17 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool return bool(has_invalid_values) def shutdown(self): - if isinstance(self.event_manager, NativeKVCacheEventManager): - self.event_manager.shutdown() - self.event_manager = None for kv_cache in self.kv_cache_map.values(): kv_cache.close() self.kv_cache_map.clear() self.impl.shutdown() + # Shut the native 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, NativeKVCacheEventManager): + self.event_manager.shutdown() if self.conversation_manager is not None: self.conversation_manager.clear() diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0835e158bfe9..86ef9f0268e9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3631,24 +3631,29 @@ class KVEventsConfig(StrictBaseModel): ) endpoint: str = Field( default="tcp://*:5557", - description="Base ZeroMQ endpoint used to publish KV cache events.") + 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, description= "Optional base ZeroMQ endpoint used to replay KV cache events.") buffer_steps: int = Field( default=10_000, - ge=0, + gt=0, description="Number of previously published batches retained for replay." ) hwm: int = Field(default=100_000, - ge=0, - description="ZeroMQ publisher socket high-water mark.") + 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, - ge=0, - description="Maximum number of batches queued for background publishing." - ) + 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." From 8f3474b3a3cef77abf63e5da85bfed41ebbe01cf Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 5 Aug 2026 14:25:01 -0700 Subject: [PATCH 08/25] fix: address native KV events review comments - Reuse truncate_sha256_hash_to_int64 for the vLLM wire hash instead of a second, divergent SHA-256->int64 truncation, keeping native and legacy event hashes consistent for the same block. - Replace logger.exception (absent on tensorrt_llm's logger; would raise AttributeError) with logger.error + traceback.format_exc() at all four call sites. Signed-off-by: tanmayv25 --- .../_torch/pyexecutor/kv_cache_events.py | 28 ++++++++++++++----- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 77821f5f9c67..eb0734216896 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -21,6 +21,7 @@ import queue import threading import time +import traceback from abc import ABC, abstractmethod from collections import deque from itertools import count @@ -32,6 +33,7 @@ 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 ExternalBlockHash = bytes | int @@ -242,7 +244,10 @@ def _publisher_thread(self) -> None: try: self._service_replay() except Exception: - logger.exception("Failed to service native KV event replay request") + logger.error( + "Failed to service native KV event replay request\n" + f"{traceback.format_exc()}" + ) try: event = self._event_queue.get(timeout=0.1) except queue.Empty: @@ -264,8 +269,9 @@ def _publisher_thread(self) -> None: self.published_batches += 1 except Exception: self._send_error_drops += 1 - logger.exception( - f"Failed to publish native KV event batch rank={self._rank} seq={seq}" + logger.error( + f"Failed to publish native KV event batch rank={self._rank} " + f"seq={seq}\n{traceback.format_exc()}" ) time.sleep(0.1) finally: @@ -338,8 +344,10 @@ def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: """Reuse an existing SHA-256 radix key as vLLM's signed integer event hash.""" if len(block_key) < 8: raise ValueError("V2 radix block keys must contain at least 8 bytes") - unsigned_hash = int.from_bytes(block_key[-8:], "big", signed=False) - # Reinterpret the low 64 bits as signed two's-complement for the wire format. + # 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 vLLM's signed wire hash. + unsigned_hash = truncate_sha256_hash_to_int64(block_key) return unsigned_hash - 2**64 if unsigned_hash >= 2**63 else unsigned_hash @@ -454,7 +462,10 @@ def _add_full_block(self, block: Any) -> None: except ValueError: self.dropped_events += 1 self._pending_entries -= 1 - logger.exception("Dropping native KV store event with unsupported token data") + logger.error( + "Dropping native KV store event with unsupported token data\n" + f"{traceback.format_exc()}" + ) return self._stored_blocks[key] = state if self._pending_events and isinstance(self._pending_events[-1], BlockStored): @@ -577,7 +588,10 @@ def flush_iteration_events(self) -> None: self.dropped_batches += 1 except Exception: self.dropped_batches += 1 - logger.exception(f"Dropping native KV event iteration batch on rank={self._rank}") + logger.error( + f"Dropping native 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]: # Native publishing pushes events out-of-band, so the pull API has From c70dc1557dcd047e30edbb579301523040c657e4 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Mon, 10 Aug 2026 13:29:25 -0700 Subject: [PATCH 09/25] test: align native KV event wire-hash assertions with truncation change 8f3474b switched the wire hash to truncate_sha256_hash_to_int64 (first 8 bytes of the radix key) but left the test asserting the old last-8-byte values, so the test failed deterministically in CI. Update the synthetic keys and expected hashes to the first-8-byte convention, keeping the signed-wraparound branch covered. Signed-off-by: tanmayv25 --- .../test_native_kv_events.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py index f21f26473d4c..be99bca11b93 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py @@ -76,11 +76,14 @@ def block( ], ) - first_hash = b"\x11" * 24 + b"\x80\x00\x00\x00\x00\x00\x00\x01" + # 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 of the wire hash -- 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"\x33" * 24 + b"\x00\x00\x00\x00\x00\x00\x00\x02" - first_wire_hash = int.from_bytes(first_hash[-8:], "big") - second_wire_hash = int.from_bytes(second_hash[-8:], "big") + second_hash = b"\x00\x00\x00\x00\x00\x00\x00\x02" + b"\x33" * 24 + first_wire_hash = int.from_bytes(first_hash[:8], "big") + second_wire_hash = int.from_bytes(second_hash[:8], "big") first_wire_hash = first_wire_hash - 2**64 if first_wire_hash >= 2**63 else first_wire_hash second_wire_hash = second_wire_hash - 2**64 if second_wire_hash >= 2**63 else second_wire_hash first = block(first_hash, [1, 2, 3, 4], root) From 6e46ee7b6bd25b2fca1f6a4032ab8d6acd8c44ee Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Mon, 10 Aug 2026 13:33:14 -0700 Subject: [PATCH 10/25] refactor: drop _NativeStoredBlockState wrapper; test config+endpoint (review) Address code-review minors: - Replace the single-int _NativeStoredBlockState wrapper with a plain dict[bytes, int] mapping radix key -> wire hash; deletes the class and a redundant tuple slot. - Add tests for KVEventsConfig publisher default resolution (None -> zmq/null) and offset_endpoint_port (base_port+rank, ipc/inproc suffix, u16 overflow, bad scheme). Signed-off-by: tanmayv25 --- .../_torch/pyexecutor/kv_cache_events.py | 29 +++++-------- .../test_native_kv_events.py | 41 ++++++++++++++++++- 2 files changed, 51 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index eb0734216896..96ed770b9e28 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -351,13 +351,6 @@ def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: return unsigned_hash - 2**64 if unsigned_hash >= 2**63 else unsigned_hash -class _NativeStoredBlockState: - __slots__ = ("block_hash",) - - def __init__(self, block_hash: int) -> None: - self.block_hash = block_hash - - class NativeKVCacheEventManager: """Scheduler-local fast path that produces vLLM wire events directly. @@ -383,7 +376,7 @@ def __init__( 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, _NativeStoredBlockState] = {} + self._stored_blocks: dict[bytes, int] = {} self._pending_events: list[BlockStored | BlockRemoved | AllBlocksCleared] = [] self._pending_entries = 0 self._closed = False @@ -458,7 +451,7 @@ def _add_full_block(self, block: Any) -> None: return try: token_ids = self._token_ids(block.tokens) - block_hash, parent_hash, state = self._block_hashes(block) + block_hash, parent_hash = self._block_hashes(block) except ValueError: self.dropped_events += 1 self._pending_entries -= 1 @@ -467,7 +460,7 @@ def _add_full_block(self, block: Any) -> None: f"{traceback.format_exc()}" ) return - self._stored_blocks[key] = state + 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: @@ -500,12 +493,12 @@ def _token_ids(tokens: Any) -> list[int]: def _block_hashes( self, block: Any, - ) -> tuple[int, int | None, _NativeStoredBlockState]: + ) -> tuple[int, int | None]: parent = block.prev is_root_child = getattr(parent, "ordinal", -1) == -1 block_hash = _vllm_wire_hash_from_radix_key(bytes(block.key)) parent_hash = None if is_root_child else _vllm_wire_hash_from_radix_key(bytes(parent.key)) - return block_hash, parent_hash, _NativeStoredBlockState(block_hash) + return block_hash, parent_hash def add_removed_event(self, block_hashes: Any) -> None: if self._closed: @@ -516,9 +509,9 @@ def add_removed_event(self, block_hashes: Any) -> None: for block_key in block_hashes: if not isinstance(block_key, bytes): continue - state = self._stored_blocks.pop(block_key, None) - if state is not None: - removed_hashes.append(state.block_hash) + 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: @@ -527,9 +520,9 @@ def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> if int(life_cycle_id) != self._target_life_cycle_id: self.non_target_life_cycles_ignored += 1 return - state = self._stored_blocks.pop(block_hash, None) - if state is not None: - self._add_removed_hashes([state.block_hash]) + 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: diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py index be99bca11b93..abf34b711dc1 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py @@ -18,9 +18,14 @@ from types import SimpleNamespace import msgspec +import pytest import zmq -from tensorrt_llm._torch.pyexecutor.kv_cache_events import BlockRemoved, NativeKVCacheEventManager +from tensorrt_llm._torch.pyexecutor.kv_cache_events import ( + BlockRemoved, + NativeKVCacheEventManager, + ZmqEventPublisher, +) from tensorrt_llm.llmapi.llm_args import KVEventsConfig @@ -184,3 +189,37 @@ def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: assert sum(len(event.block_hashes) for event in removed) == 2 manager.shutdown() + + +def test_kv_events_config_publisher_default(): + """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): + assert ZmqEventPublisher.offset_endpoint_port(endpoint, rank) == expected + + +def test_offset_endpoint_port_rejects_bad_input(): + # 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) From 3a9bf62b164f41cc3c1a65ec3a4f11458b073752 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Mon, 10 Aug 2026 15:16:06 -0700 Subject: [PATCH 11/25] refactor: rename KV events 'native/legacy' -> 'streaming/buffered' The 'native' vs 'legacy' naming was misleading: 'native' is overloaded in TRT-LLM, and 'legacy' wrongly implied the buffered gather/poll path is deprecated when it is actually the fuller-fidelity default. Rename to describe the delivery mechanism: - NativeKVCacheEventManager -> StreamingKVCacheEventManager (+ native_kv_events_enabled -> streaming_kv_events_enabled). - 'native'/'legacy' -> 'streaming (push-based)'/'buffered (gather/poll)' in log messages, comments, KVEventsConfig docstrings, and the test file name. Public config identifiers (kv_events_config, enable_kv_cache_events) are unchanged; only descriptions were updated (not captured by the golden manifest). Signed-off-by: tanmayv25 --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- .../_torch/pyexecutor/kv_cache_events.py | 34 ++++++++-------- .../_torch/pyexecutor/kv_cache_manager_v2.py | 40 +++++++++---------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 2 +- tensorrt_llm/llmapi/llm_args.py | 8 ++-- ..._events.py => test_streaming_kv_events.py} | 0 6 files changed, 43 insertions(+), 43 deletions(-) rename tests/unittest/kv_cache_manager_v2_tests/{test_native_kv_events.py => test_streaming_kv_events.py} (100%) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 1e125ad78eb6..2a0c379b7a9f 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1996,7 +1996,7 @@ def _create_kv_cache_manager( 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 native KV event " + "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): diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 96ed770b9e28..4cefca547b7d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -164,7 +164,7 @@ def __init__( ) self._thread.start() logger.info( - f"Started native KV event publisher rank={self._rank} " + f"Started streaming KV event publisher rank={self._rank} " f"endpoint={self._endpoint} topic={topic!r}" ) @@ -189,7 +189,7 @@ def publish(self, events: EventBatch) -> bool: drops = self._queue_full_drops if drops == 1 or (drops & (drops - 1) == 0): logger.warning( - f"Dropping native KV event batch on rank={self._rank} because " + f"Dropping streaming KV event batch on rank={self._rank} because " "the publisher queue is full; " f"dropped_batches={self.dropped_batches}" ) @@ -208,11 +208,11 @@ def shutdown(self) -> None: self._thread.join(timeout=self.SHUTDOWN_TIMEOUT) if self._thread.is_alive(): logger.warning( - f"Native KV event publisher rank={self._rank} did not stop " + f"Streaming KV event publisher rank={self._rank} did not stop " f"within {self.SHUTDOWN_TIMEOUT:.1f}s" ) logger.info( - f"Stopped native KV event publisher rank={self._rank} " + 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}" @@ -245,7 +245,7 @@ def _publisher_thread(self) -> None: self._service_replay() except Exception: logger.error( - "Failed to service native KV event replay request\n" + "Failed to service streaming KV event replay request\n" f"{traceback.format_exc()}" ) try: @@ -270,7 +270,7 @@ def _publisher_thread(self) -> None: except Exception: self._send_error_drops += 1 logger.error( - f"Failed to publish native KV event batch rank={self._rank} " + f"Failed to publish streaming KV event batch rank={self._rank} " f"seq={seq}\n{traceback.format_exc()}" ) time.sleep(0.1) @@ -285,7 +285,7 @@ def _service_replay(self) -> None: assert self._replay is not None frame = self._replay.recv_multipart() if len(frame) != 3: - logger.warning(f"Invalid native KV event replay request: {frame}") + 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") @@ -351,7 +351,7 @@ def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: return unsigned_hash - 2**64 if unsigned_hash >= 2**63 else unsigned_hash -class NativeKVCacheEventManager: +class StreamingKVCacheEventManager: """Scheduler-local fast path that produces vLLM wire events directly. Implements the V2 KV-cache-manager event-sink hook interface by duck @@ -403,10 +403,10 @@ def set_layer_group_window_sizes(self, window_sizes: dict[int, int]) -> None: if window_size == largest_window ] if not target_ids: - raise ValueError("Native KV events require an attention KV cache life cycle") + raise ValueError("Streaming KV events require an attention KV cache life cycle") self._target_life_cycle_id = min(target_ids) logger.info( - "Native KV event fast path selected " + "Streaming KV event fast path selected " f"lifecycle_id={self._target_life_cycle_id} " f"window_size={self._max_window_size}" ) @@ -419,7 +419,7 @@ def add_created_event( return def add_stored_event(self, *args: Any, **kwargs: Any) -> None: - # Native publishing derives stored events from the per-block hooks + # Streaming publishing derives stored events from the per-block hooks # below; the aggregate stored-event hook is intentionally unused. return @@ -456,7 +456,7 @@ def _add_full_block(self, block: Any) -> None: self.dropped_events += 1 self._pending_entries -= 1 logger.error( - "Dropping native KV store event with unsupported token data\n" + "Dropping streaming KV store event with unsupported token data\n" f"{traceback.format_exc()}" ) return @@ -557,7 +557,7 @@ def _reserve_entries(self, num_entries: int) -> bool: self.dropped_events & (self.dropped_events - 1) == 0 ): logger.warning( - "Dropping native KV events because the per-iteration safety " + "Dropping streaming KV events because the per-iteration safety " f"cap was exceeded; dropped_events={self.dropped_events}" ) return False @@ -582,14 +582,14 @@ def flush_iteration_events(self) -> None: except Exception: self.dropped_batches += 1 logger.error( - f"Dropping native KV event iteration batch on rank={self._rank}\n" + 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]: - # Native publishing pushes events out-of-band, so the pull API has + # Streaming publishing pushes events out-of-band, so the pull API has # nothing to return. Return empty instead of raising so callers of the - # legacy polling path degrade cleanly rather than erroring. + # buffered polling path degrade cleanly rather than erroring. return [] def shutdown(self) -> None: @@ -599,7 +599,7 @@ def shutdown(self) -> None: self._closed = True self._publisher.shutdown() logger.info( - "Native KV event fast path " + "Streaming KV event fast path " f"rank={self._rank} " f"stored_blocks={self.stored_blocks} " f"removed_blocks={self.removed_blocks} " diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 963a8b880066..045dbf9fddc9 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -82,7 +82,7 @@ from ...mapping import CpType, Mapping from ..utils import maybe_compile from .connectors.kv_cache_connector import KvCacheConnectorManager -from .kv_cache_events import NativeKVCacheEventManager +from .kv_cache_events import StreamingKVCacheEventManager from .kv_cache_stats import ( KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, @@ -868,32 +868,32 @@ 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 | NativeKVCacheEventManager] = None - native_events_enabled = ( + 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 native_events_enabled: + if streaming_events_enabled: if self.event_buffer_max_size > 0: logger.warning( - "Both kv_cache_config.event_buffer_max_size and native " - "kv_events_config are enabled; native publishing takes " - "precedence and the legacy get_kv_cache_events() poll path " + "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." ) if mapping.pp_size > 1: - raise ValueError("Native KV events do not support pipeline parallelism") + raise ValueError("Streaming KV events do not support pipeline parallelism") if mapping.cp_size > 1: - raise ValueError("Native KV events do not support context parallelism") + raise ValueError("Streaming KV events do not support context parallelism") assert kv_events_config is not None if mapping.enable_attention_dp or mpi_rank() == 0: event_rank = mapping.rank if mapping.enable_attention_dp else 0 - self.event_manager = NativeKVCacheEventManager( + self.event_manager = StreamingKVCacheEventManager( kv_events_config, data_parallel_rank=event_rank, block_size=self.tokens_per_block, max_window_size=event_window_size, ) - logger.info("Native KV event fast path reuses V2 radix block hashes") + logger.info("Streaming KV event fast path reuses V2 radix block hashes") elif self.event_buffer_max_size > 0: if mapping.enable_attention_dp: self.event_manager = KVCacheEventManager( @@ -1073,7 +1073,7 @@ def append_to_kv_heads_per_layer( self.kv_cache_manager_py_config = config - # The native event manager has already bound its ZMQ socket and started + # The streaming event manager has already bound its ZMQ socket and started # its background thread, so tear it down if impl construction or # event-manager setup fails here -- otherwise the socket and daemon # thread leak and an in-process retry cannot rebind the same endpoint. @@ -1105,7 +1105,7 @@ def append_to_kv_heads_per_layer( self._get_event_layer_group_ids(), ) except Exception: - if isinstance(self.event_manager, NativeKVCacheEventManager): + if isinstance(self.event_manager, StreamingKVCacheEventManager): self.event_manager.shutdown() raise @@ -1507,7 +1507,7 @@ 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): life_cycle = self.impl._life_cycles.get_life_cycle(LifeCycleId(layer_group_id)) - # Native KV events track attention prefix reuse only. Excluding SSM + # Streaming KV events track 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. @@ -2943,8 +2943,8 @@ def flush_iteration_events(self): event_manager.flush_iteration_events() def get_latest_events(self, timeout_ms: Optional[float] = None): - # Native publishing pushes events out-of-band; in that mode the event - # manager's get_latest_events returns [], so the legacy pull path + # 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 @@ -2953,8 +2953,8 @@ def get_latest_events(self, timeout_ms: Optional[float] = None): return event_manager.get_latest_events(timeout_ms) @property - def native_kv_events_enabled(self) -> bool: - return isinstance(self.event_manager, NativeKVCacheEventManager) + def streaming_kv_events_enabled(self) -> bool: + return isinstance(self.event_manager, StreamingKVCacheEventManager) def get_iteration_stats(self): if not self.enable_stats: @@ -3482,12 +3482,12 @@ def shutdown(self): kv_cache.close() self.kv_cache_map.clear() self.impl.shutdown() - # Shut the native event manager down last so removals emitted during + # 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, NativeKVCacheEventManager): + 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 500f46aa4b22..7f5f34f9438f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -647,7 +647,7 @@ def __init__( 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 or getattr( - self.kv_cache_manager, "native_kv_events_enabled", False)) + 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/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 86ef9f0268e9..d51194c09a9b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3619,11 +3619,11 @@ class MambaStateConfig(StrictBaseModel): class KVEventsConfig(StrictBaseModel): - """Configuration for native KV cache event publishing.""" + """Configuration for streaming (push-based) KV cache event publishing.""" enable_kv_cache_events: bool = Field( default=False, - description="Whether to produce and publish native KV cache events.") + description="Whether to produce and publish KV cache events over the streaming (push) path.") publisher: Optional[Literal["null", "zmq"]] = Field( default=None, description= @@ -3736,9 +3736,9 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): default=None, status="prototype", description= - "Native KV cache event publishing (KV cache manager V2 only). When set, " + "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 legacy event_buffer_max_size gather/poll path.") + "of the buffered event_buffer_max_size gather/poll path.") enable_partial_reuse: bool = Field( default=True, description= diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py similarity index 100% rename from tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py rename to tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py From 69b318a8cc9efbfd05698e2d0ca99f8e11d550f7 Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Tue, 11 Aug 2026 11:28:22 -0700 Subject: [PATCH 12/25] fix: address code-review findings for streaming KV events - ZmqEventPublisher.__init__: close the PUB/ROUTER sockets if _socket_setup raises, so a bind/scheme failure doesn't leak sockets on the shared context (shutdown() is unreachable when __init__ never returns). - offset_endpoint_port: match the scheme with startswith (consistent with _socket_setup) and reject a TCP endpoint with no port, instead of parsing the scheme colon into int() and raising an opaque error on ranks > 0. - Skip multimodal cache-key blocks (bytes token digests) via a dedicated _MultimodalBlockError so they no longer flood the log with malformed-data tracebacks; count them in multimodal_blocks_suppressed. - KVEventsConfig.replay_endpoint: add min_length=1 so an empty string is rejected up front rather than failing at bind() time. - Tests: annotate test procedures with -> None, run manager/subscriber cleanup under try/finally, assert removals reach the wire after flush, and cover the no-port endpoint rejection. Signed-off-by: tanmayv25 --- .../_torch/pyexecutor/kv_cache_events.py | 45 ++- tensorrt_llm/llmapi/llm_args.py | 1 + .../test_streaming_kv_events.py | 272 ++++++++++-------- 3 files changed, 187 insertions(+), 131 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 4cefca547b7d..b34ba5074c08 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -156,7 +156,16 @@ def __init__( self.published_batches = 0 self._queue_full_drops = 0 self._send_error_drops = 0 - self._socket_setup() + try: + self._socket_setup() + except Exception: + # __init__ never returns on failure, so shutdown() is unreachable; + # close the sockets here to avoid leaking them on the shared context. + if self._pub is not None: + self._pub.close(linger=0) + if self._replay is not None: + self._replay.close(linger=0) + raise self._thread = threading.Thread( target=self._publisher_thread, daemon=True, @@ -307,10 +316,17 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | """Apply vLLM's base-port-plus-rank endpoint convention.""" if not endpoint or data_parallel_rank == 0: 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 "inproc" in endpoint or "ipc" in endpoint: + if endpoint.startswith(("inproc://", "ipc://")): return f"{endpoint}_dp{data_parallel_rank}" - if "tcp" in endpoint and ":" in endpoint: + 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] base_port = int(endpoint[last_colon_idx + 1 :]) @@ -320,7 +336,9 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | f"KV event endpoint port exceeds 65535 for rank {data_parallel_rank}" ) return f"{base_addr}:{new_port}" - raise ValueError("Invalid endpoint: must contain 'inproc', 'ipc', or 'tcp'") + raise ValueError( + "Invalid endpoint: must start with 'inproc://', 'ipc://', or 'tcp://'" + ) def create_event_publisher(config: KVEventsConfig, data_parallel_rank: int) -> EventPublisher: @@ -351,6 +369,15 @@ def _vllm_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 vLLM-wire integer representation. Such blocks are skipped + quietly rather than routed through the malformed-data traceback path. + """ + + class StreamingKVCacheEventManager: """Scheduler-local fast path that produces vLLM wire events directly. @@ -383,6 +410,7 @@ 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 @@ -452,6 +480,12 @@ def _add_full_block(self, block: Any) -> None: 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 @@ -485,6 +519,9 @@ def _add_full_block(self, block: Any) -> None: 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("vLLM-compatible KV events require integer token IDs") token_ids.append(token) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d51194c09a9b..28b3697a6ee2 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3638,6 +3638,7 @@ class KVEventsConfig(StrictBaseModel): "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.") buffer_steps: int = Field( 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 abf34b711dc1..fee6d05729cd 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 @@ -22,8 +22,7 @@ import zmq from tensorrt_llm._torch.pyexecutor.kv_cache_events import ( - BlockRemoved, - NativeKVCacheEventManager, + StreamingKVCacheEventManager, ZmqEventPublisher, ) from tensorrt_llm.llmapi.llm_args import KVEventsConfig @@ -35,7 +34,7 @@ def _unused_tcp_port() -> int: return int(sock.getsockname()[1]) -def test_native_fast_path_publishes_only_full_max_window_blocks(): +def test_streaming_fast_path_publishes_only_full_max_window_blocks() -> None: """Protect radix hash reuse, filtering, wire format, and shutdown.""" port = _unused_tcp_port() bind_endpoint = f"tcp://*:{port}" @@ -46,7 +45,7 @@ def test_native_fast_path_publishes_only_full_max_window_blocks(): subscriber.setsockopt_string(zmq.SUBSCRIBE, topic) subscriber.connect(connect_endpoint) - manager = NativeKVCacheEventManager( + manager = StreamingKVCacheEventManager( KVEventsConfig( enable_kv_cache_events=True, publisher="zmq", @@ -58,140 +57,155 @@ def test_native_fast_path_publishes_only_full_max_window_blocks(): block_size=4, max_window_size=128, ) - manager.set_layer_group_window_sizes({0: 128, 1: 64}) - time.sleep(0.2) - - root = SimpleNamespace(ordinal=-1) - - def block( - key: bytes, - tokens: list[int], - prev: object, - ) -> SimpleNamespace: - max_window_page = object() - smaller_window_page = object() - return SimpleNamespace( - key=key, - tokens=tokens, - prev=prev, - ordinal=getattr(prev, "ordinal", -1) + 1, - storage=[ - lambda: max_window_page, - lambda: smaller_window_page, - ], - ) - - # 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 of the wire hash -- 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") - second_wire_hash = int.from_bytes(second_hash[:8], "big") - first_wire_hash = first_wire_hash - 2**64 if first_wire_hash >= 2**63 else first_wire_hash - second_wire_hash = second_wire_hash - 2**64 if second_wire_hash >= 2**63 else second_wire_hash - 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): - assert subscriber.poll(2_000) - frames.append(subscriber.recv_multipart()) - - assert [frame[0] for frame in frames] == [topic.encode(), topic.encode()] - assert [int.from_bytes(frame[1], "big") for frame in frames] == [0, 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 - - # Native publishing pushes events out-of-band, so the legacy pull API must - # degrade to an empty result rather than raising. - assert manager.get_latest_events() == [] - - manager.shutdown() - manager.shutdown() - subscriber.close(linger=0) + try: + manager.set_layer_group_window_sizes({0: 128, 1: 64}) + time.sleep(0.2) + + root = SimpleNamespace(ordinal=-1) + + def block( + key: bytes, + tokens: list[int], + prev: object, + ) -> SimpleNamespace: + max_window_page = object() + smaller_window_page = object() + return SimpleNamespace( + key=key, + tokens=tokens, + prev=prev, + ordinal=getattr(prev, "ordinal", -1) + 1, + storage=[ + lambda: max_window_page, + lambda: smaller_window_page, + ], + ) + + # 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") + second_wire_hash = int.from_bytes(second_hash[:8], "big") + first_wire_hash = first_wire_hash - 2**64 if first_wire_hash >= 2**63 else first_wire_hash + second_wire_hash = second_wire_hash - 2**64 if second_wire_hash >= 2**63 else second_wire_hash + 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): + assert subscriber.poll(2_000) + frames.append(subscriber.recv_multipart()) + + assert [frame[0] for frame in frames] == [topic.encode(), topic.encode()] + assert [int.from_bytes(frame[1], "big") for frame in frames] == [0, 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. + manager.shutdown() + manager.shutdown() + finally: + manager.shutdown() + subscriber.close(linger=0) replacement = context.socket(zmq.PUB) replacement.bind(bind_endpoint) replacement.close(linger=0) -def test_native_removals_are_never_dropped_by_the_entry_cap(): +def test_streaming_removals_are_never_dropped_by_the_entry_cap() -> None: """Removals must survive the per-iteration cap or the consumer desyncs.""" - manager = NativeKVCacheEventManager( + 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.set_layer_group_window_sizes({0: 128}) - - root = SimpleNamespace(ordinal=-1) - - def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: - page = object() - 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]) - - removed = [event for event in manager._pending_events if isinstance(event, BlockRemoved)] - assert manager.removed_blocks == 2 - assert sum(len(event.block_hashes) for event in removed) == 2 - - manager.shutdown() - - -def test_kv_events_config_publisher_default(): + 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 = object() + 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_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" @@ -212,14 +226,18 @@ def test_kv_events_config_publisher_default(): (None, 5, None), ], ) -def test_offset_endpoint_port(endpoint, rank, expected): +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(): +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) From b92128a88d13105bc6f841d4e698a3bd95e389fc Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Tue, 11 Aug 2026 17:02:28 -0700 Subject: [PATCH 13/25] refactor: name the KV event wire format in TensorRT-LLM terms The published KV-cache-event wire schema was described throughout as 'vLLM-compatible', which reads oddly in a TensorRT-LLM module. Rename the descriptive references to TensorRT-LLM's own 'KV cache event wire' terminology: - _vllm_wire_hash_from_radix_key -> _kv_event_wire_hash_from_radix_key - 'vLLM-compatible'/'vLLM wire'/'vLLM's ... hash' -> 'KV cache event wire ...' across the struct, publisher, endpoint, and manager docstrings and messages. The one retained vLLM mention is the module-header source attribution: the on-wire schema was adapted from vLLM's Apache-2.0 vllm/distributed/kv_events.py. No behavior change; the renamed helper is private to this module. Signed-off-by: tanmayv25 --- .../_torch/pyexecutor/kv_cache_events.py | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index b34ba5074c08..57202bc96689 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -13,8 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # -# The wire schema and ZeroMQ framing in this file are adapted from vLLM's -# vllm/distributed/kv_events.py. +# This module defines TensorRT-LLM's KV cache event wire format: msgpack event +# batches published over ZeroMQ. The on-wire schema and three-frame framing are +# adapted from vLLM's vllm/distributed/kv_events.py so that external +# KV-cache-aware routers can consume TensorRT-LLM events without translation. from __future__ import annotations @@ -45,7 +47,7 @@ class EventBatch( omit_defaults=True, # type: ignore[call-arg] gc=False, # type: ignore[call-arg] ): - """vLLM-compatible event batch envelope.""" + """KV cache event wire batch envelope.""" ts: float events: list[Any] @@ -58,7 +60,7 @@ class KVCacheWireEvent( gc=False, # type: ignore[call-arg] tag=True, ): - """Base class for vLLM-compatible KV cache events.""" + """Base class for KV cache event wire messages.""" class BlockStored(KVCacheWireEvent): @@ -98,7 +100,7 @@ class KVEventBatch(EventBatch): class EventPublisher(ABC): - """Publishes vLLM-compatible event batches for one cache rank.""" + """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 @@ -123,7 +125,7 @@ def shutdown(self) -> None: class ZmqEventPublisher(EventPublisher): - """Publishes event batches with vLLM's three-frame ZeroMQ protocol.""" + """Publishes event batches over the three-frame ZeroMQ wire protocol.""" SHUTDOWN_TIMEOUT = 1.0 END_SEQ = (-1).to_bytes(8, "big", signed=True) @@ -313,7 +315,7 @@ def _service_replay(self) -> None: @staticmethod def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | None: - """Apply vLLM's base-port-plus-rank endpoint convention.""" + """Apply the base-port-plus-rank endpoint convention (each rank binds base_port + rank).""" if not endpoint or data_parallel_rank == 0: return endpoint # Match the scheme with startswith so detection agrees with @@ -324,9 +326,7 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | 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}" - ) + raise ValueError(f"TCP KV event endpoint must include a port: {endpoint!r}") last_colon_idx = endpoint.rfind(":") base_addr = endpoint[:last_colon_idx] base_port = int(endpoint[last_colon_idx + 1 :]) @@ -336,9 +336,7 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | 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://'" - ) + raise ValueError("Invalid endpoint: must start with 'inproc://', 'ipc://', or 'tcp://'") def create_event_publisher(config: KVEventsConfig, data_parallel_rank: int) -> EventPublisher: @@ -358,13 +356,13 @@ def create_event_publisher(config: KVEventsConfig, data_parallel_rank: int) -> E raise ValueError(f"Unsupported KV event publisher: {config.publisher!r}") -def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: - """Reuse an existing SHA-256 radix key as vLLM's signed integer event hash.""" +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 vLLM's signed wire hash. + # 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 @@ -373,13 +371,13 @@ 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 vLLM-wire integer representation. Such blocks are skipped + 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 vLLM wire events directly. + """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 @@ -523,7 +521,7 @@ def _token_ids(tokens: Any) -> list[int]: # Multimodal cache-key digest; not representable as a wire int. raise _MultimodalBlockError if type(token) is not int: - raise ValueError("vLLM-compatible KV events require integer token IDs") + raise ValueError("KV cache event wire format requires integer token IDs") token_ids.append(token) return token_ids @@ -533,8 +531,10 @@ def _block_hashes( ) -> tuple[int, int | None]: parent = block.prev is_root_child = getattr(parent, "ordinal", -1) == -1 - block_hash = _vllm_wire_hash_from_radix_key(bytes(block.key)) - parent_hash = None if is_root_child else _vllm_wire_hash_from_radix_key(bytes(parent.key)) + 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: From e01518a36a11cd1e8c8a4b2b425b797c72a4714e Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Tue, 11 Aug 2026 17:31:13 -0700 Subject: [PATCH 14/25] fix: validate KV event endpoint port and fix test formatting - offset_endpoint_port: reject a non-numeric or out-of-range TCP port (e.g. tcp://host:abc, :0, :-5) with an endpoint-naming error instead of an opaque int()/ZeroMQ bind failure that only surfaces on ranks > 0; add tests. - test_streaming_kv_events.py: apply the ruff-format wrapping the CI pre-commit check flagged (long signed-wraparound line). Signed-off-by: tanmayv25 --- tensorrt_llm/_torch/pyexecutor/kv_cache_events.py | 9 ++++++++- .../test_streaming_kv_events.py | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 57202bc96689..4f0390f7f6cd 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -329,7 +329,14 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | raise ValueError(f"TCP KV event endpoint must include a port: {endpoint!r}") last_colon_idx = endpoint.rfind(":") base_addr = endpoint[:last_colon_idx] - base_port = int(endpoint[last_colon_idx + 1 :]) + 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}" + ) + base_port = int(port_text) new_port = base_port + data_parallel_rank if new_port > 65_535: raise ValueError( 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 fee6d05729cd..af593bf0c422 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 @@ -90,7 +90,9 @@ def block( first_wire_hash = int.from_bytes(first_hash[:8], "big") second_wire_hash = int.from_bytes(second_hash[:8], "big") first_wire_hash = first_wire_hash - 2**64 if first_wire_hash >= 2**63 else first_wire_hash - second_wire_hash = second_wire_hash - 2**64 if second_wire_hash >= 2**63 else second_wire_hash + second_wire_hash = ( + second_wire_hash - 2**64 if second_wire_hash >= 2**63 else second_wire_hash + ) 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) @@ -241,3 +243,8 @@ def test_offset_endpoint_port_rejects_bad_input() -> None: # 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) From 9a989538a9da00352c6f646a19cdfa4cfef34b9c Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Tue, 11 Aug 2026 17:34:49 -0700 Subject: [PATCH 15/25] style: fix llm_utils.py import order after merge The merge left the llm_args import block in llm_utils.py isort-dirty (KVEventsConfig sorted before KvCacheConfig). Apply isort ordering; this is the Pre-commit Check CI failure. Signed-off-by: tanmayv25 --- tensorrt_llm/llmapi/llm_utils.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index ca24db67597f..c3feb9c0f91c 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -28,11 +28,10 @@ from .llm_args import (CalibConfig, CudaGraphConfig, DecodeCudaGraphConfig, DraftTargetDecodingConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, - KVEventsConfig, 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 From 4ebfac8f0d6f90e09ed973c4482e213c1df442ae Mon Sep 17 00:00:00 2001 From: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Date: Thu, 13 Aug 2026 02:00:18 -0700 Subject: [PATCH 16/25] fix: address streaming KV event review feedback Addresses the unresolved review threads on NVIDIA/TensorRT-LLM#17023. - Gate the streaming manager on the active KV cache manager V2 backend. StreamingKVCacheEventManager is a duck-typed Python event sink, but the default backend is "cpp", whose nanobind KVCacheManager casts its event_manager to std::shared_ptr. Enabling kv_events_config there raised an opaque TypeError; it now raises an error naming TLLM_KV_CACHE_MANAGER_V2_BACKEND. Expose the active backend as kv_cache_manager_v2.BACKEND. - Scope the AttnLifeCycle filter in _get_event_window_sizes_by_layer_group() to the streaming manager. It previously also dropped SSM layer groups from the buffered KVCacheEventManager's window sizes, changing the existing path on hybrid models where every attention layer is sliding-window. - Reject publish/replay endpoint port ranges that overlap. Both apply the base_port+rank convention, so with N ranks per host a replay base within N-1 of the publish base made one rank's bind collide with another's. The span is the per-host rank count, not the total: a multi-node deployment legitimately reuses the same port numbers on each node. Document the required spacing on the replay_endpoint field. - Reserve publisher sequence numbers before enqueue instead of in the publisher thread, so a batch dropped by a full queue or a failed send leaves a detectable hole. Previously a queue-full drop consumed no sequence number, letting a consumer accept an incomplete stream as complete. - Correct the module header's wire-format claim. The events encode as maps tagged with a "type" key, which is the documented contract for custom router backends and what Dynamo's own TensorRT-LLM publisher emits, not vLLM's array_like positional encoding. Tighten ExternalBlockHash to int, since a bytes hash would fail the decode for the whole batch. - Make the ZeroMQ test deterministic: retry the publish/receive setup on a fresh port and a fresh manager rather than relying on a released port and a fixed sleep for subscription propagation. - Extract the streaming preconditions into validate_streaming_support() so they can be unit tested without building a manager, and cover the backend gate, endpoint-range validation and sequence-gap behavior with tests. - Document kv_cache_config.kv_events_config in docs/source/features/kvcache.md: endpoint convention, replay semantics, wire format, delivery guarantees and the V2-only plus parallelism constraints. Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> --- docs/source/features/kvcache.md | 67 ++++ .../_torch/pyexecutor/kv_cache_events.py | 112 +++++- .../_torch/pyexecutor/kv_cache_manager_v2.py | 42 ++- tensorrt_llm/llmapi/llm_args.py | 5 +- .../runtime/kv_cache_manager_v2/__init__.py | 5 + .../test_streaming_kv_events.py | 320 ++++++++++++------ 6 files changed, 413 insertions(+), 138 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 576927c59475..9be139bc9354 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -208,6 +208,73 @@ 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`, so `N` ranks +occupy `[base_port, base_port + N - 1]`. Co-located engines — for example disaggregated +prefill and decode on one host — must use base ports at least `N` apart, and +```replay_endpoint``` follows the same convention, so its base port must also be at least +`N` away from ```endpoint```'s. Only ranks sharing a host can collide, so `N` here is the +number of ranks per host; a multi-node deployment reuses the same port numbers on each +node. 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/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 4f0390f7f6cd..5bf54819778b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -13,10 +13,14 @@ # 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. The on-wire schema and three-frame framing are -# adapted from vLLM's vllm/distributed/kv_events.py so that external -# KV-cache-aware routers can consume TensorRT-LLM events without translation. +# 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 @@ -38,7 +42,9 @@ 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 -ExternalBlockHash = bytes | int +# Subscribers decode block hashes as 64-bit ints, so a bytes value would fail the +# decode for the entire batch. +ExternalBlockHash = int class EventBatch( @@ -125,7 +131,13 @@ def shutdown(self) -> None: class ZmqEventPublisher(EventPublisher): - """Publishes event batches over the three-frame ZeroMQ wire protocol.""" + """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) @@ -141,7 +153,7 @@ def __init__( topic: str = "", ) -> None: super().__init__(data_parallel_rank) - self._event_queue = Queue[EventBatch | None](maxsize=max_queue_size) + 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 @@ -191,8 +203,12 @@ def publish(self, events: EventBatch) -> bool: 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(events) + self._event_queue.put_nowait((seq, events)) self.enqueued_batches += 1 return True except queue.Full: @@ -201,8 +217,8 @@ def publish(self, events: EventBatch) -> bool: if drops == 1 or (drops & (drops - 1) == 0): logger.warning( f"Dropping streaming KV event batch on rank={self._rank} because " - "the publisher queue is full; " - f"dropped_batches={self.dropped_batches}" + f"the publisher queue is full; seq={seq} will be missing from the " + f"stream; dropped_batches={self.dropped_batches}" ) return False @@ -260,13 +276,13 @@ def _publisher_thread(self) -> None: f"{traceback.format_exc()}" ) try: - event = self._event_queue.get(timeout=0.1) + item = self._event_queue.get(timeout=0.1) except queue.Empty: continue - if event is None: + if item is None: self._event_queue.task_done() break - seq = next(self._seq_gen) + seq, event = item try: payload = encoder.encode(event) self._pub.send_multipart( @@ -281,8 +297,9 @@ def _publisher_thread(self) -> None: except Exception: self._send_error_drops += 1 logger.error( - f"Failed to publish streaming KV event batch rank={self._rank} " - f"seq={seq}\n{traceback.format_exc()}" + 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: @@ -346,6 +363,71 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | 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, + 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) + + +def validate_endpoint_ranges(config: KVEventsConfig, ranks_per_host: int) -> None: + """Reject configurations whose publish and replay port ranges overlap. + + Every rank binds ``base_port + rank`` on both endpoints, so intersecting spans make + one rank's publish bind collide with another's replay bind. Only ranks sharing a + host can collide, so the span is the number of ranks per host, not the total: a + multi-node deployment legitimately reuses the same port numbers on each node. + 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) + if abs(pub_base - replay_base) < span: + raise ValueError( + f"KV event endpoint {config.endpoint!r} and replay_endpoint " + f"{config.replay_endpoint!r} overlap: with {span} rank(s) per host the " + f"publish range is [{pub_base}, {pub_base + span - 1}] and the replay range " + f"is [{replay_base}, {replay_base + span - 1}]. Use base ports {span} apart." + ) + + def create_event_publisher(config: KVEventsConfig, data_parallel_rank: int) -> EventPublisher: """Create the configured publisher for one cache rank.""" if config.publisher == "null": diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 87d731cc5b00..417c457e817c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -74,6 +74,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 @@ -84,7 +85,7 @@ from ...mapping import CpType, Mapping from ..utils import maybe_compile from .connectors.kv_cache_connector import KvCacheConnectorManager -from .kv_cache_events import StreamingKVCacheEventManager +from .kv_cache_events import StreamingKVCacheEventManager, validate_streaming_support from .kv_cache_stats import ( KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, @@ -891,11 +892,17 @@ def __init__( "precedence and the buffered get_kv_cache_events() poll path " "will return no events." ) - if mapping.pp_size > 1: - raise ValueError("Streaming KV events do not support pipeline parallelism") - if mapping.cp_size > 1: - raise ValueError("Streaming KV events do not support context parallelism") 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, + # Only ranks sharing a host can collide on a port. + ranks_per_host=min(mapping.dp_size, mapping.gpus_per_node), + backend=KV_CACHE_MANAGER_V2_BACKEND, + ) if mapping.enable_attention_dp or mpi_rank() == 0: event_rank = mapping.rank if mapping.enable_attention_dp else 0 self.event_manager = StreamingKVCacheEventManager( @@ -1109,7 +1116,9 @@ def append_to_kv_heads_per_layer( raise 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(cache_tiers, tokens_per_block), @@ -1535,11 +1544,19 @@ 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] @@ -1548,13 +1565,10 @@ 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): - life_cycle = self.impl._life_cycles.get_life_cycle(LifeCycleId(layer_group_id)) - # Streaming KV events track 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. - if not isinstance(life_cycle, AttnLifeCycle): - continue + 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 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e8fd74527b61..7e106613188c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3749,7 +3749,10 @@ class KVEventsConfig(StrictBaseModel): default=None, min_length=1, description= - "Optional base ZeroMQ endpoint used to replay KV cache events.") + "Optional base ZeroMQ endpoint used to replay KV cache events. Ranks apply " + "the same base_port+rank convention as `endpoint`, so with N attention-DP " + "ranks per host the two base ports must be at least N 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, diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 51d5ea079a48..eb5e598d87ca 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 @@ -290,6 +294,7 @@ def typed_range(*args: int) -> range: __all__ = [ "AggregatedPageDesc", "AttentionLayerConfig", + "BACKEND", "BAD_PAGE_INDEX", "CACHE_LEVEL1", "BatchDesc", 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 af593bf0c422..c4d6874b3028 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 @@ -14,19 +14,29 @@ # limitations under the License. import socket -import time 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 = 8 +_RECEIVE_TIMEOUT_MS = 2_000 + + +class _NotReceived(Exception): + """No batch arrived: the subscription had not propagated yet.""" + def _unused_tcp_port() -> int: with socket.socket() as sock: @@ -34,126 +44,139 @@ def _unused_tcp_port() -> int: return int(sock.getsockname()[1]) +def _run_on_fresh_port(scenario: Callable[[int], None]) -> None: + """Retry `scenario(port)` until its sockets come up. + + A PUB socket drops messages published before a subscriber's subscription + propagates, and `_unused_tcp_port()` releases its port before the publisher binds + it. Both are transient; 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.""" - port = _unused_tcp_port() - bind_endpoint = f"tcp://*:{port}" - connect_endpoint = f"tcp://127.0.0.1:{port}" topic = "kv-events" context = zmq.Context.instance() - subscriber = context.socket(zmq.SUB) - subscriber.setsockopt_string(zmq.SUBSCRIBE, topic) - subscriber.connect(connect_endpoint) - 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, - ) - try: - manager.set_layer_group_window_sizes({0: 128, 1: 64}) - time.sleep(0.2) + # 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") - root = SimpleNamespace(ordinal=-1) - - def block( - key: bytes, - tokens: list[int], - prev: object, - ) -> SimpleNamespace: - max_window_page = object() - smaller_window_page = object() - return SimpleNamespace( - key=key, - tokens=tokens, - prev=prev, - ordinal=getattr(prev, "ordinal", -1) + 1, - storage=[ - lambda: max_window_page, - lambda: smaller_window_page, - ], + # 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.set_layer_group_window_sizes({0: 128, 1: 64}) - # 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") - second_wire_hash = int.from_bytes(second_hash[:8], "big") - first_wire_hash = first_wire_hash - 2**64 if first_wire_hash >= 2**63 else first_wire_hash - second_wire_hash = ( - second_wire_hash - 2**64 if second_wire_hash >= 2**63 else second_wire_hash - ) - 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) + root = SimpleNamespace(ordinal=-1) - 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() + def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: + max_window_page = object() + smaller_window_page = object() + return SimpleNamespace( + key=key, + tokens=tokens, + prev=prev, + ordinal=getattr(prev, "ordinal", -1) + 1, + storage=[lambda: max_window_page, lambda: smaller_window_page], + ) - frames = [] - for _ in range(2): - assert subscriber.poll(2_000) - frames.append(subscriber.recv_multipart()) - - assert [frame[0] for frame in frames] == [topic.encode(), topic.encode()] - assert [int.from_bytes(frame[1], "big") for frame in frames] == [0, 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 + 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) - # 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() == [] + 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() - # shutdown() must be idempotent. - manager.shutdown() - manager.shutdown() - finally: - manager.shutdown() - subscriber.close(linger=0) + 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()] + assert [int.from_bytes(frame[1], "big") for frame in frames] == [0, 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 - replacement = context.socket(zmq.PUB) - replacement.bind(bind_endpoint) - replacement.close(linger=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: @@ -207,6 +230,87 @@ def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: 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.""" + publisher = ZmqEventPublisher( + data_parallel_rank=0, + endpoint="inproc://kv-events-drop-test", + max_queue_size=1, + ) + try: + # Stop the publisher thread so the queue stays full and the next publish drops. + publisher._running = False + publisher._thread.join(timeout=ZmqEventPublisher.SHUTDOWN_TIMEOUT) + assert not publisher._thread.is_alive() + publisher._running = True + + 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_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, 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) + else: + validate_endpoint_ranges(config, ranks_per_host) + + 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" From e425ac1ec14ced4941b970826f8bc7724d7cb5f3 Mon Sep 17 00:00:00 2001 From: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Date: Wed, 26 Aug 2026 00:54:56 -0700 Subject: [PATCH 17/25] fix: address streaming KV event review comments - Require the target page to cover the whole radix block before publishing a BlockStored. V2 can attach a page adopted from a shorter sibling, where page.num_tokens_in_block < len(block.tokens); publishing that told the external router the engine holds a prefix it cannot fully reuse. The buffered manager already applies this rule in _life_cycle_ids_from_radix_block(). Blocks whose target page covers only part of the block are now suppressed until coverage is complete. - Defer StreamingKVCacheEventManager construction into the protected region. It binds a ZMQ socket and starts a background thread, but initialization still ran asserts and distributed collectives before reaching the cleanup region, so a rank-local failure there leaked the publisher (blocking an in-process retry from rebinding) and could strand peers in a later collective. - Guard the life-cycle hooks against a None life_cycle_id or an unconfigured target instead of calling int() on it. - Correct the endpoint documentation and range diagnostics. Ranks bind base_port+rank by global rank, so the sockets span a cluster-wide range and each rank's port is distinct; only ranks co-located on one host contend for a port, which is why the required spacing between the publish and replay base ports is the per-host rank count rather than the total. Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> --- docs/source/features/kvcache.md | 19 +++-- .../_torch/pyexecutor/kv_cache_events.py | 45 +++++++---- .../_torch/pyexecutor/kv_cache_manager_v2.py | 36 +++++---- tensorrt_llm/llmapi/llm_args.py | 7 +- .../test_streaming_kv_events.py | 77 +++++++++++++++++-- 5 files changed, 141 insertions(+), 43 deletions(-) diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index de4bdde729c5..c947895f3102 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -253,14 +253,17 @@ parallelism and context parallelism are rejected. Events are not published for d 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`, so `N` ranks -occupy `[base_port, base_port + N - 1]`. Co-located engines — for example disaggregated -prefill and decode on one host — must use base ports at least `N` apart, and -```replay_endpoint``` follows the same convention, so its base port must also be at least -`N` away from ```endpoint```'s. Only ranks sharing a host can collide, so `N` here is the -number of ranks per host; a multi-node deployment reuses the same port numbers on each -node. Overlapping ranges are rejected at startup. For `ipc://` and `inproc://` endpoints, -which have no port, each rank appends a `_dp` suffix instead. +**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 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 5bf54819778b..2fb31ac4b9bf 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -380,6 +380,7 @@ def validate_streaming_support( 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. @@ -402,29 +403,37 @@ def validate_streaming_support( "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) + validate_endpoint_ranges(config, ranks_per_host, data_parallel_size) -def validate_endpoint_ranges(config: KVEventsConfig, ranks_per_host: int) -> None: +def validate_endpoint_ranges( + config: KVEventsConfig, ranks_per_host: int, data_parallel_size: int +) -> None: """Reject configurations whose publish and replay port ranges overlap. - Every rank binds ``base_port + rank`` on both endpoints, so intersecting spans make - one rank's publish bind collide with another's replay bind. Only ranks sharing a - host can collide, so the span is the number of ranks per host, not the total: a - multi-node deployment legitimately reuses the same port numbers on each node. - Catch it before any socket is created rather than as an opaque ``EADDRINUSE``. + 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) - if abs(pub_base - replay_base) < span: + 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: with {span} rank(s) per host the " - f"publish range is [{pub_base}, {pub_base + span - 1}] and the replay range " - f"is [{replay_base}, {replay_base + span - 1}]. Use base ports {span} apart." + 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." ) @@ -545,11 +554,21 @@ def add_stored_block_event_from_block(self, block: Any) -> None: if life_cycle_id >= len(block.storage): return page_ref = block.storage[life_cycle_id] - if page_ref is None or page_ref() is None: + 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 @@ -641,7 +660,7 @@ def add_removed_event(self, block_hashes: Any) -> None: self._add_removed_hashes(removed_hashes) def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: - if self._closed: + 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 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 36e95007bb25..406e6d573a03 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -936,6 +936,7 @@ def __init__( for window_size in self.max_attention_window_vec ) self.event_manager: Optional[KVCacheEventManager | StreamingKVCacheEventManager] = None + pending_streaming_event_rank: Optional[int] = None streaming_events_enabled = ( kv_events_config is not None and kv_events_config.enable_kv_cache_events ) @@ -954,19 +955,19 @@ def __init__( kv_events_config, pp_size=mapping.pp_size, cp_size=mapping.cp_size, - # Only ranks sharing a host can collide on a port. + # 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: - 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, - ) - logger.info("Streaming KV event fast path reuses V2 radix block hashes") + # Do not construct it here: it binds a ZMQ socket and starts a + # background thread, and the initialization below still runs asserts + # and distributed collectives that can fail. A rank-local failure in + # that window would leak the publisher (blocking an in-process retry + # from rebinding) and could strand peers in a later collective. Record + # the rank and build it inside the protected region instead. + pending_streaming_event_rank = mapping.rank if mapping.enable_attention_dp else 0 elif self.event_buffer_max_size > 0: if mapping.enable_attention_dp: self.event_manager = KVCacheEventManager( @@ -1151,11 +1152,20 @@ def append_to_kv_heads_per_layer( self.kv_cache_manager_py_config = config - # The streaming event manager has already bound its ZMQ socket and started - # its background thread, so tear it down if impl construction or - # event-manager setup fails here -- otherwise the socket and daemon - # thread leak and an in-process retry cannot rebind the same endpoint. + # The streaming event manager binds a ZMQ socket and starts a background + # thread, so it is created here -- inside the cleanup region -- and torn down + # if anything below fails. Otherwise the socket and daemon thread leak and an + # in-process retry cannot rebind the same endpoint. try: + if pending_streaming_event_rank is not None: + assert kv_events_config is not None + self.event_manager = StreamingKVCacheEventManager( + kv_events_config, + data_parallel_rank=pending_streaming_event_rank, + block_size=self.tokens_per_block, + max_window_size=event_window_size, + ) + logger.info("Streaming KV event fast path reuses V2 radix block hashes") try: self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) except (CuError, KVCacheOutOfMemoryError): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0dcebda65042..90b28d512dc0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3857,9 +3857,10 @@ class KVEventsConfig(StrictBaseModel): min_length=1, description= "Optional base ZeroMQ endpoint used to replay KV cache events. Ranks apply " - "the same base_port+rank convention as `endpoint`, so with N attention-DP " - "ranks per host the two base ports must be at least N apart or a rank's replay " - "bind collides with another rank's publish bind on that host.") + "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, 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 c4d6874b3028..c4cdd9b832c0 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 @@ -103,8 +103,8 @@ def scenario(port: int) -> None: root = SimpleNamespace(ordinal=-1) def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: - max_window_page = object() - smaller_window_page = object() + 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, @@ -199,7 +199,7 @@ def test_streaming_removals_are_never_dropped_by_the_entry_cap() -> None: root = SimpleNamespace(ordinal=-1) def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: - page = object() + page = SimpleNamespace(num_tokens_in_block=len(tokens)) return SimpleNamespace( key=key, tokens=tokens, @@ -262,7 +262,7 @@ def test_dropped_batches_leave_a_sequence_gap() -> None: 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, backend="python") + 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) @@ -306,9 +306,74 @@ def test_validate_endpoint_ranges(endpoint, replay_endpoint, ranks_per_host, ove 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) + validate_endpoint_ranges(config, ranks_per_host, ranks_per_host) else: - validate_endpoint_ranges(config, ranks_per_host) + 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, + ) + 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, + ) + 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: From 44004b3864a76c01690e02cedd10678b598233dd Mon Sep 17 00:00:00 2001 From: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:41:38 -0700 Subject: [PATCH 18/25] refactor: start the streaming KV event publisher explicitly Constructing StreamingKVCacheEventManager no longer binds a socket or starts a thread. ZmqEventPublisher.__init__ only records configuration, and a new start() performs the bind and launches the publisher thread; EventPublisher grows a no-op start() so the null publisher inherits it, and shutdown() tolerates a publisher that never started. KVCacheManagerV2 builds the manager at its original position and calls start() as the last statement of __init__, so nothing is acquired until every other check has passed. That removes the need to wrap initialization in a cleanup region: a failure anywhere earlier -- including the rank-coordinated aborts, where a rank raises only because a peer failed -- leaves no socket bound and no thread running, so there is nothing to tear down. It also leaves the rank-coordinated initialization block untouched, which keeps this feature clear of the region under active refactoring. Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_events.py | 45 ++++- .../_torch/pyexecutor/kv_cache_manager_v2.py | 184 ++++++++---------- .../test_streaming_kv_events.py | 54 ++++- 3 files changed, 168 insertions(+), 115 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 2fb31ac4b9bf..aff9119182b9 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -111,6 +111,14 @@ class EventPublisher(ABC): 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.""" @@ -170,15 +178,24 @@ def __init__( 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: - # __init__ never returns on failure, so shutdown() is unreachable; - # close the sockets here to avoid leaking them on the shared context. + # 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, @@ -188,7 +205,7 @@ def __init__( self._thread.start() logger.info( f"Started streaming KV event publisher rank={self._rank} " - f"endpoint={self._endpoint} topic={topic!r}" + f"endpoint={self._endpoint} topic={self._topic!r}" ) @property @@ -232,12 +249,13 @@ def shutdown(self) -> None: except queue.Full: # The thread exits after draining the full queue. pass - 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" - ) + 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} " @@ -513,6 +531,15 @@ def __init__( 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) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 8167cf7663ca..e44ae9932186 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -956,7 +956,6 @@ def __init__( for window_size in self.max_attention_window_vec ) self.event_manager: Optional[KVCacheEventManager | StreamingKVCacheEventManager] = None - pending_streaming_event_rank: Optional[int] = None streaming_events_enabled = ( kv_events_config is not None and kv_events_config.enable_kv_cache_events ) @@ -981,13 +980,15 @@ def __init__( backend=KV_CACHE_MANAGER_V2_BACKEND, ) if mapping.enable_attention_dp or mpi_rank() == 0: - # Do not construct it here: it binds a ZMQ socket and starts a - # background thread, and the initialization below still runs asserts - # and distributed collectives that can fail. A rank-local failure in - # that window would leak the publisher (blocking an in-process retry - # from rebinding) and could strand peers in a later collective. Record - # the rank and build it inside the protected region instead. - pending_streaming_event_rank = mapping.rank if mapping.enable_attention_dp else 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( @@ -1170,109 +1171,84 @@ def append_to_kv_heads_per_layer( isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers ) - # The streaming event manager binds a ZMQ socket and starts a background - # thread, so it is created here -- inside the cleanup region -- and torn down if - # anything below fails. That includes the rank-coordinated abort paths, where - # this rank raises only because a peer failed and so has no local exception of - # its own. Otherwise the socket and daemon thread leak and an in-process retry - # cannot rebind the same endpoint. - try: - if pending_streaming_event_rank is not None: - assert kv_events_config is not None - self.event_manager = StreamingKVCacheEventManager( - kv_events_config, - data_parallel_rank=pending_streaming_event_rank, - block_size=self.tokens_per_block, - max_window_size=event_window_size, - ) - logger.info("Streaming KV event fast path reuses V2 radix block hashes") - - candidate: Optional[KVCacheManagerPy] = None - if not has_host_cache_tier: + candidate: Optional[KVCacheManagerPy] = None + if not has_host_cache_tier: + candidate = KVCacheManagerPy(config, event_manager=self.event_manager) + else: + init_error: Optional[Exception] = None + local_init_status = _KVCacheManagerInitStatus.KEEP_HOST + try: candidate = KVCacheManagerPy(config, event_manager=self.event_manager) - else: - init_error: Optional[Exception] = None - local_init_status = _KVCacheManagerInitStatus.KEEP_HOST + except Exception as error: + if isinstance(error, (CuError, KVCacheOutOfMemoryError)): + local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST + else: + init_error = error.with_traceback(None) + local_init_status = _KVCacheManagerInitStatus.ABORT + + init_status = _sync_kv_cache_manager_init_status(local_init_status, mapping) + + if init_status == _KVCacheManagerInitStatus.ABORT: + if candidate is not None: + candidate.shutdown() + if init_error is not None: + raise init_error + raise RuntimeError("KV cache manager initialization failed on another rank") + + if init_status == _KVCacheManagerInitStatus.USE_NO_HOST: + logger.warning( + "At least one rank could not use the KV cache manager host tier " + "(cuMemHostRegister may have failed). Rebuilding without the " + "host cache tier on all ranks." + ) + fallback_error: Optional[Exception] = None try: + if candidate is not None: + candidate.shutdown() + candidate = None + config = replace( + config, + cache_tiers=[ + tier + for tier in config.cache_tiers + if not isinstance(tier, HostCacheTierConfig) + ], + ) candidate = KVCacheManagerPy(config, event_manager=self.event_manager) except Exception as error: - if isinstance(error, (CuError, KVCacheOutOfMemoryError)): - local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST - else: - init_error = error.with_traceback(None) - local_init_status = _KVCacheManagerInitStatus.ABORT + fallback_error = error.with_traceback(None) - init_status = _sync_kv_cache_manager_init_status(local_init_status, mapping) + local_fallback_status = ( + _KVCacheManagerInitStatus.USE_NO_HOST + if fallback_error is None + else _KVCacheManagerInitStatus.ABORT + ) + fallback_status = _sync_kv_cache_manager_init_status(local_fallback_status, mapping) - if init_status == _KVCacheManagerInitStatus.ABORT: + if fallback_status == _KVCacheManagerInitStatus.ABORT: if candidate is not None: candidate.shutdown() - if init_error is not None: - raise init_error - raise RuntimeError("KV cache manager initialization failed on another rank") - - if init_status == _KVCacheManagerInitStatus.USE_NO_HOST: - logger.warning( - "At least one rank could not use the KV cache manager host tier " - "(cuMemHostRegister may have failed). Rebuilding without the " - "host cache tier on all ranks." - ) - fallback_error: Optional[Exception] = None - try: - if candidate is not None: - candidate.shutdown() - candidate = None - config = replace( - config, - cache_tiers=[ - tier - for tier in config.cache_tiers - if not isinstance(tier, HostCacheTierConfig) - ], - ) - candidate = KVCacheManagerPy(config, event_manager=self.event_manager) - except Exception as error: - fallback_error = error.with_traceback(None) - - local_fallback_status = ( - _KVCacheManagerInitStatus.USE_NO_HOST - if fallback_error is None - else _KVCacheManagerInitStatus.ABORT - ) - fallback_status = _sync_kv_cache_manager_init_status( - local_fallback_status, mapping + if fallback_error is not None: + raise fallback_error + raise RuntimeError( + "KV cache manager initialization without the host cache tier " + "failed on another rank" ) - if fallback_status == _KVCacheManagerInitStatus.ABORT: - if candidate is not None: - candidate.shutdown() - if fallback_error is not None: - raise fallback_error - raise RuntimeError( - "KV cache manager initialization without the host cache tier " - "failed on another rank" - ) - - assert candidate is not None - self.kv_cache_manager_py_config = config - self.impl = candidate - 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( - 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 - ), - self._get_event_layer_group_ids(), + assert candidate is not None + self.kv_cache_manager_py_config = config + self.impl = candidate + 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( + attention_only=isinstance(self.event_manager, StreamingKVCacheEventManager) ) - except Exception: - if isinstance(self.event_manager, StreamingKVCacheEventManager): - self.event_manager.shutdown() - raise + ) + self.event_manager.add_created_event( + self._get_event_num_blocks_per_cache_level(config.cache_tiers, tokens_per_block), + self._get_event_layer_group_ids(), + ) # Both backends build layer_grouping on demand, and the layer order # within a group is not part of its contract. Cache a stable physical- @@ -1387,6 +1363,14 @@ def append_to_kv_heads_per_layer( 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. 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 c4cdd9b832c0..54e5ae10b9d8 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 @@ -98,6 +98,7 @@ def scenario(port: int) -> None: block_size=4, max_window_size=128, ) + manager.start() manager.set_layer_group_window_sizes({0: 128, 1: 64}) root = SimpleNamespace(ordinal=-1) @@ -188,6 +189,7 @@ def test_streaming_removals_are_never_dropped_by_the_entry_cap() -> None: max_window_size=128, max_entries=2, ) + manager.start() try: manager.set_layer_group_window_sizes({0: 128}) @@ -232,18 +234,14 @@ def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: 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: - # Stop the publisher thread so the queue stays full and the next publish drops. - publisher._running = False - publisher._thread.join(timeout=ZmqEventPublisher.SHUTDOWN_TIMEOUT) - assert not publisher._thread.is_alive() - publisher._running = True - 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 @@ -260,6 +258,48 @@ def test_dropped_batches_leave_a_sequence_gap() -> None: 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 + # The endpoint is still free, so an unrelated socket can take it. + context = zmq.Context.instance() + squatter = context.socket(zmq.PUB) + squatter.bind(endpoint) + squatter.close(linger=0) + + 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") @@ -319,6 +359,7 @@ def test_partial_target_page_coverage_is_suppressed_until_fully_covered() -> Non block_size=4, max_window_size=128, ) + manager.start() try: manager.set_layer_group_window_sizes({0: 128}) published: list[object] = [] @@ -362,6 +403,7 @@ def test_life_cycle_hooks_ignore_none_ids() -> None: 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. From ab55d430276ac01a64b40b3152085799f2c185cb Mon Sep 17 00:00:00 2001 From: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:25:22 -0700 Subject: [PATCH 19/25] fix: synchronise the streaming KV event test on the subscription test_streaming_fast_path_publishes_only_full_max_window_blocks failed in CI. A PUB socket silently drops everything published before a subscriber's subscription has propagated, and the manager publishes both batches within a millisecond of binding, so they were always lost. Retrying on a fresh port did not help -- every attempt lost them identically, which is why the failure was deterministic rather than flaky. Publish probe batches until one is actually received, drain them, and only then produce the events under test. The probes consume sequence numbers, so assert the two batches are dense from the probe count rather than exactly [0, 1]; that still pins the counter's origin at zero, and density is the property subscribers rely on to detect loss. The fresh-port retry stays, now only for the port race it was meant to cover. Also drop the squatter socket from test_construction_binds_nothing_until_start: binding and closing the endpoint to prove it was free left it in TIME_WAIT, so the following start() could not rebind it. Asserting that the publisher holds no socket and no thread already carries that meaning. Verified in nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc24: 24 passed. Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> --- .../test_streaming_kv_events.py | 42 ++++++++++++++----- 1 file changed, 31 insertions(+), 11 deletions(-) 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 54e5ae10b9d8..c0c6775d800a 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 @@ -30,8 +30,10 @@ ) from tensorrt_llm.llmapi.llm_args import KVEventsConfig -_ZMQ_SETUP_ATTEMPTS = 8 +_ZMQ_SETUP_ATTEMPTS = 4 _RECEIVE_TIMEOUT_MS = 2_000 +_SUBSCRIBE_ATTEMPTS = 50 +_PROBE_TIMEOUT_MS = 100 class _NotReceived(Exception): @@ -44,12 +46,30 @@ def _unused_tcp_port() -> int: 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)` until its sockets come up. + """Retry `scenario(port)` on a fresh port if its sockets could not come up. - A PUB socket drops messages published before a subscriber's subscription - propagates, and `_unused_tcp_port()` releases its port before the publisher binds - it. Both are transient; assertion failures inside `scenario` are not retried. + `_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: @@ -99,6 +119,7 @@ def scenario(port: int) -> None: 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) @@ -133,7 +154,11 @@ def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: frames.append(subscriber.recv_multipart()) assert [frame[0] for frame in frames] == [topic.encode(), topic.encode()] - assert [int.from_bytes(frame[1], "big") for frame in frames] == [0, 1] + # 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 @@ -272,11 +297,6 @@ def test_construction_binds_nothing_until_start() -> None: publisher = manager._publisher assert publisher._pub is None assert publisher._thread is None - # The endpoint is still free, so an unrelated socket can take it. - context = zmq.Context.instance() - squatter = context.socket(zmq.PUB) - squatter.bind(endpoint) - squatter.close(linger=0) manager.start() assert publisher._pub is not None From 573b40e21bb5d38c9c4c8c6077019c542395148c Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Sat, 29 Aug 2026 18:29:24 -0700 Subject: [PATCH 20/25] fix early return for rank 0 Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/kv_cache_events.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index aff9119182b9..2352df07cc06 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -351,13 +351,13 @@ def _service_replay(self) -> None: @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 or data_parallel_rank == 0: + 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 f"{endpoint}_dp{data_parallel_rank}" + 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: @@ -371,6 +371,8 @@ def offset_endpoint_port(endpoint: str | None, data_parallel_rank: int) -> str | 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: From e959a2937e110f6c86c27ec3facf149e983b493b Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Sat, 29 Aug 2026 18:55:24 -0700 Subject: [PATCH 21/25] fix malformed llm args import Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/_util.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 3112bfe9a6de..e1a34d972d4e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -30,10 +30,6 @@ from tensorrt_llm.llmapi.llm_args import ( CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, KVEventsConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, - MultimodalEncoderSchedulingPolicy, PeftCacheConfig, SamplerType, - SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, - WaitingQueuePolicy) - KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, MultimodalEncoderSchedulingPolicy, PeftCacheConfig, SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy) # isort: on From 700f6038f85178c575e1f1afc54d34448b5c9401 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Tue, 1 Sep 2026 14:23:03 -0700 Subject: [PATCH 22/25] Fix KV cache manager merge resolution Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/_util.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 92b40ba6d925..0114c575e310 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2245,8 +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, kv_events_config: Optional[KVEventsConfig] = None) -> KVCacheManager: - cold_page_codec_provider: Optional[object] = None) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config @@ -2383,14 +2383,14 @@ def _create_kv_cache_manager( manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): 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__}.") - manager_extra_kwargs[ - "cold_page_codec_provider"] = cold_page_codec_provider if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg From 38bf92ae59526c31fc695a57a23f847f735affc4 Mon Sep 17 00:00:00 2001 From: Guan Luo <41310872+GuanLuo@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:08:43 -0700 Subject: [PATCH 23/25] fix: validate the KV event port span on every rank ZmqEventPublisher.__init__ resolves base_port + rank through offset_endpoint_port() before any collective runs, and that raises only on the ranks whose offset overflows the u16 range. With endpoint "tcp://*:65535" and DP size 2, rank 1 aborts there while rank 0 accepts the port and proceeds into the following all-reduce, where it waits on a peer that is already gone. Extend validate_endpoint_ranges(), which validate_streaming_support() already calls on every rank before any collective, to also reject base_port + data_parallel_size - 1 > 65535 for both the publish and replay endpoints. Every rank then fails identically, with an error naming the endpoint, the rank count and the highest base port that fits. Covered by a parameterized validation-level regression (including the reported 65535/DP-2 case, the 65534 boundary, the replay endpoint checked independently, and ipc:// exempted) plus a case pinning that rank 0 refuses what rank 1 would. Signed-off-by: Guan Luo <41310872+GuanLuo@users.noreply.github.com> --- .../_torch/pyexecutor/kv_cache_events.py | 34 ++++++++++++---- .../test_streaming_kv_events.py | 39 +++++++++++++++++++ 2 files changed, 66 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 2352df07cc06..d69096b46a6a 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -429,23 +429,43 @@ def validate_streaming_support( def validate_endpoint_ranges( config: KVEventsConfig, ranks_per_host: int, data_parallel_size: int ) -> None: - """Reject configurations whose publish and replay port ranges overlap. + """Reject endpoint configurations the per-rank binding cannot honour. 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``. + distinct cluster-wide and the sockets span ``[base, base + world - 1]``. Two things + can go wrong, and both are checked on every rank -- before any socket is created and + before the initialization collectives -- so all ranks fail identically rather than + one aborting while its peers wait in an all-reduce: + + * The span can run past port 65535. ``ZmqEventPublisher.__init__`` resolves + ``base_port + rank`` itself, but only raises on the ranks that actually overflow. + * The publish and replay spans can intersect. Only ranks co-located on one host + 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, not the total. """ + world = max(1, data_parallel_size) pub_base = _tcp_base_port(config.endpoint) replay_base = _tcp_base_port(config.replay_endpoint) + + for name, endpoint, base_port in ( + ("endpoint", config.endpoint, pub_base), + ("replay_endpoint", config.replay_endpoint, replay_base), + ): + if base_port is None: + continue + highest = base_port + world - 1 + if highest > 65_535: + raise ValueError( + f"KV event {name} {endpoint!r} does not fit {world} rank(s): ranks bind " + f"base_port+rank, so the highest would be {highest}, above the maximum " + f"port 65535. Use a base port at or below {65_535 - world + 1}." + ) + 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 " 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 c0c6775d800a..dc460b56c1eb 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 @@ -438,6 +438,45 @@ def test_life_cycle_hooks_ignore_none_ids() -> None: manager.shutdown() +@pytest.mark.parametrize( + "endpoint,replay_endpoint,dp_size,overflows", + [ + # rank 1 would bind 65536, which offset_endpoint_port rejects on that rank + # alone -- before rank 0 reaches the collective. + ("tcp://*:65535", None, 2, True), + ("tcp://*:65535", None, 1, False), + ("tcp://*:65534", None, 2, False), + ("tcp://*:65534", None, 3, True), + # The replay endpoint is checked too, not just the publish endpoint. + ("tcp://*:5557", "tcp://*:65535", 2, True), + ("tcp://*:5557", "tcp://*:60000", 2, False), + # ipc/inproc have no port, so the span does not apply. + ("ipc:///tmp/kv-events", None, 64, False), + ], +) +def test_validate_endpoint_ranges_rejects_port_overflow( + endpoint, replay_endpoint, dp_size, overflows +) -> None: + kwargs = {"replay_endpoint": replay_endpoint} if replay_endpoint else {} + config = KVEventsConfig(enable_kv_cache_events=True, endpoint=endpoint, **kwargs) + # ranks_per_host=1 isolates the span check from the overlap check. + if overflows: + with pytest.raises(ValueError, match="above the maximum port 65535"): + validate_endpoint_ranges(config, 1, dp_size) + else: + validate_endpoint_ranges(config, 1, dp_size) + + +def test_validate_streaming_support_rejects_overflowing_port_span() -> None: + """Every rank must reject the span, so none reaches the following collective.""" + config = KVEventsConfig(enable_kv_cache_events=True, endpoint="tcp://*:65535") + supported = dict(pp_size=1, cp_size=1, ranks_per_host=1, backend="python") + # One rank fits; two do not, and rank 0 must refuse it just as rank 1 would. + validate_streaming_support(config, **supported, data_parallel_size=1) + with pytest.raises(ValueError, match="above the maximum port 65535"): + validate_streaming_support(config, **supported, data_parallel_size=2) + + 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" From e40d3b271bde9c34bb975b8b0147f74dcde146c8 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 7 Sep 2026 11:11:44 -0700 Subject: [PATCH 24/25] fix: resolve KV cache manager signature merge conflict Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index cad48b1b67ff..3b330300f5e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2367,7 +2367,7 @@ def _create_kv_cache_manager( kv_cache_type=None, is_disagg: bool = False, cold_page_codec_provider: Optional[object] = None, - kv_events_config: Optional[KVEventsConfig] = None) -> KVCacheManager: + kv_events_config: Optional[KVEventsConfig] = None, joint_kv_cache_reuse: bool = False) -> KVCacheManager: """ Returns: From d5d3337ac3d30f88ffc69689b7eef9bd9728efa2 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Mon, 7 Sep 2026 11:13:15 -0700 Subject: [PATCH 25/25] fix: preserve joint KV cache reuse during merge resolution Signed-off-by: Allison Lim --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 3b330300f5e0..40c71952f04c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2508,12 +2508,12 @@ def _create_kv_cache_manager( manager_extra_kwargs[ "cold_page_codec_provider"] = cold_page_codec_provider manager_extra_kwargs["kv_events_config"] = kv_events_config + manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse 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__}.") - manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg