[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 32 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>
|
/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>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThis PR adds configurable native KV-cache event publishing. It defines event wire types, asynchronous ZeroMQ publishing, replay support, and streaming event batching. It integrates the feature with ChangesNative KV-cache events
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🔵 Low · up to The opt-in streaming KV-cache event path is mergeable, but the documentation currently lists event types that do not match the published contract, which could mislead external router implementers; align the documentation before or immediately after merge. Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant KVCacheManagerV2
participant StreamingKVCacheEventManager
participant ZmqEventPublisher
PyExecutor->>KVCacheManagerV2: pass KVEventsConfig
KVCacheManagerV2->>StreamingKVCacheEventManager: initialize streaming manager
KVCacheManagerV2->>StreamingKVCacheEventManager: forward cache lifecycle hooks
StreamingKVCacheEventManager->>ZmqEventPublisher: publish stored or removed batch
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description clearly explains the motivation, implementation, configuration, constraints, behavior changes, validation status, and related context. It is mostly complete, although it does not use the template's exact Description and Test Coverage headings or include the checklist. Full details: Docstring CoverageExplanation Docstring coverage is 21.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 60 functions across 5 files. (1 skipped: 1 unsupported.) ✨ 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
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>
|
PR_Github #69563 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #69797 [ run ] triggered by Bot. Commit: |
|
PR_Github #69797 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #70005 [ run ] triggered by Bot. Commit: |
|
PR_Github #70005 [ run ] completed with state
|
|
/bot run |
|
PR_Github #70055 [ run ] triggered by Bot. Commit: |
|
PR_Github #70055 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
All concerns from the previous round check out in the current code, each with a regression test:
- cpp-backend incompatibility →
validate_streaming_supportnow raises with an actionable message namingTLLM_KV_CACHE_MANAGER_V2_BACKEND=python, before any socket is bound. - vLLM wire-format claim → module header now states the map-vs-positional-array divergence accurately.
AttnLifeCyclefilter leaking into the buffered path →_get_event_window_sizes_by_layer_group(attention_only=...)defaults off; only the streaming manager opts in, soKVCacheEventManagerwindows are unchanged.- publish/replay port collision →
validate_endpoint_rangesrejects overlapping ranges at startup with per-host spacing semantics, and the parametrized test covers both directions plus the multi-node case. - docs →
docs/source/features/kvcache.mdnow documents both paths, the base_port+rank convention, wire format, and the gap-means-loss delivery contract.
One prior item remains open: the tracking ticket. This is a substantial new feature (user-facing config, a wire protocol, a background publisher thread) still carried under [None] — please file/link a TRTLLM JIRA in the title before merge. Not re-blocking on it, but it was asked last round.
Non-blocking note: the description's validation section still says GB300 benchmark numbers are "in progress" — worth updating with results (or dropping the promise) before merge so the merged description is accurate.
Approving — the inline comment is an optional touch-up.
|
/bot run |
|
PR_Github #70120 [ run ] triggered by Bot. Commit: |
|
PR_Github #70120 [ run ] completed with state
|
Signed-off-by: allisonlim-nv <allim@nvidia.com>
Signed-off-by: Allison Lim <allim@nvidia.com>
7755478 to
573b40e
Compare
Signed-off-by: Allison Lim <allim@nvidia.com>
|
/bot run |
|
PR_Github #70176 [ run ] triggered by Bot. Commit: |
|
PR_Github #70176 [ run ] completed with state
|
|
/bot run |
|
PR_Github #70199 [ run ] triggered by Bot. Commit: |
|
PR_Github #70199 [ run ] completed with state
|
14 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
End-to-end benchmarking on a multi-node GB300 disaggregated deployment
(DeepSeek-V4-Pro) is in progress; numbers will be added once available. (The
streaming path requires
TLLM_KV_CACHE_MANAGER_V2_BACKEND=python, since theper-block event hooks run in the V2 Python backend.)
Context
Supersedes #16869 and #16876 by @alec-flowers (original authorship preserved).
Relates to RFC #17013.
Dev Engineer Review
KVEventsConfigand public exports.msgspec.BACKENDexport.QA Engineer Review
tests/unittest/kv_cache_manager_v2_tests/test_streaming_kv_events.py.test_dropped_batches_leave_a_sequence_gaptest_validate_streaming_support_rejects_unsupported_setupstest_validate_endpoint_rangestests/integration/test_lists/coverage entry was identified.