diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index aef6f5f82e18..c195f1db0c5a 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -174,6 +174,7 @@ /scripts/generate_config_database_tests.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners /scripts/generate_config_table.py @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-doc-owners /tensorrt_llm/_torch/models @NVIDIA/trt-llm-models-devs +/tensorrt_llm/_torch/multimodal @NVIDIA/trt-llm-models-devs @NVIDIA/trt-llm-runtime-devs /tensorrt_llm/_torch/modules @NVIDIA/trt-llm-models-devs /tensorrt_llm/_torch/modules/dwdp @NVIDIA/trt-llm-runtime-devs /tensorrt_llm/_torch/modules/linear.py @NVIDIA/trt-llm-runtime-devs diff --git a/docs/source/developer-guide/multimodal.md b/docs/source/developer-guide/multimodal.md index af91b64e2e4a..a44c374dbe17 100644 --- a/docs/source/developer-guide/multimodal.md +++ b/docs/source/developer-guide/multimodal.md @@ -168,6 +168,46 @@ override selection explicitly; the compatibility selector cannot describe an arbitrary model layout. An empty partition dict means every field is shared, so it is appropriate only when selecting an item requires no payload changes. +### Encoder Rows and Token Fusion + +`MultimodalRuntimeMixin` and `MultimodalModelMixin` live in +[_torch/multimodal/model.py](source:tensorrt_llm/_torch/multimodal/model.py). +`MultimodalRuntimeMixin` is the executor's multimodal capability contract and +the encoder runtime: encoding, request-local reuse across prefill chunks, the +persistent encoder-output cache, side-stream prefetch and item scheduling. A +model that runs its own decoder, such as a Modeling V2 target, inherits it +directly. `MultimodalModelMixin` extends it for composite models whose language +model is a separate module: it delegates to that module and adds a default +`forward` and the default text embedding fusion. The previous +`_torch/models/modeling_multimodal_mixin.py` path re-exports the same classes +for compatibility. + +Every `MultimodalRuntimeMixin` subclass implements `encode_multimodal_inputs`, +`embedding_dim` and `embedding_dtype`. It can also declare its encoder's limits: + +* `max_encoder_tokens_per_item` bounds the dummy items that encoder memory + profiling builds and the largest atomic item that item scheduling budgets for. +* `supports_concurrent_encoder` permits side-stream prefetch, which runs the + encoder on an auxiliary CUDA stream. The engine rejects + `encoder_side_stream_max_ahead > 0` for a model that does not set it. + Composite models set it by default. +* `allows_breakable_prefill_cuda_graph` admits breakable prefill CUDA graphs + for a model that encodes outside the captured forward. + +`prepare_multimodal_embeddings` performs cache lookup and encoding, waits for +prefetched results, selects current-chunk rows, and runs both model hooks. +Its result contains the input token IDs returned by the full-embedding hook, +active multimodal rows, and any auxiliary embedding tensors returned by the +active-embedding hook. It does not look up text embeddings. When the caller +passes the executor's `mm_token_indices`, the result carries them with the +rows, after checking that both count the same rows, so a model with its own +embedding splice gets rows and positions together. + +`MultimodalModelMixin.prepare_multimodal_inputs` uses the same preparation and +then applies the default text embedding fusion. Model-specific encoder +implementations and token embedding operations remain in their respective +model packages. + ### Item Boundaries and Model Hooks An item is the smallest independently encodable unit, not necessarily one diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 686829541385..c5afbf464255 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -125,6 +125,8 @@ it uses only runtime interfaces, for example: * `DecoderModel`, `DecoderModelForCausalLM`, `ModelConfig` and `Mapping`. * Attention and speculative metadata, and the KV cache managers. +* `MultimodalRuntimeMixin` from `_torch/multimodal/model.py`, which provides + shared encoder execution and caching independently of model implementations. A target does not import, subclass or call Modeling V1 model code: the model classes in `_torch/models/`, the layers in `_torch/modules/`, MLA and the @@ -134,6 +136,62 @@ path does not cover, and weight loading. The [DeepSeek R1 target](models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/) is an example that uses no Modeling V1 model code. +### Multimodal targets + +A target for a checkpoint whose input processor accepts media declares how it +serves media. It inherits `MultimodalRuntimeMixin` from +`_torch/multimodal/model.py` to serve media through the shared multimodal +runtime, or `TextOnlyTarget` from [`_multimodal.py`](_multimodal.py) to serve +the language model only. The marker refuses requests that carry media at +admission, so each fails on its own rather than failing its batch in `forward`. `modeling_v2_resolve` +checks this for every target a `route` returns, in `auto` as well as +`require`, and `test_modeling_v2_multimodal_contract.py` checks it for every +routed target in CI. + +A runtime target inherits `DecoderModelForCausalLM` and +`MultimodalRuntimeMixin`. The mixin has no notion of a separate language model +and defines nothing the decoder defines, so the decoder's speculative and +runtime hooks are unaffected. Implement `encode_multimodal_inputs`, +`embedding_dim` and `embedding_dtype`. The model loader reads the last two from +the constructed target, so they must work with the checkpoint's real +configuration. Expose `dtype` for encoder memory profiling. + +Before the core's forward, call `prepare_multimodal_embeddings` with the +engine's context count, the request Params and the executor's +`mm_token_indices`. It returns the active encoder rows after request-local +reuse, persistent cache lookup, stream synchronization, chunk selection, and +model hooks, paired with their placeholder positions. Feed both to the +target's own token splice. Keep original token IDs available when other model +operations need them. + +Filter position-only Params in `select_multimodal_params`: the processor emits +them for text requests, and the executor adds them for generation requests. +Enable `supports_encoder_cache` only when the processor supplies stable item +keys and exact output row counts. Declare packed input partitions or implement +item selection so partial cache hits encode only missing items. Encoder outputs +and cached tensors are read-only. + +Declare the encoder's limits on the target: + +* `max_encoder_tokens_per_item` bounds the dummy items that encoder memory + profiling builds. Return the largest item the encoder runs in one piece when + it splits or refuses larger ones. +* `supports_concurrent_encoder` permits side-stream prefetch. Set it only when + mutable encoder state, such as kernel workspaces, is separate per CUDA + stream. +* `allows_breakable_prefill_cuda_graph` admits breakable prefill CUDA graphs. + Return True only when encoding happens outside the captured forward. + +Item scheduling is a separate capability. Set +`supports_mm_encoder_item_scheduling` only after the complete cost, selection, +and memory-profile contracts work for the target's certified encoder shapes. +The Qwen3.8 Flash Next target uses inline encoding, chunk reuse, persistent +encoder caching, prefetch and breakable prefill CUDA graphs; it leaves item +scheduling disabled because the processor's maximum atomic image exceeds its +certified per-frame patch bound. See the +[multimodal developer guide](../../../../docs/source/developer-guide/multimodal.md) +for the shared contracts. + ## Adding a family or a target 1. **New family only:** create `models//routing.py` with a `route` diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_multimodal.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_multimodal.py new file mode 100644 index 000000000000..beda02aea73c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_multimodal.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The multimodal contract every target of a multimodal checkpoint meets. + +A checkpoint is multimodal when its input processor accepts media +(``BaseMultimodalInputProcessor``). A target for such a checkpoint does one of two things: + +* It serves media through the shared multimodal runtime: it inherits + ``MultimodalRuntimeMixin`` (``_torch/multimodal/model.py``) and implements its required + members. The executor then gives it encoder-output caching, chunk reuse, prefetch and the + other runtime features. +* It serves the checkpoint's language model only: it inherits :class:`TextOnlyTarget`, + builds no encoder, and the marker refuses requests carrying media at admission. + +``modeling_v2_resolve`` checks this for every target a ``route`` returns, under ``auto`` as +well as ``require``: a target that meets neither is a bug in the target, not a deployment it +declines. + +Like the router, this module imports nothing at load time: the registries it reads are +imported when a target is checked. +""" + +#: ``multimodal_data`` entries that carry media for an encoder, or encoder output in place of it. +_MEDIA_KEYS = frozenset({"audio", "image", "video", "multimodal_embedding"}) + + +class TextOnlyTarget: + """Marker for a target that serves only the language model of a multimodal checkpoint. + + The target builds no multimodal encoder. Requests that carry media are refused at + admission, so each fails on its own instead of failing every request in its batch. + """ + + def validate_multimodal_request_data(self, mm_data: dict) -> None: + """Refuse a request whose multimodal payload carries media; accept position-only data.""" + media = sorted(key for key in _MEDIA_KEYS if mm_data.get(key) is not None) + if media: + raise ValueError( + f"{type(self).__name__} serves this checkpoint's language model only and " + f"does not accept multimodal input; the request carries {', '.join(media)}" + ) + + +def checkpoint_accepts_media(architecture: str) -> bool: + """Whether the input processor registered for ``architecture`` accepts media.""" + from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class + from tensorrt_llm.inputs.registry import INPUT_PROCESSOR_REGISTRY, BaseMultimodalInputProcessor + + model_cls = get_registered_model_class(architecture) + if model_cls is None: + return False + processor_cls = INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type.get(model_cls) + return processor_cls is not None and issubclass(processor_cls, BaseMultimodalInputProcessor) + + +def check_target_multimodal_contract(architecture: str, target: type) -> None: + """Raise ``MultimodalModelContractError`` unless ``target`` meets the contract above.""" + from tensorrt_llm._torch.multimodal.model import ( + MultimodalModelContractError, + MultimodalRuntimeMixin, + ) + + runs_runtime = issubclass(target, MultimodalRuntimeMixin) + text_only = issubclass(target, TextOnlyTarget) + if runs_runtime and text_only: + raise MultimodalModelContractError( + f"{target.__qualname__} inherits both MultimodalRuntimeMixin and TextOnlyTarget" + ) + if runs_runtime: + target.check_multimodal_class_contract() + return + if text_only or not checkpoint_accepts_media(architecture): + return + raise MultimodalModelContractError( + f"{target.__qualname__} is a target for the multimodal architecture {architecture!r}; " + "it must inherit MultimodalRuntimeMixin to serve media, or TextOnlyTarget to serve " + "the language model only" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 1d714f7dd885..7d4711f3b575 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -46,6 +46,8 @@ from tensorrt_llm.logger import logger +from ._multimodal import check_target_multimodal_contract + if TYPE_CHECKING: from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.llmapi.llm_args import TorchLlmArgs @@ -223,6 +225,7 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti ) return None if target is not None: + check_target_multimodal_contract(arch, target) _RESOLVED_MODES[target] = mode return target diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py index 926a19c8ee8f..02cf2baf7085 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py @@ -95,6 +95,7 @@ import torch from torch import nn +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import TextOnlyTarget from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kernels.k3_mla.decode_view import ( k3_mla_decode_view, @@ -3216,7 +3217,7 @@ def _check_construction(model_config: ModelConfig) -> None: # >>> route B: this target's name -class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1(KimiLinearForCausalLM): +class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp16ep1(KimiLinearForCausalLM, TextOnlyTarget): # <<< route B """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py index fb22fd8c7306..4aa4307c299c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py @@ -93,6 +93,7 @@ import torch from torch import nn +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import TextOnlyTarget from tensorrt_llm._torch._experimental.modeling_v2._native import load_family_ops from tensorrt_llm._torch._experimental.modeling_v2.models.kimi_k3_vl.kernels.k3_mla.decode_view import ( k3_mla_decode_view, @@ -3294,7 +3295,7 @@ def _check_construction(model_config: ModelConfig) -> None: ) -class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4(KimiLinearForCausalLM): +class ModelingV2KimiK3Mxfp4Sm100Tp16Moetp4ep4(KimiLinearForCausalLM, TextOnlyTarget): """The class ``route`` returns: this target's text model (`KimiLinearModel` above) behind its checks. It inherits the built-in Kimi K3 causal LM for the checkpoint load and the engine hooks (`load_weights` through diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py index c0bf5cd12b17..2606034485c9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp4/modeling.py @@ -77,6 +77,7 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import TextOnlyTarget from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( ModelingV2Mode, NoTarget, @@ -4498,7 +4499,7 @@ def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: class ModelingV2MiniMaxM3Nvfp4Sm103Tp4( - DecoderModelForCausalLM[MiniMaxM3Tp4ModelingV2, PretrainedConfig] + DecoderModelForCausalLM[MiniMaxM3Tp4ModelingV2, PretrainedConfig], TextOnlyTarget ): """MiniMax-M3 NVFP4 / sm_103 / tp4, text only, with one-model Eagle3 or no speculation. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py index 1e8c520ca9d6..e86fc4888cd8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/modeling.py @@ -79,6 +79,7 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import TextOnlyTarget from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( ModelingV2Mode, NoTarget, @@ -4167,7 +4168,7 @@ def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: class ModelingV2MiniMaxM3Nvfp4Sm103Tp8( - DecoderModelForCausalLM[MiniMaxM3Tp8ModelingV2, PretrainedConfig] + DecoderModelForCausalLM[MiniMaxM3Tp8ModelingV2, PretrainedConfig], TextOnlyTarget ): """MiniMax-M3 NVFP4 / sm_103 / tp8, text only, with one-model Eagle3 or no speculation. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py index 15926f7e5259..bd591e780358 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py @@ -88,6 +88,7 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import TextOnlyTarget from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( ModelingV2Mode, resolved_mode, @@ -6092,7 +6093,7 @@ def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: class ModelingV2MiniMaxM3Nvfp4Sm107Tp2( - DecoderModelForCausalLM[MiniMaxM3Tp2ModelingV2, PretrainedConfig] + DecoderModelForCausalLM[MiniMaxM3Tp2ModelingV2, PretrainedConfig], TextOnlyTarget ): """MiniMax-M3 NVFP4 / sm_107 / tp2, text only, with one-model Eagle3 or no speculation. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py index eb1b89fde8c1..77709f859ac5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py @@ -79,6 +79,7 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import TextOnlyTarget from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( ModelingV2Mode, resolved_mode, @@ -5566,7 +5567,7 @@ def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: class ModelingV2MiniMaxM3Nvfp4Sm107Tp4( - DecoderModelForCausalLM[MiniMaxM3Tp4ModelingV2, PretrainedConfig] + DecoderModelForCausalLM[MiniMaxM3Tp4ModelingV2, PretrainedConfig], TextOnlyTarget ): """MiniMax-M3 NVFP4 / sm_107 / tp4, text only, with one-model Eagle3 or no speculation. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py index 4c1b7189cc08..201c780e532b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py @@ -86,6 +86,7 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import TextOnlyTarget from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( ModelingV2Mode, resolved_mode, @@ -5784,7 +5785,7 @@ def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: class ModelingV2MiniMaxM3Nvfp4Sm107Tp8( - DecoderModelForCausalLM[MiniMaxM3Tp8ModelingV2, PretrainedConfig] + DecoderModelForCausalLM[MiniMaxM3Tp8ModelingV2, PretrainedConfig], TextOnlyTarget ): """MiniMax-M3 NVFP4 / sm_107 / tp8, text only, with one-model Eagle3 or no speculation. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py index 2807813eae62..cdcf68c1132f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/modeling.py @@ -140,6 +140,7 @@ import math import os import re +from collections.abc import Sequence from typing import TYPE_CHECKING, NamedTuple import torch @@ -437,12 +438,14 @@ flashinfer_cached_replay_update, ) from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.multimodal.model import MultimodalRuntimeMixin from tensorrt_llm._torch.pyexecutor.model_loader import effective_mamba_ssm_cache_dtype from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks from tensorrt_llm._torch.speculative import get_spec_worker, should_use_separate_draft_kv_cache from tensorrt_llm._torch.utils import ActivationType, is_gdn_replay_enabled from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.functional import AllReduceFusionOp +from tensorrt_llm.inputs.multimodal import MultimodalItemPartition, MultimodalParams from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantAlgo @@ -6758,7 +6761,7 @@ def _text_settings(text) -> dict: class ModelingV2Qwen38FlashNextNvfp4Sm103Default( - DecoderModelForCausalLM[Qwen4ExpDefaultModelingV2, PretrainedConfig] + DecoderModelForCausalLM[Qwen4ExpDefaultModelingV2, PretrainedConfig], MultimodalRuntimeMixin ): """Qwen3.8-Flash-Next-NVFP4 at tp1, tp4, tep4, dep4 or dep2 on sm_103.""" @@ -6768,6 +6771,63 @@ class ModelingV2Qwen38FlashNextNvfp4Sm103Default( # decode-step prepare), validated there; off at every other layout. spec_decode_fast_paths: bool = False runtime_hooks: RuntimeHooks | None = None + supports_encoder_cache = True + # The tower keeps one attention workspace per CUDA stream (`VisionTower._attention_workspace`). + supports_concurrent_encoder = True + + @property + def embedding_dim(self) -> int: + """Width of one merged vision row; the composite config has no top-level width.""" + return self.model_config.pretrained_config.vision_config.out_hidden_size + + @property + def embedding_dtype(self) -> torch.dtype: + return self.model_config.torch_dtype + + @property + def dtype(self) -> torch.dtype: + return self.embedding_dtype + + def select_multimodal_params( + self, multimodal_params: Sequence[MultimodalParams], num_context_requests: int + ) -> list[MultimodalParams]: + context_params = super().select_multimodal_params(multimodal_params, num_context_requests) + return [ + param + for param in context_params + if any( + param.multimodal_data.get(key) is not None + for key in ("image", "video", "multimodal_embedding") + ) + ] + + def encode_multimodal_inputs( + self, multimodal_params: Sequence[MultimodalParams] + ) -> torch.Tensor: + return _vision.encode_multimodal_inputs(self.vision_tower, multimodal_params) + + def get_multimodal_item_partitions( + self, param: MultimodalParams, modality: str + ) -> dict[str, MultimodalItemPartition]: + pixel_key, grid_key = _vision.MEDIA_KEYS[modality] + grids = param.multimodal_data[modality][grid_key] + return { + pixel_key: MultimodalItemPartition( + item_lengths=tuple(t * h * w for t, h, w in grids.tolist()) + ), + grid_key: MultimodalItemPartition(item_axis=0), + } + + def max_encoder_tokens_per_item(self, modality: str) -> int: + """The tower runs at most `VISION_CALL_MAX_PATCHES` patch rows per call and cuts a + larger image or video between frames; admission refuses a single frame above it. An + item of that size therefore bounds the tower's per-call activations.""" + return _vision.VISION_CALL_MAX_PATCHES + + def allows_breakable_prefill_cuda_graph(self, llm_args: TorchLlmArgs) -> bool: + """The shell encodes before the core's captured `forward`, and the prefill target + splices the rows through an eager break (`PrefillTarget._splice`).""" + return True @classmethod def get_preferred_kv_cache_manager_version(cls, pretrained_config=None) -> str: @@ -7025,6 +7085,14 @@ def forward( the sampler verifies, `spec_metadata.gather_ids`) and hands the worker the trunk's ungathered stream rows, which the MTP module reads; `position_ids` goes in the engine's layout.""" + assert input_ids is not None, "this target's PLE layer hashes token ids" + prepared_mm = self.prepare_multimodal_embeddings( + input_ids=input_ids, + multimodal_params=kwargs.get("multimodal_params"), + num_context_requests=attn_metadata.num_contexts, + mm_token_indices=kwargs.get("mm_token_indices"), + attn_metadata=attn_metadata, + ) self.model.begin_step( attn_metadata, input_ids, @@ -7032,7 +7100,7 @@ def forward( inputs_embeds=inputs_embeds, lora_params=lora_params, spec_metadata=spec_metadata, - vision=_vision.multimodal_embeds(self.vision_tower, kwargs), + vision=_vision.multimodal_embeds(prepared_mm), ) if self.spec_worker is None: assert spec_metadata is None, ( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py index 818d919c6551..021daa21fe65 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/vision.py @@ -10,7 +10,7 @@ the same rows back, since every op is a deterministic function of its operands; under attention DP a rank its own requests'. -The shell calls `multimodal_embeds` on every step, before the core's `begin_step`, and the +The shell prepares cached encoder rows before the core's `begin_step`, and the executor calls `validate_request` on every new request. Neither runs inside a captured graph: the tower is eager and reads the request's grid on the host. """ @@ -18,6 +18,7 @@ from __future__ import annotations import math +from collections.abc import Sequence from typing import NamedTuple import torch @@ -44,7 +45,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default.ops.vision.triton_pos_embed_interpolate import ( # noqa: E501 triton_pos_embed_interpolate, ) -from tensorrt_llm.inputs.multimodal import get_multimodal_item_metadata +from tensorrt_llm._torch.multimodal.model import PreparedMultimodalEmbeddings +from tensorrt_llm.inputs.multimodal import MultimodalParams, get_multimodal_item_metadata #: Tower rope base (`vision_config` has no rope field); unrelated to the text 1e7. VISION_ROPE_THETA = 10000.0 @@ -155,6 +157,7 @@ DEVICE_PATHS = ( "image.pixel_values", "video.pixel_values_videos", + "multimodal_embedding", "mrope_config.mrope_position_ids", "mrope_config.mrope_position_deltas", ) @@ -334,9 +337,18 @@ def __init__(self, w: dict, vision_config) -> None: self.grid_side = math.isqrt(vcfg.num_position_embeddings) self.dtype = w["embed"].dtype self._w = _vision_views(w, self.depth) - # Persistent scratch: `thop_attention` grows `workspace` in place; never the - # engine's own workspace. - self._ws: torch.Tensor | None = None + # Encoder prefetch and inline encoding can overlap on different streams. + # Each stream owns scratch that `thop_attention` may grow in place. + self._workspaces: dict[tuple[torch.device, int], torch.Tensor] = {} + + def _attention_workspace(self, device: torch.device) -> torch.Tensor: + stream = torch.cuda.current_stream(device) + key = (device, stream.cuda_stream) + workspace = self._workspaces.get(key) + if workspace is None: + workspace = torch.empty([0], dtype=torch.int8, device=device) + self._workspaces[key] = workspace + return workspace def warm_up(self) -> None: """Run the tower once per `WARM_UP_GRIDS` grid on zero pixels, so every JIT kernel it @@ -408,8 +420,7 @@ def _run(self, pixel_values: torch.Tensor, grid: list[tuple[int, int, int]]) -> host_types = torch.zeros([num_segments], dtype=torch.int32) host_total = torch.tensor([num_patches, 0], dtype=torch.int32) dev_lens = host_lens.to(dev) - if self._ws is None: - self._ws = torch.empty([0], dtype=torch.int8, device=dev) + workspace = self._attention_workspace(dev) # The rotary's angles, once for all blocks. cos, sin = _rope_tables(grid, self.merge, VISION_ROPE_THETA, self.head_dim, dev) @@ -433,7 +444,7 @@ def _run(self, pixel_values: torch.Tensor, grid: list[tuple[int, int, int]]) -> attn_out = thop_attention( q=qkv, output=attn_out, - workspace=self._ws, + workspace=workspace, sequence_length=dev_lens, host_past_key_value_lengths=host_lens, host_total_kv_lens=host_total, @@ -469,78 +480,43 @@ def _run(self, pixel_values: torch.Tensor, grid: list[tuple[int, int, int]]) -> # ------------------------------------------------------------ the shell's side --- -def multimodal_embeds(tower: VisionTower | None, kwargs: dict): - """Run the vision tower for this batch; returns `(rows, mm_token_indices)` (the tower's - `[R, hidden]` rows in packed-batch order and the `[R]` placeholder rows they replace, - which the prefill target splices into its embedding), or None when the step carries no - media. - - Called by the shell, because the inherited `DecoderModelForCausalLM.forward` drops - `multimodal_params` from `**kwargs` on the no-speculation path. Host-side bookkeeping: - `multimodal_params` is a subset of the context requests, so never zip it positionally; - `multimodal_runtime` (chunked prefill / KV reuse only) gives the cached and in-chunk row - counts over the request's images and videos together, in prompt order; `mm_token_indices` - names every placeholder row. - - The tower runs once per modality a request carries, and `in_prompt_order` merges the - two outputs. Positions are not derived here: every coordinate of a vision span, a - video's temporal one included, is in the payload's `mrope_config.mrope_position_ids`, - which the engine slices into `position_ids`.""" - mm_params = kwargs.get("multimodal_params") or [] - if not mm_params: - return None - rows = [] - for params in mm_params: - data = getattr(params, "multimodal_data", None) or {} - image, video = data.get("image"), data.get("video") - if image is None and video is None: - continue - mrope_ids = (data.get("mrope_config") or {}).get("mrope_position_ids") - assert ( - mrope_ids is not None and mrope_ids.dim() == 3 and mrope_ids.shape[0] == _POSITION_AXES - ), ( - f"a media request needs mrope_config.mrope_position_ids of shape " - f"[{_POSITION_AXES}, 1, L]; got {None if mrope_ids is None else tuple(mrope_ids.shape)}" - ) - runtime = getattr(params, "multimodal_runtime", None) - if runtime is not None and int(runtime.num_mm_tokens_in_chunk) == 0: - # Every placeholder row of this request lies in another chunk. - continue - assert tower is not None, "post_load_weights builds the vision tower before any step" +def encode_multimodal_inputs( + tower: VisionTower | None, multimodal_params: Sequence[MultimodalParams] +) -> torch.Tensor: + """Encode complete media items in request and prompt order. + + Cache lookup and active-chunk selection belong to the shared model runtime. + This callback also serves encoder memory profiling, whose inputs contain no + language-model position IDs. + """ + if tower is None: + raise ValueError("Raw multimodal inputs require a local vision tower") + outputs = [] + for params in multimodal_params: towers = {} - if image is not None: - towers["image"] = run_tower(tower, image, *MEDIA_KEYS["image"]) - if video is not None: - towers["video"] = run_tower(tower, video, *MEDIA_KEYS["video"]) - embeds = in_prompt_order(tower.merge, params, towers) - if runtime is not None: - # Chunked prefill or reused prefix: splice only this chunk's slice. - start = int(runtime.num_cached_mm_tokens) - count = int(runtime.num_mm_tokens_in_chunk) - total = embeds.shape[0] - assert int(runtime.total_embeds_in_request) == total, ( - f"the executor counted {runtime.total_embeds_in_request} placeholder rows for " - f"this request and the tower produced {total} over its {sorted(towers)}; the " - "two are derived from the same grids and a disagreement means the merger's " - "row count is wrong" - ) - assert 0 <= start <= start + count <= total, (start, count, total) - embeds = embeds[start : start + count] - rows.append(embeds) - if not rows: + for modality, keys in MEDIA_KEYS.items(): + media = params.multimodal_data.get(modality) + if media is not None: + towers[modality] = run_tower(tower, media, *keys) + if not towers: + raise ValueError("Encoder inputs must contain image or video data") + outputs.append(in_prompt_order(tower.merge, params, towers)) + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) + + +def multimodal_embeds( + prepared: PreparedMultimodalEmbeddings, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """The step's `(rows, mm_token_indices)` for the prefill target's splice, or None when + the step has no active encoder rows. The shared runtime checked that both count the same + rows.""" + embeddings = prepared.multimodal_embeddings + if not embeddings: return None - vision = rows[0] if len(rows) == 1 else torch.cat(rows, 0) - mm_token_indices = kwargs.get("mm_token_indices") - assert mm_token_indices is not None, ( - "the executor shipped multimodal params but no mm_token_indices; that pairing is what " - "the shell's `multimodal_token_ids` buys, and without it there is no sync-free way to " - "find the placeholder rows" - ) - assert mm_token_indices.shape[0] == vision.shape[0], ( - f"the batch has {mm_token_indices.shape[0]} placeholder rows and the tower produced " - f"{vision.shape[0]} embedding rows" - ) - return vision, mm_token_indices + if prepared.mm_token_indices is None: + raise ValueError("Active multimodal embeddings require mm_token_indices") + rows = embeddings[0] if len(embeddings) == 1 else torch.cat(tuple(embeddings), dim=0) + return rows, prepared.mm_token_indices def run_tower(tower, media: dict, pixel_key: str, grid_key: str): diff --git a/tensorrt_llm/_torch/models/modeling_cosmos3_edge.py b/tensorrt_llm/_torch/models/modeling_cosmos3_edge.py index 0b2d41241946..efdf7a8525d6 100644 --- a/tensorrt_llm/_torch/models/modeling_cosmos3_edge.py +++ b/tensorrt_llm/_torch/models/modeling_cosmos3_edge.py @@ -60,17 +60,13 @@ from ..modules.linear import TensorParallelMode from ..modules.mlp import MLP from ..modules.rms_norm import RMSNorm +from ..multimodal.model import EncoderGroup, MultimodalModelMixin, encode_multimodal_by_groups from ..speculative import SpecMetadata from ..utils import relu2 from .checkpoints.base_weight_mapper import BaseWeightMapper from .checkpoints.hf.cosmos3_edge_weight_mapper import Cosmos3EdgeHfWeightMapper from .modeling_cosmos3 import _get_cosmos3_model_paths from .modeling_multimodal_encoder import MultimodalEncoderMixin -from .modeling_multimodal_mixin import ( - EncoderGroup, - MultimodalModelMixin, - encode_multimodal_by_groups, -) from .modeling_qwen2vl import Qwen2_5_VLVisionAttention, _prepare_qwen_vl_vision_attn_metadata from .modeling_qwen3vl import ( Qwen3VLInputProcessorBase, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv41_multimodal.py b/tensorrt_llm/_torch/models/modeling_deepseekv41_multimodal.py index 3a0ad829fd0e..7a8b74622f19 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv41_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv41_multimodal.py @@ -17,10 +17,10 @@ from tensorrt_llm.inputs.multimodal import MultimodalParams from ..model_config import ModelConfig +from ..multimodal.model import MultimodalModelMixin from .deepseek_v41_processor import IMAGE_PLACEHOLDER, DeepseekV41InputProcessor from .deepseek_v41_vision import DeepseekV41Aligner, DeepseekV41VisionModel from .modeling_deepseekv41 import DeepseekV41ForCausalLM -from .modeling_multimodal_mixin import MultimodalModelMixin from .modeling_utils import register_auto_model diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 92c666a0f4e7..e46686c80a5f 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -52,15 +52,11 @@ from ...sampling_params import SamplingParams from ..modules.embedding import Embedding from ..modules.linear import Linear +from ..multimodal.model import EncoderCachePartition, MultimodalModelMixin, PreparedLlmInputs from ..peft.lora.ownership import DecoderLoraOwner from .modeling_gemma4 import Gemma4ForCausalLM from .modeling_gemma4_audio import Gemma4AudioModel from .modeling_gemma4_vision import Gemma4VisionModel -from .modeling_multimodal_mixin import ( - EncoderCachePartition, - MultimodalModelMixin, - PreparedLlmInputs, -) from .modeling_multimodal_utils import ( _MULTIMODAL_ENV_NAME, VisionNormalization, diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py index 60a1aea3e113..37491fc6471b 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py @@ -70,14 +70,10 @@ from ..custom_ops import flashinfer_apply_rope_with_cos_sin_cache_inplace else: # pragma: no cover - CPU / no-flashinfer environments flashinfer_apply_rope_with_cos_sin_cache_inplace = None +from ..multimodal.model import EncoderGroup, MultimodalModelMixin, encode_multimodal_by_groups from .modeling_auto import AutoModelForCausalLM from .modeling_glm5_next import Glm5NextForCausalLM, glm5_next_attention_mapping from .modeling_multimodal_encoder import MultimodalEncoderMixin, prepare_encoder_attn_metadata -from .modeling_multimodal_mixin import ( - EncoderGroup, - MultimodalModelMixin, - encode_multimodal_by_groups, -) from .modeling_utils import ( ModelConfig, QuantConfig, diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index c17d4f04b653..0be7ec132992 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -26,8 +26,6 @@ MistralWeightMapper from tensorrt_llm._torch.models.modeling_mistral_large3 import ( Mistral3Gate, MistralLarge3ForCausalLM) -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( - MultimodalModelMixin, PreparedLlmInputs) from tensorrt_llm._torch.models.modeling_multimodal_utils import ( _MULTIMODAL_ENV_NAME, _is_mm_disagg) from tensorrt_llm._torch.models.modeling_utils import (DecoderModel, @@ -40,6 +38,8 @@ from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.multimodal.model import (MultimodalModelMixin, + PreparedLlmInputs) from tensorrt_llm._torch.peft.lora.ownership import DecoderLoraOwner from tensorrt_llm._torch.speculative import SpecMetadata from tensorrt_llm._utils import nvtx_range diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 242c6cedda90..0c2409fd5a06 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -14,1811 +14,46 @@ # # SPDX-License-Identifier: Apache-2.0 -import contextlib -import copy -import itertools -from dataclasses import dataclass -from typing import ( - TYPE_CHECKING, - Any, - Callable, - ClassVar, - Dict, - Hashable, - Iterable, - Iterator, - List, - Optional, - Sequence, - Tuple, +"""Compatibility exports for the shared multimodal runtime.""" + +from ..multimodal.model import ( + EncoderCachePartition, + EncoderGroup, + MultimodalEncoderContractError, + MultimodalModelContractError, + MultimodalModelMixin, + MultimodalRuntimeMixin, + PreparedLlmInputs, + PreparedMultimodalEmbeddings, + _assemble_multimodal_encoder_embeddings, + _build_request_multimodal_input, + _collect_cross_iter_prefetch_candidates, + _dispatch_cross_iter_prefetch, + _get_mm_aux_stream, + _lengths_by_modality, + _reorder_embeds_by_manifest, + _run_on_aux_stream, + encode_multimodal_by_groups, + maybe_prefetch_mm_encoder_for_next_iter, ) -import torch - -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache -from tensorrt_llm._utils import prefer_pinned -from tensorrt_llm.inputs.multimodal import ( - MULTIMODAL_ENCODER_ITEM_METADATA_KEY, - MultimodalInput, - MultimodalItemMetadata, - MultimodalItemPartition, - MultimodalParams, - MultimodalRuntimeData, - select_multimodal_fields, -) -from tensorrt_llm.logger import logger - -from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs -from ..pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS, RuntimeHooks, runtime_hooks_of -from .modeling_multimodal_utils import ( - _store_chunked_prefill_embeddings, - find_input_mm_embeds, - fuse_input_embeds, - get_multimodal_embeddings, -) - - -class MultimodalEncoderContractError(ValueError): - """A request-local MM encoder input or output contract violation.""" - - -def _assemble_multimodal_encoder_embeddings( - item_tensors: Dict[int, torch.Tensor], total_items: int -) -> torch.Tensor: - """Copy prompt-ordered item embeddings into one contiguous tensor.""" - ordered = [item_tensors[i] for i in range(total_items)] - first = ordered[0] - if any( - output.shape[1:] != first.shape[1:] - or output.dtype != first.dtype - or output.device != first.device - for output in ordered[1:] - ): - raise ValueError( - "MM encoder items for one request must have matching output shape, dtype, and device" - ) - - # The copy is intentional even for one item: cache hits alias cache-owned - # storage, while this result must be independently owned by the request. - return torch.cat([output.detach() for output in ordered], dim=0) - - -@dataclass(frozen=True) -class EncoderGroup: - """Modalities that share a single encoder call. - - Batching all items in a group into one encoder invocation amortizes - fixed costs (kernel launches, dispatch) across items. The framework - splits the output back per-modality and reorders it into prompt order - via each request's `MultimodalItemMetadata.item_refs` manifest. - - Contract between `build_batched_input` and `encoder_fn`: - - * `build_batched_input` must concatenate items across requests in - `modalities` order (all items for the first modality across requests, - then all items for the second modality, etc.). - * Within a modality, items must appear in the same per-request iteration - order that `_lengths_by_modality` uses (i.e. the order of - `multimodal_params` passed in). - * `encoder_fn` must return one tensor whose rows correspond 1:1 to the - input layout produced by `build_batched_input`, so the framework can - split the output by `_lengths_by_modality` and reorder into prompt - order via each request's item manifest. - """ - - modalities: Tuple[str, ...] - """Ordered modality names that share this encoder. Defines the row - layout of the encoder output tensor: first all items of - `modalities[0]`, then all items of `modalities[1]`, etc.""" - - encoder_fn: Callable[..., torch.Tensor] - """Encoder call invoked as `encoder_fn(**build_batched_input(params))`. - Returns a single tensor with one row per embedding, laid out per the - contract above. When those rows are already in prompt order, the - framework may return a row slice of this tensor without a copy, so it - must not alias storage that a later call overwrites (for example, a CUDA - graph's static output buffer); clone such outputs before returning - them.""" - - build_batched_input: Callable[[List[MultimodalParams]], Dict[str, Any]] - """Builds the kwargs dict passed to `encoder_fn`. Responsible for - concatenating raw per-item tensors from `multimodal_data` across - requests in the order described in the class docstring.""" - - -def _lengths_by_modality( - item_layouts: Sequence[tuple[Sequence[tuple[str, int]], Sequence[int]]], - modalities: Tuple[str, ...], -) -> Dict[str, List[int]]: - """Split resolved output lengths into each group's payload-local order. - - Layouts describe whole requests, so only modalities in this encoder - group contribute lengths (for example, image/video in a vision group - while the same request's audio is handled by another group). - - Encoders batch in payload order, which can differ from prompt order. - Modality-local indices restart for each request. For example:: - - item_layouts = [ - ([("image", 1), ("video", 0), ("image", 0)], [20, 30, 10]), - ([("audio", 0), ("image", 0)], [40, 50]), - ] - modalities = ("image", "video") - # Result: {"image": [10, 20, 50], "video": [30]} - - These lengths split encoder outputs before restoring prompt order. - """ - by_modality: Dict[str, List[int]] = {m: [] for m in modalities} - for item_refs, lengths in item_layouts: - if not item_refs: - continue - modality = item_refs[0][0] - if ( - len(item_refs) == len(lengths) - and modality in by_modality - and all( - item_modality == modality and index == expected - for expected, (item_modality, index) in enumerate(item_refs) - ) - ): - by_modality[modality].extend(lengths) - continue - request_lengths: dict[str, list[tuple[int, int]]] = {m: [] for m in modalities} - for (modality, index), length in zip(item_refs, lengths, strict=True): - if modality in request_lengths: - request_lengths[modality].append((index, length)) - for modality, indexed_lengths in request_lengths.items(): - by_modality[modality].extend(length for _, length in sorted(indexed_lengths)) - return by_modality - - -def _reorder_embeds_by_manifest( - item_layouts: Sequence[tuple[Sequence[tuple[str, int]], Sequence[int]]], - per_modality_embeds: Dict[str, torch.Tensor], - per_modality_lengths: Dict[str, List[int]], -) -> torch.Tensor: - """Slice per-modality tensors item-by-item and concat in prompt order. - - If only one modality has items and every request lists them in encoder - order, that modality's tensor is returned as is, without a copy. - """ - populated_modalities = [ - (modality, embedding) - for modality, embedding in per_modality_embeds.items() - if per_modality_lengths.get(modality) - ] - if len(populated_modalities) == 1: - modality, embedding = populated_modalities[0] - prompt_order_matches_encoder = True - for item_refs, _ in item_layouts: - item_indices = [ - index for item_modality, index in item_refs if item_modality == modality - ] - if item_indices != list(range(len(item_indices))): - prompt_order_matches_encoder = False - break - if prompt_order_matches_encoder: - # A group with one populated modality already emits request and - # item rows in prompt order. Empty peer modalities (for example, - # video in an image-only Qwen batch) do not require a reorder. - return embedding - - per_modality_row_starts: Dict[str, List[int]] = { - m: list(itertools.accumulate(lens, initial=0)) for m, lens in per_modality_lengths.items() - } - - slices: List[torch.Tensor] = [] - # Item indices are per-request per-modality; advance a cursor to - # translate it into a global item index within `per_modality_embeds`. - per_modality_cursor: Dict[str, int] = {m: 0 for m in per_modality_embeds} - for item_refs, _ in item_layouts: - req_counts: Dict[str, int] = {} - for m, item_index in item_refs: - if m not in per_modality_embeds: - continue - i = per_modality_cursor[m] + item_index - starts = per_modality_row_starts[m] - slices.append(per_modality_embeds[m][starts[i] : starts[i + 1]]) - req_counts[m] = req_counts.get(m, 0) + 1 - for m, c in req_counts.items(): - per_modality_cursor[m] += c - if not slices: - # Preserve encoder dtype/device and width when no item contributes rows. - if per_modality_embeds: - return torch.cat(list(per_modality_embeds.values()), dim=0) - return torch.empty(0) - return torch.cat(slices, dim=0) - - -def encode_multimodal_by_groups( - mm_encoder_groups: Sequence["EncoderGroup"], - multimodal_params: List[MultimodalParams], -) -> torch.Tensor: - """Run each group's encoder over its batched items and reorder into - per-request prompt order. - - For each group present in the batch, one encoder call is issued over all - items across all requests belonging to that group's modalities - (arithmetic-intensity win). The output is split back per-modality using - the prompt-ordered output lengths in `MultimodalItemMetadata`, - then reordered by the same metadata's `item_refs`. - - Shared entry point for both the aggregated (`MultimodalModelMixin`) and - mm-encoder-only (`Qwen3VisionModelBase.forward`) paths so the ordering - contract lives in one place. - """ - per_modality_embeds: Dict[str, torch.Tensor] = {} - per_modality_lengths: Dict[str, List[int]] = {} - item_layouts = [mp.get_item_layout() for mp in multimodal_params] - for group in mm_encoder_groups: - group_params = [] - group_layouts = [] - for mp, layout in zip(multimodal_params, item_layouts, strict=True): - if any(mp.multimodal_data.get(m) is not None for m in group.modalities): - if not layout[0]: - raise MultimodalEncoderContractError( - "MM item metadata with output lengths is required for grouped encoding" - ) - group_params.append(mp) - group_layouts.append(layout) - if not group_params: - continue - out = group.encoder_fn(**group.build_batched_input(group_params)) - lengths = _lengths_by_modality(group_layouts, group.modalities) - cursor = 0 - for m in group.modalities: - total = sum(lengths[m]) - per_modality_embeds[m] = out[cursor : cursor + total] - cursor += total - per_modality_lengths.update(lengths) - return _reorder_embeds_by_manifest(item_layouts, per_modality_embeds, per_modality_lengths) - - -if TYPE_CHECKING: - from ..pyexecutor.guided_decoder import CapturableGuidedDecoder - from ..pyexecutor.llm_request import LlmRequest - - -_MM_DATA_INPUT_MODALITY_KEYS = frozenset({"audio", "image", "video"}) -_MM_AUX_STREAM: Optional[tuple[int, torch.cuda.Stream]] = None -_MM_ENCODER_CACHE_LOG_NAME = "mm_encoder_cache" - - -def _build_request_multimodal_input( - request: "LlmRequest", cache_enabled: bool -) -> Optional[MultimodalInput]: - """Build the encoder-cache key metadata carried by one request.""" - if not cache_enabled: - return None - params = getattr(request, "py_multimodal_params", None) - if params is not None and params.multimodal_input is not None: - return params.multimodal_input - hashes = request.multimodal_hashes - if hashes is None: - return None - multimodal_input = MultimodalInput.from_components( - hashes, - request.multimodal_positions, - request.multimodal_lengths, - mm_item_run_cu_offsets=request.multimodal_item_run_cu_offsets, - mm_run_positions=request.multimodal_run_positions, - mm_run_lengths=request.multimodal_run_lengths, - ) - if params is not None: - params.multimodal_input = multimodal_input - return multimodal_input - - -def _get_mm_aux_stream(max_prefetch_ahead: int = 0) -> Optional[torch.cuda.Stream]: - """Return the side CUDA stream used for multimodal encoder prefetch. - - Returns `None` when side-stream prefetch is disabled, CUDA is unavailable, - or the current stream is being captured. The cache intentionally keeps only - one stream because executor processes are expected to run on one current - CUDA device; if the current device changes, the cached stream is replaced. - """ - global _MM_AUX_STREAM - - if max_prefetch_ahead <= 0: - return None - if not torch.cuda.is_available(): - return None - if torch.cuda.is_current_stream_capturing(): - return None - - device = torch.cuda.current_device() - if _MM_AUX_STREAM is None or _MM_AUX_STREAM[0] != device: - _MM_AUX_STREAM = (device, torch.cuda.Stream(device=device)) - logger.warning_once( - f"Using multimodal encoder side stream on CUDA device {device} " - f"with encoder_side_stream_max_ahead={max_prefetch_ahead}. " - "This may increase peak GPU memory usage because raw multimodal " - "encoder inputs and computed embeddings can be resident before " - "request prefill.", - key=f"mm_aux_stream_used_device_{device}", - ) - return _MM_AUX_STREAM[1] - - -@contextlib.contextmanager -def _run_on_aux_stream(aux_stream: torch.cuda.Stream) -> Iterator[torch.cuda.Event]: - """Run a block on `aux_stream` independently of the caller stream. - - Yields a CUDA event recorded on `aux_stream` when the block exits. Callers - can wait on that event from another stream before consuming tensors written - in the block. - - No entrance barrier is enforced: cross-iter MM encoder prefetch operates on - data disjoint from the current iteration's batch, so serializing aux-stream - work behind the caller stream would eliminate the overlap this stream - exists for. - """ - exit_event = torch.cuda.Event() - with torch.cuda.stream(aux_stream): - try: - yield exit_event - finally: - # Keep the sync point valid even if the block raises after queuing - # aux-stream work. - exit_event.record() - - -@dataclass(frozen=True) -class PreparedLlmInputs: - """Prepared inputs returned by `MultimodalModelMixin`.""" - - input_ids: Optional[torch.Tensor] - inputs_embeds: Optional[torch.Tensor] - extra_embeds: Sequence[torch.Tensor] = () - - -@dataclass(frozen=True) -class EncoderCachePartition: - """Per-item cache partition for a single `MultimodalParams`. - - `hits` maps item index to its cached embedding row-block; `miss_indices` lists item - indices that still require encoder work; `keys` is aligned to item order so miss - embeddings can be written back after they are computed. - - `keys` always spans every item in the request -- full embedding assembly needs the total - item count, and misses are written back by item index. `hits` and `miss_indices` instead - cover only `looked_up`, the indices the caller asked about. The item scheduler passes the - subset it picked for this iteration, so items it has no budget for yet are not touched: an - LRU `get` refreshes recency, and probing an item the caller cannot encode would reorder - eviction against items that are actually in flight. - """ - - hits: Dict[int, torch.Tensor] - miss_indices: list[int] - keys: list[Hashable] - looked_up: list[int] - - @property - def is_full_hit(self) -> bool: - return bool(self.looked_up) and not self.miss_indices - - @property - def is_full_miss(self) -> bool: - return bool(self.looked_up) and not self.hits - - -class MultimodalModelMixin: - """Template-method mixin for PyTorch multimodal causal LM models. - - Concrete model forwards can call `prepare_multimodal_inputs` while keeping their explicit - language-model delegation. - - Current limitations: - - * For the time being, the persistent multimodal encoder cache stores per-item embeddings for - single-modality `MultimodalParams` objects. Mixed-modality objects bypass the cache. - * A partially cached `MultimodalParams` is handled by encoding only its miss items - and interleaving cached items back in original per-item order. The default - `build_multimodal_encoder_input` selects fields declared by - `get_multimodal_item_partitions`. Undeclared models retain legacy stacked/packed inference. - """ - - supports_vision_gpu_normalize: ClassVar[bool] = False - """Whether the model accepts processor-described raw vision pixels.""" - - supports_encoder_cache: ClassVar[bool] = False - """Whether the model's production forward path uses the persistent encoder cache.""" - - supports_mm_encoder_item_scheduling: ClassVar[bool] = False - """Whether the model supports item-level MM encoder scheduling: it implements the item-encode - forward and pairs with a processor that overrides - `BaseMultimodalInputProcessor.get_mm_item_metadata`. This is the executor-layer gate for - item scheduling.""" - - model_config: ModelConfig - _multimodal_encoder_cache: Optional[TensorLRUCache] = None - - @classmethod - def _cast_multimodal_encoder_dtype( - cls, - module: torch.nn.Module, - dtype: torch.dtype, - ) -> torch.nn.Module: - """Cast a multimodal encoder dtype without materializing meta tensors.""" - - def convert(tensor: torch.Tensor) -> torch.Tensor: - if not (tensor.is_floating_point() or tensor.is_complex()): - return tensor - if tensor.device == torch.device("meta"): - return torch.empty_like(tensor, dtype=dtype) - return tensor.to(dtype=dtype) - - return module._apply(convert) - - # Per-model registration of encoder-batching groups. Each `EncoderGroup` - # bundles a set of modalities that share one encoder call. Set as a class - # attribute or on `self` in `__init__` (when `encoder_fn` binds to instance - # methods). Consumers call the module-level `encode_multimodal_by_groups` - # with these groups; both the aggregated and mm-encoder-only paths share - # that helper so the ordering contract lives in one place. - mm_encoder_groups: Sequence[EncoderGroup] = () - - def encode_multimodal_inputs( - self, - multimodal_params: Sequence[MultimodalParams], - ) -> torch.Tensor: - """Run model-specific multimodal encoder work. - - Returns the single primary multimodal embedding tensor for the supplied params. Rows are - expected to be concatenated in request order, and special multimodal tokens occupy token - positions but do not have rows here. - """ - raise NotImplementedError - - def validate_multimodal_request_data(self, mm_data: dict) -> None: - """Admission-time validation of one request's multimodal payload. - - Called from `PyExecutor._validate_request` before activation; raise - `ValueError` to reject just this request. Conditions a model would - reject inside `forward` belong here instead: a forward-time failure - is handled batch-wide and fails every request scheduled alongside - the bad one. Default: accept everything. - """ - return None - - def prepare_multimodal_encoder_inputs( - self, - selected_items: Sequence[tuple[MultimodalParams, int]], - ) -> list[tuple[MultimodalParams, Sequence[int], str]]: - """Build selected item encoder inputs before the caller performs H2D. - - Adjacent items from the same request and modality are sliced in one - call. That is not just tidier: the packed-layout slicer splits the - request's whole pixel payload and concatenates the chosen pieces, so - slicing item-by-item re-splits the payload N times and copies each - item separately. - - Args: - selected_items: `(request params, item index)` pairs in scheduler-selected order. - - Returns: - Tuples of `(sliced encoder params, per-item embedding row counts, modality)`, in - input order. Flattening the row-count lists recovers the per-item sequence. - """ - encoder_inputs: list[tuple[MultimodalParams, Sequence[int], str]] = [] - for ( - multimodal_param, - run_indices, - modality, - item_metadata, - ) in self._runs_by_request_modality(selected_items): - item_refs = item_metadata.item_refs - # The two slicers take different index spaces. The raw-tensor - # slicer indexes the modality's own payload, so it gets the - # modality-local indices (plus the modality, which also lets it - # slice out of an interleaved request); the parallel metadata is - # prompt-ordered across modalities, so it gets the global ones. - try: - residual = self.build_multimodal_encoder_input( - multimodal_param, - [item_refs[i][1] for i in run_indices], - modality=modality, - ) - self._apply_metadata_slice( - residual, multimodal_param, run_indices, item_metadata=item_metadata - ) - except MultimodalEncoderContractError: - raise - except (KeyError, IndexError, TypeError, ValueError) as error: - raise MultimodalEncoderContractError( - f"Invalid multimodal encoder item input: {error}" - ) from error - encoder_inputs.append( - ( - residual, - residual.item_metadata.output_embedding_lengths, - modality, - ) - ) - return encoder_inputs - - @staticmethod - def _runs_by_request_modality( - selected_items: Sequence[tuple[MultimodalParams, int]], - ) -> Iterator[tuple[MultimodalParams, list[int], str, MultimodalItemMetadata]]: - """Split scheduler order into maximal same-request, same-modality runs. - - Only adjacent items merge, preserving scheduler order. Reuse the item - metadata validated at the input boundary or executor admission. - """ - run_param: Optional[MultimodalParams] = None - run_modality: Optional[str] = None - run_metadata: Optional[MultimodalItemMetadata] = None - run_indices: list[int] = [] - for multimodal_param, item_idx in selected_items: - try: - metadata = ( - run_metadata - if multimodal_param is run_param - else multimodal_param.item_metadata - ) - except (TypeError, ValueError) as error: - raise MultimodalEncoderContractError(str(error)) from error - if metadata is None or metadata.output_embedding_lengths is None: - raise MultimodalEncoderContractError( - "MM item metadata with output lengths is required for item encoding" - ) - if item_idx < 0 or item_idx >= len(metadata.item_refs): - raise MultimodalEncoderContractError( - f"MM item index {item_idx} is out of range for {len(metadata.item_refs)} item(s)" - ) - modality = metadata.item_refs[item_idx][0] - if run_indices and (multimodal_param is not run_param or modality != run_modality): - yield run_param, run_indices, run_modality, run_metadata - run_indices = [] - run_param, run_modality, run_metadata = multimodal_param, modality, metadata - run_indices.append(item_idx) - if run_indices: - yield run_param, run_indices, run_modality, run_metadata - - def forward_multimodal_encoder_items( - self, - encoder_inputs: Sequence[tuple[MultimodalParams, Sequence[int], str]], - ) -> list[torch.Tensor]: - """Forward prepared MM encoder inputs in scheduler item order. - - Args: - encoder_inputs: Tuples returned by `prepare_multimodal_encoder_inputs`. Consecutive - inputs with the same modality must be batch-compatible. - - Returns: - One encoder output tensor **per item** (not per input tuple). Each tensor has the - declared embedding row count and retains scheduler input order. - """ - outputs: list[torch.Tensor] = [] - group_params: list[MultimodalParams] = [] - group_lengths: list[int] = [] - group_modality: Optional[str] = None - - def flush_group() -> None: - if not group_params: - return - embeddings = self.encode_multimodal_inputs(group_params) - expected_length = sum(group_lengths) - if embeddings.shape[0] != expected_length: - raise MultimodalEncoderContractError( - f"MM encoder output length {embeddings.shape[0]} does not " - f"match the {expected_length} rows declared by the " - "selected items" - ) - outputs.extend(torch.split(embeddings, group_lengths, dim=0)) - group_params.clear() - group_lengths.clear() - - for multimodal_param, embedding_lengths, modality in encoder_inputs: - if group_modality is not None and modality != group_modality: - flush_group() - group_modality = modality - group_params.append(multimodal_param) - group_lengths.extend(embedding_lengths) - flush_group() - return outputs - - @property - def multimodal_token_ids(self) -> Optional[Sequence[int] | torch.Tensor]: - """Return placeholder token ids in `input_ids` replaced by MM embeds. - - These are sentinel token positions whose text embeddings are replaced - by multimodal embeddings. Return `None` to use the out-of-vocabulary - sentinel behavior in `fuse_input_embeds`. - """ - return None - - @property - def text_embedding_layer(self): - """Return the token embedding layer used by `fuse_input_embeds`.""" - raise NotImplementedError - - @property - def embedding_dim(self) -> int: - """Return the width of each cached multimodal embedding row.""" - raise NotImplementedError - - @property - def embedding_dtype(self) -> torch.dtype: - """Return the dtype of each cached multimodal embedding row.""" - raise NotImplementedError - - @property - def encoder_cache_active(self) -> bool: - """Whether the persistent encoder cache is active for this model. - - Single source of truth shared by: - - * the in-iter consume path - * the side-stream prefetch dispatch - * the engine's cache-related gate - * the KV-cache memory reservation. - - The cache is only active when the model opts in via `supports_encoder_cache` and configures - a positive capacity. - """ - if not self.supports_encoder_cache: - return False - multimodal_config = self.model_config.multimodal_config - return multimodal_config is not None and multimodal_config.encoder_cache_max_bytes > 0 - - def select_multimodal_params( - self, - multimodal_params: Sequence[MultimodalParams], - num_context_requests: int, - ) -> Sequence[MultimodalParams]: - """Select the params that participate in multimodal encoder work. - - 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() - 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 runtime_hooks(self) -> RuntimeHooks: - """Return the inner language model's step-loop hooks.""" - try: - language_model = self.language_model - except NotImplementedError: - return DEFAULT_RUNTIME_HOOKS - return runtime_hooks_of(language_model) - - @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 set_guided_decoder(self, guided_decoder: "CapturableGuidedDecoder") -> bool: - """Install a guided decoder on the inner language model. - - Returns False when the inner model does not support guided decoding, - matching the contract ModelEngine.set_guided_decoder() expects. - """ - inner = self.language_model - if not hasattr(inner, "set_guided_decoder"): - return False - return inner.set_guided_decoder(guided_decoder) - - def get_language_model_extra_forward_kwargs( - self, - *, - attn_metadata: Any, - 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. - - The attention metadata identifies context and generation requests in this batch. - """ - 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( - attn_metadata=attn_metadata, - 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, - *, - input_ids: torch.Tensor, - multimodal_params: Sequence[MultimodalParams], - embeddings: torch.Tensor, - **forward_kwargs: Any, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Optional hook after multimodal embeddings are gathered. - - Runs after cache lookup or encoder execution, before active-row model - hooks. The row layout depends on the whole batch: if any param carries - item-scheduled current-chunk rows, the gather also cuts every other - param that has runtime data down to its current-chunk rows; otherwise - each param keeps its full-request rows. Row-dependent model transforms - should use `after_active_multimodal_embeddings` instead. - """ - return input_ids, embeddings - - def after_active_multimodal_embeddings( - self, - *, - active_embeddings: list[torch.Tensor], - multimodal_params: Sequence[MultimodalParams], - **forward_kwargs: Any, - ) -> tuple[list[torch.Tensor], Sequence[torch.Tensor]]: - """Optional hook after active chunk rows are selected and before fusion. - - Models can transform or split the active multimodal embeddings here - and return additional embedding tensors to fuse alongside the primary - multimodal embeddings. - """ - # Models with packed auxiliary features (e.g. Qwen3-VL) can split them here and return - # them as extra embeds without changing the base flow. - return active_embeddings, () - - def get_multimodal_item_partitions( - self, - param: MultimodalParams, - modality: str, - ) -> dict[str, MultimodalItemPartition] | None: - """Declare how items occupy each field of one modality's payload. - - Models declare batched axes, packed item lengths, and shared values. - Unlisted fields pass through unchanged. Returning None retains legacy - layout inference for models that have not declared their fields. - """ - metadata = param.item_metadata - return metadata.input_partitions.get(modality) if metadata is not None else None - - def build_multimodal_encoder_input( - self, - param: MultimodalParams, - local_item_indices: Sequence[int], - modality: Optional[str] = None, - ) -> MultimodalParams: - """Build a forward view containing selected modality-local items. - - ``local_item_indices`` indexes this modality's payload, not positions - in the request's prompt-ordered item metadata. - Declared field layouts drive common tensor/list selection. Models with - custom preprocessing can still override this hook. The caller slices - the parallel item metadata after this returns, preserving original - request indices for scheduler targets and cache ownership. - """ - if modality is None: - modality = self._encoder_cache_modality(param) - if modality is None: - raise NotImplementedError( - "Default `build_multimodal_encoder_input` cannot infer the modality of a " - "mixed-modality param. Pass `modality` with modality-local indices, or " - "override for other layouts." - ) - modality_data = param.multimodal_data[modality] - if not isinstance(modality_data, dict): - raise TypeError( - f"multimodal_data[{modality!r}] must be a dict, got {type(modality_data).__name__}" - ) - item_partitions = self.get_multimodal_item_partitions(param, modality) - if item_partitions is None: - sliced = self._select_legacy_multimodal_fields( - modality_data, modality, local_item_indices - ) - else: - sliced = select_multimodal_fields(modality_data, item_partitions, local_item_indices) - - # Metadata slicing must not rewrite the original request's hashes. - residual_input = ( - copy.copy(param.multimodal_input) if param.multimodal_input is not None else None - ) - residual_data = { - key: value - for key, value in param.multimodal_data.items() - if key not in _MM_DATA_INPUT_MODALITY_KEYS or key == modality - } - residual_data[modality] = sliced - return MultimodalParams( - multimodal_data=residual_data, - multimodal_input=residual_input, - ) - - def _select_legacy_multimodal_fields( - self, - modality_data: Dict[str, Any], - modality: str, - item_indices: Sequence[int], - ) -> Dict[str, Any]: - """Select recognized packed/stacked layouts for undeclared models.""" - indices = list(item_indices) - grid_key = {"image": "image_grid_thw", "video": "video_grid_thw"}.get(modality) - pixel_key = {"image": "pixel_values", "video": "pixel_values_videos"}.get(modality) - - if ( - (grid_key and pixel_key) - and (grid_key in modality_data) - and (pixel_key in modality_data) - ): - # Packed layout: prefix-sum patch counts to locate each item's slab, then - # concat the requested subset in item-index order. - grids = modality_data[grid_key] - n_items = grids.shape[0] - patch_counts = [int(c) for c in torch.prod(grids, dim=1).tolist()] - row_starts = list(itertools.accumulate(patch_counts, initial=0)) - if indices == list(range(indices[0], indices[0] + len(indices))): - # Contiguous run: the concatenation is just a row range, so take - # a view instead of copying the payload. The common case for - # both callers -- a scheduler picks items in order, and cache - # misses cluster. - pixel_slice = modality_data[pixel_key][ - row_starts[indices[0]] : row_starts[indices[-1] + 1] - ] - else: - per_item = torch.split(modality_data[pixel_key], patch_counts, dim=0) - pixel_slice = torch.cat([per_item[i] for i in indices], dim=0) - sliced = { - pixel_key: pixel_slice, - grid_key: grids[indices], - } - elif ( - modality == "image" - and "pixel_values" in modality_data - and "image_sizes" in modality_data - ): - # Stacked layout: dim-0 select from `pixel_values` and list-index `image_sizes`. - n_items = modality_data["pixel_values"].shape[0] - miss_sizes = [modality_data["image_sizes"][i] for i in indices] - miss_pixel = modality_data["pixel_values"][indices] - # `pixel_values` was padded to the request-wide max H/W by the input - # processor. After keeping only the miss subset, crop the trailing H/W back - # down to that subset's own max true size -- otherwise a downstream re-batch - # step (e.g. Mistral 3's `batch_pixel_values`) that pads to - # `max(residual.image_sizes)` would compute a negative pad amount whenever - # the omitted items were the largest in the original request. - if miss_sizes and miss_pixel.dim() >= 4: - max_h = max(int(s[0]) for s in miss_sizes) - max_w = max(int(s[1]) for s in miss_sizes) - miss_pixel = miss_pixel[..., :max_h, :max_w] - sliced = { - "pixel_values": miss_pixel, - "image_sizes": miss_sizes, - } - elif modality == "audio" and "input_features" in modality_data: - # Stacked layout: `input_features [B, mel_bins, T]` sliced on dim 0. - # Per-item masks (`input_features_mask`, `feature_attention_mask`, ...) - # are handled by the sibling-slice pass below. - n_items = modality_data["input_features"].shape[0] - sliced = { - "input_features": modality_data["input_features"][indices], - } - else: - raise NotImplementedError( - f"Default `build_multimodal_encoder_input` cannot slice {modality} layout " - f"with fields {sorted(modality_data)}; override this method." - ) - - # Sibling per-item fields (e.g. `second_per_grid_ts` on Qwen2.5-VL video) - # must be sliced alongside the load-bearing keys above, or the residual - # carries a shape-mismatched encoder input. - sliced = { - **modality_data, - **sliced, - **self._slice_per_item_sibling_fields(modality_data, n_items, indices, sliced.keys()), - } - - return sliced - - @staticmethod - def _slice_per_item_sibling_fields( - modality_data: Dict[str, Any], - n_items: int, - item_indices: Sequence[int], - already_sliced: Iterable[str], - ) -> Dict[str, Any]: - """Slice modality-dict siblings whose first axis is parallel to items. - - Anything with `shape[0] == n_items` (tensor) or `len == n_items` (list) is - assumed to be per-item metadata and sliced by `item_indices`. Fields already - handled by the caller (`already_sliced`) and everything else pass through. - """ - skip = set(already_sliced) - sliced: Dict[str, Any] = {} - for key, value in modality_data.items(): - if key in skip: - continue - if isinstance(value, torch.Tensor) and value.dim() > 0 and value.shape[0] == n_items: - sliced[key] = value[item_indices] - elif isinstance(value, list) and len(value) == n_items: - sliced[key] = [value[i] for i in item_indices] - return sliced - - # A future optional mixin-owned forward can build on the same template method. - def prepare_multimodal_inputs( - self, - *, - input_ids: torch.Tensor, - positions: Optional[torch.Tensor], - multimodal_params: Optional[Sequence[MultimodalParams]], - num_context_requests: int, - **forward_kwargs: Any, - ) -> PreparedLlmInputs: - """Prepare multimodal inputs for a concrete model forward. - - This method owns the common framework sequence around a model-specific - encoder hook: retrieve/cache embeddings, select active chunk rows if - needed, run optional model hooks, and fuse rows into text embeds. - """ - context_params = list( - self.select_multimodal_params( - multimodal_params or [], - num_context_requests, - ) - ) - if not context_params: - return PreparedLlmInputs(input_ids=input_ids, inputs_embeds=None) - - embeddings = self._get_or_encode_multimodal_embeddings(context_params) - - input_ids, embeddings = self.after_full_multimodal_embeddings( - input_ids=input_ids, - multimodal_params=context_params, - embeddings=embeddings, - **forward_kwargs, - ) - - active_embeddings = find_input_mm_embeds([embeddings], list(context_params)) - active_embeddings, extra_embeds = self.after_active_multimodal_embeddings( - active_embeddings=active_embeddings, - multimodal_params=context_params, - **forward_kwargs, - ) - - fused_input_ids, inputs_embeds, fused_extra_embeds = self._fuse_multimodal_embeddings( - input_ids=input_ids, - multimodal_embeddings=active_embeddings, - mm_token_ids=self.multimodal_token_ids, - embedding_layer=self.text_embedding_layer, - extra_embeds=extra_embeds, - # `text_token_indices` / `mm_token_indices` are pre-computed by the - # executor (see DecoderInputPreparer.prepare_inputs) and must reach - # `fuse_input_embeds` to (a) preserve the active-chunk subset - # contract when MM rows are a subset of visible MM tokens and - # (b) avoid the torch.where host sync inside - # `filter_mm_token_from_input_ids`. - text_token_indices=forward_kwargs.get("text_token_indices"), - mm_token_indices=forward_kwargs.get("mm_token_indices"), - ) - return PreparedLlmInputs( - input_ids=fused_input_ids, - inputs_embeds=inputs_embeds, - extra_embeds=fused_extra_embeds, - ) - - def _get_or_encode_multimodal_embeddings( - self, - multimodal_params: Sequence[MultimodalParams], - ) -> torch.Tensor: - """Return cached multimodal embeddings or run the encoder for misses. - - Delegates cache lookup and gather behavior to `get_multimodal_embeddings`, then validates - the single tensor contract for both encoded and cached-only paths. - If item scheduling supplies current-chunk rows, the gather selects - current-chunk rows for every request before concatenating the batch. - - During side-stream prefetch, this runs with the auxiliary stream current, so the H2D copies, - the encoder, and every persistent-cache `put()` are issued on that stream. `TensorLRUCache` - records each entry's producer event on the issuing (aux) stream; the next iteration's - main-stream consumer waits on the request-level `encoder_event` for ordering. - """ - encoder_cache = self._multimodal_encoder_cache - cache_misses: list[MultimodalParams] = [] - partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = [] - if encoder_cache is not None: - for param in multimodal_params: - if param.multimodal_data.get("multimodal_embedding") is not None: - # A present embedding means either an earlier forward already wrote the - # persistent cache, or a prefetch hit attached a cache-owned tensor. Either - # way, skip re-lookup and re-write. `get_multimodal_embeddings` waits on the - # request event and records the attached tensor on the consuming stream before - # gathering it. - continue - partition = self.partition_encoder_cache(param, encoder_cache) - if partition is None or partition.is_full_miss: - cache_misses.append(param) - continue - if partition.is_full_hit: - param.multimodal_data["multimodal_embedding"] = ( - _assemble_multimodal_encoder_embeddings(partition.hits, len(partition.keys)) - ) - continue - partial_hits.append((param, partition)) - - if partial_hits: - # `encoder_cache` is non-None here because partitions are only produced when the cache - # exists. - self._encode_with_partial_cache(partial_hits, encoder_cache) - - embeddings = get_multimodal_embeddings( - encoder_forward_fn=self.encode_multimodal_inputs, - multimodal_params=list(multimodal_params), - ) - if encoder_cache is not None: - for param in cache_misses: - self._write_encoder_cache_entries(param, encoder_cache) - - # Validate post-gather so cached-only paths (KV reuse, all-cached chunked prefill) are also - # checked, not just paths that ran the encoder. - self._validate_embeddings( - embeddings, - multimodal_params, - current_chunk_only=any( - param.multimodal_data.get("multimodal_embedding_is_chunk", False) - for param in multimodal_params - ), - ) - return embeddings[0] - - def _initialize_multimodal_encoder_cache(self, max_bytes: int) -> Optional[TensorLRUCache]: - """Initialize the model-owned multimodal encoder-output cache once. - - The cache stores per-item embeddings keyed by the common item metadata. - `ModelEngine` resolves `max_bytes` as the larger of the item-scheduling - output budget and persistent-reuse capacity before runtime access. - Zero leaves the cache disabled. - """ - if self._multimodal_encoder_cache is not None: - if max_bytes != self._multimodal_encoder_cache.max_bytes: - raise ValueError( - f"{_MM_ENCODER_CACHE_LOG_NAME}: already initialized with " - f"max_bytes={self._multimodal_encoder_cache.max_bytes}; " - f"cannot reinitialize with max_bytes={max_bytes}" - ) - return self._multimodal_encoder_cache - - if max_bytes == 0: - logger.debug_once( - f"{_MM_ENCODER_CACHE_LOG_NAME}: disabled because neither item scheduling nor " - "persistent reuse requires storage.", - key="mm_encoder_cache_disabled", - ) - return None - - multimodal_config = self.model_config.multimodal_config - self._multimodal_encoder_cache = TensorLRUCache( - max_bytes, - name=_MM_ENCODER_CACHE_LOG_NAME, - cuda_stream_aware=multimodal_config.encoder_side_stream_max_ahead > 0, - ) - try: - embedding_dim = self.embedding_dim - embedding_dtype = self.embedding_dtype - except NotImplementedError: - logger.info( - f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, " - "embedding row capacity unavailable because the model does not implement " - "embedding_dim and embedding_dtype." - ) - else: - bytes_per_embedding_row = ( - embedding_dim * torch.empty((), dtype=embedding_dtype).element_size() - ) - max_embedding_rows = max_bytes // bytes_per_embedding_row - logger.info( - f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, " - f"max_embedding_rows={max_embedding_rows}, embedding_dim={embedding_dim}, " - f"embedding_dtype={embedding_dtype}" - ) - return self._multimodal_encoder_cache - - @staticmethod - def _encoder_cache_modality(param: MultimodalParams) -> Optional[str]: - """Return the single modality represented by `param`, if cacheable. - - `None` means the params either do not identify a modality or contain - multiple modality inputs. Per-item operations use the common item refs - when one request contains several modalities. - """ - mm_data = param.multimodal_data or {} - modalities = [key for key in _MM_DATA_INPUT_MODALITY_KEYS if key in mm_data] - - modality = mm_data.get("modality_type") - if isinstance(modality, str): - # Trust the explicit `modality_type` only when it agrees with the actual data keys. - # Otherwise fall through to the mixed-modality skip so an inconsistent producer (e.g. - # `modality_type="image"` while both image and audio data are present) cannot bypass the - # safety check below and have the cache serve embeddings for the wrong modality. - if modalities == [modality]: - return modality - - if len(modalities) != 1: - # A whole-param modality is ambiguous for mixed inputs. - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params with {len(modalities)} detected " - "modalities." - ) - return None - return modalities[0] - - @classmethod - def _encoder_cache_keys( - cls, - param: MultimodalParams, - ) -> Optional[list[Hashable]]: - """Build per-item keys from content hashes and the common item record.""" - mm_input = param.multimodal_input - mm_data = param.multimodal_data or {} - if mm_input is None: - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params without multimodal hashes." - ) - return None - - item_metadata = param.item_metadata - if item_metadata is None: - logger.debug(f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params without item metadata.") - return None - if item_metadata.output_embedding_lengths is None: - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params without item output lengths." - ) - return None - item_modalities = {modality for modality, _ in item_metadata.item_refs} - if any( - modality not in item_modalities - for modality in _MM_DATA_INPUT_MODALITY_KEYS - if mm_data.get(modality) is not None - ): - return None - return cls.build_encoder_cache_item_keys( - mm_input.multimodal_hashes, - item_metadata.item_refs, - item_metadata.output_embedding_lengths, - mm_data.get("mm_processor_kwargs_hash"), - ) - - @staticmethod - def _encoder_cache_item_key( - modality: str, - item_hash: Sequence[int], - embedding_length: int, - kwargs_hash: str, - ) -> Hashable: - # Sole definition of the key format, so entries written by the - # full-request path and the item-scheduling path hit from either. - return (modality, tuple(item_hash), int(embedding_length), kwargs_hash) - - @classmethod - def build_encoder_cache_item_keys( - cls, - multimodal_hashes: Optional[Sequence[Sequence[int]]], - item_refs: Sequence[tuple[str, int]], - embedding_lengths: Sequence[int], - kwargs_hash: Optional[str], - ) -> Optional[list[Hashable]]: - """Build per-item cache keys from request-level item metadata. - - The modality comes from each item's `item_refs` entry, so mixed-modality requests are - keyable per item. Returns `None` when the request cannot participate in the cache (missing - hashes or kwargs hash, or item counts that do not line up). - """ - if multimodal_hashes is None or kwargs_hash is None: - return None - if not (len(multimodal_hashes) == len(item_refs) == len(embedding_lengths)): - # Malformed metadata rather than a normal miss: the request loses - # cache participation silently, so say so once with the counts. - logger.warning_once( - f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping item keys because " - f"multimodal_hashes ({len(multimodal_hashes)}), item_refs " - f"({len(item_refs)}) and multimodal_embedding_lengths " - f"({len(embedding_lengths)}) counts disagree", - key="mm_encoder_cache_item_key_count_mismatch", - ) - return None - return [ - cls._encoder_cache_item_key(modality, item_hash, embedding_length, kwargs_hash) - for (modality, _), item_hash, embedding_length in zip( - item_refs, - multimodal_hashes, - embedding_lengths, - strict=True, - ) - ] - - @classmethod - def partition_encoder_cache( - cls, - param: MultimodalParams, - encoder_cache: TensorLRUCache, - item_indices: Optional[Sequence[int]] = None, - keys: Optional[Sequence[Hashable]] = None, - ) -> Optional[EncoderCachePartition]: - """Look up `param`'s items in the cache and return the per-item partition. - - Args: - item_indices: restrict the lookup to these items. `None` looks up every item, - which is what the full-request path wants. The item scheduler passes the - subset it selected for this iteration so unscheduled items keep their LRU - recency -- see `EncoderCachePartition`. - keys: precomputed per-item keys. `None` derives them from `param`, which needs - `multimodal_input` to carry the content hashes. The executor builds params - from `LlmRequest.py_multimodal_data` alone -- the hashes live on the - request, not the params -- so it derives keys there and passes them in. - - Returns `None` when the param is not keyable (missing item metadata, content - hashes or processor-kwargs hash, or a request-local embedding is already - attached); the caller should treat that as a full miss. - """ - if param.multimodal_data.get("multimodal_embedding") is not None: - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: request-local multimodal embedding present; " - "skipping persistent cache lookup" - ) - return None - - keys = list(keys) if keys is not None else cls._encoder_cache_keys(param) - if not keys: - return None - - looked_up = list(range(len(keys))) if item_indices is None else list(item_indices) - out_of_range = [i for i in looked_up if not 0 <= i < len(keys)] - if out_of_range: - raise MultimodalEncoderContractError( - f"item_indices {out_of_range} out of range for {len(keys)} cache keys" - ) - - hits: Dict[int, torch.Tensor] = {} - miss_indices: list[int] = [] - for i in looked_up: - cached = encoder_cache.get(keys[i]) - if cached is None: - miss_indices.append(i) - else: - hits[i] = cached - - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: partition hit_items={len(hits)}, " - f"miss_items={len(miss_indices)}, looked_up={len(looked_up)}, " - f"total_items={len(keys)}." - ) - return EncoderCachePartition( - hits=hits, miss_indices=miss_indices, keys=keys, looked_up=looked_up - ) - - @staticmethod - def _apply_metadata_slice( - residual: MultimodalParams, - source: MultimodalParams, - prompt_item_indices: Sequence[int], - *, - item_metadata: Optional[MultimodalItemMetadata] = None, - ) -> None: - """Overwrite `residual`'s per-item metadata to match the sliced items. - - Models slice raw modality tensors in `build_multimodal_encoder_input`; the mixin owns - the parallel per-item metadata slice so every model gets it identically. - ``prompt_item_indices`` indexes the source metadata's ``item_refs`` - across all modalities, rather than a modality's payload. - """ - metadata = item_metadata if item_metadata is not None else source.item_metadata - if metadata is None: - raise MultimodalEncoderContractError("Item selection requires multimodal item metadata") - residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata.select( - prompt_item_indices - ) - if residual.multimodal_input is not None and source.multimodal_input is not None: - source_hashes = source.multimodal_input.multimodal_hashes - residual.multimodal_input.multimodal_hashes = [ - source_hashes[i] for i in prompt_item_indices - ] - - def _encode_with_partial_cache( - self, - partials: Sequence[tuple[MultimodalParams, EncoderCachePartition]], - encoder_cache: TensorLRUCache, - ) -> None: - """Encode only the miss items of each partial-hit param and stitch results. - - Miss residuals from all partial-hit params in the batch are encoded in a - single call and the concatenated output is split back per param, mirroring - how `get_multimodal_embeddings` batches full-miss params. After this returns, - each param's `multimodal_embedding` has the same shape as a full encoder run - so downstream `get_multimodal_embeddings` treats it as fully cached. - """ - if not partials: - return - - # Cross-iter prefetch may have staged some params' raw MM tensors on the aux - # stream. If a prefetch encoder call then raised, the request reaches this - # iteration with an `encoder_event` but no `multimodal_embedding`; slicing - # those tensors on the main stream before the event would race the aux-stream - # H2D copy. Wait per param up front, before any raw-tensor read. - for param, _ in partials: - if param.encoder_event is not None: - torch.cuda.current_stream().wait_event(param.encoder_event) - - encoder_inputs = self.prepare_multimodal_encoder_inputs( - [(param, i) for param, partition in partials for i in partition.miss_indices] - ) - # Keep one call across requests and modalities so EncoderGroup can - # batch modalities that share an encoder. - batched_output = self.encode_multimodal_inputs([param for param, _, _ in encoder_inputs]) - output_lengths = [length for _, lengths, _ in encoder_inputs for length in lengths] - encoded_items = iter(torch.split(batched_output, output_lengths, dim=0)) - - for param, partition in partials: - miss_tensors = [next(encoded_items) for _ in partition.miss_indices] - - by_item: Dict[int, torch.Tensor] = dict(partition.hits) - for miss_idx, tensor in zip(partition.miss_indices, miss_tensors, strict=True): - by_item[miss_idx] = tensor - param.multimodal_data["multimodal_embedding"] = _assemble_multimodal_encoder_embeddings( - by_item, len(partition.keys) - ) - - inserted = 0 - rejected = 0 - for miss_idx, tensor in zip(partition.miss_indices, miss_tensors, strict=True): - if encoder_cache.put(partition.keys[miss_idx], tensor): - inserted += 1 - else: - rejected += 1 - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: partial-hit encode " - f"total_items={len(partition.keys)} " - f"hit_items={len(partition.hits)} " - f"encoded_items={len(partition.miss_indices)} " - f"cache_writes_inserted={inserted} " - f"cache_writes_rejected={rejected}" - ) - - @classmethod - def _write_encoder_cache_entries( - cls, - param: MultimodalParams, - encoder_cache: TensorLRUCache, - ) -> None: - keys = cls._encoder_cache_keys(param) - if not keys: - return - - embedding = param.multimodal_data.get("multimodal_embedding") - if isinstance(embedding, list): - embedding = torch.cat(embedding, dim=0) - param.multimodal_data["multimodal_embedding"] = embedding - if not isinstance(embedding, torch.Tensor): - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping write because no tensor embedding was " - "attached after encoder execution." - ) - return - - _, embedding_lengths = param.get_item_layout() - if sum(embedding_lengths) != embedding.shape[0]: - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping write because embedding row count " - "does not match multimodal_embedding_lengths." - ) - return - - # Encoder outputs are concatenated per params object. Splitting by item length lets future - # requests reuse matching items independently, even when their request-level item order - # differs. - inserted_entries = 0 - rejected_entries = 0 - for key, item_embedding in zip( - keys, - torch.split(embedding, embedding_lengths, dim=0), - strict=True, - ): - if encoder_cache.put(key, item_embedding): - inserted_entries += 1 - else: - rejected_entries += 1 - logger.debug( - f"{_MM_ENCODER_CACHE_LOG_NAME}: wrote {inserted_entries} item entries, " - f"rejected={rejected_entries}, rows={embedding.shape[0]}." - ) - encoder_cache.log_stats("multimodal encoder cache write.") - - 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]]: - """Thin adapter over `fuse_input_embeds`. - - The framework does not forward `prepare_multimodal_inputs` kwargs - into `fuse_input_embeds`; only inputs the helper actually consumes - are surfaced here. Models needing to bypass token filtering should - pass pre-computed `text_token_indices`/`mm_token_indices`. - """ - if mm_token_ids is not None and not isinstance(mm_token_ids, torch.Tensor): - mm_token_ids = torch.tensor( - list(mm_token_ids), dtype=input_ids.dtype, device=input_ids.device - ) - - result = fuse_input_embeds( - embedding_layer=embedding_layer, - input_ids=input_ids, - mm_embeds=multimodal_embeddings, - mm_token_ids=mm_token_ids, - text_token_indices=text_token_indices, - mm_token_indices=mm_token_indices, - extra_embeds=list(extra_embeds) if extra_embeds else None, - ) - if len(result) == 3: - fused_input_ids, inputs_embeds, fused_extra_embeds = result - return fused_input_ids, inputs_embeds, fused_extra_embeds or () - - fused_input_ids, inputs_embeds = result - return fused_input_ids, inputs_embeds, () - - @staticmethod - def _validate_embeddings( - embeddings: list[torch.Tensor], - multimodal_params: Sequence[MultimodalParams], - *, - current_chunk_only: bool = False, - ) -> None: - """Validate gathered embeddings' row count against runtime metadata. - - `current_chunk_only` matches the gather layout; it does not change the - original per-request runtime counts used for later prefill chunks. - Skipped if any param lacks `multimodal_runtime.total_embeds_in_request`, since the contract - cannot be evaluated without complete metadata. - """ - if len(embeddings) != 1: - raise ValueError( - f"MultimodalModelMixin requires a single embedding tensor, got {len(embeddings)} " - "tensors." - ) - - embeddings_tensor = embeddings[0] - expected_rows = 0 - has_runtime_metadata = [] - for param in multimodal_params: - runtime = param.multimodal_runtime - has_runtime = runtime is not None and runtime.total_embeds_in_request is not None - has_runtime_metadata.append(has_runtime) - if has_runtime: - expected_rows += ( - runtime.num_mm_tokens_in_chunk - if current_chunk_only - else runtime.total_embeds_in_request - ) - - if any(has_runtime_metadata) and not all(has_runtime_metadata): - raise ValueError( - "Multimodal runtime metadata must be present for every param or none of them." - ) - if not all(has_runtime_metadata): - logger.debug( - "Skipping multimodal embedding row-count validation: runtime metadata missing " - "for all params." - ) - return - - actual_rows = embeddings_tensor.shape[0] - if actual_rows != expected_rows: - raise ValueError( - f"Multimodal embedding row count mismatch: expected {expected_rows}, got {actual_rows}." - ) - - -def _collect_cross_iter_prefetch_candidates( - pending_requests: Sequence["LlmRequest"], - in_flight_request_ids: Iterable[int], - max_prefetch: int, - max_prefetch_ahead: int, -) -> list[tuple["LlmRequest", Dict[str, Any], torch.Tensor]]: - """Select cross-iteration prefetch candidates without touching CUDA. - - Returns up to `max_prefetch` `(request, multimodal_data, cumsum)` tuples - while keeping total not-in-flight prefetched requests at or below - `max_prefetch_ahead`. A new candidate must have a supported raw input - modality in `multimodal_data`, no cached `multimodal_embedding`, and a valid - `multimodal_embed_mask_cumsum`. Existing cached embeddings or pending encoder - events count against the ahead limit only when attached to a request with - real MM encoder work. - """ - if max_prefetch <= 0 or max_prefetch_ahead <= 0: - return [] - - in_flight = set(in_flight_request_ids) - outstanding_prefetches = 0 - candidates: list[tuple["LlmRequest", Dict[str, Any], torch.Tensor]] = [] - - for req in pending_requests: - if req.py_request_id in in_flight: - continue - mm_data = req.py_multimodal_data or {} - has_cached_embedding = mm_data.get("multimodal_embedding") is not None - has_raw_mm_input = any(key in mm_data for key in _MM_DATA_INPUT_MODALITY_KEYS) - if req.py_mm_encoder_event is not None: - if has_cached_embedding or has_raw_mm_input: - outstanding_prefetches += 1 - continue - if has_cached_embedding: - outstanding_prefetches += 1 - continue - if not has_raw_mm_input: - continue - cumsum = mm_data.get("multimodal_embed_mask_cumsum") - if cumsum is None: - continue - candidates.append((req, mm_data, cumsum)) - - available_slots = min(max_prefetch, max_prefetch_ahead - outstanding_prefetches) - if available_slots <= 0: - return [] - return candidates[:available_slots] - - -def _dispatch_cross_iter_prefetch( - model: "MultimodalModelMixin", - candidates: Sequence[tuple["LlmRequest", Dict[str, Any], torch.Tensor]], - aux_stream: "torch.cuda.Stream", -) -> None: - """H2D-copy MM data, run the encoder, and cache embeddings on `aux_stream`. - - Stamps a CUDA event on every candidate's `py_mm_encoder_event` so the - next iteration's consume site waits on it before reading cached tensors. - The event covers all work queued in the aux-stream block, so the same - event object is shared across all candidates. - """ - encoder_cache_enabled = model.encoder_cache_active - params_list = [ - MultimodalParams( - multimodal_input=_build_request_multimodal_input(req, encoder_cache_enabled), - multimodal_data=mm_data, - multimodal_runtime=MultimodalRuntimeData( - past_seen_token_num=0, - chunk_end_pos=cumsum.numel(), - embed_mask_cumsum=cumsum, - ), - ) - for req, mm_data, cumsum in candidates - ] - - # Prefetch targets requests outside the current iteration, so their - # multimodal tensors are not touched by the main stream. The caller queues - # this after the iteration's LLM kernels so aux-stream H2D copies and - # encoder work can overlap them. - # - # Request-local ordering is handled by `encoder_event`; the consume path also `record_stream`s - # attached embeddings before gathering them so post-prefill request cleanup cannot release - # storage while main-stream work is pending. Persistent-cache clones use their own producer - # events and consumer `record_stream` calls inside `TensorLRUCache`. - encoder_event = None - try: - with _run_on_aux_stream(aux_stream) as encoder_event: - encoder_cache = model._multimodal_encoder_cache if encoder_cache_enabled else None - cache_misses: list[MultimodalParams] = [] - partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = [] - if encoder_cache is None: - cache_misses = params_list - else: - for param in params_list: - partition = model.partition_encoder_cache(param, encoder_cache) - if partition is None or partition.is_full_miss: - cache_misses.append(param) - elif partition.is_full_hit: - param.multimodal_data["multimodal_embedding"] = ( - _assemble_multimodal_encoder_embeddings( - partition.hits, len(partition.keys) - ) - ) - else: - partial_hits.append((param, partition)) - - params_to_transfer = cache_misses + [param for param, _ in partial_hits] - for param in params_to_transfer: - param.to_device( - "multimodal_data", - "cuda", - pin_memory=prefer_pinned(), - target_keywords=getattr(model, "multimodal_data_device_paths", None), - ) - # `to_device` may replace `multimodal_data` with a new dict; reattach - # the (possibly new) dict to each request so the next iteration's - # `_prepare_inputs` sees the cached embedding stamped below. Mirrors - # the reassignment at the canonical to_device call site in - # DecoderInputPreparer.prepare_inputs. - for (req, _, _), p in zip(candidates, params_list): - req.py_multimodal_data = p.multimodal_data - - # Prefetch runs between executor steps, outside `ModelCaller`, so the vision tower's - # TP all-reduces can return NCCL window tensors here. The stored embeddings are read - # next iteration, so keep a clone rather than a view of a released window. - with discard_nccl_window_tensor_outputs(params_to_transfer): - if partial_hits and encoder_cache is not None: - model._encode_with_partial_cache(partial_hits, encoder_cache) - if cache_misses: - encoder_output = model.encode_multimodal_inputs(cache_misses).clone() - _store_chunked_prefill_embeddings(cache_misses, [encoder_output]) - if encoder_cache is not None: - for param in cache_misses: - model._write_encoder_cache_entries(param, encoder_cache) - - # Prefetch only needs to attach each request's embedding. Validate each request - # independently instead of gathering the unused batch output with `torch.cat`. - for param in params_list: - embedding = param.multimodal_data.get("multimodal_embedding") - if not isinstance(embedding, torch.Tensor): - raise ValueError("Multimodal encoder prefetch did not produce an embedding.") - model._validate_embeddings([embedding], [param]) - finally: - # Stash the event on every candidate's durable LlmRequest (not the - # per-iter `MultimodalParams`), since `_prepare_inputs` rebuilds the - # wrapper each iteration. The transfer to `MultimodalParams.encoder_event` - # happens in `_prepare_inputs` when the request is next scheduled. - # - # This runs in `finally` (and `_run_on_aux_stream` records the event in - # its own `finally`) so that on partial failure -- e.g. `to_device` - # mutated `req.py_multimodal_data` in place to aux-stream CUDA tensors, - # then `encode_multimodal_inputs` raised -- the consumer still has an - # event to wait on before reading those tensors on the main stream. - # Without this, the request would carry aux-stream tensors with no sync - # point, producing a cross-stream data race in the next iteration's - # in-iter encode path. - if encoder_event is not None: - for req, _, _ in candidates: - req.py_mm_encoder_event = encoder_event - - -def maybe_prefetch_mm_encoder_for_next_iter( - model: Any, - pending_requests: Sequence["LlmRequest"], - in_flight_request_ids: Iterable[int] = (), - max_prefetch: int = 1, - max_prefetch_ahead: Optional[int] = None, -) -> int: - """Speculative cross-iteration MM encoder prefetch on a side CUDA stream. - - For up to `max_prefetch` `pending_requests`, subject to the outstanding ahead cap, runs - `model.encode_multimodal_inputs` on a side CUDA stream. - - The resulting embeddings are written into `request.py_multimodal_data` so the next iteration's - `_prepare_inputs` picks them up via the standard cache path, and a CUDA event is stamped on - `request.py_mm_encoder_event` for `_prepare_inputs` to transfer onto the new `MultimodalParams`. - - The mixin consume sites (e.g. `get_multimodal_embeddings)` need to wait on the event before - reading the cached tensors. - - While the current iteration's LLM kernels run on the main stream, this queues encoder work for - an "admit-likely context" request on the aux stream. - - Mis-predictions waste GPU time, but cached embeddings remain valid until the request is admitted - or terminated. Examples include: - - - The prefetched request can be terminated before admission (client cancel / disconnect, - timeout, validation failure, ...). - - With `max_prefetch < len(pending)`, if the head is bumped by budget reasons, the next-admitted - request is one we did not prefetch. - - Gated by `MultimodalConfig.encoder_side_stream_max_ahead`: 0 disables the side stream; a - positive integer enables it and caps the total number of not-in-flight requests with - prefetched MM encoder work. - - Returns the number of requests for which an encoder kick-off was queued. - """ - if not isinstance(model, MultimodalModelMixin): - return 0 - if max_prefetch <= 0: - return 0 - if max_prefetch_ahead is None: - max_prefetch_ahead = 0 - if max_prefetch_ahead <= 0: - return 0 - - aux_stream = _get_mm_aux_stream(max_prefetch_ahead) - if aux_stream is None: - return 0 - - candidates = _collect_cross_iter_prefetch_candidates( - pending_requests, in_flight_request_ids, max_prefetch, max_prefetch_ahead - ) - if not candidates: - return 0 - _dispatch_cross_iter_prefetch(model, candidates, aux_stream) - return len(candidates) +__all__ = [ + "MultimodalEncoderContractError", + "_assemble_multimodal_encoder_embeddings", + "EncoderGroup", + "_lengths_by_modality", + "_reorder_embeds_by_manifest", + "encode_multimodal_by_groups", + "_build_request_multimodal_input", + "_get_mm_aux_stream", + "_run_on_aux_stream", + "PreparedLlmInputs", + "PreparedMultimodalEmbeddings", + "EncoderCachePartition", + "MultimodalModelContractError", + "MultimodalModelMixin", + "MultimodalRuntimeMixin", + "_collect_cross_iter_prefetch_candidates", + "_dispatch_cross_iter_prefetch", + "maybe_prefetch_mm_encoder_for_next_iter", +] diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py index fcd32b77aa4a..c58940da61ae 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py @@ -19,8 +19,7 @@ import functools import math import os -from typing import (Any, Callable, Dict, List, NamedTuple, Optional, Tuple, - TypedDict, Union, cast) +from typing import Any, NamedTuple, TypedDict import numpy as np import torch @@ -32,9 +31,62 @@ from torchvision.transforms import Normalize, Resize, ToTensor from transformers.image_processing_backends import TorchvisionBackend -from tensorrt_llm._torch.modules.embedding import Embedding from tensorrt_llm.inputs.multimodal import MultimodalParams -from tensorrt_llm.logger import logger + +from ..multimodal.embeddings import _as_copy_index +from ..multimodal.embeddings import \ + _get_uncached_multimodal_params as _get_uncached_multimodal_params +from ..multimodal.embeddings import _join_embeddings +from ..multimodal.embeddings import \ + _normalize_encoder_embeddings as _normalize_encoder_embeddings +from ..multimodal.embeddings import \ + _store_chunked_prefill_embeddings as _store_chunked_prefill_embeddings +from ..multimodal.embeddings import \ + filter_mm_token_from_input_ids as filter_mm_token_from_input_ids +from ..multimodal.embeddings import find_input_mm_embeds, fuse_input_embeds +from ..multimodal.embeddings import \ + get_attached_multimodal_embeddings as get_attached_multimodal_embeddings +from ..multimodal.embeddings import \ + get_multimodal_embeddings as get_multimodal_embeddings + +__all__ = [ + "VisionNormalization", + "vision_normalization_from_processor", + "_vision_pixel_lookup_table", + "_lookup_vision_pixels_kernel", + "normalize_vision_pixels", + "_MULTIMODAL_ENV_NAME", + "_is_mm_disagg", + "has_raw_multimodal_payload", + "_PROCESSOR_OUTPUT_KEYS", + "_get_cached_merged_typed_dict", + "_install_processor_output_validation_filter", + "preprocess_dispatch", + "find_closest_aspect_ratio", + "dynamic_preprocess", + "dynamic_preprocess_torch", + "dynamic_preprocess_dispatch", + "dynamic_s2_preprocess", + "dynamic_s2_preprocess_torch", + "dynamic_s2_preprocess_dispatch", + "s2_split_chessboard", + "s2_merge_chessboard", + "s2_batched_forward", + "multiscale_forward", + "merge_chessboard", + "split_chessboard", + "merge_features_for_dynamic_s2", + "_get_uncached_multimodal_params", + "_store_chunked_prefill_embeddings", + "_normalize_encoder_embeddings", + "_join_embeddings", + "get_multimodal_embeddings", + "get_attached_multimodal_embeddings", + "find_input_mm_embeds", + "filter_mm_token_from_input_ids", + "_as_copy_index", + "fuse_input_embeds", +] class VisionNormalization(NamedTuple): @@ -281,519 +333,6 @@ def _filtered_validate(schema, data): b.validate_typed_dict = _filtered_validate -def _get_uncached_multimodal_params( - multimodal_params: List[MultimodalParams], ) -> List[MultimodalParams]: - """ - Get uncached multimodal params that need encoder processing for chunk prefill. - """ - params_to_run = [] - - for param in multimodal_params: - # Skip if no multimodal content - if not param.has_content(): - continue - - # Check if embeddings are already cached - if (param.multimodal_data - and "multimodal_embedding" in param.multimodal_data - and param.multimodal_data["multimodal_embedding"] is not None): - logger.debug( - "Skipping encoder forward for param with cached multimodal_embedding" - ) - continue - - # This param needs encoder processing - params_to_run.append(param) - - return params_to_run - - -def _store_chunked_prefill_embeddings( - multimodal_params: List[MultimodalParams], - embeddings: List[torch.Tensor], -) -> None: - """ - Store computed multimodal embeddings back to multimodal_data to avoid recomputation. - - NOTE: this function only stores multimodal embeddings within the current request context, - mostly for chunked prefill. It does not persist embeddings across different requests or sessions. - """ - # TODO: support multiple multimodal modalities per request - if len(embeddings) > 1: - raise ValueError("Multiple modalities caching is not supported yet.") - mm_embed = embeddings[0] - - # Collect embedding lengths for each parameter - embed_lengths = [] - for param in multimodal_params: - if param.multimodal_runtime is not None: - embed_lengths.append( - param.multimodal_runtime.total_embeds_in_request) - - # Validate total length matches - total_expected = sum(embed_lengths) - assert len(mm_embed) == total_expected, \ - f"Number of mm_embeds ({len(mm_embed)}) does not match expected total ({total_expected})" - - # Use torch.split for efficient tensor splitting - split_embeddings = torch.split(mm_embed, embed_lengths, dim=0) - valid_params = [ - p for p in multimodal_params if p.multimodal_runtime is not None - ] - - # Cache split embeddings to each parameter - logger.debug( - f"Caching {len(split_embeddings)} multimodal embedding chunks in {len(multimodal_params)} params" - ) - for param, embed_chunk in zip(valid_params, split_embeddings, strict=True): - param.multimodal_data["multimodal_embedding"] = embed_chunk - - logger.debug( - f"Cached {len(split_embeddings)} multimodal embedding chunks in this iteration" - ) - - -def _normalize_encoder_embeddings( - encoder_embeddings: Union[torch.Tensor, List[torch.Tensor]], -) -> List[torch.Tensor]: - if isinstance(encoder_embeddings, torch.Tensor): - return [encoder_embeddings] - - if (not isinstance(encoder_embeddings, list) or not all( - isinstance(embedding, torch.Tensor) - for embedding in encoder_embeddings)): - raise TypeError( - "encoder_forward_fn must return a torch.Tensor or a list of torch.Tensor." - ) - - return encoder_embeddings - - -def _join_embeddings(embeddings: List[torch.Tensor], - *, - start: int = 0, - end: Optional[int] = None) -> torch.Tensor: - """Join per-item/per-request embeddings along dim 0 without a needless copy. - - ``torch.cat`` always allocates, so calling it on a one-element list - duplicates the whole embedding. That matters here because the same rows - are joined several times on the way to the LLM input (per request, per - batch, and again in ``fuse_input_embeds``), and each redundant copy is a - full extra allocation live at the prefill peak. Multimodal embeddings are - read-only downstream -- they are sliced and scattered into the input - embedding buffer, never written in place -- so returning the sole tensor - is equivalent to concatenating it. If a row range is given, slice each - segment first so rows outside the current prefill chunk are never copied. - """ - if start != 0 or end is not None: - slices = [] - offset = 0 - for embedding in embeddings: - segment_end = offset + embedding.shape[0] - slice_end = segment_end if end is None else min(end, segment_end) - if max(start, offset) < slice_end: - slices.append(embedding[max(start - offset, 0):slice_end - - offset]) - offset = segment_end - if end is not None and offset >= end: - break - if not slices: - return embeddings[0][:0] - embeddings = slices - - return embeddings[0] if len(embeddings) == 1 else torch.cat(embeddings, - dim=0) - - -def get_multimodal_embeddings( - encoder_forward_fn: Callable[ - ..., - torch.Tensor | List[torch.Tensor], - ], - multimodal_params: List[MultimodalParams], - encoder_kwargs: Optional[Dict[str, Any]] = None, -) -> List[torch.Tensor]: - """ - High-level utility to get multimodal embeddings from encoder or cached embeddings. - - This function will: - 1. Identify which parameters need encoder processing - 2. Run encoder forward only on uncached parameters - 3. Cache newly computed embeddings (if enabled) - 4. Gather all embeddings for the batch - - Args: - encoder_forward_fn: Callable that performs encoder forward pass. - Should accept List[MultimodalParams] and return either - a single torch.Tensor or List[torch.Tensor]. - multimodal_params: All multimodal parameters in the batch. - encoder_kwargs: Optional kwargs to pass to encoder_forward_fn. - Returns: - List of multimodal embeddings for all multimodal params in the batch. - """ - if not multimodal_params: - return [] - - # Wait before touching tensors produced on the MM side stream. Do not clear the event here; - # repeated stream-side waits are cheap, and leaving the event field untouched avoids races if a - # caller accidentally reuses it. - # Register attached embeddings with the consumer stream as well: the request drops its Python - # references after prefill, potentially before the asynchronous gather below has finished - # reading them. - for param in multimodal_params: - if param.encoder_event is not None: - consumer_stream = torch.cuda.current_stream() - consumer_stream.wait_event(param.encoder_event) - embeds = param.multimodal_data.get("multimodal_embedding") - if isinstance(embeds, torch.Tensor): - embeds = [embeds] - if isinstance(embeds, list): - for embed in embeds: - if isinstance(embed, torch.Tensor) and embed.is_cuda: - embed.record_stream(consumer_stream) - - # Step 1: Find uncached multimodal params that need encoder processing - uncached_multimodal_params = _get_uncached_multimodal_params( - multimodal_params) - - # Step 2: Run encoder forward only on uncached parameters - if uncached_multimodal_params: - kwargs = encoder_kwargs or {} - encoder_embeddings = encoder_forward_fn(uncached_multimodal_params, - **kwargs) - encoder_embeddings = _normalize_encoder_embeddings(encoder_embeddings) - - # TODO: support multiple multimodal modalities per request - if len(encoder_embeddings) > 1: - logger.warning( - f"Multiple modalities caching is not supported yet. " - f"encoder returned {len(encoder_embeddings)} embeddings " - f"(types: {[type(e).__name__ for e in encoder_embeddings]}, " - f"shapes: {[e.shape if hasattr(e, 'shape') else 'N/A' for e in encoder_embeddings]}) " - f"for {len(uncached_multimodal_params)} uncached params. " - f"encoder_forward_fn={encoder_forward_fn}") - return encoder_embeddings - - # Validate that multimodal_runtime has required attributes for caching - if (not hasattr(uncached_multimodal_params[0], 'multimodal_runtime') - or uncached_multimodal_params[0].multimodal_runtime is None - or uncached_multimodal_params[0].multimodal_runtime. - total_embeds_in_request is None): - logger.warning( - "Multimodal runtime data missing or incomplete, will not cache embeddings." - ) - return encoder_embeddings - - # Step 3: Cache the computed embeddings to multimodal_data["multimodal_embedding"]. - # A TP-sharded encoder can return an NCCL window tensor whose lease ends - # with the enclosing ModelCaller scope, while chunked prefill reads the - # cached rows in later iterations, so cache an owned copy instead of - # views of the encoder output. - encoder_embeddings = [encoder_embeddings[0].clone()] - _store_chunked_prefill_embeddings(uncached_multimodal_params, - encoder_embeddings) - - # Step 4: Gather after the stream wait and any encoder/cache writes. - return get_attached_multimodal_embeddings(multimodal_params) - - -def get_attached_multimodal_embeddings( - multimodal_params: List[MultimodalParams]) -> List[torch.Tensor]: - """Gather embeddings already stored on MultimodalParams. - - Use this on E/P prefill workers and cached-only paths. The encoder already ran somewhere else. - Normally `find_input_mm_embeds` selects active rows afterwards. If item - scheduling attached current-chunk embeddings, select the same chunk from - any full-request embeddings before joining the batch. This keeps mixed - batches in one row layout without copying unrelated future chunks. - - Side-stream-prefetched requests must use `get_multimodal_embeddings`, which waits on - `encoder_event` and registers attached tensors with the consuming stream before gathering them. - """ - current_chunk_only = any( - param.multimodal_data.get("multimodal_embedding_is_chunk", False) - for param in multimodal_params) - attached_embeddings = [] - for param in multimodal_params: - embeds = param.multimodal_data.get("multimodal_embedding") - # No attached embedding for this request. - if embeds is None: - continue - runtime = param.multimodal_runtime - if (current_chunk_only and runtime is not None - and not param.multimodal_data.get( - "multimodal_embedding_is_chunk", False)): - # Keep the full attached payload intact for later chunks/cache - # writes; only this forward's gathered tensor is chunk-local. - start = runtime.num_cached_mm_tokens - end = start + runtime.num_mm_tokens_in_chunk - if isinstance(embeds, list): - embeds = _join_embeddings(embeds, start=start, end=end) - else: - embeds = embeds[start:end] - elif isinstance(embeds, list): - # Legacy paths can attach a list of full-request output pieces. - embeds = _join_embeddings(embeds) - param.multimodal_data["multimodal_embedding"] = embeds - if not isinstance(embeds, torch.Tensor): - raise TypeError("multimodal_embedding must be a torch.Tensor") - attached_embeddings.append(embeds) - - if not attached_embeddings: - return [] - # Match get_multimodal_embeddings output: one concatenated tensor. - return [_join_embeddings(attached_embeddings)] - - -def find_input_mm_embeds( - mm_embeds: List[torch.Tensor], - multimodal_params: List[MultimodalParams]) -> List[torch.Tensor]: - """ - Find the multimodal mm_embeds that need processing from multimodal_params for each batch. - Supports both KV cache reuse and chunked prefill scenarios. - - Args: - - mm_embeds: List[torch.Tensor] - Multimodal embeddings for each batch - - multimodal_params: List[MultimodalParams] - Multimodal parameters with runtime data - - Returns: - - List[torch.Tensor] - Sliced mm_embeds containing only tokens that need processing: - - For KV cache reuse: tokens that are not cached - - For chunked prefill: tokens that are in the current chunk - - For mixed scenarios: both uncached and current chunk tokens - - Empty list if all tokens are cached or beyond current chunk - - Note: - - Supports both individual batching (len(mm_embeds) == len(multimodal_params)) - and pre-concatenated batching (len(mm_embeds) == 1) - - Call get_attached_multimodal_embeddings before this helper when - embeddings are already attached to multimodal_params. - - Handles chunked prefill by considering chunk boundaries and current chunk tokens - - Example: if a request has 8 MM embed rows, 2 cached rows, and 3 rows - in the current chunk, this keeps rows [2:5]. - - This function reads only CPU-resident counts from - ``MultimodalParams.multimodal_runtime`` and never touches GPU - tensors, so it does not introduce host-device synchronization. - The companion ``torch.where`` host sync sits in ``fuse_input_embeds`` - (via ``filter_mm_token_from_input_ids``) and is avoided by routing - executor-precomputed ``mm_token_indices`` / ``text_token_indices`` - through to that call — see the contract there. - """ - if not isinstance(mm_embeds, list): - raise TypeError("mm_embeds must be a list") - - # Current support two batching modes: - # 1. Pre-concatenated mm_embeds for each batch, i.e., len(mm_embeds) == 1 - # 2. Individual mm_embeds for each multimodal param, i.e., len(mm_embeds) == len(multimodal_params) - if len(mm_embeds) > 1 and len(mm_embeds) != len(multimodal_params): - raise ValueError( - f"Number of mm_embeds ({len(mm_embeds)}) does not match number of multimodal params ({len(multimodal_params)})." - ) - - if not multimodal_params or multimodal_params[0].multimodal_runtime is None: - # No slicing, return the full mm_embeds - return mm_embeds - - total_mm_tokens = sum(param.multimodal_runtime.num_mm_tokens_in_chunk - for param in multimodal_params - if param.multimodal_runtime is not None) - - if total_mm_tokens == 0: - logger.debug( - "All multimodal tokens are cached or beyond current chunk, skipping vision encoder forward" - ) - return [] - - if not mm_embeds: - raise ValueError( - "No multimodal embeddings were provided or cached for active multimodal tokens." - ) - - if total_mm_tokens == sum(mm_embed.shape[0] for mm_embed in mm_embeds): - # Item-scheduled batches were already sliced by the producer/gather - # path. Do not apply the cached-prefix offset a second time. - return mm_embeds - - current_pos = 0 - slices = [] - for param in multimodal_params: - runtime = param.multimodal_runtime - if runtime is None: - continue - local_start_pos = runtime.num_cached_mm_tokens - local_end_pos = local_start_pos + runtime.num_mm_tokens_in_chunk - slices.append( - (current_pos + local_start_pos, current_pos + local_end_pos)) - if len(mm_embeds) == 1: # pre-concatenated; advance global cursor - current_pos += runtime.total_embeds_in_request - - if len(mm_embeds) == 1: - sliced = [mm_embeds[0][start:end] for start, end in slices] - return [_join_embeddings(sliced)] - return [mm_embeds[i][start:end] for i, (start, end) in enumerate(slices)] - - -def filter_mm_token_from_input_ids( - input_ids: torch.IntTensor, - vocab_size: int, - mm_token_ids: Optional[torch.IntTensor] = None, -) -> Tuple[torch.IntTensor, torch.IntTensor]: - """ - Filter multimodal tokens from input_ids. - Args: - input_ids: shape [text_total_length + mm_total_length]. - vocab_size: size of the model's vocabulary - mm_token_ids: possible token ids for multimodal tokens, if known. If - not known and set to None, it is assumed that the multimodal tokens - are out-of-vocabulary tokens i.e. the ``input_ids`` contains tokens - >= vocab_size that represent the multimodal tokens. - Note: - Example: input_ids=[1, 55, 2, 101], vocab_size=100, and - mm_token_ids=[55] returns mm_token_indices=[1]; token 101 is text - because explicit mm_token_ids overrides the OOV fallback. - This function involves host-device synchronization due to torch.where() (= torch.nonzero) requiring - host allocation. The output indices reside on the same device as input_ids. - Returns: - text_token_indices: indices of text tokens in the input_ids - mm_token_indices: indices of multimodal tokens in the input_ids - """ - if mm_token_ids is None: - # If mm_token_ids is None, assume the multimodal tokens are out-of-vocab - # (input_ids >= vocab_size). Avoids torch.isin() over a potentially - # unbounded mm_token_ids set. - mm_token_mask = input_ids >= vocab_size - else: - mm_token_ids = mm_token_ids.to(input_ids.device, dtype=input_ids.dtype) - mm_token_mask = torch.isin(input_ids, mm_token_ids) - # NOTE: torch.where() enforces a host sync - text_token_indices = torch.where(~mm_token_mask)[0] - mm_token_indices = torch.where(mm_token_mask)[0] - return text_token_indices, mm_token_indices - - -def _as_copy_index(index: torch.Tensor, - destination: torch.Tensor) -> torch.Tensor: - """Return ``index`` as the int64 index ``index_copy_`` needs on ``destination``. - - Explicit indices can come from callers with another dtype or device; for the - executor's int64 device indices this is a no-op. - """ - return index.to(device=destination.device, - dtype=torch.long, - non_blocking=True) - - -def fuse_input_embeds( - embedding_layer: Embedding, - input_ids: torch.IntTensor, - mm_embeds: List[torch.Tensor], - mm_token_ids: Optional[torch.IntTensor] = None, - text_token_indices: Optional[torch.IntTensor] = None, - mm_token_indices: Optional[torch.IntTensor] = None, - extra_embeds: Optional[List[torch.Tensor]] = None, - **kwargs, - # TODO: make unified return type for all models -) -> Union[Tuple[Optional[torch.IntTensor], Optional[torch.FloatTensor]], - Tuple[Optional[torch.IntTensor], Optional[torch.FloatTensor], - Optional[List[torch.FloatTensor]]]]: - """ - Fuse text and multimodal embeddings. input_ids is [text_total_length + mm_total_length] and mm_embed is [mm_total_length, hidden_dim]. We just need to fuse them into [text_total_length + mm_total_length, hidden_dim] by slice-and-assign to the corresponding entries. - - Args: - embedding_layer: embedding layer of the model. - input_ids: shape [text_total_length + mm_total_length], flattened from List[(text_length1 + mm_total_length1), ..., (text_lengthi + mm_total_lengthi)]. For LLM model, the requests are inflight batched together, but the input_ids are flattened with padding removed. By the slice condition < vocab_size, we can easily separate text / multimodal tokens and naturally batched the LLM embedding lookup - mm_embeds: List[(mm_total_length1, hidden_dim), ..., (mm_total_lengthi, hidden_dim)]. - mm_token_ids: possible token ids for multimodal tokens, if known. If not known and set to None, it is assumed that the multimodal tokens are out-of-vocabulary tokens. - extra_embeds: Optional list of extra embed tensors for models that support it (e.g., Qwen3-VL/Qwen3-MoE-VL). - Returns: - - If (1) JIT test run, (2) non-multimodal run, i.e. all text-only requests, either context or generation phase (3) multimodal run, all requests in generation phase --> there is no multimodal data, return only the input_ids - - If (4) multimodal run, mixed batch of context and generation requests, each context request has a multimodal feature --> return only the fused input_embeds of shape [total length, hidden_dim]. For text tokens, LLM embedding layer has already run. - Note: - - Precedence: If kwargs provide indices (text_token_indices and mm_token_indices), those are used. If any one of them is not provided, fallback to filtering method. Sentinel-/OOV-based filtering (e.g., tokens >= vocab_size) is used only when neither index tensor and mm_token_ids is provided. - - Example: len(torch.cat(mm_embeds)) must match len(mm_token_indices); - for chunked prefill, pass only the current chunk's mm_embeds or - explicit indices for the active MM token positions. - - Sync-free contract: passing both ``text_token_indices`` and - ``mm_token_indices`` skips the GPU ``torch.where`` host sync. The - executor (``DecoderInputPreparer.prepare_inputs``) precomputes them on a CPU - ``input_ids`` copy via ``runners.prepare_multimodal_indices`` (which uses - ``filter_mm_token_from_input_ids`` against ``self.model.mm_token_ids`` - when present, else the OOV fallback ``>= vocab_size``) and ships - them as pinned async H2D tensors in the inputs dict. VLM forwards - are expected to forward these explicitly rather than relying on a - ``**kwargs`` splat — if they don't, this function falls back to the - host-syncing ``filter_mm_token_from_input_ids`` on GPU ``input_ids``. - - In-vocab fast path: when ``mm_token_ids`` is provided, the caller is - declaring that mm placeholder IDs are real vocabulary entries. In - that case the text path skips its ``index_select`` + ``torch.empty`` - + text scatter and embeds the full ``input_ids`` once, overwriting - mm rows afterwards. This is the path taken by every VLM currently - in tree EXCEPT Hyperclovax, which keeps the legacy ``vocab_size + 1`` - OOV remap in its input processor and therefore passes - ``mm_token_ids=None`` so the OOV branch is taken — embedding an OOV - id would be out-of-bounds. - """ - if len(mm_embeds) == 0: - if extra_embeds is not None and len(extra_embeds) > 0: - return input_ids, None, extra_embeds - return input_ids, None - - mm_embed = _join_embeddings(mm_embeds) - - # TODO: support the case where only one index tensor is provided, the other is derived as the complement (try to avoid implicit host-device synchronization) - if text_token_indices is None or mm_token_indices is None: - # NOTE: This function involves host-device synchronization due to torch.where() used in filter_mm_token_from_input_ids. - 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) - if mm_token_indices.shape[0] != mm_embed.shape[0]: - raise ValueError( - f"Multimodal token count mismatch: found {len(mm_token_indices)} image tokens in input_ids " - f"but received {mm_embed.shape[0]} image embeddings.") - - if mm_token_ids is not None: - # In-vocab fast path: caller declared mm tokens are real vocabulary - # IDs, so a single full embedding lookup is safe. Saves the - # ``index_select`` over text positions, the ``torch.empty`` alloc, and - # the text scatter (~2 GPU kernels per VLM forward) at the cost of - # ~mm-fraction extra embedding lookups whose result is overwritten. - input_embeds = embedding_layer(input_ids) - else: - # OOV path: input_ids holds placeholder IDs >= num_embeddings at mm - # positions (e.g. ``vocab_size + 1``), so embedding the full sequence - # would be out-of-bounds. Gather only the in-vocab text positions and - # scatter into a fresh buffer. - text_embed = embedding_layer(input_ids[text_token_indices]) - input_embeds = torch.empty(input_ids.shape[0], - mm_embed.shape[-1], - device=text_embed.device, - dtype=text_embed.dtype) - input_embeds.index_copy_( - 0, _as_copy_index(text_token_indices, input_embeds), text_embed) - if extra_embeds is not None and len(extra_embeds) > 0: - # only support single modality for deepstack features for now - for i, extra_feature in enumerate(extra_embeds): - extra_embed = torch.zeros( - input_ids.shape[0], - mm_embed.shape[-1], - device=extra_feature.device, - dtype=extra_feature.dtype, - ) - extra_embed.index_copy_( - 0, _as_copy_index(mm_token_indices, extra_embed), extra_feature) - extra_embeds[i] = extra_embed - - input_embeds.index_copy_( - 0, - _as_copy_index(mm_token_indices, input_embeds), - mm_embed.to(dtype=input_embeds.dtype, device=input_embeds.device), - ) - if extra_embeds is not None and len(extra_embeds) > 0: - return None, cast(torch.FloatTensor, input_embeds), extra_embeds - return None, cast(torch.FloatTensor, input_embeds) - - #region VILA utils # ------------------------------------------------------------------------------------------ # VILA image preprocessing utils diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py b/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py index a15b80e8ddef..4dc2757fe8c7 100644 --- a/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py +++ b/tensorrt_llm/_torch/models/modeling_musespark_open_omni.py @@ -69,8 +69,8 @@ from ...sampling_params import SamplingParams from ..model_config import ModelConfig from ..modules.embedding import Embedding +from ..multimodal.model import MultimodalModelMixin from .checkpoints.hf.musespark_open_weight_mapper import load_musespark_open_omni_weights -from .modeling_multimodal_mixin import MultimodalModelMixin from .modeling_musespark_open import MuseSparkOpenForCausalLM from .modeling_utils import register_auto_model, register_vision_encoder from .musespark_open_audio import ( diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index 9d5dd7ff74b3..7af1356440a1 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -57,15 +57,11 @@ from ...sampling_params import SamplingParams from ..attention.backends import AttentionMetadata from ..model_config import ModelConfig +from ..multimodal.model import EncoderGroup, MultimodalModelMixin, encode_multimodal_by_groups from ..peft.lora.config import LoraConfig from ..peft.lora.ownership import DecoderLoraOwner from ..speculative import SpecMetadata from .modeling_auto import AutoModelForCausalLM -from .modeling_multimodal_mixin import ( - EncoderGroup, - MultimodalModelMixin, - encode_multimodal_by_groups, -) from .modeling_multimodal_utils import ( _is_mm_disagg, filter_mm_token_from_input_ids, diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index d677e2c775aa..f366562f0708 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -70,11 +70,11 @@ from ..attention.rotary_embedding import MRotaryEmbedding, RotaryEmbedding from ..modules.gated_mlp import GatedMLP +from ..multimodal.model import MultimodalModelMixin from .modeling_auto import AutoModelForCausalLM from .modeling_multimodal_encoder import MultimodalEncoderMixin from .modeling_multimodal_encoder import \ prepare_encoder_attn_metadata as _prepare_qwen_vl_vision_attn_metadata -from .modeling_multimodal_mixin import MultimodalModelMixin from .modeling_multimodal_utils import ( _install_processor_output_validation_filter, find_input_mm_embeds, fuse_input_embeds, get_attached_multimodal_embeddings, diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 96107532b165..614476342407 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -51,16 +51,16 @@ from ..modules.layer_norm import LayerNorm from ..modules.linear import Linear, TensorParallelMode from ..modules.mlp import MLP -from .checkpoints.base_weight_mapper import BaseWeightMapper -from .checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper -from .modeling_auto import AutoModelForCausalLM -from .modeling_multimodal_encoder import MultimodalEncoderMixin, prepare_encoder_attn_metadata -from .modeling_multimodal_mixin import ( +from ..multimodal.model import ( EncoderGroup, MultimodalModelMixin, PreparedLlmInputs, encode_multimodal_by_groups, ) +from .checkpoints.base_weight_mapper import BaseWeightMapper +from .checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper +from .modeling_auto import AutoModelForCausalLM +from .modeling_multimodal_encoder import MultimodalEncoderMixin, prepare_encoder_attn_metadata from .modeling_qwen2vl import ( Qwen2_5_VLVisionAttention, Qwen2VLInputProcessorBase, diff --git a/tensorrt_llm/_torch/multimodal/__init__.py b/tensorrt_llm/_torch/multimodal/__init__.py new file mode 100644 index 000000000000..0fb201f6f880 --- /dev/null +++ b/tensorrt_llm/_torch/multimodal/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Shared multimodal model execution contracts.""" diff --git a/tensorrt_llm/_torch/multimodal/embeddings.py b/tensorrt_llm/_torch/multimodal/embeddings.py new file mode 100644 index 000000000000..00690afb3ba6 --- /dev/null +++ b/tensorrt_llm/_torch/multimodal/embeddings.py @@ -0,0 +1,547 @@ +# Copyright 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Encoder embedding storage, active-row selection, and token fusion.""" + +from __future__ import annotations + +from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple, Union, cast + +import torch + +from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.logger import logger + + +class TokenEmbedding(Protocol): + """Token lookup required by the default multimodal fusion path.""" + + num_embeddings: int + + def __call__(self, input_ids: torch.Tensor) -> torch.Tensor: ... + + +def _get_uncached_multimodal_params( + multimodal_params: List[MultimodalParams], +) -> List[MultimodalParams]: + """ + Get uncached multimodal params that need encoder processing for chunk prefill. + """ + params_to_run = [] + + for param in multimodal_params: + # Skip if no multimodal content + if not param.has_content(): + continue + + # Check if embeddings are already cached + if ( + param.multimodal_data + and "multimodal_embedding" in param.multimodal_data + and param.multimodal_data["multimodal_embedding"] is not None + ): + logger.debug("Skipping encoder forward for param with cached multimodal_embedding") + continue + + # This param needs encoder processing + params_to_run.append(param) + + return params_to_run + + +def _store_chunked_prefill_embeddings( + multimodal_params: List[MultimodalParams], + embeddings: List[torch.Tensor], +) -> None: + """ + Store computed multimodal embeddings back to multimodal_data to avoid recomputation. + + NOTE: this function only stores multimodal embeddings within the current request context, + mostly for chunked prefill. It does not persist embeddings across different requests or sessions. + """ + # TODO: support multiple multimodal modalities per request + if len(embeddings) > 1: + raise ValueError("Multiple modalities caching is not supported yet.") + mm_embed = embeddings[0] + + # Collect embedding lengths for each parameter + embed_lengths = [] + for param in multimodal_params: + if param.multimodal_runtime is not None: + embed_lengths.append(param.multimodal_runtime.total_embeds_in_request) + + # Validate total length matches + total_expected = sum(embed_lengths) + assert len(mm_embed) == total_expected, ( + f"Number of mm_embeds ({len(mm_embed)}) does not match expected total ({total_expected})" + ) + + # Use torch.split for efficient tensor splitting + split_embeddings = torch.split(mm_embed, embed_lengths, dim=0) + valid_params = [p for p in multimodal_params if p.multimodal_runtime is not None] + + # Cache split embeddings to each parameter + logger.debug( + f"Caching {len(split_embeddings)} multimodal embedding chunks in {len(multimodal_params)} params" + ) + for param, embed_chunk in zip(valid_params, split_embeddings, strict=True): + param.multimodal_data["multimodal_embedding"] = embed_chunk + + logger.debug(f"Cached {len(split_embeddings)} multimodal embedding chunks in this iteration") + + +def _normalize_encoder_embeddings( + encoder_embeddings: Union[torch.Tensor, List[torch.Tensor]], +) -> List[torch.Tensor]: + if isinstance(encoder_embeddings, torch.Tensor): + return [encoder_embeddings] + + if not isinstance(encoder_embeddings, list) or not all( + isinstance(embedding, torch.Tensor) for embedding in encoder_embeddings + ): + raise TypeError("encoder_forward_fn must return a torch.Tensor or a list of torch.Tensor.") + + return encoder_embeddings + + +def _join_embeddings( + embeddings: List[torch.Tensor], *, start: int = 0, end: Optional[int] = None +) -> torch.Tensor: + """Join per-item/per-request embeddings along dim 0 without a needless copy. + + ``torch.cat`` always allocates, so calling it on a one-element list + duplicates the whole embedding. That matters here because the same rows + are joined several times on the way to the LLM input (per request, per + batch, and again in ``fuse_input_embeds``), and each redundant copy is a + full extra allocation live at the prefill peak. Multimodal embeddings are + read-only downstream -- they are sliced and scattered into the input + embedding buffer, never written in place -- so returning the sole tensor + is equivalent to concatenating it. If a row range is given, slice each + segment first so rows outside the current prefill chunk are never copied. + """ + if start != 0 or end is not None: + slices = [] + offset = 0 + for embedding in embeddings: + segment_end = offset + embedding.shape[0] + slice_end = segment_end if end is None else min(end, segment_end) + if max(start, offset) < slice_end: + slices.append(embedding[max(start - offset, 0) : slice_end - offset]) + offset = segment_end + if end is not None and offset >= end: + break + if not slices: + return embeddings[0][:0] + embeddings = slices + + return embeddings[0] if len(embeddings) == 1 else torch.cat(embeddings, dim=0) + + +def get_multimodal_embeddings( + encoder_forward_fn: Callable[ + ..., + torch.Tensor | List[torch.Tensor], + ], + multimodal_params: List[MultimodalParams], + encoder_kwargs: Optional[Dict[str, Any]] = None, +) -> List[torch.Tensor]: + """ + High-level utility to get multimodal embeddings from encoder or cached embeddings. + + This function will: + 1. Identify which parameters need encoder processing + 2. Run encoder forward only on uncached parameters + 3. Cache newly computed embeddings (if enabled) + 4. Gather all embeddings for the batch + + Args: + encoder_forward_fn: Callable that performs encoder forward pass. + Should accept List[MultimodalParams] and return either + a single torch.Tensor or List[torch.Tensor]. + multimodal_params: All multimodal parameters in the batch. + encoder_kwargs: Optional kwargs to pass to encoder_forward_fn. + Returns: + List of multimodal embeddings for all multimodal params in the batch. + """ + if not multimodal_params: + return [] + + # Wait before touching tensors produced on the MM side stream. Do not clear the event here; + # repeated stream-side waits are cheap, and leaving the event field untouched avoids races if a + # caller accidentally reuses it. + # Register attached embeddings with the consumer stream as well: the request drops its Python + # references after prefill, potentially before the asynchronous gather below has finished + # reading them. + for param in multimodal_params: + if param.encoder_event is not None: + consumer_stream = torch.cuda.current_stream() + consumer_stream.wait_event(param.encoder_event) + embeds = param.multimodal_data.get("multimodal_embedding") + if isinstance(embeds, torch.Tensor): + embeds = [embeds] + if isinstance(embeds, list): + for embed in embeds: + if isinstance(embed, torch.Tensor) and embed.is_cuda: + embed.record_stream(consumer_stream) + + # Step 1: Find uncached multimodal params that need encoder processing + uncached_multimodal_params = _get_uncached_multimodal_params(multimodal_params) + + # Step 2: Run encoder forward only on uncached parameters + if uncached_multimodal_params: + kwargs = encoder_kwargs or {} + encoder_embeddings = encoder_forward_fn(uncached_multimodal_params, **kwargs) + encoder_embeddings = _normalize_encoder_embeddings(encoder_embeddings) + + # TODO: support multiple multimodal modalities per request + if len(encoder_embeddings) > 1: + logger.warning( + f"Multiple modalities caching is not supported yet. " + f"encoder returned {len(encoder_embeddings)} embeddings " + f"(types: {[type(e).__name__ for e in encoder_embeddings]}, " + f"shapes: {[e.shape if hasattr(e, 'shape') else 'N/A' for e in encoder_embeddings]}) " + f"for {len(uncached_multimodal_params)} uncached params. " + f"encoder_forward_fn={encoder_forward_fn}" + ) + return encoder_embeddings + + # Validate that multimodal_runtime has required attributes for caching + if ( + not hasattr(uncached_multimodal_params[0], "multimodal_runtime") + or uncached_multimodal_params[0].multimodal_runtime is None + or uncached_multimodal_params[0].multimodal_runtime.total_embeds_in_request is None + ): + logger.warning( + "Multimodal runtime data missing or incomplete, will not cache embeddings." + ) + return encoder_embeddings + + # Step 3: Cache the computed embeddings to multimodal_data["multimodal_embedding"]. + # A TP-sharded encoder can return an NCCL window tensor whose lease ends + # with the enclosing ModelCaller scope, while chunked prefill reads the + # cached rows in later iterations, so cache an owned copy instead of + # views of the encoder output. + encoder_embeddings = [encoder_embeddings[0].clone()] + _store_chunked_prefill_embeddings(uncached_multimodal_params, encoder_embeddings) + + # Step 4: Gather after the stream wait and any encoder/cache writes. + return get_attached_multimodal_embeddings(multimodal_params) + + +def get_attached_multimodal_embeddings( + multimodal_params: List[MultimodalParams], +) -> List[torch.Tensor]: + """Gather embeddings already stored on MultimodalParams. + + Use this on E/P prefill workers and cached-only paths. The encoder already ran somewhere else. + Normally `find_input_mm_embeds` selects active rows afterwards. If item + scheduling attached current-chunk embeddings, select the same chunk from + any full-request embeddings before joining the batch. This keeps mixed + batches in one row layout without copying unrelated future chunks. + + Side-stream-prefetched requests must use `get_multimodal_embeddings`, which waits on + `encoder_event` and registers attached tensors with the consuming stream before gathering them. + """ + current_chunk_only = any( + param.multimodal_data.get("multimodal_embedding_is_chunk", False) + for param in multimodal_params + ) + attached_embeddings = [] + for param in multimodal_params: + embeds = param.multimodal_data.get("multimodal_embedding") + # No attached embedding for this request. + if embeds is None: + continue + runtime = param.multimodal_runtime + if ( + current_chunk_only + and runtime is not None + and not param.multimodal_data.get("multimodal_embedding_is_chunk", False) + ): + # Keep the full attached payload intact for later chunks/cache + # writes; only this forward's gathered tensor is chunk-local. + start = runtime.num_cached_mm_tokens + end = start + runtime.num_mm_tokens_in_chunk + if isinstance(embeds, list): + embeds = _join_embeddings(embeds, start=start, end=end) + else: + embeds = embeds[start:end] + elif isinstance(embeds, list): + # Legacy paths can attach a list of full-request output pieces. + embeds = _join_embeddings(embeds) + param.multimodal_data["multimodal_embedding"] = embeds + if not isinstance(embeds, torch.Tensor): + raise TypeError("multimodal_embedding must be a torch.Tensor") + attached_embeddings.append(embeds) + + if not attached_embeddings: + return [] + # Match get_multimodal_embeddings output: one concatenated tensor. + return [_join_embeddings(attached_embeddings)] + + +def find_input_mm_embeds( + mm_embeds: List[torch.Tensor], multimodal_params: List[MultimodalParams] +) -> List[torch.Tensor]: + """ + Find the multimodal mm_embeds that need processing from multimodal_params for each batch. + Supports both KV cache reuse and chunked prefill scenarios. + + Args: + - mm_embeds: List[torch.Tensor] - Multimodal embeddings for each batch + - multimodal_params: List[MultimodalParams] - Multimodal parameters with runtime data + + Returns: + - List[torch.Tensor] - Sliced mm_embeds containing only tokens that need processing: + - For KV cache reuse: tokens that are not cached + - For chunked prefill: tokens that are in the current chunk + - For mixed scenarios: both uncached and current chunk tokens + - Empty list if all tokens are cached or beyond current chunk + + Note: + - Supports both individual batching (len(mm_embeds) == len(multimodal_params)) + and pre-concatenated batching (len(mm_embeds) == 1) + - Call get_attached_multimodal_embeddings before this helper when + embeddings are already attached to multimodal_params. + - Handles chunked prefill by considering chunk boundaries and current chunk tokens + - Example: if a request has 8 MM embed rows, 2 cached rows, and 3 rows + in the current chunk, this keeps rows [2:5]. + - This function reads only CPU-resident counts from + ``MultimodalParams.multimodal_runtime`` and never touches GPU + tensors, so it does not introduce host-device synchronization. + The companion ``torch.where`` host sync sits in ``fuse_input_embeds`` + (via ``filter_mm_token_from_input_ids``) and is avoided by routing + executor-precomputed ``mm_token_indices`` / ``text_token_indices`` + through to that call — see the contract there. + """ + if not isinstance(mm_embeds, list): + raise TypeError("mm_embeds must be a list") + + # Current support two batching modes: + # 1. Pre-concatenated mm_embeds for each batch, i.e., len(mm_embeds) == 1 + # 2. Individual mm_embeds for each multimodal param, i.e., len(mm_embeds) == len(multimodal_params) + if len(mm_embeds) > 1 and len(mm_embeds) != len(multimodal_params): + raise ValueError( + f"Number of mm_embeds ({len(mm_embeds)}) does not match number of " + f"multimodal params ({len(multimodal_params)})." + ) + + if not multimodal_params or multimodal_params[0].multimodal_runtime is None: + # No slicing, return the full mm_embeds + return mm_embeds + + total_mm_tokens = sum( + param.multimodal_runtime.num_mm_tokens_in_chunk + for param in multimodal_params + if param.multimodal_runtime is not None + ) + + if total_mm_tokens == 0: + logger.debug( + "All multimodal tokens are cached or beyond current chunk, skipping vision encoder forward" + ) + return [] + + if not mm_embeds: + raise ValueError( + "No multimodal embeddings were provided or cached for active multimodal tokens." + ) + + if total_mm_tokens == sum(mm_embed.shape[0] for mm_embed in mm_embeds): + # Item-scheduled batches were already sliced by the producer/gather + # path. Do not apply the cached-prefix offset a second time. + return mm_embeds + + current_pos = 0 + slices = [] + for param in multimodal_params: + runtime = param.multimodal_runtime + if runtime is None: + continue + local_start_pos = runtime.num_cached_mm_tokens + local_end_pos = local_start_pos + runtime.num_mm_tokens_in_chunk + slices.append((current_pos + local_start_pos, current_pos + local_end_pos)) + if len(mm_embeds) == 1: # pre-concatenated; advance global cursor + current_pos += runtime.total_embeds_in_request + + if len(mm_embeds) == 1: + sliced = [mm_embeds[0][start:end] for start, end in slices] + return [_join_embeddings(sliced)] + return [mm_embeds[i][start:end] for i, (start, end) in enumerate(slices)] + + +def filter_mm_token_from_input_ids( + input_ids: torch.IntTensor, + vocab_size: int, + mm_token_ids: Optional[torch.IntTensor] = None, +) -> Tuple[torch.IntTensor, torch.IntTensor]: + """ + Filter multimodal tokens from input_ids. + Args: + input_ids: shape [text_total_length + mm_total_length]. + vocab_size: size of the model's vocabulary + mm_token_ids: possible token ids for multimodal tokens, if known. If + not known and set to None, it is assumed that the multimodal tokens + are out-of-vocabulary tokens i.e. the ``input_ids`` contains tokens + >= vocab_size that represent the multimodal tokens. + Note: + Example: input_ids=[1, 55, 2, 101], vocab_size=100, and + mm_token_ids=[55] returns mm_token_indices=[1]; token 101 is text + because explicit mm_token_ids overrides the OOV fallback. + This function involves host-device synchronization due to torch.where() (= torch.nonzero) requiring + host allocation. The output indices reside on the same device as input_ids. + Returns: + text_token_indices: indices of text tokens in the input_ids + mm_token_indices: indices of multimodal tokens in the input_ids + """ + if mm_token_ids is None: + # If mm_token_ids is None, assume the multimodal tokens are out-of-vocab + # (input_ids >= vocab_size). Avoids torch.isin() over a potentially + # unbounded mm_token_ids set. + mm_token_mask = input_ids >= vocab_size + else: + mm_token_ids = mm_token_ids.to(input_ids.device, dtype=input_ids.dtype) + mm_token_mask = torch.isin(input_ids, mm_token_ids) + # NOTE: torch.where() enforces a host sync + text_token_indices = torch.where(~mm_token_mask)[0] + mm_token_indices = torch.where(mm_token_mask)[0] + return text_token_indices, mm_token_indices + + +def _as_copy_index(index: torch.Tensor, destination: torch.Tensor) -> torch.Tensor: + """Return ``index`` as the int64 index ``index_copy_`` needs on ``destination``. + + Explicit indices can come from callers with another dtype or device; for the + executor's int64 device indices this is a no-op. + """ + return index.to(device=destination.device, dtype=torch.long, non_blocking=True) + + +def fuse_input_embeds( + embedding_layer: TokenEmbedding, + input_ids: torch.IntTensor, + mm_embeds: List[torch.Tensor], + mm_token_ids: Optional[torch.IntTensor] = None, + text_token_indices: Optional[torch.IntTensor] = None, + mm_token_indices: Optional[torch.IntTensor] = None, + extra_embeds: Optional[List[torch.Tensor]] = None, + **kwargs, + # TODO: make unified return type for all models +) -> Union[ + Tuple[Optional[torch.IntTensor], Optional[torch.FloatTensor]], + Tuple[ + Optional[torch.IntTensor], Optional[torch.FloatTensor], Optional[List[torch.FloatTensor]] + ], +]: + """ + Fuse text and multimodal rows into [text_total_length + mm_total_length, hidden_dim]. + + Look up text embeddings and scatter multimodal rows into their token positions. + + Args: + embedding_layer: embedding layer of the model. + input_ids: Packed token IDs of shape [text_total_length + mm_total_length], + with padding removed across the in-flight batch. + mm_embeds: List[(mm_total_length1, hidden_dim), ..., (mm_total_lengthi, hidden_dim)]. + mm_token_ids: Known multimodal token IDs. None selects the out-of-vocabulary + placeholder convention. + extra_embeds: Optional list of extra embed tensors for models that support it (e.g., Qwen3-VL/Qwen3-MoE-VL). + Returns: + With no multimodal rows, return input_ids and no input embeddings. + Otherwise return no input IDs and fused embeddings of shape [total_length, hidden_dim]. + Auxiliary embeddings, when present, form the third return value. + Note: + - Use explicit indices when both are supplied. If either is missing, derive both + from mm_token_ids, or from the OOV convention when mm_token_ids is None. + - Example: len(torch.cat(mm_embeds)) must match len(mm_token_indices); + for chunked prefill, pass only the current chunk's mm_embeds or + explicit indices for the active MM token positions. + - Sync-free contract: passing both ``text_token_indices`` and + ``mm_token_indices`` skips the GPU ``torch.where`` host sync. The + executor (``DecoderInputPreparer.prepare_inputs``) precomputes them on a CPU + ``input_ids`` copy via ``runners.prepare_multimodal_indices`` (which uses + ``filter_mm_token_from_input_ids`` against ``self.model.mm_token_ids`` + when present, else the OOV fallback ``>= vocab_size``) and ships + them as pinned async H2D tensors in the inputs dict. VLM forwards + are expected to forward these explicitly rather than relying on a + ``**kwargs`` splat — if they don't, this function falls back to the + host-syncing ``filter_mm_token_from_input_ids`` on GPU ``input_ids``. + - In-vocab fast path: when ``mm_token_ids`` is provided, the caller is + declaring that mm placeholder IDs are real vocabulary entries. In + that case the text path skips its ``index_select`` + ``torch.empty`` + + text scatter and embeds the full ``input_ids`` once, overwriting + mm rows afterwards. This is the path taken by every VLM currently + in tree EXCEPT Hyperclovax, which keeps the legacy ``vocab_size + 1`` + OOV remap in its input processor and therefore passes + ``mm_token_ids=None`` so the OOV branch is taken — embedding an OOV + id would be out-of-bounds. + """ + if len(mm_embeds) == 0: + if extra_embeds is not None and len(extra_embeds) > 0: + return input_ids, None, extra_embeds + return input_ids, None + + mm_embed = _join_embeddings(mm_embeds) + + # TODO: derive a missing index tensor from the supplied complement without host synchronization. + if text_token_indices is None or mm_token_indices is None: + # Filtering uses torch.where(), which synchronizes CUDA inputs with the host. + 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 + ) + if mm_token_indices.shape[0] != mm_embed.shape[0]: + raise ValueError( + f"Multimodal token count mismatch: found {len(mm_token_indices)} image tokens in input_ids " + f"but received {mm_embed.shape[0]} image embeddings." + ) + + if mm_token_ids is not None: + # In-vocab fast path: caller declared mm tokens are real vocabulary + # IDs, so a single full embedding lookup is safe. Saves the + # ``index_select`` over text positions, the ``torch.empty`` alloc, and + # the text scatter (~2 GPU kernels per VLM forward) at the cost of + # ~mm-fraction extra embedding lookups whose result is overwritten. + input_embeds = embedding_layer(input_ids) + else: + # OOV path: input_ids holds placeholder IDs >= num_embeddings at mm + # positions (e.g. ``vocab_size + 1``), so embedding the full sequence + # would be out-of-bounds. Gather only the in-vocab text positions and + # scatter into a fresh buffer. + text_embed = embedding_layer(input_ids[text_token_indices]) + input_embeds = torch.empty( + input_ids.shape[0], mm_embed.shape[-1], device=text_embed.device, dtype=text_embed.dtype + ) + input_embeds.index_copy_(0, _as_copy_index(text_token_indices, input_embeds), text_embed) + if extra_embeds is not None and len(extra_embeds) > 0: + # only support single modality for deepstack features for now + for i, extra_feature in enumerate(extra_embeds): + extra_embed = torch.zeros( + input_ids.shape[0], + mm_embed.shape[-1], + device=extra_feature.device, + dtype=extra_feature.dtype, + ) + extra_embed.index_copy_(0, _as_copy_index(mm_token_indices, extra_embed), extra_feature) + extra_embeds[i] = extra_embed + + input_embeds.index_copy_( + 0, + _as_copy_index(mm_token_indices, input_embeds), + mm_embed.to(dtype=input_embeds.dtype, device=input_embeds.device), + ) + if extra_embeds is not None and len(extra_embeds) > 0: + return None, cast(torch.FloatTensor, input_embeds), extra_embeds + return None, cast(torch.FloatTensor, input_embeds) diff --git a/tensorrt_llm/_torch/multimodal/model.py b/tensorrt_llm/_torch/multimodal/model.py new file mode 100644 index 000000000000..20f91cec5964 --- /dev/null +++ b/tensorrt_llm/_torch/multimodal/model.py @@ -0,0 +1,2006 @@ +# Copyright 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import copy +import inspect +import itertools +from dataclasses import dataclass, replace +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ClassVar, + Dict, + Hashable, + Iterable, + Iterator, + List, + Optional, + Sequence, + Tuple, +) + +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalInput, + MultimodalItemMetadata, + MultimodalItemPartition, + MultimodalParams, + MultimodalRuntimeData, + select_multimodal_fields, +) +from tensorrt_llm.logger import logger + +from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs +from ..pyexecutor.runtime_hooks import DEFAULT_RUNTIME_HOOKS, RuntimeHooks, runtime_hooks_of +from .embeddings import ( + _store_chunked_prefill_embeddings, + find_input_mm_embeds, + fuse_input_embeds, + get_multimodal_embeddings, +) + + +class MultimodalEncoderContractError(ValueError): + """A request-local MM encoder input or output contract violation.""" + + +class MultimodalModelContractError(TypeError): + """A multimodal model class or instance does not implement the encoder runtime contract.""" + + +def _assemble_multimodal_encoder_embeddings( + item_tensors: Dict[int, torch.Tensor], total_items: int +) -> torch.Tensor: + """Copy prompt-ordered item embeddings into one contiguous tensor.""" + ordered = [item_tensors[i] for i in range(total_items)] + first = ordered[0] + if any( + output.shape[1:] != first.shape[1:] + or output.dtype != first.dtype + or output.device != first.device + for output in ordered[1:] + ): + raise ValueError( + "MM encoder items for one request must have matching output shape, dtype, and device" + ) + + # The copy is intentional even for one item: cache hits alias cache-owned + # storage, while this result must be independently owned by the request. + return torch.cat([output.detach() for output in ordered], dim=0) + + +@dataclass(frozen=True) +class EncoderGroup: + """Modalities that share a single encoder call. + + Batching all items in a group into one encoder invocation amortizes + fixed costs (kernel launches, dispatch) across items. The framework + splits the output back per-modality and reorders it into prompt order + via each request's `MultimodalItemMetadata.item_refs` manifest. + + Contract between `build_batched_input` and `encoder_fn`: + + * `build_batched_input` must concatenate items across requests in + `modalities` order (all items for the first modality across requests, + then all items for the second modality, etc.). + * Within a modality, items must appear in the same per-request iteration + order that `_lengths_by_modality` uses (i.e. the order of + `multimodal_params` passed in). + * `encoder_fn` must return one tensor whose rows correspond 1:1 to the + input layout produced by `build_batched_input`, so the framework can + split the output by `_lengths_by_modality` and reorder into prompt + order via each request's item manifest. + """ + + modalities: Tuple[str, ...] + """Ordered modality names that share this encoder. Defines the row + layout of the encoder output tensor: first all items of + `modalities[0]`, then all items of `modalities[1]`, etc.""" + + encoder_fn: Callable[..., torch.Tensor] + """Encoder call invoked as `encoder_fn(**build_batched_input(params))`. + Returns a single tensor with one row per embedding, laid out per the + contract above. When those rows are already in prompt order, the + framework may return a row slice of this tensor without a copy, so it + must not alias storage that a later call overwrites (for example, a CUDA + graph's static output buffer); clone such outputs before returning + them.""" + + build_batched_input: Callable[[List[MultimodalParams]], Dict[str, Any]] + """Builds the kwargs dict passed to `encoder_fn`. Responsible for + concatenating raw per-item tensors from `multimodal_data` across + requests in the order described in the class docstring.""" + + +def _lengths_by_modality( + item_layouts: Sequence[tuple[Sequence[tuple[str, int]], Sequence[int]]], + modalities: Tuple[str, ...], +) -> Dict[str, List[int]]: + """Split resolved output lengths into each group's payload-local order. + + Layouts describe whole requests, so only modalities in this encoder + group contribute lengths (for example, image/video in a vision group + while the same request's audio is handled by another group). + + Encoders batch in payload order, which can differ from prompt order. + Modality-local indices restart for each request. For example:: + + item_layouts = [ + ([("image", 1), ("video", 0), ("image", 0)], [20, 30, 10]), + ([("audio", 0), ("image", 0)], [40, 50]), + ] + modalities = ("image", "video") + # Result: {"image": [10, 20, 50], "video": [30]} + + These lengths split encoder outputs before restoring prompt order. + """ + by_modality: Dict[str, List[int]] = {m: [] for m in modalities} + for item_refs, lengths in item_layouts: + if not item_refs: + continue + modality = item_refs[0][0] + if ( + len(item_refs) == len(lengths) + and modality in by_modality + and all( + item_modality == modality and index == expected + for expected, (item_modality, index) in enumerate(item_refs) + ) + ): + by_modality[modality].extend(lengths) + continue + request_lengths: dict[str, list[tuple[int, int]]] = {m: [] for m in modalities} + for (modality, index), length in zip(item_refs, lengths, strict=True): + if modality in request_lengths: + request_lengths[modality].append((index, length)) + for modality, indexed_lengths in request_lengths.items(): + by_modality[modality].extend(length for _, length in sorted(indexed_lengths)) + return by_modality + + +def _reorder_embeds_by_manifest( + item_layouts: Sequence[tuple[Sequence[tuple[str, int]], Sequence[int]]], + per_modality_embeds: Dict[str, torch.Tensor], + per_modality_lengths: Dict[str, List[int]], +) -> torch.Tensor: + """Slice per-modality tensors item-by-item and concat in prompt order. + + If only one modality has items and every request lists them in encoder + order, that modality's tensor is returned as is, without a copy. + """ + populated_modalities = [ + (modality, embedding) + for modality, embedding in per_modality_embeds.items() + if per_modality_lengths.get(modality) + ] + if len(populated_modalities) == 1: + modality, embedding = populated_modalities[0] + prompt_order_matches_encoder = True + for item_refs, _ in item_layouts: + item_indices = [ + index for item_modality, index in item_refs if item_modality == modality + ] + if item_indices != list(range(len(item_indices))): + prompt_order_matches_encoder = False + break + if prompt_order_matches_encoder: + # A group with one populated modality already emits request and + # item rows in prompt order. Empty peer modalities (for example, + # video in an image-only Qwen batch) do not require a reorder. + return embedding + + per_modality_row_starts: Dict[str, List[int]] = { + m: list(itertools.accumulate(lens, initial=0)) for m, lens in per_modality_lengths.items() + } + + slices: List[torch.Tensor] = [] + # Item indices are per-request per-modality; advance a cursor to + # translate it into a global item index within `per_modality_embeds`. + per_modality_cursor: Dict[str, int] = {m: 0 for m in per_modality_embeds} + for item_refs, _ in item_layouts: + req_counts: Dict[str, int] = {} + for m, item_index in item_refs: + if m not in per_modality_embeds: + continue + i = per_modality_cursor[m] + item_index + starts = per_modality_row_starts[m] + slices.append(per_modality_embeds[m][starts[i] : starts[i + 1]]) + req_counts[m] = req_counts.get(m, 0) + 1 + for m, c in req_counts.items(): + per_modality_cursor[m] += c + if not slices: + # Preserve encoder dtype/device and width when no item contributes rows. + if per_modality_embeds: + return torch.cat(list(per_modality_embeds.values()), dim=0) + return torch.empty(0) + return torch.cat(slices, dim=0) + + +def encode_multimodal_by_groups( + mm_encoder_groups: Sequence["EncoderGroup"], + multimodal_params: List[MultimodalParams], +) -> torch.Tensor: + """Run each group's encoder over its batched items and reorder into + per-request prompt order. + + For each group present in the batch, one encoder call is issued over all + items across all requests belonging to that group's modalities + (arithmetic-intensity win). The output is split back per-modality using + the prompt-ordered output lengths in `MultimodalItemMetadata`, + then reordered by the same metadata's `item_refs`. + + Shared entry point for both the aggregated (`MultimodalModelMixin`) and + mm-encoder-only (`Qwen3VisionModelBase.forward`) paths so the ordering + contract lives in one place. + """ + per_modality_embeds: Dict[str, torch.Tensor] = {} + per_modality_lengths: Dict[str, List[int]] = {} + item_layouts = [mp.get_item_layout() for mp in multimodal_params] + for group in mm_encoder_groups: + group_params = [] + group_layouts = [] + for mp, layout in zip(multimodal_params, item_layouts, strict=True): + if any(mp.multimodal_data.get(m) is not None for m in group.modalities): + if not layout[0]: + raise MultimodalEncoderContractError( + "MM item metadata with output lengths is required for grouped encoding" + ) + group_params.append(mp) + group_layouts.append(layout) + if not group_params: + continue + out = group.encoder_fn(**group.build_batched_input(group_params)) + lengths = _lengths_by_modality(group_layouts, group.modalities) + cursor = 0 + for m in group.modalities: + total = sum(lengths[m]) + per_modality_embeds[m] = out[cursor : cursor + total] + cursor += total + per_modality_lengths.update(lengths) + return _reorder_embeds_by_manifest(item_layouts, per_modality_embeds, per_modality_lengths) + + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + from ..pyexecutor.guided_decoder import CapturableGuidedDecoder + from ..pyexecutor.llm_request import LlmRequest + + +_MM_DATA_INPUT_MODALITY_KEYS = frozenset({"audio", "image", "video"}) +_MM_AUX_STREAM: Optional[tuple[int, torch.cuda.Stream]] = None +_MM_ENCODER_CACHE_LOG_NAME = "mm_encoder_cache" + + +def _build_request_multimodal_input( + request: "LlmRequest", cache_enabled: bool +) -> Optional[MultimodalInput]: + """Build the encoder-cache key metadata carried by one request.""" + if not cache_enabled: + return None + params = getattr(request, "py_multimodal_params", None) + if params is not None and params.multimodal_input is not None: + return params.multimodal_input + hashes = request.multimodal_hashes + if hashes is None: + return None + multimodal_input = MultimodalInput.from_components( + hashes, + request.multimodal_positions, + request.multimodal_lengths, + mm_item_run_cu_offsets=request.multimodal_item_run_cu_offsets, + mm_run_positions=request.multimodal_run_positions, + mm_run_lengths=request.multimodal_run_lengths, + ) + if params is not None: + params.multimodal_input = multimodal_input + return multimodal_input + + +def _get_mm_aux_stream(max_prefetch_ahead: int = 0) -> Optional[torch.cuda.Stream]: + """Return the side CUDA stream used for multimodal encoder prefetch. + + Returns `None` when side-stream prefetch is disabled, CUDA is unavailable, + or the current stream is being captured. The cache intentionally keeps only + one stream because executor processes are expected to run on one current + CUDA device; if the current device changes, the cached stream is replaced. + """ + global _MM_AUX_STREAM + + if max_prefetch_ahead <= 0: + return None + if not torch.cuda.is_available(): + return None + if torch.cuda.is_current_stream_capturing(): + return None + + device = torch.cuda.current_device() + if _MM_AUX_STREAM is None or _MM_AUX_STREAM[0] != device: + _MM_AUX_STREAM = (device, torch.cuda.Stream(device=device)) + logger.warning_once( + f"Using multimodal encoder side stream on CUDA device {device} " + f"with encoder_side_stream_max_ahead={max_prefetch_ahead}. " + "This may increase peak GPU memory usage because raw multimodal " + "encoder inputs and computed embeddings can be resident before " + "request prefill.", + key=f"mm_aux_stream_used_device_{device}", + ) + return _MM_AUX_STREAM[1] + + +@contextlib.contextmanager +def _run_on_aux_stream(aux_stream: torch.cuda.Stream) -> Iterator[torch.cuda.Event]: + """Run a block on `aux_stream` independently of the caller stream. + + Yields a CUDA event recorded on `aux_stream` when the block exits. Callers + can wait on that event from another stream before consuming tensors written + in the block. + + No entrance barrier is enforced: cross-iter MM encoder prefetch operates on + data disjoint from the current iteration's batch, so serializing aux-stream + work behind the caller stream would eliminate the overlap this stream + exists for. + """ + exit_event = torch.cuda.Event() + with torch.cuda.stream(aux_stream): + try: + yield exit_event + finally: + # Keep the sync point valid even if the block raises after queuing + # aux-stream work. + exit_event.record() + + +@dataclass(frozen=True) +class PreparedMultimodalEmbeddings: + """Active encoder rows before a model's token-embedding fusion. + + ``multimodal_embeddings`` contains request-ordered tensors with only rows + consumed by this forward. ``extra_embeds`` carries model-specific features + selected by the same hooks as the default fusion path. + """ + + input_ids: torch.Tensor + multimodal_embeddings: Sequence[torch.Tensor] = () + extra_embeds: Sequence[torch.Tensor] = () + mm_token_indices: Optional[torch.Tensor] = None + """Placeholder positions of `multimodal_embeddings` in this forward's packed tokens, + when the caller supplied them; same row count as `multimodal_embeddings`.""" + + +@dataclass(frozen=True) +class PreparedLlmInputs: + """Prepared inputs returned by `MultimodalModelMixin`.""" + + input_ids: Optional[torch.Tensor] + inputs_embeds: Optional[torch.Tensor] + extra_embeds: Sequence[torch.Tensor] = () + + +@dataclass(frozen=True) +class EncoderCachePartition: + """Per-item cache partition for a single `MultimodalParams`. + + `hits` maps item index to its cached embedding row-block; `miss_indices` lists item + indices that still require encoder work; `keys` is aligned to item order so miss + embeddings can be written back after they are computed. + + `keys` always spans every item in the request -- full embedding assembly needs the total + item count, and misses are written back by item index. `hits` and `miss_indices` instead + cover only `looked_up`, the indices the caller asked about. The item scheduler passes the + subset it picked for this iteration, so items it has no budget for yet are not touched: an + LRU `get` refreshes recency, and probing an item the caller cannot encode would reorder + eviction against items that are actually in flight. + """ + + hits: Dict[int, torch.Tensor] + miss_indices: list[int] + keys: list[Hashable] + looked_up: list[int] + + @property + def is_full_hit(self) -> bool: + return bool(self.looked_up) and not self.miss_indices + + @property + def is_full_miss(self) -> bool: + return bool(self.looked_up) and not self.hits + + +class MultimodalRuntimeMixin: + """Multimodal encoder runtime and executor capability contract. + + The executor treats every subclass as a multimodal model. The mixin owns encoder execution, + request-local reuse across prefill chunks, the persistent encoder-output cache, side-stream + prefetch and item scheduling. `prepare_multimodal_embeddings` returns the rows consumed by + one forward; the model splices them into its own token embeddings. + + The mixin has no notion of a separate language model and defines nothing that + `DecoderModelForCausalLM` defines, so a model that runs its own decoder can inherit both in + either order. Composite models whose language model is a separate module inherit + `MultimodalModelMixin`. + + Subclasses implement `encode_multimodal_inputs`, `embedding_dim` and `embedding_dtype`. + `check_multimodal_class_contract` and `check_multimodal_instance_contract` verify them. + + Current limitations: + + * For the time being, the persistent multimodal encoder cache stores per-item embeddings for + single-modality `MultimodalParams` objects. Mixed-modality objects bypass the cache. + * A partially cached `MultimodalParams` is handled by encoding only its miss items + and interleaving cached items back in original per-item order. The default + `build_multimodal_encoder_input` selects fields declared by + `get_multimodal_item_partitions`. Undeclared models retain legacy stacked/packed inference. + """ + + supports_vision_gpu_normalize: ClassVar[bool] = False + """Whether the model accepts processor-described raw vision pixels.""" + + supports_encoder_cache: ClassVar[bool] = False + """Whether the model's production forward path uses the persistent encoder cache.""" + + supports_mm_encoder_item_scheduling: ClassVar[bool] = False + """Whether the model supports item-level MM encoder scheduling: it implements the item-encode + forward and pairs with a processor that overrides + `BaseMultimodalInputProcessor.get_mm_item_metadata`. This is the executor-layer gate for + item scheduling.""" + + supports_concurrent_encoder: ClassVar[bool] = False + """Whether `encode_multimodal_inputs` may run on an auxiliary CUDA stream while the main + stream runs the model, which side-stream prefetch requires. A model that sets it keeps every + piece of mutable encoder state, such as kernel workspaces, separate per CUDA stream.""" + + REQUIRED_MULTIMODAL_MEMBERS: ClassVar[Tuple[str, ...]] = ( + "encode_multimodal_inputs", + "embedding_dim", + "embedding_dtype", + ) + """Members every subclass overrides; checked by `check_multimodal_class_contract`.""" + + model_config: ModelConfig + _multimodal_encoder_cache: Optional[TensorLRUCache] = None + + @classmethod + def check_multimodal_class_contract(cls) -> None: + """Raise `MultimodalModelContractError` if `cls` leaves a required member undefined.""" + missing = [ + name + for name in cls.REQUIRED_MULTIMODAL_MEMBERS + if inspect.getattr_static(cls, name) + is inspect.getattr_static(MultimodalRuntimeMixin, name) + ] + if missing: + raise MultimodalModelContractError( + f"{cls.__qualname__} inherits MultimodalRuntimeMixin but does not implement " + f"{', '.join(missing)}" + ) + + def check_multimodal_instance_contract(self) -> None: + """Read the declared embedding row layout from the constructed model. + + Properties run with the model's real configuration, so a property that reads a field + the configuration lacks fails here, naming the property, rather than later inside the + executor, where `torch.nn.Module.__getattr__` would replace the error with a generic + missing-attribute message. + """ + for name in ("embedding_dim", "embedding_dtype"): + descriptor = inspect.getattr_static(type(self), name) + try: + value = ( + descriptor.__get__(self, type(self)) + if hasattr(descriptor, "__get__") + else descriptor + ) + except (AttributeError, NotImplementedError, TypeError, ValueError) as error: + raise MultimodalModelContractError( + f"{type(self).__qualname__}.{name} failed: {error!r}" + ) from error + if name == "embedding_dim" and not (isinstance(value, int) and value > 0): + raise MultimodalModelContractError( + f"{type(self).__qualname__}.embedding_dim must be a positive int, got {value!r}" + ) + if name == "embedding_dtype" and not isinstance(value, torch.dtype): + raise MultimodalModelContractError( + f"{type(self).__qualname__}.embedding_dtype must be a torch.dtype, " + f"got {value!r}" + ) + + def max_encoder_tokens_per_item(self, modality: str) -> Optional[int]: + """Return the largest item, in encoder tokens, the encoder processes in one piece. + + Encoder memory profiling never builds a larger dummy item, and item scheduling never + sizes its per-iteration budget for a larger atomic item. A model that splits larger + items internally, or refuses them at admission, declares its bound here. `None`, the + default, leaves the input processor's limit in place. + """ + return None + + def allows_breakable_prefill_cuda_graph(self, llm_args: "TorchLlmArgs") -> bool: + """Whether a breakable prefill CUDA graph may capture this model's prefill forward. + + Encoder work does not fit in a captured graph. The default refuses, which is correct for + any model that encodes inside `forward`. A model that encodes before the captured region + and splices rows through an eager break, or never encodes under `llm_args`, overrides + this. Composite models (`MultimodalModelMixin`) keep the engine's structural check. + """ + return False + + @classmethod + def _cast_multimodal_encoder_dtype( + cls, + module: torch.nn.Module, + dtype: torch.dtype, + ) -> torch.nn.Module: + """Cast a multimodal encoder dtype without materializing meta tensors.""" + + def convert(tensor: torch.Tensor) -> torch.Tensor: + if not (tensor.is_floating_point() or tensor.is_complex()): + return tensor + if tensor.device == torch.device("meta"): + return torch.empty_like(tensor, dtype=dtype) + return tensor.to(dtype=dtype) + + return module._apply(convert) + + # Per-model registration of encoder-batching groups. Each `EncoderGroup` + # bundles a set of modalities that share one encoder call. Set as a class + # attribute or on `self` in `__init__` (when `encoder_fn` binds to instance + # methods). Consumers call the module-level `encode_multimodal_by_groups` + # with these groups; both the aggregated and mm-encoder-only paths share + # that helper so the ordering contract lives in one place. + mm_encoder_groups: Sequence[EncoderGroup] = () + + def encode_multimodal_inputs( + self, + multimodal_params: Sequence[MultimodalParams], + ) -> torch.Tensor: + """Run model-specific multimodal encoder work. + + Returns the single primary multimodal embedding tensor for the supplied params. Rows are + expected to be concatenated in request order, and special multimodal tokens occupy token + positions but do not have rows here. + """ + raise NotImplementedError + + def validate_multimodal_request_data(self, mm_data: dict) -> None: + """Admission-time validation of one request's multimodal payload. + + Called from `PyExecutor._validate_request` before activation; raise + `ValueError` to reject just this request. Conditions a model would + reject inside `forward` belong here instead: a forward-time failure + is handled batch-wide and fails every request scheduled alongside + the bad one. Default: accept everything. + """ + return None + + def prepare_multimodal_encoder_inputs( + self, + selected_items: Sequence[tuple[MultimodalParams, int]], + ) -> list[tuple[MultimodalParams, Sequence[int], str]]: + """Build selected item encoder inputs before the caller performs H2D. + + Adjacent items from the same request and modality are sliced in one + call. That is not just tidier: the packed-layout slicer splits the + request's whole pixel payload and concatenates the chosen pieces, so + slicing item-by-item re-splits the payload N times and copies each + item separately. + + Args: + selected_items: `(request params, item index)` pairs in scheduler-selected order. + + Returns: + Tuples of `(sliced encoder params, per-item embedding row counts, modality)`, in + input order. Flattening the row-count lists recovers the per-item sequence. + """ + encoder_inputs: list[tuple[MultimodalParams, Sequence[int], str]] = [] + for ( + multimodal_param, + run_indices, + modality, + item_metadata, + ) in self._runs_by_request_modality(selected_items): + item_refs = item_metadata.item_refs + # The two slicers take different index spaces. The raw-tensor + # slicer indexes the modality's own payload, so it gets the + # modality-local indices (plus the modality, which also lets it + # slice out of an interleaved request); the parallel metadata is + # prompt-ordered across modalities, so it gets the global ones. + try: + residual = self.build_multimodal_encoder_input( + multimodal_param, + [item_refs[i][1] for i in run_indices], + modality=modality, + ) + self._apply_metadata_slice( + residual, multimodal_param, run_indices, item_metadata=item_metadata + ) + except MultimodalEncoderContractError: + raise + except (KeyError, IndexError, TypeError, ValueError) as error: + raise MultimodalEncoderContractError( + f"Invalid multimodal encoder item input: {error}" + ) from error + encoder_inputs.append( + ( + residual, + residual.item_metadata.output_embedding_lengths, + modality, + ) + ) + return encoder_inputs + + @staticmethod + def _runs_by_request_modality( + selected_items: Sequence[tuple[MultimodalParams, int]], + ) -> Iterator[tuple[MultimodalParams, list[int], str, MultimodalItemMetadata]]: + """Split scheduler order into maximal same-request, same-modality runs. + + Only adjacent items merge, preserving scheduler order. Reuse the item + metadata validated at the input boundary or executor admission. + """ + run_param: Optional[MultimodalParams] = None + run_modality: Optional[str] = None + run_metadata: Optional[MultimodalItemMetadata] = None + run_indices: list[int] = [] + for multimodal_param, item_idx in selected_items: + try: + metadata = ( + run_metadata + if multimodal_param is run_param + else multimodal_param.item_metadata + ) + except (TypeError, ValueError) as error: + raise MultimodalEncoderContractError(str(error)) from error + if metadata is None or metadata.output_embedding_lengths is None: + raise MultimodalEncoderContractError( + "MM item metadata with output lengths is required for item encoding" + ) + if item_idx < 0 or item_idx >= len(metadata.item_refs): + raise MultimodalEncoderContractError( + f"MM item index {item_idx} is out of range for {len(metadata.item_refs)} item(s)" + ) + modality = metadata.item_refs[item_idx][0] + if run_indices and (multimodal_param is not run_param or modality != run_modality): + yield run_param, run_indices, run_modality, run_metadata + run_indices = [] + run_param, run_modality, run_metadata = multimodal_param, modality, metadata + run_indices.append(item_idx) + if run_indices: + yield run_param, run_indices, run_modality, run_metadata + + def forward_multimodal_encoder_items( + self, + encoder_inputs: Sequence[tuple[MultimodalParams, Sequence[int], str]], + ) -> list[torch.Tensor]: + """Forward prepared MM encoder inputs in scheduler item order. + + Args: + encoder_inputs: Tuples returned by `prepare_multimodal_encoder_inputs`. Consecutive + inputs with the same modality must be batch-compatible. + + Returns: + One encoder output tensor **per item** (not per input tuple). Each tensor has the + declared embedding row count and retains scheduler input order. + """ + outputs: list[torch.Tensor] = [] + group_params: list[MultimodalParams] = [] + group_lengths: list[int] = [] + group_modality: Optional[str] = None + + def flush_group() -> None: + if not group_params: + return + embeddings = self.encode_multimodal_inputs(group_params) + expected_length = sum(group_lengths) + if embeddings.shape[0] != expected_length: + raise MultimodalEncoderContractError( + f"MM encoder output length {embeddings.shape[0]} does not " + f"match the {expected_length} rows declared by the " + "selected items" + ) + outputs.extend(torch.split(embeddings, group_lengths, dim=0)) + group_params.clear() + group_lengths.clear() + + for multimodal_param, embedding_lengths, modality in encoder_inputs: + if group_modality is not None and modality != group_modality: + flush_group() + group_modality = modality + group_params.append(multimodal_param) + group_lengths.extend(embedding_lengths) + flush_group() + return outputs + + @property + def multimodal_token_ids(self) -> Optional[Sequence[int] | torch.Tensor]: + """Return placeholder token ids in `input_ids` replaced by MM embeds. + + These are sentinel token positions whose text embeddings are replaced + by multimodal embeddings. Return `None` to use the out-of-vocabulary + sentinel behavior in `fuse_input_embeds`. + """ + return None + + @property + def embedding_dim(self) -> int: + """Return the width of each cached multimodal embedding row.""" + raise NotImplementedError + + @property + def embedding_dtype(self) -> torch.dtype: + """Return the dtype of each cached multimodal embedding row.""" + raise NotImplementedError + + @property + def encoder_cache_active(self) -> bool: + """Whether the persistent encoder cache is active for this model. + + Single source of truth shared by: + + * the in-iter consume path + * the side-stream prefetch dispatch + * the engine's cache-related gate + * the KV-cache memory reservation. + + The cache is only active when the model opts in via `supports_encoder_cache` and configures + a positive capacity. + """ + if not self.supports_encoder_cache: + return False + multimodal_config = self.model_config.multimodal_config + return multimodal_config is not None and multimodal_config.encoder_cache_max_bytes > 0 + + def select_multimodal_params( + self, + multimodal_params: Sequence[MultimodalParams], + num_context_requests: int, + ) -> Sequence[MultimodalParams]: + """Select the params that participate in multimodal encoder work. + + 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() + and ( + param.multimodal_runtime is None + or param.multimodal_runtime.num_mm_tokens_in_chunk != 0 + ) + ] + + def after_full_multimodal_embeddings( + self, + *, + input_ids: torch.Tensor, + multimodal_params: Sequence[MultimodalParams], + embeddings: torch.Tensor, + **forward_kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Optional hook after multimodal embeddings are gathered. + + Runs after cache lookup or encoder execution, before active-row model + hooks. The row layout depends on the whole batch: if any param carries + item-scheduled current-chunk rows, the gather also cuts every other + param that has runtime data down to its current-chunk rows; otherwise + each param keeps its full-request rows. Row-dependent model transforms + should use `after_active_multimodal_embeddings` instead. + """ + return input_ids, embeddings + + def after_active_multimodal_embeddings( + self, + *, + active_embeddings: list[torch.Tensor], + multimodal_params: Sequence[MultimodalParams], + **forward_kwargs: Any, + ) -> tuple[list[torch.Tensor], Sequence[torch.Tensor]]: + """Optional hook after active chunk rows are selected and before fusion. + + Models can transform or split the active multimodal embeddings here + and return additional embedding tensors to fuse alongside the primary + multimodal embeddings. + """ + # Models with packed auxiliary features (e.g. Qwen3-VL) can split them here and return + # them as extra embeds without changing the base flow. + return active_embeddings, () + + def get_multimodal_item_partitions( + self, + param: MultimodalParams, + modality: str, + ) -> dict[str, MultimodalItemPartition] | None: + """Declare how items occupy each field of one modality's payload. + + Models declare batched axes, packed item lengths, and shared values. + Unlisted fields pass through unchanged. Returning None retains legacy + layout inference for models that have not declared their fields. + """ + metadata = param.item_metadata + return metadata.input_partitions.get(modality) if metadata is not None else None + + def build_multimodal_encoder_input( + self, + param: MultimodalParams, + local_item_indices: Sequence[int], + modality: Optional[str] = None, + ) -> MultimodalParams: + """Build a forward view containing selected modality-local items. + + ``local_item_indices`` indexes this modality's payload, not positions + in the request's prompt-ordered item metadata. + Declared field layouts drive common tensor/list selection. Models with + custom preprocessing can still override this hook. The caller slices + the parallel item metadata after this returns, preserving original + request indices for scheduler targets and cache ownership. + """ + if modality is None: + modality = self._encoder_cache_modality(param) + if modality is None: + raise NotImplementedError( + "Default `build_multimodal_encoder_input` cannot infer the modality of a " + "mixed-modality param. Pass `modality` with modality-local indices, or " + "override for other layouts." + ) + modality_data = param.multimodal_data[modality] + if not isinstance(modality_data, dict): + raise TypeError( + f"multimodal_data[{modality!r}] must be a dict, got {type(modality_data).__name__}" + ) + item_partitions = self.get_multimodal_item_partitions(param, modality) + if item_partitions is None: + sliced = self._select_legacy_multimodal_fields( + modality_data, modality, local_item_indices + ) + else: + sliced = select_multimodal_fields(modality_data, item_partitions, local_item_indices) + + # Metadata slicing must not rewrite the original request's hashes. + residual_input = ( + copy.copy(param.multimodal_input) if param.multimodal_input is not None else None + ) + residual_data = { + key: value + for key, value in param.multimodal_data.items() + if key not in _MM_DATA_INPUT_MODALITY_KEYS or key == modality + } + residual_data[modality] = sliced + return MultimodalParams( + multimodal_data=residual_data, + multimodal_input=residual_input, + ) + + def _select_legacy_multimodal_fields( + self, + modality_data: Dict[str, Any], + modality: str, + item_indices: Sequence[int], + ) -> Dict[str, Any]: + """Select recognized packed/stacked layouts for undeclared models.""" + indices = list(item_indices) + grid_key = {"image": "image_grid_thw", "video": "video_grid_thw"}.get(modality) + pixel_key = {"image": "pixel_values", "video": "pixel_values_videos"}.get(modality) + + if ( + (grid_key and pixel_key) + and (grid_key in modality_data) + and (pixel_key in modality_data) + ): + # Packed layout: prefix-sum patch counts to locate each item's slab, then + # concat the requested subset in item-index order. + grids = modality_data[grid_key] + n_items = grids.shape[0] + patch_counts = [int(c) for c in torch.prod(grids, dim=1).tolist()] + row_starts = list(itertools.accumulate(patch_counts, initial=0)) + if indices == list(range(indices[0], indices[0] + len(indices))): + # Contiguous run: the concatenation is just a row range, so take + # a view instead of copying the payload. The common case for + # both callers -- a scheduler picks items in order, and cache + # misses cluster. + pixel_slice = modality_data[pixel_key][ + row_starts[indices[0]] : row_starts[indices[-1] + 1] + ] + else: + per_item = torch.split(modality_data[pixel_key], patch_counts, dim=0) + pixel_slice = torch.cat([per_item[i] for i in indices], dim=0) + sliced = { + pixel_key: pixel_slice, + grid_key: grids[indices], + } + elif ( + modality == "image" + and "pixel_values" in modality_data + and "image_sizes" in modality_data + ): + # Stacked layout: dim-0 select from `pixel_values` and list-index `image_sizes`. + n_items = modality_data["pixel_values"].shape[0] + miss_sizes = [modality_data["image_sizes"][i] for i in indices] + miss_pixel = modality_data["pixel_values"][indices] + # `pixel_values` was padded to the request-wide max H/W by the input + # processor. After keeping only the miss subset, crop the trailing H/W back + # down to that subset's own max true size -- otherwise a downstream re-batch + # step (e.g. Mistral 3's `batch_pixel_values`) that pads to + # `max(residual.image_sizes)` would compute a negative pad amount whenever + # the omitted items were the largest in the original request. + if miss_sizes and miss_pixel.dim() >= 4: + max_h = max(int(s[0]) for s in miss_sizes) + max_w = max(int(s[1]) for s in miss_sizes) + miss_pixel = miss_pixel[..., :max_h, :max_w] + sliced = { + "pixel_values": miss_pixel, + "image_sizes": miss_sizes, + } + elif modality == "audio" and "input_features" in modality_data: + # Stacked layout: `input_features [B, mel_bins, T]` sliced on dim 0. + # Per-item masks (`input_features_mask`, `feature_attention_mask`, ...) + # are handled by the sibling-slice pass below. + n_items = modality_data["input_features"].shape[0] + sliced = { + "input_features": modality_data["input_features"][indices], + } + else: + raise NotImplementedError( + f"Default `build_multimodal_encoder_input` cannot slice {modality} layout " + f"with fields {sorted(modality_data)}; override this method." + ) + + # Sibling per-item fields (e.g. `second_per_grid_ts` on Qwen2.5-VL video) + # must be sliced alongside the load-bearing keys above, or the residual + # carries a shape-mismatched encoder input. + sliced = { + **modality_data, + **sliced, + **self._slice_per_item_sibling_fields(modality_data, n_items, indices, sliced.keys()), + } + + return sliced + + @staticmethod + def _slice_per_item_sibling_fields( + modality_data: Dict[str, Any], + n_items: int, + item_indices: Sequence[int], + already_sliced: Iterable[str], + ) -> Dict[str, Any]: + """Slice modality-dict siblings whose first axis is parallel to items. + + Anything with `shape[0] == n_items` (tensor) or `len == n_items` (list) is + assumed to be per-item metadata and sliced by `item_indices`. Fields already + handled by the caller (`already_sliced`) and everything else pass through. + """ + skip = set(already_sliced) + sliced: Dict[str, Any] = {} + for key, value in modality_data.items(): + if key in skip: + continue + if isinstance(value, torch.Tensor) and value.dim() > 0 and value.shape[0] == n_items: + sliced[key] = value[item_indices] + elif isinstance(value, list) and len(value) == n_items: + sliced[key] = [value[i] for i in item_indices] + return sliced + + def prepare_multimodal_embeddings( + self, + *, + input_ids: torch.Tensor, + multimodal_params: Optional[Sequence[MultimodalParams]], + num_context_requests: int, + mm_token_indices: Optional[torch.Tensor] = None, + **forward_kwargs: Any, + ) -> PreparedMultimodalEmbeddings: + """Prepare active encoder rows for a model-owned fusion operation. + + Owns encoder/cache lookup, request-local storage, stream ordering, + active-chunk selection, and both model hooks. Targets with their own + token embedding or splice can consume these rows without computing + text embeddings or replacing their original input token IDs. + + `mm_token_indices` are the executor's placeholder positions for this forward, the + `mm_token_indices` forward argument. When given, the model hooks receive them as before, + and the result pairs them with the rows after checking that both count the same rows. + A result without rows carries no indices. + """ + if mm_token_indices is not None: + forward_kwargs["mm_token_indices"] = mm_token_indices + prepared = self._prepare_active_multimodal_embeddings( + input_ids=input_ids, + multimodal_params=multimodal_params, + num_context_requests=num_context_requests, + **forward_kwargs, + ) + if mm_token_indices is None or not prepared.multimodal_embeddings: + return prepared + num_rows = sum(rows.shape[0] for rows in prepared.multimodal_embeddings) + if mm_token_indices.shape[0] != num_rows: + raise ValueError( + f"The batch has {mm_token_indices.shape[0]} multimodal placeholder positions " + f"but {num_rows} active multimodal embedding rows" + ) + return replace(prepared, mm_token_indices=mm_token_indices) + + def _prepare_active_multimodal_embeddings( + self, + *, + input_ids: torch.Tensor, + multimodal_params: Optional[Sequence[MultimodalParams]], + num_context_requests: int, + **forward_kwargs: Any, + ) -> PreparedMultimodalEmbeddings: + """Select, encode or reuse, and post-process the rows of one forward.""" + context_params = list( + self.select_multimodal_params( + multimodal_params or [], + num_context_requests, + ) + ) + if not context_params: + return PreparedMultimodalEmbeddings(input_ids=input_ids) + + embeddings = self._get_or_encode_multimodal_embeddings(context_params) + + input_ids, embeddings = self.after_full_multimodal_embeddings( + input_ids=input_ids, + multimodal_params=context_params, + embeddings=embeddings, + **forward_kwargs, + ) + + active_embeddings = find_input_mm_embeds([embeddings], list(context_params)) + active_embeddings, extra_embeds = self.after_active_multimodal_embeddings( + active_embeddings=active_embeddings, + multimodal_params=context_params, + **forward_kwargs, + ) + return PreparedMultimodalEmbeddings( + input_ids=input_ids, + multimodal_embeddings=active_embeddings, + extra_embeds=extra_embeds, + ) + + def _get_or_encode_multimodal_embeddings( + self, + multimodal_params: Sequence[MultimodalParams], + ) -> torch.Tensor: + """Return cached multimodal embeddings or run the encoder for misses. + + Delegates cache lookup and gather behavior to `get_multimodal_embeddings`, then validates + the single tensor contract for both encoded and cached-only paths. + If item scheduling supplies current-chunk rows, the gather selects + current-chunk rows for every request before concatenating the batch. + + During side-stream prefetch, this runs with the auxiliary stream current, so the H2D copies, + the encoder, and every persistent-cache `put()` are issued on that stream. `TensorLRUCache` + records each entry's producer event on the issuing (aux) stream; the next iteration's + main-stream consumer waits on the request-level `encoder_event` for ordering. + """ + encoder_cache = self._multimodal_encoder_cache + cache_misses: list[MultimodalParams] = [] + partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = [] + if encoder_cache is not None: + for param in multimodal_params: + if param.multimodal_data.get("multimodal_embedding") is not None: + # A present embedding means either an earlier forward already wrote the + # persistent cache, or a prefetch hit attached a cache-owned tensor. Either + # way, skip re-lookup and re-write. `get_multimodal_embeddings` waits on the + # request event and records the attached tensor on the consuming stream before + # gathering it. + continue + partition = self.partition_encoder_cache(param, encoder_cache) + if partition is None or partition.is_full_miss: + cache_misses.append(param) + continue + if partition.is_full_hit: + param.multimodal_data["multimodal_embedding"] = ( + _assemble_multimodal_encoder_embeddings(partition.hits, len(partition.keys)) + ) + continue + partial_hits.append((param, partition)) + + if partial_hits: + # `encoder_cache` is non-None here because partitions are only produced when the cache + # exists. + self._encode_with_partial_cache(partial_hits, encoder_cache) + + embeddings = get_multimodal_embeddings( + encoder_forward_fn=self.encode_multimodal_inputs, + multimodal_params=list(multimodal_params), + ) + if encoder_cache is not None: + for param in cache_misses: + self._write_encoder_cache_entries(param, encoder_cache) + + # Validate post-gather so cached-only paths (KV reuse, all-cached chunked prefill) are also + # checked, not just paths that ran the encoder. + self._validate_embeddings( + embeddings, + multimodal_params, + current_chunk_only=any( + param.multimodal_data.get("multimodal_embedding_is_chunk", False) + for param in multimodal_params + ), + ) + return embeddings[0] + + def _initialize_multimodal_encoder_cache(self, max_bytes: int) -> Optional[TensorLRUCache]: + """Initialize the model-owned multimodal encoder-output cache once. + + The cache stores per-item embeddings keyed by the common item metadata. + `ModelEngine` resolves `max_bytes` as the larger of the item-scheduling + output budget and persistent-reuse capacity before runtime access. + Zero leaves the cache disabled. + """ + if self._multimodal_encoder_cache is not None: + if max_bytes != self._multimodal_encoder_cache.max_bytes: + raise ValueError( + f"{_MM_ENCODER_CACHE_LOG_NAME}: already initialized with " + f"max_bytes={self._multimodal_encoder_cache.max_bytes}; " + f"cannot reinitialize with max_bytes={max_bytes}" + ) + return self._multimodal_encoder_cache + + if max_bytes == 0: + logger.debug_once( + f"{_MM_ENCODER_CACHE_LOG_NAME}: disabled because neither item scheduling nor " + "persistent reuse requires storage.", + key="mm_encoder_cache_disabled", + ) + return None + + multimodal_config = self.model_config.multimodal_config + self._multimodal_encoder_cache = TensorLRUCache( + max_bytes, + name=_MM_ENCODER_CACHE_LOG_NAME, + cuda_stream_aware=multimodal_config.encoder_side_stream_max_ahead > 0, + ) + try: + embedding_dim = self.embedding_dim + embedding_dtype = self.embedding_dtype + except NotImplementedError: + logger.info( + f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, " + "embedding row capacity unavailable because the model does not implement " + "embedding_dim and embedding_dtype." + ) + else: + bytes_per_embedding_row = ( + embedding_dim * torch.empty((), dtype=embedding_dtype).element_size() + ) + max_embedding_rows = max_bytes // bytes_per_embedding_row + logger.info( + f"{_MM_ENCODER_CACHE_LOG_NAME}: created with max_bytes={max_bytes}, " + f"max_embedding_rows={max_embedding_rows}, embedding_dim={embedding_dim}, " + f"embedding_dtype={embedding_dtype}" + ) + return self._multimodal_encoder_cache + + @staticmethod + def _encoder_cache_modality(param: MultimodalParams) -> Optional[str]: + """Return the single modality represented by `param`, if cacheable. + + `None` means the params either do not identify a modality or contain + multiple modality inputs. Per-item operations use the common item refs + when one request contains several modalities. + """ + mm_data = param.multimodal_data or {} + modalities = [key for key in _MM_DATA_INPUT_MODALITY_KEYS if key in mm_data] + + modality = mm_data.get("modality_type") + if isinstance(modality, str): + # Trust the explicit `modality_type` only when it agrees with the actual data keys. + # Otherwise fall through to the mixed-modality skip so an inconsistent producer (e.g. + # `modality_type="image"` while both image and audio data are present) cannot bypass the + # safety check below and have the cache serve embeddings for the wrong modality. + if modalities == [modality]: + return modality + + if len(modalities) != 1: + # A whole-param modality is ambiguous for mixed inputs. + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params with {len(modalities)} detected " + "modalities." + ) + return None + return modalities[0] + + @classmethod + def _encoder_cache_keys( + cls, + param: MultimodalParams, + ) -> Optional[list[Hashable]]: + """Build per-item keys from content hashes and the common item record.""" + mm_input = param.multimodal_input + mm_data = param.multimodal_data or {} + if mm_input is None: + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params without multimodal hashes." + ) + return None + + item_metadata = param.item_metadata + if item_metadata is None: + logger.debug(f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params without item metadata.") + return None + if item_metadata.output_embedding_lengths is None: + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping params without item output lengths." + ) + return None + item_modalities = {modality for modality, _ in item_metadata.item_refs} + if any( + modality not in item_modalities + for modality in _MM_DATA_INPUT_MODALITY_KEYS + if mm_data.get(modality) is not None + ): + return None + return cls.build_encoder_cache_item_keys( + mm_input.multimodal_hashes, + item_metadata.item_refs, + item_metadata.output_embedding_lengths, + mm_data.get("mm_processor_kwargs_hash"), + ) + + @staticmethod + def _encoder_cache_item_key( + modality: str, + item_hash: Sequence[int], + embedding_length: int, + kwargs_hash: str, + ) -> Hashable: + # Sole definition of the key format, so entries written by the + # full-request path and the item-scheduling path hit from either. + return (modality, tuple(item_hash), int(embedding_length), kwargs_hash) + + @classmethod + def build_encoder_cache_item_keys( + cls, + multimodal_hashes: Optional[Sequence[Sequence[int]]], + item_refs: Sequence[tuple[str, int]], + embedding_lengths: Sequence[int], + kwargs_hash: Optional[str], + ) -> Optional[list[Hashable]]: + """Build per-item cache keys from request-level item metadata. + + The modality comes from each item's `item_refs` entry, so mixed-modality requests are + keyable per item. Returns `None` when the request cannot participate in the cache (missing + hashes or kwargs hash, or item counts that do not line up). + """ + if multimodal_hashes is None or kwargs_hash is None: + return None + if not (len(multimodal_hashes) == len(item_refs) == len(embedding_lengths)): + # Malformed metadata rather than a normal miss: the request loses + # cache participation silently, so say so once with the counts. + logger.warning_once( + f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping item keys because " + f"multimodal_hashes ({len(multimodal_hashes)}), item_refs " + f"({len(item_refs)}) and multimodal_embedding_lengths " + f"({len(embedding_lengths)}) counts disagree", + key="mm_encoder_cache_item_key_count_mismatch", + ) + return None + return [ + cls._encoder_cache_item_key(modality, item_hash, embedding_length, kwargs_hash) + for (modality, _), item_hash, embedding_length in zip( + item_refs, + multimodal_hashes, + embedding_lengths, + strict=True, + ) + ] + + @classmethod + def partition_encoder_cache( + cls, + param: MultimodalParams, + encoder_cache: TensorLRUCache, + item_indices: Optional[Sequence[int]] = None, + keys: Optional[Sequence[Hashable]] = None, + ) -> Optional[EncoderCachePartition]: + """Look up `param`'s items in the cache and return the per-item partition. + + Args: + item_indices: restrict the lookup to these items. `None` looks up every item, + which is what the full-request path wants. The item scheduler passes the + subset it selected for this iteration so unscheduled items keep their LRU + recency -- see `EncoderCachePartition`. + keys: precomputed per-item keys. `None` derives them from `param`, which needs + `multimodal_input` to carry the content hashes. The executor builds params + from `LlmRequest.py_multimodal_data` alone -- the hashes live on the + request, not the params -- so it derives keys there and passes them in. + + Returns `None` when the param is not keyable (missing item metadata, content + hashes or processor-kwargs hash, or a request-local embedding is already + attached); the caller should treat that as a full miss. + """ + if param.multimodal_data.get("multimodal_embedding") is not None: + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: request-local multimodal embedding present; " + "skipping persistent cache lookup" + ) + return None + + keys = list(keys) if keys is not None else cls._encoder_cache_keys(param) + if not keys: + return None + + looked_up = list(range(len(keys))) if item_indices is None else list(item_indices) + out_of_range = [i for i in looked_up if not 0 <= i < len(keys)] + if out_of_range: + raise MultimodalEncoderContractError( + f"item_indices {out_of_range} out of range for {len(keys)} cache keys" + ) + + hits: Dict[int, torch.Tensor] = {} + miss_indices: list[int] = [] + for i in looked_up: + cached = encoder_cache.get(keys[i]) + if cached is None: + miss_indices.append(i) + else: + hits[i] = cached + + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: partition hit_items={len(hits)}, " + f"miss_items={len(miss_indices)}, looked_up={len(looked_up)}, " + f"total_items={len(keys)}." + ) + return EncoderCachePartition( + hits=hits, miss_indices=miss_indices, keys=keys, looked_up=looked_up + ) + + @staticmethod + def _apply_metadata_slice( + residual: MultimodalParams, + source: MultimodalParams, + prompt_item_indices: Sequence[int], + *, + item_metadata: Optional[MultimodalItemMetadata] = None, + ) -> None: + """Overwrite `residual`'s per-item metadata to match the sliced items. + + Models slice raw modality tensors in `build_multimodal_encoder_input`; the mixin owns + the parallel per-item metadata slice so every model gets it identically. + ``prompt_item_indices`` indexes the source metadata's ``item_refs`` + across all modalities, rather than a modality's payload. + """ + metadata = item_metadata if item_metadata is not None else source.item_metadata + if metadata is None: + raise MultimodalEncoderContractError("Item selection requires multimodal item metadata") + residual.multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = metadata.select( + prompt_item_indices + ) + if residual.multimodal_input is not None and source.multimodal_input is not None: + source_hashes = source.multimodal_input.multimodal_hashes + residual.multimodal_input.multimodal_hashes = [ + source_hashes[i] for i in prompt_item_indices + ] + + def _encode_with_partial_cache( + self, + partials: Sequence[tuple[MultimodalParams, EncoderCachePartition]], + encoder_cache: TensorLRUCache, + ) -> None: + """Encode only the miss items of each partial-hit param and stitch results. + + Miss residuals from all partial-hit params in the batch are encoded in a + single call and the concatenated output is split back per param, mirroring + how `get_multimodal_embeddings` batches full-miss params. After this returns, + each param's `multimodal_embedding` has the same shape as a full encoder run + so downstream `get_multimodal_embeddings` treats it as fully cached. + """ + if not partials: + return + + # Cross-iter prefetch may have staged some params' raw MM tensors on the aux + # stream. If a prefetch encoder call then raised, the request reaches this + # iteration with an `encoder_event` but no `multimodal_embedding`; slicing + # those tensors on the main stream before the event would race the aux-stream + # H2D copy. Wait per param up front, before any raw-tensor read. + for param, _ in partials: + if param.encoder_event is not None: + torch.cuda.current_stream().wait_event(param.encoder_event) + + encoder_inputs = self.prepare_multimodal_encoder_inputs( + [(param, i) for param, partition in partials for i in partition.miss_indices] + ) + # Keep one call across requests and modalities so EncoderGroup can + # batch modalities that share an encoder. + batched_output = self.encode_multimodal_inputs([param for param, _, _ in encoder_inputs]) + output_lengths = [length for _, lengths, _ in encoder_inputs for length in lengths] + encoded_items = iter(torch.split(batched_output, output_lengths, dim=0)) + + for param, partition in partials: + miss_tensors = [next(encoded_items) for _ in partition.miss_indices] + + by_item: Dict[int, torch.Tensor] = dict(partition.hits) + for miss_idx, tensor in zip(partition.miss_indices, miss_tensors, strict=True): + by_item[miss_idx] = tensor + param.multimodal_data["multimodal_embedding"] = _assemble_multimodal_encoder_embeddings( + by_item, len(partition.keys) + ) + + inserted = 0 + rejected = 0 + for miss_idx, tensor in zip(partition.miss_indices, miss_tensors, strict=True): + if encoder_cache.put(partition.keys[miss_idx], tensor): + inserted += 1 + else: + rejected += 1 + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: partial-hit encode " + f"total_items={len(partition.keys)} " + f"hit_items={len(partition.hits)} " + f"encoded_items={len(partition.miss_indices)} " + f"cache_writes_inserted={inserted} " + f"cache_writes_rejected={rejected}" + ) + + @classmethod + def _write_encoder_cache_entries( + cls, + param: MultimodalParams, + encoder_cache: TensorLRUCache, + ) -> None: + keys = cls._encoder_cache_keys(param) + if not keys: + return + + embedding = param.multimodal_data.get("multimodal_embedding") + if isinstance(embedding, list): + embedding = torch.cat(embedding, dim=0) + param.multimodal_data["multimodal_embedding"] = embedding + if not isinstance(embedding, torch.Tensor): + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping write because no tensor embedding was " + "attached after encoder execution." + ) + return + + _, embedding_lengths = param.get_item_layout() + if sum(embedding_lengths) != embedding.shape[0]: + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping write because embedding row count " + "does not match multimodal_embedding_lengths." + ) + return + + # Encoder outputs are concatenated per params object. Splitting by item length lets future + # requests reuse matching items independently, even when their request-level item order + # differs. + inserted_entries = 0 + rejected_entries = 0 + for key, item_embedding in zip( + keys, + torch.split(embedding, embedding_lengths, dim=0), + strict=True, + ): + if encoder_cache.put(key, item_embedding): + inserted_entries += 1 + else: + rejected_entries += 1 + logger.debug( + f"{_MM_ENCODER_CACHE_LOG_NAME}: wrote {inserted_entries} item entries, " + f"rejected={rejected_entries}, rows={embedding.shape[0]}." + ) + encoder_cache.log_stats("multimodal encoder cache write.") + + @staticmethod + def _validate_embeddings( + embeddings: list[torch.Tensor], + multimodal_params: Sequence[MultimodalParams], + *, + current_chunk_only: bool = False, + ) -> None: + """Validate gathered embeddings' row count against runtime metadata. + + `current_chunk_only` matches the gather layout; it does not change the + original per-request runtime counts used for later prefill chunks. + Skipped if any param lacks `multimodal_runtime.total_embeds_in_request`, since the contract + cannot be evaluated without complete metadata. + """ + if len(embeddings) != 1: + raise ValueError( + f"MultimodalRuntimeMixin requires a single embedding tensor, got {len(embeddings)} " + "tensors." + ) + + embeddings_tensor = embeddings[0] + expected_rows = 0 + has_runtime_metadata = [] + for param in multimodal_params: + runtime = param.multimodal_runtime + has_runtime = runtime is not None and runtime.total_embeds_in_request is not None + has_runtime_metadata.append(has_runtime) + if has_runtime: + expected_rows += ( + runtime.num_mm_tokens_in_chunk + if current_chunk_only + else runtime.total_embeds_in_request + ) + + if any(has_runtime_metadata) and not all(has_runtime_metadata): + raise ValueError( + "Multimodal runtime metadata must be present for every param or none of them." + ) + if not all(has_runtime_metadata): + logger.debug( + "Skipping multimodal embedding row-count validation: runtime metadata missing " + "for all params." + ) + return + + actual_rows = embeddings_tensor.shape[0] + if actual_rows != expected_rows: + raise ValueError( + f"Multimodal embedding row count mismatch: expected {expected_rows}, got {actual_rows}." + ) + + +class MultimodalModelMixin(MultimodalRuntimeMixin): + """Composite multimodal causal LM: an encoder runtime wrapped around an inner language model. + + Modeling V1 multimodal models keep the language model in a separate module. This class adds + the delegation to that module (`language_model`, the step-loop hooks, vocabulary and + guided-decoding queries), a default `forward`, and the default text-embedding fusion + (`prepare_multimodal_inputs`). Models that run their own decoder, such as Modeling V2 + targets, inherit `MultimodalRuntimeMixin` directly instead. + """ + + supports_concurrent_encoder: ClassVar[bool] = True + """Composite models run side-stream prefetch; one whose encoder shares mutable state across + CUDA streams sets this to False.""" + + @property + def text_embedding_layer(self): + """Return the token embedding layer used by `fuse_input_embeds`.""" + raise NotImplementedError + + @property + def language_model(self) -> torch.nn.Module: + """Return the inner language model that receives prepared inputs.""" + raise NotImplementedError + + @property + def runtime_hooks(self) -> RuntimeHooks: + """Return the inner language model's step-loop hooks.""" + try: + language_model = self.language_model + except NotImplementedError: + return DEFAULT_RUNTIME_HOOKS + return runtime_hooks_of(language_model) + + @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 set_guided_decoder(self, guided_decoder: "CapturableGuidedDecoder") -> bool: + """Install a guided decoder on the inner language model. + + Returns False when the inner model does not support guided decoding, + matching the contract ModelEngine.set_guided_decoder() expects. + """ + inner = self.language_model + if not hasattr(inner, "set_guided_decoder"): + return False + return inner.set_guided_decoder(guided_decoder) + + def get_language_model_extra_forward_kwargs( + self, + *, + attn_metadata: Any, + 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. + + The attention metadata identifies context and generation requests in this batch. + """ + 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( + attn_metadata=attn_metadata, + 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 prepare_multimodal_inputs( + self, + *, + input_ids: torch.Tensor, + positions: Optional[torch.Tensor], + multimodal_params: Optional[Sequence[MultimodalParams]], + num_context_requests: int, + **forward_kwargs: Any, + ) -> PreparedLlmInputs: + """Prepare encoder rows and fuse them with the model's text embeddings.""" + prepared = self._prepare_active_multimodal_embeddings( + input_ids=input_ids, + multimodal_params=multimodal_params, + num_context_requests=num_context_requests, + **forward_kwargs, + ) + if not prepared.multimodal_embeddings and not prepared.extra_embeds: + return PreparedLlmInputs(input_ids=prepared.input_ids, inputs_embeds=None) + + fused_input_ids, inputs_embeds, fused_extra_embeds = self._fuse_multimodal_embeddings( + input_ids=prepared.input_ids, + multimodal_embeddings=list(prepared.multimodal_embeddings), + mm_token_ids=self.multimodal_token_ids, + embedding_layer=self.text_embedding_layer, + extra_embeds=prepared.extra_embeds, + # `text_token_indices` / `mm_token_indices` are pre-computed by the + # executor (see DecoderInputPreparer.prepare_inputs) and must reach + # `fuse_input_embeds` to (a) preserve the active-chunk subset + # contract when MM rows are a subset of visible MM tokens and + # (b) avoid the torch.where host sync inside + # `filter_mm_token_from_input_ids`. + text_token_indices=forward_kwargs.get("text_token_indices"), + mm_token_indices=forward_kwargs.get("mm_token_indices"), + ) + return PreparedLlmInputs( + input_ids=fused_input_ids, + inputs_embeds=inputs_embeds, + extra_embeds=fused_extra_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]]: + """Thin adapter over `fuse_input_embeds`. + + The framework does not forward `prepare_multimodal_inputs` kwargs + into `fuse_input_embeds`; only inputs the helper actually consumes + are surfaced here. Models needing to bypass token filtering should + pass pre-computed `text_token_indices`/`mm_token_indices`. + """ + if mm_token_ids is not None and not isinstance(mm_token_ids, torch.Tensor): + mm_token_ids = torch.tensor( + list(mm_token_ids), dtype=input_ids.dtype, device=input_ids.device + ) + + result = fuse_input_embeds( + embedding_layer=embedding_layer, + input_ids=input_ids, + mm_embeds=multimodal_embeddings, + mm_token_ids=mm_token_ids, + text_token_indices=text_token_indices, + mm_token_indices=mm_token_indices, + extra_embeds=list(extra_embeds) if extra_embeds else None, + ) + if len(result) == 3: + fused_input_ids, inputs_embeds, fused_extra_embeds = result + return fused_input_ids, inputs_embeds, fused_extra_embeds or () + + fused_input_ids, inputs_embeds = result + return fused_input_ids, inputs_embeds, () + + +def _collect_cross_iter_prefetch_candidates( + pending_requests: Sequence["LlmRequest"], + in_flight_request_ids: Iterable[int], + max_prefetch: int, + max_prefetch_ahead: int, +) -> list[tuple["LlmRequest", Dict[str, Any], torch.Tensor]]: + """Select cross-iteration prefetch candidates without touching CUDA. + + Returns up to `max_prefetch` `(request, multimodal_data, cumsum)` tuples + while keeping total not-in-flight prefetched requests at or below + `max_prefetch_ahead`. A new candidate must have a supported raw input + modality in `multimodal_data`, no cached `multimodal_embedding`, and a valid + `multimodal_embed_mask_cumsum`. Existing cached embeddings or pending encoder + events count against the ahead limit only when attached to a request with + real MM encoder work. + """ + if max_prefetch <= 0 or max_prefetch_ahead <= 0: + return [] + + in_flight = set(in_flight_request_ids) + outstanding_prefetches = 0 + candidates: list[tuple["LlmRequest", Dict[str, Any], torch.Tensor]] = [] + + for req in pending_requests: + if req.py_request_id in in_flight: + continue + mm_data = req.py_multimodal_data or {} + has_cached_embedding = mm_data.get("multimodal_embedding") is not None + has_raw_mm_input = any(key in mm_data for key in _MM_DATA_INPUT_MODALITY_KEYS) + if req.py_mm_encoder_event is not None: + if has_cached_embedding or has_raw_mm_input: + outstanding_prefetches += 1 + continue + if has_cached_embedding: + outstanding_prefetches += 1 + continue + if not has_raw_mm_input: + continue + cumsum = mm_data.get("multimodal_embed_mask_cumsum") + if cumsum is None: + continue + candidates.append((req, mm_data, cumsum)) + + available_slots = min(max_prefetch, max_prefetch_ahead - outstanding_prefetches) + if available_slots <= 0: + return [] + return candidates[:available_slots] + + +def _dispatch_cross_iter_prefetch( + model: "MultimodalRuntimeMixin", + candidates: Sequence[tuple["LlmRequest", Dict[str, Any], torch.Tensor]], + aux_stream: "torch.cuda.Stream", +) -> None: + """H2D-copy MM data, run the encoder, and cache embeddings on `aux_stream`. + + Stamps a CUDA event on every candidate's `py_mm_encoder_event` so the + next iteration's consume site waits on it before reading cached tensors. + The event covers all work queued in the aux-stream block, so the same + event object is shared across all candidates. + """ + encoder_cache_enabled = model.encoder_cache_active + params_list = [ + MultimodalParams( + multimodal_input=_build_request_multimodal_input(req, encoder_cache_enabled), + multimodal_data=mm_data, + multimodal_runtime=MultimodalRuntimeData( + past_seen_token_num=0, + chunk_end_pos=cumsum.numel(), + embed_mask_cumsum=cumsum, + ), + ) + for req, mm_data, cumsum in candidates + ] + + # Prefetch targets requests outside the current iteration, so their + # multimodal tensors are not touched by the main stream. The caller queues + # this after the iteration's LLM kernels so aux-stream H2D copies and + # encoder work can overlap them. + # + # Request-local ordering is handled by `encoder_event`; the consume path also `record_stream`s + # attached embeddings before gathering them so post-prefill request cleanup cannot release + # storage while main-stream work is pending. Persistent-cache clones use their own producer + # events and consumer `record_stream` calls inside `TensorLRUCache`. + encoder_event = None + try: + with _run_on_aux_stream(aux_stream) as encoder_event: + encoder_cache = model._multimodal_encoder_cache if encoder_cache_enabled else None + cache_misses: list[MultimodalParams] = [] + partial_hits: list[tuple[MultimodalParams, EncoderCachePartition]] = [] + if encoder_cache is None: + cache_misses = params_list + else: + for param in params_list: + partition = model.partition_encoder_cache(param, encoder_cache) + if partition is None or partition.is_full_miss: + cache_misses.append(param) + elif partition.is_full_hit: + param.multimodal_data["multimodal_embedding"] = ( + _assemble_multimodal_encoder_embeddings( + partition.hits, len(partition.keys) + ) + ) + else: + partial_hits.append((param, partition)) + + params_to_transfer = cache_misses + [param for param, _ in partial_hits] + for param in params_to_transfer: + param.to_device( + "multimodal_data", + "cuda", + pin_memory=prefer_pinned(), + target_keywords=getattr(model, "multimodal_data_device_paths", None), + ) + # `to_device` may replace `multimodal_data` with a new dict; reattach + # the (possibly new) dict to each request so the next iteration's + # `_prepare_inputs` sees the cached embedding stamped below. Mirrors + # the reassignment at the canonical to_device call site in + # DecoderInputPreparer.prepare_inputs. + for (req, _, _), p in zip(candidates, params_list): + req.py_multimodal_data = p.multimodal_data + + # Prefetch runs between executor steps, outside `ModelCaller`, so the vision tower's + # TP all-reduces can return NCCL window tensors here. The stored embeddings are read + # next iteration, so keep a clone rather than a view of a released window. + with discard_nccl_window_tensor_outputs(params_to_transfer): + if partial_hits and encoder_cache is not None: + model._encode_with_partial_cache(partial_hits, encoder_cache) + if cache_misses: + encoder_output = model.encode_multimodal_inputs(cache_misses).clone() + _store_chunked_prefill_embeddings(cache_misses, [encoder_output]) + if encoder_cache is not None: + for param in cache_misses: + model._write_encoder_cache_entries(param, encoder_cache) + + # Prefetch only needs to attach each request's embedding. Validate each request + # independently instead of gathering the unused batch output with `torch.cat`. + for param in params_list: + embedding = param.multimodal_data.get("multimodal_embedding") + if not isinstance(embedding, torch.Tensor): + raise ValueError("Multimodal encoder prefetch did not produce an embedding.") + model._validate_embeddings([embedding], [param]) + finally: + # Stash the event on every candidate's durable LlmRequest (not the + # per-iter `MultimodalParams`), since `_prepare_inputs` rebuilds the + # wrapper each iteration. The transfer to `MultimodalParams.encoder_event` + # happens in `_prepare_inputs` when the request is next scheduled. + # + # This runs in `finally` (and `_run_on_aux_stream` records the event in + # its own `finally`) so that on partial failure -- e.g. `to_device` + # mutated `req.py_multimodal_data` in place to aux-stream CUDA tensors, + # then `encode_multimodal_inputs` raised -- the consumer still has an + # event to wait on before reading those tensors on the main stream. + # Without this, the request would carry aux-stream tensors with no sync + # point, producing a cross-stream data race in the next iteration's + # in-iter encode path. + if encoder_event is not None: + for req, _, _ in candidates: + req.py_mm_encoder_event = encoder_event + + +def maybe_prefetch_mm_encoder_for_next_iter( + model: Any, + pending_requests: Sequence["LlmRequest"], + in_flight_request_ids: Iterable[int] = (), + max_prefetch: int = 1, + max_prefetch_ahead: Optional[int] = None, +) -> int: + """Speculative cross-iteration MM encoder prefetch on a side CUDA stream. + + For up to `max_prefetch` `pending_requests`, subject to the outstanding ahead cap, runs + `model.encode_multimodal_inputs` on a side CUDA stream. + + The resulting embeddings are written into `request.py_multimodal_data` so the next iteration's + `_prepare_inputs` picks them up via the standard cache path, and a CUDA event is stamped on + `request.py_mm_encoder_event` for `_prepare_inputs` to transfer onto the new `MultimodalParams`. + + The mixin consume sites (e.g. `get_multimodal_embeddings)` need to wait on the event before + reading the cached tensors. + + While the current iteration's LLM kernels run on the main stream, this queues encoder work for + an "admit-likely context" request on the aux stream. + + Mis-predictions waste GPU time, but cached embeddings remain valid until the request is admitted + or terminated. Examples include: + + - The prefetched request can be terminated before admission (client cancel / disconnect, + timeout, validation failure, ...). + - With `max_prefetch < len(pending)`, if the head is bumped by budget reasons, the next-admitted + request is one we did not prefetch. + + Gated by `MultimodalConfig.encoder_side_stream_max_ahead`: 0 disables the side stream; a + positive integer enables it and caps the total number of not-in-flight requests with + prefetched MM encoder work. + + Returns the number of requests for which an encoder kick-off was queued. + """ + if not isinstance(model, MultimodalRuntimeMixin): + return 0 + if max_prefetch <= 0: + return 0 + if max_prefetch_ahead is None: + max_prefetch_ahead = 0 + if max_prefetch_ahead <= 0: + return 0 + + aux_stream = _get_mm_aux_stream(max_prefetch_ahead) + if aux_stream is None: + return 0 + + candidates = _collect_cross_iter_prefetch_candidates( + pending_requests, in_flight_request_ids, max_prefetch, max_prefetch_ahead + ) + if not candidates: + return 0 + _dispatch_cross_iter_prefetch(model, candidates, aux_stream) + return len(candidates) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b35ed3fc40e5..0b67bb8cb0cf 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -50,8 +50,8 @@ create_kv_cache_transceiver) from ..hostfunc import set_low_latency_dispatch from ..model_config import ModelConfig, _get_kv_cache_layer_specs -from ..models.modeling_multimodal_mixin import MultimodalModelMixin from ..modules.dwdp.dense.manager import DenseDwdpManager +from ..multimodal.model import MultimodalRuntimeMixin from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from ..speculative import (draft_prompt_lookahead, get_draft_block_width, get_num_extra_kv_tokens, get_num_spec_layers, @@ -1550,8 +1550,14 @@ def _create_dummy_context_requests( return requests def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]: - """Build one processed MM encoder batch at its scheduling limits.""" - if not isinstance(self._model_engine.model, MultimodalModelMixin): + """Build one processed MM encoder batch at its scheduling limits. + + Items stay within the model's `max_encoder_tokens_per_item`. When that bound is below + the processor's largest item, the batch holds one params object per item, each built + under the bound, so the batch still fills the encoder token budget. + """ + model = self._model_engine.model + if not isinstance(model, MultimodalRuntimeMixin): return [] if isinstance( self._profiling_stage_data, @@ -1579,6 +1585,19 @@ def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]: raise ValueError( "Multimodal encoder token counts must be positive; " f"got {num_tokens} for {modality}") + # Modality -> per-item token budget under the model's encoder bound. + item_token_budgets: dict[str, int] = {} + for modality, num_tokens in list(max_tokens_per_item.items()): + bound = model.max_encoder_tokens_per_item(modality) + if bound is None or num_tokens <= bound: + continue + bounded_tokens = input_processor.get_mm_max_tokens_per_item( + max_num_encoder_tokens=bound).get(modality) + if bounded_tokens is None or bounded_tokens <= 0: + del max_tokens_per_item[modality] + continue + max_tokens_per_item[modality] = bounded_tokens + item_token_budgets[modality] = bound max_tokens_per_item = { modality: num_tokens for modality, num_tokens in max_tokens_per_item.items() @@ -1595,23 +1614,37 @@ def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]: self._model_engine.encoder_batch_size, encoder_max_num_tokens // num_tokens_per_item, ) - mm_data = input_processor.get_dummy_mm_data( - max_num_encoder_tokens=encoder_max_num_tokens, - mm_counts={modality: num_items}, - dtype=self._model_engine.model.dtype, - ) + if modality in item_token_budgets: + mm_data_list = [ + input_processor.get_dummy_mm_data( + max_num_encoder_tokens=item_token_budgets[modality], + mm_counts={modality: 1}, + dtype=model.dtype, + ) for _ in range(num_items) + ] + else: + mm_data_list = [ + input_processor.get_dummy_mm_data( + max_num_encoder_tokens=encoder_max_num_tokens, + mm_counts={modality: num_items}, + dtype=model.dtype, + ) + ] except NotImplementedError: logger.info("Multimodal memory profiling skipped: " f"{type(input_processor).__name__} does not implement " "get_dummy_mm_data().") return [] - if not mm_data: + if not all(mm_data_list): return [] - if not isinstance(mm_data, dict): + if not all(isinstance(mm_data, dict) for mm_data in mm_data_list): raise ValueError( "get_dummy_mm_data() must return a multimodal_data " "dictionary") - return [MultimodalParams(multimodal_data=mm_data)] + return [ + MultimodalParams(multimodal_data=mm_data) + for mm_data in mm_data_list + ] def _encode_dummy_inputs(self) -> Optional[torch.Tensor]: """Run the full-budget MM encoder and retain equivalent store output.""" diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index 67a9b2d017e6..fb78b7db32ec 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -16,11 +16,11 @@ from torch import nn from tensorrt_llm._torch.models.modeling_multimodal_encoder import MultimodalEncoderMixin -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( +from tensorrt_llm._torch.multimodal.embeddings import _join_embeddings +from tensorrt_llm._torch.multimodal.model import ( MultimodalEncoderContractError, - MultimodalModelMixin, + MultimodalRuntimeMixin, ) -from tensorrt_llm._torch.models.modeling_multimodal_utils import _join_embeddings from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm._utils import prefer_pinned @@ -46,11 +46,34 @@ from ..llm_request import LlmRequest, MultimodalEncoderRequestError, get_mm_items_for_chunk +def bound_mm_max_tokens_per_item( + model: MultimodalRuntimeMixin, max_tokens_per_item: Mapping[str, int] +) -> dict[str, int]: + """Clamp the input processor's largest item per modality to the model's encoder bound.""" + bounded = {} + for modality, num_tokens in max_tokens_per_item.items(): + bound = model.max_encoder_tokens_per_item(modality) + bounded[modality] = num_tokens if bound is None else min(num_tokens, bound) + return bounded + + def resolve_mm_encoder_token_budget(base_budget: int, model_max_atomic_item_tokens: int) -> int: """Keep the model's largest indivisible MM item schedulable.""" return max(base_budget, model_max_atomic_item_tokens) +def validate_mm_encoder_prefetch_compatibility(llm_args: TorchLlmArgs, model: nn.Module) -> None: + """Reject side-stream prefetch for a model whose encoder cannot run on a second stream.""" + if llm_args.multimodal_config.encoder_side_stream_max_ahead == 0: + return + if isinstance(model, MultimodalRuntimeMixin) and not model.supports_concurrent_encoder: + raise ValueError( + f"{type(model).__name__} does not support MM encoder side-stream prefetch " + "(multimodal_config.encoder_side_stream_max_ahead > 0): its encoder does not " + "declare supports_concurrent_encoder; set encoder_side_stream_max_ahead=0" + ) + + def validate_mm_encoder_scheduling_compatibility( llm_args: TorchLlmArgs, item_scheduling_enabled: bool ) -> None: @@ -89,17 +112,17 @@ def validate_mm_encoder_scheduling_compatibility( def is_multimodal(model: nn.Module, input_processor: InputProcessor) -> bool: """True iff this engine drives a multimodal model. - Primary signal: ``MultimodalModelMixin`` is the canonical marker -- + Primary signal: ``MultimodalRuntimeMixin`` is the canonical marker -- multimodal LM classes inherit from it. Until every model has migrated (Mistral done; Qwen-VL, Nemotron, Gemma, Phi-4-MM, etc. pending), fall back to whether the input processor subclasses ``BaseMultimodalInputProcessor``, which every multimodal model necessarily provides at the data boundary. TODO(TRTLLM-13542): Once all multimodal models inherit - ``MultimodalModelMixin``, drop the input-processor fallback so the model + ``MultimodalRuntimeMixin``, drop the input-processor fallback so the model class itself is the single source of truth. """ - if isinstance(model, MultimodalModelMixin): + if isinstance(model, MultimodalRuntimeMixin): return True return isinstance(input_processor, BaseMultimodalInputProcessor) @@ -107,7 +130,7 @@ class itself is the single source of truth. def mm_item_scheduling_enabled(llm_args: TorchLlmArgs, model: nn.Module) -> bool: """Whether the item-scheduling wiring is engaged this run. - Item scheduling is declared once, as a model capability (the ``MultimodalModelMixin`` + Item scheduling is declared once, as a model capability (the ``MultimodalRuntimeMixin`` ClassVar). This is the actionable flag derived from it: the engine stores it because the executor and the scheduler wrap both read it back off the engine, and the setup below, the scheduler wrap and the executor encoder step must all agree. Both @@ -127,7 +150,7 @@ def mm_item_scheduling_enabled(llm_args: TorchLlmArgs, model: nn.Module) -> bool ) enabled = ( not llm_args.disable_mm_encoder - and isinstance(model, MultimodalModelMixin) + and isinstance(model, MultimodalRuntimeMixin) and model.supports_mm_encoder_item_scheduling and policy != MultimodalEncoderSchedulingPolicy.DISABLED ) @@ -159,7 +182,7 @@ def use_mrope(model: nn.Module) -> bool: def mm_encoder_cache_enabled(model: nn.Module) -> bool: """Whether the multimodal encoder cache is active for this model.""" - return isinstance(model, MultimodalModelMixin) and model.encoder_cache_active + return isinstance(model, MultimodalRuntimeMixin) and model.encoder_cache_active def setup_mm_encoder_cuda_graphs( @@ -241,7 +264,7 @@ def setup_mm_encoder_attn_metadata( module.set_attn_max_seq_len(encoder_max_seq_len) -def resolve_bytes_per_mm_encoder_embedding(model: MultimodalModelMixin) -> int: +def resolve_bytes_per_mm_encoder_embedding(model: MultimodalRuntimeMixin) -> int: """Bytes occupied by one multimodal encoder output embedding. Prefers the mixin's explicit ``embedding_dim``/``embedding_dtype`` contract, then the @@ -283,7 +306,7 @@ def resolve_bytes_per_mm_encoder_embedding(model: MultimodalModelMixin) -> int: def resolve_mm_encoder_output_budget( input_processor: BaseMultimodalDummyInputsBuilder, encoder_max_num_tokens: int, - model: MultimodalModelMixin, + model: MultimodalRuntimeMixin, ) -> tuple[int, int]: """Resolve ``(output_budget_bytes, bytes_per_embedding)`` for resident MM encoder outputs. @@ -334,7 +357,7 @@ class MultimodalItemScheduler: def __init__( self, *, - model: MultimodalModelMixin, + model: MultimodalRuntimeMixin, input_processor: BaseMultimodalDummyInputsBuilder | None = None, attention_metadata_capacity: dict[str, int] | None = None, encoder_max_num_tokens: int | None = None, @@ -371,7 +394,7 @@ def maybe_create( validate_mm_encoder_scheduling_compatibility(llm_args, enabled) if not enabled: return None - assert isinstance(model, MultimodalModelMixin) # narrowed by the predicate + assert isinstance(model, MultimodalRuntimeMixin) # narrowed by the predicate return cls.create( model=model, input_processor=input_processor, @@ -383,7 +406,7 @@ def maybe_create( def create( cls, *, - model: MultimodalModelMixin, + model: MultimodalRuntimeMixin, input_processor: BaseMultimodalDummyInputsBuilder, encoder_max_num_tokens: int | None, configured_encoder_max_num_tokens: int | None, @@ -422,6 +445,7 @@ def create( ) if any(value <= 0 for value in max_tokens_per_item.values()): raise ValueError("get_mm_max_tokens_per_item() must return positive token counts") + max_tokens_per_item = bound_mm_max_tokens_per_item(model, max_tokens_per_item) model_max_atomic_item_tokens = max(max_tokens_per_item.values()) encoder_token_budget_base = encoder_max_num_tokens effective_encoder_token_budget = resolve_mm_encoder_token_budget( @@ -513,8 +537,8 @@ def forward_items( """Encode selected producer items into their reserved cache entries.""" if not scheduled_items: return - if not isinstance(self.model, MultimodalModelMixin): - raise TypeError("Item-level MM scheduling requires MultimodalModelMixin") + if not isinstance(self.model, MultimodalRuntimeMixin): + raise TypeError("Item-level MM scheduling requires MultimodalRuntimeMixin") encoder_cache = self.encoder_cache if encoder_cache is None: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py index 98dce6b41550..7eba998511f6 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/common.py @@ -11,7 +11,7 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.distributed import Distributed -from tensorrt_llm._torch.models.modeling_multimodal_utils import filter_mm_token_from_input_ids +from tensorrt_llm._torch.multimodal.embeddings import filter_mm_token_from_input_ids from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.utils import set_per_request_prefill_cuda_graph_flag from tensorrt_llm._utils import maybe_pin_memory @@ -196,7 +196,7 @@ def prepare_multimodal_indices( input_ids = torch.tensor(input_ids, dtype=torch.int, device="cpu") vocab_size = model.config.vocab_size # `multimodal_token_ids` is the common wrapper-model contract. Keep the legacy name as a - # fallback for models not yet migrated to `MultimodalModelMixin`. + # fallback for models not yet migrated to `MultimodalRuntimeMixin`. mm_token_ids = getattr(model, "multimodal_token_ids", None) if mm_token_ids is None: mm_token_ids = getattr(model, "mm_token_ids", None) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index df05cc46e71e..20651c8f63fb 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -16,7 +16,7 @@ from tensorrt_llm._torch.distributed import Distributed from tensorrt_llm._torch.kernels.spec_step_copies import op as spec_step_copies from tensorrt_llm._torch.metadata import KVCacheParams -from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input +from tensorrt_llm._torch.multimodal.model import _build_request_multimodal_input from tensorrt_llm._torch.pyexecutor import step_inputs from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/mm_encoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/mm_encoder.py index 28a529b3c199..c699fa144d02 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/mm_encoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/mm_encoder.py @@ -10,11 +10,11 @@ from torch import nn from typing_extensions import Self -from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( MoeLoadBalancer, MoeLoadBalancerIterContext, ) +from tensorrt_llm._torch.multimodal.model import _build_request_multimodal_input from tensorrt_llm._torch.nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from tensorrt_llm._torch.pyexecutor.llm_request import get_multimodal_embedding_lengths from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 514e9efd7399..b9209170bb73 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -30,10 +30,10 @@ from ..compilation.backend import Backend from ..distributed import Distributed from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader -from ..models.modeling_multimodal_mixin import MultimodalModelMixin from ..models.modeling_utils import DecoderModelForCausalLM, timing_metric from ..moe.fused_moe.moe_comm_timeout_guard import set_moe_comm_warmup from ..moe.fused_moe.moe_load_balancer import BaseMoeLoadBalancer +from ..multimodal.model import MultimodalModelMixin, MultimodalRuntimeMixin from ..route_capture import ROUTE_CAPTURE_ATTR, RouteCapture from ..speculative import SpecMetadata, update_spec_config_from_loaded_model from ..tensor_lru_cache import TensorLRUCache @@ -59,6 +59,7 @@ setup_mm_encoder_attn_metadata, setup_mm_encoder_cuda_graphs, use_mrope, + validate_mm_encoder_prefetch_compatibility, ) from .engine.runners import resolve_runner_type from .engine.runners.common import make_scheduled_inputs, resolve_mrope_position_deltas_cache @@ -349,6 +350,7 @@ def _init_from_model(self) -> None: # Every MM item-scheduling decision -- policy, capability, feature # validation, budget resolution -- lives in engine/multimodal.py; the # engine only copies back the three budgets that are external contract. + validate_mm_encoder_prefetch_compatibility(self.llm_args, self.model) mm_item_scheduler = MultimodalItemScheduler.maybe_create( llm_args=self.llm_args, model=self.model, @@ -366,7 +368,7 @@ def _init_from_model(self) -> None: self.encoder_max_num_tokens = mm_item_scheduler.encoder_max_num_tokens self.mm_encoder_output_budget_bytes = mm_item_scheduler.output_budget_bytes self.bytes_per_mm_encoder_embedding = mm_item_scheduler.bytes_per_embedding - if isinstance(self.model, MultimodalModelMixin) and self.mapping.is_first_pp_rank(): + if isinstance(self.model, MultimodalRuntimeMixin) and self.mapping.is_first_pp_rank(): multimodal_config = self.model.model_config.multimodal_config reuse_capacity_bytes = ( multimodal_config.encoder_cache_max_bytes if self.model.encoder_cache_active else 0 @@ -889,6 +891,17 @@ def _validate_breakable_cuda_graph_compatibility(self) -> None: if self.llm_args.prefill_cuda_graph_backend != PrefillCudaGraphBackend.BREAKABLE: return + if isinstance(self.model, MultimodalRuntimeMixin) and not isinstance( + self.model, MultimodalModelMixin + ): + # A model that runs its own decoder states where its encoder runs; composite + # models keep the structural check below. + if self.model.allows_breakable_prefill_cuda_graph(self.llm_args): + return + raise ValueError( + f"breakable prefill CUDA graph is not supported by {type(self.model).__name__}: " + "its multimodal encoder may run inside the captured prefill forward" + ) if isinstance(self.model, DecoderModelForCausalLM): return decoder_model = getattr(self.model, "llm", None) @@ -912,7 +925,7 @@ def forward_multimodal_encoder_items( if not scheduled_items: return if self._mm_item_scheduler is None: - raise TypeError("Item-level MM scheduling requires MultimodalModelMixin") + raise TypeError("Item-level MM scheduling requires MultimodalRuntimeMixin") self._mm_item_scheduler.forward_items(requests, scheduled_items) @property @@ -920,7 +933,7 @@ def mm_encoder_cache(self) -> TensorLRUCache | None: """Return the model-owned encoder-output cache used by this engine.""" # Whole-model torch.compile wraps the model after its cache is created. model = getattr(self.model, "_orig_mod", self.model) - if not isinstance(model, MultimodalModelMixin): + if not isinstance(model, MultimodalRuntimeMixin): return None if self._mm_item_scheduler is None and not model.encoder_cache_active: return None diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index b5ac7b67dbbf..c2bfbcd5f91a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -54,7 +54,6 @@ from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader from ..models.checkpoints.checkpoint_catalog import CheckpointCatalog from ..models.checkpoints.weight_load_plan import WeightLoadPlan -from ..models.modeling_multimodal_mixin import MultimodalModelMixin from ..models.modeling_utils import (DecoderModelForCausalLM, MetaInitMode, get_registered_model_class, timing_metric) from ..modules.embedding import Embedding, LMHead @@ -62,6 +61,7 @@ from ..moe.fused_moe.moe_load_balancer import (BaseMoeLoadBalancer, maybe_create_moe_load_balancer) from ..moe.fused_moe.weight_owner import is_moe_weight_owner +from ..multimodal.model import MultimodalModelMixin, MultimodalRuntimeMixin from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope from .config_utils import (is_hybrid_linear, is_minimax_m3, is_mla, @@ -1274,6 +1274,11 @@ def load( validate_lm_head_output_dtype(model, self.llm_args.lm_head_output_dtype) + # A model that runs its own decoder has no composite fallbacks for its declared + # embedding row layout, so it must read correctly from the real configuration. + if isinstance(model, MultimodalRuntimeMixin) and not isinstance( + model, MultimodalModelMixin): + model.check_multimodal_instance_contract() requires_standard_hf_loading = any( getattr(module, "_requires_standard_hf_loading", False) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c96ca67f4ac9..d8a8fced6f36 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -73,12 +73,11 @@ from ..disaggregation.orchestration.transfer_manager import AsyncTransferManager from ..distributed import Distributed from ..distributed.communicator import ReduceOp -from ..models.modeling_multimodal_mixin import \ - maybe_prefetch_mm_encoder_for_next_iter from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.decoder_layer import DecoderLayer from ..modules.dwdp.dense.manager import DenseDwdpManager from ..moe.expert_statistic import ExpertStatistic +from ..multimodal.model import maybe_prefetch_mm_encoder_for_next_iter from ..speculative.drafter import Drafter from ..speculative.spec_sampler_base import SpecSampler from ..speculative.speculation_gate import SpeculationGate @@ -10122,7 +10121,7 @@ def _maybe_prefetch_next_iter_mm_encoders( work can overlap current-iteration sampling in the non-overlap loop and previous-batch `_update_requests` in the overlap loop. No-op unless `multimodal_config.encoder_side_stream_max_ahead` is positive and the - model is a `MultimodalModelMixin` subclass. + model is a `MultimodalRuntimeMixin` subclass. Walks `active_requests` for context-init candidates that are NOT in the just-scheduled batch (and, in overlap mode, not in the diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index a5c816045ec2..4e029eb56d88 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -677,7 +677,7 @@ def unselect( class MultimodalScheduler(RequestScheduler): """Add per-item MM encoder limits and cache lookup to the LLM scheduler. - The wrapper is constructed only for ``MultimodalModelMixin`` models. It + The wrapper is constructed only for ``MultimodalRuntimeMixin`` models. It deliberately reuses the wrapped scheduler's capacity and microbatch schedulers so MM encoder costs never enter the LLM token budget. diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 875598f6e366..d57d12a7ac8c 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -102,6 +102,7 @@ l0_cpu: - unittest/_torch/peft - unittest/_torch/memory - unittest/_torch/modeling + - unittest/_torch/multimodal/test_prepared_multimodal_embeddings.py # modeling_v2 routing and routing-table checks: config shapes and a monkeypatched # compute capability, no device. The default targets' shard tests: small CPU # tensors and the CPU block_scale_interleave. The rest of unittest/_torch/modeling_v2 @@ -114,6 +115,7 @@ l0_cpu: - unittest/_torch/modeling_v2/test_modeling_v2_reference_precision.py - unittest/_torch/modeling_v2/test_modeling_v2_fp4_test_helpers.py - unittest/_torch/modeling_v2/test_modeling_v2_routing.py + - unittest/_torch/modeling_v2/test_modeling_v2_multimodal_contract.py - unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_equality.py - unittest/_torch/modeling_v2/comm/test_modeling_v2_rank_job.py # The GPU stages run with -m "not cpu_only", so this entry is what runs it. diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py b/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py index 7360ad670233..0a4956d3d79a 100644 --- a/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py +++ b/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py @@ -23,7 +23,7 @@ from tensorrt_llm._torch.models.modeling_deepseekv41_multimodal import ( DeepseekV41ForConditionalGeneration, ) -from tensorrt_llm._torch.models.modeling_multimodal_mixin import PreparedLlmInputs +from tensorrt_llm._torch.multimodal.model import PreparedLlmInputs from tensorrt_llm.inputs import create_input_processor_with_hash pytestmark = pytest.mark.cpu_only diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index abdf6689c333..6147cde46ed4 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -77,13 +77,13 @@ Gemma4MultimodalEmbedder, Gemma4MultimodalModelBase, ) -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin # noqa: E402 from tensorrt_llm._torch.models.modeling_multimodal_utils import ( # noqa: E402 VisionNormalization, find_input_mm_embeds, get_multimodal_embeddings, vision_normalization_from_processor, ) +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin # noqa: E402 from tensorrt_llm.inputs.multimodal import ( # noqa: E402 MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalInput, diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index f99cb3486a8d..dfc63d44815f 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -382,7 +382,7 @@ def test_profiling_dummy_has_encoder_item_metadata(num_images: int) -> None: Glm5NextInputProcessor, _glm5_next_build_batched_input, ) - from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( + from tensorrt_llm._torch.multimodal.model import ( EncoderGroup, MultimodalModelMixin, encode_multimodal_by_groups, @@ -400,6 +400,8 @@ def test_profiling_dummy_has_encoder_item_metadata(num_images: int) -> None: ) processor._processor = SimpleNamespace(image_processor=image_processor) model = Mock(spec=MultimodalModelMixin, dtype=torch.float16) + # The GLM-5 encoder declares no item bound of its own. + model.max_encoder_tokens_per_item.return_value = None profiler = object.__new__(KvCacheCreator) profiler._llm_args = SimpleNamespace(disable_mm_encoder=False) profiler._profiling_stage_data = {"enable_mm_reqs": True} diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py index 2ae878887807..f2713b10f012 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py @@ -1026,10 +1026,7 @@ def _make_encoder_fn(self, model): which spreads the pack dict as kwargs — so the wrappers must accept `multimodal_params` by keyword. """ - from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( - EncoderGroup, - encode_multimodal_by_groups, - ) + from tensorrt_llm._torch.multimodal.model import EncoderGroup, encode_multimodal_by_groups def _pack(params): return {"multimodal_params": params} diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index 5bea1f78846f..7882f04713ce 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -2021,7 +2021,7 @@ def _normalization_transport_spec(): @pytest.mark.parametrize("declared", [False, True]) @pytest.mark.parametrize("modality", ["image", "video"]) def test_normalization_transport_item_selection(declared: bool, modality: str) -> None: - from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin + from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalItemMetadata, @@ -2171,7 +2171,7 @@ def test_normalization_transport_shared_handles_and_process(producer_dtype) -> N @pytest.mark.cpu_only def test_normalization_transport_partial_encoder_cache() -> None: - from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin + from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 3e04985b0a00..031afaecbbd1 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -14,13 +14,13 @@ import pytest import torch -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( +from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLModelBase +from tensorrt_llm._torch.multimodal.model import ( EncoderGroup, MultimodalEncoderContractError, MultimodalModelMixin, encode_multimodal_by_groups, ) -from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLModelBase from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, @@ -876,7 +876,7 @@ def test_items_from_multiple_requests_stay_separated(): def test_request_cache_input_is_built_once_and_shared_with_forward_views(): - from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input + from tensorrt_llm._torch.multimodal.model import _build_request_multimodal_input class Request: def __init__(self): diff --git a/tests/unittest/_torch/modeling/test_musespark_open_omni_textpath.py b/tests/unittest/_torch/modeling/test_musespark_open_omni_textpath.py index 69ca486a1f28..8e87b63a678c 100644 --- a/tests/unittest/_torch/modeling/test_musespark_open_omni_textpath.py +++ b/tests/unittest/_torch/modeling/test_musespark_open_omni_textpath.py @@ -1284,12 +1284,10 @@ def test_the_wrapper_satisfies_the_encoder_profiling_gate(): hands the tower's whole activation footprint to the KV cache. It is an isinstance test, which no amount of duck typing satisfies. """ - from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( # noqa: PLC0415 - MultimodalModelMixin, - ) from tensorrt_llm._torch.models.modeling_musespark_open_omni import ( # noqa: PLC0415 MuseSparkForConditionalGeneration, ) + from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin # noqa: PLC0415 assert issubclass(MuseSparkForConditionalGeneration, MultimodalModelMixin) # The profiler calls this one directly, and the mixin's body raises @@ -1305,12 +1303,10 @@ def test_the_wrapper_uses_the_mixins_forward(): A local `forward` would fork the merge from the mixin's, silently. """ - from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( # noqa: PLC0415 - MultimodalModelMixin, - ) from tensorrt_llm._torch.models.modeling_musespark_open_omni import ( # noqa: PLC0415 MuseSparkForConditionalGeneration, ) + from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin # noqa: PLC0415 assert "forward" not in vars(MuseSparkForConditionalGeneration) assert MuseSparkForConditionalGeneration.forward is MultimodalModelMixin.forward diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_multimodal_runtime.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_multimodal_runtime.py new file mode 100644 index 000000000000..77570eb4be88 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_multimodal_runtime.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The real V2 shell and vision tower through shared embedding preparation.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import ( + check_target_multimodal_contract, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( + modeling as M, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( + vision as V, +) +from tensorrt_llm._torch.configs import Qwen4ExpConfig, Qwen4ExpVisionConfig +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin, MultimodalRuntimeMixin +from tensorrt_llm._torch.pyexecutor.model_engine import _spec_decode_fast_paths +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks, runtime_hooks_of +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalInput, + MultimodalItemMetadata, + MultimodalParams, + MultimodalRuntimeData, +) +from tensorrt_llm.llmapi.llm_args import MultimodalConfig + +pytestmark = pytest.mark.cpu_only +Shell = M.ModelingV2Qwen38FlashNextNvfp4Sm103Default + + +def make_shell(monkeypatch, cache_bytes=4096): + config = Qwen4ExpVisionConfig( + depth=0, + hidden_size=2, + num_heads=1, + in_channels=1, + patch_size=1, + temporal_patch_size=1, + spatial_merge_size=2, + num_position_embeddings=4, + out_hidden_size=2, + ) + weights = { + "embed": torch.zeros(16, 2), + "vis_patch_w": torch.zeros(2, 1, 1, 1, 1), + "vis_patch_b": torch.zeros(2), + "vis_pos": torch.zeros(4, 2), + "vis_m_n_w": torch.ones(2), + "vis_m_n_b": torch.zeros(2), + "vis_m_fc1_w": torch.zeros(8, 8), + "vis_m_fc1_b": torch.zeros(8), + "vis_m_fc2_w": torch.zeros(2, 8), + "vis_m_fc2_b": torch.zeros(2), + } + tower = V.VisionTower(weights, config) + calls = [] + + def run(pixels, grid): + calls.append(tuple(float(value) for value in pixels[:, 0].unique())) + return pixels[::4].expand(-1, 2).clone() + + monkeypatch.setattr(tower, "_run", run) + shell = Shell.__new__(Shell) + torch.nn.Module.__init__(shell) + shell.model_config = ModelConfig( + pretrained_config=Qwen4ExpConfig( + text_config={"hidden_size": 2, "vocab_size": 16}, + vision_config=config.to_dict(), + dtype="float32", + ), + multimodal_config=MultimodalConfig(encoder_cache_max_bytes=cache_bytes), + ) + shell.vision_tower = tower + shell._mm_token_ids_value = (7, 8) + shell._mm_token_ids_cache = None + shell._initialize_multimodal_encoder_cache(cache_bytes) + return shell, calls + + +def make_params(markers, lengths, start=0, end=None, kwargs_hash="same"): + total = sum(lengths) + positions = [] + offset = 0 + for length in lengths: + positions.append(offset) + offset += length + return MultimodalParams( + multimodal_input=MultimodalInput( + multimodal_hashes=[[marker] * 8 for marker in markers], + multimodal_positions=positions, + multimodal_lengths=lengths, + ), + multimodal_data={ + "image": { + "pixel_values": torch.cat( + [ + torch.full((4 * length, 1), float(marker)) + for marker, length in zip(markers, lengths, strict=True) + ] + ), + "image_grid_thw": torch.tensor([[1, 2, 2 * length] for length in lengths]), + }, + "mm_processor_kwargs_hash": kwargs_hash, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + item_refs=[("image", i) for i in range(len(markers))], + encoder_token_lengths=[4 * length for length in lengths], + output_embedding_lengths=lengths, + ), + }, + multimodal_runtime=MultimodalRuntimeData( + past_seen_token_num=start, + chunk_end_pos=total if end is None else end, + embed_mask_cumsum=torch.arange(1, total + 1), + ), + ) + + +def prepare(shell, params): + count = params.multimodal_runtime.num_mm_tokens_in_chunk + ids = torch.full((count,), 7, dtype=torch.int32) + prepared = shell.prepare_multimodal_embeddings( + input_ids=ids, + multimodal_params=[params], + num_context_requests=1, + mm_token_indices=torch.arange(count), + ) + assert prepared.input_ids is ids + assert not prepared.extra_embeds + return V.multimodal_embeds(prepared) + + +def test_chunks_and_repeated_requests_share_encoder_outputs(monkeypatch): + shell, calls = make_shell(monkeypatch) + params = make_params([10], [6], end=2) + torch.testing.assert_close(prepare(shell, params)[0], torch.full((2, 2), 10.0)) + params.multimodal_runtime = make_params([10], [6], start=2, end=5).multimodal_runtime + torch.testing.assert_close(prepare(shell, params)[0], torch.full((3, 2), 10.0)) + torch.testing.assert_close(prepare(shell, make_params([10], [6]))[0], torch.full((6, 2), 10.0)) + assert calls == [(10.0,)] + + +def test_partial_hits_select_only_missing_items_and_restore_prompt_order(monkeypatch): + shell, calls = make_shell(monkeypatch) + prepare(shell, make_params([10], [2])) + mixed = make_params([10, 20], [2, 2]) + rows = prepare(shell, mixed)[0] + torch.testing.assert_close(rows[:, 0], torch.tensor([10.0, 10.0, 20.0, 20.0])) + reordered = prepare(shell, make_params([20, 10], [2, 2]))[0] + torch.testing.assert_close(reordered[:, 0], torch.tensor([20.0, 20.0, 10.0, 10.0])) + assert calls == [(10.0,), (20.0,)] + assert mixed.multimodal_data["image"]["pixel_values"].shape == (16, 1) + + +@pytest.mark.parametrize("cache_bytes", [0, 8]) +def test_disabled_or_undersized_persistent_cache_keeps_chunk_reuse(monkeypatch, cache_bytes): + shell, calls = make_shell(monkeypatch, cache_bytes) + params = make_params([10], [6], end=2) + prepare(shell, params) + params.multimodal_runtime = make_params([10], [6], start=2, end=4).multimodal_runtime + prepare(shell, params) + assert len(calls) == 1 + prepare(shell, make_params([10], [6])) + assert len(calls) == 2 + + +def test_processor_kwargs_are_part_of_persistent_cache_key(monkeypatch): + shell, calls = make_shell(monkeypatch) + prepare(shell, make_params([10], [2])) + prepare(shell, make_params([10], [2], kwargs_hash="different")) + assert len(calls) == 2 + + +def test_position_only_params_are_filtered_and_attached_embeddings_are_consumed(monkeypatch): + shell, calls = make_shell(monkeypatch) + positions = MultimodalParams(multimodal_data={"mrope_config": {}}) + attached = MultimodalParams(multimodal_data={"multimodal_embedding": torch.ones(2, 2)}) + chosen = shell.select_multimodal_params([positions, attached], 2) + assert len(chosen) == 1 and chosen[0] is attached + result = shell.prepare_multimodal_embeddings( + input_ids=torch.tensor([7, 7]), + multimodal_params=[positions, attached], + num_context_requests=2, + ) + torch.testing.assert_close(result.multimodal_embeddings[0], torch.ones(2, 2)) + assert not calls + + +def test_decoder_contracts_and_speculative_hooks_keep_the_integrated_shell(monkeypatch): + shell, _ = make_shell(monkeypatch) + assert isinstance(shell, MultimodalRuntimeMixin) + assert not isinstance(shell, MultimodalModelMixin) + # The runtime mixin delegates nothing to a separate language model. + assert not hasattr(shell, "language_model") + assert not hasattr(shell, "set_guided_decoder") + assert Shell.vocab_size_padded is DecoderModelForCausalLM.vocab_size_padded + assert Shell.infer_max_seq_len is DecoderModelForCausalLM.infer_max_seq_len + assert Shell.forward is not MultimodalModelMixin.forward + assert shell.embedding_dim == 2 and shell.dtype == torch.float32 + shell.spec_decode_fast_paths = True + shell.runtime_hooks = RuntimeHooks(loop_on_execution_stream=True) + assert _spec_decode_fast_paths(shell) + assert runtime_hooks_of(shell).loop_on_execution_stream + + +def test_the_shell_meets_the_multimodal_runtime_contract(monkeypatch): + shell, _ = make_shell(monkeypatch) + Shell.check_multimodal_class_contract() + shell.check_multimodal_instance_contract() + check_target_multimodal_contract("Qwen4ExpForConditionalGeneration", Shell) + assert Shell.supports_concurrent_encoder + assert not Shell.supports_mm_encoder_item_scheduling + assert shell.max_encoder_tokens_per_item("image") == V.VISION_CALL_MAX_PATCHES + assert shell.max_encoder_tokens_per_item("video") == V.VISION_CALL_MAX_PATCHES + assert shell.allows_breakable_prefill_cuda_graph(None) diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py index 5c9bd6a64fc4..eee6bb6051b8 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision.py @@ -26,7 +26,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__default import ( # noqa: E501 weights as W, ) -from tensorrt_llm._torch.configs import Qwen4ExpVisionConfig +from tensorrt_llm._torch.configs import Qwen4ExpConfig, Qwen4ExpVisionConfig +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.multimodal.model import PreparedMultimodalEmbeddings from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalItemMetadata, @@ -314,15 +316,29 @@ def _params(image=None, video=None, order=None, runtime=None, mrope=True): return MultimodalParams(multimodal_data=data, multimodal_runtime=runtime) -def _splice(tower, params, num_rows, padding=0): +def _splice(tower, params, num_rows, padding=0, num_placeholders=None): """The tower's rows spliced, as the prefill target's break does, over a prompt whose tokens 1..num_rows are the placeholders and which `padding` rows follow; returns the rows - that landed there, one value per row, or None when the step carries no media.""" - kwargs = { - "multimodal_params": [params], - "mm_token_indices": torch.arange(1, 1 + num_rows, dtype=torch.int64), - } - vision = V.multimodal_embeds(tower, kwargs) + that landed there, one value per row, or None when the step carries no media. + `num_placeholders` overrides the executor's placeholder count.""" + shell = Shell.__new__(Shell) + torch.nn.Module.__init__(shell) + shell.model_config = ModelConfig( + pretrained_config=Qwen4ExpConfig( + text_config={"hidden_size": HIDDEN}, + vision_config={"out_hidden_size": HIDDEN}, + dtype="float32", + ) + ) + shell.vision_tower = tower + num_placeholders = num_rows if num_placeholders is None else num_placeholders + prepared = shell.prepare_multimodal_embeddings( + input_ids=torch.zeros(num_rows + 2, dtype=torch.int32), + multimodal_params=[params], + num_context_requests=1, + mm_token_indices=torch.arange(1, 1 + num_placeholders, dtype=torch.int64), + ) + vision = V.multimodal_embeds(prepared) if vision is None: return None total = num_rows + 2 @@ -346,10 +362,9 @@ def _item(modality, index): def test_a_text_step_runs_no_tower(): - assert V.multimodal_embeds(None, {}) is None - assert V.multimodal_embeds(None, {"multimodal_params": []}) is None + assert V.multimodal_embeds(PreparedMultimodalEmbeddings(input_ids=torch.zeros(1))) is None # Position ids alone (a text prompt through a multimodal processor) are not media. - assert V.multimodal_embeds(None, {"multimodal_params": [_params()]}) is None + assert _splice(None, _params(), 0) is None def test_the_splice_fills_the_placeholders_and_zeroes_the_padding(): @@ -402,18 +417,22 @@ def test_a_chunk_without_placeholders_skips_the_tower(): assert tower.calls == 0 -def test_a_media_request_needs_the_processors_position_ids(): - with pytest.raises(AssertionError, match="mrope_position_ids"): - _splice(_Tower(), _params(video=VIDEOS, mrope=False), 12) +def test_encoder_inputs_do_not_require_language_model_position_ids(): + embeddings = V.encode_multimodal_inputs(_Tower(), [_params(video=VIDEOS, mrope=False)]) + assert embeddings.shape == (12, HIDDEN) def test_the_placeholder_rows_must_match_the_towers_rows(): - kwargs = { - "multimodal_params": [_params(image=IMAGES)], - "mm_token_indices": torch.arange(4, dtype=torch.int64), - } - with pytest.raises(AssertionError, match="4 placeholder rows and the tower produced 5"): - V.multimodal_embeds(_Tower(), kwargs) + with pytest.raises(ValueError, match="4 multimodal placeholder positions but 5 active"): + _splice(_Tower(), _params(image=IMAGES), 5, num_placeholders=4) + + +def test_active_rows_need_the_executors_placeholder_positions(): + rows = PreparedMultimodalEmbeddings( + input_ids=torch.zeros(5), multimodal_embeddings=[torch.zeros(5, HIDDEN)] + ) + with pytest.raises(ValueError, match="require mm_token_indices"): + V.multimodal_embeds(rows) def test_the_grids_stay_on_the_host(): @@ -423,6 +442,7 @@ def test_the_grids_stay_on_the_host(): assert set(paths) == { "image.pixel_values", "video.pixel_values_videos", + "multimodal_embedding", "mrope_config.mrope_position_ids", "mrope_config.mrope_position_deltas", } diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision_warm_up.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision_warm_up.py index 73a73c4d1d78..96231f76bda2 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision_warm_up.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__default/test_vision_warm_up.py @@ -1,7 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU test of the default target's vision tower warm-up: after `VisionTower.warm_up`, which -`post_load_weights` runs, a media request of any grid compiles no Triton kernel. +"""GPU tests of the default target's vision warm-up and concurrent execution. + +After `VisionTower.warm_up`, which `post_load_weights` runs, a media request of any grid +compiles no Triton kernel. Concurrent forwards use independent attention scratch. The tower is one block deep (the block's kernels are every block's) over random weights at the checkpoint's shapes, built from the weight table's own `vis_*` declarations. @@ -92,6 +94,46 @@ def _variants() -> dict: return counts +def test_attention_workspace_is_owned_by_each_cuda_stream(): + tower = _tower() + device = torch.device("cuda", torch.cuda.current_device()) + first = torch.cuda.Stream(device=device) + second = torch.cuda.Stream(device=device) + with torch.cuda.stream(first): + first_workspace = tower._attention_workspace(device) + first_workspace.resize_(64).fill_(1) + assert tower._attention_workspace(device) is first_workspace + with torch.cuda.stream(second): + second_workspace = tower._attention_workspace(device) + second_workspace.resize_(64).fill_(2) + assert tower._attention_workspace(device) is second_workspace + first.synchronize() + second.synchronize() + assert first_workspace is not second_workspace + assert torch.equal(first_workspace.cpu(), torch.ones(64, dtype=torch.int8)) + assert torch.equal(second_workspace.cpu(), torch.full((64,), 2, dtype=torch.int8)) + + +def test_concurrent_vision_calls_match_serial_outputs(): + tower = _tower() + grids = [torch.tensor([[1, 4, 4]]), torch.tensor([[1, 8, 4]])] + pixels = [ + torch.randn(rows, tower.pixel_width, dtype=tower.dtype, device="cuda") for rows in (16, 32) + ] + expected = [tower(values, grid) for values, grid in zip(pixels, grids, strict=True)] + producer = torch.cuda.current_stream() + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + actual = [] + for stream, values, grid in zip(streams, pixels, grids, strict=True): + stream.wait_stream(producer) + with torch.cuda.stream(stream): + actual.append(tower(values, grid)) + for stream in streams: + stream.synchronize() + for output, reference in zip(actual, expected, strict=True): + torch.testing.assert_close(output, reference, rtol=1e-3, atol=1e-4) + + def test_after_warm_up_no_grid_compiles_a_kernel(): tower = _tower() tower.warm_up() diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_multimodal_contract.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_multimodal_contract.py new file mode 100644 index 000000000000..0bc058c486fd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_multimodal_contract.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Every target of a multimodal checkpoint runs the shared multimodal runtime or is text only. + +``modeling_v2_resolve`` enforces this when a deployment resolves a target. The source check +below enforces it for every target a ``route`` can return, without importing target code or +needing a GPU, so a target that would fail at resolution fails here first. +""" + +from __future__ import annotations + +import ast +import importlib.util +from pathlib import Path + +import pytest +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2 import _router_index +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import ( + TextOnlyTarget, + check_target_multimodal_contract, + checkpoint_accepts_media, +) +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.multimodal.model import ( + MultimodalModelContractError, + MultimodalRuntimeMixin, +) +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +pytestmark = pytest.mark.cpu_only + +_REPO = Path(__file__).resolve().parents[4] +_MODELS = _REPO / "tensorrt_llm/_torch/_experimental/modeling_v2/models" +_MULTIMODAL_ARCH = "Qwen4ExpForConditionalGeneration" +_TEXT_ARCH = "GptOssForCausalLM" + + +def _load_source_parser(): + path = _REPO / "jenkins/scripts/cbts/cbts/rules/modeling_v2_source.py" + spec = importlib.util.spec_from_file_location("_cbts_modeling_v2_source", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_source = _load_source_parser() + + +def _routed_architectures(tree: ast.AST) -> list[str]: + """The ``architectures[0]`` values the module's ``@routes(...)`` decorators name.""" + names = [] + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef): + continue + for decorator in node.decorator_list: + if ( + isinstance(decorator, ast.Call) + and isinstance(decorator.func, ast.Name) + and decorator.func.id == "routes" + ): + names.extend(arg.value for arg in decorator.args if isinstance(arg, ast.Constant)) + return names + + +def _routed_target_classes() -> list[tuple[str, Path, ast.ClassDef]]: + """``(architecture, modeling.py, class)`` for every target class a ``route`` returns.""" + found = [] + for routing in sorted(_MODELS.glob("*/routing.py")): + tree = ast.parse(routing.read_text()) + architectures = _routed_architectures(tree) + for node, target in _source.route_returns(tree): + if target is None or not isinstance(node.value, ast.Name): + continue + modeling = routing.parent / target / "modeling.py" + classes = { + cls.name: cls + for cls in ast.parse(modeling.read_text()).body + if isinstance(cls, ast.ClassDef) + } + found.extend( + (architecture, modeling, classes[node.value.id]) for architecture in architectures + ) + return found + + +def _base_names(cls: ast.ClassDef) -> set[str]: + names = set() + for base in cls.bases: + node = base.value if isinstance(base, ast.Subscript) else base + if isinstance(node, ast.Attribute): + names.add(node.attr) + elif isinstance(node, ast.Name): + names.add(node.id) + return names + + +def _defined_members(cls: ast.ClassDef) -> set[str]: + return { + node.name for node in cls.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + } + + +def test_route_returns_are_discovered(): + classes = {cls.name for _, _, cls in _routed_target_classes()} + assert "ModelingV2Qwen38FlashNextNvfp4Sm103Default" in classes + assert "ModelingV2GptOss120bSm10xDefault" in classes + + +def test_every_target_of_a_multimodal_checkpoint_declares_how_it_serves_media(): + problems = [] + for architecture, modeling, cls in _routed_target_classes(): + bases = _base_names(cls) + runs_runtime = MultimodalRuntimeMixin.__name__ in bases + text_only = TextOnlyTarget.__name__ in bases + where = f"{modeling.relative_to(_REPO)}::{cls.name}" + if runs_runtime and text_only: + problems.append(f"{where} inherits both MultimodalRuntimeMixin and TextOnlyTarget") + elif runs_runtime: + missing = set(MultimodalRuntimeMixin.REQUIRED_MULTIMODAL_MEMBERS) - _defined_members( + cls + ) + if missing: + problems.append(f"{where} does not implement {sorted(missing)}") + elif not text_only and checkpoint_accepts_media(architecture): + problems.append(f"{where} serves {architecture}, which accepts media") + assert not problems, "\n".join(problems) + + +def test_checkpoint_media_detection_reads_the_registered_input_processor(): + assert checkpoint_accepts_media(_MULTIMODAL_ARCH) + assert not checkpoint_accepts_media(_TEXT_ARCH) + assert not checkpoint_accepts_media("NotARegisteredArchitecture") + + +class _Plain: + pass + + +class _TextOnly(TextOnlyTarget): + pass + + +class _Runtime(MultimodalRuntimeMixin): + def encode_multimodal_inputs(self, multimodal_params): + return torch.zeros(0, 1) + + @property + def embedding_dim(self) -> int: + return 1 + + @property + def embedding_dtype(self) -> torch.dtype: + return torch.float32 + + +class _IncompleteRuntime(MultimodalRuntimeMixin): + def encode_multimodal_inputs(self, multimodal_params): + return torch.zeros(0, 1) + + +class _Both(_Runtime, TextOnlyTarget): + pass + + +@pytest.mark.parametrize( + "architecture,target", + [ + (_TEXT_ARCH, _Plain), + (_MULTIMODAL_ARCH, _TextOnly), + (_MULTIMODAL_ARCH, _Runtime), + (_TEXT_ARCH, _Runtime), + ], +) +def test_conforming_targets_pass(architecture, target): + check_target_multimodal_contract(architecture, target) + + +@pytest.mark.parametrize( + "target,match", + [ + (_Plain, "must inherit MultimodalRuntimeMixin to serve media, or TextOnlyTarget"), + (_IncompleteRuntime, "does not implement embedding_dim, embedding_dtype"), + (_Both, "inherits both"), + ], +) +def test_nonconforming_targets_fail(target, match): + with pytest.raises(MultimodalModelContractError, match=match): + check_target_multimodal_contract(_MULTIMODAL_ARCH, target) + + +@pytest.mark.parametrize("key", ["image", "video", "audio", "multimodal_embedding"]) +def test_text_only_targets_refuse_media_at_admission(key): + with pytest.raises(ValueError, match=f"language model only.*carries {key}"): + _TextOnly().validate_multimodal_request_data({key: object()}) + + +@pytest.mark.parametrize( + "mm_data", [{"mrope_config": {"mrope_position_deltas": object()}}, {"image": None}] +) +def test_text_only_targets_accept_payloads_without_media(mm_data): + _TextOnly().validate_multimodal_request_data(mm_data) + + +@pytest.mark.parametrize("mode", ["auto", "require"]) +def test_resolution_rejects_a_nonconforming_target_in_every_mode(monkeypatch, mode): + """A contract violation is a bug in the target, so ``auto`` does not fall back on it.""" + _router_index.registered_routes() + monkeypatch.setitem(_router_index._ROUTES, _MULTIMODAL_ARCH, lambda *args: _Plain) + monkeypatch.setattr(_router_index, "_device_compute_capability", lambda: (10, 3)) + config = ModelConfig(pretrained_config=PretrainedConfig(architectures=[_MULTIMODAL_ARCH])) + llm_args = TorchLlmArgs(model="/tmp/dummy_model", modeling_v2=mode) + with pytest.raises(MultimodalModelContractError, match="_Plain"): + _router_index.modeling_v2_resolve(config, llm_args) diff --git a/tests/unittest/_torch/multimodal/test_encoder_group.py b/tests/unittest/_torch/multimodal/test_encoder_group.py index b04224b4bd81..f3558ed74516 100644 --- a/tests/unittest/_torch/multimodal/test_encoder_group.py +++ b/tests/unittest/_torch/multimodal/test_encoder_group.py @@ -14,10 +14,7 @@ import pytest import torch -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( - _lengths_by_modality, - _reorder_embeds_by_manifest, -) +from tensorrt_llm._torch.multimodal.model import _lengths_by_modality, _reorder_embeds_by_manifest from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalItemMetadata, diff --git a/tests/unittest/_torch/multimodal/test_fuse_input_embeds.py b/tests/unittest/_torch/multimodal/test_fuse_input_embeds.py index 68f993deb19c..f2f1e71103d4 100644 --- a/tests/unittest/_torch/multimodal/test_fuse_input_embeds.py +++ b/tests/unittest/_torch/multimodal/test_fuse_input_embeds.py @@ -1,9 +1,25 @@ +# Copyright 2026 NVIDIA CORPORATION & AFFILIATES +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + from unittest import mock import pytest import torch -import tensorrt_llm._torch.models.modeling_multimodal_utils as multimodal_utils +import tensorrt_llm._torch.multimodal.embeddings as multimodal_utils from tensorrt_llm._torch.models.modeling_multimodal_utils import ( filter_mm_token_from_input_ids, find_input_mm_embeds, fuse_input_embeds) from tensorrt_llm._torch.modules.embedding import Embedding 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 987dee4ba9b4..b39bb6f12f10 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 @@ -26,13 +26,13 @@ import torch from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models import modeling_multimodal_mixin as mm_mixin -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( +from tensorrt_llm._torch.models.modeling_multimodal_utils import get_multimodal_embeddings +from tensorrt_llm._torch.multimodal import model as mm_mixin +from tensorrt_llm._torch.multimodal.model import ( MultimodalModelMixin, _get_mm_aux_stream, maybe_prefetch_mm_encoder_for_next_iter, ) -from tensorrt_llm._torch.models.modeling_multimodal_utils import get_multimodal_embeddings from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index 80668cb785e7..df6c4dc54f0a 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -21,11 +21,11 @@ import torch from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.multimodal.model import ( MultimodalModelMixin, _assemble_multimodal_encoder_embeddings, ) -from tensorrt_llm._torch.modules.embedding import Embedding from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalInput, @@ -382,7 +382,7 @@ def test_cache_reinitialization_requires_same_capacity(): def test_encoder_cache_creation_logs_embedding_row_capacity(): - with patch("tensorrt_llm._torch.models.modeling_multimodal_mixin.logger.info") as info: + with patch("tensorrt_llm._torch.multimodal.model.logger.info") as info: CountingEncoderMultimodalModel( make_embedding(hidden_size=4), torch.tensor([7]), @@ -403,7 +403,7 @@ def test_encoder_cache_creation_logs_byte_capacity_without_embedding_metadata(): multimodal_config=MultimodalConfig(encoder_cache_max_bytes=4096) ) - with patch("tensorrt_llm._torch.models.modeling_multimodal_mixin.logger.info") as info: + with patch("tensorrt_llm._torch.multimodal.model.logger.info") as info: model._initialize_multimodal_encoder_cache(4096) messages = [" ".join(map(str, call.args)) for call in info.call_args_list] @@ -621,7 +621,7 @@ def test_encoder_cache_partial_hit_encodes_miss_and_interleaves(): ) model._get_or_encode_multimodal_embeddings([first]) - with patch("tensorrt_llm._torch.models.modeling_multimodal_mixin.logger.debug") as debug: + with patch("tensorrt_llm._torch.multimodal.model.logger.debug") as debug: embeddings = model._get_or_encode_multimodal_embeddings([partial]) # Encoder ran twice: once for the initial miss item, once for the residual containing @@ -672,7 +672,7 @@ def test_encoder_cache_logs_rejected_oversized_write(): ) param = make_keyed_multimodal_param(embedding_lengths=[2]) - with patch("tensorrt_llm._torch.models.modeling_multimodal_mixin.logger.debug") as debug: + with patch("tensorrt_llm._torch.multimodal.model.logger.debug") as debug: model._get_or_encode_multimodal_embeddings([param]) cache = model._multimodal_encoder_cache diff --git a/tests/unittest/_torch/multimodal/test_multimodal_runtime_contract.py b/tests/unittest/_torch/multimodal/test_multimodal_runtime_contract.py new file mode 100644 index 000000000000..8ac542ea9d7e --- /dev/null +++ b/tests/unittest/_torch/multimodal/test_multimodal_runtime_contract.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The multimodal runtime base, its composite subclass, and the contract checks.""" + +from collections.abc import Sequence + +import pytest +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.multimodal.model import ( + MultimodalModelContractError, + MultimodalModelMixin, + MultimodalRuntimeMixin, +) +from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData + +pytestmark = pytest.mark.cpu_only + +_COMPOSITE_MEMBERS = ( + "language_model", + "runtime_hooks", + "vocab_size_padded", + "infer_max_seq_len", + "set_guided_decoder", + "forward", + "get_language_model_forward_kwargs", + "text_embedding_layer", + "prepare_multimodal_inputs", +) + + +class RowRuntime(torch.nn.Module, MultimodalRuntimeMixin): + """A model that runs its own decoder: rows come from `multimodal_data["image"]["rows"]`.""" + + def __init__(self, hidden_size: int = 2) -> None: + super().__init__() + self.model_config = ModelConfig(pretrained_config=PretrainedConfig(hidden_size=hidden_size)) + self.hook_kwargs: list[dict] = [] + + def encode_multimodal_inputs(self, params: Sequence[MultimodalParams]) -> torch.Tensor: + return torch.cat([param.multimodal_data["image"]["rows"] for param in params]) + + @property + def embedding_dim(self) -> int: + return self.model_config.pretrained_config.hidden_size + + @property + def embedding_dtype(self) -> torch.dtype: + return torch.float32 + + def after_active_multimodal_embeddings(self, *, active_embeddings, **kwargs): + self.hook_kwargs.append(kwargs) + return active_embeddings, () + + +class Composite(MultimodalModelMixin): + """A composite model with a real text embedding table and in-vocabulary placeholders.""" + + def __init__(self) -> None: + self.model_config = ModelConfig() + self.embedding = torch.nn.Embedding(16, 2) + + def encode_multimodal_inputs(self, params: Sequence[MultimodalParams]) -> torch.Tensor: + return torch.cat([param.multimodal_data["image"]["rows"] for param in params]) + + @property + def text_embedding_layer(self) -> torch.nn.Module: + return self.embedding + + @property + def multimodal_token_ids(self) -> Sequence[int]: + return (7,) + + +def image_params(num_rows: int, start: int = 0, end: int | None = None) -> MultimodalParams: + return MultimodalParams( + multimodal_data={"image": {"rows": torch.arange(num_rows * 2.0).reshape(num_rows, 2)}}, + multimodal_runtime=MultimodalRuntimeData( + past_seen_token_num=start, + chunk_end_pos=num_rows if end is None else end, + embed_mask_cumsum=torch.arange(1, num_rows + 1), + ), + ) + + +def test_the_runtime_base_delegates_nothing_to_a_separate_language_model(): + assert issubclass(MultimodalModelMixin, MultimodalRuntimeMixin) + for name in _COMPOSITE_MEMBERS: + assert not hasattr(MultimodalRuntimeMixin, name), name + assert hasattr(MultimodalModelMixin, name), name + assert not MultimodalRuntimeMixin.supports_concurrent_encoder + assert MultimodalModelMixin.supports_concurrent_encoder + + +def test_runtime_defaults_leave_processor_limits_and_refuse_breakable_graphs(): + model = RowRuntime() + assert model.max_encoder_tokens_per_item("image") is None + assert not model.allows_breakable_prefill_cuda_graph(None) + + +def test_rows_come_back_paired_with_their_placeholder_positions(): + model = RowRuntime() + indices = torch.tensor([1, 2, 4]) + prepared = model.prepare_multimodal_embeddings( + input_ids=torch.zeros(5, dtype=torch.int32), + multimodal_params=[image_params(3)], + num_context_requests=1, + mm_token_indices=indices, + ) + assert prepared.mm_token_indices is indices + assert sum(rows.shape[0] for rows in prepared.multimodal_embeddings) == 3 + # The model hooks keep receiving the executor's indices. + assert model.hook_kwargs[0]["mm_token_indices"] is indices + + +def test_placeholder_positions_must_count_the_active_rows(): + with pytest.raises(ValueError, match="2 multimodal placeholder positions but 3 active"): + RowRuntime().prepare_multimodal_embeddings( + input_ids=torch.zeros(5, dtype=torch.int32), + multimodal_params=[image_params(3)], + num_context_requests=1, + mm_token_indices=torch.tensor([1, 2]), + ) + + +def test_a_forward_without_active_rows_carries_no_positions(): + prepared = RowRuntime().prepare_multimodal_embeddings( + input_ids=torch.zeros(2, dtype=torch.int32), + multimodal_params=[image_params(3, start=3, end=3)], + num_context_requests=1, + mm_token_indices=torch.tensor([], dtype=torch.int64), + ) + assert not prepared.multimodal_embeddings + assert prepared.mm_token_indices is None + + +def test_composite_fusion_keeps_deriving_positions_it_was_not_given(): + """Composite models may receive executor indices that do not describe their rows; the + default fusion derives positions from `multimodal_token_ids` when one index set is + missing, so the composite path does not apply the runtime's pairing check.""" + model = Composite() + input_ids = torch.tensor([1, 7, 7, 3]) + prepared = model.prepare_multimodal_inputs( + input_ids=input_ids, + positions=None, + multimodal_params=[image_params(2)], + num_context_requests=1, + mm_token_indices=torch.tensor([], dtype=torch.int64), + ) + assert prepared.input_ids is None + torch.testing.assert_close(prepared.inputs_embeds[1:3], torch.arange(4.0).reshape(2, 2)) + torch.testing.assert_close(prepared.inputs_embeds[0], model.embedding.weight[1]) + + +def test_class_contract_names_every_missing_member(): + class Incomplete(MultimodalRuntimeMixin): + def encode_multimodal_inputs(self, params): + return torch.zeros(0, 1) + + RowRuntime.check_multimodal_class_contract() + with pytest.raises(MultimodalModelContractError, match="embedding_dim, embedding_dtype"): + Incomplete.check_multimodal_class_contract() + + +def test_instance_contract_reports_the_failing_property_with_its_cause(): + class ComposedConfig(RowRuntime): + @property + def embedding_dim(self) -> int: + return self.model_config.pretrained_config.text_width + + model = ComposedConfig() + # Through `torch.nn.Module`, the property's AttributeError loses its cause. + with pytest.raises(AttributeError, match="has no attribute 'embedding_dim'"): + _ = model.embedding_dim + with pytest.raises(MultimodalModelContractError, match=r"embedding_dim failed.*text_width"): + model.check_multimodal_instance_contract() + RowRuntime().check_multimodal_instance_contract() + + +@pytest.mark.parametrize("hidden_size", [0, 2.0]) +def test_instance_contract_rejects_invalid_row_widths(hidden_size): + with pytest.raises(MultimodalModelContractError, match="positive int"): + RowRuntime(hidden_size=hidden_size).check_multimodal_instance_contract() diff --git a/tests/unittest/_torch/multimodal/test_prepared_multimodal_embeddings.py b/tests/unittest/_torch/multimodal/test_prepared_multimodal_embeddings.py new file mode 100644 index 000000000000..303958ea8148 --- /dev/null +++ b/tests/unittest/_torch/multimodal/test_prepared_multimodal_embeddings.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared encoder-row preparation independently of text embedding fusion.""" + +from collections.abc import Sequence + +import pytest +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models import modeling_multimodal_mixin as legacy_model +from tensorrt_llm._torch.models import modeling_multimodal_utils as legacy_embeddings +from tensorrt_llm._torch.multimodal import embeddings, model +from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData + +pytestmark = pytest.mark.cpu_only + + +class RowModel(model.MultimodalModelMixin): + def __init__(self) -> None: + self.model_config = ModelConfig() + self.encode_calls = 0 + self.hooks: list[str] = [] + + @property + def text_embedding_layer(self) -> torch.nn.Module: + raise AssertionError("Row preparation must not access text embeddings") + + def encode_multimodal_inputs(self, params: Sequence[MultimodalParams]) -> torch.Tensor: + self.encode_calls += 1 + return torch.cat([param.multimodal_data["image"]["rows"] for param in params]) + + def after_full_multimodal_embeddings( + self, *, input_ids: torch.Tensor, embeddings: torch.Tensor, **kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + self.hooks.append("full") + return input_ids + 1, embeddings + 10 + + def after_active_multimodal_embeddings( + self, *, active_embeddings: list[torch.Tensor], **kwargs + ) -> tuple[list[torch.Tensor], tuple[torch.Tensor, ...]]: + self.hooks.append("active") + return [rows * 2 for rows in active_embeddings], tuple( + rows + 100 for rows in active_embeddings + ) + + +def runtime(start: int, end: int) -> MultimodalRuntimeData: + return MultimodalRuntimeData( + past_seen_token_num=start, + chunk_end_pos=end, + embed_mask_cumsum=torch.arange(1, 7, dtype=torch.int64), + ) + + +def test_row_preparation_reuses_full_embeddings_and_runs_hooks_after_chunk_selection(): + instance = RowModel() + rows = torch.arange(12, dtype=torch.float32).reshape(6, 2) + params = MultimodalParams( + multimodal_data={"image": {"rows": rows}}, multimodal_runtime=runtime(0, 2) + ) + input_ids = torch.tensor([7, 7]) + first = instance.prepare_multimodal_embeddings( + input_ids=input_ids, multimodal_params=[params], num_context_requests=1 + ) + torch.testing.assert_close(first.input_ids, input_ids + 1) + torch.testing.assert_close(first.multimodal_embeddings[0], (rows[:2] + 10) * 2) + torch.testing.assert_close(first.extra_embeds[0], rows[:2] + 110) + torch.testing.assert_close(params.multimodal_data["multimodal_embedding"], rows) + + params.multimodal_runtime = runtime(2, 5) + second = instance.prepare_multimodal_embeddings( + input_ids=torch.tensor([7, 7, 7]), multimodal_params=[params], num_context_requests=1 + ) + torch.testing.assert_close(second.multimodal_embeddings[0], (rows[2:5] + 10) * 2) + assert instance.encode_calls == 1 + assert instance.hooks == ["full", "active", "full", "active"] + + +@pytest.mark.parametrize("num_contexts,start,end", [(0, 0, 2), (1, 6, 6)]) +def test_inactive_rows_do_not_encode_or_run_hooks(num_contexts, start, end): + instance = RowModel() + params = MultimodalParams( + multimodal_data={"image": {"rows": torch.zeros(6, 2)}}, + multimodal_runtime=runtime(start, end), + ) + input_ids = torch.tensor([7]) + prepared = instance.prepare_multimodal_embeddings( + input_ids=input_ids, multimodal_params=[params], num_context_requests=num_contexts + ) + assert prepared.input_ids is input_ids + assert not prepared.multimodal_embeddings and not prepared.extra_embeds + assert instance.encode_calls == 0 and not instance.hooks + + +def test_already_chunk_local_embeddings_are_not_sliced_twice(): + instance = RowModel() + rows = torch.arange(4, dtype=torch.float32).reshape(2, 2) + params = MultimodalParams( + multimodal_data={"multimodal_embedding": rows, "multimodal_embedding_is_chunk": True}, + multimodal_runtime=runtime(2, 4), + ) + prepared = instance.prepare_multimodal_embeddings( + input_ids=torch.tensor([7, 7]), multimodal_params=[params], num_context_requests=1 + ) + torch.testing.assert_close(prepared.multimodal_embeddings[0], (rows + 10) * 2) + assert instance.encode_calls == 0 + + +def test_legacy_imports_preserve_runtime_identity(): + assert legacy_model.MultimodalModelMixin is model.MultimodalModelMixin + assert legacy_model.MultimodalRuntimeMixin is model.MultimodalRuntimeMixin + assert legacy_model.PreparedLlmInputs is model.PreparedLlmInputs + assert legacy_model.EncoderGroup is model.EncoderGroup + + +@pytest.mark.parametrize( + "name", + [ + "_get_uncached_multimodal_params", + "_store_chunked_prefill_embeddings", + "_normalize_encoder_embeddings", + "_join_embeddings", + "get_multimodal_embeddings", + "get_attached_multimodal_embeddings", + "find_input_mm_embeds", + "filter_mm_token_from_input_ids", + "_as_copy_index", + "fuse_input_embeds", + ], +) +def test_legacy_embedding_exports_preserve_runtime_identity(name): + assert getattr(legacy_embeddings, name) is getattr(embeddings, name) diff --git a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py index 15f9c543fd86..e1e851b66a24 100644 --- a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py +++ b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py @@ -14,11 +14,8 @@ import torch -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( - EncoderGroup, - encode_multimodal_by_groups, -) from tensorrt_llm._torch.models.modeling_qwen3vl import _qwen3vl_build_batched_input +from tensorrt_llm._torch.multimodal.model import EncoderGroup, encode_multimodal_by_groups from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalItemMetadata, diff --git a/tests/unittest/_torch/peft/test_lora_ownership.py b/tests/unittest/_torch/peft/test_lora_ownership.py index 64299f7db1b9..a350954d43cb 100644 --- a/tests/unittest/_torch/peft/test_lora_ownership.py +++ b/tests/unittest/_torch/peft/test_lora_ownership.py @@ -15,7 +15,7 @@ import torch from safetensors.torch import save_file -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.peft.lora.cuda_graph_lora_params import CudaGraphLoraParams from tensorrt_llm._torch.peft.lora.layer import LoraLayer, LoraModuleType diff --git a/tests/unittest/_torch/pyexecutor/engine/test_lora.py b/tests/unittest/_torch/pyexecutor/engine/test_lora.py index 61012486a990..8e187f2d3e6c 100644 --- a/tests/unittest/_torch/pyexecutor/engine/test_lora.py +++ b/tests/unittest/_torch/pyexecutor/engine/test_lora.py @@ -11,8 +11,8 @@ import pytest import torch -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.peft.lora.cuda_graph_lora_manager import CudaGraphLoraManager from tensorrt_llm._torch.peft.lora.cuda_graph_lora_params import CudaGraphLoraParams diff --git a/tests/unittest/_torch/pyexecutor/engine/test_multimodal.py b/tests/unittest/_torch/pyexecutor/engine/test_multimodal.py index 5941744adb41..5cda2d1b59a1 100644 --- a/tests/unittest/_torch/pyexecutor/engine/test_multimodal.py +++ b/tests/unittest/_torch/pyexecutor/engine/test_multimodal.py @@ -19,17 +19,20 @@ import tensorrt_llm from tensorrt_llm._torch.models.modeling_multimodal_encoder import MultimodalEncoderMixin -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( +from tensorrt_llm._torch.multimodal.model import ( MultimodalEncoderContractError, MultimodalModelMixin, + MultimodalRuntimeMixin, _build_request_multimodal_input, ) from tensorrt_llm._torch.pyexecutor.engine.multimodal import ( MultimodalItemScheduler, + bound_mm_max_tokens_per_item, resolve_bytes_per_mm_encoder_embedding, resolve_mm_encoder_output_budget, setup_mm_encoder_attn_metadata, setup_mm_encoder_cuda_graphs, + validate_mm_encoder_prefetch_compatibility, validate_mm_encoder_scheduling_compatibility, ) from tensorrt_llm._torch.pyexecutor.llm_request import ( @@ -50,6 +53,7 @@ MultimodalConfig, MultimodalEncoderCudaGraphConfig, MultimodalEncoderSchedulingPolicy, + TorchLlmArgs, ) from tensorrt_llm.mapping import Mapping @@ -465,6 +469,42 @@ def test_side_stream_compatibility_is_checked_only_for_item_scheduled_models() - validate_mm_encoder_scheduling_compatibility(args, item_scheduling_enabled=True) +class _SingleStreamEncoderModel(MultimodalRuntimeMixin): + """A runtime model whose encoder keeps state shared across CUDA streams.""" + + +class _BoundedEncoderModel(MultimodalRuntimeMixin): + def max_encoder_tokens_per_item(self, modality: str) -> int | None: + return {"image": 1024}.get(modality) + + +@pytest.mark.parametrize("max_ahead", [0, 1]) +def test_side_stream_prefetch_requires_a_concurrent_encoder(max_ahead: int) -> None: + args = TorchLlmArgs( + model="dummy", multimodal_config=MultimodalConfig(encoder_side_stream_max_ahead=max_ahead) + ) + # Composite models keep prefetch; models that are not multimodal ignore the setting. + validate_mm_encoder_prefetch_compatibility(args, MultimodalModelMixin()) + validate_mm_encoder_prefetch_compatibility(args, torch.nn.Linear(1, 1)) + if max_ahead == 0: + validate_mm_encoder_prefetch_compatibility(args, _SingleStreamEncoderModel()) + return + with pytest.raises(ValueError, match="supports_concurrent_encoder"): + validate_mm_encoder_prefetch_compatibility(args, _SingleStreamEncoderModel()) + + +def test_item_sizes_are_bounded_by_the_model_encoder() -> None: + processor_sizes = {"image": 4096, "video": 8192, "audio": 512} + assert bound_mm_max_tokens_per_item(_BoundedEncoderModel(), processor_sizes) == { + "image": 1024, + "video": 8192, + "audio": 512, + } + assert bound_mm_max_tokens_per_item(MultimodalModelMixin(), processor_sizes) == ( + processor_sizes + ) + + @pytest.mark.parametrize( ("mm_config_kwargs", "falls_back"), [ diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py index 49e2921cb335..7b21396244a1 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py @@ -22,7 +22,7 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures from tensorrt_llm._torch.model_config import KVCacheLayerSpec as ModelKVCacheLayerSpec from tensorrt_llm._torch.model_config import ModelConfig, _get_kv_cache_layer_specs -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin, MultimodalRuntimeMixin from tensorrt_llm._torch.pyexecutor._util import ( CacheCost, KvCacheCreator, @@ -380,6 +380,46 @@ def encode_multimodal_inputs(self, multimodal_params): assert creator._dummy_encoder_inputs == [] +def test_encoder_profiling_builds_items_within_the_model_bound(): + """Above the model's per-item bound, profiling asks the processor for one item per + params object under that bound, and still fills the encoder token budget.""" + + class _InputProcessor: + def __init__(self): + self.calls = [] + + def get_mm_max_tokens_per_item(self, max_num_encoder_tokens=None): + return {"image": min(4096, max_num_encoder_tokens or 4096), "audio": 512} + + def get_dummy_mm_data(self, *, max_num_encoder_tokens, mm_counts, dtype): + self.calls.append((max_num_encoder_tokens, mm_counts, dtype)) + return {"image": {"item_count": mm_counts["image"]}} + + class _Model(MultimodalRuntimeMixin): + dtype = torch.float16 + + def max_encoder_tokens_per_item(self, modality): + return 1024 if modality == "image" else None + + input_processor = _InputProcessor() + creator = object.__new__(KvCacheCreator) + creator._model_engine = SimpleNamespace( + model=_Model(), + input_processor=input_processor, + encoder_batch_size=8, + encoder_max_num_tokens=8192, + ) + creator._llm_args = TorchLlmArgs(model="dummy", checkpoint_format="HF") + creator._profiling_stage_data = {"enable_mm_reqs": True} + + encoder_inputs = creator._create_dummy_encoder_inputs() + + assert input_processor.calls == [(1024, {"image": 1}, torch.float16)] * 8 + assert [params.multimodal_data for params in encoder_inputs] == [ + {"image": {"item_count": 1}} + ] * 8 + + def _make_creator( tokens_per_block, dummy_reqs, diff --git a/tests/unittest/_torch/pyexecutor/multimodal_utils.py b/tests/unittest/_torch/pyexecutor/multimodal_utils.py index 1465a02f50a4..ae6ea192ca8c 100644 --- a/tests/unittest/_torch/pyexecutor/multimodal_utils.py +++ b/tests/unittest/_torch/pyexecutor/multimodal_utils.py @@ -12,7 +12,7 @@ import torch -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor.engine.multimodal import MultimodalItemScheduler from tensorrt_llm._torch.pyexecutor.llm_request import ( LlmRequest, diff --git a/tests/unittest/_torch/pyexecutor/test_model_engine.py b/tests/unittest/_torch/pyexecutor/test_model_engine.py index 4a3e8341d9d3..8da83fac4582 100644 --- a/tests/unittest/_torch/pyexecutor/test_model_engine.py +++ b/tests/unittest/_torch/pyexecutor/test_model_engine.py @@ -26,8 +26,8 @@ ) import tensorrt_llm._torch.pyexecutor.model_engine as model_engine_module -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin, MultimodalRuntimeMixin from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType from tensorrt_llm._torch.pyexecutor.engine import model_setup from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller @@ -715,6 +715,28 @@ def test_breakable_allows_multimodal_wrapper_in_text_only_mode(self) -> None: self.assertTrue(is_multimodal(engine.model, engine.input_processor)) engine._validate_breakable_cuda_graph_compatibility() + def test_breakable_follows_a_runtime_decoder_declaration(self) -> None: + """A decoder that runs the multimodal runtime itself decides, even under the + `DecoderModelForCausalLM` rule that admits text decoders.""" + for allows in (False, True): + with self.subTest(allows=allows): + + class RuntimeDecoder(DecoderModelForCausalLM, MultimodalRuntimeMixin): + def allows_breakable_prefill_cuda_graph(self, llm_args) -> bool: + return allows + + engine = object.__new__(ModelEngine) + engine.model = RuntimeDecoder.__new__(RuntimeDecoder) + engine.input_processor = Mock(spec=BaseMultimodalInputProcessor) + engine.llm_args = TorchLlmArgs( + model="dummy", prefill_cuda_graph_backend=PrefillCudaGraphBackend.BREAKABLE + ) + if allows: + engine._validate_breakable_cuda_graph_compatibility() + else: + with self.assertRaisesRegex(ValueError, "inside the captured prefill forward"): + engine._validate_breakable_cuda_graph_compatibility() + def test_restore_after_refit_never_touches_the_kv_cache(self) -> None: """The re-wrap hook must not run the warm-up KV-cache cleanup. diff --git a/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py b/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py index 16c8058b2ce5..fd33da0a74f2 100644 --- a/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py +++ b/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py @@ -50,7 +50,6 @@ from tensorrt_llm._torch.models import modeling_qwen3_moe as modeling_qwen3_moe_mod from tensorrt_llm._torch.models.checkpoints.hf.qwen3_moe_weight_mapper import Qwen3MoeHfWeightMapper from tensorrt_llm._torch.models.checkpoints.mx.checkpoint_loader import MXCheckpointLoader -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm._torch.modules import linear as linear_mod from tensorrt_llm._torch.modules.linear import Linear, WeightMode @@ -61,6 +60,7 @@ from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode from tensorrt_llm._torch.moe.fused_moe.moe_resolution import resolve_moe_impl from tensorrt_llm._torch.moe.fused_moe.quantization import UnquantizedFusedMoEMethod +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( diff --git a/tests/unittest/_torch/pyexecutor/test_multimodal_scheduler.py b/tests/unittest/_torch/pyexecutor/test_multimodal_scheduler.py index 12af6f2f33ae..e1a84c9c8901 100644 --- a/tests/unittest/_torch/pyexecutor/test_multimodal_scheduler.py +++ b/tests/unittest/_torch/pyexecutor/test_multimodal_scheduler.py @@ -13,7 +13,7 @@ record_output, ) -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( +from tensorrt_llm._torch.multimodal.model import ( MultimodalEncoderContractError, MultimodalModelMixin, ) diff --git a/tests/unittest/executor/test_forward_error_survival.py b/tests/unittest/executor/test_forward_error_survival.py index d5777b7ed49f..192932827168 100644 --- a/tests/unittest/executor/test_forward_error_survival.py +++ b/tests/unittest/executor/test_forward_error_survival.py @@ -25,7 +25,8 @@ import pytest -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch._experimental.modeling_v2._multimodal import TextOnlyTarget +from tensorrt_llm._torch.multimodal.model import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.py_executor import BatchState, PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager @@ -202,6 +203,16 @@ def test_validate_request_skips_hook_for_text_only_requests(): model.validate_multimodal_request_data.assert_not_called() +def test_validate_request_rejects_media_for_text_only_targets(): + class _TextOnlyModel(TextOnlyTarget): + pass + + executor = _validating_executor(_TextOnlyModel()) + with pytest.raises(ValueError, match="language model only"): + PyExecutor._validate_request(executor, _request(mm_data={"image": object()})) + PyExecutor._validate_request(executor, _request(mm_data={"mrope_config": {}})) + + def test_validate_request_accepts_models_without_the_hook(): executor = _validating_executor(object()) PyExecutor._validate_request(executor, _request(mm_data={"image": object()}))