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

.claude/skills/trtllm-model-onboard-multimodal/SKILL.md

@@ -306,16 +306,16 @@ class {Name}Model(PreTrainedModel):
 
 Subclass **both** `BaseMultimodalInputProcessor` (drives every real request) and `BaseMultimodalDummyInputsBuilder` (drives engine warmup / KV-cache profiling). Colocate in the modeling file. References: `Qwen3VLInputProcessorBase` (image+video), `Mistral3InputProcessor` (Pixtral).
 
-**Encoder KV-cache memory profiling (deterministic dummy sizing).** The KV-cache profiler sizes the encoder's memory contribution by running the encoder **once** on a worst-case dummy (held resident through the peak measurement), **decoupled** from the text-only LLM dummy. To opt in, the model exposes `encode_multimodal_inputs` (via `MultimodalModelMixin`) and the input processor implements the modality-agnostic dummy contract on `BaseMultimodalDummyInputsBuilder`:
+**Encoder KV-cache memory profiling (deterministic dummy sizing).** The KV-cache profiler sizes the encoder's memory contribution by running the encoder **once** on a worst-case dummy (held resident through the peak measurement), **decoupled** from the text-only LLM dummy. To opt in, the model exposes `encode_multimodal_inputs` (via `MultimodalModelMixin`) and the input processor implements the existing `BaseMultimodalDummyInputsBuilder` contract:
 
-- `get_mm_max_tokens_per_item() -> {modality: tokens}` — per-modality worst-case single-item encoder-attention tokens. The keys enumerate the modalities the model encodes; the profiler splits the shared `encoder_max_num_tokens` across them in proportion to these (so they share one microbatch cap, not each the whole budget). Default `{}` → no direct encoder profiling.
-- `get_dummy_mm_data_for_tokens(*, max_tokens_per_modality, dtype) -> multimodal_data` — materialize the processed encoder tensors **directly** (zeros of the exact shape the processor would emit; no PIL image + HF-processor round-trip), merged into one `multimodal_data` dict so a single `encode_multimodal_inputs` profiles the combined peak. Default raises `NotImplementedError`.
+- `get_mm_max_tokens_per_item(max_num_encoder_tokens=None) -> {modality: tokens}` — report bounded startup maxima when the argument is `None`, or the largest legal item per modality under a concrete runtime token budget.
+- `get_dummy_mm_data(*, max_num_encoder_tokens, mm_counts, dtype) -> multimodal_data` — build the profiler-selected per-modality item counts **directly** as processed encoder tensors (no PIL image + HF-processor round-trip).
 
-Vision models implement these via the size trio: `get_num_mm_tokens(*, width, height, num_frames)` (size → **pre-merger encoder-attention tokens**; the single source of truth shared with the hashing path `get_num_tokens_per_image`/`_video`), its inverse `get_size_for_max_tokens(max_tokens)` (largest aspect-bounded size whose token count ≤ budget, capped at `max_pixels`), and `get_dummy_mm_data_for_size(...)`. Qwen builds `pixel_values`/`image_grid_thw`; Mistral builds `pixel_values`/`image_sizes` and keeps its ViT patch count off `get_num_mm_tokens` (a private `_vit_tokens` helper) so the LLM-side Pixtral hashing count is unchanged. A model with neither contract falls back to a text-only dummy (encoder memory unaccounted). Don't hardcode the encoder attention workspace (`max_num_*=8192`): inherit `MultimodalEncoderMixin` and let the engine size it via `setup_attn_metadata` at load.
+Following vLLM's separation of common budgeting from model processing geometry, the profiler selects the largest legal runtime workload and computes its repetition count under `encoder_max_num_tokens` and `encoder_max_batch_size`. It passes the resulting plan (for example, `{"image": 8}`) to the concrete processor's `get_dummy_mm_data`; the processor owns only geometry and tensor layout. Qwen builds `pixel_values`/`*_grid_thw`; Mistral builds `pixel_values`/`image_sizes` and uses its private `_vit_tokens` helper so the LLM-side Pixtral hashing count is unchanged. A model without the contract falls back to a text-only dummy (encoder memory unaccounted). Don't hardcode the encoder attention workspace (`max_num_*=8192`): inherit `MultimodalEncoderMixin` and let the engine size it via `setup_attn_metadata` at load.
 
-The workspace dimensions come from `TorchLlmArgs.get_encoder_runtime_sizes()` → `(encoder_max_batch_size, encoder_max_num_tokens)` — two prototype knobs that size the encoder's `AttentionMetadata` independently of the LLM batch and fall back to the LLM-side `max_batch_size` / `max_num_tokens` when unset. `encoder_max_num_tokens` is exactly the per-iteration encoder microbatch token cap that `get_dummy_mm_data_for_tokens` saturates (and that the profiler splits across modalities), so an encoder microbatch can be sized larger than the LLM `max_num_tokens` without inflating the KV-cache budget. Read them via `get_encoder_runtime_sizes()` rather than the raw fields so the fallback is applied.
+The workspace dimension comes from `encoder_max_num_tokens`, falling back to the LLM-side `max_num_tokens` when unset. `encoder_max_batch_size` independently bounds the number of atomic items in a multimodal encoder iteration and falls back to `max_batch_size`. The profiler resolves both limits into `mm_counts` before calling `get_dummy_mm_data`. Model-internal attention sequence counts are not the item batch size; they are derived separately from the token budget and model geometry.
 
-> Mixed image+video+audio models (nemotron-nano / phi4mm) compose multiple modality dummies through the same `get_dummy_mm_data_for_tokens` (return `{"image": ..., "audio": ...}`); a per-modality `ModalityDummySizer` composition is the planned home for the shared orchestration.
+> Mixed image+video+audio models profile the supported modality with the largest legal per-item workload under the configured limits. Runtime mixed-modality requests still share the same aggregate token limit.
 
 Implement `call_with_text_prompt(inputs, sampling_params)` — the per-model text-prompt path. **Don't override `__call__`**: the base class's concrete `__call__` dispatches here for text prompts, and also detokenizes `prompt_token_ids → prompt` and falls through to here for non-fast-path VLMs. `call_with_text_prompt` does:
 
@@ -444,7 +444,7 @@ Follow `CONTRIBUTING.md`. Title `[JIRA/NVBUG/None][type] description`, `git comm
 
 **Input processor**
 - [ ] Subclasses both `BaseMultimodalInputProcessor` and `BaseMultimodalDummyInputsBuilder`.
-- [ ] Encoder KV-cache profiling: implements the deterministic dummy contract (`get_mm_max_tokens_per_item` + `get_dummy_mm_data_for_tokens`, vision via the `get_num_mm_tokens` / `get_size_for_max_tokens` / `get_dummy_mm_data_for_size` trio) and the model exposes `encode_multimodal_inputs`; encoder inherits `MultimodalEncoderMixin` (no hardcoded `max_num_*=8192` — sized by `setup_attn_metadata`). Skipping these = text-only dummy, encoder memory unaccounted.
+- [ ] Encoder KV-cache profiling: implements the deterministic dummy contract (`get_mm_max_tokens_per_item` + `get_dummy_mm_data`) and the model exposes `encode_multimodal_inputs`; encoder inherits `MultimodalEncoderMixin` (no hardcoded `max_num_*=8192` — sized by `setup_attn_metadata`). Skipping these = text-only dummy, encoder memory unaccounted.
 - [ ] `call_with_text_prompt` (not `__call__` — that's the base-class dispatcher) runs HF AutoProcessor + tokenizer, builds `multimodal_data` by modality, computes `mrope_config` on CPU, `_postprocess`-rewrites mm token ids to the OOV sentinel.
 - [ ] `mm_processor_kwargs` flow-through preserved. (Tokenized fast path is optional: set `supports_token_id_mm_expansion = True` + implement `get_text_with_mm_placeholders` / `expand_prompt_token_ids_for_mm`; otherwise the base class detokenizes token-ID inputs automatically.)
 - [ ] `_attach_multimodal_embeddings_impl` implemented (not the `attach_multimodal_embeddings` wrapper) if `@support_multimodal_disaggregated`.

docs/source/models/supported-models.md

@@ -176,8 +176,12 @@ The following optimizations are available to models that implement
   Set `multimodal_config.encoder_cache_max_bytes` to its capacity (for example, `"512MiB"`), or
   `0` to disable it. Entries are cached per multimodal item, but a request reuses cached embeddings
   only when all of its items hit the cache. At present, only single-modality requests are cacheable;
-  mixed-modality requests bypass the cache. When combined with side-stream prefetch, peak memory is
-  the cache capacity plus any in-flight prefetched encoder inputs and outputs.
+  mixed-modality requests bypass the cache. For models with item-level encoder scheduling, the
+  cache also composes with the item path: cached items skip encoder execution after selection but
+  still count against the per-iteration item and token budgets, partially cached requests re-compute
+  only the missing items, and items encoded through the item path populate the cache for later
+  requests. When combined with side-stream prefetch, peak memory is the cache capacity plus any
+  in-flight prefetched encoder inputs and outputs.
 
 # Visual Generation Models
 

tensorrt_llm/_torch/models/modeling_gemma4_vision.py

@@ -735,8 +735,8 @@ def __init__(self, model_config: ModelConfig):
 
         # SigLip-style context-only metadata (kv_cache_manager=None, no decode);
         # built by the engine via ``MultimodalEncoderMixin.setup_attn_metadata``
-        # at the encoder ``(encoder_max_batch_size, encoder_max_num_tokens)``
-        # budget, then re-prepared each forward with the actual per-image seq
+        # at the ``encoder_max_num_tokens`` budget, then re-prepared each
+        # forward with the actual per-image seq
         # lens. The vision tower runs once per LLM step across all images, so
         # the batch axis is the cross-request image count.
         self.metadata_cls = get_attention_backend(model_config.attn_backend).Metadata

tensorrt_llm/_torch/models/modeling_gemma4mm.py

@@ -610,6 +610,7 @@ def partition_encoder_cache(
                 hits={},
                 miss_indices=list(range(len(partition.keys))),
                 keys=partition.keys,
+                looked_up=partition.looked_up,
             )
         if (
             modality in ("image", "audio")
@@ -634,6 +635,7 @@ def partition_encoder_cache(
                     hits={},
                     miss_indices=list(range(len(partition.keys))),
                     keys=partition.keys,
+                    looked_up=partition.looked_up,
                 )
         return partition
 

tensorrt_llm/_torch/models/modeling_mistral.py

@@ -1,6 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import copy
 import dataclasses
-from typing import Any, Dict, List, Sequence, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
 
 import torch
 import torchvision
@@ -47,7 +50,8 @@
                                  MultimodalPlaceholderPlacement, TextPrompt,
                                  register_input_processor)
 from tensorrt_llm.inputs.multimodal import MultimodalParams
-from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY
+from tensorrt_llm.inputs.registry import (MULTIMODAL_PLACEHOLDER_REGISTRY,
+                                          MultimodalEncoderItemMetadata)
 from tensorrt_llm.inputs.utils import encode_base64_image
 from tensorrt_llm.llmapi import SamplingParams
 from tensorrt_llm.logger import logger
@@ -406,6 +410,40 @@ def processor(self) -> AutoProcessor:
     def dtype(self) -> torch.dtype:
         return self._dtype
 
+    def get_mm_encoder_item_metadata(
+        self,
+        _prompt_token_ids: List[int],
+        multimodal_data: Dict[str, Any],
+    ) -> Optional[MultimodalEncoderItemMetadata]:
+        """Return Pixtral image items and physical ViT patch counts."""
+        image_data = multimodal_data.get("image")
+        if not isinstance(image_data, dict):
+            return None
+        image_sizes = image_data.get("image_sizes")
+        if image_sizes is None:
+            return None
+        patch, merge, _, _ = self._vision_geometry()
+        encoder_token_lengths = [
+            self._vit_tokens(width=int(width), height=int(height), patch=patch)
+            for height, width in image_sizes
+        ]
+        min_tokens_per_image = merge * merge
+        if any(token_length < min_tokens_per_image or token_length %
+               min_tokens_per_image for token_length in encoder_token_lengths):
+            raise ValueError(
+                "Processed Mistral image geometry must contain a nonempty "
+                f"multiple of {min_tokens_per_image} encoder tokens")
+        item_refs = [("image", item_idx)
+                     for item_idx in range(len(encoder_token_lengths))]
+        output_embedding_lengths = [
+            token_length // (merge * merge)
+            for token_length in encoder_token_lengths
+        ]
+        return MultimodalEncoderItemMetadata(
+            item_refs=item_refs,
+            encoder_token_lengths=encoder_token_lengths,
+            output_embedding_lengths=output_embedding_lengths)
+
     @torch.inference_mode()
     def call_with_text_prompt(
         self, inputs: TextPrompt, sampling_params: SamplingParams
@@ -453,7 +491,7 @@ def call_with_text_prompt(
     # Deterministic dummy sizing for KV-cache encoder profiling.
     #
     # These power the modality-agnostic dummy contract
-    # (``get_mm_max_tokens_per_item`` / ``get_dummy_mm_data_for_tokens``) and
+    # (``get_mm_max_tokens_per_item`` / ``get_dummy_mm_data``) and
     # are deliberately kept separate from the hashing path: the hashing path
     # (``get_num_tokens_per_image``) keeps using the processor's LLM-side
     # token count (Pixtral grid + ``[IMG_BREAK]``/``[IMG_END]`` framing), while
@@ -495,65 +533,101 @@ def get_size_for_max_tokens(self, *, max_tokens: int) -> Dict[str, int]:
             edge -= unit
         return {"width": edge, "height": edge, "num_frames": 1}
 
-    def get_dummy_mm_data_for_size(
+    def _dummy_mm_data_for_size(
         self,
         *,
         width: int,
         height: int,
-        num_frames: int = 1,
         num_images: int = 1,
         dtype: torch.dtype | None = None,
     ) -> Dict[str, Any]:
-        """Processed Pixtral encoder tensors for ``num_images`` identical
-        ``(width, height)`` images: a ``[num_images, C, H, W]`` ``pixel_values``
-        zero tensor (content is irrelevant for memory profiling) plus the
-        matching ``image_sizes`` list the vision tower consumes."""
+        """Build processed Pixtral vision tensors of the requested geometry.
+
+        Image-only: Pixtral has no temporal axis, so there is no frame count
+        to take here (contrast the Qwen helper, which grids over frames).
+        """
         _, _, channels, _ = self._vision_geometry()
         num_images = max(num_images, 1)
-        pixel_values = torch.zeros((num_images, channels, height, width),
-                                   dtype=dtype or self.dtype)
-        image_sizes = [[height, width]] * num_images
+        pixel_values = torch.zeros(
+            (num_images, channels, height, width),
+            dtype=dtype or self.dtype,
+        )
         return {
             "image": {
                 "pixel_values": pixel_values,
-                "image_sizes": image_sizes,
+                "image_sizes": [[height, width]] * num_images,
             }
         }
 
-    def get_mm_max_tokens_per_item(self) -> Dict[str, int]:
-        """Largest single image's ViT patch count (the ``max_image_size``-capped
-        square), used to weight the shared-budget split. Image only -- image and
-        video share the Pixtral ViT."""
+    def get_mm_max_tokens_per_item(
+        self,
+        max_num_encoder_tokens: Optional[int] = None,
+    ) -> Dict[str, int]:
+        """Return the largest legal Pixtral image encoder item size."""
         patch, merge, _, max_size = self._vision_geometry()
         unit = patch * merge
         edge = max((max_size // unit) * unit, unit)
-        return {"image": self._vit_tokens(width=edge, height=edge, patch=patch)}
-
-    def get_dummy_mm_data_for_tokens(
+        max_image_tokens = self._vit_tokens(width=edge,
+                                            height=edge,
+                                            patch=patch)
+        token_budget = (max_num_encoder_tokens if max_num_encoder_tokens
+                        is not None else max_image_tokens)
+        size = self.get_size_for_max_tokens(max_tokens=token_budget)
+        encoder_tokens = self._vit_tokens(width=size["width"],
+                                          height=size["height"],
+                                          patch=patch)
+        if encoder_tokens > token_budget:
+            return {}
+        return {"image": encoder_tokens}
+
+    def get_max_mm_encoder_output_embeddings(
+            self, max_num_encoder_tokens: int) -> int:
+        """Bound post-merger embeddings from one Pixtral encoder iteration."""
+        _, merge, _, _ = self._vision_geometry()
+        return max_num_encoder_tokens // (merge * merge)
+
+    def get_mm_encoder_attention_metadata_capacity(
+            self, max_num_tokens: int) -> Optional[Dict[str, int]]:
+        """Bound Pixtral contexts by the physical-token budget."""
+        _, merge, _, _ = self._vision_geometry()
+        min_tokens_per_image = merge * merge
+        return {"attention": max(1, max_num_tokens // min_tokens_per_image)}
+
+    def get_dummy_mm_data(
         self,
         *,
-        max_tokens_per_modality: Dict[str, int],
+        max_num_encoder_tokens: int,
+        mm_counts: Mapping[str, int],
         dtype: torch.dtype | None = None,
     ) -> Dict[str, Any]:
-        """Vision implementation of the agnostic profiler entry: fill the
-        ``"image"`` budget with identical worst-case images. ``num_images`` is
-        derived from the realized patch count so the batch saturates the
-        budget."""
-        budget = max_tokens_per_modality.get("image")
-        if not budget:
+        """Build processed Pixtral tensors for profiler-selected images."""
+        if max_num_encoder_tokens <= 0:
+            raise ValueError("max_num_encoder_tokens must be positive")
+        unsupported_modalities = set(mm_counts) - {"image"}
+        if unsupported_modalities:
+            raise ValueError("Pixtral cannot build dummy data for modalities "
+                             f"{sorted(unsupported_modalities)}")
+        num_images = mm_counts.get("image", 0)
+        if num_images < 0:
+            raise ValueError("Multimodal item counts must be nonnegative; got "
+                             f"{num_images} for image")
+        if num_images == 0:
             return {}
+
         patch, _, _, _ = self._vision_geometry()
-        size = self.get_size_for_max_tokens(max_tokens=budget)
-        tokens_per_image = max(
-            1,
-            self._vit_tokens(width=size["width"],
-                             height=size["height"],
-                             patch=patch))
-        num_images = max(1, budget // tokens_per_image)
-        return self.get_dummy_mm_data_for_size(width=size["width"],
-                                               height=size["height"],
-                                               num_images=num_images,
-                                               dtype=dtype)
+        size = self.get_size_for_max_tokens(max_tokens=max_num_encoder_tokens)
+        tokens_per_image = self._vit_tokens(width=size["width"],
+                                            height=size["height"],
+                                            patch=patch)
+        if num_images * tokens_per_image > max_num_encoder_tokens:
+            raise ValueError("Requested multimodal dummy items exceed "
+                             f"max_num_encoder_tokens={max_num_encoder_tokens}")
+        return self._dummy_mm_data_for_size(
+            width=size["width"],
+            height=size["height"],
+            num_images=num_images,
+            dtype=dtype,
+        )
 
     def get_vocab_size(self) -> int:
         """Return the vocab size of the model."""
@@ -712,6 +786,7 @@ class Mistral3VLM(MultimodalModelMixin, PreTrainedModel):
     """
 
     supports_encoder_cache = True
+    supports_mm_encoder_item_scheduling = True
 
     def __init__(
         self,