diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 40ab2e3a64ed..1e125ad78eb6 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_cache_config.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,12 @@ 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 + 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/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py new file mode 100644 index 000000000000..eb0734216896 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -0,0 +1,618 @@ +# 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 +import traceback +from abc import ABC, abstractmethod +from collections import deque +from itertools import count +from queue import Queue +from typing import Any, Optional + +import msgspec +import zmq + +from tensorrt_llm.llmapi.llm_args import KVEventsConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime.kv_cache_hash import truncate_sha256_hash_to_int64 +from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import KVCacheEvent, KVCacheEventDiff + +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._queue_full_drops = 0 + self._send_error_drops = 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}" + ) + + @property + def dropped_batches(self) -> int: + # Two independent writers: the scheduler thread bumps _queue_full_drops + # (queue full) and the publisher thread bumps _send_error_drops (send + # failure). Each counter has a single writer, so the sum needs no lock. + return self._queue_full_drops + self._send_error_drops + + def publish(self, events: EventBatch) -> bool: + if not self._running: + return False + if events.data_parallel_rank is None: + events.data_parallel_rank = self._data_parallel_rank + try: + self._event_queue.put_nowait(events) + self.enqueued_batches += 1 + return True + except queue.Full: + self._queue_full_drops += 1 + drops = self._queue_full_drops + if drops == 1 or (drops & (drops - 1) == 0): + logger.warning( + f"Dropping 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 not self._endpoint: + raise ValueError("KV event publisher endpoint must not be empty") + if not self._endpoint.startswith(("tcp://", "ipc://", "inproc://")): + raise ValueError(f"Unsupported KV event endpoint scheme: {self._endpoint!r}") + # The publisher owns its endpoint and subscribers connect to it, so the + # PUB socket always binds -- including explicit-host TCP binds like + # tcp://0.0.0.0:5557 that the previous '*'-only heuristic wrongly + # treated as connect targets (silently dropping every event). + self._pub.bind(self._endpoint) + + if self._replay_endpoint is not None: + self._replay = self._ctx.socket(zmq.ROUTER) + self._replay.bind(self._replay_endpoint) + + def _publisher_thread(self) -> None: + encoder = msgspec.msgpack.Encoder() + assert self._pub is not None + try: + while self._running or not self._event_queue.empty(): + if self._replay is not None and self._replay.poll(0): + try: + self._service_replay() + except Exception: + logger.error( + "Failed to service native KV event replay request\n" + f"{traceback.format_exc()}" + ) + 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._send_error_drops += 1 + 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: + 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 + # 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(":") + 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', 'ipc', 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 _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") + # 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 + + +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. + + Implements the V2 KV-cache-manager event-sink hook interface by duck + typing rather than inheriting ``KVCacheEventManager``: it fully replaces + event production (reusing the radix block hashes) and shares none of the + base manager's state, so subclassing would only risk partially initialised + base attributes. + """ + + def __init__( + self, + config: KVEventsConfig, + *, + data_parallel_rank: int, + block_size: int, + max_window_size: int, + max_entries: int = 50_000, + ) -> None: + self._rank = data_parallel_rank + self._publisher = create_event_publisher(config, data_parallel_rank) + self._block_size = block_size + self._max_window_size = max_window_size + self._max_entries = max_entries + self._target_life_cycle_id: int | None = None + self._stored_blocks: dict[bytes, _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 + 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 = [ + 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_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 + 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.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): + 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 self._closed: + return + if isinstance(block_hashes, (bytes, str, int)): + block_hashes = (block_hashes,) + removed_hashes: list[ExternalBlockHash] = [] + for block_key in block_hashes: + if not isinstance(block_key, bytes): + continue + 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 self._closed: + return + 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 + # Removals are never dropped by the per-iteration cap and, unlike stores, + # do not consume the _pending_entries budget: each hash was already + # reported as stored (so removals are bounded by the stored set), and + # counting them against the store budget would starve legitimate + # BlockStored events in a removal-heavy iteration. + if self._pending_events and isinstance(self._pending_events[-1], BlockRemoved): + self._pending_events[-1].block_hashes.extend(block_hashes) + else: + self._pending_events.append(BlockRemoved(block_hashes=block_hashes, medium="GPU")) + self.removed_blocks += len(block_hashes) + + def add_updated_event( + self, + block_hash: Any, + *, + cache_level: KVCacheEventDiff | None = None, + priority: KVCacheEventDiff | None = None, + layer_group_id: int | None = None, + ) -> None: + return + + def _reserve_entries(self, num_entries: int) -> bool: + if self._pending_entries + num_entries <= self._max_entries: + self._pending_entries += num_entries + return True + self.dropped_events += num_entries + if self.dropped_events == num_entries or ( + self.dropped_events & (self.dropped_events - 1) == 0 + ): + logger.warning( + "Dropping 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 + batch = KVEventBatch( + ts=time.time(), + events=events, + data_parallel_rank=self._rank, + ) + try: + if self._publisher.publish(batch): + self.enqueued_batches += 1 + self.enqueued_events += len(events) + else: + self.dropped_batches += 1 + except Exception: + self.dropped_batches += 1 + logger.error( + f"Dropping 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 + # 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: + 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={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 518c7b711162..963a8b880066 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 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, @@ -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 NativeKVCacheEventManager 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 @@ -866,8 +868,33 @@ def __init__( self.max_seq_len if window_size is None else int(window_size) for window_size in self.max_attention_window_vec ) - self.event_manager: Optional[KVCacheEventManager] = None - if self.event_buffer_max_size > 0: + self.event_manager: Optional[KVCacheEventManager | NativeKVCacheEventManager] = None + native_events_enabled = ( + 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: + 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.event_manager = NativeKVCacheEventManager( + 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") + elif self.event_buffer_max_size > 0: if mapping.enable_attention_dp: self.event_manager = KVCacheEventManager( self.event_buffer_max_size, @@ -1046,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. @@ -1466,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) @@ -2893,13 +2938,23 @@ def get_kv_cache_stats(self): return kv_cache_stats def flush_iteration_events(self): - if self.event_manager is not None: - self.event_manager.flush_iteration_events() + event_manager = self.event_manager + if event_manager is not None: + event_manager.flush_iteration_events() def get_latest_events(self, timeout_ms: Optional[float] = None): - if self.event_manager is None: + # 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. 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: + return isinstance(self.event_manager, NativeKVCacheEventManager) def get_iteration_stats(self): if not self.enable_stats: @@ -3427,6 +3482,13 @@ 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 + # 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/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34ebd340a5ee..500f46aa4b22 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -645,7 +645,9 @@ def __init__( self._is_kv_manager_v2 = isinstance(self.kv_cache_manager, KVCacheManagerV2) self._prefetched_request_ids: set[int] = set() - self.enable_kv_cache_events = self.kv_cache_manager is not None and self.kv_cache_manager.event_buffer_max_size > 0 + self.enable_kv_cache_events = self.kv_cache_manager is not None and ( + self.kv_cache_manager.event_buffer_max_size > 0 or getattr( + self.kv_cache_manager, "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/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 1bd895dbd59b..48fd5a0e048d 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, KvCacheConfig, + KVEventsConfig, 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..86ef9f0268e9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3618,6 +3618,52 @@ 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", + 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, + gt=0, + description="Number of previously published batches retained for replay." + ) + hwm: int = Field(default=100_000, + gt=0, + description="ZeroMQ publisher socket high-water mark. " + "0 means unlimited in ZeroMQ, so it is disallowed here.") + max_queue_size: int = Field( + default=100_000, + gt=0, + description="Maximum number of batches queued for background publishing. " + "Must be positive; 0 would make the queue unbounded.") + topic: str = Field( + default="", + description="ZeroMQ subscription topic used for KV cache event batches." + ) + + def model_post_init(self, __context) -> None: + if self.publisher is None: + self.publisher = "zmq" if self.enable_kv_cache_events else "null" + + @PybindMirror.mirror_pybind_fields(_KvCacheConfig) class KvCacheConfig(StrictBaseModel, PybindMirror): """Configuration for the KV cache.""" @@ -3685,6 +3731,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= 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..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", 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..f21f26473d4c --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py @@ -0,0 +1,183 @@ +# 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 +from types import SimpleNamespace + +import msgspec +import zmq + +from tensorrt_llm._torch.pyexecutor.kv_cache_events import BlockRemoved, NativeKVCacheEventManager +from tensorrt_llm.llmapi.llm_args import KVEventsConfig + + +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_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}" + topic = "kv-events" + context = zmq.Context.instance() + subscriber = context.socket(zmq.SUB) + subscriber.setsockopt_string(zmq.SUBSCRIBE, topic) + subscriber.connect(connect_endpoint) + + manager = NativeKVCacheEventManager( + KVEventsConfig( + enable_kv_cache_events=True, + publisher="zmq", + endpoint=bind_endpoint, + topic=topic, + max_queue_size=8, + ), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + ) + manager.set_layer_group_window_sizes({0: 128, 1: 64}) + time.sleep(0.2) + + root = SimpleNamespace(ordinal=-1) + + def block( + key: bytes, + tokens: list[int], + prev: object, + ) -> SimpleNamespace: + max_window_page = object() + smaller_window_page = object() + return SimpleNamespace( + key=key, + tokens=tokens, + prev=prev, + ordinal=getattr(prev, "ordinal", -1) + 1, + storage=[ + lambda: max_window_page, + lambda: smaller_window_page, + ], + ) + + 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): + assert subscriber.poll(2_000) + frames.append(subscriber.recv_multipart()) + + assert [frame[0] for frame in frames] == [topic.encode(), topic.encode()] + assert [int.from_bytes(frame[1], "big") for frame in frames] == [0, 1] + stored_batch = msgspec.msgpack.decode(frames[0][2]) + removed_batch = msgspec.msgpack.decode(frames[1][2]) + assert stored_batch[2] == 0 + assert stored_batch[1] == [ + { + "type": "BlockStored", + "block_hashes": [first_wire_hash, second_wire_hash], + "parent_block_hash": None, + "token_ids": [1, 2, 3, 4, 5, 6, 7, 8], + "block_size": 4, + "lora_id": None, + "medium": "GPU", + "lora_name": None, + } + ] + assert removed_batch[1] == [ + { + "type": "BlockRemoved", + "block_hashes": [first_wire_hash, second_wire_hash], + "medium": "GPU", + } + ] + assert manager.stored_blocks == 2 + assert manager.removed_blocks == 2 + assert manager.partial_blocks_suppressed == 1 + assert manager.non_target_life_cycles_ignored == 1 + assert manager.dropped_events == 0 + + # Native publishing pushes events out-of-band, so the legacy pull API must + # degrade to an empty result rather than raising. + assert manager.get_latest_events() == [] + + manager.shutdown() + manager.shutdown() + subscriber.close(linger=0) + + 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.""" + manager = NativeKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=2, + max_window_size=128, + max_entries=2, + ) + manager.set_layer_group_window_sizes({0: 128}) + + root = SimpleNamespace(ordinal=-1) + + def block(key: bytes, tokens: list[int], prev: object) -> SimpleNamespace: + page = object() + return SimpleNamespace( + key=key, + tokens=tokens, + prev=prev, + ordinal=getattr(prev, "ordinal", -1) + 1, + storage=[lambda: page], + ) + + first = block(b"\x01" * 32, [1, 2], root) + second = block(b"\x02" * 32, [3, 4], first) + manager.add_stored_block_event_from_block(first) + manager.add_stored_block_event_from_block(second) + + # Both stores fill the entry cap (max_entries=2); the removals must still be + # emitted rather than dropped, or the consumer treats the blocks as resident + # forever. + manager.add_removed_event([b"\x01" * 32, b"\x02" * 32]) + + removed = [event for event in manager._pending_events if isinstance(event, BlockRemoved)] + assert manager.removed_blocks == 2 + assert sum(len(event.block_hashes) for event in removed) == 2 + + manager.shutdown()