[https://nvbugs/6550127][fix] Support Gemma4 multimodal cache partial hits

tensorrt_llm/_torch/models/modeling_gemma4mm.py

@@ -23,6 +23,7 @@
 import copy
 import dataclasses
 import math
+from collections.abc import Sequence
 from itertools import groupby
 from typing import Dict, List, Optional, Tuple
 
@@ -32,6 +33,7 @@
 from torch import nn
 
 from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper
+from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache
 
 from ..._utils import nvtx_range
 from ...inputs import (
@@ -52,7 +54,11 @@
 from .modeling_gemma4 import Gemma4ForCausalLM
 from .modeling_gemma4_audio import Gemma4AudioModel
 from .modeling_gemma4_vision import Gemma4VisionModel
-from .modeling_multimodal_mixin import MultimodalModelMixin, PreparedLlmInputs
+from .modeling_multimodal_mixin import (
+    EncoderCachePartition,
+    MultimodalModelMixin,
+    PreparedLlmInputs,
+)
 from .modeling_multimodal_utils import _MULTIMODAL_ENV_NAME, _is_mm_disagg
 from .modeling_utils import ModelConfig, filter_weights, register_auto_model
 
@@ -576,6 +582,118 @@ def multimodal_data_device_paths(self) -> List[str]:
             "audio.audio_features_mask",
         ]
 
+    # TODO(TRTLLM-14981): Implement finer-grained caching for videos.
+    @classmethod
+    def partition_encoder_cache(
+        cls,
+        param: MultimodalParams,
+        encoder_cache: TensorLRUCache,
+    ) -> Optional[EncoderCachePartition]:
+        """Treat unsupported partial Gemma4 cache hits as full misses."""
+        partition = super().partition_encoder_cache(param, encoder_cache)
+        modality = cls._encoder_cache_modality(param)
+        if (
+            modality == "video"
+            and partition is not None
+            and not partition.is_full_hit
+            and not partition.is_full_miss
+        ):
+            # Gemma4 flattens each video's frames into dim 0 before encoding, so the persistent
+            # cache cannot (yet) rebuild a partial-hit input by video item. Re-encode the complete
+            # video payload while retaining cache lookup and reuse for full hits.
+            logger.warning_once(
+                "Gemma4 video encoder cache has a partial hit, but frame slicing is not "
+                "supported; re-encoding the complete video payload.",
+                key="gemma4_video_encoder_cache_partial_hit_unsupported",
+            )
+            return EncoderCachePartition(
+                hits={},
+                miss_indices=list(range(len(partition.keys))),
+                keys=partition.keys,
+            )
+        if (
+            modality in ("image", "audio")
+            and partition is not None
+            and not partition.is_full_hit
+            and not partition.is_full_miss
+        ):
+            input_key = {"image": "pixel_values", "audio": "audio_features"}[modality]
+            modality_data = param.multimodal_data[modality]
+            input_tensor = modality_data.get(input_key) if isinstance(modality_data, dict) else None
+            if (
+                not isinstance(input_tensor, torch.Tensor)
+                or input_tensor.dim() == 0
+                or input_tensor.shape[0] != len(partition.keys)
+            ):
+                logger.warning_once(
+                    f"Gemma4 {modality} encoder cache has a partial hit, but {input_key} is not "
+                    "item-major; re-encoding the complete payload.",
+                    key=f"gemma4_{modality}_encoder_cache_partial_hit_unsupported",
+                )
+                return EncoderCachePartition(
+                    hits={},
+                    miss_indices=list(range(len(partition.keys))),
+                    keys=partition.keys,
+                )
+        return partition
+
+    def build_multimodal_encoder_input(
+        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.
+        if modality not in ("image", "audio"):
+            return super().build_multimodal_encoder_input(param, item_indices)
+        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 ())
+        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,
+        )
+
     def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) -> torch.Tensor:
         """Encode uncached Gemma4 image, video, and audio payloads."""
         modality_inputs = (

tests/integration/test_lists/waives.txt

@@ -99,7 +99,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)