Skip to content

[None][feat] Enable MM encoder cache on Qwen3.x and Gemma4 VLMs #16662

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
2ez4bz wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[None][feat] Enable MM encoder cache on Qwen3.x and Gemma4 VLMs #16662

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
63ca9aa
[None][feat] Enable MM encoder cache on Qwen3.x and Gemma4 VLMs
2ez4bz Jul 20, 2026
99db94e
trim fat
2ez4bz Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion docs/source/models/supported-models.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,11 +127,17 @@ Note:
## Multimodal Encoder Optimizations

The following optimizations are available to models that implement
`MultimodalModelMixin`. Currently, only `Mistral3ForConditionalGeneration` supports them.
`MultimodalModelMixin`.

| Model Architecture | Multimodal Encoder Side Stream | Multimodal Embeddings Cache |
| ------------------ | ------------------------------ | --------------------------- |
| `Gemma4ForConditionalGeneration` | Yes | Yes |
| `Gemma4UnifiedForConditionalGeneration` | Yes | Yes |
| `Mistral3ForConditionalGeneration` | Yes | Yes |
| `Qwen3VLForConditionalGeneration` | Yes | Yes |
| `Qwen3VLMoeForConditionalGeneration` | Yes | Yes |
| `Qwen3_5ForConditionalGeneration` | Yes | Yes |
| `Qwen3_5MoeForConditionalGeneration` | Yes | Yes |
Comment thread
2ez4bz marked this conversation as resolved.

- **Multimodal encoder side stream** prefetches encoder work for pending requests on a separate
CUDA stream, allowing it to overlap with work on the main stream. Set
Expand Down
194 changes: 9 additions & 185 deletions tensorrt_llm/_torch/models/modeling_gemma4_unified.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,15 +33,14 @@
`layer_scalar`, final-logit softcap, tied embeddings). 12B has PLE and
KV-sharing **off** (`hidden_size_per_layer_input=0`, `num_kv_shared_layers=0`).

This module reuses the existing Gemma 4 multimodal wrapper
(:class:`Gemma4ForConditionalGeneration`) for all engine plumbing
This module inherits the shared Gemma 4 multimodal wrapper base
(:class:`Gemma4MultimodalModelBase`) for all engine plumbing
(`post_config` / `get_sub_model_config` / `infer_max_seq_len` /
`vocab_size_padded` / `get_model_defaults`), and reuses
:class:`Gemma4MultimodalEmbedder` (audio) and :class:`Gemma4InputProcessor`
(HF `AutoProcessor` resolves to `Gemma4UnifiedProcessor`; the output dict
keys match). It overrides `__init__` / `forward` / `_get_image_features` /
`_get_audio_features` / `load_weights` to drop the encoder towers and use the
encoder-free projections instead.
keys match). It implements `__init__`, `_get_image_features`,
`_get_audio_features`, and `load_weights` for the encoder-free projections.

TRT-LLM provides its own `gemma4_unified` config classes
(`_torch/configs/gemma4_unified.py`) and multimodal preprocessing (the vendored
Expand All @@ -55,13 +54,12 @@
import math
import os
import re
from typing import Dict, List, Optional
from typing import Dict

import numpy as np
import torch
from torch import nn
from torchvision.transforms.v2 import functional as tvF
from transformers import PreTrainedModel
from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
from transformers.image_processing_backends import TorchvisionBackend
from transformers.image_processing_utils import BatchFeature
Expand All @@ -77,20 +75,13 @@
MultimodalPlaceholderPlacement,
register_input_processor,
)
from ...inputs.multimodal import MultimodalParams
from ..attention_backend import AttentionMetadata
from ..modules.layer_norm import LayerNorm
from ..modules.linear import Linear
from .modeling_gemma4 import Gemma4ForCausalLM
from .modeling_gemma4mm import (
Gemma4ForConditionalGeneration,
Gemma4InputProcessor,
Gemma4MultimodalEmbedder,
)
from .modeling_multimodal_utils import (
find_input_mm_embeds,
fuse_input_embeds,
get_multimodal_embeddings,
Gemma4MultimodalModelBase,
)
from .modeling_utils import ModelConfig, filter_weights, register_auto_model

Expand Down Expand Up @@ -230,16 +221,12 @@ def __init__(self, *args, **kwargs):
interleave_placeholders=True,
),
)
class Gemma4UnifiedForConditionalGeneration(Gemma4ForConditionalGeneration):
"""Gemma 4 12B Unified (encoder-free). Reuses the Gemma 4 MM wrapper for
engine plumbing + the text core (:class:`Gemma4ForCausalLM`); replaces the
vision/audio towers with encoder-free linear embedders."""
class Gemma4UnifiedForConditionalGeneration(Gemma4MultimodalModelBase):
"""Gemma 4 12B Unified model with encoder-free multimodal projections."""

def __init__(self, model_config: ModelConfig):
config = model_config.pretrained_config
# Skip Gemma4ForConditionalGeneration.__init__ (it builds vision/audio
# *towers* the unified architecture does not have) and init the HF base.
PreTrainedModel.__init__(self, config)
super().__init__(config)

# ModelConfig always has `mapping`, and Mapping always has `local_rank`.
local_rank = model_config.mapping.local_rank
Expand Down Expand Up @@ -354,169 +341,6 @@ def _get_audio_features(self, audio_features, audio_features_mask=None):
features = features.reshape(-1, features.shape[-1])
return features.contiguous()

@property
def multimodal_data_device_paths(self) -> List[str]:
return [
"image.pixel_values",
"image.image_position_ids",
"video.pixel_values",
"video.image_position_ids",
"audio.audio_features",
"audio.audio_features_mask",
]

@staticmethod
def _has_active_multimodal_tokens(multimodal_param: MultimodalParams) -> bool:
"""Whether a context parameter needs embeddings in this forward.

Mirrors `Gemma4ForConditionalGeneration._has_active_multimodal_tokens`
from #15848 (this class overrides `forward` entirely, so the parent's
version would not apply here; deduplicate once #15848 lands on main).
"""
runtime = multimodal_param.multimodal_runtime
if runtime is not None and runtime.num_mm_tokens_in_chunk == 0:
return False

multimodal_data = multimodal_param.multimodal_data
if multimodal_data.get("multimodal_embedding") is not None:
return True

payload_fields = (
("image", "pixel_values"),
("video", "pixel_values"),
("audio", "audio_features"),
)
return any(
multimodal_data.get(modality, {}).get(field) is not None
for modality, field in payload_fields
)

def _forward_multimodal_encoder(
self, multimodal_params: List[MultimodalParams]
) -> torch.Tensor:
"""Run the encoder-free projectors for all uncached multimodal payloads.

Called by `get_multimodal_embeddings`, which caches the result in
`multimodal_data["multimodal_embedding"]` so later prefill chunks reuse
the embedding without re-running the projectors (mirrors the pattern in
`Gemma4ForConditionalGeneration._forward_multimodal_encoder` from #15848).
"""
pixel_values_list, image_position_ids_list = [], []
audio_features_list, audio_mask_list = [], []
video_pixel_values_list, video_position_ids_list = [], []
for multimodal_param in multimodal_params:
image_data = multimodal_param.multimodal_data.get("image", {})
if image_data.get("pixel_values") is not None:
pixel_values_list.append(image_data["pixel_values"])
if image_data.get("image_position_ids") is not None:
image_position_ids_list.append(image_data["image_position_ids"])
audio_data = multimodal_param.multimodal_data.get("audio", {})
if audio_data.get("audio_features") is not None:
audio_features_list.append(audio_data["audio_features"])
audio_mask_list.append(audio_data.get("audio_features_mask"))
video_data = multimodal_param.multimodal_data.get("video", {})
if video_data.get("pixel_values") is not None:
video_pixel_values_list.append(video_data["pixel_values"])
if video_data.get("image_position_ids") is not None:
video_position_ids_list.append(video_data["image_position_ids"])

embeddings = []
if pixel_values_list and self.embed_vision is not None:
pixel_values = torch.cat(pixel_values_list)
image_position_ids = (
torch.cat(image_position_ids_list)
if len(image_position_ids_list) == len(pixel_values_list)
else None
)
embeddings.append(self._get_image_features(pixel_values, image_position_ids))

if video_pixel_values_list and self.embed_vision is not None:
video_pixel_values = torch.cat(video_pixel_values_list)
video_position_ids = (
torch.cat(video_position_ids_list)
if len(video_position_ids_list) == len(video_pixel_values_list)
else None
)
embeddings.append(self._get_image_features(video_pixel_values, video_position_ids))

if audio_features_list and self.embed_audio is not None:
per_clip = []
for clip_index, audio_feature in enumerate(audio_features_list):
per_clip.append(
self._get_audio_features(audio_feature, audio_mask_list[clip_index])
)
embeddings.append(torch.cat(per_clip, dim=0))

return torch.cat(embeddings, dim=0) if embeddings else torch.empty(0)

def forward(
self,
attn_metadata: AttentionMetadata,
input_ids: Optional[torch.LongTensor] = None,
position_ids: Optional[torch.LongTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
return_context_logits: bool = False,
**kwargs,
) -> torch.Tensor:
multimodal_params = kwargs.get("multimodal_params", [])

# Filter to params that have active multimodal tokens in this chunk.
# get_multimodal_embeddings caches the result after the first run, so
# later prefill chunks reuse the embedding without re-running the projectors.
active_multimodal_params = [
mp for mp in multimodal_params if self._has_active_multimodal_tokens(mp)
]

mm_embeds: List[torch.Tensor] = []
all_mm_token_ids: List[torch.Tensor] = []
mm_token_type_ids = None

if active_multimodal_params:
mm_embeds = get_multimodal_embeddings(
encoder_forward_fn=self._forward_multimodal_encoder,
multimodal_params=active_multimodal_params,
)
mm_embeds = find_input_mm_embeds(mm_embeds, active_multimodal_params)

# Collect every defined multimodal token id. On cache-hit chunks the
# raw payloads (pixel_values / audio_features) may be absent while the
# cached embedding is used, so the ids cannot be derived from payload
# presence; extra ids are harmless (they simply match no position).
for token_ids in (self.image_token_ids, self.video_token_ids, self.audio_token_ids):
if token_ids is not None:
all_mm_token_ids.append(token_ids)

# Integer mm_token_type_ids (0=text,1=image,2=video,3=audio) drive the
# inherited bidirectional-vision attention mask in Gemma4ForCausalLM.
if mm_embeds and input_ids is not None:
mm_token_type_ids = torch.zeros_like(input_ids, dtype=torch.long)
if self.image_token_ids is not None:
mm_token_type_ids[torch.isin(input_ids, self.image_token_ids)] = 1
if self.video_token_ids is not None:
mm_token_type_ids[torch.isin(input_ids, self.video_token_ids)] = 2
if self.audio_token_ids is not None:
mm_token_type_ids[torch.isin(input_ids, self.audio_token_ids)] = 3

fuse_token_ids = torch.cat(all_mm_token_ids) if all_mm_token_ids else self.image_token_ids

input_ids, inputs_embeds = fuse_input_embeds(
embedding_layer=self.llm.model.embed_tokens,
input_ids=input_ids,
mm_embeds=mm_embeds,
mm_token_ids=fuse_token_ids,
**kwargs,
)
# 12B has PLE off (hidden_size_per_layer_input=0) -> no ple_input_ids.
return self.llm.forward(
attn_metadata=attn_metadata,
input_ids=input_ids,
position_ids=position_ids,
inputs_embeds=inputs_embeds,
return_context_logits=return_context_logits,
mm_token_type_ids=mm_token_type_ids,
lora_params=kwargs.get("lora_params", None),
)

def load_weights(self, weights: Dict, weight_mapper):
# Text backbone: "model.language_model.X" -> "model.X" (same as the
# parent), then load via the reused Gemma4 text core.
Expand Down
Loading
Loading