From 63ca9aa6dd21fdb25a76a054fe102d36e0bd1edc Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Mon, 20 Jul 2026 16:26:57 -0700 Subject: [PATCH 1/2] [None][feat] Enable MM encoder cache on Qwen3.x and Gemma4 VLMs Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- docs/source/models/supported-models.md | 7 +- .../_torch/models/modeling_gemma4_unified.py | 172 +------ .../_torch/models/modeling_gemma4mm.py | 431 ++++++++---------- .../_torch/models/modeling_mistral.py | 88 +--- .../models/modeling_multimodal_mixin.py | 132 +++++- .../_torch/models/modeling_qwen3_5.py | 2 + .../_torch/models/modeling_qwen3vl.py | 287 ++++-------- .../_torch/models/modeling_qwen3vl_moe.py | 2 + .../models/modeling_qwen_image_bench.py | 12 + .../_torch/pyexecutor/model_engine.py | 7 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 4 - .../executor/test_pytorch_model_engine.py | 22 + .../_torch/modeling/test_gemma4_multimodal.py | 87 +++- .../_torch/modeling/test_modeling_qwen3vl.py | 35 ++ 14 files changed, 583 insertions(+), 705 deletions(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index b03f9961f7cc..9b4c5653a70e 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -127,11 +127,16 @@ 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 | | `Mistral3ForConditionalGeneration` | Yes | Yes | +| `Qwen3VLForConditionalGeneration` | Yes | Yes | +| `Qwen3VLMoeForConditionalGeneration` | Yes | Yes | +| `Qwen3_5ForConditionalGeneration` | Yes | Yes | +| `Qwen3_5MoeForConditionalGeneration` | Yes | Yes | - **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 diff --git a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py index f2def0290eec..d0cae7e25ffc 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py @@ -55,7 +55,7 @@ import math import os import re -from typing import Dict, List, Optional +from typing import Dict import numpy as np import torch @@ -77,8 +77,6 @@ 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 @@ -87,11 +85,6 @@ Gemma4InputProcessor, Gemma4MultimodalEmbedder, ) -from .modeling_multimodal_utils import ( - find_input_mm_embeds, - fuse_input_embeds, - get_multimodal_embeddings, -) from .modeling_utils import ModelConfig, filter_weights, register_auto_model # Raw HF nn.LayerNorm default eps (the unified vision LayerNorms use this, NOT @@ -354,169 +347,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. diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 0f5a340c0a9a..be999b51d973 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -47,18 +47,13 @@ from ...inputs.multimodal import MultimodalParams from ...logger import logger from ...sampling_params import SamplingParams -from ..attention_backend import AttentionMetadata +from ..modules.embedding import Embedding from ..modules.linear import Linear from .modeling_gemma4 import Gemma4ForCausalLM from .modeling_gemma4_audio import Gemma4AudioModel from .modeling_gemma4_vision import Gemma4VisionModel -from .modeling_multimodal_utils import ( - _MULTIMODAL_ENV_NAME, - _is_mm_disagg, - find_input_mm_embeds, - fuse_input_embeds, - get_multimodal_embeddings, -) +from .modeling_multimodal_mixin import MultimodalModelMixin, PreparedLlmInputs +from .modeling_multimodal_utils import _MULTIMODAL_ENV_NAME, _is_mm_disagg from .modeling_utils import ModelConfig, filter_weights, register_auto_model _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0" @@ -550,6 +545,146 @@ def call_with_text_prompt( # --------------------------------------------------------------------------- +class Gemma4MultimodalModelBase(MultimodalModelMixin, PreTrainedModel): + """Shared multimodal encoder flow for Gemma4 conditional generation models.""" + + supports_encoder_cache = True + + @property + def multimodal_data_device_paths(self) -> List[str]: + """Dotted multimodal-data paths that the engine transfers to the GPU.""" + return [ + "image.pixel_values", + "image.image_position_ids", + "video.pixel_values", + "video.image_position_ids", + "audio.audio_features", + "audio.audio_features_mask", + ] + + def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) -> torch.Tensor: + """Encode uncached Gemma4 image, video, and audio payloads.""" + modality_inputs = ( + ("image", "pixel_values"), + ("video", "pixel_values"), + ("audio", "audio_features"), + ) + param_modalities = [ + [ + modality + for modality, input_field in modality_inputs + if multimodal_param.multimodal_data.get(modality, {}).get(input_field) is not None + ] + for multimodal_param in multimodal_params + ] + + # The cache helper splits this tensor by request. Preserve that order while still batching + # consecutive requests with the same modality. + if ( + all(len(modalities) == 1 for modalities in param_modalities) + and len({modalities[0] for modalities in param_modalities}) > 1 + ): + embeddings = [] + for _, param_group in groupby( + zip(param_modalities, multimodal_params), key=lambda item: item[0][0] + ): + embeddings.append( + self.encode_multimodal_inputs( + [multimodal_param for _, multimodal_param in param_group] + ) + ) + return torch.cat(embeddings, dim=0) + + pixel_values_list = [] + image_position_ids_list = [] + image_seq_lens_extended: List[int] = [] + audio_features_list = [] + audio_features_mask_list: List[Optional[torch.Tensor]] = [] + video_pixel_values_list = [] + video_position_ids_list = [] + video_seq_lens_extended: List[int] = [] + + for multimodal_param in multimodal_params: + multimodal_data = multimodal_param.multimodal_data + image_data = multimodal_data.get("image", {}) + pixel_values = image_data.get("pixel_values") + if pixel_values is not None: + pixel_values_list.append(pixel_values) + image_position_ids = image_data.get("image_position_ids") + if image_position_ids is not None: + image_position_ids_list.append(image_position_ids) + image_seq_lens = image_data.get("image_seq_lens") + if image_seq_lens is not None: + image_seq_lens_extended.extend(image_seq_lens) + + audio_data = multimodal_data.get("audio", {}) + audio_features = audio_data.get("audio_features") + if audio_features is not None: + audio_features_list.append(audio_features) + audio_features_mask_list.append(audio_data.get("audio_features_mask")) + + video_data = multimodal_data.get("video", {}) + video_pixel_values = video_data.get("pixel_values") + if video_pixel_values is not None: + video_pixel_values_list.append(video_pixel_values) + video_position_ids = video_data.get("image_position_ids") + if video_position_ids is not None: + video_position_ids_list.append(video_position_ids) + video_seq_lens = video_data.get("image_seq_lens") + if video_seq_lens is not None: + video_seq_lens_extended.extend(video_seq_lens) + + multimodal_embeddings = [] + if pixel_values_list: + 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 + ) + multimodal_embeddings.append( + self._get_image_features( + pixel_values=pixel_values, + image_position_ids=image_position_ids, + image_seq_lens=( + image_seq_lens_extended + if len(image_seq_lens_extended) == pixel_values.shape[0] + else None + ), + ) + ) + + if video_pixel_values_list: + 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 + ) + multimodal_embeddings.append( + self._get_image_features( + pixel_values=video_pixel_values, + image_position_ids=video_position_ids, + image_seq_lens=( + video_seq_lens_extended + if len(video_seq_lens_extended) == video_pixel_values.shape[0] + else None + ), + ) + ) + + if audio_features_list and self.embed_audio is not None: + per_audio_embeddings = [ + self._get_audio_features(audio_features, audio_features_mask_list[index]) + for index, audio_features in enumerate(audio_features_list) + ] + multimodal_embeddings.append(torch.cat(per_audio_embeddings, dim=0)) + + if not multimodal_embeddings: + raise ValueError("Gemma4 received active multimodal parameters without encoder inputs") + return torch.cat(multimodal_embeddings, dim=0) + + @register_auto_model("Gemma4ForConditionalGeneration") @register_input_processor( Gemma4InputProcessor, @@ -578,7 +713,7 @@ def call_with_text_prompt( interleave_placeholders=True, ), ) -class Gemma4ForConditionalGeneration(PreTrainedModel): +class Gemma4ForConditionalGeneration(Gemma4MultimodalModelBase): """Gemma4 multimodal model: LLM + vision tower + multimodal embedder. Follows the Gemma3VLM pattern but adapted for Gemma4's architecture: @@ -595,13 +730,6 @@ def get_model_defaults(cls, llm_args) -> dict: "attn_backend": "FLASHINFER", } - def _check_and_adjust_experts_implementation(self, *args, **kwargs): - # transformers 5.x ``PreTrainedModel.__init__`` calls this with an - # ``experts_implementation`` argument and fails for VL wrapper models - # that do not directly contain MoE layers. TRT-LLM manages expert - # implementations independently, so skip the check. - return None - def __init__(self, model_config: ModelConfig[Gemma4Config]): if _is_mm_disagg(): raise NotImplementedError( @@ -775,30 +903,9 @@ def post_config(self): self.config = self.llm.config self.model_config.pretrained_config = self.llm.config - def infer_max_seq_len(self) -> int: - return self.llm.infer_max_seq_len() - @property - def vocab_size_padded(self) -> int: - return self.llm.vocab_size_padded - - @property - def multimodal_data_device_paths(self) -> List[str]: - """Dotted paths in ``multimodal_data`` that the engine should ship to - GPU. Anything not listed stays CPU-resident — notably - ``image.image_seq_lens`` / ``video.image_seq_lens`` (Python - ``List[int]`` carrying per-image valid-patch counts, consumed - host-side by ``Gemma4VisionModel.forward`` to populate - ``attn_metadata.prompt_lens`` without a GPU→CPU sync). - """ - return [ - "image.pixel_values", - "image.image_position_ids", - "video.pixel_values", - "video.image_position_ids", - "audio.audio_features", - "audio.audio_features_mask", - ] + def language_model(self) -> torch.nn.Module: + return self.llm @nvtx_range("[Vision] process") def _get_image_features( @@ -879,197 +986,26 @@ def _get_audio_features( projected = projected.reshape(-1, projected.shape[-1]) return projected.contiguous() - @staticmethod - def _has_active_multimodal_tokens(multimodal_param: MultimodalParams) -> bool: - """Whether a context parameter needs embeddings in this forward.""" - 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: - """Encode uncached Gemma4 image, video, and audio payloads.""" - modality_inputs = ( - ("image", "pixel_values"), - ("video", "pixel_values"), - ("audio", "audio_features"), - ) - param_modalities = [ - [ - modality - for modality, input_field in modality_inputs - if multimodal_param.multimodal_data.get(modality, {}).get(input_field) is not None - ] - for multimodal_param in multimodal_params - ] - - # The cache helper splits this tensor by request. Preserve that order while still batching - # consecutive requests with the same modality. - if ( - all(len(modalities) == 1 for modalities in param_modalities) - and len({modalities[0] for modalities in param_modalities}) > 1 - ): - embeddings = [] - for _, param_group in groupby( - zip(param_modalities, multimodal_params), key=lambda item: item[0][0] - ): - embeddings.append( - self._forward_multimodal_encoder( - [multimodal_param for _, multimodal_param in param_group] - ) - ) - return torch.cat(embeddings, dim=0) - - pixel_values_list = [] - image_position_ids_list = [] - image_seq_lens_extended: List[int] = [] - audio_features_list = [] - audio_features_mask_list: List[Optional[torch.Tensor]] = [] - video_pixel_values_list = [] - video_position_ids_list = [] - video_seq_lens_extended: List[int] = [] - - for multimodal_param in multimodal_params: - multimodal_data = multimodal_param.multimodal_data - image_data = multimodal_data.get("image", {}) - pixel_values = image_data.get("pixel_values") - if pixel_values is not None: - pixel_values_list.append(pixel_values) - image_position_ids = image_data.get("image_position_ids") - if image_position_ids is not None: - image_position_ids_list.append(image_position_ids) - image_seq_lens = image_data.get("image_seq_lens") - if image_seq_lens is not None: - image_seq_lens_extended.extend(image_seq_lens) - - audio_data = multimodal_data.get("audio", {}) - audio_features = audio_data.get("audio_features") - if audio_features is not None: - audio_features_list.append(audio_features) - audio_features_mask_list.append(audio_data.get("audio_features_mask")) - - video_data = multimodal_data.get("video", {}) - video_pixel_values = video_data.get("pixel_values") - if video_pixel_values is not None: - video_pixel_values_list.append(video_pixel_values) - video_position_ids = video_data.get("image_position_ids") - if video_position_ids is not None: - video_position_ids_list.append(video_position_ids) - video_seq_lens = video_data.get("image_seq_lens") - if video_seq_lens is not None: - video_seq_lens_extended.extend(video_seq_lens) - - multimodal_embeddings = [] - if pixel_values_list: - 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 - ) - multimodal_embeddings.append( - self._get_image_features( - pixel_values=pixel_values, - image_position_ids=image_position_ids, - image_seq_lens=( - image_seq_lens_extended - if len(image_seq_lens_extended) == pixel_values.shape[0] - else None - ), - ) - ) - - # Video frames use the same vision tower as images. - if video_pixel_values_list: - 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 - ) - multimodal_embeddings.append( - self._get_image_features( - pixel_values=video_pixel_values, - image_position_ids=video_position_ids, - image_seq_lens=( - video_seq_lens_extended - if len(video_seq_lens_extended) == video_pixel_values.shape[0] - else None - ), - ) - ) - - if audio_features_list and self.audio_tower is not None: - per_audio_embeddings = [ - self._get_audio_features(audio_features, audio_features_mask_list[i]) - for i, audio_features in enumerate(audio_features_list) - ] - multimodal_embeddings.append(torch.cat(per_audio_embeddings, dim=0)) - - if not multimodal_embeddings: - raise ValueError("Gemma4 received active multimodal parameters without encoder inputs") - return torch.cat(multimodal_embeddings, dim=0) - - @torch.inference_mode() - def forward( + def get_language_model_extra_forward_kwargs( 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: Optional[bool] = False, - **kwargs, - ) -> torch.Tensor: - multimodal_params = kwargs.get("multimodal_params", [])[: attn_metadata.num_contexts] - active_multimodal_params = [ - multimodal_param - for multimodal_param in multimodal_params - if self._has_active_multimodal_tokens(multimodal_param) - ] - - mm_embeds = [] - 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) - + *, + raw_input_ids: Optional[torch.Tensor], + position_ids: Optional[torch.Tensor], + mm_inputs: PreparedLlmInputs, + lora_params=None, + **forward_kwargs, + ) -> Dict: + """Build Gemma4-specific language-model forward arguments.""" + del position_ids, forward_kwargs mm_token_type_ids = None - if mm_embeds: - # Build integer token types only when this chunk contains active - # multimodal embeddings. Singleton comparisons avoid repeated - # torch.isin scans over the packed input. - mm_token_type_ids = torch.zeros_like(input_ids, dtype=torch.long) - mm_token_type_ids[input_ids == self.image_token_ids[0]] = 1 + if raw_input_ids is not None and mm_inputs.input_ids is None: + mm_token_type_ids = torch.zeros_like(raw_input_ids, dtype=torch.long) + mm_token_type_ids[raw_input_ids == self.image_token_ids[0]] = 1 if self.video_token_ids is not None: - mm_token_type_ids[input_ids == self.video_token_ids[0]] = 2 + mm_token_type_ids[raw_input_ids == self.video_token_ids[0]] = 2 if self.audio_token_ids is not None: - mm_token_type_ids[input_ids == self.audio_token_ids[0]] = 3 - - # Build a PLE-safe view of the original input_ids where every - # multimodal token is replaced by the text pad_token_id. The - # Gemma4 Per-Layer Embedding lookup uses this view so the PLE table - # is not consulted at audio/image/video positions (matches HF's - # Gemma4Model behaviour). Without this, multimodal requests on - # E2B/E4B (which use PLE) produce garbage output because - # ``fuse_input_embeds`` returns ``input_ids=None`` and PLE is then - # silently skipped inside ``Gemma4TextModel.forward``. + mm_token_type_ids[raw_input_ids == self.audio_token_ids[0]] = 3 + ple_input_ids = None if mm_token_type_ids is not None and self.llm.model.hidden_size_per_layer_input: text_config = getattr(self.config, "text_config", self.config) @@ -1077,30 +1013,27 @@ def forward( if pad_id is not None: ple_input_ids = torch.where( mm_token_type_ids > 0, - torch.full_like(input_ids, pad_id), - input_ids, + torch.full_like(raw_input_ids, pad_id), + raw_input_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=self.mm_token_ids, - mm_token_indices=kwargs.get("mm_token_indices"), - text_token_indices=kwargs.get("text_token_indices"), - ) - logits = 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, - ple_input_ids=ple_input_ids, - lora_params=kwargs.get("lora_params", None), - ) - return logits + return { + "mm_token_type_ids": mm_token_type_ids, + "ple_input_ids": ple_input_ids, + "lora_params": lora_params, + } @property - def mm_token_ids(self) -> torch.Tensor: + def multimodal_token_ids(self) -> torch.Tensor: return self._mm_token_ids + + @property + def text_embedding_layer(self) -> Embedding: + return self.llm.model.embed_tokens + + @property + def embedding_dim(self) -> int: + return self.text_embedding_layer.embedding_dim + + @property + def embedding_dtype(self) -> torch.dtype: + return self.text_embedding_layer.weight.dtype diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 2246a73f8fdd..9d2a71dfe3de 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -799,25 +799,6 @@ def embedding_dim(self) -> int: def embedding_dtype(self) -> torch.dtype: return self.text_embedding_layer.weight.dtype - @property - def draft_config(self): - return self.llm.draft_config - - @property - def draft_model(self): - return self.llm.draft_model - - @property - def load_draft_weights(self): - return self.llm.load_draft_weights - - @property - def vocab_size_padded(self) -> int: - return self.llm.vocab_size_padded - - def infer_max_seq_len(self) -> int: - return self.llm.infer_max_seq_len() - def encode_multimodal_inputs( self, multimodal_params: Sequence[MultimodalParams], @@ -825,81 +806,22 @@ def encode_multimodal_inputs( mm_embeds = self._vision_forward(list(multimodal_params)) return mm_embeds[0] - def get_language_model_forward_kwargs( + def get_language_model_extra_forward_kwargs( self, *, - attn_metadata: AttentionMetadata, - input_ids: torch.Tensor | None, + raw_input_ids: torch.Tensor | None, position_ids: torch.Tensor | None, - inputs_embeds: torch.Tensor | None, mm_inputs: PreparedLlmInputs, - return_context_logits: bool, spec_metadata: SpecMetadata | None, - resource_manager: Any | None, + resource_manager: Any | None = None, + **forward_kwargs: Any, ) -> dict[str, Any]: + del raw_input_ids, position_ids, mm_inputs, forward_kwargs return { - "attn_metadata": attn_metadata, - "input_ids": input_ids, - "position_ids": position_ids, - "inputs_embeds": inputs_embeds, - "return_context_logits": return_context_logits, "spec_metadata": spec_metadata, "resource_manager": resource_manager, } - @torch.inference_mode() - def forward( - self, - attn_metadata: AttentionMetadata, - input_ids: torch.LongTensor | None = None, - position_ids: torch.LongTensor | None = None, - inputs_embeds: torch.FloatTensor | None = None, - return_context_logits: bool = False, - spec_metadata: SpecMetadata | None = None, - **kwargs, - ) -> torch.Tensor: - """Forward method.""" - num_context_requests = attn_metadata.num_contexts - # multimodal_params is consumed by prepare_multimodal_inputs; remove it - # from passthrough kwargs to avoid rebinding it via **kwargs. - multimodal_params = kwargs.pop("multimodal_params", []) - - mm_inputs = self.prepare_multimodal_inputs( - input_ids=input_ids, - positions=position_ids, - multimodal_params=multimodal_params, - num_context_requests=num_context_requests, - attn_metadata=attn_metadata, - **kwargs, - ) - if inputs_embeds is not None: - if mm_inputs.inputs_embeds is not None: - # The caller supplied pre-computed inputs_embeds while the - # multimodal pipeline also produced fused embeds. Refuse to - # silently drop one or the other; let the caller resolve it. - raise ValueError( - "Mistral3VLM.forward received both caller-supplied inputs_embeds " - "and multimodal-derived inputs_embeds. These paths are mutually " - "exclusive; pass at most one.") - mm_inputs = PreparedLlmInputs( - input_ids=None, - inputs_embeds=inputs_embeds, - extra_embeds=mm_inputs.extra_embeds, - ) - - llm_kwargs = self.get_language_model_forward_kwargs( - attn_metadata=attn_metadata, - input_ids=mm_inputs.input_ids, - position_ids=position_ids, - inputs_embeds=mm_inputs.inputs_embeds, - mm_inputs=mm_inputs, - return_context_logits=return_context_logits, - spec_metadata=spec_metadata, - resource_manager=kwargs.get("resource_manager"), - ) - - return self.language_model.forward(**llm_kwargs) - @staticmethod def _get_sub_model_config( model_config: ModelConfig[MistralConfig], diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index b1a2c5df6845..59f1164a9e95 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -197,16 +197,144 @@ def select_multimodal_params( ) -> Sequence[MultimodalParams]: """Select the params that participate in multimodal encoder work. - Returns the context-slice params with multimodal content. Helpers below + Returns the context-slice params with active multimodal content. Helpers below this method (`get_multimodal_embeddings`, `find_input_mm_embeds`, `fuse_input_embeds`) operate on the returned list and therefore see only `has_content()` params. Models overriding this hook must preserve that invariant. """ return [ - param for param in list(multimodal_params)[:num_context_requests] if param.has_content() + param + for param in list(multimodal_params)[:num_context_requests] + if param.has_content() + and ( + param.multimodal_runtime is None + or param.multimodal_runtime.num_mm_tokens_in_chunk != 0 + ) ] + @property + def language_model(self) -> torch.nn.Module: + """Return the inner language model that receives prepared inputs.""" + raise NotImplementedError + + @property + def draft_config(self): + """Expose the inner language model's draft configuration.""" + return self.language_model.draft_config + + @property + def draft_model(self): + """Expose the inner language model's draft model.""" + return self.language_model.draft_model + + def load_draft_weights(self, *args, **kwargs): + """Delegate draft-weight loading to the inner language model.""" + return self.language_model.load_draft_weights(*args, **kwargs) + + @property + def vocab_size_padded(self) -> int: + """Return the inner language model's padded vocabulary size.""" + return self.language_model.vocab_size_padded + + def infer_max_seq_len(self) -> int: + """Return the inner language model's maximum sequence length.""" + return self.language_model.infer_max_seq_len() + + def _check_and_adjust_experts_implementation(self, *args, **kwargs): + """Skip Transformers' MoE validation for multimodal wrapper models.""" + return None + + def get_language_model_extra_forward_kwargs( + self, + *, + raw_input_ids: Optional[torch.Tensor], + position_ids: Optional[torch.Tensor], + mm_inputs: PreparedLlmInputs, + **forward_kwargs: Any, + ) -> dict[str, Any]: + """Return model-specific arguments for the inner language-model forward.""" + return {} + + def get_language_model_forward_kwargs( + self, + *, + attn_metadata: Any, + input_ids: Optional[torch.Tensor], + raw_input_ids: Optional[torch.Tensor], + position_ids: Optional[torch.Tensor], + inputs_embeds: Optional[torch.Tensor], + mm_inputs: PreparedLlmInputs, + return_context_logits: bool, + **forward_kwargs: Any, + ) -> dict[str, Any]: + """Build common and model-specific inner language-model forward arguments.""" + llm_kwargs = { + "attn_metadata": attn_metadata, + "input_ids": input_ids, + "position_ids": position_ids, + "inputs_embeds": inputs_embeds, + "return_context_logits": return_context_logits, + } + llm_kwargs.update( + self.get_language_model_extra_forward_kwargs( + raw_input_ids=raw_input_ids, + position_ids=position_ids, + mm_inputs=mm_inputs, + **forward_kwargs, + ) + ) + return llm_kwargs + + @torch.inference_mode() + def forward( + self, + attn_metadata: Any, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Any = None, + **kwargs: Any, + ) -> torch.Tensor: + """Prepare multimodal inputs and dispatch them to the language model.""" + multimodal_params = kwargs.pop("multimodal_params", []) + mm_inputs = self.prepare_multimodal_inputs( + input_ids=input_ids, + positions=position_ids, + multimodal_params=multimodal_params, + num_context_requests=attn_metadata.num_contexts, + attn_metadata=attn_metadata, + **kwargs, + ) + if inputs_embeds is not None: + if mm_inputs.inputs_embeds is not None: + raise ValueError( + "MultimodalModelMixin.forward received both caller-supplied inputs_embeds " + "and multimodal-derived inputs_embeds. These paths are mutually exclusive; " + "pass at most one." + ) + mm_inputs = PreparedLlmInputs( + input_ids=None, + inputs_embeds=inputs_embeds, + extra_embeds=mm_inputs.extra_embeds, + ) + + llm_kwargs = self.get_language_model_forward_kwargs( + attn_metadata=attn_metadata, + input_ids=mm_inputs.input_ids, + raw_input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=mm_inputs.inputs_embeds, + mm_inputs=mm_inputs, + return_context_logits=return_context_logits, + multimodal_params=multimodal_params, + num_generation_requests=attn_metadata.num_generations, + spec_metadata=spec_metadata, + **kwargs, + ) + return self.language_model.forward(**llm_kwargs) + def after_full_multimodal_embeddings( self, *, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 6fc43129a353..84be9d691b85 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -663,6 +663,8 @@ class _Qwen3_5VLModel(Qwen3VLModelBase): decorators (outer arch string + input-processor `model_type`). """ + supports_encoder_cache = True + @classmethod def get_model_defaults(cls, llm_args): # `ModelLoader` applies `get_model_defaults()` on the resolved outer diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 2ae083d77355..ab2e54cf14f9 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -37,6 +37,7 @@ from ..attention_backend import AttentionMetadata from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams from ..attention_backend.utils import get_attention_backend +from ..modules.embedding import Embedding from ..modules.layer_norm import LayerNorm from ..modules.linear import Linear, TensorParallelMode from ..modules.mlp import MLP @@ -45,14 +46,7 @@ from .checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper from .modeling_auto import AutoModelForCausalLM from .modeling_multimodal_encoder import MultimodalEncoderMixin -from .modeling_multimodal_mixin import MultimodalModelMixin -from .modeling_multimodal_utils import ( - filter_mm_token_from_input_ids, - find_input_mm_embeds, - fuse_input_embeds, - get_attached_multimodal_embeddings, - get_multimodal_embeddings, -) +from .modeling_multimodal_mixin import MultimodalModelMixin, PreparedLlmInputs from .modeling_qwen2vl import ( Qwen2_5_VLVisionAttention, Qwen2VLInputProcessorBase, @@ -1219,7 +1213,7 @@ def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tenso return embeds -class Qwen3VLModelBase(PreTrainedModel, MultimodalModelMixin): +class Qwen3VLModelBase(MultimodalModelMixin, PreTrainedModel): def encode_multimodal_inputs( self, multimodal_params: List[MultimodalParams], **encoder_kwargs: Any ) -> torch.Tensor: @@ -1227,25 +1221,19 @@ def encode_multimodal_inputs( Runs the vision encoder over ``multimodal_params`` and returns the embeddings as a single tensor (Qwen3-VL folds deepstack streams into - the hidden dim, so the single-tensor contract holds). Used by the - startup memory profiler to invoke the encoder directly; the model's - own ``forward`` keeps its custom deepstack fusion path. + the hidden dim, so the single-tensor contract holds). """ - mm_embeds = get_multimodal_embeddings( - encoder_forward_fn=self.mm_encoder.forward, multimodal_params=list(multimodal_params) - ) + if self.mm_encoder is None: + raise ValueError("Raw multimodal inputs require a local multimodal encoder.") + + mm_embeds = self.mm_encoder.forward(list(multimodal_params), **encoder_kwargs) + if len(mm_embeds) != 1: + raise ValueError( + "Qwen3-VL multimodal encoder must return one packed embedding tensor, " + f"but returned {len(mm_embeds)} tensors." + ) return mm_embeds[0] - def _check_and_adjust_experts_implementation(self, *args, **kwargs): - """No-op override. - - Transformers 5.x's ``PreTrainedModel.__init__`` calls this method - (with an ``experts_implementation`` argument) which fails for VL - wrapper models that do not directly contain MoE layers. TRT-LLM - manages expert implementations independently, so skip the check. - """ - return None - def __init__( self, model_config: ModelConfig[PretrainedConfig], @@ -1343,22 +1331,6 @@ def __init__( self.deepstack_num_level = ( len(config.vision_config.deepstack_visual_indexes) if self.use_deepstack else 0 ) - if self.use_deepstack: - # Pre-allocated `(L, max_num_tokens, hidden)` scratch buffer for - # per-layer deepstack embeddings; replaces `L` fresh - # `torch.zeros` + `L` scatters per prefill. - # `persistent=False` keeps it out of `state_dict`. - self.register_buffer( - "deepstack_input_embeds", - torch.zeros( - self.deepstack_num_level, - model_config.max_num_tokens, - config.text_config.hidden_size, - device="cuda", - dtype=config.text_config.torch_dtype, - ), - persistent=False, - ) # Surface the in-vocab image / video placeholder IDs to the model # engine's ``_prepare_multimodal_indices`` so it selects the @@ -1375,40 +1347,29 @@ def __init__( self.post_config() @property - def mm_token_ids(self) -> torch.Tensor: + def multimodal_token_ids(self) -> torch.Tensor: return self._mm_token_ids - def post_config(self): - # use llm.config as config for pytorch model engine - self.model_config.pretrained_config = self.llm.config - self.config = self.model_config.pretrained_config + @property + def language_model(self) -> torch.nn.Module: + return self.llm @property - def vocab_size_padded(self) -> int: - return self.llm.vocab_size_padded - - def infer_max_seq_len(self) -> int: - return self.llm.infer_max_seq_len() - - # Draft-model (two-model speculative decoding, e.g. DFlash / Eagle3) - # delegation: `ModelLoader.load` reads `draft_config` / `draft_model` and - # calls `load_draft_weights` on the *outer* model it resolved, but the - # spec-decoding wrapper (`SpecDecOneEngineForCausalLM`) is applied to the - # inner `self.llm` when this VLM composes it. Composite checkpoints - # (e.g. Qwen3.5-4B publishes text_config + vision_config) route text-only - # spec tests through this wrapper, so surface the inner LM's draft state. - # Note: `load_draft_weights` must keep an explicit signature — the loader - # dispatches kwargs via `inspect.getfullargspec`. + def text_embedding_layer(self) -> Embedding: + return self.llm.model.embed_tokens + @property - def draft_config(self): - return self.llm.draft_config + def embedding_dim(self) -> int: + return self.text_embedding_layer.embedding_dim * (self.deepstack_num_level + 1) @property - def draft_model(self): - return self.llm.draft_model + def embedding_dtype(self) -> torch.dtype: + return self.text_embedding_layer.weight.dtype - def load_draft_weights(self, weights: Dict, weight_mapper: Optional[BaseWeightMapper] = None): - return self.llm.load_draft_weights(weights, weight_mapper=weight_mapper) + def post_config(self): + # use llm.config as config for pytorch model engine + self.model_config.pretrained_config = self.llm.config + self.config = self.model_config.pretrained_config def apply_llm_torch_compile(self, *, backend: Any, fullgraph: bool) -> None: # TODO: Move this hook to MultimodalModelMixin once multimodal models @@ -1459,138 +1420,90 @@ def split_mm_embeds(self, mm_embed, deepstack_num_level): mm_embed_chunks = torch.split(mm_embed, [num_elements] * (deepstack_num_level + 1), dim=1) return mm_embed_chunks[0], list(mm_embed_chunks[1:]) - @torch.inference_mode() - def forward( + def select_multimodal_params( self, - attn_metadata: AttentionMetadata, - input_ids: Optional[torch.IntTensor] = None, - position_ids: Optional[torch.IntTensor] = None, - input_embeds: Optional[torch.Tensor] = None, - return_context_logits: bool = False, - **kwargs, - ) -> torch.Tensor: - """ - VLM forward logic with inflight batching support. - """ - num_context_requests, num_generation_requests = ( - attn_metadata.num_contexts, - attn_metadata.num_generations, + multimodal_params: List[MultimodalParams], + num_context_requests: int, + ) -> List[MultimodalParams]: + """Select requests with image/video embeddings for the current context batch.""" + context_params = super().select_multimodal_params(multimodal_params, num_context_requests) + multimodal_params, has_raw_image_or_video_data = self._get_requests_with_mm_data( + context_params ) + if not multimodal_params: + return [] + if has_raw_image_or_video_data and self.mm_encoder is None: + raise ValueError( + "Raw multimodal inputs require a local multimodal encoder on this " + "worker, or multimodal_embedding handles from an encoder handoff." + ) + if not has_raw_image_or_video_data and not getattr(self, "support_mm_disagg", False): + raise NotImplementedError( + f"{type(self)} does not support disaggregated inference yet. Please unset " + "the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'." + ) + return multimodal_params - multimodal_params = kwargs.get("multimodal_params", []) - mm_embeds = [] - mrope_config = {} - deepstack_embeds = [] + def after_active_multimodal_embeddings( + self, + *, + active_embeddings: List[torch.Tensor], + multimodal_params: List[MultimodalParams], + **forward_kwargs: Any, + ) -> tuple[List[torch.Tensor], List[torch.Tensor]]: + """Separate Qwen3-VL's packed deepstack streams from primary embeddings.""" + del multimodal_params, forward_kwargs + if not self.use_deepstack: + return active_embeddings, [] - # NOTE: Qwen*-VL series has mrope_config even on the text-only prompts, - # so we need to separate the mm_multimodal_params from the text-only prompts. - if num_context_requests > 0: - mm_multimodal_params, has_raw_image_or_video_data = self._get_requests_with_mm_data( - multimodal_params[:num_context_requests] + deepstack_embeds = [] + for index, mm_embed in enumerate(active_embeddings): + active_embeddings[index], deepstack_embed = self.split_mm_embeds( + mm_embed, self.deepstack_num_level ) - else: - mm_multimodal_params = [] - has_raw_image_or_video_data = False - if len(mm_multimodal_params) > 0: - # Raw image/video tensors: run local encoder. - if has_raw_image_or_video_data and self.mm_encoder is not None: - mm_embeds = get_multimodal_embeddings( - encoder_forward_fn=self.mm_encoder.forward, - multimodal_params=mm_multimodal_params, - ) - # Raw image/video tensors on a worker with no encoder: bad route. - elif has_raw_image_or_video_data: - raise ValueError( - "Raw multimodal inputs require a local multimodal encoder on this " - "worker, or multimodal_embedding handles from an encoder handoff." - ) - # support_mm_disagg is only set in subclasses of Qwen3VLModelBase that support EPD - elif not getattr(self, "support_mm_disagg", False): - raise NotImplementedError( - f"{type(self)} does not support disaggregated inference yet. Please unset " - "the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'." - ) - # E/P prefill: encoder already ran; use attached embeddings. - else: - mm_embeds = get_attached_multimodal_embeddings(mm_multimodal_params) - mm_embeds = find_input_mm_embeds(mm_embeds, mm_multimodal_params) - - if self.use_deepstack: - for i, mm_embed in enumerate(mm_embeds): - mm_embed, deepstack_embed = self.split_mm_embeds( - mm_embed, self.deepstack_num_level - ) - mm_embeds[i] = mm_embed - deepstack_embeds.extend(deepstack_embed) + deepstack_embeds.extend(deepstack_embed) + return active_embeddings, deepstack_embeds + def get_language_model_extra_forward_kwargs( + self, + *, + raw_input_ids: Optional[torch.Tensor], + position_ids: Optional[torch.Tensor], + mm_inputs: PreparedLlmInputs, + multimodal_params: List[MultimodalParams], + num_generation_requests: int, + spec_metadata: Any, + resource_manager: Any = None, + mrope_delta_write_seq_slots: Optional[torch.Tensor] = None, + mrope_delta_read_seq_slots: Optional[torch.Tensor] = None, + mm_token_indices: Optional[torch.Tensor] = None, + **forward_kwargs: Any, + ) -> Dict[str, Any]: + """Build Qwen3-VL-specific language-model forward arguments.""" + del forward_kwargs + mrope_config = {} if not self.model_config.pretrained_config.disable_fuse_rope: mrope_config = self.prepare_mrope_config( multimodal_params, num_generation_requests, position_ids, - mrope_delta_write_seq_slots=kwargs.get("mrope_delta_write_seq_slots"), - mrope_delta_read_seq_slots=kwargs.get("mrope_delta_read_seq_slots"), - ) - - # Prefer the indices the executor already computed (CPU-side - # `filter_mm_token_from_input_ids` + async H2D) and forwarded via - # kwargs; fall back to filtering only on engine-bypass paths - # (e.g., direct `forward` calls in unit tests). - text_token_indices = kwargs.get("text_token_indices") - mm_token_indices = kwargs.get("mm_token_indices") - if len(mm_embeds) > 0 and (text_token_indices is None or mm_token_indices is None): - text_token_indices, mm_token_indices = filter_mm_token_from_input_ids( - input_ids, - vocab_size=self.llm.model.embed_tokens.num_embeddings, - mm_token_ids=self.mm_token_ids, - ) - - # Expand the per-level deepstack mm embeddings into the pre-allocated - # `(L, max_num_tokens, H)` buffer with a single packed scatter, - # avoiding `L` fresh `torch.zeros` + `L` scatters inside - # `fuse_input_embeds`. - if self.use_deepstack and len(deepstack_embeds) > 0: - num_tokens = input_ids.shape[0] - deepstack_buffer = self.deepstack_input_embeds[:, :num_tokens, :] - deepstack_buffer.zero_() - packed_deepstack = torch.stack(deepstack_embeds, dim=0) - deepstack_buffer[:, mm_token_indices, :] = packed_deepstack.to( - dtype=deepstack_buffer.dtype, device=deepstack_buffer.device + mrope_delta_write_seq_slots=mrope_delta_write_seq_slots, + mrope_delta_read_seq_slots=mrope_delta_read_seq_slots, ) - deepstack_embeds = list(deepstack_buffer.unbind(0)) - - # Preserve the pre-fusion token IDs. `fuse_input_embeds` collapses - # input_ids -> None when MM embeddings are fused in, but spec - # decoding (MTP / Eagle) still needs the original prompt token - # IDs for drafter context preparation; pass them through as a - # dedicated kwarg consumed by `SpecDecOneEngineForCausalLM.forward`. - orig_input_ids = input_ids - input_ids, input_embeds = fuse_input_embeds( - self.llm.model.embed_tokens, - input_ids, - mm_embeds, - text_token_indices=text_token_indices, - mm_token_indices=mm_token_indices, - ) - - output_prob = self.llm.forward( - attn_metadata=attn_metadata, - input_ids=input_ids, - position_ids=position_ids, - inputs_embeds=input_embeds, - return_context_logits=return_context_logits, - deepstack_embeds=deepstack_embeds, - mrope_config=mrope_config, - spec_metadata=kwargs.get("spec_metadata"), - resource_manager=kwargs.get("resource_manager"), - orig_input_ids=orig_input_ids, - ) - # Spec-decoding (MTP / Eagle) returns a dict (accepted tokens, - # draft tokens, logits); plain forward returns a tensor. - if hasattr(output_prob, "shape"): - logger.debug(f"output shape: {output_prob.shape}") - return output_prob + deepstack_embeds = list(mm_inputs.extra_embeds) + # `prepare_multimodal_inputs` passes these through `fuse_input_embeds`, which has already + # expanded each packed deepstack feature to the full input sequence. Do not scatter them + # a second time here: their leading dimension is now `num_tokens`, not the number of + # multimodal placeholders. + + return { + "deepstack_embeds": deepstack_embeds, + "mrope_config": mrope_config, + "spec_metadata": spec_metadata, + "resource_manager": resource_manager, + "orig_input_ids": raw_input_ids, + } def _get_requests_with_mm_data(self, multimodal_params): mm_multimodal_params = [] @@ -1629,6 +1542,8 @@ def _get_requests_with_mm_data(self, multimodal_params): ), ) class Qwen3VLModel(Qwen3VLModelBase): + supports_encoder_cache = True + def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): # NOTE: HF implementation. kwargs["vision_model_class"] = Qwen3VisionModel diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py index fd7fd20f323f..41912d28a0f9 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py @@ -45,6 +45,8 @@ ), ) class Qwen3MoeVLModel(Qwen3VLModelBase): + supports_encoder_cache = True + def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): # NOTE: HF implementation. kwargs["vision_model_class"] = Qwen3VisionModel diff --git a/tensorrt_llm/_torch/models/modeling_qwen_image_bench.py b/tensorrt_llm/_torch/models/modeling_qwen_image_bench.py index 061d21bc0df0..ecbd03499641 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen_image_bench.py +++ b/tensorrt_llm/_torch/models/modeling_qwen_image_bench.py @@ -1,12 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from dataclasses import replace from typing import Dict, List, Optional import torch from transformers import PretrainedConfig from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg +from tensorrt_llm.llmapi.llm_args import MultimodalConfig from ...inputs import ( ContentFormat, @@ -95,6 +97,16 @@ def load_weights( ) class QwenImageBenchModel(_QwenImageBenchModelMixin, Qwen3VLModelBase): def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): + if model_config.multimodal_config is not None: + model_config = replace( + model_config, + multimodal_config=MultimodalConfig( + **model_config.multimodal_config.model_dump( + exclude={"encoder_cache_max_bytes"} + ), + encoder_cache_max_bytes=0, + ), + ) kwargs["vision_model_class"] = Qwen3VisionModel kwargs["disable_fuse_rope"] = kwargs.get("disable_fuse_rope", False) super().__init__(model_config, *args, **kwargs) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 50aa317e9e43..5425bb81a3d8 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2995,8 +2995,11 @@ def get_padded_piecewise_tokens(tokens): def _prepare_multimodal_indices(self, input_ids: list[int]): input_ids = torch.tensor(input_ids, dtype=torch.int, device="cpu") vocab_size = self.model.config.vocab_size - # TODO: unify naming of mm_token_ids across models - mm_token_ids = getattr(self.model, "mm_token_ids", None) + # `multimodal_token_ids` is the common wrapper-model contract. Keep the legacy name as a + # fallback for models not yet migrated to `MultimodalModelMixin`. + mm_token_ids = getattr(self.model, "multimodal_token_ids", None) + if mm_token_ids is None: + mm_token_ids = getattr(self.model, "mm_token_ids", None) text_token_indices, mm_token_indices = filter_mm_token_from_input_ids( input_ids, vocab_size=vocab_size, mm_token_ids=mm_token_ids) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 0945db4d2394..61ddc4bdbb4f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -51,9 +51,6 @@ CacheLevel = NewType("CacheLevel", int) TokenId = NewType("TokenId", int) TokenIdExt = Union[TokenId, bytes] -class PlannedDropHandle: - def drop(self) -> None: ... - class ReuseScope(NamedTuple): lora_id: int | None = None salt: int | None = None @@ -352,7 +349,6 @@ class _KVCache: def committed_tokens(self) -> list[TokenIdExt]: ... @property def reuse_scope(self) -> ReuseScope: ... - def plan_committed_block_drop(self) -> PlannedDropHandle | None: ... def stop_committing(self) -> None: ... def suspend(self) -> None: ... def resume(self, cuda_stream: CudaStream | None = None) -> bool: ... diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 51c36071274c..179c6cb7149d 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -74,6 +74,18 @@ def forward(self, *args, **kwargs) -> torch.Tensor: return {"logits": torch.randn((batch_size, 10), device='cuda')} +class DummyMultimodalIndexModel(torch.nn.Module): + + class Config: + vocab_size = 100 + + config = Config() + + @property + def multimodal_token_ids(self) -> torch.Tensor: + return torch.tensor([90, 91], dtype=torch.int32) + + class DummyModelEngine(PyTorchModelEngine): def __init__(self, llm_args: TorchLlmArgs, dtype: torch.dtype) -> None: @@ -150,6 +162,16 @@ def create_model_engine_and_kvcache(llm_args: TorchLlmArgs = None, class PyTorchModelEngineTestCase(unittest.TestCase): + def test_prepare_multimodal_indices_uses_mixin_token_ids(self) -> None: + engine = object.__new__(PyTorchModelEngine) + engine.model = DummyMultimodalIndexModel() + + text_indices, multimodal_indices = engine._prepare_multimodal_indices( + [1, 90, 2, 91, 3]) + + torch.testing.assert_close(text_indices, torch.tensor([0, 2, 4])) + torch.testing.assert_close(multimodal_indices, torch.tensor([1, 3])) + def test_build_request_multimodal_input_skips_when_cache_disabled( self) -> None: request = LlmRequest( diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index 1424a5e840a4..0080d35a079e 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -62,11 +62,17 @@ Gemma4ForConditionalGeneration, Gemma4MultimodalEmbedder, ) +from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin # noqa: E402 from tensorrt_llm._torch.models.modeling_multimodal_utils import ( # noqa: E402 find_input_mm_embeds, get_multimodal_embeddings, ) -from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData # noqa: E402 +from tensorrt_llm.inputs.multimodal import ( # noqa: E402 + MultimodalInput, + MultimodalParams, + MultimodalRuntimeData, +) +from tensorrt_llm.llmapi.llm_args import MultimodalConfig # noqa: E402 from tensorrt_llm.mapping import Mapping # noqa: E402 # --------------------------------------------------------------------------- @@ -125,6 +131,59 @@ } +class _Gemma4EncoderCacheHarness(MultimodalModelMixin): + """Lightweight Gemma4 encoder-cache harness without model weights.""" + + supports_encoder_cache = True + encode_multimodal_inputs = Gemma4ForConditionalGeneration.encode_multimodal_inputs + + def __init__(self, embedding_dim: int = 12) -> None: + self.model_config = ModelConfig( + multimodal_config=MultimodalConfig(encoder_cache_max_bytes=4096) + ) + self._embedding_dim = embedding_dim + self.encoder_calls = 0 + self.audio_tower = None + + @property + def embedding_dim(self) -> int: + return self._embedding_dim + + @property + def embedding_dtype(self) -> torch.dtype: + return torch.float32 + + def _get_image_features(self, pixel_values: torch.Tensor, **kwargs) -> torch.Tensor: + del kwargs + self.encoder_calls += 1 + return torch.full( + (pixel_values.shape[0] * 2, self.embedding_dim), + float(self.encoder_calls), + dtype=self.embedding_dtype, + ) + + +def _make_keyed_image_param() -> MultimodalParams: + embedding_lengths = [2] + return MultimodalParams( + multimodal_input=MultimodalInput( + multimodal_hashes=[[1, 2, 3, 4, 5, 6, 7, 8]], + multimodal_positions=[0], + multimodal_lengths=embedding_lengths, + ), + multimodal_data={ + "image": {"pixel_values": torch.empty(1, 1, 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), + past_seen_token_num=0, + chunk_end_pos=2, + ), + ) + + # Mirror the engine's encoder runtime sizes (``get_encoder_runtime_sizes`` -> # ``encoder_max_batch_size`` / ``encoder_max_num_tokens``, defaulting to # ``max_batch_size`` / ``max_num_tokens``). The encoder ``AttentionMetadata`` is @@ -702,6 +761,19 @@ def test_instantiation_with_vision(self): # TRT-LLM class, not ``transformers.AutoModel`` output. self.assertIsInstance(model.vision_tower, Gemma4VisionModel) + def test_encoder_cache_reuses_image_embedding_across_requests(self): + """Persistent cache reuse applies to the shared dense/MoE Gemma4 wrapper.""" + self.assertTrue(issubclass(Gemma4ForConditionalGeneration, MultimodalModelMixin)) + self.assertTrue(Gemma4ForConditionalGeneration.supports_encoder_cache) + + model = _Gemma4EncoderCacheHarness() + first = model._get_or_encode_multimodal_embeddings([_make_keyed_image_param()]) + second = model._get_or_encode_multimodal_embeddings([_make_keyed_image_param()]) + + self.assertEqual(model.encoder_calls, 1) + torch.testing.assert_close(second, first) + self.assertEqual(len(model._multimodal_encoder_cache), 1) + def test_chunked_prefill_reuses_cached_vision_embeddings(self): """Later active chunks slice cached features without rerunning vision.""" model = self._make_model() @@ -724,7 +796,7 @@ def test_chunked_prefill_reuses_cached_vision_embeddings(self): model, "_get_image_features", return_value=expected_embeddings ) as image_encoder: all_embeddings = get_multimodal_embeddings( - model._forward_multimodal_encoder, [multimodal_param] + model.encode_multimodal_inputs, [multimodal_param] ) first_chunk = find_input_mm_embeds(all_embeddings, [multimodal_param]) @@ -734,7 +806,7 @@ def test_chunked_prefill_reuses_cached_vision_embeddings(self): embed_mask_cumsum=embed_mask_cumsum, ) all_embeddings = get_multimodal_embeddings( - model._forward_multimodal_encoder, [multimodal_param] + model.encode_multimodal_inputs, [multimodal_param] ) second_chunk = find_input_mm_embeds(all_embeddings, [multimodal_param]) @@ -752,8 +824,9 @@ def test_chunk_without_multimodal_tokens_is_inactive(self): embed_mask_cumsum=torch.tensor([0, 0, 1, 2], dtype=torch.int64), ), ) - self.assertFalse( - Gemma4ForConditionalGeneration._has_active_multimodal_tokens(multimodal_param) + self.assertEqual( + Gemma4ForConditionalGeneration.select_multimodal_params(None, [multimodal_param], 1), + [], ) def test_mixed_modality_batch_preserves_request_order(self): @@ -791,7 +864,7 @@ def test_mixed_modality_batch_preserves_request_order(self): "_get_image_features", side_effect=lambda pixel_values, **_: pixel_values[:, 0], ): - embeddings = get_multimodal_embeddings(model._forward_multimodal_encoder, params) + embeddings = get_multimodal_embeddings(model.encode_multimodal_inputs, params) expected = torch.tensor([[2.0], [1.0], [3.0]]) torch.testing.assert_close(embeddings[0], expected) @@ -815,7 +888,7 @@ def test_single_request_with_multiple_modalities_is_allowed(self): "_get_image_features", side_effect=lambda pixel_values, **_: pixel_values[:, 0], ): - embeddings = model._forward_multimodal_encoder([multimodal_param]) + embeddings = model.encode_multimodal_inputs([multimodal_param]) torch.testing.assert_close(embeddings, torch.tensor([[1.0], [2.0]])) diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index b01d961655b6..7b7539cf0b60 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -4,6 +4,7 @@ import copy import os from dataclasses import dataclass +from types import SimpleNamespace from typing import List, Optional import pytest @@ -16,6 +17,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper +from tensorrt_llm._torch.models.modeling_multimodal_mixin import PreparedLlmInputs from tensorrt_llm._torch.models.modeling_qwen3vl import ( Qwen3VisionModel, Qwen3VLInputProcessorBase, @@ -84,6 +86,39 @@ } +def test_qwen3vl_forwards_fused_deepstack_embeddings_without_rescattering(): + """Deepstack features are sequence-shaped after the shared fusion step.""" + hidden_size = 4 + model = object.__new__(Qwen3VLModel) + torch.nn.Module.__init__(model) + model.use_deepstack = True + model.model_config = SimpleNamespace(pretrained_config=SimpleNamespace(disable_fuse_rope=True)) + model.llm = SimpleNamespace( + model=SimpleNamespace( + embed_tokens=SimpleNamespace(num_embeddings=32), + ) + ) + + deepstack_embeds = [torch.randn(5, hidden_size) for _ in range(3)] + result = model.get_language_model_extra_forward_kwargs( + raw_input_ids=torch.tensor([1, 7, 2, 7, 3]), + position_ids=None, + mm_inputs=PreparedLlmInputs( + input_ids=None, + inputs_embeds=torch.randn(5, hidden_size), + extra_embeds=deepstack_embeds, + ), + multimodal_params=[], + num_generation_requests=0, + spec_metadata=None, + mm_token_indices=torch.tensor([1, 3]), + ) + + assert len(result["deepstack_embeds"]) == len(deepstack_embeds) + for actual, expected in zip(result["deepstack_embeds"], deepstack_embeds): + assert actual is expected + + @dataclass(repr=False) class TestQwen3VLScenario(MultimodalScenario): disable_fuse_rope: bool = False From 99db94ea0dd541f089536729f11e300ca5359997 Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Thu, 23 Jul 2026 14:07:51 -0700 Subject: [PATCH 2/2] trim fat Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- docs/source/models/supported-models.md | 1 + .../_torch/models/modeling_gemma4_unified.py | 22 +- .../_torch/models/modeling_gemma4mm.py | 191 +++++++++--------- .../_torch/models/modeling_mistral.py | 8 + .../models/modeling_multimodal_mixin.py | 21 +- .../_torch/models/modeling_qwen3vl.py | 103 +++++++++- .../runtime/kv_cache_manager_v2/__init__.pyi | 4 + .../_torch/modeling/test_modeling_qwen3vl.py | 133 ++++++++---- .../test_mm_encoder_cross_iter_prefetch.py | 8 + 9 files changed, 327 insertions(+), 164 deletions(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 9b4c5653a70e..679e17627d80 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -132,6 +132,7 @@ The following optimizations are available to models that implement | Model Architecture | Multimodal Encoder Side Stream | Multimodal Embeddings Cache | | ------------------ | ------------------------------ | --------------------------- | | `Gemma4ForConditionalGeneration` | Yes | Yes | +| `Gemma4UnifiedForConditionalGeneration` | Yes | Yes | | `Mistral3ForConditionalGeneration` | Yes | Yes | | `Qwen3VLForConditionalGeneration` | Yes | Yes | | `Qwen3VLMoeForConditionalGeneration` | Yes | Yes | diff --git a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py index d0cae7e25ffc..5b47bd3cc0f2 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py @@ -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 @@ -61,7 +60,6 @@ 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 @@ -81,9 +79,9 @@ from ..modules.linear import Linear from .modeling_gemma4 import Gemma4ForCausalLM from .modeling_gemma4mm import ( - Gemma4ForConditionalGeneration, Gemma4InputProcessor, Gemma4MultimodalEmbedder, + Gemma4MultimodalModelBase, ) from .modeling_utils import ModelConfig, filter_weights, register_auto_model @@ -223,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 diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index be999b51d973..46d885cd60e5 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -550,6 +550,20 @@ class Gemma4MultimodalModelBase(MultimodalModelMixin, PreTrainedModel): supports_encoder_cache = True + @classmethod + def get_model_defaults(cls, llm_args) -> dict: + """Gemma4-specific defaults — see Gemma4ForCausalLM.get_model_defaults.""" + return { + "attn_backend": "FLASHINFER", + } + + def _check_and_adjust_experts_implementation(self, *args, **kwargs): + # transformers 5.x ``PreTrainedModel.__init__`` calls this with an + # ``experts_implementation`` argument and fails for VL wrapper models + # that do not directly contain MoE layers. TRT-LLM manages expert + # implementations independently, so skip the check. + return None + @property def multimodal_data_device_paths(self) -> List[str]: """Dotted multimodal-data paths that the engine transfers to the GPU.""" @@ -684,6 +698,91 @@ def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) -> raise ValueError("Gemma4 received active multimodal parameters without encoder inputs") return torch.cat(multimodal_embeddings, dim=0) + @staticmethod + def get_sub_model_config( + model_config: ModelConfig[Gemma4Config], + name: str, + ) -> ModelConfig: + assert name in ["text_config", "vision_config", "audio_config"], ( + f"Expected subconfig name to be 'text_config', 'vision_config', " + f"or 'audio_config'. Got {name} instead." + ) + pretrained_config = getattr(model_config.pretrained_config, name) + quant_config = model_config.quant_config if name == "text_config" else None + preferred_backend = "FLASHINFER" if name == "text_config" else "TRTLLM" + sub_config: ModelConfig = dataclasses.replace( + model_config, + pretrained_config=pretrained_config, + attn_backend=preferred_backend, + quant_config=quant_config, + ) + if ( + hasattr(sub_config.pretrained_config, "torch_dtype") + and sub_config.pretrained_config.torch_dtype is None + ): + sub_config.pretrained_config.torch_dtype = model_config.pretrained_config.torch_dtype + return sub_config + + def post_config(self): + self.config = self.llm.config + self.model_config.pretrained_config = self.llm.config + + @property + def language_model(self) -> torch.nn.Module: + return self.llm + + def get_language_model_extra_forward_kwargs( + self, + *, + raw_input_ids: Optional[torch.Tensor], + position_ids: Optional[torch.Tensor], + mm_inputs: PreparedLlmInputs, + lora_params=None, + **forward_kwargs, + ) -> Dict: + """Build Gemma4-specific language-model forward arguments.""" + del position_ids, forward_kwargs + mm_token_type_ids = None + if raw_input_ids is not None and mm_inputs.input_ids is None: + mm_token_type_ids = torch.zeros_like(raw_input_ids, dtype=torch.long) + mm_token_type_ids[raw_input_ids == self.image_token_ids[0]] = 1 + if self.video_token_ids is not None: + mm_token_type_ids[raw_input_ids == self.video_token_ids[0]] = 2 + if self.audio_token_ids is not None: + mm_token_type_ids[raw_input_ids == self.audio_token_ids[0]] = 3 + + ple_input_ids = None + if mm_token_type_ids is not None and self.llm.model.hidden_size_per_layer_input: + text_config = getattr(self.config, "text_config", self.config) + pad_id = getattr(text_config, "pad_token_id", None) + if pad_id is not None: + ple_input_ids = torch.where( + mm_token_type_ids > 0, + torch.full_like(raw_input_ids, pad_id), + raw_input_ids, + ) + return { + "mm_token_type_ids": mm_token_type_ids, + "ple_input_ids": ple_input_ids, + "lora_params": lora_params, + } + + @property + def multimodal_token_ids(self) -> torch.Tensor: + return self._mm_token_ids + + @property + def text_embedding_layer(self) -> Embedding: + return self.llm.model.embed_tokens + + @property + def embedding_dim(self) -> int: + return self.text_embedding_layer.embedding_dim + + @property + def embedding_dtype(self) -> torch.dtype: + return self.text_embedding_layer.weight.dtype + @register_auto_model("Gemma4ForConditionalGeneration") @register_input_processor( @@ -723,13 +822,6 @@ class Gemma4ForConditionalGeneration(Gemma4MultimodalModelBase): - mm_token_type_ids-based bidirectional masking """ - @classmethod - def get_model_defaults(cls, llm_args) -> dict: - """Gemma4-specific defaults — see Gemma4ForCausalLM.get_model_defaults.""" - return { - "attn_backend": "FLASHINFER", - } - def __init__(self, model_config: ModelConfig[Gemma4Config]): if _is_mm_disagg(): raise NotImplementedError( @@ -843,31 +935,6 @@ def __init__(self, model_config: ModelConfig[Gemma4Config]): self.post_config() self.is_loaded = True - @staticmethod - def get_sub_model_config( - model_config: ModelConfig[Gemma4Config], - name: str, - ) -> ModelConfig: - assert name in ["text_config", "vision_config", "audio_config"], ( - f"Expected subconfig name to be 'text_config', 'vision_config', " - f"or 'audio_config'. Got {name} instead." - ) - pretrained_config = getattr(model_config.pretrained_config, name) - quant_config = model_config.quant_config if name == "text_config" else None - preferred_backend = "FLASHINFER" if name == "text_config" else "TRTLLM" - sub_config: ModelConfig = dataclasses.replace( - model_config, - pretrained_config=pretrained_config, - attn_backend=preferred_backend, - quant_config=quant_config, - ) - if ( - hasattr(sub_config.pretrained_config, "torch_dtype") - and sub_config.pretrained_config.torch_dtype is None - ): - sub_config.pretrained_config.torch_dtype = model_config.pretrained_config.torch_dtype - return sub_config - def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper): # Gemma4 checkpoint keys: model.language_model.X -> need model.X for LLM # Remap: "model.language_model.layers.0..." -> "model.layers.0..." @@ -899,14 +966,6 @@ def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper): embed_a_weights = filter_weights("embed_audio", stripped) self.embed_audio.load_weights(embed_a_weights) - def post_config(self): - self.config = self.llm.config - self.model_config.pretrained_config = self.llm.config - - @property - def language_model(self) -> torch.nn.Module: - return self.llm - @nvtx_range("[Vision] process") def _get_image_features( self, @@ -985,55 +1044,3 @@ def _get_audio_features( else: projected = projected.reshape(-1, projected.shape[-1]) return projected.contiguous() - - def get_language_model_extra_forward_kwargs( - self, - *, - raw_input_ids: Optional[torch.Tensor], - position_ids: Optional[torch.Tensor], - mm_inputs: PreparedLlmInputs, - lora_params=None, - **forward_kwargs, - ) -> Dict: - """Build Gemma4-specific language-model forward arguments.""" - del position_ids, forward_kwargs - mm_token_type_ids = None - if raw_input_ids is not None and mm_inputs.input_ids is None: - mm_token_type_ids = torch.zeros_like(raw_input_ids, dtype=torch.long) - mm_token_type_ids[raw_input_ids == self.image_token_ids[0]] = 1 - if self.video_token_ids is not None: - mm_token_type_ids[raw_input_ids == self.video_token_ids[0]] = 2 - if self.audio_token_ids is not None: - mm_token_type_ids[raw_input_ids == self.audio_token_ids[0]] = 3 - - ple_input_ids = None - if mm_token_type_ids is not None and self.llm.model.hidden_size_per_layer_input: - text_config = getattr(self.config, "text_config", self.config) - pad_id = getattr(text_config, "pad_token_id", None) - if pad_id is not None: - ple_input_ids = torch.where( - mm_token_type_ids > 0, - torch.full_like(raw_input_ids, pad_id), - raw_input_ids, - ) - return { - "mm_token_type_ids": mm_token_type_ids, - "ple_input_ids": ple_input_ids, - "lora_params": lora_params, - } - - @property - def multimodal_token_ids(self) -> torch.Tensor: - return self._mm_token_ids - - @property - def text_embedding_layer(self) -> Embedding: - return self.llm.model.embed_tokens - - @property - def embedding_dim(self) -> int: - return self.text_embedding_layer.embedding_dim - - @property - def embedding_dtype(self) -> torch.dtype: - return self.text_embedding_layer.weight.dtype diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 9d2a71dfe3de..0e643518d7a2 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -799,6 +799,14 @@ def embedding_dim(self) -> int: def embedding_dtype(self) -> torch.dtype: return self.text_embedding_layer.weight.dtype + @property + def draft_config(self): + return self.llm.draft_config + + @property + def draft_model(self): + return self.llm.draft_model + def encode_multimodal_inputs( self, multimodal_params: Sequence[MultimodalParams], diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 59f1164a9e95..b984f6cafdfd 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -218,20 +218,6 @@ def language_model(self) -> torch.nn.Module: """Return the inner language model that receives prepared inputs.""" raise NotImplementedError - @property - def draft_config(self): - """Expose the inner language model's draft configuration.""" - return self.language_model.draft_config - - @property - def draft_model(self): - """Expose the inner language model's draft model.""" - return self.language_model.draft_model - - def load_draft_weights(self, *args, **kwargs): - """Delegate draft-weight loading to the inner language model.""" - return self.language_model.load_draft_weights(*args, **kwargs) - @property def vocab_size_padded(self) -> int: """Return the inner language model's padded vocabulary size.""" @@ -241,10 +227,6 @@ def infer_max_seq_len(self) -> int: """Return the inner language model's maximum sequence length.""" return self.language_model.infer_max_seq_len() - def _check_and_adjust_experts_implementation(self, *args, **kwargs): - """Skip Transformers' MoE validation for multimodal wrapper models.""" - return None - def get_language_model_extra_forward_kwargs( self, *, @@ -847,8 +829,9 @@ def _dispatch_cross_iter_prefetch( chunk_end_pos=cumsum.numel(), embed_mask_cumsum=cumsum, ), + mm_item_order=req.py_mm_item_order, ) - for _, mm_data, cumsum in candidates + for req, mm_data, cumsum in candidates ] # Prefetch targets requests outside the current iteration, so their diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index ab2e54cf14f9..dcee660b2ccf 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -5,7 +5,7 @@ import math import re from functools import lru_cache -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union import numpy as np import torch @@ -19,7 +19,10 @@ Qwen3VLVisionPatchEmbed as HFQwen3VLVisionPatchEmbed, ) -from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg +from tensorrt_llm._torch.models.modeling_multimodal_utils import ( + _is_mm_disagg, + filter_mm_token_from_input_ids, +) from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.mapping import Mapping @@ -1234,6 +1237,16 @@ def encode_multimodal_inputs( ) return mm_embeds[0] + def _check_and_adjust_experts_implementation(self, *args, **kwargs): + """No-op override. + + Transformers 5.x's ``PreTrainedModel.__init__`` calls this method + (with an ``experts_implementation`` argument) which fails for VL + wrapper models that do not directly contain MoE layers. TRT-LLM + manages expert implementations independently, so skip the check. + """ + return None + def __init__( self, model_config: ModelConfig[PretrainedConfig], @@ -1331,6 +1344,21 @@ def __init__( self.deepstack_num_level = ( len(config.vision_config.deepstack_visual_indexes) if self.use_deepstack else 0 ) + if self.deepstack_num_level > 0: + # Reuse one `(L, max_num_tokens, hidden)` scratch allocation for + # per-layer deepstack embeddings. The generic extra-embedding path + # allocates and scatters one full-sequence tensor per level. + self.register_buffer( + "deepstack_input_embeds", + torch.zeros( + self.deepstack_num_level, + model_config.max_num_tokens, + config.text_config.hidden_size, + device="cuda", + dtype=config.text_config.torch_dtype, + ), + persistent=False, + ) # Surface the in-vocab image / video placeholder IDs to the model # engine's ``_prepare_multimodal_indices`` so it selects the @@ -1354,6 +1382,26 @@ def multimodal_token_ids(self) -> torch.Tensor: def language_model(self) -> torch.nn.Module: return self.llm + # Draft-model (two-model speculative decoding, e.g. DFlash / Eagle3) + # delegation: `ModelLoader.load` reads `draft_config` / `draft_model` and + # calls `load_draft_weights` on the *outer* model it resolved, but the + # spec-decoding wrapper (`SpecDecOneEngineForCausalLM`) is applied to the + # inner `self.llm` when this VLM composes it. Composite checkpoints + # (e.g. Qwen3.5-4B publishes text_config + vision_config) route text-only + # spec tests through this wrapper, so surface the inner LM's draft state. + # Note: `load_draft_weights` must keep an explicit signature — the loader + # dispatches kwargs via `inspect.getfullargspec`. + @property + def draft_config(self): + return self.llm.draft_config + + @property + def draft_model(self): + return self.llm.draft_model + + def load_draft_weights(self, weights: Dict, weight_mapper: Optional[BaseWeightMapper] = None): + return self.llm.load_draft_weights(weights, weight_mapper=weight_mapper) + @property def text_embedding_layer(self) -> Embedding: return self.llm.model.embed_tokens @@ -1452,7 +1500,6 @@ def after_active_multimodal_embeddings( **forward_kwargs: Any, ) -> tuple[List[torch.Tensor], List[torch.Tensor]]: """Separate Qwen3-VL's packed deepstack streams from primary embeddings.""" - del multimodal_params, forward_kwargs if not self.use_deepstack: return active_embeddings, [] @@ -1464,6 +1511,55 @@ def after_active_multimodal_embeddings( deepstack_embeds.extend(deepstack_embed) return active_embeddings, deepstack_embeds + def _fuse_multimodal_embeddings( + self, + *, + input_ids: torch.Tensor, + multimodal_embeddings: List[torch.Tensor], + mm_token_ids: Optional[Sequence[int] | torch.Tensor], + embedding_layer, + extra_embeds: Sequence[torch.Tensor], + text_token_indices: Optional[torch.Tensor] = None, + mm_token_indices: Optional[torch.Tensor] = None, + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Sequence[torch.Tensor]]: + """Fuse primary embeddings and expand deepstack features in reusable scratch.""" + # Qwen only needs the explicit MM indices below when deepstack features + # must be scattered into its reusable buffer. Without extra embeds, + # `fuse_input_embeds` performs its normal index fallback inside `super()`. + if extra_embeds and (text_token_indices is None or mm_token_indices is None): + text_token_indices, mm_token_indices = filter_mm_token_from_input_ids( + input_ids, + vocab_size=embedding_layer.num_embeddings, + mm_token_ids=mm_token_ids, + ) + + fused_input_ids, inputs_embeds, _ = super()._fuse_multimodal_embeddings( + input_ids=input_ids, + multimodal_embeddings=multimodal_embeddings, + mm_token_ids=mm_token_ids, + embedding_layer=embedding_layer, + # Keep auxiliary fusion out of the generic path: passing non-empty + # `extra_embeds` would allocate and scatter one full-sequence tensor + # per deepstack level before Qwen replaces them with its buffer views. + extra_embeds=(), + text_token_indices=text_token_indices, + mm_token_indices=mm_token_indices, + ) + if not extra_embeds: + return fused_input_ids, inputs_embeds, () + + # Expand the per-level deepstack mm embeddings into the pre-allocated + # `(L, max_num_tokens, H)` buffer with a single packed scatter, avoiding `L` fresh + # `torch.zeros` + `L` scatters inside `fuse_input_embeds`. + deepstack_buffer = self.deepstack_input_embeds[:, : input_ids.shape[0], :] + deepstack_buffer.zero_() + packed_deepstack = torch.stack(tuple(extra_embeds), dim=0) + deepstack_buffer[:, mm_token_indices, :] = packed_deepstack.to( + dtype=deepstack_buffer.dtype, + device=deepstack_buffer.device, + ) + return fused_input_ids, inputs_embeds, tuple(deepstack_buffer.unbind(0)) + def get_language_model_extra_forward_kwargs( self, *, @@ -1480,7 +1576,6 @@ def get_language_model_extra_forward_kwargs( **forward_kwargs: Any, ) -> Dict[str, Any]: """Build Qwen3-VL-specific language-model forward arguments.""" - del forward_kwargs mrope_config = {} if not self.model_config.pretrained_config.disable_fuse_rope: mrope_config = self.prepare_mrope_config( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 61ddc4bdbb4f..0945db4d2394 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -51,6 +51,9 @@ CacheLevel = NewType("CacheLevel", int) TokenId = NewType("TokenId", int) TokenIdExt = Union[TokenId, bytes] +class PlannedDropHandle: + def drop(self) -> None: ... + class ReuseScope(NamedTuple): lora_id: int | None = None salt: int | None = None @@ -349,6 +352,7 @@ class _KVCache: def committed_tokens(self) -> list[TokenIdExt]: ... @property def reuse_scope(self) -> ReuseScope: ... + def plan_committed_block_drop(self) -> PlannedDropHandle | None: ... def stop_committing(self) -> None: ... def suspend(self) -> None: ... def resume(self, cuda_stream: CudaStream | None = None) -> bool: ... diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index 7b7539cf0b60..a7633100f34f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -4,7 +4,6 @@ import copy import os from dataclasses import dataclass -from types import SimpleNamespace from typing import List, Optional import pytest @@ -17,7 +16,6 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper -from tensorrt_llm._torch.models.modeling_multimodal_mixin import PreparedLlmInputs from tensorrt_llm._torch.models.modeling_qwen3vl import ( Qwen3VisionModel, Qwen3VLInputProcessorBase, @@ -86,39 +84,6 @@ } -def test_qwen3vl_forwards_fused_deepstack_embeddings_without_rescattering(): - """Deepstack features are sequence-shaped after the shared fusion step.""" - hidden_size = 4 - model = object.__new__(Qwen3VLModel) - torch.nn.Module.__init__(model) - model.use_deepstack = True - model.model_config = SimpleNamespace(pretrained_config=SimpleNamespace(disable_fuse_rope=True)) - model.llm = SimpleNamespace( - model=SimpleNamespace( - embed_tokens=SimpleNamespace(num_embeddings=32), - ) - ) - - deepstack_embeds = [torch.randn(5, hidden_size) for _ in range(3)] - result = model.get_language_model_extra_forward_kwargs( - raw_input_ids=torch.tensor([1, 7, 2, 7, 3]), - position_ids=None, - mm_inputs=PreparedLlmInputs( - input_ids=None, - inputs_embeds=torch.randn(5, hidden_size), - extra_embeds=deepstack_embeds, - ), - multimodal_params=[], - num_generation_requests=0, - spec_metadata=None, - mm_token_indices=torch.tensor([1, 3]), - ) - - assert len(result["deepstack_embeds"]) == len(deepstack_embeds) - for actual, expected in zip(result["deepstack_embeds"], deepstack_embeds): - assert actual is expected - - @dataclass(repr=False) class TestQwen3VLScenario(MultimodalScenario): disable_fuse_rope: bool = False @@ -386,6 +351,104 @@ def setup_scenario(self, scenario: TestQwen3VLScenario): ) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen3vl_deepstack_fusion_reuses_registered_buffer(): + hidden_size = 4 + num_tokens = 5 + mm_token_id = 7 + config_dict = copy.deepcopy(QWEN3_VL_8B_CONFIG) + config_dict.update( + image_token_id=mm_token_id, + video_token_id=8, + vision_start_token_id=9, + vision_end_token_id=10, + ) + config_dict["text_config"].update( + hidden_size=hidden_size, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=hidden_size, + vocab_size=32, + ) + config_dict["vision_config"].update( + deepstack_visual_indexes=[0, 1], + depth=2, + hidden_size=32, + intermediate_size=64, + num_heads=1, + out_hidden_size=hidden_size, + ) + model_config = ModelConfig( + pretrained_config=Qwen3VLConfig.from_dict(config_dict), + disable_mm_encoder=True, + max_num_tokens=num_tokens, + ) + model = Qwen3VLModel(model_config, disable_fuse_rope=True).to("cuda") + + input_ids = torch.tensor([1, mm_token_id, 2, mm_token_id, 3], device="cuda") + text_token_indices = torch.tensor([0, 2, 4], device="cuda") + mm_token_indices = torch.tensor([1, 3], device="cuda") + primary_embeds = [ + torch.full( + (2, hidden_size), + 10.0, + dtype=model.embedding_dtype, + device="cuda", + ) + ] + deepstack_embeds = [ + torch.full( + (2, hidden_size), + value, + dtype=model.embedding_dtype, + device="cuda", + ) + for value in (20.0, 30.0) + ] + + _, _, fused_deepstack = model._fuse_multimodal_embeddings( + input_ids=input_ids, + multimodal_embeddings=primary_embeds, + mm_token_ids=model.multimodal_token_ids, + embedding_layer=model.text_embedding_layer, + extra_embeds=deepstack_embeds, + text_token_indices=text_token_indices, + mm_token_indices=mm_token_indices, + ) + + assert len(fused_deepstack) == 2 + # The returned per-level tensors must be views into the model-owned scratch + # buffer, not newly allocated full-sequence tensors. + assert ( + fused_deepstack[0].untyped_storage().data_ptr() + == model.deepstack_input_embeds.untyped_storage().data_ptr() + ) + # Sharing storage alone does not prove that fusion scattered each level to + # the right token rows, so verify the multimodal positions independently. + torch.testing.assert_close( + fused_deepstack[0][mm_token_indices], + deepstack_embeds[0], + ) + torch.testing.assert_close( + fused_deepstack[1][mm_token_indices], + deepstack_embeds[1], + ) + # The scratch buffer is reused across forwards; non-multimodal positions + # must be cleared before every scatter to avoid leaking stale features. + torch.testing.assert_close( + model.deepstack_input_embeds[:, text_token_indices], + torch.zeros( + 2, + len(text_token_indices), + hidden_size, + dtype=model.embedding_dtype, + device="cuda", + ), + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_qwen3vl_init_preserves_caller_quant_config(): """Building Qwen3VLModel must not mutate the caller's quant_config.""" diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py index 01e538d16092..6e3a8e08aa94 100644 --- a/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py +++ b/tests/unittest/_torch/multimodal/test_mm_encoder_cross_iter_prefetch.py @@ -38,10 +38,12 @@ def __init__(self, hidden_size: int, tokens_per_image: int): self._tokens_per_image = tokens_per_image self.encoder_call_count = 0 self.encoder_call_stream_id = None + self.encoder_mm_item_order = None def encode_multimodal_inputs(self, multimodal_params, **kwargs) -> torch.Tensor: self.encoder_call_count += 1 self.encoder_call_stream_id = torch.cuda.current_stream().cuda_stream + self.encoder_mm_item_order = multimodal_params[0].mm_item_order pv = multimodal_params[0].multimodal_data["image"]["pixel_values"] assert pv.device.type == "cuda", ( "pixel_values should be on CUDA after to_device in the helper." @@ -229,12 +231,18 @@ def test_cross_iter_prefetch_materializes_on_side_stream(monkeypatch): model = _StubModel(hidden_size=8, tokens_per_image=4) req = _make_request(request_id=0, num_tokens=4) + expected_item_order = [ + {"modality": "image", "index": 0}, + {"modality": "video", "index": 0}, + ] + req.py_mm_item_order = expected_item_order n = maybe_prefetch_mm_encoder_for_next_iter(model, [req], max_prefetch_ahead=max_prefetch_ahead) assert n == 1 assert model.encoder_call_count == 1 assert model.encoder_call_stream_id == aux_stream.cuda_stream + assert model.encoder_mm_item_order == expected_item_order embedding = req.py_multimodal_data.get("multimodal_embedding") assert isinstance(embedding, torch.Tensor) assert embedding.shape == (4, 8)