[None][feat] Native V2 KV cache event publishing - #17023
Conversation
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
- 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 <tanmay2592@gmail.com>
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 <tanmay2592@gmail.com>
- 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 <tanmay2592@gmail.com>
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 <tanmay2592@gmail.com>
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 <tanmay2592@gmail.com>
| raise ValueError(f"Unsupported KV event publisher: {config.publisher!r}") | ||
|
|
||
|
|
||
| def _vllm_wire_hash_from_radix_key(block_key: bytes) -> int: |
There was a problem hiding this comment.
Why not use truncate_sha256_hash_to_int64 (defined in tensorrt_llm/runtime/kv_cache_hash.py)? Both return a 64 bit hash from sha-256 key, but they use different parts of the 256 bit key.
There was a problem hiding this comment.
Good call — fixed in 8f3474b. _vllm_wire_hash_from_radix_key now reuses truncate_sha256_hash_to_int64 (first 8 bytes) and applies only the signed two's-complement reinterpretation on top for the vLLM wire format. So the native path derives the same 64-bit value from a block's SHA-256 key as the shared util (modulo the wire-format sign), instead of a second, divergent truncation.
| except ValueError: | ||
| self.dropped_events += 1 | ||
| self._pending_entries -= 1 | ||
| logger.exception("Dropping native KV store event with unsupported token data") |
There was a problem hiding this comment.
logger is imported from tensorrt_llm.logger. It does not have an exception method. You should probably use logger.error instead. The current code will throw an AttributeError.
There was a problem hiding this comment.
Confirmed and fixed in 8f3474b. tensorrt_llm.logger.Logger exposes only critical/error/warning/info/debug (no exception, and no __getattr__ delegating to the underlying stdlib logger), so this would have raised AttributeError. Replaced with logger.error(...) + traceback.format_exc() here and at the three other sites (lines 245, 267, 580). Thanks for catching it.
|
/bot run |
- 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 <tanmay2592@gmail.com>
WalkthroughAdds configurable native KV-cache event publishing with vLLM-compatible MessagePack events and ZeroMQ transport. Integrates event generation with ChangesNative KV-cache event publishing
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant KVCacheManagerV2
participant NativeKVCacheEventManager
participant ZmqEventPublisher
PyExecutor->>KVCacheManagerV2: enable native KV events
KVCacheManagerV2->>NativeKVCacheEventManager: initialize with KVEventsConfig
KVCacheManagerV2->>NativeKVCacheEventManager: provide cache lifecycle hooks
NativeKVCacheEventManager->>ZmqEventPublisher: publish batched events
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (4)
tensorrt_llm/llmapi/llm_args.py (1)
3662-3664: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the
model_post_initcontext parameter.The coding guidelines require an annotation on every function parameter. Pydantic v2 declares the hook as
model_post_init(self, context: Any, /) -> None, so rename and annotate the parameter.♻️ Proposed refactor
- def model_post_init(self, __context) -> None: + def model_post_init(self, context: Any) -> None: if self.publisher is None: self.publisher = "zmq" if self.enable_kv_cache_events else "null"As per coding guidelines: "Annotate every function, use
Nonefor procedures" and "avoid unnecessary double underscores".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/llmapi/llm_args.py` around lines 3662 - 3664, Update model_post_init so its context parameter is named context and annotated with Any, while preserving the existing publisher initialization behavior.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/_util.py (1)
1145-1147: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument why this argument bypasses the resolved
kv_cache_config.Every other argument in this call uses the
kv_cache_configlocal resolved at lines 1111-1112, which honorskv_cache_config_override. This expression instead readsself._llm_args.kv_cache_config.kv_events_config. The values agree today, because each override is produced bymodel_copy()and shares the nestedKVEventsConfigobject. A short comment prevents a future maintainer from adding a per-manager override ofkv_events_configand finding it ignored.♻️ Proposed refactor
+ # Native events are a single top-level setting, deliberately not + # taken from kv_cache_config_override: only one manager per rank + # may bind the endpoint. Estimation managers are transient and + # draft managers have no prefix reuse to report, so both get None. kv_events_config=None if estimating_kv_cache or model_engine.is_draft_model else self._llm_args.kv_cache_config.kv_events_config,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 1145 - 1147, Add a concise comment next to the kv_events_config expression explaining that it intentionally reads self._llm_args.kv_cache_config.kv_events_config rather than the resolved kv_cache_config, because overrides share the nested KVEventsConfig and this argument must retain the existing behavior.tensorrt_llm/_torch/pyexecutor/kv_cache_events.py (2)
510-522: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCount removal keys that are skipped as non-bytes.
Line 517 skips any entry that is not
byteswith no counter and no log. Every other suppression path in this class increments a counter, such aspartial_blocks_suppressedornon_target_life_cycles_ignored. The current V2 call sites pass byte keys, so this branch is unreachable today. A missed removal is the one failure mode that makes a consumer treat a block as resident forever, so make a future contract change visible instead of silent.♻️ Proposed refactor
self.non_target_life_cycles_ignored = 0 + self.unsupported_removal_keys = 0 self.dropped_events = 0for block_key in block_hashes: if not isinstance(block_key, bytes): + self.unsupported_removal_keys += 1 continue state = self._stored_blocks.pop(block_key, None)Add the counter to the
shutdownsummary alongside the existing counters.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py` around lines 510 - 522, Update add_removed_event to count entries skipped because block_key is not bytes, using a dedicated counter consistent with the class’s existing suppression counters. Increment it before continuing, and include the counter in the shutdown summary alongside partial_blocks_suppressed and non_target_life_cycles_ignored.
246-246: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider narrowing the caught exception types.
The coding guidelines require the narrowest exception possible. The publisher thread must survive transport failures, so a broad catch is defensible here, but naming the expected types documents the contract and lets a genuine programming error surface. The concrete failures are
zmq.ZMQErrorfromsend_multipartandrecv_multipart, andmsgspec.EncodeErrorfromencoder.encode.If you keep the broad catch, add a short comment stating that the thread must never terminate. Ruff
BLE001is reported by static analysis, but the repository's enabled Ruff rule set does not includeBLE, so this is not a lint failure.As per coding guidelines: "Catch the narrowest exception possible" and "Catch specific exceptions instead of using broad or bare
except:handlers."Also applies to: 270-270
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py` at line 246, Update the exception handlers in the publisher thread around send_multipart, recv_multipart, and encoder.encode to catch the specific expected zmq.ZMQError and msgspec.EncodeError types instead of Exception, while preserving thread survival on transport or encoding failures. Apply the same narrowing to both affected handlers.Sources: Coding guidelines, Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py`:
- Around line 159-165: Update the initialization flow around _socket_setup to
close the already-created PUB socket whenever setup raises before the object is
fully constructed, then re-raise the original exception. Prefer moving endpoint
validation before socket creation within _socket_setup to avoid allocating
sockets for invalid or unsupported endpoints, while preserving existing bind
behavior for valid endpoints.
- Around line 305-323: Update offset_endpoint_port to detect transports using
scheme prefixes consistent with _socket_setup, so hostnames containing “ipc” or
“inproc” remain valid TCP endpoints. For TCP endpoints, validate that a port is
present and numeric before converting and offsetting it, while preserving the
existing range check and rank-zero behavior.
- Around line 343-351: Update the expected hash calculation in
test_native_kv_events.py to use the first 8 bytes of the block hash, matching
truncate_sha256_hash_to_int64() and _vllm_wire_hash_from_radix_key(). Replace
the current last-8-byte slicing while preserving the existing signed 64-bit
conversion expectations.
- Around line 491-498: Update the token conversion flow around _token_ids to
recognize blocks containing the bytes digest produced by
gen_multimodal_cache_key_tokens before native event conversion. Exclude those
multimodal blocks without raising ValueError or incrementing dropped_events
through the traceback path; otherwise, define and consistently emit a valid
vLLM-compatible token_ids representation for them.
- Around line 259-278: Update the event publishing flow around the sequence
allocation, enqueue, and exception handling so every dropped batch remains
observable to consumers. Ensure queue-full drops and encode/send failures either
reserve a sequence number and emit an explicit loss marker on the wire, or
otherwise add a corresponding marker to the replay buffer before advancing to
the next sequence; preserve ordering and ensure END_SEQ cannot make an
incomplete stream appear complete.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 3639-3642: Update the replay_endpoint Field declaration to enforce
a minimum length of 1, matching the validation applied to endpoint. Preserve
None as the allowed unset value while rejecting empty strings during
configuration validation before socket setup.
In `@tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py`:
- Around line 146-181: Update
test_native_removals_are_never_dropped_by_the_entry_cap to flush the manager via
flush_iteration_events() after queuing the removals, using a recording publisher
or ZeroMQ subscriber to capture emitted batches. Assert the flushed MessagePack
payload contains both removed block hashes, rather than only inspecting
manager._pending_events.
- Around line 33-35: Add a None return annotation to both test functions:
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py lines 33-35,
test_native_fast_path_publishes_only_full_max_window_blocks, and lines 146-147,
test_native_removals_are_never_dropped_by_the_entry_cap.
- Around line 33-143: Wrap the manager and subscriber lifecycle in
test_native_fast_path_publishes_only_full_max_window_blocks with try/finally so
manager.shutdown(), subscriber.close(), and endpoint cleanup run even when
assertions fail. Also wrap the manager lifecycle in the test spanning
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py lines 146-183
with try/finally, ensuring its shutdown executes on every failure path.
- Around line 27-30: Replace the _unused_tcp_port approach and fixed time.sleep
synchronization in the NativeKVCacheEventManager ZeroMQ setup with a retry
fixture. Have the fixture retry the publish-and-receive setup when binding fails
due to an address-in-use zmq.ZMQError, catching only that expected error and
allowing other failures to propagate. Ensure the test proceeds only after the
subscriber successfully receives the published event.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 1145-1147: Add a concise comment next to the kv_events_config
expression explaining that it intentionally reads
self._llm_args.kv_cache_config.kv_events_config rather than the resolved
kv_cache_config, because overrides share the nested KVEventsConfig and this
argument must retain the existing behavior.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py`:
- Around line 510-522: Update add_removed_event to count entries skipped because
block_key is not bytes, using a dedicated counter consistent with the class’s
existing suppression counters. Increment it before continuing, and include the
counter in the shutdown summary alongside partial_blocks_suppressed and
non_target_life_cycles_ignored.
- Line 246: Update the exception handlers in the publisher thread around
send_multipart, recv_multipart, and encoder.encode to catch the specific
expected zmq.ZMQError and msgspec.EncodeError types instead of Exception, while
preserving thread survival on transport or encoding failures. Apply the same
narrowing to both affected handlers.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 3662-3664: Update model_post_init so its context parameter is
named context and annotated with Any, while preserving the existing publisher
initialization behavior.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 20bf962b-8636-4228-8564-cb26da63e134
📒 Files selected for processing (9)
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_events.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/llmapi/llm_utils.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py
| self._socket_setup() | ||
| self._thread = threading.Thread( | ||
| target=self._publisher_thread, | ||
| daemon=True, | ||
| name=f"trtllm-kv-events-rank-{self._rank}", | ||
| ) | ||
| self._thread.start() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Close the PUB socket when _socket_setup fails.
_socket_setup creates self._pub at line 222 and can then raise at line 225, 227, or 232 (empty endpoint, unsupported scheme, or a bind failure such as EADDRINUSE). The exception propagates out of __init__, so the object is never returned and shutdown() is unreachable. The already-created socket stays open on the shared zmq.Context.instance(), which leaks a file descriptor and keeps the context alive. KVCacheManagerV2 cannot recover this case either, because it only calls event_manager.shutdown() for failures after NativeKVCacheEventManager.__init__ assigned self._publisher.
Guard the setup call so a failure closes the socket before the exception propagates.
🛠️ Proposed fix
- self._socket_setup()
+ try:
+ self._socket_setup()
+ except Exception:
+ if self._pub is not None:
+ self._pub.close(linger=0)
+ self._pub = None
+ if self._replay is not None:
+ self._replay.close(linger=0)
+ self._replay = None
+ raise
self._thread = threading.Thread(Move the endpoint validation ahead of socket creation so the common misconfiguration never allocates a socket:
def _socket_setup(self) -> None:
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}")
self._pub = self._ctx.socket(zmq.PUB)
self._pub.set_hwm(self._hwm)
self._pub.bind(self._endpoint)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self._socket_setup() | |
| self._thread = threading.Thread( | |
| target=self._publisher_thread, | |
| daemon=True, | |
| name=f"trtllm-kv-events-rank-{self._rank}", | |
| ) | |
| self._thread.start() | |
| try: | |
| self._socket_setup() | |
| except Exception: | |
| if self._pub is not None: | |
| self._pub.close(linger=0) | |
| self._pub = None | |
| if self._replay is not None: | |
| self._replay.close(linger=0) | |
| self._replay = None | |
| raise | |
| self._thread = threading.Thread( | |
| target=self._publisher_thread, | |
| daemon=True, | |
| name=f"trtllm-kv-events-rank-{self._rank}", | |
| ) | |
| self._thread.start() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py` around lines 159 - 165,
Update the initialization flow around _socket_setup to close the already-created
PUB socket whenever setup raises before the object is fully constructed, then
re-raise the original exception. Prefer moving endpoint validation before socket
creation within _socket_setup to avoid allocating sockets for invalid or
unsupported endpoints, while preserving existing bind behavior for valid
endpoints.
| 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() |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Determine whether any consumer relies on KV event sequence continuity to detect loss.
set -euo pipefail
# Find readers of the three-frame KV event protocol and the replay path.
rg -n -C 6 'END_SEQ|start_seq|_seq_gen|recv_multipart' --glob '*.py'
# Find documentation or examples describing the native KV event wire contract.
rg -n -C 4 'kv_events_config|KVEventBatch|BlockRemoved' --glob '*.md' --glob 'examples/**'
# Check whether the drop counters are exported anywhere a consumer could observe.
rg -n -C 4 'dropped_batches|dropped_events|_queue_full_drops' --glob '*.py'Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- file outline ---'
ast-grep outline tensorrt_llm/_torch/pyexecutor/kv_cache_events.py
printf '%s\n' '--- relevant implementation ---'
sed -n '1,330p' tensorrt_llm/_torch/pyexecutor/kv_cache_events.py
printf '%s\n' '--- all repository references ---'
rg -n -C 5 'KVCacheEvent|KVEvent|END_SEQ|start_seq|dropped_batches|_queue_full_drops|recv_multipart|kv_events' \
tensorrt_llm tests examples docs README.md 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,330p' tensorrt_llm/_torch/pyexecutor/kv_cache_events.py
rg -n -C 5 'KVCacheEvent|KVEvent|END_SEQ|start_seq|dropped_batches|_queue_full_drops|recv_multipart|kv_events' . 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- native event tests ---'
sed -n '1,190p' tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py
printf '%s\n' '--- focused protocol references ---'
rg -n -C 8 'END_SEQ|start_seq|replay|sequence|desync|drop|dropped|buffer' \
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py \
tensorrt_llm/_torch/pyexecutor/kv_cache_events.py \
tensorrt_llm/llmapi/llm_args.py \
docs examples README.md 2>/dev/null || true
printf '%s\n' '--- exact vLLM adaptation references ---'
rg -n -C 4 'vllm/distributed/kv_events.py|vLLM-compatible|KV events' . \
--glob '!*.json' --glob '!*.lock' 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50378
🌐 Web query:
vLLM distributed kv_events.py END_SEQ sequence replay dropped events
💡 Result:
In vLLM's distributed architecture, kv_events.py implements a reliable event publisher, ZmqEventPublisher, designed to handle KV cache events across different data-parallel ranks [1][2]. The END_SEQ constant, defined as (-1).to_bytes(8, "big", signed=True), serves as a terminal marker for sequence replay [1][3]. When a subscriber requests a replay of missed batches by sending a starting sequence number to the ROUTER endpoint, the ZmqEventPublisher iterates through its in-memory buffer (deque) [1][3]. For every event batch starting from the requested sequence number, it sends the batch to the subscriber, followed by the END_SEQ marker [1][3]. This marker explicitly signals the end of the requested replay sequence to the subscriber [1][4]. Events are typically considered "dropped" in this context if they fall outside the publisher's bounded buffer or if a subscriber fails to request a replay before the relevant batches are evicted from the buffer [1][5]. The buffer size is configurable via the buffer_steps parameter, which defaults to 10,000 entries [1]. If a subscriber cannot keep up or reconnects after the buffer has advanced past its last received sequence number, those events are effectively lost unless persisted elsewhere [1][5]. The system is designed for best-effort delivery within these memory constraints, using sequence numbers to maintain ordering and facilitate reconciliation [1][2].
Citations:
- 1: https://github.com/vllm-project/vllm/blob/main/vllm/distributed/kv_events.py
- 2: https://docs.vllm.ai/en/v0.22.0/api/vllm/distributed/kv_events/
- 3: https://docs.vllm.ai/en/v0.10.1/api/vllm/distributed/kv_events.html
- 4: https://vllm.website.cncfstack.com/api/vllm/distributed/kv_events.html
- 5: https://devfra.me/guide/streaming
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,190p' tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py
rg -n -C 8 'END_SEQ|start_seq|replay|sequence|desync|drop|dropped|buffer' \
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py \
tensorrt_llm/_torch/pyexecutor/kv_cache_events.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 44031
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("tensorrt_llm/_torch/pyexecutor/kv_cache_events.py")
tree = ast.parse(path.read_text())
class_info = {}
for node in tree.body:
if isinstance(node, ast.ClassDef) and node.name == "ZmqEventPublisher":
class_info[node.name] = {
child.name: child
for child in node.body
if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef))
}
methods = class_info["ZmqEventPublisher"]
publish = methods["publish"]
publisher = methods["_publisher_thread"]
replay = methods["_service_replay"]
def calls(node, name):
return [
n for n in ast.walk(node)
if isinstance(n, ast.Call)
and (
(isinstance(n.func, ast.Name) and n.func.id == name)
or (isinstance(n.func, ast.Attribute) and n.func.attr == name)
)
]
def source_order(node, predicate):
return [
(n.lineno, ast.unparse(n))
for n in ast.walk(node)
if predicate(n)
]
seq_assignments = source_order(
publisher,
lambda n: isinstance(n, ast.Call)
and isinstance(n.func, ast.Name)
and n.func.id == "next",
)
buffer_appends = source_order(
publisher,
lambda n: isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "append"
and isinstance(n.func.value, ast.Attribute)
and n.func.value.attr == "_buffer",
)
queue_puts = source_order(
publish,
lambda n: isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "put_nowait",
)
replay_sends = source_order(
replay,
lambda n: isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "send_multipart",
)
print("sequence allocation:", seq_assignments)
print("buffer append:", buffer_appends)
print("publish queue insertion:", queue_puts)
print("replay sends:", replay_sends)
# Model the two drop paths using the same ordering as the implementation.
published = []
next_seq = 0
queue_full_drop = True
if not queue_full_drop:
seq = next_seq
next_seq += 1
published.append(seq)
send_failed = True
seq = next_seq
next_seq += 1
if not send_failed:
published.append(seq)
print("modeled published sequence numbers after queue-full and send-failure:",
published)
print("modeled next sequence number:", next_seq)
assert published == []
assert next_seq == 1
PY
printf '%s\n' '--- upstream protocol implementation ---'
curl -L --fail --silent --show-error \
https://raw.githubusercontent.com/vllm-project/vllm/main/vllm/distributed/kv_events.py \
| rg -n -C 8 'seq|END_SEQ|replay|buffer|start_seq' | head -n 160Repository: NVIDIA/TensorRT-LLM
Length of output: 6732
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
tree = ast.parse(Path("tensorrt_llm/_torch/pyexecutor/kv_cache_events.py").read_text())
cls = next(n for n in tree.body if isinstance(n, ast.ClassDef)
and n.name == "ZmqEventPublisher")
methods = {
n.name: n for n in cls.body
if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))
}
def find_calls(node, attr=None, name=None):
result = []
for n in ast.walk(node):
if not isinstance(n, ast.Call):
continue
if attr and isinstance(n.func, ast.Attribute) and n.func.attr == attr:
result.append(n.lineno)
if name and isinstance(n.func, ast.Name) and n.func.id == name:
result.append(n.lineno)
return result
print("next() lines:", find_calls(methods["_publisher_thread"], name="next"))
print("buffer.append() lines:", find_calls(methods["_publisher_thread"], attr="append"))
print("queue.put_nowait() lines:", find_calls(methods["publish"], attr="put_nowait"))
print("replay.send_multipart() lines:", find_calls(methods["_service_replay"],
attr="send_multipart"))
# Queue-full drop: no sequence is allocated.
# Send failure: sequence is allocated, but buffer append is skipped.
published = []
next_seq = 0
queue_full = True
if not queue_full:
published.append(next_seq)
next_seq += 1
send_failed = True
failed_seq = next_seq
next_seq += 1
if not send_failed:
published.append(failed_seq)
assert published == []
assert next_seq == 1
print("modeled published sequence numbers:", published)
print("modeled next sequence number:", next_seq)
PY
curl -L --fail --silent --show-error \
https://raw.githubusercontent.com/vllm-project/vllm/main/vllm/distributed/kv_events.py \
| rg -n -C 6 'END_SEQ|start_seq|buffer|seq' | head -n 120Repository: NVIDIA/TensorRT-LLM
Length of output: 4571
Make dropped batches observable on the wire.
Queue-full drops allocate no sequence number. Encode or send failures consume a sequence number but do not enter _buffer. Replay then sends the next available batch followed by END_SEQ, so consumers can accept an incomplete stream as complete and retain stale KV-cache state. Reserve sequence numbers before enqueue, or emit an explicit loss marker.
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 270-270: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py` around lines 259 - 278,
Update the event publishing flow around the sequence allocation, enqueue, and
exception handling so every dropped batch remains observable to consumers.
Ensure queue-full drops and encode/send failures either reserve a sequence
number and emit an explicit loss marker on the wire, or otherwise add a
corresponding marker to the replay buffer before advancing to the next sequence;
preserve ordering and ensure END_SEQ cannot make an incomplete stream appear
complete.
| @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'") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Use prefix matching for the scheme and validate the port.
offset_endpoint_port detects the transport with substring tests, but _socket_setup validates it with startswith. The two disagree for a TCP endpoint whose host contains ipc or inproc. For example, tcp://ipc-host:5557 on rank 1 matches the "ipc" in endpoint branch and becomes tcp://ipc-host:5557_dp1. That value still passes the tcp:// scheme check in _socket_setup, so the failure surfaces as an opaque ZeroMQ bind error.
The port parse is also unguarded. ":" in endpoint is satisfied by the scheme colon alone, so a TCP endpoint without a port, such as tcp://host, reaches int("//host") and raises ValueError. Rank 0 returns early and never hits this, so the failure appears only on ranks above 0.
🛠️ Proposed fix
`@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:
+ if endpoint.startswith(("inproc://", "ipc://")):
return f"{endpoint}_dp{data_parallel_rank}"
- if "tcp" in endpoint and ":" in endpoint:
+ if endpoint.startswith("tcp://"):
last_colon_idx = endpoint.rfind(":")
+ if last_colon_idx <= len("tcp:"):
+ raise ValueError(
+ f"KV event TCP endpoint must specify a base port: {endpoint!r}"
+ )
base_addr = endpoint[:last_colon_idx]
- base_port = int(endpoint[last_colon_idx + 1 :])
+ port_text = endpoint[last_colon_idx + 1 :]
+ if not port_text.isdigit():
+ raise ValueError(
+ f"KV event TCP endpoint has a non-numeric base port: {endpoint!r}"
+ )
+ base_port = int(port_text)
new_port = base_port + data_parallel_rank
if new_port > 65_535:
raise ValueError(
f"KV event endpoint port exceeds 65535 for rank {data_parallel_rank}"
)
return f"{base_addr}:{new_port}"
raise ValueError("Invalid endpoint: must contain 'inproc', 'ipc', or 'tcp'")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @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'") | |
| `@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 endpoint.startswith(("inproc://", "ipc://")): | |
| return f"{endpoint}_dp{data_parallel_rank}" | |
| if endpoint.startswith("tcp://"): | |
| last_colon_idx = endpoint.rfind(":") | |
| if last_colon_idx <= len("tcp:"): | |
| raise ValueError( | |
| f"KV event TCP endpoint must specify a base port: {endpoint!r}" | |
| ) | |
| base_addr = endpoint[:last_colon_idx] | |
| port_text = endpoint[last_colon_idx + 1 :] | |
| if not port_text.isdigit(): | |
| raise ValueError( | |
| f"KV event TCP endpoint has a non-numeric base port: {endpoint!r}" | |
| ) | |
| base_port = int(port_text) | |
| new_port = base_port + data_parallel_rank | |
| if new_port > 65_535: | |
| raise ValueError( | |
| f"KV event endpoint port exceeds 65535 for rank {data_parallel_rank}" | |
| ) | |
| return f"{base_addr}:{new_port}" | |
| raise ValueError("Invalid endpoint: must contain 'inproc', 'ipc', or 'tcp'") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py` around lines 305 - 323,
Update offset_endpoint_port to detect transports using scheme prefixes
consistent with _socket_setup, so hostnames containing “ipc” or “inproc” remain
valid TCP endpoints. For TCP endpoints, validate that a port is present and
numeric before converting and offsetting it, while preserving the existing range
check and rank-zero behavior.
| 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 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect the shared SHA-256 truncation helper and its callers.
set -euo pipefail
fd -t f 'kv_cache_hash.py'
ast-grep run --pattern 'def truncate_sha256_hash_to_int64($$$):
$$$' --lang python
# Show every caller so the byte range is consistent across the event machinery.
rg -n -C 4 'truncate_sha256_hash_to_int64' --glob '*.py'Repository: NVIDIA/TensorRT-LLM
Length of output: 195
🏁 Script executed:
set -euo pipefail
file="$(fd -t f 'kv_cache_hash.py' | head -n 1)"
cat -n "$file"
printf '\nCall sites:\n'
rg -n -C 5 'truncate_sha256_hash_to_int64' --glob '*.py'
printf '\nRelevant test:\n'
fd -t f -i 'test_native_kv_events.py' -x rg -n -C 8 'first_hash|wire_hash|int.from_bytes' {}Repository: NVIDIA/TensorRT-LLM
Length of output: 4417
🏁 Script executed:
set -u
printf '%s\n' 'Tracked test candidates:'
git ls-files | rg 'test_native_kv_events\.py$|kv_cache_events\.py$'
printf '%s\n' 'Relevant test assertions:'
rg -n -C 10 'first_hash|int\.from_bytes|wire_hash|vllm_wire|BlockStored' tests tensorrt_llm 2>/dev/null | head -n 240Repository: NVIDIA/TensorRT-LLM
Length of output: 21343
🏁 Script executed:
python3 - <<'PY'
def truncate_sha256_hash_to_int64(block_hash: bytes) -> int:
return int.from_bytes(block_hash[:8], "big", signed=False)
def signed_wire_hash(value: int) -> int:
return value - 2**64 if value >= 2**63 else value
first_hash = b"\x11" * 24 + b"\x80\x00\x00\x00\x00\x00\x00\x01"
second_hash = b"\x33" * 24 + b"\x00\x00\x00\x00\x00\x00\x00\x02"
for name, value in (("first", first_hash), ("second", second_hash)):
helper_value = signed_wire_hash(truncate_sha256_hash_to_int64(value))
test_value = signed_wire_hash(int.from_bytes(value[-8:], "big"))
print(f"{name}: helper={helper_value}, test={test_value}, match={helper_value == test_value}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 287
Use the first 8 bytes in the test expectation. truncate_sha256_hash_to_int64() uses block_hash[:8], but test_native_kv_events.py derives expected hashes from the last 8 bytes. Update the test to match the shared helper contract.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py` around lines 343 - 351,
Update the expected hash calculation in test_native_kv_events.py to use the
first 8 bytes of the block hash, matching truncate_sha256_hash_to_int64() and
_vllm_wire_hash_from_radix_key(). Replace the current last-8-byte slicing while
preserving the existing signed 64-bit conversion expectations.
| @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 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Determine whether TokenIdExt subclasses int and can reach the native event store path.
set -euo pipefail
# Resolve the TokenIdExt definition.
rg -n -C 6 'TokenIdExt' --glob '*.py' -g '!**/tests/**'
# Confirm augmented tokens are committed to the radix blocks the event hooks observe.
rg -n -C 8 '_augment_tokens_for_block_reuse' --glob 'tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py'
# Find where a radix block's `tokens` attribute is populated.
rg -n -C 6 '\.tokens\b' --glob 'tensorrt_llm/runtime/kv_cache_manager_v2/**/*.py'Repository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- TokenIdExt definitions and references ---'
rg -n -C 8 'TokenIdExt' --glob '*.py' . || true
printf '%s\n' '--- augmentation references ---'
rg -n -C 10 '_augment_tokens_for_block_reuse' . || true
printf '%s\n' '--- event manager and radix block token assignments ---'
rg -n -C 8 'add_stored_block_event_from_block|_token_ids|tokens\s*=' tensorrt_llm --glob '*.py' || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- exact TokenIdExt matches ---'
rg -n -F 'TokenIdExt' --glob '*.py' . | head -200
printf '%s\n' '--- exact augmentation matches ---'
rg -n -F '_augment_tokens_for_block_reuse' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py -A 20 -B 20 || true
printf '%s\n' '--- relevant event-path matches ---'
rg -n -F 'add_stored_block_event_from_block' tensorrt_llm/_torch/pyexecutor/kv_cache_events.py -A 25 -B 15
rg -n -F '_token_ids' tensorrt_llm/_torch/pyexecutor/kv_cache_events.py -A 12 -B 8Repository: NVIDIA/TensorRT-LLM
Length of output: 19599
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- TokenIdExt aliases and multimodal token generation ---'
sed -n '45,75p' tensorrt_llm/runtime/kv_cache_manager_v2/_common.py
sed -n '2545,2605p' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
printf '%s\n' '--- radix block storage and commit path ---'
sed -n '120,150p' tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
sed -n '320,355p' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
sed -n '950,1015p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
sed -n '1035,1060p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
printf '%s\n' '--- native manager construction and event consumer contract ---'
rg -n -C 10 'NativeKVCacheEventManager|vLLM-compatible|token_ids' tensorrt_llm/_torch/pyexecutor/kv_cache_events.py tensorrt_llm --glob '*.py' | head -300Repository: NVIDIA/TensorRT-LLM
Length of output: 36557
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- multimodal augmentation helpers ---'
sed -n '485,590p' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
printf '%s\n' '--- standalone behavior probe ---'
python3 - <<'PY'
from typing import NewType
TokenId = NewType("TokenId", int)
TokenIdExt = TokenId | bytes
samples = [1, TokenId(1), b"digest"]
for value in samples:
accepted = type(value) is int
print(type(value).__name__, repr(value), "accepted_by_token_ids=", accepted)
assert type(TokenId(1)) is int
assert type(b"digest") is not int
assert not (type(b"digest") is int)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 4994
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- remaining contiguous augmentation helper ---'
sed -n '590,650p' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
printf '%s\n' '--- multimodal token generator definition and callers ---'
rg -n -C 12 'def gen_multimodal_cache_key_tokens|gen_multimodal_cache_key_tokens\(' \
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py \
tensorrt_llm/runtime/kv_cache_manager_v2 --glob '*.py'
printf '%s\n' '--- Python runtime and direct behavior probe ---'
python3 - <<'PY'
import sys
print(sys.version)
samples = [1, b"digest"]
for value in samples:
print(type(value).__name__, repr(value), "accepted_by_token_ids=", type(value) is int)
assert type(1) is int
assert type(b"digest") is not int
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 12796
Handle multimodal cache-key tokens before native event conversion.
gen_multimodal_cache_key_tokens inserts a bytes digest for the first multimodal token in each item. A full block containing that token fails _token_ids, increments dropped_events, and logs a traceback. Exclude these blocks without the traceback path, or define a valid vLLM wire representation for their token_ids.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py` around lines 491 - 498,
Update the token conversion flow around _token_ids to recognize blocks
containing the bytes digest produced by gen_multimodal_cache_key_tokens before
native event conversion. Exclude those multimodal blocks without raising
ValueError or incrementing dropped_events through the traceback path; otherwise,
define and consistently emit a valid vLLM-compatible token_ids representation
for them.
| replay_endpoint: Optional[str] = Field( | ||
| default=None, | ||
| description= | ||
| "Optional base ZeroMQ endpoint used to replay KV cache events.") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add min_length=1 to replay_endpoint.
endpoint rejects an empty string, but replay_endpoint does not. An empty string is not None, so ZmqEventPublisher.offset_endpoint_port returns it unchanged and _socket_setup reaches self._replay.bind(""), which raises a ZeroMQ error during engine startup instead of a clear configuration error.
🛠️ Proposed fix
replay_endpoint: Optional[str] = Field(
default=None,
+ min_length=1,
description=
"Optional base ZeroMQ endpoint used to replay KV cache events.")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| replay_endpoint: Optional[str] = Field( | |
| default=None, | |
| description= | |
| "Optional base ZeroMQ endpoint used to replay KV cache events.") | |
| replay_endpoint: Optional[str] = Field( | |
| default=None, | |
| min_length=1, | |
| description= | |
| "Optional base ZeroMQ endpoint used to replay KV cache events.") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/llmapi/llm_args.py` around lines 3639 - 3642, Update the
replay_endpoint Field declaration to enforce a minimum length of 1, matching the
validation applied to endpoint. Preserve None as the allowed unset value while
rejecting empty strings during configuration validation before socket setup.
| def _unused_tcp_port() -> int: | ||
| with socket.socket() as sock: | ||
| sock.bind(("127.0.0.1", 0)) | ||
| return int(sock.getsockname()[1]) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Make ZeroMQ setup deterministic.
_unused_tcp_port() releases the port before NativeKVCacheEventManager binds it. Another process can claim the port during that interval. time.sleep(0.2) also does not guarantee that the PUB socket has received the subscriber subscription.
Use a retry fixture that catches only zmq.ZMQError for address-in-use failures. Retry the publish-and-receive setup instead of depending on a released port and a fixed delay.
Also applies to: 39-57
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py` around
lines 27 - 30, Replace the _unused_tcp_port approach and fixed time.sleep
synchronization in the NativeKVCacheEventManager ZeroMQ setup with a retry
fixture. Have the fixture retry the publish-and-receive setup when binding fails
due to an address-in-use zmq.ZMQError, catching only that expected error and
allowing other failures to propagate. Ensure the test proceeds only after the
subscriber successfully receives the published event.
| def test_native_fast_path_publishes_only_full_max_window_blocks(): | ||
| """Protect radix hash reuse, filtering, wire format, and shutdown.""" | ||
| port = _unused_tcp_port() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add procedure return annotations.
The test procedures do not declare -> None.
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py#L33-L35: add-> Nonetotest_native_fast_path_publishes_only_full_max_window_blocks.tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py#L146-L147: add-> Nonetotest_native_removals_are_never_dropped_by_the_entry_cap.
As per coding guidelines, “Annotate every function.”
📍 Affects 1 file
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py#L33-L35(this comment)tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py#L146-L147
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py` around
lines 33 - 35, Add a None return annotation to both test functions:
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py lines 33-35,
test_native_fast_path_publishes_only_full_max_window_blocks, and lines 146-147,
test_native_removals_are_never_dropped_by_the_entry_cap.
Source: Coding guidelines
| 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) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Run cleanup when an assertion fails.
An assertion failure skips manager.shutdown() and socket cleanup. A background publisher can then retain a thread or endpoint and affect later tests.
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py#L33-L143: wrap the manager and subscriber lifecycle intry/finally.tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py#L146-L183: wrap the manager lifecycle intry/finally.
📍 Affects 1 file
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py#L33-L143(this comment)tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py#L146-L183
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py` around
lines 33 - 143, Wrap the manager and subscriber lifecycle in
test_native_fast_path_publishes_only_full_max_window_blocks with try/finally so
manager.shutdown(), subscriber.close(), and endpoint cleanup run even when
assertions fail. Also wrap the manager lifecycle in the test spanning
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py lines 146-183
with try/finally, ensuring its shutdown executes on every failure path.
| 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 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Verify removal delivery after flushing.
This test inspects manager._pending_events before flush_iteration_events(). It proves that BlockRemoved objects are queued, but it does not prove that the publisher emits them after the entry cap is reached.
Flush the events and assert the emitted MessagePack removal batch through a recording publisher or a ZeroMQ subscriber.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py` around
lines 146 - 181, Update test_native_removals_are_never_dropped_by_the_entry_cap
to flush the manager via flush_iteration_events() after queuing the removals,
using a recording publisher or ZeroMQ subscriber to capture emitted batches.
Assert the flushed MessagePack payload contains both removed block hashes,
rather than only inspecting manager._pending_events.
3 participants
Motivation
External KV-cache-aware routers (e.g. Dynamo) subscribe to a stream of block
stored / removed events to route requests to the engine that already holds a
prefix. The existing path builds Python
KVCacheEventobjects, buffers them,all-gathers them onto rank 0 under attention DP, and exposes them through a
per-iteration pull API (
LLM.get_kv_cache_events()).This PR adds an opt-in path where each rank publishes its own events directly
over ZeroMQ, reusing the V2 radix block hashes it already computed instead of
re-deriving events and gathering them.
What this changes
Adds an opt-in native KV-event path for KV cache manager V2 (PyTorch
backend). Off by default; the legacy path is unchanged.
NativeKVCacheEventManager(newtensorrt_llm/_torch/pyexecutor/kv_cache_events.py)implements the V2 event-sink hooks by duck typing. Per stored/removed block it
builds wire-format
msgspecstructs (BlockStored/BlockRemoved/AllBlocksCleared), reusing the low 64 bits of the radix block key as the wirehash (no re-hash) and coalescing consecutive blocks into one event.
ZmqEventPublishermsgpack-encodes each per-iteration batch and sends itfrom a background thread over a ZeroMQ
PUBsocket (3 frames:topic,seq,payload), with an optionalROUTERreplay socket. Each attention-DP rankbinds
base_port + rank.KVEventsConfig, nested askv_cache_config.kv_events_config(marked
prototype). Scope guards: V2 only (warns on a non-V2 manager);excluded for draft models and KV-cache-size estimation; raises under pipeline
or context parallelism.
[], soLLM.get_kv_cache_events()degrades cleanly instead of raising.
Before / After
Before — legacy build → buffer → gather → per-iteration pull (still the
default; unchanged by this PR):
After — native per-rank publish (opt-in via
kv_cache_config.kv_events_config):flowchart LR subgraph S["scheduler / KV-manager thread"] H["V2 radix tree<br/>store / remove hooks"] --> NM["NativeKVCacheEventManager<br/>build wire structs<br/>reuse radix hash · coalesce"] NM -->|"flush per iteration<br/>non-blocking enqueue"| Q["bounded queue"] end subgraph BG["background publisher thread"] Q --> ENC["msgpack encode"] ENC --> PUB["ZeroMQ PUB<br/>binds base_port + rank"] end PUB --> SUB["external subscriber<br/>(e.g. Dynamo)"]Each rank builds events on the scheduler thread (cheap, non-blocking enqueue) and
a background thread publishes them — no gather onto rank 0 and no per-iteration
pull path.
Context
Supersedes #16869 and #16876 by @alec-flowers (original authorship preserved).
Relates to RFC #17013.
Dev Engineer Review
KVCacheManagerV2.KVEventsConfigwith ZeroMQ, replay, queue, HWM, buffer, topic, and publisher settings.NativeKVCacheEventManagerwithmsgspecencoding, event coalescing, replay buffering, queue limits, counters, and clean shutdown.QA Engineer Review
test_native_fast_path_publishes_only_full_max_window_blocks.test_native_removals_are_never_dropped_by_the_entry_cap.tests/integration/test_lists/,test-db/, orqa/coverage entry was identified.