[TRTLLM-13419][refactor] Batch multimodal encoder inputs by item partitions by yechank-nvidia · Pull Request #21052 · NVIDIA/TensorRT-LLM

yechank-nvidia · 2026-10-27T00:28:01Z

Description

Why: When the multimodal encoder runs several requests together, each model packs its own encoder inputs across requests in its own code. These copies drift apart: some models call the tower once per request or per run of requests, and Qwen2-VL could not serve image and video requests in one batch. The input processors already declare how every item sits in its payload fields (MultimodalItemMetadata.input_partitions), and item scheduling and the encoder cache already use these declarations to cut items out. The same declarations are enough to batch the items of several requests. So this PR adds one common batching path that is driven by those partitions and is the inverse of item selection, and moves the core models to it.

What:

New batch_by_item_partitions(pieces) in tensorrt_llm/inputs/multimodal.py, next to select_multimodal_fields. It batches the pieces of one payload key (whole requests or single items) by their partitions. A shared value (item_axis=None) is returned once. Other pieces are concatenated along the item axis, and other dims that differ are padded at their end with the partition's padding_value, a new MultimodalItemPartition field (zero, False for bool, by default; item selection keeps it). The returned partition lists every batched item, so select gets each piece back.
New default builder build_batched_input_by_partitions. How a key is batched (item axis, padding value) is part of the partition that the input processor declares, so it is decided once and travels with the data. One pass over the payloads (each modality of the group, then each request) gathers every declared key, one check per key decides whether the payloads can share a call, and each key is batched once. Only declared keys reach the encoder, as keyword arguments named after the payload keys.
The input checks stay, without a helper call per payload: the payloads of a key must share item axis and padding value, packed sizes must match their item lengths, item tensors must share dtype and rank, and the item record must be a MultimodalItemMetadata.
EncoderGroup.build_batched_input is now optional. Without it, the group binds the default builder, so a model writes EncoderGroup(("image", "video"), self.encode_batched). A model whose packing is different still passes its own builder, and that path works as before.
encode_multimodal_by_groups puts payloads in separate encoder calls when a shared value (such as a normalization record) or a tensor dtype differs, or when payloads of one modality give different keys, and then puts the rows back in prompt order. Normally there is one call. For every group, the encoder must now return as many rows as the item metadata expects; otherwise the encode raises MultimodalEncoderContractError instead of slicing silently.

Models moved in this PR (one commit each): the Qwen3-VL family (Qwen3-VL, Qwen3-VL-MoE, Qwen3.5-VL, Qwen4Exp, Cosmos3) and Cosmos3 Edge; Qwen2-VL, Qwen2.5-VL and Exaone4.5; Mistral3 and Pixtral; Gemma4 and Gemma4 Unified. Each model drops its own packing code and declares every key its encoder needs (for example vision_normalization as a shared value).

Behaviour changes:

Qwen2-VL, Qwen2.5-VL and Exaone4.5 now serve batches that hold both image and video requests: one ViT call, and one output tensor in request order, then prompt order. Before, the encoder-only runner rejected such batches, and the aggregated path paired the two modality tensors with requests by position.
Gemma4 encodes all audio clips of a batch in one padded, masked call (before: one call per request). Gemma4 patch positions declare -1, the processors' padding sentinel, as their padding value, so vision payloads with different patch counts share a call and the vision encoder needs no restore pass, host copy or stream sync to mark the padding. Requests of one modality that are not next to each other in the batch share one call per modality (before: one call per run of adjacent requests).
Qwen3-VL and Gemma4 now use two calls instead of one when raw uint8 pixels and already-normalized pixels of one modality are in the same batch. This happens when GPU normalization is on and only some requests pass mm_processor_kwargs. Batches that agree on pixel dtype and normalization record, the normal case, still make one call with the same input as before.
Mistral3 keeps image_sizes on the host (multimodal_data_device_paths), so its encoder crops the pixels without a device sync.
Gemma4 Unified counts the output rows of each item in its input processor when a request has no item metadata yet (no multimodal hashing).
The row-count check also runs for the groups that keep their own builder (GLM5 Next vision and Nemotron-H). Their CPU tests are unchanged; their GPU paths were not run for this PR.

Compatibility: Python only. No C++ change and no new config field, environment variable or CLI flag. EncoderGroup keeps its fields and their order, and groups that pass their own build_batched_input work as before. MultimodalItemPartition gains the optional third field padding_value (default 0), so existing partitions keep their meaning.

Note for earlier reviewers: this PR was rewritten. The earlier version added a new EncoderGroup.input_fields mapping and moved every model at once. This version reuses the partition declarations that are already on main and moves only the core models.

Performance

Screening only, not a formal benchmark: one GB200 GPU on a shared node, bf16, TP=1, TRTLLM attention, random-weight towers with real widths and fewer layers (Qwen2.5-VL: 8 layers of width 1280; Qwen3-VL: 8 layers of width 1152; Gemma4: 26B vision widths with 4 layers and the default audio tower with 12 layers). Baseline: upstream main bae7da432c. Each round runs 3 warm-up and then 15 timed batched encodes of all requests, one at a time. The timer wraps the whole encode with a CUDA synchronize before and after, and leaves out building the payloads and moving them to the GPU. Two rounds per side (main, PR, main, PR); each cell shows the round's median time, main → PR, and the change. Qwen images have patch grid (1, 32, 32) and Qwen2.5-VL videos (4, 24, 24); Gemma4 payloads come from the real image, video and audio processors.

Model	Batch	Tower calls (main → PR)	Round 1: main → PR (ms)	Round 2: main → PR (ms)
Qwen2.5-VL	8 image + 8 video requests, alternating	2 → 1	18.79 → 15.85 (-15.6%)	18.84 → 15.87 (-15.7%)
Qwen2.5-VL	16 image requests (control)	1 → 1	10.80 → 10.84 (+0.4%)	10.80 → 10.87 (+0.7%)
Qwen3-VL	16 image requests, raw uint8 and normalized alternating	1 → 2	10.41 → 15.65 (+50.3%)	10.49 → 15.85 (+51.1%)
Qwen3-VL	16 normalized image requests (control)	1 → 1	10.08 → 10.15 (+0.6%)	10.14 → 10.20 (+0.6%)
Gemma4	8 audio requests, one 1-30 s clip each	8 → 1	192.55 → 30.88 (-84.0%)	193.32 → 30.92 (-84.0%)
Gemma4	12 requests cycling image, audio, video	12 → 3	149.63 → 41.76 (-72.1%)	151.54 → 41.07 (-72.9%)
Gemma4	8 image requests, raw uint8 and host-normalized alternating	1 → 2	15.16 → 16.44 (+8.5%)	15.21 → 16.48 (+8.3%)
Gemma4	8 image requests (control)	1 → 1	14.85 → 14.89 (+0.3%)	14.85 → 14.91 (+0.4%)

The gain comes where the number of tower calls drops: Gemma4 audio -84%, the mixed Gemma4 batch -72% and the mixed Qwen2.5-VL batch -16% (round 1). When the call count does not change (the controls), the PR adds at most +0.07 ms per batch (+0.7%). A batch that mixes raw and normalized pixels pays for one extra call. Mistral3 was not timed, because its call count does not change. End-to-end serving was not measured.

Host work of the common path, measured with the towers replaced by stubs that only allocate their outputs (120 warm encodes per side, main → PR): Qwen2.5-VL control: 14 → 32 Python calls, 0.072 → 0.125 ms (+0.053 ms); Qwen3-VL control: 116 → 37 Python calls, 0.141 → 0.166 ms (+0.025 ms); Gemma4 control: 21 → 26 Python calls, 0.081 → 0.098 ms (+0.016 ms). Kernel launches are the same as main, no encode makes a stream sync, and the Gemma4 control makes no copies (main: 1). Between two runs of this pass, main's own time for a control moved by up to 0.021 ms, so differences of that size are run-to-run variation. The input checks take 7.6 us of the 38.1 us host path for 16 Qwen image requests (CPU microbenchmark).

Test Coverage

tests/unittest/inputs/test_multimodal.py (-m cpu_only): batch_by_item_partitions next to the select tests: shared values, one row per item, packed lengths, ragged and bool padding, padding values, a non-leading item axis, one piece, lists, select returning each piece, whole requests mixed with single items, and the errors (146 tests, 18 new).
tests/unittest/_torch/multimodal/test_encoder_group.py (-m cpu_only): default builder order and declared keys, padding values, packed sizes, the item record type, one call for equal records, split calls for different shared values, dtypes and keys, own builders getting the original params in one call, and the row-count check.
Model tests, each compared with a copy of the packing it replaced, per request and batched, including selected-item views, profiling dummies and split calls: test_modeling_qwen3vl.py (Qwen3-VL, Qwen2-VL / Qwen2.5-VL and the shared Qwen requests), test_modeling_cosmos3_edge.py, test_modeling_mistral.py and test_gemma4_multimodal.py, all under tests/unittest/_torch/modeling/.
CPU suite (-m cpu_only over tests/unittest/_torch/multimodal, the multimodal model test files and tests/unittest/inputs/test_multimodal.py): 1950 passed, 77 skipped, 4 failed, against 1843 / 77 / 4 on main. No test changed status. The 4 failures are the Nemotron Nano V2 VL input-processor tests that fail on main too because the checkpoint is missing on the test machine. Every commit of the PR passes the same check on its own.
GPU, the touched unit test files with all markers on one GB200: 972 passed, 59 skipped, 8 failed, against 865 / 59 / 8 on main; no test changed status; the 8 failures (test_qwen3_vision_cuda_graph_replay_matches_eager with the FLASHINFER backend) fail the same way on main on the test machine.
GPU equivalence on one GB200 with reduced random-weight towers (bf16, TF32 off) for Qwen2.5-VL, Qwen3-VL, Mistral3 / Pixtral, Gemma4 and Gemma4 Unified (48 cases), against main: every request encoded alone is bitwise equal to main. Batched rows are bitwise equal to main's batched rows, including the Qwen mixed image and video batches and the Qwen3-VL batch split into two calls. The only exceptions are calls that now hold more items: Gemma4 batched audio is within the bf16 floating-point-order bound (rel L2 8.2e-3, and 1.6e-6 in fp32), and Gemma4 Unified non-adjacent image requests are within the same bound.
Not covered: end-to-end serving and real-checkpoint accuracy, exclusive-node formal performance, and TP > 1.

Follow-up: the other multimodal models (for example Gemma3, LLaVA-NeXT, Llama4, Kimi VL, GLM5 Next, Nemotron-H, Phi-4 multimodal, MiniCPM-V, MiniMax-M3, Step3p7, HyperCLOVA X, DeepSeek-V4.1 and MuseSpark) keep their current code in this PR and move to the default builder in later PRs.

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Problem: Cross-request batching of multimodal encoder inputs is written per model: each model packs its own payload fields for one encoder call. The input processors already declare how items occupy those fields (MultimodalItemMetadata.input_partitions), which item scheduling and the encoder cache use to cut items out. The same declarations are enough to batch the items of several requests, so the packing can be shared. The shared path has to stay cheap on the host, because a batch that a model already encodes in one call gains nothing from it. A draft that batched each key with helper calls per payload made 474 Python calls for 16 Qwen2.5-VL image requests where main makes 14, and added 0.118 ms of host time per batch, most of it before the first tower call, which the GPU waits for. Solution: - batch_by_item_partitions batches pieces of one field (whole requests or single items) by their partitions, the inverse of select_multimodal_fields. A shared value (item_axis=None) is returned once. Other pieces are concatenated along the item axis; tensors whose other dims differ are padded at the end of those dims with the partition's padding_value, a new MultimodalItemPartition field (zero, False for bool, by default; item selection keeps it). The returned partition lists every batched item, so select gets each piece back. - build_batched_input_by_partitions is the common builder. How a key is batched (item axis, padding value) is part of the partition that the input processor declares, so it is decided once and travels with the data. The builder gathers every declared key in one pass over the payloads (each modality of the group, then each request) and batches each key once. Only declared keys reach the encoder. A key that several modalities declare batches their items in modality order. - The input checks stay, without a helper call per payload: the payloads of a key must share item axis and padding value, a packed value's size must match its item lengths, item tensors must share dtype and rank (checked before torch.cat, which also accepts a 1-D empty tensor next to tensors of another rank), and the item record must be a MultimodalItemMetadata. - EncoderGroup.build_batched_input is optional. When it is None, the group binds the common builder for its modalities. A group that passes its own builder works as before. - encode_multimodal_by_groups: with the common builder, payloads whose shared value or tensor dtype for a key differs (for example raw uint8 pixels next to normalized ones), or whose keys differ within one modality, are encoded in separate calls, and the rows are put back in the group layout. Normally this is one call: one check per key over the gathered payloads finds it, and only a batch that has to split is looked at payload by payload. The prompt-order restore reads the item layouts without a helper call per request. Groups with their own builder keep one call, since their payloads may not declare partitions. For every group the encoder rows must now match the item metadata; otherwise the encode raises MultimodalEncoderContractError instead of slicing silently. No model uses the common builder yet; models move to it in the following commits. Performance: This commit alone: not measured, since no model uses the common builder yet. Host side (screening, not a formal benchmark): one GB200 GPU on a shared node; the towers are replaced by stubs that only allocate their outputs, so the time is the host work of one batched encode plus its copy kernels, with a CUDA synchronize before and after; payloads are built and moved to the GPU outside the timer; 5 warm-up and 60 timed encodes per process, 2 processes per side, median of the 120 samples; warm means no gc.collect between encodes, cold means gc.collect before each. Baseline: upstream main bae7da4. With the following model commits applied, the three controls, whose tower-call count does not change, read (main -> this change): - Qwen2.5-VL, 16 image requests: Python calls 14 -> 32; host time 0.072 -> 0.125 ms warm (+0.053 ms), 0.328 -> 0.468 ms cold (+0.140 ms) - Qwen3-VL, 16 image requests: Python calls 116 -> 37; host time 0.141 -> 0.166 ms warm (+0.025 ms), 0.484 -> 0.532 ms cold (+0.048 ms) - Gemma4, 8 image requests: Python calls 21 -> 26; host time 0.081 -> 0.098 ms warm (+0.016 ms), 0.412 -> 0.428 ms cold (+0.016 ms) Kernel launches are the same as main and no encode makes a stream synchronize. The draft above made 474, 472 and 384 Python calls for these controls and added 0.118, 0.081 and 0.184 ms warm against its baseline (upstream main 66ce99b). Between two runs of this pass, main's own warm time for a control moved by up to 0.021 ms, so differences of that size are run-to-run variation. The input checks are cheap: in a CPU microbenchmark of the host path up to the encoder call for 16 Qwen image requests, they take 7.6 us of 38.1 us. Real-tower timings are in the model commits. Validation: - New CPU tests in tests/unittest/inputs/test_multimodal.py, next to the select tests: batching shared values, one row per item, packed lengths, ragged and bool padding, a ragged middle piece, a padding value, a non-leading item axis, one piece, lists, select returning each piece, whole requests mixed with single items, item selection keeping padding values, and the errors, including a 1-D empty piece between 2-D pieces. - New CPU tests in tests/unittest/_torch/multimodal/test_encoder_group.py: the common builder's modality-then-request order, declared keys only, a key shared by image and video, a padding value reaching the encoder, payloads that disagree on how a key is batched, a payload without partitions, a packed value whose size does not match its item lengths and an item record of another type (both rejected by the builder and by the grouped encode); one call for compatible payloads with prompt order restored, one call for equal records held by different requests, separate calls for different shared values and dtypes, own builders getting the original params in one call, and the row-count check with one call and with split calls. - Multimodal CPU suite at this commit (-m cpu_only over tests/unittest/_torch/multimodal, the multimodal model test files and tests/unittest/inputs/test_multimodal.py): 1877 passed, 77 skipped, 4 failed; main bae7da4: 1843 / 77 / 4. No test changed status and none was removed; the 34 new tests pass. The 4 failures also fail on main (Nemotron Nano V2 VL input processor, checkpoint missing on the test machine). The CPU grouping tests that carry no cpu_only marker, run with all markers: no status change at any commit of the series. The test machine's native build is older than main; six operators it lacks are left unregistered, and no selected test uses them. - pre-commit hooks pass on the changed files. - Reproduce: pytest -q -m cpu_only tests/unittest/inputs/test_multimodal.py tests/unittest/_torch/multimodal/test_encoder_group.py - GPU and end-to-end serving: not run for this commit alone, since no model's encode path changes in it. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com>

Problem: Gemma4 packs its encoder inputs in encode_multimodal_inputs. It splits the batch into runs of adjacent requests with the same modality and calls itself once per run. Within a run it normalizes raw pixels per payload, batches the image or video payloads for one tower call, and runs the audio tower once per request. The processors already declare the partitions of these payload keys for item scheduling, so the common partition builder can do the packing. Vision payloads with different patch counts need padding, and patch positions use -1 to mark padding. Padding them with zeros and restoring the -1 afterwards needs the per-frame lengths on the GPU: in a draft, that restore made one pageable host-to-device copy and one stream synchronize per vision call (8 image requests: 2 copies, 1 synchronize and 5 kernel launches, against 1, 0 and 2 on its baseline, upstream main 66ce99b, with the towers stubbed). Solution: - Gemma4 and Gemma4 Unified use three EncoderGroups with the common builder: image, video and audio. The vision encoder normalizes raw uint8 patches with the processor's record after batching and keeps the rule that per-frame lengths are used only when their count matches the frames. The audio encoder takes the batched features and masks. The per-run recursion and the packing are removed; requests with more than one modality are still rejected. - The processors also declare vision_normalization as a shared value. When they build item metadata, they fill the per-item inputs a processor may omit (an all-valid audio mask, square-grid patch positions and their lengths), so payloads of different lengths can be padded into one call. Patch positions declare the processors' -1 padding sentinel as their padding value, so batching itself marks added patches as padding: the vision encoder needs no restore pass, no host-to-device copy and no synchronize for it. Declaring partitions at run time leaves the payload unchanged. - The profiling dummy carries item metadata with output rows and partitions. Gemma4 Unified counts the output rows of its items when a request has no item metadata yet (no multimodal hashing). Behaviour: - Image-only and video-only batches of adjacent requests give the tower the same input as before, so their outputs are unchanged. - Requests of one modality that are not adjacent in the batch share one call per modality instead of one call per run. - The audio clips of all requests go to one padded, masked tower call instead of one call per request. Outputs can differ from before only in floating-point order. - Raw uint8 pixels next to normalized pixels, or different normalization records, are encoded in one call per compatible set instead of one call, and the rows are put back in prompt order. This happens when some requests pass mm_processor_kwargs (host normalization) and others do not while GPU normalization is on. - Vision payloads with different patch counts can share a call (before, batching them failed). Performance: Screening, not a formal benchmark: one GB200 GPU on a shared node, bf16, TP=1, reduced random-weight Gemma4 towers (26B vision widths with 4 layers, default audio tower with 12 layers, TRTLLM attention). Baseline: upstream main bae7da4. Payloads come from the real image, video and audio processors. Each round runs 3 warm-up and then 15 timed batched encodes of all requests, one at a time; the timer wraps the whole encode with a CUDA synchronize before and after, and leaves out building the payloads and moving them to the GPU. Two rounds (main, this change, main, this change); each value is that round's median, as before -> after; difference; change in elapsed time. This is neither a kernel-only nor a serving-throughput number. - 8 audio requests of one 1-30 s clip each (8 audio calls before, one call of [8, 2707, 128] features after): round 1: 192.55 -> 30.88 ms; -161.67 ms; -84.0% round 2: 193.32 -> 30.92 ms; -162.40 ms; -84.0% - 12 requests cycling one image, one 6 s audio clip and one 4-frame video (12 calls before; after, 3 calls: images [4, 2520, 768], video frames [16, 630, 768], audio [4, 599, 128]): round 1: 149.63 -> 41.76 ms; -107.87 ms; -72.1% round 2: 151.54 -> 41.07 ms; -110.47 ms; -72.9% - 8 single-image requests, raw uint8 (GPU normalization) and host normalized alternating (one call of [8, 2520, 768] before, two calls of [4, 2520, 768] after): round 1: 15.16 -> 16.44 ms; +1.29 ms; +8.5% round 2: 15.21 -> 16.48 ms; +1.26 ms; +8.3% - 8 single-image requests (one call of [8, 2520, 768] in both): round 1: 14.85 -> 14.89 ms; +0.05 ms; +0.3% round 2: 14.85 -> 14.91 ms; +0.07 ms; +0.4% With the towers stubbed (method in the first commit), the 8-image control makes 2 kernel launches as main does, no copies (main: 1) and no stream synchronize; Python calls 21 -> 26, host time 0.081 -> 0.098 ms warm (+0.016 ms). With only the tower bodies stubbed (wrapper and embedder live, cold), the per-call work goes away with the calls: 8 audio requests take 4.063 -> 1.332 ms of host time and the 12 mixed requests 6.397 -> 2.368 ms. Validation: - New CPU tests against a copy of the previous encode path: per-request and batched equality and call counts for single and mixed batches, including non-adjacent requests, selected items, the image, video and audio profiling dummies, raw and normalized pixel splits, ragged vision with explicit or implicit positions, ragged audio with and without masks, and the Gemma4 Unified wrappers and output-row counting. Vision inputs equal the previous ones (values, dtype, device); a reduced random-weight audio tower matches per-request outputs within 1e-5 (fp32). A runtime partition lookup leaves the payload unchanged. - Multimodal CPU suite at this commit (-m cpu_only over tests/unittest/_torch/multimodal, the multimodal model test files and tests/unittest/inputs/test_multimodal.py): 1950 passed, 77 skipped, 4 failed; main bae7da4: 1843 / 77 / 4. No test changed status and none was removed; the 107 new tests pass. The 4 failures also fail on main (Nemotron Nano V2 VL input processor, checkpoint missing on the test machine). - GPU, reduced random-weight Gemma4 towers and Gemma4 Unified front end (bf16, TRTLLM attention, TF32 off), payloads from the real processors, 18 cases: every request alone equals the previous output bitwise; batched image, video, raw uint8, split and non-adjacent batches equal the previous batched output bitwise, except Unified non-adjacent image requests, which now share one call and stay within the bf16 floating-point-order bound. Batched audio is within that bound (rel L2 8.2e-3); in fp32 with TF32 off it differs from per-request encoding by rel L2 1.6e-6, so the difference is floating-point order only. - GPU, test_gemma4_multimodal.py with all markers (whole series applied): 170 passed, 20 skipped, 0 failed (main: 146 / 20 / 0); no test changed status; its ragged-vision sentinel tests run on CPU only. - pre-commit hooks pass on the changed files. - Reproduce: pytest -q -m cpu_only tests/unittest/_torch/modeling/test_gemma4_multimodal.py - End-to-end serving and real-checkpoint accuracy: not measured. Exclusive-node formal performance and TP > 1: not measured. GPU validation used reduced random-weight towers. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com>

Problem: Qwen3-VL packs its encoder inputs in a private builder, _qwen3vl_build_batched_input. It concatenates the image rows, then the video rows, of all requests and normalizes raw uint8 pixels per payload with the processor's record; Cosmos3 Edge imports the same builder. The Qwen processor already declares how items occupy these payload keys (packed patch rows, one grid row per item) for item scheduling, so the common partition builder can do the packing. Solution: - The Qwen vision partitions also declare vision_normalization as a shared value, so the record reaches the encoder and payloads with different records are encoded apart. - Qwen3VisionModelBase (Qwen3-VL, Qwen3-VL-MoE, Qwen3.5-VL, Qwen4Exp, Cosmos3) and Cosmos3 Edge use the common builder. Their encode_batched takes the payload keys, normalizes raw uint8 keys with the record and concatenates the image and video rows, image rows first, for one ViT call. The private builder is removed. Behaviour: - Batches whose payloads agree in pixel dtype and normalization record (the normal case) make one ViT call with the same input as before, so the outputs are unchanged. - A batch that mixes raw uint8 and normalized pixels of one modality, or different normalization records, is encoded in one call per compatible set instead of one call, and the rows are put back in prompt order. Its outputs can differ from before only in floating-point order. This happens when some requests pass mm_processor_kwargs (host normalization) and others do not while GPU normalization is on. - A call that holds both images and videos copies the pixel rows once more, to concatenate the two payload keys. Performance: Screening, not a formal benchmark: one GB200 GPU on a shared node, bf16, TP=1, a reduced random-weight Qwen3-VL tower (8 layers of width 1152, patch size 16, TRTLLM attention). Baseline: upstream main bae7da4. 16 image requests of one image each, patch grid (1, 32, 32), 1024 patch rows per image. Each round runs 3 warm-up and then 15 timed batched encodes of all requests, one at a time; the timer wraps the whole encode with a CUDA synchronize before and after, and leaves out building the payloads and moving them to the GPU. Two rounds (main, this change, main, this change); each value is that round's median, as before -> after; difference; change in elapsed time. This is neither a kernel-only nor a serving-throughput number. - All images normalized (one ViT call of 16384 rows in both): round 1: 10.08 -> 10.15 ms; +0.07 ms; +0.6% round 2: 10.14 -> 10.20 ms; +0.06 ms; +0.6% - Raw uint8 and normalized images alternating (one call of 16384 rows before, two calls of 8192 rows after): round 1: 10.41 -> 15.65 ms; +5.24 ms; +50.3% round 2: 10.49 -> 15.85 ms; +5.36 ms; +51.1% The alternating batch pays for the second ViT call. With the tower stubbed (method in the first commit), host time is 0.141 -> 0.166 ms warm (+0.025 ms) for the normalized control and 0.362 -> 0.378 ms warm (+0.016 ms) for the alternating batch, which makes 10 -> 6 kernel launches since raw pixels are normalized once per call instead of once per payload. Validation: - New CPU tests in test_modeling_qwen3vl.py against copies of the removed builder and the previous encode_batched, with item metadata from the Qwen processor: the ViT input and the output are equal for one image, image only, video only, mixed and video-first batches, with raw uint8 and normalized pixels, for selected-item views and for the profiling dummy; batched output equals per-request output; different records split the call; the encoders take exactly the declared keys. In test_modeling_cosmos3_edge.py, Cosmos3 Edge gets the same ViT input as before. The shared Qwen request helpers are defined in test_modeling_qwen3vl.py. - Multimodal CPU suite at this commit (-m cpu_only over tests/unittest/_torch/multimodal, the multimodal model test files and tests/unittest/inputs/test_multimodal.py): 1896 passed, 77 skipped, 4 failed; main bae7da4: 1843 / 77 / 4. No test changed status and none was removed; the 53 new tests pass. The 4 failures also fail on main (Nemotron Nano V2 VL input processor, checkpoint missing on the test machine). - GPU, reduced random-weight Qwen3-VL tower (bf16, TRTLLM attention, TF32 off), 13 cases: per-request and batched rows equal the previous code's bitwise for normalized and raw uint8 pixels, including a batch encoded in two calls. - GPU, test_modeling_qwen3vl.py and test_modeling_cosmos3_edge.py with all markers (whole series applied): 516 passed, 39 skipped, 8 failed (main: 482 / 39 / 8); no test changed status; the 8 failures (test_qwen3_vision_cuda_graph_replay_matches_eager with the FLASHINFER backend) fail the same way on main on the test machine. - pre-commit hooks pass on the changed files. - Reproduce: pytest -q -m cpu_only tests/unittest/_torch/modeling/test_modeling_qwen3vl.py tests/unittest/_torch/modeling/test_modeling_cosmos3_edge.py - End-to-end serving and real-checkpoint accuracy: not measured. Exclusive-node formal performance and TP > 1: not measured. GPU validation used reduced random-weight towers. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com>

Problem: Qwen2-VL, Qwen2.5-VL and Exaone4.5 pack their encoder inputs in a private helper that concatenates the image rows and the video rows of all requests (or passes one request's tensors through). The encoder then calls the ViT once for the images and once for the videos and returns two tensors, image rows first. A batch that holds both images and videos therefore hands its callers two modality-major tensors. The encoder-only runner rejects such a batch. The aggregated path rejects it unless it has exactly two requests, and then pairs the tensors with the requests by position, so a video request before an image request gets the image rows. Solution: The Qwen processor already declares the partitions of these payload keys, so Qwen2VisionModelBase now uses one image and video EncoderGroup with the common partition builder, as Qwen3-VL does. Its encode_batched concatenates the image and video rows, image rows first, and runs the ViT once; the ViT tells items apart only by their grid rows. forward returns the shared encode result. Exaone4.5 inherits the encoder. The private helper is removed. Behaviour: - Image-only and video-only batches give the ViT the same input as before, so their outputs are unchanged. A single request's pixels are copied once, as the other batch sizes already were. - Batches with both images and videos make one ViT call instead of two, so their outputs can differ only in floating-point order. They now reach the language model and the encoder-only runner as one tensor in request order, then prompt order, so these batches are served correctly. Performance: Screening, not a formal benchmark: one GB200 GPU on a shared node, bf16, TP=1, a reduced random-weight Qwen2.5-VL tower (8 layers of width 1280, patch size 14, TRTLLM attention). Baseline: upstream main bae7da4. 16 requests of one item each; an image has patch grid (1, 32, 32) (1024 patch rows), a video (4, 24, 24) (2304 patch rows). Each round runs 3 warm-up and then 15 timed batched encodes of all requests, one at a time; the timer wraps the whole encode with a CUDA synchronize before and after, and leaves out building the payloads and moving them to the GPU. Two rounds (main, this change, main, this change); each value is that round's median, as before -> after; difference; change in elapsed time. This is neither a kernel-only nor a serving-throughput number. - 8 image and 8 video requests alternating (two ViT calls of 8192 and 18432 rows before, one call of 26624 rows after): round 1: 18.79 -> 15.85 ms; -2.94 ms; -15.6% round 2: 18.84 -> 15.87 ms; -2.96 ms; -15.7% - 16 image requests (one call of 16384 rows in both): round 1: 10.80 -> 10.84 ms; +0.04 ms; +0.4% round 2: 10.80 -> 10.87 ms; +0.07 ms; +0.7% With the tower stubbed (method in the first commit), host time is 0.072 -> 0.125 ms warm (+0.053 ms) for the image control and 0.083 -> 0.177 ms warm (+0.093 ms) for the mixed batch: main returned the two modality tensors as they were, while the common path checks the item layouts and puts the rows in request order. Validation: - New CPU tests in test_modeling_qwen3vl.py, next to the Qwen3-VL ones and with the same requests, against copies of the removed helper and the previous forward, with item metadata from the Qwen processor: one ViT call whose input equals the previous image input followed by the previous video input, for one image, image only, video only, mixed and video-first batches and selected-item views; one output tensor equal to per-item encoding in prompt order, exactly with a stub ViT and within 1e-5 with a tiny random-weight Qwen2-VL ViT, and equal to the previous output for single-modality batches; the profiling dummy encodes through the group; the Qwen2-VL and Exaone4.5 encoders take exactly the declared keys. - Multimodal CPU suite at this commit (-m cpu_only over tests/unittest/_torch/multimodal, the multimodal model test files and tests/unittest/inputs/test_multimodal.py): 1911 passed, 77 skipped, 4 failed; main bae7da4: 1843 / 77 / 4. No test changed status and none was removed; the 68 new tests pass. The 4 failures also fail on main (Nemotron Nano V2 VL input processor, checkpoint missing on the test machine). - GPU, reduced random-weight Qwen2.5-VL tower (bf16, TRTLLM attention, TF32 off), 6 cases: per-request and batched rows equal the previous code's bitwise, also with images and videos in one call. - GPU, test_modeling_qwen3vl.py and test_modeling_qwen2vl.py with all markers (whole series applied): 515 passed, 36 skipped, 8 failed (main: 483 / 36 / 8); no test changed status; the 8 failures (test_qwen3_vision_cuda_graph_replay_matches_eager with the FLASHINFER backend) fail the same way on main on the test machine. - pre-commit hooks pass on the changed files. - Reproduce: pytest -q -m cpu_only tests/unittest/_torch/modeling/test_modeling_qwen3vl.py -k qwen2_vision - End-to-end serving and real-checkpoint accuracy: not measured. Exclusive-node formal performance and TP > 1: not measured. GPU validation used reduced random-weight towers. Exaone4.5 was not run with its own vision weights; it shares the Qwen2.5-VL encoder class. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com>

…ions Problem: Mistral3 and Pixtral pack their encoder inputs in _vision_forward: it pads every request's pixel_values to the largest image_sizes of the batch with batch_pixel_values and runs the Pixtral tower once. The processors declare no input partitions, so item scheduling and the encoder cache cut selected images with the legacy selector, which also crops the trailing height and width. Solution: - The HF and native processors (MistralPixtralEncoderContract) declare pixel_values and image_sizes with one row per image. The profiling dummy carries the same item metadata (output rows and partitions). - Mistral3VLM uses one image EncoderGroup with the common partition builder. Its encoder takes the batched pixel_values and image_sizes, crops the pixels to the largest image size of the call, as batch_pixel_values did, and runs the unchanged tower and projector. _vision_forward is removed; batch_pixel_values stays for the native processor. - image_sizes stay on the host (multimodal_data_device_paths), so the crop reads them without a device sync. - A comment in the legacy item selector no longer names Mistral, which now declares its partitions. Behaviour: - A batch gives the tower the same input as before in one call, so the outputs are unchanged, also for selected-item views. - Payloads whose pixel dtype differs are encoded in separate calls, and their rows are put back in request order. Performance: Not measured: the tower call and its input are unchanged. Validation: - New CPU tests against a copy of the previous packer, with stub tower and projector: the tower input equals the previous one (values, dtype, device) for ragged image sizes given as lists or tensors, for selected and reordered items and for the HF and native profiling dummies; batched output equals per-request output; different dtypes split the call. The profiling-dummy contract test now drives the encoder group, which also checks the output rows against the item metadata. - Multimodal CPU suite at this commit (-m cpu_only over tests/unittest/_torch/multimodal, the multimodal model test files and tests/unittest/inputs/test_multimodal.py): 1926 passed, 77 skipped, 4 failed; main bae7da4: 1843 / 77 / 4. No test changed status and none was removed; the 83 new tests pass. The 4 failures also fail on main (Nemotron Nano V2 VL input processor, checkpoint missing on the test machine). - GPU, reduced random-weight Pixtral tower and projector (bf16, TRTLLM attention, TF32 off), 8 cases: per-request and batched rows equal the previous code's bitwise, with image sizes as lists and as tensors. - GPU, test_modeling_mistral.py with all markers (whole series applied): 52 passed, 0 skipped, 0 failed (main: 37 / 0 / 0); no test changed status. - pre-commit hooks pass on the changed files. - Reproduce: pytest -q -m cpu_only tests/unittest/_torch/modeling/test_modeling_mistral.py; the dummy contract test runs without the marker: tests/unittest/_torch/modeling/test_modeling_mistral.py::test_dummy_mm_data_satisfies_the_encoder_input_contract - End-to-end serving and real-checkpoint accuracy: not measured. Exclusive-node formal performance and TP > 1: not measured. GPU validation used reduced random-weight towers. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com>

hanjingtian · 2026-11-12T04:32:31Z

Could you please rebase this branch onto the latest upstream main to include the CI hang fix in #22101 (6edbe383497a)? The current branch still contains the regression that can cause pipelines to hang for an extended period and worsen CI queueing

github-actions Bot assigned yechank-nvidia Oct 27, 2026

yechank-nvidia changed the title ~~[TRTLLM-13419][refactor] Share encoder input packing across requests~~ [TRTLLM-13419][refactor] Share encoder input batching across requests Oct 30, 2026

yechank-nvidia force-pushed the yechank/mm-common-encoder-batching branch 3 times, most recently from eb64b08 to 7832960 Compare November 5, 2026 04:28

yechank-nvidia changed the title ~~[TRTLLM-13419][refactor] Share encoder input batching across requests~~ [TRTLLM-13419][refactor] Move multimodal encoder input batching into EncoderGroup for all models Nov 5, 2026

yechank-nvidia added 5 commits November 5, 2026 06:15

yechank-nvidia force-pushed the yechank/mm-common-encoder-batching branch from 7832960 to 4c70ba0 Compare November 11, 2026 05:09

yechank-nvidia changed the title ~~[TRTLLM-13419][refactor] Move multimodal encoder input batching into EncoderGroup for all models~~ [TRTLLM-13419][refactor] Batch multimodal encoder inputs by item partitions Nov 11, 2026

yechank-nvidia force-pushed the yechank/mm-common-encoder-batching branch from 4c70ba0 to 2ea1188 Compare November 11, 2026 06:06

yechank-nvidia mentioned this pull request Nov 12, 2026

[None][refactor] define a shared multimodal model protocol #20906

Draft

1 task done

Navigation Menu

Repository navigation

[TRTLLM-13419][refactor] Batch multimodal encoder inputs by item partitions - #21052

[TRTLLM-13419][refactor] Batch multimodal encoder inputs by item partitions#21052
yechank-nvidia wants to merge 5 commits into
NVIDIA:mainfrom
yechank-nvidia:yechank/mm-common-encoder-batching

yechank-nvidia commented Oct 27, 2026 •

edited

Loading

Uh oh!

hanjingtian commented Nov 12, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

yechank-nvidia commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Performance

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

hanjingtian commented Nov 12, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yechank-nvidia commented Oct 27, 2026 •

edited

Loading