[https://nvbugs/6550127][fix] Added a fourth dispatch branch for plain dim-0 stacks, resolved per modality…

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],

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)

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"],