diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index e085367b143b..8454026f5791 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -238,6 +238,16 @@ void initBindings(nb::module_& m) } return hashes; }) + .def_prop_ro("multimodal_uuids", + [](GenLlmReq& self) + { + std::optional>> uuids = std::nullopt; + if (self.getMultimodalUuids()) + { + uuids = *self.getMultimodalUuids().value(); + } + return uuids; + }) .def_prop_ro("multimodal_positions", [](GenLlmReq& self) { diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index 3f46d098c863..684a5fd695a8 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -323,6 +323,16 @@ prompt = TextPrompt( V2 derives `mm_keys` directly from the cached token sequence. Each entry identifies a continuous multimodal segment within that block: `hash` is the item's digest, and `start_offset` is the segment's first token offset within the item. An item spanning multiple blocks retains the same digest with increasing offsets. Text may separate segments of the same item. Items are processed in prompt order; one item cannot resume after another item has started. The item digest is distinct from `block_hash`, which also depends on the preceding token sequence. +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. + ### Enable Offloading to Host Memory diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 8984856fd55e..0f76181101fa 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6126,6 +6126,16 @@ def _augment_tokens_for_block_reuse( "reuse." ) + if chunk_start == 0 and isinstance(self.event_manager, StreamingKVCacheEventManager): + uuids = req.multimodal_uuids + if uuids is not None: + uuid_pairs = [ + (_hash_to_digest(digest), uuid) + for digest, uuid in zip(req.multimodal_hashes, uuids, strict=True) + if uuid is not None + ] + self.event_manager.remember_multimodal_uuids(uuid_pairs) + # Multimodal path: materialize a Python-int list (digest bytes get spliced in below), # which flows through the per-element binding fallback. tokens may be a zero-copy numpy # int32 view (get_tokens_view) — use tolist() so elements are Python ints, not np.int32. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py index 29533afa3c37..d6063b5fe5f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_events.py @@ -29,10 +29,10 @@ import time import traceback from abc import ABC, abstractmethod -from collections import deque +from collections import OrderedDict, deque from itertools import count from queue import Queue -from typing import Any, Iterator, Optional +from typing import Any, Iterable, Iterator, Optional import msgspec import zmq @@ -80,6 +80,14 @@ class MultimodalKey( hash: str start_offset: int + uuid: str | None = None + + +class ImageRoutingMetadata(msgspec.Struct, frozen=True, gc=False): + """Versioned digest-first image normalization metadata for KV routers.""" + + version: int + token_id_offset: int class BlockStored(KVCacheWireEvent): @@ -99,6 +107,7 @@ class BlockStored(KVCacheWireEvent): locality: str | None = None # Aligned one-for-one with block_hashes when multimodal decoding is enabled. mm_keys: list[list[MultimodalKey]] | None = None + trtllm_image_routing: ImageRoutingMetadata | None = None class BlockRemoved(KVCacheWireEvent): @@ -551,6 +560,10 @@ def __init__( raise ValueError("mm_token_id_offset must be non-negative") self._block_size = block_size self._include_mm_keys = mm_token_id_offset is not None + self._mm_token_id_offset = mm_token_id_offset + self._uuid_capacity = max_entries + self._uuids: OrderedDict[bytes, str | None] = OrderedDict() + self._uuid_lock = threading.Lock() self._event_sink = kv_cache_manager_v2_runtime.StreamingEventSink( max_entries=max_entries, mm_token_id_offset=mm_token_id_offset, @@ -560,6 +573,27 @@ def __init__( def event_sink(self) -> kv_cache_manager_v2_runtime.StreamingEventSink: return self._event_sink + def remember_multimodal_uuids(self, items: Iterable[tuple[bytes, str]]) -> None: + """Retain bounded routing hints without changing native cache keys.""" + for digest, uuid in items: + if len(uuid.encode("utf-8")) > 256: + continue + 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 + else: + self._uuids[digest] = uuid + self._uuids.move_to_end(digest) + while len(self._uuids) > self._uuid_capacity: + self._uuids.popitem(last=False) + + def _multimodal_key(self, key: tuple) -> MultimodalKey: + digest = bytes(key[0]) + with self._uuid_lock: + uuid = self._uuids.get(digest) + return MultimodalKey(hash=digest.hex(), start_offset=key[1], uuid=uuid) + def set_target_life_cycle(self, life_cycle_id: int) -> None: self._event_sink.set_target_life_cycle(life_cycle_id) @@ -567,6 +601,14 @@ def drain_events(self) -> list[BlockStored | BlockRemoved | AllBlocksCleared]: result: list[BlockStored | BlockRemoved | AllBlocksCleared] = [] for event in self._event_sink.drain_iteration_events(): if isinstance(event, kv_cache_manager_v2_runtime.StreamingBlockStoredData): + mm_keys = ( + [ + [self._multimodal_key(key) for key in block_keys] + for block_keys in event.mm_keys + ] + if self._include_mm_keys + else None + ) result.append( BlockStored( block_hashes=list(event.block_hashes), @@ -576,15 +618,14 @@ def drain_events(self) -> list[BlockStored | BlockRemoved | AllBlocksCleared]: lora_id=event.lora_id, medium="GPU", lora_name=None, - mm_keys=( - [ - [ - MultimodalKey(hash=bytes(key[0]).hex(), start_offset=key[1]) - for key in block_keys - ] - for block_keys in event.mm_keys - ] - if self._include_mm_keys + mm_keys=mm_keys, + trtllm_image_routing=( + ImageRoutingMetadata( + version=1, token_id_offset=self._mm_token_id_offset + ) + if mm_keys is not None + and self._mm_token_id_offset is not None + and any(mm_keys) else None ), ) @@ -646,6 +687,10 @@ def start(self) -> None: """Bind the publisher's sockets and start its background thread.""" self._publisher.start() + def remember_multimodal_uuids(self, items: Iterable[tuple[bytes, str]]) -> None: + """Associate native MM digests with optional request identities.""" + self._event_source.remember_multimodal_uuids(items) + def get_source_descriptor(self) -> dict[str, str | int] | None: """Describe a live ZMQ publisher for OpenEngine node discovery.""" publisher = self._publisher diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 5ce68cf169c1..1ecae66b131a 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -73,7 +73,7 @@ The `Inference.Generate` RPC loads the selected model through TensorRT-LLM's PyT Clients must continuously consume the response stream. If response delivery remains stalled for 30 seconds, the server aborts the engine request and terminates the stream with a retryable overload error. -Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, multimodal media, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`. +Features without a faithful TensorRT-LLM mapping return `UNIMPLEMENTED`: prefix-cache bypass, LoRA lifecycle selection, explicit-token or all-vocabulary log-probability selection, nonzero prompt-logprob offsets, per-request grammar-backend selection, and priority metadata. Attention-DP rank hints are validated and used for first-turn placement. With conversation affinity on, the first hint establishes the conversation's rank, and later turns retain that rank even if their hints differ; the request must carry a stable conversation ID. `ServerInfo.extra.trtllm_supports_dp_rank_targeting` advertises this support. The AutoDeploy backend is rejected at startup until it supports request cancellation. `Control` implements `GetServerInfo`, `GetModelInfo`, `GetLoad`, `GetKvEventSources`, `Health` and `Abort`; its LoRA lifecycle and `SubscribeKvEvents` RPCs return `UNIMPLEMENTED`. ### Disaggregated serving @@ -118,6 +118,39 @@ Enable load reporting independently with `--grpc --grpc-protocol openengine --op For attention DP, each publisher binds `base_port + global_rank`. When event streaming is enabled, OpenEngine collects the live publisher descriptors at executor startup and advertises every rank, including remote nodes. Those hostnames and ports must be reachable from the subscriber. A missing publisher fails startup; an unavailable discovery snapshot fails the RPC rather than synthesizing socket addresses. PP and CP streaming remain unsupported. The extra placement collection is OpenEngine-only; native HTTP serving and OpenEngine with events disabled do not enable it. +### Image routing metadata + +`GetModelInfo.extra.trtllm_image_routing_version = 1` advertises image routing +metadata for aggregated Qwen3-VL with KV cache manager V2 explicitly enabled, +block reuse enabled, and ZMQ event publishing. Image passthrough separately +requires `supports_multimodal` and `extra.trtllm_supported_modalities` to advertise +image support. Set `use_kv_cache_manager_v2=True` explicitly: `"auto"` is +excluded from capability discovery even if the engine resolves it to V2. +Video, audio, and disaggregated image routing are not advertised. + +For image blocks, `BlockStored.trtllm_image_routing` contains +`version` and `token_id_offset` (the vocabulary size). The existing +`BlockStored.mm_keys` contains the keys, aligned with the stored blocks. Each key +retains its native content `hash` and `start_offset` and gains an optional `uuid`. +Keys are serialized only in `mm_keys`. The publisher emits this extension for +all multimodal keys, including video/audio keys; the name identifies the current +image-routing contract, not a modality filter. Only image routing has been +validated and advertised by OpenEngine. Native image tokens start with a digest +and continue with IDs above `token_id_offset`; a router can normalize those +tokens using identities it sent in `MediaItem.uuid`. Native block hashes, parent chains, and removals retain +their existing meaning. Cache hashing continues to include image content. + +The publisher retains at most `max_entries` digest-to-UUID associations and skips +UUIDs longer than 256 UTF-8 bytes. Conflicting hints mark a digest ambiguous until +its entry is evicted or the publisher restarts, without failing inference. +Every registration refreshes the entry's LRU position, including ambiguous entries, +so a frequently registered digest can remain ambiguous for the publisher's lifetime. +An agreeing registration does not restore a hint after a conflict. Missing, evicted, +or ambiguous hints omit the UUID while retaining the image-routing metadata. Consumers +must ignore unmapped image stores while continuing to process removals. Routing +identities must refer to stable media content. This extension does not change the +OpenEngine schema. + ### Node-local discovery OpenEngine exposes optional, versioned `ServerInfo.extra.trtllm_node` metadata without changing the protocol schema. It describes a shared engine incarnation (`engine_id`), local hostname (`node_id`), whether this is the leader, node count, KV block size, local publishing DP ranks, and the global rank-to-node ownership map. The engine ID matches `GetServerInfo.instance_id` on the serving ingress and node metadata listeners. Consumers can use this information to partition telemetry collection without introducing routing-framework-specific server options. diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 83a6f9979b51..e407448aebf7 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -67,6 +67,22 @@ def _loaded_model_type(llm: Any) -> str: return resolve_top_level_model_type(llm.input_processor.config) +def supports_image_routing(llm: Any, modalities: Collection[str]) -> bool: + """Whether the loaded model exposes the validated image KV-event contract.""" + processor_config = getattr(getattr(llm, "input_processor", None), "config", None) + if "image" not in modalities or resolve_top_level_model_type(processor_config) != "qwen3_vl": + return False + cache = getattr(llm.args, "kv_cache_config", None) + 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" + ) + + def resolve_media_capabilities( llm: Any, implemented_modalities: Collection[str] ) -> MediaCapabilities: diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index bbd5a9e824ae..287c079cc52b 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -35,7 +35,7 @@ openengine_pb2_grpc, server_pb2, ) -from .capabilities import supported_guides +from .capabilities import supported_guides, supports_image_routing from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError from .errors import AbortFailedError from .kv_events import data_parallel_size, events_config @@ -296,6 +296,8 @@ async def GetModelInfo( reasoning_parser=str(self._llm.args.reasoning_parser or ""), ) info.extra.update({"trtllm_supported_modalities": sorted(media_modalities)}) + if supports_image_routing(self._llm, media_modalities): + info.extra.update({"trtllm_image_routing_version": 1}) max_seq_len = _positive_int(self._llm.args.max_seq_len) if max_seq_len is not None: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 1bce3cdda018..1f941b2c63c3 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -69,6 +69,7 @@ l0_a10: # false advertisement by construction. The image case uses Qwen3-VL-2B. - unittest/grpc/openengine/test_capability_conformance.py - unittest/grpc/openengine/test_multimodal_conformance.py + - unittest/grpc/openengine/test_image_routing_conformance.py - unittest/disaggregated/test_coordinator_e2e.py - unittest/disaggregated/test_coordinator_worker.py - unittest/disaggregated/test_agent_multi_backends.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index b3e075c3ce14..385829969352 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -211,6 +211,7 @@ l0_cpu: - unittest/_torch/test_shared_tensor_cpu.py - unittest/grpc/openengine/test_frontend_lifecycle.py - unittest/grpc/openengine/test_control.py + - unittest/grpc/openengine/test_image_routing.py - unittest/grpc/openengine/test_multimodal.py - unittest/grpc/openengine/test_node.py - unittest/grpc/openengine/test_openengine.py diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py index 855289820381..41683b60e3a2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py @@ -52,6 +52,7 @@ def make_stub_manager( # for a one-model draft, which reads D prompt tokens past the target's end. mgr.reuse_match_backoff = reuse_match_backoff mgr.vocab_size = 32000 + mgr.event_manager = None mgr.conversation_manager = None mgr.kv_connector_manager = None mgr.kv_cache_map = {} diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py index b31abb20a596..e037230de177 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py @@ -23,6 +23,7 @@ def _make_manager(vocab_size: int): manager = KVCacheManagerV2.__new__(KVCacheManagerV2) manager.vocab_size = vocab_size + manager.event_manager = None return manager diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 23a7b27d7b15..e101b647be29 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -20,7 +20,10 @@ model_pb2, server_pb2, ) -from tensorrt_llm.grpc.openengine.capabilities import MediaCapabilities # noqa: E402 +from tensorrt_llm.grpc.openengine.capabilities import ( # noqa: E402 + MediaCapabilities, + supports_image_routing, +) from tensorrt_llm.grpc.openengine.control import OpenEngineControlServicer # noqa: E402 from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer # noqa: E402 from tensorrt_llm.llmapi import KVEventsConfig # noqa: E402 @@ -31,6 +34,40 @@ MODEL = "test-model" +@pytest.mark.parametrize("model_type", ["qwen3_vl", "qwen2_vl", "text"]) +@pytest.mark.parametrize("enabled", [False, True, "auto"]) +@pytest.mark.parametrize( + "block_reuse,events_enabled,publisher", + [(True, True, "zmq"), (False, True, "zmq"), (True, False, "zmq"), (True, True, "null")], +) +def test_image_routing_capability_requires_streaming_v2_and_validated_model( + model_type, enabled, block_reuse, events_enabled, publisher +): + llm = SimpleNamespace( + input_processor=SimpleNamespace(config=SimpleNamespace(model_type=model_type)), + args=SimpleNamespace( + kv_cache_config=SimpleNamespace( + use_kv_cache_manager_v2=enabled, + enable_block_reuse=block_reuse, + kv_events_config=SimpleNamespace( + enable_kv_cache_events=events_enabled, publisher=publisher + ), + ) + ), + ) + expected = ( + enabled is True + and model_type == "qwen3_vl" + and block_reuse + and events_enabled + and publisher == "zmq" + ) + assert supports_image_routing(llm, {"image"}) is expected + assert not supports_image_routing(llm, set()) + llm.args.cache_transceiver_config = object() + assert not supports_image_routing(llm, {"image"}) + + class _FakeHandle: def __init__(self, fail: bool = False) -> None: self.aborted = False @@ -189,6 +226,34 @@ async def test_model_info_advertises_loaded_image_support(monkeypatch): assert list(info.extra["trtllm_supported_modalities"]) == ["image"] +@pytest.mark.asyncio +@pytest.mark.parametrize("events_enabled", [False, True]) +async def test_model_info_advertises_image_routing_only_with_events(monkeypatch, events_enabled): + llm = _llm( + kv_cache_config=SimpleNamespace( + use_kv_cache_manager_v2=True, + enable_block_reuse=True, + tokens_per_block=32, + kv_events_config=SimpleNamespace( + enable_kv_cache_events=events_enabled, publisher="zmq" + ), + ) + ) + llm.input_processor = SimpleNamespace(config=SimpleNamespace(model_type="qwen3_vl")) + inference = OpenEngineInferenceServicer(llm, MODEL) + monkeypatch.setattr( + inference, "_media_capabilities", MediaCapabilities({"image": ""}) + ) + info = await OpenEngineControlServicer(llm, MODEL, inference).GetModelInfo( + model_pb2.GetModelInfoRequest(model=MODEL), FakeServicerContext() + ) + assert info.supports_multimodal + if events_enabled: + assert info.extra["trtllm_image_routing_version"] == 1 + else: + assert "trtllm_image_routing_version" not in info.extra + + @pytest.mark.asyncio async def test_model_info_leaves_the_context_window_unset_when_unknown(): """An unset limit means "unknown"; advertising zero would be a real claim.""" diff --git a/tests/unittest/grpc/openengine/test_image_routing.py b/tests/unittest/grpc/openengine/test_image_routing.py new file mode 100644 index 000000000000..cbbb321b285d --- /dev/null +++ b/tests/unittest/grpc/openengine/test_image_routing.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Image routing hints preserve engine identities and survive the event wire.""" + +from types import SimpleNamespace + +import msgspec +import pytest +from PIL import Image + +from tensorrt_llm._torch.pyexecutor import kv_cache_events as events +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.inputs.multimodal import apply_mm_hashes +from tensorrt_llm.llmapi.llm_args import KVEventsConfig + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("uuid", ["dynamo-image-v1:0123456789abcdef", "caller-owned"]) +def test_streaming_image_hints_are_bounded_and_content_sensitive(monkeypatch, uuid): + red = Image.new("RGB", (32, 32), "red") + blue = Image.new("RGB", (32, 32), "blue") + hashes = [ + apply_mm_hashes({"image": [image]}, {"image": [uuid]}).hashes["image"][0] + for image in (red, blue) + ] + assert hashes[0] != hashes[1], "routing hints must not replace content-sensitive cache keys" + + manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + source = manager._event_source + native_sink = source._event_sink + try: + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[0]), uuid)]) + + class Stored: + block_hashes = [123] + parent_block_hash = None + token_ids = [1, hashes[0], 1001, 7] + lora_id = None + mm_keys = [[(bytes.fromhex(hashes[0]), 0)]] + + monkeypatch.setattr(events.kv_cache_manager_v2_runtime, "StreamingBlockStoredData", Stored) + source._event_sink = SimpleNamespace(drain_iteration_events=lambda: [Stored()]) + wire = msgspec.msgpack.decode(msgspec.msgpack.encode(source.drain_events()[0])) + metadata = wire["trtllm_image_routing"] + assert metadata["version"] == 1 and metadata["token_id_offset"] == 1000 + assert set(metadata) == {"version", "token_id_offset"} + assert wire["mm_keys"][0][0]["uuid"] == uuid + assert wire["mm_keys"][0][0]["hash"] == hashes[0] + encoded = msgspec.msgpack.encode(source.drain_events()[0]) + assert encoded.count(hashes[0].encode()) == 2 # Token stream and mm_keys. + assert encoded.count(uuid.encode()) == 1 + assert wire["token_ids"] == Stored.token_ids + + Stored.token_ids = [1, 2, 3, 4] + Stored.mm_keys = [[]] + text_wire = msgspec.msgpack.decode(msgspec.msgpack.encode(source.drain_events()[0])) + assert "trtllm_image_routing" not in text_wire + + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[1]), uuid)]) + assert len(source._uuids) == 1 + assert source._multimodal_key((bytes.fromhex(hashes[0]), 0)).uuid is None + assert source._multimodal_key((bytes.fromhex(hashes[1]), 0)).uuid == uuid + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[0]), "x" * 257)]) + assert len(source._uuids) == 1 + manager.remember_multimodal_uuids([(bytes.fromhex(hashes[1]), "other")]) + assert source._multimodal_key((bytes.fromhex(hashes[1]), 0)).uuid is None + finally: + source._event_sink = native_sink + manager.shutdown() + + +@pytest.mark.parametrize( + "hashes,uuids,error", + [ + ([tuple(range(8)), tuple(range(8, 16))], ["image-a"], "shorter"), + ([tuple(range(8))], ["image-a", "image-b"], "longer"), + ([tuple(range(8)), (1, 2, 3)], ["image-a", "image-b"], "Expected 8 int32"), + ], +) +def test_invalid_request_uuid_pairs_leave_routing_hints_unchanged(hashes, uuids, error): + event_manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + try: + event_manager.remember_multimodal_uuids([(b"x" * 32, "existing")]) + before = event_manager._event_source._uuids.copy() + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + manager.vocab_size = 1000 + manager.event_manager = event_manager + request = SimpleNamespace( + multimodal_hashes=hashes, + multimodal_uuids=uuids, + multimodal_positions=[0] * len(hashes), + multimodal_lengths=[1] * len(hashes), + ) + with pytest.raises(ValueError, match=error): + manager._augment_tokens_for_block_reuse([1, 2], request) + assert event_manager._event_source._uuids == before + finally: + event_manager.shutdown() + + +def test_request_uuid_pairs_preserve_partial_uuid_slots(): + hashes = [tuple(range(8)), tuple(range(8, 16))] + event_manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + try: + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + manager.vocab_size = 1000 + manager.event_manager = event_manager + request = SimpleNamespace( + multimodal_hashes=hashes, + multimodal_uuids=["image-a", None], + multimodal_positions=[0, 1], + multimodal_lengths=[1, 1], + multimodal_item_run_cu_offsets=None, + multimodal_run_positions=None, + multimodal_run_lengths=None, + ) + manager._augment_tokens_for_block_reuse([1, 2], request) + digest = b"".join(value.to_bytes(4, "big", signed=True) for value in hashes[0]) + assert dict(event_manager._event_source._uuids) == {digest: "image-a"} + finally: + event_manager.shutdown() + + +def test_chunked_prefill_registers_hints_before_images_without_repeating(monkeypatch): + event_manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=2, + max_window_size=128, + mm_token_id_offset=1000, + ) + cache = KVCacheManagerV2.__new__(KVCacheManagerV2) + cache.vocab_size = 1000 + cache.event_manager = event_manager + request = SimpleNamespace( + multimodal_hashes=[tuple(range(8))], + multimodal_uuids=["image-a"], + multimodal_positions=[2], + multimodal_lengths=[3], + multimodal_item_run_cu_offsets=None, + multimodal_run_positions=None, + multimodal_run_lengths=None, + ) + registrations = [] + original = event_manager.remember_multimodal_uuids + + def remember(pairs): + registrations.append(pairs) + original(pairs) + + monkeypatch.setattr(event_manager, "remember_multimodal_uuids", remember) + digest = b"".join(value.to_bytes(4, "big", signed=True) for value in range(8)) + try: + tokens = [1, 2, 3, 4, 5, 6] + first = cache._augment_tokens_for_block_reuse(tokens, request, start=0, end=2) + assert len(registrations) == 1 + assert event_manager._event_source._multimodal_key((digest, 0)).uuid == "image-a" + # Later chunks cannot read or register the request's UUID vector again. + del request.multimodal_uuids + middle = cache._augment_tokens_for_block_reuse(tokens, request, start=2, end=4) + last = cache._augment_tokens_for_block_reuse(tokens, request, start=4, end=6) + assert len(registrations) == 1 + request.multimodal_uuids = ["image-a"] + expected = cache._augment_tokens_for_block_reuse(tokens, request) + assert first + middle + last == expected + finally: + event_manager.shutdown() + + +def test_conflicting_request_hints_preserve_inference_cache_keys(): + manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + cache = KVCacheManagerV2.__new__(KVCacheManagerV2) + cache.vocab_size = 1000 + cache.event_manager = manager + hashes = [tuple(range(8))] + digest = b"".join(value.to_bytes(4, "big", signed=True) for value in hashes[0]) + results = [] + try: + for uuid in ["image-a", "image-b", "image-a"]: + request = SimpleNamespace( + multimodal_hashes=hashes, + multimodal_uuids=[uuid], + multimodal_positions=[0], + multimodal_lengths=[1], + multimodal_item_run_cu_offsets=None, + multimodal_run_positions=None, + multimodal_run_lengths=None, + ) + results.append(cache._augment_tokens_for_block_reuse([1, 2], request)) + hint = manager._event_source._multimodal_key((digest, 0)).uuid + assert hint == ("image-a" if len(results) == 1 else None) + assert results[0] == results[1] == results[2] + finally: + manager.shutdown() + + +@pytest.mark.parametrize("hint_state", ["absent", "evicted", "ambiguous"]) +def test_unmapped_integer_image_continuations_keep_routing_metadata(monkeypatch, hint_state): + manager = events.StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + source = manager._event_source + native_sink = source._event_sink + digest = bytes(range(32)) + try: + if hint_state != "absent": + manager.remember_multimodal_uuids([(digest, "image-a")]) + if hint_state == "evicted": + manager.remember_multimodal_uuids([(bytes(reversed(digest)), "image-b")]) + else: + manager.remember_multimodal_uuids([(digest, "image-b")]) + + class Stored: + block_hashes = [123] + parent_block_hash = 122 + token_ids = [1003, 1004, 1005, 7] + lora_id = None + mm_keys = [[(digest, 3)]] + + monkeypatch.setattr(events.kv_cache_manager_v2_runtime, "StreamingBlockStoredData", Stored) + source._event_sink = SimpleNamespace(drain_iteration_events=lambda: [Stored()]) + wire = msgspec.msgpack.decode(msgspec.msgpack.encode(source.drain_events()[0])) + metadata = wire["trtllm_image_routing"] + assert metadata["version"] == 1 and metadata["token_id_offset"] == 1000 + assert set(metadata) == {"version", "token_id_offset"} + key = wire["mm_keys"][0][0] + assert key["hash"] == digest.hex() and key["start_offset"] == 3 + assert key.get("uuid") is None + assert wire["token_ids"] == Stored.token_ids + finally: + source._event_sink = native_sink + manager.shutdown() diff --git a/tests/unittest/grpc/openengine/test_image_routing_conformance.py b/tests/unittest/grpc/openengine/test_image_routing_conformance.py new file mode 100644 index 000000000000..7ffefc99f1cb --- /dev/null +++ b/tests/unittest/grpc/openengine/test_image_routing_conformance.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Real-weight image cache reuse and streaming routing-hint conformance.""" + +import asyncio +import base64 +from io import BytesIO + +import grpc +import msgspec +import pytest +import torch +import zmq +from PIL import Image +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM +from tensorrt_llm.grpc.openengine.bindings import generation_pb2 as pb +from tensorrt_llm.grpc.openengine.bindings import model_pb2, openengine_pb2_grpc +from tensorrt_llm.grpc.openengine.server import OpenEngineServer +from tensorrt_llm.llmapi import KvCacheConfig, KVEventsConfig + +pytestmark = pytest.mark.threadleak(enabled=False) + + +def test_real_image_routing_keeps_cache_identity_content_sensitive(unused_tcp_port, tmp_path): + assert torch.cuda.is_available(), "GPU conformance requires a CUDA device" + endpoint = f"tcp://127.0.0.1:{unused_tcp_port}" + context = zmq.Context() + subscriber = context.socket(zmq.SUB) + subscriber.setsockopt(zmq.SUBSCRIBE, b"") + subscriber.connect(endpoint) + llm = None + try: + llm = LLM( + model=str(llm_models_root(check=True) / "Qwen3/Qwen3-VL-2B-Instruct"), + max_batch_size=1, + max_seq_len=2048, + max_num_tokens=2048, + kv_cache_config=KvCacheConfig( + max_tokens=4096, + free_gpu_memory_fraction=0.3, + enable_block_reuse=True, + use_kv_cache_manager_v2=True, + kv_events_config=KVEventsConfig( + enable_kv_cache_events=True, publisher="zmq", endpoint=endpoint + ), + ), + ) + + async def run(): + server = OpenEngineServer(host="127.0.0.1", port=0, llm=llm, model="image-routing") + try: + await server.start() + async with grpc.aio.insecure_channel(f"127.0.0.1:{server.port}") as channel: + control = openengine_pb2_grpc.ControlStub(channel) + info = await control.GetModelInfo( + model_pb2.GetModelInfoRequest(model="image-routing") + ) + assert info.extra["trtllm_image_routing_version"] == 1 + inference = openengine_pb2_grpc.InferenceStub(channel) + + async def generate(color, request_id): + image = BytesIO() + Image.new("RGB", (224, 224), color).save(image, format="PNG") + request = pb.GenerateRequest( + model="image-routing", + request_id=request_id, + prompt=( + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Name the solid background color. Answer with one color word." + "<|im_end|>\n<|im_start|>assistant\n" + ), + sampling=pb.SamplingParams(temperature=0), + stopping=pb.StoppingOptions(max_tokens=16), + media=[ + pb.MediaItem( + modality=pb.MODALITY_IMAGE, + uuid="dynamo-image-v1:0123456789abcdef", + data_uri="data:image/png;base64," + + base64.b64encode(image.getvalue()).decode(), + ) + ], + ) + ids, usages, finished = [], [], 0 + async for response in inference.Generate(request, timeout=180): + assert not response.HasField("error"), response + if response.HasField("token"): + ids.extend(token.token_id for token in response.token.tokens) + if response.HasField("usage"): + usages.append(response.usage) + finished += int(response.HasField("finished")) + assert finished == 1 and usages + answer = llm.tokenizer.decode(ids).lower() + assert color in answer, answer + return usages[-1].cached_prompt_tokens + + first = await generate("red", "red-first") + repeated = await generate("red", "red-repeat") + changed = await generate("blue", "blue-same-hint") + assert repeated > first + assert changed < repeated, "changed content reused the first image's KV prefix" + finally: + await server.stop() + + asyncio.run(run()) + captured, hashes = [], set() + while subscriber.poll(1000): + frames = subscriber.recv_multipart() + captured.append(frames[-1]) + batch = msgspec.msgpack.decode(frames[-1]) + for event in batch[1]: + metadata = event.get("trtllm_image_routing") + if not metadata: + continue + for keys in event["mm_keys"]: + for key in keys: + assert key["uuid"] == "dynamo-image-v1:0123456789abcdef" + hashes.add(key["hash"]) + assert len(hashes) == 2, ( + "both content digests must retain the routing hint on the event wire" + ) + (tmp_path / "image-routing-events.msgpack").write_bytes(msgspec.msgpack.encode(captured)) + finally: + if llm is not None: + llm.shutdown() + subscriber.close(linger=0) + context.term() diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index 8198529e4069..90d3ec7508bc 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -393,6 +393,12 @@ def test_native_streaming_sink_preserves_multimodal_event_data(real_block_factor make_block = real_block_factory(event_sink, tokens_per_block=4) digest_a = bytes(range(32)) digest_b = bytes(reversed(range(32))) + manager.remember_multimodal_uuids( + [ + (digest_a, "dynamo-image-v1:0123456789abcdef"), + (digest_b, "dynamo-image-v1:fedcba9876543210"), + ] + ) first = make_block([1, digest_a, 1001, 1002], [4]) gap = make_block([2, 3, 4, 5], [4], parent=first) continued = make_block([1003, 7, 1004, 1005], [4], parent=gap) @@ -432,6 +438,44 @@ def test_native_streaming_sink_preserves_multimodal_event_data(real_block_factor [(digest_a.hex(), 6), (digest_b.hex(), 0)], ] assert manager.stored_blocks == 4 + assert stored.trtllm_image_routing.version == 1 + assert stored.trtllm_image_routing.token_id_offset == 1000 + assert [[key.uuid for key in keys] for keys in stored.mm_keys] == [ + ["dynamo-image-v1:0123456789abcdef"], + [], + ["dynamo-image-v1:0123456789abcdef"] * 2, + ["dynamo-image-v1:0123456789abcdef", "dynamo-image-v1:fedcba9876543210"], + ] + finally: + manager.shutdown() + + +def test_streaming_uuid_hints_are_bounded_and_do_not_replace_content_hashes(): + manager = StreamingKVCacheEventManager( + KVEventsConfig(enable_kv_cache_events=True, publisher="null"), + data_parallel_rank=0, + block_size=4, + max_window_size=128, + max_entries=1, + mm_token_id_offset=1000, + ) + a, b = bytes(range(32)), bytes(reversed(range(32))) + try: + manager.remember_multimodal_uuids([(a, "same-routing-hint"), (b, "same-routing-hint")]) + source = manager._event_source + assert len(source._uuids) == 1 + first = source._multimodal_key((a, 0)) + second = source._multimodal_key((b, 0)) + assert first.uuid is None + assert second.uuid == "same-routing-hint" + assert first.hash != second.hash + manager.remember_multimodal_uuids([(a, "x" * 257)]) + assert len(source._uuids) == 1 + manager.remember_multimodal_uuids([(b, "different-hint")]) + assert source._multimodal_key((b, 0)).uuid is None + manager.remember_multimodal_uuids([(b, "same-routing-hint")]) + assert source._multimodal_key((b, 0)).uuid is None + assert len(source._uuids) == 1 finally: manager.shutdown()