From f0feb15e7fc521544d23c2d23de0e327a509876b Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Fri, 22 May 2026 22:31:00 +0800 Subject: [PATCH] [Multimodal] Simplify ViT CUDA graph interfaces (#41234) Signed-off-by: Isotr0py --- .../generation/test_vit_cudagraph.py | 3 - tests/v1/cudagraph/test_encoder_cudagraph.py | 57 +++++++++---------- vllm/model_executor/models/interfaces.py | 28 +++------ vllm/model_executor/models/qwen2_5_vl.py | 43 +++++++------- vllm/model_executor/models/qwen2_vl.py | 26 +++++---- vllm/model_executor/models/qwen3_vl.py | 48 +++++++++------- vllm/model_executor/models/step3_vl.py | 47 ++++++--------- vllm/v1/worker/encoder_cudagraph.py | 26 +++++---- vllm/v1/worker/encoder_cudagraph_defs.py | 22 +++++++ 9 files changed, 151 insertions(+), 149 deletions(-) diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index b23b8d5057..18630e3559 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -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] diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index c67b5ec39b..2ba140707e 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -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)) diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 000de94ee6..3d46bda7ff 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -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. """ ... diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index 9e1f394f53..a4c75ab2e0 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -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) diff --git a/vllm/model_executor/models/qwen2_vl.py b/vllm/model_executor/models/qwen2_vl.py index d0c059cdaf..869e044c2d 100644 --- a/vllm/model_executor/models/qwen2_vl.py +++ b/vllm/model_executor/models/qwen2_vl.py @@ -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] diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 6f10259d13..d7765351e6 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -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) diff --git a/vllm/model_executor/models/step3_vl.py b/vllm/model_executor/models/step3_vl.py index 36e0e73e36..efe552f554 100644 --- a/vllm/model_executor/models/step3_vl.py +++ b/vllm/model_executor/models/step3_vl.py @@ -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 ] diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index 4794fa7742..6d5aa6f412 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -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) diff --git a/vllm/v1/worker/encoder_cudagraph_defs.py b/vllm/v1/worker/encoder_cudagraph_defs.py index 00ab97b3cd..70e15703f9 100644 --- a/vllm/v1/worker/encoder_cudagraph_defs.py +++ b/vllm/v1/worker/encoder_cudagraph_defs.py @@ -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: