| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Encoder CUDA graph support for the DeepSeek-V4.1 vision tower. |
| 4 | |
| 5 | Mixin implementing the ``SupportsEncoderCudaGraph`` protocol |
| 6 | (``vllm/model_executor/models/interfaces.py``) on top of the shared |
| 7 | ``DeepseekV4ViT``/``DeepseekV4Aligner``. The captured graph packs a whole |
| 8 | batch of images into one varlen run: ViT blocks attend per image via |
| 9 | ``cu_seqlens``, and the aligner's spatial merge becomes a gather+mask |
| 10 | against precomputed indices. Span assembly (IMAGE_START/IMAGE_NEW_LINE/ |
| 11 | IMAGE_END delimiters) stays eager in ``postprocess_encoder_output`` but is |
| 12 | batched over all items of the packed batch. |
| 13 | |
| 14 | Enabled via ``--compilation-config '{"cudagraph_mm_encoder": true}'``; the |
| 15 | manager (``vllm/v1/worker/encoder_cudagraph.py``) handles budget packing, DP |
| 16 | sharding and eager fallback. |
| 17 | """ |
| 18 | |
| 19 | from collections.abc import Hashable |
| 20 | from typing import TYPE_CHECKING, Any |
| 21 | |
| 22 | import torch |
| 23 | |
| 24 | from vllm.models.deepseek_v4.common.vision import ( |
| 25 | build_packed_merge_metadata, |
| 26 | build_packed_vit_metadata, |
| 27 | ) |
| 28 | from vllm.utils.torch_utils import async_tensor_h2d |
| 29 | from vllm.v1.worker.encoder_cudagraph_defs import ( |
| 30 | EncoderCudaGraphCaptureInputs, |
| 31 | EncoderCudaGraphConfig, |
| 32 | EncoderCudaGraphReplayBuffers, |
| 33 | EncoderItemSpec, |
| 34 | ) |
| 35 | |
| 36 | from .mm_preprocess import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START |
| 37 | |
| 38 | if TYPE_CHECKING: |
| 39 | from vllm.config import VllmConfig |
| 40 | from vllm.config.multimodal import MultiModalConfig |
| 41 | from vllm.models.deepseek_v4.common.vision import ( |
| 42 | DeepseekV4Aligner, |
| 43 | DeepseekV4ViT, |
| 44 | ) |
| 45 | from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config |
| 46 | |
| 47 | |
| 48 | class DeepseekV4VLEncoderCudaGraphMixin: |
| 49 | """``SupportsEncoderCudaGraph`` for ``DeepseekV41ForCausalLM``. |
| 50 | |
| 51 | Expects ``self.vision``/``self.aligner``/``self.config``/ |
| 52 | ``self.multimodal_config`` and the ``image_start``/``image_newline``/ |
| 53 | ``image_end`` parameters from the host model. |
| 54 | """ |
| 55 | |
| 56 | # Provided by the host model; declared here so mixin methods type-check. |
| 57 | vision: "DeepseekV4ViT" |
| 58 | aligner: "DeepseekV4Aligner" |
| 59 | config: "DeepseekV41Config" |
| 60 | multimodal_config: "MultiModalConfig | None" |
| 61 | image_start: torch.nn.Parameter |
| 62 | image_end: torch.nn.Parameter |
| 63 | image_newline: torch.nn.Parameter |
| 64 | |
| 65 | # -- helpers shared with the eager multimodal path -- |
| 66 | |
| 67 | def _encode_image( |
| 68 | self, |
| 69 | patches: torch.Tensor, |
| 70 | n_vit_h: int, |
| 71 | n_vit_w: int, |
| 72 | ) -> torch.Tensor: |
| 73 | # Aligner rows in reading order, one per IMAGE slot. |
| 74 | return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) |
| 75 | |
| 76 | def _build_image_span( |
| 77 | self, image_embeds: torch.Tensor, types: torch.Tensor |
| 78 | ) -> torch.Tensor: |
| 79 | """Full image span: aligner rows at IMAGE slots, the learned |
| 80 | delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" |
| 81 | types = types.to(image_embeds.device) |
| 82 | span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) |
| 83 | dtype = image_embeds.dtype |
| 84 | span[types == IMAGE_START] = self.image_start.to(dtype) |
| 85 | span[types == IMAGE_END] = self.image_end.to(dtype) |
| 86 | span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) |
| 87 | span[types == IMAGE] = image_embeds |
| 88 | return span |
| 89 | |
| 90 | @staticmethod |
| 91 | def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]: |
| 92 | grid = mm_kwargs[key] |
| 93 | if isinstance(grid, torch.Tensor): |
| 94 | grid = grid.tolist() |
| 95 | return [[int(v) for v in row] for row in grid] |
| 96 | |
| 97 | _encoder_cg_pad_totals: dict[int, int] = {} |
| 98 | |
| 99 | # -- SupportsEncoderCudaGraph protocol -- |
| 100 | |
| 101 | def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: |
| 102 | mm_cfg = self.multimodal_config |
| 103 | if ( |
| 104 | mm_cfg is not None |
| 105 | and mm_cfg.mm_encoder_attn_dtype == "fp8" |
| 106 | and mm_cfg.mm_encoder_fp8_scale_path is None |
| 107 | ): |
| 108 | raise ValueError( |
| 109 | "cudagraph_mm_encoder is incompatible with dynamic FP8 " |
| 110 | "scaling for the vision encoder (the amax history is " |
| 111 | "host-side state that cannot be captured). Provide a static " |
| 112 | "scale file via --mm-encoder-fp8-scale-path or disable " |
| 113 | "cudagraph_mm_encoder." |
| 114 | ) |
| 115 | |
| 116 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 117 | |
| 118 | backend = self.vision.blocks[0].attn.attn.attn_backend |
| 119 | # FlashInfer reads max_seqlen on the host and TORCH_SDPA calls |
| 120 | # .tolist() on CUDA cu_seqlens; neither can be captured. |
| 121 | if backend in ( |
| 122 | AttentionBackendEnum.FLASHINFER, |
| 123 | AttentionBackendEnum.TORCH_SDPA, |
| 124 | ): |
| 125 | raise ValueError( |
| 126 | f"cudagraph_mm_encoder is not supported with the {backend.name} " |
| 127 | "ViT attention backend for DeepSeek-V4.1. Set " |
| 128 | "--mm-encoder-attn-backend FLASH_ATTN or disable " |
| 129 | "cudagraph_mm_encoder." |
| 130 | ) |
| 131 | |
| 132 | pad_totals = self._encoder_cg_pad_totals |
| 133 | |
| 134 | def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: |
| 135 | # Varlen attention requires cu_seqlens[-1] to equal the number of |
| 136 | # rows actually passed in. The captured buffers are sized for the |
| 137 | # full token budget, so a smaller real batch is completed with one |
| 138 | # trailing padding sequence; declaring fewer rows than the buffer |
| 139 | # holds is undefined behaviour and returns NaN on FlashAttn. |
| 140 | total = pad_totals.get(dst.data_ptr()) |
| 141 | if total is None: |
| 142 | raise RuntimeError( |
| 143 | "cu_seqlens replay buffer was not registered at capture " |
| 144 | "time; the manager must replay into the capture-time " |
| 145 | "buffers." |
| 146 | ) |
| 147 | n = min(src.shape[0], dst.shape[0]) |
| 148 | dst[:n].copy_(src[:n]) |
| 149 | dst[n:] = total |
| 150 | |
| 151 | return EncoderCudaGraphConfig( |
| 152 | modalities=["image"], |
| 153 | buffer_keys=[ |
| 154 | "patches", |
| 155 | "vit_cos", |
| 156 | "vit_sin", |
| 157 | "cu_seqlens", |
| 158 | "max_seqlen", |
| 159 | "merge_idx", |
| 160 | "merge_mask", |
| 161 | ], |
| 162 | out_hidden_size=self.config.hidden_size, |
| 163 | padding_logics={"cu_seqlens": pad_cu_seqlens}, |
| 164 | ) |
| 165 | |
| 166 | def get_encoder_cudagraph_budget_range( |
| 167 | self, |
| 168 | vllm_config: "VllmConfig", |
| 169 | ) -> tuple[int, int]: |
| 170 | # Min: comfortably above the smallest possible image span (4 tokens |
| 171 | # for a 1x1 grid); also sets max_batch_size = 64 via the manager's |
| 172 | # auto-inference (max_batch_size <= min_budget invariant). |
| 173 | min_budget = 64 |
| 174 | max_budget = min( |
| 175 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 176 | vllm_config.model_config.max_model_len, |
| 177 | ) |
| 178 | return (min_budget, max_budget) |
| 179 | |
| 180 | def get_encoder_cudagraph_item_specs( |
| 181 | self, |
| 182 | mm_kwargs: dict[str, Any], |
| 183 | ) -> list[EncoderItemSpec]: |
| 184 | vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") |
| 185 | llm_grid = self._get_grid_list(mm_kwargs, "llm_grid") |
| 186 | return [ |
| 187 | EncoderItemSpec( |
| 188 | input_size=h * w, |
| 189 | output_tokens=lh * (lw + 1) + 2, |
| 190 | ) |
| 191 | for (h, w), (lh, lw) in zip(vit_grid, llm_grid, strict=True) |
| 192 | ] |
| 193 | |
| 194 | def select_encoder_cudagraph_items( |
| 195 | self, |
| 196 | mm_kwargs: dict[str, Any], |
| 197 | indices: list[int], |
| 198 | ) -> dict[str, Any]: |
| 199 | vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") |
| 200 | llm_grid = self._get_grid_list(mm_kwargs, "llm_grid") |
| 201 | patches = mm_kwargs["patches"] |
| 202 | types = mm_kwargs["types"] |
| 203 | |
| 204 | if len(indices) == 0: |
| 205 | return { |
| 206 | "patches": patches[:0], |
| 207 | "vit_grid": torch.zeros((0, 2), dtype=torch.int64), |
| 208 | "llm_grid": torch.zeros((0, 2), dtype=torch.int64), |
| 209 | "types": types[:0], |
| 210 | } |
| 211 | |
| 212 | cum_patches = [0] |
| 213 | for h, w in vit_grid: |
| 214 | cum_patches.append(cum_patches[-1] + h * w) |
| 215 | cum_spans = [0] |
| 216 | for lh, lw in llm_grid: |
| 217 | cum_spans.append(cum_spans[-1] + lh * (lw + 1) + 2) |
| 218 | |
| 219 | return { |
| 220 | "patches": torch.cat( |
| 221 | [patches[cum_patches[i] : cum_patches[i + 1]] for i in indices] |
| 222 | ), |
| 223 | "vit_grid": torch.tensor([vit_grid[i] for i in indices], dtype=torch.int64), |
| 224 | "llm_grid": torch.tensor([llm_grid[i] for i in indices], dtype=torch.int64), |
| 225 | "types": torch.cat( |
| 226 | [types[cum_spans[i] : cum_spans[i + 1]] for i in indices] |
| 227 | ), |
| 228 | } |
| 229 | |
| 230 | def prepare_encoder_cudagraph_capture_inputs( |
| 231 | self, |
| 232 | token_budget: int, |
| 233 | max_batch_size: int, |
| 234 | max_frames_per_batch: int, |
| 235 | device: torch.device, |
| 236 | dtype: torch.dtype, |
| 237 | path: str = "default", |
| 238 | axis_keys: tuple[Hashable, ...] | None = None, |
| 239 | ) -> EncoderCudaGraphCaptureInputs: |
| 240 | if ( |
| 241 | self.multimodal_config is not None |
| 242 | and self.multimodal_config.get_limit_per_prompt("image") == 0 |
| 243 | ): |
| 244 | raise RuntimeError( |
| 245 | "cudagraph_mm_encoder requires the vision tower, but it was " |
| 246 | "stubbed out because the image limit is 0. Disable " |
| 247 | "cudagraph_mm_encoder." |
| 248 | ) |
| 249 | |
| 250 | config = self.config |
| 251 | r = config.vision_downsample_ratio |
| 252 | # Span rows per dummy item (ceiling so the buffers tile the budget). |
| 253 | # Each dummy image is a single aligner row of per_item_out blocks: |
| 254 | # r x (r * per_item_out) patches. One real image using the whole |
| 255 | # budget needs at most r^2 * token_budget patches, which this fits. |
| 256 | per_item_out = -(-token_budget // max_batch_size) |
| 257 | grids = [[r, r * per_item_out]] * max_batch_size |
| 258 | total_patches = r * r * per_item_out * max_batch_size |
| 259 | |
| 260 | p = config.vision_patch_size |
| 261 | patches = torch.zeros(total_patches, 3, p, p, device=device, dtype=dtype) |
| 262 | |
| 263 | metadata = build_packed_vit_metadata( |
| 264 | grids, |
| 265 | rope_dim=self.vision.rope_dim, |
| 266 | rope_theta=self.vision.rope_theta, |
| 267 | device=device, |
| 268 | max_seqlen_override=total_patches, |
| 269 | cached=False, |
| 270 | ) |
| 271 | # Spare cu_seqlens slots let replay append a padding sequence that |
| 272 | # covers the rows a smaller real batch does not fill. |
| 273 | real_cu = metadata.pop("cu_seqlens") |
| 274 | cu_seqlens = torch.full( |
| 275 | (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device |
| 276 | ) |
| 277 | cu_seqlens[: real_cu.numel()] = real_cu |
| 278 | |
| 279 | merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype) |
| 280 | |
| 281 | self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches |
| 282 | |
| 283 | return EncoderCudaGraphCaptureInputs( |
| 284 | values={ |
| 285 | "patches": patches, |
| 286 | **metadata, |
| 287 | "cu_seqlens": cu_seqlens, |
| 288 | **merge, |
| 289 | } |
| 290 | ) |
| 291 | |
| 292 | def prepare_encoder_cudagraph_replay_buffers( |
| 293 | self, |
| 294 | mm_kwargs: dict[str, Any], |
| 295 | max_batch_size: int, |
| 296 | max_frames_per_batch: int, |
| 297 | path: str = "default", |
| 298 | ) -> EncoderCudaGraphReplayBuffers: |
| 299 | vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") |
| 300 | patches = mm_kwargs["patches"] |
| 301 | dtype = self.aligner.w1.weight.dtype |
| 302 | if patches.dtype != dtype: |
| 303 | patches = patches.to(dtype) |
| 304 | |
| 305 | # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and |
| 306 | # pad_cu_seqlens appends the padding sequence covering the tail rows. |
| 307 | metadata = build_packed_vit_metadata( |
| 308 | vit_grid, |
| 309 | rope_dim=self.vision.rope_dim, |
| 310 | rope_theta=self.vision.rope_theta, |
| 311 | device=patches.device, |
| 312 | ) |
| 313 | merge = build_packed_merge_metadata( |
| 314 | vit_grid, |
| 315 | self.config.vision_downsample_ratio, |
| 316 | device=patches.device, |
| 317 | dtype=dtype, |
| 318 | ) |
| 319 | return EncoderCudaGraphReplayBuffers( |
| 320 | values={"patches": patches, **metadata, **merge} |
| 321 | ) |
| 322 | |
| 323 | def encoder_cudagraph_forward( |
| 324 | self, |
| 325 | inputs: dict[str, torch.Tensor], |
| 326 | path: str = "default", |
| 327 | ) -> torch.Tensor: |
| 328 | vit_out = self.vision.forward_packed( |
| 329 | inputs["patches"], |
| 330 | inputs["vit_cos"], |
| 331 | inputs["vit_sin"], |
| 332 | inputs["cu_seqlens"], |
| 333 | inputs["max_seqlen"], |
| 334 | ) |
| 335 | return self.aligner.forward_packed( |
| 336 | vit_out, inputs["merge_idx"], inputs["merge_mask"] |
| 337 | ) |
| 338 | |
| 339 | def encoder_eager_forward( |
| 340 | self, |
| 341 | mm_kwargs: dict[str, Any], |
| 342 | path: str = "default", |
| 343 | ) -> torch.Tensor: |
| 344 | patches = mm_kwargs["patches"].to(self.aligner.w1.weight.dtype) |
| 345 | vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") |
| 346 | outs: list[torch.Tensor] = [] |
| 347 | offset = 0 |
| 348 | for n_vit_h, n_vit_w in vit_grid: |
| 349 | n_vit = n_vit_h * n_vit_w |
| 350 | outs.append( |
| 351 | self._encode_image(patches[offset : offset + n_vit], n_vit_h, n_vit_w) |
| 352 | ) |
| 353 | offset += n_vit |
| 354 | if not outs: |
| 355 | return patches.new_zeros((0, self.config.hidden_size)) |
| 356 | return torch.cat(outs, dim=0) |
| 357 | |
| 358 | def postprocess_encoder_output( |
| 359 | self, |
| 360 | outputs: dict[str, torch.Tensor], |
| 361 | indices: list[int], |
| 362 | per_item_out_tokens: list[int], |
| 363 | dest: dict[int, torch.Tensor] | list[torch.Tensor | None], |
| 364 | clone: bool = False, |
| 365 | batch_mm_kwargs: dict[str, Any] | None = None, |
| 366 | ) -> None: |
| 367 | if not indices: |
| 368 | return |
| 369 | assert batch_mm_kwargs is not None |
| 370 | aligner_out = outputs["default"] |
| 371 | r = self.config.vision_downsample_ratio |
| 372 | vit_grid = self._get_grid_list(batch_mm_kwargs, "vit_grid") |
| 373 | llm_grid = self._get_grid_list(batch_mm_kwargs, "llm_grid") |
| 374 | types = batch_mm_kwargs["types"] |
| 375 | |
| 376 | n_rows = sum(-(-h // r) * (-(-w // r)) for h, w in vit_grid) |
| 377 | span_lens = [lh * (lw + 1) + 2 for lh, lw in llm_grid] |
| 378 | |
| 379 | # Batched span assembly: one gather over [aligner rows, START, |
| 380 | # NEW_LINE, END] for all items. The index is built from the host-side |
| 381 | # ``types`` so that neither the copy nor the gather syncs. |
| 382 | is_image = types == IMAGE |
| 383 | src_idx = torch.where( |
| 384 | is_image, |
| 385 | is_image.cumsum(0) - 1, |
| 386 | n_rows + (types == IMAGE_NEW_LINE).long() + 2 * (types == IMAGE_END).long(), |
| 387 | ) |
| 388 | delims = torch.stack([self.image_start, self.image_newline, self.image_end]) |
| 389 | src = torch.cat([aligner_out[:n_rows], delims.to(aligner_out.dtype)]) |
| 390 | span = src[async_tensor_h2d(src_idx, aligner_out.device)] |
| 391 | |
| 392 | # Freshly allocated, so later replays cannot clobber the results. |
| 393 | offset = 0 |
| 394 | for idx, span_len in zip(indices, span_lens): |
| 395 | dest[idx] = span[offset : offset + span_len] |
| 396 | offset += span_len |