@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 392edee34a

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

EVS Support (Video tokens pruning) (#22980)

Signed-off-by: Eugene Khvedchenia <ekhvedchenia@nvidia.com> Signed-off-by: Eugene Khvedchenya <ekhvedchenya@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io>

Eugene Khvedchenya committed Sep 26, 2025 at 06:54 UTC 392edee34a008af2453d936cb3cdbd97842984a7

8 files changed +783 -39

tests/models/multimodal/generation/test_qwen2_5_vl.py new

+132

@@ -0,0 +1,132 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import pytest

5 +

6 +from vllm.multimodal.video import sample_frames_from_video

7 +

8 +from ....conftest import VIDEO_ASSETS

9 +

10 +models = ["Qwen/Qwen2.5-VL-3B-Instruct"]

11 +target_dtype = "bfloat16"

12 +

13 +VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"

14 +

15 +

16 +def qwen2_5_vl_chat_template(*query):

17 + return f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{''.join(query)}<|im_end|><|im_start|>assistant\n" # noqa: E501

18 +

19 +

20 +VIDEO_PROMPTS = VIDEO_ASSETS.prompts({

21 + "baby_reading":

22 + qwen2_5_vl_chat_template(

23 + VIDEO_PLACEHOLDER,

24 + "Describe this video with a short sentence ",

25 + "(no more than 20 words)",

26 + ),

27 +})

28 +

29 +

30 +@pytest.mark.core_model

31 +@pytest.mark.parametrize("model", models)

32 +@pytest.mark.parametrize("video_pruning_rate", [0.0, 0.75])

33 +@pytest.mark.parametrize("num_frames", [16])

34 +@pytest.mark.parametrize("dtype", [target_dtype])

35 +@pytest.mark.parametrize("max_tokens", [128])

36 +def test_qwen2_5_vl_evs_functionality(vllm_runner, video_assets, model,

37 + video_pruning_rate: float,

38 + num_frames: int, dtype: str,

39 + max_tokens: int) -> None:

40 + """Test EVS (Efficient Video Sampling) functionality with different

41 + pruning rates.

42 + """

43 +

44 + # Sample frames from video assets

45 + sampled_vids = [

46 + sample_frames_from_video(asset.np_ndarrays, num_frames)

47 + for asset in video_assets

48 + ]

49 +

50 + prompts = [VIDEO_PROMPTS[0]]

51 + videos = [sampled_vids[0]]

52 +

53 + # Initialize model with EVS configuration

54 + with vllm_runner(model,

55 + runner="generate",

56 + max_model_len=4000,

57 + max_num_seqs=1,

58 + dtype=dtype,

59 + limit_mm_per_prompt={"video": 1},

60 + tensor_parallel_size=1,

61 + video_pruning_rate=video_pruning_rate) as vllm_model:

62 +

63 + # Generate output - this should not crash

64 + outputs = vllm_model.generate_greedy(prompts,

65 + max_tokens,

66 + videos=videos)

67 +

68 + # Basic validation that we got a response

69 + assert len(outputs) == 1

70 + output_ids, output_text = outputs[0]

71 +

72 + # Ensure we got some output

73 + assert len(output_ids) > 0

74 + assert len(output_text) > 0

75 +

76 + # Ensure the output is a string

77 + assert isinstance(output_text, str)

78 +

79 +

80 +@pytest.mark.core_model

81 +@pytest.mark.parametrize("model", models)

82 +@pytest.mark.parametrize("video_pruning_rate", [0.0, 0.75])

83 +@pytest.mark.parametrize("num_frames", [16])

84 +@pytest.mark.parametrize("dtype", [target_dtype])

85 +@pytest.mark.parametrize("max_tokens", [128])

86 +def test_qwen2_5_vl_evs_batched_videos(vllm_runner, video_assets, model,

87 + video_pruning_rate: float,

88 + num_frames: int, dtype: str,

89 + max_tokens: int) -> None:

90 + """Test EVS functionality with batched videos.

91 +

92 + This test validates that:

93 + 1. The model handles batched video inputs correctly with EVS

94 + 2. Both pruning configurations work with multiple videos

95 + 3. The model doesn't crash when processing multiple videos simultaneously

96 + """

97 + # Sample frames from video assets

98 + sampled_vids = [

99 + sample_frames_from_video(asset.np_ndarrays, num_frames)

100 + for asset in video_assets

101 + ]

102 +

103 + # Test batched videos

104 + prompts = [VIDEO_PROMPTS[0], VIDEO_PROMPTS[0]]

105 + videos = [sampled_vids[0],

106 + sampled_vids[0]] # Use same video twice for testing

107 +

108 + # Initialize model with EVS configuration

109 + with vllm_runner(model,

110 + runner="generate",

111 + max_model_len=4000,

112 + max_num_seqs=2,

113 + dtype=dtype,

114 + limit_mm_per_prompt={"video": 2},

115 + tensor_parallel_size=1,

116 + video_pruning_rate=video_pruning_rate) as vllm_model:

117 +

118 + # Generate output - this should not crash

119 + outputs = vllm_model.generate_greedy(prompts,

120 + max_tokens,

121 + videos=videos)

122 +

123 + # Basic validation that we got responses for both videos

124 + assert len(outputs) == 2

125 +

126 + for output_ids, output_text in outputs:

127 + # Ensure we got some output for each video

128 + assert len(output_ids) > 0

129 + assert len(output_text) > 0

130 +

131 + # Ensure the output is a string

132 + assert isinstance(output_text, str)

vllm/config/model.py

+16 -11

@@ -283,6 +283,7 @@ class ModelConfig:

283 mm_encoder_tp_mode: InitVar[Optional[MMEncoderTPMode]] = None

284 interleave_mm_strings: InitVar[Optional[bool]] = None

285 skip_mm_profiling: InitVar[Optional[bool]] = None

286 + video_pruning_rate: InitVar[Optional[float]] = None

287

288 def compute_hash(self) -> str:

289 """

@@ -311,6 +312,7 @@ class ModelConfig:

312 factors.append(self.override_generation_config)

313 factors.append(self.rope_scaling)

314 factors.append(self.rope_theta)

315 + factors.append(self.video_pruning_rate)

316

317 # hf_config can control how the model looks!

318 try:

@@ -338,17 +340,19 @@ class ModelConfig:

340 return hashlib.sha256(str(factors).encode()).hexdigest()

341

342 def __post_init__(

341 - self,

342 - # Multimodal config init vars

343 - limit_mm_per_prompt: Optional[dict[str, int]],

344 - media_io_kwargs: Optional[dict[str, dict[str, Any]]],

345 - mm_processor_kwargs: Optional[dict[str, Any]],

346 - mm_processor_cache_gb: Optional[float],

347 - mm_processor_cache_type: Optional[MMCacheType],

348 - mm_shm_cache_max_object_size_mb: Optional[int],

349 - mm_encoder_tp_mode: Optional[MMEncoderTPMode],

350 - interleave_mm_strings: Optional[bool],

351 - skip_mm_profiling: Optional[bool]) -> None:

343 + self,

344 + # Multimodal config init vars

345 + limit_mm_per_prompt: Optional[dict[str, int]],

346 + media_io_kwargs: Optional[dict[str, dict[str, Any]]],

347 + mm_processor_kwargs: Optional[dict[str, Any]],

348 + mm_processor_cache_gb: Optional[float],

349 + mm_processor_cache_type: Optional[MMCacheType],

350 + mm_shm_cache_max_object_size_mb: Optional[int],

351 + mm_encoder_tp_mode: Optional[MMEncoderTPMode],

352 + interleave_mm_strings: Optional[bool],

353 + skip_mm_profiling: Optional[bool],

354 + video_pruning_rate: Optional[float],

355 + ) -> None:

356 # Set the default seed to 0 in V1.

357 # NOTE(woosuk): In V0, we set the default seed to None because the

358 # driver worker shares the same process as the user process, and thus

@@ -612,6 +616,7 @@ class ModelConfig:

616 mm_encoder_tp_mode=mm_encoder_tp_mode,

617 interleave_mm_strings=interleave_mm_strings,

618 skip_mm_profiling=skip_mm_profiling,

619 + video_pruning_rate=video_pruning_rate,

620 )

621

622 mm_config_kwargs = {

vllm/config/multimodal.py

@@ -78,6 +78,11 @@ class MultiModalConfig:

78 This reduces engine startup time but shifts the responsibility to users for

79 estimating the peak memory usage of the activation of multimodal encoder and

80 embedding cache."""

81 + video_pruning_rate: Optional[float] = None

82 + """Sets pruning rate for video pruning via Efficient Video Sampling.

83 + Value sits in range [0;1) and determines fraction of media tokens

84 + from each video to be pruned.

85 + """

87 def compute_hash(self) -> str:

88 """

@@ -118,3 +123,7 @@ class MultiModalConfig:

123 """

124 kwargs = self.mm_processor_kwargs or {}

125 return kwargs | dict(inference_kwargs)

126 +

127 + def is_multimodal_pruning_enabled(self):

128 + return (self.video_pruning_rate is not None

129 + and self.video_pruning_rate > 0)

vllm/engine/arg_utils.py

@@ -391,6 +391,7 @@ class EngineArgs:

391 mm_encoder_tp_mode: MMEncoderTPMode = MultiModalConfig.mm_encoder_tp_mode

392 io_processor_plugin: Optional[str] = None

393 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling

394 + video_pruning_rate: float = MultiModalConfig.video_pruning_rate

395 # LoRA fields

396 enable_lora: bool = False

397 enable_lora_bias: bool = LoRAConfig.bias_enabled

@@ -813,6 +814,9 @@ class EngineArgs:

814 multimodal_group.add_argument("--skip-mm-profiling",

815 **multimodal_kwargs["skip_mm_profiling"])

816

817 + multimodal_group.add_argument(

818 + "--video-pruning-rate", **multimodal_kwargs["video_pruning_rate"])

819 +

820 # LoRA related configs

821 lora_kwargs = get_kwargs(LoRAConfig)

822 lora_group = parser.add_argument_group(

@@ -1032,6 +1036,7 @@ class EngineArgs:

1036 model_impl=self.model_impl,

1037 override_attention_dtype=self.override_attention_dtype,

1038 logits_processors=self.logits_processors,

1039 + video_pruning_rate=self.video_pruning_rate,

1040 io_processor_plugin=self.io_processor_plugin,

1041 )

1042

vllm/model_executor/models/interfaces.py

+55

@@ -115,6 +115,42 @@ class SupportsMultiModal(Protocol):

115 ...

116

117

118 +@runtime_checkable

119 +class SupportsMultiModalPruning(Protocol):

120 + """The interface required for models that support returning both input

121 + embeddings and positions. Model may require custom positions for dynamic

122 + pruning of multimodal embeddings.

123 + """

124 + supports_multimodal_pruning: ClassVar[Literal[True]] = True

125 +

126 + def recompute_mrope_positions(

127 + self, input_ids: list[int],

128 + multimodal_embeddings: MultiModalEmbeddings,

129 + mrope_positions: torch.LongTensor, num_computed_tokens: int

130 + ) -> tuple[MultiModalEmbeddings, Tensor, int]:

131 + """

132 + Update part of input mrope positions (starting with

133 + num_computed_tokens index). Original mrope_positions are computed

134 + for unpruned sequence and becomes incorrect once pruning occurs,

135 + so once we prune media tokens we should reflect this in the

136 + mrope_positions before we feed it to LLM.

137 +

138 + Args:

139 + input_ids: (N,) All input tokens of the prompt containing

140 + entire sequence.

141 + multimodal_embeddings: Tuple of multimodal embeddings that

142 + fits into the prefill chunk that is being processed.

143 + mrope_positions: Existing mrope positions (3, N) for entire

144 + sequence

145 + num_computed_tokens: A number of computed tokens so far.

146 +

147 + Returns:

148 + Tuple of (multimodal_embeddings, mrope_positions,

149 + mrope_position_delta).

150 + """

151 + ...

152 +

153 +

154 @overload

155 def supports_multimodal(

156 model: type[object]) -> TypeIs[type[SupportsMultiModal]]:

@@ -142,6 +178,25 @@ def supports_multimodal_encoder_tp_data(

178 return getattr(model, "supports_encoder_tp_data", False)

179

180

181 +@overload

182 +def supports_multimodal_pruning(

183 + model: type[object]) -> TypeIs[type[SupportsMultiModalPruning]]:

184 + ...

185 +

186 +

187 +@overload

188 +def supports_multimodal_pruning(

189 + model: object) -> TypeIs[SupportsMultiModalPruning]:

190 + ...

191 +

192 +

193 +def supports_multimodal_pruning(

194 + model: Union[type[object], object],

195 +) -> Union[TypeIs[type[SupportsMultiModalPruning]],

196 + TypeIs[SupportsMultiModalPruning]]:

197 + return getattr(model, "supports_multimodal_pruning", False)

198 +

199 +

200 @runtime_checkable

201 class SupportsScoreTemplate(Protocol):

202 """The interface required for all models that support score template."""

vllm/model_executor/models/qwen2_5_vl.py

+226 -12

@@ -25,9 +25,9 @@

25 # See the License for the specific language governing permissions and

26 # limitations under the License.

27 """Inference-only Qwen2.5-VL model compatible with HuggingFace weights."""

28 -from collections.abc import Iterable, Mapping

28 +from collections.abc import Iterable, Mapping, Sequence

29 from functools import lru_cache, partial

30 -from typing import Annotated, Callable, Literal, Optional, Union

30 +from typing import Annotated, Any, Callable, Literal, Optional, Union

32 import torch

33 import torch.nn as nn

@@ -58,7 +58,13 @@ from vllm.model_executor.layers.quantization.gptq_marlin import (

58 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

59 from vllm.model_executor.models.module_mapping import MultiModelKeys

60 from vllm.multimodal import MULTIMODAL_REGISTRY

61 -from vllm.multimodal.inputs import MultiModalFieldConfig

61 +from vllm.multimodal.evs import (compute_mrope_for_media,

62 + compute_retained_tokens_count,

63 + compute_retention_mask,

64 + recompute_mrope_positions)

65 +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargs

66 +from vllm.multimodal.parse import MultiModalDataItems

67 +from vllm.multimodal.processing import PromptReplacement, PromptUpdate

68 from vllm.platforms import _Backend

69 from vllm.sequence import IntermediateTensors

70 from vllm.transformers_utils.config import uses_mrope

@@ -66,7 +72,8 @@ from vllm.utils import is_pin_memory_available

72 from vllm.utils.tensor_schema import TensorSchema, TensorShape

74 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

69 - SupportsMultiModal, SupportsPP, SupportsQuant)

75 + SupportsMultiModal, SupportsMultiModalPruning,

76 + SupportsPP, SupportsQuant)

77 from .qwen2_vl import Qwen2VLDummyInputsBuilder as Qwen2_5_VLDummyInputsBuilder

78 from .qwen2_vl import (Qwen2VLMultiModalProcessor, Qwen2VLProcessingInfo,

79 apply_rotary_pos_emb_vision)

@@ -86,9 +93,9 @@ class Qwen2_5_VLImagePixelInputs(TensorSchema):

93 - np: Number of patches

94 - ni: Number of images

95 - cps: Number of channels * patch_size * patch_size

89 -

96 +

97 Historical context:

91 - - pixel_values shape: (num_patches, num_channels * patch_size *

98 + - pixel_values shape: (num_patches, num_channels * patch_size *

99 patch_size)

100 - image_grid_thw shape: (num_images, 3) in (grid_t, grid_h, grid_w)

101 formatnum_channels * patch_size * patch_size

@@ -112,7 +119,7 @@ class Qwen2_5_VLImageEmbeddingInputs(TensorSchema):

119 - nf: Number of image features

120 - hs: Hidden size

121 - ni: Number of images

115 -

122 +

123 Historical context:

124 - image_embeds shape: (num_image_features, hidden_size)

125 - num_image_features varies based on the number and resolution of the

@@ -143,11 +150,11 @@ class Qwen2_5_VLVideoPixelInputs(TensorSchema):

150 Dimensions:

151 - np: Number of patches

152 - nv: Number of videos

146 - - ctps: Number of channels * temporal_patch_size * patch_size *

153 + - ctps: Number of channels * temporal_patch_size * patch_size *

154 patch_size

148 -

155 +

156 Historical context:

150 - - pixel_values_videos shape: (num_patches, num_channels *

157 + - pixel_values_videos shape: (num_patches, num_channels *

158 temporal_patch_size * patch_size * patch_size)

159 - video_grid_thw shape: (num_videos, 3) in (grid_t, grid_h, grid_w)

160 format

@@ -179,7 +186,7 @@ class Qwen2_5_VLVideoEmbeddingInputs(TensorSchema):

186 - nf: Number of video features

187 - hs: Hidden size

188 - nv: Number of videos

182 -

189 +

190 Historical context:

191 - video_embeds shape: (num_video_features, hidden_size)

192 - num_video_features varies based on the number and resolution of the

@@ -905,6 +912,55 @@ class Qwen2_5_VLMultiModalProcessor(Qwen2VLMultiModalProcessor):

912 second_per_grid_ts=MultiModalFieldConfig.batched("video"),

913 )

914

915 + def _get_prompt_updates(

916 + self,

917 + mm_items: MultiModalDataItems,

918 + hf_processor_mm_kwargs: Mapping[str, Any],

919 + out_mm_kwargs: MultiModalKwargs,

920 + ) -> Sequence[PromptUpdate]:

921 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

922 + image_processor = self.info.get_image_processor(

923 + **hf_processor_mm_kwargs)

924 + tokenizer = self.info.get_tokenizer()

925 + vocab = tokenizer.get_vocab()

926 +

927 + placeholder = {

928 + "image": vocab[hf_processor.image_token],

929 + "video": vocab[hf_processor.video_token],

930 + }

931 +

932 + merge_length = image_processor.merge_size**2

933 +

934 + def get_replacement_qwen2vl(item_idx: int, modality: str):

935 + out_item = out_mm_kwargs[modality][item_idx]

936 + grid_thw = out_item[f"{modality}_grid_thw"].data

937 + assert isinstance(grid_thw, torch.Tensor)

938 +

939 + num_tokens = int(grid_thw.prod()) // merge_length

940 +

941 + # EVS-specific code

942 + video_pruning_rate = self.info.ctx.get_mm_config(

943 + ).video_pruning_rate

944 + if (modality == "video" and video_pruning_rate is not None

945 + and video_pruning_rate > 0.0):

946 + num_tokens = compute_retained_tokens_count(

947 + grid_thw,

948 + image_processor.merge_size,

949 + video_pruning_rate,

950 + )

951 + # End of EVS-specific code

952 +

953 + return [placeholder[modality]] * num_tokens

954 +

955 + return [

956 + PromptReplacement(

957 + modality=modality,

958 + target=[placeholder[modality]],

959 + replacement=partial(get_replacement_qwen2vl,

960 + modality=modality),

961 + ) for modality in ("image", "video")

962 + ]

963 +

964

965 @MULTIMODAL_REGISTRY.register_processor(

966 Qwen2_5_VLMultiModalProcessor,

@@ -912,7 +968,8 @@ class Qwen2_5_VLMultiModalProcessor(Qwen2VLMultiModalProcessor):

968 dummy_inputs=Qwen2_5_VLDummyInputsBuilder)

969 class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

970 SupportsLoRA, SupportsPP,

915 - SupportsQuant):

971 + SupportsQuant,

972 + SupportsMultiModalPruning):

973

974 packed_modules_mapping = {

975 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

@@ -949,6 +1006,9 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1006 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

1007 self.config = config

1008 self.multimodal_config = multimodal_config

1009 + self.video_pruning_rate = multimodal_config.video_pruning_rate

1010 + self.is_multimodal_pruning_enabled = (

1011 + multimodal_config.is_multimodal_pruning_enabled())

1012

1013 if multimodal_config.get_limit_per_prompt("image") or \

1014 multimodal_config.get_limit_per_prompt("video"):

@@ -1090,6 +1150,36 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1150

1151 return image_embeds.split(sizes)

1152

1153 + def _postprocess_image_embeds_evs(

1154 + self, image_embeds_split: tuple[torch.Tensor, ...],

1155 + image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:

1156 + """

1157 + Append mrope positions for each for images.

1158 + This is necessary to recover correct mrope

1159 + positions after video pruning

1160 +

1161 + Args:

1162 + image_embeds_split: Tuple of image embeddings for

1163 + each image item.

1164 + image_input: Image input data.

1165 +

1166 + Returns:

1167 + Tuple of image embeddings for each image item.

1168 + Resulting embeddings will have extra 4 channels for

1169 + computed mrope positions.

1170 + """

1171 + merge_size = self.visual.spatial_merge_size

1172 + grid_thw = image_input["image_grid_thw"]

1173 + grid_thw_list = grid_thw.tolist()

1174 + image_embeds_out = []

1175 + for emb, size in zip(image_embeds_split, grid_thw_list):

1176 + positions = compute_mrope_for_media(size,

1177 + merge_size).to(emb.device)

1178 + emb = torch.cat([emb, positions], dim=1)

1179 + image_embeds_out.append(emb)

1180 + image_embeds_split = image_embeds_out

1181 + return tuple(image_embeds_split)

1182 +

1183 def _process_video_input(

1184 self,

1185 video_input: Qwen2_5_VLVideoInputs) -> tuple[torch.Tensor, ...]:

@@ -1119,6 +1209,114 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1209

1210 return video_embeds.split(sizes)

1211

1212 + def _postprocess_video_embeds_evs(

1213 + self, video_embeds_split: tuple[torch.Tensor, ...],

1214 + video_input: Qwen2_5_VLVideoInputs) -> tuple[torch.Tensor, ...]:

1215 + """

1216 + Prunes video embeddings via Efficient Video Sampling (EVS)

1217 + and then appends mrope positions for each retained embeddings

1218 +

1219 + Args:

1220 + video_embeds_split: Tuple of video embeddings for each video item.

1221 + video_input: Video input data.

1222 +

1223 + Returns:

1224 + Tuple of video embeddings for each video item.

1225 + Resulting embeddings will have extra 4 channels for

1226 + computed mrope positions.

1227 + """

1228 + grid_thw = video_input["video_grid_thw"]

1229 + assert grid_thw.ndim == 2

1230 + grid_thw_list = grid_thw.tolist()

1231 + merge_size = self.visual.spatial_merge_size

1232 +

1233 + # Cast to long to match the original code

1234 + # https://github.com/huggingface/transformers/blob/41980ce93e775f6c88500c51c8db7946fc6a2add/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py#L491 # noqa

1235 + second_per_grid_ts = video_input["second_per_grid_ts"].long()

1236 + tokens_per_second = self.config.vision_config.tokens_per_second

1237 +

1238 + video_embeds_out = []

1239 + for emb, size, video_second_per_grid_t in zip(video_embeds_split,

1240 + grid_thw_list,

1241 + second_per_grid_ts):

1242 + # For each video, we compute retention mask using EVS

1243 + retention_mask = compute_retention_mask(

1244 + emb,

1245 + size,

1246 + spatial_merge_size=self.visual.spatial_merge_size,

1247 + q=self.video_pruning_rate,

1248 + )

1249 + positions = compute_mrope_for_media(

1250 + size,

1251 + merge_size,

1252 + tokens_per_second=tokens_per_second,

1253 + video_second_per_grid=video_second_per_grid_t.item(),

1254 + ).to(emb.device)

1255 +

1256 + emb = emb[retention_mask]

1257 + positions = positions[retention_mask]

1258 + emb = torch.cat([emb, positions], dim=1)

1259 + video_embeds_out.append(emb)

1260 + return tuple(video_embeds_out)

1261 +

1262 + def recompute_mrope_positions(

1263 + self,

1264 + input_ids: list[int],

1265 + multimodal_embeddings: tuple[torch.Tensor, ...],

1266 + mrope_positions: torch.LongTensor,

1267 + num_computed_tokens: int,

1268 + ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, int]:

1269 + """

1270 + Update part of input mrope positions (starting with

1271 + num_computed_tokens index). Original mrope_positions are computed

1272 + for unpruned sequence and becomes incorrect once pruning occurs,

1273 + so once we prune media tokens we should reflect this in the

1274 + mrope_positions before we feed it to LLM.

1275 +

1276 + Args:

1277 + input_ids: (N,) All input tokens of the prompt (Containing

1278 + entire sequence).

1279 + multimodal_embeddings: Tuple of multimodal embeddings.

1280 + mrope_positions: Existing mrope positions (3, N) for entire

1281 + sequence

1282 + num_computed_tokens: A number of computed tokens so far.

1283 +

1284 + Returns:

1285 + Tuple of (multimodal_embeddings, mrope_positions,

1286 + mrope_position_delta).

1287 + """

1288 + image_token_id = self.config.image_token_id

1289 + video_token_id = self.config.video_token_id

1290 + vision_start_token_id = self.config.vision_start_token_id

1291 +

1292 + # Device

1293 + device = (multimodal_embeddings[0].device

1294 + if len(multimodal_embeddings) else mrope_positions.device)

1295 +

1296 + # Tensors

1297 + input_ids_t = torch.as_tensor(input_ids,

1298 + device=device,

1299 + dtype=torch.long)

1300 +

1301 + # fmt: off

1302 + mm_embeddings_out = [mm[:, :-4] for mm in

1303 + multimodal_embeddings]

1304 + mm_embeddings_pos = [mm[:, -4:].permute(1, 0).long() for mm in

1305 + multimodal_embeddings]

1306 + # fmt: in

1307 +

1308 + positions, mrope_positions_delta = recompute_mrope_positions(

1309 + input_ids_t,

1310 + mm_embeddings_pos,

1311 + mrope_positions,

1312 + num_computed_tokens,

1313 + vision_start_token_id,

1314 + image_token_id,

1315 + video_token_id,

1316 + )

1317 +

1318 + return tuple(mm_embeddings_out), positions, mrope_positions_delta

1319 +

1320 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

1321 mm_input_by_modality = {}

1322

@@ -1156,9 +1354,17 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1354 multimodal_input = mm_input_by_modality[modality]

1355 if modality == "image":

1356 vision_embeddings = self._process_image_input(multimodal_input)

1357 + if self.is_multimodal_pruning_enabled:

1358 + vision_embeddings = self._postprocess_image_embeds_evs(

1359 + vision_embeddings, multimodal_input

1360 + )

1361 multimodal_embeddings += vision_embeddings

1362 if modality == "video":

1363 video_embeddings = self._process_video_input(multimodal_input)

1364 + if self.is_multimodal_pruning_enabled:

1365 + video_embeddings = self._postprocess_video_embeds_evs(

1366 + video_embeddings, multimodal_input

1367 + )

1368 multimodal_embeddings += video_embeddings

1369 return multimodal_embeddings

1370

@@ -1184,6 +1390,10 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1390 inputs_embeds = self.get_input_embeddings(input_ids)

1391 if image_input is not None:

1392 image_embeds = self._process_image_input(image_input)

1393 + if self.is_multimodal_pruning_enabled:

1394 + image_embeds = self._postprocess_image_embeds_evs(

1395 + image_embeds, image_input

1396 + )

1397 inputs_embeds = merge_multimodal_embeddings(

1398 input_ids,

1399 inputs_embeds,

@@ -1193,6 +1403,10 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1403

1404 if video_input is not None:

1405 video_embeds = self._process_video_input(video_input)

1406 + if self.is_multimodal_pruning_enabled:

1407 + video_embeds = self._postprocess_video_embeds_evs(

1408 + video_embeds, video_input

1409 + )

1410 inputs_embeds = merge_multimodal_embeddings(

1411 input_ids,

1412 inputs_embeds,

vllm/multimodal/evs.py new

+273

@@ -0,0 +1,273 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +#

5 +# NVIDIA CORPORATION and its licensors retain all intellectual property

6 +# and proprietary rights in and to this software, related documentation

7 +# and any modifications thereto. Any use, reproduction, disclosure or

8 +# distribution of this software and related documentation without an express

9 +# license agreement from NVIDIA CORPORATION is strictly prohibited.

10 +

11 +import typing

12 +

13 +import torch

14 +

15 +

16 +def compute_retained_tokens_count(video_size_thw: torch.LongTensor,

17 + spatial_merge_size: int, q: float) -> int:

18 + """

19 + Compute the number of retained tokens for a given video.

20 + Method ensures that we retain all the tokens from the first frame

21 + regardless of the pruning rate.

22 +

23 + Args:

24 + video_size_thw: The size of the video in the format of (T, H, W).

25 + spatial_merge_size: The size of the spatial merge.

26 + q: The pruning rate.

27 +

28 + Returns:

29 + The number of retained tokens.

30 + """

31 + T, H, W = map(int, video_size_thw)

32 + min_num_tokens = (H // spatial_merge_size) * (W // spatial_merge_size)

33 + evs_num_tokens = int(T * min_num_tokens * (1 - q))

34 + return max(min_num_tokens, evs_num_tokens)

35 +

36 +

37 +def compute_retention_mask(

38 + video_embeds: torch.Tensor,

39 + video_size_thw: torch.LongTensor,

40 + spatial_merge_size: int,

41 + q: float,

42 +) -> torch.Tensor:

43 + """

44 + Computes the retention mask for input video embeddings.

45 +

46 + Args:

47 + video_embeds (`torch.Tensor`): The input video embeddings

48 + of shape `(T * H * W // spatial_merge_size ^ 2, hidden_size)`

49 + video_size_thw (`torch.LongTensor` of shape `(3)`):

50 + The temporal, height and width of video.

51 + spatial_merge_size: Size reduction for rows & cols dimensions.

52 + q: (`float`): Pruning rate factor [0,1)

53 +

54 + Returns:

55 + `torch.Tensor`: The retention mask for the video embeddings of

56 + `(T * H * W // spatial_merge_size ^ 2)` shape.

57 + """

58 + T, H, W = video_size_thw

59 +

60 + # Use reshape instead of einops to avoid graph breaks

61 + video_embeds = video_embeds.reshape(

62 + T,

63 + H // spatial_merge_size,

64 + W // spatial_merge_size,

65 + video_embeds.size(-1),

66 + )

67 +

68 + # Core EVS

69 + similarity = torch.nn.functional.cosine_similarity(video_embeds[1:, ...],

70 + video_embeds[:-1, ...],

71 + dim=-1)

72 + dissimilarity = 1 - similarity

73 +

74 + # Always ensure we include all tokens from the first frame

75 + dissimilarity = torch.cat(

76 + [255 * torch.ones_like(video_embeds[:1, :, :, 0]), dissimilarity],

77 + dim=0)

78 +

79 + dissimilarity_flat = dissimilarity.view(-1)

80 + order = torch.argsort(dissimilarity_flat,

81 + dim=-1,

82 + descending=True,

83 + stable=True)

84 + retain_num_tokens = compute_retained_tokens_count(video_size_thw,

85 + spatial_merge_size, q)

86 + topk_indices = order[:retain_num_tokens]

87 +

88 + retention_mask = torch.zeros_like(dissimilarity_flat, dtype=torch.bool)

89 + retention_mask[topk_indices] = True

90 + retention_mask = retention_mask.reshape(dissimilarity.size())

91 +

92 + mask = retention_mask.view(-1) # "T H W -> (T H W)"

93 + return mask

94 +

95 +

96 +def compute_mrope_for_media(

97 + video_size_thw: torch.LongTensor,

98 + spatial_merge_size: int,

99 + tokens_per_second: float = 1.0,

100 + video_second_per_grid: float = 1.0,

101 +) -> torch.Tensor:

102 + """

103 + Computes the mrope for video embeddings based on the grid dimensions.

104 + Computed mrope positions match original qwen 2.5 implementation,

105 + but positions are built for media being the first element in sequence.

106 +

107 + Args:

108 + video_size_thw: Media size (num frames, rows, cols)

109 + spatial_merge_size: Size reduction for rows & cols dimensions.

110 + tokens_per_second: Number of tokens per second.

111 + video_second_per_grid: Number of seconds per video.

112 +

113 + Returns:

114 + Tensor of shape `(T * H * W, 4)` where last dimension

115 + represents mrope positions [0:3), while the last channel

116 + contains value of llm_grid_w repeated for all positions.

117 + """

118 + llm_grid_t = video_size_thw[0]

119 + llm_grid_h = video_size_thw[1] // spatial_merge_size

120 + llm_grid_w = video_size_thw[2] // spatial_merge_size

121 +

122 + t_index = ((torch.arange(llm_grid_t).view(-1, 1).expand(

123 + -1, llm_grid_h * llm_grid_w).mul(

124 + tokens_per_second * video_second_per_grid)).long().flatten())

125 + h_index = (torch.arange(llm_grid_h).view(1, -1,

126 + 1).expand(llm_grid_t, -1,

127 + llm_grid_w).flatten())

128 + w_index = (torch.arange(llm_grid_w).view(1, 1, -1).expand(

129 + llm_grid_t, llm_grid_h, -1).flatten())

130 + llm_grid_w = (torch.tensor([llm_grid_w

131 + ]).view(1, 1,

132 + 1).expand(llm_grid_t, llm_grid_h,

133 + llm_grid_w).flatten())

134 +

135 + positions = torch.stack([t_index, h_index, w_index, llm_grid_w], dim=1)

136 + return positions

137 +

138 +

139 +def recompute_mrope_positions(

140 + input_ids: torch.LongTensor,

141 + multimodal_positions: list[torch.Tensor],

142 + mrope_positions: torch.LongTensor,

143 + num_computed_tokens: int,

144 + vision_start_token_id: int,

145 + image_token_id: int,

146 + video_token_id: int,

147 +) -> tuple[torch.LongTensor, int]:

148 + """

149 + Update part of input mrope positions.

150 + Original mrope_positions are computed incorrectly, so once we prune media

151 + tokens we should reflect this in the mrope positions for the LLM.

152 +

153 + This method supports chunked prefill approach where

154 + multimodal_embeddings are passed to LLM in chunks, so input

155 + multimodal_embeddings may contain zero, some or even some part of all

156 + multimodal_embeddings for a given prompt.

157 +

158 + Each multimodal_positions has 4 extra channels

159 + (First 3 channels corresponds to original 3 mrope positions, last channel

160 + is the maximum width of the media repeated). Provided multimodal_positions

161 + do not reflect location of media position in sequence - they are computed

162 + like the media is in the 0-th position in the sequence.

163 +

164 + Method works as follows: it recomputes mrope_positions starting from the

165 + `num_computed_tokens` for `total_len_of_multimodal_embeddings` and then

166 + shifts all text tokens that goes after total_len_of_multimodal_embeddings.

167 +

168 + It also handles case when multimodal_embeddings is partial

169 + (e.g. one media is split into two prefill stages)

170 +

171 + Args:

172 + input_ids: (N,) All input tokens of the prompt (entire sequence).

173 + multimodal_positions: List of mrope positsions for each media.

174 + mrope_positions: Existing mrope positions (4, N) for entire sequence.

175 + num_computed_tokens: A number of computed tokens so far.

176 + vision_start_token_id: Token indicating start of vision media.

177 + image_token_id: Image token id

178 + video_token_id: Video token id

179 +

180 + Returns:

181 + Tuple of (mrope_positions, mrope_position_delta).

182 + """

183 +

184 + # Tensors

185 + positions: torch.LongTensor = typing.cast(

186 + torch.LongTensor, mrope_positions.clone()) # (3, N)

187 + N = input_ids.numel()

188 +

189 + image_mask = input_ids.eq(image_token_id)

190 + video_mask = input_ids.eq(video_token_id)

191 + media_mask = image_mask | video_mask

192 + text_mask = ~media_mask

193 +

194 + # Early exit: no media in this chunk

195 + if len(multimodal_positions) == 0:

196 + delta = (int((positions.max().item() + 1) -

197 + N) if positions.numel() else -N)

198 + return positions, delta

199 +

200 + total_mm_tokens = torch.count_nonzero(media_mask)

201 + seen_mm_tokens = torch.count_nonzero(media_mask[:num_computed_tokens])

202 +

203 + # Early exit: we've updated positions for all media tokens

204 + # (and consequently - for all remaining text tokens)

205 + if seen_mm_tokens == total_mm_tokens:

206 + delta = (int((positions.max().item() + 1) -

207 + N) if positions.numel() else -N)

208 + return positions, delta

209 +

210 + vision_start_indices = (input_ids == vision_start_token_id).nonzero(

211 + as_tuple=True)[0]

212 +

213 + for mm_pos in multimodal_positions:

214 + # Each mm_pos can be a complete embedding for single media

215 + # or it can be a part of a single media (due to chunked prefill)

216 +

217 + # Cases to cover

218 + # - Current prefill chunk has no vision start indexes at all

219 + # - Vision start token appeared in previous prefill round

220 + # - Regular case

221 + seen_vision_start_indices = vision_start_indices[vision_start_indices <

222 + num_computed_tokens]

223 +

224 + if len(seen_vision_start_indices):

225 + # If we have encountered some vision start indexes,

226 + # then we should check the condition:

227 + # | --- prefill 1 ------| ---- prefill 2 ----- |

228 + # | TTTTTTTTTSVVVVVVVVVV|VVVVVVTTTTTTTTTTTTTTTT|

229 + last_vision_start_token = seen_vision_start_indices[-1]

230 + seem_mm_tokens_before_last_vision_start = torch.count_nonzero(

231 + media_mask[:last_vision_start_token])

232 + in_the_middle_of_media = (

233 + seen_mm_tokens > seem_mm_tokens_before_last_vision_start)

234 +

235 + if in_the_middle_of_media:

236 + mm_embeddings_seen = (seen_mm_tokens -

237 + seem_mm_tokens_before_last_vision_start)

238 + global_mm_start = last_vision_start_token

239 + else:

240 + # We have completed previous mm_embedding part and

241 + # ready to start a new one

242 + next_vision_start_token = vision_start_indices[

243 + vision_start_indices >= num_computed_tokens][0]

244 + mm_embeddings_seen = 0

245 + global_mm_start = next_vision_start_token

246 +

247 + else:

248 + # If there were no vision start indexes so far,

249 + # let's find first vision start index

250 + next_vision_start_token = vision_start_indices[

251 + vision_start_indices >= num_computed_tokens][0]

252 +

253 + mm_embeddings_seen = 0

254 + global_mm_start = next_vision_start_token

255 +

256 + # Offset right after vision_start_token

257 + base = positions[-1, global_mm_start] + 1

258 + local_start = global_mm_start + 1 + mm_embeddings_seen

259 + local_end = local_start + mm_pos.shape[1]

260 + positions[:, local_start:local_end] = mm_pos[0:3] + base

261 +

262 + # mm_pos[3, 0] is the max width of the media

263 + offset = mm_pos[3, 0] + base

264 +

265 + text_pos_sum = torch.cumsum(text_mask[local_end:].long(), dim=0)

266 +

267 + positions[:, local_end:N] = text_pos_sum + offset - 1

268 +

269 + # Include distance to the next vision start token

270 + num_computed_tokens += mm_pos.shape[1]

271 +

272 + mrope_positions_delta = (positions.max() + 1 - N).item()

273 + return positions, mrope_positions_delta

vllm/v1/worker/gpu_model_runner.py

+67 -16

@@ -40,11 +40,15 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

40 from vllm.model_executor.layers.mamba.abstract import MambaBase

41 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

42 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader

43 +# yapf conflicts with isort for this block

44 +# yapf: disable

45 from vllm.model_executor.models.interfaces import (SupportsMultiModal,

46 is_mixture_of_experts,

47 supports_eagle3,

48 supports_mrope,

49 + supports_multimodal_pruning,

50 supports_transcription)

51 +# yapf: enable

52 from vllm.model_executor.models.interfaces_base import (

53 VllmModelForPooling, is_pooling_model, is_text_generation_model)

54 from vllm.multimodal import MULTIMODAL_REGISTRY

@@ -206,7 +210,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

210 self.enable_prompt_embeds = model_config.enable_prompt_embeds

211 self.is_multimodal_raw_input_only_model = (

212 model_config.is_multimodal_raw_input_only_model)

209 -

213 + # This will be overridden in load_model()

214 + self.is_multimodal_pruning_enabled = False

215 self.max_model_len = model_config.max_model_len

216 self.dcp_world_size = self.parallel_config.decode_context_parallel_size

217 self.max_num_tokens = scheduler_config.max_num_batched_tokens

@@ -1530,29 +1535,47 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1535 # encoder outputs.

1536 model = cast(SupportsMultiModal, self.model)

1537 encoder_outputs = []

1533 - for _, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(

1538 + for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(

1539 mm_kwargs,

1540 device=self.device,

1541 pin_memory=self.pin_memory,

1542 merge_by_field_config=model.merge_by_field_config,

1543 ):

1539 - # Run the encoder.

1540 - # `curr_group_outputs` is either of the following:

1541 - # 1. A tensor of shape (num_items, feature_size, hidden_size)

1542 - # in case feature_size is fixed across all multimodal items.

1543 - # 2. A list or tuple (length: num_items) of tensors, each of shape

1544 - # (feature_size, hidden_size) in case the feature size is dynamic

1545 - # depending on the input multimodal items.

1546 - curr_group_outputs = model.get_multimodal_embeddings(

1547 - **mm_kwargs_group)

1544 + # (ekhvedchenia): Temporary hack to limit peak memory usage when

1545 + # processing multimodal data.This solves the issue with scheduler

1546 + # putting too many video samples into a single batch. Scheduler

1547 + # uses pruned vision tokens count to compare it versus compute

1548 + # budget which is incorrect (Either input media size or non-pruned

1549 + # output vision tokens count should be considered)

1550 + curr_group_outputs = []

1551 +

1552 + if self.is_multimodal_pruning_enabled and modality == "video":

1553 + micro_batch_size = 1

1554 + for i in range(0, num_items, micro_batch_size):

1555 + micro_batch_mm_inputs = dict(

1556 + (k, v[i:i + micro_batch_size])

1557 + for k, v in mm_kwargs_group.items())

1558 +

1559 + micro_batch_outputs = model.get_multimodal_embeddings(

1560 + **micro_batch_mm_inputs)

1561 +

1562 + curr_group_outputs.extend(micro_batch_outputs)

1563 + else:

1564 + # Run the encoder.

1565 + # `curr_group_outputs` is either of the following:

1566 + # 1. A tensor of shape (num_items, feature_size, hidden_size)

1567 + # in case feature_size is fixed across all multimodal items.

1568 + # 2. A list or tuple (length: num_items) of tensors,

1569 + # each of shape (feature_size, hidden_size) in case the feature

1570 + # size is dynamic depending on the input multimodal items.

1571 + curr_group_outputs = model.get_multimodal_embeddings(

1572 + **mm_kwargs_group)

1573

1574 sanity_check_mm_encoder_outputs(

1575 curr_group_outputs,

1576 expected_num_items=num_items,

1577 )

1553 -

1554 - for output in curr_group_outputs:

1555 - encoder_outputs.append(output)

1578 + encoder_outputs.extend(curr_group_outputs)

1579

1580 # Cache the encoder outputs by mm_hash

1581 for (mm_hash, pos_info), output in zip(mm_hashes_pos, encoder_outputs):

@@ -1566,8 +1589,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1589 scheduler_output: "SchedulerOutput",

1590 shift_computed_tokens: int = 0,

1591 ) -> list[torch.Tensor]:

1592 + should_sync_mrope_positions = False

1593 mm_embeds: list[torch.Tensor] = []

1594 for req_id in self.input_batch.req_ids:

1595 + mm_embeds_req: list[torch.Tensor] = []

1596 +

1597 num_scheduled_tokens = scheduler_output.num_scheduled_tokens[

1598 req_id]

1599 req_state = self.requests[req_id]

@@ -1609,7 +1635,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1635 encoder_output[start_idx:end_idx],

1636 is_embed=is_embed,

1637 )

1612 - mm_embeds.append(mm_embeds_item)

1638 + mm_embeds_req.append(mm_embeds_item)

1639 +

1640 + if self.is_multimodal_pruning_enabled and self.uses_mrope:

1641 + should_sync_mrope_positions = True

1642 + mm_embeds_req, new_mrope_positions, new_delta = (

1643 + self.model.recompute_mrope_positions(

1644 + input_ids=req_state.prompt_token_ids,

1645 + multimodal_embeddings=mm_embeds_req,

1646 + mrope_positions=req_state.mrope_positions,

1647 + num_computed_tokens=req_state.num_computed_tokens,

1648 + ))

1649 + assert req_state.mrope_positions is not None

1650 + req_state.mrope_positions.copy_(new_mrope_positions)

1651 + req_state.mrope_position_delta = new_delta

1652 +

1653 + mm_embeds.extend(mm_embeds_req)

1654 +

1655 + if should_sync_mrope_positions:

1656 + self._calc_mrope_positions(scheduler_output)

1657 + self.mrope_positions.copy_to_gpu(

1658 + scheduler_output.total_num_scheduled_tokens)

1659 +

1660 return mm_embeds

1661

1662 def _extract_encoder_inputs(

@@ -2589,6 +2636,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2636 time_after_load - time_before_load)

2637 prepare_communication_buffer_for_model(self.model)

2638

2639 + self.is_multimodal_pruning_enabled = (supports_multimodal_pruning(

2640 + self.model) and self.model_config.multimodal_config.

2641 + is_multimodal_pruning_enabled())

2642 +

2643 if is_mixture_of_experts(

2644 self.model) and self.parallel_config.enable_eplb:

2645 logger.info("EPLB is enabled for model %s.",

@@ -2843,7 +2894,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2894 Args:

2895 num_tokens: Number of tokens to run the dummy forward pass.

2896 cudagraph_runtime_mode: used to control the behavior.

2846 - - if not set will determine the cudagraph mode based on using

2897 + - if not set will determine the cudagraph mode based on using

2898 the self.cudagraph_dispatcher.

2899 - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run

2900 - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.