diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index fc0e1fc8bbe1..bbe5f77ec2fc 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -93,7 +93,7 @@ unset or when the safety sanitizer rejects the runtime value. | `enable_resource_governor` | `` | `value` | | | | `enable_speculative_beam_history_d2h` | `` | `value` | | | | `encode_only` | `` | `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` | `` | `value` | | | | `garbage_collection_gen0_threshold` | `` | `value` | | | diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index b03f9961f7cc..68063ec8f6b4 100644 --- a/docs/source/models/supported-models.md +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_gemma4_vision.py b/tensorrt_llm/_torch/models/modeling_gemma4_vision.py index 40c595dfe5de..39829cd77e9c 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4_vision.py +++ b/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. diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 2246a73f8fdd..75e3042d93ae 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/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, diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py index 08e38f195350..65fc4afcb519 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py +++ b/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): diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index b1a2c5df6845..863afb7ef976 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -34,6 +34,7 @@ from tensorrt_llm._torch.tensor_lru_cache import TensorLRUCache from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.inputs.multimodal import MultimodalParams, MultimodalRuntimeData +from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata from tensorrt_llm.logger import logger from .modeling_multimodal_utils import ( @@ -165,6 +166,142 @@ def encode_multimodal_inputs( """ raise NotImplementedError + @staticmethod + def _build_multimodal_encoder_input( + multimodal_param: MultimodalParams, + item_idx: int, + ) -> MultimodalParams: + """Build encoder inputs for one canonical original MM item.""" + multimodal_data = multimodal_param.multimodal_data or {} + item_metadata = get_multimodal_encoder_item_metadata(multimodal_data) + if item_metadata is None: + raise ValueError(f"Missing multimodal item reference for index {item_idx}") + item_refs = item_metadata.item_refs + if item_idx >= len(item_refs): + raise ValueError(f"Missing multimodal item reference for index {item_idx}") + modality, local_idx = item_refs[item_idx] + modality_data = multimodal_data.get(modality) + if modality_data is None: + raise ValueError(f"Missing {modality} encoder data for item {item_idx}") + if not isinstance(modality_data, dict): + raise TypeError(f"{modality} encoder data must be a dict") + + selected_data: Dict[str, Any] = {} + grid_key = ( + "image_grid_thw" + if modality == "image" + else "video_grid_thw" + if modality == "video" + else None + ) + pixel_key = ( + "pixel_values" + if modality == "image" + else "pixel_values_videos" + if modality == "video" + else None + ) + if ( + grid_key is not None + and pixel_key is not None + and grid_key in modality_data + and pixel_key in modality_data + ): + grids = modality_data[grid_key] + if local_idx >= len(grids): + raise ValueError(f"Invalid {modality} item index {local_idx}") + patch_counts = torch.prod(grids, dim=1).tolist() + patch_start = sum(int(count) for count in patch_counts[:local_idx]) + patch_end = patch_start + int(patch_counts[local_idx]) + selected_data[pixel_key] = modality_data[pixel_key][patch_start:patch_end] + selected_data[grid_key] = grids[local_idx : local_idx + 1] + elif modality == "image" and "image_sizes" in modality_data: + image_sizes = modality_data["image_sizes"] + pixel_values = modality_data["pixel_values"] + selected_data["image_sizes"] = image_sizes[local_idx : local_idx + 1] + selected_data["pixel_values"] = pixel_values[local_idx : local_idx + 1] + else: + raise TypeError(f"Unsupported multimodal item layout for {modality}") + + return MultimodalParams(multimodal_data={modality: selected_data}) + + def prepare_multimodal_encoder_inputs( + self, + selected_items: Sequence[tuple[MultimodalParams, int]], + ) -> list[tuple[MultimodalParams, int, str]]: + """Build selected item encoder inputs before the caller performs H2D. + + Args: + selected_items: ``(request params, item index)`` pairs in + scheduler-selected order. Each params object must contain + parallel item references and embedding lengths. + + Returns: + Tuples containing the single-item encoder params, its expected + encoder output row count, and its modality, in input order. + """ + encoder_inputs: list[tuple[MultimodalParams, int, str]] = [] + for multimodal_param, item_idx in selected_items: + multimodal_data = multimodal_param.multimodal_data or {} + item_metadata = get_multimodal_encoder_item_metadata(multimodal_data) + if item_metadata is None: + raise ValueError("MM item metadata is required for item encoding") + item_refs = item_metadata.item_refs + embedding_lengths = item_metadata.output_embedding_lengths + modality = item_refs[item_idx][0] + encoder_inputs.append( + ( + self._build_multimodal_encoder_input(multimodal_param, item_idx), + embedding_lengths[item_idx], + modality, + ) + ) + return encoder_inputs + + def forward_multimodal_encoder_items( + self, + encoder_inputs: Sequence[tuple[MultimodalParams, int, str]], + ) -> list[torch.Tensor]: + """Forward prepared MM encoder inputs in scheduler item order. + + Args: + encoder_inputs: Tuples returned by + :meth:`prepare_multimodal_encoder_inputs`. Consecutive inputs + with the same modality must be batch-compatible. + + Returns: + One encoder output tensor per prepared item. Each tensor has the + declared embedding row count and retains scheduler input order. + """ + outputs: list[torch.Tensor] = [] + group_params: list[MultimodalParams] = [] + group_lengths: list[int] = [] + group_modality: Optional[str] = None + + def flush_group() -> None: + if not group_params: + return + embeddings = self.encode_multimodal_inputs(group_params) + expected_length = sum(group_lengths) + if embeddings.shape[0] != expected_length: + raise ValueError( + f"MM encoder output length {embeddings.shape[0]} does not " + f"match the {expected_length} rows declared by the " + "selected items" + ) + outputs.extend(torch.split(embeddings, group_lengths, dim=0)) + group_params.clear() + group_lengths.clear() + + for multimodal_param, embedding_length, modality in encoder_inputs: + if group_modality is not None and modality != group_modality: + flush_group() + group_modality = modality + group_params.append(multimodal_param) + group_lengths.append(embedding_length) + flush_group() + return outputs + @property def multimodal_token_ids(self) -> Optional[Sequence[int] | torch.Tensor]: """Return placeholder token ids in `input_ids` replaced by MM embeds. @@ -336,10 +473,31 @@ def _get_or_encode_multimodal_embeddings( return embeddings[0] def _get_multimodal_encoder_cache(self) -> Optional[TensorLRUCache]: - """Return the per-model encoder cache, if enabled. + """Return the per-model full-request-path encoder clone cache, if enabled. The cache stores per-item embeddings for params that can be represented by one modality. See `_encoder_cache_keys` for the mixed-modality skip path and its technical limitation. + + Scope: this serves only the full-request consumers (side-stream + prefetch, `mm_encoder_only`/disagg encoding, cache-enabled models + without item scheduling). Item-scheduling models store encoder + outputs in the engine-owned `MultimodalEncoderCacheManager` instead; + for them this cache stays empty in the executor loop (their requests + arrive here with embeddings already attached). The key format is + shared (`_encoder_cache_item_key`) so the two stores can unify later. + + TODO(TRTLLM-14477): unify the full-request consumers onto the + `MultimodalEncoderCacheManager` and retire this clone cache: + (1) inject the manager into the model and switch + `_attach_encoder_cache_hit` to `get_and_hold` and + `_write_encoder_cache_entries` to a clone-adopt `put` (batch-tensor + slices must not retain the whole batch allocation) — this also puts + side-stream prefetch under the byte budget; (2) assemble partial + hits by encoding only the misses through + `prepare_multimodal_encoder_inputs` (resolves TRTLLM-13996 for all + consumers); (3) delete this getter, `_encoder_cache_keys`, + `supports_encoder_cache`, and the legacy reservation branch in + `_reserve_multimodal_encoder_cache_memory`. """ multimodal_config = self.model_config.multimodal_config if multimodal_config is None: @@ -462,12 +620,7 @@ def _encoder_cache_keys( # different request layout; the current request order is restored when cached item tensors # are concatenated below. return [ - ( - modality, - tuple(item_hash), - int(embedding_length), - kwargs_hash, - ) + cls._encoder_cache_item_key(modality, item_hash, embedding_length, kwargs_hash) for item_hash, embedding_length in zip( mm_input.multimodal_hashes, embedding_lengths, @@ -475,6 +628,56 @@ def _encoder_cache_keys( ) ] + @staticmethod + def _encoder_cache_item_key( + modality: str, + item_hash: Sequence[int], + embedding_length: int, + kwargs_hash: str, + ) -> Hashable: + """Build the cache key of one atomic item. + + The single source of the key format: the full-request path + (`_encoder_cache_keys`) and the item-scheduling path + (`build_encoder_cache_item_keys`) must produce identical keys so + entries written by either path hit from the other. + """ + return (modality, tuple(item_hash), int(embedding_length), kwargs_hash) + + @classmethod + def build_encoder_cache_item_keys( + cls, + multimodal_hashes: Optional[Sequence[Sequence[int]]], + item_refs: Sequence[tuple[str, int]], + embedding_lengths: Sequence[int], + kwargs_hash: Optional[str], + ) -> Optional[list[Hashable]]: + """Build per-item cache keys from request-level item metadata. + + Unlike `_encoder_cache_keys`, the modality comes from each item's + `item_refs` entry, so mixed-modality requests are keyable per item. + Returns `None` when the request cannot participate in the cache + (missing hashes or kwargs hash, or item counts that do not line up). + """ + 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( + f"{_MM_ENCODER_CACHE_LOG_NAME}: skipping item keys with " + "mismatched multimodal_hashes, item_refs, and " + "multimodal_embedding_lengths counts" + ) + return None + return [ + cls._encoder_cache_item_key(modality, item_hash, embedding_length, kwargs_hash) + for (modality, _), item_hash, embedding_length in zip( + item_refs, + multimodal_hashes, + embedding_lengths, + strict=True, + ) + ] + @classmethod def _attach_encoder_cache_hit( cls, @@ -497,10 +700,19 @@ def _attach_encoder_cache_hit( for key in keys: cached_embedding = encoder_cache.get(key) if cached_embedding is None: - # TODO(TRTLLM-13996): allow re-computing only the uncached items. + # TODO(TRTLLM-13996): allow re-computing only the uncached items — + # planned via storage unification (TRTLLM-14477): once this path reads the + # `MultimodalEncoderCacheManager`, misses can be encoded per + # item through `prepare_multimodal_encoder_inputs` and + # assembled as a view list (see `_get_multimodal_encoder_cache`). # `get_multimodal_embeddings` treats a param as either fully cached or uncached. # Attaching partial hits would make the later concatenated tensor ambiguous because # there is no placeholder for missing item rows inside `multimodal_embedding`. + # For models with item-level encoder scheduling the executor loop encodes + # exclusively through the item path (pre-scheduling cache-hit attachment plus + # per-item encoding of the misses), so this limitation applies to the remaining + # full-request consumers: side-stream prefetch, mm_encoder_only/disagg encoding, + # and cache-enabled models without item scheduling. logger.debug( f"{_MM_ENCODER_CACHE_LOG_NAME}: cache miss; hit_items={len(cached_embeddings)}," f" total_items={len(keys)}." diff --git a/tensorrt_llm/_torch/models/modeling_pixtral.py b/tensorrt_llm/_torch/models/modeling_pixtral.py index e2ac38dcc5f0..e1d812350d6c 100644 --- a/tensorrt_llm/_torch/models/modeling_pixtral.py +++ b/tensorrt_llm/_torch/models/modeling_pixtral.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from typing import List import torch @@ -205,6 +208,16 @@ def __init__( model_config.attn_backend ).Metadata + def get_encoder_attention_metadata_capacity( + self, max_num_items: int, max_num_tokens: int + ) -> dict[str, int]: + """Conservatively map each item to one Pixtral attention context. + + Mistral3's processor normally injects the tighter merge-tile-aware + token bound into :meth:`setup_attn_metadata`. + """ + return {"attention": max(1, min(max_num_items, max_num_tokens))} + @torch.inference_mode() def forward( self, diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index 2079bfdf3127..7d3d7c5bf8ca 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -32,6 +32,7 @@ from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs, MultimodalParams) +from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata from ..._utils import async_tensor_h2d, nvtx_range, prefer_pinned from ...inputs import (BaseMultimodalDummyInputsBuilder, @@ -69,8 +70,7 @@ from ..modules.gated_mlp import GatedMLP from ..modules.rotary_embedding import MRotaryEmbedding, RotaryEmbedding from .modeling_auto import AutoModelForCausalLM -from .modeling_multimodal_encoder import (_ENCODER_FALLBACK_MAX_NUM_REQUESTS, - MultimodalEncoderMixin) +from .modeling_multimodal_encoder import MultimodalEncoderMixin from .modeling_multimodal_mixin import MultimodalModelMixin from .modeling_multimodal_utils import ( _install_processor_output_validation_filter, find_input_mm_embeds, @@ -114,7 +114,22 @@ def _prepare_qwen_vl_vision_attn_metadata( non_blocking=True) attn_metadata.cu_q_seqlens = cu_seqlens attn_metadata.cu_kv_seqlens = cu_seqlens - attn_metadata.max_seq_len = max(seq_lens) + # Deliberately do NOT update `max_seq_len` to this batch's maximum: + # it feeds the C++ attention-op cache key (mMaxSeqLen / + # mMaxContextLength), so a per-batch value creates and permanently + # caches a new AttentionOp — each with its own cudaMalloc'd workspace + # and semaphores — for every distinct segment length the serving + # traffic produces, growing GPU usage without bound. `max_seq_len` is + # pinned once in `setup_attn_metadata` to the encoder token budget, so + # every batch reuses the single op that warmup already created and + # profiled. Kernels read the real lengths from `seq_lens`. + batch_max_seq_len = max(seq_lens) + if batch_max_seq_len > attn_metadata.max_seq_len: + raise ValueError( + f"Encoder attention segment of {batch_max_seq_len} tokens " + f"exceeds the startup maximum {attn_metadata.max_seq_len}; " + "increase encoder_max_num_tokens or reduce the input " + "max_pixels") # The vision tower runs no-cache, context-only attention and supplies its # own `cu_seqlens` above, so the heavy KV-oriented `prepare()` (kv_lens / # prompt_lens / host_request_types setup) is unnecessary host work. @@ -189,6 +204,7 @@ def _prepare_qwen_vl_mrope_config( class Qwen2VLInputProcessorBase(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): + supports_mm_encoder_item_scheduling = True def __init__(self, model_path: str, @@ -239,6 +255,143 @@ 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 prompt-ordered Qwen vision items and pre-merger costs.""" + grids_by_modality: Dict[str, torch.Tensor] = {} + for modality, grid_key in (("image", "image_grid_thw"), + ("video", "video_grid_thw")): + modality_data = multimodal_data.get(modality) + if isinstance(modality_data, + dict) and modality_data.get(grid_key) is not None: + grids_by_modality[modality] = modality_data[grid_key] + if not grids_by_modality: + return None + + # Qwen2.5's fixed attention metadata is sized from the live + # processor's startup geometry. Validate the processed grids against + # that same contract before the request reaches the scheduler. This + # also covers request-local processor overrides (for example, + # ``min_pixels`` or ``do_resize=False``) without trying to infer their + # effect from the raw kwargs. + if getattr(self.config.vision_config, "window_size", None) is not None: + self._validate_encoder_attention_geometry(grids_by_modality) + + config = self.config + image_token_id = int(config.image_token_id) + video_token_id = int(config.video_token_id) + vision_start_token_id = int(config.vision_start_token_id) + vision_end_token_id = int(config.vision_end_token_id) + + # Bounded `list.index` scans run at C speed, unlike a per-token + # Python loop over the full prompt. + modalities: List[str] = [] + search_start = 0 + prompt_length = len(prompt_token_ids) + while search_start < prompt_length: + try: + item_start = prompt_token_ids.index(vision_start_token_id, + search_start) + except ValueError: + break + try: + item_end = prompt_token_ids.index(vision_end_token_id, + item_start + 1) + except ValueError as error: + raise ValueError( + "Unclosed Qwen vision span in prompt token IDs") from error + + image_pos = None + video_pos = None + try: + image_pos = prompt_token_ids.index(image_token_id, + item_start + 1, item_end) + except ValueError: + pass + try: + video_pos = prompt_token_ids.index(video_token_id, + item_start + 1, item_end) + except ValueError: + pass + + if image_pos is not None and video_pos is not None: + raise ValueError( + "Qwen vision span contains both image and video tokens") + if image_pos is not None: + modalities.append("image") + elif video_pos is not None: + modalities.append("video") + else: + raise ValueError( + "Qwen vision span contains no image or video token") + + search_start = item_end + 1 + + video_grids = grids_by_modality.get("video") + num_video_spans = modalities.count("video") + video_spans_per_item: List[int] = [] + if video_grids is not None: + if num_video_spans == len(video_grids): + video_spans_per_item = [1] * len(video_grids) + else: + temporal_spans = [int(grid[0]) for grid in video_grids] + if num_video_spans != sum(temporal_spans): + raise ValueError( + "Prompt video spans do not match processed video grids") + video_spans_per_item = temporal_spans + + canonical_modalities: List[str] = [] + span_idx = 0 + video_idx = 0 + while span_idx < len(modalities): + modality = modalities[span_idx] + canonical_modalities.append(modality) + if modality == "image": + span_idx += 1 + continue + if video_idx >= len(video_spans_per_item): + raise ValueError( + "Prompt contains more video items than encoder grids") + spans = video_spans_per_item[video_idx] + if modalities[span_idx:span_idx + spans] != ["video"] * spans: + raise ValueError( + "One original video must occupy consecutive prompt spans") + span_idx += spans + video_idx += 1 + if video_idx != len(video_spans_per_item): + raise ValueError( + "Processed video grids contain items absent from the prompt") + + local_indices = {"image": 0, "video": 0} + item_refs: List[Tuple[str, int]] = [] + encoder_token_lengths: List[int] = [] + for modality in canonical_modalities: + local_idx = local_indices[modality] + grids = grids_by_modality.get(modality) + if grids is None or local_idx >= len(grids): + raise ValueError( + f"Prompt contains more {modality} items than encoder grids") + item_refs.append((modality, local_idx)) + encoder_token_lengths.append( + int(torch.prod(grids[local_idx]).item())) + local_indices[modality] += 1 + + expected_items = sum(len(grids) for grids in grids_by_modality.values()) + if len(item_refs) != expected_items: + raise ValueError("Prompt multimodal item order does not match " + "processed Qwen encoder grids") + output_embedding_lengths = [ + token_length // self.spatial_merge_unit + 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) + # ------------------------------------------------------------------ # Deterministic dummy-input sizing for multimodal profiling. # @@ -433,30 +586,157 @@ def get_mm_max_tokens_per_item(self) -> Dict[str, int]: self._num_vision_tokens(width=size["width"], height=size["height"]) } + def _post_merge_frame_area_bounds(self) -> Tuple[int, int]: + """Return startup frame-area bounds in merged-grid units. + + Images and videos may use distinct HF processors. Qwen2/2.5 applies + the pixel clamp independently to every temporal frame, so use the + least restrictive minimum and greatest maximum across both. + """ + cfg = self.config.vision_config + factor = cfg.patch_size * cfg.spatial_merge_size + pixel_bounds = [] + for processor_name in ("image_processor", "video_processor"): + processor = getattr(self.processor, processor_name, None) + size = getattr(processor, "size", None) + if size is None: + continue + try: + min_pixels = int(size["shortest_edge"]) + max_pixels = int(size["longest_edge"]) + except (KeyError, TypeError, ValueError): + continue + if min_pixels > 0 and max_pixels >= min_pixels: + pixel_bounds.append((min_pixels, max_pixels)) + if not pixel_bounds: + pixel_bounds.append(self._vision_pixel_bounds()) + + min_pixels = min(bounds[0] for bounds in pixel_bounds) + max_pixels = max(bounds[1] for bounds in pixel_bounds) + factor_area = factor * factor + min_area = max(1, math.ceil(min_pixels / factor_area)) + max_area = max(min_area, max_pixels // factor_area) + return min_area, max_area + + @staticmethod + @lru_cache(maxsize=32) + def _max_window_ratio(min_area: int, max_area: int, + window_side: int) -> Tuple[int, int]: + """Return ``(windows, merged_tokens)`` with maximum ratio. + + A resized Qwen frame has positive integer merged-grid dimensions + ``(h, w)`` and ``min_area <= h*w <= max_area``. Enumerating that + bounded domain gives a tight ratio upper bound for both images and + repeated video frames without assuming the runtime aspect ratio. + """ + best_windows = 1 + best_area = min_area + for grid_h in range(1, max_area + 1): + min_grid_w = max(1, math.ceil(min_area / grid_h)) + 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) * + math.ceil(grid_w / window_side)) + if windows * best_area > best_windows * area: + best_windows = windows + best_area = area + return best_windows, best_area + + def _validate_encoder_attention_geometry( + self, grids_by_modality: Dict[str, torch.Tensor]) -> None: + """Validate runtime grids against startup metadata-sizing bounds.""" + cfg = self.config.vision_config + merge_size = cfg.spatial_merge_size + min_area, max_area = self._post_merge_frame_area_bounds() + window_side = max(1, cfg.window_size // merge_size // cfg.patch_size) + max_windows, ratio_area = self._max_window_ratio( + min_area, max_area, window_side) + + for modality, grids in grids_by_modality.items(): + for grid in grids: + _, grid_h, grid_w = (int(value) for value in grid) + if grid_h % merge_size or grid_w % merge_size: + raise ValueError( + f"Processed Qwen {modality} grid ({grid_h}, {grid_w}) " + f"must be divisible by spatial_merge_size={merge_size}") + merged_h = grid_h // merge_size + merged_w = grid_w // merge_size + area = merged_h * merged_w + windows = (math.ceil(merged_h / window_side) * + math.ceil(merged_w / window_side)) + if (area < min_area + or windows * ratio_area > max_windows * area): + raise ValueError( + f"Processed Qwen {modality} frame geometry " + f"({grid_h}, {grid_w}) exceeds the encoder attention " + "metadata capacity derived from the startup processor " + "geometry. Configure the processor geometry at startup " + "instead of lowering it per request.") + + def get_mm_encoder_attention_metadata_capacity( + self, max_num_items: int, + max_num_tokens: int) -> Optional[Dict[str, int]]: + """Bound Qwen2.5 full/window segments using processor geometry. + + ``max_num_items`` is intentionally not part of the bound: one atomic + video item may contain enough temporal segments to consume the whole + token budget. + """ + cfg = self.config.vision_config + window_size = getattr(cfg, "window_size", None) + if window_size is None: + return None + + merge_unit = cfg.spatial_merge_size * cfg.spatial_merge_size + min_area, max_area = self._post_merge_frame_area_bounds() + min_frame_tokens = merge_unit * min_area + full_capacity = max(1, max_num_tokens // min_frame_tokens) + + window_side = max( + 1, window_size // cfg.spatial_merge_size // cfg.patch_size) + windows, area = self._max_window_ratio(min_area, max_area, window_side) + window_capacity = max(1, + max_num_tokens * windows // (merge_unit * area)) + return { + "full_attention": full_capacity, + "window_attention": window_capacity, + } + 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: Optional[torch.dtype] = None, ) -> Dict[str, Any]: """Vision implementation of the modality-agnostic profiler entry: for the ``"image"`` budget, pick the worst-case image size, fill it with identical copies, and materialize the encoder tensors. + ``max_items_per_modality`` selects the many-item profiling boundary. ``num_images`` is computed from the *realized* token count of the chosen - size (which the ``max_pixels`` cap may make smaller than the budget), so - the batch saturates the budget rather than assuming one image fills it. + size (which processor pixel bounds may clamp), so the batch stays within + both the item and token budgets. """ max_tokens = max_tokens_per_modality.get("image") if not max_tokens: return {} - size = self.get_size_for_max_tokens(max_tokens=max_tokens) + target_num_images = (None if max_items_per_modality is None else max( + 1, max_items_per_modality.get("image", 1))) + per_image_budget = (max_tokens if target_num_images is None else max( + 1, max_tokens // target_num_images)) + size = self.get_size_for_max_tokens(max_tokens=per_image_budget) tokens_per_image = max( 1, self._num_vision_tokens(width=size["width"], height=size["height"], num_frames=size.get("num_frames", 1))) + if tokens_per_image > max_tokens: + return {} num_images = max(1, max_tokens // 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_frames=size.get( @@ -1408,31 +1688,54 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): None, persistent=False) - def setup_attn_metadata(self, max_num_requests: int, - max_num_tokens: int) -> None: + def setup_attn_metadata( + self, + max_num_items: int, + max_num_tokens: int, + attention_metadata_capacity: Optional[Dict[str, int]] = None, + ) -> None: # Override: Qwen2/2.5-VL uses two metadata objects (full + window # attention) instead of the mixin's single ``attn_metadata``. # - # Windowed attention splits each image into many attention sequences - # (one per window grid cell), so ``max_num_requests`` here is the - # **window** count, not the image count, and can far exceed the - # LLM-side ``max_batch_size`` that ``encoder_max_batch_size`` falls back - # to. Floor it at the same legacy fallback the mixin uses (see the TODO - # there: derive from ``encoder_max_num_tokens`` once the scheduler caps - # encoder forwards at it). - max_num_requests = max(max_num_requests, - _ENCODER_FALLBACK_MAX_NUM_REQUESTS) - kwargs = dict(max_num_requests=max_num_requests, - max_num_tokens=max_num_tokens, - kv_cache_manager=None) - self.full_attn_metadata = self.metadata_cls(**kwargs) - self.window_attn_metadata = self.metadata_cls(**kwargs) + capacities = (attention_metadata_capacity + if attention_metadata_capacity is not None else + self.get_encoder_attention_metadata_capacity( + max_num_items, max_num_tokens)) + kwargs = dict(max_num_tokens=max_num_tokens, kv_cache_manager=None) + self.full_attn_metadata = self.metadata_cls( + max_num_requests=capacities["full_attention"], **kwargs) + self.window_attn_metadata = self.metadata_cls( + max_num_requests=capacities["window_attention"], **kwargs) + # Pin the no-cache `max_seq_len` to the startup token budget once. + # It participates in the C++ attention-op cache key, so keeping it + # constant means one op (created during warmup, its workspace part + # of the profiled peak) serves every runtime batch instead of a new + # cudaMalloc'd op per distinct segment length; no encoder segment + # can exceed the budget (an atomic item must fit it, enforced at + # admission). + self.full_attn_metadata.max_seq_len = max_num_tokens + self.window_attn_metadata.max_seq_len = max_num_tokens # Size the vision-block ``rope_position_ids`` scratch to the encoder # token budget; ``forward`` grows it on the rare miss above the budget. self._rope_position_ids_buffer = torch.arange(max_num_tokens, dtype=torch.int32, device=self.device) + def get_encoder_attention_metadata_capacity( + self, max_num_items: int, max_num_tokens: int) -> Dict[str, int]: + """Conservatively bound Qwen sequences without processor geometry. + + ``max_num_items`` is intentionally ignored because one atomic video + item can expand to multiple temporal and window-attention sequences. + The normal model-engine path injects the tighter live-processor bound + into :meth:`setup_attn_metadata` instead. + """ + max_num_sequences = max(1, max_num_tokens // self.spatial_merge_unit) + return { + "full_attention": max_num_sequences, + "window_attention": max_num_sequences, + } + def get_rotary_pos_emb_window_data( self, grid_rows: List[List[int]] ) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 2ae083d77355..8a95bf8144bc 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -265,6 +265,20 @@ def _build_temporal_block( # `Qwen2VLInputProcessorBase` -- the grid math and the HF `smart_resize` # it defers to are identical for Qwen3-VL. + def get_mm_encoder_attention_metadata_capacity( + self, max_num_items: int, max_num_tokens: int + ) -> Optional[Dict[str, int]]: + """Bound temporal segments using Qwen3's hard geometry minimum. + + ``max_num_items`` is intentionally not part of the bound: one atomic + video may contain enough temporal segments to consume the token + budget. Unlike Qwen2.5, Qwen3 clamps the aggregate temporal pixel + volume, so a long video can shrink each frame to one merged cell. + """ + cfg = self.config.vision_config + merge_unit = cfg.spatial_merge_size * cfg.spatial_merge_size + return {"attention": max(1, max_num_tokens // merge_unit)} + @classmethod def get_rope_index( cls, @@ -801,22 +815,36 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): def device(self) -> torch.device: return self.patch_embed.proj.weight.device - def setup_attn_metadata(self, max_num_requests: int, max_num_tokens: int) -> None: + def setup_attn_metadata( + self, + max_num_items: int, + max_num_tokens: int, + attention_metadata_capacity: Optional[Dict[str, int]] = None, + ) -> None: # Override the mixin default: each image / video frame is its own # attention segment (``seq_lens.extend([h * w] * t)`` in ``forward``), # so a single multi-image or video request can produce many more - # segments than ``max_batch_size``. The number of segments in one - # encoder forward is bounded by the token budget (every segment holds - # at least one token), NOT by the request count -- so floor the - # metadata's request capacity at ``max_num_tokens`` to keep the - # per-request buffers (prompt_lens / host_request_types / kv_lens) from - # overflowing when ``num_contexts`` is set to the segment count. - max_num_requests = max(max_num_requests, max_num_tokens) + # segments than ``max_num_items``. The number of segments in one + # encoder forward is bounded by the token budget and each segment + # contains at least ``spatial_merge_unit`` physical tokens. Use the + # processor/model capacity contract to keep the per-request buffers + # (prompt_lens / host_request_types / kv_lens) from overflowing when + # ``num_contexts`` is set to the segment count. + 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: + # it participates in the C++ attention-op cache key, so a stable + # value means one warmup-created op serves every runtime batch + # instead of a new cudaMalloc'd op per distinct segment length. + self.attn_metadata.max_seq_len = max_num_tokens # Pre-allocate the vision-block ``rope_position_ids`` as an ``arange`` # sized to the encoder's ``max_num_tokens`` (engine-driven) so per-call # code just slices ``[:seq_len]`` instead of allocating a fresh @@ -826,6 +854,19 @@ def setup_attn_metadata(self, max_num_requests: int, max_num_tokens: int) -> Non max_num_tokens, dtype=torch.int32, device=self.device ) + def get_encoder_attention_metadata_capacity( + self, max_num_items: int, max_num_tokens: int + ) -> Dict[str, int]: + """Conservatively bound temporal segments from the token budget. + + ``max_num_items`` is intentionally ignored because one atomic video + item can expand to multiple temporal attention segments. Qwen3's + processor-derived capacity intentionally resolves to this same hard + geometry bound because long videos can reach one merged cell per + segment. + """ + return {"attention": max(1, max_num_tokens // self.spatial_merge_unit)} + @staticmethod @lru_cache(maxsize=1024) def rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor: @@ -1236,6 +1277,24 @@ def encode_multimodal_inputs( ) return mm_embeds[0] + @property + def embedding_dim(self) -> int: + """Width of each encoder output row (`MultimodalModelMixin` contract). + + Qwen3-VL folds the deepstack feature maps into the hidden dim of + every embedding row, so a row is wider than the text hidden size — + byte accounting that assumed `hidden_size` would undercount 4x. + """ + # `self.config` is narrowed to the text sub-config during init, and + # `deepstack_num_level` is cached there as a plain attribute (0 when + # the checkpoint has no deepstack). + text_config = getattr(self.config, "text_config", self.config) + return text_config.hidden_size * (1 + self.deepstack_num_level) + + @property + def embedding_dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + def _check_and_adjust_experts_implementation(self, *args, **kwargs): """No-op override. diff --git a/tensorrt_llm/_torch/models/modeling_radio.py b/tensorrt_llm/_torch/models/modeling_radio.py index ab9842c4f8c0..bce4b65f8e9d 100644 --- a/tensorrt_llm/_torch/models/modeling_radio.py +++ b/tensorrt_llm/_torch/models/modeling_radio.py @@ -21,8 +21,8 @@ interface as attention_interface from tensorrt_llm._torch.attention_backend import utils as attention_utils from tensorrt_llm._torch.models import modeling_utils -from tensorrt_llm._torch.models.modeling_multimodal_encoder import ( - _ENCODER_FALLBACK_MAX_NUM_REQUESTS, MultimodalEncoderMixin) +from tensorrt_llm._torch.models.modeling_multimodal_encoder import \ + MultimodalEncoderMixin from tensorrt_llm._torch.models.multimodal_encoder_graph import ( EncoderGraphKey, EncoderGraphTensorSpec, EncoderMetadataProvider, MultimodalEncoderGraphRunner) @@ -759,7 +759,7 @@ def __init__( # 8192 requests, `model_config.max_num_tokens`) so the CUDA-graph runner # / metadata provider can read `self.attn_metadata` before the engine # re-sizes it via `setup_attn_metadata`. - self.setup_attn_metadata(max_num_requests=8192, + self.setup_attn_metadata(max_num_items=8192, max_num_tokens=model_config.max_num_tokens) # CUDA-graph runner for the block stack. None means graphs are off and @@ -769,7 +769,7 @@ def __init__( # order. self._blocks_graph_runner: Optional[MultimodalEncoderGraphRunner] = None - def setup_attn_metadata(self, max_num_requests: int, + def setup_attn_metadata(self, max_num_items: int, max_num_tokens: int) -> None: # Override the default to add `kv_layout="NHD"` for FlashInfer. # FlashInfer's original default is "NHD"; TRT-LLM switched the default @@ -778,11 +778,11 @@ def setup_attn_metadata(self, max_num_requests: int, # always in NHD format ([tokens, heads, dim]). # Floor the request capacity at the same legacy fallback as the mixin # default: one attention segment per image, which can exceed the - # LLM-side `max_batch_size` that `encoder_max_batch_size` falls back to. - max_num_requests = max(max_num_requests, - _ENCODER_FALLBACK_MAX_NUM_REQUESTS) + # atomic-item budget derived from `encoder_max_num_items`. + capacities = self.get_encoder_attention_metadata_capacity( + max_num_items, max_num_tokens) metadata_kwargs = dict( - max_num_requests=max_num_requests, + max_num_requests=capacities["attention"], max_num_tokens=max_num_tokens, kv_cache_manager=None, ) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..922873db3d16 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -68,7 +68,8 @@ from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, TRTLLMSampler) from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, - KVCacheV2Scheduler, SimpleScheduler, + KVCacheV2Scheduler, MultimodalEagerEncoderScheduler, + MultimodalScheduler, SimpleScheduler, SimpleUnifiedScheduler) from .seq_slot_manager import SeqSlotManager @@ -331,7 +332,7 @@ def __init__( self._max_batch_size = max_batch_size self._net_max_seq_len = net_max_seq_len self._dummy_reqs = None - self._dummy_encoder_inputs: List[MultimodalParams] = [] + self._dummy_encoder_inputs: List[List[MultimodalParams]] = [] self._profiling_stage_data = profiling_stage_data self._is_disagg = is_disagg self._cache_transceiver_config = llm_args.cache_transceiver_config @@ -557,10 +558,12 @@ def _create_dummy_context_requests( requests = requests * self._mapping.tp_size return requests - def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]: - """Build the worst-case dummy multimodal batch for direct encoder - profiling: the processor's worst-case dummy saturating - ``encoder_max_num_tokens`` (one batched encoder forward), staged to GPU. + def _create_dummy_encoder_inputs(self) -> List[List[MultimodalParams]]: + """Build boundary-shape MM batches for direct encoder profiling. + + The same aggregate token budget can have different peaks as one long + item or as many shorter items. Build both shapes on CPU; the encoder + profiling path stages and releases one batch at a time. Returns an empty list when the model is not a multimodal-encoder model or the processor has no dummy builder (then the encoder is not profiled @@ -587,13 +590,12 @@ def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]: if getattr(self._model_engine.model, "mm_encoder", object()) is None: return [] input_processor = self._model_engine.input_processor - _, encoder_max_num_tokens = self._llm_args.get_encoder_runtime_sizes() + encoder_max_num_tokens = self._model_engine.encoder_max_num_tokens # Modality-agnostic: the model declares each modality's per-item token # demand; split the shared ``encoder_max_num_tokens`` budget across them # in proportion to that demand (they share one encoder microbatch cap, so # the shares sum to the budget rather than each claiming all of it). The - # processor then materializes a dummy per modality, merged into one batch - # so a single encode profiles the combined peak. Empty demand / + # processor then materializes a dummy per modality. Empty demand / # NotImplementedError on the builder → text-only dummy fallback. demand = input_processor.get_mm_max_tokens_per_item() total_demand = sum(demand.values()) @@ -603,41 +605,105 @@ def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]: m: max(1, encoder_max_num_tokens * d // total_demand) for m, d in demand.items() } - try: - multimodal_data = input_processor.get_dummy_mm_data_for_tokens( - max_tokens_per_modality=max_tokens_per_modality, - dtype=self._model_engine.model.dtype) - except NotImplementedError: - return [] - if not multimodal_data: - return [] - params = MultimodalParams(multimodal_data=multimodal_data) - params.to_device("multimodal_data", - "cuda", - pin_memory=prefer_pinned(), - target_keywords=getattr( - self._model_engine.model, - "multimodal_data_device_paths", None)) - return [params] + def distribute_items(max_num_items: int) -> Dict[str, int]: + """Split one shared item budget proportionally across modalities.""" + allocations = { + modality: max_num_items * item_demand // total_demand + for modality, item_demand in demand.items() + } + remaining = max_num_items - sum(allocations.values()) + order = sorted( + demand, + key=lambda modality: (max_num_items * demand[modality] % + total_demand, demand[modality], modality), + reverse=True, + ) + for modality in order[:remaining]: + allocations[modality] += 1 + return allocations + + # Two boundary batches share the same token budget but peak + # differently: one maximal-length item stresses length-dominated + # attention scratch, while the maximal item count stresses per-item + # overheads. Profile both; skip the second when it degenerates to + # the first. + 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)) + + batches: List[List[MultimodalParams]] = [] + for max_items_per_modality in item_limits: + try: + multimodal_data = input_processor.get_dummy_mm_data_for_tokens( + max_tokens_per_modality=max_tokens_per_modality, + max_items_per_modality=max_items_per_modality, + dtype=self._model_engine.model.dtype) + except NotImplementedError: + logger.info( + "MM encoder warmup skipped: %s does not implement " + "get_dummy_mm_data_for_tokens(); memory estimation will " + "not include encoder peak usage.", + type(input_processor).__name__) + return [] + if multimodal_data: + batches.append( + [MultimodalParams(multimodal_data=multimodal_data)]) + return batches def _encode_dummy_inputs(self): - """Run the multimodal encoder(s) once on the pre-built worst-case dummy - batch (``self._dummy_encoder_inputs``) and return its output so the - embeddings stay resident while the peak is measured (the caller must hold - the returned tensors so the peak accounts for the live embeddings). - Returns ``None`` when direct profiling does not apply.""" + """Profile long-item and many-item MM boundary batches sequentially. + + Inputs are staged immediately before each forward and released before + the next. The final (many-item, when present) output remains live for + the subsequent LLM dummy forward; supported models preserve total + output size for equal aggregate encoder-token budgets. + """ if not self._dummy_encoder_inputs: return None with torch.inference_mode(): - return self._model_engine.model.encode_multimodal_inputs( - self._dummy_encoder_inputs) + retained_output = None + for batch_idx, encoder_inputs in enumerate( + self._dummy_encoder_inputs): + try: + for encoder_input in encoder_inputs: + encoder_input.to_device( + "multimodal_data", + "cuda", + pin_memory=prefer_pinned(), + target_keywords=getattr( + self._model_engine.model, + "multimodal_data_device_paths", None), + ) + output = self._model_engine.model.encode_multimodal_inputs( + encoder_inputs) + finally: + encoder_inputs.clear() + if batch_idx + 1 == len(self._dummy_encoder_inputs): + retained_output = output + else: + del output + return retained_output def _reserve_multimodal_encoder_cache_memory( self, peak_memory: int, ) -> int: - """Reserve encoder-cache capacity when the model implements that cache.""" + """Reserve encoder-output storage capacity in the estimated peak. + + Item-scheduling models store every encoder output in the + engine-owned `MultimodalEncoderCacheManager`, so reserving its + (min-clamped) byte budget bounds ALL runtime encoder-output + residency: together with the boundary-batch transient profiled by + `_run_dummy_encoder_forwards`, runtime cannot exceed the estimated + peak. Other models keep the legacy full-request clone cache + reservation. + """ + manager = getattr(self._model_engine, "mm_encoder_cache_manager", None) + if manager is not None: + return peak_memory + manager.max_bytes + model = self._model_engine.model if (not isinstance(model, MultimodalModelMixin) or not model.supports_encoder_cache): @@ -2389,6 +2455,12 @@ def create_py_executor_instance( resources[ResourceManagerType.SEQ_SLOT_MANAGER] = SeqSlotManager( max_num_sequences) + if getattr(model_engine, "mm_encoder_cache_manager", None) is not None: + # Route MM encoder-output teardown through the standard + # free_resources funnel (completion, cancellation, and error paths). + resources[ResourceManagerType.MM_ENCODER_CACHE_MANAGER] = ( + model_engine.mm_encoder_cache_manager) + # Register the compression manager (if one is configured) with the other # managers, before building ResourceManager, so it is part of the manager # set from the start. Reads its own config, not the sparse-attention one. @@ -2526,6 +2598,31 @@ def create_py_executor_instance( reorder_policy_config.policy_args.agent_inflight_seq_num) scheduler = SimpleScheduler(capacity_scheduler, mb_scheduler) + if getattr(model_engine, "supports_mm_encoder_item_scheduling", False): + # Wrap the LLM scheduler with atomic MM item budgeting. The + # side-stream conflict is checked here rather than in an args + # validator because item scheduling is a model capability + # (MultimodalModelMixin + processor opt-in) only known after model + # load; side-stream prefetch stays valid for other MM models. + multimodal_config = llm_args.multimodal_config + if multimodal_config.encoder_side_stream_max_ahead > 0: + raise NotImplementedError( + "MM side-stream prefetch is not yet compatible with " + "item-level encoder scheduling; set " + "multimodal_config.encoder_side_stream_max_ahead to 0") + scheduler_cls = MultimodalScheduler + if multimodal_config.enable_eager_encoder_scheduling: + logger.info("Eager multimodal encoder scheduling is enabled for " + "capacity-rejected active requests.") + scheduler_cls = MultimodalEagerEncoderScheduler + scheduler = scheduler_cls( + scheduler, + max_num_items=model_engine.encoder_max_num_items, + max_num_tokens=model_engine.encoder_max_num_tokens, + cache_manager=model_engine.mm_encoder_cache_manager, + embedding_row_bytes=model_engine.mm_embedding_row_bytes, + ) + config = model_engine.model.model_config.pretrained_config attention_type = AttentionTypeCpp.MLA if is_mla( config) else AttentionTypeCpp.DEFAULT diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 88513bc326d6..29c0f3ddb1af 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1,5 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Python extensions for executor requests.""" + from copy import copy, deepcopy from dataclasses import dataclass, field +from enum import Enum, auto from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast import torch @@ -9,6 +14,8 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.bindings import executor as tllm_executor from tensorrt_llm.executor.result import SimpleTokenLogprobs, TokenLogprobs +from tensorrt_llm.inputs.multimodal import strip_mm_encoder_inputs +from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata from tensorrt_llm.sampling_params import LogprobMode SamplingConfig = tensorrt_llm.bindings.SamplingConfig @@ -46,6 +53,131 @@ # request ids do not use 0; disaggregated global ids use a separate high range. ATTENTION_DP_DUMMY_REQUEST_ID = 0 + +class MultimodalEncoderProgress(Enum): + """Python-only progress derived from request-local MM item outputs.""" + + PENDING = auto() + """No MM item of the request has an encoder output yet.""" + + PARTIAL = auto() + """Some but not all items are encoded; the request must continue on the + item-scheduling path until every item has an output.""" + + READY = auto() + """No further encoder work is needed: every item is encoded, a + precomputed embedding was supplied, or the request has no MM payload.""" + + +@dataclass +class MultimodalEncoderRequestState: + """Per-item MM encoder bookkeeping owned by one request. + + Created at admission by `initialize_multimodal_encoder_request` for + requests whose raw MM payloads run through the item scheduler; the + request's `py_mm_encoder_state` is `None` otherwise. Item slots are + recorded from two sources — fresh encoder outputs adopted into the + `MultimodalEncoderCacheManager` and manager hits — through the single + `record()` writer, so validation cannot diverge between them. + + The state holds no storage of its own: every slot aliases a + manager-owned entry that the request holds until its prefill consumes + the embedding. Holds are released through the executor's teardown + funnel (`release_holds`), never by this object; held entries cannot + be evicted, so a recorded slot cannot regress. + + The state is rank-local and never crosses a serialization boundary: + schedule distribution carries request/item IDs only, and every rank + constructs its own state at admission. + """ + + embedding_lengths: List[int] + """Declared embedding row count of each atomic item, in prompt order. + `record()` validates incoming views against these.""" + + outputs: List[Optional[torch.Tensor]] + """One slot per atomic item, in prompt order. ``None`` until the item is + recorded; then an alias of a held ``MultimodalEncoderCacheManager`` + entry.""" + + @classmethod + def from_embedding_lengths( + cls, + embedding_lengths: List[int]) -> "MultimodalEncoderRequestState": + return cls(embedding_lengths=list(embedding_lengths), + outputs=[None] * len(embedding_lengths)) + + def __post_init__(self) -> None: + if len(self.embedding_lengths) != len(self.outputs): + raise ValueError("MM encoder embedding lengths must have exactly " + "one entry per item slot") + + @property + def num_items(self) -> int: + return len(self.outputs) + + @property + def progress(self) -> MultimodalEncoderProgress: + if all(output is not None for output in self.outputs): + return MultimodalEncoderProgress.READY + if any(output is not None for output in self.outputs): + return MultimodalEncoderProgress.PARTIAL + return MultimodalEncoderProgress.PENDING + + def pending_item_indices(self) -> List[int]: + """Indices of items that still need an encoder output, prompt order.""" + return [ + item_idx for item_idx, output in enumerate(self.outputs) + if output is None + ] + + def record(self, item_idx: int, output_view: torch.Tensor) -> None: + """Record one item's manager-owned output view into its slot. + + ``output_view`` must come from the `MultimodalEncoderCacheManager` + (`adopt()` for fresh encoder outputs, `get_and_hold()` for hits) so + it is already held on this request's behalf; no copy happens + here. Raises when the view does not match the item's declared row + count, when the item was already recorded, or when it disagrees + with previously recorded items on trailing shape/dtype/device + (items of one request must concatenate cleanly at fuse time). + """ + expected_rows = self.embedding_lengths[item_idx] + if output_view.shape[0] != expected_rows: + raise ValueError( + f"MM item {item_idx} produced {output_view.shape[0]} " + f"embeddings; expected {expected_rows}") + if self.outputs[item_idx] is not None: + raise ValueError( + f"MM item {item_idx} was already recorded; items are " + "encoded at most once per request") + reference = next((o for o in self.outputs if o is not None), None) + if reference is not None and ( + reference.shape[1:] != output_view.shape[1:] or reference.dtype + != output_view.dtype or reference.device != output_view.device): + raise ValueError( + "MM encoder items for one request must have matching " + "output shape, dtype, and device") + self.outputs[item_idx] = output_view + + def finalize_into(self, multimodal_data: Dict[str, Any]) -> bool: + """Publish the item views once every slot is recorded. + + Attaches the prompt-ordered list of held item views as + ``multimodal_embedding`` and drops the raw pre-encoder inputs. The + views stay manager-owned; the per-request contiguous embedding is + materialized lazily inside the prefill forward + (`get_multimodal_embeddings`), so between encode and prefill the + only GPU residency is the manager's accounted entries. No-op + returning ``False`` while any slot is still pending. + """ + if not self.outputs or any(output is None for output in self.outputs): + return False + multimodal_data["multimodal_embedding"] = list(self.outputs) + strip_mm_encoder_inputs(multimodal_data) + return True + + if TYPE_CHECKING: from .sampler.sampling_utils import Strategy @@ -678,6 +810,11 @@ def __init__( # Multimodal data self.py_multimodal_data = kwargs.pop("py_multimodal_data", None) self.py_mm_item_order = kwargs.pop("py_mm_item_order", None) + # Per-item MM encoder state, created at admission by + # `initialize_multimodal_encoder_request` when raw MM payloads must + # run through the item scheduler; `None` otherwise (also the + # request-kind signal: state presence == item-scheduling request). + self.py_mm_encoder_state: Optional[MultimodalEncoderRequestState] = None encoder_input_tokens = kwargs.get("encoder_input_tokens") encoder_output_len = kwargs.get("encoder_output_len") return_encoder_output = bool(kwargs.get("return_encoder_output", False)) @@ -1068,6 +1205,92 @@ def get_multimodal_embedding_lengths( return multimodal_embedding_lengths +def get_multimodal_encoder_token_lengths( + request: LlmRequest) -> Optional[List[int]]: + """Return per-item physical encoder attention-token costs. + + Field-level validation happens inside + `get_multimodal_encoder_item_metadata`; only the request-level + cross-check against `multimodal_embedding_lengths` lives here. + """ + item_metadata = get_multimodal_encoder_item_metadata( + request.py_multimodal_data) + if item_metadata is None: + return None + embedding_lengths = get_multimodal_embedding_lengths(request) + if embedding_lengths is None: + raise ValueError("Multimodal encoder item scheduling requires " + "multimodal_embedding_lengths") + if item_metadata.output_embedding_lengths != embedding_lengths: + raise ValueError("Multimodal encoder item output lengths must match " + "multimodal_embedding_lengths") + return item_metadata.encoder_token_lengths + + +def initialize_multimodal_encoder_request( + request: LlmRequest, + max_num_tokens: int, + *, + max_output_bytes: Optional[int] = None, + embedding_row_bytes: int = 0) -> None: + """Initialize immutable request kind and mutable per-item encoder state. + + Raises `ValueError` (failing only this request) when the request can + never execute under the startup guarantees: an atomic item larger than + the effective encoder token budget, or — when the encoder-output + storage budget is supplied — a total embedding footprint that could + never fit `multimodal_config.encoder_cache_max_bytes`. The latter is + ordinarily unreachable because prompts are bounded by `max_num_tokens`, + but LLM chunked prefill admits longer prompts whose full embedding must + stay resident across chunks. + """ + mm_data = request.py_multimodal_data + has_raw_payload = isinstance(mm_data, dict) and any( + isinstance(mm_data.get(modality), dict) + for modality in ("image", "video", "audio")) + has_full_embedding = (isinstance(mm_data, dict) + and mm_data.get("multimodal_embedding") is not None) + token_lengths = get_multimodal_encoder_token_lengths( + request) if has_raw_payload and not has_full_embedding else None + if token_lengths is not None: + largest = max(token_lengths, default=0) + if largest > max_num_tokens: + raise ValueError( + f"Multimodal item requires {largest} encoder tokens, exceeding " + f"the effective startup maximum {max_num_tokens}") + + request.py_mm_encoder_state = None + if token_lengths is not None: + embedding_lengths = get_multimodal_embedding_lengths(request) + if embedding_lengths is None: + raise ValueError("Multimodal item scheduling requires " + "multimodal_embedding_lengths") + if max_output_bytes is not None and embedding_row_bytes > 0: + total_bytes = sum(embedding_lengths) * embedding_row_bytes + if total_bytes > max_output_bytes: + raise ValueError( + f"Multimodal request needs {total_bytes} bytes of " + "encoder output storage but multimodal_config." + f"encoder_cache_max_bytes only allows {max_output_bytes}; " + "raise the cache budget to serve inputs of this size") + request.py_mm_encoder_state = ( + MultimodalEncoderRequestState.from_embedding_lengths( + embedding_lengths)) + + +def is_multimodal_encoder_ready(request: LlmRequest) -> bool: + """Return whether this request needs no further MM encoder work. + + Readiness is purely a function of the request's encoder item state: a + request without one (text-only, precomputed embedding, or a model + outside item scheduling) needs no encoder work; otherwise the request + is ready once every item slot is filled (finer-grained progress lives + on `MultimodalEncoderRequestState.progress`). + """ + state = request.py_mm_encoder_state + return state is None or state.progress is MultimodalEncoderProgress.READY + + def executor_request_to_llm_request( req_id: int, executor_request: ExecutorRequest, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 7ac471e3a852..6fd5666806ac 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -11,7 +11,7 @@ import weakref from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Hashable, List, Optional, Tuple, Union import torch import torch._dynamo.config @@ -29,7 +29,8 @@ strip_mm_data_for_generation) from tensorrt_llm.inputs.registry import (BaseMultimodalInputProcessor, create_input_processor, - create_input_processor_with_hash) + create_input_processor_with_hash, + get_multimodal_encoder_item_metadata) from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, DecodingBaseConfig, EncodeCudaGraphConfig, SeqLenAwareSparseAttentionConfig, @@ -84,6 +85,7 @@ get_multimodal_embedding_lengths) from .mamba_cache_manager import MambaHybridCacheManager from .model_loader import ModelLoader, _construct_checkpoint_loader +from .multimodal_encoder_cache_manager import MultimodalEncoderCacheManager from .resource_manager import (BaseResourceManager, KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType) @@ -92,6 +94,12 @@ from .trace_log_utils import log_mem_snapshot +def _resolve_mm_encoder_token_budget(base_budget: int, + model_max_atomic_item_tokens: int) -> int: + """Keep the model's largest indivisible MM item schedulable.""" + return max(base_budget, model_max_atomic_item_tokens) + + class ModelEngine(ABC): @abstractmethod @@ -307,9 +315,10 @@ def __init__( self.max_seq_len = max_seq_len self.max_beam_width = max_beam_width # Multimodal encoder runtime sizes; fall back to LLM-side values when - # the encoder-specific knobs are unset. + # the encoder-specific knobs are unset. The token budget may be + # raised after model load because atomic MM items cannot be split. ( - self.encoder_batch_size, + self.encoder_max_num_items, self.encoder_max_num_tokens, ) = llm_args.get_encoder_runtime_sizes() @@ -439,6 +448,67 @@ def __init__( # In case that some tests use stub models and override `_load_model`. if not hasattr(self.model, 'extra_attrs'): self.model.extra_attrs = {} + self.supports_mm_encoder_item_scheduling = ( + isinstance(self.model, MultimodalModelMixin) + and self.input_processor.supports_mm_encoder_item_scheduling) + self.mm_encoder_attention_metadata_capacity: Optional[Dict[str, + int]] = None + self.mm_encoder_cache_manager: Optional[ + MultimodalEncoderCacheManager] = None + if self.supports_mm_encoder_item_scheduling: + if self.encoder_max_num_tokens is None: + raise ValueError( + "MM encoder item scheduling requires a token budget; set " + "encoder_max_num_tokens or max_num_tokens") + max_tokens_per_item = self.input_processor.get_mm_max_tokens_per_item( + ) + if not max_tokens_per_item: + raise ValueError( + "A model with MM encoder item scheduling must implement " + "get_mm_max_tokens_per_item()") + if any(value <= 0 for value in max_tokens_per_item.values()): + raise ValueError( + "get_mm_max_tokens_per_item() must return positive token " + "counts") + model_max_atomic_item_tokens = max(max_tokens_per_item.values()) + encoder_token_budget_base = self.encoder_max_num_tokens + effective_encoder_token_budget = _resolve_mm_encoder_token_budget( + encoder_token_budget_base, model_max_atomic_item_tokens) + if effective_encoder_token_budget > self.encoder_max_num_tokens: + logger.warning_once( + f"encoder_max_num_tokens={self.encoder_max_num_tokens} " + "is smaller than the model's largest atomic " + f"multimodal item ({model_max_atomic_item_tokens}); " + f"using {model_max_atomic_item_tokens} as the " + "effective encoder runtime budget.", + key="raise_encoder_max_num_tokens_for_atomic_item", + ) + self.encoder_max_num_tokens = effective_encoder_token_budget + attention_metadata_capacity = ( + self.input_processor.get_mm_encoder_attention_metadata_capacity( + self.encoder_max_num_items, + self.encoder_max_num_tokens, + )) + if attention_metadata_capacity is not None: + if (not attention_metadata_capacity or any( + value <= 0 + for value in attention_metadata_capacity.values())): + raise ValueError( + "get_mm_encoder_attention_metadata_capacity() must " + "return nonempty positive capacities or None") + self.mm_encoder_attention_metadata_capacity = ( + attention_metadata_capacity) + logger.info( + "Multimodal encoder token budget: configured=%s, base=%d, " + "effective=%d, model_atomic_max=%d, attention_capacity=%s.", + llm_args.encoder_max_num_tokens, + encoder_token_budget_base, + self.encoder_max_num_tokens, + model_max_atomic_item_tokens, + self.mm_encoder_attention_metadata_capacity, + ) + self.mm_encoder_cache_manager = ( + self._create_mm_encoder_cache_manager(llm_args)) self._set_up_multimodal_encoder_attn_metadata() if self.llm_args.enable_layerwise_nvtx_marker: layerwise_nvtx_marker = LayerwiseNvtxMarker() @@ -2486,22 +2556,197 @@ def is_multimodal(self) -> bool: return True return isinstance(self.input_processor, BaseMultimodalInputProcessor) + @torch.inference_mode() + def forward_multimodal_encoder_items( + self, + requests: List[LlmRequest], + scheduled_items: Dict[int, List[int]], + ) -> None: + """Forward selected MM encoder items and commit request-local outputs.""" + if not scheduled_items: + return + if not isinstance(self.model, MultimodalModelMixin): + raise TypeError( + "Item-level MM scheduling requires MultimodalModelMixin") + + request_by_id = {request.request_id: request for request in requests} + selected_items = [] + selected_owners: List[Tuple[LlmRequest, int]] = [] + item_keys_by_id: Dict[int, Optional[List[Hashable]]] = {} + for request_id, item_indices in scheduled_items.items(): + request = request_by_id.get(request_id) + if request is None: + raise RuntimeError( + f"Scheduled MM request {request_id} is no longer active") + item_keys_by_id[request_id] = self.get_mm_encoder_item_keys(request) + multimodal_param = MultimodalParams( + multimodal_data=request.py_multimodal_data) + for item_idx in item_indices: + selected_items.append((multimodal_param, item_idx)) + selected_owners.append((request, item_idx)) + + encoder_inputs = self.model.prepare_multimodal_encoder_inputs( + selected_items) + for encoder_input, _, _ in encoder_inputs: + encoder_input.to_device( + "multimodal_data", + "cuda", + pin_memory=prefer_pinned(), + target_keywords=getattr(self.model, + "multimodal_data_device_paths", None), + ) + + outputs = self.model.forward_multimodal_encoder_items(encoder_inputs) + if len(outputs) != len(selected_owners): + raise RuntimeError( + "MM item encoder must return one output per item") + + manager = self.mm_encoder_cache_manager + if manager is None: + raise RuntimeError( + "MM encoder item scheduling requires the encoder cache " + "manager") + for output, (request, item_idx) in zip(outputs, + selected_owners, + strict=True): + state = request.py_mm_encoder_state + if state is None: + raise RuntimeError( + f"Scheduled MM request {request.py_request_id} has no " + "encoder item state") + item_keys = item_keys_by_id[request.request_id] + # Requests without stable content keys store under + # request-scoped temporary keys: never shared, reclaimed once + # the request's holds are released. + key = (item_keys[item_idx] if item_keys is not None else + ("mm_tmp", request.request_id, item_idx)) + # `adopt` takes ownership without cloning and holds the entry + # for this request; duplicate keys collapse to the already + # resident tensor (within-iteration dedup). + output_view = manager.adopt(key, output, request.request_id) + state.record(item_idx, output_view) + + touched_requests = { + request.request_id: request + for request, _ in selected_owners + } + for request in touched_requests.values(): + request.py_mm_encoder_state.finalize_into( + request.py_multimodal_data) + + def _create_mm_encoder_cache_manager( + self, llm_args) -> MultimodalEncoderCacheManager: + """Create the single budgeted storage for MM encoder outputs. + + The byte budget is `multimodal_config.encoder_cache_max_bytes` + clamped from below so one full prefill iteration's embeddings + always fit: without chunked MM prefill a request consumes all its + embedding rows in one iteration, and the rows co-resident for that + iteration are bounded by `max_num_tokens`. A budget below that + (including 0, which used to mean "cache off") cannot host the + embeddings the executor is guaranteed to need, so it is raised + with a warning instead of honored. + """ + row_bytes = self._resolve_mm_embedding_row_bytes() + self.mm_embedding_row_bytes = row_bytes + min_bytes = self.max_num_tokens * row_bytes + multimodal_config = getattr(llm_args, "multimodal_config", None) + configured = (multimodal_config.encoder_cache_max_bytes + if multimodal_config is not None else 0) + budget = max(configured, min_bytes) + if budget > configured: + logger.warning_once( + f"multimodal_config.encoder_cache_max_bytes={configured} is " + "smaller than one prefill iteration's embeddings " + f"(max_num_tokens={self.max_num_tokens} x " + f"row_bytes={row_bytes}); MM encoder outputs are stored " + f"exclusively in this cache, so using {budget} bytes.", + key="raise_mm_encoder_cache_budget", + ) + return MultimodalEncoderCacheManager(budget, + name="mm_encoder_cache_manager") + + def _resolve_mm_embedding_row_bytes(self) -> int: + """Bytes of one MM embedding row (hidden size x element size). + + Prefers the mixin's explicit `embedding_dim`/`embedding_dtype` + contract, then the text embedding layer's weight, then the + pretrained config's hidden size with the loaded weights' dtype — + both mixin properties are optional and most VLMs implement + neither. + """ + model = self.model + try: + return (model.embedding_dim * torch.empty( + (), dtype=model.embedding_dtype).element_size()) + except NotImplementedError: + pass + try: + weight = model.text_embedding_layer.weight + return weight.shape[-1] * weight.element_size() + except NotImplementedError: + pass + pretrained = model.model_config.pretrained_config + hidden_size = getattr(pretrained, "hidden_size", None) + if hidden_size is None: + hidden_size = getattr(getattr(pretrained, "text_config", None), + "hidden_size", None) + if hidden_size is None: + raise ValueError( + "Cannot derive the MM embedding row size: the model " + "implements neither embedding_dim/embedding_dtype nor " + "text_embedding_layer, and its pretrained config exposes " + "no (text_config.)hidden_size") + element_size = next(model.parameters()).dtype.itemsize + return hidden_size * element_size + + def get_mm_encoder_item_keys( + self, request: LlmRequest) -> Optional[List[Hashable]]: + """Return the request's per-item cache keys, or `None`. + + The single guard for content-addressed (shareable) storage keys on + the item scheduling path. `None` means the request cannot build + stable keys — it lacks item metadata, content hashes, or a + processor-kwargs hash — and its encoder outputs are stored under + request-scoped temporary keys instead (never shared, reclaimed + after its holds are released). + """ + # `getattr`: unit tests exercise partially constructed engines. + if not getattr(self, "supports_mm_encoder_item_scheduling", False): + return None + mm_data = request.py_multimodal_data + item_metadata = get_multimodal_encoder_item_metadata(mm_data) + if item_metadata is None: + return None + return self.model.build_encoder_cache_item_keys( + request.multimodal_hashes, + item_metadata.item_refs, + item_metadata.output_embedding_lengths, + mm_data.get("mm_processor_kwargs_hash"), + ) + def _set_up_multimodal_encoder_attn_metadata(self) -> None: """Construct AttentionMetadata for any multimodal encoders inside the loaded model, using the engine's encoder runtime sizes - (`encoder_max_batch_size` / `encoder_max_num_tokens`, falling back to + (`encoder_max_num_items` / `encoder_max_num_tokens`, falling back to the LLM-side `max_batch_size` / `max_num_tokens`). Mirrors `_set_up_attn_metadata` for the LLM backbone: encoders opt in by inheriting `MultimodalEncoderMixin`, and the engine drives the construction so the sizes match ``llm_args.get_encoder_runtime_sizes()`` rather - than being hardcoded inside each encoder's ``__init__``. + than being hardcoded inside each encoder's ``__init__``. The + item-count input counts atomic MM items; each encoder maps that item + capacity to its own internal attention sequence/window capacity. """ for module in self.model.modules(): if isinstance(module, MultimodalEncoderMixin): - module.setup_attn_metadata( - max_num_requests=self.encoder_batch_size, + setup_kwargs: Dict[str, Any] = dict( + max_num_items=self.encoder_max_num_items, max_num_tokens=self.encoder_max_num_tokens) + if self.mm_encoder_attention_metadata_capacity is not None: + setup_kwargs["attention_metadata_capacity"] = ( + self.mm_encoder_attention_metadata_capacity) + module.setup_attn_metadata(**setup_kwargs) def _set_up_spec_metadata( self, diff --git a/tensorrt_llm/_torch/pyexecutor/multimodal_encoder_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/multimodal_encoder_cache_manager.py new file mode 100644 index 000000000000..787c8049662e --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/multimodal_encoder_cache_manager.py @@ -0,0 +1,313 @@ +# Copyright 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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Hashable +from threading import RLock +from typing import TYPE_CHECKING, Generic, NamedTuple, TypeVar + +import torch + +from tensorrt_llm.logger import logger + +from .resource_manager import BaseResourceManager + +if TYPE_CHECKING: + from .llm_request import LlmRequest + +K = TypeVar("K", bound=Hashable) + +RequestId = int + + +class _Entry(NamedTuple): + value: torch.Tensor + size_bytes: int + + +class MultimodalEncoderCacheStats(NamedTuple): + max_bytes: int + current_bytes: int + held_bytes: int + item_count: int + held_count: int + hits: int + misses: int + adoptions: int + dedup_adoptions: int + evictions: int + + +class MultimodalEncoderCacheManager(BaseResourceManager, Generic[K]): + """Budgeted single storage for multimodal encoder outputs. + + Encoder outputs of item-scheduling models live *only* here: requests + reference entries by key, and the byte budget therefore bounds the + total GPU memory resident encoder outputs can occupy. Cross-request + reuse is a side effect of the storage being content-addressed, not a + separate cache. + + Registered as a ``BaseResourceManager`` so request teardown flows + through the executor's standard ``free_resources`` funnel; admission is + NOT capacity-gated here (the MM item scheduler enforces the byte budget + via allocate-before-compute), so the scheduler-facing resource counts + are zero. + + Semantics: + + - **Holding.** An entry is held while at least one request references + it. Held entries are never evicted; a hold is the request's guarantee + that the embedding survives until its prefill consumes it. + - **Allocate-before-compute.** Callers (the MM item scheduler) must check + `can_allocate()` before encoding an item. `adopt()` evicting only + zero-reference entries relies on that contract and raises if it cannot + make room. + - **Adopt, not copy.** `adopt()` takes ownership of the caller's tensor + without cloning when it exclusively owns its storage; views into a + larger allocation are materialized so accounted bytes always equal + physical residency. If the key already exists, the existing entry is + held and returned instead — concurrent duplicate encodes of the same + content collapse to one resident copy. + - **Zero-ref residency.** Entries whose last hold was released stay + resident in LRU order for reuse and are reclaimed only when a later + allocation needs the space. + + Returned tensors alias manager-owned storage and must be treated as + immutable by callers. + + Args: + max_bytes: Maximum logical tensor bytes held by this manager. + name: Short label used in debug log messages. + """ + + def __init__( + self, + max_bytes: int, + *, + name: str = "mm_encoder_cache_manager", + ) -> None: + if max_bytes <= 0: + raise ValueError("max_bytes must be positive") + + self._max_bytes = max_bytes + self._name = name + self._current_bytes = 0 + self._held_bytes = 0 + # Recency order over all entries; eviction walks from the LRU end + # skipping held entries. + self._entries: OrderedDict[K, _Entry] = OrderedDict() + self._key_holders: dict[K, set[RequestId]] = {} + self._request_keys: dict[RequestId, set[K]] = {} + self._lock = RLock() + self._hits = 0 + self._misses = 0 + self._adoptions = 0 + self._dedup_adoptions = 0 + self._evictions = 0 + + @property + def max_bytes(self) -> int: + return self._max_bytes + + @property + def current_bytes(self) -> int: + with self._lock: + return self._current_bytes + + @property + def held_bytes(self) -> int: + with self._lock: + return self._held_bytes + + def __len__(self) -> int: + with self._lock: + return len(self._entries) + + def available_bytes(self, *, reserved_bytes: int = 0) -> int: + """Bytes obtainable for new allocations: free plus zero-ref + reclaimable space, minus the caller's outstanding reservation.""" + with self._lock: + return max(0, self._max_bytes - self._held_bytes - reserved_bytes) + + def can_allocate(self, nbytes: int, *, reserved_bytes: int = 0) -> bool: + """Whether `nbytes` can be admitted without evicting held entries. + + `reserved_bytes` lets the scheduler set space aside for the + head-of-line request's remaining items before admitting items of + later requests (deadlock avoidance).""" + return nbytes <= self.available_bytes(reserved_bytes=reserved_bytes) + + def contains(self, key: K) -> bool: + with self._lock: + return key in self._entries + + def get_and_hold(self, key: K, request_id: RequestId) -> torch.Tensor | None: + """Return the entry for `key` held on behalf of `request_id`, + or `None` on miss. The hit is promoted to most-recently-used.""" + with self._lock: + entry = self._entries.get(key) + if entry is None: + self._misses += 1 + return None + self._hits += 1 + self._entries.move_to_end(key) + self._hold(key, entry, request_id) + return entry.value + + def adopt(self, key: K, value: torch.Tensor, request_id: RequestId) -> torch.Tensor: + """Store `value` (taking ownership) held for `request_id`. + + Tensors that exclusively own their storage are stored without a + copy. Views into a larger backing allocation (e.g. `torch.split` + outputs of a grouped encoder forward) are materialized first: + entries from one batch would otherwise share storage, so evicting + one while a sibling survives frees accounted bytes without freeing + physical memory, breaking the budget's residency guarantee. + + If `key` is already resident the existing tensor is held and + returned and `value` is dropped. Evicts zero-reference entries as + needed; raises `RuntimeError` when space cannot be made — callers + must have passed `can_allocate()` first (allocate-before-compute). + """ + size_bytes = value.numel() * value.element_size() + if value.untyped_storage().nbytes() != size_bytes: + value = value.clone() + + with self._lock: + existing = self._entries.get(key) + if existing is not None: + self._dedup_adoptions += 1 + self._entries.move_to_end(key) + self._hold(key, existing, request_id) + return existing.value + + if not self.can_allocate(size_bytes): + raise RuntimeError( + f"{self._name}: cannot adopt {size_bytes} bytes; " + f"held={self._held_bytes}, max={self._max_bytes}. " + "The MM item scheduler must reserve space via " + "can_allocate() before encoding." + ) + + self._evict_zero_ref_until(self._max_bytes - size_bytes) + entry = _Entry(value=value, size_bytes=size_bytes) + self._entries[key] = entry + self._current_bytes += size_bytes + self._adoptions += 1 + self._hold(key, entry, request_id) + return entry.value + + def release_holds(self, request_id: RequestId) -> None: + """Release every hold of `request_id` (idempotent). + + This is the single teardown funnel: prefill completion, cancellation, + and error paths all route here. Entries left with zero references + stay resident for reuse until eviction needs their space. + """ + with self._lock: + keys = self._request_keys.pop(request_id, None) + if not keys: + return + for key in keys: + holders = self._key_holders.get(key) + if holders is None: + continue + holders.discard(request_id) + if not holders: + del self._key_holders[key] + entry = self._entries.get(key) + if entry is not None: + self._held_bytes -= entry.size_bytes + + def free_resources(self, request: "LlmRequest") -> None: + """`BaseResourceManager` teardown hook: release the request's holds. + + Runs for every request termination (completion, cancellation, + errors) through `ResourceManager.free_resources`; the post-prefill + strip additionally calls it early so entries become reclaimable as + soon as their embedding has been consumed. + """ + self.release_holds(request.request_id) + + def get_max_resource_count(self) -> int: + """Encoder-output bytes are not a capacity-scheduler resource (the + MM item scheduler budgets them at selection); return 0 so the + executor does not gate admission on this manager.""" + return 0 + + def get_needed_resource_to_completion(self, request: "LlmRequest") -> int: + """See `get_max_resource_count`.""" + return 0 + + def shutdown(self) -> None: + self.clear() + + def clear(self) -> None: + """Drop all entries and holds. Only safe when no request is active + (e.g. warmup teardown or weight updates).""" + with self._lock: + self._entries.clear() + self._key_holders.clear() + self._request_keys.clear() + self._current_bytes = 0 + self._held_bytes = 0 + + def stats(self) -> MultimodalEncoderCacheStats: + with self._lock: + return MultimodalEncoderCacheStats( + max_bytes=self._max_bytes, + current_bytes=self._current_bytes, + held_bytes=self._held_bytes, + item_count=len(self._entries), + held_count=len(self._key_holders), + hits=self._hits, + misses=self._misses, + adoptions=self._adoptions, + dedup_adoptions=self._dedup_adoptions, + evictions=self._evictions, + ) + + def log_stats(self, reason: str) -> None: + stats = self.stats() + logger.debug( + f"{self._name}: stats after {reason}: items={stats.item_count} " + f"(held={stats.held_count}), bytes={stats.current_bytes}" + f"/{stats.max_bytes} (held={stats.held_bytes}), " + f"hits={stats.hits}, misses={stats.misses}, " + f"adoptions={stats.adoptions} (dedup={stats.dedup_adoptions}), " + f"evictions={stats.evictions}" + ) + + def _hold(self, key: K, entry: _Entry, request_id: RequestId) -> None: + holders = self._key_holders.setdefault(key, set()) + if not holders: + self._held_bytes += entry.size_bytes + holders.add(request_id) + self._request_keys.setdefault(request_id, set()).add(key) + + def _evict_zero_ref_until(self, target_bytes: int) -> None: + if self._current_bytes <= target_bytes: + return + for key in list(self._entries): + if self._current_bytes <= target_bytes: + break + if key in self._key_holders: + continue + entry = self._entries.pop(key) + self._current_bytes -= entry.size_bytes + self._evictions += 1 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4f4fc081a1d1..74530611b08d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -39,6 +39,7 @@ ReqIdsSet) from tensorrt_llm.executor.request import TruncateKVCacheRequest from tensorrt_llm.inputs.multimodal import strip_mm_data_for_generation +from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata from tensorrt_llm.llmapi.llm_args import PeftCacheConfig, WaitingQueuePolicy from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType @@ -74,7 +75,9 @@ from .llm_request import (ATTENTION_DP_DUMMY_REQUEST_ID, MAX_SPEC_DECODE_POSITIONS, ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, - get_draft_token_length) + get_draft_token_length, + initialize_multimodal_encoder_request, + is_multimodal_encoder_ready) from .mamba_cache_manager import (BaseMambaCacheManager, MixedMambaHybridCacheManager) from .model_engine import ModelEngine @@ -234,7 +237,7 @@ def _load_iteration_indexes(env_var: str): def _strip_py_multimodal_data_post_prefill(request: LlmRequest) -> None: - """Drop pinned encoder cache and raw pre-encoder tensors after prefill completes. + """Drop encoder outputs and raw pre-encoder tensors after prefill completes. Wraps `strip_mm_data_for_generation` and mutates the shared `request.py_multimodal_data` in-place so the `LlmRequest`'s multimodal tensors actually get freed (unlike @@ -245,6 +248,11 @@ def _strip_py_multimodal_data_post_prefill(request: LlmRequest) -> None: if not mm_data: return strip_mm_data_for_generation(mm_data) + # Drop the per-item encoder state alongside the dict clear above: both + # alias manager-owned entries, so every alias source disappears at the + # same moment the caller releases the request's holds — after this, + # eviction of the zero-ref entries actually frees the GPU memory. + request.py_mm_encoder_state = None @dataclasses.dataclass @@ -564,10 +572,15 @@ def __init__( # related modules self.resource_manager = resource_manager - self.scheduler = scheduler self.model_engine = model_engine self._enable_dsv4_adp_dummy_fixes = getattr( model_engine, "_enable_dsv4_adp_dummy_fixes", False) + # Compatibility checks and MultimodalScheduler wrapping live in + # `create_py_executor_instance`; the executor only keeps the flag + # its loop paths branch on. + self._supports_mm_encoder_item_scheduling = getattr( + model_engine, "supports_mm_encoder_item_scheduling", False) + self.scheduler = scheduler self.enable_attention_dp = model_engine.enable_attention_dp self.dist = dist self.sampler = sampler @@ -2544,6 +2557,17 @@ def _executor_loop_pp(self): if self.dist.rank != 0: # Retry until current rank can run first PP's schedule result. self._pp_retry_until_can_schedule(scheduled_batch) + # Replay the cache-hit attach (mirrors `_schedule()` on + # the scheduling rank): every rank encodes the same + # broadcast items into its own cache, so the replay makes + # identical hit decisions. Without it, leader-side hits + # stay pending here and the model forward falls back to + # the unbudgeted legacy inline encode. + # TODO: interim — proper VLM PP support should scope MM + # work to the embedding-consuming stages (fuse + + # deepstack-injection) so other ranks skip it entirely. + if self._supports_mm_encoder_item_scheduling: + self._attach_mm_encoder_cache_hits() # Run scheduler locally because scheduler may change llm requests' state. if hasattr(self.kv_cache_manager, "prepare_expect_snapshot_points"): @@ -2559,6 +2583,10 @@ def _executor_loop_pp(self): local_disagg_candidates, fitting_disagg_gen_init_requests) + if (self._supports_mm_encoder_item_scheduling + and scheduled_batch.scheduled_mm_encoder_items): + self._forward_multimodal_encoder_step(scheduled_batch) + # For requests that are fitting disagg gen init, also prepare resources for KV cache manager if self.kv_cache_transceiver: self._prepare_disagg_gen_init( @@ -3979,6 +4007,10 @@ def _executor_loop(self): self._finalize_adp_dummy_allocation(False) continue + if (self._supports_mm_encoder_item_scheduling + and scheduled_batch.scheduled_mm_encoder_items): + self._forward_multimodal_encoder_step(scheduled_batch) + if not self._is_kv_manager_v2: self._terminate_requests(scheduled_batch.paused_requests) self._pause_requests(scheduled_batch.paused_requests) @@ -4455,6 +4487,10 @@ def _executor_loop_overlap(self): self._finalize_adp_dummy_allocation(False) continue + if (self._supports_mm_encoder_item_scheduling + and scheduled_batch.scheduled_mm_encoder_items): + self._forward_multimodal_encoder_step(scheduled_batch) + if not self._is_kv_manager_v2: self._terminate_requests(scheduled_batch.paused_requests) @@ -4921,6 +4957,75 @@ def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, waiting_queue.add_requests(new_requests) + def _apply_mm_encoder_admission( + self, waiting_queue: WaitingQueue, + new_requests: List[RequestQueueItem]) -> List[RequestQueueItem]: + """Apply MM encoder item/token budgets to capacity-admitted requests. + + Active requests consume this iteration's encoder budget first. New + requests then preserve waiting-queue FCFS order: once the head request + cannot make encoder progress, all later requests are deferred and + prepended to the waiting queue. + """ + remaining_items = self.model_engine.encoder_max_num_items + remaining_tokens = self.model_engine.encoder_max_num_tokens + + def consume_pending(costs, pending_indices=None): + nonlocal remaining_items, remaining_tokens + if pending_indices is None: + pending_indices = range(len(costs)) + progressed = False + for item_idx in pending_indices: + cost = costs[item_idx] + if remaining_items == 0 or cost > remaining_tokens: + break + remaining_items -= 1 + remaining_tokens -= cost + progressed = True + return progressed + + for request in self.active_requests: + state = request.py_mm_encoder_state + if state is None or is_multimodal_encoder_ready(request): + continue + mm_data = request.py_multimodal_data or {} + item_metadata = get_multimodal_encoder_item_metadata(mm_data) + if item_metadata is None: + continue + consume_pending(item_metadata.encoder_token_lengths, + state.pending_item_indices()) + + admitted = [] + deferred = [] + blocked = False + for queue_item in new_requests: + if blocked: + deferred.append(queue_item) + continue + mm_data = getattr(queue_item.request, "py_multimodal_data", None) + item_metadata = (get_multimodal_encoder_item_metadata(mm_data) + if isinstance(mm_data, dict) else None) + costs = (item_metadata.encoder_token_lengths + if item_metadata is not None else None) + has_full_embedding = (isinstance(mm_data, dict) + and mm_data.get("multimodal_embedding") + is not None) + if costs and any(cost > self.model_engine.encoder_max_num_tokens + for cost in costs): + # Admit the invalid request so normal request validation can + # return an error instead of leaving the FCFS queue blocked. + admitted.append(queue_item) + continue + if costs and not has_full_embedding and not consume_pending(costs): + blocked = True + deferred.append(queue_item) + continue + admitted.append(queue_item) + + if deferred: + waiting_queue.prepend_requests(deferred) + return admitted + def _pop_from_waiting_queue( self, waiting_queue: WaitingQueue, @@ -4944,12 +5049,17 @@ def _pop_from_waiting_queue( self._fill_admit_cap = min(self._fill_admit_cap * 2, total_max) max_new_requests = min(max_new_requests, self._fill_admit_cap) - return get_from_waiting_queue( + new_requests = get_from_waiting_queue( waiting_queue, max_new_requests, enable_attention_dp=self.enable_attention_dp, max_num_active_requests=self.max_num_active_requests, all_ranks_num_active_requests=all_ranks_num_active_requests) + if (not new_requests or + not getattr(self, "_supports_mm_encoder_item_scheduling", False) + or self.enable_attention_dp): + return new_requests + return self._apply_mm_encoder_admission(waiting_queue, new_requests) @nvtx_range("_fetch_new_requests") def _fetch_new_requests( @@ -5085,6 +5195,16 @@ def _respond_if_invalid(request: LlmRequest) -> bool: """ try: self._validate_request(request) + if self._supports_mm_encoder_item_scheduling: + manager = self.model_engine.mm_encoder_cache_manager + initialize_multimodal_encoder_request( + request, + max_num_tokens=self.model_engine.encoder_max_num_tokens, + max_output_bytes=(manager.max_bytes + if manager is not None else None), + embedding_row_bytes=getattr(self.model_engine, + "mm_embedding_row_bytes", + 0)) return False except Exception as e: self._handle_errors(str(e), @@ -5228,12 +5348,45 @@ def _waiting_requests(self, context_requests: list[LlmRequest], self.batch_wait_iters_count = 0 return context_requests + def _attach_mm_encoder_cache_hits(self) -> None: + """Complete pending MM items from resident cache entries before + scheduling. + + Hits are held for the request and recorded into its item slots + ahead of scheduler selection, so they neither consume the encoder + budgets nor get re-encoded. A request with remaining misses becomes + PARTIAL and keeps following the item path, which re-computes only + those misses; a fully attached request graduates to READY without + any encoder work. Requests without stable content keys can never + hit (their outputs live under request-scoped temporary keys), so + they are skipped. + """ + manager = self.model_engine.mm_encoder_cache_manager + if manager is None: + return + for request in self.active_requests: + state = request.py_mm_encoder_state + if state is None or is_multimodal_encoder_ready(request): + continue + item_keys = self.model_engine.get_mm_encoder_item_keys(request) + if item_keys is None: + continue + for item_idx in state.pending_item_indices(): + cached_output = manager.get_and_hold(item_keys[item_idx], + request.request_id) + if cached_output is None: + continue + state.record(item_idx, cached_output) + state.finalize_into(request.py_multimodal_data) + @nvtx_range("_schedule") def _schedule(self): if hasattr(self.kv_cache_manager, "prepare_expect_snapshot_points"): self.kv_cache_manager.prepare_expect_snapshot_points( self.active_requests) + if self._supports_mm_encoder_item_scheduling: + self._attach_mm_encoder_cache_hits() scheduler_output = self.scheduler.schedule_request( self.active_requests, self.inflight_req_ids) @@ -5269,9 +5422,20 @@ def _schedule(self): scheduled_requests.reset_context_requests(scheduled_context_requests) scheduled_requests.generation_requests = scheduler_output.generation_requests scheduled_requests.paused_requests = scheduler_output.paused_requests + scheduled_requests.scheduled_mm_encoder_items = ( + scheduler_output.scheduled_mm_encoder_items) return scheduled_requests, scheduler_output.fitting_disagg_gen_init_requests, num_fitting + def _forward_multimodal_encoder_step( + self, scheduled_requests: ScheduledRequests) -> None: + """Run scheduler-selected MM encoder work before LLM resources.""" + scheduled_items = scheduled_requests.scheduled_mm_encoder_items + if not scheduled_items: + return + self.model_engine.forward_multimodal_encoder_items( + self.active_requests, scheduled_items) + # --------------------------------------------------------------- # Encoder-decoder support: encoder iteration in the executor loop. # @@ -6315,6 +6479,7 @@ def _update_request_states_tp(self, scheduled_requests: ScheduledRequests): # requests stay pinned on GPU through the full decode lifetime and can lead to OOMs # at high concurrency. _strip_py_multimodal_data_post_prefill(request) + self._release_mm_encoder_holds(request) if not self.disable_overlap_scheduler and request.will_complete_next_iteration( ): request.set_exclude_last_generation_logits(False) @@ -6542,7 +6707,22 @@ def _terminate_request(self, request: LlmRequest): else: self._do_terminate_request(request) + def _release_mm_encoder_holds(self, request: LlmRequest) -> None: + """Release the request's MM encoder cache holds early (idempotent). + + The registered MM_ENCODER_CACHE_MANAGER resource manager already + releases holds on every termination path via + `resource_manager.free_resources`; this early call at the + post-prefill strip makes entries reclaimable as soon as their + embedding has been consumed instead of at request end. + """ + manager = getattr(self.model_engine, "mm_encoder_cache_manager", None) + if manager is not None: + manager.free_resources(request) + def _do_terminate_request(self, request: LlmRequest): + # MM encoder holds release inside free_resources via the registered + # MM_ENCODER_CACHE_MANAGER resource manager. self.resource_manager.free_resources(request) self._prefetched_request_ids.discard(request.py_request_id) self._disagg_timed_out_ctx_cancelled_ids.discard(request.py_request_id) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 157e2d7d720a..acebe081499e 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -99,6 +99,7 @@ class ResourceManagerType(enum.Enum): SEQ_SLOT_MANAGER = "SEQ_SLOT_MANAGER" SPEC_RESOURCE_MANAGER = "SPEC_RESOURCE_MANAGER" KV_CACHE_COMPRESSION_MANAGER = "KV_CACHE_COMPRESSION_MANAGER" + MM_ENCODER_CACHE_MANAGER = "MM_ENCODER_CACHE_MANAGER" def compute_page_count(token_count: int, tokens_per_page: int) -> int: diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/__init__.py b/tensorrt_llm/_torch/pyexecutor/scheduler/__init__.py index 0c549f89f92b..ef2927d24713 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -27,6 +27,8 @@ BindMicroBatchScheduler, CapacityScheduler, MicroBatchScheduler, + MultimodalEagerEncoderScheduler, + MultimodalScheduler, PyCapacityScheduler, PyMicroBatchScheduler, RequestList, @@ -54,6 +56,8 @@ "CapacityScheduler", "KVCacheV2Scheduler", "MicroBatchScheduler", + "MultimodalEagerEncoderScheduler", + "MultimodalScheduler", "PyCapacityScheduler", "PyMicroBatchScheduler", "RequestList", diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index caa8e3cb3de1..7e2215c76a74 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -1,3 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Request schedulers used by the PyTorch executor.""" + import dataclasses import inspect from abc import ABC, abstractmethod @@ -13,7 +18,12 @@ from tensorrt_llm.logger import logger # Assuming these imports exist in your environment -from ..llm_request import LlmRequest, LlmRequestState +from ..llm_request import ( + LlmRequest, + LlmRequestState, + get_multimodal_encoder_token_lengths, + is_multimodal_encoder_ready, +) RequestList = list[LlmRequest] PrefixReuseSummary: TypeAlias = tb_internal.batch_manager.PrefixReuseSummary @@ -62,7 +72,11 @@ def _call_with_optional_summary( "paused_requests", "fitting_disagg_gen_init_requests", "num_fitting_requests", + # request id -> prompt-ordered indices of its MM items selected for + # encoder execution this iteration + "scheduled_mm_encoder_items", ], + defaults=[None], ) @@ -152,6 +166,10 @@ class ScheduledRequests: """Requests that are in the generation phase.""" paused_requests: RequestList """Requests that are paused.""" + scheduled_mm_encoder_items: dict[int, list[int]] | None + """Maps a request id to the prompt-ordered indices of its multimodal items + selected for encoder execution this iteration (only items whose encoder + outputs are still missing). ``None`` when no items were scheduled.""" def __init__(self): self.encoder_requests: RequestList = [] @@ -159,6 +177,7 @@ def __init__(self): self.context_requests_last_chunk: RequestList = [] self.generation_requests: RequestList = [] self.paused_requests: RequestList = [] + self.scheduled_mm_encoder_items: dict[int, list[int]] | None = None @property def is_generation_only(self) -> bool: @@ -257,6 +276,9 @@ class SerializableSchedulerOutput: ] # request ids of fitting disaggregated generation initialization requests num_fitting_requests: int # number of fitting requests wait_for_disagg_gen_transfer_progress: bool = False + # request id -> prompt-ordered indices of its MM items selected for + # encoder execution this iteration + scheduled_mm_encoder_items: dict[int, list[int]] | None = None @classmethod def from_scheduler_result( @@ -281,6 +303,7 @@ def from_scheduler_result( ], num_fitting_requests=num_fitting_requests, wait_for_disagg_gen_transfer_progress=wait_for_disagg_gen_transfer_progress, + scheduled_mm_encoder_items=scheduled_requests.scheduled_mm_encoder_items, ) def to_scheduler_result( @@ -303,6 +326,7 @@ def to_scheduler_result( scheduled_requests.paused_requests = [ id_to_request[req_id] for req_id in self.paused_requests ] + scheduled_requests.scheduled_mm_encoder_items = self.scheduled_mm_encoder_items fitting_disagg_gen_init_requests = [ id_to_request[req_id] for req_id in self.fitting_disagg_gen_init_requests ] @@ -453,6 +477,268 @@ def can_schedule(self, requests: RequestList) -> bool: return len(fitting_requests) == len(requests) +class MultimodalScheduler(RequestScheduler): + """Add atomic multimodal item budgeting around the existing scheduler. + + The wrapper is constructed only for ``MultimodalModelMixin`` models. It + deliberately reuses the wrapped scheduler's capacity and microbatch + schedulers so MM encoder costs never enter the LLM token budget. + + ``max_num_items`` is the resolved user-facing + ``encoder_max_num_items``. It counts atomic MM items across all requests + and modalities, not LLM requests or model-internal attention segments; + item *size* is budgeted separately by ``max_num_tokens``, and encoders + that split one item into multiple attention sequences derive their own + workspace capacity from the token budget. + + When a ``MultimodalEncoderCacheManager`` is attached, selection also + enforces its byte budget (allocate-before-compute): encoder outputs are + stored exclusively in the manager, so an item is only selected when its + embedding bytes fit alongside everything already held or claimed. + """ + + def __init__( + self, + scheduler: SimpleScheduler, + max_num_items: int, + max_num_tokens: int, + *, + cache_manager=None, + embedding_row_bytes: int = 0, + ) -> None: + self.scheduler = scheduler + self.max_num_items = max_num_items + self.max_num_tokens = max_num_tokens + # Optional `MultimodalEncoderCacheManager`: when present, item + # selection additionally performs allocate-before-compute against + # the manager's byte budget (encoder outputs are stored exclusively + # there). `embedding_row_bytes` converts declared embedding rows to + # bytes and must be positive alongside a manager. + self.cache_manager = cache_manager + self.embedding_row_bytes = embedding_row_bytes + if cache_manager is not None and embedding_row_bytes <= 0: + raise ValueError( + "embedding_row_bytes must be positive when a cache manager " + "budgets MM encoder output storage" + ) + self.has_separate_stages = hasattr(scheduler, "capacity_scheduler") and hasattr( + scheduler, "micro_batch_scheduler" + ) + + def _select_items(self, requests: RequestList) -> tuple[dict[int, list[int]], RequestList]: + """Greedily select pending MM items under the encoder budgets. + + Requests are visited in the wrapped capacity scheduler's FCFS order + with no explicit `MultimodalEncoderProgress`-based priority: a + request left `PARTIAL` by a budget split necessarily sits ahead of + anything admitted later, so its remaining items resume before newer + work by order alone. A request that arrives `PARTIAL` through cache + hits gets no boost over older `PENDING` requests — a deliberate + fairness default. + + When a cache manager is attached, selection also performs + allocate-before-compute: an item is only selected if the manager can + host its embedding bytes, counting bytes claimed earlier in this + pass (selected items are adopted only later this iteration). The + first request left incomplete additionally reserves its remaining + bytes so requests behind it cannot squat the space it needs across + iterations — without this head-of-line reservation, several + partially-stored requests could hold fragments of the budget and + deadlock waiting on one another. + + Returns the selected item indices per request id, plus the requests + eligible for LLM microbatch scheduling this iteration (encoder + outputs already ready, or every pending item selected above). + """ + remaining_items = self.max_num_items + remaining_tokens = self.max_num_tokens + manager = self.cache_manager + reserved_bytes = 0 + head_of_line_reserved = False + selected: dict[int, list[int]] = {} + llm_eligible: RequestList = [] + + for request in requests: + state = request.py_mm_encoder_state + if state is None: + llm_eligible.append(request) + continue + if is_multimodal_encoder_ready(request): + llm_eligible.append(request) + continue + + token_lengths = get_multimodal_encoder_token_lengths(request) + if token_lengths is None: + raise ValueError( + f"Multimodal request {request.py_request_id} is missing " + "multimodal_encoder_item_metadata" + ) + if len(token_lengths) != state.num_items: + raise ValueError( + f"Multimodal request {request.py_request_id} has " + f"{state.num_items} item slots but {len(token_lengths)} " + "encoder token lengths" + ) + + pending = state.pending_item_indices() + request_items: list[int] = [] + for item_idx in pending: + cost = token_lengths[item_idx] + if remaining_items == 0 or cost > remaining_tokens: + break + if manager is not None: + item_bytes = state.embedding_lengths[item_idx] * self.embedding_row_bytes + if not manager.can_allocate(item_bytes, reserved_bytes=reserved_bytes): + break + reserved_bytes += item_bytes + request_items.append(item_idx) + remaining_items -= 1 + remaining_tokens -= cost + + if request_items: + selected[request.request_id] = request_items + + if pending and len(request_items) == len(pending): + llm_eligible.append(request) + elif manager is not None and not head_of_line_reserved: + # Head-of-line reservation: `request_items` is a prefix of + # `pending`, so the unselected suffix is what this request + # still needs in future iterations. + remaining_request_bytes = ( + sum( + state.embedding_lengths[item_idx] + for item_idx in pending[len(request_items) :] + ) + * self.embedding_row_bytes + ) + total_request_bytes = sum(state.embedding_lengths) * self.embedding_row_bytes + if total_request_bytes > manager.max_bytes: + # Liveness backstop: admission + # (`initialize_multimodal_encoder_request`) already + # rejects requests whose holds can never coexist within + # the budget, so reaching this means an accounting bug + # rather than a user input. + raise RuntimeError( + f"Multimodal request {request.py_request_id} needs " + f"{total_request_bytes} bytes of encoder output " + "storage but multimodal_config." + f"encoder_cache_max_bytes only allows " + f"{manager.max_bytes}; raise the cache budget" + ) + reserved_bytes += remaining_request_bytes + head_of_line_reserved = True + + return selected, llm_eligible + + def _schedule_micro_batch( + self, + fitting_requests: RequestList, + fitting_disagg_gen_init_requests: RequestList, + paused_requests: RequestList, + inflight_request_ids: set[int], + *, + llm_eligible: RequestList, + selected_items: dict[int, list[int]] | None = None, + ) -> SchedulerOutput: + encoder_requests, context_requests, generation_requests = ( + self.scheduler.micro_batch_scheduler.schedule(llm_eligible, inflight_request_ids) + ) + return SchedulerOutput( + encoder_requests=encoder_requests, + context_requests=context_requests, + generation_requests=generation_requests, + paused_requests=list(paused_requests), + fitting_disagg_gen_init_requests=list(fitting_disagg_gen_init_requests), + num_fitting_requests=len(fitting_requests), + scheduled_mm_encoder_items=selected_items or None, + ) + + def schedule_request( + self, active_requests: RequestList, inflight_request_ids: set[int] + ) -> SchedulerOutput: + """Apply the default LLM-capacity-coupled MM scheduling policy. + + First use the wrapped scheduler to determine which requests fit + LLM/KV capacity, then select their pending atomic MM items under the + encoder budgets. The executor's encoder step is the single site that + runs MM encoders: an in-budget batch is simply the case where every + pending item gets selected. Only requests whose encoder outputs are + ready, or become ready this iteration, enter LLM microbatch + scheduling. + """ + if not self.has_separate_stages: + # Compatibility path for schedulers exposing only a combined API: + # schedule the LLM batch first, then enforce MM budgets on its + # context requests, withholding contexts that will still lack MM + # embeddings after this iteration. + scheduler_output = self.scheduler.schedule_request( + active_requests, inflight_request_ids + ) + selected_items, llm_eligible = self._select_items( + list(scheduler_output.context_requests) + ) + return scheduler_output._replace( + context_requests=llm_eligible, + scheduled_mm_encoder_items=selected_items or None, + ) + + # Only requests admitted by ordinary LLM/KV capacity may consume MM + # encoder budget this iteration. + fitting_requests, fitting_disagg_gen_init_requests, paused_requests = ( + self.scheduler.capacity_scheduler.schedule_request(active_requests) + ) + selected_items, llm_eligible = self._select_items(list(fitting_requests)) + # Preserve the capacity scheduler's decisions while attaching the MM + # item plan that the executor must run before the selected LLM + # microbatch. + return self._schedule_micro_batch( + fitting_requests, + fitting_disagg_gen_init_requests, + paused_requests, + inflight_request_ids, + llm_eligible=llm_eligible, + selected_items=selected_items, + ) + + def can_schedule(self, requests: RequestList) -> bool: + return self.scheduler.can_schedule(requests) + + +class MultimodalEagerEncoderScheduler(MultimodalScheduler): + """Eagerly schedule encoder work for already-active MM requests. + + Unlike the default coupled policy, this policy selects encoder items before + LLM capacity scheduling. An active request may therefore make encoder + progress even when it is not selected for the current LLM batch. It does + not admit waiting requests or bypass LLM capacity for decoder execution. + """ + + def schedule_request( + self, active_requests: RequestList, inflight_request_ids: set[int] + ) -> SchedulerOutput: + selected_items, llm_eligible = self._select_items(active_requests) + + if not self.has_separate_stages: + scheduler_output = self.scheduler.schedule_request(llm_eligible, inflight_request_ids) + return scheduler_output._replace(scheduled_mm_encoder_items=selected_items or None) + + llm_eligible_ids = {request.request_id for request in llm_eligible} + fitting_requests, fitting_disagg_gen_init_requests, paused_requests = ( + self.scheduler.capacity_scheduler.schedule_request(active_requests) + ) + fitting_llm_eligible = [ + request for request in fitting_requests if request.request_id in llm_eligible_ids + ] + return self._schedule_micro_batch( + fitting_requests, + fitting_disagg_gen_init_requests, + paused_requests, + inflight_request_ids, + llm_eligible=fitting_llm_eligible, + selected_items=selected_items, + ) + + class ChunkingPolicy(Enum): EQUAL_PROGRESS = 1 FIRST_COME_FIRST_SERVED = 2 diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py index 3b9a5d51d052..0330b9d509c6 100644 --- a/tensorrt_llm/inputs/__init__.py +++ b/tensorrt_llm/inputs/__init__.py @@ -12,7 +12,8 @@ # yapf: disable from .registry import (BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor, ExtraProcessedInputs, - InputProcessor, MultimodalPlaceholderMetadata, + InputProcessor, MultimodalEncoderItemMetadata, + MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, create_input_processor, create_input_processor_with_hash, maybe_compute_mm_embed_cumsum, register_input_processor, @@ -51,6 +52,7 @@ "ExtraProcessedInputs", "BaseMultimodalDummyInputsBuilder", "BaseMultimodalInputProcessor", + "MultimodalEncoderItemMetadata", "MultimodalPlaceholderMetadata", "MultimodalPlaceholderPlacement", "ConversationMessage", diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index fb4b33779d32..0cd7e5d7c90d 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -17,6 +17,7 @@ # Default hasher default_hasher = blake3 _INT32_MAX = 2**31 - 1 +MULTIMODAL_ENCODER_ITEM_METADATA_KEY = "multimodal_encoder_item_metadata" # Versioned tag prefixed to every content hash so the canonical, self-describing # serialization scheme can evolve without silently reusing stale cache keys. @@ -167,6 +168,18 @@ def strip_mm_data_for_generation(mm_data: Dict[str, Any]) -> None: mm_data['mrope_config'] = {'mrope_position_deltas': mrope_deltas} +def strip_mm_encoder_inputs(mm_data: Dict[str, Any]) -> None: + """Drop raw encoder payloads while preserving cached MM embeddings. + + Item-level encoder scheduling keeps the original request payload on CPU + and transfers only selected items to GPU. Once every item is encoded, the + raw modality dictionaries are no longer needed and must not be moved to + GPU by the subsequent LLM input-preparation path. + """ + for modality in ("image", "video", "audio"): + mm_data.pop(modality, None) + + @dataclass class MultimodalInput: """Per-logical-unit multimodal metadata for KV-cache hashing (C++ layer). @@ -449,6 +462,7 @@ def __post_init__(self): _CPU_ONLY_MULTIMODAL_DATA_KEYS = frozenset({ "multimodal_embed_mask_cumsum", "multimodal_embedding_lengths", + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, }) @@ -1066,6 +1080,7 @@ def find_mm_token_lengths( "mrope_config", "multimodal_embed_mask_cumsum", "multimodal_embedding_lengths", + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, "special_token_offsets", "layout_metadata", }) diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 6712c40196c4..225c75bc1ab2 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -13,13 +13,14 @@ # limitations under the License. # # SPDX-License-Identifier: Apache-2.0 +"""Input processor registry and multimodal preprocessing helpers.""" import enum import traceback from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import (Any, Callable, ClassVar, Dict, List, Optional, Protocol, - Tuple, Type, TypeVar, Union) +from typing import (Any, Callable, ClassVar, Dict, List, NamedTuple, Optional, + Protocol, Tuple, Type, TypeVar, Union) import torch from PIL import Image @@ -32,7 +33,8 @@ from ..sampling_params import SamplingParams from .content_format import ContentFormat from .data import TextPrompt -from .multimodal import (MultimodalInput, _as_cpu_tensor, _compute_mm_masks, +from .multimodal import (MULTIMODAL_ENCODER_ITEM_METADATA_KEY, MultimodalInput, + _as_cpu_tensor, _compute_mm_masks, _find_mm_embedding_lengths_from_masks, _find_mm_token_runs_from_mask, _find_mm_token_start_pos_from_masks, apply_mm_hashes, @@ -55,6 +57,81 @@ def _hash_mm_processor_kwargs(mm_processor_kwargs: Dict[str, Any], return hasher.hexdigest() +class MultimodalEncoderItemMetadata(NamedTuple): + """Prompt-ordered metadata for atomic multimodal encoder items.""" + + item_refs: List[Tuple[str, int]] + """``(modality, local index)`` per atomic item, in prompt order. + + ``modality`` is ``"image"``/``"video"``/``"audio"`` and the local index + identifies the item within that modality's processed payload (for + example, its row range in ``pixel_values``), so an item can be sliced + out for a single-item encoder call. + + TODO(TRTLLM-14477): `MultimodalParams.mm_item_order` carries the same prompt-order + manifest for the mixed-modality full-request encode path. Single-source + the two representations at the input processor so one manifest serves + both the item-scheduling path and the full-request interleave. + """ + + encoder_token_lengths: List[int] + """Physical encoder attention-token cost of each item. + + This is what one encoder forward actually attends over (for example, + pre-merger patch tokens for Qwen ViTs) and is the unit the scheduler + charges against ``encoder_max_num_tokens``. + """ + + output_embedding_lengths: List[int] + """Embedding rows each item contributes to the LLM prompt. + + Must equal the item's placeholder span in the prompt (the + ``multimodal_embedding_lengths`` contract); used to size and split the + encoder output back into per-item tensors. + """ + + def validate(self) -> None: + """Enforce the cross-field contract on producer-supplied metadata.""" + if not all(isinstance(values, list) for values in self): + raise TypeError( + "Multimodal encoder item metadata fields must be lists") + if not all( + isinstance(length, int) + for lengths in (self.encoder_token_lengths, + self.output_embedding_lengths) + for length in lengths): + raise TypeError( + "Multimodal encoder item lengths must contain only integers") + if not (len(self.item_refs) == len(self.encoder_token_lengths) == len( + self.output_embedding_lengths)): + raise ValueError( + "Multimodal encoder item references and lengths must align") + if any(length <= 0 for length in self.encoder_token_lengths): + raise ValueError( + "Multimodal encoder token lengths must be positive") + if any(length <= 0 for length in self.output_embedding_lengths): + raise ValueError( + "Multimodal encoder output embedding lengths must be positive") + + +def get_multimodal_encoder_item_metadata( + multimodal_data: Optional[Dict[str, Any]], +) -> Optional[MultimodalEncoderItemMetadata]: + """Return typed atomic-item metadata from Python multimodal data.""" + if multimodal_data is None: + return None + if not isinstance(multimodal_data, dict): + raise TypeError("multimodal_data must be a dict") + metadata = multimodal_data.get(MULTIMODAL_ENCODER_ITEM_METADATA_KEY) + if metadata is None: + return None + if not isinstance(metadata, MultimodalEncoderItemMetadata): + raise TypeError(f"{MULTIMODAL_ENCODER_ITEM_METADATA_KEY} must be a " + "MultimodalEncoderItemMetadata") + metadata.validate() + return metadata + + class InputProcessor(Protocol): """ Protocol for InputProcessor classes. @@ -177,6 +254,22 @@ class BaseMultimodalInputProcessor(ABC): # inputs to `call_with_token_ids` instead of detokenizing upstream. supports_token_id_mm_expansion: ClassVar[bool] = False + # Whether this processor provides complete atomic-item metadata for the + # multimodal encoder runtime scheduler. + supports_mm_encoder_item_scheduling: ClassVar[bool] = False + + def get_mm_encoder_item_metadata( + self, + prompt_token_ids: List[int], + multimodal_data: Dict[str, Any], + ) -> Optional[MultimodalEncoderItemMetadata]: + """Return item refs, physical costs, and encoder output lengths. + + Models opting into runtime MM encoder scheduling override this hook. + The default keeps legacy multimodal processors unchanged. + """ + return None + def __init__(self, model_path, config, @@ -641,6 +734,22 @@ def get_mm_max_tokens_per_item(self) -> Dict[str, int]: """ return {} + def get_mm_encoder_attention_metadata_capacity( + self, max_num_items: int, + max_num_tokens: int) -> Optional[Dict[str, int]]: + """Return a processor-geometry-aware encoder sequence capacity. + + The keys identify model-specific attention metadata objects (for + example, Qwen2.5-VL has separate ``full_attention`` and + ``window_attention`` entries). ``None`` keeps the encoder model's + conservative fallback mapping. Concrete processors should return a + positive upper bound derived from the startup item/token budgets and + the same geometry constraints used to normalize runtime media. + + The default intentionally ignores both inputs. + """ + return None + def get_preferred_media_io_kwargs(self) -> Dict[str, Dict[str, Any]]: """Per-modality media-IO decode defaults for this model. @@ -654,12 +763,17 @@ 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: Optional[torch.dtype] = None, ) -> Dict[str, Any]: """Build the worst-case dummy ``multimodal_data`` per modality budget. The modality-agnostic entry the KV-cache encoder profiler calls, sizing - each modality to saturate its share of the token budget. + each modality to saturate its share of the token budget. When + ``max_items_per_modality`` is provided, concrete builders should spread + that budget across up to the requested number of equal-sized items; + this lets profiling cover the many-item boundary in addition to the + longest-item boundary. Returns the ``multimodal_data`` dict the model's encoder consumes (e.g. ``{"image": {"pixel_values": ..., "image_grid_thw": ...}}`` for Qwen-VL). @@ -1306,6 +1420,47 @@ def input_processor_wrapper( maybe_compute_mm_embed_cumsum(prompt_token_ids, extra_processed_inputs, input_processor) + if extra_processed_inputs is None: + return prompt_token_ids, extra_processed_inputs + + multimodal_data = extra_processed_inputs.get("multimodal_data") + if (not isinstance(multimodal_data, dict) + or multimodal_data.get("multimodal_embedding") is not None): + return prompt_token_ids, extra_processed_inputs + + # `getattr` because unregistered/text-only models wrap a + # `DefaultInputProcessor`, which is not a + # `BaseMultimodalInputProcessor` and lacks the class var. + supports_item_scheduling = getattr( + input_processor, "supports_mm_encoder_item_scheduling", False) + has_raw_payload = any( + isinstance(multimodal_data.get(modality), dict) + for modality in ("image", "video", "audio")) + if not (supports_item_scheduling and has_raw_payload): + return prompt_token_ids, extra_processed_inputs + + item_metadata = input_processor.get_mm_encoder_item_metadata( + prompt_token_ids, multimodal_data) + if not isinstance(item_metadata, MultimodalEncoderItemMetadata): + raise TypeError( + "get_mm_encoder_item_metadata() must return item metadata " + "(a MultimodalEncoderItemMetadata) for raw multimodal " + f"payloads, got {type(item_metadata).__name__}") + item_metadata.validate() + + multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] = item_metadata + existing_embedding_lengths = multimodal_data.get( + "multimodal_embedding_lengths") + if existing_embedding_lengths is not None: + if not isinstance(existing_embedding_lengths, list): + raise TypeError("multimodal_embedding_lengths must be a list") + if (existing_embedding_lengths + != item_metadata.output_embedding_lengths): + raise ValueError( + "Computed multimodal encoder embedding lengths " + "do not match the existing prompt metadata") + multimodal_data[ + "multimodal_embedding_lengths"] = item_metadata.output_embedding_lengths return prompt_token_ids, extra_processed_inputs return input_processor_wrapper diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a11f79a6ec07..c8d8dd77572e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -571,17 +571,33 @@ class MultimodalConfig(StrictBaseModel): encoder_cache_max_bytes: NonNegativeInt = Field( default=134_217_728, # 128 MiB. description= - ("Maximum bytes for the per-model cross-request multimodal encoder embedding cache. " - "Set to 0 to disable. String values such as '512MB' and '1GiB' use binary units. " - "Cache entries are per multimodal item, but reuse is all-or-nothing for each request: " - "every item in the request must hit the cache before cached embeddings are reused. " - "Only single-modality requests are cacheable for the time being. " - "For the time being, this is incompatible with encoder_side_stream_max_ahead > 0. " + ("Maximum bytes for the per-model multimodal encoder embedding cache. " + "String values such as '512MB' and '1GiB' use binary units. " + "For models with item-level encoder scheduling this cache is the " + "exclusive storage for encoder outputs: its budget bounds encoder-" + "output GPU residency (reserved during memory profiling), entries " + "are reused per item across requests, and values below one prefill " + "iteration's embeddings (including 0) are raised to that minimum " + "with a warning. For other multimodal models it is a cross-request " + "cache with all-or-nothing per-request reuse (every item must hit " + "before cached embeddings are used); 0 disables it, and only " + "single-modality requests are cacheable for the time being. " + "For the time being, this is incompatible with " + "encoder_side_stream_max_ahead > 0. " "NOTE: This is only valid for child implementations of the `MultimodalModelMixin`." ), status="prototype", ) + enable_eager_encoder_scheduling: bool = Field( + default=False, + description=( + "Schedule encoder work for active multimodal requests before LLM " + "capacity filtering. This may advance encoder work for requests " + "that are not selected for the current LLM batch."), + status="prototype", + ) + video_pruning_rate: Optional[float] = Field( default=None, ge=0.0, @@ -4748,32 +4764,54 @@ def init_multimodal_config(cls, v): description="DWDP (Distributed Weight Data Parallelism) config.", status="prototype") - encoder_max_batch_size: Optional[int] = Field( + encoder_max_num_items: Optional[int] = Field( default=None, - description=( - "Maximum batch size for the multimodal encoder's AttentionMetadata. " - "Falls back to `max_batch_size` when unset. This budget is shared " - "proportionately across all modalities the model encodes, not set " - "per modality; per-modality knobs may be added later."), + description= + ("Maximum number of atomic multimodal items (e.g. one image or one " + "video) scheduled for encoder execution in one iteration, shared " + "across requests and modalities; per-modality limits may be added " + "later. Item size is budgeted separately by `encoder_max_num_tokens`. " + "This count also defines the many-item warmup boundary and the " + "default encoder attention capacity for encoders whose items map " + "1:1 to attention sequences; encoders that split one item into " + "multiple sequences derive their capacity from the token budget " + "instead. Falls back to `max_batch_size` when unset."), status="prototype") encoder_max_num_tokens: Optional[int] = Field( default=None, description=( - "Maximum number of tokens for the multimodal encoder's " - "AttentionMetadata. Falls back to `max_num_tokens` when unset. This " - "budget is shared proportionately across all modalities the model " - "encodes, not set per modality; per-modality knobs may be added " - "later."), + "Maximum number of encoder attention tokens scheduled in one " + "multimodal encoder iteration. It falls back to `max_num_tokens` " + "when unset. Because an atomic multimodal item cannot be split, " + "the effective budget is raised to the model's largest atomic item " + "when necessary. This budget is shared proportionately across all " + "encoded modalities; per-modality knobs may be added later."), status="prototype") - @field_validator("encoder_max_batch_size", "encoder_max_num_tokens") + @field_validator("encoder_max_num_items", "encoder_max_num_tokens") @classmethod def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: if v is not None and v <= 0: raise ValueError("must be a positive integer when set") return v + @model_validator(mode="after") + def validate_eager_encoder_scheduling_compatibility(self) -> 'TorchLlmArgs': + if not self.multimodal_config.enable_eager_encoder_scheduling: + return self + if self.enable_attention_dp: + raise ValueError( + "multimodal_config.enable_eager_encoder_scheduling does not " + "yet support attention DP (enable_attention_dp=True)") + if (self.cache_transceiver_config is not None + and self.cache_transceiver_config.backend is not None): + raise ValueError( + "multimodal_config.enable_eager_encoder_scheduling does not " + "yet support disaggregated serving " + "(cache_transceiver_config)") + return self + attn_backend: str = Field( default='TRTLLM', description="Attention backend to use.", @@ -5085,15 +5123,15 @@ def quant_config(self, value: QuantConfig): self._quant_config = value def get_encoder_runtime_sizes(self) -> Tuple[int, int]: - """Return encoder runtime batch and token limits. + """Return encoder runtime item-count and token limits. - Returns `(encoder_max_batch_size, encoder_max_num_tokens)`, falling + Returns `(encoder_max_num_items, encoder_max_num_tokens)`, falling back to the LLM-side `max_batch_size` / `max_num_tokens` when the encoder-specific knobs are not set. """ return ( - self.encoder_max_batch_size - if self.encoder_max_batch_size is not None else self.max_batch_size, + self.encoder_max_num_items + if self.encoder_max_num_items is not None else self.max_batch_size, self.encoder_max_num_tokens if self.encoder_max_num_tokens is not None else self.max_num_tokens, ) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 739db09bfb23..5c533bbde9fc 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -479,7 +479,7 @@ "annotation": "Optional[int]", "converter": "", "kind": "value", - "path": "encoder_max_batch_size" + "path": "encoder_max_num_items" }, { "allowed_values": [], @@ -973,6 +973,13 @@ "kind": "value", "path": "moe_tensor_parallel_size" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "multimodal_config.enable_eager_encoder_scheduling" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 8a5be984f312..337159ec1114 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Tests for KV cache token estimation in KvCacheCreator._get_token_num_for_estimation. Guards the ADP (Attention Data Parallelism) cache-block reduction: when @@ -11,11 +14,13 @@ from unittest.mock import Mock, patch import pytest +import torch from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MultimodalConfig # --------------------------------------------------------------------------- @@ -49,6 +54,56 @@ class _EncoderCacheMultimodalModel(_MultimodalModel): supports_encoder_cache = True +def test_encoder_profiling_builds_and_runs_long_and_many_item_boundaries( + monkeypatch, +): + class _InputProcessor: + def get_mm_max_tokens_per_item(self): + return {"image": 65536} + + def get_dummy_mm_data_for_tokens( + self, + *, + max_tokens_per_modality, + max_items_per_modality, + dtype, + ): + del max_tokens_per_modality, dtype + return {"image": {"num_items": max_items_per_modality["image"]}} + + class _Model(MultimodalModelMixin): + dtype = torch.float32 + + def __init__(self): + self.forwarded_item_counts = [] + + def encode_multimodal_inputs(self, multimodal_params): + count = multimodal_params[0].multimodal_data["image"]["num_items"] + self.forwarded_item_counts.append(count) + return torch.tensor([count]) + + model = _Model() + creator = object.__new__(KvCacheCreator) + creator._model_engine = SimpleNamespace( + model=model, + input_processor=_InputProcessor(), + encoder_max_num_items=8, + encoder_max_num_tokens=8192, + ) + creator._profiling_stage_data = {"enable_mm_reqs": True} + + batches = creator._create_dummy_encoder_inputs() + assert [batch[0].multimodal_data["image"]["num_items"] for batch in batches] == [1, 8] + + monkeypatch.setattr(MultimodalParams, "to_device", lambda self, *args, **kwargs: self) + creator._dummy_encoder_inputs = batches + output = creator._encode_dummy_inputs() + + assert model.forwarded_item_counts == [1, 8] + assert output.tolist() == [8] + assert all(not batch for batch in creator._dummy_encoder_inputs) + + def _make_creator( tokens_per_block, dummy_reqs, @@ -216,6 +271,21 @@ def test_kv_cache_estimation_reserves_multimodal_encoder_cache( assert creator._reserve_multimodal_encoder_cache_memory(1000) == expected +def test_reserve_uses_manager_budget_for_item_scheduling_models(): + from tensorrt_llm._torch.pyexecutor.multimodal_encoder_cache_manager import ( + MultimodalEncoderCacheManager, + ) + + creator = object.__new__(KvCacheCreator) + # The min-clamped manager budget (not the raw config bytes) bounds all + # encoder-output residency for item-scheduling models. + creator._model_engine = SimpleNamespace( + mm_encoder_cache_manager=MultimodalEncoderCacheManager(512, name="t") + ) + + assert creator._reserve_multimodal_encoder_cache_memory(1000) == 1512 + + # --------------------------------------------------------------------------- # VSWA hybrid attention pool-group scaling (Gemma4 hybrid MMMU Pro hang fix) # --------------------------------------------------------------------------- @@ -601,6 +671,9 @@ def test_estimation_temporarily_uses_inferred_pool_sizing() -> None: ) model_engine = Mock() model_engine.model.model_config.attn_backend = "TRTLLM" + # A bare Mock would auto-create the manager attribute; real engines set + # it to None unless the model opted into MM item scheduling. + model_engine.mm_encoder_cache_manager = None llm_args = Mock(cache_transceiver_config=None) with patch.object( diff --git a/tests/unittest/_torch/executor/test_multimodal_encoder_cache_manager.py b/tests/unittest/_torch/executor/test_multimodal_encoder_cache_manager.py new file mode 100644 index 000000000000..829f12821c13 --- /dev/null +++ b/tests/unittest/_torch/executor/test_multimodal_encoder_cache_manager.py @@ -0,0 +1,159 @@ +# Copyright 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. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.multimodal_encoder_cache_manager import ( + MultimodalEncoderCacheManager, +) + + +def _tensor(rows: int) -> torch.Tensor: + # 1 row == 8 bytes (2 fp32 elements) for easy byte math. + return torch.full((rows, 2), 1.0, dtype=torch.float32) + + +ROW_BYTES = 8 + + +def test_adopt_takes_ownership_without_clone(): + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + value = _tensor(4) + stored = manager.adopt("k", value, request_id=1) + assert stored.untyped_storage().data_ptr() == value.untyped_storage().data_ptr() + assert manager.current_bytes == 4 * ROW_BYTES + assert manager.held_bytes == 4 * ROW_BYTES + + +def test_adopt_materializes_views_of_shared_storage(): + # Grouped encoder forwards split one batch tensor into per-item views; + # storing the views directly would let a surviving sibling entry keep + # the whole batch allocation resident after another entry is evicted, + # decoupling accounted bytes from physical memory. + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + batch = torch.full((4, 2), 1.0, dtype=torch.float32) + view_a, view_b = torch.split(batch, 2, dim=0) + + stored_a = manager.adopt("a", view_a, request_id=1) + stored_b = manager.adopt("b", view_b, request_id=2) + + batch_ptr = batch.untyped_storage().data_ptr() + assert stored_a.untyped_storage().data_ptr() != batch_ptr + assert stored_b.untyped_storage().data_ptr() != batch_ptr + assert stored_a.untyped_storage().data_ptr() != stored_b.untyped_storage().data_ptr() + # Accounted bytes equal each entry's exclusive storage. + assert manager.current_bytes == 4 * ROW_BYTES + + +def test_duplicate_adopt_collapses_to_one_entry(): + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + first = manager.adopt("k", _tensor(4), request_id=1) + second = manager.adopt("k", _tensor(4), request_id=2) + assert second is first # single resident copy, both requests holding + assert manager.current_bytes == 4 * ROW_BYTES + assert len(manager) == 1 + assert manager.stats().dedup_adoptions == 1 + + # Entry stays held until the *last* referencing request releases it. + manager.release_holds(1) + assert manager.held_bytes == 4 * ROW_BYTES + manager.release_holds(2) + assert manager.held_bytes == 0 + assert manager.current_bytes == 4 * ROW_BYTES # still resident for reuse + + +def test_held_entries_are_never_evicted(): + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + manager.adopt("a", _tensor(6), request_id=1) + manager.adopt("b", _tensor(4), request_id=1) + + assert not manager.can_allocate(1 * ROW_BYTES) + with pytest.raises(RuntimeError, match="can_allocate"): + manager.adopt("c", _tensor(1), request_id=2) + assert manager.contains("a") and manager.contains("b") + + +def test_zero_ref_lru_eviction_makes_room(): + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + manager.adopt("old", _tensor(6), request_id=1) + manager.adopt("new", _tensor(4), request_id=2) + manager.release_holds(1) + + assert manager.can_allocate(6 * ROW_BYTES) + manager.adopt("c", _tensor(6), request_id=3) + assert not manager.contains("old") # zero-ref LRU victim + assert manager.contains("new") and manager.contains("c") + assert manager.stats().evictions == 1 + + +def test_get_and_hold_revives_freeable_entry(): + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + manager.adopt("k", _tensor(6), request_id=1) + manager.release_holds(1) + assert manager.held_bytes == 0 + + hit = manager.get_and_hold("k", request_id=2) + assert hit is not None + assert manager.held_bytes == 6 * ROW_BYTES + # Re-held entry is protected again. + with pytest.raises(RuntimeError): + manager.adopt("big", _tensor(5), request_id=3) + + +def test_get_and_hold_miss_returns_none(): + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + assert manager.get_and_hold("absent", request_id=1) is None + assert manager.stats().misses == 1 + + +def test_release_holds_is_idempotent_and_scoped(): + manager = MultimodalEncoderCacheManager(20 * ROW_BYTES) + manager.adopt("a", _tensor(4), request_id=1) + manager.adopt("b", _tensor(4), request_id=1) + manager.adopt("c", _tensor(4), request_id=2) + + manager.release_holds(1) + manager.release_holds(1) # idempotent: cancel path may race normal strip + manager.release_holds(99) # unknown request is a no-op + assert manager.held_bytes == 4 * ROW_BYTES # only request 2's entry + + +def test_reserved_bytes_shrink_allocatable_space(): + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + assert manager.can_allocate(4 * ROW_BYTES, reserved_bytes=6 * ROW_BYTES) + assert not manager.can_allocate(5 * ROW_BYTES, reserved_bytes=6 * ROW_BYTES) + + +def test_free_resources_is_the_resource_manager_teardown_hook(): + from types import SimpleNamespace + + manager = MultimodalEncoderCacheManager(10 * ROW_BYTES) + manager.adopt("k", _tensor(4), request_id=7) + + # BaseResourceManager surface: teardown by LlmRequest, zero + # scheduler-facing resource counts (budgeting happens at MM item + # selection, not capacity admission). + request = SimpleNamespace(request_id=7) + assert manager.get_max_resource_count() == 0 + assert manager.get_needed_resource_to_completion(request) == 0 + manager.free_resources(request) + assert manager.held_bytes == 0 + manager.free_resources(request) # idempotent like release_holds + + +def test_rejects_nonpositive_budget(): + with pytest.raises(ValueError): + MultimodalEncoderCacheManager(0) diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py new file mode 100644 index 000000000000..dfc28a727214 --- /dev/null +++ b/tests/unittest/_torch/executor/test_multimodal_scheduler.py @@ -0,0 +1,816 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.models.modeling_mistral import Mistral3InputProcessor +from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2VLInputProcessorBase +from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem +from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequest, + MultimodalEncoderProgress, + MultimodalEncoderRequestState, + get_multimodal_encoder_token_lengths, + initialize_multimodal_encoder_request, + is_multimodal_encoder_ready, +) +from tensorrt_llm._torch.pyexecutor.model_engine import ( + PyTorchModelEngine, + _resolve_mm_encoder_token_budget, +) +from tensorrt_llm._torch.pyexecutor.multimodal_encoder_cache_manager import ( + MultimodalEncoderCacheManager, +) +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ( + MultimodalEagerEncoderScheduler, + MultimodalScheduler, +) +from tensorrt_llm._torch.pyexecutor.scheduler.waiting_queue import FCFSWaitingQueue +from tensorrt_llm.bindings import SamplingConfig +from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, + MultimodalParams, + strip_mm_encoder_inputs, +) +from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata + + +class _CapacityScheduler: + def schedule_request(self, requests): + return list(requests), [], [] + + +class _RejectMultimodalCapacityScheduler: + def schedule_request(self, requests): + fitting = [request for request in requests if request.py_mm_encoder_state is None] + return fitting, [], [] + + +class _MicroBatchScheduler: + def schedule(self, requests, inflight_request_ids): + del inflight_request_ids + return [], list(requests), [] + + +class _BaseScheduler: + def __init__(self): + self.capacity_scheduler = _CapacityScheduler() + self.micro_batch_scheduler = _MicroBatchScheduler() + + def can_schedule(self, requests): + return bool(requests) + + +def _llm_request(request_id, multimodal_data=None): + return LlmRequest( + request_id=request_id, + max_new_tokens=1, + input_tokens=[1, 2, 3], + sampling_config=SamplingConfig(), + is_streaming=False, + py_multimodal_data=multimodal_data, + ) + + +def _request(request_id, costs, *, ready=()): + request = _llm_request( + request_id, + multimodal_data={ + "image": {"pixel_values": torch.empty(len(costs), 1)}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + item_refs=[("image", item_idx) for item_idx in range(len(costs))], + encoder_token_lengths=costs, + output_embedding_lengths=[1] * len(costs), + ), + "multimodal_embedding_lengths": [1] * len(costs), + }, + ) + initialize_multimodal_encoder_request(request, max_num_tokens=1 << 30) + for item_idx in ready: + request.py_mm_encoder_state.outputs[item_idx] = torch.empty(1) + return request + + +def test_mm_encoder_token_lengths_distinguishes_missing_and_invalid_data(): + request = _llm_request(1) + + assert get_multimodal_encoder_token_lengths(request) is None + + request.py_multimodal_data = [] + with pytest.raises(TypeError, match="multimodal_data must be a dict"): + get_multimodal_encoder_token_lengths(request) + + +def test_mm_encoder_readiness_is_derived_from_request_local_outputs(): + request = _request(1, [4, 4]) + assert request.py_mm_encoder_state.progress is MultimodalEncoderProgress.PENDING + assert not is_multimodal_encoder_ready(request) + + request.py_mm_encoder_state.outputs[0] = torch.empty(1) + assert request.py_mm_encoder_state.progress is MultimodalEncoderProgress.PARTIAL + assert not is_multimodal_encoder_ready(request) + + request.py_mm_encoder_state.outputs[1] = torch.empty(1) + assert is_multimodal_encoder_ready(request) + + # A precomputed-embedding request never gets item state in the first + # place (initialize skips it), and post-prefill strip drops the state: + # both report ready through the state-absence branch. + request.py_mm_encoder_state = None + assert is_multimodal_encoder_ready(request) + + +def test_multimodal_scheduler_keeps_items_atomic_and_backfills_requests(): + scheduler = MultimodalScheduler(_BaseScheduler(), max_num_items=2, max_num_tokens=10) + first = _request(1, [7, 7]) + second = _request(2, [3]) + + output = scheduler.schedule_request([first, second], set()) + + assert output.scheduled_mm_encoder_items == {1: [0], 2: [0]} + assert output.context_requests == [second] + + +def test_scheduler_defers_items_beyond_cache_byte_budget(): + # Budget hosts exactly one 1-row item (4 bytes): the second request's + # item must wait even though the token/item budgets would admit it + # (allocate-before-compute). + manager = MultimodalEncoderCacheManager(4, name="test") + scheduler = MultimodalScheduler( + _BaseScheduler(), + max_num_items=8, + max_num_tokens=1 << 20, + cache_manager=manager, + embedding_row_bytes=4, + ) + first = _request(1, [3]) + second = _request(2, [3]) + + output = scheduler.schedule_request([first, second], set()) + + assert output.scheduled_mm_encoder_items == {1: [0]} + assert output.context_requests == [first] + + +def test_head_of_line_reservation_blocks_later_requests(): + # Token budget splits the head request across iterations; its unencoded + # remainder reserves manager bytes, so the request behind it cannot + # squat the space the head needs to ever complete (deadlock avoidance). + manager = MultimodalEncoderCacheManager(8, name="test") + scheduler = MultimodalScheduler( + _BaseScheduler(), + max_num_items=8, + max_num_tokens=5, + cache_manager=manager, + embedding_row_bytes=4, + ) + head = _request(1, [5, 5]) # second item exceeds this iteration's tokens + follower = _request(2, [3]) + + output = scheduler.schedule_request([head, follower], set()) + + # 8 bytes total: head's item 0 claims 4, its pending item 1 reserves the + # other 4, leaving nothing for the follower despite free space. + assert output.scheduled_mm_encoder_items == {1: [0]} + assert output.context_requests == [] + + +def test_admission_rejects_requests_larger_than_output_budget(): + # A long-video request whose total embedding footprint can never fit + # the storage budget fails at admission (failing only that request), + # with guidance to raise encoder_cache_max_bytes. Reachable once LLM + # chunked prefill admits prompts longer than max_num_tokens. + request = _llm_request( + 1, + multimodal_data={ + "video": {"pixel_values_videos": torch.empty(3, 1)}, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + item_refs=[("video", 0)], + encoder_token_lengths=[12], + output_embedding_lengths=[3], + ), + "multimodal_embedding_lengths": [3], + }, + ) + with pytest.raises(ValueError, match="raise the cache budget"): + initialize_multimodal_encoder_request( + request, + max_num_tokens=1 << 30, + max_output_bytes=2 * 4, # fits 2 rows; the video needs 3 + embedding_row_bytes=4, + ) + + +def test_oversized_request_fails_fast_instead_of_starving(): + manager = MultimodalEncoderCacheManager(4, name="test") + scheduler = MultimodalScheduler( + _BaseScheduler(), + max_num_items=8, + max_num_tokens=1 << 20, + cache_manager=manager, + embedding_row_bytes=4, + ) + request = _request(1, [3, 3]) # 2 rows = 8 bytes > 4-byte budget + + with pytest.raises(RuntimeError, match="raise the cache budget"): + scheduler.schedule_request([request], set()) + + +def test_scheduler_requires_row_bytes_alongside_manager(): + with pytest.raises(ValueError, match="embedding_row_bytes"): + MultimodalScheduler( + _BaseScheduler(), + max_num_items=1, + max_num_tokens=1, + cache_manager=MultimodalEncoderCacheManager(4, name="test"), + ) + + +def test_multimodal_scheduler_selects_all_items_and_admits_request_when_batch_fits(): + scheduler = MultimodalScheduler(_BaseScheduler(), max_num_items=2, max_num_tokens=10) + request = _request(1, [6, 4]) + + output = scheduler.schedule_request([request], set()) + + # The encoder step is the single encode site: an in-budget batch simply + # has every pending item selected, and the request still enters the LLM + # batch in the same iteration (encode runs before the LLM forward). + assert output.scheduled_mm_encoder_items == {1: [0, 1]} + assert output.context_requests == [request] + + +def test_multimodal_scheduler_withholds_request_on_budget_overflow(): + scheduler = MultimodalScheduler(_BaseScheduler(), max_num_items=3, max_num_tokens=10) + request = _request(1, [6, 4, 1]) + + output = scheduler.schedule_request([request], set()) + + assert output.scheduled_mm_encoder_items == {1: [0, 1]} + assert output.context_requests == [] + + +def test_multimodal_scheduler_preserves_non_multimodal_requests(): + scheduler = MultimodalScheduler(_BaseScheduler(), max_num_items=1, max_num_tokens=1) + request = _llm_request(1) + initialize_multimodal_encoder_request(request, max_num_tokens=1) + + output = scheduler.schedule_request([request], set()) + + assert output.scheduled_mm_encoder_items is None + assert output.context_requests == [request] + + +def test_encoder_token_budget_auto_raises_for_atomic_item(): + assert _resolve_mm_encoder_token_budget(8192, 65536) == 65536 + + +def test_request_rejects_item_above_effective_startup_maximum(): + request = _request(1, [9]) + request.py_multimodal_data["image"] = {"pixel_values": torch.empty(1)} + + with pytest.raises(ValueError, match="exceeding the effective startup maximum 8"): + initialize_multimodal_encoder_request(request, max_num_tokens=8) + + +def test_eager_scheduler_encodes_request_rejected_by_llm_capacity(): + base_scheduler = _BaseScheduler() + base_scheduler.capacity_scheduler = _RejectMultimodalCapacityScheduler() + scheduler = MultimodalEagerEncoderScheduler(base_scheduler, max_num_items=1, max_num_tokens=8) + multimodal_request = _request(1, [8]) + text_request = _llm_request(2) + initialize_multimodal_encoder_request(text_request, max_num_tokens=8) + + output = scheduler.schedule_request([multimodal_request, text_request], set()) + + assert output.scheduled_mm_encoder_items == {1: [0]} + assert output.context_requests == [text_request] + + +def test_forward_multimodal_encoder_step_delegates_to_model_engine(): + calls = [] + executor = object.__new__(PyExecutor) + executor.active_requests = [SimpleNamespace(request_id=1)] + executor.model_engine = SimpleNamespace( + forward_multimodal_encoder_items=lambda requests, items: calls.append((requests, items)) + ) + scheduled_items = {1: [0]} + scheduled_requests = SimpleNamespace(scheduled_mm_encoder_items=scheduled_items) + + executor._forward_multimodal_encoder_step(scheduled_requests) + + assert calls == [(executor.active_requests, scheduled_items)] + + +def _executor_for_mm_admission(active_requests, *, max_num_tokens=8): + executor = object.__new__(PyExecutor) + executor.enable_attention_dp = False + executor.dist = SimpleNamespace(tp_size=1) + executor.max_num_active_requests = 8 + executor.is_benchmark_disagg = False + executor._supports_mm_encoder_item_scheduling = True + executor.model_engine = SimpleNamespace( + encoder_max_num_items=8, + encoder_max_num_tokens=max_num_tokens, + ) + executor.active_requests = active_requests + return executor + + +def _waiting_item(request_id, costs=None): + multimodal_data = None + if costs is not None: + multimodal_data = { + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + item_refs=[("image", item_idx) for item_idx in range(len(costs))], + encoder_token_lengths=costs, + output_embedding_lengths=[1] * len(costs), + ), + "multimodal_embedding_lengths": [1] * len(costs), + } + return RequestQueueItem( + request_id, + _llm_request(request_id, multimodal_data=multimodal_data), + ) + + +def test_mm_admission_does_not_charge_ready_active_request(): + active = _request(1, [8], ready=(0,)) + waiting = FCFSWaitingQueue([_waiting_item(2, [8])]) + executor = _executor_for_mm_admission([active]) + + admitted = executor._pop_from_waiting_queue(waiting, 1) + + assert [item.id for item in admitted] == [2] + assert not waiting + + +def test_mm_admission_passes_oversized_request_to_validation(): + waiting = FCFSWaitingQueue([_waiting_item(1, [9]), _waiting_item(2, None)]) + executor = _executor_for_mm_admission([], max_num_tokens=8) + + admitted = executor._pop_from_waiting_queue(waiting, 0) + + assert [item.id for item in admitted] == [1, 2] + assert not waiting + + +def test_item_encoder_slices_and_restores_selected_item_order(): + class _Model(MultimodalModelMixin): + def encode_multimodal_inputs(self, multimodal_params): + return torch.cat( + [param.multimodal_data["image"]["pixel_values"] for param in multimodal_params] + ) + + multimodal_param = MultimodalParams( + multimodal_data={ + "image": { + "pixel_values": torch.arange(5).unsqueeze(1), + "image_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 3]]), + }, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + item_refs=[("image", 0), ("image", 1)], + encoder_token_lengths=[2, 3], + output_embedding_lengths=[2, 3], + ), + "multimodal_embedding_lengths": [2, 3], + } + ) + + model = _Model() + encoder_inputs = model.prepare_multimodal_encoder_inputs( + [(multimodal_param, 1), (multimodal_param, 0)] + ) + outputs = model.forward_multimodal_encoder_items(encoder_inputs) + + assert [output.squeeze(1).tolist() for output in outputs] == [ + [2, 3, 4], + [0, 1], + ] + + +def test_prepare_multimodal_encoder_inputs_slices_before_device_transfer(): + multimodal_param = MultimodalParams( + multimodal_data={ + "image": { + "pixel_values": torch.arange(5).unsqueeze(1), + "image_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 3]]), + }, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + item_refs=[("image", 0), ("image", 1)], + encoder_token_lengths=[2, 3], + output_embedding_lengths=[2, 3], + ), + "multimodal_embedding_lengths": [2, 3], + } + ) + + encoder_inputs = MultimodalModelMixin.prepare_multimodal_encoder_inputs( + MultimodalModelMixin(), [(multimodal_param, 1)] + ) + + item_param, embedding_length, modality = encoder_inputs[0] + assert modality == "image" + assert embedding_length == 3 + assert item_param.multimodal_data["image"]["pixel_values"].squeeze(1).tolist() == [2, 3, 4] + assert multimodal_param.multimodal_data["image"]["pixel_values"].shape[0] == 5 + + +def test_prepare_multimodal_encoder_inputs_rejects_invalid_metadata_types(): + multimodal_param = MultimodalParams( + multimodal_data={ + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: ("image", 0), + "multimodal_embedding_lengths": [1], + } + ) + + with pytest.raises(TypeError, match="must be a MultimodalEncoderItemMetadata"): + MultimodalModelMixin().prepare_multimodal_encoder_inputs([(multimodal_param, 0)]) + + +def test_strip_mm_encoder_inputs_preserves_embedding_and_runtime_metadata(): + embedding = torch.empty(3, 4) + mm_data = { + "image": {"pixel_values": torch.empty(2, 3)}, + "video": {"pixel_values_videos": torch.empty(2, 3)}, + "multimodal_embedding": embedding, + "multimodal_embed_mask_cumsum": torch.tensor([0, 1]), + } + + strip_mm_encoder_inputs(mm_data) + + assert "image" not in mm_data + assert "video" not in mm_data + assert mm_data["multimodal_embedding"] is embedding + assert "multimodal_embed_mask_cumsum" in mm_data + + +def test_item_outputs_accumulate_in_manager_and_release_raw_data(monkeypatch): + class _Model(MultimodalModelMixin): + def forward_multimodal_encoder_items(self, encoder_inputs): + return [ + torch.full((embedding_length, 2), float(embedding_length)) + for _, embedding_length, _ in encoder_inputs + ] + + monkeypatch.setattr(MultimodalParams, "to_device", lambda self, *args, **kwargs: self) + engine = object.__new__(PyTorchModelEngine) + engine.model = _Model() + engine.supports_mm_encoder_item_scheduling = True + engine.mm_encoder_cache_manager = MultimodalEncoderCacheManager(1 << 20, name="test") + request = _llm_request( + 1, + multimodal_data={ + "image": { + "pixel_values": torch.arange(5).unsqueeze(1), + "image_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 3]]), + }, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + item_refs=[("image", 0), ("image", 1)], + encoder_token_lengths=[2, 3], + output_embedding_lengths=[2, 3], + ), + "multimodal_embedding_lengths": [2, 3], + }, + ) + initialize_multimodal_encoder_request(request, max_num_tokens=8) + assert request.py_mm_encoder_state.embedding_lengths == [2, 3] + + engine.forward_multimodal_encoder_items([request], {1: [0]}) + + # Items encoded across iterations accumulate as pinned manager entries; + # raw inputs stay until the request completes. + assert request.py_mm_encoder_state.outputs[0].shape == (2, 2) + assert request.py_mm_encoder_state.outputs[1] is None + assert engine.mm_encoder_cache_manager.held_bytes == 2 * 2 * 4 + assert "image" in request.py_multimodal_data + + engine.forward_multimodal_encoder_items([request], {1: [1]}) + + published = request.py_multimodal_data["multimodal_embedding"] + assert published == request.py_mm_encoder_state.outputs + assert [slot.tolist() for slot in published] == [ + [[2.0, 2.0], [2.0, 2.0]], + [[3.0, 3.0], [3.0, 3.0], [3.0, 3.0]], + ] + assert "image" not in request.py_multimodal_data + + +def test_qwen_item_metadata_uses_prompt_order_and_pre_merger_costs(): + processor = object.__new__(Qwen2VLInputProcessorBase) + processor._config = SimpleNamespace( + image_token_id=11, + video_token_id=12, + vision_start_token_id=10, + vision_end_token_id=13, + vision_config=SimpleNamespace(spatial_merge_size=2), + ) + prompt_token_ids = [10, 12, 12, 13, 1, 10, 11, 11, 13] + multimodal_data = { + "image": {"image_grid_thw": torch.tensor([[1, 4, 4]])}, + "video": {"video_grid_thw": torch.tensor([[2, 4, 4]])}, + } + + metadata = processor.get_mm_encoder_item_metadata(prompt_token_ids, multimodal_data) + + assert isinstance(metadata, MultimodalEncoderItemMetadata) + assert metadata.item_refs == [("video", 0), ("image", 0)] + assert metadata.encoder_token_lengths == [32, 16] + assert metadata.output_embedding_lengths == [8, 4] + + +def test_qwen_item_metadata_collapses_frame_spans_into_original_video(): + processor = object.__new__(Qwen2VLInputProcessorBase) + processor._config = SimpleNamespace( + image_token_id=11, + video_token_id=12, + vision_start_token_id=10, + vision_end_token_id=13, + vision_config=SimpleNamespace(spatial_merge_size=2), + ) + prompt_token_ids = [10, 12, 13, 100, 10, 12, 13] + multimodal_data = { + "video": {"video_grid_thw": torch.tensor([[2, 4, 4]])}, + } + + metadata = processor.get_mm_encoder_item_metadata(prompt_token_ids, multimodal_data) + + assert metadata.item_refs == [("video", 0)] + assert metadata.encoder_token_lengths == [32] + assert metadata.output_embedding_lengths == [8] + + +def test_mistral_item_metadata_separates_patch_and_embedding_units(): + processor = object.__new__(Mistral3InputProcessor) + processor._vision_geometry = lambda: (14, 2, 3, 1024) + + metadata = processor.get_mm_encoder_item_metadata( + [], {"image": {"image_sizes": [[28, 56], [56, 56]]}} + ) + + assert metadata.item_refs == [("image", 0), ("image", 1)] + assert metadata.encoder_token_lengths == [8, 16] + assert metadata.output_embedding_lengths == [2, 4] + + +# --------------------------------------------------------------------------- +# Encoder cache x item scheduling integration +# --------------------------------------------------------------------------- + + +def _cache_request(request_id, *, hashes, embedding_lengths, kwargs_hash="kw"): + """A cache-keyable item-scheduling request with raw image payload.""" + num_items = len(embedding_lengths) + multimodal_data = { + "image": { + "pixel_values": torch.arange(sum(embedding_lengths)).unsqueeze(1), + "image_grid_thw": torch.tensor([[1, 1, length] for length in embedding_lengths]), + }, + MULTIMODAL_ENCODER_ITEM_METADATA_KEY: MultimodalEncoderItemMetadata( + item_refs=[("image", item_idx) for item_idx in range(num_items)], + encoder_token_lengths=list(embedding_lengths), + output_embedding_lengths=list(embedding_lengths), + ), + "multimodal_embedding_lengths": list(embedding_lengths), + } + if kwargs_hash is not None: + multimodal_data["mm_processor_kwargs_hash"] = kwargs_hash + request = LlmRequest( + request_id=request_id, + max_new_tokens=1, + input_tokens=[1, 2, 3], + sampling_config=SamplingConfig(), + is_streaming=False, + py_multimodal_data=multimodal_data, + multimodal_hashes=hashes, + ) + initialize_multimodal_encoder_request(request, max_num_tokens=1 << 30) + return request + + +def _cache_engine(manager, monkeypatch): + class _Model(MultimodalModelMixin): + def __init__(self): + self.encoded_item_counts = [] + + def forward_multimodal_encoder_items(self, encoder_inputs): + self.encoded_item_counts.append(len(encoder_inputs)) + return [ + torch.full((embedding_length, 2), float(embedding_length)) + for _, embedding_length, _ in encoder_inputs + ] + + monkeypatch.setattr(MultimodalParams, "to_device", lambda self, *args, **kwargs: self) + engine = object.__new__(PyTorchModelEngine) + engine.model = _Model() + engine.supports_mm_encoder_item_scheduling = True + engine.mm_encoder_cache_manager = manager + return engine + + +def _sweep_executor(engine, active_requests): + executor = object.__new__(PyExecutor) + executor.active_requests = active_requests + executor.model_engine = engine + executor._supports_mm_encoder_item_scheduling = True + return executor + + +def test_item_cache_keys_pin_the_full_request_path_format(): + hashes = [[1, 2], [3, 4]] + + keys = MultimodalModelMixin.build_encoder_cache_item_keys( + hashes, [("image", 0), ("video", 0)], [2, 3], "kw" + ) + + # Deliberately identical to `_encoder_cache_item_key` / + # `_encoder_cache_keys`: the manager and the legacy full-request clone + # cache are separate stores today, but keeping one key format lets the + # remaining full-request consumers unify onto the manager later without + # invalidating entries. + assert keys == [("image", (1, 2), 2, "kw"), ("video", (3, 4), 3, "kw")] + + +def test_item_encode_adopts_outputs_into_manager(monkeypatch): + manager = MultimodalEncoderCacheManager(1 << 20, name="test") + engine = _cache_engine(manager, monkeypatch) + request = _cache_request(1, hashes=[[1, 2], [3, 4]], embedding_lengths=[2, 3]) + + engine.forward_multimodal_encoder_items([request], {1: [0, 1]}) + + key0, key1 = MultimodalModelMixin.build_encoder_cache_item_keys( + [[1, 2], [3, 4]], [("image", 0), ("image", 1)], [2, 3], "kw" + ) + # Slots alias the manager-owned entries (single storage, no clones), and + # the entries are held on the request's behalf. + assert manager.get_and_hold(key0, 99) is request.py_mm_encoder_state.outputs[0] + assert manager.get_and_hold(key1, 99) is request.py_mm_encoder_state.outputs[1] + assert manager.held_bytes == manager.current_bytes > 0 + + +def test_sweep_completes_duplicate_request_without_encoding_or_budget(monkeypatch): + manager = MultimodalEncoderCacheManager(1 << 20, name="test") + engine = _cache_engine(manager, monkeypatch) + first = _cache_request(1, hashes=[[1, 2], [3, 4]], embedding_lengths=[2, 3]) + engine.forward_multimodal_encoder_items( + [ + first, + ], + {1: [0, 1]}, + ) + assert engine.model.encoded_item_counts == [2] + + second = _cache_request(2, hashes=[[1, 2], [3, 4]], embedding_lengths=[2, 3]) + executor = _sweep_executor(engine, [second]) + executor._attach_mm_encoder_cache_hits() + + assert engine.model.encoded_item_counts == [2] # no further encoding + assert is_multimodal_encoder_ready(second) + # Published embedding is the prompt-ordered list of held views, and the + # duplicate request shares the first request's storage (dedup). + published = second.py_multimodal_data["multimodal_embedding"] + assert published == second.py_mm_encoder_state.outputs + assert published[0] is first.py_mm_encoder_state.outputs[0] + assert published[1] is first.py_mm_encoder_state.outputs[1] + assert "image" not in second.py_multimodal_data + + scheduler = MultimodalScheduler(_BaseScheduler(), max_num_items=8, max_num_tokens=1 << 20) + output = scheduler.schedule_request([second], set()) + assert output.scheduled_mm_encoder_items is None + assert output.context_requests == [second] + + +def test_partial_hit_routes_to_item_path_and_schedules_only_misses(monkeypatch): + manager = MultimodalEncoderCacheManager(1 << 20, name="test") + engine = _cache_engine(manager, monkeypatch) + request = _cache_request(1, hashes=[[1, 2], [3, 4]], embedding_lengths=[2, 3]) + key0, _ = MultimodalModelMixin.build_encoder_cache_item_keys( + [[1, 2], [3, 4]], [("image", 0), ("image", 1)], [2, 3], "kw" + ) + # Resident zero-ref entry left behind by an earlier request. + manager.adopt(key0, torch.full((2, 2), 7.0), request_id=999) + manager.release_holds(999) + + executor = _sweep_executor(engine, [request]) + executor._attach_mm_encoder_cache_hits() + + assert request.py_mm_encoder_state.progress is MultimodalEncoderProgress.PARTIAL + + # A PARTIAL request must take the item path even though the whole batch + # would fit the budgets, and only the miss is scheduled. + scheduler = MultimodalScheduler(_BaseScheduler(), max_num_items=8, max_num_tokens=1 << 20) + output = scheduler.schedule_request([request], set()) + assert output.scheduled_mm_encoder_items == {1: [1]} + + engine.forward_multimodal_encoder_items([request], {1: [1]}) + assert engine.model.encoded_item_counts == [1] + assert is_multimodal_encoder_ready(request) + + +def test_held_entries_survive_allocation_pressure(monkeypatch): + # Budget fits the pinned request plus one extra row, but not another + # 2-row item: allocation pressure must fail loudly rather than evict a + # held entry, so progress cannot regress. + manager = MultimodalEncoderCacheManager(3 * 2 * 4, name="test") + engine = _cache_engine(manager, monkeypatch) + request = _cache_request(1, hashes=[[1, 2]], embedding_lengths=[2]) + engine.forward_multimodal_encoder_items([request], {1: [0]}) + + assert not manager.can_allocate(2 * 2 * 4) + with pytest.raises(RuntimeError, match="can_allocate"): + manager.adopt("intruder", torch.full((2, 2), 5.0), request_id=2) + + assert is_multimodal_encoder_ready(request) + published = request.py_multimodal_data["multimodal_embedding"] + assert published[0].shape == (2, 2) + + # After the holds are released (post-prefill / termination funnel), the + # same allocation succeeds by evicting the now zero-ref entry. + manager.release_holds(request.request_id) + assert manager.can_allocate(2 * 2 * 4) + + +@pytest.mark.parametrize( + "case", + ["no_hashes", "no_kwargs_hash", "count_mismatch", "flag_off"], +) +def test_key_guards_fall_back_to_request_scoped_storage(case, monkeypatch): + manager = MultimodalEncoderCacheManager(1 << 20, name="test") + engine = _cache_engine(manager, monkeypatch) + if case == "flag_off": + engine.supports_mm_encoder_item_scheduling = False + request = _cache_request( + 1, + hashes=None + if case == "no_hashes" + else ([[1, 2]] if case == "count_mismatch" else [[1, 2], [3, 4]]), + embedding_lengths=[2, 3], + kwargs_hash=None if case == "no_kwargs_hash" else "kw", + ) + + assert engine.get_mm_encoder_item_keys(request) is None + + # The request still runs the item path; outputs land under temporary + # request-scoped keys, never shareable across requests. + engine.forward_multimodal_encoder_items([request], {1: [0, 1]}) + assert is_multimodal_encoder_ready(request) + assert manager.contains(("mm_tmp", 1, 0)) + assert manager.contains(("mm_tmp", 1, 1)) + + +# --------------------------------------------------------------------------- +# MultimodalEncoderRequestState unit behavior +# --------------------------------------------------------------------------- + + +def test_mm_encoder_state_enforces_lengths_slot_invariant(): + with pytest.raises(ValueError, match="one entry per item slot"): + MultimodalEncoderRequestState(embedding_lengths=[2], outputs=[None, None]) + + +def test_mm_encoder_state_progress_and_pending_transitions(): + state = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) + + assert state.progress is MultimodalEncoderProgress.PENDING + assert state.pending_item_indices() == [0, 1] + + state.record(1, torch.ones(3, 2)) + assert state.progress is MultimodalEncoderProgress.PARTIAL + assert state.pending_item_indices() == [0] + + state.record(0, torch.zeros(2, 2)) + assert state.progress is MultimodalEncoderProgress.READY + assert [slot.tolist() for slot in state.outputs] == [ + [[0, 0], [0, 0]], + [[1, 1], [1, 1], [1, 1]], + ] + + +def test_mm_encoder_state_record_rejects_mismatched_outputs(): + state = MultimodalEncoderRequestState.from_embedding_lengths([2, 3]) + + with pytest.raises(ValueError, match="expected 2"): + state.record(0, torch.ones(5, 2)) + + state.record(0, torch.ones(2, 2)) + with pytest.raises(ValueError, match="matching"): + state.record(1, torch.ones(3, 4)) # hidden dim mismatch vs items + with pytest.raises(ValueError, match="already recorded"): + state.record(0, torch.ones(2, 2)) # items encode at most once + + +def test_mm_encoder_state_finalize_into_is_a_conditional_no_op(): + state = MultimodalEncoderRequestState.from_embedding_lengths([2]) + multimodal_data = {"image": {"pixel_values": torch.empty(2, 1)}} + + assert state.finalize_into(multimodal_data) is False + assert "multimodal_embedding" not in multimodal_data + + state.record(0, torch.ones(2, 2)) + assert state.finalize_into(multimodal_data) is True + assert multimodal_data["multimodal_embedding"] == state.outputs + assert "image" not in multimodal_data diff --git a/tests/unittest/_torch/executor/test_scheduler_serializable_output.py b/tests/unittest/_torch/executor/test_scheduler_serializable_output.py index 5f0763869487..b73a367803d3 100644 --- a/tests/unittest/_torch/executor/test_scheduler_serializable_output.py +++ b/tests/unittest/_torch/executor/test_scheduler_serializable_output.py @@ -28,6 +28,7 @@ def test_serializable_scheduler_output_round_trip(): scheduled_requests.context_requests_last_chunk = [request_pool[1], request_pool[2]] scheduled_requests.generation_requests = [request_pool[3]] scheduled_requests.paused_requests = [request_pool[4]] + scheduled_requests.scheduled_mm_encoder_items = {1: [0, 2], 7: [1]} fitting_disagg_gen_init_requests = [request_pool[5], request_pool[6]] num_fitting_requests = 3 @@ -67,4 +68,8 @@ def test_serializable_scheduler_output_round_trip(): assert _request_ids(restored_schedule.paused_requests) == _request_ids( scheduled_requests.paused_requests ) + assert ( + restored_schedule.scheduled_mm_encoder_items + == scheduled_requests.scheduled_mm_encoder_items + ) assert _request_ids(restored_fitting) == _request_ids(fitting_disagg_gen_init_requests) diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py index 1424a5e840a4..59293c9ece57 100644 --- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py +++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py @@ -126,12 +126,12 @@ # Mirror the engine's encoder runtime sizes (``get_encoder_runtime_sizes`` -> -# ``encoder_max_batch_size`` / ``encoder_max_num_tokens``, defaulting to +# ``encoder_max_num_items`` / ``encoder_max_num_tokens``, defaulting to # ``max_batch_size`` / ``max_num_tokens``). The encoder ``AttentionMetadata`` is # sized once at load to this max budget; each forward re-preps it with the real # per-image seq lens. Two distinct axes: requests = image/sequence count budget, # tokens = total patch budget. -_ENCODER_TEST_MAX_NUM_REQUESTS = 2048 +_ENCODER_TEST_MAX_NUM_ITEMS = 2048 _ENCODER_TEST_MAX_NUM_TOKENS = 8192 @@ -177,7 +177,7 @@ def _build_trt_vision_tower(vision_cfg, dtype=torch.float32, device="cuda"): # `_set_up_multimodal_encoder_attn_metadata`; standalone tests must mirror # that before the encoder forward. tower.setup_attn_metadata( - max_num_requests=_ENCODER_TEST_MAX_NUM_REQUESTS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS + max_num_items=_ENCODER_TEST_MAX_NUM_ITEMS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS ) return tower diff --git a/tests/unittest/_torch/modeling/test_modeling_clip.py b/tests/unittest/_torch/modeling/test_modeling_clip.py index 5095defe6a58..7c3ff5d87516 100644 --- a/tests/unittest/_torch/modeling/test_modeling_clip.py +++ b/tests/unittest/_torch/modeling/test_modeling_clip.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import unittest from copy import deepcopy from dataclasses import dataclass @@ -33,12 +36,12 @@ } # Mirror the engine's encoder runtime sizes (``get_encoder_runtime_sizes`` -> -# ``encoder_max_batch_size`` / ``encoder_max_num_tokens``, defaulting to +# ``encoder_max_num_items`` / ``encoder_max_num_tokens``, defaulting to # ``max_batch_size`` / ``max_num_tokens``). The encoder ``AttentionMetadata`` is # sized once at load to this max budget; each forward re-preps it with the real # per-image seq lens. Two distinct axes: requests = image/sequence count budget, # tokens = total patch budget. -_ENCODER_TEST_MAX_NUM_REQUESTS = 2048 +_ENCODER_TEST_MAX_NUM_ITEMS = 2048 _ENCODER_TEST_MAX_NUM_TOKENS = 8192 @@ -87,7 +90,7 @@ def test_clip_vision_allclose_to_hf(self, scenario: Scenario): tllm_model = CLIPVisionModel(model_config).to(dtype).to(device) # Engine normally calls this after model load; standalone tests must do it themselves. tllm_model.setup_attn_metadata( - max_num_requests=_ENCODER_TEST_MAX_NUM_REQUESTS, + max_num_items=_ENCODER_TEST_MAX_NUM_ITEMS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS) # Use the load_weights method we are testing tllm_model.load_weights(hf_model.state_dict()) diff --git a/tests/unittest/_torch/modeling/test_modeling_mistral.py b/tests/unittest/_torch/modeling/test_modeling_mistral.py index 793b70455057..1a5a715de75f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_mistral.py +++ b/tests/unittest/_torch/modeling/test_modeling_mistral.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import contextlib import os import re @@ -669,6 +672,17 @@ def test_dummy_mm_max_tokens_per_item_is_image_only(): assert demand["image"] == (1540 // 14) ** 2 == 110**2 +def test_attention_metadata_capacity_uses_item_and_token_budgets(): + proc = _make_dummy_processor(spatial_merge_size=2) + + assert proc.get_mm_encoder_attention_metadata_capacity(max_num_items=8, max_num_tokens=100) == { + "attention": 8 + } + assert proc.get_mm_encoder_attention_metadata_capacity( + max_num_items=100, max_num_tokens=12 + ) == {"attention": 3} + + @pytest.mark.parametrize("budget", [1024, 4096, 8192]) def test_dummy_get_size_for_max_tokens_fits_and_aligns(budget): proc = _make_dummy_processor() @@ -707,6 +721,22 @@ def test_dummy_get_dummy_mm_data_for_tokens_shapes_and_saturation(budget): assert n * per_image + per_image > budget +def test_dummy_get_dummy_mm_data_for_tokens_covers_many_item_boundary(): + proc = _make_dummy_processor(num_channels=3) + data = proc.get_dummy_mm_data_for_tokens( + max_tokens_per_modality={"image": 8192}, + max_items_per_modality={"image": 8}, + dtype=torch.float16, + ) + + pixel_values = data["image"]["pixel_values"] + num_images, _, height, width = pixel_values.shape + tokens_per_image = (height // 14) * (width // 14) + assert num_images == 8 + assert tokens_per_image == 1024 + assert num_images * tokens_per_image == 8192 + + def test_dummy_for_tokens_empty_without_image_budget(): proc = _make_dummy_processor() assert proc.get_dummy_mm_data_for_tokens(max_tokens_per_modality={"audio": 1024}) == {} diff --git a/tests/unittest/_torch/modeling/test_modeling_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_multimodal.py index 112b665dfb8c..e8329611083c 100644 --- a/tests/unittest/_torch/modeling/test_modeling_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_multimodal.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import gc import os import unittest @@ -214,11 +217,10 @@ def create_trtllm_model( for module in model.modules(): if isinstance(module, MultimodalEncoderMixin): module.setup_attn_metadata( - # Encoder batch axis (image/sequence count) is a distinct - # budget from the token axis; mirror the engine's - # max_batch_size default. Qwen-family subclasses floor this - # up to max_num_tokens internally for windowed fan-out. - max_num_requests=2048, + # Atomic item count is distinct from the token budget. + # Each encoder maps both inputs to its internal attention + # sequence/window capacity. + max_num_items=2048, max_num_tokens=model_config.max_num_tokens, ) diff --git a/tests/unittest/_torch/modeling/test_modeling_pixtral.py b/tests/unittest/_torch/modeling/test_modeling_pixtral.py index 74b7217538db..1b08de1d1756 100644 --- a/tests/unittest/_torch/modeling/test_modeling_pixtral.py +++ b/tests/unittest/_torch/modeling/test_modeling_pixtral.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import gc import os import pathlib @@ -28,12 +31,12 @@ pytestmark = pytest.mark.threadleak(enabled=False) # Mirror the engine's encoder runtime sizes (``get_encoder_runtime_sizes`` -> -# ``encoder_max_batch_size`` / ``encoder_max_num_tokens``, defaulting to +# ``encoder_max_num_items`` / ``encoder_max_num_tokens``, defaulting to # ``max_batch_size`` / ``max_num_tokens``). The encoder ``AttentionMetadata`` is # sized once at load to this max budget; each forward re-preps it with the real # per-image seq lens. Two distinct axes: requests = image/sequence count budget, # tokens = total patch budget. -_ENCODER_TEST_MAX_NUM_REQUESTS = 2048 +_ENCODER_TEST_MAX_NUM_ITEMS = 2048 _ENCODER_TEST_MAX_NUM_TOKENS = 8192 @@ -89,7 +92,7 @@ def test_pixtral_vision_model_vs_hf(): ) # Engine normally calls this after model load; standalone tests must do it themselves. pixtral_model.setup_attn_metadata( - max_num_requests=_ENCODER_TEST_MAX_NUM_REQUESTS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS + max_num_items=_ENCODER_TEST_MAX_NUM_ITEMS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS ) hf_pixtral_model = init_hf_model( cls=hf_modeling_pixtral.PixtralVisionModel, @@ -148,7 +151,7 @@ def test_tensor_parallelism(mpi_pool_executor, tmp_path): modeling_pixtral.PixtralVisionModel(model_config=pixtral_vision_config).eval().to("cuda") ) pixtral_model.setup_attn_metadata( - max_num_requests=_ENCODER_TEST_MAX_NUM_REQUESTS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS + max_num_items=_ENCODER_TEST_MAX_NUM_ITEMS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS ) pixtral_model.load_weights(state_dict) # Save the number of params to check that the model gets shared in the workers. @@ -221,7 +224,7 @@ def _run_pixtral_and_compare_against_ref( modeling_pixtral.PixtralVisionModel(model_config=pixtral_vision_config).eval().to("cuda") ) pixtral_model.setup_attn_metadata( - max_num_requests=_ENCODER_TEST_MAX_NUM_REQUESTS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS + max_num_items=_ENCODER_TEST_MAX_NUM_ITEMS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS ) state_dict = torch.load(hf_weights_path, map_location="cuda") pixtral_model.load_weights(state_dict) diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py index 3f6bc50706ac..b5625d3c9fe4 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen2_5vl.py @@ -694,6 +694,58 @@ def test_mm_max_tokens_per_item_is_image_only(processor_cls): assert demand["image"] > 0 +def test_qwen2_5_attention_metadata_capacity_uses_processor_geometry(): + proc = _make_dummy_processor( + Qwen2VLInputProcessorBase, + patch_size=16, + spatial_merge_size=2, + min_pixels=32 * 32 * 4, + max_pixels=32 * 32 * 64, + ) + proc.config.vision_config.window_size = 128 + + capacities = proc.get_mm_encoder_attention_metadata_capacity( + max_num_items=8, max_num_tokens=160) + + # One full-attention frame has at least 4 merged cells * 4 physical + # patches. For 4x4 windows, the maximum window/area ratio is 2/5 (a + # 1x5 merged grid), or one window segment per 10 physical patches. + assert capacities == { + "full_attention": 10, + "window_attention": 16, + } + + +def test_qwen2_5_rejects_runtime_grid_below_startup_geometry(): + proc = _make_dummy_processor( + Qwen2VLInputProcessorBase, + patch_size=16, + spatial_merge_size=2, + min_pixels=32 * 32 * 4, + max_pixels=32 * 32 * 64, + ) + proc.config.vision_config.window_size = 128 + + proc._validate_encoder_attention_geometry( + {"image": torch.tensor([[1, 4, 4]])}) + with pytest.raises(ValueError, match="startup processor geometry"): + proc._validate_encoder_attention_geometry( + {"image": torch.tensor([[1, 2, 2]])}) + + +def test_qwen3_attention_capacity_keeps_long_video_safe(): + proc = _make_dummy_processor(Qwen3VLInputProcessorBase, + spatial_merge_size=2) + + capacities = proc.get_mm_encoder_attention_metadata_capacity( + max_num_items=1, max_num_tokens=160) + + # Qwen3's video resize clamp applies to aggregate temporal pixels. A long + # atomic video can therefore reach the hard one-merged-cell minimum for + # each temporal segment, irrespective of the item-count budget. + assert capacities == {"attention": 40} + + @pytest.mark.parametrize("processor_cls", _DUMMY_PROCESSORS) @pytest.mark.parametrize("budget", [1024, 4096, 16384]) def test_get_dummy_mm_data_for_tokens_saturates_budget(processor_cls, budget): @@ -709,3 +761,19 @@ def test_get_dummy_mm_data_for_tokens_saturates_budget(processor_cls, budget): # Saturates: within the budget, and adding one more image would exceed it. assert total_patches <= budget assert total_patches + per_image > budget + + +@pytest.mark.parametrize("processor_cls", _DUMMY_PROCESSORS) +def test_get_dummy_mm_data_for_tokens_covers_many_item_boundary(processor_cls): + proc = _make_dummy_processor(processor_cls) + data = proc.get_dummy_mm_data_for_tokens( + max_tokens_per_modality={"image": 8192}, + max_items_per_modality={"image": 8}, + dtype=torch.float32, + ) + + grid = data["image"]["image_grid_thw"] + token_lengths = grid.prod(dim=1) + assert grid.shape[0] == 8 + assert token_lengths.tolist() == [1024] * 8 + assert int(token_lengths.sum().item()) == 8192 diff --git a/tests/unittest/_torch/modeling/test_modeling_radio.py b/tests/unittest/_torch/modeling/test_modeling_radio.py index ddf8e98bd6a5..90dcee7787c7 100644 --- a/tests/unittest/_torch/modeling/test_modeling_radio.py +++ b/tests/unittest/_torch/modeling/test_modeling_radio.py @@ -22,12 +22,12 @@ ) # Mirror the engine's encoder runtime sizes (``get_encoder_runtime_sizes`` -> -# ``encoder_max_batch_size`` / ``encoder_max_num_tokens``, defaulting to +# ``encoder_max_num_items`` / ``encoder_max_num_tokens``, defaulting to # ``max_batch_size`` / ``max_num_tokens``). The encoder ``AttentionMetadata`` is # sized once at load to this max budget; each forward re-preps it with the real # per-image seq lens. Two distinct axes: requests = image/sequence count budget, # tokens = total patch budget. -_ENCODER_TEST_MAX_NUM_REQUESTS = 2048 +_ENCODER_TEST_MAX_NUM_ITEMS = 2048 _ENCODER_TEST_MAX_NUM_TOKENS = 8192 @@ -98,7 +98,7 @@ def test_radio_fp8_parent_kv_cache_does_not_leak_into_vit(tiny_vit_config): for module in vision_model.modules(): if isinstance(module, MultimodalEncoderMixin): module.setup_attn_metadata( - max_num_requests=_ENCODER_TEST_MAX_NUM_REQUESTS, + max_num_items=_ENCODER_TEST_MAX_NUM_ITEMS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS, ) diff --git a/tests/unittest/_torch/modeling/test_modeling_siglip.py b/tests/unittest/_torch/modeling/test_modeling_siglip.py index 74b2fd2d6b40..db0684cbc124 100644 --- a/tests/unittest/_torch/modeling/test_modeling_siglip.py +++ b/tests/unittest/_torch/modeling/test_modeling_siglip.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import unittest from copy import deepcopy from dataclasses import dataclass @@ -44,12 +47,12 @@ } # Mirror the engine's encoder runtime sizes (``get_encoder_runtime_sizes`` -> -# ``encoder_max_batch_size`` / ``encoder_max_num_tokens``, defaulting to +# ``encoder_max_num_items`` / ``encoder_max_num_tokens``, defaulting to # ``max_batch_size`` / ``max_num_tokens``). The encoder ``AttentionMetadata`` is # sized once at load to this max budget; each forward re-preps it with the real # per-image seq lens. Two distinct axes: requests = image/sequence count budget, # tokens = total patch budget. -_ENCODER_TEST_MAX_NUM_REQUESTS = 2048 +_ENCODER_TEST_MAX_NUM_ITEMS = 2048 _ENCODER_TEST_MAX_NUM_TOKENS = 8192 @@ -119,7 +122,7 @@ def test_siglip_vision_allclose_to_hf(self, scenario: Scenario): model_config, use_post_layernorm=True).to(dtype).to(device) # Engine normally calls this after model load; standalone tests must do it themselves. tllm_model.setup_attn_metadata( - max_num_requests=_ENCODER_TEST_MAX_NUM_REQUESTS, + max_num_items=_ENCODER_TEST_MAX_NUM_ITEMS, max_num_tokens=_ENCODER_TEST_MAX_NUM_TOKENS) tllm_model.load_weights(hf_model.state_dict()) diff --git a/tests/unittest/_torch/modeling/test_modeling_step3p7vl.py b/tests/unittest/_torch/modeling/test_modeling_step3p7vl.py index c05b872b0aef..bb7c5b61eb76 100644 --- a/tests/unittest/_torch/modeling/test_modeling_step3p7vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_step3p7vl.py @@ -88,10 +88,10 @@ def _load_config(checkpoint_dir: str) -> dict: _GPU_DTYPE = torch.bfloat16 # Mirror the engine's encoder runtime sizes (``get_encoder_runtime_sizes`` -> -# ``encoder_max_batch_size`` / ``encoder_max_num_tokens``, defaulting to +# ``encoder_max_num_items`` / ``encoder_max_num_tokens``, defaulting to # ``max_batch_size`` / ``max_num_tokens``). Two distinct axes: requests = # image/sequence count budget, tokens = total patch budget. -_ENCODER_TEST_MAX_NUM_REQUESTS = 2048 +_ENCODER_TEST_MAX_NUM_ITEMS = 2048 _ENCODER_TEST_MAX_NUM_TOKENS = 8192 @@ -199,7 +199,7 @@ def _setup_encoder_attn_metadata(module, max_num_tokens: int = _ENCODER_TEST_MAX for m in module.modules(): if isinstance(m, MultimodalEncoderMixin): m.setup_attn_metadata( - max_num_requests=_ENCODER_TEST_MAX_NUM_REQUESTS, max_num_tokens=max_num_tokens + max_num_items=_ENCODER_TEST_MAX_NUM_ITEMS, max_num_tokens=max_num_tokens ) return module diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_mixin.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_mixin.py index 5366d0ad02ad..97ce118fbfd1 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_mixin.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_mixin.py @@ -1,19 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 """Unit tests for `MultimodalEncoderMixin`. Verifies that the mixin's default `setup_attn_metadata` builds the encoder's -AttentionMetadata exactly once with the runtime sizes injected by the engine -(`max_num_requests`, `max_num_tokens`) and `kv_cache_manager=None`, flooring -`max_num_requests` at the encoder fallback (one attention segment per vision -tile can exceed the request count). These tests use a stub `metadata_cls` to +AttentionMetadata exactly once with the item/token budgets injected by the +engine and `kv_cache_manager=None`. The default mapping floors +`max_num_requests` at the encoder fallback because one attention segment per +vision tile can exceed the item count. These tests use a stub `metadata_cls` to avoid pulling in any real attention backend (and any GPU/CUDA dependency). """ +from types import SimpleNamespace + import torch.nn as nn from tensorrt_llm._torch.models.modeling_multimodal_encoder import ( _ENCODER_FALLBACK_MAX_NUM_REQUESTS, MultimodalEncoderMixin, ) +from tensorrt_llm._torch.models.modeling_pixtral import PixtralVisionModel +from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2_5_VisionModel +from tensorrt_llm._torch.models.modeling_qwen3vl import Qwen3VisionModel class _StubMetadata: @@ -35,10 +42,10 @@ def test_attn_metadata_is_none_before_setup(): def test_setup_attn_metadata_builds_with_engine_sizes(): - """A request count above the encoder fallback passes through unchanged.""" + """An item count above the encoder fallback passes through unchanged.""" encoder = _DummyEncoder() big = _ENCODER_FALLBACK_MAX_NUM_REQUESTS + 100 - encoder.setup_attn_metadata(max_num_requests=big, max_num_tokens=1234) + encoder.setup_attn_metadata(max_num_items=big, max_num_tokens=1234) assert isinstance(encoder.attn_metadata, _StubMetadata) assert encoder.attn_metadata.kwargs == { @@ -48,12 +55,12 @@ def test_setup_attn_metadata_builds_with_engine_sizes(): } -def test_setup_attn_metadata_floors_small_request_count(): - """A request count below the encoder fallback is floored to it: the encoder - runs one attention segment per vision tile, which can exceed the request +def test_setup_attn_metadata_floors_small_item_count(): + """An item count below the encoder fallback is floored to it: the encoder + runs one attention segment per vision tile, which can exceed the item count, so the per-segment buffers must not undersize.""" encoder = _DummyEncoder() - encoder.setup_attn_metadata(max_num_requests=8, max_num_tokens=100) + encoder.setup_attn_metadata(max_num_items=8, max_num_tokens=100) assert encoder.attn_metadata.kwargs == { "max_num_requests": _ENCODER_FALLBACK_MAX_NUM_REQUESTS, @@ -62,6 +69,21 @@ def test_setup_attn_metadata_floors_small_request_count(): } +def test_setup_attn_metadata_accepts_processor_capacity(): + encoder = _DummyEncoder() + encoder.setup_attn_metadata( + max_num_items=8, + max_num_tokens=100, + attention_metadata_capacity={"attention": 7}, + ) + + assert encoder.attn_metadata.kwargs == { + "max_num_requests": 7, + "max_num_tokens": 100, + "kv_cache_manager": None, + } + + def test_setup_attn_metadata_is_idempotent_per_call(): """Subsequent calls overwrite the previous AttentionMetadata, which matches how the engine drives `_set_up_multimodal_encoder_attn_metadata` (called @@ -70,10 +92,10 @@ def test_setup_attn_metadata_is_idempotent_per_call(): encoder = _DummyEncoder() big1 = _ENCODER_FALLBACK_MAX_NUM_REQUESTS + 8 big2 = _ENCODER_FALLBACK_MAX_NUM_REQUESTS + 16 - encoder.setup_attn_metadata(max_num_requests=big1, max_num_tokens=100) + encoder.setup_attn_metadata(max_num_items=big1, max_num_tokens=100) first = encoder.attn_metadata - encoder.setup_attn_metadata(max_num_requests=big2, max_num_tokens=200) + encoder.setup_attn_metadata(max_num_items=big2, max_num_tokens=200) second = encoder.attn_metadata assert first is not second @@ -84,6 +106,48 @@ def test_setup_attn_metadata_is_idempotent_per_call(): } +def test_pixtral_maps_each_atomic_item_to_one_attention_context(): + class _PixtralCapacityEncoder(_DummyEncoder): + get_encoder_attention_metadata_capacity = ( + PixtralVisionModel.get_encoder_attention_metadata_capacity + ) + + encoder = _PixtralCapacityEncoder() + encoder.setup_attn_metadata(max_num_items=3, max_num_tokens=100) + + assert encoder.attn_metadata.kwargs == { + "max_num_requests": 3, + "max_num_tokens": 100, + "kv_cache_manager": None, + } + + encoder.setup_attn_metadata(max_num_items=100, max_num_tokens=3) + assert encoder.attn_metadata.kwargs["max_num_requests"] == 3 + + +def test_qwen2_maps_token_budget_to_full_and_window_attention_contexts(): + encoder = SimpleNamespace(spatial_merge_unit=4) + + capacities = Qwen2_5_VisionModel.get_encoder_attention_metadata_capacity( + encoder, max_num_items=8, max_num_tokens=100 + ) + + assert capacities == { + "full_attention": 25, + "window_attention": 25, + } + + +def test_qwen3_maps_token_budget_to_temporal_attention_contexts(): + encoder = SimpleNamespace(spatial_merge_unit=4) + + capacities = Qwen3VisionModel.get_encoder_attention_metadata_capacity( + encoder, max_num_items=8, max_num_tokens=100 + ) + + assert capacities == {"attention": 25} + + def test_subclass_can_override_setup_attn_metadata(): """Special encoders (e.g. RADIO with `kv_layout="NHD"`, Qwen2-VL with multiple metadata objects) override `setup_attn_metadata`. Verify that @@ -95,14 +159,14 @@ def __init__(self): super().__init__() self.metadata_cls = _StubMetadata - def setup_attn_metadata(self, max_num_requests, max_num_tokens): - captured["max_num_requests"] = max_num_requests + def setup_attn_metadata(self, max_num_items, max_num_tokens): + captured["max_num_items"] = max_num_items captured["max_num_tokens"] = max_num_tokens # Intentionally skip building self.attn_metadata to confirm # the default impl wasn't called. encoder = _OverridingEncoder() - encoder.setup_attn_metadata(max_num_requests=5, max_num_tokens=50) + encoder.setup_attn_metadata(max_num_items=5, max_num_tokens=50) assert encoder.attn_metadata is None # default impl not invoked - assert captured == {"max_num_requests": 5, "max_num_tokens": 50} + assert captured == {"max_num_items": 5, "max_num_tokens": 50} diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 7ff208c41ce3..397498d71dee 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -127,7 +127,7 @@ methods: annotation: bool default: False status: prototype - encoder_max_batch_size: + encoder_max_num_items: annotation: Optional[int] default: null status: prototype diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index d7d36e19fff4..b26851643598 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -8,6 +8,7 @@ import torch from tensorrt_llm.inputs.multimodal import ( + MULTIMODAL_ENCODER_ITEM_METADATA_KEY, DisaggPrefillMultimodalInputs, MultimodalInput, MultimodalRuntimeData, @@ -15,11 +16,81 @@ find_mm_token_lengths, ) from tensorrt_llm.inputs.registry import ( + MultimodalEncoderItemMetadata, create_input_processor_with_hash, maybe_compute_mm_embed_cumsum, ) +class _ItemMetadataFakeProcessor: + multimodal_hashing_supported = False + supports_mm_encoder_item_scheduling = True + + def __init__(self, existing_embedding_lengths): + self.existing_embedding_lengths = existing_embedding_lengths + + def __call__(self, inputs, sampling_params): + return [101, 102], { + "multimodal_data": { + "image": {}, + "multimodal_embedding_lengths": self.existing_embedding_lengths, + } + } + + def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): + assert prompt_token_ids == [101, 102] + assert "image" in multimodal_data + return MultimodalEncoderItemMetadata( + item_refs=[("image", 0)], + encoder_token_lengths=[4], + output_embedding_lengths=[2], + ) + + def get_vocab_size(self): + return 100 + + def get_mm_token_ids(self): + return None + + def get_mm_special_token_ids(self): + return None + + +def test_mm_item_metadata_is_materialized_when_embedding_lengths_match(): + input_processor = create_input_processor_with_hash(_ItemMetadataFakeProcessor([2])) + + _, extra = input_processor({"prompt": "unused"}, sampling_params=None) + + multimodal_data = extra["multimodal_data"] + item_metadata = multimodal_data[MULTIMODAL_ENCODER_ITEM_METADATA_KEY] + assert item_metadata == MultimodalEncoderItemMetadata( + item_refs=[("image", 0)], + encoder_token_lengths=[4], + output_embedding_lengths=[2], + ) + assert multimodal_data["multimodal_embedding_lengths"] == [2] + assert "multimodal_item_refs" not in multimodal_data + assert "multimodal_encoder_token_lengths" not in multimodal_data + + +def test_mm_item_metadata_rejects_mismatched_embedding_lengths(): + input_processor = create_input_processor_with_hash(_ItemMetadataFakeProcessor([1])) + + with pytest.raises(ValueError, match="do not match"): + input_processor({"prompt": "unused"}, sampling_params=None) + + +def test_mm_item_scheduling_contract_requires_metadata_for_raw_payload(): + class MissingMetadataProcessor(_ItemMetadataFakeProcessor): + def get_mm_encoder_item_metadata(self, prompt_token_ids, multimodal_data): + return None + + input_processor = create_input_processor_with_hash(MissingMetadataProcessor([2])) + + with pytest.raises(TypeError, match="must return item metadata"): + input_processor({"prompt": "unused"}, sampling_params=None) + + def test_maybe_compute_mm_embed_cumsum_populates_py_multimodal_data(): """Producer writes a flat int64 cumsum tensor at py_multimodal_data[multimodal_embed_mask_cumsum].""" diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index b5410a8edfdb..4805aa47f9d3 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -208,16 +208,16 @@ def test_llm_args_with_model_kwargs(self, llm_args_cls): class TestEncoderRuntimeSizes: """Cover encoder runtime size fields and fallback to LLM limits. - `encoder_max_batch_size` / `encoder_max_num_tokens` are user-facing - knobs that size multimodal encoder AttentionMetadata; when unset they - fall back to the LLM-side `max_batch_size` / `max_num_tokens`. They are - PyTorch-backend only (the multimodal encoder profiling path), so they - live on `TorchLlmArgs` rather than the shared `BaseLlmArgs`. + `encoder_max_num_items` / `encoder_max_num_tokens` are user-facing + knobs for multimodal encoder scheduling and AttentionMetadata. When unset, + they fall back to the corresponding LLM-side values before model-aware + atomic-item compatibility resolution. They are PyTorch-backend only, so + they live on `TorchLlmArgs` rather than the shared `BaseLlmArgs`. """ def test_defaults_are_none(self, llm_args_cls): llm_args = llm_args_cls(model=llama_model_path) - assert llm_args.encoder_max_batch_size is None + assert llm_args.encoder_max_num_items is None assert llm_args.encoder_max_num_tokens is None @pytest.mark.parametrize( @@ -225,10 +225,10 @@ def test_defaults_are_none(self, llm_args_cls): [ # Neither encoder knob set -- falls back to LLM limits. (dict(max_batch_size=64, max_num_tokens=2048), (64, 2048)), - # Only encoder_max_batch_size overrides. + # Only encoder_max_num_items overrides. (dict(max_batch_size=64, max_num_tokens=2048, - encoder_max_batch_size=512), (512, 2048)), + encoder_max_num_items=512), (512, 2048)), # Only encoder_max_num_tokens overrides. (dict(max_batch_size=64, max_num_tokens=2048, @@ -236,10 +236,10 @@ def test_defaults_are_none(self, llm_args_cls): # Both encoder knobs override. (dict(max_batch_size=64, max_num_tokens=2048, - encoder_max_batch_size=512, + encoder_max_num_items=512, encoder_max_num_tokens=32768), (512, 32768)), ], - ids=["fallback", "only_batch", "only_tokens", "both"], + ids=["fallback", "only_items", "only_tokens", "both"], ) def test_get_encoder_runtime_sizes(self, llm_args_cls, kwargs, expected_runtime_sizes): @@ -249,8 +249,8 @@ def test_get_encoder_runtime_sizes(self, llm_args_cls, kwargs, @pytest.mark.parametrize( "field_name, invalid_value", [ - ("encoder_max_batch_size", 0), - ("encoder_max_batch_size", -1), + ("encoder_max_num_items", 0), + ("encoder_max_num_items", -1), ("encoder_max_num_tokens", 0), ("encoder_max_num_tokens", -1), ], @@ -261,6 +261,31 @@ def test_rejects_non_positive(self, llm_args_cls, field_name, llm_args_cls(model=llama_model_path, **{field_name: invalid_value}) +class TestEagerEncoderSchedulingCompatibility: + """Eager MM encoder scheduling rejects unsupported feature combinations.""" + + def test_eager_alone_is_accepted(self): + args = TorchLlmArgs(model=llama_model_path, + multimodal_config=MultimodalConfig( + enable_eager_encoder_scheduling=True)) + assert args.multimodal_config.enable_eager_encoder_scheduling + + def test_eager_rejects_attention_dp(self): + with pytest.raises(ValidationError, match="attention DP"): + TorchLlmArgs(model=llama_model_path, + enable_attention_dp=True, + multimodal_config=MultimodalConfig( + enable_eager_encoder_scheduling=True)) + + def test_eager_rejects_disaggregated_serving(self): + with pytest.raises(ValidationError, match="disaggregated"): + TorchLlmArgs( + model=llama_model_path, + cache_transceiver_config=CacheTransceiverConfig(backend="NIXL"), + multimodal_config=MultimodalConfig( + enable_eager_encoder_scheduling=True)) + + def test_decoding_type_eagle3_parses_to_eagle3_decoding_config(): adapter = TypeAdapter(SpeculativeConfig) spec_cfg = adapter.validate_python( @@ -713,6 +738,7 @@ def test_default_encoder_cuda_graph_is_none(self): assert MultimodalConfig().encoder_cuda_graph is None assert MultimodalConfig().encoder_side_stream_max_ahead == 0 assert MultimodalConfig().encoder_cache_max_bytes == 128 * 1024**2 + assert not MultimodalConfig().enable_eager_encoder_scheduling assert MultimodalConfig().video_pruning_rate is None def test_torch_llm_args_default_multimodal_config(self): @@ -721,6 +747,7 @@ def test_torch_llm_args_default_multimodal_config(self): assert args.multimodal_config.encoder_cuda_graph is None assert args.multimodal_config.encoder_side_stream_max_ahead == 0 assert args.multimodal_config.encoder_cache_max_bytes == 128 * 1024**2 + assert not args.multimodal_config.enable_eager_encoder_scheduling assert args.multimodal_config.video_pruning_rate is None @pytest.mark.parametrize( @@ -752,6 +779,12 @@ def test_torch_llm_args_with_encoder_side_stream_max_ahead(self): )) assert args.multimodal_config.encoder_side_stream_max_ahead == 2 + def test_torch_llm_args_with_eager_encoder_scheduling(self): + args = TorchLlmArgs(model=llama_model_path, + multimodal_config=MultimodalConfig( + enable_eager_encoder_scheduling=True)) + assert args.multimodal_config.enable_eager_encoder_scheduling + def test_torch_llm_args_with_multimodal_video_pruning_rate(self): args = TorchLlmArgs( model=llama_model_path, @@ -790,10 +823,12 @@ def test_torch_llm_args_with_encoder_side_stream_max_ahead_yaml(self): multimodal_config={ "encoder_side_stream_max_ahead": 2, "encoder_cache_max_bytes": 0, + "enable_eager_encoder_scheduling": True, "video_pruning_rate": 0.5, }, ) assert args.multimodal_config.encoder_side_stream_max_ahead == 2 + assert args.multimodal_config.enable_eager_encoder_scheduling assert args.multimodal_config.video_pruning_rate == 0.5 def test_encoder_cuda_graph_and_side_stream_max_ahead_are_exclusive(self):