[None][feat] Streaming (push-based) KV cache event publishing (V2) - #17023
[None][feat] Streaming (push-based) KV cache event publishing (V2)#17023tanmayv25 wants to merge 16 commits into
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>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThis PR adds configurable native KV-cache event publishing. It defines msgspec event types, asynchronous ZeroMQ publishers, replay support, and streaming event batching. It integrates streaming events with ChangesNative KV-cache events
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant KVCacheManagerV2
participant StreamingKVCacheEventManager
participant ZmqEventPublisher
PyExecutor->>KVCacheManagerV2: enable streaming KV events
KVCacheManagerV2->>StreamingKVCacheEventManager: initialize per-rank manager
KVCacheManagerV2->>StreamingKVCacheEventManager: forward cache lifecycle hooks
StreamingKVCacheEventManager->>ZmqEventPublisher: publish stored or removed batch
ZmqEventPublisher-->>StreamingKVCacheEventManager: report queue or publish status
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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
| 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.
| 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.
| 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
| 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.
8f3474b switched the wire hash to truncate_sha256_hash_to_int64 (first 8 bytes of the radix key) but left the test asserting the old last-8-byte values, so the test failed deterministically in CI. Update the synthetic keys and expected hashes to the first-8-byte convention, keeping the signed-wraparound branch covered. Signed-off-by: tanmayv25 <tanmay2592@gmail.com>
…(review) Address code-review minors: - Replace the single-int _NativeStoredBlockState wrapper with a plain dict[bytes, int] mapping radix key -> wire hash; deletes the class and a redundant tuple slot. - Add tests for KVEventsConfig publisher default resolution (None -> zmq/null) and offset_endpoint_port (base_port+rank, ipc/inproc suffix, u16 overflow, bad scheme). Signed-off-by: tanmayv25 <tanmay2592@gmail.com>
The 'native' vs 'legacy' naming was misleading: 'native' is overloaded in TRT-LLM, and 'legacy' wrongly implied the buffered gather/poll path is deprecated when it is actually the fuller-fidelity default. Rename to describe the delivery mechanism: - NativeKVCacheEventManager -> StreamingKVCacheEventManager (+ native_kv_events_enabled -> streaming_kv_events_enabled). - 'native'/'legacy' -> 'streaming (push-based)'/'buffered (gather/poll)' in log messages, comments, KVEventsConfig docstrings, and the test file name. Public config identifiers (kv_events_config, enable_kv_cache_events) are unchanged; only descriptions were updated (not captured by the golden manifest). Signed-off-by: tanmayv25 <tanmay2592@gmail.com>
tanmayv25
changed the title
- ZmqEventPublisher.__init__: close the PUB/ROUTER sockets if _socket_setup raises, so a bind/scheme failure doesn't leak sockets on the shared context (shutdown() is unreachable when __init__ never returns). - offset_endpoint_port: match the scheme with startswith (consistent with _socket_setup) and reject a TCP endpoint with no port, instead of parsing the scheme colon into int() and raising an opaque error on ranks > 0. - Skip multimodal cache-key blocks (bytes token digests) via a dedicated _MultimodalBlockError so they no longer flood the log with malformed-data tracebacks; count them in multimodal_blocks_suppressed. - KVEventsConfig.replay_endpoint: add min_length=1 so an empty string is rejected up front rather than failing at bind() time. - Tests: annotate test procedures with -> None, run manager/subscriber cleanup under try/finally, assert removals reach the wire after flush, and cover the no-port endpoint rejection. Signed-off-by: tanmayv25 <tanmay2592@gmail.com>
|
Thanks @coderabbitai — addressed in Fixed
Already addressed
Deferred (with rationale)
|
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
echo '== Commit availability =='
git show --no-patch --oneline 69b318a8cc
echo '== Relevant file outlines =='
ast-grep outline tensorrt_llm/_torch/pyexecutor/kv_cache_events.py --items all --type class,function --match 'ZmqEventPublisher|offset_endpoint_port|NativeKVCacheEventManager|_MultimodalBlockError'
ast-grep outline tensorrt_llm/llmapi/llm_args.py --items all --type class --match 'KVEventsConfig'
ast-grep outline tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.py --items all --type function --match 'endpoint|removal|native|zmq'
echo '== Target implementation regions =='
rg -n -A20 -B12 'def __init__|def _socket_setup|def offset_endpoint_port|_MultimodalBlockError|multimodal_blocks_suppressed|replay_endpoint' \
tensorrt_llm/_torch/pyexecutor/kv_cache_events.py \
tensorrt_llm/llmapi/llm_args.py
echo '== Target test assertions =='
rg -n -A12 -B8 'no.port|no-port|flush_iteration_events|try:|finally:|-> None|remov' \
tests/unittest/kv_cache_manager_v2_tests/test_native_kv_events.pyLength of output: 35434 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo '== Test files changed by the fix commit =='
git diff-tree --no-commit-id --name-status -r 69b318a8cc -- 'tests/**'
echo '== Candidate native KV event tests =='
fd -t f -i '.*(native.*kv.*event|kv.*event.*native|kv_cache_event).*' tests 2>/dev/null || true
echo '== Relevant diff from the fix commit =='
git show --format= --unified=8 69b318a8cc -- \
tensorrt_llm/_torch/pyexecutor/kv_cache_events.py \
tensorrt_llm/llmapi/llm_args.py \
'tests/**'
echo '== Relevant assertions in discovered candidate tests =='
fd -t f -i '.*(native.*kv.*event|kv.*event.*native|kv_cache_event).*' tests 2>/dev/null \
| while IFS= read -r test_file; do
echo "--- ${test_file} ---"
rg -n -A12 -B8 'no.port|no-port|flush_iteration_events|try:|finally:|-> None|remov' "$test_file" || true
doneLength of output: 60710
The socket cleanup, endpoint validation, multimodal-token suppression, and
I accept the deferred best-effort delivery behavior and the current ZeroMQ test timing rationale. The tests were inspected but not executed. 🐇 ✅ ✏️ Learnings added
You are interacting with an AI system. |
…ents-simplification Signed-off-by: tanmayv25 <tanmay2592@gmail.com> # Conflicts: # tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py # tensorrt_llm/llmapi/__init__.py # tensorrt_llm/llmapi/llm_args.py
The published KV-cache-event wire schema was described throughout as 'vLLM-compatible', which reads oddly in a TensorRT-LLM module. Rename the descriptive references to TensorRT-LLM's own 'KV cache event wire' terminology: - _vllm_wire_hash_from_radix_key -> _kv_event_wire_hash_from_radix_key - 'vLLM-compatible'/'vLLM wire'/'vLLM's ... hash' -> 'KV cache event wire ...' across the struct, publisher, endpoint, and manager docstrings and messages. The one retained vLLM mention is the module-header source attribution: the on-wire schema was adapted from vLLM's Apache-2.0 vllm/distributed/kv_events.py. No behavior change; the renamed helper is private to this module. Signed-off-by: tanmayv25 <tanmay2592@gmail.com>
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
tensorrt_llm/_torch/pyexecutor/kv_cache_events.py (2)
638-648: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInclude the remaining counters in the shutdown summary.
The summary omits
multimodal_blocks_suppressed,enqueued_events, and the publisher-side drop count.dropped_batcheson this class counts only local publish failures;self._publisher.dropped_batchescounts queue-full and send-error drops insideZmqEventPublisher. Without the publisher counter, a run that drops batches at the ZeroMQ layer looks clean in this log line.♻️ Proposed change
logger.info( "Streaming 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"multimodal_blocks_suppressed={self.multimodal_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"enqueued_events={self.enqueued_events} " f"dropped_batches={self.dropped_batches}" )🤖 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 638 - 648, Extend the shutdown summary log in the relevant KV cache events class to include multimodal_blocks_suppressed, enqueued_events, and the publisher-side drop count from self._publisher.dropped_batches, while retaining the existing local dropped_batches counter.
147-148: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
X | Nonefor consistency with the rest of the file.Lines 147 and 148 use
Optional[zmq.Socket], while every other annotation in this file uses the|union form. The coding guidelines ask for built-in generic types and|.from __future__ import annotationsis already present, so the change is safe.♻️ Proposed change
- self._pub: Optional[zmq.Socket] = None - self._replay: Optional[zmq.Socket] = None + self._pub: zmq.Socket | None = None + self._replay: zmq.Socket | None = NoneThen drop
Optionalfrom thetypingimport on line 31 if it becomes unused.🤖 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 147 - 148, Update the _pub and _replay attribute annotations to use zmq.Socket | None, matching the file’s existing union-style annotations. Remove Optional from the typing import if it is no longer used.Source: Coding guidelines
tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py (2)
152-154: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueEndpoint-release check runs only on the success path.
Lines 152-154 sit outside the
try/finallyblock. If an assertion fails, this check does not run. That is acceptable, because the check itself is an assertion about cleanup. No change required if that is intended.🤖 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_streaming_kv_events.py` around lines 152 - 154, The endpoint-release assertion around the replacement PUB socket is intentionally outside the try/finally cleanup block, so leave this code unchanged and preserve the current success-path-only check.
1-243: 📐 Maintainability & Code Quality | 🔵 TrivialTest coverage summary.
Test functions added (all new, none modified or removed):
test_streaming_fast_path_publishes_only_full_max_window_blocks— covers radix-hash reuse, partial-block suppression, non-target lifecycle filtering, the three-frame ZeroMQ wire format, sequence numbering, the empty buffered-pull degradation, idempotentshutdown(), and endpoint release.test_streaming_removals_are_never_dropped_by_the_entry_cap— covers removal emission pastmax_entriesand the msgpackBlockRemovedpayload.test_kv_events_config_publisher_default— coversKVEventsConfig.model_post_initpublisher resolution.test_offset_endpoint_port— covers rank offsetting for tcp, ipc, inproc, andNone.test_offset_endpoint_port_rejects_bad_input— covers port overflow, unknown scheme, and a missing TCP port.Test list registration: these are unit tests under
tests/unittest/, so they are collected bypytest tests/unittest/. No entry was added undertests/integration/test_lists/test-db/ortests/integration/test_lists/qa/. That matches the unit-test path convention.Coverage verdict: needs follow-up. The wire-format, filtering, and endpoint-validation paths are well covered. Two gaps remain in the changed production surface:
KVCacheManagerV2streaming integration is untested: the parallelism guards atkv_cache_manager_v2.pylines 894-897, the failure cleanup at lines 1118-1121, the attention-only lifecycle filter at lines 1549-1559, and thestreaming_kv_events_enabledproperty at lines 3026-3028.StreamingKVCacheEventManagermultimodal suppression (multimodal_blocks_suppressed) and the queue-full drop counter have no assertions.Add a manager-level test for the
ValueErrorraised onpp_size > 1andcp_size > 1, and an assertion onmultimodal_blocks_suppressedfor a bytes-valued token.Note: I inspected the tests but did not execute them.
As per path instructions, "Always produce a test coverage summary, even if no issues are found" and the summary must state which test functions changed, whether they are listed in the appropriate test list files, and a coverage verdict.
Do you want me to generate the manager-level guard tests and the multimodal-suppression assertion?
🤖 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_streaming_kv_events.py` around lines 1 - 243, Add coverage for the missing production paths: create a manager-level test exercising the KVCacheManagerV2 parallelism guard and asserting ValueError when both pp_size and cp_size exceed one, and extend StreamingKVCacheEventManager coverage to pass a bytes-valued token and assert multimodal_blocks_suppressed increments. Use existing test fixtures/configuration patterns and preserve the current test-list conventions.Source: Path instructions
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)
1549-1559: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPrefer the public layer config over
impl._life_cycles.Line 1551 reads the private
self.impl._life_cyclesattribute. The same file already derives the attention/SSM distinction from public data in_stats_life_cycle_metadata(lines 2686-2695) by checkingisinstance(layer, AttentionLayerConfig)onself.kv_cache_manager_py_config.layers. Reusing that public path keeps the manager independent of the impl's internals and avoids a second, divergent classification rule.♻️ Proposed refactor using the public layer config
window_sizes: Dict[int, int] = {} for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping): - life_cycle = self.impl._life_cycles.get_life_cycle(LifeCycleId(layer_group_id)) # Streaming KV events track attention prefix reuse only. Excluding SSM # and other non-attention life cycles prevents a state life cycle # (which reports max_seq_len as its window) from tying with the # attention life cycle and being selected as the event target. - if not isinstance(life_cycle, AttnLifeCycle): + layer_config = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] + if not isinstance(layer_config, AttentionLayerConfig): continue window_sizes[int(layer_group_id)] = get_event_window_size(int(layer_ids[0])) return window_sizesIf you keep the current form, the
LifeCycleIdimport at line 63 stays required; the refactor removes that need.🤖 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_manager_v2.py` around lines 1549 - 1559, Update the window-size collection around the layer-group loop to classify attention groups using self.kv_cache_manager_py_config.layers and AttentionLayerConfig, matching _stats_life_cycle_metadata, instead of accessing self.impl._life_cycles via LifeCycleId. Preserve inclusion of only attention layer groups and remove the now-unused LifeCycleId dependency if applicable.
🤖 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 326-338: Update the TCP endpoint handling around the port
conversion in the endpoint rank-adjustment function to validate the port text
before calling int(). Reject non-numeric, zero, and negative base ports with a
ValueError that includes the full endpoint, then preserve the existing rank
offset and upper-bound validation for valid ports.
In `@tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py`:
- Around line 90-94: Run ruff format on the test block around first_wire_hash
and second_wire_hash, applying its wrapping to the overlong conditional
assignments, then commit the formatter’s output without changing behavior.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_events.py`:
- Around line 638-648: Extend the shutdown summary log in the relevant KV cache
events class to include multimodal_blocks_suppressed, enqueued_events, and the
publisher-side drop count from self._publisher.dropped_batches, while retaining
the existing local dropped_batches counter.
- Around line 147-148: Update the _pub and _replay attribute annotations to use
zmq.Socket | None, matching the file’s existing union-style annotations. Remove
Optional from the typing import if it is no longer used.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 1549-1559: Update the window-size collection around the
layer-group loop to classify attention groups using
self.kv_cache_manager_py_config.layers and AttentionLayerConfig, matching
_stats_life_cycle_metadata, instead of accessing self.impl._life_cycles via
LifeCycleId. Preserve inclusion of only attention layer groups and remove the
now-unused LifeCycleId dependency if applicable.
In `@tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py`:
- Around line 152-154: The endpoint-release assertion around the replacement PUB
socket is intentionally outside the try/finally cleanup block, so leave this
code unchanged and preserve the current success-path-only check.
- Around line 1-243: Add coverage for the missing production paths: create a
manager-level test exercising the KVCacheManagerV2 parallelism guard and
asserting ValueError when both pp_size and cp_size exceed one, and extend
StreamingKVCacheEventManager coverage to pass a bytes-valued token and assert
multimodal_blocks_suppressed increments. Use existing test
fixtures/configuration patterns and preserve the current test-list conventions.
🪄 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: 9a8faf0f-d510-43a7-8185-321537a3343e
📒 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_streaming_kv_events.py
🚧 Files skipped from review as they are similar to previous changes (6)
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/llmapi/llm_utils.py
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tensorrt_llm/llmapi/init.py
- tensorrt_llm/llmapi/llm_args.py
- tensorrt_llm/_torch/pyexecutor/_util.py
- offset_endpoint_port: reject a non-numeric or out-of-range TCP port (e.g. tcp://host:abc, :0, :-5) with an endpoint-naming error instead of an opaque int()/ZeroMQ bind failure that only surfaces on ranks > 0; add tests. - test_streaming_kv_events.py: apply the ruff-format wrapping the CI pre-commit check flagged (long signed-wraparound line). Signed-off-by: tanmayv25 <tanmay2592@gmail.com>
|
Thanks @coderabbitai — both addressed in
|
The merge left the llm_args import block in llm_utils.py isort-dirty (KVEventsConfig sorted before KvCacheConfig). Apply isort ordering; this is the Pre-commit Check CI failure. Signed-off-by: tanmayv25 <tanmay2592@gmail.com>
tanmayv25
added
the
api-compatible
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
echo '== Commit =='
git show --no-patch --oneline e01518a36a
echo '== Focused diff =='
git show --format= --unified=12 e01518a36a -- \
tensorrt_llm/_torch/pyexecutor/kv_cache_events.py \
tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py
echo '== Publisher and endpoint implementation outline =='
ast-grep outline tensorrt_llm/_torch/pyexecutor/kv_cache_events.py \
--items all --type class,function --match 'ZmqEventPublisher|offset_endpoint_port'
echo '== Focused implementation and tests =='
rg -n -A35 -B12 \
'def offset_endpoint_port|invalid TCP endpoint|out of range|non-numeric|tcp://host:abc|tcp://host:0|tcp://host:-5|65535|65536|signed' \
tensorrt_llm/_torch/pyexecutor/kv_cache_events.py \
tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.pyLength of output: 30095
The signed-wraparound expression in The tests were inspected but not executed. 🐇 ✅ 🧠 Learnings usedYou are interacting with an AI system. |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Two whole-PR items, plus a pointer:
- Tracking ticket: this is a substantial new feature (new user-facing config, a wire protocol, a background publisher thread) carried under
[None]. Please file/link a TRTLLM JIRA — the RFC (#17013) reference is useful context but isn't a tracking ticket. - Docs:
kv_cache_config.kv_events_configis user-facing (prototype status notwithstanding), and nothing indocs/sourcecovers it —docs/source/features/kvcache.mddocuments only the pull API. External consumers (Dynamo) need the endpoint/base-port-plus-rank convention, replay semantics, wire format, and the V2-only + parallelism constraints written down.
The main mergeability question is backend support — see the comment on [kv_cache_manager_v2.py:901](https://github.com/NVIDIA/TensorRT-LLM/pull/17023/files#diff-38c4fe2f7a5f883b2fd9b2d894d2e186e846fbc0814bd41ce1cdf33b1be6bba6R901): as written the streaming manager can only be consumed by the pure-Python V2 backend, while the default backend is cpp.
| 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 = StreamingKVCacheEventManager( |
There was a problem hiding this comment.
This duck-typed manager cannot be consumed by the default V2 backend. tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py:23 selects the backend via TLLM_KV_CACHE_MANAGER_V2_BACKEND with default "cpp", where KVCacheManagerPy resolves to the nanobind kv::KvCacheManager, whose constructor does nb::cast<std::shared_ptr<kv::EventManager>>(eventManager) (cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp:1992). A plain Python StreamingKVCacheEventManager fails that cast, so enabling kv_events_config under the default backend raises an opaque TypeError at manager construction — and even if the cast were relaxed, the C++ radix tree invokes kv::EventManager methods natively, so the Python hooks would never fire. As written the feature only works with TLLM_KV_CACHE_MANAGER_V2_BACKEND=python. Either add C++-side support, or gate this branch on the active backend with a clear error/warning (like the V1 warning in _util.py). Relatedly, no test constructs KVCacheManagerV2 with kv_events_config set — an integration test on the default backend would have caught this.
| data_parallel_rank: int | None = None | ||
|
|
||
|
|
||
| class KVCacheWireEvent( |
There was a problem hiding this comment.
The module header says the schema is adapted from vLLM's kv_events.py so routers can consume events "without translation", but vLLM's KVCacheEvent base sets array_like=True (events encode as tagged positional arrays, [tag, field1, ...]), while KVCacheWireEvent here omits it — events encode as maps with a "type" key (the new test asserts exactly that dict form). A consumer decoding with vLLM's msgspec schema will fail on these batches. Please verify against the actual Dynamo subscriber; if vLLM wire compatibility is the goal, add array_like=True here (which also makes field order part of the wire contract), and if divergence is intentional, fix the header comment so consumers know translation is required.
| # 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): |
There was a problem hiding this comment.
This AttnLifeCycle filter also changes the existing buffered path: _get_event_window_sizes_by_layer_group() feeds the buffered KVCacheEventManager.set_layer_group_window_sizes() too, and SSM/state layer groups are now absent from the dict, so their events fall back to the manager-level default window (_get_window_size → _window_size = max attention window) instead of the max_seq_len they previously reported. That only diverges on hybrid models where every attention layer is sliding-window, but the PR description claims the buffered path is unchanged. Either scope the filter to the streaming manager (e.g. filter inside StreamingKVCacheEventManager.set_layer_group_window_sizes) or state the buffered-path change as intended.
| "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( |
There was a problem hiding this comment.
replay_endpoint gets the same base-port-plus-rank offset as endpoint, so with N attention-DP ranks the pub range is [base, base+N-1] and the replay range is [replay_base, replay_base+N-1]. If a user picks a replay base within N-1 ports of the pub base (e.g. endpoint :5557, replay :5558, 2 ranks), rank 1's pub bind collides with rank 0's replay bind and fails at startup. Document the required spacing here (the endpoint description already warns about co-located engines) or validate the two ranges don't overlap given the DP size.
4 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 streaming (push-based) KV-event path for KV cache manager
V2 (PyTorch backend). Off by default; the buffered gather/poll path is
unchanged.
StreamingKVCacheEventManager(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 — buffered build → buffer → gather → per-iteration pull (still the
default; unchanged by this PR):
After — streaming 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["StreamingKVCacheEventManager<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.
Validation
Early internal validation on a multi-node GB300 disaggregated deployment
(DeepSeek-V4-Pro, high-reuse agentic workload) shows the streaming path
delivering meaningfully higher end-to-end throughput than the buffered
gather/poll path — consistent with the mechanism above (no rank-0 gather, no
per-iteration pull). Broader benchmarking is ongoing.
Context
Supersedes #16869 and #16876 by @alec-flowers (original authorship preserved).
Relates to RFC #17013.
Dev Engineer Review
KVEventsConfigwith ZeroMQ, replay, queue, HWM, buffer, topic, and publisher settings.StreamingKVCacheEventManagerwithmsgspecevents, event coalescing, replay buffering, queue handling, background publishing, shutdown, and counters.KVEventsConfig.QA Engineer Review
Added test functions:
test_streaming_fast_path_publishes_only_full_max_window_blockstest_streaming_removals_are_never_dropped_by_the_entry_captest_kv_events_config_publisher_defaulttest_offset_endpoint_porttest_offset_endpoint_port_rejects_bad_inputNo corresponding
test-db/orqa/coverage entries were identified. Tests were inspected but not executed. Verdict: needs follow-up.