From 919732796ecc537f37f90b3e1675f0626a5c26fd Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Wed, 14 Oct 2026 22:45:57 -0700
Subject: [PATCH 1/6] [None][refactor] share multimodal input metadata helpers

Signed-off-by: yechank <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 86f52f71da3c..6f0713b1069c 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 7c2c2a1b0e30..3f705cca4570 100644
--- a/tensorrt_llm/inputs/registry.py
+++ b/tensorrt_llm/inputs/registry.py
@@ -42,7 +42,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
 
 N = TypeVar("N", bound=Type[nn.Module])
@@ -1280,6 +1280,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],
@@ -1303,15 +1311,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)),
@@ -1336,8 +1337,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(
@@ -1451,9 +1451,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))
@@ -1512,13 +1510,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:
@@ -1588,7 +1584,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 14f87fafa4a81dcb843737a102e64a2be3d62d23 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Wed, 14 Oct 2026 23:51:38 -0700
Subject: [PATCH 2/6] [TRTLLM-14477][refactor] unify multimodal item ordering

Signed-off-by: yechank <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 a224714198b2..f0f66642f4c5 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,
@@ -86,7 +91,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`:
 
@@ -99,7 +105,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, ...]
@@ -126,9 +132,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:
 
@@ -147,26 +152,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
 
 
@@ -189,10 +183,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
@@ -207,19 +201,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
@@ -242,16 +233,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(
@@ -264,9 +275,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
@@ -1409,7 +1420,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 3f705cca4570..e600fa6e007b 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]
@@ -1600,6 +1599,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 ecac218840420cfe21c95e441893f72711a85b9f Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 15 Oct 2026 19:50:25 -0700
Subject: [PATCH 3/6] [TRTLLM-14477][refactor] centralize multimodal item
 metadata

Signed-off-by: yechank <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 1d9e2cf33ed7..77eb3309919e 100644
--- a/tensorrt_llm/_torch/models/modeling_mistral.py
+++ b/tensorrt_llm/_torch/models/modeling_mistral.py
@@ -49,9 +49,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 f0f66642f4c5..8a459b61d45f 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 .modeling_multimodal_utils import (
@@ -129,30 +125,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}
@@ -165,7 +148,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:
@@ -182,8 +165,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
             ]
@@ -204,8 +186,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:
@@ -233,38 +214,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],
@@ -285,23 +234,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:
@@ -567,9 +518,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
@@ -1261,7 +1210,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,
@@ -1420,28 +1369,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 48354e65d4ca..640f4e7d539d 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, prefer_pinned
 from ...inputs import (BaseMultimodalDummyInputsBuilder,
diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py
index e25aab3d17b5..8d5786c7b390 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py
@@ -26,13 +26,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 MultimodalEncoderSchedulingPolicy, TorchLlmArgs
 from tensorrt_llm.logger import logger
diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index c7d38e701ef8..fcc71e6c5f3b 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 e600fa6e007b..d70dcb387391 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
@@ -35,13 +35,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
 
@@ -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 b972d7e15de0d1b5714ad871f7af785779d28283 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 15 Oct 2026 20:23:05 -0700
Subject: [PATCH 4/6] [TRTLLM-14477][refactor] declare multimodal item field
 layouts

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_gemma4mm.py        |  60 +++-------
 .../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 |  23 ++++
 .../test_multimodal_encoder_item_groups.py    |  72 ++++++++++++
 .../multimodal/test_multimodal_mixin.py       |  34 ++++++
 tests/unittest/inputs/test_multimodal.py      |  74 ++++++++++++
 9 files changed, 403 insertions(+), 101 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index 05476526e185..bd189be9a746 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -23,7 +23,6 @@
 import copy
 import dataclasses
 import math
-from collections.abc import Sequence
 from itertools import groupby
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple
 
@@ -47,7 +46,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
@@ -702,62 +701,31 @@ def partition_encoder_cache(
                 )
         return partition
 
-    def build_multimodal_encoder_input(
+    def get_multimodal_field_layout(
         self,
         param: MultimodalParams,
-        item_indices: Sequence[int],
-    ) -> MultimodalParams:
-        """Build a Gemma4 image 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` and uses `audio_features`.
-
-        Partial encoder-cache hits therefore need this override to slice those tensors, together
-        with their per-item position, length, and mask fields.
-        """
-        modality = self._encoder_cache_modality(param)
-        # Partial video partitions are converted to full misses above, so they cannot call this
-        # hook. Delegate unexpected direct calls 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)
+            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)
-
-        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}
 
     def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) -> torch.Tensor:
         """Encode uncached Gemma4 image, video, and audio payloads."""
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
index 8a459b61d45f..8fd40cbc3a74 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
 
@@ -390,8 +392,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
@@ -801,44 +803,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)
@@ -853,7 +842,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)
@@ -930,26 +946,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 640f4e7d539d..9fd856d69d23 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, prefer_pinned
@@ -2153,9 +2154,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 2d84557faf49..95b3c93c740d 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py
@@ -37,7 +37,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 @@
     Qwen2VLInputProcessorBase,
     _prepare_qwen_vl_mrope_config,
     _prepare_qwen_vl_vision_attn_metadata,
+    _qwen_vision_field_layout,
 )
 from .modeling_utils import (
     ModelConfig,
@@ -1250,6 +1255,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 f454bf7c2caa..d4321fc82c0b 100644
--- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
+++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
@@ -70,6 +70,8 @@
     get_multimodal_embeddings,
 )
 from tensorrt_llm.inputs.multimodal import (  # noqa: E402
+    MULTIMODAL_ENCODER_ITEM_METADATA_KEY,
+    MultimodalEncoderItemMetadata,
     MultimodalInput,
     MultimodalParams,
     MultimodalRuntimeData,
@@ -905,11 +907,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],
             }
@@ -920,6 +924,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_encoder_cache_partial_hit_slices_gemma4_audio_input(self):
         """A partial hit encodes only the missing Gemma4 audio item."""
@@ -936,6 +941,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 dae1bb3a070ca3c09bce8ad8c5f57f87a7bd9558 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Fri, 16 Oct 2026 01:39:31 -0700
Subject: [PATCH 5/6] refactor: name multimodal item partitions and fix
 negative selection

Signed-off-by: yechank <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 bd189be9a746..a53e90afc91f 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -46,7 +46,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
@@ -701,11 +701,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
@@ -725,7 +725,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}
 
     def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) -> torch.Tensor:
         """Encode uncached Gemma4 image, video, and audio payloads."""
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
index 8fd40cbc3a74..529373a25c76 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,
@@ -393,7 +393,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
@@ -803,11 +803,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.
@@ -842,11 +842,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 9fd856d69d23..3cce0f356b22 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, prefer_pinned
@@ -2154,13 +2154,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
@@ -2184,7 +2184,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,
     }
@@ -2193,12 +2193,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 95b3c93c740d..3dfb8737127e 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py
@@ -39,7 +39,7 @@
 )
 from ...inputs.multimodal import (
     DisaggPrefillMultimodalInputs,
-    MultimodalFieldLayout,
+    MultimodalItemPartition,
     MultimodalParams,
 )
 from ...logger import logger
@@ -66,7 +66,7 @@
     Qwen2VLInputProcessorBase,
     _prepare_qwen_vl_mrope_config,
     _prepare_qwen_vl_vision_attn_metadata,
-    _qwen_vision_field_layout,
+    _qwen_vision_item_partitions,
 )
 from .modeling_utils import (
     ModelConfig,
@@ -1255,12 +1255,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 08c7a29bcadd0d0a73c23b75fd6163e8933d49d7 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Fri, 16 Oct 2026 02:10:10 -0700
Subject: [PATCH 6/6] refactor: group multimodal request data and encoder state

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_gemma4mm.py        |  58 +--
 .../models/modeling_multimodal_mixin.py       |  40 ++-
 .../models/modeling_nemotron_h_multimodal.py  |  24 ++
 .../_torch/models/modeling_qwen2vl.py         |  57 ++-
 .../_torch/pyexecutor/engine/multimodal.py    |  18 +-
 .../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    |  93 ++++-
 .../_torch/multimodal/test_encoder_group.py   |  20 +-
 tests/unittest/inputs/test_multimodal.py      | 115 +++++-
 16 files changed, 982 insertions(+), 316 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index a53e90afc91f..1b9d2d62e9c9 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -46,7 +46,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
@@ -199,6 +199,26 @@ def _normalize_audio_inputs(audios, target_sr: int = 16000):
 # ---------------------------------------------------------------------------
 
 
+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.
 
@@ -207,6 +227,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
@@ -706,26 +737,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)
 
     def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) -> torch.Tensor:
         """Encode uncached Gemma4 image, video, and audio payloads."""
diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
index 529373a25c76..62b306f46570 100644
--- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
+++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
@@ -40,8 +40,8 @@
 from tensorrt_llm._utils import prefer_pinned
 from tensorrt_llm.inputs.multimodal import (
     MULTIMODAL_ENCODER_ITEM_METADATA_KEY,
-    MultimodalEncoderItemMetadata,
     MultimodalInput,
+    MultimodalItemMetadata,
     MultimodalItemPartition,
     MultimodalParams,
     MultimodalRuntimeData,
@@ -89,7 +89,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`:
@@ -226,7 +226,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`.
 
@@ -502,18 +502,17 @@ 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, "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.
+        carries the request's item metadata, already validated by the producer
+        and executor admission. The caller reuses it for selection and outputs.
         """
         run_param: Optional[MultimodalParams] = None
         run_modality: Optional[str] = None
-        run_metadata: Optional[MultimodalEncoderItemMetadata] = None
+        run_metadata: Optional[MultimodalItemMetadata] = None
         run_indices: list[int] = []
         for multimodal_param, item_idx in selected_items:
             try:
@@ -524,9 +523,9 @@ def _runs_by_request_modality(
                 )
             except (TypeError, ValueError) as error:
                 raise MultimodalEncoderContractError(str(error)) from error
-            if metadata is None:
+            if metadata is None or metadata.output_embedding_lengths is None:
                 raise MultimodalEncoderContractError(
-                    "MM item metadata is required for item encoding"
+                    "MM item metadata with output lengths is required for item encoding"
                 )
             if item_idx < 0 or item_idx >= len(metadata.item_refs):
                 raise MultimodalEncoderContractError(
@@ -814,7 +813,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,
@@ -1209,6 +1209,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,
@@ -1370,9 +1381,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 e6b362f553d5..f20fddc74bed 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,
@@ -1159,6 +1161,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 3cce0f356b22..9a709482c879 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)
 
@@ -384,6 +385,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.
     #
@@ -2161,26 +2178,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 8d5786c7b390..ed4e9a73cdc3 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py
@@ -26,7 +26,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 (
@@ -37,12 +38,7 @@
 from tensorrt_llm.llmapi.llm_args import MultimodalEncoderSchedulingPolicy, TorchLlmArgs
 from tensorrt_llm.logger import logger
 
-from ..llm_request import (
-    LlmRequest,
-    MultimodalEncoderProgress,
-    MultimodalEncoderRequestError,
-    _Unset,
-)
+from ..llm_request import LlmRequest, MultimodalEncoderProgress, MultimodalEncoderRequestError
 
 
 def resolve_mm_encoder_token_budget(base_budget: int, model_max_atomic_item_tokens: int) -> int:
@@ -429,12 +425,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,
@@ -493,7 +489,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 2aeffa719254..bf1edba7395f 100644
--- a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py
+++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py
@@ -137,7 +137,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 3efea91cec6f..ed6f9cacc4ab 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 (Any, Container, Dict, Hashable, Iterable, List, Optional,
-                    Union, cast)
+from typing import 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.
-
-    Only `READY` gates control flow today (`is_multimodal_encoder_ready`);
-    `PENDING`/`PARTIAL` are the intermediate states the item-level chunked
-    MM prefill follow-up will consume (per-chunk readiness) and document the
-    FCFS priority in the meantime."""
-
-    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
-
-
 @dataclass(slots=True)
 class PerfTimingInfo:
     """Stores performance timing information for a request."""
@@ -880,9 +729,14 @@ def __init__(
                                                     None)
         self.py_end_id: Optional[int] = kwargs.pop("end_id", None)
         self.py_lora_path: str | None = kwargs.pop("py_lora_path", 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
@@ -1302,7 +1156,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)
 
@@ -1381,13 +1293,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 "
@@ -1434,6 +1348,9 @@ def initialize_multimodal_encoder_request(
     remain resident until the request becomes LLM-eligible.
     """
     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 c4b14391194b..0f574ee477f5 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -42,8 +42,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,
                                           WaitingQueuePolicy)
@@ -6592,8 +6592,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
@@ -8489,6 +8491,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 c4d604a7e0dc..297746bfbbab 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
@@ -177,6 +178,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]:
@@ -188,6 +341,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")
@@ -195,6 +355,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):
@@ -643,7 +953,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.
@@ -704,15 +1014,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.
@@ -723,6 +1040,10 @@ def get_item_layout(self) -> tuple[list[tuple[str, int]], list[int]]:
         """
         metadata = self.item_metadata
         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 d70dcb387391..934c854fbfda 100644
--- a/tensorrt_llm/inputs/registry.py
+++ b/tensorrt_llm/inputs/registry.py
@@ -38,14 +38,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
 
@@ -199,6 +201,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,
@@ -1503,55 +1530,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 7c8ec0a787c1..8c8f62f10444 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 (
@@ -38,6 +35,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,
 )
@@ -50,6 +48,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]
@@ -100,9 +181,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():
@@ -606,17 +686,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 e04064c1aa9b..bd284586f961 100644
--- a/tests/unittest/_torch/executor/test_request_utils.py
+++ b/tests/unittest/_torch/executor/test_request_utils.py
@@ -29,11 +29,64 @@
 from tensorrt_llm._torch.pyexecutor.scheduler import FCFSWaitingQueue
 from tensorrt_llm.bindings import executor as trtllm
 from tensorrt_llm.conversation_params import ConversationParams
+from tensorrt_llm.inputs.multimodal import (
+    MULTIMODAL_ENCODER_ITEM_METADATA_KEY,
+    MultimodalItemMetadata,
+    MultimodalItemPartition,
+)
 from tensorrt_llm.mapping import CpType
 
 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 ca7554b470cd..3ec66dd89f26 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
@@ -23,6 +24,8 @@
 from tensorrt_llm.inputs.multimodal import (
     MULTIMODAL_ENCODER_ITEM_METADATA_KEY,
     MultimodalEncoderItemMetadata,
+    MultimodalEncoderRequestState,
+    MultimodalItemMetadata,
     MultimodalItemPartition,
     MultimodalParams,
 )
@@ -34,6 +37,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]
@@ -198,7 +279,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)]
@@ -209,11 +290,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
@@ -284,7 +371,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)
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 6590fb99ab39..9578e89a3f1c 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,
@@ -166,7 +167,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],
@@ -179,7 +180,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,
     )
 
 
@@ -191,9 +192,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
@@ -202,7 +202,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():
@@ -268,6 +271,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]."""
 
@@ -623,7 +724,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: