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: