[None][feat] Add ZMQ KV-cache event tee with gap replay

requirements.txt

@@ -46,6 +46,7 @@ click
 click_option_group
 aenum
 pyzmq
+msgspec
 fastapi>=0.120.1,<=0.121.3
 starlette>=0.49.1
 uvicorn

tensorrt_llm/commands/serve.py

@@ -30,6 +30,7 @@
 from tensorrt_llm.commands.utils import get_is_diffusion_model
 from tensorrt_llm.executor.utils import LlmLauncherEnvs
 from tensorrt_llm.inputs.multimodal import MultimodalServerConfig
+from tensorrt_llm.serve.kv_events_config import KVEventsConfig
 from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy,
                                  DynamicBatchConfig, KvCacheConfig,
                                  SchedulerConfig)
@@ -372,7 +373,8 @@ def launch_server(
         server_role: Optional[ServerRole] = None,
         disagg_cluster_config: Optional[DisaggClusterConfig] = None,
         multimodal_server_config: Optional[MultimodalServerConfig] = None,
-        served_model_name: Optional[str] = None):
+        served_model_name: Optional[str] = None,
+        kv_events_config: Optional[KVEventsConfig] = None):
 
     # Install the SIGUSR1 signal handler for debugging
     signal.signal(signal.SIGUSR1, print_stack_trace)
@@ -418,7 +420,8 @@ def launch_server(
                               metadata_server_cfg=metadata_server_cfg,
                               disagg_cluster_config=disagg_cluster_config,
                               multimodal_server_config=multimodal_server_config,
-                              chat_template=chat_template)
+                              chat_template=chat_template,
+                              kv_events_config=kv_events_config)
         _apply_fastapi_middlewares(server.app, middleware)
 
         # Optionally disable GC (default: not disabled)
@@ -923,6 +926,19 @@ def convert(self, value: Any, param: Optional["click.Parameter"],
     default=None,
     help=
     "Types of agents to schedule. Now Only Support Open Deep Research agent.")
+@click.option(
+    "--kv_events_config",
+    type=str,
+    default=None,
+    help="JSON config for teeing KV-cache events (including token_ids) over "
+    "ZMQ in vLLM wire format, mirroring vLLM's KVEventsConfig. Read-only: does "
+    "not affect the SSE /kv_cache_events path. Keys: enable_kv_cache_events "
+    "(bool, master switch), endpoint (PUB, e.g. tcp://*:5557), replay_endpoint "
+    "(ROUTER for gap recovery, e.g. tcp://*:5558; null disables replay), "
+    "buffer_steps, hwm, max_queue_size, topic. Disabled unless "
+    "enable_kv_cache_events is true. Example: "
+    '\'{"enable_kv_cache_events": true, "endpoint": "tcp://*:5557", '
+    '"replay_endpoint": "tcp://*:5558"}\'')
 def serve(
         model: str, tokenizer: Optional[str], custom_tokenizer: Optional[str],
         host: str, port: int, log_level: str, backend: str, max_beam_width: int,
@@ -942,13 +958,18 @@ def serve(
         agent_types: Optional[str], video_pruning_rate: Optional[float],
         telemetry: bool, custom_module_dirs: list[Path],
         chat_template: Optional[str], middleware: tuple[str, ...], grpc: bool,
-        served_model_name: Optional[str], visual_gen_args: Optional[str]):
+        served_model_name: Optional[str], visual_gen_args: Optional[str],
+        kv_events_config: Optional[str]):
     """Running an OpenAI API compatible server
 
     MODEL: model name | HF checkpoint path | TensorRT engine path
     """
     logger.set_level(log_level)
 
+    # Parse the optional KV-events tee config (vLLM-shaped JSON) once, at the
+    # CLI boundary, so a malformed value fails fast before the engine loads.
+    kv_events_cfg = KVEventsConfig.from_cli(kv_events_config)
+
     if moe_cluster_parallel_size is not None:
         logger.warning(
             "--moe_cluster_parallel_size / --cluster_size is deprecated and "
@@ -1112,7 +1133,8 @@ def _serve_llm():
                           server_role,
                           disagg_cluster_config,
                           multimodal_server_config,
-                          served_model_name=served_model_name)
+                          served_model_name=served_model_name,
+                          kv_events_config=kv_events_cfg)
 
     def _serve_visual_gen():
         parsed_visual_gen_args = (VisualGenArgs.from_yaml(visual_gen_args)

tensorrt_llm/serve/kv_events_config.py

@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Serving-layer config for the read-only KV-cache event ZMQ tee.
+
+Field names and defaults mirror vLLM's ``vllm.config.KVEventsConfig`` so the
+operator-facing knobs line up across engines (and the wire format already
+matches). A single ``--kv_events_config`` JSON object replaces the older pair of
+``--kv_events_zmq_endpoint`` / ``--kv_events_replay_endpoint`` flags.
+
+This module is dependency-light (stdlib only) on purpose: it is imported at CLI
+parse time, before the heavier ``KvZmqPublisher`` (which pulls in ``zmq``).
+"""
+
+import dataclasses
+import json
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+
+@dataclass
+class KVEventsConfig:
+    """Configuration for KV-cache event publishing over ZMQ.
+
+    Mirrors vLLM's ``KVEventsConfig``:
+
+    - ``enable_kv_cache_events``: master switch; the tee is off unless this is
+      true (matches vLLM, where presence of an endpoint alone does nothing).
+    - ``publisher``: ``"zmq"`` to publish, ``"null"`` to disable. Set
+      automatically from ``enable_kv_cache_events`` when left at the default.
+    - ``endpoint``: PUB address. ``tcp://*:5557`` to bind (broadcast).
+    - ``replay_endpoint``: optional ROUTER address for gap recovery. ``None``
+      disables replay (vLLM default). Set it explicitly, e.g. ``tcp://*:5558``.
+    - ``buffer_steps``: number of recent batches kept for replay.
+    - ``hwm``: ZeroMQ high-water-mark for the PUB socket.
+    - ``max_queue_size``: max events buffered in memory before drops.
+    - ``topic``: ZMQ topic prefix on every published frame.
+    """
+
+    enable_kv_cache_events: bool = False
+    publisher: Optional[Literal["null", "zmq"]] = None
+    endpoint: str = "tcp://*:5557"
+    replay_endpoint: Optional[str] = None
+    buffer_steps: int = 10000
+    hwm: int = 100000
+    max_queue_size: int = 100000
+    topic: str = ""
+
+    def __post_init__(self):
+        # Mirror vLLM exactly: an unset (None) publisher follows the master
+        # switch; an explicitly-set publisher (incl. "null") is left untouched.
+        if self.publisher is None:
+            self.publisher = "zmq" if self.enable_kv_cache_events else "null"
+
+    @classmethod
+    def from_cli(cls, value: Optional[str]) -> Optional["KVEventsConfig"]:
+        """Parse the ``--kv_events_config`` JSON string into a config.
+
+        Returns ``None`` when no value is given (tee disabled). Raises
+        ``ValueError`` on malformed JSON or unknown keys, so a typo fails fast
+        at startup instead of silently disabling the tee.
+        """
+        if not value:
+            return None
+        try:
+            data = json.loads(value)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"--kv_events_config is not valid JSON: {e}") from e
+        if not isinstance(data, dict):
+            raise ValueError("--kv_events_config must be a JSON object")
+        known = {f.name for f in dataclasses.fields(cls)}
+        unknown = set(data) - known
+        if unknown:
+            raise ValueError(
+                f"--kv_events_config has unknown keys {sorted(unknown)}; "
+                f"valid keys are {sorted(known)}")
+        return cls(**data)