From 053e7daa79208fa33ec5fb1801520c4f5da4d9ca Mon Sep 17 00:00:00 2001 From: Yi Zhong <207368749+vincentzed@users.noreply.github.com> Date: Fri, 12 Jun 2026 09:17:26 -0700 Subject: [PATCH] [Model] Add encoder CUDA graph support to Lfm2VL (#44930) Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com> --- vllm/model_executor/models/lfm2_vl.py | 434 ++++++++++++++++++++++++-- 1 file changed, 416 insertions(+), 18 deletions(-) diff --git a/vllm/model_executor/models/lfm2_vl.py b/vllm/model_executor/models/lfm2_vl.py index 9be8c5c1e5..7062884b7e 100644 --- a/vllm/model_executor/models/lfm2_vl.py +++ b/vllm/model_executor/models/lfm2_vl.py @@ -4,7 +4,7 @@ import itertools import math from collections.abc import Iterable, Mapping, Sequence -from typing import Annotated, Literal +from typing import Annotated, Any, Literal import torch import torch.nn as nn @@ -49,6 +49,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( IsHybrid, MultiModalEmbeddings, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMultiModal, SupportsPP, @@ -63,6 +64,17 @@ from .utils import ( from .vision import is_vit_use_data_parallel +def _pad_cumulative_seqlens_buffer( + dst: torch.Tensor, + src: torch.Tensor, +) -> None: + n = src.shape[0] + dst.zero_() + dst[:n].copy_(src) + if n < dst.shape[0]: + dst[n:] = src[-1] + + class Lfm2VLImagePixelInputs(TensorSchema): """ Dimensions: @@ -558,13 +570,26 @@ class Lfm2VLMultiModalProjector(nn.Module): if gather_idx_parts: gather_idx = torch.cat(gather_idx_parts).to(device=device) - gathered = vision_features_packed.index_select(0, gather_idx) - unshuffled = gathered.reshape(-1, factor * factor * hidden_size) + return self.forward_with_gather_idx(vision_features_packed, gather_idx) else: unshuffled = vision_features_packed.new_empty( (0, factor * factor * hidden_size) ) + return self.forward_from_unshuffled(unshuffled) + + def forward_with_gather_idx( + self, + vision_features_packed: torch.Tensor, + gather_idx: torch.Tensor, + ) -> torch.Tensor: + hidden_size = vision_features_packed.shape[-1] + factor = self.factor + gathered = vision_features_packed.index_select(0, gather_idx) + unshuffled = gathered.reshape(-1, factor * factor * hidden_size) + return self.forward_from_unshuffled(unshuffled) + + def forward_from_unshuffled(self, unshuffled: torch.Tensor) -> torch.Tensor: if self.projector_use_layernorm: unshuffled = self.layer_norm(unshuffled) hidden_states = self.linear_1(unshuffled) @@ -579,7 +604,12 @@ class Lfm2VLMultiModalProjector(nn.Module): dummy_inputs=Lfm2VLDummyInputsBuilder, ) class Lfm2VLForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP, IsHybrid + nn.Module, + SupportsMultiModal, + SupportsEncoderCudaGraph, + SupportsLoRA, + SupportsPP, + IsHybrid, ): merge_by_field_config = True @@ -645,6 +675,7 @@ class Lfm2VLForConditionalGeneration( self.config = config self.vllm_config = vllm_config + self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" @@ -697,7 +728,7 @@ class Lfm2VLForConditionalGeneration( self, pixel_values: torch.FloatTensor, spatial_shapes: torch.Tensor, - ) -> torch.Tensor: + ) -> list[torch.Tensor]: assert spatial_shapes.device.type == "cpu", ( "Expected `spatial_shapes` on CPU to avoid device-to-host sync in " "variable-length packing." @@ -759,23 +790,13 @@ class Lfm2VLForConditionalGeneration( ) vision_features_packed = image_outputs_packed[0] - factor = self.multi_modal_projector.factor - projected_lengths_list: list[int] = [] - for (height, width), length in zip(spatial_shapes_list, lengths_list): - if length <= 0: - projected_lengths_list.append(0) - continue - if height % factor != 0 or width % factor != 0: - raise ValueError( - "spatial_shapes must be divisible by downsample_factor: " - f"got ({height}, {width}) with factor={factor}." - ) - projected_lengths_list.append((height // factor) * (width // factor)) - projected_packed = self.multi_modal_projector( vision_features_packed=vision_features_packed, spatial_shapes=spatial_shapes, ) + projected_lengths_list = self._get_lfm2vl_tile_output_lengths( + spatial_shapes_list + ) image_features: list[torch.Tensor] = [] offset = 0 @@ -819,6 +840,383 @@ class Lfm2VLForConditionalGeneration( return self._process_image_input(image_input) + def get_encoder_cudagraph_config(self): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphConfig, + ) + + return EncoderCudaGraphConfig( + modalities=["image"], + buffer_keys=[ + "pixel_values_packed", + "pos_embeds", + "cu_seqlens", + "max_seqlen", + "gather_idx", + ], + out_hidden_size=self.config.text_config.hidden_size, + padding_logics={ + "cu_seqlens": _pad_cumulative_seqlens_buffer, + }, + ) + + def get_max_frames_per_video(self) -> int: + return 0 + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: VllmConfig, + ) -> tuple[int, int]: + min_budget = self._get_lfm2vl_min_image_tokens() + max_budget = min( + vllm_config.scheduler_config.max_num_batched_tokens, + self.model_config.max_model_len, + ) + return min_budget, max_budget + + def _get_spatial_shapes_list( + self, + spatial_shapes: torch.Tensor, + ) -> list[list[int]]: + assert spatial_shapes.device.type == "cpu", ( + "Expected `spatial_shapes` on CPU to avoid device-to-host sync in " + "variable-length packing." + ) + return spatial_shapes.tolist() + + @staticmethod + def _get_lfm2vl_tile_input_lengths( + spatial_shapes_list: list[list[int]], + ) -> list[int]: + return [height * width for height, width in spatial_shapes_list] + + def _get_lfm2vl_tile_output_lengths( + self, + spatial_shapes_list: list[list[int]], + ) -> list[int]: + factor = self.multi_modal_projector.factor + output_lengths: list[int] = [] + for height, width in spatial_shapes_list: + if height % factor != 0 or width % factor != 0: + raise ValueError( + "spatial_shapes must be divisible by downsample_factor: " + f"got ({height}, {width}) with factor={factor}." + ) + output_lengths.append((height // factor) * (width // factor)) + return output_lengths + + def _get_lfm2vl_mm_processor_kwargs(self) -> Mapping[str, object]: + return self.multimodal_config.mm_processor_kwargs or {} + + def _get_lfm2vl_min_image_tokens(self) -> int: + value = self._get_lfm2vl_mm_processor_kwargs().get( + "min_image_tokens", + getattr(self.config, "min_image_tokens", None) or 64, + ) + return max(1, int(value)) + + def _get_lfm2vl_item_tile_slices( + self, + num_patches: torch.Tensor, + ) -> list[tuple[int, int]]: + num_patches_list = [int(x) for x in num_patches.tolist()] + starts = [0] + for count in num_patches_list: + starts.append(starts[-1] + count) + return list(zip(starts[:-1], starts[1:])) + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ): + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + spatial_shapes = mm_kwargs["spatial_shapes"] + num_patches = mm_kwargs["num_patches"] + spatial_shapes_list = self._get_spatial_shapes_list(spatial_shapes) + input_lengths = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list) + output_lengths = self._get_lfm2vl_tile_output_lengths(spatial_shapes_list) + + return [ + EncoderItemSpec( + input_size=sum(input_lengths[start:end]), + output_tokens=sum(output_lengths[start:end]), + ) + for start, end in self._get_lfm2vl_item_tile_slices(num_patches) + ] + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + pixel_values = mm_kwargs["pixel_values"] + spatial_shapes = mm_kwargs["spatial_shapes"] + num_patches = mm_kwargs["num_patches"] + + tile_slices = self._get_lfm2vl_item_tile_slices(num_patches) + + if len(indices) == 0: + return { + "pixel_values": pixel_values[:0], + "spatial_shapes": spatial_shapes[:0], + "num_patches": num_patches[:0], + } + + tile_indices: list[int] = [] + for image_idx in indices: + start, end = tile_slices[image_idx] + tile_indices.extend(range(start, end)) + + return { + "pixel_values": pixel_values[tile_indices], + "spatial_shapes": spatial_shapes[tile_indices], + "num_patches": num_patches[indices], + } + + def _pack_lfm2vl_pixel_values( + self, + pixel_values: torch.Tensor, + spatial_shapes_list: list[list[int]], + ) -> torch.Tensor: + input_lengths = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list) + total_tokens = sum(input_lengths) + packed = pixel_values.new_empty((total_tokens, pixel_values.shape[-1])) + + offset = 0 + for i, length in enumerate(input_lengths): + if length <= 0: + continue + packed[offset : offset + length].copy_(pixel_values[i, :length]) + offset += length + return packed + + def _get_lfm2vl_pos_embeds( + self, + spatial_shapes: torch.Tensor, + spatial_shapes_list: list[list[int]], + ) -> torch.Tensor: + embeddings = self.vision_tower.vision_model.embeddings + positional_embeddings = embeddings.position_embedding.weight.reshape( + embeddings.position_embedding_size, + embeddings.position_embedding_size, + -1, + ) + lengths_list = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list) + return embeddings.resize_positional_embeddings_packed( + positional_embeddings, + spatial_shapes, + lengths_list=lengths_list, + ) + + def _get_lfm2vl_cu_seqlens( + self, + spatial_shapes_list: list[list[int]], + device: torch.device, + ) -> torch.Tensor: + lengths = torch.tensor( + self._get_lfm2vl_tile_input_lengths(spatial_shapes_list), + dtype=torch.int32, + device=device, + ) + cu_seqlens = torch.zeros( + lengths.shape[0] + 1, + dtype=torch.int32, + device=device, + ) + if lengths.numel() > 0: + cu_seqlens[1:] = torch.cumsum(lengths, dim=0) + return cu_seqlens + + def _get_lfm2vl_max_seqlen( + self, + spatial_shapes_list: list[list[int]], + ) -> torch.Tensor: + input_lengths = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list) + max_seqlen = max(input_lengths) if input_lengths else 0 + return torch.tensor(max_seqlen, dtype=torch.int32) + + def _get_lfm2vl_projector_gather_idx( + self, + spatial_shapes_list: list[list[int]], + device: torch.device, + ) -> torch.Tensor: + factor = self.multi_modal_projector.factor + dh = torch.arange(factor, dtype=torch.int64) + dw = torch.arange(factor, dtype=torch.int64) + dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij") + dh_flat = dh_grid.reshape(-1) + dw_flat = dw_grid.reshape(-1) + + gather_idx_parts: list[torch.Tensor] = [] + offset = 0 + for height, width in spatial_shapes_list: + length = height * width + if length <= 0: + continue + if height % factor != 0 or width % factor != 0: + raise ValueError( + "spatial_shapes must be divisible by downsample_factor: " + f"got ({height}, {width}) with factor={factor}." + ) + + rows_out = torch.arange(height // factor, dtype=torch.int64) + cols_out = torch.arange(width // factor, dtype=torch.int64) + rr, cc = torch.meshgrid(rows_out, cols_out, indexing="ij") + rr = rr.reshape(-1) + cc = cc.reshape(-1) + token_idx = (rr[:, None] * factor + dh_flat[None, :]) * width + ( + cc[:, None] * factor + dw_flat[None, :] + ) + gather_idx_parts.append(token_idx.reshape(-1) + offset) + offset += length + + if not gather_idx_parts: + return torch.empty(0, dtype=torch.int64, device=device) + return torch.cat(gather_idx_parts).to(device=device) + + def _prepare_lfm2vl_cudagraph_values( + self, + pixel_values: torch.Tensor, + spatial_shapes: torch.Tensor, + ) -> dict[str, torch.Tensor]: + spatial_shapes_list = self._get_spatial_shapes_list(spatial_shapes) + pixel_values_packed = self._pack_lfm2vl_pixel_values( + pixel_values, + spatial_shapes_list, + ) + pos_embeds = self._get_lfm2vl_pos_embeds(spatial_shapes, spatial_shapes_list) + device = pixel_values.device + + return { + "pixel_values_packed": pixel_values_packed, + "pos_embeds": pos_embeds, + "cu_seqlens": self._get_lfm2vl_cu_seqlens(spatial_shapes_list, device), + "max_seqlen": self._get_lfm2vl_max_seqlen(spatial_shapes_list), + "gather_idx": self._get_lfm2vl_projector_gather_idx( + spatial_shapes_list, + device, + ), + } + + def _get_lfm2vl_capture_spatial_shapes( + self, + token_budget: int, + ) -> torch.Tensor: + factor = self.multi_modal_projector.factor + min_image_tokens = self._get_lfm2vl_min_image_tokens() + remaining = token_budget + shapes: list[list[int]] = [] + + while remaining > 0: + out_tokens = min(remaining, min_image_tokens) + shapes.append([factor, out_tokens * factor]) + remaining -= out_tokens + + return torch.tensor(shapes, dtype=torch.int64) + + 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, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + spatial_shapes = self._get_lfm2vl_capture_spatial_shapes(token_budget) + spatial_shapes_list = self._get_spatial_shapes_list(spatial_shapes) + input_lengths = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list) + total_input_tokens = sum(input_lengths) + + patch_dim = ( + self.vision_tower.vision_model.embeddings.patch_embedding.weight.shape[1] + ) + dummy_pixel_values = torch.randn( + total_input_tokens, + patch_dim, + device=device, + dtype=dtype, + ) + pos_embeds = self._get_lfm2vl_pos_embeds( + spatial_shapes, + spatial_shapes_list, + ).to(device=device, dtype=dtype) + + # max_seqlen.item() is baked into the captured ViT attention graph, so + # capture with a budget-level upper bound that covers any replay item. + max_tile_input_tokens = token_budget * self.multi_modal_projector.factor**2 + values = { + "pixel_values_packed": dummy_pixel_values, + "pos_embeds": pos_embeds, + "cu_seqlens": self._get_lfm2vl_cu_seqlens(spatial_shapes_list, device), + "max_seqlen": torch.tensor(max_tile_input_tokens, dtype=torch.int32), + "gather_idx": self._get_lfm2vl_projector_gather_idx( + spatial_shapes_list, + device, + ), + } + + 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, + ): + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + values = self._prepare_lfm2vl_cudagraph_values( + mm_kwargs["pixel_values"], + mm_kwargs["spatial_shapes"], + ) + return EncoderCudaGraphReplayBuffers(values=values) + + def encoder_cudagraph_forward( + self, + values: dict[str, torch.Tensor], + ) -> torch.Tensor: + embeddings = self.vision_tower.vision_model.embeddings + pixel_values = values["pixel_values_packed"].to( + dtype=embeddings.patch_embedding.weight.dtype + ) + patch_embeds = embeddings.patch_embedding(pixel_values) + hidden_states = (patch_embeds + values["pos_embeds"]).unsqueeze(0) + + with set_forward_context(None, self.vllm_config): + encoder_outputs = self.vision_tower.vision_model.encoder( + inputs_embeds=hidden_states, + cu_seqlens=values["cu_seqlens"], + max_seqlen=values["max_seqlen"], + ) + + post_layernorm = self.vision_tower.vision_model.post_layernorm + if post_layernorm is not None: + encoder_outputs = post_layernorm(encoder_outputs) + + return self.multi_modal_projector.forward_with_gather_idx( + vision_features_packed=encoder_outputs[0], + gather_idx=values["gather_idx"], + ) + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + ) -> torch.Tensor: + image_input = LFM2VLImageInputs( + type="pixel_values", + pixel_values=mm_kwargs["pixel_values"], + spatial_shapes=mm_kwargs["spatial_shapes"], + num_patches=mm_kwargs["num_patches"], + ) + return torch.cat(self._process_image_input(image_input), dim=0) + def forward( self, input_ids: torch.Tensor | None,