iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Multimodal] Simplify ViT CUDA graph interfaces (#41234)

Browse Source 
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
This commit is contained in:
Isotr0py
2026-05-22 22:31:00 +08:00
committed by GitHub
parent fb21d8b4f9
commit f0feb15e7f
9 changed files with 151 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/multimodal/generation/test_vit_cudagraph.py
-3
View File
@@ -10,7 +10,6 @@ from vllm.multimodal.video import sample_frames_from_video
from vllm.platforms import current_platform
from ....conftest import IMAGE_ASSETS, VIDEO_ASSETS
from ....utils import create_new_process_for_each_test
from ...utils import dummy_hf_overrides
from .vlm_utils.builders import sample_frames_with_video_metadata
@@ -142,7 +141,6 @@ def get_compilation_config(config: VitCudagraphTestConfig):
@pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS))
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
@create_new_process_for_each_test()
def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):
config = MODEL_CONFIGS[model_id]
@@ -184,7 +182,6 @@ def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):
@pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS))
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
@create_new_process_for_each_test()
def test_vit_cudagraph_video(model_id, vllm_runner, video_assets):
config = MODEL_CONFIGS[model_id]
tests/v1/cudagraph/test_encoder_cudagraph.py
+28 -29
View File
@@ -26,6 +26,7 @@ from vllm.v1.worker.encoder_cudagraph_defs import (
EncoderCudaGraphCaptureInputs,
EncoderCudaGraphConfig,
EncoderCudaGraphReplayBuffers,
EncoderItemSpec,
)
# ---------------------------------------------------------------------------
@@ -288,24 +289,18 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
# For tests: min=4, max=128 (small values for fast capture)
return (4, 128)
def get_encoder_cudagraph_num_items(
def get_encoder_cudagraph_item_specs(
self,
mm_kwargs: dict[str, Any],
) -> int:
return len(mm_kwargs["image_grid_thw"])
def get_encoder_cudagraph_per_item_output_tokens(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
) -> list[EncoderItemSpec]:
m = _SPATIAL_MERGE
return [t * (h // m) * (w // m) for t, h, w in mm_kwargs["image_grid_thw"]]
def get_encoder_cudagraph_per_item_input_sizes(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
return [t * h * w for t, h, w in mm_kwargs["image_grid_thw"]]
return [
EncoderItemSpec(
input_size=t * h * w,
output_tokens=t * (h // m) * (w // m),
)
for t, h, w in mm_kwargs["image_grid_thw"]
]
def select_encoder_cudagraph_items(
self,
@@ -591,19 +586,18 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
# Protocol overrides that depend on modality keys
# ------------------------------------------------------------------
def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int:
return len(self._get_grid_thw(mm_kwargs))
def get_encoder_cudagraph_per_item_output_tokens(
self, mm_kwargs: dict[str, Any]
) -> list[int]:
def get_encoder_cudagraph_item_specs(
self,
mm_kwargs: dict[str, Any],
) -> list[EncoderItemSpec]:
m = _SPATIAL_MERGE
return [t * (h // m) * (w // m) for t, h, w in self._get_grid_thw(mm_kwargs)]
def get_encoder_cudagraph_per_item_input_sizes(
self, mm_kwargs: dict[str, Any]
) -> list[int]:
return [t * h * w for t, h, w in self._get_grid_thw(mm_kwargs)]
return [
EncoderItemSpec(
input_size=t * h * w,
output_tokens=t * (h // m) * (w // m),
)
for t, h, w in self._get_grid_thw(mm_kwargs)
]
def select_encoder_cudagraph_items(
self, mm_kwargs: dict[str, Any], indices: list[int]
@@ -679,11 +673,16 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
def encoder_cudagraph_forward(
self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor]
self,
mm_kwargs: dict[str, Any],
buffers: dict[str, torch.Tensor],
) -> torch.Tensor:
return self._forward(self._get_pixel_values(mm_kwargs))
def encoder_eager_forward(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:
def encoder_eager_forward(
self,
mm_kwargs: dict[str, Any],
) -> torch.Tensor:
return self._forward(self._get_pixel_values(mm_kwargs))
vllm/model_executor/models/interfaces.py
+7 -21
View File
@@ -50,6 +50,7 @@ if TYPE_CHECKING:
EncoderCudaGraphCaptureInputs,
EncoderCudaGraphConfig,
EncoderCudaGraphReplayBuffers,
EncoderItemSpec,
)
else:
VllmConfig = object
@@ -1549,30 +1550,15 @@ class SupportsEncoderCudaGraph(Protocol):
"""
...
def get_encoder_cudagraph_num_items(
def get_encoder_cudagraph_item_specs(
self,
mm_kwargs: dict[str, Any],
) -> int:
"""Return the number of items (e.g. images) in the batch."""
...
) -> list["EncoderItemSpec"]:
"""Return specs describing each item in the batch.
def get_encoder_cudagraph_per_item_output_tokens(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
"""Return output token count for each item.
Used for greedy packing and DP load balancing.
"""
...
def get_encoder_cudagraph_per_item_input_sizes(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
"""Return input size (e.g. patch count) for each item.
Used for input tensor slicing offsets.
Replaces the former separate methods for num_items,
per_item_output_tokens, and per_item_input_sizes.
The manager derives all three from this single return value.
"""
...
vllm/model_executor/models/qwen2_5_vl.py
+24 -19
View File
@@ -1578,6 +1578,7 @@ class Qwen2_5_VLForConditionalGeneration(
# embedding formats between eager and cudagraph paths.
modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"]
max_frames = self.get_max_frames_per_video() if "video" in modalities else 1
return EncoderCudaGraphConfig(
modalities=modalities,
input_key_by_modality={
@@ -1595,6 +1596,7 @@ class Qwen2_5_VLForConditionalGeneration(
"max_seqlen_window",
],
out_hidden_size=self.visual.out_hidden_size,
max_frames_per_video=max_frames,
)
def get_input_modality(
@@ -1603,7 +1605,9 @@ class Qwen2_5_VLForConditionalGeneration(
) -> str:
if "image_grid_thw" in mm_kwargs:
return "image"
return "video"
elif "video_grid_thw" in mm_kwargs:
return "video"
raise AssertionError("This line should be unreachable.")
def get_max_frames_per_video(self) -> int:
mm_registry = MULTIMODAL_REGISTRY
@@ -1649,26 +1653,21 @@ class Qwen2_5_VLForConditionalGeneration(
grid_thw = grid_thw.tolist()
return grid_thw
def get_encoder_cudagraph_num_items(
def get_encoder_cudagraph_item_specs(
self,
mm_kwargs: dict[str, Any],
) -> int:
return len(self._get_grid_thw_by_modality(mm_kwargs))
):
from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
def get_encoder_cudagraph_per_item_output_tokens(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
m = self.visual.spatial_merge_size
grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
return [t * (h // m) * (w // m) for t, h, w in grid_thw]
def get_encoder_cudagraph_per_item_input_sizes(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
return [t * h * w for t, h, w in grid_thw]
return [
EncoderItemSpec(
input_size=t * h * w,
output_tokens=t * (h // m) * (w // m),
)
for t, h, w in grid_thw
]
def select_encoder_cudagraph_items(
self,
@@ -1684,11 +1683,13 @@ class Qwen2_5_VLForConditionalGeneration(
"pixel_values": pixel_values[:0],
"image_grid_thw": [],
}
else:
elif self.get_input_modality(mm_kwargs) == "video":
return {
"pixel_values_videos": pixel_values[:0],
"video_grid_thw": [],
}
else:
raise AssertionError("This line should be unreachable.")
# Compute cumulative patch offsets for slicing pixel_values
patches_per_item = [t * h * w for t, h, w in grid_thw]
@@ -1706,11 +1707,13 @@ class Qwen2_5_VLForConditionalGeneration(
"pixel_values": selected_pv,
"image_grid_thw": selected_grid,
}
else:
elif self.get_input_modality(mm_kwargs) == "video":
return {
"pixel_values_videos": selected_pv,
"video_grid_thw": selected_grid,
}
else:
raise AssertionError("This line should be unreachable.")
def prepare_encoder_cudagraph_capture_inputs(
self,
@@ -1831,7 +1834,7 @@ class Qwen2_5_VLForConditionalGeneration(
self.model_config.max_model_len,
),
)
else:
elif modality == "video":
buffers = self.visual.prepare_encoder_metadata(
grid_thw_list,
max_frames_per_batch=max_frames_per_batch,
@@ -1840,6 +1843,8 @@ class Qwen2_5_VLForConditionalGeneration(
self.model_config.max_model_len,
),
)
else:
raise AssertionError("This line should be unreachable.")
return EncoderCudaGraphReplayBuffers(buffers=buffers)
vllm/model_executor/models/qwen2_vl.py
+14 -12
View File
@@ -1457,6 +1457,7 @@ class Qwen2VLForConditionalGeneration(
EncoderCudaGraphConfig,
)
max_frames = self.get_max_frames_per_video()
return EncoderCudaGraphConfig(
modalities=["image", "video"],
input_key_by_modality={
@@ -1470,6 +1471,7 @@ class Qwen2VLForConditionalGeneration(
"max_seqlen",
],
out_hidden_size=self.visual.out_hidden_size,
max_frames_per_video=max_frames,
)
def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:
@@ -1515,21 +1517,21 @@ class Qwen2VLForConditionalGeneration(
grid_thw = grid_thw.tolist()
return grid_thw
def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int:
return len(self._get_grid_thw_by_modality(mm_kwargs))
def get_encoder_cudagraph_item_specs(
self,
mm_kwargs: dict[str, Any],
):
from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
def get_encoder_cudagraph_per_item_output_tokens(
self, mm_kwargs: dict[str, Any]
) -> list[int]:
m = self.visual.spatial_merge_size
grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
return [t * (h // m) * (w // m) for t, h, w in grid_thw]
def get_encoder_cudagraph_per_item_input_sizes(
self, mm_kwargs: dict[str, Any]
) -> list[int]:
grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
return [t * h * w for t, h, w in grid_thw]
return [
EncoderItemSpec(
input_size=t * h * w,
output_tokens=t * (h // m) * (w // m),
)
for t, h, w in grid_thw
]
def select_encoder_cudagraph_items(
self, mm_kwargs: dict[str, Any], indices: list[int]
vllm/model_executor/models/qwen3_vl.py
+26 -22
View File
@@ -1776,6 +1776,9 @@ class Qwen3VLForConditionalGeneration(
# disable CUDA graph for all modalities when pruning is on.
modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"]
# Compute max_frames_per_video for budget sizing.
max_frames = self.get_max_frames_per_video() if "video" in modalities else 1
return EncoderCudaGraphConfig(
modalities=modalities,
input_key_by_modality={
@@ -1791,6 +1794,7 @@ class Qwen3VLForConditionalGeneration(
"sequence_lengths",
],
out_hidden_size=self.visual.out_hidden_size,
max_frames_per_video=max_frames,
)
def get_input_modality(
@@ -1799,7 +1803,9 @@ class Qwen3VLForConditionalGeneration(
) -> str:
if "image_grid_thw" in mm_kwargs:
return "image"
return "video"
elif "video_grid_thw" in mm_kwargs:
return "video"
raise AssertionError("This line should be unreachable.")
def get_max_frames_per_video(self) -> int:
mm_registry = MULTIMODAL_REGISTRY
@@ -1830,11 +1836,12 @@ class Qwen3VLForConditionalGeneration(
self,
mm_kwargs: dict[str, Any],
) -> torch.Tensor:
if self.get_input_modality(mm_kwargs) == "image":
pixel_values = mm_kwargs["pixel_values"]
else:
pixel_values = mm_kwargs["pixel_values_videos"]
return pixel_values
modality = self.get_input_modality(mm_kwargs)
if modality == "image":
return mm_kwargs["pixel_values"]
elif modality == "video":
return mm_kwargs["pixel_values_videos"]
raise AssertionError("This line should be unreachable.")
def _get_grid_thw_by_modality(
self,
@@ -1846,26 +1853,21 @@ class Qwen3VLForConditionalGeneration(
grid_thw = grid_thw.tolist()
return grid_thw
def get_encoder_cudagraph_num_items(
def get_encoder_cudagraph_item_specs(
self,
mm_kwargs: dict[str, Any],
) -> int:
return len(self._get_grid_thw_by_modality(mm_kwargs))
):
from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
def get_encoder_cudagraph_per_item_output_tokens(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
m = self.visual.spatial_merge_size
grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
return [t * (h // m) * (w // m) for t, h, w in grid_thw]
def get_encoder_cudagraph_per_item_input_sizes(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
return [t * h * w for t, h, w in grid_thw]
return [
EncoderItemSpec(
input_size=t * h * w,
output_tokens=t * (h // m) * (w // m),
)
for t, h, w in grid_thw
]
def select_encoder_cudagraph_items(
self,
@@ -2007,11 +2009,13 @@ class Qwen3VLForConditionalGeneration(
grid_thw_list,
max_batch_size=max_batch_size,
)
else:
elif modality == "video":
buffers = self.visual.prepare_encoder_metadata(
grid_thw_list,
max_frames_per_batch=max_frames_per_batch,
)
else:
raise AssertionError("This line should be unreachable.")
return EncoderCudaGraphReplayBuffers(buffers=buffers)
vllm/model_executor/models/step3_vl.py
+16 -31
View File
@@ -713,11 +713,6 @@ class Step3VLForConditionalGeneration(
) -> str:
return "image"
def get_max_frames_per_video(
self,
) -> int:
return 0
def get_encoder_cudagraph_budget_range(
self,
vllm_config: "VllmConfig",
@@ -734,52 +729,42 @@ class Step3VLForConditionalGeneration(
)
return min_budget, max_budget
def get_encoder_cudagraph_num_items(
def get_encoder_cudagraph_item_specs(
self,
mm_kwargs: dict[str, Any],
) -> int:
return len(mm_kwargs.get("pixel_values", []))
):
from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
def get_encoder_cudagraph_per_item_output_tokens(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
num_patches = mm_kwargs.get("num_patches")
img_output_tokens = self._compute_spatial_tokens(
self.config.vision_config.image_size,
self.config.vision_config.patch_size,
self.config.understand_projector_stride,
)
patch_output_tokens = self._compute_spatial_tokens(
504,
self.config.vision_config.patch_size,
self.config.understand_projector_stride,
)
return [
img_output_tokens + num_patch * patch_output_tokens
for num_patch in num_patches
]
def get_encoder_cudagraph_per_item_input_sizes(
self,
mm_kwargs: dict[str, Any],
) -> list[int]:
img_grid = (
self.config.vision_config.image_size // self.config.vision_config.patch_size
)
# NOTE: 504 is the hard coded size for each patch after processing
# by the vision model, which is determined by the current architecture
# of the vision model and may need to be updated if the architecture changes.
# The number of tokens for each patch is calculated based on this
# size and the patch size.
patch_output_tokens = self._compute_spatial_tokens(
504,
self.config.vision_config.patch_size,
self.config.understand_projector_stride,
)
img_grid = (
self.config.vision_config.image_size // self.config.vision_config.patch_size
)
patch_grid = 504 // self.config.vision_config.patch_size
total_image_pixel = img_grid * img_grid
total_patch_pixel = patch_grid * patch_grid
num_patches = mm_kwargs.get("num_patches")
return [
total_image_pixel + num_patch * total_patch_pixel
EncoderItemSpec(
input_size=(total_image_pixel + num_patch * total_patch_pixel),
output_tokens=(img_output_tokens + num_patch * patch_output_tokens),
)
for num_patch in num_patches
]
vllm/v1/worker/encoder_cudagraph.py
+14 -12
View File
@@ -21,6 +21,7 @@ from vllm.model_executor.models.utils import scatter_output_slices
from vllm.model_executor.models.vision import get_load_balance_assignment
from vllm.v1.worker.encoder_cudagraph_defs import (
EncoderCudaGraphConfig,
EncoderItemSpec,
)
logger = init_logger(__name__)
@@ -141,8 +142,8 @@ class EncoderCudaGraphManager:
elif user_max_frames is not None:
self.max_frames_per_batch = user_max_frames
else:
# Set it to the model-specific value according to its `processing_info`.
max_frames_per_video = self.model.get_max_frames_per_video()
# Set it to the model-specific value from config.
max_frames_per_video = self.config.max_frames_per_video
self.max_frames_per_batch = self.max_batch_size * max_frames_per_video
mm_config = vllm_config.model_config.multimodal_config
@@ -250,12 +251,13 @@ class EncoderCudaGraphManager:
return budget
return None
def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]:
"""Get item specs from the model."""
return self.model.get_encoder_cudagraph_item_specs(mm_kwargs)
def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]:
"""Get per-item output token counts as plain ints."""
return [
int(t)
for t in self.model.get_encoder_cudagraph_per_item_output_tokens(mm_kwargs)
]
return [spec.output_tokens for spec in self._get_item_specs(mm_kwargs)]
def _run_budget_graph(
self,
@@ -274,7 +276,7 @@ class EncoderCudaGraphManager:
Returns:
Encoder outputs, or None if graph not captured.
"""
num_items = self.model.get_encoder_cudagraph_num_items(mm_kwargs)
num_items = len(self._get_item_specs(mm_kwargs))
if token_budget not in self.budget_graphs:
self.graph_misses += num_items
return None
@@ -332,10 +334,11 @@ class EncoderCudaGraphManager:
always satisfy total_tokens <= max_budget and therefore
always find a valid budget (no miss).
"""
num_items = self.model.get_encoder_cudagraph_num_items(mm_kwargs)
item_specs = self._get_item_specs(mm_kwargs)
num_items = len(item_specs)
max_budget = self.token_budgets[-1]
per_item_out_tokens = self._get_per_item_out_tokens(mm_kwargs)
per_item_out_tokens = [spec.output_tokens for spec in item_specs]
# Sort ascending by output token count (smallest first)
sorted_indices = sorted(range(num_items), key=lambda i: per_item_out_tokens[i])
@@ -459,9 +462,8 @@ class EncoderCudaGraphManager:
tp_size = get_tensor_model_parallel_world_size()
current_rank = get_tensor_model_parallel_rank()
per_item_input_sizes = self.model.get_encoder_cudagraph_per_item_input_sizes(
mm_kwargs
)
item_specs = self._get_item_specs(mm_kwargs)
per_item_input_sizes = [spec.input_size for spec in item_specs]
(image_rank_assignment, images_per_rank, input_patches_per_rank) = (
get_load_balance_assignment(per_item_input_sizes, tp_size)
vllm/v1/worker/encoder_cudagraph_defs.py
+22
View File
@@ -8,6 +8,23 @@ from typing import Any
import torch
@dataclass
class EncoderItemSpec:
"""Description of a single encoder input item.
Returned by ``get_encoder_cudagraph_item_specs()`` to describe each
image or video in a batch without the manager needing to understand
model-specific input formats.
"""
input_size: int
"""Number of input patches/rows for this item."""
output_tokens: int
"""Number of output tokens after encoder processing (e.g. after
spatial merge)."""
@dataclass
class EncoderCudaGraphConfig:
"""Configuration for encoder CUDA graph management.
@@ -34,6 +51,11 @@ class EncoderCudaGraphConfig:
"""Output hidden dim of the vision encoder.
Used for DP gather buffer allocation."""
max_frames_per_video: int = 1
"""Maximum number of frames per video.
Only relevant when "video" is in ``modalities``.
Image-only models can use the default of 1."""
@dataclass
class EncoderCudaGraphCaptureInputs: