Conversation
c095803 to
4fc3763
Compare
|
Understand this PR’s impact Explore downstream dependencies and potential security impact with Blast Radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (15)
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review. WalkthroughThe PR adds native C++ streaming events for KV-cache blocks. It exposes the sink through bindings, integrates native event draining with Python publication, supports both V2 backends, and adds multimodal event handling and tests. ChangesNative streaming KV-cache events
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~45 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant KVCacheManager
participant StreamingEventSink
participant StreamingKVCacheEventManager
participant WirePublisher
KVCacheManager->>StreamingEventSink: submit block lifecycle events
StreamingEventSink->>StreamingKVCacheEventManager: provide drained native DTOs
StreamingKVCacheEventManager->>WirePublisher: publish stored or removed wire events
StreamingKVCacheEventManager->>StreamingEventSink: synchronize native statistics
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 17.44% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 86 functions across 15 files. (2 skipped: 2 unsupported.)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
9850306 to
003d83f
Compare
Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: Guan Luo <gluo@nvidia.com>
003d83f to
3e0a0df
Compare
| class StreamingEventSink final : public EventSink | ||
| { | ||
| public: | ||
| StreamingEventSink(int tokensPerBlock, int maxEntries, std::optional<int> mmTokenIdOffset = std::nullopt); |
There was a problem hiding this comment.
tokensPerBlock is stored and validated here but never read anywhere in StreamingEventSink — Block::isFull() already derives the block size from the block itself — so a wrong value silently does nothing; please drop the parameter. Relatedly, max_entries and mm_token_id_offset now have to be configured identically on both the native sink and StreamingKVCacheEventManager (which keeps its own inert _max_entries/_pending_entries in native mode), and a mismatch fails silently — e.g. a sink with an offset but a facade without one drops the computed mm_keys. Could StreamingKVCacheEventManager construct the native sink itself from its own parameters instead?
| block_hashes: list[ExternalBlockHash] | ||
| parent_block_hash: ExternalBlockHash | None | ||
| token_ids: list[int] | ||
| token_ids: list[EventTokenId] |
There was a problem hiding this comment.
Streaming token_ids was previously guaranteed to be list[int] (multimodal blocks were skipped via multimodal_blocks_suppressed) and can now contain hex digest strings, so existing subscribers that hash it as ints will break on multimodal payloads. This is a separate feature from C++-backend support, and the description's "preserving the existing serialization" no longer holds once mm_keys is added and token_ids is widened — could multimodal enablement be split out, or at least called out in the description?
| } | ||
| if (block.prev == nullptr) | ||
| { | ||
| throw std::logic_error("Cannot publish an orphan KV cache block"); |
There was a problem hiding this comment.
The orphan check throws after reserveEntryUnlocked() has already consumed an entry from the per-iteration budget, and unlike the Python path (_add_full_block, which rolls back _pending_entries) nothing releases it. Move the block.prev == nullptr guard above reserveEntryUnlocked().
| if mm_token_id_offset is not None and mm_token_id_offset < 0: | ||
| raise ValueError("mm_token_id_offset must be non-negative") | ||
| self._mm_token_id_offset = mm_token_id_offset | ||
| self._native_event_sink = native_event_sink |
There was a problem hiding this comment.
In native mode _stored_blocks, _pending_events, _pending_entries, _max_entries and the whole _add_full_block/_decode_block/add_removed_event sink surface are unreachable, and in Python mode _drain_native_events/_sync_native_stats are — consider splitting the two sink implementations and keeping only the shared publisher lifecycle, batching and counters on the facade. In particular, the new Python _decode_block is a hand-maintained twin of C++ decodeEventBlock that only serves the backend this PR describes as being removed. Also worth documenting that stored_blocks/removed_blocks only refresh at flush_iteration_events()/shutdown(), unlike the Python-owned batch counters.
|
|
||
| window_sizes: Dict[int, int] = {} | ||
| for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping): | ||
| layer_config = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] |
There was a problem hiding this comment.
layer_config is looked up unconditionally here and then looked up again by the same index inside get_event_window_size. Pass layer_config into the helper instead.
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Approve with nits.
2 participants
Dev Engineer Review
The PR adds native C++ streaming KV-cache events with lifecycle filtering, event coalescing, bounded buffering, statistics, and multimodal metadata. Nanobind exposes the native sink and DTOs. Python converts DTOs into existing
BlockStoredandBlockRemovedevents.The C++ backend uses the native sink. The Python backend keeps its existing path. Verify event ordering, overflow handling, shutdown flushing, lifecycle selection, multimodal key alignment, and API compatibility.
QA Engineer Review
Tests cover native stored and removed events, event coalescing, payload conversion, lifecycle filtering, statistics, multimodal metadata, backend validation, and attention-layer window filtering. Pipeline- and context-parallelism rejection checks remain covered.
No changed test-list files are present. CI or manual-QA list registration is therefore not required by the changed files.
Coverage verdict: sufficient based on the supplied tests and reported checks.
Per-File QA Perspective
CMakeLists.txt: Adds native event sources to the build. Verify compilation and linking.streamingEventSink.cpp: Adds event filtering, coalescing, buffering limits, removal handling, and statistics. Verify invalid, duplicate, multimodal, orphan, and overflow cases.streamingEventSink.h: Defines native DTOs and sink APIs. Verify binding compatibility.eventData.cpp: Decodes tokens, digests, and multimodal metadata. Verify hexadecimal encoding and block context.eventData.h: Adds shared event-token and multimodal-key types. Verify parity with Python structures.eventManager.cpp: Uses shared event decoding for stored blocks. Verify cache metadata and event payloads.eventManager.h: Removes duplicated event-data declarations. Verify dependent C++ interfaces.kvCacheManagerV2.cpp: Exposes sink types, DTOs, statistics, helpers, and sink inputs. Verify bindings and defaults.kv_cache_manager_v2.py: Selects the native sink for the C++ backend. Verify host-tier fallback and attention-layer filtering.kv_cache_events.py: Supports multimodal metadata and native DTO conversion. Verify flushing, validation, digest handling, and removal behavior.runtime/kv_cache_manager_v2/__init__.py: Exports native event types for non-Python backends. Verify imports and__all__.runtime/kv_cache_manager_v2/__init__.pyi: Declares native event types. Verify stub and binding signature parity._introspection.py: Adds native sink test helpers. Verify forwarding and unavailable-backend errors.kvcache.md: Documents C++ backend support and multimodal events. Verify exclusions and parallelism constraints.test_kv_cache_manager_v2.py: Covers attention-layer window filtering without private backend state. No changed test-list entry applies.test_kv_cache_event_manager.py: Covers native events, payloads, lifecycle filtering, counters, and multimodal metadata. No changed test-list entry applies.test_streaming_kv_events.py: Covers backend support, multimodal events, and invalid parallelism or backend configurations. No changed test-list entry applies.Description
The Python KV cache manager backend is being removed, but the streaming KV-event publisher introduced by #17023 currently depends on Python-side KV manager events.
This change keeps streaming KV events available with the C++ KV cache manager V2 backend:
StreamingEventSinkthat captures stored and removed block events as compact C++ DTOs.BlockStoredandBlockRemovedmsgspec structures, preserving the existing serialization and publisher path.The boundary intentionally remains semantic rather than serialized:
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions).
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, comment
/bot help.