[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.

…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-10T02:34:20Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-10T02:39:59Z

PR_Github #64923 [ run ] triggered by Bot. Commit: 79c5886 Link to invocation

tensorrt-cicd · 2026-08-10T06:46:10Z

PR_Github #64923 [ run ] completed with state SUCCESS. Commit: 79c5886
/LLM/main/L0_MergeRequest_PR pipeline #52761 completed with status: 'FAILURE'

CI Report

⚠️ 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

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

yechank-nvidia · 2026-08-11T07:41:26Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-11T07:47:10Z

PR_Github #65236 [ run ] triggered by Bot. Commit: 95106c7 Link to invocation

tensorrt-cicd · 2026-08-11T13:54:06Z

PR_Github #65236 [ run ] completed with state SUCCESS. Commit: 95106c7
/LLM/main/L0_MergeRequest_PR pipeline #53018 completed with status: 'FAILURE'

CI Report

⚠️ 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 21 commits August 10, 2026 00:45

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

bcb391e

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

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

db7f9dd

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

[None][fix] address focused multimodal review feedback

21e95b1

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

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

cbc6915

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

[None][fix] address multimodal CI and review feedback

5a5aebe

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

[None][docs] sync multimodal cache and profiling guidance

02a2f28

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

[None][fix] reject incompatible multimodal encoder scheduling

d457ece

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

[None][fix] contain multimodal encoder item failures

f6b0b7b

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

[None][fix] isolate multimodal admission metadata errors

54b427a

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

[None][fix] classify multimodal encoder request failures

e367447

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

[None][refactor] tighten multimodal scheduling boundaries

0022e3e

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

[None][fix] contain multimodal model contract errors

eab583c

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

[None][fix] release terminated multimodal encoder state

6b0cbc8

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

[None][test] cover empty attention DP encoder flush

79c5886

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

yechank-nvidia force-pushed the multimodal-encoder-runtime-scheduling branch from 0c73853 to 79c5886 Compare August 10, 2026 02:33

yechank-nvidia added 2 commits August 11, 2026 01:05

[None][refactor] remove redundant multimodal encoder item limit

e01a312

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

[None][refactor] align multimodal encoder profiling with item limits

95106c7

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

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 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

yechank-nvidia commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

yechank-nvidia commented Jul 7, 2026 •

edited by coderabbitai Bot

Loading