[DSV4.1] Add encoder cuda graph support for deepseek-v4.1-flash (#56625)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

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｜></think>"
+
+
 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"],

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)

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,

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,

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

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,