diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 379e5f16b52..1db82ffa688 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/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` | ✅︎ | ✅︎ | ❌︎ | diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 48521c52482..1b3741a3e42 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/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", diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 0496031988f..52b28ca8600 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/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( diff --git a/vllm/model_executor/models/kimi_vl.py b/vllm/model_executor/models/kimi_vl.py index e3bc08c654d..2b08fc6c1fd 100644 --- a/vllm/model_executor/models/kimi_vl.py +++ b/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: diff --git a/vllm/model_executor/models/moonvit.py b/vllm/model_executor/models/moonvit.py index 8c699865618..73e17cb9fb6 100644 --- a/vllm/model_executor/models/moonvit.py +++ b/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(