From 442d36031ce710aa2a777c353becb660c57ab2bd Mon Sep 17 00:00:00 2001 From: Yunzhu Lu Date: Thu, 10 Sep 2026 14:28:48 +0800 Subject: [PATCH] [MM][CG] Enable encoder CUDA Graph for MiniCPM-V (#42785) Signed-off-by: YunzhuLu Signed-off-by: Isotr0py Co-authored-by: Isotr0py --- docs/design/cuda_graphs_multimodal.md | 3 + .../multimodal/vision_language_offline.py | 3 + .../generation/test_vit_cudagraph.py | 65 ++ tests/v1/cudagraph/test_encoder_cudagraph.py | 10 + vllm/model_executor/models/cohere_compass.py | 3 +- vllm/model_executor/models/deepseek_ocr.py | 3 +- vllm/model_executor/models/ernie45_vl.py | 3 +- vllm/model_executor/models/gemma3_mm.py | 3 +- vllm/model_executor/models/gemma4_mm.py | 3 +- vllm/model_executor/models/glm4_1v.py | 3 +- .../models/idefics2_vision_model.py | 89 +-- vllm/model_executor/models/idefics3.py | 3 +- vllm/model_executor/models/interfaces.py | 16 +- vllm/model_executor/models/internvl.py | 3 +- vllm/model_executor/models/kimi_k25.py | 3 +- vllm/model_executor/models/kimi_vl.py | 3 +- vllm/model_executor/models/lfm2_vl.py | 3 +- vllm/model_executor/models/minicpmv.py | 570 +++++++++++++++++- vllm/model_executor/models/mllama4.py | 3 +- vllm/model_executor/models/qwen2_5_vl.py | 3 +- vllm/model_executor/models/qwen2_vl.py | 3 +- vllm/model_executor/models/qwen3_vl.py | 3 +- vllm/model_executor/models/step3_vl.py | 3 +- vllm/models/kimi_k3/nvidia/model.py | 3 +- vllm/v1/worker/encoder_cudagraph.py | 123 +++- vllm/v1/worker/encoder_cudagraph_defs.py | 16 +- 26 files changed, 843 insertions(+), 103 deletions(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 0fbb6801e4..dca200e2de 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -35,6 +35,8 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili | `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ | | `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ | | `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ | +| `MiniCPMV` | `MiniCPMV2.5` | ✅︎ | ❌︎ | ❌︎ | +| `MiniCPMV` | `MiniCPMV2.6`, `MiniCPMV4.0` | ✅︎ | ✅︎ | ❌︎ | ### Model x Hardware @@ -54,6 +56,7 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili | `Qwen3_5ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | | `Qwen3_5MoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | | `Step3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `MiniCPMV` | ✅︎ | ✅︎ | ❔ | ❔ | !!! note Encoder CUDA Graph has currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 97053ba0c2..33e1a8a36b 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -2276,6 +2276,9 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [ "glm4_1v", "deepseek_ocr", "ernie45_vl", + "minicpmv2_5_vl", + "minicpmv2_6_vl", + "minicpmv4_vl", ] diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 10dbd0c837..f8136ae8f3 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -74,6 +74,20 @@ def ernie45_vl_chat_template(content: str) -> str: ) +def minicpmv_25_chat_template(content: str) -> str: + """Llama3-style chat template used by MiniCPM-V 2.5.""" + return ( + f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n" + f"{content}" + f"<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + ) + + +def minicpmv_chat_template(content: str) -> str: + """ChatML template used by MiniCPM-V 2.6 / 4.0 / 4.5.""" + return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n" + + MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { "gemma3": VitCudagraphTestConfig( model="google/gemma-3-4b-it", @@ -316,6 +330,57 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { needs_video_metadata=True, marks=[pytest.mark.core_model], ), + "minicpmv_25": VitCudagraphTestConfig( + model="openbmb/MiniCPM-Llama3-V-2_5", + modalities=["image"], + image_prompt=minicpmv_25_chat_template( + "(./)\nWhat is in this image?" + ), + # CI runs on 35GB MIG slices: every budget captures one graph per + # patch-grid bucket, so the default budgets through max_model_len + # OOM there. A small budget set covers the test images (each item + # is at most (max_slice_num + 1) * query_num = 640 tokens). + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [64, 1024], + }, + vllm_runner_kwargs={"trust_remote_code": True}, + marks=[pytest.mark.core_model], + ), + "minicpmv_26": VitCudagraphTestConfig( + model="openbmb/MiniCPM-V-2_6", + image_prompt=minicpmv_chat_template( + "(./)\nWhat is in this image?" + ), + video_prompt=minicpmv_chat_template( + "()\nDescribe this video in one sentence." + ), + max_model_len=2048, + # Fewer frames keep the video item within the 1024 budget. + num_video_frames=4, + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [64, 1024], + "encoder_cudagraph_max_frames_per_batch": 4, + }, + vllm_runner_kwargs={"trust_remote_code": True}, + marks=[pytest.mark.core_model], + ), + "minicpmv_40": VitCudagraphTestConfig( + model="openbmb/MiniCPM-V-4", + image_prompt=minicpmv_chat_template( + "(./)\nWhat is in this image?" + ), + video_prompt=minicpmv_chat_template( + "()\nDescribe this video in one sentence." + ), + max_model_len=2048, + num_video_frames=4, + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [64, 1024], + "encoder_cudagraph_max_frames_per_batch": 4, + }, + vllm_runner_kwargs={"trust_remote_code": True}, + marks=[pytest.mark.core_model], + ), } diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index a01214a40d..096f0e4075 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -12,6 +12,7 @@ Test organization: - TestEncoderCudaGraphVideoReplay — video modality capture, replay """ +from collections.abc import Hashable from typing import Any import pytest @@ -116,9 +117,15 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager: ) mgr.budget_graphs = {"default": {}} mgr.graph_pool = None + mgr._capture_axes = () mgr.graph_hits = 0 mgr.graph_misses = 0 mgr.log_stats_interval = 100 + mgr.config = EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[], + out_hidden_size=32, + ) return mgr @@ -348,6 +355,7 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph): device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ) -> EncoderCudaGraphCaptureInputs: per_image_output = token_budget // max_batch_size grid_config = [ @@ -426,6 +434,7 @@ def _make_manager_for_gpu( mgr.use_dp = False mgr.budget_graphs = {"default": {}} mgr.graph_pool = None + mgr._capture_axes = () mgr.graph_hits = 0 mgr.graph_misses = 0 mgr.log_stats_interval = 100 @@ -656,6 +665,7 @@ class SimpleMockViTVideoModel(SimpleMockViTModel): device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ) -> EncoderCudaGraphCaptureInputs: per_item_output = token_budget // max_batch_size frames_per_item = max_frames_per_batch // max_batch_size diff --git a/vllm/model_executor/models/cohere_compass.py b/vllm/model_executor/models/cohere_compass.py index cc86d5f928..0256f747c9 100644 --- a/vllm/model_executor/models/cohere_compass.py +++ b/vllm/model_executor/models/cohere_compass.py @@ -24,7 +24,7 @@ """Inference-only Cohere Compass model compatible with HuggingFace weights.""" import math -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence from functools import lru_cache, partial from itertools import islice from typing import Annotated, Any, Literal, TypeAlias @@ -1598,6 +1598,7 @@ class CohereCompassForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 8365c9069e..114498764c 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -3,7 +3,7 @@ """Inference-only Deepseek-OCR model compatible with HuggingFace weights.""" import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from typing import Annotated, Any, Literal import torch @@ -804,6 +804,7 @@ class DeepseekOCRForCausalLM( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): assert path in ("global", "local") diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index a421eea976..716aeaa3b7 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -24,7 +24,7 @@ """Inference-only Ernie VL model compatible with HuggingFace weights.""" import math -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Any, Literal @@ -1681,6 +1681,7 @@ class Ernie4_5_VLMoeForConditionalGeneration( device, dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index 114085921e..b45327fbeb 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from typing import Annotated, Any, Literal import torch @@ -798,6 +798,7 @@ class Gemma3ForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 2a15dacb64..326d179cc2 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -15,7 +15,7 @@ reason about temporal order. """ import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal import numpy as np @@ -1812,6 +1812,7 @@ class Gemma4ForConditionalGeneration( device: torch.device | str = "cpu", dtype: torch.dtype | None = None, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, **kwargs: Any, ) -> "EncoderCudaGraphCaptureInputs": from vllm.v1.worker.encoder_cudagraph_defs import ( diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index ec72210a10..163115ff3d 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -28,7 +28,7 @@ compatible with HuggingFace weights.""" import math -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Any, Literal, TypeAlias @@ -2077,6 +2077,7 @@ class Glm4vForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py index 86b05f7478..3021e66f75 100644 --- a/vllm/model_executor/models/idefics2_vision_model.py +++ b/vllm/model_executor/models/idefics2_vision_model.py @@ -38,8 +38,7 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.utils.gpu_sync_debug import gpu_sync_allowed -from vllm.utils.torch_utils import PIN_MEMORY +from vllm.utils.torch_utils import async_tensor_h2d from .utils import AutoWeightsLoader, WeightsMapper from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model @@ -81,41 +80,61 @@ class Idefics2VisionEmbeddings(nn.Module): patch_attention_mask: torch.BoolTensor, tgt_sizes: torch.IntTensor | None = None, ) -> torch.Tensor: - batch_size, max_nb_patches_h, max_nb_patches_w = patch_attention_mask.shape + batch_size = patch_attention_mask.shape[0] + device = patch_attention_mask.device boundaries = torch.arange( - 1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side - ) - position_ids = torch.full( - size=(batch_size, max_nb_patches_h * max_nb_patches_w), - fill_value=0, - pin_memory=PIN_MEMORY, + 1 / self.num_patches_per_side, + 1.0, + 1 / self.num_patches_per_side, + device=device, ) - with gpu_sync_allowed(): - patch_attention_mask_cpu = patch_attention_mask.cpu() - tgt_sizes_cpu = tgt_sizes.cpu() if tgt_sizes is not None else None + flat_patch_attention_mask = patch_attention_mask.view(batch_size, -1) + max_patches = flat_patch_attention_mask.shape[-1] - for batch_idx, p_attn_mask in enumerate(patch_attention_mask_cpu): - if tgt_sizes_cpu is not None: - nb_patches_h = tgt_sizes_cpu[batch_idx][0] - nb_patches_w = tgt_sizes_cpu[batch_idx][1] + if tgt_sizes is not None: + # tgt_sizes may come from CPU-side mm_kwargs; use a pinned async + # copy to stay clean under VLLM_GPU_SYNC_CHECK. + if tgt_sizes.device != device: + tgt_sizes = async_tensor_h2d(tgt_sizes, device) + nb_patches_h = tgt_sizes[:, 0] + nb_patches_w = tgt_sizes[:, 1] + else: + if patch_attention_mask.dim() == 3: + nb_patches_h = patch_attention_mask[:, :, 0].sum(dim=1) + nb_patches_w = patch_attention_mask[:, 0, :].sum(dim=1) else: - nb_patches_h = p_attn_mask[:, 0].sum() - nb_patches_w = p_attn_mask[0].sum() - fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h) - fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w) - bucket_coords_h = torch.bucketize( - fractional_coords_h, boundaries, right=True - ) - bucket_coords_w = torch.bucketize( - fractional_coords_w, boundaries, right=True - ) - pos_ids = ( - bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w - ).flatten() - position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids + # Fallback for unexpected 2D masks without tgt_sizes + nb_patches_h = torch.ones(batch_size, device=device) + nb_patches_w = patch_attention_mask.sum(dim=1) - return position_ids.to(self.position_embedding.weight.device, non_blocking=True) + nb_h = nb_patches_h.clamp(min=1).unsqueeze(1) + nb_w = nb_patches_w.clamp(min=1).unsqueeze(1) + + i = torch.arange(max_patches, device=device) + + if patch_attention_mask.dim() == 3 and patch_attention_mask.shape[1] > 1: + stride_w = patch_attention_mask.shape[2] # padded 2D grid width + else: + stride_w = nb_w # (batch, 1) — per-item actual patch width + + h_idx = i.unsqueeze(0) // stride_w + w_idx = i.unsqueeze(0) % stride_w + + fractional_h = h_idx / nb_h + fractional_w = w_idx / nb_w + + bucket_h = torch.bucketize(fractional_h, boundaries, right=True) + bucket_w = torch.bucketize(fractional_w, boundaries, right=True) + + pos_ids = bucket_h * self.num_patches_per_side + bucket_w + + position_ids = torch.where(flat_patch_attention_mask, pos_ids.to(torch.long), 0) + + target_device = self.position_embedding.weight.device + if position_ids.device != target_device: + position_ids = async_tensor_h2d(position_ids, target_device) + return position_ids def forward( self, @@ -461,13 +480,7 @@ class Idefics2VisionTransformer(nn.Module): # - if apply_encoder_attention_mask is False, skip (not all models # sharing this encoder apply masking in attention, e.g. Aria, Phi4) # - if patch_attention_mask was None, skip attention masking - # - if any padding exists, create an additive 4D mask and pass it - # to attention; else skip mask for performance. - if ( - not self.apply_encoder_attention_mask - or flat_patch_mask is None - or not torch.any(~flat_patch_mask) - ): + if not self.apply_encoder_attention_mask or flat_patch_mask is None: attention_mask = None else: # Additive mask: masked positions receive a large negative value. diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index 80e223e4c1..baaa4106cb 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -16,7 +16,7 @@ # limitations under the License. """Inference-only Idefics3 model compatible with HuggingFace weights.""" -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from typing import Annotated, Any, Literal, TypeAlias import torch @@ -722,6 +722,7 @@ class Idefics3ForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 83e8aaa9b6..254f2c635b 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -7,6 +7,7 @@ from bisect import bisect_right from collections.abc import ( AsyncGenerator, Callable, + Hashable, Mapping, MutableSequence, Sequence, @@ -1875,6 +1876,11 @@ class SupportsEncoderCudaGraph(Protocol): - Qwen-family: slice concatenated pixel_values by cumulative patch offsets, subset grid_thw by indices. - Batched models (CLIP): index pixel_values along dim 0. + + Models that configure ``EncoderCudaGraphConfig.capture_axes`` must + additionally store the resolved per-axis keys (one key per axis, in + order) under ``ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG`` in the returned + dict; the manager pops it before the kwargs are used elsewhere. """ ... @@ -1909,8 +1915,16 @@ class SupportsEncoderCudaGraph(Protocol): device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ) -> "EncoderCudaGraphCaptureInputs": - """Create dummy inputs and buffers for CUDA graph capture.""" + """Create dummy inputs and buffers for CUDA graph capture. + + Args: + axis_keys: The resolved capture-axis keys (one per axis of + ``EncoderCudaGraphConfig.capture_axes``) this capture is for. + None or empty when no capture axes are configured; models + without capture axes ignore it. + """ ... def prepare_encoder_cudagraph_replay_buffers( diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index c18428e26f..1fcb6b8681 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -8,7 +8,7 @@ # Licensed under The MIT License [see LICENSE for details] # -------------------------------------------------------- from abc import abstractmethod -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from functools import cached_property from typing import Annotated, Any, Literal, TypeAlias, TypedDict, TypeVar @@ -1028,6 +1028,7 @@ class InternVLChatModel( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 7f313f44ea..5b0061e3cb 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -6,7 +6,7 @@ Kimi-K2.5 Model Implementation for vLLM. Kimi-K2.5 extends Kimi-K2 with vision support. """ -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal @@ -594,6 +594,7 @@ class KimiK25ForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ) -> "EncoderCudaGraphCaptureInputs": from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphCaptureInputs diff --git a/vllm/model_executor/models/kimi_vl.py b/vllm/model_executor/models/kimi_vl.py index fc7533b192..218c257e53 100644 --- a/vllm/model_executor/models/kimi_vl.py +++ b/vllm/model_executor/models/kimi_vl.py @@ -43,7 +43,7 @@ # SOFTWARE. import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from dataclasses import dataclass from typing import Annotated, Any, Literal @@ -450,6 +450,7 @@ class KimiVLForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index d27c36d828..a8b57e7c28 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -4,7 +4,7 @@ import itertools import math import typing -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from typing import Annotated, Any, Literal import torch @@ -1146,6 +1146,7 @@ class Lfm2VLForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index f8602983e5..dac5e8367f 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -26,10 +26,10 @@ import math from collections import defaultdict -from collections.abc import Callable, Iterable, Mapping, Sequence, Set +from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence, Set from functools import cached_property, partial from itertools import chain -from typing import Annotated, Any, Literal, TypeAlias +from typing import Annotated, Any, ClassVar, Literal, TypeAlias, cast import numpy as np import torch @@ -100,11 +100,19 @@ from vllm.transformers_utils.utils import convert_model_repo_to_path from vllm.utils.collection_utils import flatten_2d_lists, is_list_of from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape -from vllm.utils.torch_utils import set_default_torch_dtype +from vllm.utils.torch_utils import async_tensor_h2d, set_default_torch_dtype +from vllm.v1.worker.encoder_cudagraph_defs import ( + ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG, + EncoderCudaGraphCaptureInputs, + EncoderCudaGraphConfig, + EncoderCudaGraphReplayBuffers, + EncoderItemSpec, +) from .idefics2_vision_model import Idefics2VisionTransformer from .interfaces import ( MultiModalEmbeddings, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMultiModal, SupportsPP, @@ -218,29 +226,36 @@ class Resampler2_5(BaseResampler): device = x.device dtype = x.dtype + # Only adjust the pos cache in eager mode; during CUDA graph capture + # the cache was already expanded by the preceding warmup run + if not torch.cuda.is_current_stream_capturing(): + self._adjust_pos_cache(tgt_sizes, device=device) + + # tgt_sizes is pinned-copied only when it comes from the CPU-side + # mm_kwargs; during capture/replay it is already a GPU buffer. + if tgt_sizes.device != x.device: + tgt_sizes = async_tensor_h2d(tgt_sizes, x.device) patch_len = tgt_sizes[:, 0] * tgt_sizes[:, 1] + max_patch_len = x.shape[1] - self._adjust_pos_cache(tgt_sizes, device=device) + seq_idx = torch.arange(max_patch_len, device=device).unsqueeze(0) - max_patch_len = patch_len.max().item() - assert isinstance(max_patch_len, int) + tgt_w = tgt_sizes[:, 1].unsqueeze(1) + tgt_w = torch.clamp(tgt_w, min=1) - key_padding_mask = torch.zeros( - (bs, max_patch_len), dtype=torch.bool, device=device - ) + h_idx = seq_idx // tgt_w + w_idx = seq_idx % tgt_w - pos_embed = [] - for i in range(bs): - tgt_h, tgt_w = tgt_sizes[i].tolist() - pos_embed.append( - self.pos_embed[:tgt_h, :tgt_w, :] - .reshape((tgt_h * tgt_w, -1)) - .to(device=device, dtype=dtype) - ) # patches * D - key_padding_mask[i, patch_len[i] :] = True - pos_embed = torch.nn.utils.rnn.pad_sequence( - pos_embed, batch_first=True, padding_value=0.0 + h_idx = torch.clamp(h_idx, max=self.pos_embed.shape[0] - 1) + w_idx = torch.clamp(w_idx, max=self.pos_embed.shape[1] - 1) + + key_padding_mask = seq_idx >= patch_len.unsqueeze(1) + + pos_embed = self.pos_embed[h_idx, w_idx].to(dtype) + pos_embed = torch.where( + key_padding_mask.unsqueeze(-1), torch.zeros_like(pos_embed), pos_embed ).permute(1, 0, 2) # BLD => L * B * D + x, _ = self.kv_proj(x) # B * L * D x = self.ln_kv(x).permute(1, 0, 2) # L * B * D @@ -808,6 +823,16 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo): return max(max_frames_per_video, 1) + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + max_tokens = {"image": self.get_max_image_tokens()} + if "video" in self.get_supported_mm_limits(): + max_tokens["video"] = self.get_max_video_tokens(seq_len, mm_counts) + return max_tokens + _I = TypeVar("_I", bound=MiniCPMVProcessingInfo, default=MiniCPMVProcessingInfo) @@ -1223,6 +1248,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): # and config class self.config = config self.multimodal_config = multimodal_config + self.vllm_config = vllm_config self.version = get_version_by_config(self.config) @@ -1428,6 +1454,504 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): raise NotImplementedError +# --- Encoder CUDA graph (MiniCPM-V 2.5 / 2.6 / 4.0; mixed into subclasses) --- +# Buffer keys +_MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES = "minicpmv_tgt_sizes" +_MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK = "minicpmv_patch_attn_mask" +_MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL = "minicpmv_pixel_values" + +# mm_kwargs keys for the flat pixel-value tensor +_MINICPMV_CUDAGRAPH_FLAT_KEY_IMAGE = "minicpmv_encoder_input_flat" +_MINICPMV_CUDAGRAPH_FLAT_KEY_VIDEO = "minicpmv_video_encoder_input_flat" +# Resolved patch-grid capture-axis key ``(nb_h, nb_w)`` of the sub-batch; +# pixel dims and max patches are derived from it and the patch size. +_MINICPMV_CUDAGRAPH_PATCH_GRID_KEY = "minicpmv_encoder_patch_grid" + +_MINICPMV_BASE_PATCH_BUCKETS: tuple[tuple[int, int], ...] = ( + (32, 32), + (32, 64), + (64, 32), +) + +_ENCODER_CUDAGRAPH_MM_KWARGS_SKIP_KEYS = frozenset( + { + _MINICPMV_CUDAGRAPH_FLAT_KEY_IMAGE, + _MINICPMV_CUDAGRAPH_FLAT_KEY_VIDEO, + _MINICPMV_CUDAGRAPH_PATCH_GRID_KEY, + } +) + + +def _mcpmv_tgt_sizes_tensor( + mm_kwargs: dict[str, Any], *, video: bool +) -> torch.Tensor | list[torch.Tensor]: + key = "video_tgt_sizes" if video else "tgt_sizes" + return mm_kwargs[key] + + +def _mcpmv_normalize_tgt_sizes( + tgt_sizes: torch.Tensor | list[torch.Tensor], + slice_counts: list[int], +) -> torch.Tensor: + """Normalize tgt_sizes to shape ``(total_slices, 2)``.""" + + if isinstance(tgt_sizes, list): + if not tgt_sizes: + tgt_sizes = torch.zeros((0, 2), dtype=torch.long) + else: + tgt_sizes = torch.cat(tgt_sizes, dim=0) + + total_slices = sum(slice_counts) + if tgt_sizes.dim() == 2 and tgt_sizes.shape[0] == total_slices: + return tgt_sizes + if tgt_sizes.dim() == 3: + return torch.cat( + [tgt_sizes[i, : slice_counts[i], :] for i in range(len(slice_counts))], + dim=0, + ) + return torch.repeat_interleave( + tgt_sizes, + torch.tensor(slice_counts, device=tgt_sizes.device), + dim=0, + ) + + +def _mcpmv_pack_flat_pixels( + slices: list[torch.Tensor], + *, + pixel_height: int, + pixel_width: int, + max_num_slices: int, + device: torch.device, + dtype: torch.dtype, + patch_size: int, + tgt_sizes: torch.Tensor, +) -> torch.Tensor: + """Pack slice tensors into a fixed ``(max_num_slices, 3*H*W)`` buffer.""" + import torch.nn.functional as F + + flat_dim = 3 * pixel_height * pixel_width + grid_h = pixel_height // patch_size + grid_w = pixel_width // patch_size + max_grid_patches = grid_h * grid_w + + packed = torch.zeros((max_num_slices, flat_dim), device=device, dtype=dtype) + n = min(len(slices), max_num_slices) + tgt_list = tgt_sizes.tolist() + for i, slc in enumerate(slices[:n]): + slc = slc.to(dtype=dtype, device=device) + C, h, w = slc.shape + if h == patch_size: + nb_h, nb_w = int(tgt_list[i][0]), int(tgt_list[i][1]) + num_patches = nb_h * nb_w + + patches = ( + slc.view(C, patch_size, nb_h, nb_w, patch_size) + .permute(0, 2, 3, 1, 4) + .reshape(C, num_patches, patch_size, patch_size) + ) + + col = patches.permute(0, 2, 3, 1).reshape( + 1, C * patch_size * patch_size, num_patches + ) + col_padded = col.new_zeros(1, C * patch_size * patch_size, max_grid_patches) + col_padded[..., :num_patches] = col + + folded = F.fold( + col_padded, + output_size=(pixel_height, pixel_width), + kernel_size=patch_size, + stride=patch_size, + ) + packed[i] = folded.reshape(-1) + else: + packed[i].view(3, pixel_height, pixel_width)[:, :h, :w] = slc + return packed + + +class _MiniCPMVEncoderCudaGraphMixin(MiniCPMVBaseModel, SupportsEncoderCudaGraph): + """SupportsEncoderCudaGraph for MiniCPM-V Idefics2 + resampler (not 2.0).""" + + supports_encoder_cudagraph: ClassVar[Literal[True]] = True + + def _mcpmv_slice_pixel_size(self) -> tuple[int, int]: + image_size = int(self.vpm.embeddings.image_size) + return image_size, image_size + + def _mcpmv_patch_grid_pixel_hw( + self, patch_grid_key: tuple[int, int] + ) -> tuple[int, int]: + patch_size = int(self.vpm.embeddings.patch_size) + th, tw = patch_grid_key + return th * patch_size, tw * patch_size + + def _mcpmv_patch_grid_num_patches(self, patch_grid_key: tuple[int, int]) -> int: + th, tw = patch_grid_key + return th * tw + + def _mcpmv_patch_grid_keys(self) -> tuple[tuple[int, int], ...]: + """Ordered patch-grid keys ``(nb_h, nb_w)`` for the capture axis.""" + max_side = int(self.vpm.embeddings.image_size) // int( + self.vpm.embeddings.patch_size + ) + keys: list[tuple[int, int]] = [] + seen: set[tuple[int, int]] = set() + for th, tw in _MINICPMV_BASE_PATCH_BUCKETS: + if th <= max_side and tw <= max_side: + keys.append((th, tw)) + seen.add((th, tw)) + full = (max_side, max_side) + if full not in seen: + keys.append(full) + return tuple(keys) + + def _mcpmv_resolve_patch_grid( + self, + tgt_sizes: torch.Tensor, + slice_counts: list[int], + indices: list[int], + ) -> tuple[int, int]: + """Smallest patch-grid key covering all selected items.""" + keys = self._mcpmv_patch_grid_keys() + if not indices: + return keys[0] + tgt_groups = torch.split(tgt_sizes, slice_counts) + selected = torch.cat([tgt_groups[i] for i in indices], dim=0) + need_h = int(selected[:, 0].max().item()) + need_w = int(selected[:, 1].max().item()) + for th, tw in keys: + if th >= need_h and tw >= need_w: + return (th, tw) + return keys[-1] + + def _mcpmv_max_slices_cap( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + ) -> int: + max_slice_num = int(getattr(self.config, "max_slice_num", 9)) + query_num = max(1, int(self.config.query_num)) + max_slices_by_token_budget = max(1, token_budget // query_num) + max_slices_by_content = max_batch_size * (max_slice_num + 1) + if self.version in {(2, 6), (4, 0)} and max_frames_per_batch > 0: + max_slices_by_content = max( + max_slices_by_content, + max_frames_per_batch * (max_slice_num + 1), + ) + return max(1, min(max_slices_by_token_budget, max_slices_by_content)) + + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: + buffer_keys = [ + _MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL, + _MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES, + _MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK, + ] + # Video is only supported from 2.6 onward. + modalities = ["image"] + if self.version in {(2, 6), (4, 0)}: + modalities.append("video") + + max_frames = self.get_max_frames_per_video() if "video" in modalities else 1 + + return EncoderCudaGraphConfig( + modalities=modalities, + buffer_keys=buffer_keys, + out_hidden_size=int(self.embed_dim), + max_frames_per_video=max_frames, + capture_axes=(self._mcpmv_patch_grid_keys(),), + ) + + def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str: + if "video_pixel_values" in mm_kwargs: + return "video" + return "image" + + def get_max_frames_per_video(self) -> int: + info = MULTIMODAL_REGISTRY.get_processing_info(self.vllm_config.model_config) + assert isinstance(info, MiniCPMVProcessingInfo) + return int( + info.get_num_frames_with_most_features( + seq_len=self.vllm_config.model_config.max_model_len, + mm_counts={ + "video": self.multimodal_config.get_limit_per_prompt("video") + }, + ) + ) + + def get_encoder_cudagraph_budget_range( + self, vllm_config: VllmConfig + ) -> tuple[int, int]: + # Each slice produces exactly query_num resampler output tokens. + # A thumbnail-only image has 1 slice, so query_num is the smallest + # possible encoder output and the natural minimum budget. + min_budget = int(self.config.query_num) + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.model_config.max_model_len, + ) + return (min_budget, max_budget) + + def get_encoder_cudagraph_item_specs( + self, mm_kwargs: dict[str, Any] + ) -> list[EncoderItemSpec]: + video = self.get_input_modality(mm_kwargs) == "video" + pixel_values_key = "video_pixel_values" if video else "pixel_values" + pixel_values: list[list[torch.Tensor]] = mm_kwargs[pixel_values_key] + slice_counts = [len(img) for img in pixel_values] + tgt_sizes = _mcpmv_tgt_sizes_tensor(mm_kwargs, video=video) + tgt_sizes = _mcpmv_normalize_tgt_sizes(tgt_sizes, slice_counts) + patch_sums = tgt_sizes.prod(-1) + input_sizes = [ + int(group.sum().item()) for group in torch.split(patch_sums, slice_counts) + ] + query_num = int(self.config.query_num) + return [ + EncoderItemSpec( + input_size=input_sizes[i], + output_tokens=slice_counts[i] * query_num, + ) + for i in range(len(pixel_values)) + ] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + subset, patch_grid = self._mcpmv_select_items(mm_kwargs, indices) + subset[ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG] = (patch_grid,) + return subset + + def _mcpmv_select_items( + self, mm_kwargs: dict[str, Any], indices: list[int] + ) -> tuple[dict[str, Any], tuple[int, int]]: + """Slice mm_kwargs for `indices`; also returns the patch-grid key.""" + video = self.get_input_modality(mm_kwargs) == "video" + pixel_values_key = "video_pixel_values" if video else "pixel_values" + tgt_key = "video_tgt_sizes" if video else "tgt_sizes" + flat_key = ( + _MINICPMV_CUDAGRAPH_FLAT_KEY_VIDEO + if video + else _MINICPMV_CUDAGRAPH_FLAT_KEY_IMAGE + ) + device = next(self.vpm.parameters()).device + pixel_values: list[list[torch.Tensor]] = mm_kwargs[pixel_values_key] + tgt_sizes = _mcpmv_tgt_sizes_tensor(mm_kwargs, video=video) + + subset = { + k: v + for k, v in mm_kwargs.items() + if k not in _ENCODER_CUDAGRAPH_MM_KWARGS_SKIP_KEYS + } + + if not indices: + pixel_h, pixel_w = self._mcpmv_slice_pixel_size() + vpm_dtype = next(self.vpm.parameters()).dtype + subset.update( + { + pixel_values_key: [], + tgt_key: torch.zeros((0, 2), dtype=torch.long, device=device), + flat_key: torch.zeros( + (0, 3 * pixel_h * pixel_w), device=device, dtype=vpm_dtype + ), + } + ) + return subset, self._mcpmv_patch_grid_keys()[0] + + slice_counts = [len(item_slices) for item_slices in pixel_values] + tgt_sizes = _mcpmv_normalize_tgt_sizes(tgt_sizes, slice_counts) + tgt_groups = torch.split(tgt_sizes, slice_counts) + + patch_grid = self._mcpmv_resolve_patch_grid(tgt_sizes, slice_counts, indices) + pixel_h, pixel_w = self._mcpmv_patch_grid_pixel_hw(patch_grid) + + selected_pixel_values = [pixel_values[i] for i in indices] + selected_tgt_sizes_list = [tgt_groups[i] for i in indices] + selected_tgt_sizes = torch.cat(selected_tgt_sizes_list, dim=0) + + selected_slices = flatten_2d_lists(selected_pixel_values) + packed_flat_pixels = _mcpmv_pack_flat_pixels( + selected_slices, + pixel_height=pixel_h, + pixel_width=pixel_w, + max_num_slices=len(selected_slices), + device=selected_slices[0].device, + dtype=selected_slices[0].dtype, + patch_size=int(self.vpm.embeddings.patch_size), + tgt_sizes=selected_tgt_sizes, + ) + + subset.update( + { + pixel_values_key: selected_pixel_values, + tgt_key: selected_tgt_sizes_list, + flat_key: packed_flat_pixels, + _MINICPMV_CUDAGRAPH_PATCH_GRID_KEY: patch_grid, + } + ) + return subset, patch_grid + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, + ): + patch_grid = ( + cast("tuple[int, int]", axis_keys[0]) + if axis_keys + # Without capture-axis context, use the largest (full-resolution) + # patch grid. + else self._mcpmv_patch_grid_keys()[-1] + ) + return self._mcpmv_capture_inputs( + token_budget, + max_batch_size, + max_frames_per_batch, + device, + dtype, + patch_grid=patch_grid, + ) + + def _mcpmv_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + patch_grid: tuple[int, int], + ): + th, tw = patch_grid + pixel_h, pixel_w = self._mcpmv_patch_grid_pixel_hw(patch_grid) + max_patches = self._mcpmv_patch_grid_num_patches(patch_grid) + max_num_slices = self._mcpmv_max_slices_cap( + token_budget, + max_batch_size, + max_frames_per_batch, + ) + pixel_buffer = torch.zeros( + (max_num_slices, 3, pixel_h, pixel_w), device=device, dtype=dtype + ) + dummy_tgt_sizes = torch.zeros( + (max_num_slices, 2), dtype=torch.long, device=device + ) + dummy_tgt_sizes[:, 0] = th + dummy_tgt_sizes[:, 1] = tw + dummy_patch_mask = torch.ones( + (max_num_slices, max_patches), dtype=torch.bool, device=device + ) + values: dict[str, torch.Tensor] = { + _MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL: pixel_buffer, + _MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES: dummy_tgt_sizes, + _MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK: dummy_patch_mask, + } + return EncoderCudaGraphCaptureInputs(values=values) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ): + _ = max_batch_size + _ = max_frames_per_batch + video = self.get_input_modality(mm_kwargs) == "video" + flat_key = ( + _MINICPMV_CUDAGRAPH_FLAT_KEY_VIDEO + if video + else _MINICPMV_CUDAGRAPH_FLAT_KEY_IMAGE + ) + flat_pixels = mm_kwargs[flat_key] # (num_actual_slices, 3*pixel_h*pixel_w) + patch_grid = mm_kwargs[_MINICPMV_CUDAGRAPH_PATCH_GRID_KEY] + pixel_h, pixel_w = self._mcpmv_patch_grid_pixel_hw(patch_grid) + max_patches = self._mcpmv_patch_grid_num_patches(patch_grid) + pixel_buffer = flat_pixels.reshape(-1, 3, pixel_h, pixel_w) + + device = next(self.vpm.parameters()).device + tgt_sizes_raw = _mcpmv_tgt_sizes_tensor(mm_kwargs, video=video) + if isinstance(tgt_sizes_raw, list): + tgt_sizes_raw = torch.cat(tgt_sizes_raw, dim=0) + # tgt_sizes arrives from CPU-side mm_kwargs; use a pinned async copy + # to stay clean under VLLM_GPU_SYNC_CHECK. + tgt_sizes = async_tensor_h2d(tgt_sizes_raw, device, dtype=torch.long) + + patches_per_slice = tgt_sizes.prod(-1).clamp(max=max_patches) + col_idx = torch.arange(max_patches, device=device) + patch_attention_mask = col_idx.unsqueeze(0) < patches_per_slice.unsqueeze(1) + + values: dict[str, torch.Tensor] = { + _MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL: pixel_buffer, + _MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES: tgt_sizes, + _MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK: patch_attention_mask, + } + return EncoderCudaGraphReplayBuffers(values=values) + + def encoder_cudagraph_forward( + self, + values: dict[str, torch.Tensor], + path: str = "default", + ) -> torch.Tensor: + all_pixel_values = values[_MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL] + tgt_sizes = values[_MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES] + patch_mask = values[_MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK] + patch_attention_mask = patch_mask.unsqueeze(1) + + max_num_slices = all_pixel_values.shape[0] + # v2.5 infers patch layout from the attention mask; pass tgt_sizes=None. + vpm_tgt_sizes = None if self.version == (2, 5) else tgt_sizes + vision_embedding = self.vpm( + all_pixel_values, + patch_attention_mask=patch_attention_mask, + tgt_sizes=vpm_tgt_sizes, + ) + + resampler_out = self.resampler(vision_embedding, tgt_sizes) + + query_num = int(self.config.query_num) + return resampler_out.reshape(max_num_slices * query_num, int(self.embed_dim)) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + ) -> torch.Tensor: + """Eager encoder path; returns ``(total_tokens, embed_dim)`` like + ``encoder_cudagraph_forward``. + """ + mm_kwargs_no_flat = { + k: v + for k, v in mm_kwargs.items() + if k not in _ENCODER_CUDAGRAPH_MM_KWARGS_SKIP_KEYS + } + modalities = self._parse_and_validate_multimodal_inputs(**mm_kwargs_no_flat) + segments: list[torch.Tensor] = [] + embed_dim = self.embed_dim + for modality in modalities: + if modality == "images": + image_input = modalities["images"] + assert isinstance(image_input, MiniCPMVImagePixelInputs) + image_embeddings = self.get_vision_hidden_states(image_input) + segments.append(image_embeddings.reshape(-1, embed_dim)) + elif modality == "videos": + video_input = modalities["videos"] + assert isinstance(video_input, MiniCPMVImagePixelInputs) + video_embeddings = self.get_vision_hidden_states(video_input) + segments.append(video_embeddings.reshape(-1, embed_dim)) + if not segments: + raise RuntimeError( + "MiniCPM-V encoder cudagraph eager path expects pixel_values " + "or video_pixel_values" + ) + return torch.cat(segments, dim=0) + + class MiniCPMV2_0(MiniCPMVBaseModel): supports_encoder_tp_data = False @@ -1519,7 +2043,7 @@ class MiniCPMV2_0(MiniCPMVBaseModel): return torch.vstack(res) -class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA): +class MiniCPMV2_5(_MiniCPMVEncoderCudaGraphMixin, SupportsLoRA): packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -1611,7 +2135,7 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA): return self.resampler(vision_embedding, tgt_sizes) -class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA): +class MiniCPMV2_6(_MiniCPMVEncoderCudaGraphMixin, SupportsLoRA): hf_to_vllm_mapper = WeightsMapper( # The vision-only models have no audio tower or TTS head to load weights into. orig_to_new_prefix={"apm.": None, "audio": None, "tts": None} @@ -1715,7 +2239,7 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA): return loaded -class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): +class MiniCPMV4_0(_MiniCPMVEncoderCudaGraphMixin, SupportsLoRA): hf_to_vllm_mapper = MiniCPMV2_6.hf_to_vllm_mapper packed_modules_mapping = { "qkv_proj": [ diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 6b40cd0c86..5a6e654417 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -17,7 +17,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import math -from collections.abc import Iterable, Mapping +from collections.abc import Hashable, Iterable, Mapping from itertools import tee from typing import Annotated, Any, Literal @@ -924,6 +924,7 @@ class Llama4ForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index f409f84427..61d0153f86 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -26,7 +26,7 @@ # limitations under the License. """Inference-only Qwen2.5-VL model compatible with HuggingFace weights.""" -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Any, Literal, TypeAlias @@ -1856,6 +1856,7 @@ class Qwen2_5_VLForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index b6058a9e05..80abdc62b5 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -26,7 +26,7 @@ """Inference-only Qwen2-VL model compatible with HuggingFace weights.""" import math -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Any, Literal, TypeAlias @@ -1589,6 +1589,7 @@ class Qwen2VLForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 74e5402877..912e02e10d 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -24,7 +24,7 @@ # limitations under the License. """Inference-only Qwen3VL model compatible with HuggingFace weights.""" -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence from functools import lru_cache, partial from itertools import islice from typing import Any @@ -2120,6 +2120,7 @@ class Qwen3VLForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index a06c221968..9f738b7c64 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Hashable, Iterable, Mapping, Sequence from math import sqrt from typing import Annotated, Any, Literal, TypeAlias @@ -828,6 +828,7 @@ class Step3VLForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 3e7645d8d2..693795da38 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -3,7 +3,7 @@ """Kimi-K3 multimodal model implementation for vLLM.""" import math -from collections.abc import Iterable +from collections.abc import Hashable, Iterable from typing import Any, cast import torch @@ -1948,6 +1948,7 @@ class KimiK3ForConditionalGeneration( device: torch.device, dtype: torch.dtype, path: str = "default", + axis_keys: tuple[Hashable, ...] | None = None, ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index a1b955a741..a60e0d1892 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -2,8 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """CUDA graph manager for vision encoder budget-batch execution.""" +import itertools +import math +from collections.abc import Hashable from dataclasses import dataclass -from typing import Any +from typing import Any, TypeAlias import torch @@ -22,12 +25,20 @@ from vllm.model_executor.models.vision import get_load_balance_assignment from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import current_stream from vllm.v1.worker.encoder_cudagraph_defs import ( + ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG, EncoderCudaGraphConfig, EncoderItemSpec, ) logger = init_logger(__name__) +CaptureAxisKeys: TypeAlias = tuple[Hashable, ...] + +BudgetGraphMapKey: TypeAlias = int | tuple[int, CaptureAxisKeys] + +# Warn when capture axes multiply the captured graph count past this point. +_CAPTURE_GRAPH_COUNT_WARNING_THRESHOLD = 64 + @dataclass class BudgetGraphMetadata: @@ -50,6 +61,8 @@ class BudgetGraphMetadata: input_buffers: dict[str, torch.Tensor] # Output written by graph, read after replay output_buffer: torch.Tensor + # Per-axis capture keys, empty when no extra capture axis is used. + axis_keys: CaptureAxisKeys = () class EncoderCudaGraphManager: @@ -77,6 +90,18 @@ class EncoderCudaGraphManager: multimodal_config = vllm_config.model_config.multimodal_config + self._capture_axes: tuple[CaptureAxisKeys, ...] = tuple( + tuple(axis) for axis in self.config.capture_axes + ) + if any(len(axis) == 0 for axis in self._capture_axes): + raise ValueError("Encoder cudagraph capture axes must be non-empty.") + + if len(self.config.paths) > 1 and self._capture_axes: + raise NotImplementedError( + "Combining dual-path encoder CUDA graphs with capture axes " + "is not supported yet." + ) + # Invariant: max_batch_size <= min_token_budget. # This ensures per_image_output = budget // max_batch_size >= 1 # for every captured budget, preventing reshape crashes on empty @@ -154,7 +179,7 @@ class EncoderCudaGraphManager: and vllm_config.parallel_config.tensor_parallel_size > 1 ) - self.budget_graphs: dict[str, dict[int, BudgetGraphMetadata]] = {} + self.budget_graphs: dict[str, dict[BudgetGraphMapKey, BudgetGraphMetadata]] = {} self.graph_pool: Any | None = None self.graph_hits = 0 self.graph_misses = 0 @@ -181,11 +206,13 @@ class EncoderCudaGraphManager: logger.info( "EncoderCudaGraphManager initialized with paths=%s, " - "max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s", + "max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s, " + "capture_axes=%s", self.path_token_budgets, self.max_batch_size, self.max_frames_per_batch, self.use_dp, + self._capture_axes, ) @staticmethod @@ -219,38 +246,63 @@ class EncoderCudaGraphManager: """Capture CUDA graphs for every configured path and token budget.""" self.graph_pool = graph_pool + num_graphs = self.get_num_graphs_to_capture() + if num_graphs > _CAPTURE_GRAPH_COUNT_WARNING_THRESHOLD: + logger.warning( + "Capturing %d encoder CUDA graphs; capture axes multiply the " + "per-budget graph count, consider pruning axis keys.", + num_graphs, + ) + for path, budgets in self.path_token_budgets.items(): for token_budget in sorted(budgets, reverse=True): - if token_budget > 0: - self._capture_budget_graph(token_budget, path=path) + if token_budget == 0: + continue + for axis_keys in itertools.product(*self._capture_axes): + self._capture_budget_graph( + token_budget, + path=path, + axis_keys=axis_keys, + ) logger.info( "Encoder CUDA graph capture complete. Captured %d graphs across %d paths.", - self.get_num_graphs_to_capture(), + num_graphs, len(self.path_token_budgets), ) + def _num_graphs_for_budgets(self, budgets: list[int]) -> int: + num_budgets = sum(1 for budget in budgets if budget != 0) + return num_budgets * math.prod(len(axis) for axis in self._capture_axes) + def get_num_graphs_to_capture(self) -> int: return sum( - token_budget > 0 + self._num_graphs_for_budgets(budgets) for budgets in self.path_token_budgets.values() - for token_budget in budgets ) - def _get_graph_set(self, path: str = "default") -> dict[int, BudgetGraphMetadata]: + def _get_graph_set( + self, path: str = "default" + ) -> dict[BudgetGraphMapKey, BudgetGraphMetadata]: """Return the captured graphs for one encoder path.""" if path not in self.budget_graphs: self.budget_graphs[path] = {} return self.budget_graphs[path] - def _capture_budget_graph(self, token_budget: int, path: str = "default"): + def _capture_budget_graph( + self, + token_budget: int, + path: str = "default", + axis_keys: CaptureAxisKeys = (), + ): """Capture CUDA graph for a single token budget.""" logger.debug( "Capturing encoder cudagraph for budget=%d, max_batch_size=%d, " - "max_frames_per_batch=%d", + "max_frames_per_batch=%d, axis_keys=%s", token_budget, self.max_batch_size, self.max_frames_per_batch, + axis_keys, ) graph_set = self._get_graph_set(path) @@ -262,6 +314,7 @@ class EncoderCudaGraphManager: self.device, self.dtype, path, + axis_keys, ) values = capture_inputs.values @@ -278,13 +331,17 @@ class EncoderCudaGraphManager: output = self.model.encoder_cudagraph_forward({**values}, path=path) output_buffer.copy_(output) - graph_set[token_budget] = BudgetGraphMetadata( + graph_map_key: BudgetGraphMapKey = ( + token_budget if not axis_keys else (token_budget, axis_keys) + ) + graph_set[graph_map_key] = BudgetGraphMetadata( token_budget=token_budget, max_batch_size=self.max_batch_size, max_frames_per_batch=self.max_frames_per_batch, graph=graph, input_buffers=values, output_buffer=output_buffer, + axis_keys=axis_keys, ) def _find_smallest_fitting_budget_given_tokens( @@ -310,12 +367,24 @@ class EncoderCudaGraphManager: def _select_items( self, mm_kwargs: dict[str, Any], indices: list[int] - ) -> dict[str, Any]: - """Select the mm kwargs for `indices` from the model.""" + ) -> tuple[dict[str, Any], CaptureAxisKeys]: + """Select the mm kwargs for `indices` from the model. + + Returns the sliced mm_kwargs together with the resolved capture-axis + keys (popped out of the returned dict; empty without capture axes). + """ # Same as `_get_item_specs`: implementations re-read the per-item # grid/patch counts to slice the batch, so the D2H is inherent. with gpu_sync_allowed(): - return self.model.select_encoder_cudagraph_items(mm_kwargs, indices) + selected = self.model.select_encoder_cudagraph_items(mm_kwargs, indices) + axis_keys = selected.pop(ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG, ()) + if len(axis_keys) != len(self._capture_axes): + raise ValueError( + f"Model returned {len(axis_keys)} capture axis keys, " + f"expected {len(self._capture_axes)} " + f"(one per axis of capture_axes)." + ) + return selected, axis_keys def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]: """Get per-item output token counts as plain ints.""" @@ -334,6 +403,7 @@ class EncoderCudaGraphManager: mm_kwargs: dict[str, Any], token_budget: int, path: str = "default", + axis_keys: CaptureAxisKeys = (), ) -> torch.Tensor | None: """Execute budget graph. @@ -347,11 +417,14 @@ class EncoderCudaGraphManager: graph_set = self._get_graph_set(path) num_items = len(self._get_item_specs(mm_kwargs)) - if token_budget not in graph_set: + graph_map_key: BudgetGraphMapKey = ( + token_budget if not axis_keys else (token_budget, axis_keys) + ) + if graph_map_key not in graph_set: self.graph_misses += num_items return None - graph_meta = graph_set[token_budget] + graph_meta = graph_set[graph_map_key] replay = self.model.prepare_encoder_cudagraph_replay_buffers( mm_kwargs, @@ -432,7 +505,7 @@ class EncoderCudaGraphManager: outputs_by_orig_idx: dict[int, torch.Tensor] = {} for batch_indices, path_budgets in batches: - batch_mm_kwargs = self._select_items(mm_kwargs, batch_indices) + batch_mm_kwargs, axis_keys = self._select_items(mm_kwargs, batch_indices) graph_outputs: dict[str, torch.Tensor] = {} all_eager = True @@ -449,10 +522,14 @@ class EncoderCudaGraphManager: ) else: all_eager = False - output = self._run_budget_graph( - batch_mm_kwargs, token_budget, path=path + graph_output = self._run_budget_graph( + batch_mm_kwargs, + token_budget, + path=path, + axis_keys=axis_keys, ) - assert output is not None + assert graph_output is not None + output = graph_output graph_outputs[path] = output if all_eager: @@ -506,9 +583,9 @@ class EncoderCudaGraphManager: ] if len(local_indices) > 0: - local_mm_kwargs = self._select_items(mm_kwargs, local_indices) + local_mm_kwargs = self._select_items(mm_kwargs, local_indices)[0] else: - local_mm_kwargs = self._select_items(mm_kwargs, []) + local_mm_kwargs = self._select_items(mm_kwargs, [])[0] max_output_tokens_per_rank = ( max( diff --git a/vllm/v1/worker/encoder_cudagraph_defs.py b/vllm/v1/worker/encoder_cudagraph_defs.py index f86f43db8e..d19a4133fe 100644 --- a/vllm/v1/worker/encoder_cudagraph_defs.py +++ b/vllm/v1/worker/encoder_cudagraph_defs.py @@ -2,13 +2,19 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Data transfer objects for encoder CUDA graph management.""" -from collections.abc import Callable +from collections.abc import Callable, Hashable from dataclasses import dataclass, field import torch EncoderCudaGraphPaddingLogic = Callable[[torch.Tensor, torch.Tensor], None] +# Reserved mm_kwargs key: models with ``EncoderCudaGraphConfig.capture_axes`` +# put the resolved per-axis keys (one per axis, in order) into the dict +# returned by ``select_encoder_cudagraph_items()``; the manager pops it before +# the kwargs are used anywhere else. +ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG = "encoder_cudagraph_axis_keys" + @dataclass class EncoderItemSpec: @@ -88,6 +94,14 @@ class EncoderCudaGraphConfig: ) """Independently captured encoder paths keyed by their forward name.""" + capture_axes: tuple[tuple[Hashable, ...], ...] = () + """Extra capture axes beyond the token budget; empty to disable. + + Each entry is the ordered key set of one axis. When non-empty, one graph + is captured per token budget per combination of axis keys (cartesian + product), so the total number of captured graphs is + ``num_budgets * prod(len(axis) for axis in capture_axes)``.""" + @dataclass class EncoderCudaGraphCaptureInputs: