@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 442d36031c

[MM][CG] Enable encoder CUDA Graph for MiniCPM-V (#42785)

Signed-off-by: YunzhuLu <lucia.yunzhu@gmail.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>

Yunzhu Lu committed Sep 10, 2026 at 14:28 UTC 442d36031ce710aa2a777c353becb660c57ab2bd
26 files changed +843 -103
docs/design/cuda_graphs_multimodal.md
+3
@@ -35,6 +35,8 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili
35 | `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ |
36 | `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ |
37 | `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ |
38 +| `MiniCPMV` | `MiniCPMV2.5` | ✅︎ | ❌︎ | ❌︎ |
39 +| `MiniCPMV` | `MiniCPMV2.6`, `MiniCPMV4.0` | ✅︎ | ✅︎ | ❌︎ |
40
41 ### Model x Hardware
42
@@ -54,6 +56,7 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili
56 | `Qwen3_5ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
57 | `Qwen3_5MoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
58 | `Step3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ |
59 +| `MiniCPMV` | ✅︎ | ✅︎ | ❔ | ❔ |
60
61 !!! note
62 Encoder CUDA Graph has currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
examples/generate/multimodal/vision_language_offline.py
+3
@@ -2276,6 +2276,9 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2276 "glm4_1v",
2277 "deepseek_ocr",
2278 "ernie45_vl",
2279 + "minicpmv2_5_vl",
2280 + "minicpmv2_6_vl",
2281 + "minicpmv4_vl",
2282 ]
2283
2284
tests/models/multimodal/generation/test_vit_cudagraph.py
+65
@@ -74,6 +74,20 @@ def ernie45_vl_chat_template(content: str) -> str:
74 )
75
76
77 +def minicpmv_25_chat_template(content: str) -> str:
78 + """Llama3-style chat template used by MiniCPM-V 2.5."""
79 + return (
80 + f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n"
81 + f"{content}"
82 + f"<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
83 + )
84 +
85 +
86 +def minicpmv_chat_template(content: str) -> str:
87 + """ChatML template used by MiniCPM-V 2.6 / 4.0 / 4.5."""
88 + return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
89 +
90 +
91 MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
92 "gemma3": VitCudagraphTestConfig(
93 model="google/gemma-3-4b-it",
@@ -316,6 +330,57 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
330 needs_video_metadata=True,
331 marks=[pytest.mark.core_model],
332 ),
333 + "minicpmv_25": VitCudagraphTestConfig(
334 + model="openbmb/MiniCPM-Llama3-V-2_5",
335 + modalities=["image"],
336 + image_prompt=minicpmv_25_chat_template(
337 + "(<image>./</image>)\nWhat is in this image?"
338 + ),
339 + # CI runs on 35GB MIG slices: every budget captures one graph per
340 + # patch-grid bucket, so the default budgets through max_model_len
341 + # OOM there. A small budget set covers the test images (each item
342 + # is at most (max_slice_num + 1) * query_num = 640 tokens).
343 + compilation_config_overrides={
344 + "encoder_cudagraph_token_budgets": [64, 1024],
345 + },
346 + vllm_runner_kwargs={"trust_remote_code": True},
347 + marks=[pytest.mark.core_model],
348 + ),
349 + "minicpmv_26": VitCudagraphTestConfig(
350 + model="openbmb/MiniCPM-V-2_6",
351 + image_prompt=minicpmv_chat_template(
352 + "(<image>./</image>)\nWhat is in this image?"
353 + ),
354 + video_prompt=minicpmv_chat_template(
355 + "(<video>./</video>)\nDescribe this video in one sentence."
356 + ),
357 + max_model_len=2048,
358 + # Fewer frames keep the video item within the 1024 budget.
359 + num_video_frames=4,
360 + compilation_config_overrides={
361 + "encoder_cudagraph_token_budgets": [64, 1024],
362 + "encoder_cudagraph_max_frames_per_batch": 4,
363 + },
364 + vllm_runner_kwargs={"trust_remote_code": True},
365 + marks=[pytest.mark.core_model],
366 + ),
367 + "minicpmv_40": VitCudagraphTestConfig(
368 + model="openbmb/MiniCPM-V-4",
369 + image_prompt=minicpmv_chat_template(
370 + "(<image>./</image>)\nWhat is in this image?"
371 + ),
372 + video_prompt=minicpmv_chat_template(
373 + "(<video>./</video>)\nDescribe this video in one sentence."
374 + ),
375 + max_model_len=2048,
376 + num_video_frames=4,
377 + compilation_config_overrides={
378 + "encoder_cudagraph_token_budgets": [64, 1024],
379 + "encoder_cudagraph_max_frames_per_batch": 4,
380 + },
381 + vllm_runner_kwargs={"trust_remote_code": True},
382 + marks=[pytest.mark.core_model],
383 + ),
384 }
385
386
tests/v1/cudagraph/test_encoder_cudagraph.py
+10
@@ -12,6 +12,7 @@ Test organization:
12 - TestEncoderCudaGraphVideoReplay — video modality capture, replay
13 """
14
15 +from collections.abc import Hashable
16 from typing import Any
17
18 import pytest
@@ -116,9 +117,15 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
117 )
118 mgr.budget_graphs = {"default": {}}
119 mgr.graph_pool = None
120 + mgr._capture_axes = ()
121 mgr.graph_hits = 0
122 mgr.graph_misses = 0
123 mgr.log_stats_interval = 100
124 + mgr.config = EncoderCudaGraphConfig(
125 + modalities=["image"],
126 + buffer_keys=[],
127 + out_hidden_size=32,
128 + )
129 return mgr
130
131
@@ -348,6 +355,7 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
355 device: torch.device,
356 dtype: torch.dtype,
357 path: str = "default",
358 + axis_keys: tuple[Hashable, ...] | None = None,
359 ) -> EncoderCudaGraphCaptureInputs:
360 per_image_output = token_budget // max_batch_size
361 grid_config = [
@@ -426,6 +434,7 @@ def _make_manager_for_gpu(
434 mgr.use_dp = False
435 mgr.budget_graphs = {"default": {}}
436 mgr.graph_pool = None
437 + mgr._capture_axes = ()
438 mgr.graph_hits = 0
439 mgr.graph_misses = 0
440 mgr.log_stats_interval = 100
@@ -656,6 +665,7 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
665 device: torch.device,
666 dtype: torch.dtype,
667 path: str = "default",
668 + axis_keys: tuple[Hashable, ...] | None = None,
669 ) -> EncoderCudaGraphCaptureInputs:
670 per_item_output = token_budget // max_batch_size
671 frames_per_item = max_frames_per_batch // max_batch_size
vllm/model_executor/models/cohere_compass.py
+2 -1
@@ -24,7 +24,7 @@
24 """Inference-only Cohere Compass model compatible with HuggingFace weights."""
25
26 import math
27 -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
27 +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
28 from functools import lru_cache, partial
29 from itertools import islice
30 from typing import Annotated, Any, Literal, TypeAlias
@@ -1598,6 +1598,7 @@ class CohereCompassForConditionalGeneration(
1598 device: torch.device,
1599 dtype: torch.dtype,
1600 path: str = "default",
1601 + axis_keys: tuple[Hashable, ...] | None = None,
1602 ):
1603 from vllm.v1.worker.encoder_cudagraph_defs import (
1604 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/deepseek_ocr.py
+2 -1
@@ -3,7 +3,7 @@
3 """Inference-only Deepseek-OCR model compatible with HuggingFace weights."""
4
5 import math
6 -from collections.abc import Iterable, Mapping, Sequence
6 +from collections.abc import Hashable, Iterable, Mapping, Sequence
7 from typing import Annotated, Any, Literal
8
9 import torch
@@ -804,6 +804,7 @@ class DeepseekOCRForCausalLM(
804 device: torch.device,
805 dtype: torch.dtype,
806 path: str = "default",
807 + axis_keys: tuple[Hashable, ...] | None = None,
808 ):
809 assert path in ("global", "local")
810
vllm/model_executor/models/ernie45_vl.py
+2 -1
@@ -24,7 +24,7 @@
24 """Inference-only Ernie VL model compatible with HuggingFace weights."""
25
26 import math
27 -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
27 +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
28 from functools import partial
29 from typing import Annotated, Any, Literal
30
@@ -1681,6 +1681,7 @@ class Ernie4_5_VLMoeForConditionalGeneration(
1681 device,
1682 dtype,
1683 path: str = "default",
1684 + axis_keys: tuple[Hashable, ...] | None = None,
1685 ):
1686 from vllm.v1.worker.encoder_cudagraph_defs import (
1687 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/gemma3_mm.py
+2 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import math
4 -from collections.abc import Iterable, Mapping, Sequence
4 +from collections.abc import Hashable, Iterable, Mapping, Sequence
5 from typing import Annotated, Any, Literal
6
7 import torch
@@ -798,6 +798,7 @@ class Gemma3ForConditionalGeneration(
798 device: torch.device,
799 dtype: torch.dtype,
800 path: str = "default",
801 + axis_keys: tuple[Hashable, ...] | None = None,
802 ):
803 from vllm.v1.worker.encoder_cudagraph_defs import (
804 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/gemma4_mm.py
+2 -1
@@ -15,7 +15,7 @@ reason about temporal order.
15 """
16
17 import math
18 -from collections.abc import Iterable, Mapping, Sequence
18 +from collections.abc import Hashable, Iterable, Mapping, Sequence
19 from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal
20
21 import numpy as np
@@ -1812,6 +1812,7 @@ class Gemma4ForConditionalGeneration(
1812 device: torch.device | str = "cpu",
1813 dtype: torch.dtype | None = None,
1814 path: str = "default",
1815 + axis_keys: tuple[Hashable, ...] | None = None,
1816 **kwargs: Any,
1817 ) -> "EncoderCudaGraphCaptureInputs":
1818 from vllm.v1.worker.encoder_cudagraph_defs import (
vllm/model_executor/models/glm4_1v.py
+2 -1
@@ -28,7 +28,7 @@
28 compatible with HuggingFace weights."""
29
30 import math
31 -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
31 +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
32 from functools import partial
33 from typing import Annotated, Any, Literal, TypeAlias
34
@@ -2077,6 +2077,7 @@ class Glm4vForConditionalGeneration(
2077 device: torch.device,
2078 dtype: torch.dtype,
2079 path: str = "default",
2080 + axis_keys: tuple[Hashable, ...] | None = None,
2081 ):
2082 from vllm.v1.worker.encoder_cudagraph_defs import (
2083 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/idefics2_vision_model.py
+51 -38
@@ -38,8 +38,7 @@ from vllm.model_executor.layers.linear import (
38 RowParallelLinear,
39 )
40 from vllm.model_executor.layers.quantization import QuantizationConfig
41 -from vllm.utils.gpu_sync_debug import gpu_sync_allowed
42 -from vllm.utils.torch_utils import PIN_MEMORY
41 +from vllm.utils.torch_utils import async_tensor_h2d
42
43 from .utils import AutoWeightsLoader, WeightsMapper
44 from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model
@@ -81,41 +80,61 @@ class Idefics2VisionEmbeddings(nn.Module):
80 patch_attention_mask: torch.BoolTensor,
81 tgt_sizes: torch.IntTensor | None = None,
82 ) -> torch.Tensor:
84 - batch_size, max_nb_patches_h, max_nb_patches_w = patch_attention_mask.shape
83 + batch_size = patch_attention_mask.shape[0]
84 + device = patch_attention_mask.device
85 boundaries = torch.arange(
86 - 1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side
87 - )
88 - position_ids = torch.full(
89 - size=(batch_size, max_nb_patches_h * max_nb_patches_w),
90 - fill_value=0,
91 - pin_memory=PIN_MEMORY,
86 + 1 / self.num_patches_per_side,
87 + 1.0,
88 + 1 / self.num_patches_per_side,
89 + device=device,
90 )
91
94 - with gpu_sync_allowed():
95 - patch_attention_mask_cpu = patch_attention_mask.cpu()
96 - tgt_sizes_cpu = tgt_sizes.cpu() if tgt_sizes is not None else None
92 + flat_patch_attention_mask = patch_attention_mask.view(batch_size, -1)
93 + max_patches = flat_patch_attention_mask.shape[-1]
94
98 - for batch_idx, p_attn_mask in enumerate(patch_attention_mask_cpu):
99 - if tgt_sizes_cpu is not None:
100 - nb_patches_h = tgt_sizes_cpu[batch_idx][0]
101 - nb_patches_w = tgt_sizes_cpu[batch_idx][1]
95 + if tgt_sizes is not None:
96 + # tgt_sizes may come from CPU-side mm_kwargs; use a pinned async
97 + # copy to stay clean under VLLM_GPU_SYNC_CHECK.
98 + if tgt_sizes.device != device:
99 + tgt_sizes = async_tensor_h2d(tgt_sizes, device)
100 + nb_patches_h = tgt_sizes[:, 0]
101 + nb_patches_w = tgt_sizes[:, 1]
102 + else:
103 + if patch_attention_mask.dim() == 3:
104 + nb_patches_h = patch_attention_mask[:, :, 0].sum(dim=1)
105 + nb_patches_w = patch_attention_mask[:, 0, :].sum(dim=1)
106 else:
103 - nb_patches_h = p_attn_mask[:, 0].sum()
104 - nb_patches_w = p_attn_mask[0].sum()
105 - fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h)
106 - fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w)
107 - bucket_coords_h = torch.bucketize(
108 - fractional_coords_h, boundaries, right=True
109 - )
110 - bucket_coords_w = torch.bucketize(
111 - fractional_coords_w, boundaries, right=True
112 - )
113 - pos_ids = (
114 - bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w
115 - ).flatten()
116 - position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids
107 + # Fallback for unexpected 2D masks without tgt_sizes
108 + nb_patches_h = torch.ones(batch_size, device=device)
109 + nb_patches_w = patch_attention_mask.sum(dim=1)
110 +
111 + nb_h = nb_patches_h.clamp(min=1).unsqueeze(1)
112 + nb_w = nb_patches_w.clamp(min=1).unsqueeze(1)
113 +
114 + i = torch.arange(max_patches, device=device)
115 +
116 + if patch_attention_mask.dim() == 3 and patch_attention_mask.shape[1] > 1:
117 + stride_w = patch_attention_mask.shape[2] # padded 2D grid width
118 + else:
119 + stride_w = nb_w # (batch, 1) — per-item actual patch width
120 +
121 + h_idx = i.unsqueeze(0) // stride_w
122 + w_idx = i.unsqueeze(0) % stride_w
123 +
124 + fractional_h = h_idx / nb_h
125 + fractional_w = w_idx / nb_w
126 +
127 + bucket_h = torch.bucketize(fractional_h, boundaries, right=True)
128 + bucket_w = torch.bucketize(fractional_w, boundaries, right=True)
129 +
130 + pos_ids = bucket_h * self.num_patches_per_side + bucket_w
131 +
132 + position_ids = torch.where(flat_patch_attention_mask, pos_ids.to(torch.long), 0)
133
118 - return position_ids.to(self.position_embedding.weight.device, non_blocking=True)
134 + target_device = self.position_embedding.weight.device
135 + if position_ids.device != target_device:
136 + position_ids = async_tensor_h2d(position_ids, target_device)
137 + return position_ids
138
139 def forward(
140 self,
@@ -461,13 +480,7 @@ class Idefics2VisionTransformer(nn.Module):
480 # - if apply_encoder_attention_mask is False, skip (not all models
481 # sharing this encoder apply masking in attention, e.g. Aria, Phi4)
482 # - if patch_attention_mask was None, skip attention masking
464 - # - if any padding exists, create an additive 4D mask and pass it
465 - # to attention; else skip mask for performance.
466 - if (
467 - not self.apply_encoder_attention_mask
468 - or flat_patch_mask is None
469 - or not torch.any(~flat_patch_mask)
470 - ):
483 + if not self.apply_encoder_attention_mask or flat_patch_mask is None:
484 attention_mask = None
485 else:
486 # Additive mask: masked positions receive a large negative value.
vllm/model_executor/models/idefics3.py
+2 -1
@@ -16,7 +16,7 @@
16 # limitations under the License.
17 """Inference-only Idefics3 model compatible with HuggingFace weights."""
18
19 -from collections.abc import Iterable, Mapping, Sequence
19 +from collections.abc import Hashable, Iterable, Mapping, Sequence
20 from typing import Annotated, Any, Literal, TypeAlias
21
22 import torch
@@ -722,6 +722,7 @@ class Idefics3ForConditionalGeneration(
722 device: torch.device,
723 dtype: torch.dtype,
724 path: str = "default",
725 + axis_keys: tuple[Hashable, ...] | None = None,
726 ):
727 from vllm.v1.worker.encoder_cudagraph_defs import (
728 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/interfaces.py
+15 -1
@@ -7,6 +7,7 @@ from bisect import bisect_right
7 from collections.abc import (
8 AsyncGenerator,
9 Callable,
10 + Hashable,
11 Mapping,
12 MutableSequence,
13 Sequence,
@@ -1875,6 +1876,11 @@ class SupportsEncoderCudaGraph(Protocol):
1876 - Qwen-family: slice concatenated pixel_values by cumulative
1877 patch offsets, subset grid_thw by indices.
1878 - Batched models (CLIP): index pixel_values along dim 0.
1879 +
1880 + Models that configure ``EncoderCudaGraphConfig.capture_axes`` must
1881 + additionally store the resolved per-axis keys (one key per axis, in
1882 + order) under ``ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG`` in the returned
1883 + dict; the manager pops it before the kwargs are used elsewhere.
1884 """
1885 ...
1886
@@ -1909,8 +1915,16 @@ class SupportsEncoderCudaGraph(Protocol):
1915 device: torch.device,
1916 dtype: torch.dtype,
1917 path: str = "default",
1918 + axis_keys: tuple[Hashable, ...] | None = None,
1919 ) -> "EncoderCudaGraphCaptureInputs":
1913 - """Create dummy inputs and buffers for CUDA graph capture."""
1920 + """Create dummy inputs and buffers for CUDA graph capture.
1921 +
1922 + Args:
1923 + axis_keys: The resolved capture-axis keys (one per axis of
1924 + ``EncoderCudaGraphConfig.capture_axes``) this capture is for.
1925 + None or empty when no capture axes are configured; models
1926 + without capture axes ignore it.
1927 + """
1928 ...
1929
1930 def prepare_encoder_cudagraph_replay_buffers(
vllm/model_executor/models/internvl.py
+2 -1
@@ -8,7 +8,7 @@
8 # Licensed under The MIT License [see LICENSE for details]
9 # --------------------------------------------------------
10 from abc import abstractmethod
11 -from collections.abc import Iterable, Mapping, Sequence
11 +from collections.abc import Hashable, Iterable, Mapping, Sequence
12 from functools import cached_property
13 from typing import Annotated, Any, Literal, TypeAlias, TypedDict, TypeVar
14
@@ -1028,6 +1028,7 @@ class InternVLChatModel(
1028 device: torch.device,
1029 dtype: torch.dtype,
1030 path: str = "default",
1031 + axis_keys: tuple[Hashable, ...] | None = None,
1032 ):
1033 from vllm.v1.worker.encoder_cudagraph_defs import (
1034 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/kimi_k25.py
+2 -1
@@ -6,7 +6,7 @@ Kimi-K2.5 Model Implementation for vLLM.
6 Kimi-K2.5 extends Kimi-K2 with vision support.
7 """
8
9 -from collections.abc import Iterable, Mapping, Sequence
9 +from collections.abc import Hashable, Iterable, Mapping, Sequence
10 from dataclasses import dataclass
11 from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal
12
@@ -594,6 +594,7 @@ class KimiK25ForConditionalGeneration(
594 device: torch.device,
595 dtype: torch.dtype,
596 path: str = "default",
597 + axis_keys: tuple[Hashable, ...] | None = None,
598 ) -> "EncoderCudaGraphCaptureInputs":
599 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphCaptureInputs
600
vllm/model_executor/models/kimi_vl.py
+2 -1
@@ -43,7 +43,7 @@
43 # SOFTWARE.
44
45 import math
46 -from collections.abc import Iterable, Mapping, Sequence
46 +from collections.abc import Hashable, Iterable, Mapping, Sequence
47 from dataclasses import dataclass
48 from typing import Annotated, Any, Literal
49
@@ -450,6 +450,7 @@ class KimiVLForConditionalGeneration(
450 device: torch.device,
451 dtype: torch.dtype,
452 path: str = "default",
453 + axis_keys: tuple[Hashable, ...] | None = None,
454 ):
455 from vllm.v1.worker.encoder_cudagraph_defs import (
456 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/lfm2_vl.py
+2 -1
@@ -4,7 +4,7 @@
4 import itertools
5 import math
6 import typing
7 -from collections.abc import Iterable, Mapping, Sequence
7 +from collections.abc import Hashable, Iterable, Mapping, Sequence
8 from typing import Annotated, Any, Literal
9
10 import torch
@@ -1146,6 +1146,7 @@ class Lfm2VLForConditionalGeneration(
1146 device: torch.device,
1147 dtype: torch.dtype,
1148 path: str = "default",
1149 + axis_keys: tuple[Hashable, ...] | None = None,
1150 ):
1151 from vllm.v1.worker.encoder_cudagraph_defs import (
1152 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/minicpmv.py
+547 -23
@@ -26,10 +26,10 @@
26
27 import math
28 from collections import defaultdict
29 -from collections.abc import Callable, Iterable, Mapping, Sequence, Set
29 +from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence, Set
30 from functools import cached_property, partial
31 from itertools import chain
32 -from typing import Annotated, Any, Literal, TypeAlias
32 +from typing import Annotated, Any, ClassVar, Literal, TypeAlias, cast
33
34 import numpy as np
35 import torch
@@ -100,11 +100,19 @@ from vllm.transformers_utils.utils import convert_model_repo_to_path
100 from vllm.utils.collection_utils import flatten_2d_lists, is_list_of
101 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
102 from vllm.utils.tensor_schema import TensorSchema, TensorShape
103 -from vllm.utils.torch_utils import set_default_torch_dtype
103 +from vllm.utils.torch_utils import async_tensor_h2d, set_default_torch_dtype
104 +from vllm.v1.worker.encoder_cudagraph_defs import (
105 + ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG,
106 + EncoderCudaGraphCaptureInputs,
107 + EncoderCudaGraphConfig,
108 + EncoderCudaGraphReplayBuffers,
109 + EncoderItemSpec,
110 +)
111
112 from .idefics2_vision_model import Idefics2VisionTransformer
113 from .interfaces import (
114 MultiModalEmbeddings,
115 + SupportsEncoderCudaGraph,
116 SupportsLoRA,
117 SupportsMultiModal,
118 SupportsPP,
@@ -218,29 +226,36 @@ class Resampler2_5(BaseResampler):
226 device = x.device
227 dtype = x.dtype
228
229 + # Only adjust the pos cache in eager mode; during CUDA graph capture
230 + # the cache was already expanded by the preceding warmup run
231 + if not torch.cuda.is_current_stream_capturing():
232 + self._adjust_pos_cache(tgt_sizes, device=device)
233 +
234 + # tgt_sizes is pinned-copied only when it comes from the CPU-side
235 + # mm_kwargs; during capture/replay it is already a GPU buffer.
236 + if tgt_sizes.device != x.device:
237 + tgt_sizes = async_tensor_h2d(tgt_sizes, x.device)
238 patch_len = tgt_sizes[:, 0] * tgt_sizes[:, 1]
239 + max_patch_len = x.shape[1]
240
223 - self._adjust_pos_cache(tgt_sizes, device=device)
241 + seq_idx = torch.arange(max_patch_len, device=device).unsqueeze(0)
242
225 - max_patch_len = patch_len.max().item()
226 - assert isinstance(max_patch_len, int)
243 + tgt_w = tgt_sizes[:, 1].unsqueeze(1)
244 + tgt_w = torch.clamp(tgt_w, min=1)
245
228 - key_padding_mask = torch.zeros(
229 - (bs, max_patch_len), dtype=torch.bool, device=device
230 - )
246 + h_idx = seq_idx // tgt_w
247 + w_idx = seq_idx % tgt_w
248
232 - pos_embed = []
233 - for i in range(bs):
234 - tgt_h, tgt_w = tgt_sizes[i].tolist()
235 - pos_embed.append(
236 - self.pos_embed[:tgt_h, :tgt_w, :]
237 - .reshape((tgt_h * tgt_w, -1))
238 - .to(device=device, dtype=dtype)
239 - ) # patches * D
240 - key_padding_mask[i, patch_len[i] :] = True
241 - pos_embed = torch.nn.utils.rnn.pad_sequence(
242 - pos_embed, batch_first=True, padding_value=0.0
249 + h_idx = torch.clamp(h_idx, max=self.pos_embed.shape[0] - 1)
250 + w_idx = torch.clamp(w_idx, max=self.pos_embed.shape[1] - 1)
251 +
252 + key_padding_mask = seq_idx >= patch_len.unsqueeze(1)
253 +
254 + pos_embed = self.pos_embed[h_idx, w_idx].to(dtype)
255 + pos_embed = torch.where(
256 + key_padding_mask.unsqueeze(-1), torch.zeros_like(pos_embed), pos_embed
257 ).permute(1, 0, 2) # BLD => L * B * D
258 +
259 x, _ = self.kv_proj(x) # B * L * D
260 x = self.ln_kv(x).permute(1, 0, 2) # L * B * D
261
@@ -808,6 +823,16 @@ class MiniCPMVProcessingInfo(BaseProcessingInfo):
823
824 return max(max_frames_per_video, 1)
825
826 + def get_mm_max_tokens_per_item(
827 + self,
828 + seq_len: int,
829 + mm_counts: Mapping[str, int],
830 + ) -> Mapping[str, int]:
831 + max_tokens = {"image": self.get_max_image_tokens()}
832 + if "video" in self.get_supported_mm_limits():
833 + max_tokens["video"] = self.get_max_video_tokens(seq_len, mm_counts)
834 + return max_tokens
835 +
836
837 _I = TypeVar("_I", bound=MiniCPMVProcessingInfo, default=MiniCPMVProcessingInfo)
838
@@ -1223,6 +1248,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):
1248 # and config class
1249 self.config = config
1250 self.multimodal_config = multimodal_config
1251 + self.vllm_config = vllm_config
1252
1253 self.version = get_version_by_config(self.config)
1254
@@ -1428,6 +1454,504 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):
1454 raise NotImplementedError
1455
1456
1457 +# --- Encoder CUDA graph (MiniCPM-V 2.5 / 2.6 / 4.0; mixed into subclasses) ---
1458 +# Buffer keys
1459 +_MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES = "minicpmv_tgt_sizes"
1460 +_MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK = "minicpmv_patch_attn_mask"
1461 +_MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL = "minicpmv_pixel_values"
1462 +
1463 +# mm_kwargs keys for the flat pixel-value tensor
1464 +_MINICPMV_CUDAGRAPH_FLAT_KEY_IMAGE = "minicpmv_encoder_input_flat"
1465 +_MINICPMV_CUDAGRAPH_FLAT_KEY_VIDEO = "minicpmv_video_encoder_input_flat"
1466 +# Resolved patch-grid capture-axis key ``(nb_h, nb_w)`` of the sub-batch;
1467 +# pixel dims and max patches are derived from it and the patch size.
1468 +_MINICPMV_CUDAGRAPH_PATCH_GRID_KEY = "minicpmv_encoder_patch_grid"
1469 +
1470 +_MINICPMV_BASE_PATCH_BUCKETS: tuple[tuple[int, int], ...] = (
1471 + (32, 32),
1472 + (32, 64),
1473 + (64, 32),
1474 +)
1475 +
1476 +_ENCODER_CUDAGRAPH_MM_KWARGS_SKIP_KEYS = frozenset(
1477 + {
1478 + _MINICPMV_CUDAGRAPH_FLAT_KEY_IMAGE,
1479 + _MINICPMV_CUDAGRAPH_FLAT_KEY_VIDEO,
1480 + _MINICPMV_CUDAGRAPH_PATCH_GRID_KEY,
1481 + }
1482 +)
1483 +
1484 +
1485 +def _mcpmv_tgt_sizes_tensor(
1486 + mm_kwargs: dict[str, Any], *, video: bool
1487 +) -> torch.Tensor | list[torch.Tensor]:
1488 + key = "video_tgt_sizes" if video else "tgt_sizes"
1489 + return mm_kwargs[key]
1490 +
1491 +
1492 +def _mcpmv_normalize_tgt_sizes(
1493 + tgt_sizes: torch.Tensor | list[torch.Tensor],
1494 + slice_counts: list[int],
1495 +) -> torch.Tensor:
1496 + """Normalize tgt_sizes to shape ``(total_slices, 2)``."""
1497 +
1498 + if isinstance(tgt_sizes, list):
1499 + if not tgt_sizes:
1500 + tgt_sizes = torch.zeros((0, 2), dtype=torch.long)
1501 + else:
1502 + tgt_sizes = torch.cat(tgt_sizes, dim=0)
1503 +
1504 + total_slices = sum(slice_counts)
1505 + if tgt_sizes.dim() == 2 and tgt_sizes.shape[0] == total_slices:
1506 + return tgt_sizes
1507 + if tgt_sizes.dim() == 3:
1508 + return torch.cat(
1509 + [tgt_sizes[i, : slice_counts[i], :] for i in range(len(slice_counts))],
1510 + dim=0,
1511 + )
1512 + return torch.repeat_interleave(
1513 + tgt_sizes,
1514 + torch.tensor(slice_counts, device=tgt_sizes.device),
1515 + dim=0,
1516 + )
1517 +
1518 +
1519 +def _mcpmv_pack_flat_pixels(
1520 + slices: list[torch.Tensor],
1521 + *,
1522 + pixel_height: int,
1523 + pixel_width: int,
1524 + max_num_slices: int,
1525 + device: torch.device,
1526 + dtype: torch.dtype,
1527 + patch_size: int,
1528 + tgt_sizes: torch.Tensor,
1529 +) -> torch.Tensor:
1530 + """Pack slice tensors into a fixed ``(max_num_slices, 3*H*W)`` buffer."""
1531 + import torch.nn.functional as F
1532 +
1533 + flat_dim = 3 * pixel_height * pixel_width
1534 + grid_h = pixel_height // patch_size
1535 + grid_w = pixel_width // patch_size
1536 + max_grid_patches = grid_h * grid_w
1537 +
1538 + packed = torch.zeros((max_num_slices, flat_dim), device=device, dtype=dtype)
1539 + n = min(len(slices), max_num_slices)
1540 + tgt_list = tgt_sizes.tolist()
1541 + for i, slc in enumerate(slices[:n]):
1542 + slc = slc.to(dtype=dtype, device=device)
1543 + C, h, w = slc.shape
1544 + if h == patch_size:
1545 + nb_h, nb_w = int(tgt_list[i][0]), int(tgt_list[i][1])
1546 + num_patches = nb_h * nb_w
1547 +
1548 + patches = (
1549 + slc.view(C, patch_size, nb_h, nb_w, patch_size)
1550 + .permute(0, 2, 3, 1, 4)
1551 + .reshape(C, num_patches, patch_size, patch_size)
1552 + )
1553 +
1554 + col = patches.permute(0, 2, 3, 1).reshape(
1555 + 1, C * patch_size * patch_size, num_patches
1556 + )
1557 + col_padded = col.new_zeros(1, C * patch_size * patch_size, max_grid_patches)
1558 + col_padded[..., :num_patches] = col
1559 +
1560 + folded = F.fold(
1561 + col_padded,
1562 + output_size=(pixel_height, pixel_width),
1563 + kernel_size=patch_size,
1564 + stride=patch_size,
1565 + )
1566 + packed[i] = folded.reshape(-1)
1567 + else:
1568 + packed[i].view(3, pixel_height, pixel_width)[:, :h, :w] = slc
1569 + return packed
1570 +
1571 +
1572 +class _MiniCPMVEncoderCudaGraphMixin(MiniCPMVBaseModel, SupportsEncoderCudaGraph):
1573 + """SupportsEncoderCudaGraph for MiniCPM-V Idefics2 + resampler (not 2.0)."""
1574 +
1575 + supports_encoder_cudagraph: ClassVar[Literal[True]] = True
1576 +
1577 + def _mcpmv_slice_pixel_size(self) -> tuple[int, int]:
1578 + image_size = int(self.vpm.embeddings.image_size)
1579 + return image_size, image_size
1580 +
1581 + def _mcpmv_patch_grid_pixel_hw(
1582 + self, patch_grid_key: tuple[int, int]
1583 + ) -> tuple[int, int]:
1584 + patch_size = int(self.vpm.embeddings.patch_size)
1585 + th, tw = patch_grid_key
1586 + return th * patch_size, tw * patch_size
1587 +
1588 + def _mcpmv_patch_grid_num_patches(self, patch_grid_key: tuple[int, int]) -> int:
1589 + th, tw = patch_grid_key
1590 + return th * tw
1591 +
1592 + def _mcpmv_patch_grid_keys(self) -> tuple[tuple[int, int], ...]:
1593 + """Ordered patch-grid keys ``(nb_h, nb_w)`` for the capture axis."""
1594 + max_side = int(self.vpm.embeddings.image_size) // int(
1595 + self.vpm.embeddings.patch_size
1596 + )
1597 + keys: list[tuple[int, int]] = []
1598 + seen: set[tuple[int, int]] = set()
1599 + for th, tw in _MINICPMV_BASE_PATCH_BUCKETS:
1600 + if th <= max_side and tw <= max_side:
1601 + keys.append((th, tw))
1602 + seen.add((th, tw))
1603 + full = (max_side, max_side)
1604 + if full not in seen:
1605 + keys.append(full)
1606 + return tuple(keys)
1607 +
1608 + def _mcpmv_resolve_patch_grid(
1609 + self,
1610 + tgt_sizes: torch.Tensor,
1611 + slice_counts: list[int],
1612 + indices: list[int],
1613 + ) -> tuple[int, int]:
1614 + """Smallest patch-grid key covering all selected items."""
1615 + keys = self._mcpmv_patch_grid_keys()
1616 + if not indices:
1617 + return keys[0]
1618 + tgt_groups = torch.split(tgt_sizes, slice_counts)
1619 + selected = torch.cat([tgt_groups[i] for i in indices], dim=0)
1620 + need_h = int(selected[:, 0].max().item())
1621 + need_w = int(selected[:, 1].max().item())
1622 + for th, tw in keys:
1623 + if th >= need_h and tw >= need_w:
1624 + return (th, tw)
1625 + return keys[-1]
1626 +
1627 + def _mcpmv_max_slices_cap(
1628 + self,
1629 + token_budget: int,
1630 + max_batch_size: int,
1631 + max_frames_per_batch: int,
1632 + ) -> int:
1633 + max_slice_num = int(getattr(self.config, "max_slice_num", 9))
1634 + query_num = max(1, int(self.config.query_num))
1635 + max_slices_by_token_budget = max(1, token_budget // query_num)
1636 + max_slices_by_content = max_batch_size * (max_slice_num + 1)
1637 + if self.version in {(2, 6), (4, 0)} and max_frames_per_batch > 0:
1638 + max_slices_by_content = max(
1639 + max_slices_by_content,
1640 + max_frames_per_batch * (max_slice_num + 1),
1641 + )
1642 + return max(1, min(max_slices_by_token_budget, max_slices_by_content))
1643 +
1644 + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
1645 + buffer_keys = [
1646 + _MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL,
1647 + _MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES,
1648 + _MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK,
1649 + ]
1650 + # Video is only supported from 2.6 onward.
1651 + modalities = ["image"]
1652 + if self.version in {(2, 6), (4, 0)}:
1653 + modalities.append("video")
1654 +
1655 + max_frames = self.get_max_frames_per_video() if "video" in modalities else 1
1656 +
1657 + return EncoderCudaGraphConfig(
1658 + modalities=modalities,
1659 + buffer_keys=buffer_keys,
1660 + out_hidden_size=int(self.embed_dim),
1661 + max_frames_per_video=max_frames,
1662 + capture_axes=(self._mcpmv_patch_grid_keys(),),
1663 + )
1664 +
1665 + def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:
1666 + if "video_pixel_values" in mm_kwargs:
1667 + return "video"
1668 + return "image"
1669 +
1670 + def get_max_frames_per_video(self) -> int:
1671 + info = MULTIMODAL_REGISTRY.get_processing_info(self.vllm_config.model_config)
1672 + assert isinstance(info, MiniCPMVProcessingInfo)
1673 + return int(
1674 + info.get_num_frames_with_most_features(
1675 + seq_len=self.vllm_config.model_config.max_model_len,
1676 + mm_counts={
1677 + "video": self.multimodal_config.get_limit_per_prompt("video")
1678 + },
1679 + )
1680 + )
1681 +
1682 + def get_encoder_cudagraph_budget_range(
1683 + self, vllm_config: VllmConfig
1684 + ) -> tuple[int, int]:
1685 + # Each slice produces exactly query_num resampler output tokens.
1686 + # A thumbnail-only image has 1 slice, so query_num is the smallest
1687 + # possible encoder output and the natural minimum budget.
1688 + min_budget = int(self.config.query_num)
1689 + max_budget = min(
1690 + vllm_config.scheduler_config.max_num_batched_tokens,
1691 + vllm_config.model_config.max_model_len,
1692 + )
1693 + return (min_budget, max_budget)
1694 +
1695 + def get_encoder_cudagraph_item_specs(
1696 + self, mm_kwargs: dict[str, Any]
1697 + ) -> list[EncoderItemSpec]:
1698 + video = self.get_input_modality(mm_kwargs) == "video"
1699 + pixel_values_key = "video_pixel_values" if video else "pixel_values"
1700 + pixel_values: list[list[torch.Tensor]] = mm_kwargs[pixel_values_key]
1701 + slice_counts = [len(img) for img in pixel_values]
1702 + tgt_sizes = _mcpmv_tgt_sizes_tensor(mm_kwargs, video=video)
1703 + tgt_sizes = _mcpmv_normalize_tgt_sizes(tgt_sizes, slice_counts)
1704 + patch_sums = tgt_sizes.prod(-1)
1705 + input_sizes = [
1706 + int(group.sum().item()) for group in torch.split(patch_sums, slice_counts)
1707 + ]
1708 + query_num = int(self.config.query_num)
1709 + return [
1710 + EncoderItemSpec(
1711 + input_size=input_sizes[i],
1712 + output_tokens=slice_counts[i] * query_num,
1713 + )
1714 + for i in range(len(pixel_values))
1715 + ]
1716 +
1717 + def select_encoder_cudagraph_items(
1718 + self,
1719 + mm_kwargs: dict[str, Any],
1720 + indices: list[int],
1721 + ) -> dict[str, Any]:
1722 + subset, patch_grid = self._mcpmv_select_items(mm_kwargs, indices)
1723 + subset[ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG] = (patch_grid,)
1724 + return subset
1725 +
1726 + def _mcpmv_select_items(
1727 + self, mm_kwargs: dict[str, Any], indices: list[int]
1728 + ) -> tuple[dict[str, Any], tuple[int, int]]:
1729 + """Slice mm_kwargs for `indices`; also returns the patch-grid key."""
1730 + video = self.get_input_modality(mm_kwargs) == "video"
1731 + pixel_values_key = "video_pixel_values" if video else "pixel_values"
1732 + tgt_key = "video_tgt_sizes" if video else "tgt_sizes"
1733 + flat_key = (
1734 + _MINICPMV_CUDAGRAPH_FLAT_KEY_VIDEO
1735 + if video
1736 + else _MINICPMV_CUDAGRAPH_FLAT_KEY_IMAGE
1737 + )
1738 + device = next(self.vpm.parameters()).device
1739 + pixel_values: list[list[torch.Tensor]] = mm_kwargs[pixel_values_key]
1740 + tgt_sizes = _mcpmv_tgt_sizes_tensor(mm_kwargs, video=video)
1741 +
1742 + subset = {
1743 + k: v
1744 + for k, v in mm_kwargs.items()
1745 + if k not in _ENCODER_CUDAGRAPH_MM_KWARGS_SKIP_KEYS
1746 + }
1747 +
1748 + if not indices:
1749 + pixel_h, pixel_w = self._mcpmv_slice_pixel_size()
1750 + vpm_dtype = next(self.vpm.parameters()).dtype
1751 + subset.update(
1752 + {
1753 + pixel_values_key: [],
1754 + tgt_key: torch.zeros((0, 2), dtype=torch.long, device=device),
1755 + flat_key: torch.zeros(
1756 + (0, 3 * pixel_h * pixel_w), device=device, dtype=vpm_dtype
1757 + ),
1758 + }
1759 + )
1760 + return subset, self._mcpmv_patch_grid_keys()[0]
1761 +
1762 + slice_counts = [len(item_slices) for item_slices in pixel_values]
1763 + tgt_sizes = _mcpmv_normalize_tgt_sizes(tgt_sizes, slice_counts)
1764 + tgt_groups = torch.split(tgt_sizes, slice_counts)
1765 +
1766 + patch_grid = self._mcpmv_resolve_patch_grid(tgt_sizes, slice_counts, indices)
1767 + pixel_h, pixel_w = self._mcpmv_patch_grid_pixel_hw(patch_grid)
1768 +
1769 + selected_pixel_values = [pixel_values[i] for i in indices]
1770 + selected_tgt_sizes_list = [tgt_groups[i] for i in indices]
1771 + selected_tgt_sizes = torch.cat(selected_tgt_sizes_list, dim=0)
1772 +
1773 + selected_slices = flatten_2d_lists(selected_pixel_values)
1774 + packed_flat_pixels = _mcpmv_pack_flat_pixels(
1775 + selected_slices,
1776 + pixel_height=pixel_h,
1777 + pixel_width=pixel_w,
1778 + max_num_slices=len(selected_slices),
1779 + device=selected_slices[0].device,
1780 + dtype=selected_slices[0].dtype,
1781 + patch_size=int(self.vpm.embeddings.patch_size),
1782 + tgt_sizes=selected_tgt_sizes,
1783 + )
1784 +
1785 + subset.update(
1786 + {
1787 + pixel_values_key: selected_pixel_values,
1788 + tgt_key: selected_tgt_sizes_list,
1789 + flat_key: packed_flat_pixels,
1790 + _MINICPMV_CUDAGRAPH_PATCH_GRID_KEY: patch_grid,
1791 + }
1792 + )
1793 + return subset, patch_grid
1794 +
1795 + def prepare_encoder_cudagraph_capture_inputs(
1796 + self,
1797 + token_budget: int,
1798 + max_batch_size: int,
1799 + max_frames_per_batch: int,
1800 + device: torch.device,
1801 + dtype: torch.dtype,
1802 + path: str = "default",
1803 + axis_keys: tuple[Hashable, ...] | None = None,
1804 + ):
1805 + patch_grid = (
1806 + cast("tuple[int, int]", axis_keys[0])
1807 + if axis_keys
1808 + # Without capture-axis context, use the largest (full-resolution)
1809 + # patch grid.
1810 + else self._mcpmv_patch_grid_keys()[-1]
1811 + )
1812 + return self._mcpmv_capture_inputs(
1813 + token_budget,
1814 + max_batch_size,
1815 + max_frames_per_batch,
1816 + device,
1817 + dtype,
1818 + patch_grid=patch_grid,
1819 + )
1820 +
1821 + def _mcpmv_capture_inputs(
1822 + self,
1823 + token_budget: int,
1824 + max_batch_size: int,
1825 + max_frames_per_batch: int,
1826 + device: torch.device,
1827 + dtype: torch.dtype,
1828 + patch_grid: tuple[int, int],
1829 + ):
1830 + th, tw = patch_grid
1831 + pixel_h, pixel_w = self._mcpmv_patch_grid_pixel_hw(patch_grid)
1832 + max_patches = self._mcpmv_patch_grid_num_patches(patch_grid)
1833 + max_num_slices = self._mcpmv_max_slices_cap(
1834 + token_budget,
1835 + max_batch_size,
1836 + max_frames_per_batch,
1837 + )
1838 + pixel_buffer = torch.zeros(
1839 + (max_num_slices, 3, pixel_h, pixel_w), device=device, dtype=dtype
1840 + )
1841 + dummy_tgt_sizes = torch.zeros(
1842 + (max_num_slices, 2), dtype=torch.long, device=device
1843 + )
1844 + dummy_tgt_sizes[:, 0] = th
1845 + dummy_tgt_sizes[:, 1] = tw
1846 + dummy_patch_mask = torch.ones(
1847 + (max_num_slices, max_patches), dtype=torch.bool, device=device
1848 + )
1849 + values: dict[str, torch.Tensor] = {
1850 + _MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL: pixel_buffer,
1851 + _MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES: dummy_tgt_sizes,
1852 + _MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK: dummy_patch_mask,
1853 + }
1854 + return EncoderCudaGraphCaptureInputs(values=values)
1855 +
1856 + def prepare_encoder_cudagraph_replay_buffers(
1857 + self,
1858 + mm_kwargs: dict[str, Any],
1859 + max_batch_size: int,
1860 + max_frames_per_batch: int,
1861 + path: str = "default",
1862 + ):
1863 + _ = max_batch_size
1864 + _ = max_frames_per_batch
1865 + video = self.get_input_modality(mm_kwargs) == "video"
1866 + flat_key = (
1867 + _MINICPMV_CUDAGRAPH_FLAT_KEY_VIDEO
1868 + if video
1869 + else _MINICPMV_CUDAGRAPH_FLAT_KEY_IMAGE
1870 + )
1871 + flat_pixels = mm_kwargs[flat_key] # (num_actual_slices, 3*pixel_h*pixel_w)
1872 + patch_grid = mm_kwargs[_MINICPMV_CUDAGRAPH_PATCH_GRID_KEY]
1873 + pixel_h, pixel_w = self._mcpmv_patch_grid_pixel_hw(patch_grid)
1874 + max_patches = self._mcpmv_patch_grid_num_patches(patch_grid)
1875 + pixel_buffer = flat_pixels.reshape(-1, 3, pixel_h, pixel_w)
1876 +
1877 + device = next(self.vpm.parameters()).device
1878 + tgt_sizes_raw = _mcpmv_tgt_sizes_tensor(mm_kwargs, video=video)
1879 + if isinstance(tgt_sizes_raw, list):
1880 + tgt_sizes_raw = torch.cat(tgt_sizes_raw, dim=0)
1881 + # tgt_sizes arrives from CPU-side mm_kwargs; use a pinned async copy
1882 + # to stay clean under VLLM_GPU_SYNC_CHECK.
1883 + tgt_sizes = async_tensor_h2d(tgt_sizes_raw, device, dtype=torch.long)
1884 +
1885 + patches_per_slice = tgt_sizes.prod(-1).clamp(max=max_patches)
1886 + col_idx = torch.arange(max_patches, device=device)
1887 + patch_attention_mask = col_idx.unsqueeze(0) < patches_per_slice.unsqueeze(1)
1888 +
1889 + values: dict[str, torch.Tensor] = {
1890 + _MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL: pixel_buffer,
1891 + _MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES: tgt_sizes,
1892 + _MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK: patch_attention_mask,
1893 + }
1894 + return EncoderCudaGraphReplayBuffers(values=values)
1895 +
1896 + def encoder_cudagraph_forward(
1897 + self,
1898 + values: dict[str, torch.Tensor],
1899 + path: str = "default",
1900 + ) -> torch.Tensor:
1901 + all_pixel_values = values[_MINICPMV_CUDAGRAPH_BUF_KEY_PIXEL]
1902 + tgt_sizes = values[_MINICPMV_CUDAGRAPH_BUF_KEY_TGT_SIZES]
1903 + patch_mask = values[_MINICPMV_CUDAGRAPH_BUF_KEY_PATCH_MASK]
1904 + patch_attention_mask = patch_mask.unsqueeze(1)
1905 +
1906 + max_num_slices = all_pixel_values.shape[0]
1907 + # v2.5 infers patch layout from the attention mask; pass tgt_sizes=None.
1908 + vpm_tgt_sizes = None if self.version == (2, 5) else tgt_sizes
1909 + vision_embedding = self.vpm(
1910 + all_pixel_values,
1911 + patch_attention_mask=patch_attention_mask,
1912 + tgt_sizes=vpm_tgt_sizes,
1913 + )
1914 +
1915 + resampler_out = self.resampler(vision_embedding, tgt_sizes)
1916 +
1917 + query_num = int(self.config.query_num)
1918 + return resampler_out.reshape(max_num_slices * query_num, int(self.embed_dim))
1919 +
1920 + def encoder_eager_forward(
1921 + self,
1922 + mm_kwargs: dict[str, Any],
1923 + path: str = "default",
1924 + ) -> torch.Tensor:
1925 + """Eager encoder path; returns ``(total_tokens, embed_dim)`` like
1926 + ``encoder_cudagraph_forward``.
1927 + """
1928 + mm_kwargs_no_flat = {
1929 + k: v
1930 + for k, v in mm_kwargs.items()
1931 + if k not in _ENCODER_CUDAGRAPH_MM_KWARGS_SKIP_KEYS
1932 + }
1933 + modalities = self._parse_and_validate_multimodal_inputs(**mm_kwargs_no_flat)
1934 + segments: list[torch.Tensor] = []
1935 + embed_dim = self.embed_dim
1936 + for modality in modalities:
1937 + if modality == "images":
1938 + image_input = modalities["images"]
1939 + assert isinstance(image_input, MiniCPMVImagePixelInputs)
1940 + image_embeddings = self.get_vision_hidden_states(image_input)
1941 + segments.append(image_embeddings.reshape(-1, embed_dim))
1942 + elif modality == "videos":
1943 + video_input = modalities["videos"]
1944 + assert isinstance(video_input, MiniCPMVImagePixelInputs)
1945 + video_embeddings = self.get_vision_hidden_states(video_input)
1946 + segments.append(video_embeddings.reshape(-1, embed_dim))
1947 + if not segments:
1948 + raise RuntimeError(
1949 + "MiniCPM-V encoder cudagraph eager path expects pixel_values "
1950 + "or video_pixel_values"
1951 + )
1952 + return torch.cat(segments, dim=0)
1953 +
1954 +
1955 class MiniCPMV2_0(MiniCPMVBaseModel):
1956 supports_encoder_tp_data = False
1957
@@ -1519,7 +2043,7 @@ class MiniCPMV2_0(MiniCPMVBaseModel):
2043 return torch.vstack(res)
2044
2045
1522 -class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):
2046 +class MiniCPMV2_5(_MiniCPMVEncoderCudaGraphMixin, SupportsLoRA):
2047 packed_modules_mapping = {
2048 "qkv_proj": [
2049 "q_proj",
@@ -1611,7 +2135,7 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):
2135 return self.resampler(vision_embedding, tgt_sizes)
2136
2137
1614 -class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):
2138 +class MiniCPMV2_6(_MiniCPMVEncoderCudaGraphMixin, SupportsLoRA):
2139 hf_to_vllm_mapper = WeightsMapper(
2140 # The vision-only models have no audio tower or TTS head to load weights into.
2141 orig_to_new_prefix={"apm.": None, "audio": None, "tts": None}
@@ -1715,7 +2239,7 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):
2239 return loaded
2240
2241
1718 -class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
2242 +class MiniCPMV4_0(_MiniCPMVEncoderCudaGraphMixin, SupportsLoRA):
2243 hf_to_vllm_mapper = MiniCPMV2_6.hf_to_vllm_mapper
2244 packed_modules_mapping = {
2245 "qkv_proj": [
vllm/model_executor/models/mllama4.py
+2 -1
@@ -17,7 +17,7 @@
17 # See the License for the specific language governing permissions and
18 # limitations under the License.
19 import math
20 -from collections.abc import Iterable, Mapping
20 +from collections.abc import Hashable, Iterable, Mapping
21 from itertools import tee
22 from typing import Annotated, Any, Literal
23
@@ -924,6 +924,7 @@ class Llama4ForConditionalGeneration(
924 device: torch.device,
925 dtype: torch.dtype,
926 path: str = "default",
927 + axis_keys: tuple[Hashable, ...] | None = None,
928 ):
929 from vllm.v1.worker.encoder_cudagraph_defs import (
930 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/qwen2_5_vl.py
+2 -1
@@ -26,7 +26,7 @@
26 # limitations under the License.
27 """Inference-only Qwen2.5-VL model compatible with HuggingFace weights."""
28
29 -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
29 +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
30 from functools import partial
31 from typing import Annotated, Any, Literal, TypeAlias
32
@@ -1856,6 +1856,7 @@ class Qwen2_5_VLForConditionalGeneration(
1856 device: torch.device,
1857 dtype: torch.dtype,
1858 path: str = "default",
1859 + axis_keys: tuple[Hashable, ...] | None = None,
1860 ):
1861 from vllm.v1.worker.encoder_cudagraph_defs import (
1862 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/qwen2_vl.py
+2 -1
@@ -26,7 +26,7 @@
26 """Inference-only Qwen2-VL model compatible with HuggingFace weights."""
27
28 import math
29 -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
29 +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
30 from functools import partial
31 from typing import Annotated, Any, Literal, TypeAlias
32
@@ -1589,6 +1589,7 @@ class Qwen2VLForConditionalGeneration(
1589 device: torch.device,
1590 dtype: torch.dtype,
1591 path: str = "default",
1592 + axis_keys: tuple[Hashable, ...] | None = None,
1593 ):
1594 from vllm.v1.worker.encoder_cudagraph_defs import (
1595 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/qwen3_vl.py
+2 -1
@@ -24,7 +24,7 @@
24 # limitations under the License.
25 """Inference-only Qwen3VL model compatible with HuggingFace weights."""
26
27 -from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
27 +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
28 from functools import lru_cache, partial
29 from itertools import islice
30 from typing import Any
@@ -2120,6 +2120,7 @@ class Qwen3VLForConditionalGeneration(
2120 device: torch.device,
2121 dtype: torch.dtype,
2122 path: str = "default",
2123 + axis_keys: tuple[Hashable, ...] | None = None,
2124 ):
2125 from vllm.v1.worker.encoder_cudagraph_defs import (
2126 EncoderCudaGraphCaptureInputs,
vllm/model_executor/models/step3_vl.py
+2 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import math
4 -from collections.abc import Iterable, Mapping, Sequence
4 +from collections.abc import Hashable, Iterable, Mapping, Sequence
5 from math import sqrt
6 from typing import Annotated, Any, Literal, TypeAlias
7
@@ -828,6 +828,7 @@ class Step3VLForConditionalGeneration(
828 device: torch.device,
829 dtype: torch.dtype,
830 path: str = "default",
831 + axis_keys: tuple[Hashable, ...] | None = None,
832 ):
833 from vllm.v1.worker.encoder_cudagraph_defs import (
834 EncoderCudaGraphCaptureInputs,
vllm/models/kimi_k3/nvidia/model.py
+2 -1
@@ -3,7 +3,7 @@
3 """Kimi-K3 multimodal model implementation for vLLM."""
4
5 import math
6 -from collections.abc import Iterable
6 +from collections.abc import Hashable, Iterable
7 from typing import Any, cast
8
9 import torch
@@ -1948,6 +1948,7 @@ class KimiK3ForConditionalGeneration(
1948 device: torch.device,
1949 dtype: torch.dtype,
1950 path: str = "default",
1951 + axis_keys: tuple[Hashable, ...] | None = None,
1952 ):
1953 from vllm.v1.worker.encoder_cudagraph_defs import (
1954 EncoderCudaGraphCaptureInputs,
vllm/v1/worker/encoder_cudagraph.py
+100 -23
@@ -2,8 +2,11 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """CUDA graph manager for vision encoder budget-batch execution."""
4
5 +import itertools
6 +import math
7 +from collections.abc import Hashable
8 from dataclasses import dataclass
6 -from typing import Any
9 +from typing import Any, TypeAlias
10
11 import torch
12
@@ -22,12 +25,20 @@ from vllm.model_executor.models.vision import get_load_balance_assignment
25 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
26 from vllm.utils.torch_utils import current_stream
27 from vllm.v1.worker.encoder_cudagraph_defs import (
28 + ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG,
29 EncoderCudaGraphConfig,
30 EncoderItemSpec,
31 )
32
33 logger = init_logger(__name__)
34
35 +CaptureAxisKeys: TypeAlias = tuple[Hashable, ...]
36 +
37 +BudgetGraphMapKey: TypeAlias = int | tuple[int, CaptureAxisKeys]
38 +
39 +# Warn when capture axes multiply the captured graph count past this point.
40 +_CAPTURE_GRAPH_COUNT_WARNING_THRESHOLD = 64
41 +
42
43 @dataclass
44 class BudgetGraphMetadata:
@@ -50,6 +61,8 @@ class BudgetGraphMetadata:
61 input_buffers: dict[str, torch.Tensor]
62 # Output written by graph, read after replay
63 output_buffer: torch.Tensor
64 + # Per-axis capture keys, empty when no extra capture axis is used.
65 + axis_keys: CaptureAxisKeys = ()
66
67
68 class EncoderCudaGraphManager:
@@ -77,6 +90,18 @@ class EncoderCudaGraphManager:
90
91 multimodal_config = vllm_config.model_config.multimodal_config
92
93 + self._capture_axes: tuple[CaptureAxisKeys, ...] = tuple(
94 + tuple(axis) for axis in self.config.capture_axes
95 + )
96 + if any(len(axis) == 0 for axis in self._capture_axes):
97 + raise ValueError("Encoder cudagraph capture axes must be non-empty.")
98 +
99 + if len(self.config.paths) > 1 and self._capture_axes:
100 + raise NotImplementedError(
101 + "Combining dual-path encoder CUDA graphs with capture axes "
102 + "is not supported yet."
103 + )
104 +
105 # Invariant: max_batch_size <= min_token_budget.
106 # This ensures per_image_output = budget // max_batch_size >= 1
107 # for every captured budget, preventing reshape crashes on empty
@@ -154,7 +179,7 @@ class EncoderCudaGraphManager:
179 and vllm_config.parallel_config.tensor_parallel_size > 1
180 )
181
157 - self.budget_graphs: dict[str, dict[int, BudgetGraphMetadata]] = {}
182 + self.budget_graphs: dict[str, dict[BudgetGraphMapKey, BudgetGraphMetadata]] = {}
183 self.graph_pool: Any | None = None
184 self.graph_hits = 0
185 self.graph_misses = 0
@@ -181,11 +206,13 @@ class EncoderCudaGraphManager:
206
207 logger.info(
208 "EncoderCudaGraphManager initialized with paths=%s, "
184 - "max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s",
209 + "max_batch_size=%d, max_frames_per_batch=%s, use_dp=%s, "
210 + "capture_axes=%s",
211 self.path_token_budgets,
212 self.max_batch_size,
213 self.max_frames_per_batch,
214 self.use_dp,
215 + self._capture_axes,
216 )
217
218 @staticmethod
@@ -219,38 +246,63 @@ class EncoderCudaGraphManager:
246 """Capture CUDA graphs for every configured path and token budget."""
247 self.graph_pool = graph_pool
248
249 + num_graphs = self.get_num_graphs_to_capture()
250 + if num_graphs > _CAPTURE_GRAPH_COUNT_WARNING_THRESHOLD:
251 + logger.warning(
252 + "Capturing %d encoder CUDA graphs; capture axes multiply the "
253 + "per-budget graph count, consider pruning axis keys.",
254 + num_graphs,
255 + )
256 +
257 for path, budgets in self.path_token_budgets.items():
258 for token_budget in sorted(budgets, reverse=True):
224 - if token_budget > 0:
225 - self._capture_budget_graph(token_budget, path=path)
259 + if token_budget == 0:
260 + continue
261 + for axis_keys in itertools.product(*self._capture_axes):
262 + self._capture_budget_graph(
263 + token_budget,
264 + path=path,
265 + axis_keys=axis_keys,
266 + )
267
268 logger.info(
269 "Encoder CUDA graph capture complete. Captured %d graphs across %d paths.",
229 - self.get_num_graphs_to_capture(),
270 + num_graphs,
271 len(self.path_token_budgets),
272 )
273
274 + def _num_graphs_for_budgets(self, budgets: list[int]) -> int:
275 + num_budgets = sum(1 for budget in budgets if budget != 0)
276 + return num_budgets * math.prod(len(axis) for axis in self._capture_axes)
277 +
278 def get_num_graphs_to_capture(self) -> int:
279 return sum(
235 - token_budget > 0
280 + self._num_graphs_for_budgets(budgets)
281 for budgets in self.path_token_budgets.values()
237 - for token_budget in budgets
282 )
283
240 - def _get_graph_set(self, path: str = "default") -> dict[int, BudgetGraphMetadata]:
284 + def _get_graph_set(
285 + self, path: str = "default"
286 + ) -> dict[BudgetGraphMapKey, BudgetGraphMetadata]:
287 """Return the captured graphs for one encoder path."""
288 if path not in self.budget_graphs:
289 self.budget_graphs[path] = {}
290 return self.budget_graphs[path]
291
246 - def _capture_budget_graph(self, token_budget: int, path: str = "default"):
292 + def _capture_budget_graph(
293 + self,
294 + token_budget: int,
295 + path: str = "default",
296 + axis_keys: CaptureAxisKeys = (),
297 + ):
298 """Capture CUDA graph for a single token budget."""
299 logger.debug(
300 "Capturing encoder cudagraph for budget=%d, max_batch_size=%d, "
250 - "max_frames_per_batch=%d",
301 + "max_frames_per_batch=%d, axis_keys=%s",
302 token_budget,
303 self.max_batch_size,
304 self.max_frames_per_batch,
305 + axis_keys,
306 )
307
308 graph_set = self._get_graph_set(path)
@@ -262,6 +314,7 @@ class EncoderCudaGraphManager:
314 self.device,
315 self.dtype,
316 path,
317 + axis_keys,
318 )
319
320 values = capture_inputs.values
@@ -278,13 +331,17 @@ class EncoderCudaGraphManager:
331 output = self.model.encoder_cudagraph_forward({**values}, path=path)
332 output_buffer.copy_(output)
333
281 - graph_set[token_budget] = BudgetGraphMetadata(
334 + graph_map_key: BudgetGraphMapKey = (
335 + token_budget if not axis_keys else (token_budget, axis_keys)
336 + )
337 + graph_set[graph_map_key] = BudgetGraphMetadata(
338 token_budget=token_budget,
339 max_batch_size=self.max_batch_size,
340 max_frames_per_batch=self.max_frames_per_batch,
341 graph=graph,
342 input_buffers=values,
343 output_buffer=output_buffer,
344 + axis_keys=axis_keys,
345 )
346
347 def _find_smallest_fitting_budget_given_tokens(
@@ -310,12 +367,24 @@ class EncoderCudaGraphManager:
367
368 def _select_items(
369 self, mm_kwargs: dict[str, Any], indices: list[int]
313 - ) -> dict[str, Any]:
314 - """Select the mm kwargs for `indices` from the model."""
370 + ) -> tuple[dict[str, Any], CaptureAxisKeys]:
371 + """Select the mm kwargs for `indices` from the model.
372 +
373 + Returns the sliced mm_kwargs together with the resolved capture-axis
374 + keys (popped out of the returned dict; empty without capture axes).
375 + """
376 # Same as `_get_item_specs`: implementations re-read the per-item
377 # grid/patch counts to slice the batch, so the D2H is inherent.
378 with gpu_sync_allowed():
318 - return self.model.select_encoder_cudagraph_items(mm_kwargs, indices)
379 + selected = self.model.select_encoder_cudagraph_items(mm_kwargs, indices)
380 + axis_keys = selected.pop(ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG, ())
381 + if len(axis_keys) != len(self._capture_axes):
382 + raise ValueError(
383 + f"Model returned {len(axis_keys)} capture axis keys, "
384 + f"expected {len(self._capture_axes)} "
385 + f"(one per axis of capture_axes)."
386 + )
387 + return selected, axis_keys
388
389 def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]:
390 """Get per-item output token counts as plain ints."""
@@ -334,6 +403,7 @@ class EncoderCudaGraphManager:
403 mm_kwargs: dict[str, Any],
404 token_budget: int,
405 path: str = "default",
406 + axis_keys: CaptureAxisKeys = (),
407 ) -> torch.Tensor | None:
408 """Execute budget graph.
409
@@ -347,11 +417,14 @@ class EncoderCudaGraphManager:
417 graph_set = self._get_graph_set(path)
418 num_items = len(self._get_item_specs(mm_kwargs))
419
350 - if token_budget not in graph_set:
420 + graph_map_key: BudgetGraphMapKey = (
421 + token_budget if not axis_keys else (token_budget, axis_keys)
422 + )
423 + if graph_map_key not in graph_set:
424 self.graph_misses += num_items
425 return None
426
354 - graph_meta = graph_set[token_budget]
427 + graph_meta = graph_set[graph_map_key]
428
429 replay = self.model.prepare_encoder_cudagraph_replay_buffers(
430 mm_kwargs,
@@ -432,7 +505,7 @@ class EncoderCudaGraphManager:
505
506 outputs_by_orig_idx: dict[int, torch.Tensor] = {}
507 for batch_indices, path_budgets in batches:
435 - batch_mm_kwargs = self._select_items(mm_kwargs, batch_indices)
508 + batch_mm_kwargs, axis_keys = self._select_items(mm_kwargs, batch_indices)
509 graph_outputs: dict[str, torch.Tensor] = {}
510 all_eager = True
511
@@ -449,10 +522,14 @@ class EncoderCudaGraphManager:
522 )
523 else:
524 all_eager = False
452 - output = self._run_budget_graph(
453 - batch_mm_kwargs, token_budget, path=path
525 + graph_output = self._run_budget_graph(
526 + batch_mm_kwargs,
527 + token_budget,
528 + path=path,
529 + axis_keys=axis_keys,
530 )
455 - assert output is not None
531 + assert graph_output is not None
532 + output = graph_output
533 graph_outputs[path] = output
534
535 if all_eager:
@@ -506,9 +583,9 @@ class EncoderCudaGraphManager:
583 ]
584
585 if len(local_indices) > 0:
509 - local_mm_kwargs = self._select_items(mm_kwargs, local_indices)
586 + local_mm_kwargs = self._select_items(mm_kwargs, local_indices)[0]
587 else:
511 - local_mm_kwargs = self._select_items(mm_kwargs, [])
588 + local_mm_kwargs = self._select_items(mm_kwargs, [])[0]
589
590 max_output_tokens_per_rank = (
591 max(
vllm/v1/worker/encoder_cudagraph_defs.py
+15 -1
@@ -2,13 +2,19 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Data transfer objects for encoder CUDA graph management."""
4
5 -from collections.abc import Callable
5 +from collections.abc import Callable, Hashable
6 from dataclasses import dataclass, field
7
8 import torch
9
10 EncoderCudaGraphPaddingLogic = Callable[[torch.Tensor, torch.Tensor], None]
11
12 +# Reserved mm_kwargs key: models with ``EncoderCudaGraphConfig.capture_axes``
13 +# put the resolved per-axis keys (one per axis, in order) into the dict
14 +# returned by ``select_encoder_cudagraph_items()``; the manager pops it before
15 +# the kwargs are used anywhere else.
16 +ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG = "encoder_cudagraph_axis_keys"
17 +
18
19 @dataclass
20 class EncoderItemSpec:
@@ -88,6 +94,14 @@ class EncoderCudaGraphConfig:
94 )
95 """Independently captured encoder paths keyed by their forward name."""
96
97 + capture_axes: tuple[tuple[Hashable, ...], ...] = ()
98 + """Extra capture axes beyond the token budget; empty to disable.
99 +
100 + Each entry is the ordered key set of one axis. When non-empty, one graph
101 + is captured per token budget per combination of axis keys (cartesian
102 + product), so the total number of captured graphs is
103 + ``num_budgets * prod(len(axis) for axis in capture_axes)``."""
104 +
105
106 @dataclass
107 class EncoderCudaGraphCaptureInputs: