Skip to content

[None][feat] Add image routing metadata to OpenEngine KV events - #21125

Open
krishung5 wants to merge 57 commits into
NVIDIA:mainfrom
krishung5:openengine-image-routing
Open

krishung5 wants to merge 57 commits into
NVIDIA:mainfrom
krishung5:openengine-image-routing

Conversation

@krishung5

@krishung5 krishung5 commented Oct 27, 2026 •

Copy link
Copy Markdown
Collaborator

Description

Expose optional image routing identities alongside the native KV cache content digests so external routers can compute request-compatible image-prefix hashes. The engine continues to hash image content and caller UUIDs; a routing hint does not replace content-sensitive cache identity.

The image-input prerequisite #20446 is merged into main. This PR is independent of the image performance follow-up #21037 and the video/audio PRs, so review and validation can proceed in parallel. The companion Dynamo implementation is ai-dynamo/dynamo#17437, stacked on image passthrough ai-dynamo/dynamo#16613. End-to-end deployment requires those Dynamo changes and an official TRT-LLM runtime containing this routing implementation.

This PR targets main. The prerequisite image implementation is already on main; after the branch refresh is pushed, the Files changed view will show the routing follow-up and its processor-default regression fix.

  • Expose the existing native request UUID vector through a read-only binding.
  • Retain a bounded per-publisher digest-to-UUID mapping and include optional UUIDs in streaming KV events. Missing/evicted mappings allow consumers to ignore the affected image stores without dropping subsequent removals.
  • Add versioned optional trtllm_image_routing event metadata. Text-only stores omit this extension. The publisher does not interpret client-specific UUID prefixes.
  • Advertise ModelInfo.extra.trtllm_image_routing_version = 1 only for aggregated Qwen3-VL with explicit KV cache V2, block reuse, and ZMQ events. Disaggregated routing remains disabled.

Additive API change; a repository maintainer needs to add api-compatible because the author account cannot apply repository labels. No OpenEngine proto changes or new dependencies. Routing identities must refer to stable media content.

Test Coverage

  • 62 focused CPU cases passed for capability gating, event metadata, bounded mapping, and content-sensitive identities.
  • Real-weight Qwen3-VL-2B GPU conformance passed: correct red/blue responses; additional KV reuse for a repeated image; different image content with the same routing hint did not reuse the first image's prefix or answer.
  • Captured GPU events were replayed through the companion Dynamo decoder: two image stores indexed, none ignored.
  • Targeted pre-commit hooks passed. CPU and real-weight GPU tests are registered in the L0 CPU/A10 test lists; native event-manager regression coverage is also extended.

Local GPU evidence uses a TensorRT-LLM 1.4.0rc0 runtime with the relevant Python sources installed and explicit validation-only adapters for the new native UUID getter and release capacity discovery. This is not a matching compiled PR wheel. Validation of the new binding and the native introspection-based event suite requires the matching upstream build. Full HTTP routing between two workers and performance benchmarking remain pending. Matching-build validation and full HTTP cache-aware selection between two workers remain required before merge.

PR Checklist

  • Reviewed the applicable checklist: description, coding guidelines, tests and CI registration, additive API compatibility, documentation, dependencies, and ownership. No new dependencies or ownership changes.

Dev Engineer Review

The change adds optional image-routing UUIDs to streaming KV events. Content digests remain the cache identity. OpenEngine advertises routing version 1 only for aggregated Qwen3-VL when KV cache manager V2, block reuse, and ZMQ events are enabled. Disaggregated routing remains disabled.

The UUID mapping is bounded, ignores UUIDs over 256 UTF-8 bytes, and rejects conflicting UUIDs for the same digest. Verify that consumers handle missing or evicted hints without dropping later removals. The new read-only GenericLlmRequest.multimodal_uuids property and media-I/O defaults are additive API changes.

Reported validation includes focused CPU tests, GPU conformance, and targeted pre-commit checks. The GPU test used a different TensorRT-LLM runtime and validation-only adapters. Matching-build validation and full two-worker HTTP routing remain pending. The reported CI runs failed. No current review findings were supplied, so review severity counts are unavailable.

QA Engineer Review

Tests were added or modified in the OpenEngine, KV-cache event-manager, and executor test directories. Coverage includes capability selection, UUID propagation, content-sensitive hashes, event serialization, text-only omission, mapping eviction, oversized UUIDs, conflicting hints, and real-model cache reuse.

The CPU CI list includes test_control.py, test_image_routing.py, and test_multimodal.py. The A10 CI list includes test_image_routing_conformance.py. The KV-cache event-manager and executor tests have no test-list entries reported here. Focused CPU tests and GPU conformance were reported as passing, but the PR CI runs failed. Matching-build validation and full two-worker HTTP routing remain pending. Coverage verdict: needs follow-up.

Per-File QA Perspective

  • cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp — Verify that GenericLlmRequest.multimodal_uuids returns None when UUID data is absent and a copy when present.
  • tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py — Verify that streaming event managers receive each digest/UUID pair and that other manager types skip registration.
  • tensorrt_llm/_torch/pyexecutor/kv_cache_events.py — Verify bounded UUID retention, conflict handling, oversized UUID rejection, and optional metadata emission. Check missing or evicted hints and subsequent removals.
  • tensorrt_llm/grpc/openengine/README.md — Check that documented support limits match runtime behavior, including disaggregated routing and non-image modalities.
  • tensorrt_llm/grpc/openengine/capabilities.py — Verify the image-routing capability conditions and exclusions.
  • tensorrt_llm/grpc/openengine/control.py — Verify that GetModelInfo advertises the routing version only when the capability is supported.
  • tensorrt_llm/inputs/registry.py — Verify that the new default media-I/O kwargs preserve behavior when processors do not define decode defaults.
  • tests/integration/test_lists/test-db/l0_a10.yml — Adds the GPU conformance test to the A10 CI list.
  • tests/integration/test_lists/test-db/l0_cpu.yml — Adds the CPU image-routing test; the list also contains the modified control and multimodal tests.
  • tests/unittest/grpc/openengine/test_control.py — Covers capability conditions and exclusions. Listed in test-db/l0_cpu.yml.
  • tests/unittest/grpc/openengine/test_image_routing.py — Covers UUID propagation, content-sensitive hashes, text-only omission, eviction, oversized UUIDs, and conflicts. Listed in test-db/l0_cpu.yml.
  • tests/unittest/grpc/openengine/test_image_routing_conformance.py — Covers real-model routing, cache reuse, and captured KV events. Listed in test-db/l0_a10.yml.
  • tests/unittest/grpc/openengine/test_multimodal.py — Covers image-source mapping with absent and explicit media-I/O defaults. Listed in test-db/l0_cpu.yml.
  • tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py — Covers routing metadata, bounded retention, digest independence, oversized hints, and conflicts. No test-list entry was reported.
  • tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py — Updates the test stub’s event_manager default to None; no test-list entry was reported.
  • tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py — Updates the fixture’s event_manager default to None; no test-list entry was reported.
  • tests/unittest/_torch/executor/test_pytorch_model_engine.py — Updates two MRoPE test setup calls; no test-list entry was reported.
  • tests/unittest/_torch/test_connector_adp.py — Updates distributed test mocks and rank setup; no test-list entry was reported.

GuanLuo and others added 30 commits September 23, 2026 16:31
Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: Guan Luo <gluo@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
…ngine-image-media

Signed-off-by: krishung5 <krish@nvidia.com>

# Conflicts:
#	tensorrt_llm/grpc/openengine/capabilities.py
#	tensorrt_llm/grpc/openengine/request_mapping.py
#	tensorrt_llm/inputs/media_io.py
#	tests/unittest/grpc/openengine/test_multimodal.py
#	tests/unittest/inputs/test_url_validation.py
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>

# Conflicts:
#	tensorrt_llm/inputs/utils.py
Signed-off-by: krishung5 <krish@nvidia.com>
Signed-off-by: krishung5 <krish@nvidia.com>
@krishung5

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #81222 [ run ] triggered by Bot. Commit: d550dcf Link to invocation

@krishung5
krishung5 requested a review from JunyiXu-nv October 28, 2026 18:27

@brnguyen2 brnguyen2 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The routing-hint design keeps cache identity content-sensitive (the digest still covers content + UUID, and the UUID only rides in mm_keys[].uuid), and the CPU coverage of eviction, ambiguity, partial-UUID slots and atomic registration is thorough. Three things to resolve before merge:

  1. Unrelated changes bundled. The "Fix KV-cache, ADP and mRoPE test fixtures" commit (test_kv_cache_manager_v2_helix_superblock.py, test_pytorch_model_engine.py self arg, test_connector_adp.py MagicMock attrs) fixes tests broken by other main merges and has nothing to do with image routing. Same for the BaseMultimodalInputProcessor.get_preferred_media_io_kwargs default in [tensorrt_llm/inputs/registry.py:177](https://github.com/NVIDIA/TensorRT-LLM/pull/21125/files#diff-34567c3d44579a4e959cdc4ceaefcc344eef1adaf2c34f4f111fc1a1f06f908bR177) (a real fix: callers isinstance-check that class and call the method unconditionally, but the existing default lives on BaseMultimodalDummyInputsBuilder). Please split these into their own PR so they land independently and this PR's CI signal stays attributable. The two event_manager = None stub additions are needed here and should stay.

  2. Binding is unverified on a matching build, per the description. test_image_routing_conformance.py is the only test that exercises the new multimodal_uuids nanobind property end to end; it must pass on this PR's own wheel in the A10 stage before merge. Please drop the stale "after the branch refresh is pushed" sentence once that's done.

  3. Docs. docs/source/features/kvcache.md:324 describes the V2 mm_keys wire format (hash, start_offset) and should mention the new optional uuid field and the trtllm_image_routing extension, or link to the OpenEngine README section. Also note in that doc that the extension is emitted for any multimodal block (video/audio keys included), not only images, even though the capability flag is image-scoped.

Inline comments cover a per-chunk registration cost and a dead branch.

Comment thread tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/kv_cache_events.py Outdated
Comment thread tensorrt_llm/grpc/openengine/capabilities.py
Comment thread tensorrt_llm/_torch/pyexecutor/kv_cache_events.py
Comment thread tests/unittest/grpc/openengine/test_image_routing_conformance.py
Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com>
Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com>
@krishung5

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #81240 [ run ] triggered by Bot. Commit: 3b8f84d Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #81222 [ run ] completed with state ABORTED. Commit: d550dcf

Link to invocation

@krishung5
krishung5 requested a review from brnguyen2 October 28, 2026 20:11
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #81240 [ run ] completed with state SUCCESS. Commit: 3b8f84d
/LLM/main/L0_MergeRequest_PR pipeline #66932 completed with status: 'FAILURE'

CI Report

⚠️ Resource Maintenance — Test Coverage Reduced:

Resource maintenance in LLM/main/L0_Test-x86_64-Single-GPU #4133: L40S-PyTorch-2 is temporarily paused due to Kubernetes agent provisioning and node issues.
Affected stages:
L40S-PyTorch-2

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@brnguyen2 brnguyen2 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving — the comments below are optional touch-ups, not blockers.

The routing-hint path now matches the README in every respect but one: the README says a conflicting digest stays ambiguous "until its entry is evicted", but every registration refreshes that entry's LRU position, so a digest that keeps being requested is never evicted and never recovers. Detail inline.

Two description items need a refresh before merge:

  • The description still says GPU evidence came from a 1.4.0rc0 runtime with adapters and that "matching-build validation ... remains required", while the thread reply reports the pipeline #66913 wheel passed conformance and two-worker HTTP routing. Update the text so the merge record reflects what was actually run.
  • tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py, tests/unittest/_torch/test_connector_adp.py, and tests/unittest/_torch/executor/test_pytorch_model_engine.py are changed but unrelated to image routing and unmentioned. If they fix tests that are broken on main, split them into their own PR; if main already has the fix, they should drop out on rebase. (The two stub-manager event_manager = None additions are different: those are required by the new isinstance check and belong here.)

Since the author account cannot apply labels, a maintainer needs to add api-compatible for the additive multimodal_uuids binding and the new BlockStored / MultimodalKey fields.

Comment thread tensorrt_llm/_torch/pyexecutor/kv_cache_events.py Outdated
Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com>
Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com>
@krishung5
krishung5 requested a review from a team as a code owner October 29, 2026 01:57
@krishung5

Copy link
Copy Markdown
Collaborator Author

Following up on Brian's review:

Removed the unrelated fixture and processor-default fixes from this PR in 1daf67a, preserving them in separate local branches for follow-up PRs. The two routing-required event_manager = None fixture changes remain.

Updated the KV-cache guide with optional UUID fields, the routing extension, and its cross-modality emission. Merged latest main in 3e3cc21 without conflicts; post-merge checks passed.

The description refresh is prepared locally but hasn't been published yet. A maintainer still needs to add api-compatible.

@krishung5

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #81298 [ run ] triggered by Bot. Commit: 3e3cc21 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #81298 [ run ] completed with state FAILURE. Commit: 3e3cc21
/LLM/main/L0_MergeRequest_PR pipeline #66986 completed with status: 'FAILURE'

CI Report

⚠️ Resource Maintenance — Test Coverage Reduced:

Resource maintenance in LLM/main/L0_Test-x86_64-Single-GPU #4150: L40S-PyTorch-2 is temporarily paused due to Kubernetes agent provisioning and node issues.
Affected stages:
L40S-PyTorch-2

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@lori-ren lori-ren left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This overlaps #19529, which adds the same multimodal_uuids binding and MultimodalKey.uuid field but carries the UUID through native block state (key[2] in drain_events). Agree on one mechanism with that PR before merging, so the streaming wire has a single source for uuid.

Comment on lines +326 to +334
Streaming V2 events can also include an optional `uuid` in each `mm_keys` entry.
It is a bounded routing hint; it does not replace the content-sensitive digest.
Missing, evicted, or conflicting hints omit `uuid`.
For multimodal blocks, the optional `BlockStored.trtllm_image_routing` extension
provides the routing format version and token-ID offset. The extension can accompany
video and audio keys as well as image keys; the OpenEngine capability flag advertises
only the currently validated image-routing support. See the
[OpenEngine image routing contract](source:tensorrt_llm/grpc/openengine/README.md#image-routing-metadata)
for capability requirements and consumer handling.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This contradicts the streaming "Multimodal payloads" paragraph further down (L444–450), which still says mm_keys uses only hash/start_offset and that multimodal-aware routing awaits #19529. Update that paragraph to match the new uuid and trtllm_image_routing contract.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated both sections to describe the digest hash, optional UUID, and version/token-offset routing metadata consistently. Removed the stale wording about revisiting the contract after the dependency merges.

return resolve_top_level_model_type(llm.input_processor.config)


def supports_image_routing(llm: Any, modalities: Collection[str]) -> bool:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add supports_image_routing to this module's __all__; it is public and imported by control.py.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added supports_image_routing to all.

Comment on lines +76 to +82
events = getattr(cache, "kv_events_config", None)
return (
getattr(llm.args, "cache_transceiver_config", None) is None
and getattr(cache, "use_kv_cache_manager_v2", False) is True
and getattr(cache, "enable_block_reuse", False) is True
and getattr(events, "enable_kv_cache_events", False) is True
and getattr(events, "publisher", None) == "zmq"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use events_config(llm) is not None from .kv_events here instead of re-deriving the ZMQ-events check, so the capability gate and event discovery cannot drift apart.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The capability gate now uses events_config(llm) is not None, sharing the same predicate as event discovery.

with self._uuid_lock:
if digest in self._uuids and self._uuids[digest] != uuid:
# Keep conflicting hints ambiguous until their bounded entry is evicted.
self._uuids[digest] = None

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The digest is BLAKE3 over <uuid>…</uuid><content>… (apply_mm_hashes), and only items with a UUID are registered, so one digest cannot carry two UUIDs short of a hash collision. Drop the ambiguity state, or describe it as collision defence only — including the README paragraph at grpc/openengine/README.md:144–147.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed the Python UUID table and its ambiguity state. Streaming UUIDs now come directly from #19529’s native context, while the engine digest remains the cache identity. Updated the README accordingly.

@krishung5
krishung5 requested a review from a team as a code owner October 29, 2026 19:58
@krishung5

Copy link
Copy Markdown
Collaborator Author

/bot run

@krishung5

Copy link
Copy Markdown
Collaborator Author

This branch now uses #19529’s native UUID transport instead of a parallel Python lookup table. The OpenEngine capability and normalization metadata remain on top. #19529 itself is unchanged, and this PR depends on it merging first. Post-merge validation passed 402 focused cases and 8/8 unpinned two-worker routing repeats; all 14 responses and usage counts matched the engines. Full merged-wheel CI remains pending.

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #81504 [ run ] triggered by Bot. Commit: e894e65 Link to invocation

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

8 participants