Skip to content

[None][feat] Enforce multimodal encoder runtime budgets with budgeted output storage - #16051

Open
yechank-nvidia wants to merge 46 commits into
NVIDIA:mainfrom
yechank-nvidia:multimodal-encoder-runtime-scheduling
Open

[None][feat] Enforce multimodal encoder runtime budgets with budgeted output storage#16051
yechank-nvidia wants to merge 46 commits into
NVIDIA:mainfrom
yechank-nvidia:multimodal-encoder-runtime-scheduling

Conversation

@yechank-nvidia

@yechank-nvidia yechank-nvidia commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

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

@yechank-nvidia yechank-nvidia changed the title feat: enforce multimodal encoder runtime budgets [feat][None] enforce multimodal encoder runtime budgets Jul 7, 2026
@yechank-nvidia yechank-nvidia changed the title [feat][None] enforce multimodal encoder runtime budgets [None][feat] enforce multimodal encoder runtime budgets Jul 7, 2026
Comment thread docs/source/torch/multimodal_encoder_runtime_scheduling_design.md Outdated
Comment thread tensorrt_llm/_torch/models/modeling_mistral.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/model_engine.py Outdated
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tensorrt_llm/inputs/registry.py
Comment thread tensorrt_llm/_torch/models/modeling_multimodal_mixin.py Outdated
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_mistral.py Outdated
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_multimodal_mixin.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_multimodal_mixin.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_multimodal_utils.py Outdated
Comment thread tensorrt_llm/_torch/models/modeling_qwen2vl.py
Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated
@yechank-nvidia
yechank-nvidia force-pushed the multimodal-encoder-runtime-scheduling branch from cb85bc0 to f2242dd Compare July 10, 2026 09:14
Comment thread docs/source/torch/multimodal_encoder_runtime_scheduling_design_KR.md Outdated
Comment thread tensorrt_llm/inputs/registry.py
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/llm_request.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/llm_request.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/llm_request.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py Outdated
@yechank-nvidia
yechank-nvidia force-pushed the multimodal-encoder-runtime-scheduling branch from c552067 to 9b9ee3f Compare July 15, 2026 09:17
Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
Comment thread tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
Comment thread tensorrt_llm/_torch/pyexecutor/_util.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/_util.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tensorrt_llm/inputs/registry.py Outdated
Comment thread tests/unittest/_torch/executor/test_multimodal_scheduler.py Outdated
Comment thread tests/unittest/_torch/executor/test_multimodal_scheduler.py Outdated
@yechank-nvidia
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 yechank-nvidia added the api-compatible Accepted LLM API contract change that is backwards-compatible label Jul 20, 2026
@yechank-nvidia 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
yechank-nvidia marked this pull request as ready for review July 20, 2026 08:26
@yechank-nvidia
yechank-nvidia requested review from a team as code owners July 20, 2026 08:26
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>
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>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Add tests/unittest/_torch/executor/test_multimodal_scheduler.py to the test lists

Changed/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_optest_mm_encoder_state_finalize_is_a_conditional_no_op
  • Updated test_mm_encoder_state_publishes_the_buffer_without_copying to call finalize(...)

Coverage verdict: insufficient — no entries for this module or these test names were found under tests/integration/test_lists/test-db/ or tests/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 win

Defer MM encoder budget checks until after _init_model_capacity(). self.encoder_max_num_tokens still falls back to self.max_num_tokens here, 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_tokens path 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

📥 Commits

Reviewing files that changed from the base of the PR and between bad576a and a996d61.

📒 Files selected for processing (8)
  • tensorrt_llm/_torch/models/modeling_multimodal_encoder.py
  • tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
  • tensorrt_llm/_torch/models/modeling_qwen2vl.py
  • tensorrt_llm/_torch/pyexecutor/_util.py
  • tensorrt_llm/_torch/pyexecutor/llm_request.py
  • tensorrt_llm/_torch/pyexecutor/model_engine.py
  • tests/unittest/_torch/executor/test_multimodal_scheduler.py
  • tests/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 VALLIS-NERIA left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

approve for tests/unittest/_torch/executor/test_kv_cache_estimation.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

api-compatible Accepted LLM API contract change that is backwards-compatible

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants