diff --git a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md index 6cbf73dad66a..bef01b2e670b 100644 --- a/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md +++ b/.claude/skills/trtllm-model-onboard-multimodal/SKILL.md @@ -83,7 +83,7 @@ metadata: When `@support_multimodal_disaggregated` is set and the deployment uses `TLLM_MULTIMODAL_DISAGGREGATED=1`: - **Encoder worker:** runs as a standalone `MultimodalEncoder` (`mm_encoder_only=True`). It executes only the multimodal encoder and ships `mm_embeddings` (+ mRoPE position ids/deltas) to prefill+decode workers as shared-tensor handles. -- **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_disagg()` is true; the input processor's `_attach_multimodal_embeddings_impl()` override binds the encoder handles into the request (the base `attach_multimodal_embeddings` wrapper detokenizes tokenized inputs for non-fast-path VLMs, then delegates to your impl). For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors. +- **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_mm_disagg()` is true; the input processor's `attach_multimodal_embeddings()` override binds the encoder handles into the request. For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors. ### Templates to study @@ -229,7 +229,7 @@ class {Name}Model(PreTrainedModel): if hasattr(self, "llm"): return # idempotency guard — re-entry from `post_config` etc. - if not _is_disagg(): + if not _is_mm_disagg(): self.mm_encoder = {Name}VisionModel(model_config) else: self.mm_encoder = None @@ -269,7 +269,7 @@ class {Name}Model(PreTrainedModel): multimodal_params = kwargs.get("multimodal_params", []) mm_embeds = [] - if len(multimodal_params) > 0 and not _is_disagg(): + if len(multimodal_params) > 0 and not _is_mm_disagg(): mm_embeds = get_multimodal_embeddings( encoder_forward_fn=self.mm_encoder.forward, multimodal_params=multimodal_params[:num_context_requests], @@ -341,7 +341,7 @@ class {Name}Model(PreTrainedModel): ... ```python def load_weights(self, weights, weight_mapper): - if not _is_disagg(): + if not _is_mm_disagg(): self.mm_encoder.load_weights(weights) # Release mmap pages backing the encoder weights as soon as we're done. if hasattr(weights, "mark_consumed"): diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 73e075d949ac..5c825d20e4cf 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -81,6 +81,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^7]: Text-only support via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. [^8]: Supports text and image inputs. The vision tower runs in BF16 even when the text decoder is quantized (FP8 block-scale or NVFP4). The text decoder is also usable standalone (text-only) via the `Step3p5ForCausalLM` architecture. [^9]: Audio modality only supported on E2B/E4B variants. +[^10]: Audio requires a checkpoint with a `sound_config` and is supported only on the full (non-disaggregated) model path, not the EPD disaggregated path. # Multimodal Feature Support Matrix (PyTorch Backend) @@ -94,7 +95,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `LlavaNextForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | L + I | | `Llama4ForConditionalGeneration` | Yes | Yes | No | Yes | Yes | No | Yes | No | L + I | | `Mistral3ForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | L + I | -| `NemotronH_Nano_VL_V2` | Yes | Yes | Yes | Yes | Yes | N/A | Yes | No | L + I + V | +| `NemotronH_Nano_VL_V2` | Yes | Yes | Yes | Yes | Yes | N/A | Yes | Yes | L + I + V + A [^10] | | `Phi4MMForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | L + I + A | | `Qwen2VLForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | L + I + V | | `Qwen2_5_VLForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | L + I + V | diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 743c24042909..285a8b6fe0dd 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -66,6 +66,8 @@ { "layout_metadata", "mm_bidirectional_blocks", + "multimodal_embedding", + "multimodal_embedding_lengths", "special_token_offsets", "multimodal_embed_mask_cumsum", } diff --git a/tensorrt_llm/_torch/models/modeling_exaone4_5.py b/tensorrt_llm/_torch/models/modeling_exaone4_5.py index 1506c4fce23e..a617f7c5d34c 100644 --- a/tensorrt_llm/_torch/models/modeling_exaone4_5.py +++ b/tensorrt_llm/_torch/models/modeling_exaone4_5.py @@ -9,7 +9,7 @@ from transformers.models.auto import CONFIG_MAPPING from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper -from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_disagg +from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg from ...inputs import ( ContentFormat, @@ -198,7 +198,7 @@ def __init__( llm_model_config.pretrained_config = llm_model_config.pretrained_config.text_config self.llm = AutoModelForCausalLM.from_config(llm_model_config) - if not _is_disagg(): + if not _is_mm_disagg(): mm_encoder_config = copy.deepcopy(model_config) self.mm_encoder = Exaone4_5_VisionModel(mm_encoder_config, Qwen2_5_VisionModel) else: @@ -231,7 +231,7 @@ def forward( mm_multimodal_params = self._get_requests_with_mm_data(multimodal_params) if len(mm_multimodal_params) > 0: - if not _is_disagg(): + if not _is_mm_disagg(): mm_embeds = get_multimodal_embeddings( encoder_forward_fn=self.mm_encoder.forward, multimodal_params=mm_multimodal_params, @@ -262,6 +262,6 @@ def forward( def load_weights(self, weights, weight_mapper: BaseWeightMapper): assert isinstance(weight_mapper, Exaone4_5HfWeightMapper) weights = weight_mapper.preprocess_weights(weights) - if not _is_disagg(): + if not _is_mm_disagg(): self.mm_encoder.load_weights(weights) self.llm.load_weights(weights, weight_mapper) diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py index 5f327c82c026..58303b01e5ae 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma3vl.py +++ b/tensorrt_llm/_torch/models/modeling_gemma3vl.py @@ -1,6 +1,5 @@ import copy import dataclasses -import os from typing import List, Optional, Tuple import torch @@ -23,17 +22,11 @@ from ..modules.linear import Linear from ..modules.rms_norm import RMSNorm from .modeling_gemma3 import Gemma3ForCausalLM -from .modeling_multimodal_utils import fuse_input_embeds +from .modeling_multimodal_utils import (_MULTIMODAL_ENV_NAME, _is_mm_disagg, + fuse_input_embeds) from .modeling_siglip import SiglipVisionModel from .modeling_utils import ModelConfig, filter_weights, register_auto_model -_MULTIMODAL_ENV_NAME = "TLLM_MULTIMODAL_DISAGGREGATED" - - -# Make this a runtime lookup rather than a module-wide constant for easier unit testing. -def _is_disagg() -> bool: - return os.getenv(_MULTIMODAL_ENV_NAME, "0") == "1" - class Gemma3InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): @@ -185,7 +178,7 @@ def forward(self, vision_outputs: torch.Tensor): class Gemma3VLM(PreTrainedModel): def __init__(self, model_config: ModelConfig[Gemma3Config]): - if _is_disagg(): + if _is_mm_disagg(): raise NotImplementedError( "Gemma3VLM does not support disaggregated inference yet. Please unset " f"the {_MULTIMODAL_ENV_NAME} environment variable, or set it to '0'." diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 66ef1b03cfcc..38e33d11069a 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -23,7 +23,6 @@ import copy import dataclasses import math -import os from typing import Dict, List, Optional, Tuple import torch @@ -51,7 +50,12 @@ from .modeling_gemma4 import Gemma4ForCausalLM from .modeling_gemma4_audio import Gemma4AudioModel from .modeling_gemma4_vision import Gemma4VisionModel -from .modeling_multimodal_utils import find_input_mm_embeds, fuse_input_embeds +from .modeling_multimodal_utils import ( + _MULTIMODAL_ENV_NAME, + _is_mm_disagg, + find_input_mm_embeds, + fuse_input_embeds, +) from .modeling_utils import ModelConfig, filter_weights, register_auto_model _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0" @@ -69,12 +73,6 @@ PreTrainedModel, ) -_MULTIMODAL_ENV_NAME = "TLLM_MULTIMODAL_DISAGGREGATED" - - -def _is_disagg() -> bool: - return os.getenv(_MULTIMODAL_ENV_NAME, "0") == "1" - class RMSNormNoScale(nn.Module): """RMSNorm without learnable scale (for multimodal embedder pre-projection).""" @@ -602,7 +600,7 @@ def _check_and_adjust_experts_implementation(self, *args, **kwargs): return None def __init__(self, model_config: ModelConfig[Gemma4Config]): - if _is_disagg(): + if _is_mm_disagg(): raise NotImplementedError( "Gemma4ForConditionalGeneration does not support " "disaggregated inference yet. Please unset the " diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index d130210376ba..6595e0b405f0 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -49,7 +49,7 @@ PreTrainedTokenizerBase, ) -from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.inputs.multimodal import DisaggPrefillMultimodalInputs, MultimodalParams from tensorrt_llm.mapping import Mapping from ..._utils import prefer_pinned @@ -1492,12 +1492,12 @@ def call_with_text_prompt( "multimodal_data": multimodal_data, } - def get_prompt_token_ids( + def build_disagg_prefill_multimodal_inputs( self, inputs: TextPrompt, mm_handles: List[Dict[str, Any]], - ) -> Tuple[List[int], List[int], List[int]]: - """Build token IDs with multimodal placeholders expanded for disaggregated serving. + ) -> DisaggPrefillMultimodalInputs: + """Build disaggregated prefill inputs from multimodal embedding handles. Args: inputs: Text prompt input container. @@ -1505,7 +1505,9 @@ def get_prompt_token_ids( context phase, each containing ``tensor_size``. Returns: - Tuple of (expanded_ids, mm_token_lengths, mm_token_offsets). + DisaggPrefillMultimodalInputs containing expanded token IDs, + prompt-side MM positions/lengths, exact runs, and encoder-output + embedding lengths. """ text_prompt = inputs.get("prompt") if not text_prompt: @@ -1556,7 +1558,15 @@ def get_prompt_token_ids( expanded_ids[write_pos] = input_ids[read_pos] write_pos += 1 - return (expanded_ids.to(torch.int32).tolist(), mm_token_length, mm_token_offsets) + return DisaggPrefillMultimodalInputs( + prompt_token_ids=expanded_ids.to(torch.int32).tolist(), + multimodal_lengths=mm_token_length, + multimodal_positions=mm_token_offsets, + multimodal_embedding_lengths=[mm_handle["tensor_size"][0] for mm_handle in mm_handles], + multimodal_item_run_cu_offsets=list(range(len(mm_token_length) + 1)), + multimodal_run_positions=mm_token_offsets, + multimodal_run_lengths=mm_token_length, + ) # --------------------------------------------------------------------------- diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py index 910bc3c82a18..5b09b28a577c 100644 --- a/tensorrt_llm/_torch/models/modeling_llava_next.py +++ b/tensorrt_llm/_torch/models/modeling_llava_next.py @@ -1,5 +1,4 @@ import copy -import os from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union import numpy as np @@ -15,7 +14,9 @@ BaseWeightMapper from tensorrt_llm._torch.models.checkpoints.hf.llava_next_weight_mapper import \ LlavaNextHfWeightMapper -from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg +from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, + MultimodalParams) from ...inputs import (BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor, ContentFormat, @@ -29,11 +30,10 @@ from .modeling_auto import AutoModelForCausalLM from .modeling_clip import CLIPVisionModel from .modeling_multimodal_utils import (find_input_mm_embeds, fuse_input_embeds, + get_attached_multimodal_embeddings, get_multimodal_embeddings) from .modeling_utils import register_auto_model, register_vision_encoder -DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1' - class LlavaNextInputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): @@ -111,7 +111,8 @@ def _expand_image_placeholders_in_token_ids( num_mm_tokens_per_placeholder: List[int], ) -> Tuple[List[int], List[int], List[int]]: """ - Shared logic (called by expand_prompt_token_ids_for_mm and get_prompt_token_ids): + Shared logic (called by expand_prompt_token_ids_for_mm and + build_disagg_prefill_multimodal_inputs): replace each image placeholder token in prompt_token_ids with placeholder_id repeated num_mm_tokens_per_placeholder[i] times. @@ -268,12 +269,11 @@ def _postprocess( mm_features = mm_features.view(-1, mm_features.shape[-1]) return fused_input_ids, mm_features - def get_prompt_token_ids( - self, inputs: Union[TextPrompt, TokensPrompt], - mm_handles: List[Dict[str, - Any]]) -> Tuple[List[int], List[int], List[int]]: + def build_disagg_prefill_multimodal_inputs( + self, inputs: Union[TextPrompt, TokensPrompt], + mm_handles: List[Dict[str, Any]]) -> DisaggPrefillMultimodalInputs: """ - Build input token ids with multimodal placeholders expanded to the number of MM tokens. + Build disaggregated prefill inputs from multimodal embedding handles. Uses an already tokenized prompt or tokenizes the txt prompt first. @@ -282,10 +282,9 @@ def get_prompt_token_ids( mm_handles: List of multimodal embedding handles. Returns: - Tuple[List[int], List[int], List[int]]: - - expanded_ids: token ids with each image token expanded to a placeholder repeated per MM token - - mm_token_length: per-image MM token lengths - - mm_token_offsets: start offsets (positions) for each image's MM tokens within expanded_ids + DisaggPrefillMultimodalInputs containing expanded token IDs, + prompt-side MM positions/lengths, exact runs, and encoder-output + embedding lengths. """ # TODO: Move this function to the base input processor class when extending for more models text_prompt = inputs.get("prompt") @@ -327,7 +326,18 @@ def get_prompt_token_ids( f"({mm_token_length[-1] + mm_token_offsets[-1]}) should be less " f"than or equal to final_length ({final_length})") - return expanded_ids, mm_token_length, mm_token_offsets + return DisaggPrefillMultimodalInputs( + prompt_token_ids=expanded_ids, + multimodal_lengths=mm_token_length, + multimodal_positions=mm_token_offsets, + multimodal_embedding_lengths=[ + mm_handle["tensor_size"][0] for mm_handle in mm_handles + ], + multimodal_item_run_cu_offsets=list(range(len(mm_token_length) + + 1)), + multimodal_run_positions=mm_token_offsets, + multimodal_run_lengths=mm_token_length, + ) def _attach_multimodal_embeddings_impl( self, inputs: TextPrompt, @@ -619,7 +629,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, super().__init__(config) if hasattr(self, "llm"): return - if not DISAGG: + if not _is_mm_disagg(): self.mm_encoder = LlavaNextVisionModel(model_config) else: self.mm_encoder = None @@ -694,15 +704,14 @@ def forward( multimodal_params = kwargs.get("multimodal_params", []) mm_embeds = [] if len(multimodal_params) > 0: - if not DISAGG: + if self.mm_encoder is not None: mm_embeds = get_multimodal_embeddings( encoder_forward_fn=self.mm_encoder.forward, multimodal_params=multimodal_params[:num_context_requests]) else: - raise NotImplementedError( - "LlavaNextModel does not support disaggregated inference yet. Please unset " - f"the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'." - ) + # E/P prefill: encoder already ran; use attached embeddings. + mm_embeds = get_attached_multimodal_embeddings( + multimodal_params[:num_context_requests]) mm_embeds = find_input_mm_embeds( mm_embeds, multimodal_params[:num_context_requests]) input_ids, inputs_embeds = fuse_input_embeds( diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 7de6f2523d5f..7391920cf41c 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -25,7 +25,7 @@ from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( MultimodalEncoderOutput, MultimodalModelMixin, PreparedLlmInputs) from tensorrt_llm._torch.models.modeling_multimodal_utils import ( - _MULTIMODAL_ENV_NAME, _is_disagg) + _MULTIMODAL_ENV_NAME, _is_mm_disagg) from tensorrt_llm._torch.models.modeling_utils import (DecoderModel, DecoderModelForCausalLM, _load_weights_impl, @@ -567,7 +567,8 @@ def __init__( self, model_config: ModelConfig[Mistral3Config], ): - if _is_disagg(): + # No MM E/P handoff here yet. Fail before partial model setup. + if _is_mm_disagg(): raise NotImplementedError( "Mistral3VLM does not support disaggregated inference yet. Please unset " f"the {_MULTIMODAL_ENV_NAME} environment variable, or set it to '0'." diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py index 5615f33c1ec2..81d35a6132b9 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_utils.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_utils.py @@ -35,10 +35,18 @@ # Make this a runtime lookup rather than a module-wide constant for easier unit testing. -def _is_disagg() -> bool: +# MM E/P split flag. Not generic disaggregated serving. +def _is_mm_disagg() -> bool: return os.getenv(_MULTIMODAL_ENV_NAME, "0") == "1" +def has_raw_multimodal_payload(param: MultimodalParams) -> bool: + multimodal_data = param.multimodal_data or {} + modality_type = multimodal_data.get("modality_type") + return (modality_type in ("image", "video", "audio") + and multimodal_data.get(modality_type) is not None) + + # Processor *output* keys that transformers 5.x's # ``ProcessorMixin._merge_kwargs`` strictly rejects when they leak into # ``output_kwargs[]`` and reach ``validate_typed_dict``. They @@ -270,6 +278,34 @@ def get_multimodal_embeddings( return [all_embeddings] +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. This only makes the tensor list that + find_input_mm_embeds slices. + """ + 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 + # Some paths stash chunks. Slicer expects one tensor. + if isinstance(embeds, list): + embeds = torch.cat(embeds, dim=0) + 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 [torch.cat(attached_embeddings, dim=0)] + + def find_input_mm_embeds( mm_embeds: List[torch.Tensor], multimodal_params: List[MultimodalParams]) -> List[torch.Tensor]: @@ -291,10 +327,15 @@ def find_input_mm_embeds( 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]. """ + 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) @@ -317,6 +358,11 @@ def find_input_mm_embeds( ) 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): return mm_embeds diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_nano.py b/tensorrt_llm/_torch/models/modeling_nemotron_nano.py index c81407e77809..1ab1be554094 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_nano.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_nano.py @@ -1,7 +1,6 @@ # Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. import copy import math -import os import re from dataclasses import dataclass from typing import Any, ClassVar, Dict, List, Optional, Sequence, Tuple, Union @@ -14,7 +13,15 @@ from PIL import Image from tensorrt_llm._torch.models.checkpoints import NemotronHHfWeightMapper -from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.inputs.multimodal import ( + DisaggPrefillMultimodalInputs, + MultimodalParams, + _as_cpu_tensor, + _compute_mm_masks, + _find_mm_token_runs_from_mask, + _find_mm_token_start_pos_from_masks, + find_mm_token_lengths, +) from ...inputs import ( AudioData, @@ -24,10 +31,12 @@ MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, TextPrompt, + TokensPrompt, compute_retained_tokens_count, compute_retained_tokens_from_tubelet_budget, compute_retention_mask, register_input_processor, + support_multimodal_disaggregated, ) from ...logger import logger from ...sampling_params import SamplingParams @@ -35,13 +44,16 @@ from ..model_config import ModelConfig from .modeling_auto import AutoModelForCausalLM from .modeling_multimodal_utils import ( + _is_mm_disagg, find_input_mm_embeds, fuse_input_embeds, + get_attached_multimodal_embeddings, get_multimodal_embeddings, + has_raw_multimodal_payload, ) from .modeling_parakeet import ParakeetExtractor, ProjectedParakeet from .modeling_radio import RADIOVisionModel, calc_seq_lens -from .modeling_utils import register_auto_model +from .modeling_utils import register_auto_model, register_vision_encoder # Set max_num_tiles to 1 for video modality, to match the training behavior. VIDEO_MAX_NUM_TILES = 1 @@ -390,10 +402,6 @@ def stack(images: List[torch.Tensor], patch_size: int) -> torch.Tensor: # Make this a runtime lookup rather than a module-wide constant for easier unit testing. -def _is_disagg() -> bool: - return os.getenv("TLLM_MULTIMODAL_DISAGGREGATED", "0") == "1" - - class SquaredReLU(nn.Module): def forward(self, x): return torch.pow(torch.nn.functional.relu(x), 2) @@ -406,6 +414,7 @@ class NanoV2VLVisionEncoder(transformers.PreTrainedModel): def __init__(self, model_config: ModelConfig[transformers.PretrainedConfig]): config = model_config.pretrained_config super().__init__(config) + self.model_config = model_config self.image_size = config.force_image_size self.patch_size = config.patch_size self.num_image_token = int( @@ -878,6 +887,39 @@ def _video_tubelet_geometry(self, t: int, T: int, ih: int, iw: int) -> Tuple[int return num_tubelets, wh +class NanoV2VLMultimodalEncoder(NanoV2VLVisionEncoder): + """EPD-only encoder wrapper for Nano VL image/video handoff. + + Full Nano V3 can support more modalities through the full model path. + This wrapper is only for the mm_encoder_only EPD worker. It returns one + vision embedding tensor for image/video inputs and does not run Nano audio + or video-audio interleave logic. + """ + + def __init__(self, model_config: ModelConfig[transformers.PretrainedConfig], *args, **kwargs): + super().__init__(model_config) + + def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tensor]: + for param in multimodal_params: + modality_type = param.multimodal_data["modality_type"] + if modality_type == "audio": + # EPD encoder-only handoff does not own the Nano audio encoder. + raise NotImplementedError( + "NanoV2VL MultimodalEncoder currently supports image/video inputs, not audio." + ) + audio_data = param.multimodal_data[modality_type].get("audio") + if audio_data is not None: + # TODO(TRTLLM-13129): Add audio support for encoder handoff. + raise NotImplementedError( + "NanoV2VL MultimodalEncoder does not yet encode audio extracted from video." + ) + + mm_embeddings, _ = super().forward(multimodal_params) + if not mm_embeddings: + return [] + return [torch.cat(mm_embeddings, dim=0)] + + class NanoV2VLInputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): supports_token_id_mm_expansion: ClassVar[bool] = True @@ -2229,6 +2271,93 @@ def call_with_text_prompt( "multimodal_data": multimodal_data, } + def build_disagg_prefill_multimodal_inputs( + self, inputs: Union[TextPrompt, TokensPrompt], mm_handles: List[Dict[str, Any]] + ) -> DisaggPrefillMultimodalInputs: + text_prompt = inputs.get("prompt") + prompt_token_ids = inputs.get("prompt_token_ids") + if prompt_token_ids is None and not text_prompt: + raise ValueError("Either prompt_token_ids or text prompt is required") + if not isinstance(mm_handles, list): + raise TypeError("mm_handles must be a list") + + mm_data = inputs.get("multi_modal_data") or {} + if not mm_data: + raise ValueError("multi_modal_data is required for NanoV2VL multimodal handoff") + modalities = [name for name, value in mm_data.items() if value is not None] + if len(modalities) != 1: + raise ValueError( + "NanoV2VL multimodal handoff supports exactly one modality per request" + ) + if modalities[0] == "audio": + raise NotImplementedError( + "NanoV2VL multimodal handoff does not support audio-only inputs" + ) + + num_mm_tokens_by_key = find_mm_token_lengths(mm_data, self) + num_mm_tokens = [length for lengths in num_mm_tokens_by_key.values() for length in lengths] + if len(num_mm_tokens) != len(mm_handles): + raise RuntimeError( + f"Expected {len(num_mm_tokens)} multimodal handles, got {len(mm_handles)}." + ) + + expected_hidden_size = self.config.llm_config.hidden_size + multimodal_embedding_lengths: List[int] = [] + for i, mm_handle in enumerate(mm_handles): + tensor_size = mm_handle["tensor_size"] + if len(tensor_size) != 2: + raise RuntimeError( + f"Expected multimodal embedding {i} to be rank 2, got tensor_size={tensor_size}." + ) + if tensor_size[1] != expected_hidden_size: + raise RuntimeError( + f"Expected multimodal embedding {i} to have hidden size " + f"{expected_hidden_size}, got {tensor_size[1]}." + ) + multimodal_embedding_lengths.append(tensor_size[0]) + + if prompt_token_ids is None: + prompt_token_ids = self.tokenizer.encode(text_prompt, add_special_tokens=False) + prompt_token_ids = list(prompt_token_ids) + + expanded_ids, _ = self.expand_prompt_token_ids_for_mm( + prompt_token_ids, + num_mm_tokens, + hf_processor_mm_kwargs=inputs.get("mm_processor_kwargs"), + mm_data=mm_data, + ) + + input_ids_tensor = _as_cpu_tensor(expanded_ids) + mm_mask, embed_mask, special_mask = _compute_mm_masks( + input_ids_tensor, + vocab_size=self.get_vocab_size(), + mm_token_ids=self.get_mm_token_ids(), + mm_special_token_ids=self.get_mm_special_token_ids(), + ) + if int(embed_mask.sum().item()) != sum(multimodal_embedding_lengths): + raise RuntimeError( + "Multimodal embedding length mismatch: " + f"prompt has {int(embed_mask.sum().item())} embedding slots, " + f"handles provide {sum(multimodal_embedding_lengths)}." + ) + mm_token_offsets, special_token_offsets = _find_mm_token_start_pos_from_masks( + mm_mask, special_mask, num_mm_tokens + ) + item_run_cu_offsets, run_positions, run_lengths = _find_mm_token_runs_from_mask( + mm_mask, num_mm_tokens + ) + + return DisaggPrefillMultimodalInputs( + prompt_token_ids=expanded_ids, + multimodal_lengths=num_mm_tokens, + multimodal_positions=mm_token_offsets, + multimodal_embedding_lengths=multimodal_embedding_lengths, + multimodal_item_run_cu_offsets=item_run_cu_offsets, + multimodal_run_positions=run_positions, + multimodal_run_lengths=run_lengths, + special_token_offsets=special_token_offsets, + ) + def _prepare_audio_features( self, text: str, @@ -2423,6 +2552,8 @@ def _resample_audios( ) +@support_multimodal_disaggregated +@register_vision_encoder(NanoV2VLMultimodalEncoder) @register_auto_model("NemotronH_Nano_Omni_Reasoning_V3") @register_auto_model("NemotronH_Nano_VL_V2") @register_input_processor( @@ -2439,9 +2570,6 @@ class NemotronH_Nano_VL_V2(transformers.PreTrainedModel): _supports_flash_attn = True def __init__(self, model_config: ModelConfig): - if _is_disagg(): - raise ValueError("NanoV2VL does not support disaggregated inference yet.") - config = model_config.pretrained_config super().__init__(config) @@ -2493,10 +2621,12 @@ def load_weights(self, weights): # to be the LLM-only config and no longer has vision_config / # sound_config / force_image_size / etc. mm_pretrained = self._mm_model_config.pretrained_config - if self.vision_encoder is None and not _is_disagg(): + # Normal workers own encoders. MM E/P handoff uses attached embeddings. + is_multimodal_encoder_worker = not _is_mm_disagg() + if self.vision_encoder is None and is_multimodal_encoder_worker: self.vision_encoder = NanoV2VLVisionEncoder(self._mm_model_config).eval().to("cuda") sound_config = getattr(mm_pretrained, "sound_config", None) - if self.sound_encoder is None and sound_config is not None: + if self.sound_encoder is None and sound_config is not None and is_multimodal_encoder_worker: self.sound_encoder = ( ProjectedParakeet( sound_config, @@ -2681,6 +2811,24 @@ def _validate_evs_context_batch( "multimodal context chunks form a contiguous input_ids prefix." ) + def _check_encoders_exist(self, raw_ctx_params: List[MultimodalParams]) -> None: + """Check encoders needed by raw inputs exist. + + Raw image/video needs vision encoder; raw audio needs sound encoder. + Encoder-only EPD worker may have only one. Reject early with clear + message, not deep encoder failure. + """ + needs_vision_encoder = any( + param.multimodal_data["modality_type"] in ("image", "video") for param in raw_ctx_params + ) + if needs_vision_encoder and self.vision_encoder is None: + raise ValueError("Raw image/video inputs require a local NanoV2VL vision encoder.") + needs_sound_encoder = any( + param.multimodal_data["modality_type"] == "audio" for param in raw_ctx_params + ) + if needs_sound_encoder and self.sound_encoder is None: + raise ValueError("Raw audio inputs require a local NanoV2VL sound encoder.") + def merge_evs_mm_embeds( self, num_tokens_in_videos: List[int], @@ -2996,16 +3144,27 @@ def forward( ctx_params = multimodal_params[:num_context_requests] if self.video_pruning_rate > 0: self._validate_evs_context_batch(ctx_params, num_context_requests) - if not _is_disagg(): + raw_ctx_params = [param for param in ctx_params if has_raw_multimodal_payload(param)] + # Raw image/video/audio tensors: run local encoder. + if raw_ctx_params: + self._check_encoders_exist(raw_ctx_params) mm_embedding = get_multimodal_embeddings( encoder_forward_fn=self._encode_multimodal, multimodal_params=ctx_params, ) + # E/P prefill: encoder already ran; use attached embeddings. else: - raise NotImplementedError( - "Nano-V2-VLM does not support disaggregated inference yet. Please unset " - "the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'." - ) + if self.video_pruning_rate > 0 and any( + param.has_content() and param.multimodal_data.get("modality_type") == "video" + for param in ctx_params + ): + # TODO(TRTLLM-12534): Carry EVS retained-token counts through + # encoder handoff before enabling video pruning for E/P. + raise ValueError( + "EVS video pruning is not supported with attached " + "multimodal embeddings yet." + ) + mm_embedding = get_attached_multimodal_embeddings(ctx_params) # Adjust input_ids in videos if EVS is applied. if self.video_pruning_rate > 0: # Retrieve per-video count stashed by `_encode_multimodal`. diff --git a/tensorrt_llm/_torch/models/modeling_phi4mm.py b/tensorrt_llm/_torch/models/modeling_phi4mm.py index abcc0c1f5ff2..df9afdf34508 100644 --- a/tensorrt_llm/_torch/models/modeling_phi4mm.py +++ b/tensorrt_llm/_torch/models/modeling_phi4mm.py @@ -42,7 +42,8 @@ from ..attention_backend import AttentionMetadata from ..model_config import ModelConfig from .modeling_auto import AutoModelForCausalLM -from .modeling_multimodal_utils import (find_input_mm_embeds, fuse_input_embeds, +from .modeling_multimodal_utils import (_is_mm_disagg, find_input_mm_embeds, + fuse_input_embeds, get_multimodal_embeddings) from .modeling_utils import register_auto_model @@ -73,10 +74,6 @@ def _is_torch_compile() -> bool: return os.getenv("TLLM_MULTIMODAL_ENCODER_TORCH_COMPILE", "0") == "1" -def _is_disagg() -> bool: - return os.getenv("TLLM_MULTIMODAL_DISAGGREGATED", "0") == "1" - - # Load the Phi4MM classes from HuggingFace Phi-4-multimodal-instruct repo. # Remove this function by using the transformers version of Phi4Multimodal when weights/configs are converted to transformers format. def _load_phi4mm_classes(local_path): @@ -957,7 +954,7 @@ class Phi4MMForCausalLM(transformers.PreTrainedModel): _supports_flash_attn = True def __init__(self, model_config: ModelConfig): - if _is_disagg(): + if _is_mm_disagg(): raise ValueError( "Phi4MM does not support disaggregated inference yet.") @@ -968,7 +965,7 @@ def __init__(self, model_config: ModelConfig): if hasattr(self, "llm"): return - if not _is_disagg(): + if not _is_mm_disagg(): _load_phi4mm_classes(config._name_or_path) self.hf_phi4mm_model = HFPhi4MultimodalEncoder(config).eval() @@ -989,7 +986,7 @@ def __init__(self, model_config: ModelConfig): def load_weights(self, weights): # Load weights into HFPhi4MultimodalEncoder. - if not _is_disagg(): + if not _is_mm_disagg(): filtered_weights = {} for k, v in weights.items(): # Skip image_embed head weights since we set it as NoOp. @@ -1076,7 +1073,7 @@ def forward( multimodal_params = kwargs.get("multimodal_params", []) mm_embedding = [] if len(multimodal_params) > 0: - if not _is_disagg(): + if not _is_mm_disagg(): encoder_kwargs = { "mm_token_ids": self.mm_token_ids, } diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 7db60f085dfd..05fa290d105c 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -22,12 +22,13 @@ BaseWeightMapper from tensorrt_llm._torch.models.checkpoints.hf.qwen2vl_weight_mapper import \ Qwen2VLHfWeightMapper -from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_disagg +from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg from tensorrt_llm._torch.modules.attention import Attention from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, + MultimodalParams) from ..._utils import nvtx_range, prefer_pinned from ...inputs import (BaseMultimodalDummyInputsBuilder, @@ -55,6 +56,7 @@ from .modeling_auto import AutoModelForCausalLM from .modeling_multimodal_utils import (bypass_processor_output_validation, find_input_mm_embeds, fuse_input_embeds, + get_attached_multimodal_embeddings, get_multimodal_embeddings) from .modeling_utils import (ModelConfig, QuantConfig, _load_weights_impl, filter_weights, register_auto_model, @@ -1075,7 +1077,8 @@ def __init__( llm_model_config.pretrained_config.architectures = ["Qwen2ForCausalLM"] self.llm = AutoModelForCausalLM.from_config(llm_model_config) - if not _is_disagg(): + # Normal worker owns encoder. MM E/P prefill worker gets attached embeddings. + if not _is_mm_disagg(): mm_encoder_config = copy.deepcopy(model_config) self.mm_encoder = Qwen2VisionModelBase( mm_encoder_config, kwargs.get('vision_model_class', None)) @@ -1191,7 +1194,8 @@ def forward( mm_multimodal_params = self._get_requests_with_mm_data( multimodal_params) if len(mm_multimodal_params) > 0: - if not _is_disagg(): + # Local encoder present: raw pixels/videos become embeddings here. + if self.mm_encoder is not None: mm_embeds = get_multimodal_embeddings( encoder_forward_fn=self.mm_encoder.forward, multimodal_params=mm_multimodal_params) @@ -1200,6 +1204,10 @@ def forward( "Qwen2VLModel does not support disaggregated inference yet. Please unset " f"the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'." ) + # E/P prefill: encoder already ran; use attached embeddings. + else: + mm_embeds = get_attached_multimodal_embeddings( + mm_multimodal_params) mm_embeds = find_input_mm_embeds(mm_embeds, mm_multimodal_params) if not self.model_config.pretrained_config.disable_fuse_rope: @@ -1267,7 +1275,7 @@ def multimodal_data_device_paths(self) -> List[str]: ] def load_weights(self, weights, weight_mapper: BaseWeightMapper): - if not _is_disagg(): + if self.mm_encoder is not None: self.mm_encoder.load_weights(weights) self.llm.load_weights(weights, weight_mapper) @@ -1275,22 +1283,20 @@ def load_weights(self, weights, weight_mapper: BaseWeightMapper): class Qwen2_5VLInputProcessorBase(Qwen2VLInputProcessorBase): - def get_prompt_token_ids( - self, inputs: TextPrompt, - mm_handles: List[Dict[str, - Any]]) -> Tuple[List[int], List[int], List[int]]: + def build_disagg_prefill_multimodal_inputs( + self, inputs: TextPrompt, + mm_handles: List[Dict[str, Any]]) -> DisaggPrefillMultimodalInputs: """ - Build input token ids with multimodal placeholders expanded to the number of MM tokens. + Build disaggregated prefill inputs from multimodal embedding handles. Args: inputs: Text prompt input container. Must contain a non-empty prompt string. mm_handles: List of multimodal embedding handles. Returns: - Tuple[List[int], List[int], List[int]]: - - expanded_ids: token ids with each image token expanded to a placeholder repeated per MM token - - mm_token_length: per-image MM token lengths - - mm_token_offsets: start offsets (positions) for each image's MM tokens within expanded_ids + DisaggPrefillMultimodalInputs containing expanded token IDs, + prompt-side MM positions/lengths, exact runs, and encoder-output + embedding lengths. """ # TODO: Move this function to the base input processor class when extending for more models text_prompt = inputs.get("prompt") @@ -1347,8 +1353,18 @@ def get_prompt_token_ids( assert write_pos == final_length, f"Write position mismatch: {write_pos} != {final_length}" assert mm_token_length[-1] + mm_token_offsets[ -1] <= final_length, f"mm_token_length[-1] + mm_token_offsets[-1] ({mm_token_length[-1] + mm_token_offsets[-1]}) should be less than or equal to final_length ({final_length})" - return expanded_ids.to( - torch.int32).tolist(), mm_token_length, mm_token_offsets + return DisaggPrefillMultimodalInputs( + prompt_token_ids=expanded_ids.to(torch.int32).tolist(), + multimodal_lengths=mm_token_length, + multimodal_positions=mm_token_offsets, + multimodal_embedding_lengths=[ + mm_handle["tensor_size"][0] for mm_handle in mm_handles + ], + multimodal_item_run_cu_offsets=list(range(len(mm_token_length) + + 1)), + multimodal_run_positions=mm_token_offsets, + multimodal_run_lengths=mm_token_length, + ) @support_multimodal_disaggregated @@ -1387,7 +1403,7 @@ def load_weights(self, weights, weight_mapper: BaseWeightMapper): if isinstance(weight_mapper, Qwen2VLHfWeightMapper): weights = weight_mapper.preprocess_weights(weights) - if not _is_disagg(): + if self.mm_encoder is not None: self.mm_encoder.load_weights(weights) self.llm.load_weights(weights) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 38205de69687..f2553a2e2b0d 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -13,7 +13,7 @@ Qwen3VLVisionPatchEmbed as HFQwen3VLVisionPatchEmbed, ) -from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_disagg +from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.mapping import Mapping @@ -29,7 +29,7 @@ register_input_processor, support_multimodal_disaggregated, ) -from ...inputs.multimodal import MultimodalParams +from ...inputs.multimodal import DisaggPrefillMultimodalInputs, MultimodalParams from ...logger import logger from ...sampling_params import SamplingParams from ..attention_backend import AttentionMetadata @@ -46,6 +46,7 @@ bypass_processor_output_validation, find_input_mm_embeds, fuse_input_embeds, + get_attached_multimodal_embeddings, get_multimodal_embeddings, ) from .modeling_qwen2vl import Qwen2_5_VLVisionAttention @@ -59,6 +60,99 @@ ) +def _expand_prompt_token_ids_for_mm_handoff( + input_ids: torch.Tensor, + mm_handles: List[Dict[str, Any]], + *, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, + placeholder_id: int, +) -> DisaggPrefillMultimodalInputs: + """Expand Qwen3-VL image/video placeholders and emit sparse MM layout. + + Qwen handoff has one coarse or token per item. + This helper expands that one token to the number of embedding rows in the + handoff handle, then returns the sparse layout metadata. + + Agg gets this expansion from Qwen's HF processor taking raw images/videos + as inputs. Reusing that would be wasteful here, hence this helper that + expands based on the embedding handles row count. + + """ + placeholder_positions = [ + pos + for pos, token in enumerate(input_ids.tolist()) + if token in (image_token_id, video_token_id) + ] + if len(placeholder_positions) != len(mm_handles): + raise ValueError( + "Number of multimodal placeholders must match number of mm_handles: " + f"placeholders={len(placeholder_positions)}, " + f"mm_handles={len(mm_handles)}" + ) + + total_mm_embed_tokens = sum(mm_handle["tensor_size"][0] for mm_handle in mm_handles) + final_length = len(input_ids) - len(placeholder_positions) + total_mm_embed_tokens + expanded_ids = torch.empty(final_length, dtype=input_ids.dtype) + + mm_token_lengths: List[int] = [] + mm_token_offsets: List[int] = [] + item_types: List[int] = [] + item_run_cu_offsets: List[int] = [0] + run_positions: List[int] = [] + run_lengths: List[int] = [] + multimodal_embedding_lengths: List[int] = [] + special_token_offsets: List[int] = [] + + write_pos = 0 + mm_handle_idx = 0 + flat_mm_offset = 0 + for read_pos, token_id in enumerate(input_ids.tolist()): + if token_id not in (image_token_id, video_token_id): + expanded_ids[write_pos] = token_id + write_pos += 1 + continue + + mm_token_num = mm_handles[mm_handle_idx]["tensor_size"][0] + has_leading_special = ( + read_pos > 0 and int(input_ids[read_pos - 1].item()) == vision_start_token_id + ) + run_start = write_pos - 1 if has_leading_special else write_pos + prompt_mm_length = mm_token_num + int(has_leading_special) + + expanded_ids[write_pos : write_pos + mm_token_num] = placeholder_id + mm_token_offsets.append(run_start) + mm_token_lengths.append(prompt_mm_length) + multimodal_embedding_lengths.append(mm_token_num) + item_types.append(0 if token_id == image_token_id else 1) + run_positions.append(run_start) + run_lengths.append(prompt_mm_length) + item_run_cu_offsets.append(len(run_positions)) + + if has_leading_special: + special_token_offsets.append(flat_mm_offset) + + write_pos += mm_token_num + flat_mm_offset += prompt_mm_length + mm_handle_idx += 1 + + if write_pos != final_length: + raise RuntimeError(f"Write position mismatch: {write_pos} != {final_length}") + + return DisaggPrefillMultimodalInputs( + prompt_token_ids=expanded_ids.to(torch.int32).tolist(), + multimodal_lengths=mm_token_lengths, + multimodal_positions=mm_token_offsets, + multimodal_embedding_lengths=multimodal_embedding_lengths, + multimodal_item_run_cu_offsets=item_run_cu_offsets, + multimodal_run_positions=run_positions, + multimodal_run_lengths=run_lengths, + special_token_offsets=special_token_offsets, + item_types=item_types, + ) + + class Qwen3VLInputProcessorBase(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): def __init__( self, @@ -264,31 +358,16 @@ def get_num_tokens_per_video( video_grid_thw: Optional[torch.Tensor] = None, **kwargs, ) -> int: + if video_grid_thw is None: + raise ValueError( + "Qwen3-VL video token count requires processor-produced video_grid_thw" + ) + merge = self.config.vision_config.spatial_merge_size - if video_grid_thw is not None: - t, h, w = (int(x) for x in video_grid_thw) - return t * (h // merge) * (w // merge) - - # Must run the full processor: HF's Qwen3VLProcessor._get_num_multimodal_tokens - # (what the base class default delegates to) raises on video-only calls - # and returns a wrong-formula fallback that would break chunked prefill. - do_rescale = not (video and isinstance(video[0], torch.Tensor)) - processed = self._processor( - text=["<|vision_start|><|video_pad|><|vision_end|>"], - videos=[video], - padding=True, - do_rescale=do_rescale, - return_tensors="pt", - **kwargs, + token_counts = ( + video_grid_thw[:, 0] * (video_grid_thw[:, 1] // merge) * (video_grid_thw[:, 2] // merge) ) - vgt = processed.get("video_grid_thw") - if vgt is None or len(vgt) == 0: - raise RuntimeError( - "get_num_tokens_per_video: HF processor returned no " - "video_grid_thw for the provided video." - ) - t, h, w = (int(x) for x in vgt[0].tolist()) - return t * (h // merge) * (w // merge) + return int(token_counts.sum().item()) def _preprocess( self, text: Dict[str, Any], mm_data: Dict[str, Any], mm_processor_kwargs: Dict[str, Any] @@ -396,21 +475,20 @@ def call_with_text_prompt( "multimodal_data": multimodal_data, } - def get_prompt_token_ids( + def build_disagg_prefill_multimodal_inputs( self, inputs: TextPrompt, mm_handles: List[Dict[str, Any]] - ) -> Tuple[List[int], List[int], List[int]]: + ) -> DisaggPrefillMultimodalInputs: """ - Build input token ids with multimodal placeholders expanded to the number of MM tokens. + Build disaggregated prefill inputs from multimodal embedding handles. Args: inputs: Text prompt input container. Must contain a non-empty prompt string. mm_handles: List of multimodal embedding handles. Returns: - Tuple[List[int], List[int], List[int]]: - - expanded_ids: token ids with each image token expanded to a placeholder repeated per MM token - - mm_token_length: per-image MM token lengths - - mm_token_offsets: start offsets (positions) for each image's MM tokens within expanded_ids + DisaggPrefillMultimodalInputs containing expanded token IDs, + prompt-side MM positions/lengths, exact runs, and encoder-output + embedding lengths. """ # TODO: Move this function to the base input processor class when extending for more models text_prompt = inputs.get("prompt") @@ -433,44 +511,14 @@ def get_prompt_token_ids( input_ids = self.tokenizer(text_prompt, return_tensors="pt").input_ids[0] - # TODO: what about `video_token_id`? - image_token_index = self.config.image_token_id - - image_mask = input_ids == image_token_index - image_positions = torch.where(image_mask)[0] - num_images = len(image_positions) - assert num_images == len(mm_handles), "Number of images must match number of mm_handles" - total_mm_tokens = sum(mm_handle["tensor_size"][0] for mm_handle in mm_handles) - final_length = len(input_ids) - num_images + total_mm_tokens - # Create output tensor - expanded_ids = torch.empty(final_length, dtype=input_ids.dtype) - placeholder_id = self.tllm_multimodal_token_id - - # Fill the expanded sequence - write_pos = 0 - image_cnt = 0 - mm_token_length = [] - mm_token_offsets = [] - for read_pos in range(len(input_ids)): - if input_ids[read_pos] == image_token_index: - # Replace with placeholder id - mm_token_num = mm_handles[image_cnt]["tensor_size"][0] - expanded_ids[write_pos : write_pos + mm_token_num] = placeholder_id - mm_token_offsets.append(write_pos) - mm_token_length.append(mm_token_num) - write_pos += mm_token_num - image_cnt += 1 - else: - # Copy text token as-is - expanded_ids[write_pos] = input_ids[read_pos] - write_pos += 1 - - assert write_pos == final_length, f"Write position mismatch: {write_pos} != {final_length}" - assert mm_token_length[-1] + mm_token_offsets[-1] <= final_length, ( - f"mm_token_length[-1] + mm_token_offsets[-1] ({mm_token_length[-1] + mm_token_offsets[-1]}) should be less " - f"than or equal to final_length ({final_length})" + return _expand_prompt_token_ids_for_mm_handoff( + input_ids, + mm_handles, + image_token_id=self.config.image_token_id, + video_token_id=self.config.video_token_id, + vision_start_token_id=self.config.vision_start_token_id, + placeholder_id=self.tllm_multimodal_token_id, ) - return expanded_ids.to(torch.int32).tolist(), mm_token_length, mm_token_offsets class Qwen3VLVisionAttention(Qwen2_5_VLVisionAttention): @@ -1058,7 +1106,9 @@ def __init__( # Qwen3ForCausalLM. self.llm = AutoModelForCausalLM.from_config(llm_model_config) - if not _is_disagg(): + self.mm_encoder = None + # Normal workers own the encoder. MM E/P handoff uses attached embeddings. + if not _is_mm_disagg(): self.mm_encoder = Qwen3VisionModelBase( copy.deepcopy(model_config), kwargs.get("vision_model_class", None) ).eval() @@ -1188,18 +1238,31 @@ def forward( # NOTE: Qwen*-VL series has mrope_config even on the text-only prompts, # so we need to separate the mm_multimodal_params from the text-only prompts. - mm_multimodal_params = self._get_requests_with_mm_data(multimodal_params) + mm_multimodal_params, has_raw_image_or_video_data = self._get_requests_with_mm_data( + multimodal_params + ) if len(mm_multimodal_params) > 0: - if not _is_disagg(): + # Raw image/video tensors: run local encoder. + if has_raw_image_or_video_data and self.mm_encoder is not None: mm_embeds = get_multimodal_embeddings( encoder_forward_fn=self.mm_encoder.forward, multimodal_params=mm_multimodal_params, ) + # Raw image/video tensors on a worker with no encoder: bad route. + elif has_raw_image_or_video_data: + raise ValueError( + "Raw multimodal inputs require a local multimodal encoder on this " + "worker, or multimodal_embedding handles from an encoder handoff." + ) + # support_mm_disagg is only set in subclasses of Qwen3VLModelBase that support EPD elif not getattr(self, "support_mm_disagg", False): raise NotImplementedError( f"{type(self)} does not support disaggregated inference yet. Please unset " "the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'." ) + # E/P prefill: encoder already ran; use attached embeddings. + else: + mm_embeds = get_attached_multimodal_embeddings(mm_multimodal_params) mm_embeds = find_input_mm_embeds(mm_embeds, mm_multimodal_params) if self.use_deepstack: @@ -1239,19 +1302,22 @@ def forward( def _get_requests_with_mm_data(self, multimodal_params): mm_multimodal_params = [] + # TODO: This returns one batch-wide "has raw pixels/video" flag. That is + # safe only when a batch is all raw-MM or all attached embeddings. If a + # scheduler can mix both, split raw requests from attached-embedding + # requests and merge outputs back by request index. + has_raw_image_or_video_data = False for multimodal_param in multimodal_params: data = multimodal_param.multimodal_data - if ( - # The first 2 conditions check whether there is input on which inference should be run. + has_raw_data = ( data.get("image", {}).get("pixel_values") is not None or data.get("video", {}).get("pixel_values_videos") is not None - # This condition corresponds to when the embeddings are already populated, as is e.g. - # the case in EPD disagg in the prefill worker. - or data.get("multimodal_embedding") is not None - ): + ) + has_raw_image_or_video_data |= has_raw_data + if has_raw_data or data.get("multimodal_embedding") is not None: mm_multimodal_params.append(multimodal_param) - return mm_multimodal_params + return mm_multimodal_params, has_raw_image_or_video_data @support_multimodal_disaggregated @@ -1284,7 +1350,7 @@ def multimodal_data_device_paths(self) -> List[str]: return ["image.pixel_values", "video.pixel_values_videos", "multimodal_embedding"] def load_weights(self, weights: Dict[str, torch.Tensor], weight_mapper: BaseWeightMapper): - if not _is_disagg(): + if self.mm_encoder is not None: self.mm_encoder.load_weights(weights) weight_mapper = Qwen3VLHfWeightMapper() diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py index 4736960cf594..e26a9c317cc2 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl_moe.py @@ -3,8 +3,6 @@ import torch from transformers import PretrainedConfig -from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_disagg - from ...inputs import ( ContentFormat, MultimodalPlaceholderMetadata, @@ -64,7 +62,7 @@ def multimodal_data_device_paths(self) -> List[str]: ] def load_weights(self, weights: Dict[str, torch.Tensor], weight_mapper: BaseWeightMapper): - if not _is_disagg(): + if self.mm_encoder is not None: self.mm_encoder.load_weights(weights) weight_mapper = Qwen3VLMoeHfWeightMapper() diff --git a/tensorrt_llm/_torch/models/modeling_step3p7vl.py b/tensorrt_llm/_torch/models/modeling_step3p7vl.py index 7f4027b309f4..9f74b634244d 100644 --- a/tensorrt_llm/_torch/models/modeling_step3p7vl.py +++ b/tensorrt_llm/_torch/models/modeling_step3p7vl.py @@ -57,7 +57,7 @@ from ..modules.layer_norm import LayerNorm from ..speculative import SpecMetadata from .modeling_multimodal_utils import ( - _is_disagg, + _is_mm_disagg, find_input_mm_embeds, fuse_input_embeds, get_multimodal_embeddings, @@ -946,7 +946,7 @@ def load_weights( allow_partial_loading: bool = False, ): """Split vision/text weights and delegate to the inner LM loader.""" - if self.mm_encoder is None and not _is_disagg() and hasattr(weights, "items"): + if self.mm_encoder is None and not _is_mm_disagg() and hasattr(weights, "items"): # Construct the vision tower here, outside MetaInitMode, so its # PerceptionEncoder / HF submodules allocate real tensors. Move it # straight to CUDA (model_loader already ran model.to("cuda") for diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 69256521d364..03dbf36d0bc0 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -755,12 +755,13 @@ class SomeVLModel(...): """ def wrapper(model_cls: Type[nn.Module]) -> Type[nn.Module]: + registered = False for arch_name, registered_cls in MODEL_CLASS_MAPPING.items(): - if registered_cls.__name__ == model_cls.__name__: + if registered_cls is model_cls: MODEL_CLASS_VISION_ENCODER_MAPPING[arch_name] = ( vision_encoder_cls, vlm_base_model) - break - else: + registered = True + if not registered: raise ValueError( f"register_vision_encoder: model class {model_cls.__name__} is not registered " f"via register_auto_model; decorator order must ensure registration occurs first." diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 31b0c8d28328..1e96f88621b5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -7,6 +7,7 @@ import tensorrt_llm import tensorrt_llm.bindings.executor as trtllm +from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg from tensorrt_llm._torch.models.modeling_utils import \ MODEL_CLASS_VISION_ENCODER_MAPPING from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, @@ -392,18 +393,8 @@ def _create_dummy_mm_context_request( multimodal_input = extra_processed_inputs.get( 'multimodal_input') multimodal_data = extra_processed_inputs.get('multimodal_data') - req_mm_input = trtllm.MultimodalInput( - multimodal_hashes=multimodal_input.multimodal_hashes, - multimodal_positions=multimodal_input.multimodal_positions, - multimodal_lengths=multimodal_input.multimodal_lengths, - multimodal_uuids=multimodal_input.multimodal_uuids, - multimodal_item_run_cu_offsets=multimodal_input. - multimodal_item_run_cu_offsets, - multimodal_run_positions=multimodal_input. - multimodal_run_positions, - multimodal_run_lengths=multimodal_input. - multimodal_run_lengths, - ) if multimodal_input else None + req_mm_input = multimodal_input.to_binding( + trtllm) if multimodal_input else None request = trtllm.Request(prompt_token_ids, max_tokens=1, @@ -445,9 +436,12 @@ def _create_dummy_mm_context_request( def _create_dummy_context_requests( self, input_seq_len: int) -> List[trtllm.Request]: requests = [] - if hasattr(self._model_engine.model, - "original_arch") and MODEL_CLASS_VISION_ENCODER_MAPPING.get( - self._model_engine.model.original_arch, None): + # Disaggregated workers receive multimodal embeddings instead of raw + # pixel inputs, so capacity probing must use the text-only fallback. + if (not _is_mm_disagg() + and hasattr(self._model_engine.model, "original_arch") + and MODEL_CLASS_VISION_ENCODER_MAPPING.get( + self._model_engine.model.original_arch, None)): requests = self._create_dummy_mm_context_request(input_seq_len) # if succeed profiling with multimodal requests then return, otherwise profile # with default case diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 79cff4fde8af..16db68c79f9d 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -395,21 +395,18 @@ def append_log_probs(self, self._log_probs.append(log_probs, cum_log_probs) def append_mm_embeddings(self, mm_embeddings: torch.Tensor, - multimodal_lengths: List[int]): + mm_embedding_lengths: List[int]): """Split concatenated embeddings by per-item lengths and create handles. Args: mm_embeddings: Concatenated multimodal embeddings tensor of shape [total_tokens, hidden_dim]. - multimodal_lengths: Current per-item split lengths. + mm_embedding_lengths: Per-item encoder-output embedding lengths. """ - # TODO(TRTLLM-12175): callers currently pass request.multimodal_lengths, - # a prompt-side MM-token count that may include non-embedding - # special/framing tokens. This split needs per-item encoder-output - # embedding lengths instead. - split_embeddings = torch.split(mm_embeddings, multimodal_lengths, dim=0) + split_embeddings = torch.split(mm_embeddings, + mm_embedding_lengths, + dim=0) - # Create a SharedTensorContainer handle for each split self._mm_embeddings = [ SharedTensorContainer.from_tensor(emb).dump_to_dict() for emb in split_embeddings @@ -421,10 +418,10 @@ def set_mrope_position( mrope_position_ids: torch.Tensor, mrope_position_deltas: torch.Tensor, ): - self._mrope_position_ids = (SharedTensorContainer.from_tensor( - mrope_position_ids).dump_to_dict()) - self._mrope_position_deltas = (SharedTensorContainer.from_tensor( - mrope_position_deltas).dump_to_dict()) + self._mrope_position_ids = SharedTensorContainer.from_tensor( + mrope_position_ids).dump_to_dict() + self._mrope_position_deltas = SharedTensorContainer.from_tensor( + mrope_position_deltas).dump_to_dict() self.diff.mrope_position_ids = self._mrope_position_ids self.diff.mrope_position_deltas = self._mrope_position_deltas @@ -660,7 +657,8 @@ def __init__( return_perf_metrics=return_perf_metrics, stop_words_list=torch.tensor(stop_words_list, dtype=torch.int32) if stop_words_list else None, - **kwargs) + **kwargs, + ) self.py_client_id = client_id self.py_request_id = self.request_id self.py_llm_request_type = self.llm_request_type @@ -941,6 +939,50 @@ def convert_wordlist(word_list) -> List[List[int]]: return [tokens, offsets] +def _validate_optional_int_list(values: Any, + field_name: str) -> Optional[List[int]]: + if values is None: + return None + if not isinstance(values, list): + raise TypeError(f"{field_name} must be a list") + if not all(isinstance(value, int) for value in values): + raise TypeError(f"{field_name} must contain only integers") + return values + + +def get_multimodal_embedding_lengths( + request: LlmRequest) -> Optional[List[int]]: + """Return explicit per-item encoder-output lengths for a multimodal request.""" + py_multimodal_data = request.py_multimodal_data + if py_multimodal_data is not None and not isinstance( + py_multimodal_data, dict): + raise TypeError("py_multimodal_data must be a dict") + # `multimodal_embedding_lengths` is Python-side layout metadata, not a + # nanobind request field, so validate the flat handoff contract here. + multimodal_embedding_lengths = _validate_optional_int_list( + py_multimodal_data.get("multimodal_embedding_lengths") + if py_multimodal_data is not None else None, + "multimodal_embedding_lengths") + if multimodal_embedding_lengths is None: + return None + + if any(length < 0 for length in multimodal_embedding_lengths): + raise ValueError("multimodal_embedding_lengths must be non-negative") + multimodal_lengths = request.multimodal_lengths + if multimodal_lengths is not None: + if len(multimodal_embedding_lengths) != len(multimodal_lengths): + raise ValueError("multimodal_embedding_lengths length must match " + "multimodal_lengths") + for item_idx, (embedding_length, prompt_length) in enumerate( + zip(multimodal_embedding_lengths, multimodal_lengths)): + if embedding_length > prompt_length: + raise ValueError( + f"multimodal_embedding_lengths[{item_idx}] exceeds " + f"multimodal_lengths[{item_idx}]") + + return multimodal_embedding_lengths + + def executor_request_to_llm_request( req_id: int, executor_request: ExecutorRequest, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 789e2155da97..8c0ee9d0edf4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -20,6 +20,7 @@ from tensorrt_llm.bindings.internal.runtime import TaskLayerModuleConfig from tensorrt_llm.inputs.multimodal import (MultimodalParams, MultimodalRuntimeData, + _has_mm_payload_keys, check_mm_embed_cumsum_if_needed) from tensorrt_llm.inputs.registry import (create_input_processor, create_input_processor_with_hash) @@ -67,7 +68,8 @@ EncoderCUDAGraphRunnerConfig) from .guided_decoder import CapturableGuidedDecoder from .layerwise_nvtx_marker import LayerwiseNvtxMarker -from .llm_request import LlmRequest, get_draft_token_length +from .llm_request import (LlmRequest, get_draft_token_length, + get_multimodal_embedding_lengths) from .mamba_cache_manager import MambaHybridCacheManager from .model_loader import ModelLoader, _construct_checkpoint_loader from .resource_manager import (BaseResourceManager, KVCacheManager, @@ -4694,41 +4696,63 @@ def _forward_step_mm_encoder_only( multimodal_params = inputs.get("multimodal_params", []) if not multimodal_params or len(multimodal_params) == 0: # Return empty embeddings if no multimodal data - return {'mm_embeddings': []} - # TODO(TRTLLM-12175): split encoder outputs by explicit per-request - # encoder-output embedding lengths. multimodal_lengths is a - # prompt-side MM-token count and may include non-embedding - # special/framing tokens. - if getattr(scheduled_requests.context_requests[0], 'multimodal_lengths', - None) is None: - multimodal_chunks = None - else: - multimodal_chunks = [ - sum(request.multimodal_lengths) - for request in scheduled_requests.context_requests - if request.multimodal_lengths is not None - ] + return { + 'mm_embeddings': [], + 'mm_embedding_request_indices': [], + 'mm_embedding_lengths': [], + } + # Some ctx requests carry only mrope metadata (no actual vision + # content). Skip them so the encoder only runs on real image payloads. + mm_context_requests = [(request_idx, request) for request_idx, request + in enumerate(scheduled_requests.context_requests) + if request.py_multimodal_data is not None] + if len(mm_context_requests) != len(multimodal_params): + raise ValueError( + "mm_encoder_only expects one multimodal payload per context " + "request carrying py_multimodal_data") + mm_request_indices_with_payload = [] + mm_params_with_payload = [] + mm_embedding_lengths = [] + for (request_idx, + request), multimodal_param in zip(mm_context_requests, + multimodal_params): + if not _has_mm_payload_keys(request.py_multimodal_data): + # mrope-only warmup request (no actual vision content) -> skip. + continue + multimodal_embedding_lengths = get_multimodal_embedding_lengths( + request) + if multimodal_embedding_lengths is None: + # Vision payload keys present but no pre-computed embedding + # lengths — skip to avoid a downstream sum(None) TypeError. + continue + mm_request_indices_with_payload.append(request_idx) + mm_params_with_payload.append(multimodal_param) + mm_embedding_lengths.append(multimodal_embedding_lengths) + if not mm_params_with_payload: + return { + 'mm_embeddings': [], + 'mm_embedding_request_indices': [], + 'mm_embedding_lengths': [], + } # For mm_encoder_only mode, we only run the vision encoder part # The model should be a vision encoder (e.g., Qwen2VisionModelBase) - mm_embeddings = self.model.forward(multimodal_params) + mm_embeddings = self.model.forward(mm_params_with_payload) assert len( mm_embeddings ) == 1, "mm_embeddings should be a 1-element list, mix modality (video+image) is not supported" - if multimodal_chunks is None or len(multimodal_chunks) != len( - multimodal_params): - mm_embeddings = list( - torch.chunk(mm_embeddings[0], - scheduled_requests.num_context_requests, - dim=0)) - else: - mm_embeddings = list( - torch.split(mm_embeddings[0], multimodal_chunks, dim=0)) + split_lengths = [sum(lengths) for lengths in mm_embedding_lengths] + mm_embeddings = list(torch.split(mm_embeddings[0], split_lengths, + dim=0)) + if len(mm_embeddings) != len(mm_embedding_lengths): + raise ValueError( + "mm_encoder_only produced an embedding batch that does not " + "match mm_embedding_lengths") # Extract mrope position data from multimodal_params if available mrope_position_ids_list = [] mrope_position_deltas_list = [] - for multimodal_param in multimodal_params: + for multimodal_param in mm_params_with_payload: mrope_config = multimodal_param.multimodal_data.get( 'mrope_config', {}) mrope_position_ids = mrope_config.get('mrope_position_ids') @@ -4738,7 +4762,21 @@ def _forward_step_mm_encoder_only( if mrope_position_deltas is not None: mrope_position_deltas_list.append(mrope_position_deltas) - result = {'mm_embeddings': mm_embeddings, 'logits': None} + # mrope lists must align 1:1 with multimodal_params (or be empty); + # the sampler indexes them by per-MM-result position into mm_embeddings. + assert (len(mrope_position_ids_list) == len(mrope_position_deltas_list) + and len(mrope_position_ids_list) + in (0, len(mm_params_with_payload))), ( + f"mrope alignment: got {len(mrope_position_ids_list)} ids, " + f"{len(mrope_position_deltas_list)} deltas, " + f"{len(mm_params_with_payload)} mm params") + + result = { + 'mm_embeddings': mm_embeddings, + 'logits': None, + 'mm_embedding_request_indices': mm_request_indices_with_payload, + 'mm_embedding_lengths': mm_embedding_lengths, + } if mrope_position_ids_list: result['mrope_position_ids'] = mrope_position_ids_list if mrope_position_deltas_list: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index f4bc390060e3..fbe5bfe0d56d 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -161,6 +161,8 @@ def _ensure_int64_cpu_tensor( def _resolve_multimodal_run_metadata( req: LlmRequest) -> Optional[_MmRunMetadata]: + # TODO(perf): cache per request; block-reuse invokes this once per block, + # repeatedly rebuilding identical tensors for the same request metadata. # Worked example for one logical multimodal item split by text: # # prompt index: 0 1 2 3 4 5 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index fc8e76b923e1..3c823888c363 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -291,9 +291,65 @@ def is_generation_model(self) -> bool: @dataclass(kw_only=True) class MultimodalResult: mm_embeddings: List[torch.Tensor] + # needed to torch.split the mm_embeddings into item-wise chunks + mm_embedding_lengths: List[List[int]] + # needed when requests mix text-only and multimodal ones + mm_embedding_request_indices: List[int] + # number of context requests in the batch + num_context_requests: int # Can be used to include e.g. `mrope_position_ids`, etc. extra_data: Optional[Dict[str, Any]] = None + def __post_init__(self) -> None: + num_embeddings = len(self.mm_embeddings) + num_lengths = len(self.mm_embedding_lengths) + if num_lengths != num_embeddings: + raise ValueError( + "mm_embedding_lengths batch size does not match mm_embeddings: " + f"{num_lengths} != {num_embeddings}" + ) + num_request_indices = len(self.mm_embedding_request_indices) + if num_request_indices != num_embeddings: + raise ValueError( + "mm_embedding_request_indices batch size does not match " + f"mm_embeddings: {num_request_indices} != {num_embeddings}" + ) + for result_index, (mm_embedding, mm_embedding_lengths) in enumerate( + zip(self.mm_embeddings, self.mm_embedding_lengths, strict=True) + ): + actual_rows = len(mm_embedding) + expected_rows = sum(mm_embedding_lengths) + if actual_rows != expected_rows: + raise ValueError( + f"mm_embedding shape mismatch for result {result_index}: " + f"{actual_rows} != {expected_rows}" + ) + for request_index in self.mm_embedding_request_indices: + if request_index < 0 or request_index >= self.num_context_requests: + raise ValueError( + "mm_embedding_request_indices contains an invalid request " + f"index: {request_index} not in [0, {self.num_context_requests})" + ) + + @classmethod + def from_model_outputs( + cls, model_outputs: Dict[str, Any], num_context_requests: int + ) -> "MultimodalResult": + result_keys = { + "mm_embeddings", + "mm_embedding_lengths", + "mm_embedding_request_indices", + } + return cls( + mm_embeddings=model_outputs["mm_embeddings"], + mm_embedding_lengths=model_outputs["mm_embedding_lengths"], + mm_embedding_request_indices=model_outputs["mm_embedding_request_indices"], + num_context_requests=num_context_requests, + extra_data={ + key: value for key, value in model_outputs.items() if key not in result_keys + }, + ) + @dataclass(kw_only=True) class SampleStateWithMMResult(SampleState[SampleStateTensors, SampleStateTensors]): @@ -336,11 +392,10 @@ def sample_async( resource_manager: Optional[ResourceManager] = None, ) -> SampleState: # from model_outputs to MultimodalResult - data = MultimodalResult( - mm_embeddings=model_outputs.pop("mm_embeddings"), - extra_data={**model_outputs}, - ) assert not scheduled_requests.generation_requests + data = MultimodalResult.from_model_outputs( + model_outputs, scheduled_requests.num_context_requests + ) return self.SampleState(requests=scheduled_requests.context_requests, data=data) @override @@ -356,26 +411,28 @@ def update_requests( extra_data = state.data.extra_data or {} mrope_position_ids = extra_data.get("mrope_position_ids", None) mrope_position_deltas = extra_data.get("mrope_position_deltas", None) - for i, (request, mm_embedding) in enumerate(zip(requests, mm_embeddings)): + for request in requests: request.state = LlmRequestState.GENERATION_COMPLETE # NOTE: This is a hack: set finish reason manually and set the beam 0 request.set_finished_reason(FinishReason.LENGTH, 0) - assert request.multimodal_lengths is not None - # TODO(TRTLLM-12175): request.multimodal_lengths is a - # prompt-side MM-token count and may include non-embedding - # special/framing tokens. This validation needs per-item - # encoder-output embedding lengths instead. - if len(mm_embedding) != sum(request.multimodal_lengths): - raise ValueError( - f"mm_embedding shape mismatch: {len(mm_embedding)} != {sum(request.multimodal_lengths)}" - ) - request.py_result.append_mm_embeddings(mm_embedding, request.multimodal_lengths) + request_indices = state.data.mm_embedding_request_indices + for result_index, (request_index, mm_embedding) in enumerate( + zip(request_indices, mm_embeddings, strict=True) + ): + request = requests[request_index] + mm_embedding_lengths = state.data.mm_embedding_lengths[result_index] + + request.py_result.append_mm_embeddings(mm_embedding, mm_embedding_lengths) # Store mrope data if available if mrope_position_ids is not None and mrope_position_deltas is not None: + mrope_index = ( + request_index if len(mrope_position_ids) == len(requests) else result_index + ) request.py_result.set_mrope_position( - mrope_position_ids[i], mrope_position_deltas[i] + mrope_position_ids[mrope_index], + mrope_position_deltas[mrope_index], ) @override diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 18e8a182f210..2c6b19cb247a 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -468,21 +468,8 @@ def _enqueue_request(self, if request.multimodal_params is not None and request.multimodal_params.has_content( ): if request.multimodal_params.multimodal_input is not None: - multimodal_input = tllm.MultimodalInput( - multimodal_hashes=request.multimodal_params. - multimodal_input.multimodal_hashes, - multimodal_positions=request.multimodal_params. - multimodal_input.multimodal_positions, - multimodal_lengths=request.multimodal_params. - multimodal_input.multimodal_lengths, - multimodal_uuids=request.multimodal_params.multimodal_input. - multimodal_uuids, - multimodal_item_run_cu_offsets=request.multimodal_params. - multimodal_input.multimodal_item_run_cu_offsets, - multimodal_run_positions=request.multimodal_params. - multimodal_input.multimodal_run_positions, - multimodal_run_lengths=request.multimodal_params. - multimodal_input.multimodal_run_lengths) + multimodal_input = request.multimodal_params.multimodal_input.to_binding( + tllm) # NOTE: Setting to None here to avoid sending multimodal_input again through the 'py_multimodal_data' field request.multimodal_params.multimodal_input = None @@ -622,7 +609,8 @@ def _deduce_max_tokens(request: GenerationRequest, executor_request.py_disaggregated_params = request.disaggregated_params if self._is_pytorch_backend and request.multimodal_params is not None: if request.multimodal_params.multimodal_data is not None: - # NOTE: Deserialize SharedTensor handle to actual tensor + # Resolve SharedTensorContainer dicts inside multimodal_data, including + # E/P handoff embedding handles parked under "multimodal_embedding". request.multimodal_params.to_tensor("multimodal_data") executor_request.py_multimodal_data = request.multimodal_params.multimodal_data diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 2649f7aa5516..8cae95503b4b 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -546,7 +546,6 @@ def _handle_response(self, if hasattr(response_result, "mm_embedding_handles" ) and response_result.mm_embedding_handles is not None: - # mm_embedding_handles is a list of handles (one per multimodal item). mm_embedding_handles = response_result.mm_embedding_handles if self._disaggregated_params is not None: self._disaggregated_params.multimodal_embedding_handles = mm_embedding_handles diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index f19c0af72007..35c7e7413341 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -23,6 +23,131 @@ _HASH_SCHEME_TAG = b"trtllm.mm.hash.v1" +def _validate_int_list(values: Any, field_name: str) -> None: + """Boundary metadata must be owned Python list[int]. No tensor/tuple.""" + if not isinstance(values, list): + raise TypeError(f"{field_name} must be a list") + if not all(isinstance(value, int) for value in values): + raise TypeError(f"{field_name} must contain only integers") + + +def _validate_multimodal_positions_and_lengths( + multimodal_positions: List[int], + multimodal_lengths: List[int], + expected_num_items: int, + expected_num_items_name: str, +) -> None: + """Validate one prompt span per MM item. + + expected_num_items is owner count: hashes for MultimodalInput, + embedding lengths for E/P handoff. Positions are prompt offsets. Lengths + are prompt token counts. + """ + _validate_int_list(multimodal_positions, "multimodal_positions") + _validate_int_list(multimodal_lengths, "multimodal_lengths") + + if len(multimodal_positions) != len(multimodal_lengths): + raise ValueError(f"Position and length arrays must match in size: " + f"positions={len(multimodal_positions)}, " + f"lengths={len(multimodal_lengths)}") + if len(multimodal_positions) != expected_num_items: + raise ValueError( + f"{expected_num_items_name}, multimodal_positions, and " + "multimodal_lengths must all have the same length") + + if any(position < 0 for position in multimodal_positions): + raise ValueError("multimodal_positions must be non-negative") + if any(length <= 0 for length in multimodal_lengths): + raise ValueError("multimodal_lengths must be positive") + + +def _validate_multimodal_runs( + num_items: int, + multimodal_lengths: List[int], + multimodal_item_run_cu_offsets: Optional[List[int]], + multimodal_run_positions: Optional[List[int]], + multimodal_run_lengths: Optional[List[int]], + item_count_name: str, +) -> None: + """Validate exact runs when they are present. + + Either no run fields, or all three. Offsets length is num_items + 1. + Runs for each item must sum to multimodal_lengths[i]. Values must fit + int32 for executor/KV-cache code. + """ + run_fields = ( + multimodal_item_run_cu_offsets, + multimodal_run_positions, + multimodal_run_lengths, + ) + if all(field is None for field in run_fields): + return + if any(field is None for field in run_fields): + raise ValueError( + "multimodal_item_run_cu_offsets, multimodal_run_positions, " + "and multimodal_run_lengths must be provided together") + + assert multimodal_item_run_cu_offsets is not None + assert multimodal_run_positions is not None + assert multimodal_run_lengths is not None + + for field_name, values in ( + ("multimodal_item_run_cu_offsets", multimodal_item_run_cu_offsets), + ("multimodal_run_positions", multimodal_run_positions), + ("multimodal_run_lengths", multimodal_run_lengths), + ): + _validate_int_list(values, field_name) + if any(value > _INT32_MAX for value in values): + raise ValueError(f"{field_name} values must fit in int32") + + if len(multimodal_item_run_cu_offsets) != num_items + 1: + raise ValueError("multimodal_item_run_cu_offsets length must be " + f"len({item_count_name}) + 1") + if multimodal_item_run_cu_offsets[0] != 0: + raise ValueError("multimodal_item_run_cu_offsets must start at 0") + if len(multimodal_run_positions) != len(multimodal_run_lengths): + raise ValueError( + "multimodal_run_positions and multimodal_run_lengths must " + "have the same length") + if multimodal_item_run_cu_offsets[-1] != len(multimodal_run_positions): + raise ValueError( + "multimodal_item_run_cu_offsets[-1] must equal the number of " + "flat multimodal runs") + if not all(multimodal_item_run_cu_offsets[i] <= + multimodal_item_run_cu_offsets[i + 1] + for i in range(len(multimodal_item_run_cu_offsets) - 1)): + raise ValueError( + "multimodal_item_run_cu_offsets must be non-decreasing") + if any(pos < 0 for pos in multimodal_run_positions): + raise ValueError("multimodal_run_positions must be non-negative") + if any(length <= 0 for length in multimodal_run_lengths): + raise ValueError("multimodal_run_lengths must be positive") + for run_idx, (position, length) in enumerate( + zip(multimodal_run_positions, multimodal_run_lengths)): + if position + length > _INT32_MAX: + raise ValueError( + f"multimodal run {run_idx} end position exceeds int32 " + f"range: position={position}, length={length}, " + f"max={_INT32_MAX}") + + for item_idx, expected_length in enumerate(multimodal_lengths): + run_begin = multimodal_item_run_cu_offsets[item_idx] + run_end = multimodal_item_run_cu_offsets[item_idx + 1] + actual_length = sum(multimodal_run_lengths[run_begin:run_end]) + if actual_length != expected_length: + raise ValueError( + f"multimodal run lengths for item {item_idx} sum to " + f"{actual_length}, expected {expected_length}") + item_positions = multimodal_run_positions[run_begin:run_end] + item_lengths = multimodal_run_lengths[run_begin:run_end] + for prev_pos, prev_len, pos in zip(item_positions, item_lengths, + item_positions[1:]): + if pos < prev_pos + prev_len: + raise ValueError( + "multimodal runs must be ordered and non-overlapping " + "within each item") + + def strip_mm_data_for_generation(mm_data: Dict[str, Any]) -> None: """Clear `mm_data` in place, retaining only `mrope_config.mrope_position_deltas`. @@ -123,23 +248,12 @@ def __post_init__(self): f"All hash arrays must have the same length, got lengths: {hash_lengths}" ) - # Check that positions and lengths are valid - if not all(isinstance(x, int) for x in self.multimodal_positions): - raise TypeError("multimodal_positions must contain only integers") - - if not all(isinstance(x, int) for x in self.multimodal_lengths): - raise TypeError("multimodal_lengths must contain only integers") - - # Check position and length arrays match in size - if len(self.multimodal_positions) != len(self.multimodal_lengths): - raise ValueError( - f"Position and length arrays must match in size: " - f"positions={len(self.multimodal_positions)}, lengths={len(self.multimodal_lengths)}" - ) - if len(self.multimodal_hashes) != len(self.multimodal_positions): - raise ValueError( - "multimodal_hashes, multimodal_positions, and multimodal_lengths " - "must all have the same length") + _validate_multimodal_positions_and_lengths( + self.multimodal_positions, + self.multimodal_lengths, + len(self.multimodal_hashes), + "multimodal_hashes", + ) # Validate multimodal_uuids if provided if self.multimodal_uuids is not None: @@ -155,90 +269,14 @@ def __post_init__(self): f"multimodal_uuids[{i}] must be a string or None, got {type(uuid)}" ) - self._validate_multimodal_runs() - - def _validate_multimodal_runs(self) -> None: - run_fields = ( + _validate_multimodal_runs( + len(self.multimodal_hashes), + self.multimodal_lengths, self.multimodal_item_run_cu_offsets, self.multimodal_run_positions, self.multimodal_run_lengths, + "multimodal_hashes", ) - if all(field is None for field in run_fields): - return - if any(field is None for field in run_fields): - raise ValueError( - "multimodal_item_run_cu_offsets, multimodal_run_positions, " - "and multimodal_run_lengths must be provided together") - - assert self.multimodal_item_run_cu_offsets is not None - assert self.multimodal_run_positions is not None - assert self.multimodal_run_lengths is not None - - if len(self.multimodal_item_run_cu_offsets) != len( - self.multimodal_hashes) + 1: - raise ValueError("multimodal_item_run_cu_offsets length must be " - "len(multimodal_hashes) + 1") - if self.multimodal_item_run_cu_offsets[0] != 0: - raise ValueError("multimodal_item_run_cu_offsets must start at 0") - if len(self.multimodal_run_positions) != len( - self.multimodal_run_lengths): - raise ValueError( - "multimodal_run_positions and multimodal_run_lengths must " - "have the same length") - if self.multimodal_item_run_cu_offsets[-1] != len( - self.multimodal_run_positions): - raise ValueError( - "multimodal_item_run_cu_offsets[-1] must equal the number of " - "flat multimodal runs") - - for field_name, values in ( - ("multimodal_item_run_cu_offsets", - self.multimodal_item_run_cu_offsets), - ("multimodal_run_positions", self.multimodal_run_positions), - ("multimodal_run_lengths", self.multimodal_run_lengths), - ): - if not isinstance(values, list): - raise TypeError(f"{field_name} must be a list") - if not all(isinstance(x, int) for x in values): - raise TypeError(f"{field_name} must contain only integers") - if any(value > _INT32_MAX for value in values): - raise ValueError(f"{field_name} values must fit in int32") - - if not all( - self.multimodal_item_run_cu_offsets[i] <= - self.multimodal_item_run_cu_offsets[i + 1] - for i in range(len(self.multimodal_item_run_cu_offsets) - 1)): - raise ValueError( - "multimodal_item_run_cu_offsets must be non-decreasing") - if any(pos < 0 for pos in self.multimodal_run_positions): - raise ValueError("multimodal_run_positions must be non-negative") - if any(length <= 0 for length in self.multimodal_run_lengths): - raise ValueError("multimodal_run_lengths must be positive") - for run_idx, (position, length) in enumerate( - zip(self.multimodal_run_positions, - self.multimodal_run_lengths)): - if position + length > _INT32_MAX: - raise ValueError( - f"multimodal run {run_idx} end position exceeds int32 " - f"range: position={position}, length={length}, " - f"max={_INT32_MAX}") - - for item_idx, expected_length in enumerate(self.multimodal_lengths): - run_begin = self.multimodal_item_run_cu_offsets[item_idx] - run_end = self.multimodal_item_run_cu_offsets[item_idx + 1] - actual_length = sum(self.multimodal_run_lengths[run_begin:run_end]) - if actual_length != expected_length: - raise ValueError( - f"multimodal run lengths for item {item_idx} sum to " - f"{actual_length}, expected {expected_length}") - item_positions = self.multimodal_run_positions[run_begin:run_end] - item_lengths = self.multimodal_run_lengths[run_begin:run_end] - for prev_pos, prev_len, pos in zip(item_positions, item_lengths, - item_positions[1:]): - if pos < prev_pos + prev_len: - raise ValueError( - "multimodal runs must be ordered and non-overlapping " - "within each item") @classmethod def from_components( @@ -267,6 +305,87 @@ def to_tensor(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: torch.tensor(self.multimodal_positions, dtype=torch.int32), torch.tensor(self.multimodal_lengths, dtype=torch.int32)) + def run_metadata(self) -> Dict[str, List[int]]: + metadata = {} + if self.multimodal_item_run_cu_offsets is not None: + metadata[ + "multimodal_item_run_cu_offsets"] = self.multimodal_item_run_cu_offsets + if self.multimodal_run_positions is not None: + metadata["multimodal_run_positions"] = self.multimodal_run_positions + if self.multimodal_run_lengths is not None: + metadata["multimodal_run_lengths"] = self.multimodal_run_lengths + return metadata + + def to_binding(self, executor_module: Any) -> Any: + kwargs = dict(multimodal_hashes=self.multimodal_hashes, + multimodal_positions=self.multimodal_positions, + multimodal_lengths=self.multimodal_lengths, + multimodal_uuids=self.multimodal_uuids) + kwargs.update(self.run_metadata()) + return executor_module.MultimodalInput(**kwargs) + + +@dataclass +class DisaggPrefillMultimodalInputs: + """Typed multimodal metadata returned by E/P disagg prefill processors.""" + + prompt_token_ids: List[int] + multimodal_lengths: List[int] + multimodal_positions: List[int] + multimodal_embedding_lengths: List[int] + multimodal_item_run_cu_offsets: Optional[List[int]] = None + multimodal_run_positions: Optional[List[int]] = None + multimodal_run_lengths: Optional[List[int]] = None + special_token_offsets: Optional[List[int]] = None + item_types: Optional[List[int]] = None + + def __post_init__(self) -> None: + _validate_int_list(self.prompt_token_ids, "prompt_token_ids") + _validate_int_list(self.multimodal_embedding_lengths, + "multimodal_embedding_lengths") + _validate_multimodal_positions_and_lengths( + self.multimodal_positions, + self.multimodal_lengths, + len(self.multimodal_embedding_lengths), + "multimodal_embedding_lengths", + ) + + if any(length <= 0 for length in self.multimodal_embedding_lengths): + raise ValueError("multimodal_embedding_lengths must be positive") + + _validate_multimodal_runs( + len(self.multimodal_lengths), + self.multimodal_lengths, + self.multimodal_item_run_cu_offsets, + self.multimodal_run_positions, + self.multimodal_run_lengths, + "multimodal_lengths", + ) + self._validate_optional_metadata() + + def _validate_optional_metadata(self) -> None: + if self.special_token_offsets is not None: + _validate_int_list(self.special_token_offsets, + "special_token_offsets") + if any(offset < 0 for offset in self.special_token_offsets): + raise ValueError("special_token_offsets must be non-negative") + if self.item_types is not None: + _validate_int_list(self.item_types, "item_types") + if len(self.item_types) != len(self.multimodal_lengths): + raise ValueError("item_types length must match " + "multimodal_lengths") + + def to_multimodal_input(self, + mm_hashes: List[List[int]]) -> MultimodalInput: + return MultimodalInput.from_components( + mm_hashes, + self.multimodal_positions, + self.multimodal_lengths, + mm_item_run_cu_offsets=self.multimodal_item_run_cu_offsets, + mm_run_positions=self.multimodal_run_positions, + mm_run_lengths=self.multimodal_run_lengths, + ) + @dataclass class MultimodalRuntimeData: @@ -329,6 +448,7 @@ def __post_init__(self): # Extend only after auditing each key's consumers. _CPU_ONLY_MULTIMODAL_DATA_KEYS = frozenset({ "multimodal_embed_mask_cumsum", + "multimodal_embedding_lengths", }) @@ -354,7 +474,10 @@ class MultimodalParams: "mrope_rotary_cos_sin": torch.Tensor, # Rotary embeddings (Qwen2/2.5-VL) "mrope_position_deltas": torch.Tensor, # Position deltas (Qwen2/2.5-VL) }, - "multimodal_embedding": torch.Tensor, # Pre-computed vision embeddings + "multimodal_embedding": torch.Tensor | List[SharedTensor handle dict], + # Pre-computed embeddings. In E/P handoff this may temporarily hold + # SharedTensorContainer dicts; BaseWorker restores them to tensors with + # to_tensor("multimodal_data") before PyTorch forward. "image": { "pixel_values": torch.Tensor, "image_height": torch.Tensor | List[int], @@ -827,6 +950,13 @@ def find_mm_token_lengths( mm_video_dict = (multimodal_data or {}).get("video") or {} video_grid_thw = mm_video_dict.get("video_grid_thw") + if video_grid_thw is not None: + video_grid_thw = torch.as_tensor(video_grid_thw) + assert video_grid_thw.device.type == "cpu", ( + "video_grid_thw must be CPU-resident when computing " + f"multimodal metadata, got {video_grid_thw.device}.") + if video_grid_thw.ndim != 2 or video_grid_thw.shape[-1] != 3: + raise ValueError("video_grid_thw must have shape [num_segments, 3]") for modality, items in mm_items.items(): if not hasattr(input_processor, f"get_num_tokens_per_{modality}"): @@ -836,12 +966,14 @@ def find_mm_token_lengths( video_grid_thw_for_items = None if modality == "video" and video_grid_thw is not None: - if len(video_grid_thw) == len(items): + if len(items) == 1: + video_grid_thw_for_items = video_grid_thw + elif video_grid_thw.shape[0] == len(items): video_grid_thw_for_items = video_grid_thw else: logger.warning( "find_mm_token_lengths: video_grid_thw row count " - f"({len(video_grid_thw)}) does not match number of " + f"({video_grid_thw.shape[0]}) does not match number of " f"videos in mm_data ({len(items)}); falling back to " "per-item recompute without video_grid_thw.") @@ -872,8 +1004,9 @@ def find_mm_token_lengths( # metadata route. Keep this for now: Qwen3-VL needs the # processor-produced video_grid_thw for correct video token # counts. - call_kwargs["video_grid_thw"] = video_grid_thw_for_items[ - idx] + call_kwargs["video_grid_thw"] = ( + video_grid_thw_for_items if len(items) == 1 else + video_grid_thw_for_items[idx:idx + 1]) num_tokens = input_processor.get_num_tokens_per_video( **call_kwargs) modality_token_lengths.append(num_tokens) @@ -896,6 +1029,7 @@ def find_mm_token_lengths( _MM_METADATA_ONLY_KEYS = frozenset({ "mrope_config", "multimodal_embed_mask_cumsum", + "multimodal_embedding_lengths", "special_token_offsets", "layout_metadata", }) @@ -1132,6 +1266,32 @@ def _find_mm_token_runs_from_mask( return item_run_cu_offsets, run_positions, run_lengths +def _find_mm_embedding_lengths_from_masks( + mm_mask: torch.Tensor, + embed_mask: torch.Tensor, + num_mm_tokens: List[int], +) -> List[int]: + """Compute embedding-slot counts per logical multimodal item.""" + if not torch.any(mm_mask): + return [] + + mm_positions = torch.where(mm_mask)[0] + lengths_t = torch.tensor(num_mm_tokens) + assert mm_positions.numel() == lengths_t.sum().item(), ( + f"Number of multimodal tokens ({mm_positions.numel()}) does not match " + f"sum of per-unit lengths ({lengths_t.sum().item()}): " + f"num_mm_tokens={num_mm_tokens}") + + embedding_lengths: List[int] = [] + offset = 0 + for item_length in num_mm_tokens: + item_positions = mm_positions[offset:offset + item_length] + offset += item_length + embedding_lengths.append(int(embed_mask[item_positions].sum().item())) + + return embedding_lengths + + def validate_mm_inputs(prompt_token_ids: Union[torch.Tensor, List[int], np.ndarray], mm_hashes: List[List[int]], start_positions: List[int], diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 801e27b977eb..e9ed709010ba 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -20,6 +20,7 @@ from .content_format import ContentFormat from .data import TextPrompt from .multimodal import (MultimodalInput, _as_cpu_tensor, _compute_mm_masks, + _find_mm_embedding_lengths_from_masks, _find_mm_token_runs_from_mask, _find_mm_token_start_pos_from_masks, apply_mm_hashes, default_hasher, find_mm_token_lengths, @@ -818,10 +819,11 @@ def support_multimodal_disaggregated(model_cls: Type[nn.Module]): raise TypeError( f"{processor_cls.__name__} must inherit from BaseMultimodalInputProcessor to support multimodal disagg" ) - method = getattr(processor_cls, "get_prompt_token_ids", None) + method = getattr(processor_cls, "build_disagg_prefill_multimodal_inputs", + None) if method is None or not callable(method): raise TypeError( - f"{processor_cls.__name__} must implement a callable method `get_prompt_token_ids` to support multimodal disagg" + f"{processor_cls.__name__} must implement a callable method `build_disagg_prefill_multimodal_inputs` to support multimodal disagg" ) setattr(processor_cls, "support_mm_disagg", True) @@ -1116,6 +1118,7 @@ def multimodal_hashing_process( if input_ids_tensor.numel() == 0: start_positions, start_special_token_positions = [], [] item_run_cu_offsets, run_positions, run_lengths = [0], [], [] + multimodal_embedding_lengths = [] else: mm_mask, embed_mask, special_mask = _compute_mm_masks( input_ids_tensor, @@ -1131,6 +1134,11 @@ def multimodal_hashing_process( num_mm_tokens)) item_run_cu_offsets, run_positions, run_lengths = ( _find_mm_token_runs_from_mask(mm_mask, num_mm_tokens)) + multimodal_embedding_lengths = ( + _find_mm_embedding_lengths_from_masks(mm_mask, embed_mask, + num_mm_tokens)) + extra_processed_inputs["multimodal_data"][ + "multimodal_embedding_lengths"] = multimodal_embedding_lengths # Store special token offsets if available if len(start_special_token_positions ) > 0 and mm_special_token_ids is not None: diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 3f387d97863e..d8c471d8624e 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -17,7 +17,8 @@ from transformers import PreTrainedTokenizerBase from tensorrt_llm._utils import mpi_disabled -from tensorrt_llm.inputs.multimodal import MultimodalInput, MultimodalParams +from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, + MultimodalParams) from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.llmapi import tracing from tensorrt_llm.metrics.enums import MetricNames @@ -566,13 +567,29 @@ def _preprocess( # This branch is applicable for Encode --> Prefill handoff scenario, # in E/P/D/ and E/PD settings. Prefill worker executes this code path. if is_mm_disagg: + if self.args.backend == "_autodeploy": + raise ValueError( + "Multimodal disaggregated inference (encode -> prefill " + "embedding handoff) is not supported with the AutoDeploy " + "backend. AutoDeploy runs the multimodal encoder in-prefill " + "on raw inputs and does not consume precomputed multimodal " + "embeddings.") if not getattr(self.input_processor, "support_mm_disagg", False): raise ValueError( "Multimodal disaggregated inference is not supported for this model" ) mm_handles = disaggregated_params.multimodal_embedding_handles - prompt_token_ids, mm_token_length, mm_token_positions = self.input_processor.get_prompt_token_ids( - inputs, mm_handles) + # TODO(TRTLLM-12869): Pass encoder-side MM layout through + # DisaggregatedParams so prefill does not rebuild prompt tokens, + # positions, lengths, runs, special offsets, and cumsum here. + disagg_mm_inputs = ( + self.input_processor.build_disagg_prefill_multimodal_inputs( + inputs, mm_handles)) + if not isinstance(disagg_mm_inputs, DisaggPrefillMultimodalInputs): + raise TypeError( + "build_disagg_prefill_multimodal_inputs must return " + "DisaggPrefillMultimodalInputs") + prompt_token_ids = disagg_mm_inputs.prompt_token_ids prompt = inputs.get("prompt", None) query_token_ids = inputs.get("query_token_ids", None) if is_gen_only: @@ -581,9 +598,25 @@ def _preprocess( ) else: mm_hashes = disaggregated_params.multimodal_hashes - multimodal_input = MultimodalInput.from_components( - mm_hashes, mm_token_positions, mm_token_length) - multimodal_data = {"multimodal_embedding": mm_handles} + multimodal_input = disagg_mm_inputs.to_multimodal_input( + mm_hashes) + # E/P handoff carries SharedTensorContainer dicts. Park them under the + # embedding key so BaseWorker's recursive to_tensor("multimodal_data") + # restores local tensor views before PyTorch forward. Until then this + # key holds handles, not tensors. + multimodal_data = { + "multimodal_embedding": + mm_handles, + "multimodal_embedding_lengths": + (disagg_mm_inputs.multimodal_embedding_lengths), + } + if disagg_mm_inputs.special_token_offsets is not None: + multimodal_data["special_token_offsets"] = ( + disagg_mm_inputs.special_token_offsets) + if disagg_mm_inputs.item_types is not None: + multimodal_data["layout_metadata"] = { + "item_types": disagg_mm_inputs.item_types + } if disaggregated_params.mrope_position_ids_handle is not None: # NOTE: `PyTorchModelEngine` assumes both are present when using mrope. assert disaggregated_params.mrope_position_deltas_handle is not None diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 96316d2eb27d..59cc2f2d7295 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -158,6 +158,10 @@ class DisaggregatedParams(OpenAIBaseModel): schedule_style: Optional[DisaggScheduleStyle] = None conversation_id: Optional[str] = None ctx_usage: Optional[UsageInfo] = None + # TODO(TRTLLM-12407): Multimodal E/PD over trtllm-serve needs these protocol fields too: + # encoder embedding handles, multimodal hashes, and optional mRoPE handles. + # Add them here and in to_disaggregated_params()/to_llm_disaggregated_params() + # before routing MM encoder -> context -> generation through OpenAI protocol. class ErrorResponse(OpenAIBaseModel): diff --git a/tests/integration/defs/accuracy/references/videomme.yaml b/tests/integration/defs/accuracy/references/videomme.yaml index 2bee39c8c2b6..186971bc05b7 100644 --- a/tests/integration/defs/accuracy/references/videomme.yaml +++ b/tests/integration/defs/accuracy/references/videomme.yaml @@ -3,6 +3,9 @@ # Initial Video-MME short-shard guardrail for E2E video QA. Update these values # after collecting stable model baselines on the generated 300-question shard. +Qwen/Qwen3-VL-2B-Instruct: + - accuracy: 54.5 + num_samples: 300 nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-FP8: - quant_algo: FP8 kv_cache_quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/test_epd_disagg_multimodal.py b/tests/integration/defs/accuracy/test_epd_disagg_multimodal.py new file mode 100644 index 000000000000..20659f9b7ca9 --- /dev/null +++ b/tests/integration/defs/accuracy/test_epd_disagg_multimodal.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""VideoMME accuracy over llmapi encode / prefill-decode (E/PD) disaggregation. + +Separated from test_disaggregated_serving.py: the EPD-multimodal path uses an +in-process MultimodalEncoder plus a combined prefill/decode LLM, which is a +different mechanism from the trtllm-serve subprocess disaggregation exercised by +the other tests in that file. +""" + +# NOTE: +# The encoder and PD are resident on the same physical GPU in the current test +# harness. Placing them on different physical GPUs silently corrupts the +# embeddings (garbage output, no error raised) in TRT-LLM's current state because +# the consumer (PD worker) rebuilds the encoder's embedding from a CUDA-IPC handle +# that currently never copies the tensor onto the PD's own compute device. +# Real cross-GPU E/PD therefore requires a real cross-device transfer +# (CPU staging or NIXL/RDMA) that is currently not natively supported in TRT-LLM. + +import contextlib +import os +from dataclasses import dataclass +from typing import Any, Dict, Iterator, Mapping, Optional, Protocol +from unittest import mock + +import pytest + +from tensorrt_llm import LLM, MultimodalEncoder +from tensorrt_llm.llmapi import KvCacheConfig, RequestOutput, SamplingParams +from tensorrt_llm.quantization import QuantAlgo + +from ..conftest import llm_models_root, skip_pre_blackwell, skip_pre_hopper +from .accuracy_core import LlmapiAccuracyTestHarness, VideoMME +from .test_disaggregated_serving import DEFAULT_TEST_TIMEOUT, MyThreadPoolExecutor + + +class VideoMMECompatibleLLM(Protocol): + """LLM surface consumed by the VideoMME evaluator.""" + + args: Any + model: str + _hf_model_dir: str + tokenizer: Any + input_processor: Any + + def generate_async( + self, + inputs: Dict[str, Any], + sampling_params: Optional[SamplingParams] = None, + streaming: bool = False, + ) -> Any: ... + + +class _MultimodalEncoderPDAdapter: + """Adapter that runs VideoMME dict inputs through llmapi E/PD.""" + + def __init__( + self, encoder: MultimodalEncoder, pd_llm: LLM, thread_pool: MyThreadPoolExecutor + ) -> None: + self._encoder = encoder + self._pd_llm = pd_llm + self._thread_pool = thread_pool + self.args = pd_llm.args + self.model = pd_llm._hf_model_dir + self._hf_model_dir = pd_llm._hf_model_dir + self.tokenizer = pd_llm.tokenizer + self.input_processor = pd_llm.input_processor + + def _generate( + self, inputs: Dict[str, Any], sampling_params: Optional[SamplingParams], streaming: bool + ) -> RequestOutput: + if not isinstance(inputs, dict): + raise TypeError(f"Unsupported E/PD request input type: {type(inputs)}") + + encoder_output = self._encoder.generate_async(inputs).result() + disaggregated_params = encoder_output.disaggregated_params + if disaggregated_params is None: + raise RuntimeError("Multimodal encoder did not return disaggregated params.") + if disaggregated_params.multimodal_embedding_handles is None: + raise RuntimeError("Multimodal encoder did not return embedding handles.") + + disaggregated_params.request_type = "context_and_generation" + return self._pd_llm.generate_async( + inputs, + sampling_params=sampling_params, + streaming=streaming, + disaggregated_params=disaggregated_params, + ).result() + + def generate_async( + self, + inputs: Dict[str, Any], + sampling_params: Optional[SamplingParams] = None, + streaming: bool = False, + ): + future = self._thread_pool.submit(self._generate, inputs, sampling_params, streaming) + self._thread_pool.futures.append(future) + return future + + +@contextlib.contextmanager +def launch_multimodal_encoder_pd_llm( + encoder_llm_config: Dict[str, Any], + pd_llm_config: Dict[str, Any], + model_name: str, + max_workers: int = 16, +) -> Iterator[VideoMMECompatibleLLM]: + """Launch separate encoder and combined prefill/decode llmapi instances.""" + with contextlib.ExitStack() as stack: + stack.enter_context(mock.patch.dict(os.environ, {"TLLM_MULTIMODAL_DISAGGREGATED": "1"})) + thread_pool = stack.enter_context(MyThreadPoolExecutor(max_workers=max_workers)) + encoder = MultimodalEncoder(model=model_name, **encoder_llm_config) + pd_llm = LLM(model=model_name, **pd_llm_config) + with encoder, pd_llm: + yield _MultimodalEncoderPDAdapter(encoder, pd_llm, thread_pool) + + +@dataclass(frozen=True) +class EPDVariant: + """Immutable per-variant config for a VideoMME E/PD run.""" + + model_name: str + model_path: str + encoder_config: Mapping[str, Any] + pd_config: Mapping[str, Any] + expected_quant_algo: Optional[QuantAlgo] + max_workers: int + + @classmethod + def _build( + cls, + *, + model_name: str, + model_path: str, + kv_cache_config: KvCacheConfig, + max_batch_size: int, + expected_quant_algo: Optional[QuantAlgo], + max_num_tokens: int = 512, + attn_backend: Optional[str] = None, + max_workers: Optional[int] = None, + ) -> "EPDVariant": + """Fill shared encoder/PD defaults for one variant. + + Optional overrides are applied before construction so the frozen + instance never needs post-hoc mutation. + """ + # Optional attn_backend override, applied to both configs via a spread + # so the frozen instance never needs post-hoc mutation. + attn_override = {"attn_backend": attn_backend} if attn_backend is not None else {} + encoder_config = { + "trust_remote_code": True, + "max_batch_size": max_batch_size, + "cuda_graph_config": None, + **attn_override, + } + pd_config = { + "backend": "pytorch", + "disable_overlap_scheduler": True, + "trust_remote_code": True, + "kv_cache_config": kv_cache_config, + "enable_chunked_prefill": True, + "max_num_tokens": max_num_tokens, + "max_batch_size": max_batch_size, + "cuda_graph_config": None, + **attn_override, + } + + return cls( + model_name=model_name, + model_path=model_path, + encoder_config=encoder_config, + pd_config=pd_config, + expected_quant_algo=expected_quant_algo, + max_workers=max_workers if max_workers is not None else VideoMME.MAX_BATCH_SIZE, + ) + + @classmethod + def qwen3vl_2b(cls) -> "EPDVariant": + return cls._build( + model_name="Qwen/Qwen3-VL-2B-Instruct", + model_path=f"{llm_models_root()}/Qwen3/Qwen3-VL-2B-Instruct", + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.8, + enable_block_reuse=False, + dtype="auto", + ), + max_batch_size=16, + expected_quant_algo=None, + max_workers=16, + attn_backend="VANILLA", + # Qwen3-VL VideoMME prompts can exceed 1024 tokens after visual + # expansion; avoid splitting a single context across vanilla + # SDPA chunks in the E/P handoff path. + max_num_tokens=2048, + ) + + @classmethod + def nano_omni_fp8(cls) -> "EPDVariant": + return cls._build( + model_name="nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-FP8", + model_path=f"{llm_models_root()}/NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-FP8", + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.8, + mamba_ssm_cache_dtype="float32", + enable_block_reuse=False, + dtype="fp8", + ), + max_batch_size=64, + expected_quant_algo=QuantAlgo.FP8, + ) + + @classmethod + def nano_omni_nvfp4(cls) -> "EPDVariant": + return cls._build( + model_name="nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4", + model_path=f"{llm_models_root()}/NVIDIA-Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4", + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.8, + mamba_ssm_cache_dtype="float32", + enable_block_reuse=False, + dtype="fp8", + ), + max_batch_size=128, + expected_quant_algo=QuantAlgo.MIXED_PRECISION, + ) + + +class TestVideoMMEEPD(LlmapiAccuracyTestHarness): + """VideoMME accuracy over llmapi encode / prefill-decode (E/PD) disaggregation.""" + + SAMPLING_PARAMS = SamplingParams( + max_tokens=VideoMME.MAX_OUTPUT_LEN, + truncate_prompt_tokens=VideoMME.MAX_INPUT_LEN, + temperature=0.0, + top_k=1, + ) + + # Identical across all variants today; lifted to a class constant to mirror + # agg no_thinking_evaluator_kwargs. + NO_THINKING_EVALUATOR_KWARGS = { + "chat_template_kwargs": { + "enable_thinking": False, + }, + } + + def _launch_epd(self, variant: EPDVariant): + """Context manager: encoder + combined PD llmapi.""" + return launch_multimodal_encoder_pd_llm( + variant.encoder_config, + variant.pd_config, + variant.model_path, + max_workers=variant.max_workers, + ) + + def _run_videomme(self, llm, variant: EPDVariant) -> None: + actual_quant_algo = ( + llm.args.quant_config.quant_algo if llm.args.quant_config is not None else None + ) + assert actual_quant_algo == variant.expected_quant_algo + VideoMME(variant.model_name).evaluate( + llm, + sampling_params=self.SAMPLING_PARAMS, + extra_evaluator_kwargs=self.NO_THINKING_EVALUATOR_KWARGS, + ) + + @pytest.mark.timeout(DEFAULT_TEST_TIMEOUT) + @skip_pre_hopper + @pytest.mark.skip_less_device_memory(80000) + @pytest.mark.parametrize( + "variant", + [ + pytest.param( + EPDVariant.qwen3vl_2b(), marks=skip_pre_blackwell, id="qwen3vl_2b_instruct" + ), + pytest.param( + EPDVariant.nano_omni_fp8(), marks=skip_pre_hopper, id="nemotron_nano_v3_omni_fp8" + ), + pytest.param( + EPDVariant.nano_omni_nvfp4(), + marks=skip_pre_blackwell, + id="nemotron_nano_v3_omni_nvfp4", + ), + ], + ) + def test_disaggregated_videomme(self, variant: EPDVariant) -> None: + """Run VideoMME shard through a model-specific llmapi E/PD config.""" + with self._launch_epd(variant) as llm: + self._run_videomme(llm, variant) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index be5a78639f9f..5bed3dfa9db5 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -98,6 +98,9 @@ accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_chunked_prefill accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[noadp-ctx_tp2pp1-gen_tp1pp1] accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[adp-ctx_tp2pp1-gen_tp2pp1] accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_nixl_backend +accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] +accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] +accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_nvfp4] accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=False] accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 5ac2ad5f115f..a636a2715c90 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -74,6 +74,8 @@ l0_b200: - accuracy/test_llm_api_pytorch.py::TestQwen3_6_27B::test_fp8 - accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_nvfp4[tp1_block_reuse-cutlass] - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[nvfp4] + - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] + - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_nvfp4] - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[mtp_on] - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp[mtp_off] - accuracy/test_llm_api_pytorch.py::TestQwen3_5_9B::test_bf16[mtp_on] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 4f568c96fbb5..4dbb7eefc8ce 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -42,7 +42,7 @@ l0_h100: - unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" # ------------- MoE: test_single_gpu (by backend) --------------- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS" - - unittest/_torch/multimodal + - unittest/_torch/multimodal -k "not nemotron_nano_v2_vl_fp8" - unittest/_torch/sampler - unittest/_torch/speculative/test_eagle3.py - unittest/_torch/speculative/hw_agnostic @@ -122,6 +122,7 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_dummy_load_format - accuracy/test_llm_api_pytorch_multimodal.py::TestGemma3_27BInstruct::test_fp8_prequantized - accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8] + - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_fp8] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales_early_first_token_response - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_dummy_load_format diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_nano_v2_vl.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_nano_v2_vl.py index 3b89fe8ca204..f215fc208ebe 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_nano_v2_vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_nano_v2_vl.py @@ -3,6 +3,7 @@ import os from pathlib import Path +from types import SimpleNamespace from unittest import mock from unittest.mock import MagicMock @@ -14,15 +15,19 @@ from test_modeling_nemotron_h import extract_decode_logprobs from tensorrt_llm import LLM +from tensorrt_llm._torch.models import modeling_nemotron_nano as nemotron_nano from tensorrt_llm._torch.models.modeling_multimodal_utils import get_multimodal_embeddings from tensorrt_llm._torch.models.modeling_nemotron_nano import ( NanoV2VLInputProcessor, + NanoV2VLMultimodalEncoder, NanoV2VLVisionEncoder, NemotronH_Nano_VL_V2, ) from tensorrt_llm._torch.models.modeling_parakeet import ProjectedParakeet +from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_VISION_ENCODER_MAPPING from tensorrt_llm.inputs import ( AudioData, + VideoData, create_input_processor, create_input_processor_with_hash, default_multimodal_input_loader, @@ -36,6 +41,169 @@ MODEL_PATH = str(os.path.join(llm_models_root(), "NVIDIA-Nemotron-Nano-12B-v2-VL-BF16")) +def _make_minimal_nano_model_config(): + llm_config = SimpleNamespace(vocab_size=128) + pretrained_config = SimpleNamespace( + llm_config=llm_config, + torch_dtype=torch.bfloat16, + img_context_token_id=20, + video_context_token_id=21, + sound_context_token_id=None, + sound_config=None, + ) + return SimpleNamespace( + pretrained_config=pretrained_config, + quant_config=SimpleNamespace(exclude_modules=None), + quant_config_dict=None, + video_pruning_rate=None, + ) + + +def test_nemotron_nano_registers_native_multimodal_epd_components(): + """Native Nano VL/Omni classes advertise MM EPD support.""" + for arch in ("NemotronH_Nano_VL_V2", "NemotronH_Nano_Omni_Reasoning_V3"): + vision_encoder_cls, vlm_base_model = MODEL_CLASS_VISION_ENCODER_MAPPING[arch] + assert vision_encoder_cls is NanoV2VLMultimodalEncoder + assert vlm_base_model is None + assert NanoV2VLInputProcessor.support_mm_disagg is True + assert NemotronH_Nano_VL_V2.support_mm_disagg is True + + +def _assert_nano_video_handoff(handoff): + """Shared assertions for the EPD video handoff: split runs stay grouped under one MM item.""" + assert handoff.prompt_token_ids == [101, 30, 20, 20, 31, 55, 30, 20, 20, 31, 102] + assert handoff.multimodal_lengths == [8] + assert handoff.multimodal_positions == [1] + assert handoff.multimodal_embedding_lengths == [4] + assert handoff.multimodal_item_run_cu_offsets == [0, 2] + assert handoff.multimodal_run_positions == [1, 6] + assert handoff.multimodal_run_lengths == [4, 4] + assert handoff.special_token_offsets == [0, 3, 4, 7] + + +@pytest.mark.parametrize( + "input_field, input_value, asserts_encode_not_called", + [ + # Detokenized prompt text path: the tokenizer may encode the prompt. + ("prompt", "Question