[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 was previously unbounded at runtime: the scheduler admitted requests by LLM/KV capacity only, so one iteration could submit more encoder items — or more physical encoder attention tokens — than the configured workspace was sized for, and encoder outputs awaiting prefill accumulated without a limit or profiling visibility. This PR makes MM encoder execution budgeted, schedulable, and memory-safe end to end:

Scheduling (per-iteration budgets over atomic items)

encoder_max_num_items × encoder_max_num_tokens are now enforced by an item-granular scheduler wrapper: an MM item (one image/video) is atomic, costed by pre-merger encoder attention tokens, and selected greedily in FCFS order; a request left partial resumes its remaining items in later iterations.
The executor's encoder step is the single encode site; an in-budget batch is simply the case where every pending item is selected, and such requests still prefill in the same iteration.
Warmup profiles boundary batches (largest atomic item / max item count), so encoder attention workspaces are sized deterministically and the token budget auto-raises to keep the largest valid atomic item schedulable.

Memory (single budgeted storage for encoder outputs)

Encoder outputs live exclusively in a MultimodalEncoderCacheManager (a BaseResourceManager): requests hold entries until prefill consumes them (held entries are never evicted; releasing a hold keeps the entry resident for reuse), the item scheduler performs allocate-before-compute against the byte budget with a head-of-line reservation (prevents partial-allocation deadlock), and teardown flows through the standard free_resources funnel on every
termination path.
KV-cache estimation reserves the manager budget on top of the profiled boundary batches, so runtime encoder transients + residency cannot exceed the estimated peak.
Per-item, content-addressed storage doubles as a cross-request cache: partial hits re-encode only the misses, and duplicate items within an iteration collapse to one resident copy. encoder_cache_max_bytes is the budget, min-clamped to one prefill iteration's embeddings.

Initial model coverage: Qwen2-VL/Qwen2.5-VL, Qwen3-VL (including deepstack-widened embedding rows via a model-declared embedding_dim), Mistral3/Pixtral.
An experimental multimodal_config.enable_eager_encoder_scheduling policy (default off) lets capacity-rejected active requests make encoder progress.

Performance

A/B on Qwen3-VL-8B-Instruct, 1×H200, aiperf, 3 repetitions with distinct seeds, vs the merge-base. KV block reuse and the embeddings cache were disabled
on both sides for the scheduling comparisons (aiperf samples from a shared image pool; cache hits would contaminate the A/B).

No regression on light multimodal load (1×512px image/req, ISL 1000/OSL 100, concurrency 1/8/32):

Metric	Δ (this PR vs main)
Request throughput	±0.2%
TTFT p50 / p99	within ±3%
ITL mean / p99	within ±2%

Bursty image-heavy load (4×1024px images/req, concurrency 8/32): bounding per-iteration encoder work stops large encoder batches from stalling
in-flight decodes.

Metric	conc 8	conc 32
Request throughput	+0.7%	+0.3%
ITL mean	−9.2% (13.9→12.6 ms)	−10.7% (38.5→34.3 ms)
TTFT p50	+12.7%	+24.9%
TTFT p90/p99	parity	parity
Request latency p50	−0.7%	−2.6%

The TTFT cost is concentrated in the median (items past the budget wait a few iterations); worst-case TTFT does not get worse.

Per-item cache reuse (caches enabled; 3 shared 1408px images + 1 unique per request): warm TTFT 647→559 ms (−13.6%, ±5 ms across reps). An all-or-nothing full-request key cannot hit when any item differs; the per-item path can.

Memory-safety stress (free_gpu_memory_fraction=0.98, 7×1024px images/req, concurrency 32): an earlier revision of this branch OOM-crashed here (76 allocation failures) because encode-ahead residency was unbounded and unprofiled — the same failure class behind free_gpu_memory_fraction workarounds like #16315. With the budgeted storage this completes 128/128 requests with zero allocation failures at throughput parity with main (2.47 vs 2.50 req/s, ITL −4.7%). Re-running the full matrix after the storage rework shows all metrics within ±1%.

Test Coverage

tests/unittest/_torch/executor/test_multimodal_encoder_cache_manager.py: hold/adopt/eviction semantics, clone-free adoption, dedup, idempotent release, BaseResourceManager surface.
tests/unittest/_torch/executor/test_multimodal_scheduler.py: atomic packing/backfill, byte-budget allocate-before-compute, head-of-line reservation, oversized-request fail-fast, cache attach (partial hits, key guards with temporary-key fallback), request-state transitions, admission gating, Qwen/Mistral token-unit metadata.
tests/unittest/_torch/executor/test_kv_cache_estimation.py: manager-budget reservation in estimation.
Existing suites (test_multimodal_mixin.py, test_multimodal_embedding_lengths.py, inputs/test_multimodal.py, llmapi/test_llm_args.py) pass unchanged in behavior.

Follow-ups

TRTLLM-14477: unify the remaining full-request consumers (side-stream prefetch, mm_encoder_only/disagg, non-item models) onto the budgeted storage — resolving TRTLLM-13996 (partial-hit assembly) for all consumers — and single-source the prompt-order item manifest. TODO markers are anchored in code at the exact migration sites.

Summary by CodeRabbit

New Features
- Added item-level scheduling for multimodal encoder workloads, improving handling of image and video requests.
- Added caching and reuse of multimodal encoder outputs, including support for partial cache hits.
- Added optional eager encoder scheduling for earlier multimodal processing.
- Added support for per-item multimodal metadata across supported vision models.
Documentation
- Updated configuration and telemetry documentation to use encoder_max_num_items.
- Clarified multimodal embedding cache behavior.
Bug Fixes
- Improved encoder capacity estimation and workload budgeting for multi-item requests.

2ez4bz · 2026-07-14T23:29:08Z

        # Multimodal data
        self.py_multimodal_data = kwargs.pop("py_multimodal_data", None)
+        self.py_is_multimodal_encoder_request = False
+        self.py_mm_encoder_outputs: List[Optional[torch.Tensor]] = []


Nit: could you leave some comments for what these will hold / how they are intended to be used?

For example, it's a bit unclear just from looking at these how output_buffer and outputs will differ / what they will be distinctly used for.

Also, I don't have a great suggestion as an alternative (yet), but isn't it weird for these things to be part of the LlmRequest object? I guess they're kind of this awkward in-between since they're not part of the request, nor output, and MultimodalRuntimeData gets instantiated after scheduling - so too late.

Could they be encapsulated in a e.g. MultimodalEncoderRequestState ?

@dataclass class MultimodalEncoderRequestState: metadata: MultimodalEncoderItemMetadata output_offsets: tuple[int, ...] completed: list[bool] output_buffer: torch.Tensor | None = None @property def progress(self) -> MultimodalEncoderProgress: ... def pending_items(self) -> Iterable[tuple[int, int]]: ... def completes_with(self, item_indices: Collection[int]) -> bool: ... def commit(self, item_idx: int, output: torch.Tensor) -> torch.Tensor | None: ... def release(self) -> None: ...

and then we have

py_mm_encoder_state: MultimodalEncoderRequestState | None

in the request class.

Validation can be owned by __post_init__ or other (private) methods. Basically, we can encapsulate as much as possible in one place, and the rest of the code can go about its business using the exposed fields / methods.

could you leave some comments for what these will hold / how they are intended to be used?

Added comments.

I agree that this variable is a bit awkward. But the reason is that I didn't want to touch C++ code and ReqeustState.

Thinking about adding your suggestion, but I think that will bring some lots of code changes.

Let me think about that further.

Thx. for bringing this up.

Added the state class.

2ez4bz · 2026-07-15T17:52:19Z

+        item_limits = [distribute_items(1)]
+        if self._model_engine.encoder_max_num_items > 1:
+            item_limits.append(
+                distribute_items(self._model_engine.encoder_max_num_items))


I'm not sure I follow why we need the append here? As in, why do we need both:

item_limits = [ distribute_items(1), distributed_tiems(encoder_max_num_items), ]

?

They're two different boundary workloads under the same token budget: distribute_items(1) profiles one maximal-length item (length-dominated peak, e.g. attention scratch), while distribute_items(encoder_max_num_items) profiles the maximal item count (count-dominated peak, e.g. per-item overheads).

Which one dominates depends on the model, so warmup runs both; the guard just skips the second when it degenerates to the first (encoder_max_num_items == 1).

Added a comment in code explaining this.

2ez4bz · 2026-07-15T18:25:58Z

+                            raise TypeError(
+                                "multimodal_embedding_lengths must be a list")
+                        if (existing_embedding_lengths
+                                != item_metadata.output_embedding_lengths):


Ok but why are we checking for equality, then...?

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

- Move MultimodalEncoderItemMetadata field validation into a validate() method invoked by the metadata reader so every consumer receives validated metadata; drop the now-redundant checks in the mixin and get_multimodal_encoder_token_lengths - Flatten the input-processor wrapper metadata block into guard clauses and fold the None check into the isinstance check with a typed error - Add field-level docstrings to MultimodalEncoderItemMetadata, MultimodalEncoderProgress, the LlmRequest MM encoder fields, and scheduled_mm_encoder_items declarations - Scan Qwen vision spans with bounded list.index calls (~6x faster) and reject unclosed, mixed, and empty spans explicitly - Require an explicit encoder token budget for item-scheduling models at engine init and make initialize_multimodal_encoder_request take a required max_num_tokens - Compute the largest item cost once, build output offsets with itertools.accumulate, include both lengths in the encoder output mismatch error, and remove test leftovers Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

- Add a high-level docstring and per-branch comments to MultimodalScheduler.schedule_request describing the LLM-capacity-coupled policy and its full-request vs item paths - Log when MM encoder warmup is skipped because the input processor does not implement get_dummy_mm_data_for_tokens(), including the consequence for memory estimation - Explain why warmup profiles both the maximal-length-item and maximal-item-count boundary batches Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

The knob counts atomic multimodal items per encoder iteration, not LLM requests or attention sequences, so name it accordingly end-to-end (config field, runtime attribute, tests, api-stability reference, telemetry doc, golden manifest). The field is prototype status. Also document the contract in the field description and scheduler docstring (item count vs token-size budgets; encoders that split one item into multiple attention sequences derive their workspace capacity from the token budget instead), and fold the init-only derived attributes (configured/base/effective token budgets, model atomic max) into locals now that they only feed one startup log line. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

- Move the eager-scheduling incompatibility checks (attention DP, disaggregated serving) into a TorchLlmArgs model validator so they fail at args validation; the transceiver condition reduces to cache_transceiver_config.backend being set, which is exactly its creation condition - Move the side-stream-prefetch conflict check and the MultimodalScheduler wrapping from PyExecutor.__init__ into create_py_executor_instance: the conflict depends on the loaded model's item-scheduling capability, so it cannot be a config validator, and the composition root is the earliest point with full information - Rename PyExecutor.is_multimodal_model to the private and accurate _supports_mm_encoder_item_scheduling and drop the intermediate local - Add validator unit tests Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Replace the SimpleNamespace request fakes with real LlmRequest objects whose per-item encoder state is populated by initialize_multimodal_encoder_request, so the tests cover the same admission path as production. Drop the namedtuple pickling test, which exercised the stdlib rather than our code. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

The design documents were review-time working notes; the durable contract now lives in the llm args field descriptions and the scheduler/model-engine docstrings. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Integrate the item-scheduling path with the TensorLRUCache introduced by the MM encoder cache feature, and make the executor's encoder step the single encode site for item-scheduling models. - Extract the request-side item state transitions into single-writer helpers: fill_mm_encoder_output_into_request (validate, lazily allocate the contiguous buffer, copy into the item's reserved rows, record the slot view) and finalize_multimodal_encoder_request (publish multimodal_embedding and strip raw encoder inputs once every slot is filled). Both the encode path and the cache path go through them. - Build per-item cache keys from request metadata (build_encoder_cache_item_keys), single-sourcing the key format with the full-request path in _encoder_cache_item_key so entries written by either path hit from the other; per-item modality comes from item_refs, which also makes mixed-modality requests keyable on the item path. - Attach cache hits before scheduling (PyExecutor._attach_mm_encoder_cache_hits): hits are filled into the request's slots ahead of item selection, so they consume no encoder budget and are never re-encoded; partially attached requests re-compute only the misses. Participation is gated per request by ModelEngine.get_mm_encoder_cache_and_keys; requests that cannot build keys run exactly as without a cache. - Write freshly encoded items through to the cache after filling the request (put clones, so entries neither alias nor pin request buffers and eviction can never regress request progress). - Remove the full-request fast path (_can_schedule_full_request_mm_batch): an in-budget batch now simply selects every pending item and still prefills in the same iteration. This closes the attach/schedule window with the cache structurally and makes the all-or-nothing cache limitation (TRTLLM-13996) unreachable in the executor loop for item-scheduling models. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Restore the review-time design documents (still slated for removal before merging) and bring them and the workflow diagrams up to date with the current diff: the encoder_max_num_items rename, the single encode site (full-request fast path removal), and the embeddings-cache integration (pre-scheduling hit attachment, write-through, shared per-item key format, and the resulting TRTLLM-13996 scoping). The diagrams gain an embeddings-cache integration region and updated scheduler/executor labels. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Replace the four flat request fields (py_is_multimodal_encoder_request, py_mm_encoder_outputs, py_mm_encoder_output_buffer, py_mm_encoder_output_offsets) with one py_mm_encoder_state: Optional[MultimodalEncoderRequestState]. State presence doubles as the request-kind signal, so the flag cannot drift from the slots it used to describe. The state owns the item slots, the contiguous output buffer, and the offsets, and carries the transitions as methods: fill() is the single writer shared by the encode path and the embeddings-cache attach, and finalize_into() publishes multimodal_embedding and strips the raw encoder inputs. The offsets/slots invariant is enforced at construction, and scheduler/admission consumers iterate pending_item_indices() instead of zipping slot lists. Deliberate differences from the shape suggested in review: no metadata field (py_multimodal_data stays the only executor-boundary carrier), tensor slots instead of completed booleans (the owned copies double as eviction safety for the embeddings cache), and the state never touches the request dict except through the finalize_into(multimodal_data) argument. The state is rank-local and never serialized; each rank constructs its own at admission. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Converge the MM encoder item-scheduling path onto a single budgeted storage (the vLLM EncoderCacheManager shape) so encoder-output GPU residency is bounded and schedulable, closing the unaccounted encode-ahead residency that could OOM at high KV-cache fractions: - MultimodalEncoderCacheManager: byte-budgeted storage with per-request pinning (pinned entries are never evicted), clone-free adopt() that collapses duplicate keys to one resident copy (within-iteration dedup), zero-ref LRU eviction, and can_allocate() for allocate-before-compute. - MultimodalEncoderRequestState drops its private output buffer: slots alias pinned manager entries via the record() writer; finalize_into() publishes the item-view list and the per-request contiguous embedding materializes lazily inside the prefill forward. - The MM item scheduler performs allocate-before-compute against the manager budget with a head-of-line reservation (prevents partial-allocation deadlock) and fails fast on requests that can never fit the budget. - Pins release through a single teardown funnel: the post-prefill strip and _do_terminate_request (cancel/error/finish) both route to the idempotent unpin. - encoder_cache_max_bytes is clamped from below to one prefill iteration's embeddings (max_num_tokens x embedding row bytes) since the cache is now the only storage; requests without stable content hashes store under request-scoped temporary keys. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

The estimator already profiles boundary encoder batches with the last output held across the LLM dummy forward; pointing the additive cache reservation at the manager's min-clamped budget (instead of the legacy clone-cache bytes) completes the coexistence scene: runtime encoder transients are bounded by the profiled boundary batches and all encoder-output residency is bounded by the reserved manager budget. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

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>

The vision path updated attn_metadata.max_seq_len to each batch's largest segment, and that value participates in the C++ attention-op cache key (mMaxSeqLen/mMaxContextLength): every distinct segment length in serving traffic created and permanently cached a new AttentionOp, each with its own cudaMalloc'd workspace and semaphores, growing GPU usage without bound — the exact mechanism behind the Qwen encoder OOMs on memory-tight GPUs (observed as a cudaMalloc failure in reserveSemaphoreArray on 1xH100 during MMMU at ~430/900 requests, where weights leave ~13GB headroom). Pin max_seq_len once in setup_attn_metadata to the encoder token budget (no encoder segment can exceed it: atomic items are admission-checked against the budget, now also validated per batch) so a single op — created during warmup's boundary batches and therefore part of the profiled peak — serves every runtime batch. This matches the KV-cache attention convention where max_seq_len is a stable manager-owned value rather than a per-batch one. 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>

chienchunhung

One additional correctness concern:

chienchunhung · 2026-07-21T17:43:51Z

+            max_grid_w = max_area // grid_h
+            for grid_w in range(min_grid_w, max_grid_w + 1):
+                area = grid_h * grid_w
+                windows = (math.ceil(grid_h / window_side) *


I think the window-attention capacity can be underestimated when a grid dimension is exactly divisible by the window size.

For example, consider a merged grid of 1×4 with a 4×4 window:

The capacity calculation counts ceil(1 / 4) * ceil(4 / 4) = 1 window.

At runtime, get_window_index_by_thw computes padding as window_size - dimension % window_size. The width is exactly divisible, so it adds 4 columns instead of 0, producing a padded 4×8 grid with 2 windows.

The second window has a sequence length of zero, but it is still retained in window_seq_lens and counted by num_contexts.

Each such frame consumes 16 physical tokens, so a 160-token batch can contain 10 frames and produce 20 runtime contexts. However, the current calculation allocates window_attention for only 16 contexts in the corresponding test case.

Could we either exclude zero-length windows at runtime or use the same padding behavior when calculating and validating the metadata capacity? It would also be helpful to add this exactly-divisible case to the test.

2ez4bz · 2026-07-21T16:41:24Z

+    def get_encoder_attention_metadata_capacity(
+            self, max_num_items: int, max_num_tokens: int) -> dict[str, int]:
+        """Map item/token budgets to model-internal attention sequences."""
+        del max_num_tokens


Remove this line.

2ez4bz · 2026-07-21T16:47:22Z

+        """Map item/token budgets to model-internal attention sequences."""
+        del max_num_tokens
+        return {
+            "attention": max(max_num_items, _ENCODER_FALLBACK_MAX_NUM_REQUESTS)


Out of curiosity, is there a maximal superset of possible keys this should return? If so, should we make it slightly more structured with either a NamedTuple, or a TypedDict? The latter can be used the same way as the current free-from dict, but has the advantage of being self-documenting.

2ez4bz · 2026-07-21T16:54:00Z

+    ) -> Hashable:
+        """Build the cache key of one atomic item.
+
+        The single source of the key format: the full-request path


Nit: I think this docstring can be omitted.

2ez4bz · 2026-07-21T16:58:53Z

+    ) -> Optional[list[Hashable]]:
+        """Build per-item cache keys from request-level item metadata.
+
+        Unlike `_encoder_cache_keys`, the modality comes from each item's


Nit: I don't think we should mention private / protected attributes in public docstrings.

Could be inline code comments, though.

2ez4bz · 2026-07-21T16:59:34Z

+        if multimodal_hashes is None or kwargs_hash is None:
+            return None
+        if not (len(multimodal_hashes) == len(item_refs) == len(embedding_lengths)):
+            logger.debug(


Should this be debug or warning?

2ez4bz · 2026-07-21T18:44:43Z

+                if batch_idx + 1 == len(self._dummy_encoder_inputs):
+                    retained_output = output
+                else:
+                    del output


Nit: can you leave a comment for why this is here?

2ez4bz · 2026-07-21T20:27:36Z

+                "output shape, dtype, and device")
+        self.outputs[item_idx] = output_view
+
+    def finalize_into(self, multimodal_data: Dict[str, Any]) -> bool:


Nit: is the _into necessary?

2ez4bz · 2026-07-21T21:19:06Z

+                model_max_atomic_item_tokens,
+                self.mm_encoder_attention_metadata_capacity,
+            )
+            self.mm_encoder_cache_manager = (


A few questions:

Out of an abundance of caution, should we disable the TensorLRUCache that the mixin may have internally? Both at the config-level, and at the attribute level? So we don't cache twice.

What is our plan to consolidate these 2 levels of caching? I scheduled a meeting to discuss this later, but leaving this question so we can think about it more before then.

2ez4bz · 2026-07-21T21:20:29Z

+                self.mm_encoder_attention_metadata_capacity = (
+                    attention_metadata_capacity)
+            logger.info(
+                "Multimodal encoder token budget: configured=%s, base=%d, "


Same comment re printf style formatting vs using f-strings.

2ez4bz · 2026-07-21T21:30:01Z

+
+        manager = self.mm_encoder_cache_manager
+        if manager is None:
+            raise RuntimeError(


Out of curiosity, why?

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 27 commits July 21, 2026 06:00

refactor: clarify multimodal encoder execution path

23f05d2

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

feat: add eager multimodal encoder scheduling policy

e10b022

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

fix: honor explicit multimodal encoder runtime budgets

8283bb3

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

docs: sync multimodal encoder scheduling design

8c4424c

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

fix: align multimodal encoder runtime sizing

7d4d895

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

yechank-nvidia force-pushed the multimodal-encoder-runtime-scheduling branch from 2049786 to 591b549 Compare July 21, 2026 06:10

chienchunhung reviewed Jul 21, 2026

View reviewed changes

2ez4bz reviewed Jul 21, 2026

View reviewed changes

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

Performance

Test Coverage

Follow-ups

Summary by CodeRabbit

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

chienchunhung left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yechank-nvidia commented Jul 7, 2026 •

edited by coderabbitai Bot

Loading