| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Encoder CUDA graph support for the MiniMax M3 vision tower. |
| 4 | |
| 5 | Mixin implementing the SupportsEncoderCudaGraph protocol for any host class |
| 6 | that holds the vision tower as ``self.vision_tower`` (a |
| 7 | ``MiniMaxVLVisionModel``). The manager routes image items through this mixin; |
| 8 | video items use the eager encoder path. |
| 9 | """ |
| 10 | |
| 11 | from collections.abc import Hashable |
| 12 | from typing import TYPE_CHECKING, Any |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.config.multimodal import MultiModalConfig |
| 21 | from vllm.models.minimax_m3.common.vision_tower import MiniMaxVLVisionModel |
| 22 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 23 | EncoderCudaGraphCaptureInputs, |
| 24 | EncoderCudaGraphConfig, |
| 25 | EncoderCudaGraphReplayBuffers, |
| 26 | EncoderItemSpec, |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | class MiniMaxM3EncoderCudaGraphMixin(SupportsEncoderCudaGraph): |
| 31 | """SupportsEncoderCudaGraph for hosts with ``self.vision_tower``.""" |
| 32 | |
| 33 | vision_tower: "MiniMaxVLVisionModel" |
| 34 | multimodal_config: "MultiModalConfig | None" |
| 35 | _encoder_cg_pad_totals: dict[int, int] = {} |
| 36 | |
| 37 | def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": |
| 38 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 39 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig |
| 40 | |
| 41 | if ( |
| 42 | self.multimodal_config is not None |
| 43 | and self.multimodal_config.get_limit_per_prompt("image") == 0 |
| 44 | ): |
| 45 | raise ValueError( |
| 46 | "cudagraph_mm_encoder requires the MiniMax M3 vision tower, " |
| 47 | "but the image limit is 0. Disable cudagraph_mm_encoder." |
| 48 | ) |
| 49 | |
| 50 | mm_cfg = self.multimodal_config |
| 51 | if ( |
| 52 | mm_cfg is not None |
| 53 | and mm_cfg.mm_encoder_attn_dtype == "fp8" |
| 54 | and mm_cfg.mm_encoder_fp8_scale_path is None |
| 55 | ): |
| 56 | raise ValueError( |
| 57 | "cudagraph_mm_encoder requires static FP8 scales for the " |
| 58 | "MiniMax M3 vision encoder. Set --mm-encoder-fp8-scale-path " |
| 59 | "or disable cudagraph_mm_encoder." |
| 60 | ) |
| 61 | |
| 62 | # FlashInfer reads max_seqlen on the host and TORCH_SDPA calls |
| 63 | # .tolist() on CUDA cu_seqlens; neither can be captured. |
| 64 | backend = self.vision_tower.vision_model.attn_backend |
| 65 | if backend not in ( |
| 66 | AttentionBackendEnum.FLASH_ATTN, |
| 67 | AttentionBackendEnum.ROCM_AITER_FA, |
| 68 | ): |
| 69 | raise ValueError( |
| 70 | f"cudagraph_mm_encoder is not supported with the " |
| 71 | f"{backend.name} ViT attention backend for MiniMax M3. Set " |
| 72 | "--mm-encoder-attn-backend FLASH_ATTN or ROCM_AITER_FA " |
| 73 | "(ROCm), or disable cudagraph_mm_encoder." |
| 74 | ) |
| 75 | |
| 76 | pad_totals = self._encoder_cg_pad_totals |
| 77 | |
| 78 | def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: |
| 79 | # Varlen attention requires cu_seqlens[-1] to equal the number of |
| 80 | # rows actually passed in. The captured buffers are sized for the |
| 81 | # full token budget, so a smaller real batch has to be completed |
| 82 | # with one trailing padding sequence; declaring fewer rows than the |
| 83 | # buffer holds is undefined behaviour and returns NaN on FlashAttn. |
| 84 | total = pad_totals.get(dst.data_ptr()) |
| 85 | if total is None: |
| 86 | raise RuntimeError("cu_seqlens replay buffer was not registered") |
| 87 | n = src.shape[0] |
| 88 | dst[:n].copy_(src[:n]) |
| 89 | dst[n:] = total |
| 90 | |
| 91 | def pad_rope_planes(dst: torch.Tensor, src: torch.Tensor) -> None: |
| 92 | # cos/sin planes are (3, N, half_rot_dim): the row axis is dim 1, |
| 93 | # not dim 0 as the manager's default slice-copy assumes. |
| 94 | dst.zero_() |
| 95 | dst[:, : src.shape[1]].copy_(src) |
| 96 | |
| 97 | return EncoderCudaGraphConfig( |
| 98 | modalities=["image"], |
| 99 | buffer_keys=[ |
| 100 | "pixel_values", |
| 101 | "rotary_cos", |
| 102 | "rotary_sin", |
| 103 | "cu_seqlens", |
| 104 | "max_seqlen", |
| 105 | ], |
| 106 | out_hidden_size=self.vision_tower.out_hidden_size, |
| 107 | padding_logics={ |
| 108 | "cu_seqlens": pad_cu_seqlens, |
| 109 | "rotary_cos": pad_rope_planes, |
| 110 | "rotary_sin": pad_rope_planes, |
| 111 | }, |
| 112 | ) |
| 113 | |
| 114 | def get_encoder_cudagraph_budget_range( |
| 115 | self, |
| 116 | vllm_config: "VllmConfig", |
| 117 | ) -> tuple[int, int]: |
| 118 | min_budget = 64 |
| 119 | max_budget = min( |
| 120 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 121 | vllm_config.model_config.max_model_len, |
| 122 | ) |
| 123 | return (min_budget, max_budget) |
| 124 | |
| 125 | @staticmethod |
| 126 | def _get_grid_thw_list(mm_kwargs: dict[str, Any]) -> list[list[int]]: |
| 127 | grid_thw = mm_kwargs["image_grid_thw"] |
| 128 | if isinstance(grid_thw, torch.Tensor): |
| 129 | return [[int(x) for x in row] for row in grid_thw.tolist()] |
| 130 | return [[int(x) for x in row] for row in grid_thw] |
| 131 | |
| 132 | @staticmethod |
| 133 | def _get_pixel_values(mm_kwargs: dict[str, Any]) -> torch.Tensor: |
| 134 | pixel_values = mm_kwargs["pixel_values"] |
| 135 | if isinstance(pixel_values, list): |
| 136 | pixel_values = torch.cat(pixel_values) |
| 137 | return pixel_values |
| 138 | |
| 139 | def get_encoder_cudagraph_item_specs( |
| 140 | self, mm_kwargs: dict[str, Any] |
| 141 | ) -> list["EncoderItemSpec"]: |
| 142 | from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec |
| 143 | |
| 144 | merge = self.vision_tower.spatial_merge_size |
| 145 | return [ |
| 146 | EncoderItemSpec( |
| 147 | input_size=t * h * w, |
| 148 | output_tokens=t * h * w // (merge * merge), |
| 149 | ) |
| 150 | for t, h, w in self._get_grid_thw_list(mm_kwargs) |
| 151 | ] |
| 152 | |
| 153 | def select_encoder_cudagraph_items( |
| 154 | self, mm_kwargs: dict[str, Any], indices: list[int] |
| 155 | ) -> dict[str, Any]: |
| 156 | grid_thw_list = self._get_grid_thw_list(mm_kwargs) |
| 157 | pixel_values = self._get_pixel_values(mm_kwargs) |
| 158 | source_grid = mm_kwargs["image_grid_thw"] |
| 159 | |
| 160 | if len(indices) == 0: |
| 161 | empty_grid = ( |
| 162 | source_grid[:0] |
| 163 | if isinstance(source_grid, torch.Tensor) |
| 164 | else pixel_values.new_zeros((0, 3), dtype=torch.long) |
| 165 | ) |
| 166 | return {"pixel_values": pixel_values[:0], "image_grid_thw": empty_grid} |
| 167 | |
| 168 | patch_counts = [t * h * w for t, h, w in grid_thw_list] |
| 169 | cum = [0] |
| 170 | for pc in patch_counts: |
| 171 | cum.append(cum[-1] + pc) |
| 172 | selected_pv = torch.cat( |
| 173 | [pixel_values[cum[i] : cum[i + 1]] for i in indices], dim=0 |
| 174 | ) |
| 175 | grid_device = ( |
| 176 | source_grid.device if isinstance(source_grid, torch.Tensor) else None |
| 177 | ) |
| 178 | selected_grid = torch.tensor( |
| 179 | [grid_thw_list[i] for i in indices], |
| 180 | dtype=torch.long, |
| 181 | device=grid_device, |
| 182 | ) |
| 183 | return {"pixel_values": selected_pv, "image_grid_thw": selected_grid} |
| 184 | |
| 185 | def prepare_encoder_cudagraph_capture_inputs( |
| 186 | self, |
| 187 | token_budget: int, |
| 188 | max_batch_size: int, |
| 189 | max_frames_per_batch: int, |
| 190 | device: torch.device, |
| 191 | dtype: torch.dtype, |
| 192 | path: str = "default", |
| 193 | axis_keys: tuple[Hashable, ...] | None = None, |
| 194 | ) -> "EncoderCudaGraphCaptureInputs": |
| 195 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 196 | EncoderCudaGraphCaptureInputs, |
| 197 | ) |
| 198 | |
| 199 | merge = self.vision_tower.spatial_merge_size |
| 200 | # Output tokens per item in the dummy grid (ceiling so total >= budget). |
| 201 | per_item_out = (token_budget + max_batch_size - 1) // max_batch_size |
| 202 | grid_thw_list = [ |
| 203 | [1, merge, per_item_out * merge] for _ in range(max_batch_size) |
| 204 | ] |
| 205 | |
| 206 | embeddings = self.vision_tower.vision_model.embeddings |
| 207 | patch_dim = ( |
| 208 | embeddings.num_channels |
| 209 | * embeddings.temporal_patch_size |
| 210 | * embeddings.patch_size**2 |
| 211 | ) |
| 212 | total_patches = sum(t * h * w for t, h, w in grid_thw_list) |
| 213 | # Match the tower dtype exactly: MiniMaxVLPatchEmbed.forward mutates |
| 214 | # the module on dtype mismatch, which must not happen during capture. |
| 215 | dummy_pixel_values = torch.zeros( |
| 216 | total_patches, |
| 217 | patch_dim, |
| 218 | device=device, |
| 219 | dtype=self.vision_tower.dtype, |
| 220 | ) |
| 221 | |
| 222 | # The padding sequence can contain nearly all captured rows when a |
| 223 | # small image replays this graph. Leave one spare cu_seqlens slot. |
| 224 | metadata = self.vision_tower.vision_model.prepare_encoder_metadata( |
| 225 | grid_thw_list, |
| 226 | device=device, |
| 227 | max_batch_size=max_batch_size + 1, |
| 228 | max_seqlen_override=total_patches, |
| 229 | ) |
| 230 | |
| 231 | values: dict[str, torch.Tensor] = {"pixel_values": dummy_pixel_values} |
| 232 | values.update({k: v for k, v in metadata.items() if v is not None}) |
| 233 | self._encoder_cg_pad_totals[values["cu_seqlens"].data_ptr()] = total_patches |
| 234 | return EncoderCudaGraphCaptureInputs(values=values) |
| 235 | |
| 236 | def prepare_encoder_cudagraph_replay_buffers( |
| 237 | self, |
| 238 | mm_kwargs: dict[str, Any], |
| 239 | max_batch_size: int, |
| 240 | max_frames_per_batch: int, |
| 241 | path: str = "default", |
| 242 | ) -> "EncoderCudaGraphReplayBuffers": |
| 243 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 244 | EncoderCudaGraphReplayBuffers, |
| 245 | ) |
| 246 | |
| 247 | pixel_values = self._get_pixel_values(mm_kwargs) |
| 248 | # Unpadded: pad_cu_seqlens completes the tail with the padding sequence |
| 249 | # so cu_seqlens[-1] matches the captured buffer's row count. |
| 250 | metadata = self.vision_tower.vision_model.prepare_encoder_metadata( |
| 251 | self._get_grid_thw_list(mm_kwargs), |
| 252 | device=pixel_values.device, |
| 253 | ) |
| 254 | # FlashAttention reads this CPU scalar during capture, so the captured |
| 255 | # launch must retain the bound for the padded buffer. |
| 256 | metadata.pop("max_seqlen") |
| 257 | |
| 258 | values: dict[str, torch.Tensor | None] = {"pixel_values": pixel_values} |
| 259 | values.update(metadata) |
| 260 | return EncoderCudaGraphReplayBuffers(values=values) |
| 261 | |
| 262 | def encoder_cudagraph_forward( |
| 263 | self, |
| 264 | inputs: dict[str, torch.Tensor], |
| 265 | path: str = "default", |
| 266 | ) -> torch.Tensor: |
| 267 | pixel_values = inputs.pop("pixel_values") |
| 268 | # Remaining keys are consumed as encoder_metadata. |
| 269 | return self.vision_tower(pixel_values, grid_thw=None, encoder_metadata=inputs) |
| 270 | |
| 271 | def encoder_eager_forward( |
| 272 | self, |
| 273 | mm_kwargs: dict[str, Any], |
| 274 | path: str = "default", |
| 275 | ) -> torch.Tensor: |
| 276 | return self.vision_tower( |
| 277 | pixel_values=self._get_pixel_values(mm_kwargs).type( |
| 278 | self.vision_tower.dtype |
| 279 | ), |
| 280 | grid_thw=self._get_grid_thw_list(mm_kwargs), |
| 281 | ) |