From 70009fb9344d2a7ba642e68369e0a64a6252e8bc Mon Sep 17 00:00:00 2001 From: anthonsu <50185138+anthonsu@users.noreply.github.com> Date: Sat, 25 Jul 2026 13:02:09 -0700 Subject: [PATCH] [MM][CG] Support ViT CUDA Graph for Gemma-4 (#46837) Signed-off-by: Anthony Su Co-authored-by: Linkun Chen --- docs/design/cuda_graphs_multimodal.md | 1 + .../multimodal/vision_language_offline.py | 41 ++ .../generation/test_vit_cudagraph.py | 13 + vllm/model_executor/models/gemma4_mm.py | 422 +++++++++++++++++- 4 files changed, 476 insertions(+), 1 deletion(-) diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 7eab425d6e2..7502186ae46 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` | ✅︎ | ❌︎ | ✅︎ | | `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | | `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | +| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | | `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | | `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | | `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 661401046bd..d8e08835197 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -503,6 +503,45 @@ def run_gemma3n(questions: list[str], modality: str) -> ModelRequestData: ) +# Gemma 4 +def run_gemma4(questions: list[str], modality: str) -> ModelRequestData: + assert modality in ("image", "video") + model_name = "google/gemma-4-31B-it" + + # NOTE: Gemma-4-31B is a large model. Users running into Out-Of-Memory (OOM) + # errors might need to set `tensor_parallel_size` to > 1. + engine_args = EngineArgs( + model=model_name, + max_model_len=4096, + max_num_seqs=2, + limit_mm_per_prompt={modality: 1}, + ) + + if modality == "image": + prompts = [ + ( + "user\n" + f"<|image|>\n{question}\n" + "model\n" + ) + for question in questions + ] + else: # video + prompts = [ + ( + "user\n" + f"<|video|>\n{question}\n" + "model\n" + ) + for question in questions + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + # GLM-4v def run_glm4v(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -2303,6 +2342,7 @@ model_example_map = { "exaone4_5": run_exaone4_5, "gemma3": run_gemma3, "gemma3n": run_gemma3n, + "gemma4": run_gemma4, "glm4v": run_glm4v, "glm4_1v": run_glm4_1v, "glm4_5v": run_glm4_5v, @@ -2374,6 +2414,7 @@ MODELS_NEED_VIDEO_METADATA = [ MODELS_SUPPORT_VIT_CUDA_GRAPH = [ "llama4", + "gemma4", "qwen2_vl", "qwen2_5_vl", "qwen3_vl", diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 387046f3e74..13fbe26880d 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -249,6 +249,19 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { }, skip=True, # TODO: Re-enable this once OOM issues are resolved on CI. ), + "gemma4": VitCudagraphTestConfig( + model="google/gemma-4-E2B-it", + image_prompt=( + "user\n<|image|>\nWhat is in this image?\n" + "model\n" + ), + video_prompt=( + "user\n<|video|>\nDescribe this video in one sentence." + "\nmodel\n" + ), + needs_video_metadata=True, + marks=[pytest.mark.core_model], + ), } diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index b0169cbc6dd..e158a785016 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -16,7 +16,7 @@ reason about temporal order. import math from collections.abc import Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, Annotated, Any, Literal +from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal import numpy as np import torch @@ -72,6 +72,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( MultiModalEmbeddings, SupportsEagle3, + SupportsEncoderCudaGraph, SupportsLoRA, SupportsMultiModal, SupportsPP, @@ -86,6 +87,12 @@ from .utils import ( if TYPE_CHECKING: from vllm.model_executor.layers.quantization import QuantizationConfig + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + EncoderCudaGraphConfig, + EncoderCudaGraphReplayBuffers, + EncoderItemSpec, + ) logger = init_logger(__name__) @@ -980,7 +987,9 @@ class Gemma4ForConditionalGeneration( SupportsPP, SupportsLoRA, SupportsEagle3, + SupportsEncoderCudaGraph, ): + supports_encoder_cudagraph: ClassVar[Literal[True]] = True # Gemma4 clamps mm_prefix bidirectional ranges to the sliding window # in-kernel (HF's (causal OR blockwise) AND sliding_window). The model # runner reads this to keep image bidirectional ranges that exceed the @@ -1025,6 +1034,7 @@ class Gemma4ForConditionalGeneration( self.quant_config = quant_config self.multimodal_config = multimodal_config self.model_dtype = vllm_config.model_config.dtype + self.vllm_config = vllm_config # Only quantize towers when the quant method supports their # dimensions. BNB/torchao handle arbitrary sizes; other methods @@ -1525,6 +1535,416 @@ class Gemma4ForConditionalGeneration( return multimodal_embeddings + # ------------------------------------------------------------------ # + # EncoderCudaGraph protocol methods + # ------------------------------------------------------------------ # + + def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": + from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig + + def pad_pixel_values(dst: torch.Tensor, src: torch.Tensor) -> None: + dst.zero_() + batch_size, num_patches = src.shape[0], src.shape[1] + dst[:batch_size, :num_patches].copy_(src) + + def pad_pixel_position_ids(dst: torch.Tensor, src: torch.Tensor) -> None: + dst.fill_(-1) + batch_size, num_patches = src.shape[0], src.shape[1] + dst[:batch_size, :num_patches].copy_(src) + + return EncoderCudaGraphConfig( + modalities=["image", "video"], + buffer_keys=[ + "pixel_values", + "pixel_position_ids", + "gather_indices", + ], + out_hidden_size=self.config.text_config.hidden_size, + max_frames_per_video=_VIDEO_MAX_FRAMES, + padding_logics={ + "pixel_values": pad_pixel_values, + "pixel_position_ids": pad_pixel_position_ids, + }, + ) + + def get_encoder_cudagraph_budget_range( + self, + vllm_config: VllmConfig, + ) -> tuple[int, int]: + min_budget = _SUPPORTED_SOFT_TOKENS[0] + 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_input_modality(self, mm_kwargs: dict[str, Any]) -> str: + if "pixel_values" in mm_kwargs: + return "image" + elif "pixel_values_videos" in mm_kwargs: + return "video" + raise ValueError("Unsupported modality in mm_kwargs") + + def get_max_frames_per_video(self) -> int: + return _VIDEO_MAX_FRAMES + + def get_encoder_cudagraph_item_specs( + self, + mm_kwargs: dict[str, Any], + ) -> list["EncoderItemSpec"]: + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec + + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + if isinstance(pixel_values, list): + return [ + EncoderItemSpec( + input_size=pv.shape[0], + output_tokens=pv.shape[0] // pool_ratio, + ) + for pv in pixel_values + ] + else: + return [ + EncoderItemSpec( + input_size=pixel_values.shape[1], + output_tokens=pixel_values.shape[1] // pool_ratio, + ) + for _ in range(pixel_values.shape[0]) + ] + elif modality == "video": + pixel_values_videos = mm_kwargs["pixel_values_videos"] + video_frame_counts = mm_kwargs["video_frame_counts"] + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + np_patches = pixel_values_videos.shape[1] + return [ + EncoderItemSpec( + input_size=fc * np_patches, + output_tokens=fc * (np_patches // pool_ratio), + ) + for fc in fc_list + ] + raise ValueError(f"Unknown modality: {modality}") + + def select_encoder_cudagraph_items( + self, + mm_kwargs: dict[str, Any], + indices: list[int], + ) -> dict[str, Any]: + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + pixel_position_ids = mm_kwargs["pixel_position_ids"] + if len(indices) == 0: + is_pv_list = isinstance(pixel_values, list) + is_pp_list = isinstance(pixel_position_ids, list) + return { + "pixel_values": ([] if is_pv_list else pixel_values[:0]), + "pixel_position_ids": ( + [] if is_pp_list else pixel_position_ids[:0] + ), + } + if isinstance(pixel_values, list): + return { + "pixel_values": [pixel_values[i] for i in indices], + "pixel_position_ids": [pixel_position_ids[i] for i in indices], + } + return { + "pixel_values": pixel_values[indices], + "pixel_position_ids": pixel_position_ids[indices], + } + elif modality == "video": + pixel_values_videos = mm_kwargs["pixel_values_videos"] + pixel_position_ids_videos = mm_kwargs["pixel_position_ids_videos"] + video_frame_counts = mm_kwargs["video_frame_counts"] + + if len(indices) == 0: + is_fc_tensor = isinstance(video_frame_counts, torch.Tensor) + return { + "pixel_values_videos": pixel_values_videos[:0], + "pixel_position_ids_videos": pixel_position_ids_videos[:0], + "video_frame_counts": ( + video_frame_counts[:0] if is_fc_tensor else [] + ), + } + + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + cum_frames = [0] + for fc in fc_list: + cum_frames.append(cum_frames[-1] + fc) + + selected_pv = torch.cat( + [ + pixel_values_videos[cum_frames[i] : cum_frames[i + 1]] + for i in indices + ], + dim=0, + ) + selected_pp = torch.cat( + [ + pixel_position_ids_videos[cum_frames[i] : cum_frames[i + 1]] + for i in indices + ], + dim=0, + ) + selected_fc = ( + video_frame_counts[indices] + if isinstance(video_frame_counts, torch.Tensor) + else [video_frame_counts[i] for i in indices] + ) + return { + "pixel_values_videos": selected_pv, + "pixel_position_ids_videos": selected_pp, + "video_frame_counts": selected_fc, + } + raise ValueError(f"Unknown modality: {modality}") + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int = 256, + max_batch_size: int = 4, + max_frames_per_batch: int = 1, + device: torch.device | str = "cpu", + dtype: torch.dtype | None = None, + path: str = "default", + **kwargs: Any, + ) -> "EncoderCudaGraphCaptureInputs": + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphCaptureInputs, + ) + + dtype = dtype or torch.float32 + + max_size = max(max_batch_size, max_frames_per_batch) + + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + + # Retrieve the model's actual configured maximum tokens: + configured_max_tokens = getattr( + self.config.vision_config, + "num_soft_tokens", + _SUPPORTED_SOFT_TOKENS[2], + ) + # Dynamically compute the slot capacity per item bounded by both the + # current graph budget and the user's maximum config: + per_item_output = min(token_budget, configured_max_tokens) + # Satisfy k^2 * per_item_output = per_item_patches + per_item_patches = per_item_output * pool_ratio + + patch_size = self.vision_tower.config.patch_size + num_channels = getattr(self.vision_tower.config, "num_channels", 3) + patch_pixels = (patch_size**2) * num_channels + + dummy_pixel_values = torch.zeros( + (max_size, per_item_patches, patch_pixels), + device=device, + dtype=dtype, + ) + dummy_pixel_position_ids = torch.full( + (max_size, per_item_patches, 2), + -1, + device=device, + dtype=torch.long, + ) + dummy_gather_indices = torch.zeros( + (token_budget,), + device=device, + dtype=torch.long, + ) + + return EncoderCudaGraphCaptureInputs( + values={ + "pixel_values": dummy_pixel_values, + "pixel_position_ids": dummy_pixel_position_ids, + "gather_indices": dummy_gather_indices, + } + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int = 4, + max_frames_per_batch: int = 1, + path: str = "default", + **kwargs: Any, + ) -> "EncoderCudaGraphReplayBuffers": + from vllm.v1.worker.encoder_cudagraph_defs import ( + EncoderCudaGraphReplayBuffers, + ) + + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + pixel_values = mm_kwargs["pixel_values"] + pixel_position_ids = mm_kwargs["pixel_position_ids"] + elif modality == "video": + pixel_values = mm_kwargs["pixel_values_videos"] + pixel_position_ids = mm_kwargs["pixel_position_ids_videos"] + else: + raise ValueError(f"Unsupported modality: {modality}") + + if isinstance(pixel_values, list): + max_patches = max(pv.shape[0] for pv in pixel_values) + batch_size = len(pixel_values) + pv_tensor = torch.zeros( + (batch_size, max_patches, pixel_values[0].shape[1]), + dtype=pixel_values[0].dtype, + device=pixel_values[0].device, + ) + pp_tensor = torch.full( + (batch_size, max_patches, 2), + -1, + dtype=pixel_position_ids[0].dtype, + device=pixel_position_ids[0].device, + ) + for i, (pv, pp) in enumerate(zip(pixel_values, pixel_position_ids)): + pv_tensor[i, : pv.shape[0]].copy_(pv) + pp_tensor[i, : pp.shape[0]].copy_(pp) + pixel_values = pv_tensor + pixel_position_ids = pp_tensor + + item_specs = self.get_encoder_cudagraph_item_specs(mm_kwargs) + per_item_out_tokens = [spec.output_tokens for spec in item_specs] + total_tokens = sum(per_item_out_tokens) + + device = pixel_values.device + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + per_item_output = pixel_values.shape[1] // pool_ratio + + # ONLY allocate an array of exact size `total_tokens`. + # DO NOT pad it. The upstream Graph Manager handles the padding securely. + gather_indices = torch.zeros((total_tokens,), dtype=torch.long, device=device) + + if modality == "image": + dst_offset = 0 + for i, n_tok in enumerate(per_item_out_tokens): + safe_n_tok = min(n_tok, per_item_output) + src_start = i * per_item_output + src_end = src_start + safe_n_tok + gather_indices[dst_offset : dst_offset + safe_n_tok] = torch.arange( + src_start, src_end, dtype=torch.long, device=device + ) + dst_offset += safe_n_tok + elif modality == "video": + video_frame_counts = mm_kwargs["video_frame_counts"] + fc_list = ( + video_frame_counts.tolist() + if isinstance(video_frame_counts, torch.Tensor) + else list(video_frame_counts) + ) + vision_cfg = self.vision_tower.config + pool_ratio = getattr(vision_cfg, "pooling_kernel_size", 2) ** 2 + np_patches = pixel_values.shape[1] + frame_output_tokens = np_patches // pool_ratio + safe_frame_output_tokens = min(frame_output_tokens, per_item_output) + + dst_offset = 0 + frame_idx = 0 + for fc in fc_list: + for f in range(fc): + src_start = (frame_idx + f) * per_item_output + src_end = src_start + safe_frame_output_tokens + gather_indices[ + dst_offset : dst_offset + safe_frame_output_tokens + ] = torch.arange( + src_start, src_end, dtype=torch.long, device=device + ) + dst_offset += safe_frame_output_tokens + frame_idx += fc + + return EncoderCudaGraphReplayBuffers( + values={ + "pixel_values": pixel_values, + "pixel_position_ids": pixel_position_ids, + "gather_indices": gather_indices, + } + ) + + def encoder_cudagraph_forward( + self, + inputs: dict[str, torch.Tensor], + path: str = "default", + **kwargs: Any, + ) -> torch.Tensor: + pixel_values = inputs["pixel_values"] + pixel_position_ids = inputs["pixel_position_ids"] + gather_indices = inputs["gather_indices"] + + pad_tensor = (pixel_position_ids == -1).all(dim=-1) + + vt = self.vision_tower + inputs_embeds = vt.patch_embedder( + pixel_values, + pixel_position_ids, + pad_tensor, + ).to(self.model_dtype) + + encoder_outputs = vt.encoder( + inputs_embeds=inputs_embeds, + attention_mask=~pad_tensor, + pixel_position_ids=pixel_position_ids, + ) + hidden_states = encoder_outputs.last_hidden_state + + pool_ratio = getattr(vt.config, "pooling_kernel_size", 2) ** 2 + per_item_output = pixel_values.shape[1] // pool_ratio + + pooled_states, _ = vt.pooler( + hidden_states=hidden_states, + pixel_position_ids=pixel_position_ids, + padding_positions=pad_tensor, + output_length=per_item_output, + ) + + if getattr(vt.config, "standardize", False): + pooled_states = (pooled_states - vt.std_bias) * vt.std_scale + + flat_pooled = pooled_states.reshape(-1, pooled_states.shape[-1]) + gathered_states = flat_pooled[gather_indices] + + # Cast to the projection layer's dtype to resolve mixed-precision crash + target_dtype = self.embed_vision.embedding_projection.weight.dtype + gathered_states = gathered_states.to(target_dtype) + + flat_proj_embs = self.embed_vision( + inputs_embeds=gathered_states.unsqueeze(0) + ).squeeze(0) + + return flat_proj_embs + + def encoder_eager_forward( + self, + mm_kwargs: dict[str, Any], + path: str = "default", + **kwargs: Any, + ) -> torch.Tensor: + modality = self.get_input_modality(mm_kwargs) + if modality == "image": + image_input = self._parse_and_validate_image_input(**mm_kwargs) + assert image_input is not None + embeddings = self._process_image_input(image_input) + elif modality == "video": + video_input = self._parse_and_validate_video_input(**mm_kwargs) + assert video_input is not None + embeddings = self._process_video_input(video_input) + else: + raise ValueError(f"Unsupported modality: {modality}") + + return torch.cat(embeddings, dim=0) + def embed_input_ids( self, input_ids: torch.Tensor,