[None][feat] Enforce multimodal encoder runtime budgets with budgeted output storage by yechank-nvidia · Pull Request #16051 · NVIDIA/TensorRT-LLM

yechank-nvidia · 2026-07-07T08:11:28Z

Description

Multimodal encoder work is unbounded at runtime today: the scheduler admits on
LLM/KV capacity alone, so one iteration can submit more encoder items — or more
encoder attention tokens — than the workspace was sized for, and the embeddings
held between encode and prefill accumulate with no cap. KV-cache sizing at
startup does not know any of it exists.

This PR makes encoder execution a scheduled, budgeted resource:

Scheduling. encoder_max_num_items × encoder_max_num_tokens are
enforced per iteration over atomic items (one image or video), selected FCFS;
a request left partial resumes later. The executor's encoder step is the
single encode site.
Memory. A byte budget bounds the embeddings held between encode and
prefill. A request's first scheduled item allocates one contiguous buffer for
all of its items, so a started request can always finish; KV estimation
reserves that budget at startup. A request that could never fit is rejected at
admission with a message naming the knob, instead of becoming a CUDA OOM.
multimodal_config.encoder_scheduling_policy selects DEFAULT, EAGER
(advance encoder work for capacity-rejected requests), or DISABLED (legacy
inline encode).

Models: Qwen2-VL / Qwen2.5-VL, Qwen3-VL (deepstack-widened rows), Mistral3 /
Pixtral.

Measurements

Qwen3-VL-8B-Instruct, 1×H200, aiperf, n=8 per policy, policy the only
difference and kv_cache_config.max_tokens pinned equal. Caches off.

Availability — past the headroom, removing the cap does not degrade: it
loses essentially every request and the server process dies.

KV pinned	`DISABLED`	`DEFAULT`
101 GiB	128/128	128/128
105 GiB	one instance served, an independent one died	128/128
108 GiB	0/128, process dead	128/128, no OOM

Reproduced on two hosts. At 105 GiB DISABLED is non-deterministic.

Peak memory — DEFAULT has zero variance across repetitions and barely
moves as the workload gets 7× heavier (136.1 / 135.9 / 135.1 GiB at 1 / 4 / 7
images per request); DISABLED tracks the traffic (131.7 → 138.0 → 138.1).

Latency — a trade, and only under heavy load:

Workload	TTFT avg	ITL avg
4 images, conc 32	+2.0%	−0.5%
7 images, conc 32	−12.3%	+12.2%

At 7 images both are real (non-overlapping ranges). DEFAULT raises the fast
ITL percentiles and lowers the slow ones (p50 +17%, p99 −2.0%) — spreading
encoder work across iterations. Quote both halves or neither.

Test coverage

test_multimodal_scheduler.py (atomic packing, byte-budget
allocate-before-compute, whole-request charging, admission fail-fast, per-item
cache read-through, contiguous-buffer ownership, per-rank resolution under PP),
test_kv_cache_estimation.py (encoder profiled at its own budget; reservation
of unmaterialized capacity), test_modeling_qwen2_5vl.py /
test_modeling_mistral.py (capacity from processor geometry, window counts
matching the encoder's padding, dummy tensors satisfying the encoder contract),
test_scheduler_serializable_output.py (item schedule survives rank
distribution).

Follow-ups

Unify the remaining full-request consumers (side-stream prefetch,
mm_encoder_only / disagg, non-item models) onto the item path and
single-source the item manifest; TODO markers are anchored at the migration
sites.

Dev Engineer Review

Adds FCFS multimodal encoder item scheduling with item, token, and resident-output byte budgets.
Adds partial request progress, contiguous embedding storage, cache reuse, admission rejection, and centralized executor encoding.
Adds DEFAULT, EAGER, and DISABLED scheduling policies.
Updates Qwen2-VL, Qwen2.5-VL, Qwen3-VL, Mistral3, Pixtral, and shared encoder capacity handling.
Replaces encoder_max_batch_size with encoder_max_num_items across APIs, configuration, telemetry, manifests, and documentation.
Adds typed item metadata, cache-key helpers, raw-input cleanup, and public exports.
Review focus: verify setup_attn_metadata callers, budget calculations, admission behavior, cache reuse, and partial-progress handling.
No test-list files were changed.

QA Engineer Review

Added and updated tests for scheduler selection, FCFS ordering, eager scheduling, admission, token limits, byte limits, request progress, storage, output validation, cache reuse, serialization, metadata handling, input cleanup, KV-cache estimation, and encoder capacity calculations.
Added model-specific coverage for Qwen2-VL, Qwen2.5-VL, Qwen3-VL, Mistral3, Pixtral, Gemma4, RADIO, and SigLIP.
Replaced obsolete get_dummy_mm_data_for_tokens tests with coverage for get_dummy_mm_data.
No files under tests/integration/test_lists/ were modified.
The changed tests are not registered in test-db/ or qa/ based on the provided changes.
Verdict: needs follow-up.

In the PP executor loop only the scheduling rank runs _schedule() and therefore _attach_mm_encoder_cache_hits(); follower ranks replayed the scheduler for its state effects but never the attach, so items the leader resolved from its cache stayed permanently pending in follower item slots (never encoded either, being absent from the broadcast schedule). Replay the attach before the followers' local scheduler run: encoder execution of broadcast items is already replayed on every rank, so each rank-local cache receives the identical adopt sequence and the attach makes identical hit decisions, keeping slots, holds, and LRU recency in lockstep across ranks. Reported by chienchunhung on PR NVIDIA#16051. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Replace the MultimodalEncoderCacheManager (owning store, holds/eviction/PP replay) with a stateless design where each request owns its encoder outputs and the scheduler derives the byte budget from live request states each tick. - Delete MultimodalEncoderCacheManager and its resource-manager registration. MultimodalEncoderRequestState.record() clones each item output; the scheduler sums resident bytes over live states (clearing state at strip IS the release, no hold/evict/replay bookkeeping). - Declare item-scheduling capability on the model (MultimodalModelMixin.supports_mm_encoder_item_scheduling); a processor participates by overriding get_mm_encoder_item_metadata (no processor flag). - Bound resident encoder outputs by an internal floor (max_num_tokens x embedding-row bytes), mirroring vLLM's default; no new user-facing arg. Over-budget requests are rejected explicitly at admission. - Read-through reuse against the model's existing TensorLRUCache for cache-enabled models only (supports_encoder_cache); hits record a clone and skip the encode, misses populate the cache. Qwen stays cache-free. - KV estimation reserves the output budget plus the cache for cache-enabled item models. Restore encoder_cache_max_bytes to its cache-only meaning. TensorLRUCache is left untouched. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Replace the boolean enable_eager_encoder_scheduling with a three-valued MultimodalEncoderSchedulingPolicy (DISABLED/DEFAULT/EAGER), mirroring the capacity_scheduler_policy pattern, so users can opt out of item-level MM encoder scheduling without disabling multimodal serving. Keep supports_mm_encoder_item_scheduling a pure model capability (always True when the model can do item scheduling) and introduce a separate mm_encoder_item_scheduling_enabled = capability AND policy != DISABLED that gates the actual wiring (setup, scheduler wrap, executor encoder step). A DISABLED policy therefore keeps the capability but runs only the base LLM scheduler with legacy inline encode. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

The encoder output byte budget accounted one copy of a request's embeddings, but a large request could hold up to four at the prefill peak: the per-item tensors, the per-request concatenation, the batch concatenation, and a further concatenation inside fuse_input_embeds. The unaccounted copies scale with the request, so it could exceed the reservation made during KV-capacity estimation. Transfer ownership at publish: finalize_into() now moves the item tensors into multimodal_data and clears its slots. Publishing itself never copied - the list holds the same tensors - but the prefill path replaces that list with its contiguous form, and a slot still referencing the per-item tensors kept them alive past that point. A finalized flag keeps readiness and byte accounting intact, so the budget still charges the full footprint until the request is stripped; without it, emptied slots would read as PENDING and re-open the request for encoding. Replace the torch.cat calls at each join point with _join_embeddings(), which returns the sole tensor for a one-element list instead of allocating a duplicate. Multimodal embeddings are read-only downstream, so this is value-equivalent. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

KV-cache estimation built its multimodal profiling request from raw PIL images and pushed it through the input processor and the executor request path. That coupled the encoder budget to the LLM one: the request had to fit max_num_tokens, so a shrink loop lowered encoder_max_num_tokens until it did and silently profiled a smaller encoder than the runtime uses. Build the processed encoder tensors directly instead. get_dummy_mm_data() now returns what the encoder consumes, the LLM dummy stays text-only so it can fill max_num_tokens, and the encoder runs once at its own token budget with its output retained so the peak covers both. Only capacity the run did not materialize is reserved on top. The processors now restate the encoder's input layout rather than inheriting it from the input processor, so add a guard per model that drives the real parsing/batching step and fails on drift. Also release the retained profiling output on the error path, and make the shared geometry helper private now that nothing calls it polymorphically. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Per-item slots meant the prefill path had to concatenate them into the contiguous form it consumes, so the request briefly held both. The byte budget accounted for the slots only, and the extra copy scaled with the request. Size the buffer from the declared item lengths and let record() copy each item into its own row range, so the buffer is the request's final storage: finalize_into() publishes it by reference and nothing downstream rebuilds it. This also removes two special cases. The finalized flag is gone -- buffer presence is the state, so emptied slots can no longer read as PENDING -- and the head-of-line reservation is gone, because a request's first scheduled item now allocates storage for all of them, which makes the whole footprint the natural unit to charge and lets a started request always finish. Storage for a partially encoded request is therefore its full footprint, which admission already requires to fit the budget. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

upstream renamed Mistral3InputProcessor to MistralHFInputProcessor and added MistralNativeInputProcessor alongside it; the scheduler test still referred to the old name. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

The encoder attention capacity rounded window counts up with ceil(), but get_window_index_by_thw pads by 'window - dim % window', which is a whole extra window when a merged-grid dimension is an exact multiple rather than none. That window holds only padding, so its sequence length is zero, yet it still reaches window_seq_lens and still counts as a context. Sizing therefore came up short exactly on the divisible cases: a 4x4 merged grid produces 4 windows where ceil() predicted 1, and an 8x8 grid produces 9 where it predicted 4. Mirror the encoder's padding in _windows_along() and use it for both capacity selection and runtime validation, which raises the worst-ratio frame in the existing test from 16 to 20 contexts. The encoder's padding itself is left alone: window_index depends on it, and the fix belongs on the side that has to predict it. Reported-by: chienchunhung Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Fold the None check into the isinstance check, which already covers it, and name the offending type in the error instead. Escalate the item-key count mismatch from debug to warning_once and print the three counts: it means the metadata disagrees with itself, and the visible effect is that the request silently stops participating in the encoder cache. Drop the docstring from a private helper, keeping only the invariant that matters as a comment, and stop naming a private helper from a public docstring. Single backticks, no sphinx roles. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Matches the ruff line-length the file is linted at. Only lines this branch touched are rewrapped; the one pre-existing over-length line is left alone. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

tensorrt_llm.logger does not do %-interpolation, so these printed their own format string and appended the arguments: Multimodal encoder token budget: configured=%s, base=%d, effective=%d, model_atomic_max=%d, attention_capacity=%s. 65536 65536 65536 65536 {...} Both were observed that way in server logs while investigating the encoder budget. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Schedule distribution carries request and item ids only, so a follower rank never sees the leader's encoder outputs or its cache lookups. If hits were applied only where scheduling happened, followers would encode work the leader skipped and the ranks would disagree. Drive one schedule through two independent engines, each with its own rank-local cache, and require identical encoder call counts, readiness and embeddings -- while asserting the storage is not shared, so agreeing on values cannot be mistaken for sharing them. Verified the assertion bites: making the lookup leader-only turns the encoder call counts into [2] vs [2, 2]. Requested-by: chienchunhung Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

…ding an arg Drop the 'del max_num_tokens' line and say in the docstring why the default ignores it: it sizes from the item budget, and encoders that split one item across several sequences override this and do use the token budget. Also record that the returned keys name each encoder's own attention metadata objects -- one 'attention' here, 'full_attention' and 'window_attention' for a windowed encoder -- so there is no fixed superset to promote to a TypedDict. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

The argument already names the destination at every call site, so the suffix was not carrying its weight: request.py_mm_encoder_state.finalize(request.py_multimodal_data) Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Per-item partial cache hits (NVIDIA#16817) landed the same primitives this branch grew independently. Converge on the upstream ones and generalize them where item scheduling needs more than the full-request path does. Cache lookup: the executor's hand-rolled hit/miss loop becomes `partition_encoder_cache`, which gains two parameters so it can serve both callers without regressing either: - `item_indices` scopes the lookup. Probing an item the budget cannot encode this iteration would still refresh its LRU recency and reorder eviction against items actually in flight, so the scheduler passes only what it selected. `EncoderCachePartition.looked_up` records the scope, and the full-hit/full-miss predicates read it rather than `keys`. - `keys` lets the caller supply precomputed keys. The executor builds params from `py_multimodal_data` alone while the content hashes live on the `LlmRequest`, so param-derived keys are not available there. Key derivation: `_encoder_cache_keys` now delegates to `build_encoder_cache_item_keys` when the request carries atomic-item metadata, taking each key's modality from its `item_refs` entry. This removes the "mixed-modality params are not cacheable" limitation for both paths -- that metadata is exactly what the comment said was missing. Slicing: drop this branch's `_build_multimodal_encoder_input` for upstream's `build_multimodal_encoder_input` (three layouts including audio, sibling-field slicing) plus `_apply_metadata_slice`. It takes an optional `modality` so a caller that already knows an item's modality can slice one item out of an interleaved request; the residual now keeps only that modality's payload, since leaving a sibling's unsliced tensors on it makes the residual look mixed-modality to `_lengths_by_modality`. Adds junction tests. The two features each had coverage but nothing ran them together, and the pairing has real edges: the metadata re-slice is load-bearing for `encode_multimodal_by_groups`, which splits encoder output by `multimodal_embedding_lengths` and yields zero rows without it. `MultimodalEncoderRequestState.record` deliberately does not move to `assemble_full_embedding`. That helper builds the buffer from a complete item dict; items here arrive across iterations, and holding per-item tensors until the request completes would pin each contributing encoder batch instead of copying into the single pre-sized buffer. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Two entry points were building the same thing: a request's per-item encoder outputs copied into one exact-size contiguous buffer. `assemble_full_embedding` took every item at once for the full-request path; `MultimodalEncoderRequestState` filled the same buffer incrementally for item scheduling, which receives items across iterations. Same buffer, same validation, two APIs. Keep the incremental owner, since it is the one that cannot be expressed in terms of the other, and add a classmethod for the all-at-once case that routes through `record`. Callers keep the name they had; the mixin's copy is deleted. This also drops the single-item shortcut that returned the item tensor unchanged. Its inputs are cache entries, and `TensorLRUCache.get` documents its return as an alias of the cache-owned tensor, so a one-item request ended up sharing storage with the cache: eviction could no longer free those bytes, and the cache's accounting said otherwise. `record` already copied for exactly this reason, so the unified path inherits it. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

`record()` located each item's row range with `sum(self.embedding_lengths[:item_idx])`, re-adding the prefix on every call and making a request's assembly quadratic in its item count. The declared lengths are fixed at construction, so accumulate the offsets once in `__post_init__`. This matters more than it did: folding `assemble_full_embedding` into this class routed the full-request encoder path through `record` too, so both callers paid it. `resident_output_bytes` reuses the same total rather than summing again. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

`prepare_multimodal_encoder_inputs` invoked the raw-tensor slicer once per selected item. For the packed layout that slicer splits the request's entire pixel payload and concatenates the chosen pieces, so per-item calls re-split the payload N times and `torch.cat` copies each item separately -- 67 ms per 8-item request at Qwen2.5-VL 1024x1024 shapes, against 0.03 ms for the direct row-range slice this replaced. Batch adjacent same-request, same-modality items into one call, which is what the plural `item_indices` parameter is for, and take a view instead of concatenating when the indices are a contiguous run. Both callers benefit: a scheduler picks items in order, and cache misses cluster. Back to 0.032 ms per request, with one slicer call instead of eight and no payload copy. `prepare_multimodal_encoder_inputs` now returns per-tuple length lists; `forward_multimodal_encoder_items` still emits one tensor per item, so the engine contract is unchanged. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Two host-side costs repeated work whose inputs are fixed at admission. `_runs_by_request_modality` called `get_multimodal_encoder_item_metadata` once per selected item, and each call re-validates the whole record. At 32 requests x 8 items that is 256 calls, ~460 us per iteration. Fetch it when the request changes instead, and yield it with the run so `prepare_multimodal_encoder_inputs` reuses it rather than fetching again -- 256 calls become 32, no new state. `get_mm_encoder_item_keys` rebuilt a request's keys on every iteration that scheduled any of its items. The keys derive from the content hashes and item metadata, both fixed at admission, so memoize them on the encoder request state, which has exactly that lifetime. `None` is a real result -- the request cannot participate in the cache -- so an `_UNSET` sentinel separates it from "not computed yet"; without that an unkeyable request would retry every iteration, the case the second test pins. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

yechank-nvidia · 2026-08-05T04:37:34Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-05T04:42:59Z

PR_Github #63938 [ run ] triggered by Bot. Commit: ae18a32 Link to invocation

tensorrt-cicd · 2026-08-05T08:42:52Z

PR_Github #63938 [ run ] completed with state SUCCESS. Commit: ae18a32
/LLM/main/L0_MergeRequest_PR pipeline #51873 completed with status: 'FAILURE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

github-actions Bot assigned yechank-nvidia Jul 7, 2026

yechank-nvidia changed the title ~~feat: enforce multimodal encoder runtime budgets~~ [feat][None] enforce multimodal encoder runtime budgets Jul 7, 2026

yechank-nvidia changed the title ~~[feat][None] enforce multimodal encoder runtime budgets~~ [None][feat] enforce multimodal encoder runtime budgets Jul 7, 2026

yechank-nvidia commented Jul 9, 2026

View reviewed changes

Comment thread docs/source/torch/multimodal_encoder_runtime_scheduling_design.md Outdated

2ez4bz reviewed Jul 9, 2026

View reviewed changes

2ez4bz reviewed Jul 10, 2026

View reviewed changes

yechank-nvidia force-pushed the multimodal-encoder-runtime-scheduling branch from cb85bc0 to f2242dd Compare July 10, 2026 09:14

2ez4bz reviewed Jul 15, 2026

View reviewed changes

yechank-nvidia force-pushed the multimodal-encoder-runtime-scheduling branch from c552067 to 9b9ee3f Compare July 15, 2026 09:17

2ez4bz reviewed Jul 15, 2026

View reviewed changes

This was referenced Jul 15, 2026

[Bug]: Preserve per-item processed metadata when computing multimodal token lengths #16459

Open

[Bug]: find_input_mm_embeds is order-dependent for mixed full-prefill and partial VLM requests #16460

Open

yechank-nvidia force-pushed the multimodal-encoder-runtime-scheduling branch 3 times, most recently from 8daa248 to 602b402 Compare July 20, 2026 08:15

yechank-nvidia added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 20, 2026

yechank-nvidia changed the title ~~[None][feat] enforce multimodal encoder runtime budgets~~ [None][feat] Enforce multimodal encoder runtime budgets with budgeted output storage Jul 20, 2026

yechank-nvidia marked this pull request as ready for review July 20, 2026 08:26

yechank-nvidia requested review from a team as code owners July 20, 2026 08:26

yechank-nvidia requested review from QiJune, arysef, chienchunhung and pcastonguay July 20, 2026 08:26

yechank-nvidia added 26 commits August 5, 2026 04:29

fix: align multimodal encoder runtime budgets

86673ab

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

test: consolidate multimodal encoder coverage

4e799be

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

fix: profile multimodal memory through runtime requests

185bff8

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

[None][fix] address multimodal scheduling pre-commit failures

cfd9327

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

[None][fix] repair cross-iteration encoder cache hits

1aa46f0

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

[None][fix] address focused multimodal review feedback

fc441dc

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

[None][test] release Qwen model memory between tests

ae18a32

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

yechank-nvidia force-pushed the multimodal-encoder-runtime-scheduling branch from f6528d1 to ae18a32 Compare August 5, 2026 04:37

Uh oh!

Conversation

yechank-nvidia commented Jul 7, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Measurements

Test coverage

Follow-ups

Dev Engineer Review

QA Engineer Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yechank-nvidia commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

tensorrt-cicd commented Aug 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

yechank-nvidia commented Jul 7, 2026 •

edited by coderabbitai Bot

Loading