diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 71132d01513f..05518df6bf77 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -58,8 +58,10 @@ from .modeling_gemma4_vision import Gemma4VisionModel from .modeling_multimodal_mixin import ( EncoderCachePartition, + EncoderInputField, MultimodalModelMixin, PreparedLlmInputs, + batch_encoder_inputs, ) from .modeling_multimodal_utils import _MULTIMODAL_ENV_NAME, _is_mm_disagg from .modeling_utils import ModelConfig, filter_weights, register_auto_model @@ -85,6 +87,12 @@ get_aspect_ratio_preserving_size, ) +_GEMMA4_PIXEL_FIELD = {"pixel_values": EncoderInputField(("pixel_values",))} +_GEMMA4_VISION_FIELDS = { + **_GEMMA4_PIXEL_FIELD, + "image_position_ids": EncoderInputField(("image_position_ids",)), +} + class RMSNormNoScale(nn.Module): """RMSNorm without learnable scale (for multimodal embedder pre-projection).""" @@ -1043,107 +1051,91 @@ def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) -> "Gemma4 requests containing multiple multimodal input types are not supported." ) - # The cache helper splits this tensor by request. Preserve that order while still batching - # consecutive requests with the same modality. - if ( - all(len(modalities) == 1 for modalities in param_modalities) - and len({modalities[0] for modalities in param_modalities}) > 1 - ): - embeddings = [] - for _, param_group in groupby( - zip(param_modalities, multimodal_params), key=lambda item: item[0][0] - ): - embeddings.append( - self.encode_multimodal_inputs( - [multimodal_param for _, multimodal_param in param_group] - ) + # The cache helper splits this tensor by request. Preserve consecutive + # modality runs and their encoder call composition without rechecking + # each run through a recursive encoder entry. + single_modality_inputs = all(len(modalities) == 1 for modalities in param_modalities) + mixed_modality_runs = ( + single_modality_inputs and len({modalities[0] for modalities in param_modalities}) > 1 + ) + if mixed_modality_runs: + input_groups = [ + (modality, [param.multimodal_data[modality] for _, param in param_group]) + for modality, param_group in groupby( + zip(param_modalities, multimodal_params), key=lambda item: item[0][0] + ) + ] + elif single_modality_inputs and multimodal_params: + modality = param_modalities[0][0] + input_groups = [ + (modality, [param.multimodal_data[modality] for param in multimodal_params]) + ] + else: + input_groups = [ + ( + modality, + [ + data + for param in multimodal_params + if (data := param.multimodal_data.get(modality, {})).get(input_field) + is not None + ], ) - return torch.cat(embeddings, dim=0) - - pixel_values_list = [] - image_position_ids_list = [] - image_seq_lens_extended: List[int] = [] - audio_features_list = [] - audio_features_mask_list: List[Optional[torch.Tensor]] = [] - video_pixel_values_list = [] - video_position_ids_list = [] - video_seq_lens_extended: List[int] = [] - - for multimodal_param in multimodal_params: - multimodal_data = multimodal_param.multimodal_data - image_data = multimodal_data.get("image", {}) - pixel_values = image_data.get("pixel_values") - if pixel_values is not None: - pixel_values_list.append(pixel_values) - image_position_ids = image_data.get("image_position_ids") - if image_position_ids is not None: - image_position_ids_list.append(image_position_ids) - image_seq_lens = image_data.get("image_seq_lens") - if image_seq_lens is not None: - image_seq_lens_extended.extend(image_seq_lens) - - audio_data = multimodal_data.get("audio", {}) - audio_features = audio_data.get("audio_features") - if audio_features is not None: - audio_features_list.append(audio_features) - audio_features_mask_list.append(audio_data.get("audio_features_mask")) - - video_data = multimodal_data.get("video", {}) - video_pixel_values = video_data.get("pixel_values") - if video_pixel_values is not None: - video_pixel_values_list.append(video_pixel_values) - video_position_ids = video_data.get("image_position_ids") - if video_position_ids is not None: - video_position_ids_list.append(video_position_ids) - video_seq_lens = video_data.get("image_seq_lens") - if video_seq_lens is not None: - video_seq_lens_extended.extend(video_seq_lens) + for modality, input_field in modality_inputs + ] multimodal_embeddings = [] - if pixel_values_list: - pixel_values = torch.cat(pixel_values_list) - image_position_ids = ( - torch.cat(image_position_ids_list) - if len(image_position_ids_list) == len(pixel_values_list) - else None + for modality, inputs in input_groups: + if not inputs: + continue + if modality == "audio": + if self.embed_audio is not None: + multimodal_embeddings.append( + torch.cat( + [ + self._get_audio_features( + data["audio_features"], data.get("audio_features_mask") + ) + for data in inputs + ], + dim=0, + ) + ) + if mixed_modality_runs: + multimodal_embeddings[-1] = multimodal_embeddings[-1].clone( + memory_format=torch.contiguous_format + ) + elif mixed_modality_runs: + raise ValueError( + "Gemma4 received active multimodal parameters without encoder inputs" + ) + continue + fields = ( + _GEMMA4_VISION_FIELDS + if all(data.get("image_position_ids") is not None for data in inputs) + else _GEMMA4_PIXEL_FIELD ) + batched = batch_encoder_inputs(inputs, fields) + image_seq_lens = [] + for data in inputs: + if data.get("image_seq_lens") is not None: + image_seq_lens.extend(data["image_seq_lens"]) multimodal_embeddings.append( self._get_image_features( - pixel_values=pixel_values, - image_position_ids=image_position_ids, + pixel_values=batched["pixel_values"], + image_position_ids=batched.get("image_position_ids"), image_seq_lens=( - image_seq_lens_extended - if len(image_seq_lens_extended) == pixel_values.shape[0] + image_seq_lens + if len(image_seq_lens) == batched["pixel_values"].shape[0] else None ), ) ) - - if video_pixel_values_list: - video_pixel_values = torch.cat(video_pixel_values_list) - video_position_ids = ( - torch.cat(video_position_ids_list) - if len(video_position_ids_list) == len(video_pixel_values_list) - else None - ) - multimodal_embeddings.append( - self._get_image_features( - pixel_values=video_pixel_values, - image_position_ids=video_position_ids, - image_seq_lens=( - video_seq_lens_extended - if len(video_seq_lens_extended) == video_pixel_values.shape[0] - else None - ), + if mixed_modality_runs: + # Own this run's storage before another call can replay a graph. + multimodal_embeddings[-1] = multimodal_embeddings[-1].clone( + memory_format=torch.contiguous_format ) - ) - - if audio_features_list and self.embed_audio is not None: - per_audio_embeddings = [ - self._get_audio_features(audio_features, audio_features_mask_list[index]) - for index, audio_features in enumerate(audio_features_list) - ] - multimodal_embeddings.append(torch.cat(per_audio_embeddings, dim=0)) if not multimodal_embeddings: raise ValueError("Gemma4 received active multimodal parameters without encoder inputs") diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 3d3bc305e537..85d26bafadda 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -17,7 +17,10 @@ import contextlib import copy import itertools +from collections.abc import Mapping from dataclasses import dataclass +from dataclasses import field as dataclass_field +from functools import partial from typing import ( TYPE_CHECKING, Any, @@ -62,6 +65,119 @@ class MultimodalEncoderContractError(ValueError): """A request-local MM encoder input or output contract violation.""" +def _batch_encoder_input( + values: Sequence[torch.Tensor], *, dim: int, pad_dims: tuple[int, ...], pad_value: float +) -> torch.Tensor: + if pad_dims and len(values) > 1: + maxima = {axis: max(value.shape[axis] for value in values) for axis in pad_dims} + padded = [] + for value in values: + padding = [0] * (2 * value.ndim) + for axis, maximum in maxima.items(): + padding[2 * (value.ndim - axis) - 1] = maximum - value.shape[axis] + padded.append( + torch.nn.functional.pad(value, padding, value=pad_value) if any(padding) else value + ) + values = padded + return torch.cat(values, dim=dim) + + +@dataclass(frozen=True) +class EncoderInputField: + """How one encoder argument combines already ordered input payloads. + + ``source_names`` lists aliases in lookup order; the first non-None value + is used. Only ``pad_dims`` may vary besides the concatenation dimension. + Padding extends the right side with ``pad_value``. Callers supply compatible + tensor ranks, dtypes, devices and remaining dimensions, and retain ownership + of modality order, optional-field policy and attention-mask semantics. + ``batch`` applies the same padding and owned concatenation to a gathered column. + """ + + source_names: tuple[str, ...] + dim: int = 0 + pad_dims: tuple[int, ...] = () + pad_value: float = 0.0 + batch: Callable[[Sequence[torch.Tensor]], torch.Tensor] = dataclass_field( + init=False, repr=False, compare=False + ) + + def __post_init__(self) -> None: + if not self.source_names: + raise ValueError("Encoder input fields require at least one source name") + if self.dim < 0 or any(axis < 0 or axis == self.dim for axis in self.pad_dims): + raise ValueError("Encoder padding axes must be nonnegative and exclude the concat axis") + object.__setattr__( + self, + "batch", + partial( + _batch_encoder_input, + dim=self.dim, + pad_dims=self.pad_dims, + pad_value=self.pad_value, + ) + if self.pad_dims or self.dim + else torch.cat, + ) + + +def batch_encoder_inputs( + inputs: Sequence[Mapping[str, Any]], + fields: Mapping[str, EncoderInputField], +) -> dict[str, torch.Tensor]: + """Gather, explicitly pad and concatenate encoder kwargs without reordering. + + Inputs are read-only and the result always owns concatenated storage, even + for one input. Models may use this inside ``EncoderGroup.build_batched_input`` + while existing grouped output restoration continues to use item metadata. + Lengths, grids and masks are never inferred from tensor shapes. + """ + if not inputs: + raise ValueError("Cannot batch empty encoder inputs") + result: dict[str, torch.Tensor] = {} + only = inputs[0] if len(inputs) == 1 else None + for name, field in fields.items(): + source_name = field.source_names[0] + if len(field.source_names) == 1: + try: + values = ( + [only[source_name]] + if only is not None + else [payload[source_name] for payload in inputs] + ) + except KeyError as exc: + raise ValueError( + f"Missing encoder input {name!r}: expected {field.source_names}" + ) from exc + elif len(field.source_names) == 2: + alternate = field.source_names[1] + try: + values = [ + value if (value := payload.get(source_name)) is not None else payload[alternate] + for payload in inputs + ] + except KeyError as exc: + raise ValueError( + f"Missing encoder input {name!r}: expected {field.source_names}" + ) from exc + else: + values = [] + for payload in inputs: + value = payload.get(source_name) + if value is None: + for alias in field.source_names[1:]: + value = payload.get(alias) + if value is not None: + break + if value is None: + raise ValueError( + f"Missing encoder input {name!r}: expected {field.source_names}" + ) + values.append(value) + result[name] = field.batch(values) + return result + + def _assemble_multimodal_encoder_embeddings( item_tensors: Dict[int, torch.Tensor], total_items: int ) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index ce2e57941655..a82c2b62c53c 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -54,8 +54,10 @@ from .modeling_multimodal_encoder import MultimodalEncoderMixin, prepare_encoder_attn_metadata from .modeling_multimodal_mixin import ( EncoderGroup, + EncoderInputField, MultimodalModelMixin, PreparedLlmInputs, + batch_encoder_inputs, encode_multimodal_by_groups, ) from .modeling_qwen2vl import ( @@ -1297,26 +1299,67 @@ def refresh_in_place( self._vision.prepare_attn_metadata(list(padded_seq_lengths), metadata) +_QWEN3VL_ENCODER_INPUT_FIELDS = { + "pixel_values": EncoderInputField(("pixel_values", "pixel_values_videos")), + "grid_thw": EncoderInputField(("image_grid_thw", "video_grid_thw")), +} +_QWEN3VL_SINGLE_MODALITY_FIELDS = tuple( + { + name: EncoderInputField((field.source_names[index],)) + for name, field in _QWEN3VL_ENCODER_INPUT_FIELDS.items() + } + for index in (0, 1) +) + + def _qwen3vl_build_batched_input( multimodal_params: list[MultimodalParams], ) -> dict[str, Any]: """Cat image items then video items across requests (matching the ``EncoderGroup.modalities`` order) into one ViT input.""" - pixels: list[torch.Tensor] = [] - grids: list[torch.Tensor] = [] - for m, pv_key, thw_key in ( - ("image", "pixel_values", "image_grid_thw"), - ("video", "pixel_values_videos", "video_grid_thw"), - ): - for mp in multimodal_params: - bucket = mp.multimodal_data.get(m) + pixels, grids = [], [] + try: + for param in multimodal_params: + bucket = param.multimodal_data.get("image") if bucket is not None: - pixels.append(bucket[pv_key]) - grids.append(bucket[thw_key]) - return { - "pixel_values": torch.cat(pixels, dim=0), - "grid_thw": torch.cat(grids, dim=0), - } + pixels.append(bucket["pixel_values"]) + grids.append(bucket["image_grid_thw"]) + has_images = bool(pixels) + for param in multimodal_params: + bucket = param.multimodal_data.get("video") + if bucket is not None: + # Two canonical fields leave no room for a preferred image alias. + if ( + len(bucket) > 2 + and has_images + and ("pixel_values" in bucket or "image_grid_thw" in bucket) + ): + break + pixels.append(bucket["pixel_values_videos"]) + grids.append(bucket["video_grid_thw"]) + else: + if pixels: + return { + "pixel_values": _QWEN3VL_ENCODER_INPUT_FIELDS["pixel_values"].batch(pixels), + "grid_thw": _QWEN3VL_ENCODER_INPUT_FIELDS["grid_thw"].batch(grids), + } + except (KeyError, TypeError, ValueError): + pass + + # Preserve field-major alias lookup and error priority for noncanonical inputs. + inputs_by_modality = [ + [ + bucket + for param in multimodal_params + if (bucket := param.multimodal_data.get(modality)) is not None + ] + for modality in ("image", "video") + ] + images, videos = inputs_by_modality + if images and videos: + return batch_encoder_inputs(images + videos, _QWEN3VL_ENCODER_INPUT_FIELDS) + index = 0 if images else 1 + return batch_encoder_inputs(inputs_by_modality[index], _QWEN3VL_SINGLE_MODALITY_FIELDS[index]) class Qwen3VisionModelBase(nn.Module): diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index b89f9ea2681e..7df6af2c44dc 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -55,8 +55,8 @@ from transformers.models.gemma4.modeling_gemma4 import ( # noqa: E402 Gemma4MultimodalEmbedder as HFGemma4MultimodalEmbedder, ) -from transformers.models.gemma4.modeling_gemma4 import ( # noqa: E402 - Gemma4VisionModel as HFGemma4VisionModel, +from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4VisionModel as HFGemma4VisionModel, # noqa: E402 ) from transformers.models.gemma4.video_processing_gemma4 import Gemma4VideoProcessor # noqa: E402 @@ -2297,3 +2297,90 @@ def test_forward_parity_head_dim_72(self): if __name__ == "__main__": unittest.main() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("modality", ["image", "video"]) +@pytest.mark.parametrize("positions", ["all", "partial", "none"]) +def test_common_vision_batch_preserves_optional_fields(modality, positions): + model = _Gemma4EncoderCacheHarness() + inputs = [] + for index in range(2): + data = { + "pixel_values": torch.full((index + 1, 3, 2), float(index + 1)), + "image_seq_lens": [3] * (index + 1), + } + if positions == "all" or (positions == "partial" and index == 0): + data["image_position_ids"] = torch.full((index + 1, 3, 2), index) + inputs.append(data) + params = [MultimodalParams(multimodal_data={modality: data}) for data in inputs] + with mock.patch.object( + _Gemma4EncoderCacheHarness, + "_get_image_features", + side_effect=lambda pixel_values, **_: pixel_values[:, 0], + ) as encoder: + output = model.encode_multimodal_inputs(params) + kwargs = encoder.call_args.kwargs + assert torch.equal(kwargs["pixel_values"], torch.cat([data["pixel_values"] for data in inputs])) + assert kwargs["image_seq_lens"] == [3, 3, 3] + if positions == "all": + assert torch.equal( + kwargs["image_position_ids"], torch.cat([data["image_position_ids"] for data in inputs]) + ) + else: + assert kwargs["image_position_ids"] is None + assert torch.equal(output, torch.tensor([[1.0, 1.0], [2.0, 2.0], [2.0, 2.0]])) + + +@pytest.mark.cpu_only +def test_common_vision_batch_matches_serial_with_reordered_items_and_audio_masks(): + model = _Gemma4EncoderCacheHarness() + params = [] + for modality, value in (("image", 1.0), ("image", 2.0), ("video", 3.0), ("audio", 4.0)): + if modality == "audio": + data = {"audio_features": torch.full((1, 3, 1), value), "audio_features_mask": None} + else: + data = {"pixel_values": torch.tensor([[[value]], [[value + 10]]])} + params.append(MultimodalParams(multimodal_data={modality: data})) + params[0].multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = MultimodalItemMetadata( + [("image", 0), ("image", 1)], output_embedding_lengths=[1, 1] + ) + params[0] = model.build_multimodal_encoder_input(params[0], [1, 0], "image") + with ( + mock.patch.object( + _Gemma4EncoderCacheHarness, + "_get_image_features", + side_effect=lambda pixel_values, **_: pixel_values[:, 0], + ), + mock.patch.object( + _Gemma4EncoderCacheHarness, + "_get_audio_features", + side_effect=lambda features, mask: features[:, 0], + ) as audio, + ): + batched = model.encode_multimodal_inputs(params) + serial = torch.cat([model.encode_multimodal_inputs([param]) for param in params]) + assert torch.equal(batched, serial) + assert torch.equal(batched[:2, 0], torch.tensor([11.0, 1.0])) + assert all(call.args[1] is None for call in audio.call_args_list) + + +@pytest.mark.cpu_only +def test_common_vision_runs_copy_outputs_before_static_buffer_reuse(): + model = _Gemma4EncoderCacheHarness() + params = [ + MultimodalParams(multimodal_data={modality: {"pixel_values": torch.full((1, 1, 1), value)}}) + for modality, value in (("image", 1.0), ("video", 2.0), ("image", 3.0)) + ] + static_output = torch.empty(1, 1) + + def encode(pixel_values, **kwargs): + assert kwargs["image_seq_lens"] is None + return static_output.copy_(pixel_values[:, 0]) + + with mock.patch.object( + _Gemma4EncoderCacheHarness, "_get_image_features", side_effect=encode + ) as encoder: + output = model.encode_multimodal_inputs(params) + assert encoder.call_count == 3 + assert torch.equal(output[:, 0], torch.tensor([1.0, 2.0, 3.0])) diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index c430a47081d8..47efb31e0d18 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -18,6 +18,11 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_qwen3vl from tensorrt_llm._torch.models.checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper +from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( + EncoderGroup, + MultimodalModelMixin, + encode_multimodal_by_groups, +) from tensorrt_llm._torch.models.modeling_qwen3vl import ( Qwen3VisionModel, Qwen3VLInputProcessorBase, @@ -25,7 +30,15 @@ _triton_pos_embed_interpolate, ) from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType +from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalInput, + MultimodalItemMetadata, + MultimodalItemPartition, + MultimodalParams, +) from tensorrt_llm.llmapi.llm_args import MultimodalConfig, MultimodalEncoderCudaGraphConfig from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -355,6 +368,177 @@ def setup_scenario(self, scenario: TestQwen3VLScenario): ) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "selection", ["full", "reversed_requests", "reordered_items", "partial", "partial_cache"] +) +def test_qwen3vl_common_batching_preserves_packing_and_deepstack_order(selection: str) -> None: + def make_params( + image_markers: list[int], video_marker: int, refs: list[tuple[str, int]] + ) -> MultimodalParams: + return MultimodalParams( + multimodal_data={ + "image": { + "pixel_values": torch.tensor(image_markers, dtype=torch.float32) + .repeat_interleave(4) + .reshape(-1, 1), + "image_grid_thw": torch.tensor([[1, 2, 2]] * len(image_markers)), + }, + "video": { + "pixel_values_videos": torch.full((8, 1), float(video_marker)), + "video_grid_thw": torch.tensor([[2, 2, 2]]), + }, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalItemMetadata( + refs, + encoder_token_lengths=[8 if modality == "video" else 4 for modality, _ in refs], + output_embedding_lengths=[ + 2 if modality == "video" else 1 for modality, _ in refs + ], + input_partitions={ + "image": { + "pixel_values": MultimodalItemPartition( + item_lengths=(4,) * len(image_markers) + ), + "image_grid_thw": MultimodalItemPartition(), + }, + "video": { + "pixel_values_videos": MultimodalItemPartition(item_lengths=(8,)), + "video_grid_thw": MultimodalItemPartition(), + }, + }, + ), + } + ) + + first = make_params([10, 20], 30, [("video", 0), ("image", 1), ("image", 0)]) + second = make_params([40], 50, [("image", 0), ("video", 0)]) + if selection in ("full", "partial_cache"): + params, markers = [first, second], [30, 30, 20, 10, 40, 50, 50] + elif selection == "reversed_requests": + params, markers = [second, first], [40, 50, 50, 30, 30, 20, 10] + else: + selected = [(first, 1), (second, 1)] + markers = [20, 50, 50] + if selection == "reordered_items": + selected.insert(1, (first, 2)) + markers.insert(1, 10) + params = [ + prepared + for prepared, _, _ in MultimodalModelMixin().prepare_multimodal_encoder_inputs(selected) + ] + + # Keep the previous builder's literal packing as an independent reference. + pixels, grids = [], [] + for modality, pixel_key, grid_key in ( + ("image", "pixel_values", "image_grid_thw"), + ("video", "pixel_values_videos", "video_grid_thw"), + ): + for param in params: + bucket = param.multimodal_data.get(modality) + if bucket is not None: + pixels.append(bucket[pixel_key]) + grids.append(bucket[grid_key]) + packed = modeling_qwen3vl._qwen3vl_build_batched_input(params) + assert torch.equal(packed["pixel_values"], torch.cat(pixels)) + assert torch.equal(packed["grid_thw"], torch.cat(grids)) + + encoder_inputs_seen = [] + + def visual( + pixel_values: torch.Tensor, grid_thw: torch.Tensor + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + primary = pixel_values[::4] + encoder_inputs_seen.append(primary.clone()) + return primary, [primary + 100, primary + 200] + + encoder = SimpleNamespace(model_dtype=torch.float32, visual=visual) + group = EncoderGroup( + ("image", "video"), + lambda **kwargs: modeling_qwen3vl.Qwen3VisionModelBase.encode_batched(encoder, **kwargs), + modeling_qwen3vl._qwen3vl_build_batched_input, + ) + actual = encode_multimodal_by_groups((group,), params) + primary = torch.tensor(markers, dtype=torch.float32).reshape(-1, 1) + expected = torch.cat([primary, primary + 100, primary + 200], dim=1) + assert torch.equal(actual, expected) + assert torch.equal( + actual, + torch.cat([encode_multimodal_by_groups((group,), [param]) for param in params]), + ) + assert first.item_metadata.item_refs == (("video", 0), ("image", 1), ("image", 0)) + if selection == "partial_cache": + + class CachedEncoder(MultimodalModelMixin): + def encode_multimodal_inputs( + self, multimodal_params: list[MultimodalParams] + ) -> torch.Tensor: + return encode_multimodal_by_groups((group,), multimodal_params) + + model = CachedEncoder() + model._multimodal_encoder_cache = TensorLRUCache(4096) + expected_items = [] + for request_index, param in enumerate(params): + lengths = param.item_metadata.output_embedding_lengths + param.multimodal_input = MultimodalInput.from_components( + [[request_index * 10 + index + 1] * 8 for index in range(len(lengths))], + [sum(lengths[:index]) for index in range(len(lengths))], + list(lengths), + ) + param.multimodal_data["mm_processor_kwargs_hash"] = "test" + item_outputs = torch.split(encode_multimodal_by_groups((group,), [param]), lengths) + keys = model._encoder_cache_keys(param) + model._multimodal_encoder_cache.put(keys[0], item_outputs[0]) + expected_items.extend(zip(keys, item_outputs)) + encoder_inputs_seen.clear() + assert torch.equal(model._get_or_encode_multimodal_embeddings(params), expected) + assert len(encoder_inputs_seen) == 1 + assert encoder_inputs_seen[0].flatten().tolist() == [20, 10, 50, 50] + for key, item_output in expected_items: + assert torch.equal(model._multimodal_encoder_cache.get(key), item_output) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("preferred_pixels", "preferred_grid", "expected_pixels", "expected_grid"), + [ + ( + torch.tensor([[5.0]]), + torch.tensor([[7, 8, 9]]), + [[1.0], [5.0]], + [[1, 2, 3], [7, 8, 9]], + ), + (None, None, [[1.0], [3.0]], [[1, 2, 3], [4, 5, 6]]), + ], + ids=["tensor", "none"], +) +def test_qwen3vl_common_batching_preserves_video_alias_priority( + preferred_pixels: Optional[torch.Tensor], + preferred_grid: Optional[torch.Tensor], + expected_pixels: list[list[float]], + expected_grid: list[list[int]], +) -> None: + image = { + "pixel_values": torch.tensor([[1.0]]), + "image_grid_thw": torch.tensor([[1, 2, 3]]), + } + video = { + "pixel_values_videos": torch.tensor([[3.0]]), + "video_grid_thw": torch.tensor([[4, 5, 6]]), + "pixel_values": preferred_pixels, + "image_grid_thw": preferred_grid, + } + params = [MultimodalParams(multimodal_data={"video": video, "image": image})] + + packed = modeling_qwen3vl._qwen3vl_build_batched_input(params) + + assert torch.equal(packed["pixel_values"], torch.tensor(expected_pixels)) + assert torch.equal(packed["grid_thw"], torch.tensor(expected_grid)) + for result in packed.values(): + for source in (*image.values(), *video.values()): + if source is not None: + assert result.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_qwen3vl_deepstack_fusion_reuses_registered_buffer(): hidden_size = 4 diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py index 3e04985b0a00..eec1f71a52c6 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -16,8 +16,10 @@ from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( EncoderGroup, + EncoderInputField, MultimodalEncoderContractError, MultimodalModelMixin, + batch_encoder_inputs, encode_multimodal_by_groups, ) from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VLModelBase @@ -907,3 +909,73 @@ def multimodal_hashes(self): model._apply_metadata_slice(residual, request.py_multimodal_params, [1]) assert residual.multimodal_input.multimodal_hashes == [[2] * 8] assert first.multimodal_hashes == [[1] * 8, [2] * 8] + + +@pytest.mark.parametrize("axis", [1, 2]) +def test_encoder_input_batch_pads_only_declared_axes(axis): + shape = [2, 3, 4] + shorter_shape = shape.copy() + shorter_shape[axis] -= 2 + first = torch.arange(24, dtype=torch.float32).reshape(shape) + second = torch.full(shorter_shape, 7.0) + result = batch_encoder_inputs( + [{"features": first}, {"features": second}], + {"values": EncoderInputField(("features",), pad_dims=(axis,), pad_value=-1.0)}, + )["values"] + assert result.shape == (4, 3, 4) + assert torch.equal(result[:2], first) + expected = torch.full(shape, -1.0) + expected.narrow(axis, 0, shorter_shape[axis]).copy_(second) + assert torch.equal(result[2:], expected) + assert torch.equal(second, torch.full(shorter_shape, 7.0)) + + +def test_encoder_input_batch_pads_features_and_boolean_validity_together(): + features = [torch.ones(2, 3, 4), torch.full((1, 5, 4), 2.0)] + masks = [torch.tensor([[True, True, False], [True, False, False]]), torch.ones(1, 5).bool()] + result = batch_encoder_inputs( + [{"x": x, "valid": mask} for x, mask in zip(features, masks, strict=True)], + { + "features": EncoderInputField(("x",), pad_dims=(1,)), + "mask": EncoderInputField(("valid",), pad_dims=(1,)), + }, + ) + assert result["mask"].dtype == torch.bool + assert torch.equal(result["mask"].sum(1), torch.tensor([2, 1, 5])) + assert torch.count_nonzero(result["features"][:2, 3:]) == 0 + assert not result["mask"][:2, 3:].any() + assert torch.equal(result["features"][2], features[1][0]) + + +@pytest.mark.parametrize("aliases", [("primary", "alternate"), ("primary", "unused", "alternate")]) +def test_encoder_input_batch_aliases_keep_input_order_and_owned_storage(aliases): + first = torch.arange(6).reshape(2, 3).T + second = torch.full_like(first, 9) + field = EncoderInputField(aliases, dim=1) + result = batch_encoder_inputs( + [{"primary": None, "alternate": first}, {"primary": second}], {"values": field} + )["values"] + assert torch.equal(result, torch.cat([first, second], dim=1)) + single = batch_encoder_inputs([{"primary": first}], {"values": field})["values"] + assert single.data_ptr() != first.data_ptr() + assert single.dtype == first.dtype and single.device == first.device + + +def test_encoder_input_batch_rejects_missing_field_instead_of_discarding_mask(): + with pytest.raises(ValueError, match="Missing encoder input 'mask'"): + batch_encoder_inputs( + [{"valid": torch.ones(1, 2)}, {}], + {"mask": EncoderInputField(("valid",), pad_dims=(1,))}, + ) + + +def test_encoder_input_batch_does_not_infer_padding(): + with pytest.raises(RuntimeError): + batch_encoder_inputs( + [{"x": torch.ones(1, 2)}, {"x": torch.ones(1, 3)}], + {"values": EncoderInputField(("x",))}, + ) + with pytest.raises(ValueError, match="empty encoder inputs"): + batch_encoder_inputs([], {"values": EncoderInputField(("x",))}) + with pytest.raises(ValueError, match="exclude the concat axis"): + EncoderInputField(("x",), pad_dims=(0,))