From 0866d6321976b01c47eac01d624506760dcbe095 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:45:57 -0700 Subject: [PATCH 01/22] [None][refactor] share multimodal input metadata helpers Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- tensorrt_llm/inputs/media_io.py | 6 +- tensorrt_llm/inputs/multimodal.py | 5 +- tensorrt_llm/inputs/multimodal_data.py | 6 +- tensorrt_llm/inputs/registry.py | 34 ++-- tests/unittest/inputs/test_multimodal.py | 246 ++++++++++++++++++++++- 5 files changed, 269 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 4620ee864da3..e33bd4ddf3ab 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -22,7 +22,6 @@ Dict, Generic, List, - Literal, Mapping, Optional, Tuple, @@ -41,7 +40,7 @@ from packaging.version import Version from PIL import Image -from tensorrt_llm.inputs.multimodal_data import AudioData, VideoData +from tensorrt_llm.inputs.multimodal_data import AudioData, MediaModality, VideoData from tensorrt_llm.logger import logger @@ -71,9 +70,6 @@ def convert_image_mode(image: Image.Image, to_mode: str) -> Image.Image: return image.convert(to_mode) -# Canonical set of supported media modalities for Pydantic field validation. -MediaModality = Literal["image", "video", "audio"] - # Output representations supported by `ImageMediaIO`: # `"pt"` -> `torch.Tensor`, `"np"` -> uint8 HWC `numpy.ndarray`, # `"pil"` -> `PIL.Image.Image`. diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 06bae3975faf..49f512f5b632 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -10,7 +10,8 @@ from blake3 import blake3 from tensorrt_llm._utils import maybe_pin_memory -from tensorrt_llm.inputs.multimodal_data import (BaseModalityData, VideoData, +from tensorrt_llm.inputs.multimodal_data import (MEDIA_MODALITIES, + BaseModalityData, VideoData, serialize_item) from tensorrt_llm.logger import logger @@ -176,7 +177,7 @@ def strip_mm_encoder_inputs(mm_data: Dict[str, Any]) -> None: raw modality dictionaries are no longer needed and must not be moved to GPU by the subsequent LLM input-preparation path. """ - for modality in ("image", "video", "audio"): + for modality in MEDIA_MODALITIES: mm_data.pop(modality, None) diff --git a/tensorrt_llm/inputs/multimodal_data.py b/tensorrt_llm/inputs/multimodal_data.py index 65895f9fa435..216df313ae09 100644 --- a/tensorrt_llm/inputs/multimodal_data.py +++ b/tensorrt_llm/inputs/multimodal_data.py @@ -3,12 +3,16 @@ import struct from dataclasses import dataclass -from typing import Any, Protocol +from typing import Any, Literal, Protocol, get_args import numpy as np import torch from PIL import Image +# Canonical media names for input metadata and I/O validation. +MediaModality = Literal["image", "video", "audio"] +MEDIA_MODALITIES: tuple[MediaModality, ...] = get_args(MediaModality) + # Video metadata fields that participate in the cache-key hash. These describe # how frames were sampled and therefore change the model-visible content. _VIDEO_HASH_METADATA_FIELDS = ( diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 9a0fa952eb18..df4de35ac857 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -41,7 +41,7 @@ _find_mm_token_start_pos_from_masks, apply_mm_hashes, default_hasher, find_mm_token_lengths, hexdigest_to_int32, validate_mm_inputs) -from .multimodal_data import serialize_item +from .multimodal_data import MEDIA_MODALITIES, serialize_item from .prefix_token_cache import create_prefix_token_cache from .tokenization import tokenize_prompt @@ -1243,6 +1243,14 @@ def _get_single_mm_token_lengths( return num_mm_tokens +def _cache_mm_embed_cumsum(multimodal_data: Dict[str, Any], + embed_mask: torch.Tensor) -> None: + """Cache the int64 prefix sum of a flat CPU embedding mask without replacing existing data.""" + if "multimodal_embed_mask_cumsum" not in multimodal_data: + multimodal_data["multimodal_embed_mask_cumsum"] = embed_mask.cumsum( + 0, dtype=torch.int64) + + def maybe_compute_mm_embed_cumsum( prompt_token_ids: List[int], extra_processed_inputs: Optional[ExtraProcessedInputs], @@ -1266,15 +1274,8 @@ def maybe_compute_mm_embed_cumsum( if mm_data is None: return - # Always backfill the bidirectional-MM gate first (cheap; instance attr - # falls through to class attr in Python). Must happen BEFORE the cumsum - # idempotency guard below, otherwise the path where - # `multimodal_hashing_process` already setdefault()ed the cumsum (without - # this flag) leaves scheduler_v2 reading mm_bidirectional_blocks=None and - # silently skipping align — which then trips - # get_attention_mask's `cached_token_lens == 0` assert on - # 26B/31B once an image block is split across chunks. Gemma4 sets this - # per-instance from text_config.use_bidirectional_attention. + # The scheduler needs the bidirectional-block flag even when hashing has + # already cached the cumsum. Propagate it before the idempotency check. mm_data.setdefault( "mm_bidirectional_blocks", bool(getattr(input_processor, "mm_bidirectional_blocks", False)), @@ -1299,8 +1300,7 @@ def maybe_compute_mm_embed_cumsum( mm_special_token_ids=input_processor.get_mm_special_token_ids(), ) # Cache the int64 cumsum; request-invariant, read once per chunk. - mm_data["multimodal_embed_mask_cumsum"] = embed_mask.cumsum( - 0, dtype=torch.int64) + _cache_mm_embed_cumsum(mm_data, embed_mask) def create_input_processor_with_hash( @@ -1414,9 +1414,7 @@ def multimodal_hashing_process( mm_token_ids=mm_ids, mm_special_token_ids=mm_special_token_ids, ) - extra_processed_inputs["multimodal_data"].setdefault( - "multimodal_embed_mask_cumsum", - embed_mask.cumsum(0, dtype=torch.int64)) + _cache_mm_embed_cumsum(multimodal_data, embed_mask) start_positions, start_special_token_positions = ( _find_mm_token_start_pos_from_masks(mm_mask, special_mask, num_mm_tokens)) @@ -1475,13 +1473,11 @@ def process_prompt_maybe_hash( # Any subset of the supported modalities is eligible for hashing. # ``None`` payloads are skipped so an empty bucket doesn't gate off # hashing for the modalities that are actually present. - _SUPPORTED_HASHING_MODALITIES = ('image', 'video', 'audio') modalities = [ m for m, d in inputs.get('multi_modal_data', {}).items() if d is not None ] - if modalities and all(m in _SUPPORTED_HASHING_MODALITIES - for m in modalities): + if modalities and all(m in MEDIA_MODALITIES for m in modalities): if input_processor.multimodal_hashing_supported is not None: use_multimodal_hashing = input_processor.multimodal_hashing_supported else: @@ -1551,7 +1547,7 @@ def input_processor_wrapper( "get_mm_encoder_item_metadata", None) has_raw_payload = any( isinstance(multimodal_data.get(modality), dict) - for modality in ("image", "video", "audio")) + for modality in MEDIA_MODALITIES) if not has_raw_payload or get_item_metadata is None: return prompt_token_ids, extra_processed_inputs diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index c1aa7bc3befa..74f18fbdca6b 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -2,12 +2,15 @@ # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Tests for MultimodalRuntimeData cumsum math and the flat-mask producer.""" +from types import SimpleNamespace +from typing import get_args from unittest.mock import patch import numpy as np import pytest import torch +from tensorrt_llm.inputs.media_io import MEDIA_IO_REGISTRY, MediaModality from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, DisaggPrefillMultimodalInputs, @@ -17,7 +20,7 @@ _find_mm_embedding_lengths_from_masks, find_mm_token_lengths, ) -from tensorrt_llm.inputs.multimodal_data import VideoData +from tensorrt_llm.inputs.multimodal_data import MEDIA_MODALITIES, AudioData, VideoData from tensorrt_llm.inputs.registry import ( MultimodalEncoderItemMetadata, create_input_processor_with_hash, @@ -27,6 +30,11 @@ pytestmark = pytest.mark.cpu_only +def test_media_modalities_match_io_and_validation_contract() -> None: + assert MEDIA_MODALITIES == get_args(MediaModality) == ("image", "video", "audio") + assert set(MEDIA_MODALITIES) == set(MEDIA_IO_REGISTRY) + + class _ItemMetadataFakeProcessor: multimodal_hashing_supported = False @@ -145,6 +153,19 @@ def get_mm_special_token_ids(self): ) +@pytest.mark.parametrize("cached_cumsum", [None, torch.tensor([0, 1], dtype=torch.int64)]) +def test_existing_mm_cumsum_preserves_storage_and_backfills_block_alignment( + cached_cumsum: torch.Tensor | None, +) -> None: + processor = SimpleNamespace(mm_bidirectional_blocks=True) + data = {"multimodal_embed_mask_cumsum": cached_cumsum} + + maybe_compute_mm_embed_cumsum([10, 101], {"multimodal_data": data}, processor) + + assert data["multimodal_embed_mask_cumsum"] is cached_cumsum + assert data["mm_bidirectional_blocks"] is True + + def test_tokenized_multimodal_overwrites_stale_embedding_lengths(): """Final expanded prompt layout wins over dummy-placeholder metadata.""" @@ -223,6 +244,229 @@ def get_mm_special_token_ids(self): return None +_CONTENT_HASH_SNAPSHOT = { + "image": [ + 1704823964, + 362764787, + 1556099772, + -86050102, + 551081767, + 1806967272, + -1857978322, + -1074313326, + ], + "video": [ + -839685430, + -760767140, + -2055508912, + -1577466732, + 1582489614, + 1138227298, + 1889150221, + 1556646618, + ], + "audio": [ + -1693994058, + 897664400, + 754889986, + -1726904048, + -495512501, + 952543145, + -87795905, + -419978698, + ], +} +_UUID_HASH_SNAPSHOT = { + "image": [ + 1777607066, + -1240710263, + -116192886, + 1440160746, + 120140059, + -2062281779, + 960024653, + 2041693752, + ], + "video": [ + 2070136680, + -523271134, + 1570526135, + 1209057230, + 1458965967, + 334689758, + 1673654996, + -1634814920, + ], + "audio": [ + -1172540155, + -1149723152, + 1994549875, + -2028680509, + 93572480, + 1551260061, + -1949652498, + -1123551599, + ], +} + + +@pytest.mark.parametrize("hashing_enabled", [False, True], ids=["fallback", "hashing"]) +@pytest.mark.parametrize("use_uuids", [False, True], ids=["content", "uuid"]) +@pytest.mark.parametrize( + "modalities,token_ids,explicit_ids,special_ids,expected_cumsum,expected_runs", + [ + pytest.param( + ["image"], + [10, 101, 102, 20], + None, + None, + [0, 1, 2, 2], + ([0, 1], [1], [2]), + id="out-of-vocab-image", + ), + pytest.param( + ["image"], + [10, 101, 102, 20], + [101, 102], + None, + [0, 1, 2, 2], + ([0, 1], [1], [2]), + id="explicit-image-tokens", + ), + pytest.param( + ["image"], + [10, 101, 200, 102, 20], + [101, 102], + [200], + [0, 1, 1, 2, 2], + ([0, 1], [1], [3]), + id="image-framing-token", + ), + pytest.param( + ["video"], + [10, 101, 20, 102, 30], + None, + None, + [0, 1, 1, 2, 2], + ([0, 2], [1, 3], [1, 1]), + id="video-interleaved-text", + ), + pytest.param( + ["audio"], + [10, 101, 102, 20], + [101, 102], + None, + [0, 1, 2, 2], + ([0, 1], [1], [2]), + id="audio", + ), + pytest.param( + ["video", "image", "audio"], + [10, 101, 102, 20, 101, 102, 30, 101, 102, 40], + [101, 102], + None, + [0, 1, 2, 2, 3, 4, 4, 5, 6, 6], + ([0, 1, 2, 3], [1, 4, 7], [2, 2, 2]), + id="mixed-prompt-order", + ), + ], +) +def test_multimodal_input_contract_snapshot( + modalities: list[str], + token_ids: list[int], + explicit_ids: list[int] | None, + special_ids: list[int] | None, + expected_cumsum: list[int], + expected_runs: tuple[list[int], list[int], list[int]], + hashing_enabled: bool, + use_uuids: bool, +) -> None: + """Pin framework metadata independently of model-specific media encoding.""" + item_refs = [(modality, 0) for modality in modalities] + num_items = len(item_refs) + + class Processor(_KwargsHashFakeProcessor): + multimodal_hashing_supported = hashing_enabled + mm_bidirectional_blocks = True + + def __call__(self, inputs: dict, sampling_params: object) -> tuple[list[int], dict]: + return list(token_ids), { + "multimodal_data": { + **{modality: {} for modality in modalities}, + "multimodal_embedding_lengths": [2] * num_items, + } + } + + def get_num_tokens_per_image(self, *, image: np.ndarray) -> int: + return 2 + len(special_ids or []) + + def get_num_tokens_per_video(self, *, video: list[np.ndarray], **kwargs: object) -> int: + return 2 + + def get_num_tokens_per_audio(self, *, audio: AudioData) -> int: + return 2 + + def get_mm_token_ids(self) -> torch.Tensor | None: + return torch.tensor(explicit_ids) if explicit_ids is not None else None + + def get_mm_special_token_ids(self) -> torch.Tensor | None: + return torch.tensor(special_ids) if special_ids is not None else None + + def get_mm_encoder_item_metadata( + self, prompt_token_ids: list[int], multimodal_data: dict + ) -> MultimodalEncoderItemMetadata: + return MultimodalEncoderItemMetadata(item_refs, [4] * num_items, [2] * num_items) + + media = { + "image": [np.array([1], dtype=np.uint8)], + "video": [VideoData(frames=[np.full((2, 2, 3), 2, dtype=np.uint8)], metadata={})], + "audio": [AudioData(samples=np.array([0.25, -0.5], dtype=np.float32), sample_rate=16000)], + } + inputs = { + "prompt": "unused", + "multi_modal_data": {key: value for key, value in media.items() if key in modalities}, + "mm_item_order": [ + {"modality": modality, "index": 0, "placeholder": f"<{modality}>"} + for modality in modalities + ], + } + if use_uuids: + inputs["multi_modal_uuids"] = {modality: [f"{modality}-id"] for modality in modalities} + + actual_token_ids, extra = create_input_processor_with_hash(Processor())(inputs, None) + + assert actual_token_ids == token_ids + data = extra["multimodal_data"] + cumsum = data["multimodal_embed_mask_cumsum"] + assert cumsum.dtype == torch.int64 + assert cumsum.device.type == "cpu" + assert cumsum.tolist() == expected_cumsum + assert data["mm_bidirectional_blocks"] is True + assert data["multimodal_embedding_lengths"] == [2] * num_items + assert data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] == MultimodalEncoderItemMetadata( + item_refs, [4] * num_items, [2] * num_items + ) + if not hashing_enabled: + assert "multimodal_input" not in extra + return + + mm_input = extra["multimodal_input"] + hashes = _UUID_HASH_SNAPSHOT if use_uuids else _CONTENT_HASH_SNAPSHOT + assert mm_input.multimodal_hashes == [hashes[modality] for modality in modalities] + assert mm_input.multimodal_uuids == ( + [f"{modality}-id" for modality in modalities] if use_uuids else None + ) + assert mm_input.multimodal_positions == [1 + 3 * index for index in range(num_items)] + assert mm_input.multimodal_lengths == [2 + len(special_ids or [])] * num_items + assert ( + mm_input.multimodal_item_run_cu_offsets, + mm_input.multimodal_run_positions, + mm_input.multimodal_run_lengths, + ) == expected_runs + if special_ids: + assert data["special_token_offsets"] == [1] + + def test_mm_processor_kwargs_hash_is_stable_for_canonical_values(): input_processor = create_input_processor_with_hash( _KwargsHashFakeProcessor(), From e44118dc9c7512ffc5a8bb8a297708e11d37ed53 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Wed, 14 Oct 2026 23:51:38 -0700 Subject: [PATCH 02/22] [TRTLLM-14477][refactor] unify multimodal item ordering Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../models/modeling_multimodal_mixin.py | 125 +++++++++++------- tensorrt_llm/inputs/registry.py | 14 +- .../test_multimodal_encoder_item_groups.py | 93 +++++++++++++ .../_torch/multimodal/test_encoder_group.py | 40 ++++-- tests/unittest/inputs/test_multimodal.py | 49 +++++++ 5 files changed, 259 insertions(+), 62 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 1389d1fa1e11..15e28bae0fdb 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -31,7 +31,7 @@ Optional, Sequence, Tuple, - Union, + cast, ) import torch @@ -39,7 +39,12 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm._utils import prefer_pinned -from tensorrt_llm.inputs.multimodal import MultimodalInput, MultimodalParams, MultimodalRuntimeData +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalInput, + MultimodalParams, + MultimodalRuntimeData, +) from tensorrt_llm.inputs.registry import ( MultimodalEncoderItemMetadata, get_multimodal_encoder_item_metadata, @@ -87,7 +92,8 @@ class EncoderGroup: Batching all items in a group into one encoder invocation amortizes fixed costs (kernel launches, dispatch) across items. The framework splits the output back per-modality and reorders it into prompt order - via each request's `mm_item_order` manifest. + via each request's `MultimodalEncoderItemMetadata.item_refs` manifest, + falling back to `mm_item_order` for legacy processors. Contract between `build_batched_input` and `encoder_fn`: @@ -100,7 +106,7 @@ class EncoderGroup: * `encoder_fn` must return one tensor whose rows correspond 1:1 to the input layout produced by `build_batched_input`, so the framework can split the output by `_lengths_by_modality` and reorder into prompt - order via each request's `mm_item_order` manifest. + order via each request's item manifest. """ modalities: Tuple[str, ...] @@ -127,9 +133,8 @@ def _lengths_by_modality( multimodal_params: List[MultimodalParams], modalities: Tuple[str, ...], ) -> Dict[str, List[int]]: - """Invert prompt-ordered `multimodal_embedding_lengths` (number of - embedding rows per item) into per-modality per-item lengths, matching - the per-modality item order used by `EncoderGroup.build_batched_input`. + """Invert prompt-ordered output lengths into the per-modality item order + used by `EncoderGroup.build_batched_input`. Glossary used below: @@ -148,26 +153,15 @@ def _lengths_by_modality( """ by_modality: Dict[str, List[int]] = {m: [] for m in modalities} for mp in multimodal_params: - flat = mp.multimodal_data.get("multimodal_embedding_lengths") or [] - if mp.mm_item_order: - for entry, length in zip(mp.mm_item_order, flat, strict=True): - if entry["modality"] in by_modality: - by_modality[entry["modality"]].append(length) + item_refs, lengths = _get_multimodal_item_layout(mp, modalities) + if not item_refs: continue - # Raw-prompt entrypoints (non chat-parsing) do not attach a manifest, - # so this is the single enforcement point that a >1-modality request - # must carry `mm_item_order` to make prompt-order reordering possible. - # Check the *request-level* modality set so multi-group models (each - # group with a single modality) still catch cross-group mixes — the - # group-local `modalities` tuple can have length 1 and never trip. - present = [m for m in _MM_DATA_INPUT_MODALITY_KEYS if mp.multimodal_data.get(m) is not None] - if len(present) > 1: - raise ValueError( - "Request with multiple modalities present " - f"({present}) must carry mm_item_order on MultimodalParams." - ) - if present and present[0] in by_modality: - by_modality[present[0]].extend(flat) + request_lengths: dict[str, list[tuple[int, int]]] = {m: [] for m in modalities} + for (modality, index), length in zip(item_refs, lengths, strict=True): + if modality in request_lengths: + request_lengths[modality].append((index, length)) + for modality, indexed_lengths in request_lengths.items(): + by_modality[modality].extend(length for _, length in sorted(indexed_lengths)) return by_modality @@ -190,10 +184,10 @@ def _reorder_embeds_by_manifest( modality, embedding = populated_modalities[0] prompt_order_matches_encoder = True for mp in multimodal_params: - manifest = mp.mm_item_order or _synthesize_single_modality_manifest( - mp, per_modality_embeds.keys() - ) - item_indices = [entry["index"] for entry in manifest if entry["modality"] == modality] + item_refs, _ = _get_multimodal_item_layout(mp, per_modality_embeds.keys()) + item_indices = [ + index for item_modality, index in item_refs if item_modality == modality + ] if item_indices != list(range(len(item_indices))): prompt_order_matches_encoder = False break @@ -208,19 +202,16 @@ def _reorder_embeds_by_manifest( } slices: List[torch.Tensor] = [] - # `entry["index"]` is per-request per-modality; advance a cursor to + # Item indices are per-request per-modality; advance a cursor to # translate it into a global item index within `per_modality_embeds`. per_modality_cursor: Dict[str, int] = {m: 0 for m in per_modality_embeds} for mp in multimodal_params: - manifest = mp.mm_item_order or _synthesize_single_modality_manifest( - mp, per_modality_embeds.keys() - ) + item_refs, _ = _get_multimodal_item_layout(mp, per_modality_embeds.keys()) req_counts: Dict[str, int] = {} - for entry in manifest: - m = entry["modality"] + for m, item_index in item_refs: if m not in per_modality_embeds: continue - i = per_modality_cursor[m] + entry["index"] + i = per_modality_cursor[m] + item_index starts = per_modality_row_starts[m] slices.append(per_modality_embeds[m][starts[i] : starts[i + 1]]) req_counts[m] = req_counts.get(m, 0) + 1 @@ -243,16 +234,36 @@ def _reorder_embeds_by_manifest( return torch.cat(slices, dim=0) -def _synthesize_single_modality_manifest( +def _get_multimodal_item_layout( mp: MultimodalParams, modalities: Iterable[str], -) -> List[Dict[str, Union[str, int]]]: - """Trivial manifest for requests with only one modality present.""" +) -> tuple[list[tuple[str, int]], list[int]]: + """Read prompt-order refs and output lengths from the item metadata. + + Legacy processors supply `mm_item_order` and embedding lengths instead. + Single-modality requests can omit the order; mixed requests cannot. + """ + metadata = get_multimodal_encoder_item_metadata(mp.multimodal_data) + if metadata is not None: + return metadata.item_refs, metadata.output_embedding_lengths + flat = mp.multimodal_data.get("multimodal_embedding_lengths") or [] - for m in modalities: - if mp.multimodal_data.get(m) is not None: - return [{"modality": m, "index": i} for i in range(len(flat))] - return [] + if mp.mm_item_order: + item_refs = cast( + list[tuple[str, int]], + [(entry["modality"], entry["index"]) for entry in mp.mm_item_order], + ) + return item_refs, flat + + present = [m for m in _MM_DATA_INPUT_MODALITY_KEYS if mp.multimodal_data.get(m) is not None] + if len(present) > 1: + raise ValueError( + "Request with multiple modalities present " + f"({present}) must carry item metadata or mm_item_order on MultimodalParams." + ) + if present and present[0] in modalities: + return [(present[0], i) for i in range(len(flat))], flat + return [], [] def encode_multimodal_by_groups( @@ -265,9 +276,9 @@ def encode_multimodal_by_groups( For each group present in the batch, one encoder call is issued over all items across all requests belonging to that group's modalities (arithmetic-intensity win). The output is split back per-modality using - the prompt-ordered `multimodal_embedding_lengths` already stashed on - `multimodal_data`, then reordered into each request's `mm_item_order` - prompt sequence. + the prompt-ordered output lengths in `MultimodalEncoderItemMetadata`, + then reordered by the same metadata's `item_refs`. Legacy processors + fall back to `multimodal_embedding_lengths` and `mm_item_order`. Shared entry point for both the aggregated (`MultimodalModelMixin`) and mm-encoder-only (`Qwen3VisionModelBase.forward`) paths so the ordering @@ -1421,7 +1432,25 @@ def _apply_metadata_slice( Models slice raw modality tensors in `build_multimodal_encoder_input`; the mixin owns the parallel per-item metadata slice so every model gets it identically. """ - source_lengths = source.multimodal_data["multimodal_embedding_lengths"] + metadata = get_multimodal_encoder_item_metadata(source.multimodal_data) + if metadata is not None: + local_counts: dict[str, int] = {} + item_refs = [] + for i in item_indices: + modality = metadata.item_refs[i][0] + local_index = local_counts.get(modality, 0) + item_refs.append((modality, local_index)) + local_counts[modality] = local_index + 1 + residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = ( + MultimodalEncoderItemMetadata( + item_refs, + [metadata.encoder_token_lengths[i] for i in item_indices], + [metadata.output_embedding_lengths[i] for i in item_indices], + ) + ) + source_lengths = metadata.output_embedding_lengths + else: + source_lengths = source.multimodal_data["multimodal_embedding_lengths"] residual.multimodal_data["multimodal_embedding_lengths"] = [ source_lengths[i] for i in item_indices ] diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index df4de35ac857..073c0c56cadf 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -71,10 +71,9 @@ class MultimodalEncoderItemMetadata(NamedTuple): example, its row range in ``pixel_values``), so an item can be sliced out for a single-item encoder call. - TODO(TRTLLM-14477): `MultimodalParams.mm_item_order` carries the same prompt-order - manifest for the mixed-modality full-request encode path. Single-source - the two representations at the input processor so one manifest serves - both the item-scheduling path and the full-request interleave. + Shared by item scheduling and full-request encoder output assembly. + `MultimodalParams.mm_item_order` remains a fallback for legacy processors + that do not emit this metadata. """ encoder_token_lengths: List[int] @@ -1563,6 +1562,13 @@ def input_processor_wrapper( f"payloads, got {type(item_metadata).__name__}") item_metadata.validate() + mm_item_order = inputs.get("mm_item_order") + if mm_item_order and item_metadata.item_refs != [ + (entry["modality"], entry["index"]) for entry in mm_item_order + ]: + raise ValueError("Multimodal encoder item references do not match " + "the prompt item order") + multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = item_metadata existing_embedding_lengths = multimodal_data.get( "multimodal_embedding_lengths") diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 08657636b709..1bfeb03cce11 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -29,6 +29,7 @@ ) HIDDEN = 4 +pytestmark = pytest.mark.cpu_only # Rows the encoder emits per patch-grid item. Distinct values catch splits that # happen to line up only for uniform sizes. ITEM_ROWS = {"image": [2, 3], "video": [4]} @@ -181,6 +182,98 @@ def test_mixed_modality_items_keep_scheduler_order(): torch.testing.assert_close(output, _expected_rows(request, item_idx)) +@pytest.mark.parametrize("modalities", [["video", "image", "image"], ["image", "video", "image"]]) +@pytest.mark.parametrize("legacy_order", [False, True]) +def test_full_and_item_encoder_paths_share_item_metadata( + modalities: list[str], legacy_order: bool +) -> None: + model = _GroupedEncoderModel() + request = _make_request(modalities) + metadata = request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] + if legacy_order: + request.mm_item_order = [ + {"modality": modality, "index": index} for modality, index in metadata.item_refs + ] + # Full-request encoding reads both order and lengths from the same + # metadata that drives the scheduler's item selection. + request.multimodal_data.pop("multimodal_embedding_lengths") + + full_output = model.encode_multimodal_inputs([request]) + item_outputs, _ = _run_items(model, [(request, i) for i in range(len(modalities))]) + expected = torch.cat([_expected_rows(request, i) for i in range(len(modalities))]) + + torch.testing.assert_close(full_output, expected, rtol=0, atol=0) + torch.testing.assert_close(torch.cat(item_outputs), expected, rtol=0, atol=0) + + +def test_metadata_is_rebased_for_partial_and_reordered_items() -> None: + model = _GroupedEncoderModel() + request = _make_request(["image", "video", "image"]) + original = request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] + + outputs, encoder_inputs = _run_items(model, [(request, 2), (request, 0)]) + + assert len(encoder_inputs) == 1 + residual = encoder_inputs[0][0] + metadata = residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] + assert metadata == MultimodalEncoderItemMetadata([("image", 0), ("image", 1)], [3, 2], [3, 2]) + assert residual.multimodal_data["multimodal_embedding_lengths"] == [3, 2] + assert request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] is original + torch.testing.assert_close(outputs[0], _expected_rows(request, 2), rtol=0, atol=0) + torch.testing.assert_close(outputs[1], _expected_rows(request, 0), rtol=0, atol=0) + + +@pytest.mark.parametrize("use_metadata", [False, True]) +def test_full_and_item_order_can_differ_from_raw_payload_order(use_metadata: bool) -> None: + model = _GroupedEncoderModel() + request = _make_request(["image", "image"]) + request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalEncoderItemMetadata( + [("image", 1), ("image", 0)], [3, 2], [3, 2] + ) + expected = torch.cat([_expected_rows(request, 0), _expected_rows(request, 1)]) + request.multimodal_data["multimodal_embedding_lengths"] = [3, 2] + if use_metadata: + item_outputs, _ = _run_items(model, [(request, 0), (request, 1)]) + torch.testing.assert_close(torch.cat(item_outputs), expected, rtol=0, atol=0) + else: + request.multimodal_data.pop(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) + request.mm_item_order = [ + {"modality": "image", "index": 1}, + {"modality": "image", "index": 0}, + ] + + torch.testing.assert_close(model.encode_multimodal_inputs([request]), expected, rtol=0, atol=0) + + +def test_vision_and_audio_groups_share_full_and_item_ordering() -> None: + model = _GroupedEncoderModel() + + def build_audio_input(params: list[MultimodalParams]) -> dict[str, torch.Tensor]: + return { + "pixel_values": torch.cat( + [mp.multimodal_data["audio"]["input_features"].flatten(0, 1) for mp in params] + ) + } + + model.mm_encoder_groups += (EncoderGroup(("audio",), model._encode, build_audio_input),) + request = _make_request(["video", "image"]) + audio = torch.stack([torch.full((2, HIDDEN), value) for value in (10.0, 20.0)]) + request.multimodal_data["audio"] = {"input_features": audio} + request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalEncoderItemMetadata( + [("video", 0), ("audio", 0), ("image", 0), ("audio", 1)], [4, 2, 2, 2], [4, 2, 2, 2] + ) + request.multimodal_data.pop("multimodal_embedding_lengths") + expected = torch.cat( + [_expected_rows(request, 0), audio[0], _expected_rows(request, 2), audio[1]] + ) + + full_output = model.encode_multimodal_inputs([request]) + item_outputs, _ = _run_items(model, [(request, i) for i in range(4)]) + + torch.testing.assert_close(full_output, expected, rtol=0, atol=0) + torch.testing.assert_close(torch.cat(item_outputs), expected, rtol=0, atol=0) + + def test_partial_selection_only_encodes_selected_items(): """The scheduler may pick a subset (the rest being cache hits or deferred); unselected items must not leak into the encoder input.""" diff --git a/tests/unittest/_torch/multimodal/test_encoder_group.py b/tests/unittest/_torch/multimodal/test_encoder_group.py index 9cbdecf7ac34..44df0a85e212 100644 --- a/tests/unittest/_torch/multimodal/test_encoder_group.py +++ b/tests/unittest/_torch/multimodal/test_encoder_group.py @@ -6,8 +6,7 @@ * ``_lengths_by_modality`` — invert prompt-ordered ``multimodal_embedding_lengths`` into per-modality per-item lists using ``mm_item_order``. -* ``_synthesize_single_modality_manifest`` — trivial manifest for single-modality - requests that don't carry an explicit one. +* ``_get_multimodal_item_layout`` — shared item metadata with legacy fallbacks. * ``_reorder_embeds_by_manifest`` — slice per-modality tensors and concat in each request's prompt-order manifest. """ @@ -18,10 +17,14 @@ import torch from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( + _get_multimodal_item_layout, _lengths_by_modality, _reorder_embeds_by_manifest, - _synthesize_single_modality_manifest, ) +from tensorrt_llm.inputs.multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY +from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata + +pytestmark = pytest.mark.cpu_only def _mp(*, mm_item_order=None, embedding_lengths=None, buckets=None): @@ -89,19 +92,36 @@ def test_single_modality_outside_group_is_skipped(self): assert _lengths_by_modality([mp], ("image",)) == {"image": []} -class TestSynthesizeSingleModalityManifest: +class TestMultimodalItemLayout: def test_image_only(self): mp = _mp(embedding_lengths=[10, 20, 30], buckets={"image": {}}) - assert _synthesize_single_modality_manifest(mp, ("image", "video")) == [ - {"modality": "image", "index": 0}, - {"modality": "image", "index": 1}, - {"modality": "image", "index": 2}, - ] + assert _get_multimodal_item_layout(mp, ("image", "video")) == ( + [("image", 0), ("image", 1), ("image", 2)], + [10, 20, 30], + ) def test_no_group_modality_returns_empty(self): # Request has audio but the group only covers image/video. mp = _mp(embedding_lengths=[10], buckets={"audio": {}}) - assert _synthesize_single_modality_manifest(mp, ("image", "video")) == [] + assert _get_multimodal_item_layout(mp, ("image", "video")) == ([], []) + + def test_item_metadata_owns_order_and_lengths(self) -> None: + metadata = MultimodalEncoderItemMetadata( + [("video", 0), ("image", 0), ("audio", 0)], [24, 16, 8], [6, 4, 2] + ) + mp = _mp( + mm_item_order=[{"modality": "image", "index": 0}], + embedding_lengths=[99], + buckets={"image": {}, "video": {}, "audio": {}}, + ) + mp.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata + + item_refs, lengths = _get_multimodal_item_layout(mp, ("image", "video")) + + assert item_refs is metadata.item_refs + assert lengths is metadata.output_embedding_lengths + assert _lengths_by_modality([mp], ("image", "video")) == {"image": [4], "video": [6]} + assert _lengths_by_modality([mp], ("audio",)) == {"audio": [2]} class TestReorderEmbedsByManifest: diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 74f18fbdca6b..b2c368ac88f8 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -126,6 +126,17 @@ def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): input_processor({"prompt": "unused"}, sampling_params=None) +def test_mm_item_metadata_retains_three_field_tuple_contract() -> None: + metadata = MultimodalEncoderItemMetadata([("image", 0)], [4], [2]) + + item_refs, encoder_lengths, output_lengths = metadata + + assert metadata._fields == ("item_refs", "encoder_token_lengths", "output_embedding_lengths") + assert item_refs == [("image", 0)] + assert encoder_lengths == [4] + assert output_lengths == [2] + + def test_maybe_compute_mm_embed_cumsum_populates_py_multimodal_data(): """Producer writes a flat int64 cumsum tensor at py_multimodal_data[multimodal_embed_mask_cumsum].""" @@ -310,6 +321,44 @@ def get_mm_special_token_ids(self): } +@pytest.mark.parametrize("hashing_enabled", [False, True]) +def test_mm_item_metadata_rejects_conflicting_prompt_order(hashing_enabled: bool) -> None: + class Processor(_KwargsHashFakeProcessor): + multimodal_hashing_supported = hashing_enabled + + def __call__(self, inputs: dict, sampling_params: object) -> tuple[list[int], dict]: + return [10, 101, 102, 20, 101, 102, 30], { + "multimodal_data": { + "image": {}, + "video": {}, + "multimodal_embedding_lengths": [2, 2], + } + } + + def get_num_tokens_per_video(self, *, video: list[np.ndarray], **kwargs: object) -> int: + return 2 + + def get_mm_encoder_item_metadata( + self, prompt_token_ids: list[int], multimodal_data: dict + ) -> MultimodalEncoderItemMetadata: + return MultimodalEncoderItemMetadata([("video", 0), ("image", 0)], [4, 4], [2, 2]) + + processor = Processor() + inputs = { + "prompt": "unused", + "multi_modal_data": { + "image": [np.array([1], dtype=np.uint8)], + "video": [VideoData(frames=[np.full((2, 2, 3), 2, dtype=np.uint8)], metadata={})], + }, + "mm_item_order": [{"modality": "image", "index": 0}, {"modality": "video", "index": 0}], + } + + with pytest.raises(ValueError, match="item references do not match the prompt item order"): + create_input_processor_with_hash(processor)(inputs, None) + + assert processor.multimodal_hashing_supported is hashing_enabled + + @pytest.mark.parametrize("hashing_enabled", [False, True], ids=["fallback", "hashing"]) @pytest.mark.parametrize("use_uuids", [False, True], ids=["content", "uuid"]) @pytest.mark.parametrize( From d735f8f8d7cd74efc193472263cbe26a2045a9b9 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:50:25 -0700 Subject: [PATCH 03/22] [TRTLLM-14477][refactor] centralize multimodal item metadata Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_mistral.py | 6 +- .../models/modeling_multimodal_mixin.py | 121 ++++----------- .../_torch/models/modeling_qwen2vl.py | 2 +- .../_torch/pyexecutor/engine/multimodal.py | 2 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 2 +- tensorrt_llm/inputs/__init__.py | 5 +- tensorrt_llm/inputs/multimodal.py | 140 +++++++++++++++++- tensorrt_llm/inputs/registry.py | 87 ++--------- .../test_multimodal_encoder_item_groups.py | 43 ++++-- .../_torch/multimodal/test_encoder_group.py | 72 +++++---- .../test_qwen3vl_mixed_modality_encoder.py | 6 +- tests/unittest/inputs/test_multimodal.py | 42 +++++- 12 files changed, 298 insertions(+), 230 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 740e62ebc1b1..6648dc896454 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -50,9 +50,9 @@ MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) -from tensorrt_llm.inputs.multimodal import MultimodalParams -from tensorrt_llm.inputs.registry import (MULTIMODAL_PLACEHOLDER_REGISTRY, - MultimodalEncoderItemMetadata) +from tensorrt_llm.inputs.multimodal import (MultimodalEncoderItemMetadata, + MultimodalParams) +from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY from tensorrt_llm.inputs.utils import encode_base64_image from tensorrt_llm.llmapi import SamplingParams from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 15e28bae0fdb..c811f0e0b78b 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -31,7 +31,6 @@ Optional, Sequence, Tuple, - cast, ) import torch @@ -41,14 +40,11 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalEncoderItemMetadata, MultimodalInput, MultimodalParams, MultimodalRuntimeData, ) -from tensorrt_llm.inputs.registry import ( - MultimodalEncoderItemMetadata, - get_multimodal_encoder_item_metadata, -) from tensorrt_llm.logger import logger from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs @@ -130,30 +126,17 @@ class EncoderGroup: def _lengths_by_modality( - multimodal_params: List[MultimodalParams], + item_layouts: Sequence[tuple[list[tuple[str, int]], list[int]]], modalities: Tuple[str, ...], ) -> Dict[str, List[int]]: - """Invert prompt-ordered output lengths into the per-modality item order - used by `EncoderGroup.build_batched_input`. - - Glossary used below: - - * **group** — one `EncoderGroup` entry (a single tuple of modalities - that share an encoder call). E.g. Qwen3-VL registers one group - `("image", "video")`; Nano registers three groups - `("image",)`, `("video",)`, `("audio",)`. - * **group-local modalities** — the `modalities` tuple of the group - whose encoder is being invoked (this function's `modalities` arg). - * **cross-group mix** — a single request whose items span modalities - belonging to two *different* groups (e.g. one image + one audio - item for Nano, where image and audio live in separate groups). - * **multi-group model** — a model that registers more than one - `EncoderGroup`, so its group-local `modalities` tuples are each - strictly narrower than the full request-level modality set. + """Split resolved output lengths into each group's payload-local order. + + Layouts describe whole requests, so only modalities in this encoder + group contribute lengths (for example, image/video in a vision group + while the same request's audio is handled by another group). """ by_modality: Dict[str, List[int]] = {m: [] for m in modalities} - for mp in multimodal_params: - item_refs, lengths = _get_multimodal_item_layout(mp, modalities) + for item_refs, lengths in item_layouts: if not item_refs: continue request_lengths: dict[str, list[tuple[int, int]]] = {m: [] for m in modalities} @@ -166,7 +149,7 @@ def _lengths_by_modality( def _reorder_embeds_by_manifest( - multimodal_params: List[MultimodalParams], + item_layouts: Sequence[tuple[list[tuple[str, int]], list[int]]], per_modality_embeds: Dict[str, torch.Tensor], per_modality_lengths: Dict[str, List[int]], ) -> torch.Tensor: @@ -183,8 +166,7 @@ def _reorder_embeds_by_manifest( if len(populated_modalities) == 1: modality, embedding = populated_modalities[0] prompt_order_matches_encoder = True - for mp in multimodal_params: - item_refs, _ = _get_multimodal_item_layout(mp, per_modality_embeds.keys()) + for item_refs, _ in item_layouts: item_indices = [ index for item_modality, index in item_refs if item_modality == modality ] @@ -205,8 +187,7 @@ def _reorder_embeds_by_manifest( # Item indices are per-request per-modality; advance a cursor to # translate it into a global item index within `per_modality_embeds`. per_modality_cursor: Dict[str, int] = {m: 0 for m in per_modality_embeds} - for mp in multimodal_params: - item_refs, _ = _get_multimodal_item_layout(mp, per_modality_embeds.keys()) + for item_refs, _ in item_layouts: req_counts: Dict[str, int] = {} for m, item_index in item_refs: if m not in per_modality_embeds: @@ -234,38 +215,6 @@ def _reorder_embeds_by_manifest( return torch.cat(slices, dim=0) -def _get_multimodal_item_layout( - mp: MultimodalParams, - modalities: Iterable[str], -) -> tuple[list[tuple[str, int]], list[int]]: - """Read prompt-order refs and output lengths from the item metadata. - - Legacy processors supply `mm_item_order` and embedding lengths instead. - Single-modality requests can omit the order; mixed requests cannot. - """ - metadata = get_multimodal_encoder_item_metadata(mp.multimodal_data) - if metadata is not None: - return metadata.item_refs, metadata.output_embedding_lengths - - flat = mp.multimodal_data.get("multimodal_embedding_lengths") or [] - if mp.mm_item_order: - item_refs = cast( - list[tuple[str, int]], - [(entry["modality"], entry["index"]) for entry in mp.mm_item_order], - ) - return item_refs, flat - - present = [m for m in _MM_DATA_INPUT_MODALITY_KEYS if mp.multimodal_data.get(m) is not None] - if len(present) > 1: - raise ValueError( - "Request with multiple modalities present " - f"({present}) must carry item metadata or mm_item_order on MultimodalParams." - ) - if present and present[0] in modalities: - return [(present[0], i) for i in range(len(flat))], flat - return [], [] - - def encode_multimodal_by_groups( mm_encoder_groups: Sequence["EncoderGroup"], multimodal_params: List[MultimodalParams], @@ -286,23 +235,25 @@ def encode_multimodal_by_groups( """ per_modality_embeds: Dict[str, torch.Tensor] = {} per_modality_lengths: Dict[str, List[int]] = {} + item_layouts = [mp.get_item_layout() for mp in multimodal_params] for group in mm_encoder_groups: - group_params = [ - mp - for mp in multimodal_params - if any(mp.multimodal_data.get(m) is not None for m in group.modalities) - ] + group_params = [] + group_layouts = [] + for mp, layout in zip(multimodal_params, item_layouts, strict=True): + if any(mp.multimodal_data.get(m) is not None for m in group.modalities): + group_params.append(mp) + group_layouts.append(layout) if not group_params: continue out = group.encoder_fn(**group.build_batched_input(group_params)) - lengths = _lengths_by_modality(group_params, group.modalities) + lengths = _lengths_by_modality(group_layouts, group.modalities) cursor = 0 for m in group.modalities: total = sum(lengths[m]) per_modality_embeds[m] = out[cursor : cursor + total] cursor += total per_modality_lengths.update(lengths) - return _reorder_embeds_by_manifest(multimodal_params, per_modality_embeds, per_modality_lengths) + return _reorder_embeds_by_manifest(item_layouts, per_modality_embeds, per_modality_lengths) if TYPE_CHECKING: @@ -579,9 +530,7 @@ def _runs_by_request_modality( metadata = ( run_metadata if multimodal_param is run_param - else get_multimodal_encoder_item_metadata( - multimodal_param.multimodal_data or {} - ) + else multimodal_param.item_metadata ) except (TypeError, ValueError) as error: raise MultimodalEncoderContractError(str(error)) from error @@ -1273,7 +1222,7 @@ def _encoder_cache_keys( ) return None - item_metadata = get_multimodal_encoder_item_metadata(mm_data) + item_metadata = param.item_metadata if item_metadata is not None: return cls.build_encoder_cache_item_keys( mm_input.multimodal_hashes, @@ -1432,28 +1381,18 @@ def _apply_metadata_slice( Models slice raw modality tensors in `build_multimodal_encoder_input`; the mixin owns the parallel per-item metadata slice so every model gets it identically. """ - metadata = get_multimodal_encoder_item_metadata(source.multimodal_data) + metadata = source.item_metadata if metadata is not None: - local_counts: dict[str, int] = {} - item_refs = [] - for i in item_indices: - modality = metadata.item_refs[i][0] - local_index = local_counts.get(modality, 0) - item_refs.append((modality, local_index)) - local_counts[modality] = local_index + 1 - residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = ( - MultimodalEncoderItemMetadata( - item_refs, - [metadata.encoder_token_lengths[i] for i in item_indices], - [metadata.output_embedding_lengths[i] for i in item_indices], - ) + selected = metadata.select(item_indices) + residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = selected + residual.multimodal_data["multimodal_embedding_lengths"] = ( + selected.output_embedding_lengths ) - source_lengths = metadata.output_embedding_lengths else: source_lengths = source.multimodal_data["multimodal_embedding_lengths"] - residual.multimodal_data["multimodal_embedding_lengths"] = [ - source_lengths[i] for i in item_indices - ] + residual.multimodal_data["multimodal_embedding_lengths"] = [ + source_lengths[i] for i in item_indices + ] if residual.multimodal_input is not None and source.multimodal_input is not None: source_hashes = source.multimodal_input.multimodal_hashes residual.multimodal_input.multimodal_hashes = [source_hashes[i] for i in item_indices] diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index f4848ae64f06..8ba5aebb2f8f 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -31,8 +31,8 @@ from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, + MultimodalEncoderItemMetadata, MultimodalParams) -from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata from ..._utils import async_tensor_h2d, nvtx_range from ...inputs import (BaseMultimodalDummyInputsBuilder, diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index be9d642dcfd3..c7dd344fea08 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -27,13 +27,13 @@ from tensorrt_llm.inputs.multimodal import ( MultimodalParams, MultimodalRuntimeData, + get_multimodal_encoder_item_metadata, strip_mm_encoder_inputs, ) from tensorrt_llm.inputs.registry import ( BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor, InputProcessor, - get_multimodal_encoder_item_metadata, ) from tensorrt_llm.llmapi.llm_args import ( MultimodalEncoderCudaGraphConfig, diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 200ca48b37be..3c4a1706be9b 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -15,7 +15,7 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.bindings import executor as tllm_executor from tensorrt_llm.executor.result import SimpleTokenLogprobs, TokenLogprobs -from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata +from tensorrt_llm.inputs.multimodal import get_multimodal_encoder_item_metadata from tensorrt_llm.sampling_params import LogprobMode SamplingConfig = tllm_executor.SamplingConfig diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py index 4d49dd8038d6..bd6559311765 100644 --- a/tensorrt_llm/inputs/__init__.py +++ b/tensorrt_llm/inputs/__init__.py @@ -6,14 +6,13 @@ from .evs import (compute_retained_tokens_count, compute_retained_tokens_from_tubelet_budget, compute_retention_mask) -from .multimodal import MultimodalInput +from .multimodal import MultimodalEncoderItemMetadata, MultimodalInput from .multimodal_data import AudioData, BaseModalityData, VideoData # yapf and isort conflict on the following import blocks # yapf: disable from .registry import (BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor, ExtraProcessedInputs, - InputProcessor, MultimodalEncoderItemMetadata, - MultimodalPlaceholderMetadata, + InputProcessor, MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, create_input_processor, create_input_processor_with_hash, maybe_compute_mm_embed_cumsum, register_input_processor, diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 49f512f5b632..d7b201ea6e5b 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -1,9 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Multimodal utilities for handling images and other media types in TensorRT-LLM.""" +"""Common multimodal data contracts and layout operations. + +`MultimodalEncoderItemMetadata` describes request items. `MultimodalParams` +references processed payloads and the `MultimodalRuntimeData` view for the +current forward. `MultimodalInput` carries the executor boundary representation. +Processor registration and encoder execution live outside this module. +""" from dataclasses import dataclass, field -from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union +from typing import (Any, Dict, List, NamedTuple, Optional, Sequence, Tuple, + Union, cast) import numpy as np import torch @@ -25,6 +32,89 @@ _HASH_SCHEME_TAG = b"trtllm.mm.hash.v1" +class MultimodalEncoderItemMetadata(NamedTuple): + """Prompt-ordered metadata shared by item scheduling and encoder assembly.""" + + item_refs: List[Tuple[str, int]] + """``(modality, local index)`` per atomic item, in prompt order. + + The local index identifies the item within that modality's processed + payload. Legacy processors can supply `MultimodalParams.mm_item_order` + instead. + """ + + encoder_token_lengths: List[int] + """Physical attention-token cost charged to ``encoder_max_num_tokens``. + + For Qwen ViTs this counts patch tokens before the spatial merger. + """ + + output_embedding_lengths: List[int] + """Output embedding rows per item, matching ``multimodal_embedding_lengths``. + + These lengths split encoder outputs; they can differ from encoder cost + and from prompt token counts when only part of a placeholder is embedded. + """ + + def validate(self) -> None: + """Enforce the cross-field contract on producer-supplied metadata.""" + if not all(isinstance(values, list) for values in self): + raise TypeError( + "Multimodal encoder item metadata fields must be lists") + lengths = self.encoder_token_lengths + self.output_embedding_lengths + if not all(isinstance(length, int) for length in lengths): + raise TypeError( + "Multimodal encoder item lengths must contain only integers") + if not (len(self.item_refs) == len(self.encoder_token_lengths) == len( + self.output_embedding_lengths)): + raise ValueError( + "Multimodal encoder item references and lengths must align") + if any(length <= 0 for length in self.encoder_token_lengths): + raise ValueError( + "Multimodal encoder token lengths must be positive") + if any(length <= 0 for length in self.output_embedding_lengths): + raise ValueError( + "Multimodal encoder output embedding lengths must be positive") + + def select(self, + item_indices: Sequence[int]) -> "MultimodalEncoderItemMetadata": + """Select prompt-order items and rebase refs to their sliced payloads. + + Each modality's payload must be sliced in the same selected-item order. + The source metadata remains owned by the original request. + """ + local_counts: dict[str, int] = {} + item_refs = [] + for i in item_indices: + modality = self.item_refs[i][0] + local_index = local_counts.get(modality, 0) + item_refs.append((modality, local_index)) + local_counts[modality] = local_index + 1 + return MultimodalEncoderItemMetadata( + item_refs, + [self.encoder_token_lengths[i] for i in item_indices], + [self.output_embedding_lengths[i] for i in item_indices], + ) + + +def get_multimodal_encoder_item_metadata( + multimodal_data: Optional[Dict[str, Any]], +) -> Optional[MultimodalEncoderItemMetadata]: + """Return typed atomic-item metadata from Python multimodal data.""" + if multimodal_data is None: + return None + if not isinstance(multimodal_data, dict): + raise TypeError("multimodal_data must be a dict") + metadata = multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) + if metadata is None: + return None + if not isinstance(metadata, MultimodalEncoderItemMetadata): + raise TypeError(f"{MULTIMODAL_ENCODER_ITEM_METADATA_KEY} must be a " + "MultimodalEncoderItemMetadata") + metadata.validate() + return metadata + + def _validate_int_list(values: Any, field_name: str) -> None: """Boundary metadata must be owned Python list[int]. No tensor/tuple.""" if not isinstance(values, list): @@ -468,10 +558,12 @@ def __post_init__(self): @dataclass class MultimodalParams: - """Unified container for multimodal parameters. + """Processed multimodal inputs and views used by model execution. - This class encapsulates all multimodal-related data that flows through the system, - providing a clean interface for handling multimodal inputs across different models. + `item_metadata` reads the request's shared item record from processed data; + `get_item_layout` resolves its prompt order, including legacy formats. + Runtime data describes the current forward or chunk rather than encoder + readiness and cache ownership, which remain on the executor request. Attributes: multimodal_input: Multimodal input data with hashing information. @@ -537,6 +629,44 @@ def __post_init__(self): if self.multimodal_data is None: self.multimodal_data = {} + @property + def item_metadata(self) -> Optional[MultimodalEncoderItemMetadata]: + """Atomic-item metadata stored in this request's processed data.""" + return get_multimodal_encoder_item_metadata(self.multimodal_data) + + def get_item_layout(self) -> tuple[list[tuple[str, int]], list[int]]: + """Resolve prompt-order item refs and output lengths for all modalities. + + Typed metadata takes precedence. Legacy processors supply + ``mm_item_order`` and ``multimodal_embedding_lengths``; a request with + just one modality can omit the order. + """ + metadata = self.item_metadata + if metadata is not None: + return metadata.item_refs, metadata.output_embedding_lengths + + lengths = self.multimodal_data.get("multimodal_embedding_lengths") or [] + if self.mm_item_order: + item_refs = cast( + list[tuple[str, int]], + [(entry["modality"], entry["index"]) + for entry in self.mm_item_order], + ) + return item_refs, lengths + + present = [ + m for m in MEDIA_MODALITIES + if self.multimodal_data.get(m) is not None + ] + if len(present) > 1: + raise ValueError( + "Request with multiple modalities present " + f"({present}) must carry item metadata or mm_item_order on MultimodalParams." + ) + if present: + return [(present[0], i) for i in range(len(lengths))], lengths + return [], [] + def _is_shared_tensor_dict(self, obj: Any) -> bool: """Check if an object is a shared tensor dictionary. diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 073c0c56cadf..3675fa7721b8 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -20,8 +20,8 @@ import traceback from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import (Any, Callable, ClassVar, Dict, List, Mapping, NamedTuple, - Optional, Protocol, Tuple, Type, TypeVar, Union) +from typing import (Any, Callable, ClassVar, Dict, List, Mapping, Optional, + Protocol, Tuple, Type, TypeVar, Union) import numpy as np import torch @@ -34,13 +34,17 @@ from ..sampling_params import SamplingParams from .content_format import ContentFormat from .data import TextPrompt -from .multimodal import (MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalInput, - _as_cpu_tensor, _compute_mm_masks, +from .multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY +from .multimodal import \ + MultimodalEncoderItemMetadata as MultimodalEncoderItemMetadata +from .multimodal import (MultimodalInput, _as_cpu_tensor, _compute_mm_masks, _find_mm_embedding_lengths_from_masks, _find_mm_token_runs_from_mask, _find_mm_token_start_pos_from_masks, apply_mm_hashes, - default_hasher, find_mm_token_lengths, - hexdigest_to_int32, validate_mm_inputs) + default_hasher, find_mm_token_lengths) +from .multimodal import \ + get_multimodal_encoder_item_metadata as get_multimodal_encoder_item_metadata +from .multimodal import hexdigest_to_int32, validate_mm_inputs from .multimodal_data import MEDIA_MODALITIES, serialize_item from .prefix_token_cache import create_prefix_token_cache from .tokenization import tokenize_prompt @@ -60,77 +64,6 @@ def _hash_mm_processor_kwargs(mm_processor_kwargs: Dict[str, Any], return hasher.hexdigest() -class MultimodalEncoderItemMetadata(NamedTuple): - """Prompt-ordered metadata for atomic multimodal encoder items.""" - - item_refs: List[Tuple[str, int]] - """``(modality, local index)`` per atomic item, in prompt order. - - ``modality`` is ``"image"``/``"video"``/``"audio"`` and the local index - identifies the item within that modality's processed payload (for - example, its row range in ``pixel_values``), so an item can be sliced - out for a single-item encoder call. - - Shared by item scheduling and full-request encoder output assembly. - `MultimodalParams.mm_item_order` remains a fallback for legacy processors - that do not emit this metadata. - """ - - encoder_token_lengths: List[int] - """Physical encoder attention-token cost of each item. - - This is what one encoder forward actually attends over (for example, - pre-merger patch tokens for Qwen ViTs) and is the unit the scheduler - charges against ``encoder_max_num_tokens``. - """ - - output_embedding_lengths: List[int] - """Embedding rows each item contributes to the LLM prompt. - - Must equal the item's placeholder span in the prompt (the - ``multimodal_embedding_lengths`` contract); used to size and split the - encoder output back into per-item tensors. - """ - - def validate(self) -> None: - """Enforce the cross-field contract on producer-supplied metadata.""" - if not all(isinstance(values, list) for values in self): - raise TypeError( - "Multimodal encoder item metadata fields must be lists") - lengths = (self.encoder_token_lengths + self.output_embedding_lengths) - if not all(isinstance(length, int) for length in lengths): - raise TypeError( - "Multimodal encoder item lengths must contain only integers") - if not (len(self.item_refs) == len(self.encoder_token_lengths) == len( - self.output_embedding_lengths)): - raise ValueError( - "Multimodal encoder item references and lengths must align") - if any(length <= 0 for length in self.encoder_token_lengths): - raise ValueError( - "Multimodal encoder token lengths must be positive") - if any(length <= 0 for length in self.output_embedding_lengths): - raise ValueError( - "Multimodal encoder output embedding lengths must be positive") - - -def get_multimodal_encoder_item_metadata( - multimodal_data: Optional[Dict[str, Any]], -) -> Optional[MultimodalEncoderItemMetadata]: - """Return typed atomic-item metadata from Python multimodal data.""" - if multimodal_data is None: - return None - if not isinstance(multimodal_data, dict): - raise TypeError("multimodal_data must be a dict") - metadata = multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) - if metadata is None: - return None - if not isinstance(metadata, MultimodalEncoderItemMetadata): - raise TypeError(f"{MULTIMODAL_ENCODER_ITEM_METADATA_KEY} must be a " - "MultimodalEncoderItemMetadata") - metadata.validate() - return metadata - - class InputProcessor(Protocol): """ Protocol for InputProcessor classes. diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 1bfeb03cce11..b0d3f5662362 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -1,18 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Junction tests: item-level MM encoder scheduling over `EncoderGroup` models. - -The two features are developed independently and each has its own coverage, but -nothing exercised them *together* until this file. They meet at a specific -contract: `encode_multimodal_by_groups` reads `multimodal_embedding_lengths` -off each `MultimodalParams` to split the encoder output back per modality, while -item scheduling feeds it freshly sliced single-item params. A slice that omits -that key splits the encoder output to zero rows, so every scheduled item comes -back empty; the row-count check in `forward_multimodal_encoder_items` turns that -into a hard error, but only after the wasted encoder forward. +"""Full-request and selected-item encoding share prompt-order item metadata. + +Group outputs follow modality-local payload order. Selecting items must slice +the tensors and their metadata together, so output assembly returns the same +rows and order as encoding those items in the full request. """ +from functools import partial + import pytest import torch @@ -22,10 +19,10 @@ MultimodalModelMixin, encode_multimodal_by_groups, ) -from tensorrt_llm.inputs.multimodal import MultimodalParams -from tensorrt_llm.inputs.registry import ( +from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, + MultimodalParams, ) HIDDEN = 4 @@ -112,9 +109,8 @@ def _make_request(modalities: list[str], *, value_start: float = 1.0) -> Multimo encoder_token_lengths=list(lengths), output_embedding_lengths=list(lengths), ) - # Production attaches this alongside the item metadata (see - # `inputs/registry.py`); `_apply_metadata_slice` and - # `encode_multimodal_by_groups` both read it from the top level. + # Processors retain this field for legacy consumers. Encoder groups use + # the typed item metadata when available. multimodal_data["multimodal_embedding_lengths"] = list(lengths) return MultimodalParams(multimodal_data=multimodal_data) @@ -245,9 +241,24 @@ def test_full_and_item_order_can_differ_from_raw_payload_order(use_metadata: boo torch.testing.assert_close(model.encode_multimodal_inputs([request]), expected, rtol=0, atol=0) -def test_vision_and_audio_groups_share_full_and_item_ordering() -> None: +@pytest.mark.parametrize("shared_vision_encoder", [False, True]) +def test_vision_and_audio_groups_share_full_and_item_ordering(shared_vision_encoder: bool) -> None: model = _GroupedEncoderModel() + def build_vision_input( + params: list[MultimodalParams], modality: str + ) -> dict[str, torch.Tensor]: + pixel_key = "pixel_values" if modality == "image" else "pixel_values_videos" + return { + "pixel_values": torch.cat([mp.multimodal_data[modality][pixel_key] for mp in params]) + } + + if not shared_vision_encoder: + model.mm_encoder_groups = tuple( + EncoderGroup((modality,), model._encode, partial(build_vision_input, modality=modality)) + for modality in ("image", "video") + ) + def build_audio_input(params: list[MultimodalParams]) -> dict[str, torch.Tensor]: return { "pixel_values": torch.cat( diff --git a/tests/unittest/_torch/multimodal/test_encoder_group.py b/tests/unittest/_torch/multimodal/test_encoder_group.py index 44df0a85e212..41f5247b2a69 100644 --- a/tests/unittest/_torch/multimodal/test_encoder_group.py +++ b/tests/unittest/_torch/multimodal/test_encoder_group.py @@ -6,34 +6,33 @@ * ``_lengths_by_modality`` — invert prompt-ordered ``multimodal_embedding_lengths`` into per-modality per-item lists using ``mm_item_order``. -* ``_get_multimodal_item_layout`` — shared item metadata with legacy fallbacks. +* ``MultimodalParams.get_item_layout`` — item metadata with legacy fallbacks. * ``_reorder_embeds_by_manifest`` — slice per-modality tensors and concat in each request's prompt-order manifest. """ -from types import SimpleNamespace - import pytest import torch from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( - _get_multimodal_item_layout, _lengths_by_modality, _reorder_embeds_by_manifest, ) -from tensorrt_llm.inputs.multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY -from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalEncoderItemMetadata, + MultimodalParams, +) pytestmark = pytest.mark.cpu_only def _mp(*, mm_item_order=None, embedding_lengths=None, buckets=None): - """Minimal MultimodalParams stand-in — the helpers only touch - ``multimodal_data`` and ``mm_item_order``.""" + """Request parameters with a legacy layout or typed item metadata.""" data = dict(buckets or {}) if embedding_lengths is not None: data["multimodal_embedding_lengths"] = embedding_lengths - return SimpleNamespace(multimodal_data=data, mm_item_order=mm_item_order) + return MultimodalParams(multimodal_data=data, mm_item_order=mm_item_order) class TestLengthsByModality: @@ -48,7 +47,7 @@ def test_mixed_manifest_image_video_image(self): embedding_lengths=[4, 6, 2], buckets={"image": {}, "video": {}}, ) - assert _lengths_by_modality([mp], ("image", "video")) == { + assert _lengths_by_modality([mp.get_item_layout()], ("image", "video")) == { "image": [4, 2], "video": [6], } @@ -56,7 +55,7 @@ def test_mixed_manifest_image_video_image(self): def test_single_modality_no_manifest(self): # No mm_item_order; flat list is already in per-modality order. mp = _mp(embedding_lengths=[64, 64], buckets={"image": {}}) - assert _lengths_by_modality([mp], ("image", "video")) == { + assert _lengths_by_modality([mp.get_item_layout()], ("image", "video")) == { "image": [64, 64], "video": [], } @@ -72,7 +71,7 @@ def test_length_mismatch_raises(self): buckets={"image": {}, "video": {}}, ) with pytest.raises(ValueError): - _lengths_by_modality([mp], ("image", "video")) + _lengths_by_modality([mp.get_item_layout()], ("image", "video")) def test_cross_group_multi_modality_without_manifest_raises(self): # Multi-group models register one single-modality group per encoder. @@ -82,28 +81,28 @@ def test_cross_group_multi_modality_without_manifest_raises(self): # widens to the request-level modality set. mp = _mp(embedding_lengths=[4, 3], buckets={"image": {}, "audio": {}}) with pytest.raises(ValueError, match="mm_item_order"): - _lengths_by_modality([mp], ("image",)) + _lengths_by_modality([mp.get_item_layout()], ("image",)) def test_single_modality_outside_group_is_skipped(self): # Request has only audio; called on the image group. Must silently # return the group's empty bucket rather than KeyError on # ``by_modality["audio"]``. mp = _mp(embedding_lengths=[3], buckets={"audio": {}}) - assert _lengths_by_modality([mp], ("image",)) == {"image": []} + assert _lengths_by_modality([mp.get_item_layout()], ("image",)) == {"image": []} class TestMultimodalItemLayout: def test_image_only(self): mp = _mp(embedding_lengths=[10, 20, 30], buckets={"image": {}}) - assert _get_multimodal_item_layout(mp, ("image", "video")) == ( + assert mp.get_item_layout() == ( [("image", 0), ("image", 1), ("image", 2)], [10, 20, 30], ) - def test_no_group_modality_returns_empty(self): - # Request has audio but the group only covers image/video. + def test_layout_includes_all_request_modalities(self): + # Layout belongs to the request, independent of the encoder group. mp = _mp(embedding_lengths=[10], buckets={"audio": {}}) - assert _get_multimodal_item_layout(mp, ("image", "video")) == ([], []) + assert mp.get_item_layout() == ([("audio", 0)], [10]) def test_item_metadata_owns_order_and_lengths(self) -> None: metadata = MultimodalEncoderItemMetadata( @@ -116,12 +115,15 @@ def test_item_metadata_owns_order_and_lengths(self) -> None: ) mp.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata - item_refs, lengths = _get_multimodal_item_layout(mp, ("image", "video")) + item_refs, lengths = mp.get_item_layout() assert item_refs is metadata.item_refs assert lengths is metadata.output_embedding_lengths - assert _lengths_by_modality([mp], ("image", "video")) == {"image": [4], "video": [6]} - assert _lengths_by_modality([mp], ("audio",)) == {"audio": [2]} + assert _lengths_by_modality([mp.get_item_layout()], ("image", "video")) == { + "image": [4], + "video": [6], + } + assert _lengths_by_modality([mp.get_item_layout()], ("audio",)) == {"audio": [2]} class TestReorderEmbedsByManifest: @@ -147,7 +149,9 @@ def test_mixed_image_video_image_in_one_request(self): "video": self._marker_tensor(20, 6), } per_modality_lengths = {"image": [4, 2], "video": [6]} - out = _reorder_embeds_by_manifest([mp], per_modality_embeds, per_modality_lengths) + out = _reorder_embeds_by_manifest( + [mp.get_item_layout()], per_modality_embeds, per_modality_lengths + ) # Expected column-0: 10 (×4), 20 (×6), 30 (×2) in prompt order. expected = torch.tensor([10.0] * 4 + [20.0] * 6 + [30.0] * 2) assert torch.equal(out[:, 0], expected) @@ -172,7 +176,11 @@ def test_cross_request_cursors_advance_per_modality(self): "video": self._marker_tensor(2, 3), } per_modality_lengths = {"image": [2, 1], "video": [3]} - out = _reorder_embeds_by_manifest([mp_a, mp_b], per_modality_embeds, per_modality_lengths) + out = _reorder_embeds_by_manifest( + [mp_a.get_item_layout(), mp_b.get_item_layout()], + per_modality_embeds, + per_modality_lengths, + ) # A: image#0 → marker 1 (×2). B: video#0 → 2 (×3), then image#0 → 3 (×1). expected = torch.tensor([1.0] * 2 + [2.0] * 3 + [3.0] * 1) assert torch.equal(out[:, 0], expected) @@ -188,7 +196,9 @@ def test_one_populated_modality_ignores_empty_group_peer(self): "video": image_embeddings.new_empty((0, image_embeddings.shape[1])), } per_modality_lengths = {"image": [2, 3], "video": []} - out = _reorder_embeds_by_manifest([mp], per_modality_embeds, per_modality_lengths) + out = _reorder_embeds_by_manifest( + [mp.get_item_layout()], per_modality_embeds, per_modality_lengths + ) expected = torch.tensor([7.0] * 2 + [9.0] * 3) assert torch.equal(out[:, 0], expected) assert out is image_embeddings @@ -202,7 +212,9 @@ def test_single_modality_noncanonical_manifest_is_reordered(self): buckets={"image": {}}, ) image_embeddings = torch.cat([self._marker_tensor(7, 2), self._marker_tensor(9, 3)], dim=0) - out = _reorder_embeds_by_manifest([mp], {"image": image_embeddings}, {"image": [2, 3]}) + out = _reorder_embeds_by_manifest( + [mp.get_item_layout()], {"image": image_embeddings}, {"image": [2, 3]} + ) expected = torch.tensor([9.0] * 3 + [7.0] * 2) assert torch.equal(out[:, 0], expected) @@ -223,7 +235,9 @@ def test_single_modality_batch_reorders_only_out_of_order_requests(self, b_item_ dim=0, ) out = _reorder_embeds_by_manifest( - [mp_a, mp_b], {"image": image_embeddings}, {"image": [2, 3, 1]} + [mp_a.get_item_layout(), mp_b.get_item_layout()], + {"image": image_embeddings}, + {"image": [2, 3, 1]}, ) b_rows = {0: [2.0] * 3, 1: [3.0] * 1} expected = torch.tensor([1.0] * 2 + b_rows[b_item_order[0]] + b_rows[b_item_order[1]]) @@ -239,7 +253,9 @@ def test_empty_bookkeeping_returns_typed_empty(self): hidden = 5 per_modality_embeds = {"image": torch.empty((0, hidden), dtype=torch.float16)} per_modality_lengths = {"image": []} - out = _reorder_embeds_by_manifest([mp], per_modality_embeds, per_modality_lengths) + out = _reorder_embeds_by_manifest( + [mp.get_item_layout()], per_modality_embeds, per_modality_lengths + ) assert out.shape == (0, hidden) assert out.dtype == torch.float16 @@ -247,5 +263,5 @@ def test_no_embeds_returns_empty(self): # Defensive: no group produced embeddings at all — return an empty # tensor rather than crashing. mp = _mp(buckets={}) - out = _reorder_embeds_by_manifest([mp], {}, {}) + out = _reorder_embeds_by_manifest([mp.get_item_layout()], {}, {}) assert out.numel() == 0 diff --git a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py index 11ea35f3c6b7..c3c8a95ef71a 100644 --- a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py +++ b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py @@ -10,7 +10,6 @@ disagg entrypoint) delegates to, so this suite covers both paths. """ -from types import SimpleNamespace from typing import Optional import pytest @@ -21,6 +20,7 @@ encode_multimodal_by_groups, ) from tensorrt_llm._torch.models.modeling_qwen3vl import _qwen3vl_build_batched_input +from tensorrt_llm.inputs.multimodal import MultimodalParams def _encode_batched(pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: @@ -43,7 +43,7 @@ def _param( video: Optional[dict] = None, order: Optional[list] = None, lengths: Optional[list] = None, -) -> SimpleNamespace: +) -> MultimodalParams: data: dict = {} if image is not None: data["image"] = image @@ -51,7 +51,7 @@ def _param( data["video"] = video if lengths is not None: data["multimodal_embedding_lengths"] = lengths - return SimpleNamespace(multimodal_data=data, mm_item_order=order) + return MultimodalParams(multimodal_data=data, mm_item_order=order) def _patches(marker: int, n: int, dim: int = 2) -> torch.Tensor: diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index b2c368ac88f8..e072f379ace5 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -14,6 +14,7 @@ from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, DisaggPrefillMultimodalInputs, + MultimodalEncoderItemMetadata, MultimodalInput, MultimodalParams, MultimodalRuntimeData, @@ -22,7 +23,6 @@ ) from tensorrt_llm.inputs.multimodal_data import MEDIA_MODALITIES, AudioData, VideoData from tensorrt_llm.inputs.registry import ( - MultimodalEncoderItemMetadata, create_input_processor_with_hash, maybe_compute_mm_embed_cumsum, ) @@ -137,6 +137,46 @@ def test_mm_item_metadata_retains_three_field_tuple_contract() -> None: assert output_lengths == [2] +@pytest.mark.parametrize( + "selected, expected_refs, expected_costs, expected_lengths", + [ + ([3, 0, 2], [("audio", 0), ("image", 0), ("image", 1)], [40, 20, 30], [4, 2, 3]), + ([2, 0], [("image", 0), ("image", 1)], [30, 20], [3, 2]), + ([], [], [], []), + ], +) +def test_mm_item_metadata_selection_rebases_refs_without_mutating_source( + selected, expected_refs, expected_costs, expected_lengths +) -> None: + metadata = MultimodalEncoderItemMetadata( + [("image", 1), ("video", 0), ("image", 0), ("audio", 0)], + [20, 50, 30, 40], + [2, 5, 3, 4], + ) + + subset = metadata.select(selected) + + assert subset == MultimodalEncoderItemMetadata(expected_refs, expected_costs, expected_lengths) + assert metadata.item_refs == [("image", 1), ("video", 0), ("image", 0), ("audio", 0)] + assert metadata.encoder_token_lengths == [20, 50, 30, 40] + assert metadata.output_embedding_lengths == [2, 5, 3, 4] + assert all( + selected_field is not source_field for selected_field, source_field in zip(subset, metadata) + ) + + +def test_mm_item_metadata_legacy_import_and_pickle() -> None: + import pickle + + from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata as LegacyItemMetadata + + assert LegacyItemMetadata is MultimodalEncoderItemMetadata + # Existing serialized records resolve the original registry module path. + legacy_class = pickle.loads(b"ctensorrt_llm.inputs.registry\nMultimodalEncoderItemMetadata\n.") + metadata = legacy_class([("audio", 0)], [8], [2]) + assert pickle.loads(pickle.dumps(metadata)) == metadata + + def test_maybe_compute_mm_embed_cumsum_populates_py_multimodal_data(): """Producer writes a flat int64 cumsum tensor at py_multimodal_data[multimodal_embed_mask_cumsum].""" From b470512cb2e8a0fa31c748abeb8ce8c524ff4e77 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:23:05 -0700 Subject: [PATCH 04/22] [TRTLLM-14477][refactor] declare multimodal item field layouts Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_gemma4mm.py | 63 +++-------- .../models/modeling_multimodal_mixin.py | 105 +++++++++--------- .../_torch/models/modeling_qwen2vl.py | 44 ++++++++ .../_torch/models/modeling_qwen3vl.py | 14 ++- tensorrt_llm/inputs/multimodal.py | 78 +++++++++++++ .../_torch/modeling/test_gemma4_multimodal.py | 22 ++++ .../test_multimodal_encoder_item_groups.py | 72 ++++++++++++ .../multimodal/test_multimodal_mixin.py | 34 ++++++ tests/unittest/inputs/test_multimodal.py | 74 ++++++++++++ 9 files changed, 402 insertions(+), 104 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 624478bcc4dc..e7aa8ed81abe 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -48,7 +48,7 @@ TextPrompt, register_input_processor, ) -from ...inputs.multimodal import MultimodalParams +from ...inputs.multimodal import MultimodalFieldLayout, MultimodalParams from ...logger import logger from ...sampling_params import SamplingParams from ..modules.embedding import Embedding @@ -959,66 +959,31 @@ def partition_encoder_cache( ) return partition - def build_multimodal_encoder_input( + def get_multimodal_field_layout( self, param: MultimodalParams, - item_indices: Sequence[int], - modality: Optional[str] = None, - ) -> MultimodalParams: - """Build a Gemma4 image, video, or audio input containing selected items. - - The generic implementation recognizes item-major images through a parallel `image_sizes` - field and item-major audio through the `input_features` key. Gemma4 instead provides - fixed-size `pixel_values` without `image_sizes`, uses `audio_features`, and flattens the - frames of every video into dim 0. - - Partial encoder-cache hits and encoder item scheduling therefore need this override to slice - those tensors, together with their per-item position, length, and mask fields. - """ - if modality is None: - modality = self._encoder_cache_modality(param) - if modality == "video": - return self._build_video_encoder_input(param, item_indices) - # Delegate unexpected modalities to the generic validation so they fail instead of - # returning an incorrectly unsliced input. + modality: str, + ) -> dict[str, MultimodalFieldLayout] | None: + """Declare Gemma4's item-major image and audio fields.""" if modality not in ("image", "audio"): - return super().build_multimodal_encoder_input(param, item_indices, modality=modality) + return None input_key = {"image": "pixel_values", "audio": "audio_features"}[modality] - modality_data = param.multimodal_data[modality] - if not isinstance(modality_data, dict): - raise TypeError( - f"multimodal_data[{modality!r}] must be a dict, got {type(modality_data).__name__}" - ) - - item_count = len(param.multimodal_data.get("multimodal_embedding_lengths") or ()) + item_refs, _ = param.get_item_layout() + item_count = sum(item_modality == modality for item_modality, _ in item_refs) input_tensor = modality_data.get(input_key) - # Only slice dim 0 after confirming it is the per-item axis declared by the cache metadata. - # Let the generic implementation handle any other recognized layout, or reject an - # inconsistent Gemma4 payload. if ( not isinstance(input_tensor, torch.Tensor) or input_tensor.dim() == 0 or input_tensor.shape[0] != item_count ): - return super().build_multimodal_encoder_input(param, item_indices, modality=modality) - - indices = list(item_indices) - sliced = {input_key: input_tensor[indices]} - sliced = { - **modality_data, - **sliced, - **self._slice_per_item_sibling_fields( - modality_data, item_count, indices, sliced.keys() - ), - } - residual_input = ( - copy.copy(param.multimodal_input) if param.multimodal_input is not None else None - ) - return MultimodalParams( - multimodal_data={**param.multimodal_data, modality: sliced}, - multimodal_input=residual_input, + return None + fields = ( + ("pixel_values", "image_position_ids", "image_seq_lens") + if modality == "image" + else ("audio_features", "audio_features_mask") ) + return {name: MultimodalFieldLayout() for name in fields} @staticmethod def _build_video_encoder_input( diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index c811f0e0b78b..f928fc5bcb77 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -41,9 +41,11 @@ from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, + MultimodalFieldLayout, MultimodalInput, MultimodalParams, MultimodalRuntimeData, + select_multimodal_fields, ) from tensorrt_llm.logger import logger @@ -391,8 +393,8 @@ class MultimodalModelMixin: single-modality `MultimodalParams` objects. Mixed-modality objects bypass the cache. * A partially cached `MultimodalParams` is handled by encoding only its miss items and interleaving cached items back in original per-item order. The default - `build_multimodal_encoder_input` handles stacked-on-dim-0 and packed-with-grid-thw - layouts; models with other layouts override that method. + `build_multimodal_encoder_input` selects fields declared by + `get_multimodal_field_layout`. Undeclared models retain legacy stacked/packed inference. """ supports_encoder_cache: ClassVar[bool] = False @@ -813,44 +815,31 @@ def after_active_multimodal_embeddings( # them as extra embeds without changing the base flow. return active_embeddings, () + def get_multimodal_field_layout( + self, + param: MultimodalParams, + modality: str, + ) -> dict[str, MultimodalFieldLayout] | None: + """Declare how items occupy each field of one modality's payload. + + Models declare batched axes, packed item lengths, and shared values. + Unlisted fields pass through unchanged. Returning None retains legacy + layout inference for models that have not declared their fields. + """ + return None + def build_multimodal_encoder_input( self, param: MultimodalParams, item_indices: Sequence[int], modality: Optional[str] = None, ) -> MultimodalParams: - """Return a `MultimodalParams` whose raw modality inputs contain only - `item_indices` from `param`, in that order. - - `item_indices` are indices into `modality`'s own payload. When `modality` - is None it is inferred from `param`, which requires the request to hold a - single modality -- the full-request path's case, where an item's global - index and its index within its modality coincide. Callers that already - know an item's modality (item scheduling reads it from `item_refs`) pass - it explicitly along with the modality-local indices, which is what lets - a single item be sliced out of an interleaved image+video request. - - Default handles three common single-modality layouts: - - - Image, stacked on dim 0 (Mistral 3 / Pixtral / LLaVA-family): `pixel_values` - `[B, C, H, W]` with a parallel `image_sizes` list; both sliced by item. - - Image / video, packed with `*_grid_thw` offsets (Qwen2-VL family): - `pixel_values` `[total_patches, feat]` + `image_grid_thw` `[B, 3]`; - prefix-summed patch counts locate each item's slice, and `image_grid_thw` - is sliced in parallel. - - Audio, stacked on dim 0 (Whisper / Qwen2-Audio / Gemma4 audio): - `input_features` `[B, mel_bins, T]` sliced by item. - - Any additional sibling field in the modality dict whose first-axis length equals - the item count is also sliced -- covers per-item metadata such as - `second_per_grid_ts` (Qwen2.5-VL video) or `input_features_mask` / - `feature_attention_mask` (audio) without model-specific code. - - Models with a different layout (e.g. mixed-modality per param, custom packed - formats) should override this method. The parallel per-item metadata - (`multimodal_embedding_lengths`, `multimodal_hashes`) is model-agnostic and is - re-sliced by the mixin after this returns, so overrides need only handle the - modality-specific raw data. + """Build a forward view containing selected modality-local items. + + Declared field layouts drive common tensor/list selection. Models with + custom preprocessing can still override this hook. The caller slices + the parallel item metadata after this returns, preserving original + request indices for scheduler targets and cache ownership. """ if modality is None: modality = self._encoder_cache_modality(param) @@ -865,7 +854,34 @@ def build_multimodal_encoder_input( raise TypeError( f"multimodal_data[{modality!r}] must be a dict, got {type(modality_data).__name__}" ) + field_layout = self.get_multimodal_field_layout(param, modality) + if field_layout is None: + sliced = self._select_legacy_multimodal_fields(modality_data, modality, item_indices) + else: + sliced = select_multimodal_fields(modality_data, field_layout, item_indices) + + # Metadata slicing must not rewrite the original request's hashes. + residual_input = ( + copy.copy(param.multimodal_input) if param.multimodal_input is not None else None + ) + residual_data = { + key: value + for key, value in param.multimodal_data.items() + if key not in _MM_DATA_INPUT_MODALITY_KEYS or key == modality + } + residual_data[modality] = sliced + return MultimodalParams( + multimodal_data=residual_data, + multimodal_input=residual_input, + ) + def _select_legacy_multimodal_fields( + self, + modality_data: Dict[str, Any], + modality: str, + item_indices: Sequence[int], + ) -> Dict[str, Any]: + """Select recognized packed/stacked layouts for undeclared models.""" indices = list(item_indices) grid_key = {"image": "image_grid_thw", "video": "video_grid_thw"}.get(modality) pixel_key = {"image": "pixel_values", "video": "pixel_values_videos"}.get(modality) @@ -942,26 +958,7 @@ def build_multimodal_encoder_input( **self._slice_per_item_sibling_fields(modality_data, n_items, indices, sliced.keys()), } - # Shallow-copy `multimodal_input` so `_apply_metadata_slice` can rewrite - # `multimodal_hashes` on the residual without mutating the source. - residual_input = ( - copy.copy(param.multimodal_input) if param.multimodal_input is not None else None - ) - # Carry the source's other entries through, but keep only the sliced - # modality's raw payload: leaving a sibling modality's unsliced tensors - # on the residual would make it look mixed-modality to the encoder - # (`_lengths_by_modality` rejects that without an `mm_item_order` - # manifest) while its rows were never requested. - residual_data = { - key: value - for key, value in param.multimodal_data.items() - if key not in _MM_DATA_INPUT_MODALITY_KEYS or key == modality - } - residual_data[modality] = sliced - return MultimodalParams( - multimodal_data=residual_data, - multimodal_input=residual_input, - ) + return sliced @staticmethod def _slice_per_item_sibling_fields( diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 8ba5aebb2f8f..481140603980 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -32,6 +32,7 @@ from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, MultimodalEncoderItemMetadata, + MultimodalFieldLayout, MultimodalParams) from ..._utils import async_tensor_h2d, nvtx_range @@ -2101,9 +2102,52 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, return hidden_states +_QWEN_VISION_BATCHED_FIELD = MultimodalFieldLayout() + + +def _qwen_vision_field_layout( + param: MultimodalParams, + modality: str, +) -> dict[str, MultimodalFieldLayout] | None: + """Declare Qwen vision patch packing and parallel per-item grid fields.""" + if modality not in ("image", "video"): + return None + pixel_key, grid_key = (("pixel_values", + "image_grid_thw") if modality == "image" else + ("pixel_values_videos", "video_grid_thw")) + grids = param.multimodal_data[modality][grid_key] + # Host grids give packing lengths directly. Device grids reuse CPU item + # costs to avoid a device-to-host read during selection. + metadata = param.item_metadata if grids.device.type != "cpu" else None + if metadata is not None: + # Physical encoder costs are patch counts, in payload-local order. + costs_by_item = [0] * grids.shape[0] + for (item_modality, + local_index), cost in zip(metadata.item_refs, + metadata.encoder_token_lengths, + strict=True): + if item_modality == modality: + costs_by_item[local_index] = cost + patch_counts = tuple(costs_by_item) + else: + patch_counts = tuple(t * h * w for t, h, w in grids.tolist()) + return { + pixel_key: MultimodalFieldLayout(item_lengths=patch_counts), + grid_key: _QWEN_VISION_BATCHED_FIELD, + "second_per_grid_ts": _QWEN_VISION_BATCHED_FIELD, + } + + class Qwen2VLModelBase(PreTrainedModel, MultimodalModelMixin): supports_mm_encoder_item_scheduling = True + def get_multimodal_field_layout( + self, + param: MultimodalParams, + modality: str, + ) -> dict[str, MultimodalFieldLayout] | None: + return _qwen_vision_field_layout(param, modality) + def __init__( self, model_config: ModelConfig[PretrainedConfig], diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 82c35ee4534e..bdcd15b42c16 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -38,7 +38,11 @@ register_input_processor, support_multimodal_disaggregated, ) -from ...inputs.multimodal import DisaggPrefillMultimodalInputs, MultimodalParams +from ...inputs.multimodal import ( + DisaggPrefillMultimodalInputs, + MultimodalFieldLayout, + MultimodalParams, +) from ...logger import logger from ..attention.backends import AttentionMetadata from ..attention.backends.interface import PositionalEmbeddingParams, RopeParams @@ -62,6 +66,7 @@ Qwen2_5_VLVisionAttention, Qwen2VLInputProcessorBase, _prepare_qwen_vl_mrope_config, + _qwen_vision_field_layout, ) from .modeling_utils import ( ModelConfig, @@ -1417,6 +1422,13 @@ def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tenso class Qwen3VLModelBase(MultimodalModelMixin, PreTrainedModel): supports_mm_encoder_item_scheduling = True + def get_multimodal_field_layout( + self, + param: MultimodalParams, + modality: str, + ) -> dict[str, MultimodalFieldLayout] | None: + return _qwen_vision_field_layout(param, modality) + def encode_multimodal_inputs( self, multimodal_params: List[MultimodalParams], **encoder_kwargs: Any ) -> torch.Tensor: diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index d7b201ea6e5b..31f8e21ecba8 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -5,10 +5,12 @@ `MultimodalEncoderItemMetadata` describes request items. `MultimodalParams` references processed payloads and the `MultimodalRuntimeData` view for the current forward. `MultimodalInput` carries the executor boundary representation. +`MultimodalFieldLayout` describes item axes and packed ranges within payloads. Processor registration and encoder execution live outside this module. """ from dataclasses import dataclass, field +from itertools import accumulate from typing import (Any, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union, cast) @@ -32,6 +34,82 @@ _HASH_SCHEME_TAG = b"trtllm.mm.hash.v1" +class MultimodalFieldLayout(NamedTuple): + """How items occupy one processed payload field. + + ``item_axis`` identifies the item axis of a batched tensor, or is ``None`` + for a value shared by every item. Lists and tuples have item axis zero. + ``item_lengths`` partitions a packed tensor along that axis, for example + patch counts for vision or frame counts for audio. Lengths are host-side + metadata, independent of the payload tensor's device. + """ + + item_axis: int | None = 0 + item_lengths: tuple[int, ...] | None = None + + def select(self, value: Any, item_indices: Sequence[int]) -> Any: + """Select items in the requested order, preserving tensor dtype/device. + + Contiguous selections of a packed tensor return a view. Other tensor + selections copy only the selected items; shared values retain identity. + """ + if self.item_axis is None or value is None: + return value + indices = list(item_indices) + if isinstance(value, torch.Tensor): + if self.item_lengths is None: + if self.item_axis == 0: + return value[indices] + index: list[slice | list[int]] = [slice(None)] * value.ndim + index[self.item_axis] = indices + return value[tuple(index)] + + offsets = list(accumulate(self.item_lengths, initial=0)) + if offsets[-1] != value.shape[self.item_axis]: + raise ValueError( + "Packed field size does not match its item lengths") + if not indices: + return value.narrow(self.item_axis, 0, 0) + if indices == list(range(indices[0], indices[0] + len(indices))): + start = offsets[indices[0]] + if self.item_axis == 0: + return value[start:offsets[indices[-1] + 1]] + return value.narrow(self.item_axis, start, + offsets[indices[-1] + 1] - start) + if self.item_axis == 0: + pieces = [value[offsets[i]:offsets[i + 1]] for i in indices] + else: + pieces = [ + value.narrow(self.item_axis, offsets[i], + self.item_lengths[i]) for i in indices + ] + return torch.cat(pieces, dim=self.item_axis) + + if (isinstance(value, (list, tuple)) and self.item_axis == 0 + and self.item_lengths is None): + return type(value)(value[i] for i in indices) + raise TypeError( + "Item fields must be tensors or lists/tuples on item axis zero") + + +def select_multimodal_fields( + modality_data: dict[str, Any], + field_layouts: dict[str, MultimodalFieldLayout], + item_indices: Sequence[int], +) -> dict[str, Any]: + """Select declared fields; unlisted and shared fields pass through. + + Optional fields absent from the payload are ignored. A model declares + every per-item field explicitly, so matching a tensor dimension to the + item count never changes a shared field's meaning. + """ + selected = dict(modality_data) + for name, layout in field_layouts.items(): + if name in modality_data: + selected[name] = layout.select(modality_data[name], item_indices) + return selected + + class MultimodalEncoderItemMetadata(NamedTuple): """Prompt-ordered metadata shared by item scheduling and encoder assembly.""" diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index 727119d8e2d7..fec02fa2f2a1 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -79,6 +79,7 @@ ) from tensorrt_llm.inputs.multimodal import ( # noqa: E402 MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalEncoderItemMetadata, MultimodalInput, MultimodalParams, MultimodalRuntimeData, @@ -948,11 +949,13 @@ def test_build_multimodal_encoder_input_slices_gemma4_audio_input(self): """Gemma4 audio features and their mask remain item-aligned.""" audio_features = torch.arange(24, dtype=torch.float32).reshape(3, 4, 2) audio_mask = torch.arange(12).reshape(3, 4) + shared_stats = torch.arange(3) param = MultimodalParams( multimodal_data={ "audio": { "audio_features": audio_features, "audio_features_mask": audio_mask, + "shared_stats": shared_stats, }, "multimodal_embedding_lengths": [1, 1, 1], } @@ -963,6 +966,7 @@ def test_build_multimodal_encoder_input_slices_gemma4_audio_input(self): residual_audio = residual.multimodal_data["audio"] torch.testing.assert_close(residual_audio["audio_features"], audio_features[[2, 0]]) torch.testing.assert_close(residual_audio["audio_features_mask"], audio_mask[[2, 0]]) + assert residual_audio["shared_stats"] is shared_stats def test_build_multimodal_encoder_input_slices_whole_videos(self): """A video item is the run of flattened frames its `num_frames` entry declares.""" @@ -1090,6 +1094,24 @@ def test_encoder_cache_partial_hit_slices_gemma4_audio_input(self): torch.testing.assert_close(embeddings[2:], torch.full((2, model.embedding_dim), 2.0)) self.assertEqual(len(model._multimodal_encoder_cache), 2) + @pytest.mark.cpu_only + def test_build_multimodal_encoder_input_uses_typed_audio_item_count(self): + features = torch.arange(24).reshape(3, 4, 2) + param = MultimodalParams( + multimodal_data={ + "audio": {"audio_features": features}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + [("audio", 0), ("audio", 1), ("audio", 2)], [8, 8, 8], [1, 1, 1] + ), + } + ) + + residual = _Gemma4EncoderCacheHarness().build_multimodal_encoder_input(param, [2, 0]) + + torch.testing.assert_close( + residual.multimodal_data["audio"]["audio_features"], features[[2, 0]] + ) + def test_encoder_cache_partial_hit_with_unsupported_item_axis_becomes_full_miss(self): """Unexpected Gemma4 image and audio layouts bypass partial cache reuse.""" model = _Gemma4EncoderCacheHarness() diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index b0d3f5662362..f1cacd0e9430 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -19,9 +19,11 @@ MultimodalModelMixin, encode_multimodal_by_groups, ) +from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLModelBase from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, + MultimodalFieldLayout, MultimodalParams, ) @@ -47,6 +49,11 @@ class _GroupedEncoderModel(MultimodalModelMixin): supports_mm_encoder_item_scheduling = True + def get_multimodal_field_layout(self, param, modality): + if modality == "audio": + return {"input_features": MultimodalFieldLayout()} + return Qwen3VLModelBase.get_multimodal_field_layout(self, param, modality) + def __init__(self): self.encoder_calls: list[int] = [] self.mm_encoder_groups = ( @@ -132,6 +139,71 @@ def _run_items(model, selected): return model.forward_multimodal_encoder_items(encoder_inputs), encoder_inputs +@pytest.mark.parametrize("indices", [[2, 0], [1, 2]]) +def test_declared_qwen_layout_keeps_shared_fields(indices): + model = _GroupedEncoderModel() + request = _make_request(["image", "image", "image"]) + shared = torch.arange(3) + request.multimodal_data["image"]["shared_stats"] = shared + expected = torch.cat([_expected_rows(request, i) for i in indices]) + + encoder_inputs = model.prepare_multimodal_encoder_inputs([(request, i) for i in indices]) + residual = encoder_inputs[0][0] + outputs = model.forward_multimodal_encoder_items(encoder_inputs) + + torch.testing.assert_close(torch.cat(outputs), expected, rtol=0, atol=0) + assert residual.multimodal_data["image"]["shared_stats"] is shared + if indices == [1, 2]: + source_pixels = request.multimodal_data["image"]["pixel_values"] + selected_pixels = residual.multimodal_data["image"]["pixel_values"] + assert ( + selected_pixels.untyped_storage().data_ptr() + == source_pixels.untyped_storage().data_ptr() + ) + + +def test_declared_audio_layout_selects_packed_frames_on_a_nonzero_axis(): + class PackedAudioModel(_GroupedEncoderModel): + def get_multimodal_field_layout(self, param, modality): + return { + "features": MultimodalFieldLayout(item_axis=1, item_lengths=(2, 3, 1)), + "mask": MultimodalFieldLayout(item_lengths=(2, 3, 1)), + "frame_counts": MultimodalFieldLayout(), + "shared_stats": MultimodalFieldLayout(item_axis=None), + } + + features = torch.arange(12).reshape(2, 6) + mask = torch.tensor([1, 1, 1, 1, 0, 1]) + shared = torch.arange(3) + request = MultimodalParams( + multimodal_data={ + "audio": { + "features": features, + "mask": mask, + "frame_counts": [2, 3, 1], + "shared_stats": shared, + }, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + [("audio", 0), ("audio", 1), ("audio", 2)], [2, 3, 1], [1, 2, 1] + ), + } + ) + + prepared = PackedAudioModel().prepare_multimodal_encoder_inputs([(request, 2), (request, 0)]) + residual = prepared[0][0] + + torch.testing.assert_close( + residual.multimodal_data["audio"]["features"], features[:, [5, 0, 1]] + ) + torch.testing.assert_close(residual.multimodal_data["audio"]["mask"], mask[[5, 0, 1]]) + assert residual.multimodal_data["audio"]["frame_counts"] == [1, 2] + assert residual.multimodal_data["audio"]["shared_stats"] is shared + assert residual.item_metadata == MultimodalEncoderItemMetadata( + [("audio", 0), ("audio", 1)], [1, 2], [1, 1] + ) + assert request.item_metadata.item_refs == [("audio", 0), ("audio", 1), ("audio", 2)] + + def test_sliced_items_survive_the_grouped_encoder_path(): """Regression: without `multimodal_embedding_lengths` on the slice, the grouped path splits to zero rows and returns nothing for every item.""" diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index 5bfe67925607..eb2b1ca9c4ad 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -269,6 +269,40 @@ def test_prepare_multimodal_inputs_accepts_tensor_encoder_output_cuda(): _test_prepare_multimodal_inputs_accepts_tensor_encoder_output("cuda") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_declared_qwen_fields_reuse_host_costs_for_device_grids(): + from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLModelBase + from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalEncoderItemMetadata, + ) + + class QwenModel(DummyMultimodalModel): + get_multimodal_field_layout = Qwen3VLModelBase.get_multimodal_field_layout + + pixels = torch.arange(10, device="cuda").reshape(5, 2) + grids = torch.tensor([[1, 1, 2], [1, 1, 3]], device="cuda") + param = MultimodalParams( + multimodal_data={ + "image": {"pixel_values": pixels, "image_grid_thw": grids}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + [("image", 1), ("image", 0)], [3, 2], [3, 2] + ), + } + ) + model = QwenModel(make_embedding(hidden_size=1), torch.tensor([0])) + + with patch.object( + torch.Tensor, "tolist", side_effect=AssertionError("Device grids must not be read back") + ): + selected = model.build_multimodal_encoder_input(param, [1, 0]) + + torch.testing.assert_close( + selected.multimodal_data["image"]["pixel_values"], torch.cat([pixels[2:], pixels[:2]]) + ) + torch.testing.assert_close(selected.multimodal_data["image"]["image_grid_thw"], grids[[1, 0]]) + + def test_encoder_cache_first_request_writes_per_item_entries(): model = CountingEncoderMultimodalModel( make_embedding(hidden_size=4), diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index e072f379ace5..0da4bf74d256 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -15,11 +15,13 @@ MULTIMODAL_ENCODER_ITEM_METADATA_KEY, DisaggPrefillMultimodalInputs, MultimodalEncoderItemMetadata, + MultimodalFieldLayout, MultimodalInput, MultimodalParams, MultimodalRuntimeData, _find_mm_embedding_lengths_from_masks, find_mm_token_lengths, + select_multimodal_fields, ) from tensorrt_llm.inputs.multimodal_data import MEDIA_MODALITIES, AudioData, VideoData from tensorrt_llm.inputs.registry import ( @@ -30,6 +32,78 @@ pytestmark = pytest.mark.cpu_only +@pytest.mark.parametrize("axis", [0, 1, -1]) +@pytest.mark.parametrize("indices", [[2, 0], [1, 1, 0], []]) +def test_field_layout_selects_tensor_item_axis(axis: int, indices: list[int]) -> None: + values = torch.arange(9).reshape(3, 3) + selected = MultimodalFieldLayout(item_axis=axis).select(values, indices) + expected = torch.index_select(values, axis, torch.tensor(indices, dtype=torch.long)) + torch.testing.assert_close(selected, expected) + + +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("indices", [[2, 0], [1, 2], [1, 1, 0], []]) +def test_field_layout_selects_packed_items_and_preserves_contiguous_views( + axis: int, indices: list[int] +) -> None: + values = torch.arange(12).reshape(6, 2) + if axis == 1: + values = values.T + layout = MultimodalFieldLayout(item_axis=axis, item_lengths=(2, 3, 1)) + + selected = layout.select(values, indices) + + items = torch.split(values, (2, 3, 1), dim=axis) + expected = ( + torch.cat([items[i] for i in indices], dim=axis) if indices else values.narrow(axis, 0, 0) + ) + torch.testing.assert_close(selected, expected) + if indices in ([1, 2], []): + assert selected.untyped_storage().data_ptr() == values.untyped_storage().data_ptr() + else: + assert selected.untyped_storage().data_ptr() != values.untyped_storage().data_ptr() + + +@pytest.mark.parametrize("sequence_type", [list, tuple]) +def test_field_layout_selects_item_sequences(sequence_type) -> None: + values = sequence_type(["a", "b", "c"]) + assert MultimodalFieldLayout().select(values, [2, 0]) == sequence_type(["c", "a"]) + + +def test_field_layout_keeps_shared_and_unlisted_values_even_when_sizes_match() -> None: + shared = torch.arange(3) + unlisted = ["a", "b", "c"] + data = { + "frames": torch.arange(6).reshape(3, 2), + "shared": shared, + "other": unlisted, + "nullable_mask": None, + } + + selected = select_multimodal_fields( + data, + { + "frames": MultimodalFieldLayout(), + "shared": MultimodalFieldLayout(item_axis=None), + "optional_mask": MultimodalFieldLayout(), + "nullable_mask": MultimodalFieldLayout(), + }, + [2, 0], + ) + + torch.testing.assert_close(selected["frames"], data["frames"][[2, 0]]) + assert selected["shared"] is shared + assert selected["other"] is unlisted + assert "optional_mask" not in selected + assert selected["nullable_mask"] is None + assert data["frames"].shape == (3, 2) + + +def test_packed_field_layout_rejects_inconsistent_lengths() -> None: + with pytest.raises(ValueError, match="Packed field size"): + MultimodalFieldLayout(item_lengths=(2, 3)).select(torch.zeros(4, 2), [0]) + + def test_media_modalities_match_io_and_validation_contract() -> None: assert MEDIA_MODALITIES == get_args(MediaModality) == ("image", "video", "audio") assert set(MEDIA_MODALITIES) == set(MEDIA_IO_REGISTRY) From 879df7c7154d954d5653f0c38a93516c3e45b7e0 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:39:31 -0700 Subject: [PATCH 05/22] refactor: name multimodal item partitions and fix negative selection Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_gemma4mm.py | 8 ++-- .../models/modeling_multimodal_mixin.py | 14 +++--- .../_torch/models/modeling_qwen2vl.py | 16 +++---- .../_torch/models/modeling_qwen3vl.py | 10 ++-- tensorrt_llm/inputs/multimodal.py | 14 +++--- .../test_multimodal_encoder_item_groups.py | 18 +++---- .../multimodal/test_multimodal_mixin.py | 2 +- tests/unittest/inputs/test_multimodal.py | 47 +++++++++++++------ 8 files changed, 74 insertions(+), 55 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index e7aa8ed81abe..fa1c630b3543 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -48,7 +48,7 @@ TextPrompt, register_input_processor, ) -from ...inputs.multimodal import MultimodalFieldLayout, MultimodalParams +from ...inputs.multimodal import MultimodalItemPartition, MultimodalParams from ...logger import logger from ...sampling_params import SamplingParams from ..modules.embedding import Embedding @@ -959,11 +959,11 @@ def partition_encoder_cache( ) return partition - def get_multimodal_field_layout( + def get_multimodal_item_partitions( self, param: MultimodalParams, modality: str, - ) -> dict[str, MultimodalFieldLayout] | None: + ) -> dict[str, MultimodalItemPartition] | None: """Declare Gemma4's item-major image and audio fields.""" if modality not in ("image", "audio"): return None @@ -983,7 +983,7 @@ def get_multimodal_field_layout( if modality == "image" else ("audio_features", "audio_features_mask") ) - return {name: MultimodalFieldLayout() for name in fields} + return {name: MultimodalItemPartition() for name in fields} @staticmethod def _build_video_encoder_input( diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index f928fc5bcb77..45fddfbfda19 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -41,8 +41,8 @@ from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, - MultimodalFieldLayout, MultimodalInput, + MultimodalItemPartition, MultimodalParams, MultimodalRuntimeData, select_multimodal_fields, @@ -394,7 +394,7 @@ class MultimodalModelMixin: * A partially cached `MultimodalParams` is handled by encoding only its miss items and interleaving cached items back in original per-item order. The default `build_multimodal_encoder_input` selects fields declared by - `get_multimodal_field_layout`. Undeclared models retain legacy stacked/packed inference. + `get_multimodal_item_partitions`. Undeclared models retain legacy stacked/packed inference. """ supports_encoder_cache: ClassVar[bool] = False @@ -815,11 +815,11 @@ def after_active_multimodal_embeddings( # them as extra embeds without changing the base flow. return active_embeddings, () - def get_multimodal_field_layout( + def get_multimodal_item_partitions( self, param: MultimodalParams, modality: str, - ) -> dict[str, MultimodalFieldLayout] | None: + ) -> dict[str, MultimodalItemPartition] | None: """Declare how items occupy each field of one modality's payload. Models declare batched axes, packed item lengths, and shared values. @@ -854,11 +854,11 @@ def build_multimodal_encoder_input( raise TypeError( f"multimodal_data[{modality!r}] must be a dict, got {type(modality_data).__name__}" ) - field_layout = self.get_multimodal_field_layout(param, modality) - if field_layout is None: + item_partitions = self.get_multimodal_item_partitions(param, modality) + if item_partitions is None: sliced = self._select_legacy_multimodal_fields(modality_data, modality, item_indices) else: - sliced = select_multimodal_fields(modality_data, field_layout, item_indices) + sliced = select_multimodal_fields(modality_data, item_partitions, item_indices) # Metadata slicing must not rewrite the original request's hashes. residual_input = ( diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 481140603980..c9559778ce10 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -32,7 +32,7 @@ from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, MultimodalEncoderItemMetadata, - MultimodalFieldLayout, + MultimodalItemPartition, MultimodalParams) from ..._utils import async_tensor_h2d, nvtx_range @@ -2102,13 +2102,13 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, return hidden_states -_QWEN_VISION_BATCHED_FIELD = MultimodalFieldLayout() +_QWEN_VISION_BATCHED_FIELD = MultimodalItemPartition() -def _qwen_vision_field_layout( +def _qwen_vision_item_partitions( param: MultimodalParams, modality: str, -) -> dict[str, MultimodalFieldLayout] | None: +) -> dict[str, MultimodalItemPartition] | None: """Declare Qwen vision patch packing and parallel per-item grid fields.""" if modality not in ("image", "video"): return None @@ -2132,7 +2132,7 @@ def _qwen_vision_field_layout( else: patch_counts = tuple(t * h * w for t, h, w in grids.tolist()) return { - pixel_key: MultimodalFieldLayout(item_lengths=patch_counts), + pixel_key: MultimodalItemPartition(item_lengths=patch_counts), grid_key: _QWEN_VISION_BATCHED_FIELD, "second_per_grid_ts": _QWEN_VISION_BATCHED_FIELD, } @@ -2141,12 +2141,12 @@ def _qwen_vision_field_layout( class Qwen2VLModelBase(PreTrainedModel, MultimodalModelMixin): supports_mm_encoder_item_scheduling = True - def get_multimodal_field_layout( + def get_multimodal_item_partitions( self, param: MultimodalParams, modality: str, - ) -> dict[str, MultimodalFieldLayout] | None: - return _qwen_vision_field_layout(param, modality) + ) -> dict[str, MultimodalItemPartition] | None: + return _qwen_vision_item_partitions(param, modality) def __init__( self, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index bdcd15b42c16..593ae4279a42 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -40,7 +40,7 @@ ) from ...inputs.multimodal import ( DisaggPrefillMultimodalInputs, - MultimodalFieldLayout, + MultimodalItemPartition, MultimodalParams, ) from ...logger import logger @@ -66,7 +66,7 @@ Qwen2_5_VLVisionAttention, Qwen2VLInputProcessorBase, _prepare_qwen_vl_mrope_config, - _qwen_vision_field_layout, + _qwen_vision_item_partitions, ) from .modeling_utils import ( ModelConfig, @@ -1422,12 +1422,12 @@ def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tenso class Qwen3VLModelBase(MultimodalModelMixin, PreTrainedModel): supports_mm_encoder_item_scheduling = True - def get_multimodal_field_layout( + def get_multimodal_item_partitions( self, param: MultimodalParams, modality: str, - ) -> dict[str, MultimodalFieldLayout] | None: - return _qwen_vision_field_layout(param, modality) + ) -> dict[str, MultimodalItemPartition] | None: + return _qwen_vision_item_partitions(param, modality) def encode_multimodal_inputs( self, multimodal_params: List[MultimodalParams], **encoder_kwargs: Any diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 31f8e21ecba8..c4d604a7e0dc 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -5,7 +5,7 @@ `MultimodalEncoderItemMetadata` describes request items. `MultimodalParams` references processed payloads and the `MultimodalRuntimeData` view for the current forward. `MultimodalInput` carries the executor boundary representation. -`MultimodalFieldLayout` describes item axes and packed ranges within payloads. +`MultimodalItemPartition` describes item axes and packed ranges within payloads. Processor registration and encoder execution live outside this module. """ @@ -34,8 +34,8 @@ _HASH_SCHEME_TAG = b"trtllm.mm.hash.v1" -class MultimodalFieldLayout(NamedTuple): - """How items occupy one processed payload field. +class MultimodalItemPartition(NamedTuple): + """How one processed input is partitioned into multimodal items. ``item_axis`` identifies the item axis of a batched tensor, or is ``None`` for a value shared by every item. Lists and tuples have item axis zero. @@ -70,6 +70,8 @@ def select(self, value: Any, item_indices: Sequence[int]) -> Any: "Packed field size does not match its item lengths") if not indices: return value.narrow(self.item_axis, 0, 0) + item_range = range(len(self.item_lengths)) + indices = [item_range[i] for i in indices] if indices == list(range(indices[0], indices[0] + len(indices))): start = offsets[indices[0]] if self.item_axis == 0: @@ -94,7 +96,7 @@ def select(self, value: Any, item_indices: Sequence[int]) -> Any: def select_multimodal_fields( modality_data: dict[str, Any], - field_layouts: dict[str, MultimodalFieldLayout], + item_partitions: dict[str, MultimodalItemPartition], item_indices: Sequence[int], ) -> dict[str, Any]: """Select declared fields; unlisted and shared fields pass through. @@ -104,9 +106,9 @@ def select_multimodal_fields( item count never changes a shared field's meaning. """ selected = dict(modality_data) - for name, layout in field_layouts.items(): + for name, partition in item_partitions.items(): if name in modality_data: - selected[name] = layout.select(modality_data[name], item_indices) + selected[name] = partition.select(modality_data[name], item_indices) return selected diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index f1cacd0e9430..ca7554b470cd 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -23,7 +23,7 @@ from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, - MultimodalFieldLayout, + MultimodalItemPartition, MultimodalParams, ) @@ -49,10 +49,10 @@ class _GroupedEncoderModel(MultimodalModelMixin): supports_mm_encoder_item_scheduling = True - def get_multimodal_field_layout(self, param, modality): + def get_multimodal_item_partitions(self, param, modality): if modality == "audio": - return {"input_features": MultimodalFieldLayout()} - return Qwen3VLModelBase.get_multimodal_field_layout(self, param, modality) + return {"input_features": MultimodalItemPartition()} + return Qwen3VLModelBase.get_multimodal_item_partitions(self, param, modality) def __init__(self): self.encoder_calls: list[int] = [] @@ -164,12 +164,12 @@ def test_declared_qwen_layout_keeps_shared_fields(indices): def test_declared_audio_layout_selects_packed_frames_on_a_nonzero_axis(): class PackedAudioModel(_GroupedEncoderModel): - def get_multimodal_field_layout(self, param, modality): + def get_multimodal_item_partitions(self, param, modality): return { - "features": MultimodalFieldLayout(item_axis=1, item_lengths=(2, 3, 1)), - "mask": MultimodalFieldLayout(item_lengths=(2, 3, 1)), - "frame_counts": MultimodalFieldLayout(), - "shared_stats": MultimodalFieldLayout(item_axis=None), + "features": MultimodalItemPartition(item_axis=1, item_lengths=(2, 3, 1)), + "mask": MultimodalItemPartition(item_lengths=(2, 3, 1)), + "frame_counts": MultimodalItemPartition(), + "shared_stats": MultimodalItemPartition(item_axis=None), } features = torch.arange(12).reshape(2, 6) diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index eb2b1ca9c4ad..b8de957810f2 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -278,7 +278,7 @@ def test_declared_qwen_fields_reuse_host_costs_for_device_grids(): ) class QwenModel(DummyMultimodalModel): - get_multimodal_field_layout = Qwen3VLModelBase.get_multimodal_field_layout + get_multimodal_item_partitions = Qwen3VLModelBase.get_multimodal_item_partitions pixels = torch.arange(10, device="cuda").reshape(5, 2) grids = torch.tensor([[1, 1, 2], [1, 1, 3]], device="cuda") diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 0da4bf74d256..6590fb99ab39 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -15,8 +15,8 @@ MULTIMODAL_ENCODER_ITEM_METADATA_KEY, DisaggPrefillMultimodalInputs, MultimodalEncoderItemMetadata, - MultimodalFieldLayout, MultimodalInput, + MultimodalItemPartition, MultimodalParams, MultimodalRuntimeData, _find_mm_embedding_lengths_from_masks, @@ -36,20 +36,31 @@ @pytest.mark.parametrize("indices", [[2, 0], [1, 1, 0], []]) def test_field_layout_selects_tensor_item_axis(axis: int, indices: list[int]) -> None: values = torch.arange(9).reshape(3, 3) - selected = MultimodalFieldLayout(item_axis=axis).select(values, indices) + selected = MultimodalItemPartition(item_axis=axis).select(values, indices) expected = torch.index_select(values, axis, torch.tensor(indices, dtype=torch.long)) torch.testing.assert_close(selected, expected) -@pytest.mark.parametrize("axis", [0, 1]) -@pytest.mark.parametrize("indices", [[2, 0], [1, 2], [1, 1, 0], []]) -def test_field_layout_selects_packed_items_and_preserves_contiguous_views( - axis: int, indices: list[int] +@pytest.mark.parametrize("axis", [0, 1, -1]) +@pytest.mark.parametrize( + ("indices", "is_view"), + [ + ([2, 0], False), + ([1, 2], True), + ([1, 1, 0], False), + ([], True), + ([-1], True), + ([-2, -1], True), + ([-1, 0], False), + ], +) +def test_item_partition_selects_packed_items_and_preserves_contiguous_views( + axis: int, indices: list[int], is_view: bool ) -> None: values = torch.arange(12).reshape(6, 2) - if axis == 1: + if axis != 0: values = values.T - layout = MultimodalFieldLayout(item_axis=axis, item_lengths=(2, 3, 1)) + layout = MultimodalItemPartition(item_axis=axis, item_lengths=(2, 3, 1)) selected = layout.select(values, indices) @@ -58,7 +69,7 @@ def test_field_layout_selects_packed_items_and_preserves_contiguous_views( torch.cat([items[i] for i in indices], dim=axis) if indices else values.narrow(axis, 0, 0) ) torch.testing.assert_close(selected, expected) - if indices in ([1, 2], []): + if is_view: assert selected.untyped_storage().data_ptr() == values.untyped_storage().data_ptr() else: assert selected.untyped_storage().data_ptr() != values.untyped_storage().data_ptr() @@ -67,7 +78,7 @@ def test_field_layout_selects_packed_items_and_preserves_contiguous_views( @pytest.mark.parametrize("sequence_type", [list, tuple]) def test_field_layout_selects_item_sequences(sequence_type) -> None: values = sequence_type(["a", "b", "c"]) - assert MultimodalFieldLayout().select(values, [2, 0]) == sequence_type(["c", "a"]) + assert MultimodalItemPartition().select(values, [2, 0]) == sequence_type(["c", "a"]) def test_field_layout_keeps_shared_and_unlisted_values_even_when_sizes_match() -> None: @@ -83,10 +94,10 @@ def test_field_layout_keeps_shared_and_unlisted_values_even_when_sizes_match() - selected = select_multimodal_fields( data, { - "frames": MultimodalFieldLayout(), - "shared": MultimodalFieldLayout(item_axis=None), - "optional_mask": MultimodalFieldLayout(), - "nullable_mask": MultimodalFieldLayout(), + "frames": MultimodalItemPartition(), + "shared": MultimodalItemPartition(item_axis=None), + "optional_mask": MultimodalItemPartition(), + "nullable_mask": MultimodalItemPartition(), }, [2, 0], ) @@ -101,7 +112,13 @@ def test_field_layout_keeps_shared_and_unlisted_values_even_when_sizes_match() - def test_packed_field_layout_rejects_inconsistent_lengths() -> None: with pytest.raises(ValueError, match="Packed field size"): - MultimodalFieldLayout(item_lengths=(2, 3)).select(torch.zeros(4, 2), [0]) + MultimodalItemPartition(item_lengths=(2, 3)).select(torch.zeros(4, 2), [0]) + + +@pytest.mark.parametrize("index", [-4, 3]) +def test_packed_item_partition_rejects_out_of_bounds_indices(index: int) -> None: + with pytest.raises(IndexError): + MultimodalItemPartition(item_lengths=(2, 3, 1)).select(torch.zeros(6, 2), [index]) def test_media_modalities_match_io_and_validation_contract() -> None: From 53c5b1656e57fa3316f22b48c8811cdc0ec8523a Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Sun, 18 Oct 2026 21:28:11 -0700 Subject: [PATCH 06/22] Fix partial encoder cache item selection and preparation overhead Use common item selection for partial-cache misses so prompt order, modality-local payload indices, and canonical output lengths stay aligned. Reuse validated request metadata and contiguous tensor views during input preparation. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../models/modeling_multimodal_mixin.py | 77 +++++----- tensorrt_llm/inputs/multimodal.py | 90 +++++++----- .../test_multimodal_encoder_item_groups.py | 131 ++++++++++++++++++ tests/unittest/inputs/test_multimodal.py | 18 ++- 4 files changed, 244 insertions(+), 72 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 45fddfbfda19..a42b287e261b 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -141,6 +141,17 @@ def _lengths_by_modality( for item_refs, lengths in item_layouts: if not item_refs: continue + modality = item_refs[0][0] + if ( + len(item_refs) == len(lengths) + and modality in by_modality + and all( + item_modality == modality and index == expected + for expected, (item_modality, index) in enumerate(item_refs) + ) + ): + by_modality[modality].extend(lengths) + continue request_lengths: dict[str, list[tuple[int, int]]] = {m: [] for m in modalities} for (modality, index), length in zip(item_refs, lengths, strict=True): if modality in request_lengths: @@ -483,7 +494,11 @@ def prepare_multimodal_encoder_inputs( modality, item_metadata, ) in self._runs_by_request_modality(selected_items): - item_refs = item_metadata.item_refs + item_refs = ( + item_metadata.item_refs + if item_metadata is not None + else multimodal_param.get_item_layout()[0] + ) # The two slicers take different index spaces. The raw-tensor # slicer indexes the modality's own payload, so it gets the # modality-local indices (plus the modality, which also lets it @@ -495,7 +510,9 @@ def prepare_multimodal_encoder_inputs( [item_refs[i][1] for i in run_indices], modality=modality, ) - self._apply_metadata_slice(residual, multimodal_param, run_indices) + self._apply_metadata_slice( + residual, multimodal_param, run_indices, item_metadata=item_metadata + ) except MultimodalEncoderContractError: raise except (KeyError, IndexError, TypeError, ValueError) as error: @@ -505,7 +522,7 @@ def prepare_multimodal_encoder_inputs( encoder_inputs.append( ( residual, - [int(item_metadata.output_embedding_lengths[i]) for i in run_indices], + residual.multimodal_data["multimodal_embedding_lengths"], modality, ) ) @@ -514,7 +531,7 @@ def prepare_multimodal_encoder_inputs( @staticmethod def _runs_by_request_modality( selected_items: Sequence[tuple[MultimodalParams, int]], - ) -> Iterator[tuple[MultimodalParams, list[int], str, "MultimodalEncoderItemMetadata"]]: + ) -> Iterator[tuple[MultimodalParams, list[int], str, Optional[MultimodalEncoderItemMetadata]]]: """Split scheduler order into maximal same-request, same-modality runs. Only adjacent items merge, so the flattened result keeps the @@ -526,6 +543,7 @@ def _runs_by_request_modality( run_param: Optional[MultimodalParams] = None run_modality: Optional[str] = None run_metadata: Optional[MultimodalEncoderItemMetadata] = None + item_refs: list[tuple[str, int]] = [] run_indices: list[int] = [] for multimodal_param, item_idx in selected_items: try: @@ -536,16 +554,17 @@ def _runs_by_request_modality( ) except (TypeError, ValueError) as error: raise MultimodalEncoderContractError(str(error)) from error - if metadata is None: - raise MultimodalEncoderContractError( - "MM item metadata is required for item encoding" + if multimodal_param is not run_param: + item_refs = ( + metadata.item_refs + if metadata is not None + else multimodal_param.get_item_layout()[0] ) - if item_idx < 0 or item_idx >= len(metadata.item_refs): + if item_idx < 0 or item_idx >= len(item_refs): raise MultimodalEncoderContractError( - f"MM item index {item_idx} is out of range for " - f"{len(metadata.item_refs)} item(s)" + f"MM item index {item_idx} is out of range for {len(item_refs)} item(s)" ) - modality = metadata.item_refs[item_idx][0] + modality = item_refs[item_idx][0] if run_indices and (multimodal_param is not run_param or modality != run_modality): yield run_param, run_indices, run_modality, run_metadata run_indices = [] @@ -1372,13 +1391,15 @@ def _apply_metadata_slice( residual: MultimodalParams, source: MultimodalParams, item_indices: Sequence[int], + *, + item_metadata: Optional[MultimodalEncoderItemMetadata] = None, ) -> None: """Overwrite `residual`'s per-item metadata to match the sliced items. Models slice raw modality tensors in `build_multimodal_encoder_input`; the mixin owns the parallel per-item metadata slice so every model gets it identically. """ - metadata = source.item_metadata + metadata = item_metadata if item_metadata is not None else source.item_metadata if metadata is not None: selected = metadata.select(item_indices) residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = selected @@ -1419,29 +1440,17 @@ def _encode_with_partial_cache( if param.encoder_event is not None: torch.cuda.current_stream().wait_event(param.encoder_event) - # Build every residual, then run one batched encoder call over the whole set. - residuals: list[MultimodalParams] = [] - per_param_miss_lengths: list[list[int]] = [] - for param, partition in partials: - residual = self.build_multimodal_encoder_input(param, partition.miss_indices) - self._apply_metadata_slice(residual, param, partition.miss_indices) - residuals.append(residual) - per_param_miss_lengths.append( - [ - param.multimodal_data["multimodal_embedding_lengths"][i] - for i in partition.miss_indices - ] - ) - - batched_output = self.encode_multimodal_inputs(residuals) - per_param_slabs = torch.split( - batched_output, [sum(lengths) for lengths in per_param_miss_lengths], dim=0 + encoder_inputs = self.prepare_multimodal_encoder_inputs( + [(param, i) for param, partition in partials for i in partition.miss_indices] ) + # Keep one call across requests and modalities so EncoderGroup can + # batch modalities that share an encoder. + batched_output = self.encode_multimodal_inputs([param for param, _, _ in encoder_inputs]) + output_lengths = [length for _, lengths, _ in encoder_inputs for length in lengths] + encoded_items = iter(torch.split(batched_output, output_lengths, dim=0)) - for (param, partition), slab, miss_lengths in zip( - partials, per_param_slabs, per_param_miss_lengths, strict=True - ): - miss_tensors = torch.split(slab, miss_lengths, dim=0) + for param, partition in partials: + miss_tensors = [next(encoded_items) for _ in partition.miss_indices] by_item: Dict[int, torch.Tensor] = dict(partition.hits) for miss_idx, tensor in zip(partition.miss_indices, miss_tensors, strict=True): @@ -1487,7 +1496,7 @@ def _write_encoder_cache_entries( ) return - embedding_lengths = param.multimodal_data["multimodal_embedding_lengths"] + _, embedding_lengths = param.get_item_layout() if sum(embedding_lengths) != embedding.shape[0]: logger.debug( f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping write because embedding row count " diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index c4d604a7e0dc..ffce49a64080 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -50,46 +50,52 @@ class MultimodalItemPartition(NamedTuple): def select(self, value: Any, item_indices: Sequence[int]) -> Any: """Select items in the requested order, preserving tensor dtype/device. - Contiguous selections of a packed tensor return a view. Other tensor - selections copy only the selected items; shared values retain identity. + Contiguous tensor selections return views that callers must treat as + read-only. Other selections copy only the selected items; shared values + retain identity. """ - if self.item_axis is None or value is None: + axis, lengths = self + if axis is None or value is None: return value - indices = list(item_indices) if isinstance(value, torch.Tensor): - if self.item_lengths is None: - if self.item_axis == 0: - return value[indices] - index: list[slice | list[int]] = [slice(None)] * value.ndim - index[self.item_axis] = indices - return value[tuple(index)] - - offsets = list(accumulate(self.item_lengths, initial=0)) - if offsets[-1] != value.shape[self.item_axis]: - raise ValueError( - "Packed field size does not match its item lengths") + offsets = None + if lengths is not None: + offsets = list(accumulate(lengths, initial=0)) + if offsets[-1] != value.shape[axis]: + raise ValueError( + "Packed field size does not match its item lengths") + item_range = range(len(lengths)) + else: + item_range = range(value.shape[axis]) + # Python range indexing normalizes negatives and rejects out-of-range + # items before any view can silently truncate the requested rows. + indices = [item_range[i] for i in item_indices] if not indices: - return value.narrow(self.item_axis, 0, 0) - item_range = range(len(self.item_lengths)) - indices = [item_range[i] for i in indices] + return value.narrow(axis, 0, 0) if indices == list(range(indices[0], indices[0] + len(indices))): - start = offsets[indices[0]] - if self.item_axis == 0: - return value[start:offsets[indices[-1] + 1]] - return value.narrow(self.item_axis, start, - offsets[indices[-1] + 1] - start) - if self.item_axis == 0: - pieces = [value[offsets[i]:offsets[i + 1]] for i in indices] - else: - pieces = [ - value.narrow(self.item_axis, offsets[i], - self.item_lengths[i]) for i in indices - ] - return torch.cat(pieces, dim=self.item_axis) - - if (isinstance(value, (list, tuple)) and self.item_axis == 0 - and self.item_lengths is None): - return type(value)(value[i] for i in indices) + start, end = indices[0], indices[-1] + 1 + if offsets is not None: + start, end = offsets[start], offsets[end] + if axis == 0: + return value[start:end] + return value.narrow(axis, start, end - start) + if offsets is not None: + if axis == 0: + pieces = [value[offsets[i]:offsets[i + 1]] for i in indices] + else: + pieces = [ + value.narrow(axis, offsets[i], lengths[i]) + for i in indices + ] + return torch.cat(pieces, dim=axis) + if axis == 0: + return value[indices] + index: list[slice | list[int]] = [slice(None)] * value.ndim + index[axis] = indices + return value[tuple(index)] + + if (isinstance(value, (list, tuple)) and axis == 0 and lengths is None): + return type(value)(value[i] for i in item_indices) raise TypeError( "Item fields must be tensors or lists/tuples on item axis zero") @@ -163,6 +169,13 @@ def select(self, Each modality's payload must be sliced in the same selected-item order. The source metadata remains owned by the original request. """ + if len(item_indices) == 1: + i = item_indices[0] + return MultimodalEncoderItemMetadata( + [(self.item_refs[i][0], 0)], + [self.encoder_token_lengths[i]], + [self.output_embedding_lengths[i]], + ) local_counts: dict[str, int] = {} item_refs = [] for i in item_indices: @@ -721,7 +734,12 @@ def get_item_layout(self) -> tuple[list[tuple[str, int]], list[int]]: ``mm_item_order`` and ``multimodal_embedding_lengths``; a request with just one modality can omit the order. """ - metadata = self.item_metadata + # Input processors validate this record before dispatch; layout reads + # do not need to revalidate encoder costs on every forward. + metadata = cast( + Optional[MultimodalEncoderItemMetadata], + self.multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY), + ) if metadata is not None: return metadata.item_refs, metadata.output_embedding_lengths diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index ca7554b470cd..9441aefde835 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -20,11 +20,14 @@ encode_multimodal_by_groups, ) from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLModelBase +from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, + MultimodalInput, MultimodalItemPartition, MultimodalParams, + MultimodalRuntimeData, ) HIDDEN = 4 @@ -313,6 +316,134 @@ def test_full_and_item_order_can_differ_from_raw_payload_order(use_metadata: boo torch.testing.assert_close(model.encode_multimodal_inputs([request]), expected, rtol=0, atol=0) +def _attach_cache_metadata(request, item_lengths): + request.multimodal_input = MultimodalInput.from_components( + [[i + 1] * 8 for i in range(len(item_lengths))], + [sum(item_lengths[:i]) for i in range(len(item_lengths))], + list(item_lengths), + ) + request.multimodal_runtime = MultimodalRuntimeData( + 0, sum(item_lengths), torch.arange(1, sum(item_lengths) + 1, dtype=torch.int64) + ) + request.multimodal_data["mm_processor_kwargs_hash"] = "test" + + +@pytest.mark.parametrize("metadata_format", ["typed", "typed_only", "legacy"]) +@pytest.mark.parametrize("reversed_order", [False, True]) +def test_partial_cache_uses_prompt_refs_and_lengths(metadata_format, reversed_order): + model = _GroupedEncoderModel() + model._multimodal_encoder_cache = TensorLRUCache(4096) + request = _make_request(["image", "image"]) + metadata = request.item_metadata + if reversed_order: + metadata = MultimodalEncoderItemMetadata([("image", 1), ("image", 0)], [3, 2], [3, 2]) + request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata + lengths = metadata.output_embedding_lengths + request.multimodal_data["multimodal_embedding_lengths"] = lengths + if metadata_format == "typed_only": + request.multimodal_data.pop("multimodal_embedding_lengths") + elif metadata_format == "legacy": + request.mm_item_order = [ + {"modality": modality, "index": index} for modality, index in metadata.item_refs + ] + request.multimodal_data.pop(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) + _attach_cache_metadata(request, lengths) + expected = model.encode_multimodal_inputs([request]) + keys = model._encoder_cache_keys(request) + model._multimodal_encoder_cache.put(keys[0], expected[: lengths[0]]) + + actual = model._get_or_encode_multimodal_embeddings([request]) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close( + model._multimodal_encoder_cache.get(keys[1]), expected[lengths[0] :], rtol=0, atol=0 + ) + assert model.encoder_calls[-1] == lengths[1] + if metadata_format != "legacy": + assert request.item_metadata is metadata + + +def test_full_cache_write_uses_canonical_lengths_without_legacy_fields(): + model = _GroupedEncoderModel() + model._multimodal_encoder_cache = TensorLRUCache(4096) + request = _make_request(["image", "image"]) + lengths = request.item_metadata.output_embedding_lengths + request.multimodal_data.pop("multimodal_embedding_lengths") + _attach_cache_metadata(request, lengths) + expected = model.encode_multimodal_inputs([request]) + calls_before = len(model.encoder_calls) + + actual = model._get_or_encode_multimodal_embeddings([request]) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert len(model.encoder_calls) == calls_before + 1 + for key, rows in zip( + model._encoder_cache_keys(request), torch.split(expected, lengths), strict=True + ): + torch.testing.assert_close(model._multimodal_encoder_cache.get(key), rows, rtol=0, atol=0) + + +def test_partial_cache_batches_mixed_vision_and_audio_misses(): + model = _GroupedEncoderModel() + model.mm_encoder_groups += ( + EncoderGroup( + ("audio",), + model._encode, + lambda params: { + "pixel_values": torch.cat( + [p.multimodal_data["audio"]["input_features"].flatten(0, 1) for p in params] + ) + }, + ), + ) + model._multimodal_encoder_cache = TensorLRUCache(4096) + request = _make_request(["image", "video", "image"]) + request.multimodal_data["audio"] = { + "input_features": torch.stack([torch.full((1, HIDDEN), value) for value in (10.0, 20.0)]) + } + request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalEncoderItemMetadata( + [("video", 0), ("audio", 0), ("image", 1), ("image", 0), ("audio", 1)], + [4, 1, 3, 2, 1], + [4, 1, 3, 2, 1], + ) + request.multimodal_data.pop("multimodal_embedding_lengths") + lengths = request.item_metadata.output_embedding_lengths + _attach_cache_metadata(request, lengths) + expected = model.encode_multimodal_inputs([request]) + expected_items = torch.split(expected, lengths) + keys = model._encoder_cache_keys(request) + for i in (0, 3): + model._multimodal_encoder_cache.put(keys[i], expected_items[i]) + calls_before = len(model.encoder_calls) + + actual = model._get_or_encode_multimodal_embeddings([request]) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + # Image misses share one vision call; audio misses share one audio call. + assert model.encoder_calls[calls_before:] == [3, 2] + for key, rows in zip(keys, expected_items, strict=True): + torch.testing.assert_close(model._multimodal_encoder_cache.get(key), rows, rtol=0, atol=0) + + +def test_alternating_item_selection_validates_each_request_once(monkeypatch): + model = _GroupedEncoderModel() + requests = [_make_request(["image", "video"] * 4) for _ in range(2)] + validated = [] + validate = MultimodalEncoderItemMetadata.validate + + def record_validation(metadata): + validated.append(metadata) + validate(metadata) + + monkeypatch.setattr(MultimodalEncoderItemMetadata, "validate", record_validation) + + model.prepare_multimodal_encoder_inputs( + [(request, i) for request in requests for i in range(8)] + ) + + assert len(validated) == len(requests) + + @pytest.mark.parametrize("shared_vision_encoder", [False, True]) def test_vision_and_audio_groups_share_full_and_item_ordering(shared_vision_encoder: bool) -> None: model = _GroupedEncoderModel() diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 6590fb99ab39..023caef6505b 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -41,6 +41,17 @@ def test_field_layout_selects_tensor_item_axis(axis: int, indices: list[int]) -> torch.testing.assert_close(selected, expected) +@pytest.mark.parametrize("axis", [0, 1, -1]) +@pytest.mark.parametrize("indices", [[0, 1], [-2, -1], [2], []]) +def test_batched_item_partition_preserves_contiguous_views(axis, indices) -> None: + values = torch.arange(9).reshape(3, 3) + selected = MultimodalItemPartition(item_axis=axis).select(values, indices) + expected = values[indices] if axis == 0 else values[:, indices] + + torch.testing.assert_close(selected, expected) + assert selected.untyped_storage().data_ptr() == values.untyped_storage().data_ptr() + + @pytest.mark.parametrize("axis", [0, 1, -1]) @pytest.mark.parametrize( ("indices", "is_view"), @@ -116,9 +127,12 @@ def test_packed_field_layout_rejects_inconsistent_lengths() -> None: @pytest.mark.parametrize("index", [-4, 3]) -def test_packed_item_partition_rejects_out_of_bounds_indices(index: int) -> None: +@pytest.mark.parametrize("packed", [False, True]) +def test_item_partition_rejects_out_of_bounds_indices(index: int, packed: bool) -> None: + lengths = (2, 3, 1) if packed else None + values = torch.zeros(6 if packed else 3, 2) with pytest.raises(IndexError): - MultimodalItemPartition(item_lengths=(2, 3, 1)).select(torch.zeros(6, 2), [index]) + MultimodalItemPartition(item_lengths=lengths).select(values, [index]) def test_media_modalities_match_io_and_validation_contract() -> None: From aacea28e7c9f4cc1315a617e831e4636e135a6ea Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:10:10 -0700 Subject: [PATCH 07/22] refactor: group multimodal request data and encoder state Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_gemma4mm.py | 58 +-- .../models/modeling_multimodal_mixin.py | 66 ++-- .../models/modeling_nemotron_h_multimodal.py | 24 ++ .../_torch/models/modeling_qwen2vl.py | 57 ++- .../_torch/pyexecutor/engine/multimodal.py | 13 +- .../pyexecutor/engine/runners/no_kv_cache.py | 2 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 247 +++++-------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 11 +- tensorrt_llm/inputs/__init__.py | 5 +- tensorrt_llm/inputs/multimodal.py | 337 +++++++++++++++++- tensorrt_llm/inputs/registry.py | 108 +++--- .../executor/test_multimodal_scheduler.py | 110 +++++- .../_torch/executor/test_request_utils.py | 53 +++ .../test_multimodal_encoder_item_groups.py | 103 +++++- .../_torch/multimodal/test_encoder_group.py | 20 +- tests/unittest/inputs/test_multimodal.py | 115 +++++- 16 files changed, 998 insertions(+), 331 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index fa1c630b3543..f7a11b001c36 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -48,7 +48,7 @@ TextPrompt, register_input_processor, ) -from ...inputs.multimodal import MultimodalItemPartition, MultimodalParams +from ...inputs.multimodal import MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams from ...logger import logger from ...sampling_params import SamplingParams from ..modules.embedding import Embedding @@ -211,6 +211,26 @@ def _padded_patches_per_frame(processor) -> int: # --------------------------------------------------------------------------- +def _gemma4_input_partitions( + multimodal_data: dict[str, Any], + metadata: MultimodalItemMetadata, +) -> dict[str, dict[str, MultimodalItemPartition]]: + """Declare item-major fields; video frame grouping stays model-owned.""" + partitions = {} + for modality, fields in ( + ("image", ("pixel_values", "image_position_ids", "image_seq_lens")), + ("audio", ("audio_features", "audio_features_mask")), + ): + data = multimodal_data.get(modality) + if not isinstance(data, dict): + continue + item_count = sum(m == modality for m, _ in metadata.item_refs) + value = data.get(fields[0]) + if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] == item_count: + partitions[modality] = {name: MultimodalItemPartition() for name in fields} + return partitions + + class Gemma4InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): """Preprocesses image inputs for Gemma4. @@ -219,6 +239,17 @@ class Gemma4InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInpu processing using the image processor saved in the model directory. """ + def get_mm_item_metadata( + self, + prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + mm_item_order: Optional[List[Dict[str, Any]]] = None, + ) -> Optional[MultimodalItemMetadata]: + metadata = super().get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) + if metadata is not None: + metadata.input_partitions.update(_gemma4_input_partitions(multimodal_data, metadata)) + return metadata + # Default class-level fallback. Real value computed per-instance below # from text_config.use_bidirectional_attention. Only 26B/31B set # use_bidirectional_attention="vision" — their image blocks need intact @@ -964,26 +995,13 @@ def get_multimodal_item_partitions( param: MultimodalParams, modality: str, ) -> dict[str, MultimodalItemPartition] | None: - """Declare Gemma4's item-major image and audio fields.""" - if modality not in ("image", "audio"): + partitions = super().get_multimodal_item_partitions(param, modality) + if partitions is not None: + return partitions + metadata = param.item_metadata + if metadata is None: return None - input_key = {"image": "pixel_values", "audio": "audio_features"}[modality] - modality_data = param.multimodal_data[modality] - item_refs, _ = param.get_item_layout() - item_count = sum(item_modality == modality for item_modality, _ in item_refs) - input_tensor = modality_data.get(input_key) - if ( - not isinstance(input_tensor, torch.Tensor) - or input_tensor.dim() == 0 - or input_tensor.shape[0] != item_count - ): - return None - fields = ( - ("pixel_values", "image_position_ids", "image_seq_lens") - if modality == "image" - else ("audio_features", "audio_features_mask") - ) - return {name: MultimodalItemPartition() for name in fields} + return _gemma4_input_partitions(param.multimodal_data, metadata).get(modality) @staticmethod def _build_video_encoder_input( diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index a42b287e261b..e20c6cf0d76b 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -40,11 +40,12 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalEncoderItemMetadata, MultimodalInput, + MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams, MultimodalRuntimeData, + normalize_multimodal_item_metadata, select_multimodal_fields, ) from tensorrt_llm.logger import logger @@ -90,7 +91,7 @@ class EncoderGroup: Batching all items in a group into one encoder invocation amortizes fixed costs (kernel launches, dispatch) across items. The framework splits the output back per-modality and reorders it into prompt order - via each request's `MultimodalEncoderItemMetadata.item_refs` manifest, + via each request's `MultimodalItemMetadata.item_refs` manifest, falling back to `mm_item_order` for legacy processors. Contract between `build_batched_input` and `encoder_fn`: @@ -238,7 +239,7 @@ def encode_multimodal_by_groups( For each group present in the batch, one encoder call is issued over all items across all requests belonging to that group's modalities (arithmetic-intensity win). The output is split back per-modality using - the prompt-ordered output lengths in `MultimodalEncoderItemMetadata`, + the prompt-ordered output lengths in `MultimodalItemMetadata`, then reordered by the same metadata's `item_refs`. Legacy processors fall back to `multimodal_embedding_lengths` and `mm_item_order`. @@ -494,11 +495,7 @@ def prepare_multimodal_encoder_inputs( modality, item_metadata, ) in self._runs_by_request_modality(selected_items): - item_refs = ( - item_metadata.item_refs - if item_metadata is not None - else multimodal_param.get_item_layout()[0] - ) + item_refs = item_metadata.item_refs # The two slicers take different index spaces. The raw-tensor # slicer indexes the modality's own payload, so it gets the # modality-local indices (plus the modality, which also lets it @@ -531,19 +528,15 @@ def prepare_multimodal_encoder_inputs( @staticmethod def _runs_by_request_modality( selected_items: Sequence[tuple[MultimodalParams, int]], - ) -> Iterator[tuple[MultimodalParams, list[int], str, Optional[MultimodalEncoderItemMetadata]]]: + ) -> Iterator[tuple[MultimodalParams, list[int], str, MultimodalItemMetadata]]: """Split scheduler order into maximal same-request, same-modality runs. - Only adjacent items merge, so the flattened result keeps the - scheduler's order and outputs map back positionally. Each run also - carries its request's item metadata: fetching it validates the whole - record, so it is read once per request rather than once per item, and - the caller reuses it instead of fetching again. + Only adjacent items merge, preserving scheduler order. Reuse the item + metadata validated at the input boundary or executor admission. """ run_param: Optional[MultimodalParams] = None run_modality: Optional[str] = None - run_metadata: Optional[MultimodalEncoderItemMetadata] = None - item_refs: list[tuple[str, int]] = [] + run_metadata: Optional[MultimodalItemMetadata] = None run_indices: list[int] = [] for multimodal_param, item_idx in selected_items: try: @@ -552,19 +545,21 @@ def _runs_by_request_modality( if multimodal_param is run_param else multimodal_param.item_metadata ) + if metadata is None: + metadata = normalize_multimodal_item_metadata( + multimodal_param.multimodal_data, multimodal_param.mm_item_order + ) except (TypeError, ValueError) as error: raise MultimodalEncoderContractError(str(error)) from error - if multimodal_param is not run_param: - item_refs = ( - metadata.item_refs - if metadata is not None - else multimodal_param.get_item_layout()[0] + if metadata is None or metadata.output_embedding_lengths is None: + raise MultimodalEncoderContractError( + "MM item metadata with output lengths is required for item encoding" ) - if item_idx < 0 or item_idx >= len(item_refs): + if item_idx < 0 or item_idx >= len(metadata.item_refs): raise MultimodalEncoderContractError( - f"MM item index {item_idx} is out of range for {len(item_refs)} item(s)" + f"MM item index {item_idx} is out of range for {len(metadata.item_refs)} item(s)" ) - modality = item_refs[item_idx][0] + modality = metadata.item_refs[item_idx][0] if run_indices and (multimodal_param is not run_param or modality != run_modality): yield run_param, run_indices, run_modality, run_metadata run_indices = [] @@ -845,7 +840,8 @@ def get_multimodal_item_partitions( Unlisted fields pass through unchanged. Returning None retains legacy layout inference for models that have not declared their fields. """ - return None + metadata = param.item_metadata + return metadata.input_partitions.get(modality) if metadata is not None else None def build_multimodal_encoder_input( self, @@ -1240,6 +1236,17 @@ def _encoder_cache_keys( item_metadata = param.item_metadata if item_metadata is not None: + if item_metadata.output_embedding_lengths is None: + return None + item_modalities = {modality for modality, _ in item_metadata.item_refs} + if any( + modality not in item_modalities + for modality in _MM_DATA_INPUT_MODALITY_KEYS + if mm_data.get(modality) is not None + ): + # A legacy caller may add another raw modality after its + # original single-modality metadata was normalized. + return None return cls.build_encoder_cache_item_keys( mm_input.multimodal_hashes, item_metadata.item_refs, @@ -1392,7 +1399,7 @@ def _apply_metadata_slice( source: MultimodalParams, item_indices: Sequence[int], *, - item_metadata: Optional[MultimodalEncoderItemMetadata] = None, + item_metadata: Optional[MultimodalItemMetadata] = None, ) -> None: """Overwrite `residual`'s per-item metadata to match the sliced items. @@ -1403,9 +1410,10 @@ def _apply_metadata_slice( if metadata is not None: selected = metadata.select(item_indices) residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = selected - residual.multimodal_data["multimodal_embedding_lengths"] = ( - selected.output_embedding_lengths - ) + if selected.output_embedding_lengths is not None: + residual.multimodal_data["multimodal_embedding_lengths"] = ( + selected.output_embedding_lengths + ) else: source_lengths = source.multimodal_data["multimodal_embedding_lengths"] residual.multimodal_data["multimodal_embedding_lengths"] = [ diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index 38b6fb4bf048..d9b5d9021f18 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -29,6 +29,8 @@ from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper from tensorrt_llm.inputs.multimodal import ( DisaggPrefillMultimodalInputs, + MultimodalItemMetadata, + MultimodalItemPartition, MultimodalParams, _as_cpu_tensor, _compute_mm_masks, @@ -1151,6 +1153,28 @@ class NemotronHMultimodalInputProcessor( ): supports_token_id_mm_expansion: ClassVar[bool] = True + def get_mm_item_metadata( + self, + prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + mm_item_order: Optional[List[Dict[str, Any]]] = None, + ) -> Optional[MultimodalItemMetadata]: + metadata = super().get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) + audio = multimodal_data.get("audio") + if metadata is not None and isinstance(audio, dict) and audio.get("evs_ids") is None: + # One audio item can contain several clips. Video's embedded audio + # stays in its video bucket and retains the model's AV interleave. + clips = tuple(audio["audio_num_clips"].tolist()) + item_count = sum(m == "audio" for m, _ in metadata.item_refs) + if len(clips) == item_count: + packed = MultimodalItemPartition(item_lengths=clips) + metadata.input_partitions["audio"] = { + "input_audio_features": packed, + "feature_attention_mask": packed, + "audio_num_clips": MultimodalItemPartition(), + } + return metadata + def __init__( self, model_path: str, diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index c9559778ce10..f91f98ff9a0b 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -32,6 +32,7 @@ from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, MultimodalEncoderItemMetadata, + MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams) @@ -332,6 +333,22 @@ def get_mm_encoder_item_metadata( encoder_token_lengths=encoder_token_lengths, output_embedding_lengths=output_embedding_lengths) + def get_mm_item_metadata( + self, + prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + mm_item_order: Optional[List[Dict[str, Any]]] = None, + ) -> Optional[MultimodalItemMetadata]: + metadata = super().get_mm_item_metadata(prompt_token_ids, + multimodal_data, mm_item_order) + if metadata is not None: + for modality in ("image", "video"): + if isinstance(multimodal_data.get(modality), dict): + metadata.input_partitions[ + modality] = _qwen_vision_input_partitions( + multimodal_data, modality, metadata) + return metadata + # ------------------------------------------------------------------ # Deterministic dummy-input sizing for multimodal profiling. # @@ -2109,26 +2126,36 @@ def _qwen_vision_item_partitions( param: MultimodalParams, modality: str, ) -> dict[str, MultimodalItemPartition] | None: - """Declare Qwen vision patch packing and parallel per-item grid fields.""" if modality not in ("image", "video"): return None + metadata = param.item_metadata + if metadata is not None and modality in metadata.input_partitions: + return metadata.input_partitions[modality] + grid_key = "image_grid_thw" if modality == "image" else "video_grid_thw" + grids = param.multimodal_data[modality][grid_key] + return _qwen_vision_input_partitions( + param.multimodal_data, modality, + metadata if grids.device.type != "cpu" else None) + + +def _qwen_vision_input_partitions( + multimodal_data: dict[str, Any], + modality: str, + metadata: MultimodalItemMetadata | None, +) -> dict[str, MultimodalItemPartition]: + """Declare patch packing once on the producer, with a legacy fallback.""" pixel_key, grid_key = (("pixel_values", "image_grid_thw") if modality == "image" else ("pixel_values_videos", "video_grid_thw")) - grids = param.multimodal_data[modality][grid_key] - # Host grids give packing lengths directly. Device grids reuse CPU item - # costs to avoid a device-to-host read during selection. - metadata = param.item_metadata if grids.device.type != "cpu" else None - if metadata is not None: - # Physical encoder costs are patch counts, in payload-local order. - costs_by_item = [0] * grids.shape[0] - for (item_modality, - local_index), cost in zip(metadata.item_refs, - metadata.encoder_token_lengths, - strict=True): - if item_modality == modality: - costs_by_item[local_index] = cost - patch_counts = tuple(costs_by_item) + grids = multimodal_data[modality][grid_key] + if metadata is not None and metadata.encoder_token_lengths is not None: + costs = { + index: cost + for (item_modality, index), cost in zip( + metadata.item_refs, metadata.encoder_token_lengths, strict=True) + if item_modality == modality + } + patch_counts = tuple(costs[i] for i in range(grids.shape[0])) else: patch_counts = tuple(t * h * w for t, h, w in grids.tolist()) return { diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index c7dd344fea08..22a7d1ea31f5 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -27,7 +27,8 @@ from tensorrt_llm.inputs.multimodal import ( MultimodalParams, MultimodalRuntimeData, - get_multimodal_encoder_item_metadata, + _Unset, + get_multimodal_item_metadata, strip_mm_encoder_inputs, ) from tensorrt_llm.inputs.registry import ( @@ -42,7 +43,7 @@ ) from tensorrt_llm.logger import logger -from ..llm_request import LlmRequest, MultimodalEncoderRequestError, _Unset, get_mm_items_for_chunk +from ..llm_request import LlmRequest, MultimodalEncoderRequestError, get_mm_items_for_chunk def resolve_mm_encoder_token_budget(base_budget: int, model_max_atomic_item_tokens: int) -> int: @@ -484,12 +485,12 @@ def item_cache_keys(self, request: LlmRequest) -> list[Hashable] | None: return state.stable_item_cache_keys mm_data = request.py_multimodal_data try: - item_metadata = get_multimodal_encoder_item_metadata(mm_data) + item_metadata = get_multimodal_item_metadata(mm_data) except (TypeError, ValueError) as error: raise MultimodalEncoderRequestError(str(error)) from error keys = ( None - if item_metadata is None + if item_metadata is None or item_metadata.output_embedding_lengths is None else ( self.model.build_encoder_cache_item_keys( request.multimodal_hashes, @@ -548,7 +549,9 @@ def requests_using_cache_keys(cache_keys: set[Hashable]) -> set[int]: f"Scheduled MM request {request_id} has no encoder item state", request_ids={request_id}, ) - multimodal_param = MultimodalParams(multimodal_data=request.py_multimodal_data) + multimodal_param = MultimodalParams( + multimodal_data=request.py_multimodal_params.multimodal_data + ) for item_idx in item_indices: cache_key = state.item_cache_keys[item_idx] if cache_key is None: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py index 86f5bd4add1a..bdc014a561ad 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py @@ -140,7 +140,7 @@ def _prepare_inputs( request, runner_config.mm_encoder_cache_enabled, ), - multimodal_data=request.py_multimodal_data, + multimodal_data=request.py_multimodal_params.multimodal_data, mm_item_order=getattr(request, "py_mm_item_order", None), input_ids_start_offset=context_start_idx, ) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 3c4a1706be9b..ed0572e3e21c 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -4,9 +4,7 @@ from copy import copy, deepcopy from dataclasses import dataclass, field -from enum import Enum, auto -from typing import (TYPE_CHECKING, Any, Container, Dict, Hashable, Iterable, - List, Optional, Union, cast) +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union, cast import torch @@ -15,7 +13,13 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.bindings import executor as tllm_executor from tensorrt_llm.executor.result import SimpleTokenLogprobs, TokenLogprobs -from tensorrt_llm.inputs.multimodal import get_multimodal_encoder_item_metadata +from tensorrt_llm.inputs.multimodal import \ + MultimodalEncoderProgress as MultimodalEncoderProgress +from tensorrt_llm.inputs.multimodal import \ + MultimodalEncoderRequestState as MultimodalEncoderRequestState +from tensorrt_llm.inputs.multimodal import (MultimodalParams, + get_multimodal_item_metadata, + normalize_multimodal_item_metadata) from tensorrt_llm.sampling_params import LogprobMode SamplingConfig = tllm_executor.SamplingConfig @@ -65,167 +69,12 @@ def __init__(self, self.request_ids = frozenset(request_ids or ()) -class MultimodalEncoderProgress(Enum): - """Python-only progress derived from request-local MM item outputs. - - It summarizes the ready flags of all items. Scheduling is gated by the - items the remaining prefill uses instead (`get_mm_pending_items`), because - items behind confirmed prefill progress or a reusable KV prefix are - released or never encoded.""" - - PENDING = auto() - """No MM item of the request has an encoder output yet.""" - - PARTIAL = auto() - """Some but not all items are encoded; the request must continue on the - item-scheduling path until every item has an output.""" - - READY = auto() - """No further encoder work is needed: every item is encoded, a - precomputed embedding was supplied, or the request has no MM payload.""" - - -class _Unset: - """Sentinel for a memo that has not been computed yet. - - Distinct from `None`, which is a real result: the request cannot build - stable cache keys and its items always encode. - """ - - -_UNSET = _Unset() - - def make_mm_encoder_transient_cache_key(request_id: int, item_idx: int) -> tuple[str, int, int]: """Return a request-local cache key for an item without a stable key.""" return ("mm_transient", request_id, item_idx) -@dataclass -class MultimodalEncoderRequestState: - """Tracks the encoder-cache key for each MM item in a request. - - The state is created when an item-scheduled request enters the executor. - Other requests leave `py_mm_encoder_state` as `None`. - - Encoder outputs live only in the model-owned `TensorLRUCache`. This object - stores one cache key and one ready flag per item, in prompt order. On the - item-scheduled executor, each non-empty slot holds one reference. Two - identical items may point to the same cache entry, but each item still - holds and releases its own reference. The request does not store another - output tensor or combined output buffer. - - Cleanup clears the slots before releasing their references, so repeated - request cleanup is safe. - """ - - embedding_lengths: List[int] - """Expected encoder-output rows for each item, in prompt order.""" - - encoder_token_lengths: List[int] - """Validated encoder attention-token cost of each atomic item. - - Admission copies these values from the request metadata so scheduling - never has to re-validate user input inside the scheduler loop. - """ - - item_ready: List[bool] - """Whether each item's bound cache entry is ready, in prompt order.""" - - item_cache_keys: List[Optional[Hashable]] = field(default_factory=list) - """Cache keys held by this request, in prompt order. - - A non-empty slot holds one reference. `item_ready` tells whether the entry - already contains its encoder output. - """ - - stable_item_cache_keys: "list[Hashable] | _Unset | None" = _UNSET - """Cached per-item reuse keys. - - `_UNSET` means the keys have not been computed. `None` means stable keys - cannot be built, so this request uses request-local cache keys instead. - """ - - @classmethod - def from_embedding_lengths( - cls, - embedding_lengths: List[int], - *, - encoder_token_lengths: Optional[List[int]] = None - ) -> "MultimodalEncoderRequestState": - if encoder_token_lengths is None: - encoder_token_lengths = embedding_lengths - return cls(embedding_lengths=list(embedding_lengths), - encoder_token_lengths=list(encoder_token_lengths), - item_ready=[False] * len(embedding_lengths), - item_cache_keys=[None] * len(embedding_lengths)) - - def __post_init__(self) -> None: - if not self.item_cache_keys: - self.item_cache_keys = [None] * len(self.item_ready) - if not (len(self.embedding_lengths) == len(self.encoder_token_lengths) - == len(self.item_ready) == len(self.item_cache_keys)): - raise ValueError("MM encoder token and embedding lengths must have " - "exactly one cache key per item slot") - - @property - def num_items(self) -> int: - return len(self.item_ready) - - @property - def progress(self) -> MultimodalEncoderProgress: - if all(self.item_ready): - return MultimodalEncoderProgress.READY - if any(self.item_ready): - return MultimodalEncoderProgress.PARTIAL - return MultimodalEncoderProgress.PENDING - - def pending_item_indices(self) -> List[int]: - """Indices of items that still need an encoder output, prompt order.""" - return [ - item_idx for item_idx, ready in enumerate(self.item_ready) - if not ready - ] - - def set_item_cache_key(self, item_idx: int, cache_key: Hashable, *, - ready: bool) -> None: - """Assign a cache key to one item.""" - if self.item_cache_keys[item_idx] is not None: - raise RuntimeError(f"MM item {item_idx} already has a cache key") - self.item_cache_keys[item_idx] = cache_key - self.item_ready[item_idx] = ready - - def clear_item_cache_key(self, item_idx: int) -> Hashable: - """Clear and return one item's cache key.""" - cache_key = self.item_cache_keys[item_idx] - if cache_key is None: - raise RuntimeError(f"MM item {item_idx} has no cache key") - self.item_cache_keys[item_idx] = None - self.item_ready[item_idx] = False - return cache_key - - def mark_cache_key_ready(self, cache_key: Hashable) -> None: - """Mark every item using `cache_key` as ready.""" - self.mark_cache_keys_ready((cache_key, )) - - def mark_cache_keys_ready(self, cache_keys: Container[Hashable]) -> None: - """Mark every item using one of `cache_keys` as ready.""" - for item_idx, bound_cache_key in enumerate(self.item_cache_keys): - if bound_cache_key is not None and bound_cache_key in cache_keys: - self.item_ready[item_idx] = True - - def pop_all_cache_keys(self) -> List[Hashable]: - """Return and clear all cache keys, keeping prompt order and duplicates.""" - cache_keys = [ - cache_key for cache_key in self.item_cache_keys - if cache_key is not None - ] - self.item_cache_keys = [None] * len(self.item_cache_keys) - self.item_ready = [False] * len(self.item_ready) - return cache_keys - - if TYPE_CHECKING: from ..tensor_lru_cache import TensorLRUCache @@ -885,9 +734,14 @@ def __init__( self.py_lora_path: str | None = kwargs.pop("py_lora_path", None) self.py_lora_ckpt_source: str | None = kwargs.pop( "py_lora_ckpt_source", None) - # Multimodal data - self.py_multimodal_data = kwargs.pop("py_multimodal_data", None) - self.py_mm_item_order = kwargs.pop("py_mm_item_order", None) + # The request owns one MM aggregate. Compatibility properties below + # expose the existing field names without duplicate data or state. + mm_data = kwargs.pop("py_multimodal_data", None) + mm_item_order = kwargs.pop("py_mm_item_order", None) + self.py_multimodal_params: MultimodalParams | None = ( + MultimodalParams(multimodal_data=mm_data, + mm_item_order=mm_item_order) + if mm_data is not None or mm_item_order is not None else None) # Per-item MM encoder state, created at admission by # `initialize_multimodal_encoder_request` when raw MM payloads must # run through the item scheduler; `None` otherwise (also the @@ -1287,7 +1141,65 @@ def finish_by(self, reason: FinishReason, beam: int) -> None: self.state = LlmRequestState.GENERATION_COMPLETE self.set_finished_reason(reason, beam) + @property + def py_multimodal_data(self) -> dict[str, Any] | None: + params = self.py_multimodal_params + return params.multimodal_data if params is not None else None + + @py_multimodal_data.setter + def py_multimodal_data(self, data: dict[str, Any] | None) -> None: + if self.py_multimodal_params is None: + if data is not None: + self.py_multimodal_params = MultimodalParams( + multimodal_data=data) + else: + self.py_multimodal_params.multimodal_data = data + + @property + def py_mm_item_order(self) -> list[dict[str, Any]] | None: + params = self.py_multimodal_params + return params.mm_item_order if params is not None else None + + @py_mm_item_order.setter + def py_mm_item_order(self, order: list[dict[str, Any]] | None) -> None: + if self.py_multimodal_params is None: + if order is None: + return + self.py_multimodal_params = MultimodalParams() + self.py_multimodal_params.mm_item_order = order + + @property + def py_mm_encoder_state(self) -> MultimodalEncoderRequestState | None: + params = self.py_multimodal_params + return params.encoder_state if params is not None else None + + @py_mm_encoder_state.setter + def py_mm_encoder_state( + self, state: MultimodalEncoderRequestState | None) -> None: + if self.py_multimodal_params is None: + if state is None: + return + self.py_multimodal_params = MultimodalParams() + self.py_multimodal_params.encoder_state = state + + @property + def py_mm_encoder_event(self) -> torch.cuda.Event | None: + params = self.py_multimodal_params + return params.encoder_event if params is not None else None + + @py_mm_encoder_event.setter + def py_mm_encoder_event(self, event: torch.cuda.Event | None) -> None: + if self.py_multimodal_params is None: + if event is None: + return + self.py_multimodal_params = MultimodalParams() + self.py_multimodal_params.encoder_event = event + def create_child_request(self, child_id): + if (self.py_mm_encoder_state is not None + or self.py_mm_encoder_event is not None): + raise RuntimeError( + "Child requests must be created before MM encoder admission") child = super().create_child_request(child_id) py_request = LlmRequest(llm_request=child) @@ -1370,13 +1282,15 @@ def get_multimodal_encoder_token_lengths( """Return per-item physical encoder attention-token costs. Field-level validation happens inside - `get_multimodal_encoder_item_metadata`; only the request-level + `MultimodalItemMetadata.validate`; only the request-level cross-check against `multimodal_embedding_lengths` lives here. """ - item_metadata = get_multimodal_encoder_item_metadata( - request.py_multimodal_data) + item_metadata = get_multimodal_item_metadata(request.py_multimodal_data) if item_metadata is None: return None + item_metadata.validate() + if item_metadata.encoder_token_lengths is None: + return None embedding_lengths = get_multimodal_embedding_lengths(request) if embedding_lengths is None: raise ValueError("Multimodal encoder item scheduling requires " @@ -1467,6 +1381,9 @@ def initialize_multimodal_encoder_request( (see `get_mm_max_resident_embeddings`). """ mm_data = request.py_multimodal_data + if isinstance(mm_data, dict): + normalize_multimodal_item_metadata( + mm_data, getattr(request, "py_mm_item_order", None)) has_raw_payload = isinstance(mm_data, dict) and any( isinstance(mm_data.get(modality), dict) for modality in ("image", "video", "audio")) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f80b3aaed8c1..0f8723defc94 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -45,8 +45,8 @@ from tensorrt_llm.bindings.internal.batch_manager import (LlmRequestType, ReqIdsSet) from tensorrt_llm.executor.request import TruncateKVCacheRequest -from tensorrt_llm.inputs.multimodal import strip_mm_data_for_generation -from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata +from tensorrt_llm.inputs.multimodal import (get_multimodal_item_metadata, + strip_mm_data_for_generation) from tensorrt_llm.llmapi.llm_args import (CapacitySchedulerPolicy, ExecutorMemoryType, PeftCacheConfig, PrefillCudaGraphBackend, @@ -6705,8 +6705,10 @@ def consume_pending(costs, pending_indices=None): continue mm_data = getattr(queue_item.request, "py_multimodal_data", None) try: - item_metadata = (get_multimodal_encoder_item_metadata(mm_data) + item_metadata = (get_multimodal_item_metadata(mm_data) if isinstance(mm_data, dict) else None) + if item_metadata is not None: + item_metadata.validate() except (TypeError, ValueError): # Admission only estimates capacity. Pass malformed requests # through so the normal request-scoped validation path can @@ -8645,6 +8647,9 @@ def _release_multimodal_resources(self, request: LlmRequest) -> None: for cache_key in cache_keys: encoder_cache.release(cache_key) + params = getattr(request, "py_multimodal_params", None) + if params is not None: + params.encoder_event = None mm_data = getattr(request, "py_multimodal_data", None) if mm_data: strip_mm_data_for_generation(mm_data) diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py index bd6559311765..7d0eb769cbe9 100644 --- a/tensorrt_llm/inputs/__init__.py +++ b/tensorrt_llm/inputs/__init__.py @@ -6,7 +6,8 @@ from .evs import (compute_retained_tokens_count, compute_retained_tokens_from_tubelet_budget, compute_retention_mask) -from .multimodal import MultimodalEncoderItemMetadata, MultimodalInput +from .multimodal import (MultimodalEncoderItemMetadata, MultimodalInput, + MultimodalItemMetadata, MultimodalItemPartition) from .multimodal_data import AudioData, BaseModalityData, VideoData # yapf and isort conflict on the following import blocks # yapf: disable @@ -53,6 +54,8 @@ "BaseMultimodalDummyInputsBuilder", "BaseMultimodalInputProcessor", "MultimodalEncoderItemMetadata", + "MultimodalItemMetadata", + "MultimodalItemPartition", "MultimodalPlaceholderMetadata", "MultimodalPlaceholderPlacement", "ConversationMessage", diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index ffce49a64080..4137206c1275 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -2,7 +2,7 @@ # Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Common multimodal data contracts and layout operations. -`MultimodalEncoderItemMetadata` describes request items. `MultimodalParams` +`MultimodalItemMetadata` describes request items. `MultimodalParams` references processed payloads and the `MultimodalRuntimeData` view for the current forward. `MultimodalInput` carries the executor boundary representation. `MultimodalItemPartition` describes item axes and packed ranges within payloads. @@ -10,9 +10,10 @@ """ from dataclasses import dataclass, field +from enum import Enum, auto from itertools import accumulate -from typing import (Any, Dict, List, NamedTuple, Optional, Sequence, Tuple, - Union, cast) +from typing import (Any, Container, Dict, Hashable, List, NamedTuple, Optional, + Sequence, Tuple, Union, cast) import numpy as np import torch @@ -190,6 +191,158 @@ def select(self, ) +@dataclass +class MultimodalItemMetadata: + """Prompt-ordered items and their partitions in processed model inputs. + + Encoder costs and output rows are optional: their presence does not enable + item scheduling. Models keep their existing scheduling capabilities. + """ + + item_refs: list[tuple[str, int]] + encoder_token_lengths: list[int] | None = None + output_embedding_lengths: list[int] | None = None + input_partitions: dict[str, dict[str, MultimodalItemPartition]] = field( + default_factory=dict) + + @classmethod + def from_encoder_metadata( + cls, metadata: MultimodalEncoderItemMetadata + ) -> "MultimodalItemMetadata": + """Adapt a legacy producer without copying its CPU lists.""" + metadata.validate() + return cls(metadata.item_refs, metadata.encoder_token_lengths, + metadata.output_embedding_lengths) + + def validate(self) -> None: + """Validate producer data once before dispatch or executor admission.""" + if not isinstance(self.item_refs, list): + raise TypeError("Multimodal item references must be a list") + indices_by_modality: dict[str, list[int]] = {} + for ref in self.item_refs: + if (not isinstance(ref, tuple) or len(ref) != 2 + or not isinstance(ref[0], str) + or not isinstance(ref[1], int)): + raise TypeError( + "Multimodal item references must be (str, int) pairs") + indices_by_modality.setdefault(ref[0], []).append(ref[1]) + for indices in indices_by_modality.values(): + if sorted(indices) != list(range(len(indices))): + raise ValueError( + "Multimodal item references must cover each local item once" + ) + for name in ("encoder_token_lengths", "output_embedding_lengths"): + lengths = getattr(self, name) + if lengths is None: + continue + _validate_int_list(lengths, name) + if len(lengths) != len(self.item_refs): + raise ValueError( + "Multimodal item references and lengths must align") + minimum = 1 if name == "encoder_token_lengths" else 0 + if any(length < minimum for length in lengths): + raise ValueError(f"{name} values must be at least {minimum}") + for modality, partitions in self.input_partitions.items(): + if modality not in indices_by_modality: + raise ValueError( + "Input partitions must refer to a present modality") + for partition in partitions.values(): + if not isinstance(partition, MultimodalItemPartition): + raise TypeError( + "Input partitions must be MultimodalItemPartition records" + ) + if partition.item_lengths is not None: + if (partition.item_axis is None + or len(partition.item_lengths) != len( + indices_by_modality[modality]) + or any(not isinstance(n, int) or n < 0 + for n in partition.item_lengths)): + raise ValueError( + "Packed input lengths must align with modality items" + ) + + def select(self, item_indices: Sequence[int]) -> "MultimodalItemMetadata": + """Select items and rebase their refs and packed lengths to the view.""" + local_indices: dict[str, list[int]] = {} + refs = [] + for index in item_indices: + modality, local_index = self.item_refs[index] + selected = local_indices.setdefault(modality, []) + refs.append((modality, len(selected))) + selected.append(local_index) + partitions = {} + for modality, indices in local_indices.items(): + if modality in self.input_partitions: + partitions[modality] = { + name: (partition if partition.item_lengths is None else + MultimodalItemPartition( + partition.item_axis, + tuple(partition.item_lengths[i] + for i in indices))) + for name, partition in + self.input_partitions[modality].items() + } + return MultimodalItemMetadata( + refs, + ([self.encoder_token_lengths[i] for i in item_indices] + if self.encoder_token_lengths is not None else None), + ([self.output_embedding_lengths[i] for i in item_indices] + if self.output_embedding_lengths is not None else None), + partitions, + ) + + +def get_multimodal_item_metadata( + multimodal_data: dict[str, Any] | None, ) -> MultimodalItemMetadata | None: + """Read item metadata; producer/admission boundaries validate its values.""" + if multimodal_data is None: + return None + if not isinstance(multimodal_data, dict): + raise TypeError("multimodal_data must be a dict") + metadata = multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) + if isinstance(metadata, MultimodalEncoderItemMetadata): + return MultimodalItemMetadata.from_encoder_metadata(metadata) + if metadata is not None and not isinstance(metadata, + MultimodalItemMetadata): + raise TypeError( + f"{MULTIMODAL_ENCODER_ITEM_METADATA_KEY} must be a MultimodalItemMetadata" + ) + return metadata + + +def normalize_multimodal_item_metadata( + multimodal_data: dict[str, Any], + mm_item_order: list[dict[str, Any]] | None = None, +) -> MultimodalItemMetadata | None: + """Adapt legacy order at the input boundary, without guessing mixed order.""" + metadata = get_multimodal_item_metadata(multimodal_data) + if metadata is None: + lengths = multimodal_data.get("multimodal_embedding_lengths") + if lengths is None: + return None + if mm_item_order: + refs = [(entry["modality"], entry["index"]) + for entry in mm_item_order] + else: + present = [ + m for m in MEDIA_MODALITIES + if multimodal_data.get(m) is not None + ] + if len(present) != 1: + return None + refs = [(present[0], i) for i in range(len(lengths))] + metadata = MultimodalItemMetadata(refs, + output_embedding_lengths=lengths) + metadata.validate() + if mm_item_order and metadata.item_refs != [ + (entry["modality"], entry["index"]) for entry in mm_item_order + ]: + raise ValueError( + "Multimodal item references do not match the prompt item order") + multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata + return metadata + + def get_multimodal_encoder_item_metadata( multimodal_data: Optional[Dict[str, Any]], ) -> Optional[MultimodalEncoderItemMetadata]: @@ -201,6 +354,13 @@ def get_multimodal_encoder_item_metadata( metadata = multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) if metadata is None: return None + if isinstance(metadata, MultimodalItemMetadata): + if (metadata.encoder_token_lengths is None + or metadata.output_embedding_lengths is None): + return None + metadata = MultimodalEncoderItemMetadata( + metadata.item_refs, metadata.encoder_token_lengths, + metadata.output_embedding_lengths) if not isinstance(metadata, MultimodalEncoderItemMetadata): raise TypeError(f"{MULTIMODAL_ENCODER_ITEM_METADATA_KEY} must be a " "MultimodalEncoderItemMetadata") @@ -208,6 +368,156 @@ def get_multimodal_encoder_item_metadata( return metadata +class MultimodalEncoderProgress(Enum): + """Readiness derived from the request's cached item outputs.""" + + PENDING = auto() + """No MM item of the request has an encoder output yet.""" + + PARTIAL = auto() + """Some but not all items are encoded; the request must continue on the + item-scheduling path until every item has an output.""" + + READY = auto() + """No further encoder work is needed: every item is encoded, a + precomputed embedding was supplied, or the request has no MM payload.""" + + +class _Unset: + """Sentinel for a memo that has not been computed yet. + + Distinct from `None`, which is a real result: the request cannot build + stable cache keys and its items always encode. + """ + + +_UNSET = _Unset() + + +@dataclass +class MultimodalEncoderRequestState: + """Tracks the encoder-cache key for each MM item in a request. + + The state is created when an item-scheduled request enters the executor. + Other requests leave `py_mm_encoder_state` as `None`. + + Encoder outputs live only in the model-owned `TensorLRUCache`. This object + stores one cache key and one ready flag per item, in prompt order. On the + item-scheduled executor, each non-empty slot holds one reference. Two + identical items may point to the same cache entry, but each item still + holds and releases its own reference. The request does not store another + output tensor or combined output buffer. + + Cleanup clears the slots before releasing their references, so repeated + request cleanup is safe. + """ + + embedding_lengths: List[int] + """Expected encoder-output rows for each item, in prompt order.""" + + encoder_token_lengths: List[int] + """Validated encoder attention-token cost of each atomic item. + + Admission copies these values from the request metadata so scheduling + never has to re-validate user input inside the scheduler loop. + """ + + item_ready: List[bool] + """Whether each item's bound cache entry is ready, in prompt order.""" + + item_cache_keys: List[Optional[Hashable]] = field(default_factory=list) + """Cache keys held by this request, in prompt order. + + A non-empty slot holds one reference. `item_ready` tells whether the entry + already contains its encoder output. + """ + + stable_item_cache_keys: "list[Hashable] | _Unset | None" = _UNSET + """Cached per-item reuse keys. + + `_UNSET` means the keys have not been computed. `None` means stable keys + cannot be built, so this request uses request-local cache keys instead. + """ + + @classmethod + def from_embedding_lengths( + cls, + embedding_lengths: List[int], + *, + encoder_token_lengths: Optional[List[int]] = None + ) -> "MultimodalEncoderRequestState": + if encoder_token_lengths is None: + encoder_token_lengths = embedding_lengths + return cls(embedding_lengths=list(embedding_lengths), + encoder_token_lengths=list(encoder_token_lengths), + item_ready=[False] * len(embedding_lengths), + item_cache_keys=[None] * len(embedding_lengths)) + + def __post_init__(self) -> None: + if not self.item_cache_keys: + self.item_cache_keys = [None] * len(self.item_ready) + if not (len(self.embedding_lengths) == len(self.encoder_token_lengths) + == len(self.item_ready) == len(self.item_cache_keys)): + raise ValueError("MM encoder token and embedding lengths must have " + "exactly one cache key per item slot") + + @property + def num_items(self) -> int: + return len(self.item_ready) + + @property + def progress(self) -> MultimodalEncoderProgress: + if all(self.item_ready): + return MultimodalEncoderProgress.READY + if any(self.item_ready): + return MultimodalEncoderProgress.PARTIAL + return MultimodalEncoderProgress.PENDING + + def pending_item_indices(self) -> List[int]: + """Indices of items that still need an encoder output, prompt order.""" + return [ + item_idx for item_idx, ready in enumerate(self.item_ready) + if not ready + ] + + def set_item_cache_key(self, item_idx: int, cache_key: Hashable, *, + ready: bool) -> None: + """Assign a cache key to one item.""" + if self.item_cache_keys[item_idx] is not None: + raise RuntimeError(f"MM item {item_idx} already has a cache key") + self.item_cache_keys[item_idx] = cache_key + self.item_ready[item_idx] = ready + + def clear_item_cache_key(self, item_idx: int) -> Hashable: + """Clear and return one item's cache key.""" + cache_key = self.item_cache_keys[item_idx] + if cache_key is None: + raise RuntimeError(f"MM item {item_idx} has no cache key") + self.item_cache_keys[item_idx] = None + self.item_ready[item_idx] = False + return cache_key + + def mark_cache_key_ready(self, cache_key: Hashable) -> None: + """Mark every item using `cache_key` as ready.""" + self.mark_cache_keys_ready((cache_key, )) + + def mark_cache_keys_ready(self, cache_keys: Container[Hashable]) -> None: + """Mark every item using one of `cache_keys` as ready.""" + for item_idx, bound_cache_key in enumerate(self.item_cache_keys): + if bound_cache_key is not None and bound_cache_key in cache_keys: + self.item_ready[item_idx] = True + + def pop_all_cache_keys(self) -> List[Hashable]: + """Return and clear all cache keys, keeping prompt order and duplicates.""" + cache_keys = [ + cache_key for cache_key in self.item_cache_keys + if cache_key is not None + ] + self.item_cache_keys = [None] * len(self.item_cache_keys) + self.item_ready = [False] * len(self.item_ready) + return cache_keys + + def _validate_int_list(values: Any, field_name: str) -> None: """Boundary metadata must be owned Python list[int]. No tensor/tuple.""" if not isinstance(values, list): @@ -656,7 +966,7 @@ class MultimodalParams: `item_metadata` reads the request's shared item record from processed data; `get_item_layout` resolves its prompt order, including legacy formats. Runtime data describes the current forward or chunk rather than encoder - readiness and cache ownership, which remain on the executor request. + readiness and cache ownership, which belong to the request Params. Attributes: multimodal_input: Multimodal input data with hashing information. @@ -717,15 +1027,22 @@ class MultimodalParams: repr=False, compare=False) + encoder_state: MultimodalEncoderRequestState | None = field(default=None, + init=False, + repr=False, + compare=False) + """Executor-owned cache references. Forward and chunk views leave this unset.""" + def __post_init__(self): - """Ensure default values are properly set.""" if self.multimodal_data is None: self.multimodal_data = {} + normalize_multimodal_item_metadata(self.multimodal_data, + self.mm_item_order) @property - def item_metadata(self) -> Optional[MultimodalEncoderItemMetadata]: - """Atomic-item metadata stored in this request's processed data.""" - return get_multimodal_encoder_item_metadata(self.multimodal_data) + def item_metadata(self) -> MultimodalItemMetadata | None: + """The single item record stored in this request's processed data.""" + return get_multimodal_item_metadata(self.multimodal_data) def get_item_layout(self) -> tuple[list[tuple[str, int]], list[int]]: """Resolve prompt-order item refs and output lengths for all modalities. @@ -741,6 +1058,10 @@ def get_item_layout(self) -> tuple[list[tuple[str, int]], list[int]]: self.multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY), ) if metadata is not None: + if metadata.output_embedding_lengths is None: + raise ValueError( + "Encoder output assembly requires item output embedding lengths" + ) return metadata.item_refs, metadata.output_embedding_lengths lengths = self.multimodal_data.get("multimodal_embedding_lengths") or [] diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 3675fa7721b8..55f2af9fc88d 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -37,14 +37,16 @@ from .multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY from .multimodal import \ MultimodalEncoderItemMetadata as MultimodalEncoderItemMetadata -from .multimodal import (MultimodalInput, _as_cpu_tensor, _compute_mm_masks, +from .multimodal import (MultimodalInput, MultimodalItemMetadata, + _as_cpu_tensor, _compute_mm_masks, _find_mm_embedding_lengths_from_masks, _find_mm_token_runs_from_mask, _find_mm_token_start_pos_from_masks, apply_mm_hashes, default_hasher, find_mm_token_lengths) from .multimodal import \ get_multimodal_encoder_item_metadata as get_multimodal_encoder_item_metadata -from .multimodal import hexdigest_to_int32, validate_mm_inputs +from .multimodal import (hexdigest_to_int32, normalize_multimodal_item_metadata, + validate_mm_inputs) from .multimodal_data import MEDIA_MODALITIES, serialize_item from .prefix_token_cache import create_prefix_token_cache from .tokenization import tokenize_prompt @@ -162,6 +164,31 @@ def get_mm_encoder_item_metadata( """ return None + def get_mm_item_metadata( + self, + prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + mm_item_order: Optional[List[Dict[str, Any]]] = None, + ) -> Optional[MultimodalItemMetadata]: + """Describe processed items independently of encoder scheduling support. + + Keep legacy encoder-cost producers at this boundary. Precomputed + embeddings and other models can supply order and rows without costs. + """ + if (multimodal_data.get("multimodal_embedding") is None and any( + isinstance(multimodal_data.get(m), dict) + for m in MEDIA_MODALITIES)): + metadata = self.get_mm_encoder_item_metadata( + prompt_token_ids, multimodal_data) + if metadata is not None: + if not isinstance(metadata, MultimodalEncoderItemMetadata): + raise TypeError( + "get_mm_encoder_item_metadata() must return a " + "MultimodalEncoderItemMetadata or None") + return MultimodalItemMetadata.from_encoder_metadata(metadata) + return normalize_multimodal_item_metadata(multimodal_data, + mm_item_order) + def __init__(self, model_path, config, @@ -1466,55 +1493,54 @@ def input_processor_wrapper( return prompt_token_ids, extra_processed_inputs multimodal_data = extra_processed_inputs.get("multimodal_data") - if (not isinstance(multimodal_data, dict) - or multimodal_data.get("multimodal_embedding") is not None): - return prompt_token_ids, extra_processed_inputs - - # A processor participates in item scheduling by overriding - # `get_mm_encoder_item_metadata`; the base returns `None`, so its - # result gates routing without a separate capability flag. `getattr` - # because unregistered/text-only models wrap a `DefaultInputProcessor` - # that lacks the method entirely. - get_item_metadata = getattr(input_processor, - "get_mm_encoder_item_metadata", None) - has_raw_payload = any( - isinstance(multimodal_data.get(modality), dict) - for modality in MEDIA_MODALITIES) - if not has_raw_payload or get_item_metadata is None: + if not isinstance(multimodal_data, dict): return prompt_token_ids, extra_processed_inputs - item_metadata = get_item_metadata(prompt_token_ids, multimodal_data) + mm_item_order = inputs.get("mm_item_order") + get_item_metadata = getattr(input_processor, "get_mm_item_metadata", + None) + if get_item_metadata is not None: + item_metadata = get_item_metadata(prompt_token_ids, multimodal_data, + mm_item_order) + else: + # Registered processors inherit the generic hook. Accept older + # standalone processors at the same boundary. + get_encoder_metadata = getattr(input_processor, + "get_mm_encoder_item_metadata", None) + if (get_encoder_metadata is not None + and multimodal_data.get("multimodal_embedding") is None + and any( + isinstance(multimodal_data.get(m), dict) + for m in MEDIA_MODALITIES)): + legacy_metadata = get_encoder_metadata(prompt_token_ids, + multimodal_data) + if legacy_metadata is not None: + if not isinstance(legacy_metadata, + MultimodalEncoderItemMetadata): + raise TypeError( + "get_mm_encoder_item_metadata() must return a " + "MultimodalEncoderItemMetadata or None") + multimodal_data[ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = legacy_metadata + item_metadata = normalize_multimodal_item_metadata( + multimodal_data, mm_item_order) if item_metadata is None: - # Processor does not emit item metadata for this input; fall back - # to the non-item-scheduled path. return prompt_token_ids, extra_processed_inputs - if not isinstance(item_metadata, MultimodalEncoderItemMetadata): + if not isinstance(item_metadata, MultimodalItemMetadata): raise TypeError( - "get_mm_encoder_item_metadata() must return a " - "MultimodalEncoderItemMetadata or None for raw multimodal " - f"payloads, got {type(item_metadata).__name__}") + "get_mm_item_metadata() must return a MultimodalItemMetadata or None" + ) item_metadata.validate() - - mm_item_order = inputs.get("mm_item_order") - if mm_item_order and item_metadata.item_refs != [ - (entry["modality"], entry["index"]) for entry in mm_item_order - ]: - raise ValueError("Multimodal encoder item references do not match " - "the prompt item order") - multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = item_metadata - existing_embedding_lengths = multimodal_data.get( - "multimodal_embedding_lengths") - if existing_embedding_lengths is not None: - if not isinstance(existing_embedding_lengths, list): - raise TypeError("multimodal_embedding_lengths must be a list") - if (existing_embedding_lengths - != item_metadata.output_embedding_lengths): + normalize_multimodal_item_metadata(multimodal_data, mm_item_order) + embedding_lengths = item_metadata.output_embedding_lengths + if embedding_lengths is not None: + existing = multimodal_data.get("multimodal_embedding_lengths") + if existing is not None and existing != embedding_lengths: raise ValueError( "Computed multimodal encoder embedding lengths " "do not match the existing prompt metadata") - multimodal_data[ - "multimodal_embedding_lengths"] = item_metadata.output_embedding_lengths + multimodal_data["multimodal_embedding_lengths"] = embedding_lengths return prompt_token_ids, extra_processed_inputs return input_processor_wrapper diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py index 6b1830c83395..0c7d6ac2b0f5 100644 --- a/tests/unittest/_torch/executor/test_multimodal_scheduler.py +++ b/tests/unittest/_torch/executor/test_multimodal_scheduler.py @@ -13,10 +13,7 @@ record_output, ) -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( - MultimodalEncoderContractError, - MultimodalModelMixin, -) +from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( @@ -44,6 +41,7 @@ from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, MultimodalParams, strip_mm_encoder_inputs, ) @@ -64,6 +62,89 @@ def schedule_request(self, requests): return list(requests), [], [] +def test_multimodal_request_aliases_share_one_aggregate(): + request = make_mm_request(1, [4, 8]) + params = request.py_multimodal_params + assert request.py_multimodal_data is params.multimodal_data + assert request.py_mm_encoder_state is params.encoder_state + order = [{"modality": "image", "index": 0}, {"modality": "image", "index": 1}] + request.py_mm_item_order = order + assert params.mm_item_order is order + params.mm_item_order = list(reversed(order)) + assert request.py_mm_item_order is params.mm_item_order + event = object() + request.py_mm_encoder_event = event + assert params.encoder_event is event + params.encoder_event = None + assert request.py_mm_encoder_event is None + replacement = dict(request.py_multimodal_data) + request.py_multimodal_data = replacement + assert params.multimodal_data is replacement + assert request.py_mm_encoder_state is params.encoder_state + assert params.encoder_state is not None + assert params.multimodal_runtime is None + assert all( + name not in request.__dict__ + for name in ( + "py_multimodal_data", + "py_mm_item_order", + "py_mm_encoder_state", + "py_mm_encoder_event", + ) + ) + + +def test_text_requests_do_not_allocate_multimodal_params(): + request = make_llm_request(1) + request.py_mm_encoder_state = None + request.py_mm_encoder_event = None + request.py_mm_item_order = None + request.py_multimodal_data = None + assert request.py_multimodal_params is None + + +def test_generic_metadata_does_not_enable_item_scheduling(): + metadata = MultimodalItemMetadata([("image", 0)], output_embedding_lengths=[1]) + request = make_llm_request( + 1, + { + "image": {"pixel_values": torch.zeros(1, 2)}, + "multimodal_embedding_lengths": [1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, + }, + ) + assert request.py_mm_encoder_state is None + assert get_multimodal_encoder_token_lengths(request) is None + with pytest.raises(ValueError, match="when item-level encoder scheduling is enabled"): + initialize_multimodal_encoder_request(request, 32) + request.py_multimodal_data["multimodal_embedding"] = torch.zeros(1, 2) + initialize_multimodal_encoder_request(request, 32) + assert request.py_multimodal_params.encoder_state is None + + +def test_child_creation_cannot_duplicate_admitted_cache_ownership(): + request = make_mm_request(1, [4]) + with pytest.raises(RuntimeError, match="before MM encoder admission"): + request.create_child_request(2) + + +def test_cleanup_releases_duplicate_references_and_drops_state_and_event(): + request = make_mm_request(1, [4, 4]) + state = request.py_mm_encoder_state + cache_key = ("same", 0) + state.set_item_cache_key(0, cache_key, ready=True) + state.set_item_cache_key(1, cache_key, ready=True) + request.py_mm_encoder_event = object() + cache = Mock() + executor = SimpleNamespace(model_engine=SimpleNamespace(mm_encoder_cache=cache)) + PyExecutor._release_multimodal_resources(executor, request) + PyExecutor._release_multimodal_resources(executor, request) + assert cache.release.call_args_list == [((cache_key,),), ((cache_key,),)] + assert state.item_cache_keys == [None, None] + assert request.py_multimodal_params.encoder_state is None + assert request.py_multimodal_params.encoder_event is None + + class _RejectMultimodalCapacityScheduler: def schedule_request(self, requests): fitting = [request for request in requests if request.py_mm_encoder_state is None] @@ -233,9 +314,8 @@ def test_mm_encoder_token_lengths_distinguishes_missing_and_invalid_data(): assert get_multimodal_encoder_token_lengths(request) is None - request.py_multimodal_data = [] with pytest.raises(TypeError, match="multimodal_data must be a dict"): - get_multimodal_encoder_token_lengths(request) + request.py_multimodal_data = [] def test_mm_encoder_readiness_is_derived_from_item_state(): @@ -997,17 +1077,13 @@ def test_prepare_multimodal_encoder_inputs_slices_before_device_transfer(): def test_prepare_multimodal_encoder_inputs_rejects_invalid_metadata_types(): - multimodal_param = MultimodalParams( - multimodal_data={ - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: ("image", 0), - "multimodal_embedding_lengths": [1], - } - ) - - with pytest.raises( - MultimodalEncoderContractError, match="must be a MultimodalEncoderItemMetadata" - ): - MultimodalModelMixin().prepare_multimodal_encoder_inputs([(multimodal_param, 0)]) + with pytest.raises(TypeError, match="must be a MultimodalItemMetadata"): + MultimodalParams( + multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: ("image", 0), + "multimodal_embedding_lengths": [1], + } + ) def test_strip_mm_encoder_inputs_preserves_embedding_and_runtime_metadata(): diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index ea010de4bda5..26f9528d3760 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/tests/unittest/_torch/executor/test_request_utils.py @@ -33,12 +33,65 @@ from tensorrt_llm.bindings import executor as trtllm from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.executor.result import Logprob +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, + MultimodalItemPartition, +) from tensorrt_llm.mapping import CpType from tensorrt_llm.sampling_params import LogprobMode pytestmark = pytest.mark.cpu_only +def test_multimodal_payload_broadcast_and_children_keep_the_same_contract(): + import pickle + + metadata = MultimodalItemMetadata( + [("audio", 0), ("image", 0)], + output_embedding_lengths=[1, 2], + input_partitions={"audio": {"features": MultimodalItemPartition(1, (3,))}}, + ) + payload = { + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, + "audio": {"features": torch.arange(6).reshape(2, 3)}, + "image": {}, + "multimodal_embedding_lengths": [1, 2], + } + source = trtllm.Request( + input_token_ids=[1, 2, 3], + max_tokens=2, + sampling_config=trtllm.SamplingConfig(num_return_sequences=2), + ) + source.py_multimodal_data = payload + source.py_conversation_params = None + objects = RequestBroadcaster._collect_py_objects(None, [RequestQueueItem(42, source)]) + # Reproduce the existing Python-object rank handoff, without adding the + # request aggregate, CUDA event, or cache state to its wire format. + restored_objects = pickle.loads(pickle.dumps(objects)) + restored = trtllm.Request( + input_token_ids=[1, 2, 3], + max_tokens=2, + sampling_config=trtllm.SamplingConfig(num_return_sequences=2), + ) + queue_items = [RequestQueueItem(42, restored)] + attach_py_objects_to_requests(queue_items, restored_objects) + request = executor_request_to_llm_request( + 42, restored, child_req_ids=[43], exclude_last_generation_logits=False + ) + params = request.py_multimodal_params + child = request.child_requests[0].py_multimodal_params + assert params.item_metadata == child.item_metadata == metadata + assert child is not params + assert child.item_metadata is not params.item_metadata + assert params.encoder_state is child.encoder_state is None + assert params.multimodal_runtime is child.multimodal_runtime is None + assert params.item_metadata.item_refs == [("audio", 0), ("image", 0)] + torch.testing.assert_close( + child.multimodal_data["audio"]["features"], payload["audio"]["features"] + ) + + @pytest.fixture def attention_dp_config(): """Create a config dict for attention DP testing.""" diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 9441aefde835..b20149e840ed 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -9,6 +9,7 @@ """ from functools import partial +from types import SimpleNamespace import pytest import torch @@ -24,7 +25,9 @@ from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, + MultimodalEncoderRequestState, MultimodalInput, + MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams, MultimodalRuntimeData, @@ -37,6 +40,84 @@ ITEM_ROWS = {"image": [2, 3], "video": [4]} +def test_qwen_processor_produces_prompt_order_and_packed_partitions(): + from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLInputProcessorBase + + processor = object.__new__(Qwen3VLInputProcessorBase) + processor._config = SimpleNamespace( + image_token_id=10, + video_token_id=11, + vision_start_token_id=12, + vision_end_token_id=13, + vision_config=SimpleNamespace(spatial_merge_size=2), + ) + data = { + "image": {"pixel_values": torch.zeros(8, 2), "image_grid_thw": torch.tensor([[1, 2, 4]])}, + "video": { + "pixel_values_videos": torch.zeros(16, 2), + "video_grid_thw": torch.tensor([[2, 2, 4]]), + }, + } + metadata = processor.get_mm_item_metadata([12, 11, 13, 12, 10, 13], data) + metadata.validate() + assert metadata.item_refs == [("video", 0), ("image", 0)] + assert metadata.encoder_token_lengths == [16, 8] + assert metadata.output_embedding_lengths == [4, 2] + assert metadata.input_partitions["image"]["pixel_values"].item_lengths == (8,) + assert metadata.input_partitions["video"]["pixel_values_videos"].item_lengths == (16,) + + +def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + + processor = object.__new__(Gemma4InputProcessor) + data = { + "image": {"pixel_values": torch.zeros(1, 4, 6), "image_position_ids": torch.zeros(1, 4, 2)}, + "audio": {"audio_features": torch.zeros(1, 5, 8), "audio_features_mask": torch.ones(1, 5)}, + "multimodal_embedding_lengths": [2, 4], + } + order = [{"modality": "audio", "index": 0}, {"modality": "image", "index": 0}] + metadata = processor.get_mm_item_metadata([], data, order) + metadata.validate() + assert metadata.item_refs == [("audio", 0), ("image", 0)] + assert metadata.encoder_token_lengths is None + assert metadata.input_partitions["audio"]["audio_features_mask"].item_axis == 0 + assert metadata.input_partitions["image"]["image_position_ids"].item_axis == 0 + + +def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): + from tensorrt_llm._torch.models.modeling_nemotron_h_multimodal import ( + NemotronHMultimodalInputProcessor, + ) + + processor = object.__new__(NemotronHMultimodalInputProcessor) + audio = { + "input_audio_features": torch.arange(30).reshape(3, 5, 2), + "feature_attention_mask": torch.ones(3, 5), + "audio_num_clips": torch.tensor([2, 1]), + } + data = {"audio": audio, "multimodal_embedding_lengths": [4, 2]} + metadata = processor.get_mm_item_metadata([], data) + metadata.validate() + params = MultimodalParams(multimodal_data=data) + selected = MultimodalModelMixin().build_multimodal_encoder_input(params, [1, 0], "audio") + MultimodalModelMixin._apply_metadata_slice(selected, params, [1, 0]) + assert selected.item_metadata.encoder_token_lengths is None + partition = selected.item_metadata.input_partitions["audio"]["input_audio_features"] + assert partition.item_lengths == (1, 2) + torch.testing.assert_close( + selected.multimodal_data["audio"]["input_audio_features"], + audio["input_audio_features"][[2, 0, 1]], + ) + assert selected.encoder_state is None + video = {"audio": audio, "video_size": [[2, 1, 32, 32]]} + video_data = {"video": video, "multimodal_embedding_lengths": [6]} + video_metadata = processor.get_mm_item_metadata([], video_data) + assert video_metadata.item_refs == [("video", 0)] + assert "audio" not in video_metadata.input_partitions + assert video_data["video"]["audio"] is audio + + def _grid_for_rows(rows: int) -> list[int]: """A (t, h, w) grid whose product is the item's patch count.""" return [1, 1, rows] @@ -201,7 +282,7 @@ def get_multimodal_item_partitions(self, param, modality): torch.testing.assert_close(residual.multimodal_data["audio"]["mask"], mask[[5, 0, 1]]) assert residual.multimodal_data["audio"]["frame_counts"] == [1, 2] assert residual.multimodal_data["audio"]["shared_stats"] is shared - assert residual.item_metadata == MultimodalEncoderItemMetadata( + assert residual.item_metadata == MultimodalItemMetadata( [("audio", 0), ("audio", 1)], [1, 2], [1, 1] ) assert request.item_metadata.item_refs == [("audio", 0), ("audio", 1), ("audio", 2)] @@ -212,11 +293,17 @@ def test_sliced_items_survive_the_grouped_encoder_path(): grouped path splits to zero rows and returns nothing for every item.""" model = _GroupedEncoderModel() request = _make_request(["image", "image"]) + owner = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) + owner.set_item_cache_key(0, ("owned", 0), ready=True) + request.encoder_state = owner selected = [(request, 0), (request, 1)] - outputs, _ = _run_items(model, selected) + outputs, encoder_inputs = _run_items(model, selected) assert len(outputs) == 2 + assert request.encoder_state is owner + assert owner.item_cache_keys == [("owned", 0), None] + assert all(param.encoder_state is None for param, _, _ in encoder_inputs) for item_idx, output in enumerate(outputs): expected = _expected_rows(request, item_idx) assert output.shape == expected.shape @@ -287,7 +374,7 @@ def test_metadata_is_rebased_for_partial_and_reordered_items() -> None: assert len(encoder_inputs) == 1 residual = encoder_inputs[0][0] metadata = residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] - assert metadata == MultimodalEncoderItemMetadata([("image", 0), ("image", 1)], [3, 2], [3, 2]) + assert metadata == MultimodalItemMetadata([("image", 0), ("image", 1)], [3, 2], [3, 2]) assert residual.multimodal_data["multimodal_embedding_lengths"] == [3, 2] assert request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] is original torch.testing.assert_close(outputs[0], _expected_rows(request, 2), rtol=0, atol=0) @@ -336,7 +423,7 @@ def test_partial_cache_uses_prompt_refs_and_lengths(metadata_format, reversed_or request = _make_request(["image", "image"]) metadata = request.item_metadata if reversed_order: - metadata = MultimodalEncoderItemMetadata([("image", 1), ("image", 0)], [3, 2], [3, 2]) + metadata = MultimodalItemMetadata([("image", 1), ("image", 0)], [3, 2], [3, 2]) request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata lengths = metadata.output_embedding_lengths request.multimodal_data["multimodal_embedding_lengths"] = lengths @@ -425,23 +512,23 @@ def test_partial_cache_batches_mixed_vision_and_audio_misses(): torch.testing.assert_close(model._multimodal_encoder_cache.get(key), rows, rtol=0, atol=0) -def test_alternating_item_selection_validates_each_request_once(monkeypatch): +def test_alternating_item_selection_reuses_validated_request_metadata(monkeypatch): model = _GroupedEncoderModel() requests = [_make_request(["image", "video"] * 4) for _ in range(2)] validated = [] - validate = MultimodalEncoderItemMetadata.validate + validate = MultimodalItemMetadata.validate def record_validation(metadata): validated.append(metadata) validate(metadata) - monkeypatch.setattr(MultimodalEncoderItemMetadata, "validate", record_validation) + monkeypatch.setattr(MultimodalItemMetadata, "validate", record_validation) model.prepare_multimodal_encoder_inputs( [(request, i) for request in requests for i in range(8)] ) - assert len(validated) == len(requests) + assert not validated @pytest.mark.parametrize("shared_vision_encoder", [False, True]) diff --git a/tests/unittest/_torch/multimodal/test_encoder_group.py b/tests/unittest/_torch/multimodal/test_encoder_group.py index 41f5247b2a69..131559f681c5 100644 --- a/tests/unittest/_torch/multimodal/test_encoder_group.py +++ b/tests/unittest/_torch/multimodal/test_encoder_group.py @@ -61,17 +61,15 @@ def test_single_modality_no_manifest(self): } def test_length_mismatch_raises(self): - # ``strict=True`` inside the helper catches manifest/length divergence. - mp = _mp( - mm_item_order=[ - {"modality": "image", "index": 0}, - {"modality": "video", "index": 0}, - ], - embedding_lengths=[4], # one short - buckets={"image": {}, "video": {}}, - ) - with pytest.raises(ValueError): - _lengths_by_modality([mp.get_item_layout()], ("image", "video")) + with pytest.raises(ValueError, match="references and lengths must align"): + _mp( + mm_item_order=[ + {"modality": "image", "index": 0}, + {"modality": "video", "index": 0}, + ], + embedding_lengths=[4], + buckets={"image": {}, "video": {}}, + ) def test_cross_group_multi_modality_without_manifest_raises(self): # Multi-group models register one single-modality group per encoder. diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 023caef6505b..280d354b58b4 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -16,6 +16,7 @@ DisaggPrefillMultimodalInputs, MultimodalEncoderItemMetadata, MultimodalInput, + MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams, MultimodalRuntimeData, @@ -180,7 +181,7 @@ def test_mm_item_metadata_is_materialized_when_embedding_lengths_match(): multimodal_data = extra["multimodal_data"] item_metadata = multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] - assert item_metadata == MultimodalEncoderItemMetadata( + assert item_metadata == MultimodalItemMetadata( item_refs=[("image", 0)], encoder_token_lengths=[4], output_embedding_lengths=[2], @@ -193,7 +194,7 @@ def test_mm_item_metadata_is_materialized_when_embedding_lengths_match(): params.to_handle("multimodal_data") assert isinstance( params.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY], - MultimodalEncoderItemMetadata, + MultimodalItemMetadata, ) @@ -205,9 +206,8 @@ def test_mm_item_metadata_rejects_mismatched_embedding_lengths(): def test_mm_item_metadata_none_falls_back_to_non_item_path(): - # A processor whose get_mm_encoder_item_metadata returns None does not - # participate in item scheduling; routing falls through with no item - # metadata attached (no capability flag drives this). + # Generic item order does not imply that physical encoder costs or + # item scheduling are supported. class NoMetadataProcessor(_ItemMetadataFakeProcessor): def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): return None @@ -216,7 +216,10 @@ def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): _, extra = input_processor({"prompt": "unused"}, sampling_params=None) - assert MULTIMODAL_ENCODER_ITEM_METADATA_KEY not in extra["multimodal_data"] + metadata = extra["multimodal_data"][MULTIMODAL_ENCODER_ITEM_METADATA_KEY] + assert metadata.item_refs == [("image", 0)] + assert metadata.encoder_token_lengths is None + assert metadata.output_embedding_lengths == [2] def test_mm_item_metadata_rejects_wrong_return_type(): @@ -282,6 +285,104 @@ def test_mm_item_metadata_legacy_import_and_pickle() -> None: assert pickle.loads(pickle.dumps(metadata)) == metadata +@pytest.mark.parametrize("indices", [[3, 0, 2], [2, 0], [-1, 0], [], [1, 1]]) +def test_item_metadata_selection_keeps_partitions_in_payload_order(indices) -> None: + metadata = MultimodalItemMetadata( + [("image", 1), ("audio", 0), ("image", 0), ("audio", 1)], + output_embedding_lengths=[3, 1, 2, 4], + input_partitions={ + "image": {"pixels": MultimodalItemPartition(item_lengths=(2, 3))}, + "audio": { + "features": MultimodalItemPartition(1, (1, 4)), + "shared": MultimodalItemPartition(item_axis=None), + }, + }, + ) + metadata.validate() + values = {"image": torch.arange(5), "audio": torch.arange(10).reshape(2, 5)} + selected = metadata.select(indices) + selected.validate() + assert selected.encoder_token_lengths is None + assert metadata.input_partitions["image"]["pixels"].item_lengths == (2, 3) + assert metadata.input_partitions["audio"]["features"].item_lengths == (1, 4) + for modality, partitions in selected.input_partitions.items(): + local_indices = [ + metadata.item_refs[i][1] for i in indices if metadata.item_refs[i][0] == modality + ] + name = "pixels" if modality == "image" else "features" + source_partition = metadata.input_partitions[modality][name] + value = source_partition.select(values[modality], local_indices) + # A second selection uses the view's local item indices and lengths. + first = partitions[name].select(value, [0]) + expected = source_partition.select(values[modality], local_indices[:1]) + torch.testing.assert_close(first, expected) + + +def test_unknown_item_costs_and_rows_remain_explicit() -> None: + metadata = MultimodalItemMetadata([("audio", 0)]) + metadata.validate() + params = MultimodalParams(multimodal_data={MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata}) + assert params.item_metadata is metadata + assert params.item_metadata.encoder_token_lengths is None + with pytest.raises(ValueError, match="requires item output embedding lengths"): + params.get_item_layout() + + +def test_item_metadata_roundtrips_with_shared_tensor_handles() -> None: + import pickle + + values = torch.arange(10, dtype=torch.float32).reshape(2, 5) + metadata = MultimodalItemMetadata( + [("audio", 0), ("audio", 1)], + output_embedding_lengths=[2, 3], + input_partitions={"audio": {"features": MultimodalItemPartition(1, (2, 3))}}, + ) + params = MultimodalParams( + multimodal_data={ + "audio": {"features": values}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, + } + ) + params.to_handle("multimodal_data") + assert params.item_metadata is metadata + restored = pickle.loads(pickle.dumps(params)) + restored.to_tensor("multimodal_data") + torch.testing.assert_close(restored.multimodal_data["audio"]["features"], values) + assert restored.item_metadata == metadata + assert restored.encoder_state is None + + +@pytest.mark.parametrize("precomputed", [False, True]) +def test_generic_metadata_normalizes_mixed_order_without_encoder_costs(precomputed) -> None: + from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor + + class Processor(_ItemMetadataFakeProcessor): + get_mm_item_metadata = BaseMultimodalInputProcessor.get_mm_item_metadata + + def __call__(self, inputs, sampling_params): + return [101, 102], { + "multimodal_data": { + "image": {}, + "audio": {}, + "multimodal_embedding_lengths": [1, 1], + **({"multimodal_embedding": torch.zeros(2, 4)} if precomputed else {}), + } + } + + def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): + assert not precomputed + return None + + order = [{"modality": "audio", "index": 0}, {"modality": "image", "index": 0}] + _, extra = create_input_processor_with_hash(Processor(None))( + {"prompt": "unused", "mm_item_order": order}, None + ) + metadata = extra["multimodal_data"][MULTIMODAL_ENCODER_ITEM_METADATA_KEY] + assert metadata.item_refs == [("audio", 0), ("image", 0)] + assert metadata.encoder_token_lengths is None + assert metadata.output_embedding_lengths == [1, 1] + + def test_maybe_compute_mm_embed_cumsum_populates_py_multimodal_data(): """Producer writes a flat int64 cumsum tensor at py_multimodal_data[multimodal_embed_mask_cumsum].""" @@ -637,7 +738,7 @@ def get_mm_encoder_item_metadata( assert cumsum.tolist() == expected_cumsum assert data["mm_bidirectional_blocks"] is True assert data["multimodal_embedding_lengths"] == [2] * num_items - assert data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] == MultimodalEncoderItemMetadata( + assert data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] == MultimodalItemMetadata( item_refs, [4] * num_items, [2] * num_items ) if not hashing_enabled: From 9e51942a75a78f83e08fba954c465c1475a4f3de Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:34:30 -0700 Subject: [PATCH 08/22] Unify multimodal prefill handoff with request params Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_kimi_k25.py | 26 ++- .../_torch/models/modeling_llava_next.py | 34 ++-- .../models/modeling_nemotron_h_multimodal.py | 19 +- .../_torch/models/modeling_qwen2vl.py | 58 +++--- .../_torch/models/modeling_qwen3vl.py | 41 ++-- tensorrt_llm/inputs/multimodal.py | 126 ++++++------ tensorrt_llm/inputs/registry.py | 7 +- tensorrt_llm/llmapi/llm.py | 40 +--- .../test_modeling_nemotron_h_multimodal.py | 24 +-- ...est_nemotron_h_multimodal_preprocessing.py | 18 +- .../multimodal/test_qwen3vl_disagg_prompt.py | 185 ++++++++++++++++-- tests/unittest/inputs/test_multimodal.py | 88 +++++++-- 12 files changed, 444 insertions(+), 222 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index 1a8db4251764..7d34e6fc768c 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -49,7 +49,7 @@ PreTrainedTokenizerBase, ) -from tensorrt_llm.inputs.multimodal import DisaggPrefillMultimodalInputs, MultimodalParams +from tensorrt_llm.inputs.multimodal import MultimodalItemMetadata, MultimodalParams from tensorrt_llm.mapping import Mapping from ..._utils import prefer_pinned @@ -1469,18 +1469,20 @@ def build_disagg_prefill_multimodal_inputs( self, inputs: TextPrompt, mm_handles: List[Dict[str, Any]], - ) -> DisaggPrefillMultimodalInputs: + *, + mm_hashes: list[list[int]] | None = None, + ) -> tuple[list[int], MultimodalParams]: """Build disaggregated prefill inputs from multimodal embedding handles. Args: inputs: Text prompt input container. - mm_handles: List of multimodal embedding handles from the - context phase, each containing ``tensor_size``. + mm_handles: List of multimodal embedding handles, each + containing ``tensor_size``. + mm_hashes: Optional prompt-ordered hashes for KV-cache reuse. Returns: - DisaggPrefillMultimodalInputs containing expanded token IDs, - prompt-side MM positions/lengths, exact runs, and encoder-output - embedding lengths. + Expanded token IDs and common multimodal params containing + prompt layout, item metadata, and embedding handles. """ text_prompt = inputs.get("prompt") if not text_prompt: @@ -1528,11 +1530,15 @@ def build_disagg_prefill_multimodal_inputs( expanded_ids[write_pos] = input_ids[read_pos] write_pos += 1 - return DisaggPrefillMultimodalInputs( - prompt_token_ids=expanded_ids.to(torch.int32).tolist(), + return expanded_ids.to(torch.int32).tolist(), MultimodalParams.from_embeddings( + mm_handles, + MultimodalItemMetadata( + [("image", i) for i in range(len(mm_handles))], + output_embedding_lengths=[handle["tensor_size"][0] for handle in mm_handles], + ), + multimodal_hashes=mm_hashes, multimodal_lengths=mm_token_length, multimodal_positions=mm_token_offsets, - multimodal_embedding_lengths=[mm_handle["tensor_size"][0] for mm_handle in mm_handles], multimodal_item_run_cu_offsets=list(range(len(mm_token_length) + 1)), multimodal_run_positions=mm_token_offsets, multimodal_run_lengths=mm_token_length, diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py index f062bbbfd103..d6eac6adf806 100644 --- a/tensorrt_llm/_torch/models/modeling_llava_next.py +++ b/tensorrt_llm/_torch/models/modeling_llava_next.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + import copy from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union @@ -15,7 +18,7 @@ from tensorrt_llm._torch.models.checkpoints.hf.llava_next_weight_mapper import \ LlavaNextHfWeightMapper from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg -from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, +from tensorrt_llm.inputs.multimodal import (MultimodalItemMetadata, MultimodalParams) from ...inputs import (BaseMultimodalDummyInputsBuilder, @@ -283,8 +286,12 @@ def _postprocess( return fused_input_ids, mm_features def build_disagg_prefill_multimodal_inputs( - self, inputs: Union[TextPrompt, TokensPrompt], - mm_handles: List[Dict[str, Any]]) -> DisaggPrefillMultimodalInputs: + self, + inputs: Union[TextPrompt, TokensPrompt], + mm_handles: List[Dict[str, Any]], + *, + mm_hashes: list[list[int]] | None = None, + ) -> tuple[list[int], MultimodalParams]: """ Build disaggregated prefill inputs from multimodal embedding handles. @@ -293,13 +300,12 @@ def build_disagg_prefill_multimodal_inputs( Args: inputs: Inputs containing an already tokenized text prompt or a text prompt string. mm_handles: List of multimodal embedding handles. + mm_hashes: Optional prompt-ordered hashes for KV-cache reuse. Returns: - DisaggPrefillMultimodalInputs containing expanded token IDs, - prompt-side MM positions/lengths, exact runs, and encoder-output - embedding lengths. + Expanded token IDs and common multimodal params containing + prompt layout, item metadata, and embedding handles. """ - # TODO: Move this function to the base input processor class when extending for more models text_prompt = inputs.get("prompt") prompt_token_ids = inputs.get("prompt_token_ids") if text_prompt: @@ -339,13 +345,17 @@ def build_disagg_prefill_multimodal_inputs( f"({mm_token_length[-1] + mm_token_offsets[-1]}) should be less " f"than or equal to final_length ({final_length})") - return DisaggPrefillMultimodalInputs( - prompt_token_ids=expanded_ids, + return expanded_ids, MultimodalParams.from_embeddings( + mm_handles, + MultimodalItemMetadata( + [("image", i) for i in range(len(mm_handles))], + output_embedding_lengths=[ + handle["tensor_size"][0] for handle in mm_handles + ], + ), + multimodal_hashes=mm_hashes, multimodal_lengths=mm_token_length, multimodal_positions=mm_token_offsets, - multimodal_embedding_lengths=[ - mm_handle["tensor_size"][0] for mm_handle in mm_handles - ], multimodal_item_run_cu_offsets=list(range(len(mm_token_length) + 1)), multimodal_run_positions=mm_token_offsets, diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index d9b5d9021f18..9a621fe5be7c 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -28,7 +28,6 @@ from tensorrt_llm._torch.models.checkpoints import NemotronHHfWeightMapper from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper from tensorrt_llm.inputs.multimodal import ( - DisaggPrefillMultimodalInputs, MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams, @@ -2501,8 +2500,12 @@ def call_with_text_prompt( return input_ids[0].to(torch.int32).tolist(), {"multimodal_data": multimodal_data} def build_disagg_prefill_multimodal_inputs( - self, inputs: Union[TextPrompt, TokensPrompt], mm_handles: List[Dict[str, Any]] - ) -> DisaggPrefillMultimodalInputs: + self, + inputs: Union[TextPrompt, TokensPrompt], + mm_handles: List[Dict[str, Any]], + *, + mm_hashes: list[list[int]] | None = None, + ) -> tuple[list[int], MultimodalParams]: text_prompt = inputs.get("prompt") prompt_token_ids = inputs.get("prompt_token_ids") if prompt_token_ids is None and not text_prompt: @@ -2618,11 +2621,15 @@ def build_disagg_prefill_multimodal_inputs( mm_mask, num_mm_tokens ) - return DisaggPrefillMultimodalInputs( - prompt_token_ids=expanded_ids, + return expanded_ids, MultimodalParams.from_embeddings( + mm_handles, + MultimodalItemMetadata( + [(modalities[0], i) for i in range(len(mm_handles))], + output_embedding_lengths=multimodal_embedding_lengths, + ), + multimodal_hashes=mm_hashes, multimodal_lengths=num_mm_tokens, multimodal_positions=mm_token_offsets, - multimodal_embedding_lengths=multimodal_embedding_lengths, multimodal_item_run_cu_offsets=item_run_cu_offsets, multimodal_run_positions=run_positions, multimodal_run_lengths=run_lengths, diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index f91f98ff9a0b..6ae7b8694379 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -30,8 +30,7 @@ from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, - MultimodalEncoderItemMetadata, +from tensorrt_llm.inputs.multimodal import (MultimodalEncoderItemMetadata, MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams) @@ -1209,8 +1208,8 @@ def _attach_multimodal_embeddings_impl( mm_handles = [{ "tensor_size": tuple(image_embedding.shape) } for image_embedding in image_embeddings] - prompt_token_ids = build_disagg_prefill_multimodal_inputs( - inputs, mm_handles).prompt_token_ids + prompt_token_ids, multimodal_params = build_disagg_prefill_multimodal_inputs( + inputs, mm_handles) # ``build_disagg_prefill_multimodal_inputs`` already emits the in-vocab # ``image_token_id`` at mm positions (no legacy OOV remap), so @@ -1230,10 +1229,9 @@ def _attach_multimodal_embeddings_impl( mrope_config = self.get_mrope_config(mrope_input_ids, image_grid_thw, None, attention_mask, None) - multimodal_data = { - "multimodal_embedding": image_embeddings, - "mrope_config": mrope_config, - } + multimodal_data = multimodal_params.multimodal_data + multimodal_data["multimodal_embedding"] = image_embeddings + multimodal_data["mrope_config"] = mrope_config return prompt_token_ids, {"multimodal_data": multimodal_data} @nvtx_range("Qwen2VLInputProcessorBase forward()") @@ -2474,21 +2472,24 @@ def load_weights(self, weights, weight_mapper: BaseWeightMapper): class Qwen2_5VLInputProcessorBase(Qwen2VLInputProcessorBase): def build_disagg_prefill_multimodal_inputs( - self, inputs: TextPrompt, - mm_handles: List[Dict[str, Any]]) -> DisaggPrefillMultimodalInputs: + self, + inputs: TextPrompt, + mm_handles: List[Dict[str, Any]], + *, + mm_hashes: list[list[int]] | None = None, + ) -> tuple[list[int], MultimodalParams]: """ Build disaggregated prefill inputs from multimodal embedding handles. Args: inputs: Text prompt input container. Must contain a non-empty prompt string. mm_handles: List of multimodal embedding handles. + mm_hashes: Optional prompt-ordered hashes for KV-cache reuse. Returns: - DisaggPrefillMultimodalInputs containing expanded token IDs, - prompt-side MM positions/lengths, exact runs, and encoder-output - embedding lengths. + Expanded token IDs and common multimodal params containing + prompt layout, item metadata, and embedding handles. """ - # TODO: Move this function to the base input processor class when extending for more models text_prompt = inputs.get("prompt") if not text_prompt: raise ValueError("Text prompt is required but not provided") @@ -2546,18 +2547,23 @@ def build_disagg_prefill_multimodal_inputs( assert write_pos == final_length, f"Write position mismatch: {write_pos} != {final_length}" assert mm_token_length[-1] + mm_token_offsets[ -1] <= final_length, f"mm_token_length[-1] + mm_token_offsets[-1] ({mm_token_length[-1] + mm_token_offsets[-1]}) should be less than or equal to final_length ({final_length})" - return DisaggPrefillMultimodalInputs( - prompt_token_ids=expanded_ids.to(torch.int32).tolist(), - multimodal_lengths=mm_token_length, - multimodal_positions=mm_token_offsets, - multimodal_embedding_lengths=[ - mm_handle["tensor_size"][0] for mm_handle in mm_handles - ], - multimodal_item_run_cu_offsets=list(range(len(mm_token_length) + - 1)), - multimodal_run_positions=mm_token_offsets, - multimodal_run_lengths=mm_token_length, - ) + return expanded_ids.to( + torch.int32).tolist(), MultimodalParams.from_embeddings( + mm_handles, + MultimodalItemMetadata( + [("image", i) for i in range(len(mm_handles))], + output_embedding_lengths=[ + handle["tensor_size"][0] for handle in mm_handles + ], + ), + multimodal_hashes=mm_hashes, + multimodal_lengths=mm_token_length, + multimodal_positions=mm_token_offsets, + multimodal_item_run_cu_offsets=list( + range(len(mm_token_length) + 1)), + multimodal_run_positions=mm_token_offsets, + multimodal_run_lengths=mm_token_length, + ) @support_multimodal_disaggregated diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 593ae4279a42..95eb34454c31 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -38,11 +38,7 @@ register_input_processor, support_multimodal_disaggregated, ) -from ...inputs.multimodal import ( - DisaggPrefillMultimodalInputs, - MultimodalItemPartition, - MultimodalParams, -) +from ...inputs.multimodal import MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams from ...logger import logger from ..attention.backends import AttentionMetadata from ..attention.backends.interface import PositionalEmbeddingParams, RopeParams @@ -94,7 +90,8 @@ def _expand_prompt_token_ids_for_mm_handoff( image_token_id: int, video_token_id: int, vision_start_token_id: int, -) -> DisaggPrefillMultimodalInputs: + mm_hashes: list[list[int]] | None = None, +) -> tuple[list[int], MultimodalParams]: """Expand Qwen3-VL image/video placeholders and emit sparse MM layout. Qwen handoff has one coarse or token per item. @@ -125,7 +122,8 @@ def _expand_prompt_token_ids_for_mm_handoff( mm_token_lengths: List[int] = [] mm_token_offsets: List[int] = [] - item_types: List[int] = [] + item_refs: list[tuple[str, int]] = [] + local_counts = {"image": 0, "video": 0} item_run_cu_offsets: List[int] = [0] run_positions: List[int] = [] run_lengths: List[int] = [] @@ -152,7 +150,9 @@ def _expand_prompt_token_ids_for_mm_handoff( mm_token_offsets.append(run_start) mm_token_lengths.append(prompt_mm_length) multimodal_embedding_lengths.append(mm_token_num) - item_types.append(0 if token_id == image_token_id else 1) + modality = "image" if token_id == image_token_id else "video" + item_refs.append((modality, local_counts[modality])) + local_counts[modality] += 1 run_positions.append(run_start) run_lengths.append(prompt_mm_length) item_run_cu_offsets.append(len(run_positions)) @@ -167,17 +167,18 @@ def _expand_prompt_token_ids_for_mm_handoff( if write_pos != final_length: raise RuntimeError(f"Write position mismatch: {write_pos} != {final_length}") - return DisaggPrefillMultimodalInputs( - prompt_token_ids=expanded_ids.to(torch.int32).tolist(), + params = MultimodalParams.from_embeddings( + mm_handles, + MultimodalItemMetadata(item_refs, output_embedding_lengths=multimodal_embedding_lengths), + multimodal_hashes=mm_hashes, multimodal_lengths=mm_token_lengths, multimodal_positions=mm_token_offsets, - multimodal_embedding_lengths=multimodal_embedding_lengths, multimodal_item_run_cu_offsets=item_run_cu_offsets, multimodal_run_positions=run_positions, multimodal_run_lengths=run_lengths, special_token_offsets=special_token_offsets, - item_types=item_types, ) + return expanded_ids.to(torch.int32).tolist(), params def _decide_do_sample_frames( @@ -487,21 +488,24 @@ def get_preferred_media_io_kwargs(self) -> Dict[str, Dict[str, Any]]: return {"image": {"format": "np"}, "video": {"format": "np"}} def build_disagg_prefill_multimodal_inputs( - self, inputs: TextPrompt, mm_handles: List[Dict[str, Any]] - ) -> DisaggPrefillMultimodalInputs: + self, + inputs: TextPrompt, + mm_handles: List[Dict[str, Any]], + *, + mm_hashes: list[list[int]] | None = None, + ) -> tuple[list[int], MultimodalParams]: """ Build disaggregated prefill inputs from multimodal embedding handles. Args: inputs: Text prompt input container. Must contain a non-empty prompt string. mm_handles: List of multimodal embedding handles. + mm_hashes: Optional prompt-ordered hashes for KV-cache reuse. Returns: - DisaggPrefillMultimodalInputs containing expanded token IDs, - prompt-side MM positions/lengths, exact runs, and encoder-output - embedding lengths. + Expanded token IDs and common multimodal params containing + prompt layout, item metadata, and embedding handles. """ - # TODO: Move this function to the base input processor class when extending for more models text_prompt = inputs.get("prompt") if not text_prompt: raise ValueError("Text prompt is required but not provided") @@ -530,6 +534,7 @@ def build_disagg_prefill_multimodal_inputs( image_token_id=self.config.image_token_id, video_token_id=self.config.video_token_id, vision_start_token_id=self.config.vision_start_token_id, + mm_hashes=mm_hashes, ) diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 4137206c1275..5b9a283f64da 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -831,68 +831,6 @@ def to_binding(self, executor_module: Any) -> Any: return executor_module.MultimodalInput(**kwargs) -@dataclass -class DisaggPrefillMultimodalInputs: - """Typed multimodal metadata returned by E/P disagg prefill processors.""" - - prompt_token_ids: List[int] - multimodal_lengths: List[int] - multimodal_positions: List[int] - multimodal_embedding_lengths: List[int] - multimodal_item_run_cu_offsets: Optional[List[int]] = None - multimodal_run_positions: Optional[List[int]] = None - multimodal_run_lengths: Optional[List[int]] = None - special_token_offsets: Optional[List[int]] = None - item_types: Optional[List[int]] = None - - def __post_init__(self) -> None: - _validate_int_list(self.prompt_token_ids, "prompt_token_ids") - _validate_int_list(self.multimodal_embedding_lengths, - "multimodal_embedding_lengths") - _validate_multimodal_positions_and_lengths( - self.multimodal_positions, - self.multimodal_lengths, - len(self.multimodal_embedding_lengths), - "multimodal_embedding_lengths", - ) - - if any(length <= 0 for length in self.multimodal_embedding_lengths): - raise ValueError("multimodal_embedding_lengths must be positive") - - _validate_multimodal_runs( - len(self.multimodal_lengths), - self.multimodal_lengths, - self.multimodal_item_run_cu_offsets, - self.multimodal_run_positions, - self.multimodal_run_lengths, - "multimodal_lengths", - ) - self._validate_optional_metadata() - - def _validate_optional_metadata(self) -> None: - if self.special_token_offsets is not None: - _validate_int_list(self.special_token_offsets, - "special_token_offsets") - if any(offset < 0 for offset in self.special_token_offsets): - raise ValueError("special_token_offsets must be non-negative") - if self.item_types is not None: - _validate_int_list(self.item_types, "item_types") - if len(self.item_types) != len(self.multimodal_lengths): - raise ValueError("item_types length must match " - "multimodal_lengths") - - def to_multimodal_input(self, - mm_hashes: List[List[int]]) -> MultimodalInput: - return MultimodalInput.from_components( - mm_hashes, - self.multimodal_positions, - self.multimodal_lengths, - mm_item_run_cu_offsets=self.multimodal_item_run_cu_offsets, - mm_run_positions=self.multimodal_run_positions, - mm_run_lengths=self.multimodal_run_lengths, - ) - - @dataclass class MultimodalRuntimeData: """Runtime data for tracking multimodal embedding caching and reuse per request sequence. @@ -1039,6 +977,70 @@ def __post_init__(self): normalize_multimodal_item_metadata(self.multimodal_data, self.mm_item_order) + @classmethod + def from_embeddings( + cls, + embeddings: list[torch.Tensor] | list[dict[str, Any]], + item_metadata: MultimodalItemMetadata, + *, + multimodal_positions: list[int], + multimodal_lengths: list[int], + multimodal_hashes: list[list[int]] | None = None, + multimodal_item_run_cu_offsets: list[int] | None = None, + multimodal_run_positions: list[int] | None = None, + multimodal_run_lengths: list[int] | None = None, + special_token_offsets: list[int] | None = None, + ) -> "MultimodalParams": + """Attach prompt-ordered embeddings without copying their payloads. + + Each entry is a tensor of shape ``[output_rows, feature_size]`` or its + shared-tensor handle. Item output lengths describe those rows; prompt + lengths also count model-specific MM framing tokens. Hashes enable + the C++ KV-cache representation. Without hashes, only the processed + data is retained, as for external embeddings without KV-cache reuse. + """ + item_metadata.validate() + embedding_lengths = item_metadata.output_embedding_lengths + if embedding_lengths is None or len(embeddings) != len( + embedding_lengths): + raise ValueError("Embeddings and item output lengths must align") + if any(length <= 0 for length in embedding_lengths): + raise ValueError("Attached embedding lengths must be positive") + multimodal_input = None + if multimodal_hashes is not None: + if len(multimodal_hashes) != len(embedding_lengths): + raise ValueError( + "Multimodal hashes and item output lengths must align") + multimodal_input = MultimodalInput.from_components( + multimodal_hashes, + multimodal_positions, + multimodal_lengths, + mm_item_run_cu_offsets=multimodal_item_run_cu_offsets, + mm_run_positions=multimodal_run_positions, + mm_run_lengths=multimodal_run_lengths, + ) + else: + _validate_multimodal_positions_and_lengths( + multimodal_positions, multimodal_lengths, + len(embedding_lengths), "output_embedding_lengths") + _validate_multimodal_runs(len(embedding_lengths), + multimodal_lengths, + multimodal_item_run_cu_offsets, + multimodal_run_positions, + multimodal_run_lengths, + "output_embedding_lengths") + data = { + "multimodal_embedding": embeddings, + "multimodal_embedding_lengths": embedding_lengths, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: item_metadata, + } + if special_token_offsets is not None: + _validate_int_list(special_token_offsets, "special_token_offsets") + if any(offset < 0 for offset in special_token_offsets): + raise ValueError("special_token_offsets must be non-negative") + data["special_token_offsets"] = special_token_offsets + return cls(multimodal_input=multimodal_input, multimodal_data=data) + @property def item_metadata(self) -> MultimodalItemMetadata | None: """The single item record stored in this request's processed data.""" diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 55f2af9fc88d..cac942daefea 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -996,8 +996,13 @@ def support_multimodal_disaggregated(model_cls: Type[nn.Module]): Apply this to a model class AFTER its input processor has been registered via @register_input_processor. The decorator will locate the processor class, - validate requirements, and set `supports_multimodal_disagg = True` on both + validate requirements, and set `support_mm_disagg = True` on both the processor class and the model class. + + The processor implements ``build_disagg_prefill_multimodal_inputs(inputs, + mm_handles, *, mm_hashes=None)`` and returns ``(prompt_token_ids, + MultimodalParams)``. The params hold prompt layout, item metadata, and + embedding handles; the worker resolves the handles before model execution. """ processor_cls = INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type.get( model_cls) diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index d1bd36c2e982..76ff7a6743aa 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -31,8 +31,7 @@ from transformers import PreTrainedTokenizerBase from tensorrt_llm._utils import mpi_disabled -from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, - MultimodalParams) +from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.llmapi import tracing from tensorrt_llm.metrics.enums import MetricNames @@ -978,40 +977,23 @@ def _preprocess( # TODO(TRTLLM-12869): Pass encoder-side MM layout through # DisaggregatedParams so prefill does not rebuild prompt tokens, # positions, lengths, runs, special offsets, and cumsum here. - disagg_mm_inputs = ( + prompt_token_ids, multimodal_params = ( self.input_processor.build_disagg_prefill_multimodal_inputs( - inputs, mm_handles)) - if not isinstance(disagg_mm_inputs, DisaggPrefillMultimodalInputs): + inputs, + mm_handles, + mm_hashes=disaggregated_params.multimodal_hashes)) + if not isinstance(multimodal_params, MultimodalParams): raise TypeError( "build_disagg_prefill_multimodal_inputs must return " - "DisaggPrefillMultimodalInputs") - prompt_token_ids = disagg_mm_inputs.prompt_token_ids + "(prompt_token_ids, MultimodalParams)") prompt = inputs.get("prompt", None) if is_gen_only: raise ValueError( "Generation-only mode should not need multimodal parameters" ) else: - mm_hashes = disaggregated_params.multimodal_hashes - multimodal_input = disagg_mm_inputs.to_multimodal_input( - mm_hashes) - # E/P handoff carries SharedTensorContainer dicts. Park them under the - # embedding key so BaseWorker's recursive to_tensor("multimodal_data") - # restores local tensor views before PyTorch forward. Until then this - # key holds handles, not tensors. - multimodal_data = { - "multimodal_embedding": - mm_handles, - "multimodal_embedding_lengths": - (disagg_mm_inputs.multimodal_embedding_lengths), - } - if disagg_mm_inputs.special_token_offsets is not None: - multimodal_data["special_token_offsets"] = ( - disagg_mm_inputs.special_token_offsets) - if disagg_mm_inputs.item_types is not None: - multimodal_data["layout_metadata"] = { - "item_types": disagg_mm_inputs.item_types - } + # BaseWorker resolves embedding handles to local tensor views. + multimodal_data = multimodal_params.multimodal_data if disaggregated_params.mrope_position_ids_handle is not None: # NOTE: `PyTorchModelEngine` assumes both are present when using mrope. assert disaggregated_params.mrope_position_deltas_handle is not None @@ -1026,10 +1008,6 @@ def _preprocess( maybe_compute_mm_embed_cumsum( prompt_token_ids, {"multimodal_data": multimodal_data}, cast(BaseMultimodalInputProcessor, self.input_processor)) - multimodal_params = MultimodalParams( - multimodal_input=multimodal_input, - multimodal_data=multimodal_data, - ) # This condition is to ensure that this branch is not hit for models that expand # placeholder token IDs with MM data. elif ("prompt_token_ids" in inputs diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py index f059c5611f24..54a61e74208c 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py @@ -85,16 +85,17 @@ def test_nemotron_nano_registers_native_multimodal_epd_components(): assert NemotronHMultimodalModel.support_mm_disagg is True -def _assert_nano_video_handoff(handoff): +def _assert_nano_video_handoff(prompt_token_ids, params): """Shared assertions for the EPD video handoff: split runs stay grouped under one MM item.""" - assert handoff.prompt_token_ids == [101, 30, 20, 20, 31, 55, 30, 20, 20, 31, 102] - assert handoff.multimodal_lengths == [8] - assert handoff.multimodal_positions == [1] - assert handoff.multimodal_embedding_lengths == [4] - assert handoff.multimodal_item_run_cu_offsets == [0, 2] - assert handoff.multimodal_run_positions == [1, 6] - assert handoff.multimodal_run_lengths == [4, 4] - assert handoff.special_token_offsets == [0, 3, 4, 7] + assert prompt_token_ids == [101, 30, 20, 20, 31, 55, 30, 20, 20, 31, 102] + assert params.multimodal_input.multimodal_lengths == [8] + assert params.multimodal_input.multimodal_positions == [1] + assert params.item_metadata.item_refs == [("video", 0)] + assert params.item_metadata.output_embedding_lengths == [4] + assert params.multimodal_input.multimodal_item_run_cu_offsets == [0, 2] + assert params.multimodal_input.multimodal_run_positions == [1, 6] + assert params.multimodal_input.multimodal_run_lengths == [4, 4] + assert params.multimodal_data["special_token_offsets"] == [0, 3, 4, 7] @pytest.mark.parametrize( @@ -137,19 +138,20 @@ def test_nemotron_nano_epd_handoff_preserves_non_contiguous_video_runs( ) video = VideoData(frames=[object()], metadata={}, audio=None) - handoff = processor.build_disagg_prefill_multimodal_inputs( + prompt_token_ids, params = processor.build_disagg_prefill_multimodal_inputs( { input_field: input_value, "multi_modal_data": {"video": [video]}, }, [{"tensor_size": (4, 16)}], + mm_hashes=[[1] * 8], ) if asserts_encode_not_called: processor._tokenizer.encode.assert_not_called() processor.expand_prompt_token_ids_for_mm.assert_called_once() assert processor.expand_prompt_token_ids_for_mm.call_args.args[0] == [101, 98, 102] - _assert_nano_video_handoff(handoff) + _assert_nano_video_handoff(prompt_token_ids, params) @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py b/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py index 9e03b5141dd0..24825bdf2583 100644 --- a/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py +++ b/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py @@ -2523,19 +2523,25 @@ def test_handoff_rebuilds_pruned_video_layout( {"tensor_size": [sum(item), 2], "metadata": {"retained_token_counts": item}} for item in counts ] - handoff = processor.build_disagg_prefill_multimodal_inputs(inputs, handles) + prompt_token_ids, params = processor.build_disagg_prefill_multimodal_inputs( + inputs, handles, mm_hashes=[[i] * 8 for i in range(len(handles))] + ) expected = [101] for item in counts: for count in item: expected.extend([100, 500, *([20] * count), 501]) expected.append(102) - assert handoff.prompt_token_ids == expected - assert handoff.multimodal_embedding_lengths == list(map(sum, counts)) - assert handoff.multimodal_lengths == [sum(item) + 2 * len(item) for item in counts] - assert handoff.multimodal_item_run_cu_offsets == [ + assert prompt_token_ids == expected + assert params.item_metadata.output_embedding_lengths == list(map(sum, counts)) + assert params.multimodal_input.multimodal_lengths == [ + sum(item) + 2 * len(item) for item in counts + ] + assert params.multimodal_input.multimodal_item_run_cu_offsets == [ sum(len(item) for item in counts[:i]) for i in range(len(counts) + 1) ] - assert handoff.multimodal_run_lengths == [count + 2 for item in counts for count in item] + assert params.multimodal_input.multimodal_run_lengths == [ + count + 2 for item in counts for count in item + ] @pytest.mark.parametrize("item_counts", [None, [3], [-1, 4], [2, 2], [True, 2], [1.0, 2]]) def test_handoff_rejects_invalid_retained_counts( diff --git a/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py b/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py index baef8945be59..41fc606a3c93 100644 --- a/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py +++ b/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py @@ -10,6 +10,10 @@ Qwen3VLInputProcessorBase, _expand_prompt_token_ids_for_mm_handoff, ) +from tensorrt_llm.disaggregated_params import DisaggregatedParams +from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.llmapi.llm import BaseLLM +from tensorrt_llm.sampling_params import SamplingParams pytestmark = pytest.mark.cpu_only @@ -44,14 +48,15 @@ def test_qwen3vl_disagg_video_prompt_expands_video_placeholder(): [11, _VISION_START_TOKEN_ID, _VIDEO_TOKEN_ID, _VISION_END_TOKEN_ID, 12] ) - handoff = processor.build_disagg_prefill_multimodal_inputs( + prompt_token_ids, params = processor.build_disagg_prefill_multimodal_inputs( {"prompt": "video prompt"}, [{"tensor_size": (3, 16)}], + mm_hashes=[[1] * 8], ) # In-vocab contract: the video placeholder expands to the real # _VIDEO_TOKEN_ID (no legacy OOV remap). - assert handoff.prompt_token_ids == [ + assert prompt_token_ids == [ 11, _VISION_START_TOKEN_ID, _VIDEO_TOKEN_ID, @@ -60,14 +65,14 @@ def test_qwen3vl_disagg_video_prompt_expands_video_placeholder(): _VISION_END_TOKEN_ID, 12, ] - assert handoff.multimodal_lengths == [4] - assert handoff.multimodal_positions == [1] - assert handoff.multimodal_embedding_lengths == [3] - assert handoff.multimodal_item_run_cu_offsets == [0, 1] - assert handoff.multimodal_run_positions == [1] - assert handoff.multimodal_run_lengths == [4] - assert handoff.special_token_offsets == [0] - assert handoff.item_types == [1] + assert params.multimodal_input.multimodal_lengths == [4] + assert params.multimodal_input.multimodal_positions == [1] + assert params.item_metadata.item_refs == [("video", 0)] + assert params.item_metadata.output_embedding_lengths == [3] + assert params.multimodal_input.multimodal_item_run_cu_offsets == [0, 1] + assert params.multimodal_input.multimodal_run_positions == [1] + assert params.multimodal_input.multimodal_run_lengths == [4] + assert params.multimodal_data["special_token_offsets"] == [0] def test_qwen3vl_disagg_mixed_prompt_layout_preserves_item_order(): @@ -86,16 +91,17 @@ def test_qwen3vl_disagg_mixed_prompt_layout_preserves_item_order(): ] ) - handoff = _expand_prompt_token_ids_for_mm_handoff( + prompt_token_ids, params = _expand_prompt_token_ids_for_mm_handoff( input_ids, [{"tensor_size": (2, 16)}, {"tensor_size": (3, 16)}], image_token_id=_IMAGE_TOKEN_ID, video_token_id=_VIDEO_TOKEN_ID, vision_start_token_id=_VISION_START_TOKEN_ID, + mm_hashes=[[1] * 8, [2] * 8], ) # In-vocab contract: placeholders expand to the real image / video token ids. - assert handoff.prompt_token_ids == [ + assert prompt_token_ids == [ _VISION_START_TOKEN_ID, _IMAGE_TOKEN_ID, _IMAGE_TOKEN_ID, @@ -107,14 +113,14 @@ def test_qwen3vl_disagg_mixed_prompt_layout_preserves_item_order(): _VIDEO_TOKEN_ID, _VISION_END_TOKEN_ID, ] - assert handoff.multimodal_lengths == [3, 4] - assert handoff.multimodal_positions == [0, 5] - assert handoff.multimodal_item_run_cu_offsets == [0, 1, 2] - assert handoff.multimodal_run_positions == [0, 5] - assert handoff.multimodal_run_lengths == [3, 4] - assert handoff.multimodal_embedding_lengths == [2, 3] - assert handoff.special_token_offsets == [0, 3] - assert handoff.item_types == [0, 1] + assert params.multimodal_input.multimodal_lengths == [3, 4] + assert params.multimodal_input.multimodal_positions == [0, 5] + assert params.multimodal_input.multimodal_item_run_cu_offsets == [0, 1, 2] + assert params.multimodal_input.multimodal_run_positions == [0, 5] + assert params.multimodal_input.multimodal_run_lengths == [3, 4] + assert params.item_metadata.item_refs == [("image", 0), ("video", 0)] + assert params.item_metadata.output_embedding_lengths == [2, 3] + assert params.multimodal_data["special_token_offsets"] == [0, 3] def test_qwen3vl_disagg_prompt_layout_requires_handle_per_placeholder(): @@ -129,6 +135,145 @@ def test_qwen3vl_disagg_prompt_layout_requires_handle_per_placeholder(): ) +@pytest.mark.parametrize("model", ["qwen2.5", "llava", "kimi"]) +def test_disagg_image_processors_return_common_params(model): + from tensorrt_llm._torch.models.modeling_kimi_k25 import KimiK25InputProcessor + from tensorrt_llm._torch.models.modeling_llava_next import LlavaNextInputProcessor + from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2_5VLInputProcessorBase + + processor_cls = { + "qwen2.5": Qwen2_5VLInputProcessorBase, + "llava": LlavaNextInputProcessor, + "kimi": KimiK25InputProcessor, + }[model] + processor = object.__new__(processor_cls) + processor._config = SimpleNamespace( + image_token_id=_IMAGE_TOKEN_ID, + image_token_index=_IMAGE_TOKEN_ID, + text_config=SimpleNamespace(hidden_size=16), + ) + processor._media_placeholder_token_id = _IMAGE_TOKEN_ID + processor._tokenizer = _FakeTokenizer([11, _IMAGE_TOKEN_ID, 12, _IMAGE_TOKEN_ID, 13]) + handles = [{"tensor_size": (2, 16)}, {"tensor_size": (3, 16)}] + hashes = [[1] * 8, [2] * 8] + + tokens, params = processor.build_disagg_prefill_multimodal_inputs( + {"prompt": "video prompt"}, handles, mm_hashes=hashes + ) + + assert tokens == [11, _IMAGE_TOKEN_ID, _IMAGE_TOKEN_ID, 12, *([_IMAGE_TOKEN_ID] * 3), 13] + assert isinstance(params, MultimodalParams) + assert params.multimodal_data["multimodal_embedding"] is handles + assert params.get_item_layout() == ([("image", 0), ("image", 1)], [2, 3]) + assert params.item_metadata.encoder_token_lengths is None + assert params.multimodal_input.multimodal_hashes is hashes + assert params.multimodal_input.multimodal_positions == [1, 4] + assert params.multimodal_input.multimodal_lengths == [2, 3] + + +@pytest.mark.parametrize("deepstack_levels", [0, 3]) +def test_disagg_preprocess_keeps_processor_params_and_mrope(monkeypatch, deepstack_levels): + processor = object.__new__(Qwen3VLInputProcessorBase) + processor._processor = SimpleNamespace() + processor._config = SimpleNamespace( + image_token_id=_IMAGE_TOKEN_ID, + video_token_id=_VIDEO_TOKEN_ID, + vision_start_token_id=_VISION_START_TOKEN_ID, + text_config=SimpleNamespace(hidden_size=16, vocab_size=151936), + vision_config=SimpleNamespace(deepstack_visual_indexes=list(range(deepstack_levels))), + ) + processor._tokenizer = _FakeTokenizer( + [11, _VISION_START_TOKEN_ID, _VIDEO_TOKEN_ID, _VISION_END_TOKEN_ID, 12] + ) + llm = object.__new__(BaseLLM) + llm.input_processor = processor + llm._hf_model_config = SimpleNamespace(is_encoder_decoder=False) + handles = [{"tensor_size": (3, 16 * (1 + deepstack_levels))}] + mrope_ids, mrope_deltas = {"position_ids": "handle"}, {"position_deltas": "handle"} + disagg = DisaggregatedParams( + request_type="context_and_generation", + multimodal_embedding_handles=handles, + multimodal_hashes=[[1] * 8], + mrope_position_ids_handle=mrope_ids, + mrope_position_deltas_handle=mrope_deltas, + ) + build = processor.build_disagg_prefill_multimodal_inputs + built = [] + + def capture(*args, **kwargs): + result = build(*args, **kwargs) + built.append(result[1]) + return result + + monkeypatch.setattr(processor, "build_disagg_prefill_multimodal_inputs", capture) + tokens, prompt, params, encoder_tokens = llm._preprocess( + {"prompt": "video prompt"}, SamplingParams(), disagg + ) + + assert params is built[0] + assert prompt == "video prompt" + assert encoder_tokens is None + assert params.get_item_layout() == ([("video", 0)], [3]) + assert params.multimodal_data["multimodal_embedding"] is handles + assert params.multimodal_data["mrope_config"]["mrope_position_ids"] is mrope_ids + assert params.multimodal_data["mrope_config"]["mrope_position_deltas"] is mrope_deltas + cumsum = params.multimodal_data["multimodal_embed_mask_cumsum"] + assert cumsum.device.type == "cpu" + assert cumsum.numel() == len(tokens) + assert cumsum[-1].item() == 3 + assert params.encoder_state is None + assert params.encoder_event is None + + +def test_disagg_decode_rebuilds_only_mrope_params(): + llm = object.__new__(BaseLLM) + llm._hf_model_config = SimpleNamespace(is_encoder_decoder=False) + mrope_ids, mrope_deltas = {"position_ids": "handle"}, {"position_deltas": "handle"} + disagg = DisaggregatedParams( + request_type="generation_only", + mrope_position_ids_handle=mrope_ids, + mrope_position_deltas_handle=mrope_deltas, + ) + + _, _, params, _ = llm._preprocess({"prompt_token_ids": [10, 11]}, SamplingParams(), disagg) + + assert params.item_metadata is None + assert params.multimodal_input is None + assert params.multimodal_data == { + "mrope_config": {"mrope_position_ids": mrope_ids, "mrope_position_deltas": mrope_deltas} + } + + +@pytest.mark.parametrize("model", ["qwen2.5", "qwen3"]) +def test_external_embeddings_reuse_common_item_metadata(model): + from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2_5VLInputProcessorBase + + cls = Qwen3VLInputProcessorBase if model == "qwen3" else Qwen2_5VLInputProcessorBase + processor = object.__new__(cls) + processor._config = SimpleNamespace( + image_token_id=_IMAGE_TOKEN_ID, + video_token_id=_VIDEO_TOKEN_ID, + vision_start_token_id=_VISION_START_TOKEN_ID, + text_config=SimpleNamespace(hidden_size=16), + vision_config=SimpleNamespace(deepstack_visual_indexes=[], spatial_merge_size=2), + ) + processor._tokenizer = _FakeTokenizer([11, _IMAGE_TOKEN_ID, 12, _IMAGE_TOKEN_ID, 13]) + mrope_config = {"mrope_position_deltas": torch.zeros(1, dtype=torch.int32)} + processor.get_mrope_config = lambda *args: mrope_config + embeddings = [torch.ones(4, 16), torch.ones(9, 16)] + + tokens, extra = processor.attach_multimodal_embeddings( + {"prompt": "video prompt"}, {"image": embeddings}, SamplingParams() + ) + params = MultimodalParams(multimodal_data=extra["multimodal_data"]) + + assert len(tokens) == 16 + assert params.get_item_layout() == ([("image", 0), ("image", 1)], [4, 9]) + assert params.multimodal_data["multimodal_embedding"] is embeddings + assert params.multimodal_data["mrope_config"] is mrope_config + assert params.multimodal_input is None + + def test_qwen3vl_video_token_count_sums_single_video_grid_rows(): """One logical video may be represented by multiple processor grid rows.""" processor = object.__new__(Qwen3VLInputProcessorBase) diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 280d354b58b4..668031ec6588 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -13,7 +13,6 @@ from tensorrt_llm.inputs.media_io import MEDIA_IO_REGISTRY, MediaModality from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - DisaggPrefillMultimodalInputs, MultimodalEncoderItemMetadata, MultimodalInput, MultimodalItemMetadata, @@ -932,30 +931,81 @@ def get_num_tokens_per_video(self, *, video, **kwargs): ) -def test_disagg_prefill_multimodal_inputs_builds_typed_handoff(): - """Typed EPD handoff converts back to legacy MultimodalInput fields.""" - handoff = DisaggPrefillMultimodalInputs( - prompt_token_ids=[10, 1001, 1002, 2000, 1003, 20], +def test_attached_embeddings_share_item_metadata_and_prompt_layout(): + """Sparse E/P inputs use the same params as local encoder results.""" + handles = [{"tensor_size": (3, 16)}] + metadata = MultimodalItemMetadata([("video", 0)], output_embedding_lengths=[3]) + params = MultimodalParams.from_embeddings( + handles, + metadata, multimodal_lengths=[4], multimodal_positions=[1], - multimodal_embedding_lengths=[3], - multimodal_item_run_cu_offsets=[0, 1], - multimodal_run_positions=[1], - multimodal_run_lengths=[4], + multimodal_hashes=[[1, 2, 3, 4]], + multimodal_item_run_cu_offsets=[0, 2], + multimodal_run_positions=[1, 4], + multimodal_run_lengths=[2, 2], special_token_offsets=[2], - item_types=[0], ) - multimodal_input = handoff.to_multimodal_input([[1, 2, 3, 4]]) + assert params.item_metadata is metadata + assert params.get_item_layout() == ([("video", 0)], [3]) + assert params.multimodal_data["multimodal_embedding"] is handles + assert ( + params.multimodal_data["multimodal_embedding_lengths"] is metadata.output_embedding_lengths + ) + assert params.multimodal_data["special_token_offsets"] == [2] + assert params.multimodal_input.multimodal_positions == [1] + assert params.multimodal_input.multimodal_lengths == [4] + assert params.multimodal_input.multimodal_item_run_cu_offsets == [0, 2] + assert params.multimodal_input.multimodal_run_positions == [1, 4] + assert params.multimodal_input.multimodal_run_lengths == [2, 2] + assert params.encoder_state is None + assert params.encoder_event is None + + +@pytest.mark.parametrize("modality", ["image", "video", "audio"]) +def test_attached_embeddings_accept_tensors_without_cache_hashes(modality): + embedding = torch.ones(3, 16) + params = MultimodalParams.from_embeddings( + [embedding], + MultimodalItemMetadata([(modality, 0)], output_embedding_lengths=[3]), + multimodal_positions=[1], + multimodal_lengths=[3], + ) + + assert params.multimodal_input is None + assert params.item_metadata.item_refs == [(modality, 0)] + assert params.multimodal_data["multimodal_embedding"][0] is embedding + assert params.item_metadata.encoder_token_lengths is None + - assert multimodal_input.multimodal_positions == [1] - assert multimodal_input.multimodal_lengths == [4] - assert multimodal_input.multimodal_item_run_cu_offsets == [0, 1] - assert multimodal_input.multimodal_run_positions == [1] - assert multimodal_input.multimodal_run_lengths == [4] - assert handoff.multimodal_embedding_lengths == [3] - assert handoff.special_token_offsets == [2] - assert handoff.item_types == [0] +@pytest.mark.parametrize("with_hashes", [False, True]) +@pytest.mark.parametrize( + "updates,match", + [ + ({"multimodal_positions": [-1]}, "non-negative"), + ({"multimodal_positions": []}, "length|align"), + ({"multimodal_lengths": [0]}, "positive"), + ({"multimodal_run_lengths": [2]}, "sum"), + ({"special_token_offsets": [-1]}, "non-negative"), + ({"embeddings": []}, "align"), + ], +) +def test_attached_embeddings_validate_prompt_and_item_alignment(with_hashes, updates, match): + kwargs = dict( + embeddings=[{"tensor_size": (3, 16)}], + item_metadata=MultimodalItemMetadata([("audio", 0)], output_embedding_lengths=[3]), + multimodal_positions=[1], + multimodal_lengths=[3], + multimodal_item_run_cu_offsets=[0, 1], + multimodal_run_positions=[1], + multimodal_run_lengths=[3], + ) + if with_hashes: + kwargs["multimodal_hashes"] = [[1] * 8] + kwargs.update(updates) + with pytest.raises(ValueError, match=match): + MultimodalParams.from_embeddings(**kwargs) def test_multimodal_input_rejects_invalid_prompt_spans(): From 8e93b2dbe735e3206bb2700f129f1fd158763a03 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Sun, 18 Oct 2026 22:50:54 -0700 Subject: [PATCH 09/22] Document multimodal contracts and unify output-length readers Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- AGENTS.md | 19 +++++ docs/source/torch/arch_overview.md | 69 +++++++++++++++++++ tensorrt_llm/_torch/pyexecutor/llm_request.py | 25 +++++-- .../test_multimodal_embedding_lengths.py | 36 ++++++++++ 4 files changed, 143 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f4e20f4873ec..f74213cc2353 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -93,6 +93,25 @@ HuggingFace Model → LLM API → PyTorch Executor - **Disaggregated serving**: separates prefill (context) and decode (generation) across GPUs - KV cache exchange via NIXL (default), UCX, or MPI +### Multimodal request data + +Read the [multimodal data contracts](docs/source/torch/arch_overview.md#multimodal-request-data) +before changing multimodal preprocessing, item selection, or request ownership. +- Reuse `MultimodalParams` in `inputs/multimodal.py`; the executor's request aggregate is + `LlmRequest.py_multimodal_params`. Legacy request properties alias this aggregate. +- `item_metadata` reads one record in `multimodal_data`. Its `input_partitions` describe + how processed fields are selected by item; model hooks retain batching and padding policy. + Keep prompt token counts, encoder costs, and output embedding rows distinct. +- Only request Params own `encoder_state` cache references. Encoder selection and LLM chunk + Params are separate views; cleanup remains in the executor and the cache remains model-owned. +- E/P processor hooks return `(prompt_token_ids, MultimodalParams)`; use + `MultimodalParams.from_embeddings` for embedding handoffs. `DisaggregatedParams` remains + the transport envelope. Do not serialize live encoder state or CUDA events in a new wire format. +- When changing an E/P contract, check external callers such as Dynamo separately from + local processor hooks: wire compatibility and Python API compatibility are distinct. +- Validate producer/admission contracts at their boundaries. Preserve legacy adapters and + model capabilities without adding repeated full-metadata validation to selection or forward. + ## Key Files | File | Role | diff --git a/docs/source/torch/arch_overview.md b/docs/source/torch/arch_overview.md index 8bef0397312d..70173497219b 100644 --- a/docs/source/torch/arch_overview.md +++ b/docs/source/torch/arch_overview.md @@ -30,6 +30,75 @@ The single-step flow of PyExecutor involves: - Running the decoder using the model forward outputs for the scheduled requests. - Adding output tokens for each request and handling finished requests. +## Multimodal Request Data + +Multimodal preprocessing and execution use the existing `MultimodalParams` in +[inputs/multimodal.py](../../../tensorrt_llm/inputs/multimodal.py). +The executor keeps request-wide data in `LlmRequest.py_multimodal_params`, while +encoder selection and LLM forwards create separate Params views. + +| Type or member | Responsibility | +|---|---| +| `MultimodalParams.multimodal_data` | Processed model inputs, precomputed embeddings or shared-tensor handles, and model-specific auxiliary data. | +| `MultimodalParams.item_metadata` | A property reading the single `MultimodalItemMetadata` record stored in `multimodal_data`. | +| `MultimodalItemMetadata` | Prompt-ordered references to processed items, optional encoder costs and output row counts, and input partitions. | +| `MultimodalItemPartition` | A small value describing one input field's item axis and optional packed lengths. It lives inside `item_metadata.input_partitions`. | +| `MultimodalInput` | Processor-side hash, prompt-position, and exact-run data converted into the existing C++ request fields by the worker. | +| `MultimodalRuntimeData` | Derived information for the current LLM forward or prefill chunk. | +| `MultimodalEncoderRequestState` | Request-owned readiness and references to cached encoder outputs; the model still owns the `TensorLRUCache`. | + +The process boundary is: + +```text +Main / input processor + -> prompt token IDs + MultimodalParams (processed data and item metadata) + -> existing worker handoff: restore tensors; convert MultimodalInput to C++ fields +Executor / LlmRequest.py_multimodal_params + -> admission creates encoder_state where item scheduling is supported + -> selected encoder Params: selected data and metadata, no cache ownership + -> LLM forward Params: embeddings and current-chunk runtime, no cache ownership + -> executor cleanup releases request-owned cache references +``` + +After conversion, the worker clears the processor-side `multimodal_input`; +the executor uses the C++ request fields for that information. Text-only requests +do not need an aggregate. Execution views leave `encoder_state` unset, and the +request-wide Params do not carry the current chunk's runtime data. Existing +prefetch events retain their stream synchronization role. Live cache references +and CUDA events are not added to the cross-process payload. + +An item reference identifies a processed item within its modality. It does not +necessarily identify one original media object or one tensor: a video can expand +into multiple processed items, and a field can pack multiple items in one tensor. +Partitions describe selection: a shared field has no item axis, a batched field +has an item axis, and a packed field also supplies per-item lengths. Contiguous +tensor selections can share storage; reordered selections may copy. Treat shared +input storage as read-only. Model hooks and `EncoderGroup` still determine +compatible batches, padding, masks, encoder arguments, and output fusion. + +Prompt multimodal token counts, physical encoder token costs, and output embedding +row counts serve different purposes. Unknown costs or row counts remain `None`; +metadata alone does not enable item scheduling. Producers and admission boundaries +validate the required contracts rather than revalidating all items per forward. +Legacy `mm_item_order`, `multimodal_embedding_lengths`, and the three-field +`MultimodalEncoderItemMetadata` remain compatibility inputs. Migrated producers +publish the shared metadata record; the executor's output-length helper reads its +rows and rejects conflicting legacy lengths. + +For encoder/prefill disaggregation (E/P), model processor hooks return +`(prompt_token_ids, MultimodalParams)` and use `MultimodalParams.from_embeddings` +to attach embeddings or their handles. See the hook contract in +[inputs/registry.py](../../../tensorrt_llm/inputs/registry.py). +`DisaggregatedParams` remains the transport envelope; this local result does not +replace the wire protocol. The prefill worker still rebuilds model-specific prompt +layout. Shared data contracts do not change the supported model/modality matrix +or introduce a common batching policy for every model. + +External orchestrators can pass the existing `DisaggregatedParams` to the LLM API; +the prefill-side processor builds the local Params. Custom processors implementing +the E/P hook must follow its tuple return contract. Check those callers separately +from transport fields when changing this interface. + ## Model Engine The core component of `PyExecutor` is the `ModelEngine`, responsible for executing the model's forward pass efficiently on the GPU. diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index ed0572e3e21c..353bb65f2ff2 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1246,17 +1246,30 @@ def _validate_optional_int_list(values: Any, def get_multimodal_embedding_lengths( request: LlmRequest) -> Optional[List[int]]: - """Return explicit per-item encoder-output lengths for a multimodal request.""" + """Return canonical item output rows, falling back to legacy layout.""" py_multimodal_data = request.py_multimodal_data if py_multimodal_data is not None and not isinstance( py_multimodal_data, dict): raise TypeError("py_multimodal_data must be a dict") - # `multimodal_embedding_lengths` is Python-side layout metadata, not a - # nanobind request field, so validate the flat handoff contract here. - multimodal_embedding_lengths = _validate_optional_int_list( + multimodal_embedding_lengths = ( py_multimodal_data.get("multimodal_embedding_lengths") - if py_multimodal_data is not None else None, - "multimodal_embedding_lengths") + if py_multimodal_data is not None else None) + item_metadata = get_multimodal_item_metadata(py_multimodal_data) + if (item_metadata is not None + and item_metadata.output_embedding_lengths is not None): + item_lengths = item_metadata.output_embedding_lengths + if (multimodal_embedding_lengths is not None + and multimodal_embedding_lengths is not item_lengths): + _validate_optional_int_list(multimodal_embedding_lengths, + "multimodal_embedding_lengths") + if multimodal_embedding_lengths != item_lengths: + raise ValueError( + "Multimodal encoder item output lengths must match " + "multimodal_embedding_lengths") + multimodal_embedding_lengths = item_lengths + # Output rows are Python-side metadata, not a nanobind request field. + multimodal_embedding_lengths = _validate_optional_int_list( + multimodal_embedding_lengths, "multimodal_embedding_lengths") if multimodal_embedding_lengths is None: return None diff --git a/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py b/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py index 877caef4a960..5be1e50c3c9c 100644 --- a/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py +++ b/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py @@ -15,6 +15,11 @@ from tensorrt_llm._torch.pyexecutor.sampler import EarlyStopWithMMResult, MultimodalResult from tensorrt_llm._torch.shared_tensor import SharedTensorContainer from tensorrt_llm.bindings.executor import FinishReason +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, + MultimodalParams, +) from tensorrt_llm.mapping import Mapping @@ -51,6 +56,37 @@ def test_multimodal_embedding_lengths_returns_top_level_metadata(req, expected): assert get_multimodal_embedding_lengths(req) == expected +@pytest.mark.parametrize("legacy_lengths", [None, [2]]) +@pytest.mark.cpu_only +def test_multimodal_embedding_lengths_uses_item_metadata(legacy_lengths): + """Encoder output export uses the same rows as model input assembly.""" + metadata = MultimodalItemMetadata([("image", 0)], output_embedding_lengths=[2]) + data = {MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata} + if legacy_lengths is not None: + data["multimodal_embedding_lengths"] = legacy_lengths + params = MultimodalParams(multimodal_data=data) + req = SimpleNamespace(multimodal_lengths=[4], py_multimodal_data=params.multimodal_data) + + assert get_multimodal_embedding_lengths(req) is params.get_item_layout()[1] + + +@pytest.mark.cpu_only +def test_multimodal_embedding_lengths_rejects_conflicting_item_metadata(): + """A direct Params caller cannot export a different item layout.""" + params = MultimodalParams( + multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[2] + ), + "multimodal_embedding_lengths": [3], + } + ) + req = SimpleNamespace(multimodal_lengths=[4], py_multimodal_data=params.multimodal_data) + + with pytest.raises(ValueError, match="item output lengths must match"): + get_multimodal_embedding_lengths(req) + + @pytest.mark.parametrize( "req,exception,match", [ From 945273e250874eee22d882cc3519ef693ff56410 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Sun, 18 Oct 2026 23:38:53 -0700 Subject: [PATCH 10/22] Deduplicate multimodal request metadata and cache input views Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- AGENTS.md | 5 + docs/source/torch/arch_overview.md | 24 +++- .../_torch/models/modeling_mistral.py | 33 ++++- .../models/modeling_multimodal_mixin.py | 31 ++-- .../_torch/models/modeling_qwen2vl.py | 35 +++-- .../_torch/pyexecutor/engine/multimodal.py | 3 +- .../pyexecutor/engine/runners/mm_encoder.py | 2 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 30 ++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 1 + .../_torch/pyexecutor/scheduler/scheduler.py | 3 +- tensorrt_llm/inputs/multimodal.py | 135 ++++++++++++------ tensorrt_llm/inputs/registry.py | 8 -- .../_torch/executor/engine/test_multimodal.py | 4 +- .../test_multimodal_embedding_lengths.py | 25 ++-- .../executor/test_multimodal_scheduler.py | 31 ++-- .../_torch/executor/test_request_utils.py | 2 +- .../test_modeling_nemotron_h_multimodal.py | 4 +- .../test_multimodal_encoder_item_groups.py | 83 +++++++++-- ...est_nemotron_h_multimodal_preprocessing.py | 2 +- .../_torch/multimodal/test_encoder_group.py | 11 +- .../multimodal/test_multimodal_mixin.py | 2 +- .../multimodal/test_qwen3vl_disagg_prompt.py | 14 +- tests/unittest/inputs/test_multimodal.py | 56 ++++++-- 23 files changed, 370 insertions(+), 174 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f74213cc2353..e179ef8c98d1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -102,8 +102,13 @@ before changing multimodal preprocessing, item selection, or request ownership. - `item_metadata` reads one record in `multimodal_data`. Its `input_partitions` describe how processed fields are selected by item; model hooks retain batching and padding policy. Keep prompt token counts, encoder costs, and output embedding rows distinct. +- Item refs, costs, and output rows are immutable tuples. Encoder state references the + admitted record instead of copying its lengths. Normalized Params consume legacy order + and length fields; keep compatibility inputs only where normalization is unavailable. - Only request Params own `encoder_state` cache references. Encoder selection and LLM chunk Params are separate views; cleanup remains in the executor and the cache remains model-owned. +- Persistent encoder caching reuses one lazily reconstructed `MultimodalInput` per request. + Copy that wrapper before slicing its hashes; prefill/termination cleanup clears it. - E/P processor hooks return `(prompt_token_ids, MultimodalParams)`; use `MultimodalParams.from_embeddings` for embedding handoffs. `DisaggregatedParams` remains the transport envelope. Do not serialize live encoder state or CUDA events in a new wire format. diff --git a/docs/source/torch/arch_overview.md b/docs/source/torch/arch_overview.md index 70173497219b..0e9c41642319 100644 --- a/docs/source/torch/arch_overview.md +++ b/docs/source/torch/arch_overview.md @@ -45,7 +45,7 @@ encoder selection and LLM forwards create separate Params views. | `MultimodalItemPartition` | A small value describing one input field's item axis and optional packed lengths. It lives inside `item_metadata.input_partitions`. | | `MultimodalInput` | Processor-side hash, prompt-position, and exact-run data converted into the existing C++ request fields by the worker. | | `MultimodalRuntimeData` | Derived information for the current LLM forward or prefill chunk. | -| `MultimodalEncoderRequestState` | Request-owned readiness and references to cached encoder outputs; the model still owns the `TensorLRUCache`. | +| `MultimodalEncoderRequestState` | Request-owned readiness and cache references, plus a reference to the admitted item metadata; the model still owns the `TensorLRUCache`. | The process boundary is: @@ -61,7 +61,11 @@ Executor / LlmRequest.py_multimodal_params ``` After conversion, the worker clears the processor-side `multimodal_input`; -the executor uses the C++ request fields for that information. Text-only requests +the executor uses the C++ request fields for that information. When persistent +encoder caching needs a Python `MultimodalInput`, it reconstructs it once and +reuses it through the request Params. Selected inputs copy the wrapper before +changing their hashes. Prefill/termination cleanup clears the retained wrapper. +Text-only requests do not need an aggregate. Execution views leave `encoder_state` unset, and the request-wide Params do not carry the current chunk's runtime data. Existing prefetch events retain their stream synchronization role. Live cache references @@ -80,10 +84,22 @@ Prompt multimodal token counts, physical encoder token costs, and output embeddi row counts serve different purposes. Unknown costs or row counts remain `None`; metadata alone does not enable item scheduling. Producers and admission boundaries validate the required contracts rather than revalidating all items per forward. +The item refs, costs, and output rows are stored as immutable tuples; list inputs +are accepted by the metadata constructor. Encoder state references this same +record and reads its lengths without a second snapshot. Producers fill the +partition declarations before admission. Legacy `mm_item_order`, `multimodal_embedding_lengths`, and the three-field `MultimodalEncoderItemMetadata` remain compatibility inputs. Migrated producers -publish the shared metadata record; the executor's output-length helper reads its -rows and rejects conflicting legacy lengths. +publish the shared metadata record. Normalization checks and removes redundant +legacy lengths, and Params clear the order manifest once metadata is available. +Inputs without a known layout retain their compatibility path. The executor's +output-length helper reads canonical rows and rejects conflicting legacy lengths. +The encoder-result boundary still exports lengths as lists. + +Execution Params can receive model outputs. Inline encoding and prefetch retain +their existing output write-back to the request, while item-scheduled LLM chunks +keep their assembled embeddings local to the forward. CUDA event transfer remains +with the execution paths that synchronize the streams. For encoder/prefill disaggregation (E/P), model processor hooks return `(prompt_token_ids, MultimodalParams)` and use `MultimodalParams.from_embeddings` diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 6648dc896454..7f2d3a91a613 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -50,8 +50,10 @@ MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) -from tensorrt_llm.inputs.multimodal import (MultimodalEncoderItemMetadata, - MultimodalParams) +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, + MultimodalItemMetadata, MultimodalParams, + get_multimodal_encoder_item_metadata) from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY from tensorrt_llm.inputs.utils import encode_base64_image from tensorrt_llm.llmapi import SamplingParams @@ -411,11 +413,11 @@ def processor(self) -> AutoProcessor: def dtype(self) -> torch.dtype: return self._dtype - def get_mm_encoder_item_metadata( + def _build_mm_item_metadata( self, _prompt_token_ids: List[int], multimodal_data: Dict[str, Any], - ) -> Optional[MultimodalEncoderItemMetadata]: + ) -> Optional[MultimodalItemMetadata]: """Return Pixtral image items and physical ViT patch counts.""" image_data = multimodal_data.get("image") if not isinstance(image_data, dict): @@ -440,11 +442,32 @@ def get_mm_encoder_item_metadata( token_length // (merge * merge) for token_length in encoder_token_lengths ] - return MultimodalEncoderItemMetadata( + return MultimodalItemMetadata( item_refs=item_refs, encoder_token_lengths=encoder_token_lengths, output_embedding_lengths=output_embedding_lengths) + def get_mm_item_metadata( + self, + prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + mm_item_order: Optional[List[Dict[str, Any]]] = None, + ) -> Optional[MultimodalItemMetadata]: + metadata = ( + self._build_mm_item_metadata(prompt_token_ids, multimodal_data) + if multimodal_data.get("multimodal_embedding") is None else None) + return metadata if metadata is not None else super( + ).get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) + + def get_mm_encoder_item_metadata( + self, prompt_token_ids: List[int], + multimodal_data: Dict[str, + Any]) -> Optional[MultimodalEncoderItemMetadata]: + metadata = self._build_mm_item_metadata(prompt_token_ids, + multimodal_data) + return get_multimodal_encoder_item_metadata( + {MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata}) + @torch.inference_mode() def call_with_text_prompt( self, inputs: TextPrompt, sampling_params: SamplingParams diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index e20c6cf0d76b..d6faa3a46d5e 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -284,20 +284,25 @@ def _build_request_multimodal_input( request: "LlmRequest", cache_enabled: bool ) -> Optional[MultimodalInput]: """Build the encoder-cache key metadata carried by one request.""" - # Skip construction (and `from_components` validation) when no persistent cache consumes it. - if not cache_enabled or request.multimodal_hashes is None: + if not cache_enabled: return None - # `MultimodalModelMixin._encoder_cache_keys` uses UUID-aware multimodal hashes internally. - # Although the UUIDs are not exposed as an attribute, they remain in the backing C++ request - # for KV-cache block keys and cache events. - return MultimodalInput.from_components( - request.multimodal_hashes, + params = getattr(request, "py_multimodal_params", None) + if params is not None and params.multimodal_input is not None: + return params.multimodal_input + hashes = request.multimodal_hashes + if hashes is None: + return None + multimodal_input = MultimodalInput.from_components( + hashes, request.multimodal_positions, request.multimodal_lengths, mm_item_run_cu_offsets=request.multimodal_item_run_cu_offsets, mm_run_positions=request.multimodal_run_positions, mm_run_lengths=request.multimodal_run_lengths, ) + if params is not None: + params.multimodal_input = multimodal_input + return multimodal_input def _get_mm_aux_stream(max_prefetch_ahead: int = 0) -> Optional[torch.cuda.Stream]: @@ -472,7 +477,7 @@ def validate_multimodal_request_data(self, mm_data: dict) -> None: def prepare_multimodal_encoder_inputs( self, selected_items: Sequence[tuple[MultimodalParams, int]], - ) -> list[tuple[MultimodalParams, list[int], str]]: + ) -> list[tuple[MultimodalParams, Sequence[int], str]]: """Build selected item encoder inputs before the caller performs H2D. Adjacent items from the same request and modality are sliced in one @@ -488,7 +493,7 @@ def prepare_multimodal_encoder_inputs( Tuples of `(sliced encoder params, per-item embedding row counts, modality)`, in input order. Flattening the row-count lists recovers the per-item sequence. """ - encoder_inputs: list[tuple[MultimodalParams, list[int], str]] = [] + encoder_inputs: list[tuple[MultimodalParams, Sequence[int], str]] = [] for ( multimodal_param, run_indices, @@ -519,7 +524,7 @@ def prepare_multimodal_encoder_inputs( encoder_inputs.append( ( residual, - residual.multimodal_data["multimodal_embedding_lengths"], + residual.item_metadata.output_embedding_lengths, modality, ) ) @@ -570,7 +575,7 @@ def _runs_by_request_modality( def forward_multimodal_encoder_items( self, - encoder_inputs: Sequence[tuple[MultimodalParams, list[int], str]], + encoder_inputs: Sequence[tuple[MultimodalParams, Sequence[int], str]], ) -> list[torch.Tensor]: """Forward prepared MM encoder inputs in scheduler item order. @@ -1410,10 +1415,6 @@ def _apply_metadata_slice( if metadata is not None: selected = metadata.select(item_indices) residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = selected - if selected.output_embedding_lengths is not None: - residual.multimodal_data["multimodal_embedding_lengths"] = ( - selected.output_embedding_lengths - ) else: source_lengths = source.multimodal_data["multimodal_embedding_lengths"] residual.multimodal_data["multimodal_embedding_lengths"] = [ diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 6ae7b8694379..aa3cb0cebe11 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -30,10 +30,10 @@ from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.inputs.multimodal import (MultimodalEncoderItemMetadata, - MultimodalItemMetadata, - MultimodalItemPartition, - MultimodalParams) +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, + MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams, + get_multimodal_encoder_item_metadata) from ..._utils import async_tensor_h2d, nvtx_range from ...inputs import (BaseMultimodalDummyInputsBuilder, @@ -195,11 +195,11 @@ def processor(self) -> AutoProcessor: def dtype(self) -> torch.dtype: return self._dtype - def get_mm_encoder_item_metadata( + def _build_mm_item_metadata( self, prompt_token_ids: List[int], multimodal_data: Dict[str, Any], - ) -> Optional[MultimodalEncoderItemMetadata]: + ) -> Optional[MultimodalItemMetadata]: """Return prompt-ordered Qwen vision items and pre-merger costs.""" grids_by_modality: Dict[str, torch.Tensor] = {} for modality, grid_key in (("image", "image_grid_thw"), @@ -327,19 +327,36 @@ def get_mm_encoder_item_metadata( token_length // self.spatial_merge_unit for token_length in encoder_token_lengths ] - return MultimodalEncoderItemMetadata( + if any(length <= 0 for length in output_embedding_lengths): + raise ValueError( + "Processed Qwen items must produce positive embedding lengths") + return MultimodalItemMetadata( item_refs=item_refs, encoder_token_lengths=encoder_token_lengths, output_embedding_lengths=output_embedding_lengths) + def get_mm_encoder_item_metadata( + self, prompt_token_ids: List[int], + multimodal_data: Dict[str, + Any]) -> Optional[MultimodalEncoderItemMetadata]: + metadata = self._build_mm_item_metadata(prompt_token_ids, + multimodal_data) + return get_multimodal_encoder_item_metadata( + {MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata}) + def get_mm_item_metadata( self, prompt_token_ids: List[int], multimodal_data: Dict[str, Any], mm_item_order: Optional[List[Dict[str, Any]]] = None, ) -> Optional[MultimodalItemMetadata]: - metadata = super().get_mm_item_metadata(prompt_token_ids, - multimodal_data, mm_item_order) + metadata = ( + self._build_mm_item_metadata(prompt_token_ids, multimodal_data) + if multimodal_data.get("multimodal_embedding") is None else None) + if metadata is None: + metadata = super().get_mm_item_metadata(prompt_token_ids, + multimodal_data, + mm_item_order) if metadata is not None: for modality in ("image", "video"): if isinstance(multimodal_data.get(modality), dict): diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index 22a7d1ea31f5..67a9b2d017e6 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -552,6 +552,7 @@ def requests_using_cache_keys(cache_keys: set[Hashable]) -> set[int]: multimodal_param = MultimodalParams( multimodal_data=request.py_multimodal_params.multimodal_data ) + embedding_lengths = state.embedding_lengths for item_idx in item_indices: cache_key = state.item_cache_keys[item_idx] if cache_key is None: @@ -561,7 +562,7 @@ def requests_using_cache_keys(cache_keys: set[Hashable]) -> set[int]: ) scheduled_cache_keys.add(cache_key) encoder_items.append((multimodal_param, item_idx)) - output_targets.append((item_idx, cache_key, state.embedding_lengths[item_idx])) + output_targets.append((item_idx, cache_key, embedding_lengths[item_idx])) # A batch-level error cannot be traced to one output, so fail only the # scheduled producers. Requests that only share a scheduled entry keep diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/mm_encoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/mm_encoder.py index 297c74a27dcd..1aac11464cc1 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/mm_encoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/mm_encoder.py @@ -68,7 +68,7 @@ def _forward_step( continue mm_request_indices_with_payload.append(request_idx) mm_params_with_payload.append(multimodal_param) - mm_embedding_lengths.append(multimodal_embedding_lengths) + mm_embedding_lengths.append(list(multimodal_embedding_lengths)) if not mm_params_with_payload: return { "mm_embeddings": [], diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 353bb65f2ff2..d482b46998fc 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -4,7 +4,7 @@ from copy import copy, deepcopy from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Union, cast import torch @@ -1233,19 +1233,19 @@ def create_child_request(self, child_id): self.child_requests.append(py_request) -def _validate_optional_int_list(values: Any, - field_name: str) -> Optional[List[int]]: +def _validate_optional_int_sequence(values: Any, + field_name: str) -> Optional[Sequence[int]]: if values is None: return None - if not isinstance(values, list): - raise TypeError(f"{field_name} must be a list") + if not isinstance(values, (list, tuple)): + raise TypeError(f"{field_name} must be a list or tuple") if not all(isinstance(value, int) for value in values): raise TypeError(f"{field_name} must contain only integers") return values def get_multimodal_embedding_lengths( - request: LlmRequest) -> Optional[List[int]]: + request: LlmRequest) -> Optional[Sequence[int]]: """Return canonical item output rows, falling back to legacy layout.""" py_multimodal_data = request.py_multimodal_data if py_multimodal_data is not None and not isinstance( @@ -1260,15 +1260,15 @@ def get_multimodal_embedding_lengths( item_lengths = item_metadata.output_embedding_lengths if (multimodal_embedding_lengths is not None and multimodal_embedding_lengths is not item_lengths): - _validate_optional_int_list(multimodal_embedding_lengths, - "multimodal_embedding_lengths") - if multimodal_embedding_lengths != item_lengths: + _validate_optional_int_sequence(multimodal_embedding_lengths, + "multimodal_embedding_lengths") + if tuple(multimodal_embedding_lengths) != item_lengths: raise ValueError( "Multimodal encoder item output lengths must match " "multimodal_embedding_lengths") multimodal_embedding_lengths = item_lengths # Output rows are Python-side metadata, not a nanobind request field. - multimodal_embedding_lengths = _validate_optional_int_list( + multimodal_embedding_lengths = _validate_optional_int_sequence( multimodal_embedding_lengths, "multimodal_embedding_lengths") if multimodal_embedding_lengths is None: return None @@ -1291,7 +1291,7 @@ def get_multimodal_embedding_lengths( def get_multimodal_encoder_token_lengths( - request: LlmRequest) -> Optional[List[int]]: + request: LlmRequest) -> Optional[Sequence[int]]: """Return per-item physical encoder attention-token costs. Field-level validation happens inside @@ -1308,7 +1308,7 @@ def get_multimodal_encoder_token_lengths( if embedding_lengths is None: raise ValueError("Multimodal encoder item scheduling requires " "multimodal_embedding_lengths") - if item_metadata.output_embedding_lengths != embedding_lengths: + if item_metadata.output_embedding_lengths != tuple(embedding_lengths): raise ValueError("Multimodal encoder item output lengths must match " "multimodal_embedding_lengths") return item_metadata.encoder_token_lengths @@ -1450,9 +1450,9 @@ def initialize_multimodal_encoder_request( whole_request_output_resident=(context_chunk_unit_size is None), )) - request.py_mm_encoder_state = ( - MultimodalEncoderRequestState.from_embedding_lengths( - embedding_lengths, encoder_token_lengths=token_lengths)) + request.py_mm_encoder_state = MultimodalEncoderRequestState( + item_metadata=get_multimodal_item_metadata(mm_data), + item_ready=[False] * len(embedding_lengths)) def get_mm_context_chunk_start(request: LlmRequest) -> int: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 0f8723defc94..07b25a612140 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -8650,6 +8650,7 @@ def _release_multimodal_resources(self, request: LlmRequest) -> None: params = getattr(request, "py_multimodal_params", None) if params is not None: params.encoder_event = None + params.multimodal_input = None mm_data = getattr(request, "py_multimodal_data", None) if mm_data: strip_mm_data_for_generation(mm_data) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 9e870b6b3ad0..29fc9ec12af9 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -807,6 +807,7 @@ def _output_budget_blocks( ) -> bool: """Check whether unclaimed pending entries can fit beside protected cache entries.""" state = request.py_mm_encoder_state + embedding_lengths = state.embedding_lengths stable_keys = self.get_item_cache_keys(request) item_cache_keys = ( stable_keys @@ -820,7 +821,7 @@ def _output_budget_blocks( # entries without live references are evictable and do not claim space. additional_entries = { item_cache_keys[item_idx]: ( - state.embedding_lengths[item_idx] * self.bytes_per_encoder_embedding + embedding_lengths[item_idx] * self.bytes_per_encoder_embedding ) for item_idx in get_mm_pending_items(request) if item_cache_keys[item_idx] not in held_keys diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 5b9a283f64da..b8a9fee2d040 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -12,6 +12,7 @@ from dataclasses import dataclass, field from enum import Enum, auto from itertools import accumulate +from operator import itemgetter from typing import (Any, Container, Dict, Hashable, List, NamedTuple, Optional, Sequence, Tuple, Union, cast) @@ -191,7 +192,7 @@ def select(self, ) -@dataclass +@dataclass(frozen=True, slots=True, init=False) class MultimodalItemMetadata: """Prompt-ordered items and their partitions in processed model inputs. @@ -199,25 +200,43 @@ class MultimodalItemMetadata: item scheduling. Models keep their existing scheduling capabilities. """ - item_refs: list[tuple[str, int]] - encoder_token_lengths: list[int] | None = None - output_embedding_lengths: list[int] | None = None + item_refs: Sequence[tuple[str, int]] + encoder_token_lengths: Sequence[int] | None = None + output_embedding_lengths: Sequence[int] | None = None input_partitions: dict[str, dict[str, MultimodalItemPartition]] = field( default_factory=dict) + def __init__( + self, + item_refs: Sequence[tuple[str, int]], + encoder_token_lengths: Sequence[int] | None = None, + output_embedding_lengths: Sequence[int] | None = None, + input_partitions: dict[str, dict[str, MultimodalItemPartition]] + | None = None, + ) -> None: + object.__setattr__(self, "item_refs", tuple(item_refs)) + object.__setattr__( + self, "encoder_token_lengths", None + if encoder_token_lengths is None else tuple(encoder_token_lengths)) + object.__setattr__( + self, "output_embedding_lengths", None if output_embedding_lengths + is None else tuple(output_embedding_lengths)) + object.__setattr__(self, "input_partitions", + {} if input_partitions is None else input_partitions) + @classmethod def from_encoder_metadata( cls, metadata: MultimodalEncoderItemMetadata ) -> "MultimodalItemMetadata": - """Adapt a legacy producer without copying its CPU lists.""" + """Adapt the legacy producer's lists to the shared immutable record.""" metadata.validate() return cls(metadata.item_refs, metadata.encoder_token_lengths, metadata.output_embedding_lengths) def validate(self) -> None: """Validate producer data once before dispatch or executor admission.""" - if not isinstance(self.item_refs, list): - raise TypeError("Multimodal item references must be a list") + if not isinstance(self.item_refs, tuple): + raise TypeError("Multimodal item references must be a sequence") indices_by_modality: dict[str, list[int]] = {} for ref in self.item_refs: if (not isinstance(ref, tuple) or len(ref) != 2 @@ -235,7 +254,9 @@ def validate(self) -> None: lengths = getattr(self, name) if lengths is None: continue - _validate_int_list(lengths, name) + if not isinstance(lengths, tuple) or not all( + isinstance(length, int) for length in lengths): + raise TypeError(f"{name} must be a sequence of integers") if len(lengths) != len(self.item_refs): raise ValueError( "Multimodal item references and lengths must align") @@ -263,6 +284,30 @@ def validate(self) -> None: def select(self, item_indices: Sequence[int]) -> "MultimodalItemMetadata": """Select items and rebase their refs and packed lengths to the view.""" + if len(item_indices) == 0: + return MultimodalItemMetadata( + (), None if self.encoder_token_lengths is None else (), + None if self.output_embedding_lengths is None else ()) + if len(item_indices) == 1: + index = item_indices[0] + modality, local_index = self.item_refs[index] + partitions = self.input_partitions.get(modality) + return MultimodalItemMetadata( + ((modality, 0), ), + None if self.encoder_token_lengths is None else + (self.encoder_token_lengths[index], ), + None if self.output_embedding_lengths is None else + (self.output_embedding_lengths[index], ), + {} if partitions is None else { + modality: { + name: + partition if partition.item_lengths is None else + MultimodalItemPartition(partition.item_axis, ( + partition.item_lengths[local_index], )) + for name, partition in partitions.items() + } + }, + ) local_indices: dict[str, list[int]] = {} refs = [] for index in item_indices: @@ -282,11 +327,12 @@ def select(self, item_indices: Sequence[int]) -> "MultimodalItemMetadata": for name, partition in self.input_partitions[modality].items() } + select = itemgetter(*item_indices) return MultimodalItemMetadata( refs, - ([self.encoder_token_lengths[i] for i in item_indices] + (select(self.encoder_token_lengths) if self.encoder_token_lengths is not None else None), - ([self.output_embedding_lengths[i] for i in item_indices] + (select(self.output_embedding_lengths) if self.output_embedding_lengths is not None else None), partitions, ) @@ -334,11 +380,21 @@ def normalize_multimodal_item_metadata( metadata = MultimodalItemMetadata(refs, output_embedding_lengths=lengths) metadata.validate() - if mm_item_order and metadata.item_refs != [ - (entry["modality"], entry["index"]) for entry in mm_item_order - ]: + if mm_item_order and metadata.item_refs != tuple( + (entry["modality"], entry["index"]) for entry in mm_item_order): raise ValueError( "Multimodal item references do not match the prompt item order") + legacy_lengths = multimodal_data.get("multimodal_embedding_lengths") + if metadata.output_embedding_lengths is not None and legacy_lengths is not None: + if not isinstance(legacy_lengths, (list, tuple)) or not all( + isinstance(n, int) for n in legacy_lengths): + raise TypeError( + "multimodal_embedding_lengths must be a sequence of integers") + if tuple(legacy_lengths) != metadata.output_embedding_lengths: + raise ValueError( + "Multimodal item output lengths do not match multimodal_embedding_lengths" + ) + multimodal_data.pop("multimodal_embedding_lengths") multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata return metadata @@ -359,8 +415,8 @@ def get_multimodal_encoder_item_metadata( or metadata.output_embedding_lengths is None): return None metadata = MultimodalEncoderItemMetadata( - metadata.item_refs, metadata.encoder_token_lengths, - metadata.output_embedding_lengths) + list(metadata.item_refs), list(metadata.encoder_token_lengths), + list(metadata.output_embedding_lengths)) if not isinstance(metadata, MultimodalEncoderItemMetadata): raise TypeError(f"{MULTIMODAL_ENCODER_ITEM_METADATA_KEY} must be a " "MultimodalEncoderItemMetadata") @@ -412,15 +468,8 @@ class MultimodalEncoderRequestState: request cleanup is safe. """ - embedding_lengths: List[int] - """Expected encoder-output rows for each item, in prompt order.""" - - encoder_token_lengths: List[int] - """Validated encoder attention-token cost of each atomic item. - - Admission copies these values from the request metadata so scheduling - never has to re-validate user input inside the scheduler loop. - """ + item_metadata: MultimodalItemMetadata + """The admitted request record, shared with Params rather than copied.""" item_ready: List[bool] """Whether each item's bound cache entry is ready, in prompt order.""" @@ -439,27 +488,22 @@ class MultimodalEncoderRequestState: cannot be built, so this request uses request-local cache keys instead. """ - @classmethod - def from_embedding_lengths( - cls, - embedding_lengths: List[int], - *, - encoder_token_lengths: Optional[List[int]] = None - ) -> "MultimodalEncoderRequestState": - if encoder_token_lengths is None: - encoder_token_lengths = embedding_lengths - return cls(embedding_lengths=list(embedding_lengths), - encoder_token_lengths=list(encoder_token_lengths), - item_ready=[False] * len(embedding_lengths), - item_cache_keys=[None] * len(embedding_lengths)) + @property + def embedding_lengths(self) -> Sequence[int]: + return cast(Sequence[int], self.item_metadata.output_embedding_lengths) + + @property + def encoder_token_lengths(self) -> Sequence[int]: + return cast(Sequence[int], self.item_metadata.encoder_token_lengths) def __post_init__(self) -> None: if not self.item_cache_keys: self.item_cache_keys = [None] * len(self.item_ready) - if not (len(self.embedding_lengths) == len(self.encoder_token_lengths) - == len(self.item_ready) == len(self.item_cache_keys)): - raise ValueError("MM encoder token and embedding lengths must have " - "exactly one cache key per item slot") + if not (len(self.item_metadata.item_refs) == len(self.item_ready) == + len(self.item_cache_keys)): + raise ValueError( + "MM encoder state must have exactly one cache key per item slot" + ) @property def num_items(self) -> int: @@ -974,8 +1018,9 @@ class MultimodalParams: def __post_init__(self): if self.multimodal_data is None: self.multimodal_data = {} - normalize_multimodal_item_metadata(self.multimodal_data, - self.mm_item_order) + if normalize_multimodal_item_metadata(self.multimodal_data, + self.mm_item_order) is not None: + self.mm_item_order = None @classmethod def from_embeddings( @@ -1031,7 +1076,6 @@ def from_embeddings( "output_embedding_lengths") data = { "multimodal_embedding": embeddings, - "multimodal_embedding_lengths": embedding_lengths, MULTIMODAL_ENCODER_ITEM_METADATA_KEY: item_metadata, } if special_token_offsets is not None: @@ -1046,7 +1090,8 @@ def item_metadata(self) -> MultimodalItemMetadata | None: """The single item record stored in this request's processed data.""" return get_multimodal_item_metadata(self.multimodal_data) - def get_item_layout(self) -> tuple[list[tuple[str, int]], list[int]]: + def get_item_layout( + self) -> tuple[Sequence[tuple[str, int]], Sequence[int]]: """Resolve prompt-order item refs and output lengths for all modalities. Typed metadata takes precedence. Legacy processors supply diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index cac942daefea..72b29d179437 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1538,14 +1538,6 @@ def input_processor_wrapper( item_metadata.validate() multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = item_metadata normalize_multimodal_item_metadata(multimodal_data, mm_item_order) - embedding_lengths = item_metadata.output_embedding_lengths - if embedding_lengths is not None: - existing = multimodal_data.get("multimodal_embedding_lengths") - if existing is not None and existing != embedding_lengths: - raise ValueError( - "Computed multimodal encoder embedding lengths " - "do not match the existing prompt metadata") - multimodal_data["multimodal_embedding_lengths"] = embedding_lengths return prompt_token_ids, extra_processed_inputs return input_processor_wrapper diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index a30aa49dfb90..9360cca97ae0 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -373,7 +373,9 @@ def test_encoder_graph_modality_validation_allows_delegated_encoders() -> None: def _bind_items(mm_item_scheduler: MultimodalItemScheduler, request, *, row_bytes: int = 8) -> None: state = request.py_mm_encoder_state - for item_idx, rows in enumerate(state.embedding_lengths): + for item_idx, rows in enumerate( + request.py_multimodal_params.item_metadata.output_embedding_lengths + ): cache_key = make_mm_encoder_transient_cache_key(request.request_id, item_idx) assert mm_item_scheduler.encoder_cache.acquire( cache_key, rows * row_bytes, retain_after_release=False diff --git a/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py b/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py index 5be1e50c3c9c..968e80f011e7 100644 --- a/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py +++ b/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py @@ -73,18 +73,15 @@ def test_multimodal_embedding_lengths_uses_item_metadata(legacy_lengths): @pytest.mark.cpu_only def test_multimodal_embedding_lengths_rejects_conflicting_item_metadata(): """A direct Params caller cannot export a different item layout.""" - params = MultimodalParams( - multimodal_data={ - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( - [("image", 0)], output_embedding_lengths=[2] - ), - "multimodal_embedding_lengths": [3], - } - ) - req = SimpleNamespace(multimodal_lengths=[4], py_multimodal_data=params.multimodal_data) - - with pytest.raises(ValueError, match="item output lengths must match"): - get_multimodal_embedding_lengths(req) + with pytest.raises(ValueError, match="output lengths do not match"): + MultimodalParams( + multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[2] + ), + "multimodal_embedding_lengths": [3], + } + ) @pytest.mark.parametrize( @@ -127,10 +124,10 @@ def test_multimodal_embedding_lengths_rejects_conflicting_item_metadata(): ( SimpleNamespace( multimodal_lengths=[4], - py_multimodal_data={"multimodal_embedding_lengths": (4,)}, + py_multimodal_data={"multimodal_embedding_lengths": (4.0,)}, ), TypeError, - "must be a list", + "integers", ), ( SimpleNamespace( diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py index 0c7d6ac2b0f5..1e62af2d1195 100644 --- a/tests/unittest/_torch/executor/test_multimodal_scheduler.py +++ b/tests/unittest/_torch/executor/test_multimodal_scheduler.py @@ -135,6 +135,7 @@ def test_cleanup_releases_duplicate_references_and_drops_state_and_event(): state.set_item_cache_key(0, cache_key, ready=True) state.set_item_cache_key(1, cache_key, ready=True) request.py_mm_encoder_event = object() + request.py_multimodal_params.multimodal_input = object() cache = Mock() executor = SimpleNamespace(model_engine=SimpleNamespace(mm_encoder_cache=cache)) PyExecutor._release_multimodal_resources(executor, request) @@ -143,6 +144,7 @@ def test_cleanup_releases_duplicate_references_and_drops_state_and_event(): assert state.item_cache_keys == [None, None] assert request.py_multimodal_params.encoder_state is None assert request.py_multimodal_params.encoder_event is None + assert request.py_multimodal_params.multimodal_input is None class _RejectMultimodalCapacityScheduler: @@ -1071,7 +1073,7 @@ def test_prepare_multimodal_encoder_inputs_slices_before_device_transfer(): item_param, embedding_lengths, modality = encoder_inputs[0] assert modality == "image" - assert embedding_lengths == [3] + assert embedding_lengths == (3,) assert item_param.multimodal_data["image"]["pixel_values"].squeeze(1).tolist() == [2, 3, 4] assert multimodal_param.multimodal_data["image"]["pixel_values"].shape[0] == 5 @@ -1172,18 +1174,24 @@ def test_weight_invalidation_rejects_live_references(): def test_mm_encoder_state_enforces_lengths_slot_invariant(): with pytest.raises(ValueError, match="one cache key per item slot"): MultimodalEncoderRequestState( - embedding_lengths=[2], encoder_token_lengths=[4], item_ready=[False, False] + item_metadata=MultimodalItemMetadata([("image", 0), ("image", 1)], [4, 7], [2, 3]), + item_ready=[False, False], + item_cache_keys=[None], ) -def test_mm_encoder_state_copies_validated_scheduler_costs_at_admission(): - request = make_mm_request(1, [4, 7]) +def test_mm_encoder_state_uses_immutable_admitted_metadata(): + costs = [4, 7] + request = make_mm_request(1, costs) + metadata = request.py_multimodal_params.item_metadata + costs[0] = 100 - assert request.py_mm_encoder_state.encoder_token_lengths == [4, 7] - - metadata = request.py_multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] - metadata.encoder_token_lengths[0] = 100 - assert request.py_mm_encoder_state.encoder_token_lengths == [4, 7] + assert metadata.encoder_token_lengths == (4, 7) + assert request.py_mm_encoder_state.item_metadata is metadata + assert "encoder_token_lengths" not in vars(request.py_mm_encoder_state) + assert "embedding_lengths" not in vars(request.py_mm_encoder_state) + with pytest.raises(TypeError): + metadata.encoder_token_lengths[0] = 100 scheduler = _scheduler(max_batch_size=2, max_num_tokens=11) output = scheduler.schedule_request([request], set()) @@ -1191,7 +1199,10 @@ def test_mm_encoder_state_copies_validated_scheduler_costs_at_admission(): def test_mm_encoder_state_tracks_prompt_ordered_cache_key_readiness(): - state = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) + state = MultimodalEncoderRequestState( + item_metadata=MultimodalItemMetadata([("audio", 0), ("audio", 1)], [4, 7], [2, 3]), + item_ready=[False, False], + ) first_cache_key = ("cache", 0) second_cache_key = ("cache", 1) diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index 26f9528d3760..85937f91a71d 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/tests/unittest/_torch/executor/test_request_utils.py @@ -86,7 +86,7 @@ def test_multimodal_payload_broadcast_and_children_keep_the_same_contract(): assert child.item_metadata is not params.item_metadata assert params.encoder_state is child.encoder_state is None assert params.multimodal_runtime is child.multimodal_runtime is None - assert params.item_metadata.item_refs == [("audio", 0), ("image", 0)] + assert params.item_metadata.item_refs == (("audio", 0), ("image", 0)) torch.testing.assert_close( child.multimodal_data["audio"]["features"], payload["audio"]["features"] ) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py index 54a61e74208c..b58de5675f06 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py @@ -90,8 +90,8 @@ def _assert_nano_video_handoff(prompt_token_ids, params): assert prompt_token_ids == [101, 30, 20, 20, 31, 55, 30, 20, 20, 31, 102] assert params.multimodal_input.multimodal_lengths == [8] assert params.multimodal_input.multimodal_positions == [1] - assert params.item_metadata.item_refs == [("video", 0)] - assert params.item_metadata.output_embedding_lengths == [4] + assert params.item_metadata.item_refs == (("video", 0),) + assert params.item_metadata.output_embedding_lengths == (4,) assert params.multimodal_input.multimodal_item_run_cu_offsets == [0, 2] assert params.multimodal_input.multimodal_run_positions == [1, 6] assert params.multimodal_input.multimodal_run_lengths == [4, 4] diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index b20149e840ed..9b9ffdc38c9a 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -60,13 +60,29 @@ def test_qwen_processor_produces_prompt_order_and_packed_partitions(): } metadata = processor.get_mm_item_metadata([12, 11, 13, 12, 10, 13], data) metadata.validate() - assert metadata.item_refs == [("video", 0), ("image", 0)] - assert metadata.encoder_token_lengths == [16, 8] - assert metadata.output_embedding_lengths == [4, 2] + assert metadata.item_refs == (("video", 0), ("image", 0)) + assert metadata.encoder_token_lengths == (16, 8) + assert metadata.output_embedding_lengths == (4, 2) assert metadata.input_partitions["image"]["pixel_values"].item_lengths == (8,) assert metadata.input_partitions["video"]["pixel_values_videos"].item_lengths == (16,) +def test_qwen_common_metadata_preserves_positive_output_requirement(): + from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLInputProcessorBase + + processor = object.__new__(Qwen3VLInputProcessorBase) + processor._config = SimpleNamespace( + image_token_id=10, + video_token_id=11, + vision_start_token_id=12, + vision_end_token_id=13, + vision_config=SimpleNamespace(spatial_merge_size=2), + ) + data = {"image": {"image_grid_thw": torch.tensor([[1, 1, 1]])}} + with pytest.raises(ValueError, match="positive embedding lengths"): + processor.get_mm_item_metadata([12, 10, 13], data) + + def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor @@ -79,7 +95,7 @@ def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): order = [{"modality": "audio", "index": 0}, {"modality": "image", "index": 0}] metadata = processor.get_mm_item_metadata([], data, order) metadata.validate() - assert metadata.item_refs == [("audio", 0), ("image", 0)] + assert metadata.item_refs == (("audio", 0), ("image", 0)) assert metadata.encoder_token_lengths is None assert metadata.input_partitions["audio"]["audio_features_mask"].item_axis == 0 assert metadata.input_partitions["image"]["image_position_ids"].item_axis == 0 @@ -113,7 +129,7 @@ def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): video = {"audio": audio, "video_size": [[2, 1, 32, 32]]} video_data = {"video": video, "multimodal_embedding_lengths": [6]} video_metadata = processor.get_mm_item_metadata([], video_data) - assert video_metadata.item_refs == [("video", 0)] + assert video_metadata.item_refs == (("video", 0),) assert "audio" not in video_metadata.input_partitions assert video_data["video"]["audio"] is audio @@ -285,7 +301,7 @@ def get_multimodal_item_partitions(self, param, modality): assert residual.item_metadata == MultimodalItemMetadata( [("audio", 0), ("audio", 1)], [1, 2], [1, 1] ) - assert request.item_metadata.item_refs == [("audio", 0), ("audio", 1), ("audio", 2)] + assert request.item_metadata.item_refs == (("audio", 0), ("audio", 1), ("audio", 2)) def test_sliced_items_survive_the_grouped_encoder_path(): @@ -293,7 +309,9 @@ def test_sliced_items_survive_the_grouped_encoder_path(): grouped path splits to zero rows and returns nothing for every item.""" model = _GroupedEncoderModel() request = _make_request(["image", "image"]) - owner = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) + owner = MultimodalEncoderRequestState( + item_metadata=request.item_metadata, item_ready=[False, False] + ) owner.set_item_cache_key(0, ("owned", 0), ready=True) request.encoder_state = owner selected = [(request, 0), (request, 1)] @@ -319,7 +337,7 @@ def test_adjacent_same_modality_items_share_one_encoder_call(): outputs, encoder_inputs = _run_items(model, [(request, 0), (request, 1)]) assert len(encoder_inputs) == 1, "adjacent same-modality items must be sliced together" - assert encoder_inputs[0][1] == ITEM_ROWS["image"][:2] + assert encoder_inputs[0][1] == tuple(ITEM_ROWS["image"][:2]) assert model.encoder_calls == [sum(ITEM_ROWS["image"][:2])] assert len(outputs) == 2 @@ -354,7 +372,7 @@ def test_full_and_item_encoder_paths_share_item_metadata( ] # Full-request encoding reads both order and lengths from the same # metadata that drives the scheduler's item selection. - request.multimodal_data.pop("multimodal_embedding_lengths") + request.multimodal_data.pop("multimodal_embedding_lengths", None) full_output = model.encode_multimodal_inputs([request]) item_outputs, _ = _run_items(model, [(request, i) for i in range(len(modalities))]) @@ -375,7 +393,8 @@ def test_metadata_is_rebased_for_partial_and_reordered_items() -> None: residual = encoder_inputs[0][0] metadata = residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] assert metadata == MultimodalItemMetadata([("image", 0), ("image", 1)], [3, 2], [3, 2]) - assert residual.multimodal_data["multimodal_embedding_lengths"] == [3, 2] + assert residual.item_metadata.output_embedding_lengths == (3, 2) + assert "multimodal_embedding_lengths" not in residual.multimodal_data assert request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] is original torch.testing.assert_close(outputs[0], _expected_rows(request, 2), rtol=0, atol=0) torch.testing.assert_close(outputs[1], _expected_rows(request, 0), rtol=0, atol=0) @@ -426,9 +445,9 @@ def test_partial_cache_uses_prompt_refs_and_lengths(metadata_format, reversed_or metadata = MultimodalItemMetadata([("image", 1), ("image", 0)], [3, 2], [3, 2]) request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata lengths = metadata.output_embedding_lengths - request.multimodal_data["multimodal_embedding_lengths"] = lengths + request.multimodal_data["multimodal_embedding_lengths"] = list(lengths) if metadata_format == "typed_only": - request.multimodal_data.pop("multimodal_embedding_lengths") + request.multimodal_data.pop("multimodal_embedding_lengths", None) elif metadata_format == "legacy": request.mm_item_order = [ {"modality": modality, "index": index} for modality, index in metadata.item_refs @@ -455,7 +474,7 @@ def test_full_cache_write_uses_canonical_lengths_without_legacy_fields(): model._multimodal_encoder_cache = TensorLRUCache(4096) request = _make_request(["image", "image"]) lengths = request.item_metadata.output_embedding_lengths - request.multimodal_data.pop("multimodal_embedding_lengths") + request.multimodal_data.pop("multimodal_embedding_lengths", None) _attach_cache_metadata(request, lengths) expected = model.encode_multimodal_inputs([request]) calls_before = len(model.encoder_calls) @@ -493,7 +512,7 @@ def test_partial_cache_batches_mixed_vision_and_audio_misses(): [4, 1, 3, 2, 1], [4, 1, 3, 2, 1], ) - request.multimodal_data.pop("multimodal_embedding_lengths") + request.multimodal_data.pop("multimodal_embedding_lengths", None) lengths = request.item_metadata.output_embedding_lengths _attach_cache_metadata(request, lengths) expected = model.encode_multimodal_inputs([request]) @@ -563,7 +582,7 @@ def build_audio_input(params: list[MultimodalParams]) -> dict[str, torch.Tensor] request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalEncoderItemMetadata( [("video", 0), ("audio", 0), ("image", 0), ("audio", 1)], [4, 2, 2, 2], [4, 2, 2, 2] ) - request.multimodal_data.pop("multimodal_embedding_lengths") + request.multimodal_data.pop("multimodal_embedding_lengths", None) expected = torch.cat( [_expected_rows(request, 0), audio[0], _expected_rows(request, 2), audio[1]] ) @@ -613,3 +632,37 @@ def test_items_from_multiple_requests_stay_separated(): assert model.encoder_calls == [2 * ITEM_ROWS["image"][0]] torch.testing.assert_close(outputs[0], _expected_rows(first, 0)) torch.testing.assert_close(outputs[1], _expected_rows(second, 0)) + + +def test_request_cache_input_is_built_once_and_shared_with_forward_views(): + from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input + + class Request: + def __init__(self): + self.py_multimodal_params = _make_request(["image", "image"]) + self.hash_reads = 0 + self.multimodal_positions = [1, 5] + self.multimodal_lengths = [2, 3] + self.multimodal_item_run_cu_offsets = None + self.multimodal_run_positions = None + self.multimodal_run_lengths = None + + @property + def multimodal_hashes(self): + self.hash_reads += 1 + return [[1] * 8, [2] * 8] + + request = Request() + assert _build_request_multimodal_input(request, False) is None + assert request.hash_reads == 0 + + first = _build_request_multimodal_input(request, True) + second = _build_request_multimodal_input(request, True) + assert first is second is request.py_multimodal_params.multimodal_input + assert request.hash_reads == 1 + + model = _GroupedEncoderModel() + residual = model.build_multimodal_encoder_input(request.py_multimodal_params, [1], "image") + model._apply_metadata_slice(residual, request.py_multimodal_params, [1]) + assert residual.multimodal_input.multimodal_hashes == [[2] * 8] + assert first.multimodal_hashes == [[1] * 8, [2] * 8] diff --git a/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py b/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py index 24825bdf2583..c906f44ee8b9 100644 --- a/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py +++ b/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py @@ -2532,7 +2532,7 @@ def test_handoff_rebuilds_pruned_video_layout( expected.extend([100, 500, *([20] * count), 501]) expected.append(102) assert prompt_token_ids == expected - assert params.item_metadata.output_embedding_lengths == list(map(sum, counts)) + assert params.item_metadata.output_embedding_lengths == tuple(map(sum, counts)) assert params.multimodal_input.multimodal_lengths == [ sum(item) + 2 * len(item) for item in counts ] diff --git a/tests/unittest/_torch/multimodal/test_encoder_group.py b/tests/unittest/_torch/multimodal/test_encoder_group.py index 131559f681c5..79d5138ee14e 100644 --- a/tests/unittest/_torch/multimodal/test_encoder_group.py +++ b/tests/unittest/_torch/multimodal/test_encoder_group.py @@ -92,15 +92,12 @@ def test_single_modality_outside_group_is_skipped(self): class TestMultimodalItemLayout: def test_image_only(self): mp = _mp(embedding_lengths=[10, 20, 30], buckets={"image": {}}) - assert mp.get_item_layout() == ( - [("image", 0), ("image", 1), ("image", 2)], - [10, 20, 30], - ) + assert mp.get_item_layout() == ((("image", 0), ("image", 1), ("image", 2)), (10, 20, 30)) def test_layout_includes_all_request_modalities(self): # Layout belongs to the request, independent of the encoder group. mp = _mp(embedding_lengths=[10], buckets={"audio": {}}) - assert mp.get_item_layout() == ([("audio", 0)], [10]) + assert mp.get_item_layout() == ((("audio", 0),), (10,)) def test_item_metadata_owns_order_and_lengths(self) -> None: metadata = MultimodalEncoderItemMetadata( @@ -115,8 +112,8 @@ def test_item_metadata_owns_order_and_lengths(self) -> None: item_refs, lengths = mp.get_item_layout() - assert item_refs is metadata.item_refs - assert lengths is metadata.output_embedding_lengths + assert item_refs == tuple(metadata.item_refs) + assert lengths == tuple(metadata.output_embedding_lengths) assert _lengths_by_modality([mp.get_item_layout()], ("image", "video")) == { "image": [4], "video": [6], diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index b8de957810f2..7447cb974b13 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -117,7 +117,7 @@ def encode_multimodal_inputs(self, multimodal_params, **encoder_kwargs) -> torch if param.multimodal_runtime is not None: total_rows += param.multimodal_runtime.total_embeds_in_request else: - total_rows += sum(param.multimodal_data["multimodal_embedding_lengths"]) + total_rows += sum(param.get_item_layout()[1]) return torch.full( (total_rows, self.embedding.embedding_dim), float(self.encode_calls), diff --git a/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py b/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py index 41fc606a3c93..5544c7c997ab 100644 --- a/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py +++ b/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py @@ -67,8 +67,8 @@ def test_qwen3vl_disagg_video_prompt_expands_video_placeholder(): ] assert params.multimodal_input.multimodal_lengths == [4] assert params.multimodal_input.multimodal_positions == [1] - assert params.item_metadata.item_refs == [("video", 0)] - assert params.item_metadata.output_embedding_lengths == [3] + assert params.item_metadata.item_refs == (("video", 0),) + assert params.item_metadata.output_embedding_lengths == (3,) assert params.multimodal_input.multimodal_item_run_cu_offsets == [0, 1] assert params.multimodal_input.multimodal_run_positions == [1] assert params.multimodal_input.multimodal_run_lengths == [4] @@ -118,8 +118,8 @@ def test_qwen3vl_disagg_mixed_prompt_layout_preserves_item_order(): assert params.multimodal_input.multimodal_item_run_cu_offsets == [0, 1, 2] assert params.multimodal_input.multimodal_run_positions == [0, 5] assert params.multimodal_input.multimodal_run_lengths == [3, 4] - assert params.item_metadata.item_refs == [("image", 0), ("video", 0)] - assert params.item_metadata.output_embedding_lengths == [2, 3] + assert params.item_metadata.item_refs == (("image", 0), ("video", 0)) + assert params.item_metadata.output_embedding_lengths == (2, 3) assert params.multimodal_data["special_token_offsets"] == [0, 3] @@ -164,7 +164,7 @@ def test_disagg_image_processors_return_common_params(model): assert tokens == [11, _IMAGE_TOKEN_ID, _IMAGE_TOKEN_ID, 12, *([_IMAGE_TOKEN_ID] * 3), 13] assert isinstance(params, MultimodalParams) assert params.multimodal_data["multimodal_embedding"] is handles - assert params.get_item_layout() == ([("image", 0), ("image", 1)], [2, 3]) + assert params.get_item_layout() == ((("image", 0), ("image", 1)), (2, 3)) assert params.item_metadata.encoder_token_lengths is None assert params.multimodal_input.multimodal_hashes is hashes assert params.multimodal_input.multimodal_positions == [1, 4] @@ -213,7 +213,7 @@ def capture(*args, **kwargs): assert params is built[0] assert prompt == "video prompt" assert encoder_tokens is None - assert params.get_item_layout() == ([("video", 0)], [3]) + assert params.get_item_layout() == ((("video", 0),), (3,)) assert params.multimodal_data["multimodal_embedding"] is handles assert params.multimodal_data["mrope_config"]["mrope_position_ids"] is mrope_ids assert params.multimodal_data["mrope_config"]["mrope_position_deltas"] is mrope_deltas @@ -268,7 +268,7 @@ def test_external_embeddings_reuse_common_item_metadata(model): params = MultimodalParams(multimodal_data=extra["multimodal_data"]) assert len(tokens) == 16 - assert params.get_item_layout() == ([("image", 0), ("image", 1)], [4, 9]) + assert params.get_item_layout() == ((("image", 0), ("image", 1)), (4, 9)) assert params.multimodal_data["multimodal_embedding"] is embeddings assert params.multimodal_data["mrope_config"] is mrope_config assert params.multimodal_input is None diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 668031ec6588..6fc17746417b 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -185,7 +185,7 @@ def test_mm_item_metadata_is_materialized_when_embedding_lengths_match(): encoder_token_lengths=[4], output_embedding_lengths=[2], ) - assert multimodal_data["multimodal_embedding_lengths"] == [2] + assert "multimodal_embedding_lengths" not in multimodal_data assert "multimodal_item_refs" not in multimodal_data assert "multimodal_encoder_token_lengths" not in multimodal_data @@ -216,9 +216,9 @@ def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): _, extra = input_processor({"prompt": "unused"}, sampling_params=None) metadata = extra["multimodal_data"][MULTIMODAL_ENCODER_ITEM_METADATA_KEY] - assert metadata.item_refs == [("image", 0)] + assert metadata.item_refs == (("image", 0),) assert metadata.encoder_token_lengths is None - assert metadata.output_embedding_lengths == [2] + assert metadata.output_embedding_lengths == (2,) def test_mm_item_metadata_rejects_wrong_return_type(): @@ -377,9 +377,9 @@ def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): {"prompt": "unused", "mm_item_order": order}, None ) metadata = extra["multimodal_data"][MULTIMODAL_ENCODER_ITEM_METADATA_KEY] - assert metadata.item_refs == [("audio", 0), ("image", 0)] + assert metadata.item_refs == (("audio", 0), ("image", 0)) assert metadata.encoder_token_lengths is None - assert metadata.output_embedding_lengths == [1, 1] + assert metadata.output_embedding_lengths == (1, 1) def test_maybe_compute_mm_embed_cumsum_populates_py_multimodal_data(): @@ -736,7 +736,7 @@ def get_mm_encoder_item_metadata( assert cumsum.device.type == "cpu" assert cumsum.tolist() == expected_cumsum assert data["mm_bidirectional_blocks"] is True - assert data["multimodal_embedding_lengths"] == [2] * num_items + assert "multimodal_embedding_lengths" not in data assert data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] == MultimodalItemMetadata( item_refs, [4] * num_items, [2] * num_items ) @@ -948,11 +948,9 @@ def test_attached_embeddings_share_item_metadata_and_prompt_layout(): ) assert params.item_metadata is metadata - assert params.get_item_layout() == ([("video", 0)], [3]) + assert params.get_item_layout() == ((("video", 0),), (3,)) assert params.multimodal_data["multimodal_embedding"] is handles - assert ( - params.multimodal_data["multimodal_embedding_lengths"] is metadata.output_embedding_lengths - ) + assert "multimodal_embedding_lengths" not in params.multimodal_data assert params.multimodal_data["special_token_offsets"] == [2] assert params.multimodal_input.multimodal_positions == [1] assert params.multimodal_input.multimodal_lengths == [4] @@ -974,7 +972,7 @@ def test_attached_embeddings_accept_tensors_without_cache_hashes(modality): ) assert params.multimodal_input is None - assert params.item_metadata.item_refs == [(modality, 0)] + assert params.item_metadata.item_refs == ((modality, 0),) assert params.multimodal_data["multimodal_embedding"][0] is embedding assert params.item_metadata.encoder_token_lengths is None @@ -1077,3 +1075,39 @@ def test_runtime_data_requires_cumsum(): """embed_mask_cumsum is required.""" with pytest.raises(TypeError): MultimodalRuntimeData(past_seen_token_num=0, chunk_end_pos=5) + + +def test_normalized_params_keep_one_order_and_output_length_record(): + refs = [("audio", 0), ("image", 0)] + costs = [4, 8] + rows = [2, 3] + metadata = MultimodalItemMetadata(refs, costs, rows) + data = { + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, + "multimodal_embedding_lengths": list(rows), + } + params = MultimodalParams( + multimodal_data=data, + mm_item_order=[{"modality": m, "index": i} for m, i in refs], + ) + refs.reverse() + costs[0] = 100 + rows[0] = 100 + + assert params.item_metadata is metadata + assert params.mm_item_order is None + assert "multimodal_embedding_lengths" not in data + assert metadata.item_refs == (("audio", 0), ("image", 0)) + assert metadata.encoder_token_lengths == (4, 8) + assert metadata.output_embedding_lengths == (2, 3) + with pytest.raises(TypeError): + metadata.output_embedding_lengths[0] = 100 + with pytest.raises(AttributeError): + metadata.encoder_token_lengths = (100, 8) + + +def test_unknown_item_layout_preserves_legacy_order_at_input_boundary(): + order = [{"modality": "audio", "index": 0}] + params = MultimodalParams(multimodal_data={"audio": {}}, mm_item_order=order) + assert params.item_metadata is None + assert params.mm_item_order is order From 185644747eb69736591127a4ef300a5d0394a210 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:38:05 -0700 Subject: [PATCH 11/22] Use one multimodal item metadata contract Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- AGENTS.md | 10 +- docs/source/torch/arch_overview.md | 27 ++- .../_torch/models/modeling_mistral.py | 15 +- .../models/modeling_multimodal_mixin.py | 115 +++------- .../_torch/models/modeling_qwen2vl.py | 16 +- .../engine/runners/decoder/inputs.py | 1 - .../pyexecutor/engine/runners/no_kv_cache.py | 1 - tensorrt_llm/_torch/pyexecutor/llm_request.py | 68 ++---- tensorrt_llm/executor/base_worker.py | 2 - tensorrt_llm/inputs/__init__.py | 5 +- tensorrt_llm/inputs/multimodal.py | 207 +----------------- tensorrt_llm/inputs/registry.py | 177 +++++++-------- tensorrt_llm/llmapi/llm.py | 2 +- .../_torch/executor/engine/test_multimodal.py | 14 +- .../executor/engine/test_no_kv_cache.py | 8 +- .../_torch/executor/engine/test_runners.py | 30 ++- .../_torch/executor/multimodal_utils.py | 8 +- .../test_multimodal_embedding_lengths.py | 64 +++--- .../executor/test_multimodal_scheduler.py | 41 ++-- .../executor/test_pytorch_model_engine.py | 7 +- .../_torch/executor/test_request_utils.py | 1 - .../modeling/test_deepseek_v41_multimodal.py | 3 +- .../_torch/modeling/test_gemma4_multimodal.py | 23 +- .../_torch/modeling/test_modeling_mistral.py | 10 +- .../test_modeling_nemotron_h_multimodal.py | 11 +- .../test_multimodal_encoder_item_groups.py | 152 +++++++++---- ...est_nemotron_h_multimodal_preprocessing.py | 10 +- .../_torch/multimodal/test_encoder_group.py | 71 +++--- .../test_mm_encoder_cross_iter_prefetch.py | 41 ++-- .../multimodal/test_multimodal_mixin.py | 33 ++- ...st_nemotron_h_multimodal_encoder_groups.py | 2 +- .../test_qwen3vl_mixed_modality_encoder.py | 77 ++----- tests/unittest/inputs/test_multimodal.py | 152 +++++++------ 33 files changed, 586 insertions(+), 818 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e179ef8c98d1..41145c6ed6e8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -103,8 +103,8 @@ before changing multimodal preprocessing, item selection, or request ownership. how processed fields are selected by item; model hooks retain batching and padding policy. Keep prompt token counts, encoder costs, and output embedding rows distinct. - Item refs, costs, and output rows are immutable tuples. Encoder state references the - admitted record instead of copying its lengths. Normalized Params consume legacy order - and length fields; keep compatibility inputs only where normalization is unavailable. + admitted record instead of copying its lengths. Input preparation produces this record; + request and forward Params do not reconstruct it from separate order or length fields. - Only request Params own `encoder_state` cache references. Encoder selection and LLM chunk Params are separate views; cleanup remains in the executor and the cache remains model-owned. - Persistent encoder caching reuses one lazily reconstructed `MultimodalInput` per request. @@ -114,8 +114,10 @@ before changing multimodal preprocessing, item selection, or request ownership. the transport envelope. Do not serialize live encoder state or CUDA events in a new wire format. - When changing an E/P contract, check external callers such as Dynamo separately from local processor hooks: wire compatibility and Python API compatibility are distinct. -- Validate producer/admission contracts at their boundaries. Preserve legacy adapters and - model capabilities without adding repeated full-metadata validation to selection or forward. +- Validate producer/admission contracts at their boundaries, without repeated full-metadata + validation in selection or forward. Processors use `get_mm_item_metadata`; do not add a + second metadata format or old-hook dispatch. Unknown encoder costs remain `None`, and + common item metadata does not imply item scheduling support. ## Key Files diff --git a/docs/source/torch/arch_overview.md b/docs/source/torch/arch_overview.md index 0e9c41642319..c3c281e62d72 100644 --- a/docs/source/torch/arch_overview.md +++ b/docs/source/torch/arch_overview.md @@ -88,13 +88,26 @@ The item refs, costs, and output rows are stored as immutable tuples; list input are accepted by the metadata constructor. Encoder state references this same record and reads its lengths without a second snapshot. Producers fill the partition declarations before admission. -Legacy `mm_item_order`, `multimodal_embedding_lengths`, and the three-field -`MultimodalEncoderItemMetadata` remain compatibility inputs. Migrated producers -publish the shared metadata record. Normalization checks and removes redundant -legacy lengths, and Params clear the order manifest once metadata is available. -Inputs without a known layout retain their compatibility path. The executor's -output-length helper reads canonical rows and rejects conflicting legacy lengths. -The encoder-result boundary still exports lengths as lists. +Input preparation publishes the common item record directly. The frontend's +`mm_item_order` describes prompt order during preprocessing; executor Params +carry that order only in `item_metadata.item_refs`. Output row counts are stored +only in `item_metadata.output_embedding_lengths`. Request admission, selected-item +encoding, full-request encoding, cache keys and encoder output export all read +the same record. The encoder-result boundary still exports lengths as lists. + +`get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order=None)` is +the processor hook. Generic preparation supplies prompt order and known output +rows; model processors add encoder costs and partitions. Partial model metadata +retains known prompt rows, while conflicting order or rows fail at input preparation. +Gemma4 and Nemotron add partitions to this record. Qwen and Mistral also provide +physical encoder costs. A custom processor overrides this same hook and returns +`MultimodalItemMetadata`; the old three-field tuple, encoder-only metadata hook, +and Params order/length fallbacks have been removed. + +This data contract is independent of item scheduling. Processors can provide +common order and output rows with `encoder_token_lengths=None` and keep +full-request encoding. Model capability and configuration still decide whether +items are scheduled separately. Execution Params can receive model outputs. Inline encoding and prefetch retain their existing output write-back to the request, while item-scheduled LLM chunks diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 7f2d3a91a613..67acf7d93cd9 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -50,10 +50,8 @@ MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) -from tensorrt_llm.inputs.multimodal import ( - MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, - MultimodalItemMetadata, MultimodalParams, - get_multimodal_encoder_item_metadata) +from tensorrt_llm.inputs.multimodal import (MultimodalItemMetadata, + MultimodalParams) from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY from tensorrt_llm.inputs.utils import encode_base64_image from tensorrt_llm.llmapi import SamplingParams @@ -459,15 +457,6 @@ def get_mm_item_metadata( return metadata if metadata is not None else super( ).get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) - def get_mm_encoder_item_metadata( - self, prompt_token_ids: List[int], - multimodal_data: Dict[str, - Any]) -> Optional[MultimodalEncoderItemMetadata]: - metadata = self._build_mm_item_metadata(prompt_token_ids, - multimodal_data) - return get_multimodal_encoder_item_metadata( - {MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata}) - @torch.inference_mode() def call_with_text_prompt( self, inputs: TextPrompt, sampling_params: SamplingParams diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index d6faa3a46d5e..7af5b9d153aa 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -45,7 +45,6 @@ MultimodalItemPartition, MultimodalParams, MultimodalRuntimeData, - normalize_multimodal_item_metadata, select_multimodal_fields, ) from tensorrt_llm.logger import logger @@ -91,8 +90,7 @@ class EncoderGroup: Batching all items in a group into one encoder invocation amortizes fixed costs (kernel launches, dispatch) across items. The framework splits the output back per-modality and reorders it into prompt order - via each request's `MultimodalItemMetadata.item_refs` manifest, - falling back to `mm_item_order` for legacy processors. + via each request's `MultimodalItemMetadata.item_refs` manifest. Contract between `build_batched_input` and `encoder_fn`: @@ -129,7 +127,7 @@ class EncoderGroup: def _lengths_by_modality( - item_layouts: Sequence[tuple[list[tuple[str, int]], list[int]]], + item_layouts: Sequence[tuple[Sequence[tuple[str, int]], Sequence[int]]], modalities: Tuple[str, ...], ) -> Dict[str, List[int]]: """Split resolved output lengths into each group's payload-local order. @@ -163,7 +161,7 @@ def _lengths_by_modality( def _reorder_embeds_by_manifest( - item_layouts: Sequence[tuple[list[tuple[str, int]], list[int]]], + item_layouts: Sequence[tuple[Sequence[tuple[str, int]], Sequence[int]]], per_modality_embeds: Dict[str, torch.Tensor], per_modality_lengths: Dict[str, List[int]], ) -> torch.Tensor: @@ -240,8 +238,7 @@ def encode_multimodal_by_groups( items across all requests belonging to that group's modalities (arithmetic-intensity win). The output is split back per-modality using the prompt-ordered output lengths in `MultimodalItemMetadata`, - then reordered by the same metadata's `item_refs`. Legacy processors - fall back to `multimodal_embedding_lengths` and `mm_item_order`. + then reordered by the same metadata's `item_refs`. Shared entry point for both the aggregated (`MultimodalModelMixin`) and mm-encoder-only (`Qwen3VisionModelBase.forward`) paths so the ordering @@ -420,7 +417,7 @@ class MultimodalModelMixin: supports_mm_encoder_item_scheduling: ClassVar[bool] = False """Whether the model supports item-level MM encoder scheduling: it implements the item-encode forward and pairs with a processor that overrides - `BaseMultimodalInputProcessor.get_mm_encoder_item_metadata`. This is the executor-layer gate for + `BaseMultimodalInputProcessor.get_mm_item_metadata`. This is the executor-layer gate for item scheduling.""" model_config: ModelConfig @@ -550,10 +547,6 @@ def _runs_by_request_modality( if multimodal_param is run_param else multimodal_param.item_metadata ) - if metadata is None: - metadata = normalize_multimodal_item_metadata( - multimodal_param.multimodal_data, multimodal_param.mm_item_order - ) except (TypeError, ValueError) as error: raise MultimodalEncoderContractError(str(error)) from error if metadata is None or metadata.output_embedding_lengths is None: @@ -1133,8 +1126,7 @@ def _get_or_encode_multimodal_embeddings( def _initialize_multimodal_encoder_cache(self, max_bytes: int) -> Optional[TensorLRUCache]: """Initialize the model-owned multimodal encoder-output cache once. - The cache stores per-item embeddings for params that can be represented by one modality. - See `_encoder_cache_keys` for the mixed-modality skip path and its technical limitation. + The cache stores per-item embeddings keyed by the common item metadata. `ModelEngine` resolves `max_bytes` as the larger of the item-scheduling output budget and persistent-reuse capacity before runtime access. Zero leaves the cache disabled. @@ -1188,8 +1180,8 @@ def _encoder_cache_modality(param: MultimodalParams) -> Optional[str]: """Return the single modality represented by `param`, if cacheable. `None` means the params either do not identify a modality or contain - multiple modality inputs. The persistent encoder cache deliberately does - not cache mixed-modality params today. + multiple modality inputs. Per-item operations use the common item refs + when one request contains several modalities. """ mm_data = param.multimodal_data or {} modalities = [key for key in _MM_DATA_INPUT_MODALITY_KEYS if key in mm_data] @@ -1204,11 +1196,7 @@ def _encoder_cache_modality(param: MultimodalParams) -> Optional[str]: return modality if len(modalities) != 1: - # Mixed-modality params are skipped because the cache key metadata is request-item - # oriented: `multimodal_hashes` and `multimodal_embedding_lengths` are parallel per - # item, but there is no parallel per-item modality list. Without that, a cache key - # cannot unambiguously distinguish, for example, an image item from an audio item inside - # the same params object. + # A whole-param modality is ambiguous for mixed inputs. logger.debug( f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params with {len(modalities)} detected " "modalities." @@ -1221,16 +1209,7 @@ def _encoder_cache_keys( cls, param: MultimodalParams, ) -> Optional[list[Hashable]]: - """Build per-item encoder cache keys for `param`. - - The returned keys split one request's concatenated encoder output by - `multimodal_embedding_lengths`. - - When the request carries atomic-item metadata, each key takes its modality from - that item's `item_refs` entry, so mixed-modality requests are keyable per item. - Without that metadata the whole param must resolve to one modality, since there - is nothing else that says which item is which. - """ + """Build per-item keys from content hashes and the common item record.""" mm_input = param.multimodal_input mm_data = param.multimodal_data or {} if mm_input is None: @@ -1240,55 +1219,21 @@ def _encoder_cache_keys( return None item_metadata = param.item_metadata - if item_metadata is not None: - if item_metadata.output_embedding_lengths is None: - return None - item_modalities = {modality for modality, _ in item_metadata.item_refs} - if any( - modality not in item_modalities - for modality in _MM_DATA_INPUT_MODALITY_KEYS - if mm_data.get(modality) is not None - ): - # A legacy caller may add another raw modality after its - # original single-modality metadata was normalized. - return None - return cls.build_encoder_cache_item_keys( - mm_input.multimodal_hashes, - item_metadata.item_refs, - item_metadata.output_embedding_lengths, - mm_data.get("mm_processor_kwargs_hash"), - ) - - modality = cls._encoder_cache_modality(param) - embedding_lengths = mm_data.get("multimodal_embedding_lengths") - kwargs_hash = mm_data.get("mm_processor_kwargs_hash") - if modality is None or not isinstance(embedding_lengths, list) or kwargs_hash is None: - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping unkeyable params, " - f"has_modality={modality is not None}, " - f"has_embedding_lengths={isinstance(embedding_lengths, list)}, " - f"has_processor_kwargs_hash={kwargs_hash is not None}" - ) + if item_metadata is None or item_metadata.output_embedding_lengths is None: return None - if len(mm_input.multimodal_hashes) != len(embedding_lengths): - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params with mismatched " - "multimodal_hashes and multimodal_embedding_lengths counts" - ) + item_modalities = {modality for modality, _ in item_metadata.item_refs} + if any( + modality not in item_modalities + for modality in _MM_DATA_INPUT_MODALITY_KEYS + if mm_data.get(modality) is not None + ): return None - - # The item hash, embedding row count, processor kwargs, and modality fully describe a - # reusable item embedding. Request order is excluded so the same item can be reused from a - # different request layout; the current request order is restored when cached item tensors - # are concatenated below. - return [ - cls._encoder_cache_item_key(modality, item_hash, embedding_length, kwargs_hash) - for item_hash, embedding_length in zip( - mm_input.multimodal_hashes, - embedding_lengths, - strict=True, - ) - ] + return cls.build_encoder_cache_item_keys( + mm_input.multimodal_hashes, + item_metadata.item_refs, + item_metadata.output_embedding_lengths, + mm_data.get("mm_processor_kwargs_hash"), + ) @staticmethod def _encoder_cache_item_key( @@ -1412,14 +1357,11 @@ def _apply_metadata_slice( the parallel per-item metadata slice so every model gets it identically. """ metadata = item_metadata if item_metadata is not None else source.item_metadata - if metadata is not None: - selected = metadata.select(item_indices) - residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = selected - else: - source_lengths = source.multimodal_data["multimodal_embedding_lengths"] - residual.multimodal_data["multimodal_embedding_lengths"] = [ - source_lengths[i] for i in item_indices - ] + if metadata is None: + raise MultimodalEncoderContractError("Item selection requires multimodal item metadata") + residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata.select( + item_indices + ) if residual.multimodal_input is not None and source.multimodal_input is not None: source_hashes = source.multimodal_input.multimodal_hashes residual.multimodal_input.multimodal_hashes = [source_hashes[i] for i in item_indices] @@ -1695,7 +1637,6 @@ def _dispatch_cross_iter_prefetch( chunk_end_pos=cumsum.numel(), embed_mask_cumsum=cumsum, ), - mm_item_order=req.py_mm_item_order, ) for req, mm_data, cumsum in candidates ] diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index aa3cb0cebe11..9fec6eb3595a 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -30,10 +30,9 @@ from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.inputs.multimodal import ( - MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalEncoderItemMetadata, - MultimodalItemMetadata, MultimodalItemPartition, MultimodalParams, - get_multimodal_encoder_item_metadata) +from tensorrt_llm.inputs.multimodal import (MultimodalItemMetadata, + MultimodalItemPartition, + MultimodalParams) from ..._utils import async_tensor_h2d, nvtx_range from ...inputs import (BaseMultimodalDummyInputsBuilder, @@ -335,15 +334,6 @@ def _build_mm_item_metadata( encoder_token_lengths=encoder_token_lengths, output_embedding_lengths=output_embedding_lengths) - def get_mm_encoder_item_metadata( - self, prompt_token_ids: List[int], - multimodal_data: Dict[str, - Any]) -> Optional[MultimodalEncoderItemMetadata]: - metadata = self._build_mm_item_metadata(prompt_token_ids, - multimodal_data) - return get_multimodal_encoder_item_metadata( - {MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata}) - def get_mm_item_metadata( self, prompt_token_ids: List[int], diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index 578ab3e68667..6fc6b95af5dd 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -1088,7 +1088,6 @@ def _pack_context_multimodal( ), multimodal_data=self._build_multimodal_data_for_llm(request, py_multimodal_runtime), multimodal_runtime=py_multimodal_runtime, - mm_item_order=getattr(request, "py_mm_item_order", None), input_ids_start_offset=context_start_idx, ) # Transfer any cross-iter MM encoder prefetch event stamped on the request onto the diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py index bdc014a561ad..91f019205999 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py @@ -141,7 +141,6 @@ def _prepare_inputs( runner_config.mm_encoder_cache_enabled, ), multimodal_data=request.py_multimodal_params.multimodal_data, - mm_item_order=getattr(request, "py_mm_item_order", None), input_ids_start_offset=context_start_idx, ) multimodal_params.to_device( diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index d482b46998fc..106214f52428 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -18,8 +18,7 @@ from tensorrt_llm.inputs.multimodal import \ MultimodalEncoderRequestState as MultimodalEncoderRequestState from tensorrt_llm.inputs.multimodal import (MultimodalParams, - get_multimodal_item_metadata, - normalize_multimodal_item_metadata) + get_multimodal_item_metadata) from tensorrt_llm.sampling_params import LogprobMode SamplingConfig = tllm_executor.SamplingConfig @@ -737,11 +736,8 @@ def __init__( # The request owns one MM aggregate. Compatibility properties below # expose the existing field names without duplicate data or state. mm_data = kwargs.pop("py_multimodal_data", None) - mm_item_order = kwargs.pop("py_mm_item_order", None) - self.py_multimodal_params: MultimodalParams | None = ( - MultimodalParams(multimodal_data=mm_data, - mm_item_order=mm_item_order) - if mm_data is not None or mm_item_order is not None else None) + self.py_multimodal_params: MultimodalParams | None = (MultimodalParams( + multimodal_data=mm_data) if mm_data is not None else None) # Per-item MM encoder state, created at admission by # `initialize_multimodal_encoder_request` when raw MM payloads must # run through the item scheduler; `None` otherwise (also the @@ -1155,19 +1151,6 @@ def py_multimodal_data(self, data: dict[str, Any] | None) -> None: else: self.py_multimodal_params.multimodal_data = data - @property - def py_mm_item_order(self) -> list[dict[str, Any]] | None: - params = self.py_multimodal_params - return params.mm_item_order if params is not None else None - - @py_mm_item_order.setter - def py_mm_item_order(self, order: list[dict[str, Any]] | None) -> None: - if self.py_multimodal_params is None: - if order is None: - return - self.py_multimodal_params = MultimodalParams() - self.py_multimodal_params.mm_item_order = order - @property def py_mm_encoder_state(self) -> MultimodalEncoderRequestState | None: params = self.py_multimodal_params @@ -1246,45 +1229,28 @@ def _validate_optional_int_sequence(values: Any, def get_multimodal_embedding_lengths( request: LlmRequest) -> Optional[Sequence[int]]: - """Return canonical item output rows, falling back to legacy layout.""" - py_multimodal_data = request.py_multimodal_data - if py_multimodal_data is not None and not isinstance( - py_multimodal_data, dict): - raise TypeError("py_multimodal_data must be a dict") - multimodal_embedding_lengths = ( - py_multimodal_data.get("multimodal_embedding_lengths") - if py_multimodal_data is not None else None) - item_metadata = get_multimodal_item_metadata(py_multimodal_data) - if (item_metadata is not None - and item_metadata.output_embedding_lengths is not None): - item_lengths = item_metadata.output_embedding_lengths - if (multimodal_embedding_lengths is not None - and multimodal_embedding_lengths is not item_lengths): - _validate_optional_int_sequence(multimodal_embedding_lengths, - "multimodal_embedding_lengths") - if tuple(multimodal_embedding_lengths) != item_lengths: - raise ValueError( - "Multimodal encoder item output lengths must match " - "multimodal_embedding_lengths") - multimodal_embedding_lengths = item_lengths + """Read output rows from the request's common item metadata.""" + item_metadata = get_multimodal_item_metadata(request.py_multimodal_data) + multimodal_embedding_lengths = (item_metadata.output_embedding_lengths + if item_metadata is not None else None) # Output rows are Python-side metadata, not a nanobind request field. multimodal_embedding_lengths = _validate_optional_int_sequence( - multimodal_embedding_lengths, "multimodal_embedding_lengths") + multimodal_embedding_lengths, "output_embedding_lengths") if multimodal_embedding_lengths is None: return None if any(length < 0 for length in multimodal_embedding_lengths): - raise ValueError("multimodal_embedding_lengths must be non-negative") + raise ValueError("output_embedding_lengths must be non-negative") multimodal_lengths = request.multimodal_lengths if multimodal_lengths is not None: if len(multimodal_embedding_lengths) != len(multimodal_lengths): - raise ValueError("multimodal_embedding_lengths length must match " + raise ValueError("output_embedding_lengths length must match " "multimodal_lengths") for item_idx, (embedding_length, prompt_length) in enumerate( zip(multimodal_embedding_lengths, multimodal_lengths)): if embedding_length > prompt_length: raise ValueError( - f"multimodal_embedding_lengths[{item_idx}] exceeds " + f"output_embedding_lengths[{item_idx}] exceeds " f"multimodal_lengths[{item_idx}]") return multimodal_embedding_lengths @@ -1296,7 +1262,7 @@ def get_multimodal_encoder_token_lengths( Field-level validation happens inside `MultimodalItemMetadata.validate`; only the request-level - cross-check against `multimodal_embedding_lengths` lives here. + cross-check against `output_embedding_lengths` lives here. """ item_metadata = get_multimodal_item_metadata(request.py_multimodal_data) if item_metadata is None: @@ -1307,10 +1273,10 @@ def get_multimodal_encoder_token_lengths( embedding_lengths = get_multimodal_embedding_lengths(request) if embedding_lengths is None: raise ValueError("Multimodal encoder item scheduling requires " - "multimodal_embedding_lengths") + "output_embedding_lengths") if item_metadata.output_embedding_lengths != tuple(embedding_lengths): raise ValueError("Multimodal encoder item output lengths must match " - "multimodal_embedding_lengths") + "output_embedding_lengths") return item_metadata.encoder_token_lengths @@ -1394,9 +1360,6 @@ def initialize_multimodal_encoder_request( (see `get_mm_max_resident_embeddings`). """ mm_data = request.py_multimodal_data - if isinstance(mm_data, dict): - normalize_multimodal_item_metadata( - mm_data, getattr(request, "py_mm_item_order", None)) has_raw_payload = isinstance(mm_data, dict) and any( isinstance(mm_data.get(modality), dict) for modality in ("image", "video", "audio")) @@ -1425,7 +1388,7 @@ def initialize_multimodal_encoder_request( embedding_lengths = get_multimodal_embedding_lengths(request) if embedding_lengths is None: raise ValueError("Multimodal item scheduling requires " - "multimodal_embedding_lengths") + "output_embedding_lengths") # Prefill slices the cached item outputs by the prompt's embedding # slots, so the items must declare exactly one row per slot. embed_mask_cumsum = mm_data.get("multimodal_embed_mask_cumsum") @@ -1697,7 +1660,6 @@ def executor_request_to_llm_request( cache_salt=executor_request.cache_salt, arrival_time=getattr(executor_request, "py_arrival_time", None), py_multimodal_data=py_multimodal_data, - py_mm_item_order=getattr(executor_request, "py_mm_item_order", None), kv_cache_retention_config=executor_request.kv_cache_retention_config, agent_hierarchy=agent_hierarchy, logprobs_mode=getattr(executor_request, "py_logprobs_mode", diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index ea513db141ab..4cf9d669ef72 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -501,8 +501,6 @@ def _deduce_max_tokens(request: GenerationRequest) -> int: # E/P handoff embedding handles parked under "multimodal_embedding". request.multimodal_params.to_tensor("multimodal_data") executor_request.py_multimodal_data = request.multimodal_params.multimodal_data - if request.multimodal_params.mm_item_order: - executor_request.py_mm_item_order = request.multimodal_params.mm_item_order if request.sampling_params.logits_processor: # For PyTorch backend, we attach logits processors as a dynamic Python attribute diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py index 7d0eb769cbe9..36f0e23d5c72 100644 --- a/tensorrt_llm/inputs/__init__.py +++ b/tensorrt_llm/inputs/__init__.py @@ -6,8 +6,8 @@ from .evs import (compute_retained_tokens_count, compute_retained_tokens_from_tubelet_budget, compute_retention_mask) -from .multimodal import (MultimodalEncoderItemMetadata, MultimodalInput, - MultimodalItemMetadata, MultimodalItemPartition) +from .multimodal import (MultimodalInput, MultimodalItemMetadata, + MultimodalItemPartition) from .multimodal_data import AudioData, BaseModalityData, VideoData # yapf and isort conflict on the following import blocks # yapf: disable @@ -53,7 +53,6 @@ "ExtraProcessedInputs", "BaseMultimodalDummyInputsBuilder", "BaseMultimodalInputProcessor", - "MultimodalEncoderItemMetadata", "MultimodalItemMetadata", "MultimodalItemPartition", "MultimodalPlaceholderMetadata", diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index b8a9fee2d040..f0acda33a9bc 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -120,78 +120,6 @@ def select_multimodal_fields( return selected -class MultimodalEncoderItemMetadata(NamedTuple): - """Prompt-ordered metadata shared by item scheduling and encoder assembly.""" - - item_refs: List[Tuple[str, int]] - """``(modality, local index)`` per atomic item, in prompt order. - - The local index identifies the item within that modality's processed - payload. Legacy processors can supply `MultimodalParams.mm_item_order` - instead. - """ - - encoder_token_lengths: List[int] - """Physical attention-token cost charged to ``encoder_max_num_tokens``. - - For Qwen ViTs this counts patch tokens before the spatial merger. - """ - - output_embedding_lengths: List[int] - """Output embedding rows per item, matching ``multimodal_embedding_lengths``. - - These lengths split encoder outputs; they can differ from encoder cost - and from prompt token counts when only part of a placeholder is embedded. - """ - - def validate(self) -> None: - """Enforce the cross-field contract on producer-supplied metadata.""" - if not all(isinstance(values, list) for values in self): - raise TypeError( - "Multimodal encoder item metadata fields must be lists") - lengths = self.encoder_token_lengths + self.output_embedding_lengths - if not all(isinstance(length, int) for length in lengths): - raise TypeError( - "Multimodal encoder item lengths must contain only integers") - if not (len(self.item_refs) == len(self.encoder_token_lengths) == len( - self.output_embedding_lengths)): - raise ValueError( - "Multimodal encoder item references and lengths must align") - if any(length <= 0 for length in self.encoder_token_lengths): - raise ValueError( - "Multimodal encoder token lengths must be positive") - if any(length <= 0 for length in self.output_embedding_lengths): - raise ValueError( - "Multimodal encoder output embedding lengths must be positive") - - def select(self, - item_indices: Sequence[int]) -> "MultimodalEncoderItemMetadata": - """Select prompt-order items and rebase refs to their sliced payloads. - - Each modality's payload must be sliced in the same selected-item order. - The source metadata remains owned by the original request. - """ - if len(item_indices) == 1: - i = item_indices[0] - return MultimodalEncoderItemMetadata( - [(self.item_refs[i][0], 0)], - [self.encoder_token_lengths[i]], - [self.output_embedding_lengths[i]], - ) - local_counts: dict[str, int] = {} - item_refs = [] - for i in item_indices: - modality = self.item_refs[i][0] - local_index = local_counts.get(modality, 0) - item_refs.append((modality, local_index)) - local_counts[modality] = local_index + 1 - return MultimodalEncoderItemMetadata( - item_refs, - [self.encoder_token_lengths[i] for i in item_indices], - [self.output_embedding_lengths[i] for i in item_indices], - ) - - @dataclass(frozen=True, slots=True, init=False) class MultimodalItemMetadata: """Prompt-ordered items and their partitions in processed model inputs. @@ -224,15 +152,6 @@ def __init__( object.__setattr__(self, "input_partitions", {} if input_partitions is None else input_partitions) - @classmethod - def from_encoder_metadata( - cls, metadata: MultimodalEncoderItemMetadata - ) -> "MultimodalItemMetadata": - """Adapt the legacy producer's lists to the shared immutable record.""" - metadata.validate() - return cls(metadata.item_refs, metadata.encoder_token_lengths, - metadata.output_embedding_lengths) - def validate(self) -> None: """Validate producer data once before dispatch or executor admission.""" if not isinstance(self.item_refs, tuple): @@ -346,8 +265,6 @@ def get_multimodal_item_metadata( if not isinstance(multimodal_data, dict): raise TypeError("multimodal_data must be a dict") metadata = multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) - if isinstance(metadata, MultimodalEncoderItemMetadata): - return MultimodalItemMetadata.from_encoder_metadata(metadata) if metadata is not None and not isinstance(metadata, MultimodalItemMetadata): raise TypeError( @@ -356,74 +273,6 @@ def get_multimodal_item_metadata( return metadata -def normalize_multimodal_item_metadata( - multimodal_data: dict[str, Any], - mm_item_order: list[dict[str, Any]] | None = None, -) -> MultimodalItemMetadata | None: - """Adapt legacy order at the input boundary, without guessing mixed order.""" - metadata = get_multimodal_item_metadata(multimodal_data) - if metadata is None: - lengths = multimodal_data.get("multimodal_embedding_lengths") - if lengths is None: - return None - if mm_item_order: - refs = [(entry["modality"], entry["index"]) - for entry in mm_item_order] - else: - present = [ - m for m in MEDIA_MODALITIES - if multimodal_data.get(m) is not None - ] - if len(present) != 1: - return None - refs = [(present[0], i) for i in range(len(lengths))] - metadata = MultimodalItemMetadata(refs, - output_embedding_lengths=lengths) - metadata.validate() - if mm_item_order and metadata.item_refs != tuple( - (entry["modality"], entry["index"]) for entry in mm_item_order): - raise ValueError( - "Multimodal item references do not match the prompt item order") - legacy_lengths = multimodal_data.get("multimodal_embedding_lengths") - if metadata.output_embedding_lengths is not None and legacy_lengths is not None: - if not isinstance(legacy_lengths, (list, tuple)) or not all( - isinstance(n, int) for n in legacy_lengths): - raise TypeError( - "multimodal_embedding_lengths must be a sequence of integers") - if tuple(legacy_lengths) != metadata.output_embedding_lengths: - raise ValueError( - "Multimodal item output lengths do not match multimodal_embedding_lengths" - ) - multimodal_data.pop("multimodal_embedding_lengths") - multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata - return metadata - - -def get_multimodal_encoder_item_metadata( - multimodal_data: Optional[Dict[str, Any]], -) -> Optional[MultimodalEncoderItemMetadata]: - """Return typed atomic-item metadata from Python multimodal data.""" - if multimodal_data is None: - return None - if not isinstance(multimodal_data, dict): - raise TypeError("multimodal_data must be a dict") - metadata = multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) - if metadata is None: - return None - if isinstance(metadata, MultimodalItemMetadata): - if (metadata.encoder_token_lengths is None - or metadata.output_embedding_lengths is None): - return None - metadata = MultimodalEncoderItemMetadata( - list(metadata.item_refs), list(metadata.encoder_token_lengths), - list(metadata.output_embedding_lengths)) - if not isinstance(metadata, MultimodalEncoderItemMetadata): - raise TypeError(f"{MULTIMODAL_ENCODER_ITEM_METADATA_KEY} must be a " - "MultimodalEncoderItemMetadata") - metadata.validate() - return metadata - - class MultimodalEncoderProgress(Enum): """Readiness derived from the request's cached item outputs.""" @@ -936,7 +785,6 @@ def __post_init__(self): # Extend only after auditing each key's consumers. _CPU_ONLY_MULTIMODAL_DATA_KEYS = frozenset({ "multimodal_embed_mask_cumsum", - "multimodal_embedding_lengths", MULTIMODAL_ENCODER_ITEM_METADATA_KEY, }) @@ -994,12 +842,6 @@ class MultimodalParams: multimodal_data: Optional[Dict[str, Any]] = field(default_factory=dict) multimodal_runtime: Optional[MultimodalRuntimeData] = None input_ids_start_offset: int = 0 - # Prompt-order manifest of data-backed items. Each entry is - # ``{"modality": m, "index": i, "placeholder": p}`` where ``i`` indexes - # ``multi_modal_data[m]``. Encoders that interleave items across modalities - # in prompt order read ``modality``/``index``; ``placeholder`` is carried - # for parity with the tracker manifest. - mm_item_order: Optional[List[Dict[str, Union[str, int]]]] = None # CUDA event recorded on a side stream by the MM encoder prefetch path. # When set, the consume site in `get_multimodal_embeddings` issues a # `wait_event` and records attached embeddings on the current stream before @@ -1018,9 +860,6 @@ class MultimodalParams: def __post_init__(self): if self.multimodal_data is None: self.multimodal_data = {} - if normalize_multimodal_item_metadata(self.multimodal_data, - self.mm_item_order) is not None: - self.mm_item_order = None @classmethod def from_embeddings( @@ -1092,46 +931,15 @@ def item_metadata(self) -> MultimodalItemMetadata | None: def get_item_layout( self) -> tuple[Sequence[tuple[str, int]], Sequence[int]]: - """Resolve prompt-order item refs and output lengths for all modalities. - - Typed metadata takes precedence. Legacy processors supply - ``mm_item_order`` and ``multimodal_embedding_lengths``; a request with - just one modality can omit the order. - """ - # Input processors validate this record before dispatch; layout reads - # do not need to revalidate encoder costs on every forward. - metadata = cast( - Optional[MultimodalEncoderItemMetadata], - self.multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY), - ) - if metadata is not None: - if metadata.output_embedding_lengths is None: - raise ValueError( - "Encoder output assembly requires item output embedding lengths" - ) - return metadata.item_refs, metadata.output_embedding_lengths - - lengths = self.multimodal_data.get("multimodal_embedding_lengths") or [] - if self.mm_item_order: - item_refs = cast( - list[tuple[str, int]], - [(entry["modality"], entry["index"]) - for entry in self.mm_item_order], - ) - return item_refs, lengths - - present = [ - m for m in MEDIA_MODALITIES - if self.multimodal_data.get(m) is not None - ] - if len(present) > 1: + """Read prompt-order refs and output rows from the common item record.""" + metadata = self.item_metadata + if metadata is None: + return (), () + if metadata.output_embedding_lengths is None: raise ValueError( - "Request with multiple modalities present " - f"({present}) must carry item metadata or mm_item_order on MultimodalParams." + "Encoder output assembly requires item output embedding lengths" ) - if present: - return [(present[0], i) for i in range(len(lengths))], lengths - return [], [] + return metadata.item_refs, metadata.output_embedding_lengths def _is_shared_tensor_dict(self, obj: Any) -> bool: """Check if an object is a shared tensor dictionary. @@ -1681,7 +1489,6 @@ def find_mm_token_lengths( _MM_METADATA_ONLY_KEYS = frozenset({ "mrope_config", "multimodal_embed_mask_cumsum", - "multimodal_embedding_lengths", MULTIMODAL_ENCODER_ITEM_METADATA_KEY, "special_token_offsets", "layout_metadata", diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 72b29d179437..11ca6267f50d 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -19,7 +19,7 @@ import importlib import traceback from abc import ABC, abstractmethod -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import (Any, Callable, ClassVar, Dict, List, Mapping, Optional, Protocol, Tuple, Type, TypeVar, Union) @@ -34,19 +34,17 @@ from ..sampling_params import SamplingParams from .content_format import ContentFormat from .data import TextPrompt -from .multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY -from .multimodal import \ - MultimodalEncoderItemMetadata as MultimodalEncoderItemMetadata -from .multimodal import (MultimodalInput, MultimodalItemMetadata, - _as_cpu_tensor, _compute_mm_masks, +# yapf: disable +from .multimodal import (MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalInput, + MultimodalItemMetadata, _as_cpu_tensor, + _compute_mm_masks, _find_mm_embedding_lengths_from_masks, _find_mm_token_runs_from_mask, _find_mm_token_start_pos_from_masks, apply_mm_hashes, - default_hasher, find_mm_token_lengths) -from .multimodal import \ - get_multimodal_encoder_item_metadata as get_multimodal_encoder_item_metadata -from .multimodal import (hexdigest_to_int32, normalize_multimodal_item_metadata, + default_hasher, find_mm_token_lengths, + get_multimodal_item_metadata, hexdigest_to_int32, validate_mm_inputs) +# yapf: enable from .multimodal_data import MEDIA_MODALITIES, serialize_item from .prefix_token_cache import create_prefix_token_cache from .tokenization import tokenize_prompt @@ -66,6 +64,29 @@ def _hash_mm_processor_kwargs(mm_processor_kwargs: Dict[str, Any], return hasher.hexdigest() +def _merge_mm_item_metadata( + prepared: MultimodalItemMetadata | None, + metadata: MultimodalItemMetadata | None, +) -> MultimodalItemMetadata | None: + """Combine prompt layout with model costs and partitions at input preparation.""" + if metadata is None: + return prepared + if prepared is None or prepared is metadata: + return metadata + if metadata.item_refs != prepared.item_refs: + raise ValueError( + "Multimodal item references do not match the prompt item order") + known_rows = prepared.output_embedding_lengths + if known_rows is not None: + if metadata.output_embedding_lengths is None: + return replace(metadata, output_embedding_lengths=known_rows) + if metadata.output_embedding_lengths != known_rows: + raise ValueError( + "Computed item output lengths do not match the prompt embedding slots" + ) + return metadata + + class InputProcessor(Protocol): """ Protocol for InputProcessor classes. @@ -152,42 +173,18 @@ class BaseMultimodalInputProcessor(ABC): # uses `_init_tokenization_cache` and `_encode_with_tokenization_cache`. supports_tokenization_cache: ClassVar[bool] = False - def get_mm_encoder_item_metadata( - self, - prompt_token_ids: List[int], - multimodal_data: Dict[str, Any], - ) -> Optional[MultimodalEncoderItemMetadata]: - """Return item refs, physical costs, and encoder output lengths. - - Models opting into runtime MM encoder scheduling override this hook. - The default keeps legacy multimodal processors unchanged. - """ - return None - def get_mm_item_metadata( self, prompt_token_ids: List[int], multimodal_data: Dict[str, Any], mm_item_order: Optional[List[Dict[str, Any]]] = None, ) -> Optional[MultimodalItemMetadata]: - """Describe processed items independently of encoder scheduling support. + """Read the processed item record, independently of scheduling support. - Keep legacy encoder-cost producers at this boundary. Precomputed - embeddings and other models can supply order and rows without costs. + The framework supplies prompt order and output rows during input + preparation. Model processors add encoder costs and input partitions. """ - if (multimodal_data.get("multimodal_embedding") is None and any( - isinstance(multimodal_data.get(m), dict) - for m in MEDIA_MODALITIES)): - metadata = self.get_mm_encoder_item_metadata( - prompt_token_ids, multimodal_data) - if metadata is not None: - if not isinstance(metadata, MultimodalEncoderItemMetadata): - raise TypeError( - "get_mm_encoder_item_metadata() must return a " - "MultimodalEncoderItemMetadata or None") - return MultimodalItemMetadata.from_encoder_metadata(metadata) - return normalize_multimodal_item_metadata(multimodal_data, - mm_item_order) + return get_multimodal_item_metadata(multimodal_data) def __init__(self, model_path, @@ -1324,7 +1321,6 @@ def multimodal_hashing_process( "mm_processor_kwargs_hash"] = _hash_mm_processor_kwargs( inputs.get("mm_processor_kwargs") or {}, hash_lib) - # TODO: here we assume there is only one modality for now num_mm_tokens_by_key = find_mm_token_lengths( mm_data, input_processor, @@ -1344,13 +1340,25 @@ def multimodal_hashing_process( # index items at the same offsets — project into prompt order via the # manifest. mm_item_order = inputs.get("mm_item_order") + metadata = get_multimodal_item_metadata(multimodal_data) if mm_item_order: - num_mm_tokens = [ - num_mm_tokens_by_key[e["modality"]][e["index"]] - for e in mm_item_order - ] + item_refs = tuple( + (entry["modality"], entry["index"]) for entry in mm_item_order) + elif metadata is not None: + item_refs = metadata.item_refs else: - num_mm_tokens = next(iter(num_mm_tokens_by_key.values())) + present = [ + m for m, lengths in num_mm_tokens_by_key.items() if lengths + ] + if len(present) != 1: + raise ValueError( + "Mixed multimodal inputs require prompt-ordered item metadata" + ) + modality = present[0] + item_refs = tuple( + (modality, i) + for i in range(len(num_mm_tokens_by_key[modality]))) + num_mm_tokens = [num_mm_tokens_by_key[m][i] for m, i in item_refs] if len(num_mm_tokens) <= 0: raise ValueError("multimodal hashing produced an empty multimodal " "token-length list.") @@ -1387,37 +1395,21 @@ def multimodal_hashing_process( multimodal_embedding_lengths = ( _find_mm_embedding_lengths_from_masks(mm_mask, embed_mask, num_mm_tokens)) - extra_processed_inputs["multimodal_data"][ - "multimodal_embedding_lengths"] = multimodal_embedding_lengths + prepared_metadata = MultimodalItemMetadata( + item_refs, output_embedding_lengths=multimodal_embedding_lengths) + multimodal_data[ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = _merge_mm_item_metadata( + prepared_metadata, + get_multimodal_item_metadata(multimodal_data)) # Store special token offsets if available if len(start_special_token_positions ) > 0 and mm_special_token_ids is not None: extra_processed_inputs["multimodal_data"][ "special_token_offsets"] = start_special_token_positions - # Same prompt-order projection as `num_mm_tokens` above — the cache - # key indexes each item's digest by its `start_positions` offset. - mm_item_order = inputs.get("mm_item_order") - if mm_item_order: - mm_hashes_flat = [ - mm_hashes[e["modality"]][e["index"]] for e in mm_item_order - ] - else: - mm_hashes_flat = [ - h for hashes in mm_hashes.values() for h in hashes - ] - # `MultimodalInput.multimodal_uuids` must index in lockstep with - # `multimodal_hashes`, so project through the same manifest. - if mm_uuids_by_key is None: - mm_uuid_list = None - elif mm_item_order: - mm_uuid_list = [ - mm_uuids_by_key[e["modality"]][e["index"]] - for e in mm_item_order - ] - else: - mm_uuid_list = [ - u for uuids in mm_uuids_by_key.values() for u in uuids - ] + # Hashes, UUIDs and token spans use the same prompt-ordered item refs. + mm_hashes_flat = [mm_hashes[m][i] for m, i in item_refs] + mm_uuid_list = (None if mm_uuids_by_key is None else + [mm_uuids_by_key[m][i] for m, i in item_refs]) validate_mm_inputs(prompt_token_ids, mm_hashes_flat, start_positions, num_mm_tokens) mm_hashes_int32 = [hexdigest_to_int32(h) for h in mm_hashes_flat @@ -1501,43 +1493,28 @@ def input_processor_wrapper( if not isinstance(multimodal_data, dict): return prompt_token_ids, extra_processed_inputs - mm_item_order = inputs.get("mm_item_order") + prepared_metadata = get_multimodal_item_metadata(multimodal_data) get_item_metadata = getattr(input_processor, "get_mm_item_metadata", None) - if get_item_metadata is not None: - item_metadata = get_item_metadata(prompt_token_ids, multimodal_data, - mm_item_order) - else: - # Registered processors inherit the generic hook. Accept older - # standalone processors at the same boundary. - get_encoder_metadata = getattr(input_processor, - "get_mm_encoder_item_metadata", None) - if (get_encoder_metadata is not None - and multimodal_data.get("multimodal_embedding") is None - and any( - isinstance(multimodal_data.get(m), dict) - for m in MEDIA_MODALITIES)): - legacy_metadata = get_encoder_metadata(prompt_token_ids, - multimodal_data) - if legacy_metadata is not None: - if not isinstance(legacy_metadata, - MultimodalEncoderItemMetadata): - raise TypeError( - "get_mm_encoder_item_metadata() must return a " - "MultimodalEncoderItemMetadata or None") - multimodal_data[ - MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = legacy_metadata - item_metadata = normalize_multimodal_item_metadata( - multimodal_data, mm_item_order) - if item_metadata is None: - return prompt_token_ids, extra_processed_inputs - if not isinstance(item_metadata, MultimodalItemMetadata): + item_metadata = (get_item_metadata(prompt_token_ids, multimodal_data, + inputs.get("mm_item_order")) if + get_item_metadata is not None else prepared_metadata) + if item_metadata is not None and not isinstance(item_metadata, + MultimodalItemMetadata): raise TypeError( "get_mm_item_metadata() must return a MultimodalItemMetadata or None" ) + item_metadata = _merge_mm_item_metadata(prepared_metadata, + item_metadata) + if item_metadata is None: + return prompt_token_ids, extra_processed_inputs + mm_item_order = inputs.get("mm_item_order") + if mm_item_order and item_metadata.item_refs != tuple( + (entry["modality"], entry["index"]) for entry in mm_item_order): + raise ValueError( + "Multimodal item references do not match the prompt item order") item_metadata.validate() multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = item_metadata - normalize_multimodal_item_metadata(multimodal_data, mm_item_order) return prompt_token_ids, extra_processed_inputs return input_processor_wrapper diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 76ff7a6743aa..3f69af4d00f8 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1074,7 +1074,7 @@ def _preprocess( 'multimodal_input'), multimodal_data=extra_processed_inputs.get( 'multimodal_data'), - mm_item_order=inputs.get("mm_item_order")) + ) # Only pass it if it has content if not multimodal_params.has_content(): multimodal_params = None diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index 9360cca97ae0..c69db80ae663 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -38,10 +38,7 @@ MultimodalParams, MultimodalRuntimeData, ) -from tensorrt_llm.inputs.registry import ( - BaseMultimodalDummyInputsBuilder, - MultimodalEncoderItemMetadata, -) +from tensorrt_llm.inputs.registry import BaseMultimodalDummyInputsBuilder, MultimodalItemMetadata from tensorrt_llm.llmapi.llm_args import ( MultimodalConfig, MultimodalEncoderCudaGraphConfig, @@ -708,12 +705,11 @@ def forward_multimodal_encoder_items(self, encoder_inputs): "pixel_values": torch.arange(5).unsqueeze(1), "image_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 3]]), }, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( item_refs=[("image", 0), ("image", 1)], encoder_token_lengths=[2, 3], output_embedding_lengths=[2, 3], ), - "multimodal_embedding_lengths": [2, 3], "multimodal_embed_mask_cumsum": embed_mask_cumsum, }, input_tokens=list(range(8)), @@ -771,12 +767,11 @@ def make_request(embed_mask_cumsum: list[int]) -> LlmRequest: 1, multimodal_data={ "image": {"pixel_values": torch.empty(5, 1)}, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( item_refs=[("image", 0), ("image", 1)], encoder_token_lengths=[2, 3], output_embedding_lengths=[2, 3], ), - "multimodal_embedding_lengths": [2, 3], "multimodal_embed_mask_cumsum": torch.tensor(embed_mask_cumsum), }, ) @@ -801,12 +796,11 @@ class _Model(MultimodalModelMixin): "pixel_values": torch.arange(2).unsqueeze(1), "image_grid_thw": torch.tensor([[1, 1, 2]]), }, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( item_refs=[("image", 0)], encoder_token_lengths=[2], output_embedding_lengths=[2], ), - "multimodal_embedding_lengths": [2], }, ) initialize_multimodal_encoder_request(request, max_num_tokens=8) diff --git a/tests/unittest/_torch/executor/engine/test_no_kv_cache.py b/tests/unittest/_torch/executor/engine/test_no_kv_cache.py index 5f6dd2017a00..870e61d068d7 100644 --- a/tests/unittest/_torch/executor/engine/test_no_kv_cache.py +++ b/tests/unittest/_torch/executor/engine/test_no_kv_cache.py @@ -13,6 +13,7 @@ from tensorrt_llm._torch.pyexecutor.engine.runners import no_kv_cache as no_kv_cache_module from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunnerConfig from tensorrt_llm._torch.pyexecutor.engine.runners.pooling import PoolingRunner +from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend pytestmark = pytest.mark.cpu_only @@ -43,12 +44,15 @@ def _request( position_ids: list[int] | None = None, multimodal_data: dict | None = None, ) -> SimpleNamespace: + params = ( + MultimodalParams(multimodal_data=multimodal_data) if multimodal_data is not None else None + ) return SimpleNamespace( + py_multimodal_params=params, get_tokens=Mock(return_value=tokens), py_request_id=request_id, py_position_ids=position_ids, py_multimodal_data=multimodal_data, - py_mm_item_order=None, py_seq_slot=request_id, ) @@ -171,7 +175,6 @@ def test_no_kv_cache_runner_prepare_inputs_builds_and_ships_multimodal_inputs( monkeypatch: pytest.MonkeyPatch, ) -> None: request = _request(3, [1, 101], multimodal_data={"image": object()}) - request.py_mm_item_order = ["image"] multimodal_params = Mock() multimodal_params_factory = Mock(return_value=multimodal_params) multimodal_input = object() @@ -202,7 +205,6 @@ def test_no_kv_cache_runner_prepare_inputs_builds_and_ships_multimodal_inputs( multimodal_params_factory.assert_called_once_with( multimodal_input=multimodal_input, multimodal_data=request.py_multimodal_data, - mm_item_order=request.py_mm_item_order, input_ids_start_offset=0, ) multimodal_params.to_device.assert_called_once_with( diff --git a/tests/unittest/_torch/executor/engine/test_runners.py b/tests/unittest/_torch/executor/engine/test_runners.py index 6164c792e9e3..931414d0c155 100644 --- a/tests/unittest/_torch/executor/engine/test_runners.py +++ b/tests/unittest/_torch/executor/engine/test_runners.py @@ -36,6 +36,10 @@ from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.shared_tensor import SharedTensorContainer from tensorrt_llm._torch.utils import get_model_extra_attrs, model_extra_attrs +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, +) from tensorrt_llm.llmapi.llm_args import ( CudaGraphConfig, EncodeCudaGraphConfig, @@ -476,7 +480,9 @@ def test_mm_encoder_runner_forward_step_skips_metadata_only_and_missing_length_r SimpleNamespace( py_multimodal_data={ "image": object(), - "multimodal_embedding_lengths": [2], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[2] + ), }, multimodal_lengths=[2], ), @@ -524,14 +530,18 @@ def test_mm_encoder_runner_splits_embeddings_and_returns_mrope_metadata() -> Non SimpleNamespace( py_multimodal_data={ "image": object(), - "multimodal_embedding_lengths": [1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[1] + ), }, multimodal_lengths=[1], ), SimpleNamespace( py_multimodal_data={ "image": object(), - "multimodal_embedding_lengths": [2], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[2] + ), }, multimodal_lengths=[2], ), @@ -744,12 +754,22 @@ def test_model_caller_requires_streams_with_compile_backend(): def test_mm_encoder_handoff_preserves_per_item_metadata_in_sparse_batch() -> None: requests = [ SimpleNamespace( - py_multimodal_data={"image": object(), "multimodal_embedding_lengths": [1]}, + py_multimodal_data={ + "image": object(), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[1] + ), + }, multimodal_lengths=[1], ), SimpleNamespace(py_multimodal_data=None), SimpleNamespace( - py_multimodal_data={"video": object(), "multimodal_embedding_lengths": [3, 2]}, + py_multimodal_data={ + "video": object(), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("video", 0), ("video", 1)], output_embedding_lengths=[3, 2] + ), + }, multimodal_lengths=[7, 6], ), ] diff --git a/tests/unittest/_torch/executor/multimodal_utils.py b/tests/unittest/_torch/executor/multimodal_utils.py index 745b12e6f459..7cb2d40d4a2e 100644 --- a/tests/unittest/_torch/executor/multimodal_utils.py +++ b/tests/unittest/_torch/executor/multimodal_utils.py @@ -22,10 +22,7 @@ from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.bindings import SamplingConfig from tensorrt_llm.inputs.multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY -from tensorrt_llm.inputs.registry import ( - BaseMultimodalDummyInputsBuilder, - MultimodalEncoderItemMetadata, -) +from tensorrt_llm.inputs.registry import BaseMultimodalDummyInputsBuilder, MultimodalItemMetadata def bare_mm_item_scheduler( @@ -86,12 +83,11 @@ def make_mm_request( embedding_lengths = [1] * len(costs) multimodal_data = { "image": {"pixel_values": torch.empty(len(costs), 1)}, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( item_refs=[("image", item_idx) for item_idx in range(len(costs))], encoder_token_lengths=costs, output_embedding_lengths=list(embedding_lengths), ), - "multimodal_embedding_lengths": list(embedding_lengths), } input_tokens = None if interleaved: diff --git a/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py b/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py index 968e80f011e7..0dedfff93435 100644 --- a/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py +++ b/tests/unittest/_torch/executor/test_multimodal_embedding_lengths.py @@ -33,16 +33,22 @@ ( SimpleNamespace( multimodal_lengths=[6, 5], - py_multimodal_data={"multimodal_embedding_lengths": [5, 3]}, + py_multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0), ("image", 1)], output_embedding_lengths=[5, 3] + ) + }, ), - [5, 3], + (5, 3), ), ( SimpleNamespace( multimodal_lengths=[6, 5], py_multimodal_data={ "layout_metadata": { - "multimodal_embedding_lengths": [6, 4], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0), ("image", 1)], output_embedding_lengths=[6, 4] + ), }, }, ), @@ -56,41 +62,28 @@ def test_multimodal_embedding_lengths_returns_top_level_metadata(req, expected): assert get_multimodal_embedding_lengths(req) == expected -@pytest.mark.parametrize("legacy_lengths", [None, [2]]) @pytest.mark.cpu_only -def test_multimodal_embedding_lengths_uses_item_metadata(legacy_lengths): +def test_multimodal_embedding_lengths_uses_item_metadata(): """Encoder output export uses the same rows as model input assembly.""" metadata = MultimodalItemMetadata([("image", 0)], output_embedding_lengths=[2]) data = {MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata} - if legacy_lengths is not None: - data["multimodal_embedding_lengths"] = legacy_lengths params = MultimodalParams(multimodal_data=data) req = SimpleNamespace(multimodal_lengths=[4], py_multimodal_data=params.multimodal_data) assert get_multimodal_embedding_lengths(req) is params.get_item_layout()[1] -@pytest.mark.cpu_only -def test_multimodal_embedding_lengths_rejects_conflicting_item_metadata(): - """A direct Params caller cannot export a different item layout.""" - with pytest.raises(ValueError, match="output lengths do not match"): - MultimodalParams( - multimodal_data={ - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( - [("image", 0)], output_embedding_lengths=[2] - ), - "multimodal_embedding_lengths": [3], - } - ) - - @pytest.mark.parametrize( "req,exception,match", [ ( SimpleNamespace( multimodal_lengths=[4], - py_multimodal_data={"multimodal_embedding_lengths": [3, 1]}, + py_multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0), ("image", 1)], output_embedding_lengths=[3, 1] + ) + }, ), ValueError, "length must match", @@ -98,7 +91,11 @@ def test_multimodal_embedding_lengths_rejects_conflicting_item_metadata(): ( SimpleNamespace( multimodal_lengths=[4], - py_multimodal_data={"multimodal_embedding_lengths": [5]}, + py_multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[5] + ) + }, ), ValueError, "exceeds", @@ -106,7 +103,11 @@ def test_multimodal_embedding_lengths_rejects_conflicting_item_metadata(): ( SimpleNamespace( multimodal_lengths=[4], - py_multimodal_data={"multimodal_embedding_lengths": [-1]}, + py_multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[-1] + ) + }, ), ValueError, "non-negative", @@ -115,16 +116,23 @@ def test_multimodal_embedding_lengths_rejects_conflicting_item_metadata(): SimpleNamespace( multimodal_lengths=[4], py_multimodal_data={ - "multimodal_embedding_lengths": torch.tensor([4]), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", i) for i in range(len(torch.tensor([4])))], + output_embedding_lengths=torch.tensor([4]), + ), }, ), TypeError, - "must be a list", + "must contain only integers", ), ( SimpleNamespace( multimodal_lengths=[4], - py_multimodal_data={"multimodal_embedding_lengths": (4.0,)}, + py_multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=(4.0,) + ) + }, ), TypeError, "integers", @@ -135,7 +143,7 @@ def test_multimodal_embedding_lengths_rejects_conflicting_item_metadata(): py_multimodal_data=["multimodal_embedding_lengths", [4]], ), TypeError, - "py_multimodal_data must be a dict", + "multimodal_data must be a dict", ), ], ) diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py index 1e62af2d1195..60bbea141dc5 100644 --- a/tests/unittest/_torch/executor/test_multimodal_scheduler.py +++ b/tests/unittest/_torch/executor/test_multimodal_scheduler.py @@ -13,7 +13,10 @@ record_output, ) -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( + MultimodalEncoderContractError, + MultimodalModelMixin, +) from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( @@ -45,7 +48,6 @@ MultimodalParams, strip_mm_encoder_inputs, ) -from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata from tensorrt_llm.llmapi.llm_args import ( CapacitySchedulerPolicy, ContextChunkingPolicy, @@ -67,11 +69,6 @@ def test_multimodal_request_aliases_share_one_aggregate(): params = request.py_multimodal_params assert request.py_multimodal_data is params.multimodal_data assert request.py_mm_encoder_state is params.encoder_state - order = [{"modality": "image", "index": 0}, {"modality": "image", "index": 1}] - request.py_mm_item_order = order - assert params.mm_item_order is order - params.mm_item_order = list(reversed(order)) - assert request.py_mm_item_order is params.mm_item_order event = object() request.py_mm_encoder_event = event assert params.encoder_event is event @@ -87,7 +84,6 @@ def test_multimodal_request_aliases_share_one_aggregate(): name not in request.__dict__ for name in ( "py_multimodal_data", - "py_mm_item_order", "py_mm_encoder_state", "py_mm_encoder_event", ) @@ -98,7 +94,6 @@ def test_text_requests_do_not_allocate_multimodal_params(): request = make_llm_request(1) request.py_mm_encoder_state = None request.py_mm_encoder_event = None - request.py_mm_item_order = None request.py_multimodal_data = None assert request.py_multimodal_params is None @@ -109,7 +104,6 @@ def test_generic_metadata_does_not_enable_item_scheduling(): 1, { "image": {"pixel_values": torch.zeros(1, 2)}, - "multimodal_embedding_lengths": [1], MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, }, ) @@ -318,6 +312,7 @@ def test_mm_encoder_token_lengths_distinguishes_missing_and_invalid_data(): with pytest.raises(TypeError, match="multimodal_data must be a dict"): request.py_multimodal_data = [] + get_multimodal_encoder_token_lengths(request) def test_mm_encoder_readiness_is_derived_from_item_state(): @@ -447,12 +442,11 @@ def test_admission_rejects_requests_larger_than_output_budget(): 1, multimodal_data={ "video": {"pixel_values_videos": torch.empty(3, 1)}, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( item_refs=[("video", 0)], encoder_token_lengths=[12], output_embedding_lengths=[3], ), - "multimodal_embedding_lengths": [3], }, ) with pytest.raises(ValueError, match="raise encoder_max_num_tokens") as exc_info: @@ -812,7 +806,7 @@ def test_forward_multimodal_encoder_step_contains_model_contract_error(): def fail_encoder(*_): raise MultimodalEncoderRequestError( - "multimodal_encoder_item_metadata must be a MultimodalEncoderItemMetadata" + "multimodal_encoder_item_metadata must be a MultimodalItemMetadata" ) engine = SimpleNamespace(forward_multimodal_encoder_items=fail_encoder) @@ -833,7 +827,7 @@ def fail_encoder(*_): assert scheduled_requests.context_requests == [unrelated] assert scheduled_requests.scheduled_mm_encoder_items is None assert len(handled) == 1 - assert "must be a MultimodalEncoderItemMetadata" in handled[0][0] + assert "must be a MultimodalItemMetadata" in handled[0][0] assert handled[0][1] == {"requests": [failed], "charge_budget": False} @@ -914,12 +908,11 @@ def _waiting_item(request_id, costs=None): multimodal_data = None if costs is not None: multimodal_data = { - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( item_refs=[("image", item_idx) for item_idx in range(len(costs))], encoder_token_lengths=costs, output_embedding_lengths=[1] * len(costs), ), - "multimodal_embedding_lengths": [1] * len(costs), } return RequestQueueItem( request_id, @@ -1030,12 +1023,11 @@ def encode_multimodal_inputs(self, multimodal_params): "pixel_values": torch.arange(5).unsqueeze(1), "image_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 3]]), }, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( item_refs=[("image", 0), ("image", 1)], encoder_token_lengths=[2, 3], output_embedding_lengths=[2, 3], ), - "multimodal_embedding_lengths": [2, 3], } ) @@ -1058,12 +1050,11 @@ def test_prepare_multimodal_encoder_inputs_slices_before_device_transfer(): "pixel_values": torch.arange(5).unsqueeze(1), "image_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 3]]), }, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( item_refs=[("image", 0), ("image", 1)], encoder_token_lengths=[2, 3], output_embedding_lengths=[2, 3], ), - "multimodal_embedding_lengths": [2, 3], } ) @@ -1079,13 +1070,9 @@ def test_prepare_multimodal_encoder_inputs_slices_before_device_transfer(): def test_prepare_multimodal_encoder_inputs_rejects_invalid_metadata_types(): - with pytest.raises(TypeError, match="must be a MultimodalItemMetadata"): - MultimodalParams( - multimodal_data={ - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: ("image", 0), - "multimodal_embedding_lengths": [1], - } - ) + params = MultimodalParams(multimodal_data={MULTIMODAL_ENCODER_ITEM_METADATA_KEY: ("image", 0)}) + with pytest.raises(MultimodalEncoderContractError, match="must be a MultimodalItemMetadata"): + MultimodalModelMixin().prepare_multimodal_encoder_inputs([(params, 0)]) def test_strip_mm_encoder_inputs_preserves_embedding_and_runtime_metadata(): diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index b26cb976be99..2016e33240ac 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -66,6 +66,8 @@ from tensorrt_llm._torch.speculative.spec_sampler_base import \ SampleStateTensorsSpec from tensorrt_llm._torch.utils import set_per_request_prefill_cuda_graph_flag +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalItemMetadata) from tensorrt_llm.inputs.registry import (BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor) from tensorrt_llm.llmapi import (CudaGraphConfig, MTPDecodingConfig, @@ -3726,7 +3728,10 @@ def test_item_scheduled_chunk_embedding_is_not_retained(self) -> None: context.context_chunk_size = 2 context.py_seq_slot = 0 context.py_mm_encoder_state = object() - multimodal_data = {"multimodal_embedding_lengths": [2]} + multimodal_data = { + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: + MultimodalItemMetadata([('image', 0)], output_embedding_lengths=[2]) + } context.py_multimodal_data = multimodal_data model_engine._mm_item_scheduler = Mock() model_engine._runner._inputs._mm_item_scheduler = model_engine._mm_item_scheduler diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index 85937f91a71d..6acd9675f497 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/tests/unittest/_torch/executor/test_request_utils.py @@ -56,7 +56,6 @@ def test_multimodal_payload_broadcast_and_children_keep_the_same_contract(): MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, "audio": {"features": torch.arange(6).reshape(2, 3)}, "image": {}, - "multimodal_embedding_lengths": [1, 2], } source = trtllm.Request( input_token_ids=[1, 2, 3], diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py b/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py index c80425c5deb9..7360ad670233 100644 --- a/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py +++ b/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py @@ -165,7 +165,8 @@ def test_processor_preserves_interleaved_text_and_image_payload() -> None: data = extra["multimodal_data"] image = data["image"] assert image["positions"] == [1, 6] - assert image["num_tokens"] == data["multimodal_embedding_lengths"] == [4, 10] + assert image["num_tokens"] == [4, 10] + assert data["multimodal_encoder_item_metadata"].output_embedding_lengths == (4, 10) assert image["image_grid_hw"] == [(3, 3), (4, 7)] assert image["llm_grid_hw"] == [(1, 1), (2, 3)] assert image["token_types"][1].tolist() == [0, 1, 1, 1, 2, 1, 1, 1, 2, 3] diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index fec02fa2f2a1..cb2e1ab993a7 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -79,8 +79,8 @@ ) from tensorrt_llm.inputs.multimodal import ( # noqa: E402 MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalEncoderItemMetadata, MultimodalInput, + MultimodalItemMetadata, MultimodalParams, MultimodalRuntimeData, ) @@ -208,7 +208,10 @@ def _make_keyed_image_param( "image_position_ids": image_position_ids, "image_seq_lens": [1] * item_count, }, - "multimodal_embedding_lengths": embedding_lengths, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", i) for i in range(len(embedding_lengths))], + output_embedding_lengths=embedding_lengths, + ), "mm_processor_kwargs_hash": "kwargs-a", }, multimodal_runtime=MultimodalRuntimeData( @@ -230,7 +233,10 @@ def _make_keyed_video_param(item_hashes: list[list[int]]) -> MultimodalParams: ), multimodal_data={ "video": {"pixel_values": torch.arange(item_count).reshape(item_count, 1, 1)}, - "multimodal_embedding_lengths": embedding_lengths, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("video", i) for i in range(len(embedding_lengths))], + output_embedding_lengths=embedding_lengths, + ), "mm_processor_kwargs_hash": "kwargs-a", }, multimodal_runtime=MultimodalRuntimeData( @@ -255,7 +261,10 @@ def _make_keyed_audio_param(item_hashes: list[list[int]]) -> MultimodalParams: "audio_features": torch.arange(item_count * 2).reshape(item_count, 1, 2), "audio_features_mask": torch.ones(item_count, 1), }, - "multimodal_embedding_lengths": embedding_lengths, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("audio", i) for i in range(len(embedding_lengths))], + output_embedding_lengths=embedding_lengths, + ), "mm_processor_kwargs_hash": "kwargs-a", }, multimodal_runtime=MultimodalRuntimeData( @@ -957,7 +966,9 @@ def test_build_multimodal_encoder_input_slices_gemma4_audio_input(self): "audio_features_mask": audio_mask, "shared_stats": shared_stats, }, - "multimodal_embedding_lengths": [1, 1, 1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("audio", 0), ("audio", 1), ("audio", 2)], output_embedding_lengths=[1, 1, 1] + ), } ) @@ -1100,7 +1111,7 @@ def test_build_multimodal_encoder_input_uses_typed_audio_item_count(self): param = MultimodalParams( multimodal_data={ "audio": {"audio_features": features}, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( [("audio", 0), ("audio", 1), ("audio", 2)], [8, 8, 8], [1, 1, 1] ), } diff --git a/tests/unittest/_torch/modeling/test_modeling_mistral.py b/tests/unittest/_torch/modeling/test_modeling_mistral.py index 1e0b21ea44ea..90944222db3e 100644 --- a/tests/unittest/_torch/modeling/test_modeling_mistral.py +++ b/tests/unittest/_torch/modeling/test_modeling_mistral.py @@ -722,13 +722,11 @@ def test_mistral_item_metadata_separates_patch_and_embedding_units(): processor = object.__new__(MistralHFInputProcessor) processor._vision_geometry = lambda: (14, 2, 3, 1024) - metadata = processor.get_mm_encoder_item_metadata( - [], {"image": {"image_sizes": [[28, 56], [56, 56]]}} - ) + metadata = processor.get_mm_item_metadata([], {"image": {"image_sizes": [[28, 56], [56, 56]]}}) - assert metadata.item_refs == [("image", 0), ("image", 1)] - assert metadata.encoder_token_lengths == [8, 16] - assert metadata.output_embedding_lengths == [2, 4] + assert metadata.item_refs == (("image", 0), ("image", 1)) + assert metadata.encoder_token_lengths == (8, 16) + assert metadata.output_embedding_lengths == (2, 4) @pytest.mark.parametrize("budget", [1024, 4096, 8192]) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py index b58de5675f06..cf61ac09c2e1 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py @@ -38,7 +38,12 @@ default_multimodal_input_loader, prompt_inputs, ) -from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, + MultimodalParams, + MultimodalRuntimeData, +) from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.llmapi.llm import RequestOutput from tensorrt_llm.llmapi.llm_args import ( @@ -1049,7 +1054,9 @@ def _make_param_with_runtime(self, modality, num_tokens, **extra): multimodal_data={ "modality_type": modality, modality: nested, - "multimodal_embedding_lengths": [num_tokens], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [(modality, 0)], output_embedding_lengths=[num_tokens] + ), **extra, }, multimodal_runtime=runtime, diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 9b9ffdc38c9a..6f13de48f0e7 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -24,7 +24,6 @@ from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalEncoderItemMetadata, MultimodalEncoderRequestState, MultimodalInput, MultimodalItemMetadata, @@ -83,6 +82,63 @@ def test_qwen_common_metadata_preserves_positive_output_requirement(): processor.get_mm_item_metadata([12, 10, 13], data) +@pytest.mark.parametrize("model", ["qwen2", "qwen2_5", "qwen3", "mistral"]) +@pytest.mark.parametrize("custom", [False, True]) +def test_processors_use_the_common_metadata_hook(model, custom): + from dataclasses import replace + + from tensorrt_llm._torch.models.modeling_mistral import MistralHFInputProcessor + from tensorrt_llm._torch.models.modeling_qwen2vl import ( + Qwen2_5VLInputProcessorBase, + Qwen2VLInputProcessorBase, + ) + from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLInputProcessorBase + + processor_cls = { + "qwen2": Qwen2VLInputProcessorBase, + "qwen2_5": Qwen2_5VLInputProcessorBase, + "qwen3": Qwen3VLInputProcessorBase, + "mistral": MistralHFInputProcessor, + }[model] + calls = [] + + class CustomProcessor(processor_cls): + def get_mm_item_metadata(self, prompt_token_ids, multimodal_data, mm_item_order=None): + calls.append(prompt_token_ids) + metadata = super().get_mm_item_metadata( + prompt_token_ids, multimodal_data, mm_item_order + ) + return replace(metadata, encoder_token_lengths=[80]) + + processor = object.__new__(CustomProcessor if custom else processor_cls) + processor._config = SimpleNamespace( + image_token_id=10, + video_token_id=11, + vision_start_token_id=12, + vision_end_token_id=13, + vision_config=SimpleNamespace( + spatial_merge_size=2, patch_size=2, image_size=8, num_channels=3 + ), + ) + processor._processor = SimpleNamespace() + prompt = [12, 10, 10, 13] + image = {"pixel_values": torch.zeros(8, 2), "image_grid_thw": torch.tensor([[1, 2, 4]])} + if model == "mistral": + image = {"pixel_values": torch.zeros(1, 3, 4, 8), "image_sizes": [[4, 8]]} + metadata = processor.get_mm_item_metadata(prompt, {"image": image}) + assert calls == ([prompt] if custom else []) + assert metadata.item_refs == (("image", 0),) + assert metadata.output_embedding_lengths == (2,) + assert metadata.encoder_token_lengths == ((80,) if custom else (8,)) + metadata.validate() + if model != "mistral": + partition = metadata.input_partitions["image"]["pixel_values"] + assert partition.item_lengths == (8,) + torch.testing.assert_close( + partition.select(image["pixel_values"], [0]), image["pixel_values"] + ) + + def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor @@ -90,7 +146,9 @@ def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): data = { "image": {"pixel_values": torch.zeros(1, 4, 6), "image_position_ids": torch.zeros(1, 4, 2)}, "audio": {"audio_features": torch.zeros(1, 5, 8), "audio_features_mask": torch.ones(1, 5)}, - "multimodal_embedding_lengths": [2, 4], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("audio", 0), ("image", 0)], output_embedding_lengths=[2, 4] + ), } order = [{"modality": "audio", "index": 0}, {"modality": "image", "index": 0}] metadata = processor.get_mm_item_metadata([], data, order) @@ -112,7 +170,12 @@ def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): "feature_attention_mask": torch.ones(3, 5), "audio_num_clips": torch.tensor([2, 1]), } - data = {"audio": audio, "multimodal_embedding_lengths": [4, 2]} + data = { + "audio": audio, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("audio", 0), ("audio", 1)], output_embedding_lengths=[4, 2] + ), + } metadata = processor.get_mm_item_metadata([], data) metadata.validate() params = MultimodalParams(multimodal_data=data) @@ -127,7 +190,12 @@ def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): ) assert selected.encoder_state is None video = {"audio": audio, "video_size": [[2, 1, 32, 32]]} - video_data = {"video": video, "multimodal_embedding_lengths": [6]} + video_data = { + "video": video, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("video", 0)], output_embedding_lengths=[6] + ), + } video_metadata = processor.get_mm_item_metadata([], video_data) assert video_metadata.item_refs == (("video", 0),) assert "audio" not in video_metadata.input_partitions @@ -186,6 +254,28 @@ def encode_multimodal_inputs(self, multimodal_params) -> torch.Tensor: return encode_multimodal_by_groups(self.mm_encoder_groups, list(multimodal_params)) +def test_common_metadata_shares_executor_and_encoder_layout(): + from tensorrt_llm._torch.pyexecutor.llm_request import get_multimodal_embedding_lengths + + pixels = torch.arange(5 * HIDDEN).reshape(5, HIDDEN) + params = MultimodalParams( + multimodal_data={ + "image": {"pixel_values": pixels}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 1), ("image", 0)], output_embedding_lengths=[3, 2] + ), + } + ) + request = SimpleNamespace(py_multimodal_data=params.multimodal_data, multimodal_lengths=[3, 2]) + + assert get_multimodal_embedding_lengths(request) == (3, 2) + actual = _GroupedEncoderModel().encode_multimodal_inputs([params]) + + torch.testing.assert_close(actual, torch.cat((pixels[2:], pixels[:2]))) + assert params.item_metadata.encoder_token_lengths is None + assert "multimodal_embedding_lengths" not in params.multimodal_data + + def _make_request(modalities: list[str], *, value_start: float = 1.0) -> MultimodalParams: """Build a request whose items carry unique, checkable patch values.""" per_modality_rows: dict[str, list[int]] = {} @@ -211,14 +301,13 @@ def _make_request(modalities: list[str], *, value_start: float = 1.0) -> Multimo pixel_key: torch.cat(blocks, dim=0), grid_key: torch.tensor([_grid_for_rows(r) for r in rows_per_item]), } - multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalEncoderItemMetadata( + multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( item_refs=item_refs, encoder_token_lengths=list(lengths), output_embedding_lengths=list(lengths), ) # Processors retain this field for legacy consumers. Encoder groups use # the typed item metadata when available. - multimodal_data["multimodal_embedding_lengths"] = list(lengths) return MultimodalParams(multimodal_data=multimodal_data) @@ -283,7 +372,7 @@ def get_multimodal_item_partitions(self, param, modality): "frame_counts": [2, 3, 1], "shared_stats": shared, }, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( [("audio", 0), ("audio", 1), ("audio", 2)], [2, 3, 1], [1, 2, 1] ), } @@ -359,21 +448,9 @@ def test_mixed_modality_items_keep_scheduler_order(): @pytest.mark.parametrize("modalities", [["video", "image", "image"], ["image", "video", "image"]]) -@pytest.mark.parametrize("legacy_order", [False, True]) -def test_full_and_item_encoder_paths_share_item_metadata( - modalities: list[str], legacy_order: bool -) -> None: +def test_full_and_item_encoder_paths_share_item_metadata(modalities: list[str]) -> None: model = _GroupedEncoderModel() request = _make_request(modalities) - metadata = request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] - if legacy_order: - request.mm_item_order = [ - {"modality": modality, "index": index} for modality, index in metadata.item_refs - ] - # Full-request encoding reads both order and lengths from the same - # metadata that drives the scheduler's item selection. - request.multimodal_data.pop("multimodal_embedding_lengths", None) - full_output = model.encode_multimodal_inputs([request]) item_outputs, _ = _run_items(model, [(request, i) for i in range(len(modalities))]) expected = torch.cat([_expected_rows(request, i) for i in range(len(modalities))]) @@ -400,24 +477,15 @@ def test_metadata_is_rebased_for_partial_and_reordered_items() -> None: torch.testing.assert_close(outputs[1], _expected_rows(request, 0), rtol=0, atol=0) -@pytest.mark.parametrize("use_metadata", [False, True]) -def test_full_and_item_order_can_differ_from_raw_payload_order(use_metadata: bool) -> None: +def test_full_and_item_order_can_differ_from_raw_payload_order() -> None: model = _GroupedEncoderModel() request = _make_request(["image", "image"]) - request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalEncoderItemMetadata( + request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( [("image", 1), ("image", 0)], [3, 2], [3, 2] ) expected = torch.cat([_expected_rows(request, 0), _expected_rows(request, 1)]) - request.multimodal_data["multimodal_embedding_lengths"] = [3, 2] - if use_metadata: - item_outputs, _ = _run_items(model, [(request, 0), (request, 1)]) - torch.testing.assert_close(torch.cat(item_outputs), expected, rtol=0, atol=0) - else: - request.multimodal_data.pop(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) - request.mm_item_order = [ - {"modality": "image", "index": 1}, - {"modality": "image", "index": 0}, - ] + item_outputs, _ = _run_items(model, [(request, 0), (request, 1)]) + torch.testing.assert_close(torch.cat(item_outputs), expected, rtol=0, atol=0) torch.testing.assert_close(model.encode_multimodal_inputs([request]), expected, rtol=0, atol=0) @@ -434,9 +502,8 @@ def _attach_cache_metadata(request, item_lengths): request.multimodal_data["mm_processor_kwargs_hash"] = "test" -@pytest.mark.parametrize("metadata_format", ["typed", "typed_only", "legacy"]) @pytest.mark.parametrize("reversed_order", [False, True]) -def test_partial_cache_uses_prompt_refs_and_lengths(metadata_format, reversed_order): +def test_partial_cache_uses_prompt_refs_and_lengths(reversed_order): model = _GroupedEncoderModel() model._multimodal_encoder_cache = TensorLRUCache(4096) request = _make_request(["image", "image"]) @@ -445,14 +512,6 @@ def test_partial_cache_uses_prompt_refs_and_lengths(metadata_format, reversed_or metadata = MultimodalItemMetadata([("image", 1), ("image", 0)], [3, 2], [3, 2]) request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata lengths = metadata.output_embedding_lengths - request.multimodal_data["multimodal_embedding_lengths"] = list(lengths) - if metadata_format == "typed_only": - request.multimodal_data.pop("multimodal_embedding_lengths", None) - elif metadata_format == "legacy": - request.mm_item_order = [ - {"modality": modality, "index": index} for modality, index in metadata.item_refs - ] - request.multimodal_data.pop(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) _attach_cache_metadata(request, lengths) expected = model.encode_multimodal_inputs([request]) keys = model._encoder_cache_keys(request) @@ -465,8 +524,7 @@ def test_partial_cache_uses_prompt_refs_and_lengths(metadata_format, reversed_or model._multimodal_encoder_cache.get(keys[1]), expected[lengths[0] :], rtol=0, atol=0 ) assert model.encoder_calls[-1] == lengths[1] - if metadata_format != "legacy": - assert request.item_metadata is metadata + assert request.item_metadata is metadata def test_full_cache_write_uses_canonical_lengths_without_legacy_fields(): @@ -507,7 +565,7 @@ def test_partial_cache_batches_mixed_vision_and_audio_misses(): request.multimodal_data["audio"] = { "input_features": torch.stack([torch.full((1, HIDDEN), value) for value in (10.0, 20.0)]) } - request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalEncoderItemMetadata( + request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( [("video", 0), ("audio", 0), ("image", 1), ("image", 0), ("audio", 1)], [4, 1, 3, 2, 1], [4, 1, 3, 2, 1], @@ -579,7 +637,7 @@ def build_audio_input(params: list[MultimodalParams]) -> dict[str, torch.Tensor] request = _make_request(["video", "image"]) audio = torch.stack([torch.full((2, HIDDEN), value) for value in (10.0, 20.0)]) request.multimodal_data["audio"] = {"input_features": audio} - request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalEncoderItemMetadata( + request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( [("video", 0), ("audio", 0), ("image", 0), ("audio", 1)], [4, 2, 2, 2], [4, 2, 2, 2] ) request.multimodal_data.pop("multimodal_embedding_lengths", None) diff --git a/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py b/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py index c906f44ee8b9..e7bca352b491 100644 --- a/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py +++ b/tests/unittest/_torch/modeling/test_nemotron_h_multimodal_preprocessing.py @@ -28,6 +28,8 @@ video_to_pixel_values, ) from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, MultimodalParams, MultimodalRuntimeData, _compute_mm_masks, @@ -2568,13 +2570,17 @@ def test_encoder_preserves_counts_per_video_and_request(self) -> None: MultimodalParams( multimodal_data={ "video": {"video_size": [(3, 1, 32, 32), (5, 1, 32, 32)]}, - "multimodal_embedding_lengths": [3, 4], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("video", 0), ("video", 1)], output_embedding_lengths=[3, 4] + ), } ), MultimodalParams( multimodal_data={ "image": {}, - "multimodal_embedding_lengths": [1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[1] + ), } ), ] diff --git a/tests/unittest/_torch/multimodal/test_encoder_group.py b/tests/unittest/_torch/multimodal/test_encoder_group.py index 79d5138ee14e..b04224b4bd81 100644 --- a/tests/unittest/_torch/multimodal/test_encoder_group.py +++ b/tests/unittest/_torch/multimodal/test_encoder_group.py @@ -5,7 +5,7 @@ Covers the pure helpers behind ``encode_multimodal_by_groups``: * ``_lengths_by_modality`` — invert prompt-ordered ``multimodal_embedding_lengths`` - into per-modality per-item lists using ``mm_item_order``. + into per-modality per-item lists using ``item_order``. * ``MultimodalParams.get_item_layout`` — item metadata with legacy fallbacks. * ``_reorder_embeds_by_manifest`` — slice per-modality tensors and concat in each request's prompt-order manifest. @@ -20,26 +20,37 @@ ) from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalEncoderItemMetadata, + MultimodalItemMetadata, MultimodalParams, ) pytestmark = pytest.mark.cpu_only -def _mp(*, mm_item_order=None, embedding_lengths=None, buckets=None): - """Request parameters with a legacy layout or typed item metadata.""" +def _mp(*, item_order=None, embedding_lengths=None, buckets=None): + """Build the same common item record supplied by input preparation.""" data = dict(buckets or {}) - if embedding_lengths is not None: - data["multimodal_embedding_lengths"] = embedding_lengths - return MultimodalParams(multimodal_data=data, mm_item_order=mm_item_order) + refs = ( + [(entry["modality"], entry["index"]) for entry in item_order] + if item_order + else [(next(iter(data)), i) for i in range(len(embedding_lengths or ()))] + ) + metadata = MultimodalItemMetadata( + refs, + output_embedding_lengths=( + () if not refs and embedding_lengths is None else embedding_lengths + ), + ) + metadata.validate() + data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata + return MultimodalParams(multimodal_data=data) class TestLengthsByModality: def test_mixed_manifest_image_video_image(self): # Manifest order: image#0 (4 rows), video#0 (6), image#1 (2). mp = _mp( - mm_item_order=[ + item_order=[ {"modality": "image", "index": 0}, {"modality": "video", "index": 0}, {"modality": "image", "index": 1}, @@ -53,7 +64,7 @@ def test_mixed_manifest_image_video_image(self): } def test_single_modality_no_manifest(self): - # No mm_item_order; flat list is already in per-modality order. + # No item_order; flat list is already in per-modality order. mp = _mp(embedding_lengths=[64, 64], buckets={"image": {}}) assert _lengths_by_modality([mp.get_item_layout()], ("image", "video")) == { "image": [64, 64], @@ -63,7 +74,7 @@ def test_single_modality_no_manifest(self): def test_length_mismatch_raises(self): with pytest.raises(ValueError, match="references and lengths must align"): _mp( - mm_item_order=[ + item_order=[ {"modality": "image", "index": 0}, {"modality": "video", "index": 0}, ], @@ -71,15 +82,9 @@ def test_length_mismatch_raises(self): buckets={"image": {}, "video": {}}, ) - def test_cross_group_multi_modality_without_manifest_raises(self): - # Multi-group models register one single-modality group per encoder. - # A raw request (no ``mm_item_order``) with items from two of those - # groups must be rejected — the group-local ``modalities`` tuple is - # length 1 and can never trip the guard on its own, so the check - # widens to the request-level modality set. - mp = _mp(embedding_lengths=[4, 3], buckets={"image": {}, "audio": {}}) - with pytest.raises(ValueError, match="mm_item_order"): - _lengths_by_modality([mp.get_item_layout()], ("image",)) + def test_missing_item_record_does_not_infer_cross_group_order(self): + mp = MultimodalParams(multimodal_data={"image": {}, "audio": {}}) + assert mp.get_item_layout() == ((), ()) def test_single_modality_outside_group_is_skipped(self): # Request has only audio; called on the image group. Must silently @@ -100,11 +105,11 @@ def test_layout_includes_all_request_modalities(self): assert mp.get_item_layout() == ((("audio", 0),), (10,)) def test_item_metadata_owns_order_and_lengths(self) -> None: - metadata = MultimodalEncoderItemMetadata( + metadata = MultimodalItemMetadata( [("video", 0), ("image", 0), ("audio", 0)], [24, 16, 8], [6, 4, 2] ) mp = _mp( - mm_item_order=[{"modality": "image", "index": 0}], + item_order=[{"modality": "image", "index": 0}], embedding_lengths=[99], buckets={"image": {}, "video": {}, "audio": {}}, ) @@ -130,7 +135,7 @@ def _marker_tensor(marker, rows, hidden=2): def test_mixed_image_video_image_in_one_request(self): # Prompt: image#0 (4 rows, marker 10), video#0 (6, 20), image#1 (2, 30). mp = _mp( - mm_item_order=[ + item_order=[ {"modality": "image", "index": 0}, {"modality": "video", "index": 0}, {"modality": "image", "index": 1}, @@ -145,7 +150,7 @@ def test_mixed_image_video_image_in_one_request(self): } per_modality_lengths = {"image": [4, 2], "video": [6]} out = _reorder_embeds_by_manifest( - [mp.get_item_layout()], per_modality_embeds, per_modality_lengths + [(mp.item_metadata.item_refs, ())], per_modality_embeds, per_modality_lengths ) # Expected column-0: 10 (×4), 20 (×6), 30 (×2) in prompt order. expected = torch.tensor([10.0] * 4 + [20.0] * 6 + [30.0] * 2) @@ -156,11 +161,11 @@ def test_cross_request_cursors_advance_per_modality(self): # Global per-modality indexing: image=[A#0, B#0], video=[B#0]. # Manifests use per-request indices — cursor must translate them. mp_a = _mp( - mm_item_order=[{"modality": "image", "index": 0}], + item_order=[{"modality": "image", "index": 0}], buckets={"image": {}}, ) mp_b = _mp( - mm_item_order=[ + item_order=[ {"modality": "video", "index": 0}, {"modality": "image", "index": 0}, ], @@ -172,7 +177,7 @@ def test_cross_request_cursors_advance_per_modality(self): } per_modality_lengths = {"image": [2, 1], "video": [3]} out = _reorder_embeds_by_manifest( - [mp_a.get_item_layout(), mp_b.get_item_layout()], + [(mp_a.item_metadata.item_refs, ()), (mp_b.item_metadata.item_refs, ())], per_modality_embeds, per_modality_lengths, ) @@ -192,7 +197,7 @@ def test_one_populated_modality_ignores_empty_group_peer(self): } per_modality_lengths = {"image": [2, 3], "video": []} out = _reorder_embeds_by_manifest( - [mp.get_item_layout()], per_modality_embeds, per_modality_lengths + [(mp.item_metadata.item_refs, ())], per_modality_embeds, per_modality_lengths ) expected = torch.tensor([7.0] * 2 + [9.0] * 3) assert torch.equal(out[:, 0], expected) @@ -200,7 +205,7 @@ def test_one_populated_modality_ignores_empty_group_peer(self): def test_single_modality_noncanonical_manifest_is_reordered(self): mp = _mp( - mm_item_order=[ + item_order=[ {"modality": "image", "index": 1}, {"modality": "image", "index": 0}, ], @@ -208,7 +213,7 @@ def test_single_modality_noncanonical_manifest_is_reordered(self): ) image_embeddings = torch.cat([self._marker_tensor(7, 2), self._marker_tensor(9, 3)], dim=0) out = _reorder_embeds_by_manifest( - [mp.get_item_layout()], {"image": image_embeddings}, {"image": [2, 3]} + [(mp.item_metadata.item_refs, ())], {"image": image_embeddings}, {"image": [2, 3]} ) expected = torch.tensor([9.0] * 3 + [7.0] * 2) @@ -222,7 +227,7 @@ def test_single_modality_batch_reorders_only_out_of_order_requests(self, b_item_ # order unless B lists its second image first. mp_a = _mp(embedding_lengths=[2], buckets={"image": {}}) mp_b = _mp( - mm_item_order=[{"modality": "image", "index": i} for i in b_item_order], + item_order=[{"modality": "image", "index": i} for i in b_item_order], buckets={"image": {}}, ) image_embeddings = torch.cat( @@ -230,7 +235,7 @@ def test_single_modality_batch_reorders_only_out_of_order_requests(self, b_item_ dim=0, ) out = _reorder_embeds_by_manifest( - [mp_a.get_item_layout(), mp_b.get_item_layout()], + [(mp_a.item_metadata.item_refs, ()), (mp_b.item_metadata.item_refs, ())], {"image": image_embeddings}, {"image": [2, 3, 1]}, ) @@ -249,7 +254,7 @@ def test_empty_bookkeeping_returns_typed_empty(self): per_modality_embeds = {"image": torch.empty((0, hidden), dtype=torch.float16)} per_modality_lengths = {"image": []} out = _reorder_embeds_by_manifest( - [mp.get_item_layout()], per_modality_embeds, per_modality_lengths + [(mp.item_metadata.item_refs, ())], per_modality_embeds, per_modality_lengths ) assert out.shape == (0, hidden) assert out.dtype == torch.float16 @@ -258,5 +263,5 @@ def test_no_embeds_returns_empty(self): # Defensive: no group produced embeddings at all — return an empty # tensor rather than crashing. mp = _mp(buckets={}) - out = _reorder_embeds_by_manifest([mp.get_item_layout()], {}, {}) + out = _reorder_embeds_by_manifest([(mp.item_metadata.item_refs, ())], {}, {}) assert out.numel() == 0 diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py index b0587c2be6b6..987dee4ba9b4 100644 --- a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py +++ b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py @@ -34,7 +34,12 @@ ) from tensorrt_llm._torch.models.modeling_multimodal_utils import get_multimodal_embeddings from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig -from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, + MultimodalParams, + MultimodalRuntimeData, +) from tensorrt_llm.llmapi.llm_args import MultimodalConfig requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -46,12 +51,12 @@ def __init__(self, hidden_size: int, tokens_per_image: int): self._tokens_per_image = tokens_per_image self.encoder_call_count = 0 self.encoder_call_stream_id = None - self.encoder_mm_item_order = None + self.encoder_item_metadata = None def encode_multimodal_inputs(self, multimodal_params, **kwargs) -> torch.Tensor: self.encoder_call_count += 1 self.encoder_call_stream_id = torch.cuda.current_stream().cuda_stream - self.encoder_mm_item_order = multimodal_params[0].mm_item_order + self.encoder_item_metadata = multimodal_params[0].item_metadata pv = multimodal_params[0].multimodal_data["image"]["pixel_values"] assert pv.device.type == "cuda", ( "pixel_values should be on CUDA after to_device in the helper." @@ -90,15 +95,7 @@ def encode_multimodal_inputs(self, multimodal_params, **kwargs) -> torch.Tensor: self.encoder_call_count += 1 self.last_encoder_batch_size = len(multimodal_params) self.encoder_call_stream_id = torch.cuda.current_stream().cuda_stream - total_rows = sum( - sum( - param.multimodal_data.get( - "multimodal_embedding_lengths", - [self._tokens_per_image], - ) - ) - for param in multimodal_params - ) + total_rows = sum(sum(param.get_item_layout()[1]) for param in multimodal_params) return torch.full( (total_rows, self._hidden_size), float(self.encoder_call_count), @@ -120,6 +117,9 @@ def _make_request(request_id: int, num_tokens: int) -> LlmRequest: "pixel_values": pixel_values, }, "multimodal_embed_mask_cumsum": cumsum, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[num_tokens] + ), }, ) @@ -141,7 +141,9 @@ def _make_cacheable_request( py_multimodal_data={ "image": {"pixel_values": torch.randn(3, 32, 32)}, "multimodal_embed_mask_cumsum": torch.arange(1, num_tokens + 1, dtype=torch.int64), - "multimodal_embedding_lengths": [num_tokens], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[num_tokens] + ), "mm_processor_kwargs_hash": "kwargs-a", }, multimodal_hashes=[item_hash], @@ -169,7 +171,10 @@ def _make_two_item_cacheable_request( "image_sizes": [[32, 32]] * len(item_hashes), }, "multimodal_embed_mask_cumsum": torch.arange(1, num_tokens + 1, dtype=torch.int64), - "multimodal_embedding_lengths": [tokens_per_item] * len(item_hashes), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", i) for i in range(len([tokens_per_item] * len(item_hashes)))], + output_embedding_lengths=[tokens_per_item] * len(item_hashes), + ), "mm_processor_kwargs_hash": "kwargs-a", }, multimodal_hashes=item_hashes, @@ -337,18 +342,14 @@ def test_cross_iter_prefetch_materializes_on_side_stream(monkeypatch): model = _StubModel(hidden_size=8, tokens_per_image=4) req = _make_request(request_id=0, num_tokens=4) - expected_item_order = [ - {"modality": "image", "index": 0}, - {"modality": "video", "index": 0}, - ] - req.py_mm_item_order = expected_item_order + expected_item_metadata = req.py_multimodal_params.item_metadata n = maybe_prefetch_mm_encoder_for_next_iter(model, [req], max_prefetch_ahead=max_prefetch_ahead) assert n == 1 assert model.encoder_call_count == 1 assert model.encoder_call_stream_id == aux_stream.cuda_stream - assert model.encoder_mm_item_order == expected_item_order + assert model.encoder_item_metadata is expected_item_metadata embedding = req.py_multimodal_data.get("multimodal_embedding") assert isinstance(embedding, torch.Tensor) assert embedding.shape == (4, 8) diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index 7447cb974b13..80668cb785e7 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -26,7 +26,13 @@ _assemble_multimodal_encoder_embeddings, ) from tensorrt_llm._torch.modules.embedding import Embedding -from tensorrt_llm.inputs.multimodal import MultimodalInput, MultimodalParams, MultimodalRuntimeData +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalInput, + MultimodalItemMetadata, + MultimodalParams, + MultimodalRuntimeData, +) from tensorrt_llm.llmapi.llm_args import MultimodalConfig @@ -164,7 +170,10 @@ def make_keyed_multimodal_param( ), "image_sizes": [[2, 2]] * n_items, }, - "multimodal_embedding_lengths": embedding_lengths, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", i) for i in range(len(embedding_lengths))], + output_embedding_lengths=embedding_lengths, + ), "mm_processor_kwargs_hash": kwargs_hash, } if local_embedding is not None: @@ -274,7 +283,7 @@ def test_declared_qwen_fields_reuse_host_costs_for_device_grids(): from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLModelBase from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalEncoderItemMetadata, + MultimodalItemMetadata, ) class QwenModel(DummyMultimodalModel): @@ -285,7 +294,7 @@ class QwenModel(DummyMultimodalModel): param = MultimodalParams( multimodal_data={ "image": {"pixel_values": pixels, "image_grid_thw": grids}, - MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( [("image", 1), ("image", 0)], [3, 2], [3, 2] ), } @@ -710,7 +719,9 @@ def test_disabled_encoder_cache_preserves_current_behavior(): MultimodalParams( multimodal_data={ "image": {"pixel_values": torch.empty(1)}, - "multimodal_embedding_lengths": [2], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=[2] + ), "mm_processor_kwargs_hash": "kwargs-a", }, multimodal_runtime=make_runtime(2), @@ -829,7 +840,9 @@ def test_build_multimodal_encoder_input_slices_packed_grid_thw(): "per_item_list": ["a", "b", "c"], "per_request_scalar": torch.tensor(42.0), }, - "multimodal_embedding_lengths": [2, 3, 1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0), ("image", 1), ("image", 2)], output_embedding_lengths=[2, 3, 1] + ), "mm_processor_kwargs_hash": "kw", }, ) @@ -866,7 +879,9 @@ def test_build_multimodal_encoder_input_stacked_crops_padding_to_miss_max_size() "pixel_values": pixels, "image_sizes": [[3, 4], [5, 5]], }, - "multimodal_embedding_lengths": [1, 1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0), ("image", 1)], output_embedding_lengths=[1, 1] + ), "mm_processor_kwargs_hash": "kw", }, ) @@ -899,7 +914,9 @@ def test_build_multimodal_encoder_input_slices_audio_input_features(): "input_features": features, "input_features_mask": mask, }, - "multimodal_embedding_lengths": [1, 1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("audio", 0), ("audio", 1)], output_embedding_lengths=[1, 1] + ), "mm_processor_kwargs_hash": "kw", }, ) diff --git a/tests/unittest/_torch/multimodal/test_nemotron_h_multimodal_encoder_groups.py b/tests/unittest/_torch/multimodal/test_nemotron_h_multimodal_encoder_groups.py index 362bd8626f80..c43d9b132b74 100644 --- a/tests/unittest/_torch/multimodal/test_nemotron_h_multimodal_encoder_groups.py +++ b/tests/unittest/_torch/multimodal/test_nemotron_h_multimodal_encoder_groups.py @@ -45,7 +45,7 @@ def _make_multimodal_params( data["video"] = video if audio is not None: data["audio"] = audio - return MultimodalParams(multimodal_data=data, mm_item_order=None) + return MultimodalParams(multimodal_data=data) def _marker_tensor(marker: int, n: int, dim: int = 4) -> torch.Tensor: diff --git a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py index c3c8a95ef71a..15f9c543fd86 100644 --- a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py +++ b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py @@ -12,7 +12,6 @@ from typing import Optional -import pytest import torch from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( @@ -20,7 +19,11 @@ encode_multimodal_by_groups, ) from tensorrt_llm._torch.models.modeling_qwen3vl import _qwen3vl_build_batched_input -from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, + MultimodalParams, +) def _encode_batched(pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: @@ -50,8 +53,15 @@ def _param( if video is not None: data["video"] = video if lengths is not None: - data["multimodal_embedding_lengths"] = lengths - return MultimodalParams(multimodal_data=data, mm_item_order=order) + refs = ( + [(entry["modality"], entry["index"]) for entry in order] + if order + else [(next(iter(data)), i) for i in range(len(lengths))] + ) + data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( + refs, output_embedding_lengths=lengths + ) + return MultimodalParams(multimodal_data=data) def _patches(marker: int, n: int, dim: int = 2) -> torch.Tensor: @@ -105,28 +115,6 @@ def test_mixed_image_video_image_prompt_order() -> None: assert torch.equal(out[:, 0], expected) -def test_mixed_without_manifest_raises() -> None: - """Mixed request must carry mm_item_order; the shared helper enforces this - once for every model that registers a multi-modality EncoderGroup.""" - with pytest.raises(ValueError, match="mm_item_order"): - _encode( - [ - _param( - image={ - "pixel_values": _patches(1, 1), - "image_grid_thw": torch.tensor([[1, 1, 1]]), - }, - video={ - "pixel_values_videos": _patches(2, 1), - "video_grid_thw": torch.tensor([[1, 1, 1]]), - }, - lengths=[1, 1], - # order deliberately omitted - ) - ] - ) - - def test_batch_of_two_image_only_requests() -> None: out = _encode( [ @@ -230,40 +218,3 @@ def test_raw_inputs_route_through_encode_multimodal_by_groups() -> None: out = encode_multimodal_by_groups(_QWEN3VL_ENCODER_GROUPS, params) expected = torch.tensor([22.0] * 5 + [11.0] * 3) assert torch.equal(out[:, 0], expected) - - -def test_mixed_request_alongside_ordered_one_still_rejects_unordered() -> None: - """Even when a sibling request carries a valid manifest, a mixed request - without its own manifest must still raise — validation is per-request.""" - with pytest.raises(ValueError, match="mm_item_order"): - _encode( - [ - _param( - image={ - "pixel_values": _patches(10, 2), - "image_grid_thw": torch.tensor([[1, 1, 2]]), - }, - video={ - "pixel_values_videos": _patches(20, 2), - "video_grid_thw": torch.tensor([[1, 1, 2]]), - }, - order=[ - {"modality": "image", "index": 0}, - {"modality": "video", "index": 0}, - ], - lengths=[2, 2], - ), - _param( - image={ - "pixel_values": _patches(30, 1), - "image_grid_thw": torch.tensor([[1, 1, 1]]), - }, - video={ - "pixel_values_videos": _patches(40, 1), - "video_grid_thw": torch.tensor([[1, 1, 1]]), - }, - lengths=[1, 1], - # No manifest on this request even though it has both modalities. - ), - ] - ) diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 6fc17746417b..febb2773f3d5 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -13,7 +13,6 @@ from tensorrt_llm.inputs.media_io import MEDIA_IO_REGISTRY, MediaModality from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalEncoderItemMetadata, MultimodalInput, MultimodalItemMetadata, MultimodalItemPartition, @@ -150,14 +149,17 @@ def __call__(self, inputs, sampling_params): return [101, 102], { "multimodal_data": { "image": {}, - "multimodal_embedding_lengths": self.existing_embedding_lengths, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", i) for i in range(len(self.existing_embedding_lengths))], + output_embedding_lengths=self.existing_embedding_lengths, + ), } } - def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): + def get_mm_item_metadata(self, prompt_token_ids, multimodal_data, mm_item_order=None): assert prompt_token_ids == [101, 102] assert "image" in multimodal_data - return MultimodalEncoderItemMetadata( + return MultimodalItemMetadata( item_refs=[("image", 0)], encoder_token_lengths=[4], output_embedding_lengths=[2], @@ -208,7 +210,7 @@ def test_mm_item_metadata_none_falls_back_to_non_item_path(): # Generic item order does not imply that physical encoder costs or # item scheduling are supported. class NoMetadataProcessor(_ItemMetadataFakeProcessor): - def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): + def get_mm_item_metadata(self, prompt_token_ids, multimodal_data, mm_item_order=None): return None input_processor = create_input_processor_with_hash(NoMetadataProcessor([2])) @@ -224,7 +226,7 @@ def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): def test_mm_item_metadata_rejects_wrong_return_type(): # A non-None, non-metadata return is a genuine contract violation. class BadMetadataProcessor(_ItemMetadataFakeProcessor): - def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): + def get_mm_item_metadata(self, prompt_token_ids, multimodal_data, mm_item_order=None): return ("image", 0) input_processor = create_input_processor_with_hash(BadMetadataProcessor([2])) @@ -233,17 +235,6 @@ def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): input_processor({"prompt": "unused"}, sampling_params=None) -def test_mm_item_metadata_retains_three_field_tuple_contract() -> None: - metadata = MultimodalEncoderItemMetadata([("image", 0)], [4], [2]) - - item_refs, encoder_lengths, output_lengths = metadata - - assert metadata._fields == ("item_refs", "encoder_token_lengths", "output_embedding_lengths") - assert item_refs == [("image", 0)] - assert encoder_lengths == [4] - assert output_lengths == [2] - - @pytest.mark.parametrize( "selected, expected_refs, expected_costs, expected_lengths", [ @@ -255,7 +246,7 @@ def test_mm_item_metadata_retains_three_field_tuple_contract() -> None: def test_mm_item_metadata_selection_rebases_refs_without_mutating_source( selected, expected_refs, expected_costs, expected_lengths ) -> None: - metadata = MultimodalEncoderItemMetadata( + metadata = MultimodalItemMetadata( [("image", 1), ("video", 0), ("image", 0), ("audio", 0)], [20, 50, 30, 40], [2, 5, 3, 4], @@ -263,25 +254,11 @@ def test_mm_item_metadata_selection_rebases_refs_without_mutating_source( subset = metadata.select(selected) - assert subset == MultimodalEncoderItemMetadata(expected_refs, expected_costs, expected_lengths) - assert metadata.item_refs == [("image", 1), ("video", 0), ("image", 0), ("audio", 0)] - assert metadata.encoder_token_lengths == [20, 50, 30, 40] - assert metadata.output_embedding_lengths == [2, 5, 3, 4] - assert all( - selected_field is not source_field for selected_field, source_field in zip(subset, metadata) - ) - - -def test_mm_item_metadata_legacy_import_and_pickle() -> None: - import pickle - - from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata as LegacyItemMetadata - - assert LegacyItemMetadata is MultimodalEncoderItemMetadata - # Existing serialized records resolve the original registry module path. - legacy_class = pickle.loads(b"ctensorrt_llm.inputs.registry\nMultimodalEncoderItemMetadata\n.") - metadata = legacy_class([("audio", 0)], [8], [2]) - assert pickle.loads(pickle.dumps(metadata)) == metadata + assert subset == MultimodalItemMetadata(expected_refs, expected_costs, expected_lengths) + assert metadata.item_refs == (("image", 1), ("video", 0), ("image", 0), ("audio", 0)) + assert metadata.encoder_token_lengths == (20, 50, 30, 40) + assert metadata.output_embedding_lengths == (2, 5, 3, 4) + assert subset is not metadata @pytest.mark.parametrize("indices", [[3, 0, 2], [2, 0], [-1, 0], [], [1, 1]]) @@ -327,6 +304,48 @@ def test_unknown_item_costs_and_rows_remain_explicit() -> None: params.get_item_layout() +@pytest.mark.parametrize("modality", MEDIA_MODALITIES) +def test_prepared_metadata_fills_unknown_model_rows(modality) -> None: + partitions = {modality: {"features": MultimodalItemPartition()}} + original = MultimodalItemMetadata([(modality, 0)], input_partitions=partitions) + + class Processor(_ItemMetadataFakeProcessor): + def __call__(self, inputs, sampling_params): + return [101, 102], { + "multimodal_data": { + modality: {}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [(modality, 0)], output_embedding_lengths=[2] + ), + } + } + + def get_mm_item_metadata(self, prompt_token_ids, multimodal_data, mm_item_order=None): + return original + + _, extra = create_input_processor_with_hash(Processor(None))({"prompt": "unused"}, None) + params = MultimodalParams(multimodal_data=extra["multimodal_data"]) + assert params.get_item_layout() == (((modality, 0),), (2,)) + assert params.item_metadata.encoder_token_lengths is None + assert params.item_metadata.input_partitions is partitions + assert original.output_embedding_lengths is None + assert "multimodal_embedding_lengths" not in params.multimodal_data + assert ( + MultimodalParams(multimodal_data=params.multimodal_data).item_metadata + is params.item_metadata + ) + + +@pytest.mark.parametrize("lengths", [[2, 3], [-1], ["2"]]) +def test_prepared_metadata_validates_model_rows(lengths) -> None: + class Processor(_ItemMetadataFakeProcessor): + def get_mm_item_metadata(self, prompt_token_ids, multimodal_data, mm_item_order=None): + return MultimodalItemMetadata([("image", 0)], output_embedding_lengths=lengths) + + with pytest.raises((TypeError, ValueError)): + create_input_processor_with_hash(Processor([2]))({"prompt": "unused"}, None) + + def test_item_metadata_roundtrips_with_shared_tensor_handles() -> None: import pickle @@ -352,26 +371,25 @@ def test_item_metadata_roundtrips_with_shared_tensor_handles() -> None: @pytest.mark.parametrize("precomputed", [False, True]) -def test_generic_metadata_normalizes_mixed_order_without_encoder_costs(precomputed) -> None: +def test_generic_metadata_keeps_mixed_order_without_encoder_costs(precomputed) -> None: from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor class Processor(_ItemMetadataFakeProcessor): get_mm_item_metadata = BaseMultimodalInputProcessor.get_mm_item_metadata def __call__(self, inputs, sampling_params): + refs = [(entry["modality"], entry["index"]) for entry in inputs["mm_item_order"]] return [101, 102], { "multimodal_data": { "image": {}, "audio": {}, - "multimodal_embedding_lengths": [1, 1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + refs, output_embedding_lengths=[1, 1] + ), **({"multimodal_embedding": torch.zeros(2, 4)} if precomputed else {}), } } - def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): - assert not precomputed - return None - order = [{"modality": "audio", "index": 0}, {"modality": "image", "index": 0}] _, extra = create_input_processor_with_hash(Processor(None))( {"prompt": "unused", "mm_item_order": order}, None @@ -422,7 +440,7 @@ def test_existing_mm_cumsum_preserves_storage_and_backfills_block_alignment( assert data["mm_bidirectional_blocks"] is True -def test_tokenized_multimodal_overwrites_stale_embedding_lengths(): +def test_tokenized_multimodal_builds_metadata_from_final_embedding_slots(): """Final expanded prompt layout wins over dummy-placeholder metadata.""" class FakeProcessor: @@ -445,7 +463,9 @@ def __call__(self, inputs, sampling_params): return expanded_ids, { "multimodal_data": { # Stale dummy-placeholder value; recomputed to [3] downstream. - "multimodal_embedding_lengths": [999], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("image", 0)], output_embedding_lengths=None + ), } } @@ -478,7 +498,9 @@ def get_mm_special_token_ids(self): ) assert prompt_token_ids == [10, 101, 102, 200, 103, 20] - assert extra["multimodal_data"]["multimodal_embedding_lengths"] == [3] + assert extra["multimodal_data"][ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY + ].output_embedding_lengths == (3,) class _KwargsHashFakeProcessor: @@ -576,17 +598,19 @@ def __call__(self, inputs: dict, sampling_params: object) -> tuple[list[int], di "multimodal_data": { "image": {}, "video": {}, - "multimodal_embedding_lengths": [2, 2], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("video", 0), ("image", 0)], output_embedding_lengths=[2, 2] + ), } } def get_num_tokens_per_video(self, *, video: list[np.ndarray], **kwargs: object) -> int: return 2 - def get_mm_encoder_item_metadata( - self, prompt_token_ids: list[int], multimodal_data: dict - ) -> MultimodalEncoderItemMetadata: - return MultimodalEncoderItemMetadata([("video", 0), ("image", 0)], [4, 4], [2, 2]) + def get_mm_item_metadata( + self, prompt_token_ids: list[int], multimodal_data: dict, mm_item_order=None + ) -> MultimodalItemMetadata: + return MultimodalItemMetadata([("video", 0), ("image", 0)], [4, 4], [2, 2]) processor = Processor() inputs = { @@ -606,6 +630,7 @@ def get_mm_encoder_item_metadata( @pytest.mark.parametrize("hashing_enabled", [False, True], ids=["fallback", "hashing"]) @pytest.mark.parametrize("use_uuids", [False, True], ids=["content", "uuid"]) +@pytest.mark.parametrize("frontend_order", [False, True], ids=["record-order", "frontend-order"]) @pytest.mark.parametrize( "modalities,token_ids,explicit_ids,special_ids,expected_cumsum,expected_runs", [ @@ -674,6 +699,7 @@ def test_multimodal_input_contract_snapshot( expected_runs: tuple[list[int], list[int], list[int]], hashing_enabled: bool, use_uuids: bool, + frontend_order: bool, ) -> None: """Pin framework metadata independently of model-specific media encoding.""" item_refs = [(modality, 0) for modality in modalities] @@ -687,7 +713,9 @@ def __call__(self, inputs: dict, sampling_params: object) -> tuple[list[int], di return list(token_ids), { "multimodal_data": { **{modality: {} for modality in modalities}, - "multimodal_embedding_lengths": [2] * num_items, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + item_refs, output_embedding_lengths=[2] * num_items + ), } } @@ -706,10 +734,10 @@ def get_mm_token_ids(self) -> torch.Tensor | None: def get_mm_special_token_ids(self) -> torch.Tensor | None: return torch.tensor(special_ids) if special_ids is not None else None - def get_mm_encoder_item_metadata( - self, prompt_token_ids: list[int], multimodal_data: dict - ) -> MultimodalEncoderItemMetadata: - return MultimodalEncoderItemMetadata(item_refs, [4] * num_items, [2] * num_items) + def get_mm_item_metadata( + self, prompt_token_ids: list[int], multimodal_data: dict, mm_item_order=None + ) -> MultimodalItemMetadata: + return MultimodalItemMetadata(item_refs, [4] * num_items, [2] * num_items) media = { "image": [np.array([1], dtype=np.uint8)], @@ -724,6 +752,8 @@ def get_mm_encoder_item_metadata( for modality in modalities ], } + if not frontend_order: + inputs.pop("mm_item_order") if use_uuids: inputs["multi_modal_uuids"] = {modality: [f"{modality}-id"] for modality in modalities} @@ -1084,18 +1114,15 @@ def test_normalized_params_keep_one_order_and_output_length_record(): metadata = MultimodalItemMetadata(refs, costs, rows) data = { MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, - "multimodal_embedding_lengths": list(rows), } params = MultimodalParams( multimodal_data=data, - mm_item_order=[{"modality": m, "index": i} for m, i in refs], ) refs.reverse() costs[0] = 100 rows[0] = 100 assert params.item_metadata is metadata - assert params.mm_item_order is None assert "multimodal_embedding_lengths" not in data assert metadata.item_refs == (("audio", 0), ("image", 0)) assert metadata.encoder_token_lengths == (4, 8) @@ -1106,8 +1133,7 @@ def test_normalized_params_keep_one_order_and_output_length_record(): metadata.encoder_token_lengths = (100, 8) -def test_unknown_item_layout_preserves_legacy_order_at_input_boundary(): - order = [{"modality": "audio", "index": 0}] - params = MultimodalParams(multimodal_data={"audio": {}}, mm_item_order=order) +def test_params_do_not_infer_item_layout_from_raw_buckets(): + params = MultimodalParams(multimodal_data={"audio": {}, "image": {}}) assert params.item_metadata is None - assert params.mm_item_order is order + assert params.get_item_layout() == ((), ()) From 63f2002a79d5f3c771ae2ac31b06b37cdd537786 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 01:04:57 -0700 Subject: [PATCH 12/22] Remove redundant multimodal metadata checks and stale docs Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/llm_request.py | 11 +++-------- tensorrt_llm/inputs/multimodal.py | 2 +- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 106214f52428..d881d47eea41 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1260,9 +1260,8 @@ def get_multimodal_encoder_token_lengths( request: LlmRequest) -> Optional[Sequence[int]]: """Return per-item physical encoder attention-token costs. - Field-level validation happens inside - `MultimodalItemMetadata.validate`; only the request-level - cross-check against `output_embedding_lengths` lives here. + Validate item metadata and check output rows against the request's + prompt spans. """ item_metadata = get_multimodal_item_metadata(request.py_multimodal_data) if item_metadata is None: @@ -1270,13 +1269,9 @@ def get_multimodal_encoder_token_lengths( item_metadata.validate() if item_metadata.encoder_token_lengths is None: return None - embedding_lengths = get_multimodal_embedding_lengths(request) - if embedding_lengths is None: + if get_multimodal_embedding_lengths(request) is None: raise ValueError("Multimodal encoder item scheduling requires " "output_embedding_lengths") - if item_metadata.output_embedding_lengths != tuple(embedding_lengths): - raise ValueError("Multimodal encoder item output lengths must match " - "output_embedding_lengths") return item_metadata.encoder_token_lengths diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index f0acda33a9bc..a9d2d2970ea3 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -794,7 +794,7 @@ class MultimodalParams: """Processed multimodal inputs and views used by model execution. `item_metadata` reads the request's shared item record from processed data; - `get_item_layout` resolves its prompt order, including legacy formats. + `get_item_layout` reads prompt order and output rows from that record. Runtime data describes the current forward or chunk rather than encoder readiness and cache ownership, which belong to the request Params. From c9b036632f6caa126b18fb0db76e88638a83e401 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 01:31:25 -0700 Subject: [PATCH 13/22] Consolidate multimodal tests and document developer contracts Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- AGENTS.md | 26 +-- docs/source/developer-guide/multimodal.md | 109 ++++++++++ docs/source/developer-guide/overview.md | 7 + docs/source/index.rst | 1 + docs/source/torch/arch_overview.md | 97 +-------- .../test_multimodal_encoder_item_groups.py | 104 +++++----- tests/unittest/inputs/test_multimodal.py | 187 ++++++++++-------- 7 files changed, 266 insertions(+), 265 deletions(-) create mode 100644 docs/source/developer-guide/multimodal.md diff --git a/AGENTS.md b/AGENTS.md index 41145c6ed6e8..95c1c435b255 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -95,29 +95,9 @@ HuggingFace Model → LLM API → PyTorch Executor ### Multimodal request data -Read the [multimodal data contracts](docs/source/torch/arch_overview.md#multimodal-request-data) -before changing multimodal preprocessing, item selection, or request ownership. -- Reuse `MultimodalParams` in `inputs/multimodal.py`; the executor's request aggregate is - `LlmRequest.py_multimodal_params`. Legacy request properties alias this aggregate. -- `item_metadata` reads one record in `multimodal_data`. Its `input_partitions` describe - how processed fields are selected by item; model hooks retain batching and padding policy. - Keep prompt token counts, encoder costs, and output embedding rows distinct. -- Item refs, costs, and output rows are immutable tuples. Encoder state references the - admitted record instead of copying its lengths. Input preparation produces this record; - request and forward Params do not reconstruct it from separate order or length fields. -- Only request Params own `encoder_state` cache references. Encoder selection and LLM chunk - Params are separate views; cleanup remains in the executor and the cache remains model-owned. -- Persistent encoder caching reuses one lazily reconstructed `MultimodalInput` per request. - Copy that wrapper before slicing its hashes; prefill/termination cleanup clears it. -- E/P processor hooks return `(prompt_token_ids, MultimodalParams)`; use - `MultimodalParams.from_embeddings` for embedding handoffs. `DisaggregatedParams` remains - the transport envelope. Do not serialize live encoder state or CUDA events in a new wire format. -- When changing an E/P contract, check external callers such as Dynamo separately from - local processor hooks: wire compatibility and Python API compatibility are distinct. -- Validate producer/admission contracts at their boundaries, without repeated full-metadata - validation in selection or forward. Processors use `get_mm_item_metadata`; do not add a - second metadata format or old-hook dispatch. Unknown encoder costs remain `None`, and - common item metadata does not imply item scheduling support. +Read the [multimodal developer guide](docs/source/developer-guide/multimodal.md) before changing multimodal paths. +Reuse `MultimodalParams`; request Params own cache references. +Encoder and LLM forward Params are execution views. ## Key Files diff --git a/docs/source/developer-guide/multimodal.md b/docs/source/developer-guide/multimodal.md new file mode 100644 index 000000000000..671ac374bb35 --- /dev/null +++ b/docs/source/developer-guide/multimodal.md @@ -0,0 +1,109 @@ + + +# Multimodal Request Data + +Multimodal preprocessing and execution share `MultimodalParams`, defined in +[inputs/multimodal.py](source:tensorrt_llm/inputs/multimodal.py). + +| Type or member | Responsibility | +|---|---| +| `MultimodalParams.multimodal_data` | Processed model inputs, embeddings or shared-tensor handles, and model-specific auxiliary data. | +| `MultimodalParams.item_metadata` | Reads the single `MultimodalItemMetadata` record in `multimodal_data`. | +| `MultimodalItemMetadata` | Prompt-ordered item references, optional encoder costs and output row counts, and input partitions. | +| `MultimodalItemPartition` | Describes one field's item axis and optional packed lengths, inside `item_metadata.input_partitions`. | +| `MultimodalInput` | Hashes, prompt positions, and exact runs converted into existing C++ request fields. | +| `MultimodalRuntimeData` | Derived information for the current LLM forward or prefill chunk. | +| `MultimodalEncoderRequestState` | Request-owned readiness and cache references; reads lengths from the admitted item metadata. | + +## Process Boundaries and Ownership + +```text +Main / input processor + -> prompt token IDs + MultimodalParams (processed data and item metadata) +Worker handoff + -> restore tensors; convert MultimodalInput to C++ request fields +Executor / LlmRequest.py_multimodal_params + -> admission creates encoder_state where item scheduling is supported + -> encoder selection creates Params views with selected data and metadata + -> LLM forward creates Params views with embeddings and chunk runtime + -> executor cleanup releases request-owned cache references +``` + +Only the request Params own `encoder_state` cache references. Execution views +leave that state unset, and request Params leave `multimodal_runtime` unset. +The model owns the `TensorLRUCache`; the executor owns request cleanup. +Text-only requests do not need a Params aggregate. Live cache references and +CUDA events stay in the executor rather than the cross-process payload. + +The worker clears the processor-side `multimodal_input` after converting it to +C++ fields. Persistent encoder caching lazily reconstructs one Python +`MultimodalInput` per request. Selected inputs copy this wrapper before slicing +its hashes, and prefill/termination cleanup clears it. + +Inline encoding and prefetch write encoder outputs back to the request. +Item-scheduled LLM chunks keep their assembled embeddings local to the forward. +Execution paths retain their CUDA event synchronization responsibilities. + +## Item Metadata and Selection + +`item_refs` identifies processed items by modality and modality-local index, in +prompt order. A processed item is not necessarily one original media object or +one tensor. Prompt token counts, physical encoder token costs, and output +embedding rows are distinct quantities. + +Input preparation publishes the common record directly. The frontend's +`mm_item_order` is consumed during preprocessing; execution reads only +`item_metadata.item_refs`. Output row counts live in +`item_metadata.output_embedding_lengths`. Admission, full-request and selected-item +encoding, cache keys, and encoder output export all read the same record. +The encoder-result boundary exports lengths as lists. + +Item refs, costs, and output rows are immutable tuples; the constructor also +accepts lists. Encoder state references this record without copying its lengths. +Producers fill input partitions before admission and validate contracts at +producer/admission boundaries, avoiding repeated validation per forward. + +Partitions describe how to select processed fields: shared fields have no item +axis, batched fields have an item axis, and packed fields also provide host-side +per-item lengths. Contiguous tensor selections can share storage; reordered +selections may copy. Treat shared input storage as read-only. Model hooks and +`EncoderGroup` still own compatible batching, padding, masks, encoder arguments, +and output fusion. + +## Processor Hooks + +Processors implement +`get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order=None)` and +return `MultimodalItemMetadata`. Generic preparation supplies prompt order and +known output rows; model processors add encoder costs and input partitions. +Partial model records retain known prompt rows. Conflicting order or rows fail +at input preparation. Gemma4 and Nemotron add partitions; Qwen and Mistral also +provide physical encoder costs. + +Unknown costs or row counts remain `None`. Common metadata does not enable item +scheduling: models can supply order and output rows with +`encoder_token_lengths=None` and keep full-request encoding. Model capability +and configuration decide whether items are scheduled separately. + +Custom processors must use this common hook and type. The former +`get_mm_encoder_item_metadata` hook, `MultimodalEncoderItemMetadata` tuple, and +Params order/length fallbacks have been removed. + +## Encoder/Prefill Handoff + +For encoder/prefill disaggregation (E/P), model processor hooks return +`(prompt_token_ids, MultimodalParams)`. Use `MultimodalParams.from_embeddings` +to attach embeddings or shared-tensor handles. See +[inputs/registry.py](source:tensorrt_llm/inputs/registry.py) for the hook contract. + +`DisaggregatedParams` remains the transport envelope used by external +orchestrators such as Dynamo. The prefill processor still rebuilds model-specific +prompt layout; its local Params result does not replace the wire protocol. +Custom E/P hooks must follow the tuple return contract. Review Python hook +compatibility separately from transport fields when changing this interface. + +These contracts do not change the supported model/modality matrix or impose a +common batching policy on every model. diff --git a/docs/source/developer-guide/overview.md b/docs/source/developer-guide/overview.md index 8cdd2e122bc3..fb19732fdaf1 100644 --- a/docs/source/developer-guide/overview.md +++ b/docs/source/developer-guide/overview.md @@ -43,6 +43,13 @@ During each iteration of its background loop, the `PyExecutor` performs the foll - Output Handling: Updates the partial outputs for ongoing requests and finalizes the results for any requests that have reached completion, returning them to the user. +## Multimodal Requests + +Multimodal input preparation and execution share `MultimodalParams`. The executor +owns request cache references; encoder and LLM forwards use separate execution views. +See the [multimodal developer guide](multimodal.md) for the data contracts, +process boundaries, and encoder/prefill handoff. + ## Runtime Optimizations TensorRT LLM enhances inference throughput and reduces latency by integrating a suite of runtime optimizations, including CUDA Graph, [Overlap Scheduler](../features/overlap-scheduler.md), [Speculative decoding](../features/speculative-decoding.md), etc. diff --git a/docs/source/index.rst b/docs/source/index.rst index e04a5dbdf0bc..9920583eb754 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -101,6 +101,7 @@ Welcome to TensorRT LLM's Documentation! :caption: Developer Guide developer-guide/overview.md + developer-guide/multimodal.md developer-guide/perf-analysis.md developer-guide/perf-benchmarking.md developer-guide/moe-perf-guide.md diff --git a/docs/source/torch/arch_overview.md b/docs/source/torch/arch_overview.md index c3c281e62d72..b8d737d6a2a1 100644 --- a/docs/source/torch/arch_overview.md +++ b/docs/source/torch/arch_overview.md @@ -32,101 +32,8 @@ The single-step flow of PyExecutor involves: ## Multimodal Request Data -Multimodal preprocessing and execution use the existing `MultimodalParams` in -[inputs/multimodal.py](../../../tensorrt_llm/inputs/multimodal.py). -The executor keeps request-wide data in `LlmRequest.py_multimodal_params`, while -encoder selection and LLM forwards create separate Params views. - -| Type or member | Responsibility | -|---|---| -| `MultimodalParams.multimodal_data` | Processed model inputs, precomputed embeddings or shared-tensor handles, and model-specific auxiliary data. | -| `MultimodalParams.item_metadata` | A property reading the single `MultimodalItemMetadata` record stored in `multimodal_data`. | -| `MultimodalItemMetadata` | Prompt-ordered references to processed items, optional encoder costs and output row counts, and input partitions. | -| `MultimodalItemPartition` | A small value describing one input field's item axis and optional packed lengths. It lives inside `item_metadata.input_partitions`. | -| `MultimodalInput` | Processor-side hash, prompt-position, and exact-run data converted into the existing C++ request fields by the worker. | -| `MultimodalRuntimeData` | Derived information for the current LLM forward or prefill chunk. | -| `MultimodalEncoderRequestState` | Request-owned readiness and cache references, plus a reference to the admitted item metadata; the model still owns the `TensorLRUCache`. | - -The process boundary is: - -```text -Main / input processor - -> prompt token IDs + MultimodalParams (processed data and item metadata) - -> existing worker handoff: restore tensors; convert MultimodalInput to C++ fields -Executor / LlmRequest.py_multimodal_params - -> admission creates encoder_state where item scheduling is supported - -> selected encoder Params: selected data and metadata, no cache ownership - -> LLM forward Params: embeddings and current-chunk runtime, no cache ownership - -> executor cleanup releases request-owned cache references -``` - -After conversion, the worker clears the processor-side `multimodal_input`; -the executor uses the C++ request fields for that information. When persistent -encoder caching needs a Python `MultimodalInput`, it reconstructs it once and -reuses it through the request Params. Selected inputs copy the wrapper before -changing their hashes. Prefill/termination cleanup clears the retained wrapper. -Text-only requests -do not need an aggregate. Execution views leave `encoder_state` unset, and the -request-wide Params do not carry the current chunk's runtime data. Existing -prefetch events retain their stream synchronization role. Live cache references -and CUDA events are not added to the cross-process payload. - -An item reference identifies a processed item within its modality. It does not -necessarily identify one original media object or one tensor: a video can expand -into multiple processed items, and a field can pack multiple items in one tensor. -Partitions describe selection: a shared field has no item axis, a batched field -has an item axis, and a packed field also supplies per-item lengths. Contiguous -tensor selections can share storage; reordered selections may copy. Treat shared -input storage as read-only. Model hooks and `EncoderGroup` still determine -compatible batches, padding, masks, encoder arguments, and output fusion. - -Prompt multimodal token counts, physical encoder token costs, and output embedding -row counts serve different purposes. Unknown costs or row counts remain `None`; -metadata alone does not enable item scheduling. Producers and admission boundaries -validate the required contracts rather than revalidating all items per forward. -The item refs, costs, and output rows are stored as immutable tuples; list inputs -are accepted by the metadata constructor. Encoder state references this same -record and reads its lengths without a second snapshot. Producers fill the -partition declarations before admission. -Input preparation publishes the common item record directly. The frontend's -`mm_item_order` describes prompt order during preprocessing; executor Params -carry that order only in `item_metadata.item_refs`. Output row counts are stored -only in `item_metadata.output_embedding_lengths`. Request admission, selected-item -encoding, full-request encoding, cache keys and encoder output export all read -the same record. The encoder-result boundary still exports lengths as lists. - -`get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order=None)` is -the processor hook. Generic preparation supplies prompt order and known output -rows; model processors add encoder costs and partitions. Partial model metadata -retains known prompt rows, while conflicting order or rows fail at input preparation. -Gemma4 and Nemotron add partitions to this record. Qwen and Mistral also provide -physical encoder costs. A custom processor overrides this same hook and returns -`MultimodalItemMetadata`; the old three-field tuple, encoder-only metadata hook, -and Params order/length fallbacks have been removed. - -This data contract is independent of item scheduling. Processors can provide -common order and output rows with `encoder_token_lengths=None` and keep -full-request encoding. Model capability and configuration still decide whether -items are scheduled separately. - -Execution Params can receive model outputs. Inline encoding and prefetch retain -their existing output write-back to the request, while item-scheduled LLM chunks -keep their assembled embeddings local to the forward. CUDA event transfer remains -with the execution paths that synchronize the streams. - -For encoder/prefill disaggregation (E/P), model processor hooks return -`(prompt_token_ids, MultimodalParams)` and use `MultimodalParams.from_embeddings` -to attach embeddings or their handles. See the hook contract in -[inputs/registry.py](../../../tensorrt_llm/inputs/registry.py). -`DisaggregatedParams` remains the transport envelope; this local result does not -replace the wire protocol. The prefill worker still rebuilds model-specific prompt -layout. Shared data contracts do not change the supported model/modality matrix -or introduce a common batching policy for every model. - -External orchestrators can pass the existing `DisaggregatedParams` to the LLM API; -the prefill-side processor builds the local Params. Custom processors implementing -the E/P hook must follow its tuple return contract. Check those callers separately -from transport fields when changing this interface. +See the [multimodal developer guide](../developer-guide/multimodal.md) for +shared data contracts, process boundaries, and request ownership. ## Model Engine diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 6f13de48f0e7..00448ce81315 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -210,9 +210,8 @@ def _grid_for_rows(rows: int) -> list[int]: class _GroupedEncoderModel(MultimodalModelMixin): """Minimal model whose encoder serves image and video in one call. - Mirrors the Qwen3-VL shape: `encode_multimodal_inputs` delegates to - `encode_multimodal_by_groups`, which is what makes the - `multimodal_embedding_lengths` contract load-bearing. + Mirrors the Qwen3-VL grouped encoder, including its prompt-order output + assembly from the common item metadata. """ supports_mm_encoder_item_scheduling = True @@ -306,8 +305,6 @@ def _make_request(modalities: list[str], *, value_start: float = 1.0) -> Multimo encoder_token_lengths=list(lengths), output_embedding_lengths=list(lengths), ) - # Processors retain this field for legacy consumers. Encoder groups use - # the typed item metadata when available. return MultimodalParams(multimodal_data=multimodal_data) @@ -394,8 +391,7 @@ def get_multimodal_item_partitions(self, param, modality): def test_sliced_items_survive_the_grouped_encoder_path(): - """Regression: without `multimodal_embedding_lengths` on the slice, the - grouped path splits to zero rows and returns nothing for every item.""" + """Selected views preserve output rows without acquiring request cache ownership.""" model = _GroupedEncoderModel() request = _make_request(["image", "image"]) owner = MultimodalEncoderRequestState( @@ -417,46 +413,50 @@ def test_sliced_items_survive_the_grouped_encoder_path(): torch.testing.assert_close(output, expected) -def test_adjacent_same_modality_items_share_one_encoder_call(): - """`forward_multimodal_encoder_items` batches consecutive same-modality - inputs into a single grouped-encoder invocation.""" - model = _GroupedEncoderModel() - request = _make_request(["image", "image"]) - - outputs, encoder_inputs = _run_items(model, [(request, 0), (request, 1)]) - - assert len(encoder_inputs) == 1, "adjacent same-modality items must be sliced together" - assert encoder_inputs[0][1] == tuple(ITEM_ROWS["image"][:2]) - assert model.encoder_calls == [sum(ITEM_ROWS["image"][:2])] - assert len(outputs) == 2 - - -def test_mixed_modality_items_keep_scheduler_order(): - """A request interleaving modalities must still get one output per item, - in scheduler order -- this is where the per-modality reorder could bite.""" - model = _GroupedEncoderModel() - request = _make_request(["image", "video", "image"]) - selected = [(request, idx) for idx in range(3)] - - outputs, encoder_inputs = _run_items(model, selected) - - # image / video / image cannot merge across the modality change. - assert [modality for _, _, modality in encoder_inputs] == ["image", "video", "image"] - assert len(outputs) == 3 - for item_idx, output in enumerate(outputs): - torch.testing.assert_close(output, _expected_rows(request, item_idx)) - - -@pytest.mark.parametrize("modalities", [["video", "image", "image"], ["image", "video", "image"]]) -def test_full_and_item_encoder_paths_share_item_metadata(modalities: list[str]) -> None: +@pytest.mark.parametrize( + "modalities,item_order,expected_runs", + [ + pytest.param(["image", "image"], [0, 1], [("image", 5)], id="adjacent-images"), + pytest.param( + ["video", "image", "image"], + [0, 1, 2], + [("video", 4), ("image", 5)], + id="video-then-images", + ), + pytest.param( + ["image", "video", "image"], + [0, 1, 2], + [("image", 2), ("video", 4), ("image", 3)], + id="interleaved-vision", + ), + pytest.param(["image", "image"], [1, 0], [("image", 5)], id="reordered-images"), + ], +) +def test_full_and_item_encoder_paths_share_item_metadata( + modalities: list[str], item_order: list[int], expected_runs: list[tuple[str, int]] +) -> None: model = _GroupedEncoderModel() request = _make_request(modalities) - full_output = model.encode_multimodal_inputs([request]) - item_outputs, _ = _run_items(model, [(request, i) for i in range(len(modalities))]) - expected = torch.cat([_expected_rows(request, i) for i in range(len(modalities))]) + original = request.item_metadata + # Reorder prompt refs while retaining the original payload-local indices. + request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( + [original.item_refs[i] for i in item_order], + [original.encoder_token_lengths[i] for i in item_order], + [original.output_embedding_lengths[i] for i in item_order], + ) + expected = [_expected_rows(request, i) for i in range(len(item_order))] + torch.testing.assert_close( + model.encode_multimodal_inputs([request]), torch.cat(expected), rtol=0, atol=0 + ) - torch.testing.assert_close(full_output, expected, rtol=0, atol=0) - torch.testing.assert_close(torch.cat(item_outputs), expected, rtol=0, atol=0) + model.encoder_calls.clear() + item_outputs, encoder_inputs = _run_items(model, [(request, i) for i in range(len(item_order))]) + + assert [modality for _, _, modality in encoder_inputs] == [m for m, _ in expected_runs] + assert model.encoder_calls == [rows for _, rows in expected_runs] + assert len(item_outputs) == len(expected) + for actual, rows in zip(item_outputs, expected, strict=True): + torch.testing.assert_close(actual, rows, rtol=0, atol=0) def test_metadata_is_rebased_for_partial_and_reordered_items() -> None: @@ -477,19 +477,6 @@ def test_metadata_is_rebased_for_partial_and_reordered_items() -> None: torch.testing.assert_close(outputs[1], _expected_rows(request, 0), rtol=0, atol=0) -def test_full_and_item_order_can_differ_from_raw_payload_order() -> None: - model = _GroupedEncoderModel() - request = _make_request(["image", "image"]) - request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( - [("image", 1), ("image", 0)], [3, 2], [3, 2] - ) - expected = torch.cat([_expected_rows(request, 0), _expected_rows(request, 1)]) - item_outputs, _ = _run_items(model, [(request, 0), (request, 1)]) - torch.testing.assert_close(torch.cat(item_outputs), expected, rtol=0, atol=0) - - torch.testing.assert_close(model.encode_multimodal_inputs([request]), expected, rtol=0, atol=0) - - def _attach_cache_metadata(request, item_lengths): request.multimodal_input = MultimodalInput.from_components( [[i + 1] * 8 for i in range(len(item_lengths))], @@ -532,7 +519,6 @@ def test_full_cache_write_uses_canonical_lengths_without_legacy_fields(): model._multimodal_encoder_cache = TensorLRUCache(4096) request = _make_request(["image", "image"]) lengths = request.item_metadata.output_embedding_lengths - request.multimodal_data.pop("multimodal_embedding_lengths", None) _attach_cache_metadata(request, lengths) expected = model.encode_multimodal_inputs([request]) calls_before = len(model.encoder_calls) @@ -570,7 +556,6 @@ def test_partial_cache_batches_mixed_vision_and_audio_misses(): [4, 1, 3, 2, 1], [4, 1, 3, 2, 1], ) - request.multimodal_data.pop("multimodal_embedding_lengths", None) lengths = request.item_metadata.output_embedding_lengths _attach_cache_metadata(request, lengths) expected = model.encode_multimodal_inputs([request]) @@ -640,7 +625,6 @@ def build_audio_input(params: list[MultimodalParams]) -> dict[str, torch.Tensor] request.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( [("video", 0), ("audio", 0), ("image", 0), ("audio", 1)], [4, 2, 2, 2], [4, 2, 2, 2] ) - request.multimodal_data.pop("multimodal_embedding_lengths", None) expected = torch.cat( [_expected_rows(request, 0), audio[0], _expected_rows(request, 2), audio[1]] ) diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index febb2773f3d5..4b0f94e924e1 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -32,23 +32,27 @@ @pytest.mark.parametrize("axis", [0, 1, -1]) -@pytest.mark.parametrize("indices", [[2, 0], [1, 1, 0], []]) -def test_field_layout_selects_tensor_item_axis(axis: int, indices: list[int]) -> None: - values = torch.arange(9).reshape(3, 3) - selected = MultimodalItemPartition(item_axis=axis).select(values, indices) - expected = torch.index_select(values, axis, torch.tensor(indices, dtype=torch.long)) - torch.testing.assert_close(selected, expected) - - -@pytest.mark.parametrize("axis", [0, 1, -1]) -@pytest.mark.parametrize("indices", [[0, 1], [-2, -1], [2], []]) -def test_batched_item_partition_preserves_contiguous_views(axis, indices) -> None: +@pytest.mark.parametrize( + "indices,is_view", + [ + ([0, 1], True), + ([2, 0], False), + ([1, 1, 0], False), + ([], True), + ([-2, -1], True), + ([2], True), + ], +) +def test_batched_item_partition_selects_items_and_preserves_views( + axis: int, indices: list[int], is_view: bool +) -> None: values = torch.arange(9).reshape(3, 3) selected = MultimodalItemPartition(item_axis=axis).select(values, indices) expected = values[indices] if axis == 0 else values[:, indices] torch.testing.assert_close(selected, expected) - assert selected.untyped_storage().data_ptr() == values.untyped_storage().data_ptr() + shares_storage = selected.untyped_storage().data_ptr() == values.untyped_storage().data_ptr() + assert shares_storage is is_view @pytest.mark.parametrize("axis", [0, 1, -1]) @@ -628,80 +632,16 @@ def get_mm_item_metadata( assert processor.multimodal_hashing_supported is hashing_enabled -@pytest.mark.parametrize("hashing_enabled", [False, True], ids=["fallback", "hashing"]) -@pytest.mark.parametrize("use_uuids", [False, True], ids=["content", "uuid"]) -@pytest.mark.parametrize("frontend_order", [False, True], ids=["record-order", "frontend-order"]) -@pytest.mark.parametrize( - "modalities,token_ids,explicit_ids,special_ids,expected_cumsum,expected_runs", - [ - pytest.param( - ["image"], - [10, 101, 102, 20], - None, - None, - [0, 1, 2, 2], - ([0, 1], [1], [2]), - id="out-of-vocab-image", - ), - pytest.param( - ["image"], - [10, 101, 102, 20], - [101, 102], - None, - [0, 1, 2, 2], - ([0, 1], [1], [2]), - id="explicit-image-tokens", - ), - pytest.param( - ["image"], - [10, 101, 200, 102, 20], - [101, 102], - [200], - [0, 1, 1, 2, 2], - ([0, 1], [1], [3]), - id="image-framing-token", - ), - pytest.param( - ["video"], - [10, 101, 20, 102, 30], - None, - None, - [0, 1, 1, 2, 2], - ([0, 2], [1, 3], [1, 1]), - id="video-interleaved-text", - ), - pytest.param( - ["audio"], - [10, 101, 102, 20], - [101, 102], - None, - [0, 1, 2, 2], - ([0, 1], [1], [2]), - id="audio", - ), - pytest.param( - ["video", "image", "audio"], - [10, 101, 102, 20, 101, 102, 30, 101, 102, 40], - [101, 102], - None, - [0, 1, 2, 2, 3, 4, 4, 5, 6, 6], - ([0, 1, 2, 3], [1, 4, 7], [2, 2, 2]), - id="mixed-prompt-order", - ), - ], -) -def test_multimodal_input_contract_snapshot( +def _process_mm_contract( modalities: list[str], token_ids: list[int], explicit_ids: list[int] | None, special_ids: list[int] | None, - expected_cumsum: list[int], - expected_runs: tuple[list[int], list[int], list[int]], - hashing_enabled: bool, - use_uuids: bool, - frontend_order: bool, -) -> None: - """Pin framework metadata independently of model-specific media encoding.""" + *, + hashing_enabled: bool = True, + use_uuids: bool = False, + frontend_order: bool = True, +) -> tuple[list[int], dict]: item_refs = [(modality, 0) for modality in modalities] num_items = len(item_refs) @@ -757,7 +697,58 @@ def get_mm_item_metadata( if use_uuids: inputs["multi_modal_uuids"] = {modality: [f"{modality}-id"] for modality in modalities} - actual_token_ids, extra = create_input_processor_with_hash(Processor())(inputs, None) + return create_input_processor_with_hash(Processor())(inputs, None) + + +@pytest.mark.parametrize("hashing_enabled", [False, True], ids=["fallback", "hashing"]) +@pytest.mark.parametrize( + "modalities,token_ids,explicit_ids,special_ids,expected_cumsum,expected_runs", + [ + (["image"], [10, 101, 102, 20], None, None, [0, 1, 2, 2], ([0, 1], [1], [2])), + (["image"], [10, 101, 102, 20], [101, 102], None, [0, 1, 2, 2], ([0, 1], [1], [2])), + ( + ["image"], + [10, 101, 200, 102, 20], + [101, 102], + [200], + [0, 1, 1, 2, 2], + ([0, 1], [1], [3]), + ), + (["video"], [10, 101, 20, 102, 30], None, None, [0, 1, 1, 2, 2], ([0, 2], [1, 3], [1, 1])), + (["audio"], [10, 101, 102, 20], [101, 102], None, [0, 1, 2, 2], ([0, 1], [1], [2])), + ( + ["video", "image", "audio"], + [10, 101, 102, 20, 101, 102, 30, 101, 102, 40], + [101, 102], + None, + [0, 1, 2, 2, 3, 4, 4, 5, 6, 6], + ([0, 1, 2, 3], [1, 4, 7], [2, 2, 2]), + ), + ], + ids=[ + "out-of-vocab-image", + "explicit-image-tokens", + "image-framing-token", + "video-interleaved-text", + "audio", + "mixed-prompt-order", + ], +) +def test_multimodal_input_contract_snapshot( + modalities: list[str], + token_ids: list[int], + explicit_ids: list[int] | None, + special_ids: list[int] | None, + expected_cumsum: list[int], + expected_runs: tuple[list[int], list[int], list[int]], + hashing_enabled: bool, +) -> None: + """Pin framework metadata independently of model-specific media encoding.""" + actual_token_ids, extra = _process_mm_contract( + modalities, token_ids, explicit_ids, special_ids, hashing_enabled=hashing_enabled + ) + item_refs = [(modality, 0) for modality in modalities] + num_items = len(item_refs) assert actual_token_ids == token_ids data = extra["multimodal_data"] @@ -775,11 +766,9 @@ def get_mm_item_metadata( return mm_input = extra["multimodal_input"] - hashes = _UUID_HASH_SNAPSHOT if use_uuids else _CONTENT_HASH_SNAPSHOT + hashes = _CONTENT_HASH_SNAPSHOT assert mm_input.multimodal_hashes == [hashes[modality] for modality in modalities] - assert mm_input.multimodal_uuids == ( - [f"{modality}-id" for modality in modalities] if use_uuids else None - ) + assert mm_input.multimodal_uuids is None assert mm_input.multimodal_positions == [1 + 3 * index for index in range(num_items)] assert mm_input.multimodal_lengths == [2 + len(special_ids or [])] * num_items assert ( @@ -791,6 +780,30 @@ def get_mm_item_metadata( assert data["special_token_offsets"] == [1] +@pytest.mark.parametrize("use_uuids", [False, True], ids=["content", "uuid"]) +@pytest.mark.parametrize("frontend_order", [False, True], ids=["record-order", "frontend-order"]) +def test_hashes_and_uuids_follow_common_item_refs(use_uuids: bool, frontend_order: bool) -> None: + modalities = ["video", "image", "audio"] + _, extra = _process_mm_contract( + modalities, + [10, 101, 102, 20, 101, 102, 30, 101, 102, 40], + [101, 102], + None, + use_uuids=use_uuids, + frontend_order=frontend_order, + ) + mm_input = extra["multimodal_input"] + hashes = _UUID_HASH_SNAPSHOT if use_uuids else _CONTENT_HASH_SNAPSHOT + assert mm_input.multimodal_hashes == [hashes[modality] for modality in modalities] + assert mm_input.multimodal_uuids == ( + [f"{modality}-id" for modality in modalities] if use_uuids else None + ) + assert mm_input.multimodal_positions == [1, 4, 7] + assert extra["multimodal_data"][MULTIMODAL_ENCODER_ITEM_METADATA_KEY].item_refs == tuple( + (modality, 0) for modality in modalities + ) + + def test_mm_processor_kwargs_hash_is_stable_for_canonical_values(): input_processor = create_input_processor_with_hash( _KwargsHashFakeProcessor(), From 987fde0d56c4c6c40e555e0975c8a83df33d274f Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:06:44 -0700 Subject: [PATCH 14/22] [TRTLLM-14477][fix] address multimodal metadata review feedback Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- docs/source/developer-guide/multimodal.md | 90 ++++++++++++++ docs/source/release-notes.md | 8 ++ .../_torch/models/modeling_gemma4mm.py | 18 ++- .../models/modeling_multimodal_mixin.py | 44 +++++-- .../models/modeling_nemotron_h_multimodal.py | 7 +- .../_torch/models/modeling_qwen2vl.py | 12 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 23 +++- tensorrt_llm/inputs/multimodal.py | 18 ++- tensorrt_llm/inputs/registry.py | 4 + .../test_multimodal_encoder_item_groups.py | 114 +++++++++++++++++- .../multimodal/test_qwen3vl_disagg_prompt.py | 5 +- tests/unittest/inputs/test_multimodal.py | 11 ++ 12 files changed, 330 insertions(+), 24 deletions(-) diff --git a/docs/source/developer-guide/multimodal.md b/docs/source/developer-guide/multimodal.md index 671ac374bb35..6a66f32e6dc8 100644 --- a/docs/source/developer-guide/multimodal.md +++ b/docs/source/developer-guide/multimodal.md @@ -92,6 +92,96 @@ Custom processors must use this common hook and type. The former `get_mm_encoder_item_metadata` hook, `MultimodalEncoderItemMetadata` tuple, and Params order/length fallbacks have been removed. +## Onboarding a Model + +Start with full-request encoding and add item scheduling only when the model +can encode its items independently. Use these contracts rather than inventing +encoder costs or inferring item axes from matching tensor sizes. + +| Capability | Required information | +|---|---| +| Common request metadata | `item_refs` in prompt order. Unknown encoder costs and output rows stay `None`. | +| Grouped output assembly | Exact `output_embedding_lengths`; `EncoderGroup` needs them to split and reorder rows. | +| Item scheduling | Exact positive `encoder_token_lengths`, output row counts, and an item-selection implementation. | +| Persistent item caching | Output row counts and stable content/processor hashes. Partial hits also require correct item selection. | + +For example, a processor for a full-request audio encoder can return: + +```python +return MultimodalItemMetadata( + item_refs=[("audio", 0)], + encoder_token_lengths=None, + output_embedding_lengths=[num_output_rows], +) +``` + +This supports output assembly without declaring schedulable encoder costs. +If output row counts are unknown too, the model must resolve them before using +grouped assembly or item caching. A model-specific full-request encoder can +compute those counts from its result. Metadata alone never enables scheduling. + +To opt in, the model sets `supports_mm_encoder_item_scheduling = True` and +implements `encode_multimodal_inputs` and item selection. The engine also +checks `disable_mm_encoder`, the scheduling policy, and supported combinations +in `mm_item_scheduling_enabled()` and +`validate_mm_encoder_scheduling_compatibility()`. Under item scheduling, missing +costs or output rows fail admission; they do not silently switch one request to +full-request encoding. Keep the capability disabled until the model can supply +the complete scheduling contract. + +### Selecting Unusual Payloads + +Declare `input_partitions` for ordinary batched, packed, and shared fields. +For example, two audio items with two and one clips use: + +```python +input_partitions = { + "audio": { + "input_audio_features": MultimodalItemPartition(item_lengths=(2, 1)), + "feature_attention_mask": MultimodalItemPartition(item_lengths=(2, 1)), + "audio_num_clips": MultimodalItemPartition(), + }, +} +``` + +If a payload has multiple item axes, nested variable-length structures, or +selection that must rebuild sibling fields, override +`build_multimodal_encoder_input(param, local_item_indices, modality)`. Return a +new execution `MultimodalParams` view and leave the request's payload and cache +references intact. The mixin then slices the common item metadata and hashes +using prompt-global indices. Payload indices and prompt indices can differ: +for `[('video', 0), ('image', 0)]`, the image has local index `0` and prompt +index `1`. + +Returning `None` from `get_multimodal_item_partitions` uses the existing +packed/stacked compatibility selector. New models should declare partitions or +override selection explicitly; the compatibility selector cannot describe an +arbitrary model layout. An empty partition dict means every field is shared, +so it is appropriate only when selecting an item requires no payload changes. + +### Item Boundaries and Model Hooks + +An item is the smallest independently encodable unit, not necessarily one +uploaded media object. Several crops, frames, or audio clips can belong to one +item. A common item reference names one modality. Keep coupled audio/video +data in one model-owned payload under that item's modality, or keep full-request +encoding if independent item execution is not possible. Separate audio and +video references imply that either may be selected independently. + +Use `after_full_multimodal_embeddings` for request-level preparation after +cache lookup or encoding. Its tensors are not always full-request tensors: a +mixed batch containing item-scheduled chunk outputs already gathers active +rows. Use `after_active_multimodal_embeddings` for row-dependent transforms, +such as splitting Qwen3-VL deepstack features, immediately before fusion. + +Model-specific auxiliary values belong in `multimodal_data`. Declare their +selection rules alongside the primary fields, or rebuild them in the custom +selection hook. Add a field to `MultimodalItemMetadata` only when shared +preprocessing/execution consumers need it, and update validation, `select()`, +producer merging, and process-handoff tests together. Keep lengths and +partitions as host metadata; CUDA-to-host copies can synchronize even when +written as `.cpu().tolist()`. + ## Encoder/Prefill Handoff For encoder/prefill disaggregation (E/P), model processor hooks return diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 8f949a1aeb75..d7db1a83c212 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -13,6 +13,14 @@ All published functionality in the Release Notes has been fully tested and verif ### API Changes - **[BREAKING CHANGE] Removed obsolete C++ executor orchestration helpers.** The internal `RequestWithId` type, its `serializeReqWithIds` / `deserializeReqWithIds` functions, all production `insertRequestInOrder` overloads, and the unused `requestUtils.h` and `orchestratorUtils.h` helpers have been removed. These helpers belonged to the retired C++ executor orchestration path; downstream code linking their exported symbols must stop using them and rebuild. The current `Request` / `Response` interfaces and serialization remain available. Executor factory and worker constructors no longer accept the unused `engine` and `batched_logits_processor` arguments; callers should omit them and continue passing model configuration through `llm_args` and `hf_model_dir`. Unused decoder, request-cancellation, executor, legacy-model, orchestration, and speculative-decoding MPI tag names have also been removed; values 1–23, 127–129, and 1023–1025 are retired. The default, KV-cache event, and context-transfer tags retain their existing values. Internal `LogprobParams` and `LogProbsResult` tuples now contain only prompt-logprob fields; generation logprobs remain supplied by the sampler. The unused `TLLM_EXECUTOR_PERIODICAL_RESP_IN_AWAIT` environment variable has been removed. +- Multimodal processor extensions now use `get_mm_item_metadata()` and + `MultimodalItemMetadata`, replacing `get_mm_encoder_item_metadata()` and + `MultimodalEncoderItemMetadata`. The removed `DisaggPrefillMultimodalInputs` + result is replaced by `(prompt_token_ids, MultimodalParams)` from + `build_disagg_prefill_multimodal_inputs()`, which also accepts the keyword-only + `mm_hashes` argument. Custom processors must migrate these hooks; processors + defining the removed metadata hook now fail at wrapper creation. See the + [multimodal developer guide](developer-guide/multimodal.md). ### Fixed Issues diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index f7a11b001c36..fc16dfe66931 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -227,7 +227,23 @@ def _gemma4_input_partitions( item_count = sum(m == modality for m, _ in metadata.item_refs) value = data.get(fields[0]) if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] == item_count: - partitions[modality] = {name: MultimodalItemPartition() for name in fields} + modality_partitions = {} + for name in fields: + value = data.get(name) + if value is None: + continue + if isinstance(value, torch.Tensor): + size = value.shape[0] if value.ndim > 0 else None + elif isinstance(value, (list, tuple)): + size = len(value) + else: + size = None + if size != item_count: + raise ValueError( + f"{modality}.{name} must have {item_count} items on its leading axis" + ) + modality_partitions[name] = MultimodalItemPartition() + partitions[modality] = modality_partitions return partitions diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 7af5b9d153aa..70b06c306316 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -135,6 +135,18 @@ def _lengths_by_modality( Layouts describe whole requests, so only modalities in this encoder group contribute lengths (for example, image/video in a vision group while the same request's audio is handled by another group). + + Encoders batch in payload order, which can differ from prompt order. + Modality-local indices restart for each request. For example:: + + item_layouts = [ + ([("image", 1), ("video", 0), ("image", 0)], [20, 30, 10]), + ([("audio", 0), ("image", 0)], [40, 50]), + ] + modalities = ("image", "video") + # Result: {"image": [10, 20, 50], "video": [30]} + + These lengths split encoder outputs before restoring prompt order. """ by_modality: Dict[str, List[int]] = {m: [] for m in modalities} for item_refs, lengths in item_layouts: @@ -252,6 +264,10 @@ def encode_multimodal_by_groups( group_layouts = [] for mp, layout in zip(multimodal_params, item_layouts, strict=True): if any(mp.multimodal_data.get(m) is not None for m in group.modalities): + if not layout[0]: + raise MultimodalEncoderContractError( + "MM item metadata with output lengths is required for grouped encoding" + ) group_params.append(mp) group_layouts.append(layout) if not group_params: @@ -844,11 +860,13 @@ def get_multimodal_item_partitions( def build_multimodal_encoder_input( self, param: MultimodalParams, - item_indices: Sequence[int], + local_item_indices: Sequence[int], modality: Optional[str] = None, ) -> MultimodalParams: """Build a forward view containing selected modality-local items. + ``local_item_indices`` indexes this modality's payload, not positions + in the request's prompt-ordered item metadata. Declared field layouts drive common tensor/list selection. Models with custom preprocessing can still override this hook. The caller slices the parallel item metadata after this returns, preserving original @@ -869,9 +887,11 @@ def build_multimodal_encoder_input( ) item_partitions = self.get_multimodal_item_partitions(param, modality) if item_partitions is None: - sliced = self._select_legacy_multimodal_fields(modality_data, modality, item_indices) + sliced = self._select_legacy_multimodal_fields( + modality_data, modality, local_item_indices + ) else: - sliced = select_multimodal_fields(modality_data, item_partitions, item_indices) + sliced = select_multimodal_fields(modality_data, item_partitions, local_item_indices) # Metadata slicing must not rewrite the original request's hashes. residual_input = ( @@ -1219,7 +1239,13 @@ def _encoder_cache_keys( return None item_metadata = param.item_metadata - if item_metadata is None or item_metadata.output_embedding_lengths is None: + if item_metadata is None: + logger.debug(f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params without item metadata.") + return None + if item_metadata.output_embedding_lengths is None: + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params without item output lengths." + ) return None item_modalities = {modality for modality, _ in item_metadata.item_refs} if any( @@ -1347,7 +1373,7 @@ def partition_encoder_cache( def _apply_metadata_slice( residual: MultimodalParams, source: MultimodalParams, - item_indices: Sequence[int], + prompt_item_indices: Sequence[int], *, item_metadata: Optional[MultimodalItemMetadata] = None, ) -> None: @@ -1355,16 +1381,20 @@ def _apply_metadata_slice( Models slice raw modality tensors in `build_multimodal_encoder_input`; the mixin owns the parallel per-item metadata slice so every model gets it identically. + ``prompt_item_indices`` indexes the source metadata's ``item_refs`` + across all modalities, rather than a modality's payload. """ metadata = item_metadata if item_metadata is not None else source.item_metadata if metadata is None: raise MultimodalEncoderContractError("Item selection requires multimodal item metadata") residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata.select( - item_indices + prompt_item_indices ) if residual.multimodal_input is not None and source.multimodal_input is not None: source_hashes = source.multimodal_input.multimodal_hashes - residual.multimodal_input.multimodal_hashes = [source_hashes[i] for i in item_indices] + residual.multimodal_input.multimodal_hashes = [ + source_hashes[i] for i in prompt_item_indices + ] def _encode_with_partial_cache( self, diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index 9a621fe5be7c..e1708eb4eff9 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -1160,10 +1160,15 @@ def get_mm_item_metadata( ) -> Optional[MultimodalItemMetadata]: metadata = super().get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) audio = multimodal_data.get("audio") + # EVS merge also carries full prompt IDs for image/audio requests. + # Selecting individual items would require rebuilding that prompt stream. if metadata is not None and isinstance(audio, dict) and audio.get("evs_ids") is None: # One audio item can contain several clips. Video's embedded audio # stays in its video bucket and retains the model's AV interleave. - clips = tuple(audio["audio_num_clips"].tolist()) + clip_counts = audio.get("audio_num_clips") + if clip_counts is None: + return metadata + clips = tuple(clip_counts.cpu().tolist()) item_count = sum(m == "audio" for m, _ in metadata.item_refs) if len(clips) == item_count: packed = MultimodalItemPartition(item_lengths=clips) diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 9fec6eb3595a..928e8ce0e71c 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -349,7 +349,10 @@ def get_mm_item_metadata( mm_item_order) if metadata is not None: for modality in ("image", "video"): - if isinstance(multimodal_data.get(modality), dict): + modality_data = multimodal_data.get(modality) + grid_key = "image_grid_thw" if modality == "image" else "video_grid_thw" + if isinstance(modality_data, + dict) and modality_data.get(grid_key) is not None: metadata.input_partitions[ modality] = _qwen_vision_input_partitions( multimodal_data, modality, metadata) @@ -2136,11 +2139,8 @@ def _qwen_vision_item_partitions( metadata = param.item_metadata if metadata is not None and modality in metadata.input_partitions: return metadata.input_partitions[modality] - grid_key = "image_grid_thw" if modality == "image" else "video_grid_thw" - grids = param.multimodal_data[modality][grid_key] - return _qwen_vision_input_partitions( - param.multimodal_data, modality, - metadata if grids.device.type != "cpu" else None) + return _qwen_vision_input_partitions(param.multimodal_data, modality, + metadata) def _qwen_vision_input_partitions( diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index d881d47eea41..7267fee2f949 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -13,11 +13,9 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.bindings import executor as tllm_executor from tensorrt_llm.executor.result import SimpleTokenLogprobs, TokenLogprobs -from tensorrt_llm.inputs.multimodal import \ - MultimodalEncoderProgress as MultimodalEncoderProgress -from tensorrt_llm.inputs.multimodal import \ - MultimodalEncoderRequestState as MultimodalEncoderRequestState -from tensorrt_llm.inputs.multimodal import (MultimodalParams, +from tensorrt_llm.inputs.multimodal import (MultimodalEncoderProgress, + MultimodalEncoderRequestState, + MultimodalParams, get_multimodal_item_metadata) from tensorrt_llm.sampling_params import LogprobMode @@ -1139,6 +1137,11 @@ def finish_by(self, reason: FinishReason, beam: int) -> None: @property def py_multimodal_data(self) -> dict[str, Any] | None: + """Request-owned processed payload, e.g. ``{"image": {"pixel_values": tensor}}``. + + The same dict stores the common item record under + ``MULTIMODAL_ENCODER_ITEM_METADATA_KEY`` and model-specific auxiliary data. + """ params = self.py_multimodal_params return params.multimodal_data if params is not None else None @@ -1153,6 +1156,11 @@ def py_multimodal_data(self, data: dict[str, Any] | None) -> None: @property def py_mm_encoder_state(self) -> MultimodalEncoderRequestState | None: + """Item-encoding progress and cache references owned by the request Params. + + Set at encoder admission and cleared by request cleanup. Inline encoding + and requests with precomputed embeddings leave it unset. + """ params = self.py_multimodal_params return params.encoder_state if params is not None else None @@ -1167,6 +1175,11 @@ def py_mm_encoder_state( @property def py_mm_encoder_event(self) -> torch.cuda.Event | None: + """Completion event for encoder work queued on the prefetch CUDA stream. + + Consumers wait before reading the staged payload or embeddings; the + executor clears the event after that work is consumed or cleaned up. + """ params = self.py_multimodal_params return params.encoder_event if params is not None else None diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index a9d2d2970ea3..611f4be23824 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -52,6 +52,9 @@ class MultimodalItemPartition(NamedTuple): def select(self, value: Any, item_indices: Sequence[int]) -> Any: """Select items in the requested order, preserving tensor dtype/device. + ``item_indices`` contains modality-local payload indices, for example + ``[2, 0]`` selects the third item followed by the first. For a packed + field, each index selects its whole span from ``item_lengths``. Contiguous tensor selections return views that callers must treat as read-only. Other selections copy only the selected items; shared values retain identity. @@ -76,12 +79,14 @@ def select(self, value: Any, item_indices: Sequence[int]) -> Any: return value.narrow(axis, 0, 0) if indices == list(range(indices[0], indices[0] + len(indices))): start, end = indices[0], indices[-1] + 1 + # Packed items map to row spans; batched items already use row indices. if offsets is not None: start, end = offsets[start], offsets[end] if axis == 0: return value[start:end] return value.narrow(axis, start, end - start) if offsets is not None: + # Reordered packed items need concatenation; each may occupy several rows. if axis == 0: pieces = [value[offsets[i]:offsets[i + 1]] for i in indices] else: @@ -90,6 +95,7 @@ def select(self, value: Any, item_indices: Sequence[int]) -> Any: for i in indices ] return torch.cat(pieces, dim=axis) + # With no packed lengths, tensor indexing selects one row per item. if axis == 0: return value[indices] index: list[slice | list[int]] = [slice(None)] * value.ndim @@ -126,6 +132,10 @@ class MultimodalItemMetadata: Encoder costs and output rows are optional: their presence does not enable item scheduling. Models keep their existing scheduling capabilities. + + Request state shares this record, so refs and lengths cannot be reassigned. + Input partitions are completed by producers before admission. This is a + runtime data record, not a user-facing configuration model. """ item_refs: Sequence[tuple[str, int]] @@ -142,6 +152,8 @@ def __init__( input_partitions: dict[str, dict[str, MultimodalItemPartition]] | None = None, ) -> None: + """Normalize lists while initializing each frozen field once.""" + # Frozen dataclasses require object.__setattr__ during construction. object.__setattr__(self, "item_refs", tuple(item_refs)) object.__setattr__( self, "encoder_token_lengths", None @@ -202,7 +214,11 @@ def validate(self) -> None: ) def select(self, item_indices: Sequence[int]) -> "MultimodalItemMetadata": - """Select items and rebase their refs and packed lengths to the view.""" + """Select prompt-order positions in ``item_refs`` and rebase the view. + + These indices span all modalities; payload partitions instead take + modality-local indices. + """ if len(item_indices) == 0: return MultimodalItemMetadata( (), None if self.encoder_token_lengths is None else (), diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 11ca6267f50d..c14eb083b0e7 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1282,6 +1282,10 @@ def create_input_processor_with_hash( Returns: A wrapped processor that modifies prompts before processing. """ + if callable(getattr(input_processor, "get_mm_encoder_item_metadata", None)): + raise TypeError( + "get_mm_encoder_item_metadata() has been removed; implement " + "get_mm_item_metadata() returning MultimodalItemMetadata instead.") def multimodal_hashing_process( inputs: TextPrompt, sampling_params: SamplingParams diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 00448ce81315..cff43fef85ae 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -82,6 +82,47 @@ def test_qwen_common_metadata_preserves_positive_output_requirement(): processor.get_mm_item_metadata([12, 10, 13], data) +@pytest.mark.parametrize("modality", ["image", "video"]) +def test_qwen_precomputed_embeddings_skip_payloads_without_grids(modality: str) -> None: + from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2VLInputProcessorBase + + processor = object.__new__(Qwen2VLInputProcessorBase) + metadata = MultimodalItemMetadata([(modality, 0)], output_embedding_lengths=[3]) + data = { + modality: {}, + "multimodal_embedding": torch.zeros(3, HIDDEN), + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, + } + + result = processor.get_mm_item_metadata([], data) + + assert result.item_refs == metadata.item_refs + assert result.output_embedding_lengths == (3,) + assert result.input_partitions == {} + + +def test_qwen_partition_uses_metadata_costs_without_reading_grid_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.models.modeling_qwen2vl import _qwen_vision_item_partitions + + metadata = MultimodalItemMetadata([("image", 1), ("image", 0)], [32, 16], [8, 4]) + params = MultimodalParams( + multimodal_data={ + "image": {"image_grid_thw": torch.tensor([[1, 4, 4], [2, 4, 4]])}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, + } + ) + + def unexpected_grid_read(tensor: torch.Tensor) -> list: + raise AssertionError("Known encoder costs must not be recomputed from grid values") + + monkeypatch.setattr(torch.Tensor, "tolist", unexpected_grid_read) + partitions = _qwen_vision_item_partitions(params, "image") + + assert partitions["pixel_values"].item_lengths == (16, 32) + + @pytest.mark.parametrize("model", ["qwen2", "qwen2_5", "qwen3", "mistral"]) @pytest.mark.parametrize("custom", [False, True]) def test_processors_use_the_common_metadata_hook(model, custom): @@ -144,7 +185,11 @@ def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): processor = object.__new__(Gemma4InputProcessor) data = { - "image": {"pixel_values": torch.zeros(1, 4, 6), "image_position_ids": torch.zeros(1, 4, 2)}, + "image": { + "pixel_values": torch.zeros(1, 4, 6), + "image_position_ids": torch.zeros(1, 4, 2), + "image_seq_lens": [4], + }, "audio": {"audio_features": torch.zeros(1, 5, 8), "audio_features_mask": torch.ones(1, 5)}, MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( [("audio", 0), ("image", 0)], output_embedding_lengths=[2, 4] @@ -159,6 +204,32 @@ def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): assert metadata.input_partitions["image"]["image_position_ids"].item_axis == 0 +@pytest.mark.parametrize( + "modality,field,value", + [ + ("image", "image_position_ids", torch.zeros(2, 4, 2)), + ("image", "image_seq_lens", [4, 4]), + ("audio", "audio_features_mask", torch.ones(2, 5)), + ], +) +def test_gemma_processor_rejects_misaligned_sibling_items( + modality: str, field: str, value: torch.Tensor | list[int] +) -> None: + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + + processor = object.__new__(Gemma4InputProcessor) + primary = "pixel_values" if modality == "image" else "audio_features" + data = { + modality: {primary: torch.zeros(1, 4, 6), field: value}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [(modality, 0)], output_embedding_lengths=[4] + ), + } + + with pytest.raises(ValueError, match=rf"{modality}\.{field} must have 1 items"): + processor.get_mm_item_metadata([], data) + + def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): from tensorrt_llm._torch.models.modeling_nemotron_h_multimodal import ( NemotronHMultimodalInputProcessor, @@ -202,6 +273,31 @@ def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): assert video_data["video"]["audio"] is audio +@pytest.mark.parametrize("full_prompt_ids", [False, True]) +def test_nemotron_audio_skips_incomplete_or_prompt_coupled_partitions( + full_prompt_ids: bool, +) -> None: + from tensorrt_llm._torch.models.modeling_nemotron_h_multimodal import ( + NemotronHMultimodalInputProcessor, + ) + + processor = object.__new__(NemotronHMultimodalInputProcessor) + audio = {"input_audio_features": torch.zeros(1, 5, 2)} + if full_prompt_ids: + audio.update(audio_num_clips=torch.tensor([1]), evs_ids=torch.tensor([1, 2, 3])) + data = { + "audio": audio, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("audio", 0)], output_embedding_lengths=[2] + ), + } + + metadata = processor.get_mm_item_metadata([], data) + + assert metadata.item_refs == (("audio", 0),) + assert "audio" not in metadata.input_partitions + + def _grid_for_rows(rows: int) -> list[int]: """A (t, h, w) grid whose product is the item's patch count.""" return [1, 1, rows] @@ -253,6 +349,22 @@ def encode_multimodal_inputs(self, multimodal_params) -> torch.Tensor: return encode_multimodal_by_groups(self.mm_encoder_groups, list(multimodal_params)) +@pytest.mark.parametrize("empty_record", [False, True]) +def test_grouped_encoder_rejects_raw_payload_without_item_layout(empty_record: bool) -> None: + model = _GroupedEncoderModel() + data = {"image": {"pixel_values": torch.ones(2, HIDDEN)}} + if empty_record: + data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( + [], output_embedding_lengths=[] + ) + params = MultimodalParams(multimodal_data=data) + + with pytest.raises(MultimodalEncoderContractError, match="metadata.*grouped encoding"): + model.encode_multimodal_inputs([params]) + + assert model.encoder_calls == [] + + def test_common_metadata_shares_executor_and_encoder_layout(): from tensorrt_llm._torch.pyexecutor.llm_request import get_multimodal_embedding_lengths diff --git a/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py b/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py index 5544c7c997ab..3a3241e708c5 100644 --- a/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py +++ b/tests/unittest/_torch/multimodal/test_qwen3vl_disagg_prompt.py @@ -219,8 +219,9 @@ def capture(*args, **kwargs): assert params.multimodal_data["mrope_config"]["mrope_position_deltas"] is mrope_deltas cumsum = params.multimodal_data["multimodal_embed_mask_cumsum"] assert cumsum.device.type == "cpu" + # Only video embedding positions advance the mask; framing and text do not. + assert cumsum.tolist() == [0, 0, 1, 2, 3, 3, 3] assert cumsum.numel() == len(tokens) - assert cumsum[-1].item() == 3 assert params.encoder_state is None assert params.encoder_event is None @@ -267,7 +268,7 @@ def test_external_embeddings_reuse_common_item_metadata(model): ) params = MultimodalParams(multimodal_data=extra["multimodal_data"]) - assert len(tokens) == 16 + assert tokens == [11, *([_IMAGE_TOKEN_ID] * 4), 12, *([_IMAGE_TOKEN_ID] * 9), 13] assert params.get_item_layout() == ((("image", 0), ("image", 1)), (4, 9)) assert params.multimodal_data["multimodal_embedding"] is embeddings assert params.multimodal_data["mrope_config"] is mrope_config diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 4b0f94e924e1..448f96d8e98e 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -210,6 +210,17 @@ def test_mm_item_metadata_rejects_mismatched_embedding_lengths(): input_processor({"prompt": "unused"}, sampling_params=None) +def test_removed_encoder_metadata_hook_fails_before_processing() -> None: + class OldMetadataProcessor(_ItemMetadataFakeProcessor): + def get_mm_encoder_item_metadata( + self, prompt_token_ids: list[int], multimodal_data: dict + ) -> None: + raise AssertionError("The removed hook must not be called") + + with pytest.raises(TypeError, match=r"implement get_mm_item_metadata\(\)"): + create_input_processor_with_hash(OldMetadataProcessor([2])) + + def test_mm_item_metadata_none_falls_back_to_non_item_path(): # Generic item order does not imply that physical encoder costs or # item scheduling are supported. From f4e903938a72b382cf9d4523f396df95328ed81a Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:43:09 -0700 Subject: [PATCH 15/22] [TRTLLM-14477][docs] align multimodal onboarding with common item metadata Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../trtllm-model-onboard-multimodal/SKILL.md | 77 +++++++++++++------ docs/source/developer-guide/multimodal.md | 4 + docs/source/features/multi-modality.md | 10 +++ docs/source/models/adding-new-model.md | 9 +++ docs/source/torch/adding_new_model.md | 9 +++ 5 files changed, 86 insertions(+), 23 deletions(-) diff --git a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md index 94dbcc2e62f5..87b20c55ff6a 100644 --- a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md +++ b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md @@ -10,10 +10,22 @@ metadata: author: NVIDIA Corporation --- + + # TensorRT-LLM Multimodal Model Onboarding (PyTorch backend) > **Scope.** PyTorch backend only (`tensorrt_llm/_torch/`) — the default for `LLM(..., backend="pytorch")`, `trtllm-serve`, and `trtllm-bench`. +Read the [multimodal developer guide](../../../docs/source/developer-guide/multimodal.md) +for request ownership and process boundaries. Its +[model onboarding section](../../../docs/source/developer-guide/multimodal.md#onboarding-a-model) +defines minimum metadata, scheduling requirements, unusual payload selection, +coupled media items, and model-specific extensions. Use that contract alongside +the implementation steps below. + **Output:** - `tensorrt_llm/_torch/models/modeling_{name}.py` — wrapper class (multimodal encoder + LLM) decorated with `@register_auto_model`, `@register_vision_encoder`, `@register_input_processor` (and `@support_multimodal_disaggregated` if EPD is supported), plus a `BaseMultimodalInputProcessor` (+ `BaseMultimodalDummyInputsBuilder`) subclass. - `_torch/models/checkpoints/hf/{name}_weight_mapper.py` if HF prefixes need surgery. @@ -40,7 +52,9 @@ metadata: (mRoPE) mrope_position_ids + deltas computed on CPU _postprocess: HF mm token ids → tllm_multimodal_token_id (OOV sentinel) The framework wrapper around your processor computes blake3 content hashes - for KV-cache reuse. + for KV-cache reuse and merges get_mm_item_metadata into one + MultimodalItemMetadata record inside multimodal_data. The frontend's + mm_item_order is consumed here; execution reads item_metadata.item_refs. MultimodalParams.to_handle("multimodal_data") at the end → each tensor in multimodal_data is replaced by a small dict pointing at its CUDA-IPC / shm handle, so the broadcast in [3] carries pointers, not megabytes of pixels. @@ -48,24 +62,26 @@ metadata: [3] Worker fan-out (TP / PP / CP) Each worker rebuilds local tensor views via to_tensor("multimodal_data"). multimodal_input (hashes / positions / lengths) is forwarded to the C++ - executor to drive KV-cache hash matching. + executor to drive KV-cache hash matching, then cleared from the request's + Params. Persistent encoder caching reconstructs it lazily when needed. + Live encoder_state cache references and CUDA events stay in the executor. [4] Per-iteration staging (model engine) Context: build MultimodalRuntimeData (positions / lengths / chunk bounds). For item-scheduled models, reserve space in the model's one TensorLRUCache. Reuse ready outputs and encode each missing item only once, within the encoder item / token / output-byte limits. Then - join the request's cached item outputs in prompt order, keeping only - the current chunk's rows, into a multimodal_embedding flagged - multimodal_embedding_is_chunk. - Legacy models stage the request payload here as before. H2D obeys + assemble the current chunk's rows in a forward Params view with + multimodal_embedding_is_chunk. The request keeps cache references, + not a combined embedding buffer. + Full-request models stage the request payload here. H2D obeys multimodal_data_device_paths and uses pinned, non-blocking copies. Generation (mRoPE only): strip everything except mrope_position_deltas. Post-prefill / termination: release request-held cache references, then strip raw MM fields so they don't ride along in decode. [5] Model.forward(attn_metadata, input_ids, position_ids, multimodal_params=…) - MultimodalModelMixin.prepare_multimodal_inputs: the legacy path looks up + MultimodalModelMixin.prepare_multimodal_inputs: full-request encoding looks up the persistent cache and encodes partial-hit misses; get_multimodal_embeddings then encodes the remaining full misses and joins every request's rows into one batch tensor. If any request is @@ -99,7 +115,7 @@ metadata: When `@support_multimodal_disaggregated` is set and the deployment uses `TLLM_MULTIMODAL_DISAGGREGATED=1`: - **Encoder worker:** runs as a standalone `MultimodalEncoder` (`mm_encoder_only=True`). It executes only the multimodal encoder and ships `mm_embeddings` (+ mRoPE position ids/deltas) to prefill+decode workers as shared-tensor handles. -- **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_mm_disagg()` is true; the input processor's `attach_multimodal_embeddings()` override binds the encoder handles into the request. For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors. +- **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_mm_disagg()` is true. For supplied embeddings, override `_attach_multimodal_embeddings_impl`, not the `attach_multimodal_embeddings` wrapper. Handle-based E/P input uses `build_disagg_prefill_multimodal_inputs(inputs, mm_handles, *, mm_hashes=None)` and returns `(prompt_token_ids, MultimodalParams)`; build the Params with `MultimodalParams.from_embeddings`. `DisaggregatedParams` remains the transport envelope, separate from this local return value. For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors. This item-scheduled cache path currently applies only when the encoder and LLM run in the same process. `mm_encoder_only` / EPD keeps its existing @@ -218,7 +234,7 @@ A 1024×1024 fp32 patch tensor is ~12 MB; a video clip can be hundreds of MB. Na ### Contract 4 — Batch the multimodal encoder across requests Implement `MultimodalModelMixin.encode_multimodal_inputs` as one batched forward over its -`MultimodalParams` list. The legacy path batches uncached requests through +`MultimodalParams` list. Full-request encoding batches uncached requests through `get_multimodal_embeddings`. The item-scheduled path instead passes sliced inputs for selected producer items and calls the same hook once per compatible modality run. For example, `image -> video -> image` can require three calls, even if the model supports both modalities @@ -229,10 +245,14 @@ boundaries, and run the encoder blocks once. Never loop over requests and call t per request. **Pattern (Qwen3-VL).** `encode_multimodal_by_groups` concatenates items for compatible -modalities, runs one encoder call, and restores request/prompt order from `mm_item_order`. +modalities, runs one encoder call, and restores request/prompt order from +`param.item_metadata.item_refs` and `output_embedding_lengths`. `forward_multimodal_encoder_items` then splits the result into one tensor per atomic item using -the processor-declared output lengths. Override `build_multimodal_encoder_input` only when the -default packed-grid, stacked-image, or stacked-audio slicers cannot represent the model's input. +the processor-declared output lengths. Declare `item_metadata.input_partitions` +with `MultimodalItemPartition` for batched, packed, or shared fields. Override +`build_multimodal_encoder_input(param, local_item_indices, modality)` when those +partitions cannot express the payload. The mixin separately selects metadata +using prompt-global indices. New models should not rely on legacy layout inference. **Audit.** Under load, compatible requests should share one wide encoder block per hook call, not one block per request. Item scheduling may produce multiple blocks per iteration when @@ -271,8 +291,8 @@ class {Name}VisionModel(nn.Module): class {Name}Model(MultimodalModelMixin, PreTrainedModel): config_class = {Name}Config supports_encoder_cache = True - # Enable only after implementing the Phase 3 atomic-item contracts. - supports_mm_encoder_item_scheduling = True + # Enable after implementing and validating the Phase 3 atomic-item contracts. + supports_mm_encoder_item_scheduling = False def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): config = model_config.pretrained_config @@ -337,7 +357,7 @@ class {Name}Model(MultimodalModelMixin, PreTrainedModel): - Implement `encode_multimodal_inputs`, `language_model`, `text_embedding_layer`, `embedding_dim`, and `embedding_dtype`. The encoder must return one tensor whose first dimension equals the processor-declared aggregate output length. The item path validates and - splits it per item; the legacy path retains its existing full-request behavior. + splits it per item; full-request encoding uses the same common metadata record. `embedding_dim` and `embedding_dtype` describe one cached encoder-output row, not the text embedding: include any extra features folded into each row (Qwen3-VL deepstack returns `hidden * (deepstack_num_level + 1)`). Item scheduling sizes cache entries from them and @@ -374,12 +394,21 @@ The workspace dimension comes from `encoder_max_num_tokens`, falling back to the > Mixed image+video+audio models profile the supported modality with the largest legal per-item workload under the configured limits. Runtime mixed-modality requests still share the same aggregate token limit. -**Atomic-item scheduling.** Override `get_mm_encoder_item_metadata` to return -prompt-ordered `item_refs`, the encoder-token cost of each item, and each -item's output length. Also implement `get_max_mm_encoder_output_embeddings` so -the engine can size the cache for one legal encoder iteration. The declared -output lengths must match both the MM placeholder spans and the tensor splits -from `forward_multimodal_encoder_items`. +**Common item metadata.** Implement +`get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order=None)` +and return `MultimodalItemMetadata`. The framework supplies prompt order and +known output rows; the processor adds encoder costs and input partitions. +Full-request models can leave unknown costs as `None`; grouped assembly still +requires exact output row counts. Do not implement the removed +`get_mm_encoder_item_metadata` hook: the registry rejects it at wrapper creation. + +**Atomic-item scheduling (optional).** Enable +`supports_mm_encoder_item_scheduling` only when items can be encoded independently. +Supply exact positive `encoder_token_lengths`, `output_embedding_lengths`, and +item selection, and implement `get_max_mm_encoder_output_embeddings` to size +the cache for one legal iteration. Missing scheduling metadata fails admission; +it does not trigger a per-request fallback. Keep lengths and partitions on the +host: `.cpu().tolist()` can synchronize and does not belong in a forward path. Implement `call_with_text_prompt(inputs, sampling_params)` — the per-model text-prompt path. **Don't override `__call__`**: the base class's concrete `__call__` dispatches here for text prompts, and also detokenizes `prompt_token_ids → prompt` and falls through to here for non-fast-path VLMs. `call_with_text_prompt` does: @@ -392,7 +421,7 @@ Implement `call_with_text_prompt(inputs, sampling_params)` — the per-model tex **Optional tokenized+MM fast path (skip unless needed):** set `supports_token_id_mm_expansion = True` (`ClassVar`) and implement `get_text_with_mm_placeholders(mm_counts)` + `expand_prompt_token_ids_for_mm(prompt_token_ids, num_mm_tokens, ...)`. The base-class `call_with_token_ids` then builds dummy placeholder text, runs `call_with_text_prompt` on it, expands the real token IDs, and merges any returned `mm_data_updates` (e.g. video `evs_ids`) into `multimodal_data`. Leave the flag `False` and the base class just detokenizes token-ID inputs and re-runs `call_with_text_prompt`. Only LlavaNext + Nemotron H multimodal opt in today. -**EPD override (if `@support_multimodal_disaggregated`):** override `_attach_multimodal_embeddings_impl(inputs, multimodal_embedding, sampling_params)` — **not** the `attach_multimodal_embeddings` wrapper — to consume encoder outputs in the prefill+decode worker. The base wrapper detokenizes tokenized inputs for non-fast-path VLMs before delegating to your impl. +**EPD overrides (if `@support_multimodal_disaggregated`):** override `_attach_multimodal_embeddings_impl(inputs, multimodal_embedding, sampling_params)` — **not** the `attach_multimodal_embeddings` wrapper — for supplied embeddings. The base wrapper detokenizes tokenized inputs for non-fast-path VLMs before delegating to your impl. If the model supports handle-based E/P handoff, implement `build_disagg_prefill_multimodal_inputs(inputs, mm_handles, *, mm_hashes=None) -> tuple[list[int], MultimodalParams]`; preserve prompt order, hashes, output rows, and model-specific auxiliary data in the returned Params. See Qwen3-VL's implementation and the developer guide's [handoff contract](../../../docs/source/developer-guide/multimodal.md#encoderprefill-handoff). **Decorator stack** — bottom-up application; `register_vision_encoder` requires `register_auto_model` to have run first: @@ -508,10 +537,12 @@ Follow `CONTRIBUTING.md`. Title `[JIRA/NVBUG/None][type] description`, `git comm **Input processor** - [ ] Subclasses both `BaseMultimodalInputProcessor` and `BaseMultimodalDummyInputsBuilder`. -- [ ] Encoder profiling implements `get_mm_max_tokens_per_item` + `get_dummy_mm_data`; item scheduling additionally implements `get_mm_encoder_item_metadata` + `get_max_mm_encoder_output_embeddings`. The model exposes batched `encode_multimodal_inputs`; encoder metadata capacity comes from `setup_attn_metadata`, not hardcoded `max_num_*` values. +- [ ] Encoder profiling implements `get_mm_max_tokens_per_item` + `get_dummy_mm_data`. The processor returns `MultimodalItemMetadata` from `get_mm_item_metadata`; item scheduling additionally supplies exact costs/output rows, item selection, and `get_max_mm_encoder_output_embeddings`. The model exposes batched `encode_multimodal_inputs`; encoder metadata capacity comes from `setup_attn_metadata`, not hardcoded `max_num_*` values. +- [ ] Batched/packed/shared fields have explicit `input_partitions`, or the model overrides selection for its layout. Payload indices are modality-local; metadata indices are prompt-global. Execution reads `item_metadata.item_refs`, not frontend `mm_item_order`. - [ ] `call_with_text_prompt` (not `__call__` — that's the base-class dispatcher) runs HF AutoProcessor + tokenizer, builds `multimodal_data` by modality, computes `mrope_config` on CPU, `_postprocess`-rewrites mm token ids to the OOV sentinel. - [ ] `mm_processor_kwargs` flow-through preserved. (Tokenized fast path is optional: set `supports_token_id_mm_expansion = True` + implement `get_text_with_mm_placeholders` / `expand_prompt_token_ids_for_mm`; otherwise the base class detokenizes token-ID inputs automatically.) - [ ] `_attach_multimodal_embeddings_impl` implemented (not the `attach_multimodal_embeddings` wrapper) if `@support_multimodal_disaggregated`. +- [ ] Handle-based E/P hooks accept keyword-only `mm_hashes` and return `(prompt_token_ids, MultimodalParams)`. Cache references stay on request Params; execution views do not own them. **Performance contracts** - [ ] Grep clean for Contract 1 bans (`.item()` / `.cpu()` / `.tolist()` / `torch.nonzero` / single-arg `torch.where` / value-dependent `if`, etc.) in modeling `forward` paths — elementwise `torch.where(cond, x, y)` with GPU-only `cond` is fine. diff --git a/docs/source/developer-guide/multimodal.md b/docs/source/developer-guide/multimodal.md index 6a66f32e6dc8..d9fbe5e6d064 100644 --- a/docs/source/developer-guide/multimodal.md +++ b/docs/source/developer-guide/multimodal.md @@ -94,6 +94,10 @@ Params order/length fallbacks have been removed. ## Onboarding a Model +Follow [Adding a New Model](../models/adding-new-model.md) for model definition, +weight loading, and registration. This section covers the additional multimodal +processor and execution contracts. + Start with full-request encoding and add item scheduling only when the model can encode its items independently. Use these contracts rather than inventing encoder costs or inferring item axes from matching tensor sizes. diff --git a/docs/source/features/multi-modality.md b/docs/source/features/multi-modality.md index a63fe14a8c47..dc272451c7a8 100644 --- a/docs/source/features/multi-modality.md +++ b/docs/source/features/multi-modality.md @@ -1,3 +1,8 @@ + + # Multimodal Support in TensorRT LLM TensorRT LLM supports a variety of multimodal models, enabling efficient inference with inputs beyond just text. @@ -26,6 +31,11 @@ Further optimizations are under development and will be updated as they become a Please refer to the latest multimodal [support matrix](../models/supported-models.md#multimodal-feature-support-matrix-pytorch-backend). +## Adding a Model + +See the [multimodal onboarding guide](../developer-guide/multimodal.md#onboarding-a-model) +for processor metadata, payload selection, model hooks, and encoder/prefill handoff. + ## Optional dependencies Video input decoding uses OpenCV (`cv2`) and is not installed by default. Install it only for models with video inputs: diff --git a/docs/source/models/adding-new-model.md b/docs/source/models/adding-new-model.md index 3a045e29198b..67c1371fa125 100644 --- a/docs/source/models/adding-new-model.md +++ b/docs/source/models/adding-new-model.md @@ -1,3 +1,8 @@ + + # Adding a New Model ## Table of Contents @@ -15,6 +20,10 @@ This guide provides a step-by-step process for adding a new model in PyTorch Backend. +For multimodal models, also follow the +[multimodal onboarding guide](../developer-guide/multimodal.md#onboarding-a-model) +for processor metadata, payload selection, encoder scheduling, and E/P handoff. + ## Prerequisites Before you begin, ensure you have the following: diff --git a/docs/source/torch/adding_new_model.md b/docs/source/torch/adding_new_model.md index 1028a1596a28..006941981f8e 100644 --- a/docs/source/torch/adding_new_model.md +++ b/docs/source/torch/adding_new_model.md @@ -1,3 +1,8 @@ + + # Adding a New Model in PyTorch Backend ## Table of Contents @@ -15,6 +20,10 @@ This guide provides a step-by-step process for adding a new model in PyTorch Backend. +For multimodal models, also follow the +[multimodal onboarding guide](../developer-guide/multimodal.md#onboarding-a-model) +for processor metadata, payload selection, encoder scheduling, and E/P handoff. + ## Prerequisites Before you begin, ensure you have the following: From 49bcbb6e986a2e1b0f82672804af0d0362070d97 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:51:30 -0700 Subject: [PATCH 16/22] [TRTLLM-14477][fix] preserve multimodal metadata during profiling and hashing Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../trtllm-model-onboard-multimodal/SKILL.md | 4 ++ docs/source/developer-guide/multimodal.md | 5 ++ .../models/modeling_multimodal_mixin.py | 11 +--- .../_torch/models/modeling_qwen2vl.py | 27 ++++++++-- tensorrt_llm/inputs/registry.py | 25 +++++---- .../test_multimodal_encoder_item_groups.py | 53 +++++++++++++++++++ tests/unittest/inputs/test_multimodal.py | 44 +++++++++++++++ 7 files changed, 146 insertions(+), 23 deletions(-) diff --git a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md index 87b20c55ff6a..83f00f710c6b 100644 --- a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md +++ b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md @@ -388,6 +388,10 @@ Subclass **both** `BaseMultimodalInputProcessor` (drives every real request) and - `get_mm_max_tokens_per_item(max_num_encoder_tokens=None) -> {modality: tokens}` — report bounded startup maxima when the argument is `None`, or the largest legal item per modality under a concrete runtime token budget. - `get_dummy_mm_data(*, max_num_encoder_tokens, mm_counts, dtype) -> multimodal_data` — build the profiler-selected per-modality item counts **directly** as processed encoder tensors (no PIL image + HF-processor round-trip). +For grouped encoders, include `MultimodalItemMetadata` with item references and +exact output rows in the returned dummy data. Profiling bypasses the normal +processor wrapper, so it does not call `get_mm_item_metadata` for you. + Following vLLM's separation of common budgeting from model processing geometry, the profiler selects the largest legal runtime workload and computes its repetition count under `encoder_max_num_tokens` and `encoder_max_batch_size`. It passes the resulting plan (for example, `{"image": 8}`) to the concrete processor's `get_dummy_mm_data`; the processor owns only geometry and tensor layout. Qwen builds `pixel_values`/`*_grid_thw`; Mistral builds `pixel_values`/`image_sizes` and uses its private `_vit_tokens` helper so the LLM-side Pixtral hashing count is unchanged. A model without the contract falls back to a text-only dummy (encoder memory unaccounted). Don't hardcode the encoder attention workspace (`max_num_*=8192`): inherit `MultimodalEncoderMixin` and let the engine size it via `setup_attn_metadata` at load. The workspace dimension comes from `encoder_max_num_tokens`, falling back to the LLM-side `max_num_tokens` when unset. `encoder_max_batch_size` independently bounds the number of atomic items in a multimodal encoder iteration and falls back to `max_batch_size`. The profiler resolves both limits into `mm_counts` before calling `get_dummy_mm_data`. Model-internal attention sequence counts are not the item batch size; they are derived separately from the token budget and model geometry. diff --git a/docs/source/developer-guide/multimodal.md b/docs/source/developer-guide/multimodal.md index d9fbe5e6d064..af91b64e2e4a 100644 --- a/docs/source/developer-guide/multimodal.md +++ b/docs/source/developer-guide/multimodal.md @@ -109,6 +109,11 @@ encoder costs or inferring item axes from matching tensor sizes. | Item scheduling | Exact positive `encoder_token_lengths`, output row counts, and an item-selection implementation. | | Persistent item caching | Output row counts and stable content/processor hashes. Partial hits also require correct item selection. | +Startup profiling calls `get_dummy_mm_data` directly, bypassing the normal +processor wrapper. For grouped encoders, its returned data must include the +common `MultimodalItemMetadata` record with item references and exact output +row counts. + For example, a processor for a full-request audio encoder can return: ```python diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 70b06c306316..3d3bc305e537 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -223,16 +223,7 @@ def _reorder_embeds_by_manifest( for m, c in req_counts.items(): per_modality_cursor[m] += c if not slices: - # No items resolved for any request. This happens on the executor's - # KV-cache profiling pass: `_encode_dummy_inputs` runs the encoder on a - # worst-case dummy batch that carries the encoder tensors but no - # `multimodal_embedding_lengths`, so the per-modality lengths (and thus - # the sliced `per_modality_embeds`) come back empty. The encoder forward - # still ran (its activation is what peak-memory profiling captures), so - # return a correctly-typed empty embedding tensor instead of crashing on - # `torch.cat([])`. `per_modality_embeds` values are already zero-row - # slices of the encoder output, so their concat preserves dtype/device - # and the hidden dim. + # Preserve encoder dtype/device and width when no item contributes rows. if per_modality_embeds: return torch.cat(list(per_modality_embeds.values()), dim=0) return torch.empty(0) diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 928e8ce0e71c..7cda5d539670 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -30,9 +30,9 @@ from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.inputs.multimodal import (MultimodalItemMetadata, - MultimodalItemPartition, - MultimodalParams) +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalItemMetadata, + MultimodalItemPartition, MultimodalParams) from ..._utils import async_tensor_h2d, nvtx_range from ...inputs import (BaseMultimodalDummyInputsBuilder, @@ -810,6 +810,8 @@ def get_dummy_mm_data( mm_data: Dict[str, Any] = {} total_tokens = 0 + item_refs: list[tuple[str, int]] = [] + encoder_token_lengths: list[int] = [] for modality, num_items in mm_counts.items(): if modality not in ("image", "video"): raise ValueError(f"Unsupported vision modality: {modality}") @@ -826,7 +828,8 @@ def get_dummy_mm_data( raise ValueError( f"Qwen {modality} dummy does not fit within " f"max_num_encoder_tokens={max_num_encoder_tokens}") - total_tokens += num_items * math.prod(grid) + item_tokens = math.prod(grid) + total_tokens += num_items * item_tokens if total_tokens > max_num_encoder_tokens: raise ValueError( "Requested multimodal dummy items exceed " @@ -838,6 +841,22 @@ def get_dummy_mm_data( num_items=num_items, dtype=dtype, )) + item_refs.extend((modality, i) for i in range(num_items)) + encoder_token_lengths.extend([item_tokens] * num_items) + if item_refs: + metadata = MultimodalItemMetadata( + item_refs, + encoder_token_lengths, + [ + tokens // self.spatial_merge_unit + for tokens in encoder_token_lengths + ], + ) + for modality in mm_data: + metadata.input_partitions[ + modality] = _qwen_vision_input_partitions( + mm_data, modality, metadata) + mm_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata return mm_data def _dummy_mm_data_for_size( diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index c14eb083b0e7..2c1f17e2463b 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1355,13 +1355,21 @@ def multimodal_hashing_process( m for m, lengths in num_mm_tokens_by_key.items() if lengths ] if len(present) != 1: - raise ValueError( - "Mixed multimodal inputs require prompt-ordered item metadata" - ) - modality = present[0] - item_refs = tuple( - (modality, i) - for i in range(len(num_mm_tokens_by_key[modality]))) + get_item_metadata = getattr(input_processor, + "get_mm_item_metadata", None) + metadata = (get_item_metadata(prompt_token_ids, multimodal_data, + None) + if get_item_metadata is not None else None) + if metadata is None: + raise ValueError( + "Mixed multimodal inputs require prompt-ordered item metadata" + ) + item_refs = metadata.item_refs + else: + modality = present[0] + item_refs = tuple( + (modality, i) + for i in range(len(num_mm_tokens_by_key[modality]))) num_mm_tokens = [num_mm_tokens_by_key[m][i] for m, i in item_refs] if len(num_mm_tokens) <= 0: raise ValueError("multimodal hashing produced an empty multimodal " @@ -1403,8 +1411,7 @@ def multimodal_hashing_process( item_refs, output_embedding_lengths=multimodal_embedding_lengths) multimodal_data[ MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = _merge_mm_item_metadata( - prepared_metadata, - get_multimodal_item_metadata(multimodal_data)) + prepared_metadata, metadata) # Store special token offsets if available if len(start_special_token_positions ) > 0 and mm_special_token_ids is not None: diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index cff43fef85ae..1f668b3dd6f1 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -365,6 +365,59 @@ def test_grouped_encoder_rejects_raw_payload_without_item_layout(empty_record: b assert model.encoder_calls == [] +@pytest.mark.parametrize( + "mm_counts,expected_refs,expected_costs,expected_rows", + [ + ({"image": 2}, (("image", 0), ("image", 1)), (8, 8), (2, 2)), + ({"video": 2}, (("video", 0), ("video", 1)), (16, 16), (4, 4)), + ( + {"video": 1, "image": 2}, + (("video", 0), ("image", 0), ("image", 1)), + (16, 8, 8), + (4, 2, 2), + ), + ], + ids=["images", "videos", "mixed"], +) +def test_qwen_profiling_dummy_supports_grouped_encoding( + mm_counts: dict[str, int], + expected_refs: tuple[tuple[str, int], ...], + expected_costs: tuple[int, ...], + expected_rows: tuple[int, ...], +) -> None: + from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLInputProcessorBase + + class Processor(Qwen3VLInputProcessorBase): + def _get_dummy_grid_for_modality( + self, modality: str, max_num_encoder_tokens: int + ) -> tuple[int, int, int]: + return {"image": (1, 2, 4), "video": (2, 2, 4)}[modality] + + processor = object.__new__(Processor) + processor._config = SimpleNamespace( + vision_config=SimpleNamespace(spatial_merge_size=2, patch_size=1, in_channels=HIDDEN) + ) + data = processor.get_dummy_mm_data( + max_num_encoder_tokens=64, mm_counts=mm_counts, dtype=torch.float16 + ) + params = MultimodalParams(multimodal_data=data) + metadata = params.item_metadata + assert metadata is not None + metadata.validate() + assert metadata.item_refs == expected_refs + assert metadata.encoder_token_lengths == expected_costs + assert metadata.output_embedding_lengths == expected_rows + + group = EncoderGroup( + modalities=("image", "video"), + encoder_fn=lambda pixel_values: pixel_values[::4], + build_batched_input=_GroupedEncoderModel._build_batched_input, + ) + output = encode_multimodal_by_groups([group], [params]) + assert output.shape == (sum(expected_rows), HIDDEN) + assert output.dtype == torch.float16 + + def test_common_metadata_shares_executor_and_encoder_layout(): from tensorrt_llm._torch.pyexecutor.llm_request import get_multimodal_embedding_lengths diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index 448f96d8e98e..cfa53738d627 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -603,6 +603,50 @@ def get_mm_special_token_ids(self): } +def test_mixed_hashing_uses_processor_metadata_without_frontend_order() -> None: + class Processor(_KwargsHashFakeProcessor): + multimodal_hashing_supported = None + + def __call__(self, inputs: dict, sampling_params: object) -> tuple[list[int], dict]: + return [10, 101, 102, 20, 101, 102, 30], {"multimodal_data": {"image": {}, "video": {}}} + + def get_num_tokens_per_video(self, *, video: VideoData, **kwargs: object) -> int: + return 2 + + def get_mm_item_metadata( + self, + prompt_token_ids: list[int], + multimodal_data: dict, + mm_item_order: list[dict[str, object]] | None = None, + ) -> MultimodalItemMetadata: + return MultimodalItemMetadata([("video", 0), ("image", 0)], [16, 8], [2, 2]) + + processor = Processor() + wrapped = create_input_processor_with_hash(processor) + inputs = { + "prompt": "unused", + "multi_modal_data": { + "image": [np.array([1], dtype=np.uint8)], + "video": [VideoData(frames=[np.full((2, 2, 3), 2, dtype=np.uint8)], metadata={})], + }, + } + + # Exercise both the first hashing probe and a subsequent enabled request. + for _ in range(2): + _, extra = wrapped(inputs, None) + assert processor.multimodal_hashing_supported is True + mm_input = extra["multimodal_input"] + assert mm_input.multimodal_hashes == [ + _CONTENT_HASH_SNAPSHOT["video"], + _CONTENT_HASH_SNAPSHOT["image"], + ] + assert mm_input.multimodal_positions == [1, 4] + metadata = extra["multimodal_data"][MULTIMODAL_ENCODER_ITEM_METADATA_KEY] + assert metadata.item_refs == (("video", 0), ("image", 0)) + assert metadata.encoder_token_lengths == (16, 8) + assert metadata.output_embedding_lengths == (2, 2) + + @pytest.mark.parametrize("hashing_enabled", [False, True]) def test_mm_item_metadata_rejects_conflicting_prompt_order(hashing_enabled: bool) -> None: class Processor(_KwargsHashFakeProcessor): From 30ddec4d74907db7ca2824b036de6d7c8c2536a9 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:19:05 -0700 Subject: [PATCH 17/22] [TRTLLM-14477][perf] avoid multimodal metadata host syncs Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_gemma4mm.py | 28 ++++++++++++--- .../models/modeling_nemotron_h_multimodal.py | 21 +++++++++--- .../_torch/models/modeling_parakeet.py | 20 +++++++++-- tensorrt_llm/inputs/multimodal.py | 7 ++++ .../test_modeling_nemotron_h_multimodal.py | 10 +++++- .../_torch/modeling/test_modeling_parakeet.py | 29 +++++++++++++--- .../test_multimodal_encoder_item_groups.py | 34 +++++++++++++++++++ tests/unittest/inputs/test_multimodal.py | 19 +++++++++++ 8 files changed, 150 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index fc16dfe66931..d6a886d4a38b 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -946,7 +946,6 @@ def multimodal_data_device_paths(self) -> List[str]: "video.pixel_values", "video.image_position_ids", "audio.audio_features", - "audio.audio_features_mask", ] # TODO(TRTLLM-14981): Implement finer-grained caching for videos. @@ -1536,11 +1535,21 @@ def _get_audio_features( that it can be concatenated with image features before being fused into the LLM input embeddings. """ + num_audio_tokens = None + if ( + audio_features_mask is not None + and audio_features_mask.device.type == "cpu" + and audio_features_mask.dtype == torch.bool + ): + # The audio tower samples the mask through two stride-2 stages. + # Read its output size before H2D, avoiding CUDA nonzero's host sync. + num_audio_tokens = int(audio_features_mask[:, ::4].sum().item()) target_dtype = self.embed_audio.embedding_projection.weight.dtype with torch.autocast(device_type="cuda", dtype=self.model_dtype): if audio_features_mask is not None: output = self.audio_tower( - audio_features, attention_mask=audio_features_mask.to(audio_features.device) + audio_features, + attention_mask=audio_features_mask.to(audio_features.device, non_blocking=True), ) else: output = self.audio_tower(audio_features) @@ -1548,8 +1557,19 @@ def _get_audio_features( projected = self.embed_audio(hidden.to(target_dtype)) # (B, T, H_text) # Optionally drop padding frames reported by the audio tower. tower_mask = getattr(output, "attention_mask", None) - if tower_mask is not None and tower_mask.dtype == torch.bool: - projected = projected[tower_mask] + if ( + tower_mask is not None + and tower_mask.dtype == torch.bool + and audio_features_mask is not None + ): + if num_audio_tokens is None: + # Direct callers may already have placed the mask on CUDA. + projected = projected[tower_mask] + else: + indices = torch.nonzero_static( + tower_mask.reshape(-1), size=num_audio_tokens + ).flatten() + projected = projected.reshape(-1, projected.shape[-1]).index_select(0, indices) else: projected = projected.reshape(-1, projected.shape[-1]) return projected.contiguous() diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index e1708eb4eff9..ef0608859a6f 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -1168,13 +1168,16 @@ def get_mm_item_metadata( clip_counts = audio.get("audio_num_clips") if clip_counts is None: return metadata - clips = tuple(clip_counts.cpu().tolist()) + clip_counts = clip_counts.cpu() + audio["audio_num_clips"] = clip_counts + clips = tuple(clip_counts.tolist()) item_count = sum(m == "audio" for m, _ in metadata.item_refs) if len(clips) == item_count: packed = MultimodalItemPartition(item_lengths=clips) metadata.input_partitions["audio"] = { "input_audio_features": packed, "feature_attention_mask": packed, + "audio_output_lengths": packed, "audio_num_clips": MultimodalItemPartition(), } return metadata @@ -3284,9 +3287,15 @@ def _encode_audio(self, audio_data_list: List[dict]) -> List[Tuple[torch.Tensor, padded_features: List[torch.Tensor] = [] padded_masks: List[torch.Tensor] = [] clips_per_input: List[int] = [] + valid_output_lens: List[int] = [] + has_output_lengths = all( + ad.get("audio_output_lengths") is not None for ad in audio_data_list + ) for audio_data in audio_data_list: features = audio_data["input_audio_features"].to(dtype=self.model_dtype) mask = audio_data["feature_attention_mask"] + if has_output_lengths: + valid_output_lens.extend(audio_data["audio_output_lengths"].tolist()) pad_amount = max_time - features.shape[1] if pad_amount > 0: features = torch.nn.functional.pad(features, (0, 0, 0, pad_amount)) @@ -3300,10 +3309,12 @@ def _encode_audio(self, audio_data_list: List[dict]) -> List[Tuple[torch.Tensor, sound_embeds = self.sound_encoder(all_features, all_masks) - valid_input_lens = all_masks.sum(dim=1) - valid_output_lens = self.sound_encoder.encoder._get_subsampling_output_length( - valid_input_lens - ).tolist() + if not has_output_lengths: + # Compatibility for feature dicts not built by ParakeetExtractor. + # Keep one batched readback rather than synchronizing per input. + valid_output_lens = self.sound_encoder.encoder._get_subsampling_output_length( + all_masks.sum(dim=1) + ).tolist() per_clip_embeds: List[torch.Tensor] = [] per_clip_counts: List[int] = [] diff --git a/tensorrt_llm/_torch/models/modeling_parakeet.py b/tensorrt_llm/_torch/models/modeling_parakeet.py index 704e7c792911..2ed758b5848e 100644 --- a/tensorrt_llm/_torch/models/modeling_parakeet.py +++ b/tensorrt_llm/_torch/models/modeling_parakeet.py @@ -250,11 +250,11 @@ def __call__( audio_clips.extend(clips) audio_num_clips.append(len(clips)) - audio_lengths = torch.tensor( + audio_lengths_cpu = torch.tensor( [len(speech) for speech in audio_clips], dtype=torch.long, - device=device, ) + audio_lengths = audio_lengths_cpu.to(device, non_blocking=True) max_length = max(len(speech) for speech in audio_clips) input_features = self._pad_raw_speech(audio_clips, max_length, device) if self.config.preemphasis is not None: @@ -262,10 +262,24 @@ def __call__( input_features = self._torch_extract_fbank_features(input_features, device) input_features, attention_mask = self._normalize_mel_features(input_features, audio_lengths) + # Keep clip geometry on the host: the forward slices encoder outputs + # using these lengths, even when feature extraction runs on CUDA. + feature_lengths = torch.floor_divide( + audio_lengths_cpu + self.config.n_fft // 2 * 2 - self.config.n_fft, + self.config.hop_length, + ).clamp_(min=0, max=input_features.shape[1]) + output_lengths = _subsampling_output_length( + feature_lengths, + self.config.subsampling_factor, + self.config.subsampling_conv_kernel_size, + self.config.subsampling_conv_stride, + ) + return { "input_audio_features": input_features, "feature_attention_mask": attention_mask, - "audio_num_clips": torch.tensor(audio_num_clips, dtype=torch.long, device=device), + "audio_num_clips": torch.tensor(audio_num_clips, dtype=torch.long), + "audio_output_lengths": output_lengths, } def audio_length(self, audio_tokens: int) -> int: diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 611f4be23824..933d43f2b1c7 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -95,6 +95,11 @@ def select(self, value: Any, item_indices: Sequence[int]) -> Any: for i in indices ] return torch.cat(pieces, dim=axis) + if value.is_cuda: + # Python-list indexing copies an index tensor to CUDA and + # synchronizes the host. These host-known spans need no indices. + return torch.cat([value.narrow(axis, i, 1) for i in indices], + dim=axis) # With no packed lengths, tensor indexing selects one row per item. if axis == 0: return value[indices] @@ -801,6 +806,8 @@ def __post_init__(self): # Extend only after auditing each key's consumers. _CPU_ONLY_MULTIMODAL_DATA_KEYS = frozenset({ "multimodal_embed_mask_cumsum", + "audio_num_clips", + "audio_output_lengths", MULTIMODAL_ENCODER_ITEM_METADATA_KEY, }) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py index cf61ac09c2e1..b9a7d45b535c 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py @@ -982,7 +982,8 @@ def _make_audio_data(self, num_clips, time_len, valid_lens): mask[i, :vl] = 1 return {"input_audio_features": features, "feature_attention_mask": mask} - def test_batched_matches_per_input(self): + @pytest.mark.parametrize("precomputed_lengths", [False, True]) + def test_batched_matches_per_input(self, precomputed_lengths): """Bucket output equals N singleton `_encode_audio` calls. Compares ``encode([a1, a2])`` against ``[encode([a1])[0], encode([a2])[0]]`` @@ -998,6 +999,12 @@ def test_batched_matches_per_input(self): # Two inputs with different time / clip counts. a1 = self._make_audio_data(num_clips=2, time_len=10, valid_lens=[10, 6]) a2 = self._make_audio_data(num_clips=1, time_len=14, valid_lens=[12]) + if precomputed_lengths: + a1["audio_output_lengths"] = torch.tensor([5, 3]) + a2["audio_output_lengths"] = torch.tensor([6]) + model.sound_encoder.encoder._get_subsampling_output_length = mock.Mock( + side_effect=AssertionError("Forward must reuse host output lengths") + ) per_input_results = [ NemotronHMultimodalModel._encode_audio(model, [a1])[0], @@ -1006,6 +1013,7 @@ def test_batched_matches_per_input(self): batched_results = NemotronHMultimodalModel._encode_audio(model, [a1, a2]) assert len(batched_results) == 2 + assert [counts for _, counts in batched_results] == [[5, 3], [6]] for (b_emb, b_counts), (s_emb, s_counts) in zip(batched_results, per_input_results): assert b_counts == s_counts assert torch.allclose(b_emb, s_emb, atol=1e-6, rtol=1e-6) diff --git a/tests/unittest/_torch/modeling/test_modeling_parakeet.py b/tests/unittest/_torch/modeling/test_modeling_parakeet.py index 7fe4326c5149..1dd523d8b2e3 100644 --- a/tests/unittest/_torch/modeling/test_modeling_parakeet.py +++ b/tests/unittest/_torch/modeling/test_modeling_parakeet.py @@ -14,7 +14,11 @@ import pytest import torch -from tensorrt_llm._torch.models.modeling_parakeet import ParakeetExtractor, ProjectedParakeet +from tensorrt_llm._torch.models.modeling_parakeet import ( + ParakeetExtractor, + ProjectedParakeet, + _subsampling_output_length, +) def _make_extractor_config(**overrides): @@ -101,17 +105,32 @@ def test_split_audio_into_clips_preserves_total_samples(self): @pytest.mark.threadleak(enabled=False) def test_call_returns_expected_keys(self): ext = _make_extractor() - audios = [np.random.randn(16000).astype(np.float32)] + audios = [ + np.zeros(length, dtype=np.float32) for length in (16000, ext._clip_target_samples + 1) + ] result = ext(audios) assert "input_audio_features" in result assert "feature_attention_mask" in result assert "audio_num_clips" in result + assert "audio_output_lengths" in result assert result["input_audio_features"].ndim == 3 - assert result["input_audio_features"].shape[0] == 1 + assert result["input_audio_features"].shape[0] == sum(result["audio_num_clips"].tolist()) assert result["input_audio_features"].shape[2] == ext.config.feature_size assert result["feature_attention_mask"].shape == result["input_audio_features"].shape[:2] - assert result["audio_num_clips"].shape == (1,) - assert result["audio_num_clips"].item() >= 1 + assert result["audio_num_clips"].shape == (2,) + assert result["audio_num_clips"].tolist() == [1, 2] + assert result["audio_num_clips"].device.type == "cpu" + assert result["audio_output_lengths"].device.type == "cpu" + expected_lengths = _subsampling_output_length( + result["feature_attention_mask"].sum(dim=1), + ext.config.subsampling_factor, + ext.config.subsampling_conv_kernel_size, + ext.config.subsampling_conv_stride, + ) + torch.testing.assert_close(result["audio_output_lengths"], expected_lengths) + assert result["audio_output_lengths"].sum().item() == sum( + ext.audio_token_count(len(audio)) for audio in audios + ) def test_call_accepts_stereo_audio(self): ext = _make_extractor() diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 1f668b3dd6f1..88c033f87d80 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -240,6 +240,7 @@ def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): "input_audio_features": torch.arange(30).reshape(3, 5, 2), "feature_attention_mask": torch.ones(3, 5), "audio_num_clips": torch.tensor([2, 1]), + "audio_output_lengths": torch.tensor([1, 3, 2]), } data = { "audio": audio, @@ -259,6 +260,8 @@ def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): selected.multimodal_data["audio"]["input_audio_features"], audio["input_audio_features"][[2, 0, 1]], ) + assert selected.multimodal_data["audio"]["audio_num_clips"].tolist() == [1, 2] + assert selected.multimodal_data["audio"]["audio_output_lengths"].tolist() == [2, 1, 3] assert selected.encoder_state is None video = {"audio": audio, "video_size": [[2, 1, 32, 32]]} video_data = { @@ -273,6 +276,37 @@ def test_nemotron_audio_keeps_clip_groups_and_video_audio_relationship(): assert video_data["video"]["audio"] is audio +@pytest.mark.parametrize("with_mask", [False, True]) +def test_gemma_audio_compacts_valid_rows_in_input_order(with_mask: bool): + from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4ForConditionalGeneration + + hidden = torch.arange(40, dtype=torch.float32).reshape(2, 5, 4) + mask = torch.tensor( + [[True, False, True, True, False], [False, True, True, False, True]] + ).repeat_interleave(4, dim=1)[:, :17] + + class Projection(torch.nn.Identity): + embedding_projection = SimpleNamespace(weight=torch.zeros(1)) + + def tower(features, attention_mask=None): + del features + return SimpleNamespace( + last_hidden_state=hidden, + attention_mask=( + attention_mask[:, ::4] + if attention_mask is not None + else torch.ones(2, 5, dtype=torch.bool) + ), + ) + + model = SimpleNamespace(model_dtype=torch.bfloat16, audio_tower=tower, embed_audio=Projection()) + actual = Gemma4ForConditionalGeneration._get_audio_features( + model, torch.zeros(2, 17, 4), mask if with_mask else None + ) + expected = hidden[mask[:, ::4]] if with_mask else hidden.reshape(-1, 4) + torch.testing.assert_close(actual, expected) + + @pytest.mark.parametrize("full_prompt_ids", [False, True]) def test_nemotron_audio_skips_incomplete_or_prompt_coupled_partitions( full_prompt_ids: bool, diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index cfa53738d627..2420f1ebd0ff 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -55,6 +55,25 @@ def test_batched_item_partition_selects_items_and_preserves_views( assert shares_storage is is_view +@pytest.mark.parametrize("targeted", [False, True]) +def test_audio_geometry_stays_on_cpu_during_nested_device_staging(targeted: bool) -> None: + audio = { + "input_audio_features": torch.zeros(3, 4, 2), + "audio_num_clips": torch.tensor([2, 1]), + "audio_output_lengths": torch.tensor([4, 2, 3]), + } + params = MultimodalParams(multimodal_data={"video": {"audio": audio}}) + targets = [f"video.audio.{name}" for name in audio] if targeted else None + + params.to_device("multimodal_data", "meta", target_keywords=targets) + + staged = params.multimodal_data["video"]["audio"] + assert staged["input_audio_features"].device.type == "meta" + for name in ("audio_num_clips", "audio_output_lengths"): + assert staged[name] is audio[name] + assert staged[name].device.type == "cpu" + + @pytest.mark.parametrize("axis", [0, 1, -1]) @pytest.mark.parametrize( ("indices", "is_view"), From 87681f42208f3043335e373170170b3f420fc809 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:24:51 -0700 Subject: [PATCH 18/22] [TRTLLM-14477][chore] align multimodal docs and imports after rebase Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/llm_request.py | 6 +++--- tensorrt_llm/inputs/multimodal.py | 8 +++++++- tests/unittest/_torch/executor/multimodal_utils.py | 7 +++++-- .../unittest/_torch/executor/test_multimodal_scheduler.py | 4 ++-- 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 7267fee2f949..a101ce4e9abb 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -4,7 +4,8 @@ from copy import copy, deepcopy from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Union, cast +from typing import (TYPE_CHECKING, Any, Dict, Iterable, List, Optional, + Sequence, Union, cast) import torch @@ -13,8 +14,7 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.bindings import executor as tllm_executor from tensorrt_llm.executor.result import SimpleTokenLogprobs, TokenLogprobs -from tensorrt_llm.inputs.multimodal import (MultimodalEncoderProgress, - MultimodalEncoderRequestState, +from tensorrt_llm.inputs.multimodal import (MultimodalEncoderRequestState, MultimodalParams, get_multimodal_item_metadata) from tensorrt_llm.sampling_params import LogprobMode diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 933d43f2b1c7..b6358a6a5184 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -295,7 +295,13 @@ def get_multimodal_item_metadata( class MultimodalEncoderProgress(Enum): - """Readiness derived from the request's cached item outputs.""" + """Python-only progress derived from request-local MM item outputs. + + It summarizes the ready flags of all items. Scheduling is gated by the + items the remaining prefill uses instead (`get_mm_pending_items`), because + items behind confirmed prefill progress or a reusable KV prefix are + released or never encoded. + """ PENDING = auto() """No MM item of the request has an encoder output yet.""" diff --git a/tests/unittest/_torch/executor/multimodal_utils.py b/tests/unittest/_torch/executor/multimodal_utils.py index 7cb2d40d4a2e..1465a02f50a4 100644 --- a/tests/unittest/_torch/executor/multimodal_utils.py +++ b/tests/unittest/_torch/executor/multimodal_utils.py @@ -21,8 +21,11 @@ ) from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.bindings import SamplingConfig -from tensorrt_llm.inputs.multimodal import MULTIMODAL_ENCODER_ITEM_METADATA_KEY -from tensorrt_llm.inputs.registry import BaseMultimodalDummyInputsBuilder, MultimodalItemMetadata +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, +) +from tensorrt_llm.inputs.registry import BaseMultimodalDummyInputsBuilder def bare_mm_item_scheduler( diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py index 60bbea141dc5..c8c72c7df438 100644 --- a/tests/unittest/_torch/executor/test_multimodal_scheduler.py +++ b/tests/unittest/_torch/executor/test_multimodal_scheduler.py @@ -20,7 +20,6 @@ from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( - MultimodalEncoderProgress, MultimodalEncoderRequestError, MultimodalEncoderRequestState, get_mm_items_for_chunk, @@ -44,6 +43,7 @@ from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalEncoderProgress, MultimodalItemMetadata, MultimodalParams, strip_mm_encoder_inputs, @@ -728,7 +728,7 @@ def test_admission_checks_only_outputs_resident_at_once(context_chunk_unit_size, if admitted: initialize_multimodal_encoder_request(request, **budget) - assert request.py_mm_encoder_state.embedding_lengths == [1, 1] + assert request.py_mm_encoder_state.embedding_lengths == (1, 1) else: with pytest.raises(ValueError, match="raise encoder_max_num_tokens") as error: initialize_multimodal_encoder_request(request, **budget) From 110a6bfd39599b7b1050974c19da22f5071d0ec5 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:27:42 -0700 Subject: [PATCH 19/22] [TRTLLM-14477][fix] preserve GLM profiling item metadata Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- docs/source/release-notes.md | 11 +++- .../models/modeling_glm5_next_vision.py | 17 +++++- .../modeling/test_glm5_next_contracts.py | 54 +++++++++++++++++++ 3 files changed, 78 insertions(+), 4 deletions(-) diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index d7db1a83c212..ab9b2a3366e9 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -13,13 +13,20 @@ All published functionality in the Release Notes has been fully tested and verif ### API Changes - **[BREAKING CHANGE] Removed obsolete C++ executor orchestration helpers.** The internal `RequestWithId` type, its `serializeReqWithIds` / `deserializeReqWithIds` functions, all production `insertRequestInOrder` overloads, and the unused `requestUtils.h` and `orchestratorUtils.h` helpers have been removed. These helpers belonged to the retired C++ executor orchestration path; downstream code linking their exported symbols must stop using them and rebuild. The current `Request` / `Response` interfaces and serialization remain available. Executor factory and worker constructors no longer accept the unused `engine` and `batched_logits_processor` arguments; callers should omit them and continue passing model configuration through `llm_args` and `hf_model_dir`. Unused decoder, request-cancellation, executor, legacy-model, orchestration, and speculative-decoding MPI tag names have also been removed; values 1–23, 127–129, and 1023–1025 are retired. The default, KV-cache event, and context-transfer tags retain their existing values. Internal `LogprobParams` and `LogProbsResult` tuples now contain only prompt-logprob fields; generation logprobs remain supplied by the sampler. The unused `TLLM_EXECUTOR_PERIODICAL_RESP_IN_AWAIT` environment variable has been removed. -- Multimodal processor extensions now use `get_mm_item_metadata()` and +- **[BREAKING CHANGE] Unified multimodal processor metadata.** + Multimodal processor extensions now use `get_mm_item_metadata()` and `MultimodalItemMetadata`, replacing `get_mm_encoder_item_metadata()` and `MultimodalEncoderItemMetadata`. The removed `DisaggPrefillMultimodalInputs` result is replaced by `(prompt_token_ids, MultimodalParams)` from `build_disagg_prefill_multimodal_inputs()`, which also accepts the keyword-only `mm_hashes` argument. Custom processors must migrate these hooks; processors - defining the removed metadata hook now fail at wrapper creation. See the + defining the removed metadata hook now fail at wrapper creation. + Also removed: `MultimodalParams(mm_item_order=...)`, the + `multimodal_embedding_lengths` and `layout_metadata` keys in `multimodal_data`, + `registry.get_multimodal_encoder_item_metadata`, and `LlmRequest.py_mm_item_order`. + Read item order and output lengths through `MultimodalParams.item_metadata`. + Grouped encoders must also include item metadata with exact output row counts + in `get_dummy_mm_data()` results used for startup profiling. See the [multimodal developer guide](developer-guide/multimodal.md). ### Fixed Issues diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py index c059c628a00c..74ee8e7b25b9 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py @@ -44,7 +44,11 @@ TextPrompt, register_input_processor, ) -from ...inputs.multimodal import MultimodalParams +from ...inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, + MultimodalParams, +) from ...inputs.registry import BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor from ...logger import logger from ...sampling_params import SamplingParams @@ -904,11 +908,20 @@ def get_dummy_mm_data( pixels_side = side * self._patch_size image = Image.new("RGB", (pixels_side, pixels_side), (127, 127, 127)) out = self.processor.image_processor(images=[image] * num_images, return_tensors="pt") + encoder_token_lengths = out["image_grid_thw"].prod(dim=-1).tolist() + metadata = MultimodalItemMetadata( + item_refs=[("image", i) for i in range(num_images)], + encoder_token_lengths=encoder_token_lengths, + output_embedding_lengths=[ + length // self.spatial_merge_unit for length in encoder_token_lengths + ], + ) return { "image": { "pixel_values": out["pixel_values"].to(dtype or self._dtype), "image_grid_thw": out["image_grid_thw"], - } + }, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: metadata, } # -- request processing ---------------------------------------------------- diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index e929b6b09790..17612840506c 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -235,6 +235,60 @@ def test_mixed_processor_matches_separate_inputs(tensor_images): ) +@pytest.mark.cpu_only +@pytest.mark.parametrize("num_images", [1, 2]) +def test_profiling_dummy_has_encoder_item_metadata(num_images: int) -> None: + from tensorrt_llm._torch.models.modeling_glm5_next_vision import ( + Glm5NextInputProcessor, + _glm5_next_build_batched_input, + ) + from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( + EncoderGroup, + MultimodalModelMixin, + encode_multimodal_by_groups, + ) + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + + processor = Glm5NextInputProcessor("unused", _config(), tokenizer=Mock()) + image_processor = Mock( + return_value={ + "pixel_values": torch.zeros(num_images * 16, 12), + "image_grid_thw": torch.tensor([[1, 4, 4]] * num_images), + }, + max_image_tokens=4, + merge_size=2, + ) + processor._processor = SimpleNamespace(image_processor=image_processor) + model = Mock(spec=MultimodalModelMixin, dtype=torch.float16) + profiler = object.__new__(KvCacheCreator) + profiler._llm_args = SimpleNamespace(disable_mm_encoder=False) + profiler._profiling_stage_data = {"enable_mm_reqs": True} + profiler._model_engine = SimpleNamespace( + model=model, + input_processor=processor, + encoder_max_num_tokens=num_images * 16, + encoder_batch_size=num_images, + ) + + params = profiler._create_dummy_encoder_inputs() + + def encode(pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + assert grid_thw.shape == (num_images, 3) + return pixel_values.new_zeros((pixel_values.shape[0] // 4, 2)) + + output = encode_multimodal_by_groups( + (EncoderGroup(("image", "video"), encode, _glm5_next_build_batched_input),), params + ) + assert output.shape == (num_images * 4, 2) + assert output.dtype == torch.float16 + metadata = params[0].item_metadata + assert metadata is not None + metadata.validate() + assert metadata.item_refs == tuple(("image", i) for i in range(num_images)) + assert metadata.encoder_token_lengths == (16,) * num_images + assert metadata.output_embedding_lengths == (4,) * num_images + + def _materialize_meta(module: torch.nn.Module) -> None: """Allocate meta tensors on CUDA as ModelLoader does after MetaInitMode construction.""" module._apply(lambda t: torch.empty_like(t, device="cuda") if t.is_meta else t) From 30df1a39e7a3d40835982b0202b2b752d47e3a8c Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Fri, 23 Oct 2026 23:12:36 -0700 Subject: [PATCH 20/22] fix: preserve Gemma4 item scheduling after metadata unification Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_gemma4_unified.py | 21 ++++- .../_torch/models/modeling_gemma4mm.py | 91 ++++++------------- tensorrt_llm/inputs/multimodal.py | 15 ++- .../_torch/modeling/test_gemma4_multimodal.py | 82 ++++++++++------- .../test_multimodal_encoder_item_groups.py | 14 +-- 5 files changed, 113 insertions(+), 110 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py index e04b3a8d44c7..00148e5efd4b 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py @@ -71,12 +71,13 @@ from ...inputs import ( BaseMultimodalDummyInputsBuilder, + BaseMultimodalInputProcessor, ContentFormat, - MultimodalEncoderItemMetadata, MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, register_input_processor, ) +from ...inputs.multimodal import MultimodalItemMetadata from ..modules.layer_norm import LayerNorm from ..modules.linear import Linear from .modeling_gemma4 import Gemma4ForCausalLM @@ -84,6 +85,7 @@ Gemma4InputProcessor, Gemma4MultimodalEmbedder, Gemma4MultimodalModelBase, + _gemma4_input_partitions, ) from .modeling_utils import ModelConfig, filter_weights, register_auto_model @@ -217,10 +219,19 @@ def get_mm_max_tokens_per_item( self, max_num_encoder_tokens ) - def get_mm_encoder_item_metadata( - self, prompt_token_ids: list[int], multimodal_data: dict - ) -> MultimodalEncoderItemMetadata | None: - return None + def get_mm_item_metadata( + self, + prompt_token_ids: list[int], + multimodal_data: dict, + mm_item_order: list[dict] | None = None, + ) -> MultimodalItemMetadata | None: + """Keep request metadata without the tower-specific encoder costs.""" + metadata = BaseMultimodalInputProcessor.get_mm_item_metadata( + self, prompt_token_ids, multimodal_data, mm_item_order + ) + if metadata is not None: + metadata.input_partitions.update(_gemma4_input_partitions(multimodal_data, metadata)) + return metadata @register_auto_model("Gemma4UnifiedForConditionalGeneration") diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index d6a886d4a38b..71132d01513f 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -23,7 +23,7 @@ import copy import dataclasses import math -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from itertools import accumulate, groupby from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple @@ -42,7 +42,6 @@ BaseMultimodalInputProcessor, ContentFormat, ExtraProcessedInputs, - MultimodalEncoderItemMetadata, MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, TextPrompt, @@ -215,18 +214,22 @@ def _gemma4_input_partitions( multimodal_data: dict[str, Any], metadata: MultimodalItemMetadata, ) -> dict[str, dict[str, MultimodalItemPartition]]: - """Declare item-major fields; video frame grouping stays model-owned.""" + """Declare item-major inputs and the packed frame spans of whole videos.""" partitions = {} for modality, fields in ( ("image", ("pixel_values", "image_position_ids", "image_seq_lens")), + ("video", ("pixel_values", "image_position_ids", "image_seq_lens")), ("audio", ("audio_features", "audio_features_mask")), ): data = multimodal_data.get(modality) if not isinstance(data, dict): continue item_count = sum(m == modality for m, _ in metadata.item_refs) + frame_counts = data.get("num_frames") if modality == "video" else None + item_lengths = tuple(frame_counts) if frame_counts is not None else None + field_count = sum(item_lengths) if item_lengths is not None else item_count value = data.get(fields[0]) - if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] == item_count: + if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] == field_count: modality_partitions = {} for name in fields: value = data.get(name) @@ -238,11 +241,13 @@ def _gemma4_input_partitions( size = len(value) else: size = None - if size != item_count: + if size != field_count: raise ValueError( - f"{modality}.{name} must have {item_count} items on its leading axis" + f"{modality}.{name} must have {field_count} entries on its leading axis" ) - modality_partitions[name] = MultimodalItemPartition() + modality_partitions[name] = MultimodalItemPartition(item_lengths=item_lengths) + if frame_counts is not None: + modality_partitions["num_frames"] = MultimodalItemPartition() partitions[modality] = modality_partitions return partitions @@ -255,17 +260,6 @@ class Gemma4InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInpu processing using the image processor saved in the model directory. """ - def get_mm_item_metadata( - self, - prompt_token_ids: List[int], - multimodal_data: Dict[str, Any], - mm_item_order: Optional[List[Dict[str, Any]]] = None, - ) -> Optional[MultimodalItemMetadata]: - metadata = super().get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) - if metadata is not None: - metadata.input_partitions.update(_gemma4_input_partitions(multimodal_data, metadata)) - return metadata - # Default class-level fallback. Real value computed per-instance below # from text_config.use_bidirectional_attention. Only 26B/31B set # use_bidirectional_attention="vision" — their image blocks need intact @@ -443,25 +437,26 @@ def _audio_tokens_per_item(self) -> Optional[int]: audio_seq_length = getattr(self._processor, "audio_seq_length", None) return None if audio_seq_length is None else int(audio_seq_length) - def get_mm_encoder_item_metadata( + def get_mm_item_metadata( self, prompt_token_ids: List[int], multimodal_data: Dict[str, Any], - ) -> Optional[MultimodalEncoderItemMetadata]: + mm_item_order: Optional[List[Dict[str, Any]]] = None, + ) -> Optional[MultimodalItemMetadata]: """Return one encoder item per image, video, or audio clip. An item costs its padded encoder input, because the vision tower's per-image tensors grow with the padded patch count rather than the valid one. It yields one embedding row per pooled soft token, or per subsampled audio frame, that covers valid input. """ - del prompt_token_ids + prepared = super().get_mm_item_metadata(prompt_token_ids, multimodal_data, mm_item_order) modalities = [ modality for modality in ("image", "video", "audio") if isinstance(multimodal_data.get(modality), dict) ] if not modalities: - return None + return prepared if len(modalities) > 1: raise ValueError( "Gemma4 requests containing multiple multimodal input types are not supported." @@ -475,11 +470,20 @@ def get_mm_encoder_item_metadata( encoder_token_lengths, output_embedding_lengths = self._vision_item_lengths( multimodal_data[modality] ) - return MultimodalEncoderItemMetadata( - item_refs=[(modality, item_idx) for item_idx in range(len(encoder_token_lengths))], - encoder_token_lengths=encoder_token_lengths, - output_embedding_lengths=output_embedding_lengths, + if mm_item_order: + item_refs = [(entry["modality"], entry["index"]) for entry in mm_item_order] + elif prepared is not None: + item_refs = prepared.item_refs + else: + item_refs = [(modality, item_idx) for item_idx in range(len(encoder_token_lengths))] + metadata = MultimodalItemMetadata( + item_refs=item_refs, + encoder_token_lengths=[encoder_token_lengths[i] for _, i in item_refs], + output_embedding_lengths=[output_embedding_lengths[i] for _, i in item_refs], + input_partitions=prepared.input_partitions if prepared is not None else None, ) + metadata.input_partitions.update(_gemma4_input_partitions(multimodal_data, metadata)) + return metadata def _vision_item_lengths(self, modality_data: Dict[str, Any]) -> Tuple[List[int], List[int]]: """Return the patch cost and soft-token count of every image or video. @@ -1018,41 +1022,6 @@ def get_multimodal_item_partitions( return None return _gemma4_input_partitions(param.multimodal_data, metadata).get(modality) - @staticmethod - def _build_video_encoder_input( - param: MultimodalParams, item_indices: Sequence[int] - ) -> MultimodalParams: - """Slice whole videos: each one is the run of frames that its `num_frames` entry declares.""" - video_data = param.multimodal_data["video"] - if not isinstance(video_data, dict): - raise TypeError( - f"multimodal_data['video'] must be a dict, got {type(video_data).__name__}" - ) - frames_per_item = video_data["num_frames"] - frame_starts = list(accumulate(frames_per_item, initial=0)) - frame_indices = [ - frame_idx - for item_idx in item_indices - for frame_idx in range(frame_starts[item_idx], frame_starts[item_idx + 1]) - ] - sliced = { - "pixel_values": video_data["pixel_values"][frame_indices], - "num_frames": [frames_per_item[item_idx] for item_idx in item_indices], - } - position_ids = video_data.get("image_position_ids") - if position_ids is not None: - sliced["image_position_ids"] = position_ids[frame_indices] - seq_lens = video_data.get("image_seq_lens") - if seq_lens is not None: - sliced["image_seq_lens"] = [seq_lens[frame_idx] for frame_idx in frame_indices] - residual_input = ( - copy.copy(param.multimodal_input) if param.multimodal_input is not None else None - ) - return MultimodalParams( - multimodal_data={**param.multimodal_data, "video": {**video_data, **sliced}}, - multimodal_input=residual_input, - ) - def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) -> torch.Tensor: """Encode uncached Gemma4 image, video, and audio payloads.""" modality_inputs = ( diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index b6358a6a5184..f762a33f9db2 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -41,7 +41,7 @@ class MultimodalItemPartition(NamedTuple): ``item_axis`` identifies the item axis of a batched tensor, or is ``None`` for a value shared by every item. Lists and tuples have item axis zero. - ``item_lengths`` partitions a packed tensor along that axis, for example + ``item_lengths`` partitions a packed input along that axis, for example patch counts for vision or frame counts for audio. Lengths are host-side metadata, independent of the payload tensor's device. """ @@ -107,8 +107,17 @@ def select(self, value: Any, item_indices: Sequence[int]) -> Any: index[axis] = indices return value[tuple(index)] - if (isinstance(value, (list, tuple)) and axis == 0 and lengths is None): - return type(value)(value[i] for i in item_indices) + if isinstance(value, (list, tuple)) and axis == 0: + if lengths is None: + return type(value)(value[i] for i in item_indices) + offsets = list(accumulate(lengths, initial=0)) + if offsets[-1] != len(value): + raise ValueError( + "Packed field size does not match its item lengths") + item_range = range(len(lengths)) + indices = [item_range[i] for i in item_indices] + return type(value)(element for i in indices + for element in value[offsets[i]:offsets[i + 1]]) raise TypeError( "Item fields must be tensors or lists/tuples on item axis zero") diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index cb2e1ab993a7..fa054b7c7fca 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -303,9 +303,8 @@ def _make_item_input_processor( def _with_item_metadata(multimodal_data: dict, processor: Gemma4InputProcessor) -> MultimodalParams: """Attach encoder item metadata the way the input-processor wrapper does.""" - metadata = processor.get_mm_encoder_item_metadata([], multimodal_data) + metadata = processor.get_mm_item_metadata([], multimodal_data) multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata - multimodal_data["multimodal_embedding_lengths"] = metadata.output_embedding_lengths return MultimodalParams(multimodal_data=multimodal_data) @@ -989,7 +988,10 @@ def test_build_multimodal_encoder_input_slices_whole_videos(self): "image_seq_lens": [10, 11, 12, 13, 14, 15], "num_frames": [2, 1, 3], }, - "multimodal_embedding_lengths": [1, 1, 1], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + [("video", 0), ("video", 1), ("video", 2)], + output_embedding_lengths=[1, 1, 1], + ), } ) @@ -1077,7 +1079,7 @@ def test_item_encoding_matches_full_request_encoding(self): }, processor, ) - lengths = param.multimodal_data["multimodal_embedding_lengths"] + lengths = param.get_item_layout()[1] expected = torch.split(model.encode_multimodal_inputs([param]), lengths) selected = [len(lengths) - 1, 0] @@ -1329,7 +1331,7 @@ def test_image_items_cost_padded_patches_and_yield_valid_soft_tokens(self): position_ids = torch.zeros(2, 2520, 2, dtype=torch.long) position_ids[1, 1260:] = -1 # The second image fills half of its patch slots. - metadata = processor.get_mm_encoder_item_metadata( + metadata = processor.get_mm_item_metadata( [], { "image": { @@ -1339,27 +1341,37 @@ def test_image_items_cost_padded_patches_and_yield_valid_soft_tokens(self): }, ) - self.assertEqual(metadata.item_refs, [("image", 0), ("image", 1)]) - self.assertEqual(metadata.encoder_token_lengths, [2520, 2520]) - self.assertEqual(metadata.output_embedding_lengths, [280, 140]) + self.assertEqual(list(metadata.item_refs), [("image", 0), ("image", 1)]) + self.assertEqual(list(metadata.encoder_token_lengths), [2520, 2520]) + self.assertEqual(list(metadata.output_embedding_lengths), [280, 140]) def test_video_items_span_their_frames(self): processor = _make_item_input_processor() + data = { + "video": { + "pixel_values": torch.zeros(5, 630, 1), + "image_seq_lens": [630, 630, 315, 315, 315], + "num_frames": [2, 3], + } + } + metadata = processor.get_mm_item_metadata([], data) - metadata = processor.get_mm_encoder_item_metadata( - [], - { - "video": { - "pixel_values": torch.zeros(5, 630, 1), - "image_seq_lens": [630, 630, 315, 315, 315], - "num_frames": [2, 3], - } - }, + self.assertEqual(list(metadata.item_refs), [("video", 0), ("video", 1)]) + self.assertEqual(list(metadata.encoder_token_lengths), [1260, 1890]) + self.assertEqual(list(metadata.output_embedding_lengths), [140, 105]) + self.assertEqual(metadata.input_partitions["video"]["image_seq_lens"].item_lengths, (2, 3)) + self.assertEqual( + metadata.input_partitions["video"]["image_seq_lens"].select( + data["video"]["image_seq_lens"], [1, 0] + ), + [315, 315, 315, 630, 630], ) - - self.assertEqual(metadata.item_refs, [("video", 0), ("video", 1)]) - self.assertEqual(metadata.encoder_token_lengths, [1260, 1890]) - self.assertEqual(metadata.output_embedding_lengths, [140, 105]) + reordered = processor.get_mm_item_metadata( + [], data, [{"modality": "video", "index": 1}, {"modality": "video", "index": 0}] + ) + self.assertEqual(reordered.item_refs, (("video", 1), ("video", 0))) + self.assertEqual(reordered.encoder_token_lengths, (1890, 1260)) + self.assertEqual(reordered.output_embedding_lengths, (105, 140)) def test_audio_items_count_subsampled_frames(self): """Each clip yields ceil(valid mel frames / 4) rows, the count HF's processor inserts.""" @@ -1368,19 +1380,19 @@ def test_audio_items_count_subsampled_frames(self): mask[0] = True mask[1, :5] = True - metadata = processor.get_mm_encoder_item_metadata( + metadata = processor.get_mm_item_metadata( [], {"audio": {"audio_features": torch.zeros(2, 10, 128), "audio_features_mask": mask}}, ) - self.assertEqual(metadata.item_refs, [("audio", 0), ("audio", 1)]) - self.assertEqual(metadata.encoder_token_lengths, [3, 3]) - self.assertEqual(metadata.output_embedding_lengths, [3, 2]) + self.assertEqual(list(metadata.item_refs), [("audio", 0), ("audio", 1)]) + self.assertEqual(list(metadata.encoder_token_lengths), [3, 3]) + self.assertEqual(list(metadata.output_embedding_lengths), [3, 2]) def test_mixed_modality_request_is_rejected(self): processor = _make_item_input_processor(audio_seq_length=750) with self.assertRaisesRegex(ValueError, "multiple multimodal input types"): - processor.get_mm_encoder_item_metadata( + processor.get_mm_item_metadata( [], { "image": {"pixel_values": torch.zeros(1, 2520, 1)}, @@ -1415,8 +1427,8 @@ def test_video_beyond_the_default_frame_count_fills_the_budget(self): ) dummy = processor.get_dummy_mm_data(max_num_encoder_tokens=budget, mm_counts={"video": 1}) self.assertEqual(dummy["video"]["num_frames"], [64]) - metadata = processor.get_mm_encoder_item_metadata([], dummy) - self.assertEqual(metadata.encoder_token_lengths, [budget]) + metadata = processor.get_mm_item_metadata([], dummy) + self.assertEqual(list(metadata.encoder_token_lengths), [budget]) def test_encoder_iteration_bounds(self): processor = _make_item_input_processor() @@ -1442,12 +1454,14 @@ def test_dummy_items_have_the_largest_legal_size(self): max_num_encoder_tokens=budget, mm_counts={modality: num_items} ) - metadata = processor.get_mm_encoder_item_metadata([], dummy) + metadata = processor.get_mm_item_metadata([], dummy) - self.assertEqual(metadata.encoder_token_lengths, [item_tokens] * num_items) + self.assertEqual(list(metadata.encoder_token_lengths), [item_tokens] * num_items) # Every patch slot or mel frame holds valid input. rows_per_item = item_tokens if modality == "audio" else item_tokens // 9 - self.assertEqual(metadata.output_embedding_lengths, [rows_per_item] * num_items) + self.assertEqual( + list(metadata.output_embedding_lengths), [rows_per_item] * num_items + ) def test_dummy_rejects_items_beyond_the_budget(self): processor = _make_item_input_processor() @@ -1492,7 +1506,7 @@ def test_only_the_encoder_model_opts_into_item_scheduling(self): with self.assertRaises(NotImplementedError): unified.get_mm_max_tokens_per_item(max_num_encoder_tokens=8192) self.assertIsNone( - unified.get_mm_encoder_item_metadata([], {"image": {"pixel_values": torch.zeros(1)}}) + unified.get_mm_item_metadata([], {"image": {"pixel_values": torch.zeros(1)}}) ) @@ -1706,7 +1720,7 @@ def test_item_metadata_matches_prompt_soft_tokens(self): SamplingParams(max_tokens=10), ) - metadata = proc.get_mm_encoder_item_metadata(input_ids, mm_data["multimodal_data"]) + metadata = proc.get_mm_item_metadata(input_ids, mm_data["multimodal_data"]) self.assertEqual( metadata.output_embedding_lengths, @@ -1908,7 +1922,7 @@ def test_video_processing_returns_pixel_values(self): self.assertEqual(video_data["pixel_values"].dim(), 3) self.assertEqual(video_data["pixel_values"].shape[0], len(frames)) self.assertEqual(video_data["num_frames"], [len(frames)]) - metadata = proc.get_mm_encoder_item_metadata(input_ids, mm_data["multimodal_data"]) + metadata = proc.get_mm_item_metadata(input_ids, mm_data["multimodal_data"]) self.assertEqual(metadata.output_embedding_lengths, [predicted_video_tokens]) video_token_id = self._video_token_id() diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 88c033f87d80..3e04985b0a00 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -180,10 +180,10 @@ def get_mm_item_metadata(self, prompt_token_ids, multimodal_data, mm_item_order= ) -def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): - from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor +def test_gemma_unified_processor_describes_image_and_audio_without_encoder_costs(): + from tensorrt_llm._torch.models.modeling_gemma4_unified import Gemma4UnifiedInputProcessor - processor = object.__new__(Gemma4InputProcessor) + processor = object.__new__(Gemma4UnifiedInputProcessor) data = { "image": { "pixel_values": torch.zeros(1, 4, 6), @@ -212,12 +212,12 @@ def test_gemma_processor_describes_image_and_audio_without_encoder_costs(): ("audio", "audio_features_mask", torch.ones(2, 5)), ], ) -def test_gemma_processor_rejects_misaligned_sibling_items( +def test_gemma_unified_processor_rejects_misaligned_sibling_items( modality: str, field: str, value: torch.Tensor | list[int] ) -> None: - from tensorrt_llm._torch.models.modeling_gemma4mm import Gemma4InputProcessor + from tensorrt_llm._torch.models.modeling_gemma4_unified import Gemma4UnifiedInputProcessor - processor = object.__new__(Gemma4InputProcessor) + processor = object.__new__(Gemma4UnifiedInputProcessor) primary = "pixel_values" if modality == "image" else "audio_features" data = { modality: {primary: torch.zeros(1, 4, 6), field: value}, @@ -226,7 +226,7 @@ def test_gemma_processor_rejects_misaligned_sibling_items( ), } - with pytest.raises(ValueError, match=rf"{modality}\.{field} must have 1 items"): + with pytest.raises(ValueError, match=rf"{modality}\.{field} must have 1 "): processor.get_mm_item_metadata([], data) From 6938334c6dc97a3f5c2630aea5871557273039d0 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Sun, 25 Oct 2026 15:38:03 +0000 Subject: [PATCH 21/22] [TRTLLM-14477][fix] migrate Gemma4 and Cosmos3 Edge metadata tests Compare Gemma4 output lengths as sequences after MultimodalItemMetadata normalizes them to immutable tuples. The image and video assertions were still comparing tuples directly with lists. Update Cosmos3 Edge's cache fixture to publish common item metadata, unpack the inherited (prompt_token_ids, MultimodalParams) prefill result, and carry those Params through embedding fusion. Use the current encoder cache initializer directly. The model already inherits the unified processor and encoder contracts, so it needs no separate compatibility implementation. Validation: the exact two Gemma4 processor tests fail on the rebased code and pass on clean main and with these fixes, using small synthetic tokenizer and processor assets with real image/video preprocessing and no weights. The Cosmos3 Edge CPU suite changes from 7 passed / 1 failed to 8 passed; clean main also passes all 8. Full checkpoint-based accuracy is unverified. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/modeling/test_gemma4_multimodal.py | 4 +-- .../modeling/test_modeling_cosmos3_edge.py | 32 +++++++++++-------- 2 files changed, 21 insertions(+), 15 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index fa054b7c7fca..b89f9ea2681e 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -1723,7 +1723,7 @@ def test_item_metadata_matches_prompt_soft_tokens(self): metadata = proc.get_mm_item_metadata(input_ids, mm_data["multimodal_data"]) self.assertEqual( - metadata.output_embedding_lengths, + list(metadata.output_embedding_lengths), [proc.get_num_tokens_per_image(image=image) for image in images], ) self.assertEqual( @@ -1923,7 +1923,7 @@ def test_video_processing_returns_pixel_values(self): self.assertEqual(video_data["pixel_values"].shape[0], len(frames)) self.assertEqual(video_data["num_frames"], [len(frames)]) metadata = proc.get_mm_item_metadata(input_ids, mm_data["multimodal_data"]) - self.assertEqual(metadata.output_embedding_lengths, [predicted_video_tokens]) + self.assertEqual(list(metadata.output_embedding_lengths), [predicted_video_tokens]) video_token_id = self._video_token_id() if video_token_id is not None: diff --git a/tests/unittest/_torch/modeling/test_modeling_cosmos3_edge.py b/tests/unittest/_torch/modeling/test_modeling_cosmos3_edge.py index 9cfe5a64e354..2652cc51194e 100644 --- a/tests/unittest/_torch/modeling/test_modeling_cosmos3_edge.py +++ b/tests/unittest/_torch/modeling/test_modeling_cosmos3_edge.py @@ -29,7 +29,13 @@ ) from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLVisionPatchMerger from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType -from tensorrt_llm.inputs.multimodal import MultimodalInput, MultimodalParams, MultimodalRuntimeData +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalInput, + MultimodalItemMetadata, + MultimodalParams, + MultimodalRuntimeData, +) from tensorrt_llm.llmapi.llm_args import MultimodalConfig from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -155,12 +161,12 @@ def test_edge_disagg_handoff_without_deepstack(): ) processor._tokenizer = _FakeTokenizer([11, vision_start, image_token_id, vision_end, 12]) - handoff = processor.build_disagg_prefill_multimodal_inputs( + prompt_token_ids, params = processor.build_disagg_prefill_multimodal_inputs( {"prompt": "image prompt"}, [{"tensor_size": (2, 16)}], ) - assert handoff.prompt_token_ids == [ + assert prompt_token_ids == [ 11, vision_start, image_token_id, @@ -168,7 +174,7 @@ def test_edge_disagg_handoff_without_deepstack(): vision_end, 12, ] - assert handoff.multimodal_embedding_lengths == [2] + assert params.item_metadata.output_embedding_lengths == (2,) @pytest.mark.cpu_only @@ -240,10 +246,7 @@ def _enable_encoder_cache(model, max_bytes: int): model.model_config._frozen = False model.model_config.multimodal_config = MultimodalConfig(encoder_cache_max_bytes=max_bytes) assert model.encoder_cache_active - if hasattr(model, "_initialize_multimodal_encoder_cache"): - cache = model._initialize_multimodal_encoder_cache(max_bytes) - else: - cache = model._get_multimodal_encoder_cache() + cache = model._initialize_multimodal_encoder_cache(max_bytes) assert cache is not None return cache @@ -390,7 +393,9 @@ def image_param(): ), multimodal_data={ "image": {"pixel_values": pixels, "image_grid_thw": grid}, - "multimodal_embedding_lengths": [rows], + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + item_refs=[("image", 0)], output_embedding_lengths=[rows] + ), "mm_processor_kwargs_hash": "edge-review", }, multimodal_runtime=MultimodalRuntimeData( @@ -415,19 +420,20 @@ def image_param(): processor = object.__new__(Cosmos3EdgeInputProcessor) processor._config = self.hf_config processor._tokenizer = _FakeTokenizer([11, vision_start, image_token_id, vision_end, 12]) - handoff = processor.build_disagg_prefill_multimodal_inputs( + prompt_token_ids, params = processor.build_disagg_prefill_multimodal_inputs( {"prompt": "image"}, [{"tensor_size": tuple(first.shape)}], ) - self.assertEqual(handoff.multimodal_embedding_lengths, [rows]) - input_ids = torch.tensor(handoff.prompt_token_ids, dtype=torch.long, device="cuda") + self.assertEqual(params.item_metadata.output_embedding_lengths, (rows,)) + params.multimodal_data["multimodal_embedding"] = first + input_ids = torch.tensor(prompt_token_ids, dtype=torch.long, device="cuda") image_positions = torch.nonzero(input_ids == image_token_id, as_tuple=False).flatten() text_positions = torch.nonzero(input_ids != image_token_id, as_tuple=False).flatten() self.assertEqual(image_positions.numel(), rows) prepared = model.prepare_multimodal_inputs( input_ids=input_ids, positions=None, - multimodal_params=[MultimodalParams(multimodal_data={"multimodal_embedding": first})], + multimodal_params=[params], num_context_requests=1, text_token_indices=text_positions, mm_token_indices=image_positions, From 87c4ee4b80823ff9ca184d8c17df3a33a9d7ad9e Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Sun, 25 Oct 2026 15:38:03 +0000 Subject: [PATCH 22/22] [TRTLLM-14477][fix] read common multimodal item order in modeling_v2 Read Qwen4 modeling_v2 request order from the shared item metadata during admission and from MultimodalParams.item_metadata during row assembly. Both paths still read removed order attributes, rejecting mixed image/video requests even when preprocessing had supplied a valid common record. Preserve grid-derived per-item row spans and all existing payload checks. Migrate the existing admission and splice fixtures to the same contract. Keep test lists unchanged. Validation: all Qwen4 modeling_v2 CPU suites pass (65 tests), including admission, splice, chunk slicing and video frame grouping. The named B300 accuracy test sends mixed media in both orders, but its SM103 whole-model accuracy and MTP acceptance length remain unverified. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../modeling.py | 39 ++++++++++--------- .../_torch/models/modeling_qwen3vl.py | 2 +- ...modeling_v2_qwen4_exp_multimodal_embeds.py | 15 +++++-- ..._modeling_v2_qwen4_exp_validate_request.py | 16 +++++--- 4 files changed, 44 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 14ace1194c5a..17b73c9eabd2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -137,6 +137,7 @@ from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache from tensorrt_llm._torch.utils import ActivationType +from tensorrt_llm.inputs.multimodal import get_multimodal_item_metadata from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo @@ -6476,7 +6477,7 @@ def validate_request(self, request) -> None: Refused: a frame past `_VISION_CALL_MAX_PATCHES` patches (the one bound the tower keeps), a malformed grid or pixel payload, media whose rows disagree with the prompt's placeholder count, missing or mis-sized `mrope_position_ids`, and a - request with both images and videos whose `mm_item_order` is missing or does not + request with both images and videos whose item metadata is missing or does not list its items.""" data = request.py_multimodal_data if not data: @@ -6504,7 +6505,8 @@ def validate_request(self, request) -> None: ) if len(grids) > 1: - self._check_item_order(grids, getattr(request, "py_mm_item_order", None)) + metadata = get_multimodal_item_metadata(data) + self._check_item_order(grids, metadata.item_refs if metadata is not None else None) cumsum = data.get("multimodal_embed_mask_cumsum") if cumsum is not None: @@ -6558,22 +6560,23 @@ def _check_item_order(grids, order) -> None: """A request with several modalities lists its items in prompt order.""" if not order: raise ValueError( - f"a request carrying {sorted(grids)} needs mm_item_order to splice them in " + f"a request carrying {sorted(grids)} needs item metadata to splice them in " "prompt order; the payload holds each modality's rows apart" ) listed = {modality: [] for modality in grids} for entry in order: try: - listed[entry["modality"]].append(int(entry["index"])) + modality, index = entry + listed[modality].append(int(index)) except (KeyError, TypeError, ValueError) as error: raise ValueError( - f"mm_item_order entry {entry!r} is not an item of this request, which " - f"carries {sorted(grids)}; entries are {{'modality', 'index'}}" + f"item_refs entry {entry!r} is not an item of this request, which " + f"carries {sorted(grids)}; entries are (modality, index) pairs" ) from error for modality, grid in grids.items(): if sorted(listed[modality]) != list(range(len(grid))): raise ValueError( - f"mm_item_order lists {modality} items {listed[modality]}; the payload " + f"item_refs lists {modality} items {listed[modality]}; the payload " f"carries {len(grid)}" ) @@ -6672,15 +6675,16 @@ def _in_prompt_order(self, params, towers): `towers` maps each modality the request carries to `(rows, grid_thw)`; each grid row is one item of `t*h*w / merge**2` rows. One modality is already in - prompt order. Two are reassembled item by item from `mm_item_order` - (`{"modality", "index"}` per item in prompt order, `index` counting that + prompt order. Two are reassembled item by item from `item_metadata.item_refs` + (`(modality, index)` per item in prompt order, `index` counting that modality's items); without it there is no host-side way to tell which placeholder run is an image and which a video.""" if len(towers) == 1: return next(iter(towers.values()))[0] - order = getattr(params, "mm_item_order", None) + metadata = params.item_metadata + order = metadata.item_refs if metadata is not None else None assert order, ( - f"a request carrying {sorted(towers)} needs mm_item_order to splice them " + f"a request carrying {sorted(towers)} needs item metadata to splice them " "in prompt order; the payload holds each modality's rows apart" ) merge_unit = self.model.v_merge * self.model.v_merge @@ -6693,18 +6697,17 @@ def _in_prompt_order(self, params, towers): ) items[modality] = torch.split(embeds, counts, 0) listed = {modality: [] for modality in items} - for entry in order: - assert entry["modality"] in items, ( - f"mm_item_order lists a {entry['modality']!r} item; this request " - f"carries {sorted(items)}" + for modality, index in order: + assert modality in items, ( + f"item_refs lists a {modality!r} item; this request carries {sorted(items)}" ) - listed[entry["modality"]].append(int(entry["index"])) + listed[modality].append(int(index)) for modality, parts in items.items(): assert sorted(listed[modality]) == list(range(len(parts))), ( - f"mm_item_order lists {modality} items {listed[modality]}; the " + f"item_refs lists {modality} items {listed[modality]}; the " f"payload carries {len(parts)}" ) - return torch.cat([items[entry["modality"]][int(entry["index"])] for entry in order], 0) + return torch.cat([items[modality][int(index)] for modality, index in order], 0) def forward( self, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 95eb34454c31..ce2e57941655 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -1419,7 +1419,7 @@ def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tenso `split_lengths`, so the rows must be in request-then-prompt order. Delegating to `encode_multimodal_by_groups` runs the same modality-batched ViT the aggregated path uses and applies the - per-request `mm_item_order` reorder before returning. + per-request `item_metadata.item_refs` reorder before returning. """ return [encode_multimodal_by_groups(self.mm_encoder_groups, multimodal_params)] diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py index 5a371f24bf36..bf95c7a4f166 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_multimodal_embeds.py @@ -5,7 +5,7 @@ in calls of whole frames under a row ceiling. The tower is a stub whose rows name the item they came from, so the bookkeeping -- -modality dispatch, `mm_item_order`, the chunk slice, the frame grouping and the order of +modality dispatch, item metadata, the chunk slice, the frame grouping and the order of the concatenated calls -- runs on the CPU without the vision ops or an engine. """ @@ -22,6 +22,11 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 modeling as M, ) +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, + MultimodalParams, +) pytestmark = pytest.mark.cpu_only @@ -77,7 +82,9 @@ def _params(image=None, video=None, order=None, runtime=None, mrope=True): data["image"] = _media(IMAGE, image) if video is not None: data["video"] = _media(VIDEO, video) - return SimpleNamespace(multimodal_data=data, mm_item_order=order, multimodal_runtime=runtime) + if order is not None: + data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order) + return MultimodalParams(multimodal_data=data, multimodal_runtime=runtime) def _splice(tower, params, num_rows, padding=0): @@ -101,7 +108,7 @@ def _splice(tower, params, num_rows, padding=0): def _item(modality, index): - return {"modality": modality, "index": index} + return modality, index # Two videos of 4 and 8 rows, two images of 4 and 1. @@ -137,7 +144,7 @@ def test_a_mixed_request_is_reassembled_in_prompt_order(): def test_a_mixed_request_needs_its_manifest(): - with pytest.raises(AssertionError, match="needs mm_item_order"): + with pytest.raises(AssertionError, match="needs item metadata"): _splice(_Tower(), _params(image=IMAGES, video=VIDEOS), len(MIXED_ROWS)) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py index 447ca2875a90..cc089cefa5c1 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_validate_request.py @@ -20,6 +20,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 modeling as M, ) +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalItemMetadata, +) pytestmark = pytest.mark.cpu_only @@ -67,7 +71,9 @@ def _request(image=None, video=None, order=None, **overrides): data["image"] = _media("image", image) if video is not None: data["video"] = _media("video", video) - fields = {"py_multimodal_data": data, "py_mm_item_order": order, "prompt_len": prompt_len} + if order is not None: + data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata(item_refs=order) + fields = {"py_multimodal_data": data, "prompt_len": prompt_len} fields.update(overrides) return SimpleNamespace(**fields) @@ -77,7 +83,7 @@ def _admit(request): def _item(modality, index): - return {"modality": modality, "index": index} + return modality, index IMAGES = [(1, 4, 4), (1, 2, 2)] @@ -173,9 +179,9 @@ def test_missing_or_mis_sized_position_ids_fail_the_request(): def test_a_mixed_request_needs_its_manifest(): - with pytest.raises(ValueError, match="needs mm_item_order"): + with pytest.raises(ValueError, match="needs item metadata"): _admit(_request(image=IMAGES, video=VIDEOS)) - with pytest.raises(ValueError, match="needs mm_item_order"): + with pytest.raises(ValueError, match="needs item metadata"): _admit(_request(image=IMAGES, video=VIDEOS, order=[])) @@ -187,7 +193,7 @@ def test_a_manifest_that_disagrees_with_the_payload_is_refused(): with pytest.raises(ValueError, match="is not an item of this request"): _admit(_request(image=IMAGES, video=VIDEOS, order=MIXED_ORDER + [_item("audio", 0)])) with pytest.raises(ValueError, match="is not an item of this request"): - _admit(_request(image=IMAGES, video=VIDEOS, order=[{"modality": "image"}])) + _admit(_request(image=IMAGES, video=VIDEOS, order=[("image",)])) def test_the_placeholder_count_is_the_sum_over_both_modalities():