diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index bec4e802c16e..47f3e9abc3c0 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -221,6 +221,10 @@ def encode_multimodal_by_groups( _MM_DATA_INPUT_MODALITY_KEYS = frozenset({"audio", "image", "video"}) +# Raw-input keys that may hold a plain per-item dim-0 stack, tried in this order by +# `MultimodalModelMixin._stacked_item_key`. Which key belongs to which modality is not +# encoded here: the leading-axis check in that helper is what selects the right one. +_MM_STACKED_ITEM_KEYS = ("pixel_values_videos", "pixel_values", "audio_features") _MM_AUX_STREAM: Optional[tuple[int, torch.cuda.Stream]] = None _MM_ENCODER_CACHE_LOG_NAME = "mm_encoder_cache" @@ -601,7 +605,7 @@ def build_multimodal_encoder_input( """Return a `MultimodalParams` whose raw modality inputs contain only `item_indices` from `param`, in that order. - Default handles three common single-modality layouts: + Default handles four 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. @@ -611,6 +615,14 @@ def build_multimodal_encoder_input( is sliced in parallel. - Audio, stacked on dim 0 (Whisper / Qwen2-Audio / Gemma4 audio): `input_features` `[B, mel_bins, T]` sliced by item. + - Any modality, plain stack on dim 0 with no companion size/grid field + (Gemma4 image / video / audio): the raw input's first axis is already one + row per item, so a dim-0 index is the whole slice. Unlike the + `image_sizes` case there is nothing to crop, because such layouts pad to + a model constant (e.g. Gemma4's fixed per-image patch count, with `-1` + sentinels in `image_position_ids`) rather than to a request-wide max. Only + taken when the leading axis matches the declared item count -- see + `_stacked_item_key`. 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 @@ -638,6 +650,7 @@ def build_multimodal_encoder_input( 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) + declared_items = len(param.multimodal_data.get("multimodal_embedding_lengths") or ()) if ( (grid_key and pixel_key) @@ -685,6 +698,13 @@ def build_multimodal_encoder_input( sliced = { "input_features": modality_data["input_features"][indices], } + elif (stacked_key := self._stacked_item_key(modality_data, declared_items)) is not None: + # Plain per-item dim-0 stack with no companion size/grid field. Every + # per-item sibling (Gemma4's `image_position_ids` / `image_seq_lens`, + # an audio mask, ...) is picked up by the sibling-slice pass below. + # `_stacked_item_key` already proved `shape[0] == declared_items`. + n_items = declared_items + sliced = {stacked_key: modality_data[stacked_key][indices]} else: raise NotImplementedError( f"Default `build_multimodal_encoder_input` cannot slice {modality} layout " @@ -710,6 +730,27 @@ def build_multimodal_encoder_input( multimodal_input=residual_input, ) + @staticmethod + def _stacked_item_key(modality_data: Dict[str, Any], item_count: int) -> Optional[str]: + """Return the raw-input key holding a plain per-item dim-0 stack, if any. + + A dim-0 index is only a per-item slice when the first axis really is one row + per item, so this requires `shape[0]` to equal `item_count` (the request's + declared item count). That check is what makes the branch safe for a model + whose raw tensor is flattened along a sub-item axis instead -- Gemma4 video + reshapes `(B, frames, patches, C)` to `(B * frames, patches, C)`, so its first + axis counts frames, not items, and slicing it by item index would silently + encode the wrong frames. Such layouts fall through to the `NotImplementedError` + and must override `build_multimodal_encoder_input`. + """ + if item_count == 0: + return None + for key in _MM_STACKED_ITEM_KEYS: + value = modality_data.get(key) + if isinstance(value, torch.Tensor) and value.dim() > 0 and value.shape[0] == item_count: + return key + return None + @staticmethod def _slice_per_item_sibling_fields( modality_data: Dict[str, Any], diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index acc7b260a7cb..26f7d165148e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -100,7 +100,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[tep4_trtllm] accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] SKIP (https://nvbugs/6418830) accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4_mtp3_gdn_replay_tep4 SKIP (https://nvbugs/6535779) accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] SKIP (https://nvbugs/6412098) -accuracy/test_llm_api_pytorch_multimodal.py::TestGemma4_26B_A4B::test_nvfp4 SKIP (https://nvbugs/6550127) accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6248827) accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray SKIP (https://nvbugs/6427411) cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mooncake_kvcache-90] SKIP (https://nvbugs/5838199) diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index 71aa44539766..eb0db6d0c96e 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -704,15 +704,81 @@ def test_build_multimodal_encoder_input_slices_audio_input_features(): torch.testing.assert_close(residual_audio["input_features_mask"], mask[[1, 0]]) +@pytest.mark.parametrize("modality", ["image", "video", "audio"]) +def test_build_multimodal_encoder_input_slices_plain_dim0_stack(modality): + # Gemma4 layout: a plain per-item dim-0 stack with no companion `image_sizes` or + # `*_grid_thw` field, padded to a model constant (fixed patch count) rather than to + # a request-wide max -- so slicing needs no trailing crop. `image_position_ids` and + # `image_seq_lens` must slice in parallel with the raw input, which is what the + # partial-hit encode path depends on (nvbugs/6550127). + pixel_key = "audio_features" if modality == "audio" else "pixel_values" + # 3 items x 4 padded rows x 2 feats; item i has i + 2 valid rows. + raw = torch.arange(3 * 4 * 2, dtype=torch.float32).reshape(3, 4, 2) + position_ids = torch.arange(3 * 4 * 2).reshape(3, 4, 2) + seq_lens = [2, 3, 4] + param = MultimodalParams( + multimodal_input=MultimodalInput( + multimodal_hashes=[[i] * 8 for i in range(3)], + multimodal_positions=[0, 0, 0], + multimodal_lengths=[1, 1, 1], + ), + multimodal_data={ + modality: { + pixel_key: raw, + "image_position_ids": position_ids, + "image_seq_lens": seq_lens, + "per_request_scalar": torch.tensor(7.0), + }, + "multimodal_embedding_lengths": [1, 1, 1], + "mm_processor_kwargs_hash": "kw", + }, + ) + model = DummyMultimodalModel(make_embedding(hidden_size=1), torch.tensor([0])) + + residual = model.build_multimodal_encoder_input(param, [2, 0]) + + residual_data = residual.multimodal_data[modality] + torch.testing.assert_close(residual_data[pixel_key], raw[[2, 0]]) + # Both per-item siblings follow the raw input, in the requested item order. + torch.testing.assert_close(residual_data["image_position_ids"], position_ids[[2, 0]]) + assert residual_data["image_seq_lens"] == [4, 2] + # Non-per-item siblings pass through untouched. + torch.testing.assert_close(residual_data["per_request_scalar"], torch.tensor(7.0)) + + +def test_build_multimodal_encoder_input_frame_flattened_video_raises(): + # Gemma4 video reshapes `(B, frames, patches, C)` -> `(B * frames, patches, C)`, so + # dim 0 counts frames rather than items. Indexing it by item would silently encode + # the wrong frames, so this must fall through to the override request instead of + # being caught by the plain dim-0 branch. + param = MultimodalParams( + multimodal_input=MultimodalInput( + multimodal_hashes=[[i] * 8 for i in range(2)], + multimodal_positions=[0, 0], + multimodal_lengths=[1, 1], + ), + multimodal_data={ + # 2 videos x 3 frames flattened to 6 rows against an item count of 2. + "video": {"pixel_values": torch.zeros(6, 4, 2)}, + "multimodal_embedding_lengths": [1, 1], + "mm_processor_kwargs_hash": "kw", + }, + ) + model = DummyMultimodalModel(make_embedding(hidden_size=1), torch.tensor([0])) + with pytest.raises(NotImplementedError, match="cannot slice video layout"): + model.build_multimodal_encoder_input(param, [0]) + + @pytest.mark.parametrize( "mm_data, expected_match", [ # `_encoder_cache_modality` returns None -> single-modality guard fires. ({}, "only supports single-modality"), - # Modality present but layout is neither pattern A (image_sizes) nor pattern B - # (grid_thw); default has nothing to dispatch on. + # Modality present but no layout matches: no `image_sizes`, no `grid_thw`, and + # without `multimodal_embedding_lengths` the plain dim-0 branch cannot confirm + # that the first axis is per-item either. ({"image": {"pixel_values": torch.zeros(2)}}, "cannot slice image layout"), - # Audio modality but no `input_features`; default falls through. + # Audio modality with no recognized raw-input key at all. ({"audio": {"nonsense": torch.zeros(2)}}, "cannot slice audio layout"), ], ids=["no_modality", "unhandled_image_layout", "unhandled_audio_layout"],