diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 7502186ae46..04fee73af60 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -101,7 +101,7 @@ When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks Following (ViT full CUDA graph support for image inference), extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager. !!! note - Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture. + Video CUDA graphs are automatically disabled when video token pruning (EVS or VidCom2) is enabled, since pruning makes the token count data-dependent and incompatible with CUDA graph capture. Mixed inputs (image+video) per prompt are also supported now. diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index 7f6c1aee760..4cb624030e9 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -350,6 +350,31 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py) +#### Video Token Pruning + +For supported models, vLLM can prune video tokens after the vision encoder to +reduce prefill time and KV cache usage, at some cost in accuracy. Set +`--video-pruning-rate ` to prune the fraction `q` of video tokens from each +video, and `--video-pruning-method` to choose the training-free algorithm: + +- **`evs`** (Efficient Video Sampling, default): drops the tokens with the + lowest temporal dissimilarity to the previous frame. The first frame is + always fully retained. +- **`vidcom2`** (Video Compression Commander): scores tokens by similarity to + video-level and frame-level feature centers and gives distinctive frames a + larger share of the budget. At least one token per frame is retained. + +```bash +vllm serve Qwen/Qwen3-VL-8B-Instruct \ + --video-pruning-rate 0.75 --video-pruning-method vidcom2 +``` + +!!! note + `evs` is supported by all models implementing multimodal pruning; + `vidcom2` is currently supported by Qwen3-VL only. Unsupported combinations + are rejected at startup. Enabling video pruning also disables encoder CUDA + graphs, since the retained token count becomes data-dependent. + ### Audio Inputs You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary. diff --git a/tests/multimodal/test_vidcom2.py b/tests/multimodal/test_vidcom2.py new file mode 100644 index 00000000000..4a62bbd87cf --- /dev/null +++ b/tests/multimodal/test_vidcom2.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retained_tokens_count, + compute_retention_mask, +) + + +def _fake_video_embeds( + num_frames: int, + rows: int, + cols: int, + hidden: int = 64, + seed: int = 0, +) -> torch.Tensor: + """Deterministic fake ViT output with a distinct mean per frame.""" + g = torch.Generator().manual_seed(seed) + frames = [] + for f in range(num_frames): + base = torch.randn(hidden, generator=g) * (0.1 + 0.05 * f) + frames.append( + base[None, :].expand(rows * cols, hidden) + + 0.01 * torch.randn(rows * cols, hidden, generator=g) + ) + return torch.cat(frames, dim=0) + + +@pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9]) +@pytest.mark.parametrize("num_frames", [1, 4, 16]) +def test_mask_shape_and_dtype(q: float, num_frames: int) -> None: + merge = 2 + rows, cols = 6, 8 + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + assert mask.dtype == torch.bool + assert mask.shape == (num_frames * rows * cols,) + + +def test_retained_count_floors_at_one_token_per_frame() -> None: + """The global minimum is one token per frame (not a full first frame).""" + assert ( + compute_retained_tokens_count(tokens_per_frame=48, num_frames=4, q=0.999) == 4 + ) + assert ( + compute_retained_tokens_count(tokens_per_frame=48, num_frames=4, q=0.0) + == 48 * 4 + ) + + +@pytest.mark.parametrize("q", [0.25, 0.5, 0.75, 0.9]) +@pytest.mark.parametrize("num_frames", [1, 4, 16]) +def test_total_retained_matches_target(q: float, num_frames: int) -> None: + """Mask total must equal the placeholder-sizing helper.""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + expected = compute_retained_tokens_count( + tokens_per_frame=tpf, num_frames=num_frames, q=q + ) + assert int(mask.sum().item()) == expected + + +def test_per_frame_min_one_when_budget_allows() -> None: + """No frame is fully dropped when the budget allows.""" + merge = 2 + rows, cols = 6, 8 + num_frames = 8 + embeds = _fake_video_embeds(num_frames, rows, cols) + mask = compute_retention_mask( + embeds, + (num_frames, rows * merge, cols * merge), + spatial_merge_size=merge, + q=0.25, + ) + per_frame = mask.view(num_frames, rows * cols).sum(dim=1) + assert (per_frame >= 1).all(), f"zero-token frame detected: {per_frame.tolist()}" + + +def test_dynamic_per_frame_budget() -> None: + """A distinctive frame gets more retained tokens than bland ones.""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + hidden = 64 + torch.manual_seed(0) + bland = 0.01 * torch.randn(tpf, hidden) + frames = [torch.randn(tpf, hidden) * 1.0] + for _ in range(7): + frames.append(bland + 0.001 * torch.randn(tpf, hidden)) + embeds = torch.cat(frames, dim=0) + mask = compute_retention_mask( + embeds, + (8, rows * merge, cols * merge), + spatial_merge_size=merge, + q=0.5, + ) + per_frame = mask.view(8, tpf).sum(dim=1) + assert per_frame[0].item() > per_frame[1:].float().mean().item() + + +def test_empty_input_safe() -> None: + embeds = torch.zeros(0, 32) + mask = compute_retention_mask(embeds, (0, 0, 0), spatial_merge_size=2, q=0.25) + assert mask.numel() == 0 + + +@pytest.mark.parametrize("q", [0.0, 0.25, 0.5, 0.75]) +def test_first_frame_not_privileged(q: float) -> None: + """A bland first frame is not force-retained (unlike EVS).""" + merge = 2 + rows, cols = 6, 8 + tpf = rows * cols + torch.manual_seed(1) + bland = 0.01 * torch.randn(tpf, 64) + frames = [bland] + for f in range(7): + frames.append(torch.randn(tpf, 64) * (1.0 + 0.1 * f)) + embeds = torch.cat(frames, dim=0) + mask = compute_retention_mask( + embeds, + (8, rows * merge, cols * merge), + spatial_merge_size=merge, + q=q, + ) + per_frame = mask.view(8, tpf).sum(dim=1) + assert per_frame[0].item() <= tpf + if q > 0.0: + assert per_frame[0].item() < int(mask.sum().item()) diff --git a/vllm/config/model.py b/vllm/config/model.py index d64bd57e87d..c2aa650f88d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -376,6 +376,7 @@ class ModelConfig: interleave_mm_strings: InitVar[bool | None] = None skip_mm_profiling: InitVar[bool | None] = None video_pruning_rate: InitVar[float | None] = None + video_pruning_method: InitVar[str | None] = None mm_tensor_ipc: InitVar[MMTensorIPC] = None mm_ipc_gpu_memory_gb: InitVar[float | None] = None @@ -504,6 +505,7 @@ class ModelConfig: interleave_mm_strings: bool | None, skip_mm_profiling: bool | None, video_pruning_rate: float | None, + video_pruning_method: str | None, mm_tensor_ipc: MMTensorIPC, mm_ipc_gpu_memory_gb: float | None, ) -> None: @@ -735,6 +737,7 @@ class ModelConfig: interleave_mm_strings=interleave_mm_strings, skip_mm_profiling=skip_mm_profiling, video_pruning_rate=video_pruning_rate, + video_pruning_method=video_pruning_method, mm_tensor_ipc=mm_tensor_ipc, mm_ipc_gpu_memory_gb=mm_ipc_gpu_memory_gb, ) @@ -745,6 +748,19 @@ class ModelConfig: self.multimodal_config = MultiModalConfig(**mm_config_kwargs) # type: ignore[arg-type] + pruning_spec = self.multimodal_config.get_video_pruning_spec() + supported_pruning = self._model_info.supported_video_pruning_methods + if ( + pruning_spec is not None + and supported_pruning + and pruning_spec[0] not in supported_pruning + ): + raise ValueError( + f"Video pruning method '{pruning_spec[0]}' is not " + f"supported by {self._model_info.architecture} " + f"(supported methods: {supported_pruning})." + ) + if ( self.renderer_num_workers > 1 and self.multimodal_config.mm_processor_cache_gb > 0 diff --git a/vllm/config/multimodal.py b/vllm/config/multimodal.py index 865615dff79..ab5cf50fc3a 100644 --- a/vllm/config/multimodal.py +++ b/vllm/config/multimodal.py @@ -61,6 +61,7 @@ class MultiModalDummyOptionsBuiltins(TypedDict, total=False): MMEncoderTPMode = Literal["weights", "data"] MMCacheType = Literal["shm", "lru"] +VideoPruningMethod = Literal["evs", "vidcom2"] MMTensorIPC = Literal["direct_rpc", "torch_shm"] MMDummyOptions: TypeAlias = dict[str, BaseDummyOptions] """ @@ -189,9 +190,14 @@ class MultiModalConfig: estimating the peak memory usage of the activation of multimodal encoder and embedding cache.""" video_pruning_rate: float | None = Field(default=None, ge=0.0, lt=1.0) - """Sets pruning rate for video pruning via Efficient Video Sampling. - Value sits in range [0;1) and determines fraction of media tokens - from each video to be pruned. + """Fraction of video tokens to prune from each video. Value sits in range + [0;1); pruning is enabled when it is greater than 0. The pruning algorithm + is selected by `video_pruning_method`. + """ + video_pruning_method: VideoPruningMethod = "evs" + """Video token pruning algorithm applied when `video_pruning_rate` > 0: + - "evs": Efficient Video Sampling. + - "vidcom2": Video Compression Commander. """ mm_tensor_ipc: MMTensorIPC = "direct_rpc" """IPC (inter-process communication) method for multimodal tensors. @@ -360,4 +366,11 @@ class MultiModalConfig: ) def is_multimodal_pruning_enabled(self): - return self.video_pruning_rate is not None and self.video_pruning_rate > 0 + return self.get_video_pruning_spec() is not None + + def get_video_pruning_spec(self) -> tuple[VideoPruningMethod, float] | None: + """Return `(method, rate)` when video pruning is enabled, else None. + `rate` is the fraction of video tokens to prune.""" + if self.video_pruning_rate is not None and self.video_pruning_rate > 0: + return (self.video_pruning_method, float(self.video_pruning_rate)) + return None diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index b7c5f746545..fc9349a0bea 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -586,6 +586,7 @@ class EngineArgs: renderer_num_workers: int = 1 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling video_pruning_rate: float | None = MultiModalConfig.video_pruning_rate + video_pruning_method: str = MultiModalConfig.video_pruning_method mm_tensor_ipc: MMTensorIPC = MultiModalConfig.mm_tensor_ipc mm_ipc_gpu_memory_gb: float = MultiModalConfig.mm_ipc_gpu_memory_gb # LoRA fields @@ -1333,6 +1334,10 @@ class EngineArgs: multimodal_group.add_argument( "--video-pruning-rate", **multimodal_kwargs["video_pruning_rate"] ) + multimodal_group.add_argument( + "--video-pruning-method", + **multimodal_kwargs["video_pruning_method"], + ) multimodal_group.add_argument( "--mm-tensor-ipc", **multimodal_kwargs["mm_tensor_ipc"] ) @@ -1715,6 +1720,7 @@ class EngineArgs: override_attention_dtype=self.override_attention_dtype, logits_processors=self.logits_processors, video_pruning_rate=self.video_pruning_rate, + video_pruning_method=self.video_pruning_method, mm_tensor_ipc=self.mm_tensor_ipc, mm_ipc_gpu_memory_gb=self.mm_ipc_gpu_memory_gb, io_processor_plugin=self.io_processor_plugin, diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index df9c18ce534..71229439f66 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -41,6 +41,7 @@ if TYPE_CHECKING: SpeechToTextParams, VllmConfig, ) + from vllm.config.multimodal import VideoPruningMethod from vllm.inputs import PromptType, TokensPrompt from vllm.lora.model_manager import LoRAModelManager from vllm.model_executor.layers.fused_moe import MoERunner @@ -424,6 +425,13 @@ class SupportsMultiModalPruning(Protocol): supports_multimodal_pruning: ClassVar[Literal[True]] = True + supported_video_pruning_methods: ClassVar[tuple["VideoPruningMethod", ...]] = ( + "evs", + ) + """Video pruning methods (as reported by + `MultiModalConfig.get_video_pruning_spec`) implemented by this model. + Models supporting methods beyond EVS should override this.""" + def recompute_mrope_positions( self, input_ids: list[int] | torch.Tensor, diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index c04b4729454..ef7d3d58343 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -570,10 +570,7 @@ class InternS1ProForConditionalGeneration( self.config = config self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) with self._mark_tower_model(vllm_config, {"image", "video"}): self.visual = Qwen3_VisionTransformer( diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 64667503d57..5b4233b0781 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -42,10 +42,6 @@ from vllm.model_executor.models.utils import ( maybe_prefix, ) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_retained_tokens_count, - compute_retention_mask, -) from vllm.multimodal.inputs import ( AudioItem, BatchedTensorInputs, @@ -74,6 +70,10 @@ from vllm.multimodal.processing.processor import ( PromptReplacement, PromptUpdate, ) +from vllm.multimodal.video_prune.evs import ( + compute_retained_tokens_count, + compute_retention_mask, +) from vllm.renderers import TokenizeParams from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index c987e07b43d..7957ca805ff 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -67,12 +67,6 @@ from vllm.model_executor.layers.rotary_embedding.common import ( ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_mrope_for_media, - compute_retained_tokens_count, - compute_retention_mask, - recompute_mrope_positions, -) from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -80,6 +74,12 @@ from vllm.multimodal.inputs import ( ) from vllm.multimodal.parse import MultiModalDataItems from vllm.multimodal.processing import PromptReplacement, PromptUpdate +from vllm.multimodal.video_prune.evs import ( + compute_mrope_for_media, + compute_retained_tokens_count, + compute_retention_mask, + recompute_mrope_positions, +) from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index f86560e5f4e..baf75fa2dd5 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -50,7 +50,12 @@ from transformers.video_utils import VideoMetadata from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions +from vllm.config.multimodal import ( + BaseDummyOptions, + MultiModalConfig, + VideoDummyOptions, + VideoPruningMethod, +) from vllm.distributed import get_pp_group, parallel_state from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger @@ -69,12 +74,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.evs import ( - compute_mrope_for_media, - compute_retained_tokens_count, - compute_retention_mask, - recompute_mrope_positions, -) from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -92,6 +91,18 @@ from vllm.multimodal.processing import ( PromptUpdate, PromptUpdateDetails, ) +from vllm.multimodal.video_prune.evs import ( + compute_mrope_for_media, + compute_retained_tokens_count, + compute_retention_mask, + recompute_mrope_positions, +) +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retained_tokens_count as vidcom2_compute_retained_tokens_count, +) +from vllm.multimodal.video_prune.vidcom2 import ( + compute_retention_mask as vidcom2_compute_retention_mask, +) from vllm.sequence import IntermediateTensors from vllm.tokenizers.protocol import TokenizerLike from vllm.tokenizers.registry import cached_tokenizer_from_config @@ -1256,7 +1267,7 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) hf_config = self.info.get_hf_config() tokenizer = self.info.get_tokenizer() merge_size = hf_config.vision_config.spatial_merge_size - video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate + pruning_spec = self.info.ctx.get_mm_config().get_video_pruning_spec() vision_start_token_id = hf_config.vision_start_token_id vision_end_token_id = hf_config.vision_end_token_id video_token_id = hf_config.video_token_id @@ -1339,11 +1350,18 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) merge_size**2 ) - if video_pruning_rate is not None and video_pruning_rate > 0.0: - num_tokens = compute_retained_tokens_count( + # Apply video pruning (EVS or VidCom2) if enabled. + if pruning_spec is not None: + method, prune_q = pruning_spec + count_fn = ( + vidcom2_compute_retained_tokens_count + if method == "vidcom2" + else compute_retained_tokens_count + ) + num_tokens = count_fn( tokens_per_frame=tokens_per_frame_base, num_frames=num_frames, - q=video_pruning_rate, + q=prune_q, ) tokens_per_frame = [num_tokens] + [0] * (num_frames - 1) select_token_id = False @@ -1459,16 +1477,22 @@ class Qwen3VLMultiModalProcessor(BaseMultiModalProcessor[Qwen3VLProcessingInfo]) f"video length ({grid_thw[0]})." ) - # Compute tokens per frame, with EVS support + # Compute tokens per frame, with EVS / VidCom2 support num_frames = int(grid_thw[0]) tokens_per_frame_base = int(grid_thw[1:].prod()) // merge_length - video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate - if video_pruning_rate is not None and video_pruning_rate > 0.0: - num_tokens = compute_retained_tokens_count( + pruning_spec = self.info.ctx.get_mm_config().get_video_pruning_spec() + if pruning_spec is not None: + method, prune_q = pruning_spec + count_fn = ( + vidcom2_compute_retained_tokens_count + if method == "vidcom2" + else compute_retained_tokens_count + ) + num_tokens = count_fn( tokens_per_frame=tokens_per_frame_base, num_frames=num_frames, - q=video_pruning_rate, + q=prune_q, ) tokens_per_frame = [num_tokens] + [0] * (num_frames - 1) select_token_id = False @@ -1701,6 +1725,8 @@ class Qwen3VLForConditionalGeneration( supports_encoder_tp_data = True + supported_video_pruning_methods = ("evs", "vidcom2") + # To ensure correct weight loading and mapping. hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ @@ -1719,6 +1745,17 @@ class Qwen3VLForConditionalGeneration( raise ValueError("Only image or video modality is supported") + def _init_video_pruning(self, multimodal_config: MultiModalConfig) -> None: + pruning_spec = multimodal_config.get_video_pruning_spec() + if pruning_spec is None: + self.video_pruning_method: VideoPruningMethod | None = None + self.video_pruning_rate = multimodal_config.video_pruning_rate + else: + self.video_pruning_method, self.video_pruning_rate = pruning_spec + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): super().__init__() config: Qwen3VLConfig = vllm_config.model_config.hf_config @@ -1730,10 +1767,7 @@ class Qwen3VLForConditionalGeneration( self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") self.deepstack_num_level = ( @@ -1848,7 +1882,7 @@ class Qwen3VLForConditionalGeneration( EncoderCudaGraphConfig, ) - # When EVS pruning is enabled, embed_multimodal post-processes both + # When video pruning is enabled, embed_multimodal post-processes both # image and video embeddings (mrope positions are appended for image, # prune+append for video). The encoder CUDA graph path bypasses that # post-process, producing inconsistent embedding formats vs eager. So @@ -2291,9 +2325,12 @@ class Qwen3VLForConditionalGeneration( t, h, w = size if self.is_multimodal_pruning_enabled: - # For each video, compute retention mask using EVS. - # retention_mask: [11424]. - retention_mask = compute_retention_mask( + # Compute the retention mask for each video (EVS or VidCom2). + if self.video_pruning_method == "vidcom2": + mask_fn = vidcom2_compute_retention_mask + else: + mask_fn = compute_retention_mask + retention_mask = mask_fn( emb, size, spatial_merge_size=self.visual.spatial_merge_size, diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 4413e1213bb..f1409d23399 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -220,10 +220,7 @@ class Qwen3VLMoeForConditionalGeneration( self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) self.multimodal_config = multimodal_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" - self.video_pruning_rate = multimodal_config.video_pruning_rate - self.is_multimodal_pruning_enabled = ( - multimodal_config.is_multimodal_pruning_enabled() - ) + self._init_video_pruning(multimodal_config) self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes") self.deepstack_num_level = ( diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 1dcacc7936e..697f55b3727 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -797,6 +797,7 @@ class _ModelInfo: supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool + supported_video_pruning_methods: tuple[str, ...] @staticmethod def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": @@ -827,6 +828,9 @@ class _ModelInfo: supports_transcription(model) and model.supports_transcription_only ), has_noops=has_noops(model), + supported_video_pruning_methods=getattr( + model, "supported_video_pruning_methods", () + ), ) diff --git a/vllm/multimodal/video_prune/__init__.py b/vllm/multimodal/video_prune/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/multimodal/video_prune/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/multimodal/evs.py b/vllm/multimodal/video_prune/evs.py similarity index 100% rename from vllm/multimodal/evs.py rename to vllm/multimodal/video_prune/evs.py diff --git a/vllm/multimodal/video_prune/vidcom2.py b/vllm/multimodal/video_prune/vidcom2.py new file mode 100644 index 00000000000..830e47bc055 --- /dev/null +++ b/vllm/multimodal/video_prune/vidcom2.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# VidCom2 (Video Compression Commander) video token pruning. +# Liu et al., EMNLP 2025 — https://arxiv.org/abs/2505.14454 +# Adapted from the reference implementation: +# https://github.com/xuyang-liu16/VidCom2 (Apache-2.0, +# Copyright (c) 2025 the VidCom2 authors). + +import torch +import torch.nn.functional as F + +# Multi-scale Gaussian bandwidths from the reference implementation. +_ALPHAS: tuple[float, ...] = tuple(2.0**k for k in range(-3, 2)) +_LOW_VAR_CHANNEL_RATIO: float = 0.5 +_SOFTMAX_TEMPERATURE: float = 0.01 + + +def compute_retained_tokens_count( + tokens_per_frame: int, num_frames: int, q: float +) -> int: + """Number of video tokens retained after VidCom2 pruning. + + The target is `(1 - q) * total_tokens`, i.e. a retention ratio of + `1 - q` averaged across frames. Because the per-frame budget is floored + at one token, the global minimum is `num_frames` (one token per frame). + """ + total_tokens = tokens_per_frame * num_frames + base_num = int(total_tokens * (1.0 - q)) + return max(num_frames, min(base_num, total_tokens)) + + +def compute_retention_mask( + video_embeds: torch.Tensor, + video_size_thw: torch.LongTensor | tuple[int, int, int], + spatial_merge_size: int, + q: float, +) -> torch.Tensor: + """Compute the VidCom2 retention mask for a single video. + + Args: + video_embeds: `(T*H*W/merge^2, hidden_size)` post-ViT token features. + video_size_thw: `(T, H, W)` grid dimensions. + spatial_merge_size: ViT spatial merge factor (e.g. 2). + q: Pruning fraction in `[0, 1)`; retention ratio is `1 - q`. + + Returns: + Flat bool tensor of shape `(T*H*W/merge^2,)`, True for retained + tokens. The True count equals `compute_retained_tokens_count` so + placeholders sized at prompt-processing time match exactly. + """ + T, H, W = map(int, video_size_thw) + rows = H // spatial_merge_size + cols = W // spatial_merge_size + tokens_per_frame = rows * cols + total_tokens = T * tokens_per_frame + + device = video_embeds.device + if tokens_per_frame == 0 or total_tokens == 0: + return torch.ones(0, dtype=torch.bool, device=device) + + target_retained = compute_retained_tokens_count( + tokens_per_frame=tokens_per_frame, num_frames=T, q=q + ) + target_retained = min(target_retained, total_tokens) + + # 1. Score in the lowest-variance half of channels. + variances = video_embeds.var(dim=0, unbiased=False) + k_channels = max(1, int(video_embeds.size(-1) * _LOW_VAR_CHANNEL_RATIO)) + _, low_var_idx = torch.topk(variances, k=k_channels, largest=False) + sel = video_embeds.index_select(-1, low_var_idx) + + # 2. Multi-scale Gaussian similarity to video and per-frame centers. + frames = sel.view(T, tokens_per_frame, sel.size(-1)) + frames = F.normalize(frames, dim=-1) + vid_center = frames.mean(dim=(0, 1), keepdim=True) # (1, 1, C) + frame_center = frames.mean(dim=1, keepdim=True) # (T, 1, C) + v_score = _multi_scale_gaussian(frames, vid_center) + f_score = _multi_scale_gaussian(frames, frame_center) + # Higher similarity = more redundant; lowest-similarity tokens are kept. + similarity = v_score + f_score # (T, tpf) + + # 3. Per-frame dynamic budget: distinctive frames get a larger share. + base = 1.0 - q + frame_scores = -v_score.mean(dim=-1) # (T,) + probs = F.softmax((frame_scores - frame_scores.max()) / _SOFTMAX_TEMPERATURE, dim=0) + scales = (base * (1.0 + probs - probs.mean())).clamp(max=1.0) + ks = (scales * tokens_per_frame).round().long().clamp(min=1, max=tokens_per_frame) + + # 4. Retain the smallest-similarity tokens per frame. + mask_2d = torch.zeros(T, tokens_per_frame, dtype=torch.bool, device=device) + for i in range(T): + k_i = int(ks[i].item()) + if k_i <= 0: + continue + _, idx = torch.topk(similarity[i], k=k_i, largest=False, sorted=False) + mask_2d[i].scatter_(0, idx, True) + + # 5. Reconcile rounding/clamp drift to the exact target count by score. + flat_mask = mask_2d.view(-1) + flat_sim = similarity.view(-1) + current = int(flat_mask.sum().item()) + if current > target_retained: + drop_n = current - target_retained + retained_idx = flat_mask.nonzero(as_tuple=False).squeeze(-1) + retained_sim = flat_sim[retained_idx] + _, worst = torch.topk(retained_sim, k=drop_n, largest=True, sorted=False) + flat_mask[retained_idx[worst]] = False + elif current < target_retained: + add_n = target_retained - current + available_idx = (~flat_mask).nonzero(as_tuple=False).squeeze(-1) + if available_idx.numel() > 0: + available_sim = flat_sim[available_idx] + add_n = min(add_n, available_idx.numel()) + _, best = torch.topk(available_sim, k=add_n, largest=False, sorted=False) + flat_mask[available_idx[best]] = True + + return flat_mask + + +def _multi_scale_gaussian(x: torch.Tensor, center: torch.Tensor) -> torch.Tensor: + """Sum Gaussian kernels over `_ALPHAS`; `(T, N, C) -> (T, N)` scores.""" + dist_sq = ((x - center) ** 2).sum(dim=-1) + return sum(torch.exp(-dist_sq / (2.0 * a)) for a in _ALPHAS) diff --git a/vllm/transformers_utils/processors/nano_nemotron_vl.py b/vllm/transformers_utils/processors/nano_nemotron_vl.py index d48a29d6b43..028d207a25a 100644 --- a/vllm/transformers_utils/processors/nano_nemotron_vl.py +++ b/vllm/transformers_utils/processors/nano_nemotron_vl.py @@ -23,9 +23,9 @@ from PIL import Image from transformers import BatchFeature, PretrainedConfig, TensorType from vllm.model_executor.models.parakeet import ParakeetExtractor -from vllm.multimodal.evs import compute_retained_tokens_count from vllm.multimodal.inputs import AudioItem from vllm.multimodal.processing.processor import PromptUpdateDetails +from vllm.multimodal.video_prune.evs import compute_retained_tokens_count from vllm.tokenizers.hf import HfTokenizer from .internvl import calculate_internvl_targets, get_internvl_target_ratios