karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[MM][CG] Support ViT CUDA Graph for Gemma-4 (#46837)

Browse Source 
Signed-off-by: Anthony Su <xsuanthony@gmail.com>
Co-authored-by: Linkun Chen <github@lkchen.net>
This commit is contained in:
anthonsu
2026-07-25 15:02:09 -05:00
committed by GitHub
co-authored by Linkun Chen
parent ee1d996367
commit 70009fb934
4 changed files with 476 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/cuda_graphs_multimodal.md
+1
View File
@@ -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` | ✅︎ | ❌︎ | ❌︎ |
examples/generate/multimodal/vision_language_offline.py
+41
View File
@@ -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 = [
(
"<bos><start_of_turn>user\n"
f"<|image|>\n{question}<end_of_turn>\n"
"<start_of_turn>model\n"
)
for question in questions
]
else: # video
prompts = [
(
"<bos><start_of_turn>user\n"
f"<|video|>\n{question}<end_of_turn>\n"
"<start_of_turn>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",
tests/models/multimodal/generation/test_vit_cudagraph.py
+13
View File
@@ -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=(
"<bos><start_of_turn>user\n<|image|>\nWhat is in this image?<end_of_turn>\n"
"<start_of_turn>model\n"
),
video_prompt=(
"<bos><start_of_turn>user\n<|video|>\nDescribe this video in one sentence."
"<end_of_turn>\n<start_of_turn>model\n"
),
needs_video_metadata=True,
marks=[pytest.mark.core_model],
),
}
vllm/model_executor/models/gemma4_mm.py
+421 -1
View File
@@ -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,