[Bugfix] Fix Gemma 4 encoder CUDA graph capture OOM at startup (#61365)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Vadim Gimpelson
co-authored by
Claude
Claude
parent
2b85d9cb65
commit
ec8b20fbf0
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
@@ -17,6 +18,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.cache import MultiModalProcessorOnlyCache
|
||||
from vllm.multimodal.inputs import MultiModalFieldConfig
|
||||
from vllm.utils.mem_constants import GiB_bytes
|
||||
from vllm.v1.worker.encoder_cudagraph_defs import ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG
|
||||
|
||||
from ....conftest import ImageTestAssets
|
||||
from ...utils import build_model_context
|
||||
@@ -348,3 +350,56 @@ def test_encoder_chunk_no_free_memory_falls_back_to_one():
|
||||
)
|
||||
== 1
|
||||
)
|
||||
|
||||
|
||||
def _gemma4_for_encoder_cudagraph() -> Gemma4ForConditionalGeneration:
|
||||
model = object.__new__(Gemma4ForConditionalGeneration)
|
||||
vision_config = SimpleNamespace(
|
||||
patch_size=16, pooling_kernel_size=3, default_output_length=280
|
||||
)
|
||||
model.vision_tower = SimpleNamespace(config=vision_config)
|
||||
model.config = SimpleNamespace(vision_config=vision_config)
|
||||
return model
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("token_budget", "row_tokens"), [(4096, 70), (4096, 280), (70, 280)]
|
||||
)
|
||||
def test_encoder_cudagraph_capture_rows_fit_token_budget(
|
||||
token_budget: int, row_tokens: int
|
||||
):
|
||||
# Regression for #61045: rows were max_frames_per_batch (58 * 32) at
|
||||
# 280 tokens each, so the 4096-token buffer alone was 6.69 GiB.
|
||||
model = _gemma4_for_encoder_cudagraph()
|
||||
inputs = model.prepare_encoder_cudagraph_capture_inputs(
|
||||
token_budget=token_budget,
|
||||
max_batch_size=58,
|
||||
max_frames_per_batch=58 * 32,
|
||||
device="meta",
|
||||
dtype=torch.bfloat16,
|
||||
axis_keys=(row_tokens,),
|
||||
)
|
||||
|
||||
num_rows, row_patches, _ = inputs.values["pixel_values"].shape
|
||||
assert row_patches == row_tokens * 9
|
||||
assert num_rows == max(1, token_budget // row_tokens)
|
||||
|
||||
|
||||
def test_encoder_cudagraph_mixed_width_images_pack_by_padded_width():
|
||||
model = _gemma4_for_encoder_cudagraph()
|
||||
mm_kwargs = {
|
||||
"pixel_values": [torch.zeros(2520, 768), torch.zeros(630, 768)],
|
||||
"pixel_position_ids": [
|
||||
torch.zeros(2520, 2, dtype=torch.long),
|
||||
torch.zeros(630, 2, dtype=torch.long),
|
||||
],
|
||||
}
|
||||
|
||||
specs = model.get_encoder_cudagraph_item_specs(mm_kwargs)
|
||||
assert [spec.output_tokens for spec in specs] == [280, 70]
|
||||
assert [spec.get_path_output_tokens("default") for spec in specs] == [280, 280]
|
||||
|
||||
selected = model.select_encoder_cudagraph_items(mm_kwargs, [1])
|
||||
assert selected[ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG] == (70,)
|
||||
selected = model.select_encoder_cudagraph_items(mm_kwargs, [0, 1])
|
||||
assert selected[ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG] == (280,)
|
||||
|
||||
@@ -1663,6 +1663,13 @@ class Gemma4ForConditionalGeneration(
|
||||
batch_size, num_patches = src.shape[0], src.shape[1]
|
||||
dst[:batch_size, :num_patches].copy_(src)
|
||||
|
||||
configured_max_tokens = self.config.vision_config.default_output_length
|
||||
# One graph per row width (soft tokens of one image or video frame),
|
||||
# so a budget's rows are not all padded to the widest image.
|
||||
row_widths = tuple(
|
||||
t for t in _SUPPORTED_SOFT_TOKENS if t <= configured_max_tokens
|
||||
)
|
||||
|
||||
return EncoderCudaGraphConfig(
|
||||
modalities=["image", "video"],
|
||||
buffer_keys=[
|
||||
@@ -1676,6 +1683,7 @@ class Gemma4ForConditionalGeneration(
|
||||
"pixel_values": pad_pixel_values,
|
||||
"pixel_position_ids": pad_pixel_position_ids,
|
||||
},
|
||||
capture_axes=(row_widths,),
|
||||
)
|
||||
|
||||
def get_encoder_cudagraph_budget_range(
|
||||
@@ -1712,10 +1720,15 @@ class Gemma4ForConditionalGeneration(
|
||||
if modality == "image":
|
||||
pixel_values = mm_kwargs["pixel_values"]
|
||||
if isinstance(pixel_values, list):
|
||||
# A batch pads its images to the widest one, so pack by that.
|
||||
row_tokens = (
|
||||
max((pv.shape[0] for pv in pixel_values), default=0) // pool_ratio
|
||||
)
|
||||
return [
|
||||
EncoderItemSpec(
|
||||
input_size=pv.shape[0],
|
||||
output_tokens=pv.shape[0] // pool_ratio,
|
||||
path_output_tokens={"default": row_tokens},
|
||||
)
|
||||
for pv in pixel_values
|
||||
]
|
||||
@@ -1750,10 +1763,25 @@ class Gemma4ForConditionalGeneration(
|
||||
mm_kwargs: dict[str, Any],
|
||||
indices: list[int],
|
||||
) -> dict[str, Any]:
|
||||
from vllm.v1.worker.encoder_cudagraph_defs import (
|
||||
ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG,
|
||||
)
|
||||
|
||||
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"]
|
||||
pixel_position_ids = mm_kwargs["pixel_position_ids"]
|
||||
if isinstance(pixel_values, list):
|
||||
row_patches = max(
|
||||
(pixel_values[i].shape[0] for i in indices), default=0
|
||||
)
|
||||
else:
|
||||
row_patches = pixel_values.shape[1]
|
||||
axis_keys = (row_patches // pool_ratio,)
|
||||
|
||||
if len(indices) == 0:
|
||||
is_pv_list = isinstance(pixel_values, list)
|
||||
is_pp_list = isinstance(pixel_position_ids, list)
|
||||
@@ -1762,20 +1790,24 @@ class Gemma4ForConditionalGeneration(
|
||||
"pixel_position_ids": (
|
||||
[] if is_pp_list else pixel_position_ids[:0]
|
||||
),
|
||||
ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
|
||||
}
|
||||
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],
|
||||
ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
|
||||
}
|
||||
return {
|
||||
"pixel_values": pixel_values[indices],
|
||||
"pixel_position_ids": pixel_position_ids[indices],
|
||||
ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
|
||||
}
|
||||
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"]
|
||||
axis_keys = (pixel_values_videos.shape[1] // pool_ratio,)
|
||||
|
||||
if len(indices) == 0:
|
||||
is_fc_tensor = isinstance(video_frame_counts, torch.Tensor)
|
||||
@@ -1785,6 +1817,7 @@ class Gemma4ForConditionalGeneration(
|
||||
"video_frame_counts": (
|
||||
video_frame_counts[:0] if is_fc_tensor else []
|
||||
),
|
||||
ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
|
||||
}
|
||||
|
||||
fc_list = (
|
||||
@@ -1819,6 +1852,7 @@ class Gemma4ForConditionalGeneration(
|
||||
"pixel_values_videos": selected_pv,
|
||||
"pixel_position_ids_videos": selected_pp,
|
||||
"video_frame_counts": selected_fc,
|
||||
ENCODER_CUDAGRAPH_AXIS_KEYS_KWARG: axis_keys,
|
||||
}
|
||||
raise ValueError(f"Unknown modality: {modality}")
|
||||
|
||||
@@ -1839,20 +1873,21 @@ class Gemma4ForConditionalGeneration(
|
||||
|
||||
dtype = dtype or torch.float32
|
||||
|
||||
max_size = max(max_batch_size, max_frames_per_batch)
|
||||
assert axis_keys is not None
|
||||
(per_item_output,) = axis_keys
|
||||
assert isinstance(per_item_output, int)
|
||||
|
||||
# Every row of a batch replayed on this graph is an image or video
|
||||
# frame of per_item_output tokens, so a batch that fits token_budget
|
||||
# never needs more rows than this. Capture needs at least one row.
|
||||
max_size = min(
|
||||
max(max_batch_size, max_frames_per_batch),
|
||||
max(1, token_budget // per_item_output),
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user