From eb64b087778eb3d489a91f39421d1a34740db7a2 Mon Sep 17 00:00:00 2001 From: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> Date: Tue, 3 Nov 2026 03:07:52 +0000 Subject: [PATCH] [TRTLLM-13419][refactor] Declare encoder inputs on EncoderGroup Cross-request packing of multimodal encoder inputs is written per model. Qwen3-VL has a private builder that concatenates image rows, then video rows, and Cosmos3 Edge imports that builder from another model file. The builder only maps encoder arguments to payload fields, so a declared mapping lets the framework do the packing. EncoderGroup gets an optional input_fields mapping: encoder kwarg -> {modality: payload field}. encode_multimodal_by_groups concatenates those fields along dim 0, modality first, then request order, with one torch.cat per kwarg, like the removed builder. build_batched_input stays for packing that is model policy. A group sets exactly one of the two, otherwise construction raises TypeError. The EncoderGroup docstring now describes both ways. Qwen3VisionModelBase (Qwen3-VL, Qwen3-VL-MoE, Qwen3.5-VL, Qwen4Exp, Cosmos3) and Cosmos3 Edge declare the mapping inline on their EncoderGroup, and _qwen3vl_build_batched_input is deleted, so Cosmos3 Edge no longer imports a private helper from the Qwen3-VL file. The encoder gets the same tensors in the same order, device and dtype as before. Packing that stays model-owned: - GLM5 Next expands each video grid into per-frame rows and moves grids to the pixel device before the concat. - Nemotron-H passes the params objects through, because its encoders route by payload and write per-request metadata back onto them. Validation: - New CPU tests check the packed inputs against a copy of the removed builder (one image, video only, mixed, video first, selected item views), the exactly-one TypeError, one cat per kwarg with a copy for one request, and that builders get the original params objects. - Touched and related multimodal test files on CPU: no new failures against main. The CUDA pixel case passes on GB200. - GB200, reduced Qwen3-VL encoder with random weights: outputs are bit-identical to main for image-only, video-only and mixed batches in the aggregated, encoder-only and item paths. - pre-commit hooks pass on the changed files. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_cosmos3_edge.py | 7 +- .../models/modeling_multimodal_mixin.py | 74 +++++++++--- .../_torch/models/modeling_qwen3vl.py | 28 +---- .../_torch/modeling/test_modeling_qwen3vl.py | 106 ++++++++++++++++++ .../test_multimodal_encoder_item_groups.py | 19 +--- .../_torch/multimodal/test_encoder_group.py | 101 +++++++++++++++++ .../test_qwen3vl_mixed_modality_encoder.py | 18 +-- 7 files changed, 287 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_cosmos3_edge.py b/tensorrt_llm/_torch/models/modeling_cosmos3_edge.py index 0b2d41241946..034ad8a882a7 100644 --- a/tensorrt_llm/_torch/models/modeling_cosmos3_edge.py +++ b/tensorrt_llm/_torch/models/modeling_cosmos3_edge.py @@ -77,7 +77,6 @@ Qwen3VLModelBase, Qwen3VLVisionMLP, Qwen3VLVisionPatchMerger, - _qwen3vl_build_batched_input, ) from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import ( @@ -468,7 +467,11 @@ def mm_encoder_groups(self) -> Tuple[EncoderGroup, ...]: EncoderGroup( modalities=("image", "video"), encoder_fn=self.encode_batched, - build_batched_input=_qwen3vl_build_batched_input, + # One ViT input carries both modalities, image rows first. + input_fields={ + "pixel_values": {"image": "pixel_values", "video": "pixel_values_videos"}, + "grid_thw": {"image": "image_grid_thw", "video": "video_grid_thw"}, + }, ), ) diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index cdf7da0c801f..98adb2a52aec 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -28,6 +28,7 @@ Iterable, Iterator, List, + Mapping, Optional, Sequence, Tuple, @@ -93,18 +94,26 @@ class EncoderGroup: splits the output back per-modality and reorders it into prompt order via each request's `MultimodalItemMetadata.item_refs` manifest. - Contract between `build_batched_input` and `encoder_fn`: + Contract between the encoder inputs and `encoder_fn`: - * `build_batched_input` must concatenate items across requests in - `modalities` order (all items for the first modality across requests, - then all items for the second modality, etc.). - * Within a modality, items must appear in the same per-request iteration + * Encoder inputs hold items across requests in `modalities` order (all + items for the first modality across requests, then all items for the + second modality, etc.). + * Within a modality, items appear in the same per-request iteration order that `_lengths_by_modality` uses (i.e. the order of `multimodal_params` passed in). - * `encoder_fn` must return one tensor whose rows correspond 1:1 to the - input layout produced by `build_batched_input`, so the framework can - split the output by `_lengths_by_modality` and reorder into prompt - order via each request's item manifest. + * `encoder_fn` must return one tensor whose rows correspond 1:1 to that + input layout, so the framework can split the output by + `_lengths_by_modality` and reorder into prompt order via each + request's item manifest. + + A group declares its encoder inputs in exactly one way: + + * `input_fields`: the framework concatenates the named payload fields + across requests in the layout above. + * `build_batched_input`: the model builds the `encoder_fn` kwargs itself, + for packing that is model policy (padding, per-item transforms, or + passing the params through so `encoder_fn` packs them). """ modalities: Tuple[str, ...] @@ -113,7 +122,8 @@ class EncoderGroup: `modalities[0]`, then all items of `modalities[1]`, etc.""" encoder_fn: Callable[..., torch.Tensor] - """Encoder call invoked as `encoder_fn(**build_batched_input(params))`. + """Encoder call invoked once per group as `encoder_fn(**inputs)`, with + the inputs built from `input_fields` or `build_batched_input`. Returns a single tensor with one row per embedding, laid out per the contract above. When those rows are already in prompt order, the framework may return a row slice of this tensor without a copy, so it @@ -121,10 +131,40 @@ class EncoderGroup: graph's static output buffer); clone such outputs before returning them.""" - build_batched_input: Callable[[List[MultimodalParams]], Dict[str, Any]] - """Builds the kwargs dict passed to `encoder_fn`. Responsible for - concatenating raw per-item tensors from `multimodal_data` across - requests in the order described in the class docstring.""" + build_batched_input: Optional[Callable[[List[MultimodalParams]], Dict[str, Any]]] = None + """Model-owned packing. Receives the group's original `MultimodalParams` + objects in order (not copies) and returns the kwargs dict passed to + `encoder_fn`, laid out per the contract above.""" + + input_fields: Optional[Mapping[str, Mapping[str, str]]] = None + """Common packing: `encoder_fn` kwarg -> {modality: field of + `multimodal_data[modality]`}. Each kwarg names a field for every modality + in `modalities`. Every listed field must be a tensor whose items lie + along dim 0 (`MultimodalItemPartition(item_axis=0)`); shared values, + lists and padded fields belong in `build_batched_input`. Values are + concatenated along dim 0 on their own device and dtype, even when only + one request contributes; casting belongs to `encoder_fn`. A declared + field missing from a present payload raises KeyError.""" + + def __post_init__(self) -> None: + if (self.input_fields is None) == (self.build_batched_input is None): + raise TypeError("EncoderGroup takes exactly one of input_fields or build_batched_input") + + +def _concat_input_fields( + group: EncoderGroup, multimodal_params: Sequence[MultimodalParams] +) -> Dict[str, torch.Tensor]: + """Concatenate `group.input_fields` across requests in the group layout.""" + buckets = [] + for modality in group.modalities: + for param in multimodal_params: + data = param.multimodal_data.get(modality) + if data is not None: + buckets.append((modality, data)) + return { + name: torch.cat([data[sources[modality]] for modality, data in buckets], dim=0) + for name, sources in group.input_fields.items() + } def _lengths_by_modality( @@ -264,7 +304,11 @@ def encode_multimodal_by_groups( group_layouts.append(layout) if not group_params: continue - out = group.encoder_fn(**group.build_batched_input(group_params)) + if group.input_fields is not None: + encoder_inputs = _concat_input_fields(group, group_params) + else: + encoder_inputs = group.build_batched_input(group_params) + out = group.encoder_fn(**encoder_inputs) lengths = _lengths_by_modality(group_layouts, group.modalities) cursor = 0 for m in group.modalities: diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 4d1acbebe8a1..c5ffb768a639 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -1297,28 +1297,6 @@ def refresh_in_place( self._vision.prepare_attn_metadata(list(padded_seq_lengths), metadata) -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) - 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), - } - - class Qwen3VisionModelBase(nn.Module): def __init__( self, @@ -1407,7 +1385,11 @@ def mm_encoder_groups(self) -> Tuple[EncoderGroup, ...]: EncoderGroup( modalities=("image", "video"), encoder_fn=self.encode_batched, - build_batched_input=_qwen3vl_build_batched_input, + # One ViT input carries both modalities, image rows first. + input_fields={ + "pixel_values": {"image": "pixel_values", "video": "pixel_values_videos"}, + "grid_thw": {"image": "image_grid_thw", "video": "video_grid_thw"}, + }, ), ) diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index 8ea278c27080..7aad2e44e5f4 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -20,14 +20,17 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_qwen2vl, modeling_qwen3vl from tensorrt_llm._torch.models.checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper +from tensorrt_llm._torch.models.modeling_multimodal_mixin import _concat_input_fields from tensorrt_llm._torch.models.modeling_qwen3vl import ( Qwen3VisionModel, + Qwen3VisionModelBase, Qwen3VLInputProcessorBase, Qwen3VLModel, _triton_pos_embed_interpolate, ) from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.inputs.multimodal import 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 @@ -651,6 +654,109 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: assert events == ["block_0", "deepstack_merger", "block_1", "final_merger"] +def _main_qwen3vl_build_batched_input(multimodal_params: List[MultimodalParams]) -> dict: + """Verbatim copy of the Qwen3-VL builder that `input_fields` replaced.""" + 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) + 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), + } + + +def _qwen_vision_request(images=(), videos=(), pixel_device="cpu") -> MultimodalParams: + """Build one request payload; each item is `(marker, (t, h, w))`.""" + data = {} + for modality, items, pixel_key, grid_key in ( + ("image", images, "pixel_values", "image_grid_thw"), + ("video", videos, "pixel_values_videos", "video_grid_thw"), + ): + if items: + pixels = [torch.full((t * h * w, 3), float(marker)) for marker, (t, h, w) in items] + data[modality] = { + pixel_key: torch.cat(pixels).to(pixel_device), + grid_key: torch.tensor([grid for _, grid in items]), + } + return MultimodalParams(multimodal_data=data) + + +def _assert_qwen3_vision_inputs_match_main(params: List[MultimodalParams]) -> None: + (group,) = Qwen3VisionModelBase.mm_encoder_groups.fget( + SimpleNamespace(encode_batched=lambda **inputs: inputs) + ) + actual = _concat_input_fields(group, params) + expected = _main_qwen3vl_build_batched_input(params) + assert actual.keys() == expected.keys() + for name, value in expected.items(): + assert torch.equal(actual[name], value), name + assert actual[name].dtype == value.dtype, name + assert actual[name].device == value.device, name + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "requests", + [ + [{"images": [(1, (1, 2, 2))]}], + [{"videos": [(1, (2, 2, 2))]}, {"videos": [(2, (1, 2, 4))]}], + [ + {"images": [(1, (1, 2, 2)), (2, (1, 2, 4))], "videos": [(3, (2, 2, 2))]}, + {"videos": [(4, (1, 2, 2))]}, + {"images": [(5, (1, 4, 2))]}, + ], + [{"videos": [(1, (2, 2, 2))]}, {"images": [(2, (1, 2, 2))]}], + ], + ids=["image_r1", "video_only", "mixed_r3", "video_first"], +) +def test_qwen3_vision_encoder_inputs_match_main_builder(requests: list) -> None: + _assert_qwen3_vision_inputs_match_main([_qwen_vision_request(**r) for r in requests]) + + +@pytest.mark.cpu_only +def test_qwen3_vision_encoder_inputs_match_main_builder_on_selected_items() -> None: + # Item scheduling and partial cache hits encode residual payloads whose + # fields are views of the request payload. + request = _qwen_vision_request( + images=[(1, (1, 2, 2)), (2, (1, 2, 4)), (3, (1, 2, 2))], videos=[(4, (2, 2, 2))] + ) + image = request.multimodal_data["image"] + residual_image = { + "pixel_values": MultimodalItemPartition(item_lengths=(4, 8, 4)).select( + image["pixel_values"], [1, 2] + ), + "image_grid_thw": MultimodalItemPartition().select(image["image_grid_thw"], [1, 2]), + } + assert residual_image["pixel_values"].data_ptr() != image["pixel_values"].data_ptr() + assert ( + residual_image["pixel_values"].untyped_storage().data_ptr() + == image["pixel_values"].untyped_storage().data_ptr() + ) + residual = MultimodalParams( + multimodal_data={"image": residual_image, "video": request.multimodal_data["video"]} + ) + _assert_qwen3_vision_inputs_match_main( + [residual, _qwen_vision_request(images=[(5, (1, 2, 2))])] + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen3_vision_encoder_inputs_keep_cpu_grids_with_cuda_pixels() -> None: + params = [ + _qwen_vision_request(images=[(1, (1, 2, 2))], videos=[(2, (2, 2, 2))], pixel_device="cuda"), + _qwen_vision_request(images=[(3, (1, 2, 4))], pixel_device="cuda"), + ] + _assert_qwen3_vision_inputs_match_main(params) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("attn_backend", ["TRTLLM", "FLASHINFER"]) @pytest.mark.parametrize("deepstack_visual_indexes", [[0], [1, 0]]) 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..e15a4bc4921d 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py @@ -337,6 +337,9 @@ def _grid_for_rows(rows: int) -> list[int]: return [1, 1, rows] +_VISION_INPUT_FIELDS = {"pixel_values": {"image": "pixel_values", "video": "pixel_values_videos"}} + + class _GroupedEncoderModel(MultimodalModelMixin): """Minimal model whose encoder serves image and video in one call. @@ -357,22 +360,10 @@ def __init__(self): EncoderGroup( modalities=("image", "video"), encoder_fn=self._encode, - build_batched_input=self._build_batched_input, + input_fields=_VISION_INPUT_FIELDS, ), ) - @staticmethod - def _build_batched_input(params: list[MultimodalParams]) -> dict: - # Per the EncoderGroup contract: all items of the first modality across - # requests, then all items of the second. - pixels = [] - for modality, pixel_key in (("image", "pixel_values"), ("video", "pixel_values_videos")): - for mp in params: - data = mp.multimodal_data.get(modality) - if data is not None: - pixels.append(data[pixel_key]) - return {"pixel_values": torch.cat(pixels, dim=0)} - def _encode(self, pixel_values: torch.Tensor) -> torch.Tensor: self.encoder_calls.append(pixel_values.shape[0]) # One output row per input patch row, carrying the input value so the @@ -445,7 +436,7 @@ def _get_dummy_grid_for_modality( group = EncoderGroup( modalities=("image", "video"), encoder_fn=lambda pixel_values: pixel_values[::4], - build_batched_input=_GroupedEncoderModel._build_batched_input, + input_fields=_VISION_INPUT_FIELDS, ) output = encode_multimodal_by_groups([group], [params]) assert output.shape == (sum(expected_rows), HIDDEN) diff --git a/tests/unittest/_torch/multimodal/test_encoder_group.py b/tests/unittest/_torch/multimodal/test_encoder_group.py index b04224b4bd81..d8b3c9bab067 100644 --- a/tests/unittest/_torch/multimodal/test_encoder_group.py +++ b/tests/unittest/_torch/multimodal/test_encoder_group.py @@ -9,14 +9,19 @@ * ``MultimodalParams.get_item_layout`` — item metadata with legacy fallbacks. * ``_reorder_embeds_by_manifest`` — slice per-modality tensors and concat in each request's prompt-order manifest. +* ``EncoderGroup`` inputs — ``input_fields`` packing and ``build_batched_input`` + pass-through. """ import pytest import torch from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( + EncoderGroup, + _concat_input_fields, _lengths_by_modality, _reorder_embeds_by_manifest, + encode_multimodal_by_groups, ) from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, @@ -265,3 +270,99 @@ def test_no_embeds_returns_empty(self): mp = _mp(buckets={}) out = _reorder_embeds_by_manifest([(mp.item_metadata.item_refs, ())], {}, {}) assert out.numel() == 0 + + +class TestEncoderGroupInputs: + _INPUT_FIELDS = { + "values": {"image": "image_values", "video": "video_values"}, + "sizes": {"image": "image_sizes", "video": "video_sizes"}, + } + + @staticmethod + def _encode(values, sizes): + return values + + @staticmethod + def _payload(modality, marker, rows): + return { + f"{modality}_values": torch.full((rows, 2), float(marker)), + f"{modality}_sizes": torch.tensor([[rows]]), + } + + def _group(self): + return EncoderGroup(("image", "video"), self._encode, input_fields=self._INPUT_FIELDS) + + def test_requires_exactly_one_input_declaration(self): + with pytest.raises(TypeError, match="exactly one"): + EncoderGroup(("image",), self._encode) + with pytest.raises(TypeError, match="exactly one"): + EncoderGroup( + ("image",), + self._encode, + lambda params: {}, + input_fields={"values": {"image": "image_values"}}, + ) + + def test_concat_orders_modality_then_request_and_skips_absent_payloads(self): + params = [ + MultimodalParams(multimodal_data={"video": self._payload("video", 1, 3)}), + MultimodalParams(multimodal_data={"image": None, "audio": {}}), + MultimodalParams( + multimodal_data={ + "image": self._payload("image", 2, 1), + "video": self._payload("video", 3, 2), + } + ), + MultimodalParams(multimodal_data={"image": self._payload("image", 4, 2)}), + ] + inputs = _concat_input_fields(self._group(), params) + assert list(inputs) == ["values", "sizes"] + assert torch.equal(inputs["values"][:, 0], torch.tensor([2.0, 4, 4, 1, 1, 1, 3, 3])) + assert torch.equal(inputs["sizes"], torch.tensor([[1], [2], [3], [2]])) + + def test_concat_copies_once_per_field_even_for_one_request(self, monkeypatch): + payload = self._payload("image", 1, 2) + calls = [] + cat = torch.cat + + def counting_cat(tensors, *args, **kwargs): + calls.append(len(tensors)) + return cat(tensors, *args, **kwargs) + + monkeypatch.setattr(torch, "cat", counting_cat) + inputs = _concat_input_fields( + self._group(), [MultimodalParams(multimodal_data={"image": payload})] + ) + monkeypatch.undo() + assert calls == [1, 1] + assert torch.equal(inputs["values"], payload["image_values"]) + assert inputs["values"].data_ptr() != payload["image_values"].data_ptr() + + def test_missing_declared_field_raises(self): + payload = self._payload("image", 1, 2) + del payload["image_sizes"] + with pytest.raises(KeyError, match="image_sizes"): + _concat_input_fields( + self._group(), [MultimodalParams(multimodal_data={"image": payload})] + ) + + def test_builder_receives_original_params(self): + received = [] + + def build(params): + received.extend(params) + return { + "values": torch.cat([p.multimodal_data["image"]["image_values"] for p in params]) + } + + group = EncoderGroup(("image",), lambda values: values, build) + assert group.build_batched_input is build + assert group.input_fields is None + params = [ + _mp(embedding_lengths=[rows], buckets={"image": self._payload("image", marker, rows)}) + for marker, rows in ((1, 2), (2, 1)) + ] + out = encode_multimodal_by_groups([group], params) + assert len(received) == len(params) + assert all(actual is expected for actual, expected in zip(received, params)) + assert torch.equal(out[:, 0], torch.tensor([1.0, 1, 2])) diff --git a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py index 15f9c543fd86..e7eab724fcb2 100644 --- a/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py +++ b/tests/unittest/_torch/multimodal/test_qwen3vl_mixed_modality_encoder.py @@ -10,15 +10,13 @@ disagg entrypoint) delegates to, so this suite covers both paths. """ +from types import SimpleNamespace from typing import Optional import torch -from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( - EncoderGroup, - encode_multimodal_by_groups, -) -from tensorrt_llm._torch.models.modeling_qwen3vl import _qwen3vl_build_batched_input +from tensorrt_llm._torch.models.modeling_multimodal_mixin import encode_multimodal_by_groups +from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VisionModelBase from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalItemMetadata, @@ -32,12 +30,8 @@ def _encode_batched(pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch return pixel_values.clone().repeat(1, 2) # widen to hidden=marker_dim*2 -_QWEN3VL_ENCODER_GROUPS = ( - EncoderGroup( - modalities=("image", "video"), - encoder_fn=_encode_batched, - build_batched_input=_qwen3vl_build_batched_input, - ), +_QWEN3VL_ENCODER_GROUPS = Qwen3VisionModelBase.mm_encoder_groups.fget( + SimpleNamespace(encode_batched=_encode_batched) ) @@ -166,7 +160,7 @@ def test_heterogeneous_batch_image_only_and_video_only() -> None: def test_heterogeneous_batch_video_first_then_image() -> None: """Reviewer's mm-encoder-only regression case: `[video-only req0, - image-only req1]`. `_qwen3vl_build_batched_input` concatenates images + image-only req1]`. The Qwen3-VL encoder group concatenates images across requests first, then videos, so without prompt-order reordering the encoder output would be `[img(req1), vid(req0)]` — request-ordered downstream splits would then hand each request the other's rows.