From 5b90a930ea54fe615a7c41cfe5879562c25c13d5 Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Tue, 13 Oct 2026 01:09:07 -0700 Subject: [PATCH 1/8] [None][perf] share multimodal encoder CUDA graph capture Capture configured encoders through MultimodalEncoderMixin after weights and metadata are ready, before KV-cache profiling. Reuse MultimodalEncoderGraphRunner while keeping tensor specs, metadata providers, preprocessing, and mergers in each model. Add Qwen3-VL block graphs with eager deepstack mergers and share fixed-buffer packed attention metadata refresh with Qwen2-VL and GLM. Cover capture ordering, duplicate startup, padded replay, and TRTLLM/FlashInfer backends. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- .../models/modeling_glm5_next_vision.py | 52 ++--- .../models/modeling_multimodal_encoder.py | 115 +++++++++- .../models/modeling_nemotron_h_multimodal.py | 1 - .../_torch/models/modeling_qwen2vl.py | 58 +---- .../_torch/models/modeling_qwen3vl.py | 205 ++++++++++++++++-- tensorrt_llm/_torch/models/modeling_radio.py | 22 +- .../_torch/models/multimodal_encoder_graph.py | 6 +- .../_torch/pyexecutor/engine/multimodal.py | 7 + .../_torch/pyexecutor/model_engine.py | 7 +- .../test_lists/test-db/l0_h100.yml | 1 + .../_torch/executor/engine/test_multimodal.py | 59 ++++- .../modeling/test_glm5_next_contracts.py | 30 ++- .../_torch/modeling/test_modeling_qwen3vl.py | 130 ++++++++++- .../modeling/test_multimodal_encoder_graph.py | 21 ++ 14 files changed, 574 insertions(+), 140 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py index 8e9efa4abc61..5b0514b8b87c 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py @@ -68,13 +68,12 @@ flashinfer_apply_rope_with_cos_sin_cache_inplace = None from .modeling_auto import AutoModelForCausalLM from .modeling_glm5_next import Glm5NextForCausalLM, glm5_next_attention_mapping -from .modeling_multimodal_encoder import MultimodalEncoderMixin +from .modeling_multimodal_encoder import MultimodalEncoderMixin, prepare_encoder_attn_metadata from .modeling_multimodal_mixin import ( EncoderGroup, MultimodalModelMixin, encode_multimodal_by_groups, ) -from .modeling_qwen2vl import _prepare_qwen_vl_vision_attn_metadata from .modeling_utils import ( ModelConfig, QuantConfig, @@ -451,24 +450,11 @@ def device(self) -> torch.device: return self.patch_embed.proj.weight.device # -- block-stack CUDA graphs --------------------------------------------- - def enable_blocks_cuda_graph( - self, - config: Optional["MultimodalEncoderCudaGraphConfig"] = None, - *, - device: Optional[torch.device] = None, - ) -> None: - """Capture the 24-block loop for every configured bucket. - - Weights must already live on the CUDA device. The eager prelude - (patch embed, rotary tables) and tail (post norm, downsample, merger) - stay outside the graph; a request whose (num_images, total_patches) - matches no bucket falls back to the eager block loop. - """ - config = config if config is not None else self._encoder_cuda_graph_config - if config is None or self._blocks_graph_runner is not None: - return + def _build_blocks_graph_runner( + self, config: "MultimodalEncoderCudaGraphConfig" + ) -> MultimodalEncoderGraphRunner: dtype = self.patch_embed.proj.weight.dtype - runner = MultimodalEncoderGraphRunner( + return MultimodalEncoderGraphRunner( encoder_fn=self._encoder_graph_fn, metadata_provider=_Glm5NextVisionGraphMetadataProvider(self), input_specs={ @@ -479,8 +465,6 @@ def enable_blocks_cuda_graph( output_specs={"x": 0}, config=config, ) - runner.capture_all(device if device is not None else self.device) - self._blocks_graph_runner = runner def clear_blocks_cuda_graph(self, *, release_nccl_window_owners: bool = True) -> None: """Release the block graphs at the engine coordinated teardown boundary.""" @@ -525,7 +509,7 @@ def _run_blocks( ) if out is not None: return out["x"] - self.attn_metadata = _prepare_qwen_vl_vision_attn_metadata( + self.attn_metadata = prepare_encoder_attn_metadata( seq_lens, self.attn_metadata, max_seq_len=self._fixed_max_seq_len ) return self._run_blocks_eager(hidden_states, position_embeddings, self.attn_metadata) @@ -634,8 +618,8 @@ class _Glm5NextVisionGraphMetadataProvider(EncoderMetadataProvider): Per bucket one metadata instance with ``is_cuda_graph=True`` so the ``seq_lens`` setter copies into the captured ``_seq_lens_cuda`` buffer instead of reallocating; the packed cumulative lengths the FMHA reads - live in a fixed ``cu_q_seqlens`` buffer refreshed in place. Host-side - views are rebound per refresh — the C++ op reads them only at capture. + live in a fixed ``cu_q_seqlens`` buffer refreshed in place. Host lengths + and runtime views are bound once per bucket and updated before replay. """ graph_critical_attrs: Sequence[str] = ("_seq_lens_cuda", "cu_q_seqlens") @@ -655,22 +639,18 @@ def build(self, key: EncoderGraphKey) -> AttentionMetadata: key.num_contexts + 1, dtype=torch.int32, device=vit.device ) metadata.cu_kv_seqlens = metadata.cu_q_seqlens + metadata.seq_lens = torch.zeros( + key.num_contexts, dtype=torch.int32, pin_memory=prefer_pinned() + ) + metadata.bind_encoder_cuda_graph_seq_lens(metadata.seq_lens, key.num_contexts) return metadata def refresh_in_place( self, metadata: AttentionMetadata, padded_seq_lengths: Sequence[int] ) -> None: - n = len(padded_seq_lengths) - seq_lens = torch.tensor(padded_seq_lengths, dtype=torch.int32, pin_memory=prefer_pinned()) - metadata.num_contexts = n - metadata.request_ids = list(range(1, n + 1)) - metadata.seq_lens = seq_lens - metadata.bind_encoder_cuda_graph_seq_lens(metadata.seq_lens, n) - cu_seqlens = torch.zeros(n + 1, dtype=torch.int32, pin_memory=prefer_pinned()) - torch.cumsum(seq_lens, dim=0, out=cu_seqlens[1:]) - metadata.cu_q_seqlens.copy_(cu_seqlens, non_blocking=True) - metadata.max_seq_len = self._vit._fixed_max_seq_len - metadata.prepare_encoder_only() + prepare_encoder_attn_metadata( + list(padded_seq_lengths), metadata, max_seq_len=self._vit._fixed_max_seq_len + ) class Glm5NextVisionModelBase(nn.Module): @@ -1211,6 +1191,4 @@ def load_weights(self, weights: Dict[str, torch.Tensor], **_ignored: Any) -> Non "tower (serve text-only with disable_mm_encoder instead)" ) self.mm_encoder.load_weights(weights) - # Weights are on device now; capture the (opt-in) encoder graphs. - self.mm_encoder.enable_blocks_cuda_graph() self.llm.load_weights(weights) diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py index 15913f0670f9..59965f3f2594 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py @@ -16,14 +16,20 @@ # This file is based on official VILA: https://github.com/NVlabs/VILA/blob/main/llava/model/multimodal_encoder/ import os -from typing import Optional, Type +from typing import TYPE_CHECKING, Optional, Type, cast +import torch import torch.nn as nn from transformers import AutoConfig, AutoModel +from ..._utils import prefer_pinned from ..attention.backends.interface import AttentionMetadata from .modeling_multimodal_utils import multiscale_forward +if TYPE_CHECKING: + from ...llmapi.llm_args import MultimodalEncoderCudaGraphConfig + from .multimodal_encoder_graph import MultimodalEncoderGraphRunner + # 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 @@ -38,6 +44,80 @@ _ENCODER_FALLBACK_MAX_NUM_REQUESTS = 8192 +def prepare_encoder_attn_metadata(seq_lens: list[int], + attn_metadata: AttentionMetadata, *, + max_seq_len: int) -> AttentionMetadata: + """Refresh packed no-KV encoder lengths while preserving captured buffers.""" + seq_lens = [int(seq_len) for seq_len in seq_lens] + if not seq_lens: + raise ValueError("Encoder attention requires at least one segment") + + min_seq_len = min(seq_lens) + if min_seq_len < 0: + raise ValueError( + f"Encoder attention segment length must be nonnegative, got {min_seq_len}" + ) + + if max_seq_len <= 0: + raise ValueError( + f"Encoder attention max_seq_len must be positive, got {max_seq_len}" + ) + actual_max_seq_len = max(seq_lens) + if actual_max_seq_len > max_seq_len: + raise ValueError( + "Encoder attention segment length exceeds the configured " + f"maximum: actual={actual_max_seq_len}, maximum={max_seq_len}. " + "Increase encoder_max_num_tokens or reduce the input max_pixels.") + + num_segments = len(seq_lens) + seq_lens_torch = torch.tensor(seq_lens, + dtype=torch.int, + pin_memory=prefer_pinned()) + cu_seqlens = torch.empty(num_segments + 1, + dtype=torch.int32, + pin_memory=prefer_pinned()) + cu_seqlens[0] = 0 + torch.cumsum(seq_lens_torch, dim=0, out=cu_seqlens[1:]) + + attn_metadata.num_contexts = num_segments + attn_metadata.request_ids = list(range(1, num_segments + 1)) + prepare_graph_replay = getattr(attn_metadata, + "prepare_encoder_cuda_graph_replay", None) + is_cuda_graph = getattr(attn_metadata, "is_cuda_graph", False) + if is_cuda_graph and callable(prepare_graph_replay): + prepare_graph_replay(seq_lens, sum(seq_lens)) + # The replay helper refreshes only the host lengths, but the captured + # attention op also reads the device lengths; update them in place. + attn_metadata.seq_lens_cuda[:num_segments].copy_(seq_lens_torch, + non_blocking=True) + else: + attn_metadata.seq_lens = seq_lens_torch + cu_seqlens = cu_seqlens.to(device=attn_metadata.seq_lens_cuda.device, + non_blocking=True) + if (is_cuda_graph + and getattr(attn_metadata, "cu_q_seqlens", None) is not None + and attn_metadata.cu_q_seqlens.shape == cu_seqlens.shape): + # CUDA graph capture bakes these addresses into attention kernels. + # Refresh their contents without replacing the captured buffers. + attn_metadata.cu_q_seqlens.copy_(cu_seqlens, non_blocking=True) + if (attn_metadata.cu_kv_seqlens.data_ptr() + != attn_metadata.cu_q_seqlens.data_ptr()): + attn_metadata.cu_kv_seqlens.copy_(cu_seqlens, non_blocking=True) + else: + attn_metadata.cu_q_seqlens = cu_seqlens + attn_metadata.cu_kv_seqlens = cu_seqlens + # Keep the native attention-op cache key stable across image resolutions. + # Actual segment lengths remain available through seq_lens / cu_seqlens. + attn_metadata.max_seq_len = max_seq_len + # 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. + # `prepare_encoder_only` runs the lean path on backends that have one + # (TRTLLM) and falls back to the full `prepare()` elsewhere. + attn_metadata.prepare_encoder_only() + return attn_metadata + + class MultimodalEncoderMixin: """Encoder-side counterpart to ``MultimodalModelMixin``. @@ -52,6 +132,39 @@ class MultimodalEncoderMixin: """ metadata_cls: Type[AttentionMetadata] attn_metadata: Optional[AttentionMetadata] = None + _encoder_cuda_graph_config: "MultimodalEncoderCudaGraphConfig | None" = None + _blocks_graph_runner: "MultimodalEncoderGraphRunner | None" = None + + def _build_blocks_graph_runner( + self, config: "MultimodalEncoderCudaGraphConfig" + ) -> "MultimodalEncoderGraphRunner": + """Describe the block function, tensor specs, and metadata provider. + + Encoders opting into graphs override this factory. Model-specific + preprocessing and output assembly stay in the encoder forward. + """ + raise NotImplementedError( + f"{type(self).__name__} does not support encoder CUDA graphs") + + def enable_blocks_cuda_graph( + self, + config: "MultimodalEncoderCudaGraphConfig | None" = None, + *, + device: torch.device | None = None, + ) -> None: + """Capture configured block graphs once, after weights and metadata are ready. + + The engine drives this lifecycle for every local multimodal encoder. + Standalone callers may provide an explicit config and device. + """ + config = config if config is not None else self._encoder_cuda_graph_config + if config is None or self._blocks_graph_runner is not None: + return + runner = self._build_blocks_graph_runner(config) + if device is None: + device = next(cast(nn.Module, self).parameters()).device + runner.capture_all(device) + self._blocks_graph_runner = runner def get_encoder_attention_metadata_capacity( self, max_num_tokens: int) -> dict[str, int]: diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index 20c79d17d59f..be992ef8f9a1 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -2949,7 +2949,6 @@ def load_weights(self, weights): # Load vision encoder weights. if self.vision_encoder is not None: self.vision_encoder.load_weights(weights) - self.vision_encoder.enable_radio_cuda_graph() # Free vision encoder weights from the dict so the backing mmap pages can be released. if hasattr(weights, "mark_consumed"): diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py index f4236bd04023..f4848ae64f06 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py @@ -34,7 +34,7 @@ MultimodalParams) from tensorrt_llm.inputs.registry import MultimodalEncoderItemMetadata -from ..._utils import async_tensor_h2d, nvtx_range, prefer_pinned +from ..._utils import async_tensor_h2d, nvtx_range from ...inputs import (BaseMultimodalDummyInputsBuilder, BaseMultimodalInputProcessor, ContentFormat, ExtraProcessedInputs, MultimodalPlaceholderMetadata, @@ -71,6 +71,8 @@ from ..modules.gated_mlp import GatedMLP from .modeling_auto import AutoModelForCausalLM from .modeling_multimodal_encoder import MultimodalEncoderMixin +from .modeling_multimodal_encoder import \ + prepare_encoder_attn_metadata as _prepare_qwen_vl_vision_attn_metadata from .modeling_multimodal_mixin import MultimodalModelMixin from .modeling_multimodal_utils import ( _install_processor_output_validation_filter, find_input_mm_embeds, @@ -83,60 +85,6 @@ PAD_INDEX = -100 # NOTE: refer to https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py#L269 -def _prepare_qwen_vl_vision_attn_metadata( - seq_lens: List[int], attn_metadata: AttentionMetadata, *, - max_seq_len: int) -> AttentionMetadata: - seq_lens = [int(seq_len) for seq_len in seq_lens] - if not seq_lens: - raise ValueError( - "Qwen VL vision attention requires at least one segment") - - min_seq_len = min(seq_lens) - if min_seq_len < 0: - raise ValueError( - f"Qwen VL vision attention segment length must be nonnegative, got {min_seq_len}" - ) - - if max_seq_len <= 0: - raise ValueError( - f"Qwen VL vision attention max_seq_len must be positive, got {max_seq_len}" - ) - actual_max_seq_len = max(seq_lens) - if actual_max_seq_len > max_seq_len: - raise ValueError( - "Qwen VL vision attention segment length exceeds the configured " - f"maximum: actual={actual_max_seq_len}, maximum={max_seq_len}. " - "Increase encoder_max_num_tokens or reduce the input max_pixels.") - - num_segments = len(seq_lens) - seq_lens_torch = torch.tensor(seq_lens, - dtype=torch.int, - pin_memory=prefer_pinned()) - cu_seqlens = torch.empty(num_segments + 1, - dtype=torch.int32, - pin_memory=prefer_pinned()) - cu_seqlens[0] = 0 - torch.cumsum(seq_lens_torch, dim=0, out=cu_seqlens[1:]) - - attn_metadata.num_contexts = num_segments - attn_metadata.request_ids = list(range(1, num_segments + 1)) - attn_metadata.seq_lens = seq_lens_torch - cu_seqlens = cu_seqlens.to(device=attn_metadata.seq_lens_cuda.device, - non_blocking=True) - attn_metadata.cu_q_seqlens = cu_seqlens - attn_metadata.cu_kv_seqlens = cu_seqlens - # Keep the native attention-op cache key stable across image resolutions. - # Actual segment lengths remain available through seq_lens / cu_seqlens. - attn_metadata.max_seq_len = max_seq_len - # 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. - # `prepare_encoder_only` runs the lean path on backends that have one - # (TRTLLM) and falls back to the full `prepare()` elsewhere. - attn_metadata.prepare_encoder_only() - return attn_metadata - - def _prepare_qwen_vl_mrope_config( *, multimodal_params: List[MultimodalParams], diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 2b6d5163e2f3..479f60105a94 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -6,7 +6,8 @@ import math import re from functools import lru_cache -from typing import Any, Dict, List, Optional, Sequence, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union +from typing import Mapping as TypingMapping import numpy as np import torch @@ -50,7 +51,7 @@ from .checkpoints.base_weight_mapper import BaseWeightMapper from .checkpoints.hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper from .modeling_auto import AutoModelForCausalLM -from .modeling_multimodal_encoder import MultimodalEncoderMixin +from .modeling_multimodal_encoder import MultimodalEncoderMixin, prepare_encoder_attn_metadata from .modeling_multimodal_mixin import ( EncoderGroup, MultimodalModelMixin, @@ -61,7 +62,6 @@ Qwen2_5_VLVisionAttention, Qwen2VLInputProcessorBase, _prepare_qwen_vl_mrope_config, - _prepare_qwen_vl_vision_attn_metadata, ) from .modeling_utils import ( ModelConfig, @@ -71,6 +71,15 @@ register_auto_model, register_vision_encoder, ) +from .multimodal_encoder_graph import ( + EncoderGraphKey, + EncoderGraphTensorSpec, + EncoderMetadataProvider, + MultimodalEncoderGraphRunner, +) + +if TYPE_CHECKING: + from ...llmapi.llm_args import MultimodalEncoderCudaGraphConfig def _expand_prompt_token_ids_for_mm_handoff( @@ -880,6 +889,18 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): self.attn_metadata: Optional[AttentionMetadata] = None self._fixed_max_seq_len = self.model_config.max_num_tokens + self._blocks_graph_runner: MultimodalEncoderGraphRunner | None = None + + mm_config = self.model_config.multimodal_config + self._encoder_cuda_graph_config: "MultimodalEncoderCudaGraphConfig | None" = None + if mm_config is not None and mm_config.encoder_cuda_graph is not None: + unknown_modalities = set(mm_config.encoder_cuda_graph) - {"vision"} + if unknown_modalities: + raise ValueError( + "Unsupported Qwen3-VL encoder CUDA graph modalities: " + f"{sorted(unknown_modalities)}. Supported modalities: ['vision']." + ) + self._encoder_cuda_graph_config = mm_config.encoder_cuda_graph.get("vision") # Vision block's `rope_position_ids` scratch. Registered empty here; # `setup_attn_metadata` allocates it as an `arange` (see there). @@ -1064,10 +1085,122 @@ def prepare_attn_metadata( "Vision encoder AttentionMetadata is not initialized. " "It must be set up before the encoder forward runs." ) - return _prepare_qwen_vl_vision_attn_metadata( + return prepare_encoder_attn_metadata( seq_lens, attn_metadata, max_seq_len=self._fixed_max_seq_len ) + def _build_blocks_graph_runner( + self, config: "MultimodalEncoderCudaGraphConfig" + ) -> MultimodalEncoderGraphRunner: + rotary_cos_sin = self.rotary_pos_emb.rotary_cos_sin + rotary_dim = rotary_cos_sin.shape[-1] * 2 + # Match the dtypes of the eager tensors so the static-buffer copy never casts. + input_specs = { + "hidden_states": EncoderGraphTensorSpec( + shape=(self.config.hidden_size,), dtype=self.patch_embed.proj.weight.dtype + ), + "cos": EncoderGraphTensorSpec(shape=(rotary_dim,), dtype=rotary_cos_sin.dtype), + "sin": EncoderGraphTensorSpec(shape=(rotary_dim,), dtype=rotary_cos_sin.dtype), + "rope_position_ids": EncoderGraphTensorSpec(shape=(), dtype=torch.int32), + } + output_specs = {"hidden_states": 0} + output_specs.update( + {f"deepstack_{idx}": 0 for idx in range(len(self.deepstack_visual_indexes))} + ) + return MultimodalEncoderGraphRunner( + encoder_fn=self._encoder_graph_fn, + metadata_provider=_Qwen3VisionEncoderMetadataProvider(self), + input_specs=input_specs, + output_specs=output_specs, + config=config, + ) + + def _encoder_graph_fn( + self, + inputs: TypingMapping[str, torch.Tensor], + attn_metadata: AttentionMetadata, + ) -> dict[str, torch.Tensor]: + hidden_states, deepstack_hidden_states = self._run_encoder_blocks( + inputs["hidden_states"], + inputs["cos"], + inputs["sin"], + inputs["rope_position_ids"], + attn_metadata, + merge_deepstack=False, + ) + outputs = {"hidden_states": hidden_states} + outputs.update( + { + f"deepstack_{idx}": hidden_state + for idx, hidden_state in enumerate(deepstack_hidden_states) + } + ) + return outputs + + def _run_blocks( + self, + hidden_states: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rope_position_ids: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + if self._blocks_graph_runner is not None: + graph_outputs = self._blocks_graph_runner.maybe_run( + seq_lengths=attn_metadata.seq_lens.tolist(), + inputs={ + "hidden_states": hidden_states, + "cos": cos, + "sin": sin, + "rope_position_ids": rope_position_ids, + }, + ) + if graph_outputs is not None: + # Replay outputs borrow graph-owned buffers; the mergers run + # eagerly so the returned features own their storage. + return self.merger(graph_outputs["hidden_states"]), [ + merger(graph_outputs[f"deepstack_{idx}"]) + for idx, merger in enumerate(self.deepstack_merger_list) + ] + + hidden_states, deepstack_features = self._run_encoder_blocks( + hidden_states, + cos, + sin, + rope_position_ids, + attn_metadata, + merge_deepstack=True, + ) + return self.merger(hidden_states), deepstack_features + + def _run_encoder_blocks( + self, + hidden_states: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rope_position_ids: torch.Tensor, + attn_metadata: AttentionMetadata, + *, + merge_deepstack: bool, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + position_embeddings = (cos, sin) + deepstack_outputs = [] + for layer_num, block in enumerate(self.blocks): + hidden_states = block( + hidden_states, + position_ids=rope_position_ids, + attn_metadata=attn_metadata, + position_embeddings=position_embeddings, + ) + merger_idx = self._deepstack_layer_to_merger_idx.get(layer_num) + if merger_idx is not None: + deepstack_outputs.append( + self.deepstack_merger_list[merger_idx](hidden_states) + if merge_deepstack + else hidden_states + ) + return hidden_states, deepstack_outputs + @torch.inference_mode() def forward( self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs @@ -1114,23 +1247,55 @@ def forward( seq_len, dtype=torch.int32, device=self.device ) rope_position_ids = self._rope_position_ids_buffer[:seq_len] - position_embeddings = (cos, sin) + return self._run_blocks(hidden_states, cos, sin, rope_position_ids, self.attn_metadata) - deepstack_feature_lists = [] - for layer_num, block in enumerate(self.blocks): - hidden_states = block( - hidden_states, - position_ids=rope_position_ids, - attn_metadata=self.attn_metadata, - position_embeddings=position_embeddings, + +class _Qwen3VisionEncoderMetadataProvider(EncoderMetadataProvider): + """Own fixed-address attention metadata for Qwen3-VL graph buckets.""" + + def __init__(self, vision: Qwen3VisionModel) -> None: + self._vision = vision + self.graph_critical_attrs: Sequence[str] = ( + "_seq_lens_cuda", + "cu_q_seqlens", + "cu_kv_seqlens", + ) + if vision.model_config.attn_backend == "FLASHINFER": + self.graph_critical_attrs += ( + "workspace_buffer", + "_ragged_qo_indptr_buf", + "_ragged_kv_indptr_buf", ) - merger_idx = self._deepstack_layer_to_merger_idx.get(layer_num) - if merger_idx is not None: - deepstack_feature = self.deepstack_merger_list[merger_idx](hidden_states) - deepstack_feature_lists.append(deepstack_feature) - hidden_states = self.merger(hidden_states) - return hidden_states, deepstack_feature_lists + def build(self, key: EncoderGraphKey) -> AttentionMetadata: + vision = self._vision + assert vision.attn_metadata is not None + if ( + key.num_contexts > vision.attn_metadata.max_num_requests + or key.total_tokens > vision.attn_metadata.max_num_tokens + ): + raise ValueError( + f"Qwen3-VL encoder CUDA graph needs {key.total_tokens} tokens and " + f"{key.num_contexts} contexts, which exceeds the encoder attention metadata " + f"capacity (max_num_tokens={vision.attn_metadata.max_num_tokens}, " + f"max_num_requests={vision.attn_metadata.max_num_requests}). With " + "enable_padding=True, every bucket reserves one extra token and one padding " + "context. Raise encoder_max_num_tokens or use a smaller bucket." + ) + metadata = vision.attn_metadata.create_cuda_graph_metadata( + key.num_contexts, encode_only=True + ) + bind_seq_lens = getattr(metadata, "bind_encoder_cuda_graph_seq_lens", None) + if callable(bind_seq_lens): + bind_seq_lens(metadata.seq_lens, key.num_contexts) + return metadata + + def refresh_in_place( + self, + metadata: AttentionMetadata, + padded_seq_lengths: Sequence[int], + ) -> None: + self._vision.prepare_attn_metadata(list(padded_seq_lengths), metadata) def _qwen3vl_build_batched_input( @@ -1179,6 +1344,10 @@ def __init__( def post_config(self): self.config = self.model_config.pretrained_config.vision_config + def enable_cuda_graph(self) -> None: + """Capture the vision tower's opt-in encoder CUDA graphs.""" + self.visual.enable_blocks_cuda_graph() + def load_weights( self, weights: Dict[str, torch.Tensor], diff --git a/tensorrt_llm/_torch/models/modeling_radio.py b/tensorrt_llm/_torch/models/modeling_radio.py index 6e6e676f9ffa..2e3c29420d47 100644 --- a/tensorrt_llm/_torch/models/modeling_radio.py +++ b/tensorrt_llm/_torch/models/modeling_radio.py @@ -788,27 +788,6 @@ def setup_attn_metadata(self, max_num_tokens: int) -> None: metadata_kwargs["kv_layout"] = "NHD" self.attn_metadata = self.metadata_cls(**metadata_kwargs) - def enable_blocks_cuda_graph( - self, - config: "MultimodalEncoderCudaGraphConfig", - *, - device: Optional[torch.device] = None, - ) -> None: - """Build the block-stack runner from `config` and capture startup graphs. - - The caller must ensure the model parameters already live on `device` - (or on a CUDA device discoverable from `self.parameters()`); captures - are recorded immediately. - """ - if self._blocks_graph_runner is not None: - return - - graph_runner = self._build_blocks_graph_runner(config) - if device is None: - device = next(self.parameters()).device - graph_runner.capture_all(device) - self._blocks_graph_runner = graph_runner - def clear_blocks_cuda_graph(self, *, release_nccl_window_owners: bool = True @@ -1308,6 +1287,7 @@ def __init__(self, inter_feature_normalizer=inter_feature_normalizer, model_config=self.model_config, ) + self.radio_model.model._encoder_cuda_graph_config = encoder_cuda_graph_config def load_weights(self, weights): # Load radio_model weights for pytorch modules. diff --git a/tensorrt_llm/_torch/models/multimodal_encoder_graph.py b/tensorrt_llm/_torch/models/multimodal_encoder_graph.py index ac3ad36713eb..7632e5ec3cc8 100644 --- a/tensorrt_llm/_torch/models/multimodal_encoder_graph.py +++ b/tensorrt_llm/_torch/models/multimodal_encoder_graph.py @@ -213,7 +213,11 @@ def capture_all(self, device: torch.device) -> None: `enable_padding=True` captures the with-dummy-context variant for each bucket; `enable_padding=False` captures the exact-fit variant. """ - for bucket in self._buckets: + # Capture larger workloads first. Attention backends may share a graph workspace across + # metadata variants; warming up a later, larger workload can resize that workspace and + # invalidate pointers captured by an earlier graph. This matches the ordering used by the + # LLM CUDA-graph runners and also lets smaller graphs reuse the larger graph's memory pool. + for bucket in reversed(self._buckets): key = self._padded_key_for_bucket(bucket) if key in self._captured: continue diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index 6f3db93bf47f..6b54bb978abc 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -157,6 +157,13 @@ def mm_encoder_cache_enabled(model: nn.Module) -> bool: return isinstance(model, MultimodalModelMixin) and model.encoder_cache_active +def setup_mm_encoder_cuda_graphs(model: nn.Module) -> None: + """Capture local encoder graphs after metadata setup and before KV sizing.""" + for module in model.modules(): + if isinstance(module, MultimodalEncoderMixin): + module.enable_blocks_cuda_graph() + + def setup_mm_encoder_attn_metadata( model: nn.Module, input_processor: InputProcessor, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index a224c4989fb6..49345bb6df11 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -55,7 +55,8 @@ from .engine.model_call import ModelCaller from .engine.multimodal import (MultimodalItemScheduler, is_multimodal, mm_encoder_cache_enabled, - setup_mm_encoder_attn_metadata, use_mrope) + setup_mm_encoder_attn_metadata, + setup_mm_encoder_cuda_graphs, use_mrope) from .engine.runners import resolve_runner_type from .engine.runners.common import (make_scheduled_inputs, resolve_mrope_position_deltas_cache) @@ -538,6 +539,10 @@ def __init__( set_torch_compiling(False) self._nccl_window_tensor_scope_hooks = ( install_eager_nccl_window_tensor_scopes(self.model)) + + # Capture local encoders after their metadata is sized and before + # KV-cache estimation, so graph memory is included in the profile. + setup_mm_encoder_cuda_graphs(self.model) except Exception as e: import traceback traceback.print_exception(Exception, e, e.__traceback__) diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index d7c3a419ccc0..ab77e67f4b8e 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -103,6 +103,7 @@ l0_h100: # Breakable prefill CUDA graphs (BCG) on the hybrid Mamba path; mpi_ray_parity (Ray id runs in the ray stage below). - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph[tp1] TIMEOUT (90) - unittest/_torch/modeling/test_multimodal_encoder_graph.py + - unittest/_torch/modeling/test_modeling_qwen3vl.py::test_qwen3_vision_cuda_graph_replay_matches_eager # Register the full EXAONE suite on Hopper to include the SM90-only # chunked-prefill and KV-cache-reuse scenarios. - unittest/_torch/modeling/test_modeling_exaone4_5.py diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index e1299d2b5e88..cfc38b3b2fed 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -12,6 +12,7 @@ make_mm_request, ) +from tensorrt_llm._torch.models.modeling_multimodal_encoder import MultimodalEncoderMixin from tensorrt_llm._torch.models.modeling_multimodal_mixin import ( MultimodalEncoderContractError, MultimodalModelMixin, @@ -21,6 +22,7 @@ resolve_bytes_per_mm_encoder_embedding, resolve_mm_encoder_output_budget, setup_mm_encoder_attn_metadata, + setup_mm_encoder_cuda_graphs, validate_mm_encoder_scheduling_compatibility, ) from tensorrt_llm._torch.pyexecutor.llm_request import ( @@ -40,7 +42,11 @@ BaseMultimodalDummyInputsBuilder, MultimodalEncoderItemMetadata, ) -from tensorrt_llm.llmapi.llm_args import MultimodalConfig, MultimodalEncoderSchedulingPolicy +from tensorrt_llm.llmapi.llm_args import ( + MultimodalConfig, + MultimodalEncoderCudaGraphConfig, + MultimodalEncoderSchedulingPolicy, +) # The item-scheduling surface is pure logic: no kernels, no device transfers. pytestmark = pytest.mark.cpu_only @@ -55,6 +61,57 @@ def test_disabled_encoder_does_not_initialize_processor(): processor.get_mm_max_tokens_per_item.assert_not_called() +def test_encoder_graph_capture_follows_metadata_setup_and_runs_once() -> None: + events = [] + + class Encoder(torch.nn.Module, MultimodalEncoderMixin): + def __init__(self, name: str, config: MultimodalEncoderCudaGraphConfig | None) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(1)) + self.name = name + self._encoder_cuda_graph_config = config + self._blocks_graph_runner = None + self.runner = Mock() + self.runner.capture_all.side_effect = lambda device: events.append( + (self.name, "capture", device) + ) + + def setup_attn_metadata( + self, + max_num_tokens: int, + attention_metadata_capacity: dict[str, int] | None = None, + ) -> None: + events.append((self.name, "metadata", max_num_tokens)) + + def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) -> Mock: + assert events.index((self.name, "metadata", 1024)) >= 0 + assert config is self._encoder_cuda_graph_config + return self.runner + + config = MultimodalEncoderCudaGraphConfig(buckets=[(256, 1)]) + model = torch.nn.Module() + model.first = Encoder("first", config) + model.second = Encoder("second", config) + model.disabled = Encoder("disabled", None) + model.shared = model.first + model.unrelated = torch.nn.Linear(1, 1) + + setup_mm_encoder_attn_metadata(model, object(), 1024, None) + setup_mm_encoder_cuda_graphs(model) + setup_mm_encoder_cuda_graphs(model) + + assert events == [ + ("first", "metadata", 1024), + ("second", "metadata", 1024), + ("disabled", "metadata", 1024), + ("first", "capture", torch.device("cpu")), + ("second", "capture", torch.device("cpu")), + ] + assert model.first._blocks_graph_runner is model.first.runner + assert model.second._blocks_graph_runner is model.second.runner + model.disabled.runner.capture_all.assert_not_called() + + def _bind_items(mm_item_scheduler: MultimodalItemScheduler, request, *, row_bytes: int = 8) -> None: state = request.py_mm_encoder_state for item_idx, rows in enumerate(state.embedding_lengths): diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index 5763a76c2ed8..e929b6b09790 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -33,6 +33,8 @@ Glm5NextVLM, ) from tensorrt_llm._torch.pyexecutor.config_utils import get_glm5_next_layer_masks, is_glm5_next +from tensorrt_llm._torch.pyexecutor.engine.multimodal import setup_mm_encoder_cuda_graphs +from tensorrt_llm.llmapi.llm_args import MultimodalConfig, MultimodalEncoderCudaGraphConfig from tensorrt_llm.mapping import Mapping @@ -880,14 +882,24 @@ def test_kda_shards_preserve_fp32_gate_parameters(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("encoder_cuda_graph", [False, True], ids=["eager", "cuda-graph"]) @torch.inference_mode() -def test_vision_attention_dp_is_local_and_matches_single_rank(): +def test_vision_attention_dp_is_local_and_matches_single_rank(encoder_cuda_graph: bool) -> None: config = _config() reference = Glm5NextVisionModelBase(ModelConfig(pretrained_config=config)).cuda().eval() reference.visual.setup_attn_metadata(max_num_tokens=128) torch.manual_seed(7) for parameter in reference.parameters(): parameter.normal_(0, 0.05) + multimodal_config = ( + MultimodalConfig( + encoder_cuda_graph={ + "image": MultimodalEncoderCudaGraphConfig(buckets=[(16, 1), (32, 1)]) + } + ) + if encoder_cuda_graph + else None + ) # Rank 0 has no image; the other ranks see different token counts and # pixels. Any TP collective on the encoder path is an error, even when @@ -900,18 +912,30 @@ def local_allreduce(**kwargs): for rank, width in enumerate((0, 4, 6, 8)): mapping = Mapping(world_size=4, tp_size=4, rank=rank, enable_attention_dp=True) local = ( - Glm5NextVisionModelBase(ModelConfig(pretrained_config=config, mapping=mapping)) + Glm5NextVisionModelBase( + ModelConfig( + pretrained_config=config, + mapping=mapping, + multimodal_config=multimodal_config, + ) + ) .cuda() .eval() ) local.load_state_dict(reference.state_dict()) local.visual.setup_attn_metadata(max_num_tokens=128) + setup_mm_encoder_cuda_graphs(local) if not width: continue pixels = torch.randn(4 * width, 3 * 2 * 2 * 2, device="cuda", dtype=torch.bfloat16) grid = torch.tensor([[1, 4, width]]) expected = reference.encode_batched(pixels, grid) - actual = local.encode_batched(pixels, grid) + with patch.object( + local.visual, "_run_blocks_eager", wraps=local.visual._run_blocks_eager + ) as eager_blocks: + actual = local.encode_batched(pixels, grid) + if encoder_cuda_graph: + eager_blocks.assert_not_called() torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index 0d285f6a3da6..365bf6300bd4 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -26,6 +26,7 @@ ) from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.llmapi.llm_args import MultimodalConfig, MultimodalEncoderCudaGraphConfig from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -495,7 +496,7 @@ def capture_prepare_attn_metadata( monkeypatch.setattr( modeling_qwen3vl, - "_prepare_qwen_vl_vision_attn_metadata", + "prepare_encoder_attn_metadata", capture_prepare_attn_metadata, ) model = Qwen3VisionModel.__new__(Qwen3VisionModel) @@ -530,6 +531,7 @@ def test_qwen3_vision_patch_projection_matches_conv3d() -> None: torch.nn.Module.__init__(vision) vision.patch_embed = patch_embed vision.attn_metadata = object() + vision._blocks_graph_runner = None vision._rope_position_ids_buffer = None vision.rot_pos_emb = lambda grid: ( torch.empty(6, 0), @@ -547,6 +549,132 @@ def test_qwen3_vision_patch_projection_matches_conv3d() -> None: assert deepstack == [] +def test_qwen3_vision_run_blocks_uses_encoder_graph() -> None: + class FakeGraphRunner: + def __init__(self) -> None: + self.calls = [] + + def maybe_run(self, **kwargs): + self.calls.append(kwargs) + return { + "hidden_states": torch.full((4, 8), 2.0), + "deepstack_0": torch.full((4, 8), 3.0), + } + + vision = Qwen3VisionModel.__new__(Qwen3VisionModel) + torch.nn.Module.__init__(vision) + vision._blocks_graph_runner = FakeGraphRunner() + vision.deepstack_visual_indexes = [0] + vision.deepstack_merger_list = torch.nn.ModuleList([torch.nn.Identity()]) + vision.merger = torch.nn.Identity() + metadata = SimpleNamespace(seq_lens=torch.tensor([4])) + hidden_states = torch.ones(4, 8) + cos = torch.ones(4, 2) + sin = torch.zeros(4, 2) + position_ids = torch.arange(4, dtype=torch.int32) + + output, deepstack = vision._run_blocks(hidden_states, cos, sin, position_ids, metadata) + + assert torch.equal(output, torch.full((4, 8), 2.0)) + assert len(deepstack) == 1 + assert torch.equal(deepstack[0], torch.full((4, 8), 3.0)) + assert vision._blocks_graph_runner.calls == [ + { + "seq_lengths": [4], + "inputs": { + "hidden_states": hidden_states, + "cos": cos, + "sin": sin, + "rope_position_ids": position_ids, + }, + } + ] + + +def test_qwen3_vision_run_blocks_preserves_eager_merger_order() -> None: + events = [] + + class RecordingBlock(torch.nn.Module): + def __init__(self, name: str) -> None: + super().__init__() + self.name = name + + def forward(self, hidden_states, **kwargs): + events.append(self.name) + return hidden_states + + class RecordingMerger(torch.nn.Module): + def __init__(self, name: str) -> None: + super().__init__() + self.name = name + + def forward(self, hidden_states): + events.append(self.name) + return hidden_states + + vision = Qwen3VisionModel.__new__(Qwen3VisionModel) + torch.nn.Module.__init__(vision) + vision._blocks_graph_runner = None + vision.blocks = torch.nn.ModuleList([RecordingBlock("block_0"), RecordingBlock("block_1")]) + vision._deepstack_layer_to_merger_idx = {0: 0} + vision.deepstack_merger_list = torch.nn.ModuleList([RecordingMerger("deepstack_merger")]) + vision.merger = RecordingMerger("final_merger") + hidden_states = torch.ones(4, 8) + + vision._run_blocks( + hidden_states, + torch.ones(4, 2), + torch.zeros(4, 2), + torch.arange(4, dtype=torch.int32), + SimpleNamespace(seq_lens=torch.tensor([4])), + ) + + assert events == ["block_0", "deepstack_merger", "block_1", "final_merger"] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("attn_backend", ["TRTLLM", "FLASHINFER"]) +def test_qwen3_vision_cuda_graph_replay_matches_eager( + monkeypatch: pytest.MonkeyPatch, attn_backend: str +) -> None: + config_dict = copy.deepcopy(QWEN3_VL_8B_CONFIG) + config_dict["vision_config"].update(depth=2, deepstack_visual_indexes=[0]) + model_config = ModelConfig( + pretrained_config=Qwen3VLConfig.from_dict(config_dict), + attn_backend=attn_backend, + multimodal_config=MultimodalConfig( + encoder_cuda_graph={"vision": MultimodalEncoderCudaGraphConfig(buckets=[(256, 1)])} + ), + ) + torch.manual_seed(0) + encoder = modeling_qwen3vl.Qwen3VisionModelBase(model_config, Qwen3VisionModel).cuda().eval() + with torch.no_grad(): + for name, param in encoder.named_parameters(): + if "norm" in name and name.endswith("weight"): + param.fill_(1.0) + else: + param.normal_(std=0.02) + vision = encoder.visual + vision.setup_attn_metadata(max_num_tokens=1024) + # 192 real tokens replay the (256, 1) bucket with a 65-token padding context. + grid_thw = torch.tensor([[1, 12, 16]]) + pixel_values = torch.randn(192, 3 * 2 * 16 * 16, device="cuda") + eager = encoder.encode_batched(pixel_values, grid_thw) + + vision.enable_blocks_cuda_graph() + + def fail_eager_fallback(*args, **kwargs): + raise AssertionError("Qwen3-VL vision CUDA graph replay fell back to eager.") + + monkeypatch.setattr(vision, "_run_encoder_blocks", fail_eager_fallback) + graph = encoder.encode_batched(pixel_values, grid_thw) + + torch.testing.assert_close(graph, eager, rtol=5e-2, atol=5e-2) + (captured,) = vision._blocks_graph_runner._captured.values() + # Replay must refresh the device lengths read by attention, not only the host copy. + assert captured.metadata.seq_lens_cuda.tolist() == [192, 65] + + def test_qwen3_processor_max_pixels_maps_to_fixed_attention_capacity() -> None: max_pixels = 16_777_216 expected_max_tokens_per_item = {"image": 65_536} diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py index 0d2eea9b1510..23d7a237dae7 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py @@ -300,6 +300,27 @@ def test_empty_buckets_rejected(): ) +def test_capture_all_processes_largest_bucket_first() -> None: + runner = _make_logic_runner( + buckets=[ + EncoderGraphKey(num_contexts=1, total_tokens=256), + EncoderGraphKey(num_contexts=1, total_tokens=1024), + EncoderGraphKey(num_contexts=1, total_tokens=512), + ] + ) + captured_keys = [] + + def record_capture( + key: EncoderGraphKey, padded_seq_lengths: list[int], device: torch.device + ) -> None: + captured_keys.append(key) + + with mock.patch.object(runner, "_capture_key", side_effect=record_capture): + runner.capture_all(torch.device("cpu")) + + assert [key.total_tokens for key in captured_keys] == [1025, 513, 257] + + def test_metadata_buffer_snapshot_records_declared_attrs(): md = _ToyMetadata( max_contexts=2, From 18d091634cfc30694fa7e4a880b03b379f9fe4c9 Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Tue, 13 Oct 2026 01:39:36 -0700 Subject: [PATCH 2/8] [None][fix] preserve deepstack merger order in encoder graphs Select each eager merger by its checkpoint index while preserving the layer order of captured deepstack states. Cover reordered indexes on TRTLLM and FlashInfer, and mark graph logic regressions for CPU CI collection. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_qwen3vl.py | 9 ++++- .../_torch/modeling/test_modeling_qwen3vl.py | 40 ++++++++++++++----- .../modeling/test_multimodal_encoder_graph.py | 1 + 3 files changed, 37 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 479f60105a94..ae7442865551 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -1158,9 +1158,14 @@ def _run_blocks( if graph_outputs is not None: # Replay outputs borrow graph-owned buffers; the mergers run # eagerly so the returned features own their storage. + # Blocks emit deepstack states in layer order; each merger's + # checkpoint index follows deepstack_visual_indexes instead. + mergers = [ + self.deepstack_merger_list[merger_idx] + for _, merger_idx in sorted(self._deepstack_layer_to_merger_idx.items()) + ] return self.merger(graph_outputs["hidden_states"]), [ - merger(graph_outputs[f"deepstack_{idx}"]) - for idx, merger in enumerate(self.deepstack_merger_list) + merger(graph_outputs[f"deepstack_{idx}"]) for idx, merger in enumerate(mergers) ] hidden_states, deepstack_features = self._run_encoder_blocks( diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index 365bf6300bd4..0039586220ed 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -549,23 +549,38 @@ def test_qwen3_vision_patch_projection_matches_conv3d() -> None: assert deepstack == [] -def test_qwen3_vision_run_blocks_uses_encoder_graph() -> None: +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("deepstack_visual_indexes", "expected_features"), [([0], [3.0]), ([1, 0], [6.0, 4.0])] +) +def test_qwen3_vision_run_blocks_uses_encoder_graph( + deepstack_visual_indexes: list[int], expected_features: list[float] +) -> None: class FakeGraphRunner: def __init__(self) -> None: self.calls = [] - def maybe_run(self, **kwargs): + def maybe_run(self, **kwargs: object) -> dict[str, torch.Tensor]: self.calls.append(kwargs) return { "hidden_states": torch.full((4, 8), 2.0), "deepstack_0": torch.full((4, 8), 3.0), + "deepstack_1": torch.full((4, 8), 4.0), } vision = Qwen3VisionModel.__new__(Qwen3VisionModel) torch.nn.Module.__init__(vision) vision._blocks_graph_runner = FakeGraphRunner() - vision.deepstack_visual_indexes = [0] - vision.deepstack_merger_list = torch.nn.ModuleList([torch.nn.Identity()]) + vision.deepstack_visual_indexes = deepstack_visual_indexes + vision._deepstack_layer_to_merger_idx = { + layer_idx: idx for idx, layer_idx in enumerate(deepstack_visual_indexes) + } + vision.deepstack_merger_list = torch.nn.ModuleList( + [torch.nn.Linear(8, 8, bias=False) for _ in deepstack_visual_indexes] + ) + with torch.no_grad(): + for idx, merger in enumerate(vision.deepstack_merger_list): + merger.weight.copy_(torch.eye(8) * (idx + 1)) vision.merger = torch.nn.Identity() metadata = SimpleNamespace(seq_lens=torch.tensor([4])) hidden_states = torch.ones(4, 8) @@ -576,8 +591,9 @@ def maybe_run(self, **kwargs): output, deepstack = vision._run_blocks(hidden_states, cos, sin, position_ids, metadata) assert torch.equal(output, torch.full((4, 8), 2.0)) - assert len(deepstack) == 1 - assert torch.equal(deepstack[0], torch.full((4, 8), 3.0)) + assert len(deepstack) == len(expected_features) + for features, expected in zip(deepstack, expected_features): + assert torch.equal(features, torch.full((4, 8), expected)) assert vision._blocks_graph_runner.calls == [ { "seq_lengths": [4], @@ -591,6 +607,7 @@ def maybe_run(self, **kwargs): ] +@pytest.mark.cpu_only def test_qwen3_vision_run_blocks_preserves_eager_merger_order() -> None: events = [] @@ -599,7 +616,7 @@ def __init__(self, name: str) -> None: super().__init__() self.name = name - def forward(self, hidden_states, **kwargs): + def forward(self, hidden_states: torch.Tensor, **kwargs: object) -> torch.Tensor: events.append(self.name) return hidden_states @@ -608,7 +625,7 @@ def __init__(self, name: str) -> None: super().__init__() self.name = name - def forward(self, hidden_states): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: events.append(self.name) return hidden_states @@ -634,11 +651,12 @@ def forward(self, hidden_states): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("attn_backend", ["TRTLLM", "FLASHINFER"]) +@pytest.mark.parametrize("deepstack_visual_indexes", [[0], [1, 0]]) def test_qwen3_vision_cuda_graph_replay_matches_eager( - monkeypatch: pytest.MonkeyPatch, attn_backend: str + monkeypatch: pytest.MonkeyPatch, attn_backend: str, deepstack_visual_indexes: list[int] ) -> None: config_dict = copy.deepcopy(QWEN3_VL_8B_CONFIG) - config_dict["vision_config"].update(depth=2, deepstack_visual_indexes=[0]) + config_dict["vision_config"].update(depth=2, deepstack_visual_indexes=deepstack_visual_indexes) model_config = ModelConfig( pretrained_config=Qwen3VLConfig.from_dict(config_dict), attn_backend=attn_backend, @@ -663,7 +681,7 @@ def test_qwen3_vision_cuda_graph_replay_matches_eager( vision.enable_blocks_cuda_graph() - def fail_eager_fallback(*args, **kwargs): + def fail_eager_fallback(*args: object, **kwargs: object) -> None: raise AssertionError("Qwen3-VL vision CUDA graph replay fell back to eager.") monkeypatch.setattr(vision, "_run_encoder_blocks", fail_eager_fallback) diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py index 23d7a237dae7..debb044f0f29 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py @@ -300,6 +300,7 @@ def test_empty_buckets_rejected(): ) +@pytest.mark.cpu_only def test_capture_all_processes_largest_bucket_first() -> None: runner = _make_logic_runner( buckets=[ From 3aa1e30ab4f72893f065afabc34550029f1adab4 Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Wed, 14 Oct 2026 01:08:43 -0700 Subject: [PATCH 3/8] [None][refactor] generalize multimodal encoder graph interfaces Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- .../models/modeling_glm5_next_vision.py | 21 +- .../models/modeling_multimodal_encoder.py | 34 +- .../models/modeling_nemotron_h_multimodal.py | 11 +- .../_torch/models/modeling_qwen3vl.py | 15 +- tensorrt_llm/_torch/models/modeling_radio.py | 5 +- .../_torch/models/multimodal_encoder_graph.py | 320 ++++++++++-------- .../_torch/pyexecutor/engine/multimodal.py | 39 ++- .../_torch/pyexecutor/model_engine.py | 5 +- .../_torch/executor/engine/test_multimodal.py | 52 +++ .../test_modeling_nemotron_h_multimodal.py | 12 +- .../modeling/test_multimodal_encoder_graph.py | 148 +++++++- 11 files changed, 472 insertions(+), 190 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py index 5b0514b8b87c..cefa2a654ec2 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py @@ -95,21 +95,6 @@ _VISION_WEIGHT_PREFIX = "model.visual" -def _image_encoder_cuda_graph_config( - model_config: ModelConfig[PretrainedConfig], -) -> Optional["MultimodalEncoderCudaGraphConfig"]: - mm_config = model_config.multimodal_config - if mm_config is None or mm_config.encoder_cuda_graph is None: - return None - unknown = set(mm_config.encoder_cuda_graph) - {"image"} - if unknown: - raise ValueError( - "glm5_next: unsupported multimodal encoder CUDA graph modalities " - f"{sorted(unknown)}; only 'image' is supported" - ) - return mm_config.encoder_cuda_graph.get("image") - - def _text_dtype(model_config: ModelConfig[PretrainedConfig]) -> torch.dtype: text_config = unwrap_glm5_next_text_config(model_config.pretrained_config) dtype = getattr(text_config, "dtype", None) @@ -382,6 +367,8 @@ class Glm5NextVisionModel(nn.Module, MultimodalEncoderMixin): """The GLM vision tower: patch embed → 24 blocks → post-norm → downsample → projector. Context-only, ``kv_cache_manager=None``.""" + encoder_cuda_graph_modality = "image" + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: super().__init__() _require_trtllm_vision_backend(model_config, type(self).__name__) @@ -442,7 +429,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: # Optional CUDA-graph replay of the block stack (see # `enable_blocks_cuda_graph`); opted in through # `multimodal_config.encoder_cuda_graph["image"]`. - self._encoder_cuda_graph_config = _image_encoder_cuda_graph_config(model_config) + self.configure_encoder_cuda_graph(model_config.multimodal_config) self._blocks_graph_runner: Optional[MultimodalEncoderGraphRunner] = None @property @@ -612,7 +599,7 @@ def project_merged(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.merger(hidden_states) -class _Glm5NextVisionGraphMetadataProvider(EncoderMetadataProvider): +class _Glm5NextVisionGraphMetadataProvider(EncoderMetadataProvider[AttentionMetadata]): """Graph-owned TRTLLM metadata for the block-stack runner. Per bucket one metadata instance with ``is_cuda_graph=True`` so the diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py index 59965f3f2594..de594226305e 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py @@ -27,7 +27,8 @@ from .modeling_multimodal_utils import multiscale_forward if TYPE_CHECKING: - from ...llmapi.llm_args import MultimodalEncoderCudaGraphConfig + from ...llmapi.llm_args import (MultimodalConfig, + MultimodalEncoderCudaGraphConfig) from .multimodal_encoder_graph import MultimodalEncoderGraphRunner # Legacy fallback for encoders that do not yet opt into atomic-item scheduling @@ -47,7 +48,17 @@ def prepare_encoder_attn_metadata(seq_lens: list[int], attn_metadata: AttentionMetadata, *, max_seq_len: int) -> AttentionMetadata: - """Refresh packed no-KV encoder lengths while preserving captured buffers.""" + """Refresh packed no-KV encoder lengths while preserving captured buffers. + + Args: + seq_lens: Nonempty list of nonnegative lengths for independent attention segments, + including the isolated padding segment when using a padded graph. + attn_metadata: Initialized encoder metadata. Captured length buffers are updated in place. + max_seq_len: Positive, stable per-segment capacity, at least the largest live length. + + Returns: + The supplied metadata, prepared for context-only attention without a KV cache. + """ seq_lens = [int(seq_len) for seq_len in seq_lens] if not seq_lens: raise ValueError("Encoder attention requires at least one segment") @@ -132,9 +143,28 @@ class MultimodalEncoderMixin: """ metadata_cls: Type[AttentionMetadata] attn_metadata: Optional[AttentionMetadata] = None + encoder_cuda_graph_modality: str | None = None _encoder_cuda_graph_config: "MultimodalEncoderCudaGraphConfig | None" = None _blocks_graph_runner: "MultimodalEncoderGraphRunner | None" = None + def configure_encoder_cuda_graph( + self, multimodal_config: "MultimodalConfig | None") -> None: + """Select this tower's configuration without rejecting other towers' modalities. + + Subclasses declare ``encoder_cuda_graph_modality`` when their graph factory is supported. + The engine validates the requested keys against all local encoder capabilities before + capturing any graphs. Explicit configs for standalone encoders remain supported. + + Args: + multimodal_config: Model-wide configuration shared by its encoder towers. + """ + graph_configs = (multimodal_config.encoder_cuda_graph + if multimodal_config is not None else None) + self._encoder_cuda_graph_config = ( + graph_configs.get(self.encoder_cuda_graph_modality) + if graph_configs is not None + and self.encoder_cuda_graph_modality is not None else None) + def _build_blocks_graph_runner( self, config: "MultimodalEncoderCudaGraphConfig" ) -> "MultimodalEncoderGraphRunner": diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py index be992ef8f9a1..38b6fb4bf048 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h_multimodal.py @@ -283,18 +283,9 @@ class DynamicResolutionParams: def _get_vision_encoder_cuda_graph_config( mm_config: Optional["MultimodalConfig"], ) -> Optional["MultimodalEncoderCudaGraphConfig"]: + """Select the vision tower's graph config from the model-wide modality mapping.""" if mm_config is None or mm_config.encoder_cuda_graph is None: return None - - # TODO: Add support for audio encoder CUDA graphs. - supported_modalities = {"vision"} - unknown_modalities = set(mm_config.encoder_cuda_graph) - supported_modalities - if unknown_modalities: - raise ValueError( - f"Unsupported multimodal encoder CUDA graph modalities: {sorted(unknown_modalities)}. " - f"Supported modalities: {sorted(supported_modalities)}." - ) - return mm_config.encoder_cuda_graph.get("vision") diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index ae7442865551..c4a370588018 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -816,6 +816,8 @@ def _triton_pos_embed_interpolate( class Qwen3VisionModel(torch.nn.Module, MultimodalEncoderMixin): + encoder_cuda_graph_modality = "vision" + def __init__(self, model_config: ModelConfig[PretrainedConfig]): super().__init__() self.model_config = model_config @@ -891,16 +893,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): self._fixed_max_seq_len = self.model_config.max_num_tokens self._blocks_graph_runner: MultimodalEncoderGraphRunner | None = None - mm_config = self.model_config.multimodal_config - self._encoder_cuda_graph_config: "MultimodalEncoderCudaGraphConfig | None" = None - if mm_config is not None and mm_config.encoder_cuda_graph is not None: - unknown_modalities = set(mm_config.encoder_cuda_graph) - {"vision"} - if unknown_modalities: - raise ValueError( - "Unsupported Qwen3-VL encoder CUDA graph modalities: " - f"{sorted(unknown_modalities)}. Supported modalities: ['vision']." - ) - self._encoder_cuda_graph_config = mm_config.encoder_cuda_graph.get("vision") + self.configure_encoder_cuda_graph(self.model_config.multimodal_config) # Vision block's `rope_position_ids` scratch. Registered empty here; # `setup_attn_metadata` allocates it as an `arange` (see there). @@ -1255,7 +1248,7 @@ def forward( return self._run_blocks(hidden_states, cos, sin, rope_position_ids, self.attn_metadata) -class _Qwen3VisionEncoderMetadataProvider(EncoderMetadataProvider): +class _Qwen3VisionEncoderMetadataProvider(EncoderMetadataProvider[AttentionMetadata]): """Own fixed-address attention metadata for Qwen3-VL graph buckets.""" def __init__(self, vision: Qwen3VisionModel) -> None: diff --git a/tensorrt_llm/_torch/models/modeling_radio.py b/tensorrt_llm/_torch/models/modeling_radio.py index 2e3c29420d47..b069d8374763 100644 --- a/tensorrt_llm/_torch/models/modeling_radio.py +++ b/tensorrt_llm/_torch/models/modeling_radio.py @@ -565,6 +565,8 @@ class VisionTransformer(nn.Module, MultimodalEncoderMixin): Modified from https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py. """ + encoder_cuda_graph_modality = "vision" + def __init__( self, img_size: int = 224, @@ -927,7 +929,8 @@ def _run_blocks_eager(self, x: torch.Tensor, return x -class _VisionEncoderMetadataProvider(EncoderMetadataProvider): +class _VisionEncoderMetadataProvider(EncoderMetadataProvider[AttentionMetadata] + ): """Bridges the multimodal encoder graph runner to the model's eager metadata path. `build` returns a fresh `AttentionMetadata` with `is_cuda_graph=True`; that flag flips diff --git a/tensorrt_llm/_torch/models/multimodal_encoder_graph.py b/tensorrt_llm/_torch/models/multimodal_encoder_graph.py index 7632e5ec3cc8..2ad203e2a760 100644 --- a/tensorrt_llm/_torch/models/multimodal_encoder_graph.py +++ b/tensorrt_llm/_torch/models/multimodal_encoder_graph.py @@ -9,7 +9,8 @@ Design highlights: -* Padding uses an isolated dummy context, never an extension of the last real context. +* The default packed adapter uses an isolated dummy padding context, never an extension of the + last real context. Other encoder layouts supply their own shape, mask, and padding rules. Full-attention vision encoders would otherwise let real tokens attend to padding rows. * When `enable_padding=True`, the runner reserves one extra token beyond each bucket's real-token budget so the dummy padding context is always non-empty (the attention backend rejects @@ -22,10 +23,13 @@ from __future__ import annotations +from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import ( TYPE_CHECKING, Callable, Dict, + Generic, List, Mapping, NamedTuple, @@ -33,6 +37,7 @@ Protocol, Sequence, Tuple, + TypeVar, ) import torch @@ -44,7 +49,8 @@ if TYPE_CHECKING: from ...llmapi.llm_args import MultimodalEncoderCudaGraphConfig - from ..attention.backends import AttentionMetadata + +_MetadataT = TypeVar("_MetadataT") # Single token kept aside for the dummy padding context when `enable_padding=True`. # The attention backend rejects zero-length contexts, so the dummy must contain at least one token @@ -62,9 +68,8 @@ class EncoderGraphKey(NamedTuple): * As a configured bucket, `total_tokens` is the maximum number of real (post-encode) tokens the bucket is meant to cover, and `num_contexts` is the number of real contexts. Neither includes any padding reservation; the runner adds that internally. - * As a captured graph key, both fields reflect the padded state seen by the captured kernels. - When `enable_padding=True` the key's `total_tokens` is `bucket.total_tokens + 1` and - `num_contexts` is `bucket.num_contexts + 1`. + * As a captured graph key, both fields describe the adapter's fixed capture capacity. + The default packed adapter adds one token and one context when padding is enabled. """ total_tokens: int @@ -89,12 +94,146 @@ def materialize(self, total_tokens: int, device: torch.device) -> torch.Tensor: return torch.zeros(tuple(full_shape), dtype=self.dtype, device=device) -class _CapturedGraph(NamedTuple): +class EncoderGraphIOAdapter(ABC): + """Define an encoder's captured layout, padding, and tensor transfers. + + The default layout packs independent contexts along one token axis and isolates padding in a + dummy context. Encoders with batched time axes or masks can override ``graph_key``, + ``capture_seq_lengths``, and ``padded_seq_lengths`` together. ``can_replay`` can reject a + workload that fits the token budget but does not fit that layout. + + Input buffers must retain their addresses across replays. Output tensors are borrowed views + into graph-owned buffers; collection may trim different axes or output lengths without + assuming that the encoder preserves the input token count. + """ + + @abstractmethod + def allocate_inputs( + self, key: EncoderGraphKey, device: torch.device + ) -> dict[str, torch.Tensor]: + """Allocate fixed buffers initialized with safe inputs for warmup and capture.""" + + @abstractmethod + def copy_inputs( + self, + static_inputs: Mapping[str, torch.Tensor], + inputs: Mapping[str, torch.Tensor], + seq_lengths: Sequence[int], + ) -> None: + """Copy live inputs and refresh masks/padding without replacing buffers.""" + + @abstractmethod + def collect_outputs( + self, outputs: Mapping[str, torch.Tensor], seq_lengths: Sequence[int] + ) -> dict[str, torch.Tensor]: + """Return output views for the live workload, including any masks.""" + + def can_replay(self, bucket: EncoderGraphKey, seq_lengths: Sequence[int]) -> bool: + """Whether a workload fitting ``bucket`` also fits this encoder's layout.""" + return True + + def graph_key(self, bucket: EncoderGraphKey, *, enable_padding: bool) -> EncoderGraphKey: + """Map a configured bucket to its fixed capture capacity.""" + if enable_padding: + return EncoderGraphKey( + total_tokens=bucket.total_tokens + _PADDING_TOKEN_RESERVE, + num_contexts=bucket.num_contexts + 1, + ) + return bucket + + def padded_seq_lengths( + self, seq_lengths: Sequence[int], *, bucket: EncoderGraphKey, enable_padding: bool + ) -> list[int]: + """Describe the contexts passed to the metadata provider before replay.""" + if not enable_padding: + return list(seq_lengths) + padding = bucket.total_tokens - sum(seq_lengths) + _PADDING_TOKEN_RESERVE + return [*seq_lengths, padding] + + def capture_seq_lengths(self, bucket: EncoderGraphKey, key: EncoderGraphKey) -> list[int]: + """Use worst-case segment skew so capture covers every selectable workload.""" + if key.num_contexts == bucket.num_contexts: + return [bucket.total_tokens - (bucket.num_contexts - 1)] + [1] * ( + bucket.num_contexts - 1 + ) + real = [1] * bucket.num_contexts + padding = key.total_tokens - bucket.num_contexts + if padding <= 0: + raise ValueError( + f"Bucket {bucket} cannot host {bucket.num_contexts} real contexts " + f"plus a dummy padding context in {key.total_tokens} tokens." + ) + return real + [padding] + + +class PackedEncoderGraphIOAdapter(EncoderGraphIOAdapter): + """Use token-axis tensor specs for packed, length-preserving encoder blocks. + + Args: + input_specs: Non-token dimensions, dtype, and token axis of each graph input. + output_specs: Token axis of each graph output; outputs preserve the real input token count. + """ + + def __init__( + self, + input_specs: Mapping[str, EncoderGraphTensorSpec], + output_specs: Mapping[str, int], + ) -> None: + self._input_specs = dict(input_specs) + self._output_specs = dict(output_specs) + + def allocate_inputs( + self, key: EncoderGraphKey, device: torch.device + ) -> dict[str, torch.Tensor]: + return { + name: spec.materialize(key.total_tokens, device) + for name, spec in self._input_specs.items() + } + + def copy_inputs( + self, + static_inputs: Mapping[str, torch.Tensor], + inputs: Mapping[str, torch.Tensor], + seq_lengths: Sequence[int], + ) -> None: + real_tokens = sum(seq_lengths) + for name, spec in self._input_specs.items(): + if name not in inputs: + raise KeyError(f"Missing graph input '{name}' for multimodal encoder runner.") + src = inputs[name] + dst = static_inputs[name] + for dim, (d_size, s_size) in enumerate(zip(dst.shape, src.shape, strict=True)): + if dim != spec.token_dim and d_size != s_size: + raise ValueError( + f"Encoder graph input mismatch at {dim=}: static {d_size} vs runtime {s_size}." + ) + src_token_len = src.shape[spec.token_dim] + if src_token_len < real_tokens: + raise ValueError( + f"Encoder graph input has only {src_token_len} tokens on axis {spec.token_dim}, " + f"but {real_tokens} were promised." + ) + dst.narrow(spec.token_dim, 0, real_tokens).copy_( + src.narrow(spec.token_dim, 0, real_tokens) + ) + + def collect_outputs( + self, outputs: Mapping[str, torch.Tensor], seq_lengths: Sequence[int] + ) -> dict[str, torch.Tensor]: + real_tokens = sum(seq_lengths) + return { + name: tensor.narrow(self._output_specs[name], 0, real_tokens) + for name, tensor in outputs.items() + } + + +@dataclass(frozen=True) +class _CapturedGraph(Generic[_MetadataT]): """Bookkeeping for a single captured graph variant.""" graph: torch.cuda.CUDAGraph inputs: Dict[str, torch.Tensor] - metadata: AttentionMetadata + metadata: _MetadataT # Each entry is a tensor view aliased to the captured output's memory via `make_weak_ref`; it # stays valid as long as the captured graph holds the underlying allocation alive. outputs: Dict[str, torch.Tensor] @@ -105,8 +244,8 @@ class _CapturedGraph(NamedTuple): metadata_tensor_ptrs: Mapping[str, int] -class EncoderMetadataProvider(Protocol): - """Builds and refreshes the graph-owned `AttentionMetadata`. +class EncoderMetadataProvider(Protocol[_MetadataT]): + """Builds and refreshes encoder-specific graph-owned metadata. The runner takes a single provider object instead of two separate callbacks so the contract between `build` and `refresh_in_place` lives in one place: @@ -125,18 +264,18 @@ class EncoderMetadataProvider(Protocol): # this list to anchor and verify `data_ptr()` stability across replays. graph_critical_attrs: Sequence[str] - def build(self, key: "EncoderGraphKey") -> "AttentionMetadata": + def build(self, key: EncoderGraphKey) -> _MetadataT: """Return a fresh metadata instance sized to `key`.""" def refresh_in_place( self, - metadata: "AttentionMetadata", + metadata: _MetadataT, padded_seq_lengths: Sequence[int], ) -> None: """Mutate `metadata` in place for the next capture or replay.""" -class MultimodalEncoderGraphRunner: +class MultimodalEncoderGraphRunner(Generic[_MetadataT]): """Captures and replays CUDA graphs for a multimodal encoder block stack. The runner is intentionally narrow: it does not know how to build attention metadata, what @@ -156,29 +295,36 @@ class MultimodalEncoderGraphRunner: object. See `EncoderMetadataProvider` for the contract; the runner enforces the "no reallocation" half of that contract via a `data_ptr()` check after every replay's refresh call. - input_specs: Static shape/dtype descriptors for the graph inputs. + input_specs: Static shape/dtype descriptors for packed graph inputs. Used when + ``io_adapter`` is omitted. output_specs: Token axis (`token_dim`) for each graph output, keyed by output name. Outputs are not materialized by the runner - they are aliased from `encoder_fn`'s return - so only the token axis is needed, to slice each output back to the real token extent. config: Configuration object containing buckets and padding policy. See `MultimodalEncoderCudaGraphConfig`. + io_adapter: Encoder-specific layout, padding, input copying, and output collection. Supply + either this adapter or the packed ``input_specs``/``output_specs`` pair. """ def __init__( self, *, - encoder_fn: Callable[ - [Mapping[str, torch.Tensor], AttentionMetadata], Mapping[str, torch.Tensor] - ], - metadata_provider: EncoderMetadataProvider, - input_specs: Mapping[str, EncoderGraphTensorSpec], - output_specs: Mapping[str, int], + encoder_fn: Callable[[Mapping[str, torch.Tensor], _MetadataT], Mapping[str, torch.Tensor]], + metadata_provider: EncoderMetadataProvider[_MetadataT], config: MultimodalEncoderCudaGraphConfig, + input_specs: Mapping[str, EncoderGraphTensorSpec] | None = None, + output_specs: Mapping[str, int] | None = None, + io_adapter: EncoderGraphIOAdapter | None = None, ) -> None: self._encoder_fn = encoder_fn self._metadata_provider = metadata_provider - self._input_specs: Dict[str, EncoderGraphTensorSpec] = dict(input_specs) - self._output_specs: Dict[str, int] = dict(output_specs) + if io_adapter is None: + if input_specs is None or output_specs is None: + raise ValueError("Provide io_adapter or both input_specs and output_specs.") + io_adapter = PackedEncoderGraphIOAdapter(input_specs, output_specs) + elif input_specs is not None or output_specs is not None: + raise ValueError("io_adapter cannot be combined with input_specs or output_specs.") + self._io_adapter = io_adapter self._config = config # Adopted from the first captured graph so subsequent bucket captures share the pool. self._memory_pool = None @@ -199,7 +345,7 @@ def __init__( # One captured graph per internal key. Populated by `capture_all` and fixed for the runner's # lifetime; `maybe_run` never adds entries. - self._captured: Dict[EncoderGraphKey, _CapturedGraph] = {} + self._captured: dict[EncoderGraphKey, _CapturedGraph[_MetadataT]] = {} self._warned_no_bucket_match: bool = False self._replay_stats_enabled: bool = config.enable_replay_stats @@ -247,7 +393,7 @@ def maybe_run( real_contexts = len(seq_lengths) real_tokens = int(sum(seq_lengths)) - bucket = self._select_bucket(real_contexts, real_tokens) + bucket = self._select_bucket(real_contexts, real_tokens, seq_lengths) if bucket is None: if self._replay_stats_enabled: logger.info( @@ -256,9 +402,7 @@ def maybe_run( ) return None - padded_seq_lengths = self._padded_seq_lengths( - list(seq_lengths), bucket=bucket, real_tokens=real_tokens - ) + padded_seq_lengths = self._padded_seq_lengths(list(seq_lengths), bucket=bucket) key = self._padded_key_for_bucket(bucket) # Every selectable bucket is captured by `capture_all` before the runner is installed, so a @@ -271,7 +415,7 @@ def maybe_run( f"{bucket=} with {key=}, but no graph has been captured." ) return None - self._copy_inputs_into_static(captured, inputs, real_tokens=real_tokens) + self._io_adapter.copy_inputs(captured.inputs, inputs, seq_lengths) self._metadata_provider.refresh_in_place(captured.metadata, padded_seq_lengths) self._assert_metadata_buffers_stable(captured.metadata, captured.metadata_tensor_ptrs) captured.graph.replay() @@ -281,7 +425,7 @@ def maybe_run( f"{bucket=}, {real_contexts=}, {real_tokens=}, " f"pad_slack={bucket.total_tokens - real_tokens}." ) - return self._collect_outputs(captured, real_tokens=real_tokens) + return self._io_adapter.collect_outputs(captured.outputs, seq_lengths) def clear(self, *, release_nccl_window_owners: bool = True) -> None: """Reset captures, releasing their pool owner only at coordinated teardown.""" @@ -297,12 +441,19 @@ def clear(self, *, release_nccl_window_owners: bool = True) -> None: # Bucket selection / padding # ------------------------------------------------------------------ - def _select_bucket(self, real_contexts: int, real_tokens: int) -> Optional[EncoderGraphKey]: + def _select_bucket( + self, + real_contexts: int, + real_tokens: int, + seq_lengths: Sequence[int] | None = None, + ) -> Optional[EncoderGraphKey]: for bucket in self._buckets: if bucket.num_contexts != real_contexts: continue if real_tokens > bucket.total_tokens: continue + if seq_lengths is not None and not self._io_adapter.can_replay(bucket, seq_lengths): + continue if self._config.enable_padding: # Padding path: exact num_contexts match and real_tokens <= bucket budget. return bucket @@ -323,29 +474,14 @@ def _padded_seq_lengths( seq_lengths: List[int], *, bucket: EncoderGraphKey, - real_tokens: int, ) -> List[int]: - if not self._config.enable_padding: - # Exact-fit only: workload tokens must equal the bucket's real-token - # budget. _select_bucket already enforces this. - return seq_lengths - # Append one dummy context that absorbs both the unused real-token slack and the reserved - # padding token. The reservation ensures the dummy is always non-empty even when the - # workload exactly fills the bucket. - padding = (bucket.total_tokens - real_tokens) + _PADDING_TOKEN_RESERVE - return seq_lengths + [padding] + return self._io_adapter.padded_seq_lengths( + seq_lengths, bucket=bucket, enable_padding=self._config.enable_padding + ) def _padded_key_for_bucket(self, bucket: EncoderGraphKey) -> EncoderGraphKey: """Graph key for the captured variant of a bucket.""" - if self._config.enable_padding: - return EncoderGraphKey( - total_tokens=bucket.total_tokens + _PADDING_TOKEN_RESERVE, - num_contexts=bucket.num_contexts + 1, - ) - return EncoderGraphKey( - total_tokens=bucket.total_tokens, - num_contexts=bucket.num_contexts, - ) + return self._io_adapter.graph_key(bucket, enable_padding=self._config.enable_padding) def _dummy_padded_seq_lengths(self, bucket: EncoderGraphKey, key: EncoderGraphKey) -> List[int]: """Representative padded seq_lengths used during capture only. @@ -353,25 +489,7 @@ def _dummy_padded_seq_lengths(self, bucket: EncoderGraphKey, key: EncoderGraphKe Backends may derive a graph-captured work shape from metadata built with these lengths, so the layout must cover every replay shape that can select the same bucket. """ - if key.num_contexts == bucket.num_contexts: - # No padding context. Capture the worst-case skew so plan/run attention backends whose - # launch shape depends on max_token_per_sequence cover every exact-fit replay for this - # bucket. - return [bucket.total_tokens - (bucket.num_contexts - 1)] + [1] * ( - bucket.num_contexts - 1 - ) - # One dummy padding context: give each real context a single token and put the rest in the - # dummy. The captured buffers are sized to `key.total_tokens` - # (== bucket.total_tokens + _PADDING_TOKEN_RESERVE), - # so the dummy absorbs everything not assigned to a real context. - real = [1] * bucket.num_contexts - padding = key.total_tokens - bucket.num_contexts - if padding <= 0: - raise ValueError( - f"Bucket {bucket} cannot host {bucket.num_contexts} real contexts " - f"plus a dummy padding context in {key.total_tokens} tokens." - ) - return real + [padding] + return self._io_adapter.capture_seq_lengths(bucket, key) # ------------------------------------------------------------------ # Capture & replay primitives @@ -386,10 +504,7 @@ def _capture_key( padded_seq_lengths = list(padded_seq_lengths) logger.info(f"Capturing multimodal encoder CUDA graph for key={key} on device={device}.") - static_inputs = { - name: spec.materialize(key.total_tokens, device) - for name, spec in self._input_specs.items() - } + static_inputs = self._io_adapter.allocate_inputs(key, device) metadata = self._metadata_provider.build(key) self._metadata_provider.refresh_in_place(metadata, padded_seq_lengths) @@ -428,9 +543,7 @@ def _capture_key( self._memory_pool = graph_pool @staticmethod - def _snapshot_metadata_tensor_ptrs( - metadata: AttentionMetadata, attrs: Sequence[str] - ) -> Dict[str, int]: + def _snapshot_metadata_tensor_ptrs(metadata: object, attrs: Sequence[str]) -> Dict[str, int]: """Record the `data_ptr()` of each named attribute on `metadata`. Used after capture to anchor the addresses the captured graph reads from. `attrs` comes from @@ -449,9 +562,7 @@ def _snapshot_metadata_tensor_ptrs( return snapshot @staticmethod - def _assert_metadata_buffers_stable( - metadata: AttentionMetadata, expected_ptrs: Mapping[str, int] - ) -> None: + def _assert_metadata_buffers_stable(metadata: object, expected_ptrs: Mapping[str, int]) -> None: """Raise if any captured-time tensor has been rebound or moved. Catches the realistic failure mode where `refresh_in_place` accidentally reallocates a @@ -476,61 +587,6 @@ def _assert_metadata_buffers_stable( "CUDA graph." ) - def _copy_inputs_into_static( - self, - captured: _CapturedGraph, - inputs: Mapping[str, torch.Tensor], - *, - real_tokens: int, - ) -> None: - for name, spec in self._input_specs.items(): - if name not in inputs: - raise KeyError(f"Missing graph input '{name}' for multimodal encoder runner.") - src = inputs[name] - dst = captured.inputs[name] - self._copy_token_axis(dst, src, spec.token_dim, real_tokens) - - def _collect_outputs( - self, captured: _CapturedGraph, *, real_tokens: int - ) -> Dict[str, torch.Tensor]: - result: Dict[str, torch.Tensor] = {} - for name, tensor in captured.outputs.items(): - token_dim = self._output_specs[name] - result[name] = self._slice_token_axis(tensor, token_dim, real_tokens) - return result - - @staticmethod - def _copy_token_axis( - dst: torch.Tensor, - src: torch.Tensor, - token_dim: int, - real_tokens: int, - ) -> None: - # Pre-flight shape checks: every non-token dim must match exactly so the captured graph's - # loads stay in bounds. - for dim, (d_size, s_size) in enumerate(zip(dst.shape, src.shape, strict=True)): - if dim == token_dim: - continue - if d_size != s_size: - raise ValueError( - f"Encoder graph input mismatch at {dim=}: static {d_size} vs runtime {s_size}." - ) - src_token_len = src.shape[token_dim] - if src_token_len < real_tokens: - raise ValueError( - f"Encoder graph input has only {src_token_len} tokens on axis {token_dim}, but " - f"{real_tokens} were promised." - ) - idx = [slice(None)] * dst.ndim - idx[token_dim] = slice(0, real_tokens) - dst[tuple(idx)].copy_(src.narrow(token_dim, 0, real_tokens)) - - @staticmethod - def _slice_token_axis(tensor: torch.Tensor, token_dim: int, real_tokens: int) -> torch.Tensor: - if tensor.shape[token_dim] == real_tokens: - return tensor - return tensor.narrow(token_dim, 0, real_tokens) - def _log_cuda_graph_memory_warning(self) -> None: max_total_tokens = max(bucket.total_tokens for bucket in self._buckets) sum_total_tokens = sum(bucket.total_tokens for bucket in self._buckets) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index 6b54bb978abc..3e5ed82b9834 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -9,7 +9,7 @@ the runner family axis -- it has no warmup or graph-capture phase. """ -from collections.abc import Hashable +from collections.abc import Hashable, Mapping from typing import Any import torch @@ -35,7 +35,11 @@ InputProcessor, get_multimodal_encoder_item_metadata, ) -from tensorrt_llm.llmapi.llm_args import MultimodalEncoderSchedulingPolicy, TorchLlmArgs +from tensorrt_llm.llmapi.llm_args import ( + MultimodalEncoderCudaGraphConfig, + MultimodalEncoderSchedulingPolicy, + TorchLlmArgs, +) from tensorrt_llm.logger import logger from ..llm_request import LlmRequest, MultimodalEncoderRequestError, _Unset, get_mm_items_for_chunk @@ -157,11 +161,32 @@ def mm_encoder_cache_enabled(model: nn.Module) -> bool: return isinstance(model, MultimodalModelMixin) and model.encoder_cache_active -def setup_mm_encoder_cuda_graphs(model: nn.Module) -> None: - """Capture local encoder graphs after metadata setup and before KV sizing.""" - for module in model.modules(): - if isinstance(module, MultimodalEncoderMixin): - module.enable_blocks_cuda_graph() +def setup_mm_encoder_cuda_graphs( + model: nn.Module, + graph_configs: Mapping[str, MultimodalEncoderCudaGraphConfig] | None = None, +) -> None: + """Validate local encoder capabilities, then capture graphs before KV sizing. + + Args: + model: Loaded model whose local encoders have already initialized their metadata. + graph_configs: Model-wide modality configs. Validate against all local towers together; + models without local encoders may delegate encoding to another process. + """ + encoders = [module for module in model.modules() if isinstance(module, MultimodalEncoderMixin)] + if graph_configs is not None and encoders: + supported_modalities = { + module.encoder_cuda_graph_modality + for module in encoders + if module.encoder_cuda_graph_modality is not None + } + unsupported = set(graph_configs) - supported_modalities + if unsupported: + raise ValueError( + f"Unsupported multimodal encoder CUDA graph modalities: {sorted(unsupported)}. " + f"Supported modalities: {sorted(supported_modalities)}." + ) + for module in encoders: + module.enable_blocks_cuda_graph() def setup_mm_encoder_attn_metadata( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 49345bb6df11..3b3e5a30619b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -542,7 +542,10 @@ def __init__( # Capture local encoders after their metadata is sized and before # KV-cache estimation, so graph memory is included in the profile. - setup_mm_encoder_cuda_graphs(self.model) + mm_config = self.llm_args.multimodal_config + setup_mm_encoder_cuda_graphs( + self.model, + mm_config.encoder_cuda_graph if mm_config is not None else None) except Exception as e: import traceback traceback.print_exception(Exception, e, e.__traceback__) diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index cfc38b3b2fed..3d3a33c2f51a 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -112,6 +112,58 @@ def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) - model.disabled.runner.capture_all.assert_not_called() +def test_encoder_graph_configs_are_selected_per_tower_and_validated_together() -> None: + class Encoder(torch.nn.Module, MultimodalEncoderMixin): + def __init__(self, modality: str, mm_config: MultimodalConfig) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.encoder_cuda_graph_modality = modality + self.configure_encoder_cuda_graph(mm_config) + self.runner = Mock() + + def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) -> Mock: + assert config is self._encoder_cuda_graph_config + return self.runner + + vision_config = MultimodalEncoderCudaGraphConfig(buckets=[(256, 1)]) + audio_config = MultimodalEncoderCudaGraphConfig(buckets=[(1024, 1)]) + mm_config = MultimodalConfig( + encoder_cuda_graph={"vision": vision_config, "audio": audio_config} + ) + model = torch.nn.Module() + model.vision = Encoder("vision", mm_config) + model.audio = Encoder("audio", mm_config) + assert model.vision._encoder_cuda_graph_config is vision_config + assert model.audio._encoder_cuda_graph_config is audio_config + + setup_mm_encoder_cuda_graphs(model, mm_config.encoder_cuda_graph) + setup_mm_encoder_cuda_graphs(model, mm_config.encoder_cuda_graph) + model.vision.runner.capture_all.assert_called_once_with(torch.device("cpu")) + model.audio.runner.capture_all.assert_called_once_with(torch.device("cpu")) + + +def test_unsupported_encoder_graph_modality_fails_before_any_capture() -> None: + class VisionEncoder(torch.nn.Module, MultimodalEncoderMixin): + encoder_cuda_graph_modality = "vision" + + def enable_blocks_cuda_graph(self) -> None: + raise AssertionError("Validation must happen before graph capture") + + model = torch.nn.Module() + model.encoder = VisionEncoder() + graph_configs = {"audio": MultimodalEncoderCudaGraphConfig(buckets=[(1024, 1)])} + with pytest.raises(ValueError, match="Unsupported.*audio.*Supported.*vision"): + setup_mm_encoder_cuda_graphs(model, graph_configs) + + +def test_encoder_graph_modality_validation_allows_delegated_encoders() -> None: + model = torch.nn.Module() + model.mm_encoder = None + setup_mm_encoder_cuda_graphs( + model, {"audio": MultimodalEncoderCudaGraphConfig(buckets=[(1024, 1)])} + ) + + def _bind_items(mm_item_scheduler: MultimodalItemScheduler, request, *, row_bytes: int = 8) -> None: state = request.py_mm_encoder_state for item_idx, rows in enumerate(state.embedding_lengths): diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py index fcbf754f598b..f059c5611f24 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py @@ -297,13 +297,15 @@ def test_get_vision_encoder_cuda_graph_config(): assert _get_vision_encoder_cuda_graph_config(mm_config) is config -def test_get_vision_encoder_cuda_graph_config_rejects_unknown_modalities(): +def test_get_vision_encoder_cuda_graph_config_ignores_other_towers(): + vision_config = MultimodalEncoderCudaGraphConfig(buckets=[(1280, 1)]) mm_config = MultimodalConfig( - encoder_cuda_graph={"audio": MultimodalEncoderCudaGraphConfig(buckets=[(1024, 1)])} + encoder_cuda_graph={ + "vision": vision_config, + "audio": MultimodalEncoderCudaGraphConfig(buckets=[(1024, 1)]), + } ) - - with pytest.raises(ValueError, match="Unsupported multimodal encoder CUDA graph modalities"): - _get_vision_encoder_cuda_graph_config(mm_config) + assert _get_vision_encoder_cuda_graph_config(mm_config) is vision_config @pytest.fixture(scope="function") diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py index debb044f0f29..83c1599f2606 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py @@ -11,7 +11,7 @@ import contextlib from dataclasses import dataclass -from typing import Callable, Dict, List, Optional, Sequence +from typing import Callable, Dict, List, Mapping, Optional, Sequence from unittest import mock import pytest @@ -19,6 +19,7 @@ import tensorrt_llm._torch.models.multimodal_encoder_graph as encoder_graph_module from tensorrt_llm._torch.models.multimodal_encoder_graph import ( + EncoderGraphIOAdapter, EncoderGraphKey, EncoderGraphTensorSpec, MultimodalEncoderGraphRunner, @@ -93,6 +94,68 @@ def refresh_in_place(self, metadata, padded_seq_lengths: Sequence[int]) -> None: return None +class _BatchedFramesIOAdapter(EncoderGraphIOAdapter): + """Test layout with a batch/time axis, validity masks, and downsampled outputs.""" + + def allocate_inputs( + self, key: EncoderGraphKey, device: torch.device + ) -> dict[str, torch.Tensor]: + frames = key.total_tokens // key.num_contexts + return { + "features": torch.zeros(key.num_contexts, frames, 2, device=device), + "mask": torch.zeros(key.num_contexts, frames, dtype=torch.bool, device=device), + } + + def copy_inputs( + self, + static_inputs: Mapping[str, torch.Tensor], + inputs: Mapping[str, torch.Tensor], + seq_lengths: Sequence[int], + ) -> None: + for name, dst in static_inputs.items(): + dst.zero_() + src = inputs[name] + dst[:, : src.shape[1]].copy_(src) + + def collect_outputs( + self, outputs: Mapping[str, torch.Tensor], seq_lengths: Sequence[int] + ) -> dict[str, torch.Tensor]: + output_frames = (max(seq_lengths) + 1) // 2 + return {name: tensor[:, :output_frames] for name, tensor in outputs.items()} + + def can_replay(self, bucket: EncoderGraphKey, seq_lengths: Sequence[int]) -> bool: + return max(seq_lengths) <= bucket.total_tokens // bucket.num_contexts + + def graph_key(self, bucket: EncoderGraphKey, *, enable_padding: bool) -> EncoderGraphKey: + return bucket + + def padded_seq_lengths( + self, seq_lengths: Sequence[int], *, bucket: EncoderGraphKey, enable_padding: bool + ) -> list[int]: + return list(seq_lengths) + + def capture_seq_lengths(self, bucket: EncoderGraphKey, key: EncoderGraphKey) -> list[int]: + return [key.total_tokens // key.num_contexts] * key.num_contexts + + +def _batched_frames_encoder( + inputs: Mapping[str, torch.Tensor], metadata: _ToyMetadata +) -> dict[str, torch.Tensor]: + mask = inputs["mask"][:, ::2] + bias = metadata.seq_lens_cuda.sum().to(inputs["features"].dtype) + features = (inputs["features"][:, ::2] + bias) * mask.unsqueeze(-1) + return {"features": features, "mask": mask} + + +def _make_batched_frames_runner(device: torch.device) -> MultimodalEncoderGraphRunner[_ToyMetadata]: + return MultimodalEncoderGraphRunner( + encoder_fn=_batched_frames_encoder, + metadata_provider=_ToyMetadataProvider(device), + io_adapter=_BatchedFramesIOAdapter(), + config=MultimodalEncoderCudaGraphConfig(buckets=[(8, 2), (12, 2)]), + ) + + def _make_toy_encoder_fn(scale: float): """A deterministic, side-effect-free encoder for replay equivalence checks. @@ -201,7 +264,7 @@ def test_select_bucket(enable_padding, real_contexts, real_tokens, expected_buck def test_padded_seq_lengths_appends_dummy_context(): runner = _make_logic_runner() bucket = EncoderGraphKey(num_contexts=1, total_tokens=512) - padded = runner._padded_seq_lengths([300], bucket=bucket, real_tokens=300) + padded = runner._padded_seq_lengths([300], bucket=bucket) # Slack (512 - 300) plus the runner's 1-token padding reservation. assert padded == [300, 213] @@ -211,14 +274,14 @@ def test_padded_seq_lengths_reserves_one_token_on_exact_fit(): # backend rejects empty contexts. runner = _make_logic_runner() bucket = EncoderGraphKey(num_contexts=1, total_tokens=512) - padded = runner._padded_seq_lengths([512], bucket=bucket, real_tokens=512) + padded = runner._padded_seq_lengths([512], bucket=bucket) assert padded == [512, 1] def test_padded_seq_lengths_no_dummy_when_padding_disabled(): runner = _make_logic_runner(enable_padding=False) bucket = EncoderGraphKey(num_contexts=1, total_tokens=512) - padded = runner._padded_seq_lengths([512], bucket=bucket, real_tokens=512) + padded = runner._padded_seq_lengths([512], bucket=bucket) assert padded == [512] @@ -483,6 +546,63 @@ def encoder_fn(inputs, metadata): assert all(inputs is static_inputs for inputs in scoped_inputs) +@pytest.mark.cpu_only +def test_custom_io_adapter_updates_masks_and_collects_downsampled_outputs() -> None: + device = torch.device("cpu") + runner = _make_batched_frames_runner(device) + bucket = EncoderGraphKey(total_tokens=8, num_contexts=2) + key = runner._padded_key_for_bucket(bucket) + assert key == bucket + assert runner._dummy_padded_seq_lengths(bucket, key) == [4, 4] + + static_inputs = runner._io_adapter.allocate_inputs(key, device) + metadata = runner._metadata_provider.build(key) + outputs = _batched_frames_encoder(static_inputs, metadata) + graph = mock.Mock() + + def replay() -> None: + for name, tensor in _batched_frames_encoder(static_inputs, metadata).items(): + outputs[name].copy_(tensor) + + graph.replay.side_effect = replay + runner._captured[key] = _CapturedGraph( + graph=graph, + inputs=static_inputs, + metadata=metadata, + outputs=outputs, + metadata_tensor_ptrs=runner._snapshot_metadata_tensor_ptrs( + metadata, runner._metadata_provider.graph_critical_attrs + ), + ) + input_ptrs = {name: tensor.data_ptr() for name, tensor in static_inputs.items()} + for lengths, frames in [([4, 2], 4), ([1, 2], 2)]: + features = torch.arange(2 * frames * 2, dtype=torch.float32).reshape(2, frames, 2) + mask = torch.arange(frames).unsqueeze(0) < torch.tensor(lengths).unsqueeze(1) + result = runner.maybe_run(seq_lengths=lengths, inputs={"features": features, "mask": mask}) + assert result is not None + expected = torch.zeros(2, (max(lengths) + 1) // 2, 2) + for batch_idx, length in enumerate(lengths): + for frame_idx in range(0, length, 2): + expected[batch_idx, frame_idx // 2] = features[batch_idx, frame_idx] + sum(lengths) + torch.testing.assert_close(result["features"], expected) + torch.testing.assert_close(result["mask"], mask[:, ::2]) + assert {name: tensor.data_ptr() for name, tensor in static_inputs.items()} == input_ptrs + assert metadata.seq_lens_cuda.tolist() == lengths + + assert not static_inputs["mask"][:, 2:].any() + assert not static_inputs["features"][:, 2:].any() + assert graph.replay.call_count == 2 + + +@pytest.mark.cpu_only +def test_custom_layout_skips_bucket_whose_time_axis_is_too_small() -> None: + runner = _make_batched_frames_runner(torch.device("cpu")) + lengths = [5, 1] + assert runner._select_bucket(2, sum(lengths), lengths) == EncoderGraphKey(12, 2) + runner._buckets = [EncoderGraphKey(8, 2)] + assert runner.maybe_run(seq_lengths=lengths, inputs={}) is None + + pytestmark_cuda = pytest.mark.skipif( not torch.cuda.is_available(), reason="MultimodalEncoderGraphRunner requires CUDA", @@ -496,6 +616,26 @@ def cuda_device(): return torch.device("cuda") +@pytestmark_cuda +def test_cuda_custom_io_adapter_replays_masked_downsampled_outputs(cuda_device) -> None: + runner = _make_batched_frames_runner(cuda_device) + runner.capture_all(cuda_device) + for lengths, frames in [([4, 2], 4), ([3, 1], 3), ([1, 2], 2)]: + features = torch.arange(2 * frames * 2, dtype=torch.float32).reshape(2, frames, 2) + mask = torch.arange(frames).unsqueeze(0) < torch.tensor(lengths).unsqueeze(1) + result = runner.maybe_run( + seq_lengths=lengths, + inputs={"features": features.to(cuda_device), "mask": mask.to(cuda_device)}, + ) + assert result is not None + expected = torch.zeros(2, (max(lengths) + 1) // 2, 2) + for batch_idx, length in enumerate(lengths): + for frame_idx in range(0, length, 2): + expected[batch_idx, frame_idx // 2] = features[batch_idx, frame_idx] + sum(lengths) + torch.testing.assert_close(result["features"].cpu(), expected, rtol=0, atol=0) + torch.testing.assert_close(result["mask"].cpu(), mask[:, ::2]) + + @pytest.fixture def make_cuda_runner(cuda_device): """Factory fixture for CUDA end-to-end tests.""" From d935f0506585c414854b6b0a49151e6a363e1f13 Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:42:11 -0700 Subject: [PATCH 4/8] Add automatic multimodal encoder CUDA graph shapes Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- docs/source/models/supported-models.md | 48 ++++++++++++ .../models/modeling_multimodal_encoder.py | 11 ++- tensorrt_llm/_torch/models/modeling_radio.py | 12 ++- .../_torch/models/multimodal_encoder_graph.py | 42 ++++++++-- .../_torch/pyexecutor/engine/multimodal.py | 29 ++++++- .../_torch/pyexecutor/model_engine.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 35 ++++++++- .../_torch/executor/engine/test_multimodal.py | 69 +++++++++++++++++ .../_torch/modeling/test_modeling_qwen3vl.py | 31 ++++++-- .../_torch/modeling/test_modeling_radio.py | 19 +++-- .../modeling/test_multimodal_encoder_graph.py | 77 +++++++++++++++++-- tests/unittest/llmapi/test_llm_args.py | 52 +++++++++++++ 12 files changed, 397 insertions(+), 32 deletions(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 5a3eb3b1e61d..654295a8bd14 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -162,6 +162,54 @@ Note: The following optimizations are available to models that implement `MultimodalModelMixin`. +### Encoder CUDA graphs + +For encoders with CUDA graph support, set a maximum batch size in the model's +modality configuration. For example, pass this YAML through `trtllm-serve --config`: + +```yaml +multimodal_config: + encoder_cuda_graph: + vision: + max_batch_size: 8 +``` + +The encoder adapter automatically generates capture shapes for batch sizes 1 through 8, +within `encoder_max_num_tokens` (or `max_num_tokens` when unset). Use the optional +per-modality `max_num_tokens` to limit capture memory and startup work. Batch size counts +independent encoder contexts: images, video frame groups, or audio chunks as defined by +the encoder. It does not count concurrent API requests or the scheduler's atomic items. +Larger or unsupported shapes fall back to eager execution. Graphs remain disabled by default. + +Each supported tower reads its own modality entry. An audio adapter uses the same settings, +with its own shape, mask, and output rules: + +```yaml +multimodal_config: + encoder_cuda_graph: + vision: + max_batch_size: 8 + audio: + max_batch_size: 4 + max_num_tokens: 2048 +``` + +The second example requires a model with graph-enabled vision and audio towers. The current +production integrations are Qwen3-VL and related Qwen vision encoders, RADIO (`vision`), and +GLM5-Next (`image` instead of `vision`). The shared audio adapter interface is available, +but production audio towers still need an implementation and model validation. Requesting an +unsupported modality fails before any graphs are captured. + +Advanced users can supply `buckets: [[256, 1], [512, 2]]` instead of `max_batch_size` +and `max_num_tokens`. Each pair specifies total real encoder tokens and an exact context count. +Automatic shapes use per-context capacities starting at 256 tokens (or the available budget +when smaller), doubling up to each batch size's token budget. Adapters can override these +shapes for additional alignment constraints. Padded replay is enabled by default; disabling it +requires an exact token-count match. Encoder graphs and encoder side-stream prefetch are +mutually exclusive. + +### Side streams and embeddings cache + | Model Architecture | Multimodal Encoder Side Stream | Multimodal Embeddings Cache | | ------------------ | ------------------------------ | --------------------------- | | `Gemma4ForConditionalGeneration` | Yes | Yes | diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py index de594226305e..33ba8cbbdfc9 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py @@ -181,15 +181,24 @@ def enable_blocks_cuda_graph( config: "MultimodalEncoderCudaGraphConfig | None" = None, *, device: torch.device | None = None, + max_num_tokens: int | None = None, ) -> None: """Capture configured block graphs once, after weights and metadata are ready. The engine drives this lifecycle for every local multimodal encoder. - Standalone callers may provide an explicit config and device. + Standalone callers may provide an explicit config and device. Automatic + configs use their own token cap or the engine's original encoder budget; + raising the scheduling budget for an oversized item does not enlarge captures. """ config = config if config is not None else self._encoder_cuda_graph_config if config is None or self._blocks_graph_runner is not None: return + if config.buckets is None and config.max_num_tokens is None: + if max_num_tokens is None or max_num_tokens <= 0: + raise ValueError("Automatic encoder graphs need max_num_tokens " + "from the config or engine.") + config = config.model_copy( + update={"max_num_tokens": max_num_tokens}) runner = self._build_blocks_graph_runner(config) if device is None: device = next(cast(nn.Module, self).parameters()).device diff --git a/tensorrt_llm/_torch/models/modeling_radio.py b/tensorrt_llm/_torch/models/modeling_radio.py index b069d8374763..d01fc0898d3e 100644 --- a/tensorrt_llm/_torch/models/modeling_radio.py +++ b/tensorrt_llm/_torch/models/modeling_radio.py @@ -771,7 +771,11 @@ def __init__( # order. self._blocks_graph_runner: Optional[MultimodalEncoderGraphRunner] = None - def setup_attn_metadata(self, max_num_tokens: int) -> None: + def setup_attn_metadata( + self, + max_num_tokens: int, + attention_metadata_capacity: Optional[Dict[str, int]] = None, + ) -> None: # Override the default to add `kv_layout="NHD"` for FlashInfer. # FlashInfer's original default is "NHD"; TRT-LLM switched the default # to "HND" for paged KV cache paths (see PR #6917). Ragged prefill @@ -779,8 +783,10 @@ def setup_attn_metadata(self, max_num_tokens: int) -> None: # always in NHD format ([tokens, heads, dim]). # Floor the request capacity at the same legacy fallback as the mixin # default while allowing one attention segment per encoder token. - capacities = self.get_encoder_attention_metadata_capacity( - max_num_tokens) + capacities = ( + attention_metadata_capacity + if attention_metadata_capacity is not None else + self.get_encoder_attention_metadata_capacity(max_num_tokens)) metadata_kwargs = dict( max_num_requests=capacities["attention"], max_num_tokens=max_num_tokens, diff --git a/tensorrt_llm/_torch/models/multimodal_encoder_graph.py b/tensorrt_llm/_torch/models/multimodal_encoder_graph.py index 2ad203e2a760..0873d7122fe7 100644 --- a/tensorrt_llm/_torch/models/multimodal_encoder_graph.py +++ b/tensorrt_llm/_torch/models/multimodal_encoder_graph.py @@ -65,7 +65,7 @@ class EncoderGraphKey(NamedTuple): The fields' meaning depends on how the key is used: - * As a configured bucket, `total_tokens` is the maximum number of real (post-encode) tokens + * As a configured bucket, `total_tokens` is the maximum number of real block-input tokens the bucket is meant to cover, and `num_contexts` is the number of real contexts. Neither includes any padding reservation; the runner adds that internally. * As a captured graph key, both fields describe the adapter's fixed capture capacity. @@ -107,6 +107,30 @@ class EncoderGraphIOAdapter(ABC): assuming that the encoder preserves the input token count. """ + def capture_buckets(self, max_batch_size: int, max_num_tokens: int) -> list[EncoderGraphKey]: + """Generate shapes for every batch size within the total token budget. + + Each context count uses per-context capacities starting at 256 tokens and doubling up + to its share of the budget, followed by the full token budget. Dense ``[batch, time, ...]`` + layouts can round down their time dimension when allocating buffers and check it in + ``can_replay``. Adapters with additional alignment or minimum-length constraints can + override this. + + Args: + max_batch_size: Maximum number of independent encoder contexts. + max_num_tokens: Maximum total tokens, including layout padding but excluding any + adapter-specific reservation added by ``graph_key``. + """ + buckets = [] + for num_contexts in range(1, min(max_batch_size, max_num_tokens) + 1): + max_seq_len = max_num_tokens // num_contexts + seq_len = min(256, max_seq_len) + while seq_len < max_seq_len: + buckets.append(EncoderGraphKey(seq_len * num_contexts, num_contexts)) + seq_len *= 2 + buckets.append(EncoderGraphKey(max_num_tokens, num_contexts)) + return sorted(buckets) + @abstractmethod def allocate_inputs( self, key: EncoderGraphKey, device: torch.device @@ -329,17 +353,25 @@ def __init__( # Adopted from the first captured graph so subsequent bucket captures share the pool. self._memory_pool = None - if not config.buckets: + buckets = config.buckets + if buckets is None: + if config.max_batch_size is None or config.max_num_tokens is None: + raise ValueError( + "Automatic encoder graphs require max_batch_size and a resolved " + "max_num_tokens budget." + ) + buckets = io_adapter.capture_buckets(config.max_batch_size, config.max_num_tokens) + if not buckets: raise ValueError( - "MultimodalEncoderCudaGraphConfig.buckets is empty; the runner has nothing to " - "capture." + "MultimodalEncoderCudaGraphConfig.buckets is empty; the runner has " + "nothing to capture." ) # Field order is (total_tokens, num_contexts), matching the desired bucket ordering for # smallest-fit selection in `maybe_run()`. self._buckets: List[EncoderGraphKey] = sorted( EncoderGraphKey(total_tokens=total_tokens, num_contexts=num_contexts) - for total_tokens, num_contexts in config.buckets + for total_tokens, num_contexts in buckets ) self._log_cuda_graph_memory_warning() diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index 3e5ed82b9834..71758517a71f 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -164,6 +164,8 @@ def mm_encoder_cache_enabled(model: nn.Module) -> bool: def setup_mm_encoder_cuda_graphs( model: nn.Module, graph_configs: Mapping[str, MultimodalEncoderCudaGraphConfig] | None = None, + *, + max_num_tokens: int | None = None, ) -> None: """Validate local encoder capabilities, then capture graphs before KV sizing. @@ -171,6 +173,8 @@ def setup_mm_encoder_cuda_graphs( model: Loaded model whose local encoders have already initialized their metadata. graph_configs: Model-wide modality configs. Validate against all local towers together; models without local encoders may delegate encoding to another process. + max_num_tokens: Encoder token budget before any scheduling-only increase for large items. + Supplies the default capture capacity for automatic configs in each tower. """ encoders = [module for module in model.modules() if isinstance(module, MultimodalEncoderMixin)] if graph_configs is not None and encoders: @@ -186,7 +190,7 @@ def setup_mm_encoder_cuda_graphs( f"Supported modalities: {sorted(supported_modalities)}." ) for module in encoders: - module.enable_blocks_cuda_graph() + module.enable_blocks_cuda_graph(max_num_tokens=max_num_tokens) def setup_mm_encoder_attn_metadata( @@ -220,10 +224,31 @@ def setup_mm_encoder_attn_metadata( for module in encoders: setup_kwargs: dict[str, Any] = dict(max_num_tokens=encoder_max_num_tokens) + encoder_max_seq_len = max_seq_len if attention_metadata_capacity is not None: setup_kwargs["attention_metadata_capacity"] = attention_metadata_capacity + graph_config = module._encoder_cuda_graph_config + if graph_config is not None and graph_config.buckets is None: + # The default packed adapter isolates padding in one extra context/token. + # Reserve this small margin for automatic configs before any graph is captured. + reserve = int(graph_config.enable_padding) + graph_tokens = graph_config.max_num_tokens or encoder_max_num_tokens + max_batch_size = graph_config.max_batch_size + assert max_batch_size is not None + metadata_tokens = max(encoder_max_num_tokens, graph_tokens) + reserve + encoder_max_seq_len = max(max_seq_len, graph_tokens) + capacities = ( + attention_metadata_capacity + if attention_metadata_capacity is not None + else module.get_encoder_attention_metadata_capacity(metadata_tokens) + ) + setup_kwargs["max_num_tokens"] = metadata_tokens + setup_kwargs["attention_metadata_capacity"] = { + name: max(capacity, max_batch_size + reserve) + for name, capacity in capacities.items() + } module.setup_attn_metadata(**setup_kwargs) - module.set_attn_max_seq_len(max_seq_len) + module.set_attn_max_seq_len(encoder_max_seq_len) def resolve_bytes_per_mm_encoder_embedding(model: MultimodalModelMixin) -> int: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 3b3e5a30619b..9cc1e49d69b1 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -545,7 +545,9 @@ def __init__( mm_config = self.llm_args.multimodal_config setup_mm_encoder_cuda_graphs( self.model, - mm_config.encoder_cuda_graph if mm_config is not None else None) + mm_config.encoder_cuda_graph if mm_config is not None else None, + max_num_tokens=(self.llm_args.encoder_max_num_tokens + or self.max_num_tokens)) except Exception as e: import traceback traceback.print_exception(Exception, e, e.__traceback__) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index c6caaf0c274b..355c77c0b319 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -680,14 +680,36 @@ def validate_token_encoder_bucket_config( class MultimodalEncoderCudaGraphConfig(StrictBaseModel): """CUDA graph capture for multimodal vision / audio encoders. - One graph is captured per `buckets` entry; replay falls back to eager when the live workload's - padded shape does not match any bucket. + Set `max_batch_size` to generate capture shapes automatically. Advanced users can instead + specify `buckets`. Replay falls back to eager when no captured shape fits the live workload. NOTE: enabling this will lead to higher GPU memory usage. It is usually more beneficial for `buckets` corresponding to lesser compute. """ - buckets: list[tuple[PositiveInt, PositiveInt]] = Field( + max_batch_size: Optional[PositiveInt] = Field( + default=None, + description= + ("Maximum number of encoder contexts per CUDA graph. Automatically captures every " + "context count from 1 to this limit, with token capacities derived from the encoder " + "token budget. Contexts are encoder sequences, such as image or video frame groups " + "and audio chunks; they are not concurrent requests. Mutually exclusive with buckets." + ), + status="prototype", + ) + + max_num_tokens: Optional[PositiveInt] = Field( + default=None, + description= + ("Optional total token capacity for automatically generated graphs. Defaults to " + "encoder_max_num_tokens, or max_num_tokens when the encoder limit is unset. The " + "encoder adapter defines token units and layout. Larger workloads run eagerly. " + "Mutually exclusive with buckets."), + status="prototype", + ) + + buckets: Optional[list[tuple[PositiveInt, PositiveInt]]] = Field( + default=None, min_length=1, description= ("Explicit encoder graph buckets as `(total_tokens, num_contexts)` pairs. Each pair " @@ -723,6 +745,13 @@ class MultimodalEncoderCudaGraphConfig(StrictBaseModel): @model_validator(mode='after') def _normalize(self) -> 'MultimodalEncoderCudaGraphConfig': + if self.buckets is None: + if self.max_batch_size is None: + raise ValueError("Set max_batch_size or explicit buckets.") + return self + if self.max_batch_size is not None or self.max_num_tokens is not None: + raise ValueError("Explicit buckets cannot be combined with " + "max_batch_size or max_num_tokens.") for total_tokens, num_contexts in self.buckets: if total_tokens < num_contexts: raise ValueError("`buckets` entries must have total_tokens >= " diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index 3d3a33c2f51a..000ee78bc7ae 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -156,6 +156,75 @@ def enable_blocks_cuda_graph(self) -> None: setup_mm_encoder_cuda_graphs(model, graph_configs) +def test_automatic_encoder_graph_budgets_are_resolved_per_tower_without_mutating_config() -> None: + class Encoder(torch.nn.Module, MultimodalEncoderMixin): + def __init__(self, modality: str, mm_config: MultimodalConfig) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.encoder_cuda_graph_modality = modality + self.configure_encoder_cuda_graph(mm_config) + self.build_runner = Mock(return_value=Mock()) + + def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) -> Mock: + return self.build_runner(config) + + mm_config = MultimodalConfig( + encoder_cuda_graph={ + "vision": {"max_batch_size": 8}, + "audio": {"max_batch_size": 4, "max_num_tokens": 512}, + } + ) + model = torch.nn.Module() + model.vision = Encoder("vision", mm_config) + model.audio = Encoder("audio", mm_config) + setup_mm_encoder_cuda_graphs(model, mm_config.encoder_cuda_graph, max_num_tokens=2048) + setup_mm_encoder_cuda_graphs(model, mm_config.encoder_cuda_graph, max_num_tokens=2048) + + vision_config = model.vision.build_runner.call_args.args[0] + audio_config = model.audio.build_runner.call_args.args[0] + assert vision_config.max_batch_size == 8 + assert vision_config.max_num_tokens == 2048 + assert audio_config.max_batch_size == 4 + assert audio_config.max_num_tokens == 512 + assert mm_config.encoder_cuda_graph["vision"].max_num_tokens is None + assert mm_config.encoder_cuda_graph["audio"].max_num_tokens == 512 + model.vision.build_runner.assert_called_once() + model.audio.build_runner.assert_called_once() + model.vision._blocks_graph_runner.capture_all.assert_called_once_with(torch.device("cpu")) + model.audio._blocks_graph_runner.capture_all.assert_called_once_with(torch.device("cpu")) + + +def test_automatic_encoder_graphs_need_a_token_budget_before_capture() -> None: + encoder = MultimodalEncoderMixin() + encoder._encoder_cuda_graph_config = MultimodalEncoderCudaGraphConfig(max_batch_size=8) + with pytest.raises(ValueError, match="from the config or engine"): + encoder.enable_blocks_cuda_graph() + + +@pytest.mark.parametrize("enable_padding", [False, True]) +def test_automatic_encoder_graph_metadata_reserves_padding_capacity(enable_padding) -> None: + class Encoder(torch.nn.Module, MultimodalEncoderMixin): + def __init__(self) -> None: + super().__init__() + self._encoder_cuda_graph_config = MultimodalEncoderCudaGraphConfig( + max_batch_size=8, max_num_tokens=64, enable_padding=enable_padding + ) + self.setup_attn_metadata = Mock() + self.set_attn_max_seq_len = Mock() + + model = torch.nn.Module() + model.encoder = Encoder() + capacities = {"attention": 4} + setup_mm_encoder_attn_metadata(model, object(), 32, capacities) + reserve = int(enable_padding) + model.encoder.setup_attn_metadata.assert_called_once_with( + max_num_tokens=64 + reserve, + attention_metadata_capacity={"attention": 8 + reserve}, + ) + model.encoder.set_attn_max_seq_len.assert_called_once_with(64) + assert capacities == {"attention": 4} + + def test_encoder_graph_modality_validation_allows_delegated_encoders() -> None: model = torch.nn.Module() model.mm_encoder = None diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index 0039586220ed..716f1130906f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -652,16 +652,31 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("attn_backend", ["TRTLLM", "FLASHINFER"]) @pytest.mark.parametrize("deepstack_visual_indexes", [[0], [1, 0]]) +@pytest.mark.parametrize("automatic", [False, True], ids=["explicit", "automatic"]) def test_qwen3_vision_cuda_graph_replay_matches_eager( - monkeypatch: pytest.MonkeyPatch, attn_backend: str, deepstack_visual_indexes: list[int] + monkeypatch: pytest.MonkeyPatch, + attn_backend: str, + deepstack_visual_indexes: list[int], + automatic: bool, ) -> None: + from tensorrt_llm._torch.pyexecutor.engine.multimodal import ( + setup_mm_encoder_attn_metadata, + setup_mm_encoder_cuda_graphs, + ) + config_dict = copy.deepcopy(QWEN3_VL_8B_CONFIG) config_dict["vision_config"].update(depth=2, deepstack_visual_indexes=deepstack_visual_indexes) model_config = ModelConfig( pretrained_config=Qwen3VLConfig.from_dict(config_dict), attn_backend=attn_backend, multimodal_config=MultimodalConfig( - encoder_cuda_graph={"vision": MultimodalEncoderCudaGraphConfig(buckets=[(256, 1)])} + encoder_cuda_graph={ + "vision": ( + MultimodalEncoderCudaGraphConfig(max_batch_size=2) + if automatic + else MultimodalEncoderCudaGraphConfig(buckets=[(256, 1)]) + ) + } ), ) torch.manual_seed(0) @@ -673,13 +688,19 @@ def test_qwen3_vision_cuda_graph_replay_matches_eager( else: param.normal_(std=0.02) vision = encoder.visual - vision.setup_attn_metadata(max_num_tokens=1024) + if automatic: + setup_mm_encoder_attn_metadata(encoder, object(), 256, None) + else: + vision.setup_attn_metadata(max_num_tokens=1024) # 192 real tokens replay the (256, 1) bucket with a 65-token padding context. grid_thw = torch.tensor([[1, 12, 16]]) pixel_values = torch.randn(192, 3 * 2 * 16 * 16, device="cuda") eager = encoder.encode_batched(pixel_values, grid_thw) - vision.enable_blocks_cuda_graph() + if automatic: + setup_mm_encoder_cuda_graphs(encoder, max_num_tokens=256) + else: + vision.enable_blocks_cuda_graph() def fail_eager_fallback(*args: object, **kwargs: object) -> None: raise AssertionError("Qwen3-VL vision CUDA graph replay fell back to eager.") @@ -688,7 +709,7 @@ def fail_eager_fallback(*args: object, **kwargs: object) -> None: graph = encoder.encode_batched(pixel_values, grid_thw) torch.testing.assert_close(graph, eager, rtol=5e-2, atol=5e-2) - (captured,) = vision._blocks_graph_runner._captured.values() + captured = vision._blocks_graph_runner._captured[(257, 2)] # Replay must refresh the device lengths read by attention, not only the host copy. assert captured.metadata.seq_lens_cuda.tolist() == [192, 65] diff --git a/tests/unittest/_torch/modeling/test_modeling_radio.py b/tests/unittest/_torch/modeling/test_modeling_radio.py index 2c303f8ebc90..ac660f6155f8 100644 --- a/tests/unittest/_torch/modeling/test_modeling_radio.py +++ b/tests/unittest/_torch/modeling/test_modeling_radio.py @@ -171,7 +171,8 @@ def test_clear_blocks_cuda_graph_releases_only_at_explicit_boundary(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_radio_blocks_cuda_graph_matches_eager(tiny_vit_config): +@pytest.mark.parametrize("automatic", [False, True], ids=["explicit", "automatic"]) +def test_radio_blocks_cuda_graph_matches_eager(tiny_vit_config, automatic): """Block-loop CUDA graph wiring must produce the same output as eager. Runs the tiny ViT once in eager mode to capture a reference, then enables the encoder graph @@ -188,10 +189,14 @@ def test_radio_blocks_cuda_graph_matches_eager(tiny_vit_config): RADIOVisionModel( _make_bf16_model_config(), disable_quantization=True, - encoder_cuda_graph_config=MultimodalEncoderCudaGraphConfig( - buckets=[(16, 1)], - enable_padding=True, - warmup_steps=2, + encoder_cuda_graph_config=( + MultimodalEncoderCudaGraphConfig(max_batch_size=2, max_num_tokens=16) + if automatic + else MultimodalEncoderCudaGraphConfig( + buckets=[(16, 1)], + enable_padding=True, + warmup_steps=2, + ) ), ) .to(device) @@ -200,6 +205,10 @@ def test_radio_blocks_cuda_graph_matches_eager(tiny_vit_config): model.eval() _init_finite_weights(model) assert model.model_config.attn_backend == "FLASHINFER" + if automatic: + from tensorrt_llm._torch.pyexecutor.engine.multimodal import setup_mm_encoder_attn_metadata + + setup_mm_encoder_attn_metadata(model, object(), 16, None) with torch.no_grad(): eager_out = model.forward(pixel_values).clone() diff --git a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py index 83c1599f2606..5e724619286c 100644 --- a/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py +++ b/tests/unittest/_torch/modeling/test_multimodal_encoder_graph.py @@ -147,12 +147,18 @@ def _batched_frames_encoder( return {"features": features, "mask": mask} -def _make_batched_frames_runner(device: torch.device) -> MultimodalEncoderGraphRunner[_ToyMetadata]: +def _make_batched_frames_runner( + device: torch.device, *, automatic: bool = False +) -> MultimodalEncoderGraphRunner[_ToyMetadata]: return MultimodalEncoderGraphRunner( encoder_fn=_batched_frames_encoder, metadata_provider=_ToyMetadataProvider(device), io_adapter=_BatchedFramesIOAdapter(), - config=MultimodalEncoderCudaGraphConfig(buckets=[(8, 2), (12, 2)]), + config=( + MultimodalEncoderCudaGraphConfig(max_batch_size=3, max_num_tokens=12) + if automatic + else MultimodalEncoderCudaGraphConfig(buckets=[(8, 2), (12, 2)]) + ), ) @@ -363,6 +369,58 @@ def test_empty_buckets_rejected(): ) +@pytest.mark.cpu_only +def test_automatic_buckets_cover_every_batch_with_bounded_layout_shapes() -> None: + runner = _make_batched_frames_runner(torch.device("cpu"), automatic=True) + assert runner._buckets == [EncoderGraphKey(12, batch) for batch in (1, 2, 3)] + for lengths in ([7], [5, 1], [4, 2, 1]): + bucket = runner._select_bucket(len(lengths), sum(lengths), lengths) + assert bucket is not None + assert bucket.total_tokens <= 12 + assert bucket.total_tokens % bucket.num_contexts == 0 + assert runner._select_bucket(4, 4, [1] * 4) is None + assert runner._select_bucket(2, 8, [7, 1]) is None + + +@pytest.mark.cpu_only +def test_automatic_buckets_include_intermediate_token_capacities() -> None: + buckets = _BatchedFramesIOAdapter().capture_buckets(3, 4096) + assert [bucket.total_tokens for bucket in buckets if bucket.num_contexts == 1] == [ + 256, + 512, + 1024, + 2048, + 4096, + ] + assert EncoderGraphKey(4096, 3) in buckets + assert all(bucket.total_tokens <= 4096 for bucket in buckets) + inputs = _BatchedFramesIOAdapter().allocate_inputs( + EncoderGraphKey(4096, 3), torch.device("cpu") + ) + assert inputs["features"].shape == (3, 1365, 2) + assert _BatchedFramesIOAdapter().capture_buckets(8, 2) == [ + EncoderGraphKey(2, 1), + EncoderGraphKey(2, 2), + ] + + +@pytest.mark.cpu_only +def test_automatic_buckets_use_encoder_adapter_shape_rules() -> None: + class AlignedAdapter(_BatchedFramesIOAdapter): + def capture_buckets( + self, max_batch_size: int, max_num_tokens: int + ) -> list[EncoderGraphKey]: + return [EncoderGraphKey(8, 2)] + + runner = MultimodalEncoderGraphRunner( + encoder_fn=_batched_frames_encoder, + metadata_provider=_ToyMetadataProvider(torch.device("cpu")), + io_adapter=AlignedAdapter(), + config=MultimodalEncoderCudaGraphConfig(max_batch_size=2, max_num_tokens=12), + ) + assert runner._buckets == [EncoderGraphKey(8, 2)] + + @pytest.mark.cpu_only def test_capture_all_processes_largest_bucket_first() -> None: runner = _make_logic_runner( @@ -617,18 +675,23 @@ def cuda_device(): @pytestmark_cuda -def test_cuda_custom_io_adapter_replays_masked_downsampled_outputs(cuda_device) -> None: - runner = _make_batched_frames_runner(cuda_device) +@pytest.mark.parametrize("automatic", [False, True], ids=["explicit", "automatic"]) +def test_cuda_custom_io_adapter_replays_masked_downsampled_outputs(cuda_device, automatic) -> None: + runner = _make_batched_frames_runner(cuda_device, automatic=automatic) runner.capture_all(cuda_device) - for lengths, frames in [([4, 2], 4), ([3, 1], 3), ([1, 2], 2)]: - features = torch.arange(2 * frames * 2, dtype=torch.float32).reshape(2, frames, 2) + cases = [([4, 2], 4), ([3, 1], 3), ([1, 2], 2)] + if automatic: + cases += [([7], 7), ([4, 2, 1], 4), ([1, 1, 2], 2)] + for lengths, frames in cases: + batch = len(lengths) + features = torch.arange(batch * frames * 2, dtype=torch.float32).reshape(batch, frames, 2) mask = torch.arange(frames).unsqueeze(0) < torch.tensor(lengths).unsqueeze(1) result = runner.maybe_run( seq_lengths=lengths, inputs={"features": features.to(cuda_device), "mask": mask.to(cuda_device)}, ) assert result is not None - expected = torch.zeros(2, (max(lengths) + 1) // 2, 2) + expected = torch.zeros(batch, (max(lengths) + 1) // 2, 2) for batch_idx, length in enumerate(lengths): for frame_idx in range(0, length, 2): expected[batch_idx, frame_idx // 2] = features[batch_idx, frame_idx] + sum(lengths) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index fbbafa3e7d68..cf1869e4d68e 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1717,6 +1717,58 @@ def test_KvCacheConfig_disk_cache_validation(tmp_path): @pytest.mark.cpu_only class TestMultimodalEncoderCudaGraphConfig: + def test_automatic_buckets_need_only_max_batch_size(self): + config = MultimodalEncoderCudaGraphConfig(max_batch_size=8) + assert config.max_batch_size == 8 + assert config.max_num_tokens is None + assert config.buckets is None + + def test_automatic_buckets_accept_token_cap(self): + config = MultimodalEncoderCudaGraphConfig(max_batch_size=8, + max_num_tokens=1024) + assert config.max_num_tokens == 1024 + + @pytest.mark.parametrize("kwargs", [ + { + "max_batch_size": 0 + }, + { + "max_batch_size": -1 + }, + { + "max_batch_size": 8, + "max_num_tokens": 0 + }, + { + "max_num_tokens": 1024 + }, + { + "buckets": [(256, 1)], + "max_batch_size": 8 + }, + { + "buckets": [(256, 1)], + "max_num_tokens": 1024 + }, + ]) + def test_rejects_invalid_automatic_config(self, kwargs): + with pytest.raises(ValidationError): + MultimodalEncoderCudaGraphConfig(**kwargs) + + def test_both_modalities_accept_automatic_config(self): + config = MultimodalConfig( + encoder_cuda_graph={ + "vision": { + "max_batch_size": 8 + }, + "audio": { + "max_batch_size": 4, + "max_num_tokens": 2048 + }, + }) + assert config.encoder_cuda_graph["vision"].max_batch_size == 8 + assert config.encoder_cuda_graph["audio"].max_num_tokens == 2048 + def test_minimal_required_fields(self): config = MultimodalEncoderCudaGraphConfig(buckets=[(1035, 1)]) assert config.buckets == [(1035, 1)] From f737f568d7355a89ed4b3729bc84b3d68927e9a8 Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:11:31 -0700 Subject: [PATCH 5/8] [None][fix] size encoder graph metadata from capture keys Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/models/modeling_qwen3vl.py | 25 +++++-------------- .../_torch/pyexecutor/engine/multimodal.py | 24 +++++------------- .../_torch/executor/engine/test_multimodal.py | 18 ++++++++----- .../_torch/modeling/test_modeling_qwen3vl.py | 18 +++++++------ 4 files changed, 34 insertions(+), 51 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index c4a370588018..82c35ee4534e 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -1267,22 +1267,13 @@ def __init__(self, vision: Qwen3VisionModel) -> None: def build(self, key: EncoderGraphKey) -> AttentionMetadata: vision = self._vision - assert vision.attn_metadata is not None - if ( - key.num_contexts > vision.attn_metadata.max_num_requests - or key.total_tokens > vision.attn_metadata.max_num_tokens - ): - raise ValueError( - f"Qwen3-VL encoder CUDA graph needs {key.total_tokens} tokens and " - f"{key.num_contexts} contexts, which exceeds the encoder attention metadata " - f"capacity (max_num_tokens={vision.attn_metadata.max_num_tokens}, " - f"max_num_requests={vision.attn_metadata.max_num_requests}). With " - "enable_padding=True, every bucket reserves one extra token and one padding " - "context. Raise encoder_max_num_tokens or use a smaller bucket." - ) - metadata = vision.attn_metadata.create_cuda_graph_metadata( - key.num_contexts, encode_only=True + metadata = vision.metadata_cls( + max_num_requests=key.num_contexts, + max_num_tokens=key.total_tokens, + kv_cache_manager=None, ) + metadata.is_cuda_graph = True + metadata.seq_lens = torch.ones(key.num_contexts, dtype=torch.int32) bind_seq_lens = getattr(metadata, "bind_encoder_cuda_graph_seq_lens", None) if callable(bind_seq_lens): bind_seq_lens(metadata.seq_lens, key.num_contexts) @@ -1342,10 +1333,6 @@ def __init__( def post_config(self): self.config = self.model_config.pretrained_config.vision_config - def enable_cuda_graph(self) -> None: - """Capture the vision tower's opt-in encoder CUDA graphs.""" - self.visual.enable_blocks_cuda_graph() - def load_weights( self, weights: Dict[str, torch.Tensor], diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index 71758517a71f..2a894a9f2172 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -228,25 +228,13 @@ def setup_mm_encoder_attn_metadata( if attention_metadata_capacity is not None: setup_kwargs["attention_metadata_capacity"] = attention_metadata_capacity graph_config = module._encoder_cuda_graph_config - if graph_config is not None and graph_config.buckets is None: - # The default packed adapter isolates padding in one extra context/token. - # Reserve this small margin for automatic configs before any graph is captured. - reserve = int(graph_config.enable_padding) - graph_tokens = graph_config.max_num_tokens or encoder_max_num_tokens - max_batch_size = graph_config.max_batch_size - assert max_batch_size is not None - metadata_tokens = max(encoder_max_num_tokens, graph_tokens) + reserve - encoder_max_seq_len = max(max_seq_len, graph_tokens) - capacities = ( - attention_metadata_capacity - if attention_metadata_capacity is not None - else module.get_encoder_attention_metadata_capacity(metadata_tokens) + if graph_config is not None: + graph_tokens = ( + max(tokens for tokens, _ in graph_config.buckets) + if graph_config.buckets is not None + else graph_config.max_num_tokens or encoder_max_num_tokens ) - setup_kwargs["max_num_tokens"] = metadata_tokens - setup_kwargs["attention_metadata_capacity"] = { - name: max(capacity, max_batch_size + reserve) - for name, capacity in capacities.items() - } + encoder_max_seq_len = max(max_seq_len, graph_tokens) module.setup_attn_metadata(**setup_kwargs) module.set_attn_max_seq_len(encoder_max_seq_len) diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index 000ee78bc7ae..f65b34095b43 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -202,12 +202,19 @@ def test_automatic_encoder_graphs_need_a_token_budget_before_capture() -> None: @pytest.mark.parametrize("enable_padding", [False, True]) -def test_automatic_encoder_graph_metadata_reserves_padding_capacity(enable_padding) -> None: +@pytest.mark.parametrize("automatic", [False, True], ids=["explicit", "automatic"]) +def test_encoder_graphs_preserve_eager_metadata_capacity(enable_padding, automatic) -> None: class Encoder(torch.nn.Module, MultimodalEncoderMixin): def __init__(self) -> None: super().__init__() - self._encoder_cuda_graph_config = MultimodalEncoderCudaGraphConfig( - max_batch_size=8, max_num_tokens=64, enable_padding=enable_padding + self._encoder_cuda_graph_config = ( + MultimodalEncoderCudaGraphConfig( + max_batch_size=8, max_num_tokens=64, enable_padding=enable_padding + ) + if automatic + else MultimodalEncoderCudaGraphConfig( + buckets=[(64, 8)], enable_padding=enable_padding + ) ) self.setup_attn_metadata = Mock() self.set_attn_max_seq_len = Mock() @@ -216,10 +223,9 @@ def __init__(self) -> None: model.encoder = Encoder() capacities = {"attention": 4} setup_mm_encoder_attn_metadata(model, object(), 32, capacities) - reserve = int(enable_padding) model.encoder.setup_attn_metadata.assert_called_once_with( - max_num_tokens=64 + reserve, - attention_metadata_capacity={"attention": 8 + reserve}, + max_num_tokens=32, + attention_metadata_capacity={"attention": 4}, ) model.encoder.set_attn_max_seq_len.assert_called_once_with(64) assert capacities == {"attention": 4} diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index 716f1130906f..c430a47081d8 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -653,11 +653,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @pytest.mark.parametrize("attn_backend", ["TRTLLM", "FLASHINFER"]) @pytest.mark.parametrize("deepstack_visual_indexes", [[0], [1, 0]]) @pytest.mark.parametrize("automatic", [False, True], ids=["explicit", "automatic"]) +@pytest.mark.parametrize("num_tokens", [192, 256]) def test_qwen3_vision_cuda_graph_replay_matches_eager( monkeypatch: pytest.MonkeyPatch, attn_backend: str, deepstack_visual_indexes: list[int], automatic: bool, + num_tokens: int, ) -> None: from tensorrt_llm._torch.pyexecutor.engine.multimodal import ( setup_mm_encoder_attn_metadata, @@ -688,13 +690,9 @@ def test_qwen3_vision_cuda_graph_replay_matches_eager( else: param.normal_(std=0.02) vision = encoder.visual - if automatic: - setup_mm_encoder_attn_metadata(encoder, object(), 256, None) - else: - vision.setup_attn_metadata(max_num_tokens=1024) - # 192 real tokens replay the (256, 1) bucket with a 65-token padding context. - grid_thw = torch.tensor([[1, 12, 16]]) - pixel_values = torch.randn(192, 3 * 2 * 16 * 16, device="cuda") + setup_mm_encoder_attn_metadata(encoder, object(), 256, {"attention": 1}) + grid_thw = torch.tensor([[1, num_tokens // 16, 16]]) + pixel_values = torch.randn(num_tokens, 3 * 2 * 16 * 16, device="cuda") eager = encoder.encode_batched(pixel_values, grid_thw) if automatic: @@ -711,7 +709,11 @@ def fail_eager_fallback(*args: object, **kwargs: object) -> None: torch.testing.assert_close(graph, eager, rtol=5e-2, atol=5e-2) captured = vision._blocks_graph_runner._captured[(257, 2)] # Replay must refresh the device lengths read by attention, not only the host copy. - assert captured.metadata.seq_lens_cuda.tolist() == [192, 65] + assert captured.metadata.seq_lens_cuda.tolist() == [num_tokens, 257 - num_tokens] + assert vision.attn_metadata.max_num_tokens == 256 + assert vision.attn_metadata.max_num_requests == 1 + assert captured.metadata.max_num_tokens == 257 + assert captured.metadata.max_num_requests == 2 def test_qwen3_processor_max_pixels_maps_to_fixed_attention_capacity() -> None: From 363773f0497657d0c0e0643281d45f2a7cb54223 Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:41:53 -0700 Subject: [PATCH 6/8] [None][fix] preserve multimodal encoder graph setup contracts Expose encoder graph settings through the mixin property and remove the redundant metadata token-budget fallback. Keep the originally configured encoder budget through scheduler expansion and LLM token clamping, and access declared attention metadata fields directly. Cover explicit and default encoder budgets with item scheduling enabled and disabled in the real engine constructor. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- .../models/modeling_multimodal_encoder.py | 16 ++- tensorrt_llm/_torch/models/modeling_radio.py | 2 +- .../_torch/pyexecutor/engine/multimodal.py | 7 +- .../_torch/pyexecutor/model_engine.py | 10 +- .../_torch/executor/engine/test_multimodal.py | 128 +++++++++++++++++- 5 files changed, 144 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py index 33ba8cbbdfc9..80101b864658 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py @@ -94,7 +94,7 @@ def prepare_encoder_attn_metadata(seq_lens: list[int], attn_metadata.request_ids = list(range(1, num_segments + 1)) prepare_graph_replay = getattr(attn_metadata, "prepare_encoder_cuda_graph_replay", None) - is_cuda_graph = getattr(attn_metadata, "is_cuda_graph", False) + is_cuda_graph = attn_metadata.is_cuda_graph if is_cuda_graph and callable(prepare_graph_replay): prepare_graph_replay(seq_lens, sum(seq_lens)) # The replay helper refreshes only the host lengths, but the captured @@ -105,8 +105,7 @@ def prepare_encoder_attn_metadata(seq_lens: list[int], attn_metadata.seq_lens = seq_lens_torch cu_seqlens = cu_seqlens.to(device=attn_metadata.seq_lens_cuda.device, non_blocking=True) - if (is_cuda_graph - and getattr(attn_metadata, "cu_q_seqlens", None) is not None + if (is_cuda_graph and attn_metadata.cu_q_seqlens is not None and attn_metadata.cu_q_seqlens.shape == cu_seqlens.shape): # CUDA graph capture bakes these addresses into attention kernels. # Refresh their contents without replacing the captured buffers. @@ -147,6 +146,17 @@ class MultimodalEncoderMixin: _encoder_cuda_graph_config: "MultimodalEncoderCudaGraphConfig | None" = None _blocks_graph_runner: "MultimodalEncoderGraphRunner | None" = None + @property + def encoder_cuda_graph_config( + self) -> "MultimodalEncoderCudaGraphConfig | None": + """CUDA graph configuration selected for this encoder tower.""" + return self._encoder_cuda_graph_config + + @encoder_cuda_graph_config.setter + def encoder_cuda_graph_config( + self, config: "MultimodalEncoderCudaGraphConfig | None") -> None: + self._encoder_cuda_graph_config = config + def configure_encoder_cuda_graph( self, multimodal_config: "MultimodalConfig | None") -> None: """Select this tower's configuration without rejecting other towers' modalities. diff --git a/tensorrt_llm/_torch/models/modeling_radio.py b/tensorrt_llm/_torch/models/modeling_radio.py index d01fc0898d3e..50bed2a138a1 100644 --- a/tensorrt_llm/_torch/models/modeling_radio.py +++ b/tensorrt_llm/_torch/models/modeling_radio.py @@ -1296,7 +1296,7 @@ def __init__(self, inter_feature_normalizer=inter_feature_normalizer, model_config=self.model_config, ) - self.radio_model.model._encoder_cuda_graph_config = encoder_cuda_graph_config + self.radio_model.model.encoder_cuda_graph_config = encoder_cuda_graph_config def load_weights(self, weights): # Load radio_model weights for pytorch modules. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py index 2a894a9f2172..be9d642dcfd3 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/multimodal.py @@ -227,14 +227,15 @@ def setup_mm_encoder_attn_metadata( encoder_max_seq_len = max_seq_len if attention_metadata_capacity is not None: setup_kwargs["attention_metadata_capacity"] = attention_metadata_capacity - graph_config = module._encoder_cuda_graph_config + graph_config = module.encoder_cuda_graph_config if graph_config is not None: graph_tokens = ( max(tokens for tokens, _ in graph_config.buckets) if graph_config.buckets is not None - else graph_config.max_num_tokens or encoder_max_num_tokens + else graph_config.max_num_tokens ) - encoder_max_seq_len = max(max_seq_len, graph_tokens) + if graph_tokens is not None: + encoder_max_seq_len = max(max_seq_len, graph_tokens) module.setup_attn_metadata(**setup_kwargs) module.set_attn_max_seq_len(encoder_max_seq_len) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 9cc1e49d69b1..b76bd860e968 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -226,9 +226,10 @@ def __init__( # The multimodal encoder token budget falls back to the LLM-side value # when unset. It may be raised after model load because atomic MM items # cannot be split. - self.encoder_max_num_tokens = (llm_args.encoder_max_num_tokens - if llm_args.encoder_max_num_tokens - is not None else self.max_num_tokens) + configured_encoder_max_num_tokens = ( + llm_args.encoder_max_num_tokens if llm_args.encoder_max_num_tokens + is not None else self.max_num_tokens) + self.encoder_max_num_tokens = configured_encoder_max_num_tokens if checkpoint_loader is None: checkpoint_loader = _construct_checkpoint_loader( @@ -546,8 +547,7 @@ def __init__( setup_mm_encoder_cuda_graphs( self.model, mm_config.encoder_cuda_graph if mm_config is not None else None, - max_num_tokens=(self.llm_args.encoder_max_num_tokens - or self.max_num_tokens)) + max_num_tokens=configured_encoder_max_num_tokens) except Exception as e: import traceback traceback.print_exception(Exception, e, e.__traceback__) diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index f65b34095b43..182f80a53961 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -69,7 +69,7 @@ def __init__(self, name: str, config: MultimodalEncoderCudaGraphConfig | None) - super().__init__() self.weight = torch.nn.Parameter(torch.empty(1)) self.name = name - self._encoder_cuda_graph_config = config + self.encoder_cuda_graph_config = config self._blocks_graph_runner = None self.runner = Mock() self.runner.capture_all.side_effect = lambda device: events.append( @@ -85,7 +85,7 @@ def setup_attn_metadata( def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) -> Mock: assert events.index((self.name, "metadata", 1024)) >= 0 - assert config is self._encoder_cuda_graph_config + assert config is self.encoder_cuda_graph_config return self.runner config = MultimodalEncoderCudaGraphConfig(buckets=[(256, 1)]) @@ -122,7 +122,7 @@ def __init__(self, modality: str, mm_config: MultimodalConfig) -> None: self.runner = Mock() def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) -> Mock: - assert config is self._encoder_cuda_graph_config + assert config is self.encoder_cuda_graph_config return self.runner vision_config = MultimodalEncoderCudaGraphConfig(buckets=[(256, 1)]) @@ -133,8 +133,8 @@ def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) - model = torch.nn.Module() model.vision = Encoder("vision", mm_config) model.audio = Encoder("audio", mm_config) - assert model.vision._encoder_cuda_graph_config is vision_config - assert model.audio._encoder_cuda_graph_config is audio_config + assert model.vision.encoder_cuda_graph_config is vision_config + assert model.audio.encoder_cuda_graph_config is audio_config setup_mm_encoder_cuda_graphs(model, mm_config.encoder_cuda_graph) setup_mm_encoder_cuda_graphs(model, mm_config.encoder_cuda_graph) @@ -196,18 +196,132 @@ def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) - def test_automatic_encoder_graphs_need_a_token_budget_before_capture() -> None: encoder = MultimodalEncoderMixin() - encoder._encoder_cuda_graph_config = MultimodalEncoderCudaGraphConfig(max_batch_size=8) + encoder.encoder_cuda_graph_config = MultimodalEncoderCudaGraphConfig(max_batch_size=8) with pytest.raises(ValueError, match="from the config or engine"): encoder.enable_blocks_cuda_graph() +@pytest.mark.parametrize("encoder_max_num_tokens", [None, 1024]) +@pytest.mark.parametrize("item_scheduling", [False, True]) +def test_engine_captures_encoder_graphs_with_original_token_budget( + monkeypatch: pytest.MonkeyPatch, + encoder_max_num_tokens: int | None, + item_scheduling: bool, +) -> None: + from tensorrt_llm._torch.pyexecutor import _util, model_engine + from tensorrt_llm.mapping import Mapping + + class Encoder(torch.nn.Module, MultimodalEncoderMixin): + encoder_cuda_graph_modality = "vision" + + def __init__(self, mm_config: MultimodalConfig) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.configure_encoder_cuda_graph(mm_config) + self.setup_attn_metadata = Mock() + self.set_attn_max_seq_len = Mock() + self.build_runner = Mock(return_value=Mock()) + + def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) -> Mock: + return self.build_runner(config) + + mm_config = MultimodalConfig(encoder_cuda_graph={"vision": {"max_batch_size": 8}}) + mapping = Mapping() + model = torch.nn.Module() + model.encoder = Encoder(mm_config) + model.config = SimpleNamespace(torch_dtype=torch.float32) + model.model_config = SimpleNamespace( + pretrained_config=model.config, mapping=mapping, sparse_attention_config=None + ) + llm_args = SimpleNamespace( + encode_only=False, + mm_encoder_only=False, + get_runtime_sizes=lambda: (1, 4096, 128, 4), + encoder_max_batch_size=None, + encoder_max_num_tokens=encoder_max_num_tokens, + enable_in_graph_sampling=False, + enable_return_routed_experts=False, + multimodal_config=mm_config, + checkpoint_format="HF", + trust_remote_code=False, + disable_overlap_scheduler=True, + kv_cache_config=None, + sleep_config=None, + enable_layerwise_nvtx_marker=False, + cuda_graph_config=None, + torch_compile_config=None, + prefill_cuda_graph_backend=None, + prefill_capture_num_tokens=None, + allreduce_strategy="AUTO", + attn_backend="TRTLLM", + sparse_attention_config=None, + ) + for name in ( + "should_enable_adp_dummy_fixes", + "should_enable_scheduler_aware_adp_dummy", + "should_enable_non_overlap_adp_forward_intent", + "should_enable_overlap_headroom", + "resolved_kv_cache_manager_is_v2", + ): + monkeypatch.setattr(_util, name, Mock(return_value=False)) + monkeypatch.setattr(_util, "compute_max_num_sequences", Mock(return_value=4)) + monkeypatch.setattr(model_engine, "_configure_deep_gemm_pdl", Mock()) + monkeypatch.setattr(model_engine, "create_input_processor", Mock(return_value=object())) + monkeypatch.setattr( + model_engine, "resolve_mrope_position_deltas_cache", Mock(return_value=None) + ) + monkeypatch.setattr(model_engine, "is_hybrid_linear", Mock(return_value=False)) + scheduler = ( + SimpleNamespace( + encoder_max_num_tokens=65536, + output_budget_bytes=1024, + bytes_per_embedding=4, + attention_metadata_capacity=None, + ) + if item_scheduling + else None + ) + monkeypatch.setattr( + model_engine.MultimodalItemScheduler, "maybe_create", Mock(return_value=scheduler) + ) + engine_cls = model_engine.PyTorchModelEngine + monkeypatch.setattr(engine_cls, "_validate_breakable_cuda_graph_compatibility", Mock()) + monkeypatch.setattr(engine_cls, "_init_model_capacity", engine_cls._init_max_num_tokens) + monkeypatch.setattr(engine_cls, "__del__", lambda self: None) + monkeypatch.setattr( + model_engine, + "get_attention_backend", + Mock(side_effect=RuntimeError("encoder graph setup complete")), + ) + engine = engine_cls.__new__(engine_cls) + with pytest.raises(RuntimeError, match="encoder graph setup complete"): + engine_cls.__init__( + engine, + model_path="dummy", + mapping=mapping, + model=model, + llm_args=llm_args, + checkpoint_loader=Mock(), + ) + + configured_budget = encoder_max_num_tokens if encoder_max_num_tokens is not None else 4096 + assert engine.max_num_tokens == 512 + assert engine.encoder_max_num_tokens == (65536 if item_scheduling else configured_budget) + model.encoder.set_attn_max_seq_len.assert_called_once_with(engine.encoder_max_num_tokens) + capture_config = model.encoder.build_runner.call_args.args[0] + assert capture_config.max_num_tokens == configured_budget + assert mm_config.encoder_cuda_graph["vision"].max_num_tokens is None + model.encoder.build_runner.assert_called_once() + model.encoder._blocks_graph_runner.capture_all.assert_called_once_with(torch.device("cpu")) + + @pytest.mark.parametrize("enable_padding", [False, True]) @pytest.mark.parametrize("automatic", [False, True], ids=["explicit", "automatic"]) def test_encoder_graphs_preserve_eager_metadata_capacity(enable_padding, automatic) -> None: class Encoder(torch.nn.Module, MultimodalEncoderMixin): def __init__(self) -> None: super().__init__() - self._encoder_cuda_graph_config = ( + self.encoder_cuda_graph_config = ( MultimodalEncoderCudaGraphConfig( max_batch_size=8, max_num_tokens=64, enable_padding=enable_padding ) From 111e0d73fbf239d2dfbfbe3d26dfa09efaccda92 Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Mon, 19 Oct 2026 04:48:51 -0700 Subject: [PATCH 7/8] [None][fix] use multimodal config defaults in model engine test Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- .../_torch/executor/test_pytorch_model_engine_warmup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index acb8fec7c41c..65b0ff226762 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -53,6 +53,7 @@ from tensorrt_llm.llmapi.llm_args import ( DecodingBaseConfig, DraftTargetDecodingConfig, + MultimodalConfig, PARDDecodingConfig, PrefillCudaGraphBackend, TorchLlmArgs, @@ -113,7 +114,7 @@ def test_pcg_fx_fallback_policy_is_model_specific( encoder_max_num_tokens=None, enable_in_graph_sampling=False, enable_return_routed_experts=False, - multimodal_config=SimpleNamespace(video_pruning_rate=None), + multimodal_config=MultimodalConfig(), checkpoint_format="HF", trust_remote_code=False, disable_overlap_scheduler=True, From d55177984101d64fa720aed66994790713def15f Mon Sep 17 00:00:00 2001 From: yechank <161688079+yechank-nvidia@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:12:21 -0700 Subject: [PATCH 8/8] [None][fix] preserve shared encoder graph teardown after rebase Reuse the existing graph teardown in MultimodalEncoderMixin so Qwen, RADIO and GLM honor the engine NCCL window ownership policy. Cover repeated teardown and capture in the existing lifecycle test. Regenerate serving schemas for automatic capture settings and validate their YAML bounds. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com> --- .../models/modeling_glm5_next_vision.py | 6 -- .../models/modeling_multimodal_encoder.py | 10 ++ tensorrt_llm/_torch/models/modeling_radio.py | 10 -- .../schemas/trtllm-serve-config.schema.json | 90 ++++++++++++---- .../trtllm-serve-disagg-config.schema.json | 101 ++++++++++++++---- .../_torch/executor/engine/test_multimodal.py | 20 +++- .../test_trtllm_serve_config_schemas.py | 33 ++++++ 7 files changed, 209 insertions(+), 61 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py index cefa2a654ec2..c059c628a00c 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next_vision.py @@ -453,12 +453,6 @@ def _build_blocks_graph_runner( config=config, ) - def clear_blocks_cuda_graph(self, *, release_nccl_window_owners: bool = True) -> None: - """Release the block graphs at the engine coordinated teardown boundary.""" - if self._blocks_graph_runner is not None: - self._blocks_graph_runner.clear(release_nccl_window_owners=release_nccl_window_owners) - self._blocks_graph_runner = None - def _encoder_graph_fn( self, inputs: Mapping[str, torch.Tensor], attn_metadata: AttentionMetadata ) -> Dict[str, torch.Tensor]: diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py index 80101b864658..18cfc50749ed 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_encoder.py @@ -215,6 +215,16 @@ def enable_blocks_cuda_graph( runner.capture_all(device) self._blocks_graph_runner = runner + def clear_blocks_cuda_graph(self, + *, + release_nccl_window_owners: bool = True + ) -> None: + """Release block graphs at the engine's coordinated teardown boundary.""" + if self._blocks_graph_runner is not None: + self._blocks_graph_runner.clear( + release_nccl_window_owners=release_nccl_window_owners) + self._blocks_graph_runner = None + def get_encoder_attention_metadata_capacity( self, max_num_tokens: int) -> dict[str, int]: """Map the token budget to model-internal attention sequences. diff --git a/tensorrt_llm/_torch/models/modeling_radio.py b/tensorrt_llm/_torch/models/modeling_radio.py index 50bed2a138a1..090092f5c786 100644 --- a/tensorrt_llm/_torch/models/modeling_radio.py +++ b/tensorrt_llm/_torch/models/modeling_radio.py @@ -796,16 +796,6 @@ def setup_attn_metadata( metadata_kwargs["kv_layout"] = "NHD" self.attn_metadata = self.metadata_cls(**metadata_kwargs) - def clear_blocks_cuda_graph(self, - *, - release_nccl_window_owners: bool = True - ) -> None: - """Release the block graphs at the engine coordinated teardown boundary.""" - if self._blocks_graph_runner is not None: - self._blocks_graph_runner.clear( - release_nccl_window_owners=release_nccl_window_owners) - self._blocks_graph_runner = None - def prepare_attn_metadata( self, batch_size: int, seq_lengths: List[int], attn_metadata: AttentionMetadata) -> AttentionMetadata: diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 1ebeb8105b0a..3eac7f79bf0c 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4282,29 +4282,37 @@ }, "MultimodalEncoderCudaGraphConfig": { "additionalProperties": false, - "description": "CUDA graph capture for multimodal vision / audio encoders.\n\nOne graph is captured per `buckets` entry; replay falls back to eager when the live workload's\npadded shape does not match any bucket.\n\nNOTE: enabling this will lead to higher GPU memory usage. It is usually more beneficial for\n`buckets` corresponding to lesser compute.", + "description": "CUDA graph capture for multimodal vision / audio encoders.\n\nSet `max_batch_size` to generate capture shapes automatically. Advanced users can instead\nspecify `buckets`. Replay falls back to eager when no captured shape fits the live workload.\n\nNOTE: enabling this will lead to higher GPU memory usage. It is usually more beneficial for\n`buckets` corresponding to lesser compute.", "properties": { "buckets": { - "description": "Explicit encoder graph buckets as `(total_tokens, num_contexts)` pairs. Each pair captures one graph sized for up to `total_tokens` real encoder tokens across exactly `num_contexts` real contexts.", - "items": { - "maxItems": 2, - "minItems": 2, - "prefixItems": [ - { - "exclusiveMinimum": 0, - "type": "integer" + "anyOf": [ + { + "items": { + "maxItems": 2, + "minItems": 2, + "prefixItems": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "exclusiveMinimum": 0, + "type": "integer" + } + ], + "type": "array" }, - { - "exclusiveMinimum": 0, - "type": "integer" - } - ], - "type": "array" - }, - "minItems": 1, + "minItems": 1, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Explicit encoder graph buckets as `(total_tokens, num_contexts)` pairs. Each pair captures one graph sized for up to `total_tokens` real encoder tokens across exactly `num_contexts` real contexts.", "status": "prototype", - "title": "Buckets", - "type": "array" + "title": "Buckets" }, "enable_padding": { "default": true, @@ -4320,6 +4328,36 @@ "title": "Enable Replay Stats", "type": "boolean" }, + "max_batch_size": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Maximum number of encoder contexts per CUDA graph. Automatically captures every context count from 1 to this limit, with token capacities derived from the encoder token budget. Contexts are encoder sequences, such as image or video frame groups and audio chunks; they are not concurrent requests. Mutually exclusive with buckets.", + "status": "prototype", + "title": "Max Batch Size" + }, + "max_num_tokens": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional total token capacity for automatically generated graphs. Defaults to encoder_max_num_tokens, or max_num_tokens when the encoder limit is unset. The encoder adapter defines token units and layout. Larger workloads run eagerly. Mutually exclusive with buckets.", + "status": "prototype", + "title": "Max Num Tokens" + }, "warmup_steps": { "default": 2, "description": "Warmup iterations to run per bucket before capturing.", @@ -4329,9 +4367,6 @@ "type": "integer" } }, - "required": [ - "buckets" - ], "title": "MultimodalEncoderCudaGraphConfig", "type": "object" }, @@ -6967,6 +7002,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 6d1a3705e212..5e9ffbc0a526 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -2317,6 +2317,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { @@ -5835,29 +5846,37 @@ }, "MultimodalEncoderCudaGraphConfig": { "additionalProperties": false, - "description": "CUDA graph capture for multimodal vision / audio encoders.\n\nOne graph is captured per `buckets` entry; replay falls back to eager when the live workload's\npadded shape does not match any bucket.\n\nNOTE: enabling this will lead to higher GPU memory usage. It is usually more beneficial for\n`buckets` corresponding to lesser compute.", + "description": "CUDA graph capture for multimodal vision / audio encoders.\n\nSet `max_batch_size` to generate capture shapes automatically. Advanced users can instead\nspecify `buckets`. Replay falls back to eager when no captured shape fits the live workload.\n\nNOTE: enabling this will lead to higher GPU memory usage. It is usually more beneficial for\n`buckets` corresponding to lesser compute.", "properties": { "buckets": { - "description": "Explicit encoder graph buckets as `(total_tokens, num_contexts)` pairs. Each pair captures one graph sized for up to `total_tokens` real encoder tokens across exactly `num_contexts` real contexts.", - "items": { - "maxItems": 2, - "minItems": 2, - "prefixItems": [ - { - "exclusiveMinimum": 0, - "type": "integer" + "anyOf": [ + { + "items": { + "maxItems": 2, + "minItems": 2, + "prefixItems": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "exclusiveMinimum": 0, + "type": "integer" + } + ], + "type": "array" }, - { - "exclusiveMinimum": 0, - "type": "integer" - } - ], - "type": "array" - }, - "minItems": 1, + "minItems": 1, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Explicit encoder graph buckets as `(total_tokens, num_contexts)` pairs. Each pair captures one graph sized for up to `total_tokens` real encoder tokens across exactly `num_contexts` real contexts.", "status": "prototype", - "title": "Buckets", - "type": "array" + "title": "Buckets" }, "enable_padding": { "default": true, @@ -5873,6 +5892,36 @@ "title": "Enable Replay Stats", "type": "boolean" }, + "max_batch_size": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Maximum number of encoder contexts per CUDA graph. Automatically captures every context count from 1 to this limit, with token capacities derived from the encoder token budget. Contexts are encoder sequences, such as image or video frame groups and audio chunks; they are not concurrent requests. Mutually exclusive with buckets.", + "status": "prototype", + "title": "Max Batch Size" + }, + "max_num_tokens": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional total token capacity for automatically generated graphs. Defaults to encoder_max_num_tokens, or max_num_tokens when the encoder limit is unset. The encoder adapter defines token units and layout. Larger workloads run eagerly. Mutually exclusive with buckets.", + "status": "prototype", + "title": "Max Num Tokens" + }, "warmup_steps": { "default": 2, "description": "Warmup iterations to run per bucket before capturing.", @@ -5882,9 +5931,6 @@ "type": "integer" } }, - "required": [ - "buckets" - ], "title": "MultimodalEncoderCudaGraphConfig", "type": "object" }, @@ -8640,6 +8686,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { diff --git a/tests/unittest/_torch/executor/engine/test_multimodal.py b/tests/unittest/_torch/executor/engine/test_multimodal.py index 182f80a53961..a30aa49dfb90 100644 --- a/tests/unittest/_torch/executor/engine/test_multimodal.py +++ b/tests/unittest/_torch/executor/engine/test_multimodal.py @@ -61,7 +61,10 @@ def test_disabled_encoder_does_not_initialize_processor(): processor.get_mm_max_tokens_per_item.assert_not_called() -def test_encoder_graph_capture_follows_metadata_setup_and_runs_once() -> None: +@pytest.mark.parametrize("release_nccl_window_owners", [False, True]) +def test_encoder_graph_capture_follows_metadata_setup_and_runs_once( + release_nccl_window_owners: bool, +) -> None: events = [] class Encoder(torch.nn.Module, MultimodalEncoderMixin): @@ -111,6 +114,21 @@ def _build_blocks_graph_runner(self, config: MultimodalEncoderCudaGraphConfig) - assert model.second._blocks_graph_runner is model.second.runner model.disabled.runner.capture_all.assert_not_called() + for encoder in (model.first, model.second, model.disabled): + encoder.clear_blocks_cuda_graph(release_nccl_window_owners=release_nccl_window_owners) + encoder.clear_blocks_cuda_graph(release_nccl_window_owners=release_nccl_window_owners) + assert encoder._blocks_graph_runner is None + for encoder in (model.first, model.second): + encoder.runner.clear.assert_called_once_with( + release_nccl_window_owners=release_nccl_window_owners + ) + model.disabled.runner.clear.assert_not_called() + + setup_mm_encoder_cuda_graphs(model) + assert model.first.runner.capture_all.call_count == 2 + assert model.second.runner.capture_all.call_count == 2 + model.disabled.runner.capture_all.assert_not_called() + def test_encoder_graph_configs_are_selected_per_tower_and_validated_together() -> None: class Encoder(torch.nn.Module, MultimodalEncoderMixin): diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 97441e854cf4..251a3223aa9b 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -80,6 +80,39 @@ def schemas(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]: (generator.SERVE_SCHEMA, {"env_overrides": {"FLAG": 1, "OTHER": True}}, True), (generator.SERVE_SCHEMA, {"telemetry_config": {"disabled": True}}, True), (generator.SERVE_SCHEMA, {"multimodal_config": None}, True), + ( + generator.SERVE_SCHEMA, + {"multimodal_config": {"encoder_cuda_graph": {"vision": {"max_batch_size": 8}}}}, + True, + ), + ( + generator.SERVE_SCHEMA, + {"multimodal_config": {"encoder_cuda_graph": {"vision": {"max_batch_size": 0}}}}, + False, + ), + ( + generator.DISAGG_SCHEMA, + { + "multimodal_config": { + "encoder_cuda_graph": {"audio": {"max_batch_size": 4, "max_num_tokens": 2048}} + } + }, + True, + ), + ( + generator.SERVE_SCHEMA, + {"multimodal_config": {"encoder_cuda_graph": {"vision": {"buckets": [[256, 1]]}}}}, + True, + ), + ( + generator.SERVE_SCHEMA, + { + "multimodal_config": { + "encoder_cuda_graph": {"vision": {"max_batch_size": 8, "max_num_tokens": 0}} + } + }, + False, + ), (generator.SERVE_SCHEMA, {"num_serve_frontends": 0}, False), (generator.SERVE_SCHEMA, {"cuda_graph_config": {}}, True), (generator.SERVE_SCHEMA, {"cuda_graph_config": {"batch_sizes": [1, 2, 4]}}, True),