diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 1ed84ebc52..9bda7ec85d 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -55,6 +55,10 @@ def kimi_vl_chat_template(content: str) -> str: ) +def deepseek_v41_chat_template(content: str) -> str: + return f"<｜begin▁of▁sentence｜><｜User｜>{content}<｜Assistant｜>" + + def step3_vl_chat_template(content: str) -> str: return ( "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n " @@ -221,6 +225,25 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { vllm_runner_kwargs={"gpu_memory_utilization": 0.80}, marks=[pytest.mark.core_model], ), + "deepseek_v41": VitCudagraphTestConfig( + model="deepseek-ai/DeepSeek-V4.1-Flash", + modalities=["image"], + image_prompt=deepseek_v41_chat_template( + "<｜deepseek_image｜>\n\nWhat is in this image?" + ), + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [1024], + }, + vllm_runner_kwargs={ + "load_format": "dummy", + "attention_backend": "FLASHMLA_SPARSE_DSV41", + "hf_overrides": partial( + dummy_hf_overrides, + model_arch="DeepseekV41ForCausalLM", + exist_overrides={"vision_n_layers": 1}, + ), + }, + ), "step3_vl": VitCudagraphTestConfig( model="stepfun-ai/Step3-VL-10B", modalities=["image"], diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 403a01a9e2..1e1722537d 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -218,8 +218,14 @@ def test_model_tensor_schema(model_id: str): factories = model_cls._processor_factory + # Capture helpers return capture-buffer containers (e.g. + # EncoderCudaGraphCaptureInputs), not TensorSchema mm inputs. + capture_helpers = {"prepare_encoder_cudagraph_capture_inputs"} + inputs_parse_methods = [] for attr_name in dir(model_cls): + if attr_name in capture_helpers: + continue attr = getattr(model_cls, attr_name) if hasattr(attr, "__annotations__"): return_type = attr.__annotations__.get("return", None) diff --git a/vllm/models/deepseek_v4/common/vision.py b/vllm/models/deepseek_v4/common/vision.py index 4c633da6c5..bfc02768c0 100644 --- a/vllm/models/deepseek_v4/common/vision.py +++ b/vllm/models/deepseek_v4/common/vision.py @@ -40,8 +40,7 @@ from vllm.model_executor.models.vision import ( ) -@lru_cache(8) -def get_vision_cos_sin( +def _compute_vision_cos_sin( n_h: int, n_w: int, dim: int, theta: float ) -> tuple[torch.Tensor, torch.Tensor]: inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) @@ -52,6 +51,13 @@ def get_vision_cos_sin( return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) +@lru_cache(8) +def get_vision_cos_sin( + n_h: int, n_w: int, dim: int, theta: float +) -> tuple[torch.Tensor, torch.Tensor]: + return _compute_vision_cos_sin(n_h, n_w, dim, theta) + + def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: dtype = x.dtype x1, x2 = x.float().chunk(2, dim=-1) @@ -123,15 +129,23 @@ class DeepseekV4VisionAttention(nn.Module): ) def forward( - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: n = x.size(0) qkv, _ = self.wqkv(x) q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1)) q = apply_rotary(q, cos, sin).unsqueeze(0) # (b=1, n, h, d) k = apply_rotary(k, cos, sin).unsqueeze(0) - # One image per call: a dense batch, no varlen packing metadata. - o = self.attn(q, k, v.unsqueeze(0)) + # Dense (cu_seqlens=None): one image per call. Varlen (cu_seqlens + # set): packed multi-image batch for encoder CUDA graph replay. + o = self.attn( + q, k, v.unsqueeze(0), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen + ) out, _ = self.wo(o.reshape(n, -1)) return out @@ -173,9 +187,14 @@ class DeepseekV4VisionBlock(nn.Module): self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp") def forward( - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: - x = x + self.attn(self.norm1(x), cos, sin) + x = x + self.attn(self.norm1(x), cos, sin, cu_seqlens, max_seqlen) return x + self.mlp(self.norm2(x)) @@ -206,6 +225,115 @@ class DeepseekV4ViT(nn.Module): x = block(x, cos, sin) return self.norm(x) + def forward_packed( + self, + patches: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + ) -> torch.Tensor: + """Varlen path for encoder CUDA graphs: multiple images packed along + rows, per-image attention via ``cu_seqlens``, RoPE tables precomputed + on device. Pure-tensor (no host reads, no H2D), safe to capture.""" + x = self.patch_embed(patches) + for block in self.blocks: + x = block(x, cos, sin, cu_seqlens, max_seqlen) + return self.norm(x) + + +def build_packed_vit_metadata( + grids: list[list[int]] | list[tuple[int, int]], + *, + rope_dim: int, + rope_theta: float, + device: torch.device, + max_seqlen_override: int | None = None, + cached: bool = True, +) -> dict[str, torch.Tensor]: + """Precompute packed-batch ViT metadata for ``DeepseekV4ViT.forward_packed``. + + Args: + grids: ``[n_vit_h, n_vit_w]`` per image, in packing order. + rope_dim: RoPE embedding dimension of the ViT. + rope_theta: RoPE base frequency of the ViT. + device: Device for the returned tensors (except ``max_seqlen``, + which stays on the host). + max_seqlen_override: Worst-case value baked in at CUDA graph capture + (the attention wrapper reads ``max_seqlen`` on the host, so the + capture-time value becomes a graph constant). + cached: Use the shared ``lru_cache`` for RoPE tables. Capture-time + dummy grids pass ``False`` to avoid evicting real entries. + + """ + cos_sin = get_vision_cos_sin if cached else _compute_vision_cos_sin + cos_list: list[torch.Tensor] = [] + sin_list: list[torch.Tensor] = [] + cu_seqlens = [0] + max_seqlen = 0 + for n_h, n_w in grids: + cos, sin = cos_sin(n_h, n_w, rope_dim, rope_theta) + cos_list.append(cos) + sin_list.append(sin) + n = n_h * n_w + cu_seqlens.append(cu_seqlens[-1] + n) + max_seqlen = max(max_seqlen, n) + if cos_list: + cos = torch.cat(cos_list).to(device) + sin = torch.cat(sin_list).to(device) + else: + cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device) + sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device) + if max_seqlen_override is not None: + max_seqlen = max_seqlen_override + return { + "vit_cos": cos, + "vit_sin": sin, + "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32, device=device), + # Read on the host by the attention wrapper; keep on CPU. + "max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32), + } + + +def build_packed_merge_metadata( + grids: list[list[int]] | list[tuple[int, int]], + downsample_ratio: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + """Gather indices/mask for ``DeepseekV4Aligner.forward_packed``. + + For each aligner output row, the ``r x r`` source patch rows in the + packed ViT output; positions past the image edge (the eager path pads + them with zeros) get mask 0 and a clamped in-range index. + """ + r = downsample_ratio + idx_list: list[torch.Tensor] = [] + mask_list: list[torch.Tensor] = [] + offset = 0 + for n_h, n_w in grids: + rows = -(-n_h // r) + cols = -(-n_w // r) + bi = torch.arange(rows).repeat_interleave(cols) + bj = torch.arange(cols).repeat(rows) + di = torch.arange(r).repeat_interleave(r) + dj = torch.arange(r).repeat(r) + hi = bi[:, None] * r + di[None, :] + wj = bj[:, None] * r + dj[None, :] + valid = (hi < n_h) & (wj < n_w) + idx = offset + hi.clamp(max=n_h - 1) * n_w + wj.clamp(max=n_w - 1) + idx_list.append(idx) + mask_list.append(valid) + offset += n_h * n_w + if idx_list: + merge_idx = torch.cat(idx_list).to(device) + merge_mask = torch.cat(mask_list).unsqueeze(-1).to(device=device, dtype=dtype) + else: + merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device) + merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device) + return {"merge_idx": merge_idx, "merge_mask": merge_mask} + class DeepseekV4Aligner(nn.Module): """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.""" @@ -240,6 +368,27 @@ class DeepseekV4Aligner(nn.Module): out, _ = self.w2(F.gelu(hidden)) return out + def forward_packed( + self, + x: torch.Tensor, + merge_idx: torch.Tensor, + merge_mask: torch.Tensor, + ) -> torch.Tensor: + """Packed multi-image path for encoder CUDA graphs. + + ``merge_idx``/``merge_mask`` come from ``build_packed_merge_metadata``; + gather+mask reproduces the eager ``F.pad`` + ``F.unfold`` merge + exactly (unfold rows are channel-major, so the gathered patch rows + are transposed before flattening). + """ + r2 = self.downsample_ratio**2 + m = merge_idx.shape[0] + gathered = (x[merge_idx] * merge_mask).view(m, r2, -1) + gathered = gathered.permute(0, 2, 1).reshape(m, -1) + hidden, _ = self.w1(gathered) + out, _ = self.w2(F.gelu(hidden)) + return out + def run_dp_sharded_vision_tower( vision_model: DeepseekV4ViT, diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 3227b50a4e..3c57df2ad2 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -26,6 +26,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, ) @@ -44,15 +45,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.utils.tensor_schema import TensorSchema, TensorShape from ..common.mm_preprocess import ( - IMAGE, - IMAGE_END, - IMAGE_NEW_LINE, IMAGE_PLACEHOLDER, - IMAGE_START, DeepseekV4VLDummyInputsBuilder, DeepseekV4VLMultiModalProcessor, DeepseekV4VLProcessingInfo, ) +from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper( info=DeepseekV4VLProcessingInfo, dummy_inputs=DeepseekV4VLDummyInputsBuilder, ) -class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3): +class DeepseekV41ForCausalLM( + nn.Module, + DeepseekV4VLEncoderCudaGraphMixin, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsEagle3, +): """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower. ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) delegates through ``language_model`` via the protocol defaults. + ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is + implemented by ``DeepseekV4VLEncoderCudaGraphMixin``. """ supports_encoder_tp_data = True @@ -196,29 +203,6 @@ class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports resolve_bindings={"p": self.config.vision_patch_size}, ) - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - ) -> torch.Tensor: - # Aligner rows in reading order, one per IMAGE slot. - return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) - - def _build_image_span( - self, image_embeds: torch.Tensor, types: torch.Tensor - ) -> torch.Tensor: - """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" - types = types.to(image_embeds.device) - span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) - dtype = image_embeds.dtype - span[types == IMAGE_START] = self.image_start.to(dtype) - span[types == IMAGE_END] = self.image_end.to(dtype) - span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) - span[types == IMAGE] = image_embeds - return span - def _process_image_input( self, image_input: DeepseekV4VLImagePixelInputs, diff --git a/vllm/models/deepseek_v41/common/vl_cudagraph.py b/vllm/models/deepseek_v41/common/vl_cudagraph.py new file mode 100644 index 0000000000..e2a718b745 --- /dev/null +++ b/vllm/models/deepseek_v41/common/vl_cudagraph.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Encoder CUDA graph support for the DeepSeek-V4.1 vision tower. + +Mixin implementing the ``SupportsEncoderCudaGraph`` protocol +(``vllm/model_executor/models/interfaces.py``) on top of the shared +``DeepseekV4ViT``/``DeepseekV4Aligner``. The captured graph packs a whole +batch of images into one varlen run: ViT blocks attend per image via +``cu_seqlens``, and the aligner's spatial merge becomes a gather+mask +against precomputed indices. Span assembly (IMAGE_START/IMAGE_NEW_LINE/ +IMAGE_END delimiters) stays eager in ``postprocess_encoder_output`` but is +batched over all items of the packed batch. + +Enabled via ``--compilation-config '{"cudagraph_mm_encoder": true}'``; the +manager (``vllm/v1/worker/encoder_cudagraph.py``) handles budget packing, DP +sharding and eager fallback. +""" + +from collections.abc import Hashable +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.models.deepseek_v4.common.vision import ( + build_packed_merge_metadata, + build_packed_vit_metadata, +) +from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + EncoderCudaGraphConfig, + EncoderCudaGraphReplayBuffers, + EncoderItemSpec, +) + +from .mm_preprocess import IMAGE, IMAGE_END, IMAGE_NEW_LINE, IMAGE_START + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.config.multimodal import MultiModalConfig + from vllm.models.deepseek_v4.common.vision import ( + DeepseekV4Aligner, + DeepseekV4ViT, + ) + from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config + + +class DeepseekV4VLEncoderCudaGraphMixin: + """``SupportsEncoderCudaGraph`` for ``DeepseekV41ForCausalLM``. + + Expects ``self.vision``/``self.aligner``/``self.config``/ + ``self.multimodal_config`` and the ``image_start``/``image_newline``/ + ``image_end`` parameters from the host model. + """ + + # Provided by the host model; declared here so mixin methods type-check. + vision: "DeepseekV4ViT" + aligner: "DeepseekV4Aligner" + config: "DeepseekV41Config" + multimodal_config: "MultiModalConfig | None" + image_start: torch.nn.Parameter + image_end: torch.nn.Parameter + image_newline: torch.nn.Parameter + + # -- helpers shared with the eager multimodal path -- + + def _encode_image( + self, + patches: torch.Tensor, + n_vit_h: int, + n_vit_w: int, + ) -> torch.Tensor: + # Aligner rows in reading order, one per IMAGE slot. + return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) + + def _build_image_span( + self, image_embeds: torch.Tensor, types: torch.Tensor + ) -> torch.Tensor: + """Full image span: aligner rows at IMAGE slots, the learned + delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" + types = types.to(image_embeds.device) + span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) + dtype = image_embeds.dtype + span[types == IMAGE_START] = self.image_start.to(dtype) + span[types == IMAGE_END] = self.image_end.to(dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) + span[types == IMAGE] = image_embeds + return span + + @staticmethod + def _get_grid_list(mm_kwargs: dict[str, Any], key: str) -> list[list[int]]: + grid = mm_kwargs[key] + if isinstance(grid, torch.Tensor): + grid = grid.tolist() + return [[int(v) for v in row] for row in grid] + + _encoder_cg_pad_totals: dict[int, int] = {} + + # -- SupportsEncoderCudaGraph protocol -- + + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: + mm_cfg = self.multimodal_config + if ( + mm_cfg is not None + and mm_cfg.mm_encoder_attn_dtype == "fp8" + and mm_cfg.mm_encoder_fp8_scale_path is None + ): + raise ValueError( + "cudagraph_mm_encoder is incompatible with dynamic FP8 " + "scaling for the vision encoder (the amax history is " + "host-side state that cannot be captured). Provide a static " + "scale file via --mm-encoder-fp8-scale-path or disable " + "cudagraph_mm_encoder." + ) + + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + backend = self.vision.blocks[0].attn.attn.attn_backend + # FlashInfer reads max_seqlen on the host and TORCH_SDPA calls + # .tolist() on CUDA cu_seqlens; neither can be captured. + if backend in ( + AttentionBackendEnum.FLASHINFER, + AttentionBackendEnum.TORCH_SDPA, + ): + raise ValueError( + f"cudagraph_mm_encoder is not supported with the {backend.name} " + "ViT attention backend for DeepSeek-V4.1. Set " + "--mm-encoder-attn-backend FLASH_ATTN or disable " + "cudagraph_mm_encoder." + ) + + pad_totals = self._encoder_cg_pad_totals + + def pad_cu_seqlens(dst: torch.Tensor, src: torch.Tensor) -> None: + # Varlen attention requires cu_seqlens[-1] to equal the number of + # rows actually passed in. The captured buffers are sized for the + # full token budget, so a smaller real batch is completed with one + # trailing padding sequence; declaring fewer rows than the buffer + # holds is undefined behaviour and returns NaN on FlashAttn. + total = pad_totals.get(dst.data_ptr()) + if total is None: + raise RuntimeError( + "cu_seqlens replay buffer was not registered at capture " + "time; the manager must replay into the capture-time " + "buffers." + ) + n = min(src.shape[0], dst.shape[0]) + dst[:n].copy_(src[:n]) + dst[n:] = total + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[ + "patches", + "vit_cos", + "vit_sin", + "cu_seqlens", + "max_seqlen", + "merge_idx", + "merge_mask", + ], + out_hidden_size=self.config.hidden_size, + padding_logics={"cu_seqlens": pad_cu_seqlens}, + ) + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: "VllmConfig", + ) -> tuple[int, int]: + # Min: comfortably above the smallest possible image span (4 tokens + # for a 1x1 grid); also sets max_batch_size = 64 via the manager's + # auto-inference (max_batch_size <= min_budget invariant). + min_budget = 64 + 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]: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(mm_kwargs, "llm_grid") + return [ + EncoderItemSpec( + input_size=h * w, + output_tokens=lh * (lw + 1) + 2, + ) + for (h, w), (lh, lw) in zip(vit_grid, llm_grid, strict=True) + ] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(mm_kwargs, "llm_grid") + patches = mm_kwargs["patches"] + types = mm_kwargs["types"] + + if len(indices) == 0: + return { + "patches": patches[:0], + "vit_grid": torch.zeros((0, 2), dtype=torch.int64), + "llm_grid": torch.zeros((0, 2), dtype=torch.int64), + "types": types[:0], + } + + cum_patches = [0] + for h, w in vit_grid: + cum_patches.append(cum_patches[-1] + h * w) + cum_spans = [0] + for lh, lw in llm_grid: + cum_spans.append(cum_spans[-1] + lh * (lw + 1) + 2) + + return { + "patches": torch.cat( + [patches[cum_patches[i] : cum_patches[i + 1]] for i in indices] + ), + "vit_grid": torch.tensor([vit_grid[i] for i in indices], dtype=torch.int64), + "llm_grid": torch.tensor([llm_grid[i] for i in indices], dtype=torch.int64), + "types": torch.cat( + [types[cum_spans[i] : cum_spans[i + 1]] for i in indices] + ), + } + + 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, + ) -> EncoderCudaGraphCaptureInputs: + if ( + self.multimodal_config is not None + and self.multimodal_config.get_limit_per_prompt("image") == 0 + ): + raise RuntimeError( + "cudagraph_mm_encoder requires the vision tower, but it was " + "stubbed out because the image limit is 0. Disable " + "cudagraph_mm_encoder." + ) + + config = self.config + r = config.vision_downsample_ratio + # Span rows per dummy item (ceiling so the buffers tile the budget). + # Each dummy image is a single aligner row of per_item_out blocks: + # r x (r * per_item_out) patches. One real image using the whole + # budget needs at most r^2 * token_budget patches, which this fits. + per_item_out = -(-token_budget // max_batch_size) + grids = [[r, r * per_item_out]] * max_batch_size + total_patches = r * r * per_item_out * max_batch_size + + p = config.vision_patch_size + patches = torch.zeros(total_patches, 3, p, p, device=device, dtype=dtype) + + metadata = build_packed_vit_metadata( + grids, + rope_dim=self.vision.rope_dim, + rope_theta=self.vision.rope_theta, + device=device, + max_seqlen_override=total_patches, + cached=False, + ) + # Spare cu_seqlens slots let replay append a padding sequence that + # covers the rows a smaller real batch does not fill. + real_cu = metadata.pop("cu_seqlens") + cu_seqlens = torch.full( + (max_batch_size + 2,), total_patches, dtype=torch.int32, device=device + ) + cu_seqlens[: real_cu.numel()] = real_cu + + merge = build_packed_merge_metadata(grids, r, device=device, dtype=dtype) + + self._encoder_cg_pad_totals[cu_seqlens.data_ptr()] = total_patches + + return EncoderCudaGraphCaptureInputs( + values={ + "patches": patches, + **metadata, + "cu_seqlens": cu_seqlens, + **merge, + } + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + path: str = "default", + ) -> EncoderCudaGraphReplayBuffers: + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + patches = mm_kwargs["patches"] + dtype = self.aligner.w1.weight.dtype + if patches.dtype != dtype: + patches = patches.to(dtype) + + # Unpadded: the manager zero-pads patches/cos/sin/merge buffers, and + # pad_cu_seqlens appends the padding sequence covering the tail rows. + metadata = build_packed_vit_metadata( + vit_grid, + rope_dim=self.vision.rope_dim, + rope_theta=self.vision.rope_theta, + device=patches.device, + ) + merge = build_packed_merge_metadata( + vit_grid, + self.config.vision_downsample_ratio, + device=patches.device, + dtype=dtype, + ) + return EncoderCudaGraphReplayBuffers( + values={"patches": patches, **metadata, **merge} + ) + + def encoder_cudagraph_forward( + self, + inputs: dict[str, torch.Tensor], + path: str = "default", + ) -> torch.Tensor: + vit_out = self.vision.forward_packed( + inputs["patches"], + inputs["vit_cos"], + inputs["vit_sin"], + inputs["cu_seqlens"], + inputs["max_seqlen"], + ) + return self.aligner.forward_packed( + vit_out, inputs["merge_idx"], inputs["merge_mask"] + ) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + ) -> torch.Tensor: + patches = mm_kwargs["patches"].to(self.aligner.w1.weight.dtype) + vit_grid = self._get_grid_list(mm_kwargs, "vit_grid") + outs: list[torch.Tensor] = [] + offset = 0 + for n_vit_h, n_vit_w in vit_grid: + n_vit = n_vit_h * n_vit_w + outs.append( + self._encode_image(patches[offset : offset + n_vit], n_vit_h, n_vit_w) + ) + offset += n_vit + if not outs: + return patches.new_zeros((0, self.config.hidden_size)) + return torch.cat(outs, dim=0) + + def postprocess_encoder_output( + self, + outputs: dict[str, torch.Tensor], + indices: list[int], + per_item_out_tokens: list[int], + dest: dict[int, torch.Tensor] | list[torch.Tensor | None], + clone: bool = False, + batch_mm_kwargs: dict[str, Any] | None = None, + ) -> None: + if not indices: + return + assert batch_mm_kwargs is not None + aligner_out = outputs["default"] + r = self.config.vision_downsample_ratio + vit_grid = self._get_grid_list(batch_mm_kwargs, "vit_grid") + llm_grid = self._get_grid_list(batch_mm_kwargs, "llm_grid") + types = batch_mm_kwargs["types"].to(aligner_out.device) + + n_rows = sum(-(-h // r) * (-(-w // r)) for h, w in vit_grid) + span_lens = [lh * (lw + 1) + 2 for lh, lw in llm_grid] + + # Batched span assembly: one masked fill per role across all items. + dtype = aligner_out.dtype + span = aligner_out.new_empty(sum(span_lens), self.config.hidden_size) + span[types == IMAGE] = aligner_out[:n_rows] + span[types == IMAGE_START] = self.image_start.to(dtype) + span[types == IMAGE_END] = self.image_end.to(dtype) + span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) + + # Freshly allocated, so later replays cannot clobber the results. + offset = 0 + for idx, span_len in zip(indices, span_lens): + dest[idx] = span[offset : offset + span_len] + offset += span_len diff --git a/vllm/models/deepseek_v41/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py index bfbcaa89a1..901985bdd7 100644 --- a/vllm/models/deepseek_v41/nvidia/vl_model.py +++ b/vllm/models/deepseek_v41/nvidia/vl_model.py @@ -26,6 +26,7 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, ) @@ -44,15 +45,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.utils.tensor_schema import TensorSchema, TensorShape from ..common.mm_preprocess import ( - IMAGE, - IMAGE_END, - IMAGE_NEW_LINE, IMAGE_PLACEHOLDER, - IMAGE_START, DeepseekV4VLDummyInputsBuilder, DeepseekV4VLMultiModalProcessor, DeepseekV4VLProcessingInfo, ) +from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -115,11 +113,20 @@ def _make_deepseek_v4_vl_weights_mapper( info=DeepseekV4VLProcessingInfo, dummy_inputs=DeepseekV4VLDummyInputsBuilder, ) -class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3): +class DeepseekV41ForCausalLM( + nn.Module, + DeepseekV4VLEncoderCudaGraphMixin, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsPP, + SupportsEagle3, +): """Multimodal entry point for DeepSeek-V4.1 checkpoints with a vision tower. ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) delegates through ``language_model`` via the protocol defaults. + ``SupportsEncoderCudaGraph`` (ViT CUDA graph capture/replay) is + implemented by ``DeepseekV4VLEncoderCudaGraphMixin``. """ supports_encoder_tp_data = True @@ -196,29 +203,6 @@ class DeepseekV41ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Supports resolve_bindings={"p": self.config.vision_patch_size}, ) - def _encode_image( - self, - patches: torch.Tensor, - n_vit_h: int, - n_vit_w: int, - ) -> torch.Tensor: - # Aligner rows in reading order, one per IMAGE slot. - return self.aligner(self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w) - - def _build_image_span( - self, image_embeds: torch.Tensor, types: torch.Tensor - ) -> torch.Tensor: - """Full image span: aligner rows at IMAGE slots, the learned - delimiter vectors at IMAGE_START/IMAGE_NEW_LINE/IMAGE_END.""" - types = types.to(image_embeds.device) - span = image_embeds.new_empty(types.numel(), image_embeds.shape[-1]) - dtype = image_embeds.dtype - span[types == IMAGE_START] = self.image_start.to(dtype) - span[types == IMAGE_END] = self.image_end.to(dtype) - span[types == IMAGE_NEW_LINE] = self.image_newline.to(dtype) - span[types == IMAGE] = image_embeds - return span - def _process_image_input( self, image_input: DeepseekV4VLImagePixelInputs,