From 9579e3572d48bd94ef8e95a0fb4f93dccbd9c093 Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:46:25 -0700 Subject: [PATCH 1/2] [https://nvbugs/6550127] Support Gemma4 multimodal cache partial hits * Why? The generic multimodal cache path cannot slice Gemma4 image and audio layouts, causing partial encoder-cache hits to fail during input construction. * What? Override partial-hit input construction for Gemma4 images and audio, keeping their per-item metadata aligned. Bypass persistent caching for videos with a warning until frame-level slicing is supported, and remove the obsolete accuracy waiver. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- .../_torch/models/modeling_gemma4mm.py | 78 ++++++++ tests/integration/test_lists/waives.txt | 1 - .../_torch/modeling/test_gemma4_multimodal.py | 168 +++++++++++++++++- 3 files changed, 238 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index a5d08d4d06b8..6905c598797a 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/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 @@ -576,6 +577,83 @@ def multimodal_data_device_paths(self) -> List[str]: "audio.audio_features_mask", ] + # TODO(TRTLLM-14981): Implement finer-grained caching for videos. + @staticmethod + def _encoder_cache_modality(param: MultimodalParams) -> Optional[str]: + """Return the cacheable Gemma4 modality, excluding video.""" + modality = MultimodalModelMixin._encoder_cache_modality(param) + # 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. + # Treat video as uncacheable to retain the pre-cache behavior: encode the complete video + # payload and reuse only its request-local embedding for chunked prefill. + if modality == "video": + logger.warning_once( + "Gemma4 video inputs currently bypass the persistent multimodal encoder " + "cache because partial-hit frame slicing is not supported; video inputs " + "will be re-encoded for each request.", + key="gemma4_video_encoder_cache_unsupported", + ) + return None + return modality + + 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) + # Video is marked uncacheable above, so it cannot produce the partial-hit partition that + # calls this hook; the normal path re-encodes its complete payload. 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 = ( diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 2180bc6f7bb3..db5093fbfa32 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index c51b328951ff..29c949f0e97f 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -61,6 +61,7 @@ from tensorrt_llm._torch.models.modeling_gemma4mm import ( # noqa: E402 Gemma4ForConditionalGeneration, Gemma4MultimodalEmbedder, + Gemma4MultimodalModelBase, ) from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin # noqa: E402 from tensorrt_llm._torch.models.modeling_multimodal_utils import ( # noqa: E402 @@ -131,7 +132,7 @@ } -class _Gemma4EncoderCacheHarness(MultimodalModelMixin): +class _Gemma4EncoderCacheHarness(Gemma4MultimodalModelBase): """Lightweight Gemma4 encoder-cache harness without model weights.""" supports_encoder_cache = True @@ -144,6 +145,7 @@ def __init__(self, embedding_dim: int = 12) -> None: self._embedding_dim = embedding_dim self.encoder_calls = 0 self.audio_tower = None + self.embed_audio = object() @property def embedding_dim(self) -> int: @@ -162,24 +164,92 @@ def _get_image_features(self, pixel_values: torch.Tensor, **kwargs) -> torch.Ten dtype=self.embedding_dtype, ) + def _get_audio_features( + self, audio_features: torch.Tensor, audio_features_mask: torch.Tensor | None + ) -> torch.Tensor: + self.encoder_calls += 1 + return torch.full( + (audio_features.shape[0] * 2, self.embedding_dim), + float(self.encoder_calls), + dtype=self.embedding_dtype, + ) + + +def _make_keyed_image_param( + item_hashes: list[list[int]] | None = None, +) -> MultimodalParams: + if item_hashes is None: + item_hashes = [[1, 2, 3, 4, 5, 6, 7, 8]] + item_count = len(item_hashes) + embedding_lengths = [2] * item_count + pixel_values = torch.arange(item_count, dtype=torch.float32).reshape(item_count, 1, 1) + image_position_ids = torch.arange(item_count * 2).reshape(item_count, 1, 2) + return MultimodalParams( + multimodal_input=MultimodalInput( + multimodal_hashes=item_hashes, + multimodal_positions=[0] * item_count, + multimodal_lengths=embedding_lengths, + ), + multimodal_data={ + "image": { + "pixel_values": pixel_values, + "image_position_ids": image_position_ids, + "image_seq_lens": [1] * item_count, + }, + "multimodal_embedding_lengths": embedding_lengths, + "mm_processor_kwargs_hash": "kwargs-a", + }, + multimodal_runtime=MultimodalRuntimeData( + embed_mask_cumsum=torch.arange(1, sum(embedding_lengths) + 1, dtype=torch.int64), + past_seen_token_num=0, + chunk_end_pos=sum(embedding_lengths), + ), + ) + + +def _make_keyed_video_param(item_hashes: list[list[int]]) -> MultimodalParams: + item_count = len(item_hashes) + embedding_lengths = [2] * item_count + return MultimodalParams( + multimodal_input=MultimodalInput( + multimodal_hashes=item_hashes, + multimodal_positions=[0] * item_count, + multimodal_lengths=embedding_lengths, + ), + multimodal_data={ + "video": {"pixel_values": torch.arange(item_count).reshape(item_count, 1, 1)}, + "multimodal_embedding_lengths": embedding_lengths, + "mm_processor_kwargs_hash": "kwargs-a", + }, + multimodal_runtime=MultimodalRuntimeData( + embed_mask_cumsum=torch.arange(1, sum(embedding_lengths) + 1, dtype=torch.int64), + past_seen_token_num=0, + chunk_end_pos=sum(embedding_lengths), + ), + ) + -def _make_keyed_image_param() -> MultimodalParams: - embedding_lengths = [2] +def _make_keyed_audio_param(item_hashes: list[list[int]]) -> MultimodalParams: + item_count = len(item_hashes) + embedding_lengths = [2] * item_count return MultimodalParams( multimodal_input=MultimodalInput( - multimodal_hashes=[[1, 2, 3, 4, 5, 6, 7, 8]], - multimodal_positions=[0], + multimodal_hashes=item_hashes, + multimodal_positions=[0] * item_count, multimodal_lengths=embedding_lengths, ), multimodal_data={ - "image": {"pixel_values": torch.empty(1, 1, 1)}, + "audio": { + "audio_features": torch.arange(item_count * 2).reshape(item_count, 1, 2), + "audio_features_mask": torch.ones(item_count, 1), + }, "multimodal_embedding_lengths": embedding_lengths, "mm_processor_kwargs_hash": "kwargs-a", }, multimodal_runtime=MultimodalRuntimeData( - embed_mask_cumsum=torch.arange(1, 3, dtype=torch.int64), + embed_mask_cumsum=torch.arange(1, sum(embedding_lengths) + 1, dtype=torch.int64), past_seen_token_num=0, - chunk_end_pos=2, + chunk_end_pos=sum(embedding_lengths), ), ) @@ -774,6 +844,88 @@ def test_encoder_cache_reuses_image_embedding_across_requests(self): torch.testing.assert_close(second, first) self.assertEqual(len(model._multimodal_encoder_cache), 1) + def test_encoder_cache_partial_hit_slices_gemma4_image_input(self): + """A partial hit encodes only the missing Gemma4 image.""" + model = _Gemma4EncoderCacheHarness() + shared_hash = [1] * 8 + model._get_or_encode_multimodal_embeddings( + [_make_keyed_image_param(item_hashes=[shared_hash])] + ) + + embeddings = model._get_or_encode_multimodal_embeddings( + [_make_keyed_image_param(item_hashes=[shared_hash, [2] * 8])] + ) + + self.assertEqual(model.encoder_calls, 2) + torch.testing.assert_close(embeddings[:2], torch.ones(2, model.embedding_dim)) + torch.testing.assert_close(embeddings[2:], torch.full((2, model.embedding_dim), 2.0)) + self.assertEqual(len(model._multimodal_encoder_cache), 2) + + def test_build_multimodal_encoder_input_slices_gemma4_image_metadata(self): + """Gemma4 image position metadata follows the selected images.""" + param = _make_keyed_image_param(item_hashes=[[0] * 8, [1] * 8, [2] * 8]) + source_image = param.multimodal_data["image"] + + residual = _Gemma4EncoderCacheHarness().build_multimodal_encoder_input(param, [2, 0]) + + residual_image = residual.multimodal_data["image"] + torch.testing.assert_close( + residual_image["pixel_values"], source_image["pixel_values"][[2, 0]] + ) + torch.testing.assert_close( + residual_image["image_position_ids"], source_image["image_position_ids"][[2, 0]] + ) + self.assertEqual(residual_image["image_seq_lens"], [1, 1]) + + 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) + param = MultimodalParams( + multimodal_data={ + "audio": { + "audio_features": audio_features, + "audio_features_mask": audio_mask, + }, + "multimodal_embedding_lengths": [1, 1, 1], + } + ) + + residual = _Gemma4EncoderCacheHarness().build_multimodal_encoder_input(param, [2, 0]) + + 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]]) + + def test_encoder_cache_partial_hit_slices_gemma4_audio_input(self): + """A partial hit encodes only the missing Gemma4 audio item.""" + model = _Gemma4EncoderCacheHarness() + shared_hash = [1] * 8 + model._get_or_encode_multimodal_embeddings([_make_keyed_audio_param([shared_hash])]) + + embeddings = model._get_or_encode_multimodal_embeddings( + [_make_keyed_audio_param([shared_hash, [2] * 8])] + ) + + self.assertEqual(model.encoder_calls, 2) + torch.testing.assert_close(embeddings[:2], torch.ones(2, model.embedding_dim)) + torch.testing.assert_close(embeddings[2:], torch.full((2, model.embedding_dim), 2.0)) + self.assertEqual(len(model._multimodal_encoder_cache), 2) + + def test_encoder_cache_reencodes_gemma4_video(self): + """Gemma4 video bypasses the persistent cache until frame slicing is supported.""" + model = _Gemma4EncoderCacheHarness() + shared_hash = [1] * 8 + first = model._get_or_encode_multimodal_embeddings([_make_keyed_video_param([shared_hash])]) + second = model._get_or_encode_multimodal_embeddings( + [_make_keyed_video_param([shared_hash, [2] * 8])] + ) + + self.assertEqual(model.encoder_calls, 2) + torch.testing.assert_close(first, torch.ones(2, model.embedding_dim)) + torch.testing.assert_close(second, torch.full((4, model.embedding_dim), 2.0)) + self.assertEqual(len(model._multimodal_encoder_cache), 0) + def test_chunked_prefill_reuses_cached_vision_embeddings(self): """Later active chunks slice cached features without rerunning vision.""" model = self._make_model() From d3b295b4fcfc32f87d9c32b6a10eebdd0ea21bb1 Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Tue, 4 Aug 2026 09:45:59 -0700 Subject: [PATCH 2/2] address comments Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- .../_torch/models/modeling_gemma4mm.py | 80 ++++++++++++++----- .../_torch/modeling/test_gemma4_multimodal.py | 52 ++++++++++-- 2 files changed, 107 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 6905c598797a..ab211c1a81a9 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -33,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 ( @@ -53,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 @@ -578,23 +583,59 @@ def multimodal_data_device_paths(self) -> List[str]: ] # TODO(TRTLLM-14981): Implement finer-grained caching for videos. - @staticmethod - def _encoder_cache_modality(param: MultimodalParams) -> Optional[str]: - """Return the cacheable Gemma4 modality, excluding video.""" - modality = MultimodalModelMixin._encoder_cache_modality(param) - # 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. - # Treat video as uncacheable to retain the pre-cache behavior: encode the complete video - # payload and reuse only its request-local embedding for chunked prefill. - if modality == "video": + @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 inputs currently bypass the persistent multimodal encoder " - "cache because partial-hit frame slicing is not supported; video inputs " - "will be re-encoded for each request.", - key="gemma4_video_encoder_cache_unsupported", + "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, ) - return None - return modality + 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, @@ -611,10 +652,9 @@ def build_multimodal_encoder_input( with their per-item position, length, and mask fields. """ modality = self._encoder_cache_modality(param) - # Video is marked uncacheable above, so it cannot produce the partial-hit partition that - # calls this hook; the normal path re-encodes its complete payload. Delegate unexpected - # direct calls to the generic validation so they fail instead of returning an incorrectly - # unsliced input. + # 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] diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index 29c949f0e97f..4567e1717f17 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -912,19 +912,61 @@ 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) - def test_encoder_cache_reencodes_gemma4_video(self): - """Gemma4 video bypasses the persistent cache until frame slicing is supported.""" + 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() + shared_hash = [1] * 8 + cases = ( + ("image", "pixel_values", _make_keyed_image_param), + ("audio", "audio_features", _make_keyed_audio_param), + ) + + for modality, input_key, make_param in cases: + with self.subTest(modality=modality): + model._get_or_encode_multimodal_embeddings([make_param([shared_hash])]) + param = make_param([shared_hash, [2] * 8]) + input_tensor = param.multimodal_data[modality][input_key] + param.multimodal_data[modality][input_key] = torch.cat( + (input_tensor, input_tensor[:1]), dim=0 + ) + + with unittest.mock.patch( + "tensorrt_llm._torch.models.modeling_gemma4mm.logger.warning_once" + ) as warning_once: + partition = model.partition_encoder_cache( + param, model._multimodal_encoder_cache + ) + + self.assertTrue(partition.is_full_miss) + self.assertEqual(partition.hits, {}) + self.assertEqual(partition.miss_indices, [0, 1]) + warning_once.assert_called_once() + + def test_encoder_cache_reuses_full_video_hits_and_reencodes_partial_hits(self): + """Gemma4 video reuses full hits and re-encodes complete partial-hit requests.""" model = _Gemma4EncoderCacheHarness() shared_hash = [1] * 8 first = model._get_or_encode_multimodal_embeddings([_make_keyed_video_param([shared_hash])]) - second = model._get_or_encode_multimodal_embeddings( + full_hit = model._get_or_encode_multimodal_embeddings( + [_make_keyed_video_param([shared_hash])] + ) + with unittest.mock.patch( + "tensorrt_llm._torch.models.modeling_gemma4mm.logger.warning_once" + ) as warning_once: + partial_hit = model._get_or_encode_multimodal_embeddings( + [_make_keyed_video_param([shared_hash, [2] * 8])] + ) + repeated = model._get_or_encode_multimodal_embeddings( [_make_keyed_video_param([shared_hash, [2] * 8])] ) + warning_once.assert_called_once() self.assertEqual(model.encoder_calls, 2) + torch.testing.assert_close(full_hit, first) torch.testing.assert_close(first, torch.ones(2, model.embedding_dim)) - torch.testing.assert_close(second, torch.full((4, model.embedding_dim), 2.0)) - self.assertEqual(len(model._multimodal_encoder_cache), 0) + torch.testing.assert_close(partial_hit, torch.full((4, model.embedding_dim), 2.0)) + torch.testing.assert_close(repeated, partial_hit) + self.assertEqual(len(model._multimodal_encoder_cache), 2) def test_chunked_prefill_reuses_cached_vision_embeddings(self): """Later active chunks slice cached features without rerunning vision."""