Repository navigation
[TRTLLM-13419][refactor] Share encoder input packing across requests - #21052
Draft
yechank-nvidia wants to merge 3 commits into
Draft
yechank-nvidia wants to merge 3 commits into
yechank-nvidia wants to merge 3 commits into
Conversation
Problem Qwen3-VL/Cosmos and Gemma4 repeat field gathering and concatenation around the existing EncoderGroup/item metadata contract. Nemotron separately repeats feature/mask padding. This duplicates host preparation policy and makes request/item ordering and output ownership harder to review. Synthetic CPU preparation measurements identified an initial common-helper regression: repeated alias lookups and recursive Gemma run preparation added host work. This is measured preparation cost, not evidence of a GPU kernel or end-to-end serving bottleneck. The rejected Nemotron audio extraction increased CPU cost by about 11.5-12.2%; even a compact prototype remained about 5.3-6.5% slower. Solution Add frozen EncoderInputField declarations and batch_encoder_inputs beside the existing grouped encoder mechanism. Source aliases, concat axis, allowed right-padding axes and fill value are explicit. Reuse existing item selection, output lengths, prompt/request restore and cache state. Required fields fail instead of dropping masks; inputs remain read-only and packed outputs own storage. Migrate Qwen3-VL/Cosmos packing and Gemma4 image/video packing, including inherited Unified inputs. Use static single-source declarations for uniform Qwen batches and direct aliases for mixed batches. Build Gemma consecutive runs once, extend CPU lengths, and copy each mixed-run output into contiguous owned storage before another graph replay. Keep encoder call composition, DeepStack folding, mRoPE, optional positions/lengths and per-input audio calls unchanged. Defer Nemotron: its production source and tests remain byte-identical to the baseline. Temporal/AV/EVS and partial-cache behavior stay intact. Generic declared padding and boolean masks are CPU-tested for future consumers. No public API, schema, transport, user-facing configuration or environment switch changes. Performance Baseline is main 9ffa6f4. All numbers are shared-system screening on NVIDIA GB200 with an aarch64 CPU, TP1, Python 3.12.3, Torch 2.14.0a0+b2c75dd062.nv26.09 and CUDA 13.4; they are not exclusive-node/formal results. CPU: BF16, one Torch thread, four source items/request; select [0,1] for contiguous/alternating and [3,1] for reordered. Qwen has 16/32 packed patches/item of width 1536 and CPU grids; Gemma has 36 patches/item of width 768 plus positions/CPU lengths. Actual item selection is included only in selection+packing. Encoder math is replaced by a recording feature stub. Use perf_counter_ns, disabled GC, five warmups and 13 interleaved rounds of 100 calls/path. Report median of round means. Delta is candidate minus baseline; negative improvement means slower. Qwen3-VL, 64 requests, contiguous, CPU packing: baseline 0.31960447999999997 ms -> candidate 0.32304415999999997 ms; delta 0.0034396800000000005 ms; improvement -1.0762302205526033%. Qwen3-VL, 64 requests, contiguous, CPU selection+packing: baseline 0.60707071 ms -> candidate 0.60828287 ms; delta 0.0012121600000000177 ms; improvement -0.19967360968543807%. Qwen3-VL, 64 requests, reordered, CPU packing: baseline 0.32040031 ms -> candidate 0.32219423999999997 ms; delta 0.0017939299999999436 ms; improvement -0.5599027042139701%. Qwen3-VL, 64 requests, reordered, CPU selection+packing: baseline 1.1874258999999998 ms -> candidate 1.18874398 ms; delta 0.0013180800000001103 ms; improvement -0.11100313712208151%. Qwen3-VL, 64 requests, alternating, CPU packing: baseline 0.47950527000000004 ms -> candidate 0.48323391 ms; delta 0.0037286399999999498 ms; improvement -0.7776014641090284%. Qwen3-VL, 64 requests, alternating, CPU selection+packing: baseline 0.7713702299999999 ms -> candidate 0.7729654300000001 ms; delta 0.0015952000000001298 ms; improvement -0.20680082507204484%. Qwen3-VL, 128 requests, contiguous, CPU packing: baseline 0.64554911 ms -> candidate 0.64660799 ms; delta 0.0010588799999999843 ms; improvement -0.1640277995271319%. Qwen3-VL, 128 requests, contiguous, CPU selection+packing: baseline 1.21258622 ms -> candidate 1.21446782 ms; delta 0.0018816000000001498 ms; improvement -0.15517247095222225%. Qwen3-VL, 128 requests, reordered, CPU packing: baseline 0.65870879 ms -> candidate 0.66049343 ms; delta 0.001784640000000004 ms; improvement -0.27093004178675134%. Qwen3-VL, 128 requests, reordered, CPU selection+packing: baseline 2.4085113199999997 ms -> candidate 2.43397819 ms; delta 0.025466870000000252 ms; improvement -1.057369744892885%. Qwen3-VL, 128 requests, alternating, CPU packing: baseline 0.9452195 ms -> candidate 0.94485662 ms; delta -0.00036287999999995435 ms; improvement 0.038391082706181405%. Qwen3-VL, 128 requests, alternating, CPU selection+packing: baseline 1.57062621 ms -> candidate 1.57470493 ms; delta 0.004078720000000091 ms; improvement -0.25968750387783807%. Gemma4, 64 requests, contiguous, CPU packing: baseline 0.41641599 ms -> candidate 0.42101151000000003 ms; delta 0.00459552000000002 ms; improvement -1.1035887454754125%. Gemma4, 64 requests, contiguous, CPU selection+packing: baseline 0.7349667 ms -> candidate 0.73521823 ms; delta 0.0002515299999999998 ms; improvement -0.034223319233374766%. Gemma4, 64 requests, reordered, CPU packing: baseline 0.41737759 ms -> candidate 0.42032 ms; delta 0.0029424100000000064 ms; improvement -0.7049755594209086%. Gemma4, 64 requests, reordered, CPU selection+packing: baseline 2.42623611 ms -> candidate 2.43761788 ms; delta 0.011381769999999847 ms; improvement -0.4691122167825557%. Gemma4, 64 requests, alternating, CPU packing: baseline 0.66435486 ms -> candidate 0.64009663 ms; delta -0.024258229999999936 ms; improvement 3.651396484101875%. Gemma4, 64 requests, alternating, CPU selection+packing: baseline 1.00810078 ms -> candidate 0.99134782 ms; delta -0.016752959999999928 ms; improvement 1.6618338495879281%. Gemma4, 128 requests, contiguous, CPU packing: baseline 0.83675262 ms -> candidate 0.83803742 ms; delta 0.0012848000000000859 ms; improvement -0.1535459787386248%. Gemma4, 128 requests, contiguous, CPU selection+packing: baseline 1.47961085 ms -> candidate 1.47587262 ms; delta -0.0037382299999999535 ms; improvement 0.25264953957318936%. Gemma4, 128 requests, reordered, CPU packing: baseline 0.84495231 ms -> candidate 0.84477278 ms; delta -0.00017953000000003883 ms; improvement 0.02124735300150121%. Gemma4, 128 requests, reordered, CPU selection+packing: baseline 4.842278309999999 ms -> candidate 4.89149943 ms; delta 0.04922112000000034 ms; improvement -1.0164868032957064%. Gemma4, 128 requests, alternating, CPU packing: baseline 1.35763518 ms -> candidate 1.30376542 ms; delta -0.05386975999999999 ms; improvement 3.96791132062444%. Gemma4, 128 requests, alternating, CPU selection+packing: baseline 2.04518524 ms -> candidate 1.97633916 ms; delta -0.06884607999999992 ms; improvement 3.3662515577317547%. GPU: BF16 reduced real towers, seed 13419, fixed identical weights across paths. Qwen vision hidden128/depth2/heads2/intermediate256/patch16, temporal patch2, spatial merge2, DeepStack layers [0,1]; output width384. Gemma vision hidden128/layers2/heads4/head_dim32/intermediate256/patch16/pooling3, 36 patches and four output rows/item. CUDA-event encoder-forward timing excludes preparation: five warmups, nine interleaved rounds of 20 calls/path, median of round means. Qwen3-VL, 1 requests, contiguous, encoder forward: baseline 1.760223960876465 ms -> candidate 1.7692352294921876 ms; delta 0.009011268615722656 ms; improvement -0.5119387541591975%. Qwen3-VL, 4 requests, contiguous, encoder forward: baseline 1.8926143646240234 ms -> candidate 1.8980384826660157 ms; delta 0.005424118041992276 ms; improvement -0.2865939381723864%. Qwen3-VL, 4 requests, alternating, encoder forward: baseline 1.8831136703491211 ms -> candidate 1.8758783340454102 ms; delta -0.007235336303710982 ms; improvement 0.3842219626800108%. Qwen3-VL, 8 requests, alternating reversed, encoder forward: baseline 1.9655263900756836 ms -> candidate 1.9682687759399413 ms; delta 0.0027423858642576793 ms; improvement -0.13952424541865766%. Gemma4 vision, 1 requests, contiguous, encoder forward: baseline 3.140609550476074 ms -> candidate 3.1271728515625 ms; delta -0.01343669891357413 ms; improvement 0.4278372939271393%. Gemma4 vision, 4 requests, contiguous, encoder forward: baseline 3.148052787780762 ms -> candidate 3.158030319213867 ms; delta 0.009977531433105202 ms; improvement -0.3169429519045302%. Gemma4 vision, 4 requests, alternating, encoder forward: baseline 12.463820648193359 ms -> candidate 12.474774169921876 ms; delta 0.010953521728517046 ms; improvement -0.08788253648454712%. Gemma4 vision, 8 requests, alternating reversed, encoder forward: baseline 24.823648071289064 ms -> candidate 24.86558380126953 ms; delta 0.04193572998046591 ms; improvement -0.16893459760641955%. The encoder implementation and prepacked tensors are unchanged. GPU median variation overlaps round dispersion. Gemma alternating CPU preparation improves by 1.6618338495879281% and 3.3662515577317547%. Small slower CPU medians remain and are not relabeled as improvements. The strict CPU zero/no-regression bound is unproven. A fixed A/A identity control also gives apparent differences: identical Qwen128 alternating callables show +0.5264131402440978%, including a naive bootstrap interval excluding zero. Gemma128 reordered candidate and A/A paired-delta intervals overlap. Shared-system drift makes small residual CPU effects inconclusive; there is no guaranteed zero-regression claim. End-to-end serving latency/throughput, formal performance and peak memory: unmeasured (미측정). Module/CPU preparation results are not serving speedups. Validation Final validation counts: CPU1772 passed/4 existing failures/8 existing errors/77 skipped; extras314 passed/81 skipped; GPU182 passed/1 skipped; new failures0. Existing missing-asset and NIXL/render-fixture failures are reproduced by the baseline; corrected render-fixture coverage passes. New CPU cases cover aliases, padding/masks, ownership, mixed/reordered items, DeepStack, actual partial LRU hits, Gemma optional fields and graph-buffer reuse. All 20 cases are selected by existing CPU CI parent coverage; no test database edit is required. All eight final reduced BF16 GPU compositions are bit-identical against main and serial execution, max absolute/relative difference 0.0, including DeepStack width384. No non-bit-identical case was observed. Source/harness pins match CPU and GPU start/end snapshots. The same item-1 suites are repeated before/after; overlapping phase totals are not unique test counts. Validation uses new Python sources over native artifact 032eb45, not a fresh native build. Requested trained Qwen3-VL-2B/Nemotron12B checkpoints were absent locally, so conditional real-weight E/P reruns were not executed. Trained accuracy, transport/MPI/topology and untested shapes remain unqualified. Repository-pinned pre-commit checks pass. The type-check hook has no applicable configured roots and is skipped. No Qwen2-VL test, C++ transport, serving schema or public LLM API changes. Strict CPU no-regression qualification remains an explicit follow-up before that claim can be made. Reproduce changed CPU cases in a matching installed build: PYTHONPATH=tests/unittest:tests pytest --noconftest -c /dev/null \ -p test_common.magic_import_hooks -m cpu_only \ tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py \ tests/unittest/_torch/modeling/test_modeling_qwen3vl.py \ tests/unittest/_torch/modeling/test_gemma4_multimodal.py Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com>
Problem The common Gemma4 input path still constructs groupby(zip(...)) after it has already established that every request has one modality and that all modalities match. cProfile records 64 groupby key-lambda calls per 64-request preparation (8192 calls across 128 preparations). Line tracing locates the uniform payload gather and its key calls in encode_multimodal_inputs. The unchanged tensor concatenation count is three per preparation; no extra tensor copy was found. Instrumented times identify dispatch work, not an uninstrumented speedup or a GPU bottleneck. Solution Use the existing mixed_modality_runs result to reserve groupby for mixed runs. Build the one uniform input group directly from request payloads. Keep the existing inactive/empty-input fallback, required fields, audio calls, optional positions and lengths, alias handling, consecutive mixed-run order, graph-buffer snapshots and final owned concatenation. Reuse the current common packing helper; add no configuration, public API, tensor copy or new abstraction. The production diff is six added lines and one changed line. Performance Compare the previous common-batching head cceab0d with the direct-group variant. These are CPU preparation screening results on a shared NVIDIA GB200 system with an aarch64 CPU, not formal or end-to-end results. Python 3.12.3, Torch 2.14.0a0+b2c75dd062.nv26.09, BF16, TP1, one Torch/OMP/OpenBLAS thread, fixed CPU affinity. Each request has four source items of shape [36,768]; select [0,1] for contiguous and [3,1] for reordered. The encoder is replaced by the same recording stub. Use 16 in-process screening rounds, 128 discarded warmups per label, eight calls per block, randomized five-label order followed by its reverse. Report geometric means of per-round arithmetic means and paired log-ratio descriptive t95 intervals. GC stays enabled at executor defaults (20000,10,10); no timed collection occurred. Each pair uses exactly the same inputs. These screening rounds do not count toward the pre-registered qualification gate. gemma4/128/reordered/packing_only: baseline 0.8251086213312362 ms -> candidate 0.8135045522430572 ms; absolute change -0.011604069088179014 ms; candidate/baseline change -1.4063686632502044%; paired ratio 95% CI [0.9782534554037668, 0.9936795097907201]. gemma4/64/contiguous/selection_and_packing: baseline 0.719085895637334 ms -> candidate 0.7144605346095466 ms; absolute change -0.0046253610277874335 ms; candidate/baseline change -0.6432278891644707%; paired ratio 95% CI [0.9915208547207646, 0.995618812986503]. gemma4/64/contiguous/packing_only: baseline 0.4091835288816613 ms -> candidate 0.40469958631734565 ms; absolute change -0.004483942564315668 ms; candidate/baseline change -1.0958267495690197%; paired ratio 95% CI [0.9845894969257152, 0.993514100738909]. gemma4/128/reordered/selection_and_packing: baseline 4.784406223026791 ms -> candidate 4.813567107043298 ms; absolute change 0.029160884016506294 ms; candidate/baseline change 0.6094984969326768%; paired ratio 95% CI [0.9954199054660114, 1.0168845461318659]. Packing improves in both screened uniform cases; the 128-request reordered selection+packing point estimate is slower with an interval that includes no change, so it is inconclusive. Do not interpret a packing-only improvement as overall serving improvement. The full 24-case CPU qualification is unmeasured at commit time and will use 344 fresh independent round processes per case, compared with main 9ffa6f4, after the excluded pilot. Its fixed upper candidate/baseline t95 bound is 1.005 with both A/B and A/A CI radii below 0.005. GPU/module speedup, formal performance, trained-model accuracy, peak memory and end-to-end latency/throughput improvements are unmeasured. Validation Repeat the original 59-file marked CPU scan and all five compatible extra phases on this source: 1772 passed, four existing missing-asset failures, eight existing environment/harness errors and 77 skips; extras 314 passed and 81 skips. All 20 original added cases pass. Exact testcase names/statuses match the prior candidate, with no new failure, removed test or timeout. Counts overlap and are not unique-test totals. Repeat all six original GPU suites on one locked GB200: 182 passed and one existing skip. All eight reduced BF16 compositions are bitwise equal between baseline/candidate and batched/serial, maximum absolute and relative differences 0.0; Qwen DeepStack width remains 384. The GPU lease is released. Native artifact 032eb45 is reused with new Python sources; this is an overlay check, not a fresh native build. Requested trained checkpoints remain unavailable, so real-weight handoff and full serving are unqualified. Repository-pinned pre-commit checks pass. The type-check hook has no applicable configured roots and is skipped. Qwen2-VL tests and the test database remain untouched. Reproduce focused CPU coverage in a matching installed build: PYTHONPATH=tests/unittest:tests pytest --noconftest -c /dev/null \ -p test_common.magic_import_hooks -m cpu_only \ tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py \ tests/unittest/_torch/modeling/test_modeling_qwen3vl.py \ tests/unittest/_torch/modeling/test_gemma4_multimodal.py Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com>
Problem The common batching change added Python work to Qwen3-VL encoder input preparation. Uniform requests built five temporary lists instead of two, and mixed requests doubled dictionary lookups through alias fallback. A pre-registered CPU gate against the base failed all six Qwen packing cases and both alternating selection+packing cases. Solution Gather the canonical pixel and grid columns together, in image-then-video and request order, and reuse the existing field declarations through EncoderInputField.batch. Keep the previous field-major path as a cold fallback for missing, invalid or aliased payloads, so empty inputs and error order are unchanged. For the default field (dim 0, no padding) call torch.cat directly. Other axes and padding use a metadata-bound partial of one private batching function, which keeps declared-axis padding and owned output storage for a single input. The derived callable is excluded from dataclass init, repr, comparison and hash. Performance CPU-thread time of encoder input preparation (selection + packing), Qwen3-VL and Gemma4, 64/128 requests, contiguous / reordered / alternating image-video. Pre-registered gate: 344 fresh processes, randomized ABBA/BAAB order, A/A control, pass if the upper 95% bound of candidate/base is <= 1.005. Result on this head: 24/24 PASS, worst upper bound 1.0018; largest A/A half-width 0.24%. Gemma4 alternating requests are cheaper than base. Screening on a shared node; serving throughput and latency were not measured. Validation - New CPU test: preferred and explicit-None aliases, image-before-video order, separate output storage (2 cases). - CPU: 1774 passed; the 4 failures and 8 errors are the same on base (missing assets, NIXL, render fixture). - GPU: 182 passed, 1 skipped. Eight reduced BF16 encoder setups (Qwen3-VL image, video, alternating, reordered, DeepStack; Gemma4 image, video, alternating) are bitwise equal to base and to serial execution. Signed-off-by: Yechan Kim <161688079+yechank-nvidia@users.noreply.github.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@coderabbitai summary
Description
Why: When the multimodal encoder runs several requests together, each model gathers its input fields, concatenates them and pads them in its own code. Qwen3-VL and Gemma4 each have their own copy of this logic. Copies drift apart, so item order, padding and output storage are hard to keep correct across models, and every new model has to write it again.
What: This PR adds one small shared helper for packing encoder inputs across requests, and moves Qwen3-VL and Gemma4 to it.
EncoderInputField(frozen dataclass) declares one input field: the source keys (with aliases), the concatenation axis, and optional padding axes with a fill value. It has no model or modality names.batch_encoder_inputsgathers the declared fields from the selected items in order, pads only the declared axes, and concatenates. Missing required fields raise an error. A single input still gets its own output storage.EncoderGroupdispatch, item selection, cache keys, prompt-order restore and E/P handoff from [TRTLLM-14477][refactor] unify multimodal request data and item metadata #20391 are unchanged. No public API, config, environment variable or serving schema changes.Nemotron audio is not moved in this PR. Moving it made input preparation about 12% slower, so it stays as is and is left for a follow-up. Qwen2/2.5-VL are unchanged.
Performance
CPU time to prepare encoder inputs on the host (one thread, GB200 host CPU, BF16 inputs, 4 items per request, 2 selected). Each case was measured over 344 separate processes in random order, with a same-code A/A control. The pass rule, fixed before measuring, was: the upper 95% bound of new/old time must be at most +0.5%. All 24 cases pass; the worst upper bound is +0.18%.
So the shared helper costs nothing measurable, and Gemma4 requests that mix images and videos get 5-8% cheaper to prepare. GPU encoder outputs are bit-identical to main (see Test Coverage). Serving throughput was not measured.
Test Coverage
tests/unittest/_torch/modeling/test_multimodal_encoder_item_groups.py: the helper itself: declared padding axes and fill values, boolean masks, alias lookup, missing-field errors, single-input storage, reordered and mixed items.tests/unittest/_torch/modeling/test_modeling_qwen3vl.py: Qwen3-VL packing order (image before video, request order), aliases, DeepStack order, partial encoder-cache hits.tests/unittest/_torch/modeling/test_gemma4_multimodal.py: Gemma4 optional fields and CUDA-graph buffer snapshots.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.