From 92b8fd38a4fb162d490c124c9402af1fc6d7e540 Mon Sep 17 00:00:00 2001 From: Alec Flowers Date: Fri, 24 Jul 2026 23:26:56 -0700 Subject: [PATCH 01/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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/11] 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