[None][feat] Enforce multimodal encoder runtime budgets with budgeted output storage - #16051
[None][feat] Enforce multimodal encoder runtime budgets with budgeted output storage#16051yechank-nvidia wants to merge 46 commits into
Conversation
cb85bc0 to
f2242dd
Compare
c552067 to
9b9ee3f
Compare
8daa248 to
602b402
Compare
The manager budget and the scheduler's byte accounting derived row bytes from the text hidden size, but Qwen3-VL folds deepstack feature maps into every embedding row (4x wider), so reservations undercounted 4x and adopt() hit the budget ceiling mid-flight. Resolve row bytes through the mixin's embedding_dim/embedding_dtype contract first and implement it for Qwen3-VL; the engine falls back to the embedding layer's weight, then to config hidden size, for models that declare neither. Validated: the 0.98 free_gpu_memory_fraction repro (7x1024px images per request, concurrency 32) that OOM-crashed the previous design now completes 128/128 requests with zero allocation failures at throughput parity with main. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Rewrite the memory-ownership section as-built (budgeted MultimodalEncoderCacheManager, allocate-before-compute with head-of-line reservation, unpin funnel, profiling guarantee with the 0.98 repro result), update the request-state and workflow sections to the record()/adopt()/lazy-cat semantics in both EN and KR, refresh diagram region labels, and describe the new exclusive-storage semantics of encoder_cache_max_bytes in its field description. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
'unpin' misread as deallocation when the entry actually stays resident for reuse; 'hold/release' carries the right intuition (letting go of a grip does not destroy the object). Renames get_and_pin->get_and_hold, unpin_request->release_holds, pinned_bytes->held_bytes across the manager, executor funnel, state docstrings, tests, design docs, and diagrams. No behavior change. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Audit pass after the cache-manager convergence: fix the post-prefill strip comment that still described the removed output buffer and claimed the published embedding survives the dict clear (all alias sources drop together so eviction actually frees memory), disambiguate the manager's log name from the legacy full-request clone cache, scope that cache's docstring to its remaining consumers, state the byte budget in the MultimodalScheduler docstring, and correct the key-format-parity test comment (the two stores are separate today; the shared format is kept for future unification). Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Fold the manager into the executor's standard resource-manager idiom: move it under pyexecutor/, subclass BaseResourceManager, register it as ResourceManagerType.MM_ENCODER_CACHE_MANAGER so request teardown flows through the existing free_resources funnel on every termination path (the post-prefill strip still calls it early so entries become reclaimable as soon as their embedding is consumed), and report zero scheduler-facing resource counts since the byte budget is enforced at MM item selection rather than capacity admission. Also drop external framing from the class docstring and design docs in favor of describing the mechanism itself. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Anchor the follow-up plan in code where the work will happen: the legacy clone-cache getter carries the migration steps onto the MultimodalEncoderCacheManager (read/write swap, partial-hit assembly resolving TRTLLM-13996, retirement), the TRTLLM-13996 comment points at that resolution, and MultimodalEncoderItemMetadata notes the item_refs/mm_item_order manifest duplication to single-source at the input processor. Mirror the plan in the design docs' deferred list and mark the resident-byte-limit question resolved. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
These were kept in-tree during review by agreement and are not meant to merge; the as-built content is preserved outside the tree. Code docstrings and the TRTLLM-14477 TODO markers remain the in-tree source of truth. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
The scheduler's liveness guard for requests whose encoder-output bytes can never fit the storage budget raises inside the executor loop, which kills the server instead of the request. It was unreachable while prompts were bounded by max_num_tokens, but LLM chunked prefill admits longer prompts (e.g. a single long video) whose full embedding stays resident across chunks. Move the check to admission (initialize_multimodal_encoder_request), where it fails only the offending request with guidance to raise encoder_cache_max_bytes; the scheduler guard remains as an accounting-bug backstop. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Grouped encoder forwards emit torch.split views of one batch tensor; adopting the views directly let a surviving sibling entry keep the whole batch allocation resident after another entry was evicted, decoupling accounted bytes from physical memory and breaking the budget's residency guarantee. adopt() now clones values that do not exclusively own their storage (tensors that do are still stored without a copy) — the same hazard the legacy clone cache documented, reapplied at the single storage boundary. Reported by coderabbitai on PR NVIDIA#16051. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
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>
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>
591b549 to
bad576a
Compare
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>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/unittest/_torch/executor/test_multimodal_scheduler.py (1)
727-993: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd
tests/unittest/_torch/executor/test_multimodal_scheduler.pyto the test listsChanged/added test functions in this segment:
- Added
test_pipeline_ranks_reach_the_same_encoder_state_independently- Renamed
test_mm_encoder_state_finalize_into_is_a_conditional_no_op→test_mm_encoder_state_finalize_is_a_conditional_no_op- Updated
test_mm_encoder_state_publishes_the_buffer_without_copyingto callfinalize(...)Coverage verdict: insufficient — no entries for this module or these test names were found under
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/executor/test_multimodal_scheduler.py` around lines 727 - 993, Add tests/unittest/_torch/executor/test_multimodal_scheduler.py and its changed test names to the appropriate test-list entries under the test-list configuration, including test_pipeline_ranks_reach_the_same_encoder_state_independently, test_mm_encoder_state_finalize_is_a_conditional_no_op, and test_mm_encoder_state_publishes_the_buffer_without_copying.Source: Path instructions
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
507-558: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDefer MM encoder budget checks until after
_init_model_capacity().self.encoder_max_num_tokensstill falls back toself.max_num_tokenshere, and that value is only populated later in__init__. With MM encoder item scheduling enabled, the current guard can reject the common unset-max_num_tokenspath before the default budget is derived.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/model_engine.py` around lines 507 - 558, Move the MM encoder item scheduling budget validation and dependent calculations out of the current __init__ section and run them only after _init_model_capacity() has populated self.max_num_tokens and the fallback self.encoder_max_num_tokens. Preserve the existing validation, effective-budget adjustment, metadata-capacity handling, logging, and output-budget resolution once initialization is complete.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 507-558: Move the MM encoder item scheduling budget validation and
dependent calculations out of the current __init__ section and run them only
after _init_model_capacity() has populated self.max_num_tokens and the fallback
self.encoder_max_num_tokens. Preserve the existing validation, effective-budget
adjustment, metadata-capacity handling, logging, and output-budget resolution
once initialization is complete.
In `@tests/unittest/_torch/executor/test_multimodal_scheduler.py`:
- Around line 727-993: Add
tests/unittest/_torch/executor/test_multimodal_scheduler.py and its changed test
names to the appropriate test-list entries under the test-list configuration,
including test_pipeline_ranks_reach_the_same_encoder_state_independently,
test_mm_encoder_state_finalize_is_a_conditional_no_op, and
test_mm_encoder_state_publishes_the_buffer_without_copying.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 159c2227-07b1-464e-bd97-47d8d954481c
📒 Files selected for processing (8)
tensorrt_llm/_torch/models/modeling_multimodal_encoder.pytensorrt_llm/_torch/models/modeling_multimodal_mixin.pytensorrt_llm/_torch/models/modeling_qwen2vl.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_multimodal_scheduler.pytests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/models/modeling_multimodal_encoder.py
- tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
- tensorrt_llm/_torch/pyexecutor/llm_request.py
VALLIS-NERIA
left a comment
There was a problem hiding this comment.
approve for tests/unittest/_torch/executor/test_kv_cache_estimation.py
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:
encoder_max_num_items×encoder_max_num_tokensareenforced 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.
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_policyselectsDEFAULT,EAGER(advance encoder work for capacity-rejected requests), or
DISABLED(legacyinline 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 onlydifference and
kv_cache_config.max_tokenspinned equal. Caches off.Availability — past the headroom, removing the cap does not degrade: it
loses essentially every request and the server process dies.
DISABLEDDEFAULTReproduced on two hosts. At 105 GiB
DISABLEDis non-deterministic.Peak memory —
DEFAULThas zero variance across repetitions and barelymoves as the workload gets 7× heavier (136.1 / 135.9 / 135.1 GiB at 1 / 4 / 7
images per request);
DISABLEDtracks the traffic (131.7 → 138.0 → 138.1).Latency — a trade, and only under heavy load:
At 7 images both are real (non-overlapping ranges).
DEFAULTraises the fastITL 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-budgetallocate-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; reservationof unmaterialized capacity),
test_modeling_qwen2_5vl.py/test_modeling_mistral.py(capacity from processor geometry, window countsmatching the encoder's padding, dummy tensors satisfying the encoder contract),
test_scheduler_serializable_output.py(item schedule survives rankdistribution).
Follow-ups
Unify the remaining full-request consumers (side-stream prefetch,
mm_encoder_only/ disagg, non-item models) onto the item path andsingle-source the item manifest; TODO markers are anchored at the migration
sites.