[MM][Perf][CG] Support ViT full CUDA graph for Kimi-VL (#41992)

Signed-off-by: oguz <[email protected]>

docs/design/cuda_graphs_multimodal.md

@@ -129,6 +129,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 | `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
 | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
 | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
+| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
 | `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ |
 | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ |
 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ |

examples/generate/multimodal/vision_language_offline.py

@@ -2537,6 +2537,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
     "qwen2_5_vl",
     "qwen3_vl",
     "qwen3_vl_moe",
+    "kimi_vl",
     "qwen3_5",
     "qwen3_5_moe",
     "internvl_chat",

tests/models/multimodal/generation/test_vit_cudagraph.py

@@ -48,6 +48,13 @@ def internvl_chat_template(content: str) -> str:
     return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
 
 
+def kimi_vl_chat_template(content: str) -> str:
+    return (
+        f"<|im_user|>user<|im_middle|>{content}<|im_end|>"
+        "<|im_assistant|>assistant<|im_middle|>"
+    )
+
+
 def step3_vl_chat_template(content: str) -> str:
     return (
         "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n "
@@ -100,6 +107,34 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
         needs_video_metadata=False,
         marks=[pytest.mark.core_model],
     ),
+    "kimi_vl": VitCudagraphTestConfig(
+        model="moonshotai/Kimi-VL-A3B-Instruct",
+        modalities=["image"],
+        image_prompt=kimi_vl_chat_template(
+            "<|media_start|>image<|media_content|><|media_pad|><|media_end|>"
+            "What is in this image?"
+        ),
+        needs_video_metadata=False,
+        # Single bucket sized to cover the test images' output tokens.
+        # The default auto-inferred range fans out into multiple power-of-2
+        # buckets, each holding a full ViT capture pool.
+        compilation_config_overrides={
+            "encoder_cudagraph_token_budgets": [1024],
+        },
+        # Shrink to 1 text + 1 vision layer with random weights so the
+        # test runs on any CI GPU (incl. L4) and skips the multi-GiB
+        # weight download. The test only validates that encoder CG
+        # capture/replay functions correctly, not output quality.
+        vllm_runner_kwargs={
+            "trust_remote_code": True,
+            "load_format": "dummy",
+            "hf_overrides": partial(
+                dummy_hf_overrides,
+                model_arch="KimiVLForConditionalGeneration",
+            ),
+        },
+        marks=[pytest.mark.core_model],
+    ),
     "qwen3_vl": VitCudagraphTestConfig(
         model="Qwen/Qwen3-VL-2B-Instruct",
         image_prompt=qwen_vl_chat_template(

vllm/model_executor/models/kimi_vl.py

@@ -56,7 +56,11 @@ from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.interfaces import (
+    SupportsEncoderCudaGraph,
+    SupportsMultiModal,
+    SupportsPP,
+)
 from vllm.model_executor.models.moonvit import MoonVitPretrainedModel
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
@@ -79,6 +83,7 @@ from vllm.multimodal.processing import (
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.kimi_vl import KimiVLConfig, MoonViTConfig
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
+from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
 
 from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
 from .vision import is_vit_use_data_parallel, run_dp_sharded_mrope_vision_model
@@ -287,7 +292,9 @@ class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]):
     info=KimiVLProcessingInfo,
     dummy_inputs=KimiVLDummyInputsBuilder,
 )
-class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
+class KimiVLForConditionalGeneration(
+    nn.Module, SupportsMultiModal, SupportsEncoderCudaGraph, SupportsPP
+):
     supports_encoder_tp_data = True
 
     @classmethod
@@ -340,6 +347,192 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
 
         self.media_placeholder: int = self.config.media_placeholder_token_id
 
+        self.model_config = model_config
+
+    # -- SupportsEncoderCudaGraph protocol methods --
+
+    def get_encoder_cudagraph_config(self):
+        from vllm.v1.worker.encoder_cudagraph_defs import (
+            EncoderCudaGraphConfig,
+        )
+
+        return EncoderCudaGraphConfig(
+            modalities=["image"],
+            buffer_keys=[
+                "pixel_values",
+                "pos_embeds",
+                "rope_freqs_cis",
+                "cu_seqlens",
+                "max_seqlen",
+                "merge_gather_idx",
+            ],
+            out_hidden_size=self.hidden_size,
+        )
+
+    def get_encoder_cudagraph_budget_range(
+        self,
+        vllm_config,
+    ) -> tuple[int, int]:
+        # Min: estimated smallest possible encoder input.
+        # 224x224 image with patch_size=14 -> 16x16 patches, then merge
+        # kernel (2,2) -> 8x8 = 64 output tokens.
+        min_budget = 64
+        max_budget = min(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            self.model_config.max_model_len,
+        )
+        return (min_budget, max_budget)
+
+    def _get_grid_hws(
+        self,
+        mm_kwargs: dict[str, Any],
+    ) -> list[tuple[int, int]]:
+        grid_hws = mm_kwargs["image_grid_hws"]
+        if not isinstance(grid_hws, list):
+            grid_hws = grid_hws.tolist()
+        return grid_hws
+
+    def get_encoder_cudagraph_item_specs(
+        self,
+        mm_kwargs: dict[str, Any],
+    ):
+        from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
+
+        kh, kw = self.config.vision_config.merge_kernel_size
+        return [
+            EncoderItemSpec(
+                input_size=h * w,
+                output_tokens=(h // kh) * (w // kw),
+            )
+            for h, w in self._get_grid_hws(mm_kwargs)
+        ]
+
+    def select_encoder_cudagraph_items(
+        self,
+        mm_kwargs: dict[str, Any],
+        indices: list[int],
+    ) -> dict[str, Any]:
+        grid_hws = self._get_grid_hws(mm_kwargs)
+        pixel_values = mm_kwargs["pixel_values"]
+
+        if len(indices) == 0:
+            return {
+                "pixel_values": pixel_values[:0],
+                "image_grid_hws": pixel_values.new_zeros((0, 2), dtype=torch.long),
+            }
+
+        patches_per_item = [h * w for h, w in grid_hws]
+        cum_patches = [0]
+        for p in patches_per_item:
+            cum_patches.append(cum_patches[-1] + p)
+
+        selected_pv = torch.cat(
+            [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
+        )
+        selected_grid = torch.tensor(
+            [grid_hws[i] for i in indices],
+            dtype=torch.long,
+            device=pixel_values.device,
+        )
+        return {
+            "pixel_values": selected_pv,
+            "image_grid_hws": selected_grid,
+        }
+
+    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",
+    ):
+        from vllm.v1.worker.encoder_cudagraph_defs import (
+            EncoderCudaGraphCaptureInputs,
+        )
+
+        kh, kw = self.config.vision_config.merge_kernel_size
+        # Ceil so the buffer fits the worst case of one item using the full
+        # budget. Floor under-allocates when budget is not a multiple of
+        # max_batch_size.
+        per_mm_item_output = (token_budget + max_batch_size - 1) // max_batch_size
+
+        # Shape the synthetic grid so neither dimension exceeds Rope2DPosEmb's
+        # precomputed range. Pack as wide a row as fits, then add rows.
+        rope = self.vision_tower.encoder.rope_2d
+        max_wo = rope.max_width // kw
+        wo = min(per_mm_item_output, max_wo)
+        ho = (per_mm_item_output + wo - 1) // wo
+        assert ho * kh <= rope.max_height, (
+            f"per_mm_item_output={per_mm_item_output} exceeds RoPE grid capacity "
+            f"(max {(rope.max_height // kh) * (rope.max_width // kw)} tokens)"
+        )
+        grid_hws_list = [(ho * kh, wo * kw) for _ in range(max_batch_size)]
+
+        patch_size = self.config.vision_config.patch_size
+        if isinstance(patch_size, int):
+            patch_size = (patch_size, patch_size)
+
+        total_patches = sum(h * w for h, w in grid_hws_list)
+        in_channels = 3
+        dummy_pixel_values = torch.randn(
+            total_patches,
+            in_channels,
+            patch_size[0],
+            patch_size[1],
+            device=device,
+            dtype=dtype,
+        )
+
+        buffers = self.vision_tower.prepare_encoder_metadata(
+            grid_hws_list,
+            max_batch_size=max_batch_size,
+            max_seqlen_override=token_budget,
+            device=device,
+        )
+        values = buffers | {"pixel_values": dummy_pixel_values}
+
+        return EncoderCudaGraphCaptureInputs(values=values)
+
+    def prepare_encoder_cudagraph_replay_buffers(
+        self,
+        mm_kwargs: dict[str, Any],
+        max_batch_size: int,
+        max_frames_per_batch: int,
+        path: str = "default",
+    ):
+        grid_hws_list = self._get_grid_hws(mm_kwargs)
+        buffers = self.vision_tower.prepare_encoder_metadata(
+            grid_hws_list,
+            max_batch_size=max_batch_size,
+            device=mm_kwargs["pixel_values"].device,
+        )
+        values = buffers | {"pixel_values": mm_kwargs["pixel_values"]}
+        return EncoderCudaGraphReplayBuffers(values=values)
+
+    def encoder_cudagraph_forward(
+        self,
+        values: dict[str, torch.Tensor],
+        path: str = "default",
+    ) -> torch.Tensor:
+        pixel_values = values.pop("pixel_values")
+        metadata = values
+        image_features = self.vision_tower(
+            pixel_values, grid_hw=None, encoder_metadata=metadata
+        )
+        return self.multi_modal_projector(image_features)
+
+    def encoder_eager_forward(
+        self,
+        mm_kwargs: dict[str, Any],
+        path: str = "default",
+    ) -> torch.Tensor:
+        pixel_values = mm_kwargs["pixel_values"]
+        image_grid_hws = mm_kwargs["image_grid_hws"]
+        image_features = self.vision_tower(pixel_values, image_grid_hws)
+        return self.multi_modal_projector(torch.cat(image_features))
+
     def _parse_and_validate_image_input(
         self, **kwargs: object
     ) -> KimiVLImageInputs | None:

vllm/model_executor/models/moonvit.py

@@ -45,7 +45,9 @@
 from collections.abc import Sequence
 from copy import deepcopy
 from functools import cached_property
+from typing import Any
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -110,23 +112,42 @@ class Learnable2DInterpPosEmb(nn.Module):
     def reset_parameters(self):
         nn.init.normal_(self.weight)
 
-    def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
-        pos_embs = []
-        for shape in grid_hws.tolist():
-            if shape == self.weight.shape[:-1]:
+    def get_pos_embeds(
+        self,
+        grid_hws_list: list[list[int]] | list[tuple[int, int]],
+    ) -> torch.Tensor:
+        """Build packed per-token positional embeddings for a list of grids.
+
+        Returns a tensor of shape ``(sum(h * w), dim)`` formed by interpolating
+        the learned ``(height, width, dim)`` weight to each ``(h, w)`` grid and
+        concatenating the flattened results in the same order as
+        ``grid_hws_list``. Lives outside the captured CUDA graph so the
+        per-grid Python iteration is safe.
+        """
+        weight_shape = list(self.weight.shape[:-1])
+        pos_embs: list[torch.Tensor] = []
+        for shape in grid_hws_list:
+            shape_list = [int(shape[0]), int(shape[1])]
+            if shape_list == weight_shape:
                 pos_embs.append(self.weight.flatten(end_dim=1))
             else:
                 pos_embs.append(
                     F.interpolate(
                         self.weight.permute((2, 0, 1)).unsqueeze(0),
-                        size=shape,
+                        size=tuple(shape_list),
                         mode=self.interpolation_mode,
                     )
                     .squeeze(0)
                     .permute((1, 2, 0))
                     .flatten(end_dim=1)
                 )
-        out = x + torch.cat(pos_embs)
+        if not pos_embs:
+            return self.weight.new_zeros((0, self.weight.shape[-1]))
+        return torch.cat(pos_embs)
+
+    def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
+        pos_embs = self.get_pos_embeds(grid_hws.tolist())
+        out = x + pos_embs
         return out
 
 
@@ -158,19 +179,29 @@ class MoonVisionPatchEmbed(nn.Module):
             height=pos_emb_height, width=pos_emb_width, dim=out_dim
         )
 
-    def forward(self, x: torch.Tensor, grid_hw: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        x: torch.Tensor,
+        grid_hw: torch.Tensor | None = None,
+        *,
+        pos_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         """
         Args:
             x (L, Channels): input tensor
             grid_hw (N, 2): grid height and width
+            pos_embeds: precomputed positional embeddings of shape
+                ``(L, Cout)``. When provided, ``grid_hw`` is unused and the
+                CUDA-graph-incompatible interpolation in ``self.pos_emb`` is
+                skipped.
 
         Returns:
             (L, Cout) tensor
         """
         x = self.proj(x).view(x.size(0), -1)
-        # apply positional embedding
-        x = self.pos_emb(x, grid_hw)
-        return x
+        if pos_embeds is not None:
+            return x + pos_embeds
+        return self.pos_emb(x, grid_hw)
 
 
 class Rope2DPosEmb(nn.Module):
@@ -243,6 +274,35 @@ class Rope2DPosEmb(nn.Module):
         freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1)
         return freqs_cis
 
+    def get_freqs_cis_by_seqlens_list(
+        self,
+        grid_hws_list: list[list[int]] | list[tuple[int, int]],
+    ) -> torch.Tensor:
+        """List-based variant of :meth:`get_freqs_cis_by_seqlens`.
+
+        Accepts a Python list of ``(h, w)`` pairs so callers that already
+        operate outside the captured CUDA graph can avoid materializing a
+        tensor + ``.tolist()`` round-trip.
+        """
+        assert all(
+            1 <= h <= self.max_height and 1 <= w <= self.max_width
+            for h, w in grid_hws_list
+        ), (
+            grid_hws_list,
+            self.max_height,
+            self.max_width,
+        )
+        if not grid_hws_list:
+            return self.precomputed_freqs_cis.new_zeros((0, self.dim // 2))
+        freqs_cis = torch.cat(
+            [
+                self.precomputed_freqs_cis[:h, :w].reshape(-1, self.dim // 2)
+                for h, w in grid_hws_list
+            ],
+            dim=0,
+        )
+        return freqs_cis
+
     def get_freqs_cis_by_seqlens(self, grid_hws: torch.Tensor) -> torch.Tensor:
         """
         Args:
@@ -250,22 +310,7 @@ class Rope2DPosEmb(nn.Module):
         Returns:
             freqs_cis: tensor of shape (sum(t * height * width), dim//2)
         """
-        shapes = grid_hws.tolist()
-        assert all(
-            1 <= h <= self.max_height and 1 <= w <= self.max_width for h, w in shapes
-        ), (
-            shapes,
-            self.max_height,
-            self.max_width,
-        )
-        freqs_cis = torch.cat(
-            [
-                self.precomputed_freqs_cis[:h, :w].reshape(-1, self.dim // 2)
-                for h, w in shapes
-            ],
-            dim=0,
-        )
-        return freqs_cis
+        return self.get_freqs_cis_by_seqlens_list(grid_hws.tolist())
 
     def get_freqs_cis_by_idx(
         self, pos_idx: torch.Tensor, pos_idx_mask: torch.Tensor
@@ -392,11 +437,15 @@ class MoonVitEncoderLayer(nn.Module):
         x: torch.Tensor,
         cu_seqlens: torch.Tensor,
         rope_freqs_cis: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
     ):
         """
         Args:
             x (torch.Tensor): (seqlen, hidden_dim)
             cu_seqlens (torch.Tensor):
+            max_seqlen: Optional precomputed scalar tensor. When omitted it
+                is derived from ``cu_seqlens``, which produces a GPU scalar
+                that breaks CUDA graph capture.
         """
         seq_length = x.size(0)
         xqkv, _ = self.wqkv(x)
@@ -412,7 +461,8 @@ class MoonVitEncoderLayer(nn.Module):
 
         xq, xk = apply_rope(xq, xk, rope_freqs_cis)
 
-        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+        if max_seqlen is None:
+            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
         attn_out = self.attn(
             xq.unsqueeze(0),
             xk.unsqueeze(0),
@@ -433,10 +483,12 @@ class MoonVitEncoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor,
         rope_freqs_cis: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """
         Args:
             hidden_states: non-packed (B, N, D) or packed (L, D). if non-packed, seqlens should be None, if packed, seqlens should be set
+            max_seqlen: optional precomputed max-sequence-length scalar.
 
         Returns:
             output: same shape of input, non-packed (B, N, D) for non-packed input, (L, D) for packed input
@@ -444,7 +496,10 @@ class MoonVitEncoderLayer(nn.Module):
         residual = hidden_states
         hidden_states = self.norm0(hidden_states)
         attn_out = self.attention_qkvpacked(
-            hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
+            hidden_states,
+            cu_seqlens,
+            rope_freqs_cis=rope_freqs_cis,
+            max_seqlen=max_seqlen,
         )
         hidden_states = residual + attn_out
 
@@ -478,22 +533,39 @@ class MoonVitEncoder(nn.Module):
         )
         self.final_layernorm = nn.LayerNorm(hidden_dim)
 
-    def forward(
-        self, hidden_states: torch.Tensor, grid_hw: torch.Tensor
+    def get_rope_freqs_cis(
+        self,
+        grid_hws_list: list[list[int]] | list[tuple[int, int]],
     ) -> torch.Tensor:
-        rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(grid_hws=grid_hw)
+        return self.rope_2d.get_freqs_cis_by_seqlens_list(grid_hws_list)
 
-        lengths = torch.cat(
-            (
-                torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype),
-                (grid_hw[:, 0] * grid_hw[:, 1]).to(hidden_states.device),
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        grid_hw: torch.Tensor | None = None,
+        *,
+        cu_seqlens: torch.Tensor | None = None,
+        rope_freqs_cis: torch.Tensor | None = None,
+        max_seqlen: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if rope_freqs_cis is None:
+            rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(grid_hws=grid_hw)
+
+        if cu_seqlens is None:
+            lengths = torch.cat(
+                (
+                    torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype),
+                    (grid_hw[:, 0] * grid_hw[:, 1]).to(hidden_states.device),
+                )
             )
-        )
-        cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32)
+            cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32)
 
         for _, block in enumerate(self.blocks):
             hidden_states = block(
-                hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
+                hidden_states,
+                cu_seqlens,
+                rope_freqs_cis=rope_freqs_cis,
+                max_seqlen=max_seqlen,
             )
 
         hidden_states = self.final_layernorm(hidden_states)
@@ -530,6 +602,54 @@ def patch_merger(
     return outputs
 
 
+def patch_merger_packed(
+    x: torch.Tensor,
+    gather_idx: torch.Tensor,
+    merge_kernel_size: tuple[int, int],
+) -> torch.Tensor:
+    """CUDA-graph-safe equivalent of :func:`patch_merger`.
+
+    Uses a precomputed index tensor to gather the per-token reshape +
+    permute that ``patch_merger`` does inside a Python loop. The output is
+    the concatenated 3D tensor ``(sum(new_h * new_w), kh * kw, d_model)``,
+    matching what ``torch.cat(patch_merger(...))`` would produce.
+    """
+    kh, kw = merge_kernel_size
+    d_model = x.size(-1)
+    return x.index_select(0, gather_idx).view(-1, kh * kw, d_model)
+
+
+def _build_merge_gather_idx(
+    grid_hws_list: list[list[int]] | list[tuple[int, int]],
+    merge_kernel_size: tuple[int, int],
+) -> np.ndarray:
+    """Build the per-token gather indices used by :func:`patch_merger_packed`.
+
+    For each item with grid (h, w) and merge kernel (kh, kw), the output
+    block at position (nh, nw) gathers the kh*kw input tokens at rows
+    (nh*kh + ih, nw*kw + iw) of that item, in (ih, iw) row-major order.
+    """
+    kh, kw = merge_kernel_size
+    parts: list[np.ndarray] = []
+    pre_sum = 0
+    for h, w in grid_hws_list:
+        new_h, new_w = h // kh, w // kw
+        nh = np.arange(new_h, dtype=np.int64).reshape(new_h, 1, 1, 1)
+        nw = np.arange(new_w, dtype=np.int64).reshape(1, new_w, 1, 1)
+        ih = np.arange(kh, dtype=np.int64).reshape(1, 1, kh, 1)
+        iw = np.arange(kw, dtype=np.int64).reshape(1, 1, 1, kw)
+        # Linearized input row = (nh*kh + ih) * w + (nw*kw + iw), offset by
+        # the per-item base ``pre_sum``. Output is laid out as
+        # (new_h, new_w, kh, kw) which patch_merger flattens to
+        # (new_h*new_w, kh*kw).
+        idx = pre_sum + (nh * kh + ih) * w + (nw * kw + iw)
+        parts.append(idx.reshape(-1))
+        pre_sum += h * w
+    if not parts:
+        return np.zeros(0, dtype=np.int64)
+    return np.concatenate(parts)
+
+
 class MoonVitPretrainedModel(PreTrainedModel):
     config_class = MoonViTConfig
     model_type = "moonvit"
@@ -570,17 +690,126 @@ class MoonVitPretrainedModel(PreTrainedModel):
             prefix=f"{prefix}.encoder",
         )
 
+    def prepare_encoder_metadata(
+        self,
+        grid_hws_list: list[list[int]] | list[tuple[int, int]],
+        *,
+        max_batch_size: int | None = None,
+        max_seqlen_override: int | None = None,
+        device: torch.device | None = None,
+    ) -> dict[str, Any]:
+        """Precompute every grid-dependent input the encoder needs.
+
+        Used by the CUDA graph capture and replay paths to precompute
+        every grid-dependent input outside the captured graph, so per-grid
+        Python iteration and ``.tolist()`` round-trips are fine; the
+        values are then copied into fixed-shape buffers for replay.
+
+        Args:
+            grid_hws_list: List of ``(h, w)`` patch-grid sizes per image.
+            max_batch_size: When set, ``cu_seqlens`` is right-padded with
+                its last value so the buffer covers up to this many
+                sequences. Required at CUDA graph capture/replay so the
+                buffer shape matches what was recorded; padding entries
+                are zero-length sequences and are ignored by varlen
+                attention.
+            max_seqlen_override: Override the per-replay max sequence
+                length scalar. At capture this must be a safe upper bound
+                (worst case: a single image consuming the full token
+                budget) because the value is baked into the captured
+                graph.
+            device: Device for the metadata tensors. Defaults to the
+                model's parameter device.
+        """
+        if device is None:
+            device = next(self.parameters()).device
+
+        # Normalize to a list of plain Python int pairs so the helpers
+        # below never need ``.tolist()`` on a tensor.
+        grid_pairs: list[tuple[int, int]] = [(int(h), int(w)) for h, w in grid_hws_list]
+
+        metadata: dict[str, Any] = {}
+
+        pos_embeds = self.patch_embed.pos_emb.get_pos_embeds(grid_pairs)
+        metadata["pos_embeds"] = pos_embeds.to(device=device)
+
+        rope_freqs_cis = self.encoder.get_rope_freqs_cis(grid_pairs)
+        metadata["rope_freqs_cis"] = rope_freqs_cis.to(device=device)
+
+        grid_arr = np.array(grid_pairs, dtype=np.int64)
+        seq_lens = (grid_arr[:, 0] * grid_arr[:, 1]).astype(np.int32)
+        cu_seqlens_np = np.concatenate(
+            [
+                np.zeros(1, dtype=np.int32),
+                seq_lens.cumsum(dtype=np.int32),
+            ]
+        )
+
+        if max_batch_size is not None:
+            num_seqs = len(cu_seqlens_np) - 1
+            if num_seqs < max_batch_size:
+                cu_seqlens_np = np.concatenate(
+                    [
+                        cu_seqlens_np,
+                        np.full(
+                            max_batch_size - num_seqs,
+                            cu_seqlens_np[-1],
+                            dtype=np.int32,
+                        ),
+                    ]
+                )
+        metadata["cu_seqlens"] = torch.from_numpy(cu_seqlens_np).to(device)
+
+        if max_seqlen_override is not None:
+            max_seqlen_val = int(max_seqlen_override)
+        else:
+            max_seqlen_val = int(seq_lens.max()) if len(seq_lens) > 0 else 0
+        # Keep on CPU: attention wrappers may call .item() on this scalar
+        # and we want that materialization to happen outside the captured
+        # graph (the value is constant per capture anyway).
+        metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)
+
+        gather_idx_np = _build_merge_gather_idx(grid_pairs, self.merge_kernel_size)
+        metadata["merge_gather_idx"] = torch.from_numpy(gather_idx_np).to(device)
+
+        return metadata
+
     def forward(
-        self, pixel_values: torch.Tensor, grid_hw: torch.Tensor
-    ) -> torch.Tensor:
+        self,
+        pixel_values: torch.Tensor,
+        grid_hw: torch.Tensor,
+        *,
+        encoder_metadata: dict[str, Any] | None = None,
+    ) -> torch.Tensor | list[torch.Tensor]:
         """
         Args:
             pixel_values (torch.Tensor): The input pixel values.
             grid_hw (torch.Tensor): The grid height and width.
-
-        Returns:
-            torch.Tensor: The output tokens.
+            encoder_metadata: Optional precomputed metadata produced by
+                :meth:`prepare_encoder_metadata`. When provided every
+                ``.tolist()`` call in the forward path is skipped, the
+                returned tensor is the packed
+                ``(sum(new_h*new_w), kh*kw, hidden_size)`` form (suitable
+                for CUDA graph capture/replay), and ``grid_hw`` is unused.
+                When ``None`` the legacy path runs and returns a list of
+                per-image tensors.
         """
+        if encoder_metadata is not None:
+            hidden_states = self.patch_embed(
+                pixel_values, pos_embeds=encoder_metadata["pos_embeds"]
+            )
+            hidden_states = self.encoder(
+                hidden_states,
+                cu_seqlens=encoder_metadata["cu_seqlens"],
+                rope_freqs_cis=encoder_metadata["rope_freqs_cis"],
+                max_seqlen=encoder_metadata["max_seqlen"],
+            )
+            return patch_merger_packed(
+                hidden_states,
+                encoder_metadata["merge_gather_idx"],
+                merge_kernel_size=self.merge_kernel_size,
+            )
+
         hidden_states = self.patch_embed(pixel_values, grid_hw)
         hidden_states = self.encoder(hidden_states, grid_hw)
         hidden_states = patch_merger(