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

docs/source/developer-guide/telemetry.md

@@ -93,7 +93,7 @@ unset or when the safety sanitizer rejects the runtime value.
 | `enable_resource_governor` | `<class 'bool'>` | `value` |  |  |
 | `enable_speculative_beam_history_d2h` | `<class 'bool'>` | `value` |  |  |
 | `encode_only` | `<class 'bool'>` | `value` |  |  |
-| `encoder_max_batch_size` | `Optional[int]` | `value` |  |  |
+| `encoder_max_num_items` | `Optional[int]` | `value` |  |  |
 | `encoder_max_num_tokens` | `Optional[int]` | `value` |  |  |
 | `force_dynamic_quantization` | `<class 'bool'>` | `value` |  |  |
 | `garbage_collection_gen0_threshold` | `<class 'int'>` | `value` |  |  |

docs/source/models/supported-models.md

@@ -142,7 +142,10 @@ 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.
+  mixed-modality requests bypass the cache. For models with item-level encoder scheduling, the
+  cache also composes with the item path: cached items complete without consuming the encoder
+  runtime budgets, partially cached requests re-compute only the missing items, and items encoded
+  through the item path populate the cache for later requests.
 
 # Visual Generation Models
 

tensorrt_llm/_torch/models/modeling_gemma4_vision.py

@@ -735,7 +735,7 @@ 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)``
+        # at the encoder ``(encoder_max_num_items, 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.

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, Optional, Sequence, Tuple
 
 import torch
 import torchvision
@@ -47,6 +50,7 @@
                                  MultimodalPlaceholderPlacement, TextPrompt,
                                  register_input_processor)
 from tensorrt_llm.inputs.multimodal import MultimodalParams
+from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata
 from tensorrt_llm.inputs.utils import encode_base64_image
 from tensorrt_llm.llmapi import SamplingParams
 from tensorrt_llm.logger import logger
@@ -354,6 +358,7 @@ def __call__(self, text, images=None, **kwargs):
 
 class Mistral3InputProcessor(BaseMultimodalInputProcessor,
                              BaseMultimodalDummyInputsBuilder):
+    supports_mm_encoder_item_scheduling = True
 
     def __init__(
         self,
@@ -413,6 +418,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
@@ -543,27 +582,48 @@ def get_mm_max_tokens_per_item(self) -> Dict[str, int]:
         edge = max((max_size // unit) * unit, unit)
         return {"image": self._vit_tokens(width=edge, height=edge, patch=patch)}
 
+    def get_mm_encoder_attention_metadata_capacity(
+            self, max_num_items: int,
+            max_num_tokens: int) -> Optional[Dict[str, int]]:
+        """Bound Pixtral contexts by item and physical-token budgets."""
+        _, merge, _, _ = self._vision_geometry()
+        min_tokens_per_image = merge * merge
+        return {
+            "attention":
+            max(1, min(max_num_items, max_num_tokens // min_tokens_per_image))
+        }
+
     def get_dummy_mm_data_for_tokens(
         self,
         *,
         max_tokens_per_modality: Dict[str, int],
+        max_items_per_modality: Optional[Dict[str, int]] = None,
         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."""
+        derived from the realized patch count so the batch stays within the
+        token budget. ``max_items_per_modality`` selects the many-item
+        profiling boundary when provided."""
         budget = max_tokens_per_modality.get("image")
         if not budget:
             return {}
+        target_num_images = (None if max_items_per_modality is None else max(
+            1, max_items_per_modality.get("image", 1)))
+        per_image_budget = (budget if target_num_images is None else max(
+            1, budget // target_num_images))
         patch, _, _, _ = self._vision_geometry()
-        size = self.get_size_for_max_tokens(max_tokens=budget)
+        size = self.get_size_for_max_tokens(max_tokens=per_image_budget)
         tokens_per_image = max(
             1,
             self._vit_tokens(width=size["width"],
                              height=size["height"],
                              patch=patch))
+        if tokens_per_image > budget:
+            return {}
         num_images = max(1, budget // tokens_per_image)
+        if target_num_images is not None:
+            num_images = min(target_num_images, num_images)
         return self.get_dummy_mm_data_for_size(width=size["width"],
                                                height=size["height"],
                                                num_images=num_images,

tensorrt_llm/_torch/models/modeling_multimodal_encoder.py

@@ -1,4 +1,4 @@
-# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+# Copyright 2024-2026 NVIDIA CORPORATION & AFFILIATES
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -24,18 +24,17 @@
 from ..attention_backend.interface import AttentionMetadata
 from .modeling_multimodal_utils import multiscale_forward
 
-# Fallback capacity for the encoder ``AttentionMetadata``'s per-segment buffers
-# (``max_num_requests``). The encoder runs one attention segment per vision tile
-# / window, so the real bound is the segment count, which can far exceed the
-# request count (one image can expand into many tiles/windows).
+# Legacy fallback for encoders that do not yet opt into atomic-item scheduling
+# or provide a model-specific item/token-to-segment mapping. Their encoder
+# forward is not runtime-budgeted, so configured item/token limits cannot
+# safely size fixed ``AttentionMetadata`` per-segment buffers.
 #
-# TODO: Once the scheduler caps an encoder forward at ``encoder_max_num_tokens``,
-# derive this from the token budget instead -- ``encoder_max_num_tokens //
-# min_tokens_per_segment`` (an exact segment bound). We cannot do that today:
-# ``encoder_max_num_tokens`` is not yet enforced (the attention workspace grows
-# past it), so a low value would under-size these fixed buffers, which -- unlike
-# the token workspace -- cannot grow. Until then, fall back to the legacy
-# worst-case capacity.
+# TODO: Replace this fallback as the remaining encoders provide model-specific
+# conversions from atomic item/token budgets to attention segments. The exact
+# conversion requires each encoder's minimum tokens per tile/window/segment, so
+# retain this worst-case capacity for models without that contract.
+# Runtime-scheduled Qwen and Mistral/Pixtral encoders override
+# ``get_encoder_attention_metadata_capacity`` and do not consume this fallback.
 _ENCODER_FALLBACK_MAX_NUM_REQUESTS = 8192
 
 
@@ -44,7 +43,7 @@ class MultimodalEncoderMixin:
 
     Marker + default ``setup_attn_metadata`` for multimodal encoders whose
     ``AttentionMetadata`` is built by ``PyTorchModelEngine`` after model load
-    using runtime sizes (``max_batch_size``, ``max_num_tokens``).
+    using runtime sizes (``max_num_items``, ``max_num_tokens``).
 
     Subclasses set ``metadata_cls`` in their own ``__init__`` (typically from
     ``get_attention_backend(model_config.attn_backend).Metadata``) and either
@@ -54,15 +53,33 @@ class MultimodalEncoderMixin:
     metadata_cls: Type[AttentionMetadata]
     attn_metadata: Optional[AttentionMetadata] = None
 
-    def setup_attn_metadata(self, max_num_requests: int,
-                            max_num_tokens: int) -> None:
-        max_num_requests = max(max_num_requests,
-                               _ENCODER_FALLBACK_MAX_NUM_REQUESTS)
+    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
+        return {
+            "attention": max(max_num_items, _ENCODER_FALLBACK_MAX_NUM_REQUESTS)
+        }
+
+    def setup_attn_metadata(
+        self,
+        max_num_items: int,
+        max_num_tokens: int,
+        attention_metadata_capacity: Optional[dict[str, int]] = None,
+    ) -> None:
+        """Map encoder item/token budgets to attention metadata capacity."""
+        capacities = (attention_metadata_capacity
+                      if attention_metadata_capacity is not None else
+                      self.get_encoder_attention_metadata_capacity(
+                          max_num_items, max_num_tokens))
         self.attn_metadata = self.metadata_cls(
-            max_num_requests=max_num_requests,
+            max_num_requests=capacities["attention"],
             max_num_tokens=max_num_tokens,
             kv_cache_manager=None,
         )
+        # Pin the no-cache ``max_seq_len`` to the startup token budget once
+        # (stable C++ attention-op cache key; see the Qwen setup overrides).
+        self.attn_metadata.max_seq_len = max_num_tokens
 
 
 class VisionTower(nn.Module):