| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 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( |
| 17 | tokens_per_frame: int, num_frames: int, q: float |
| 18 | ) -> int: |
| 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 | tokens_per_frame: The number of tokens per frame. |
| 25 | num_frames: The total number of frames. |
| 26 | q: The pruning rate. |
| 27 | |
| 28 | Returns: |
| 29 | The number of retained tokens. |
| 30 | |
| 31 | """ |
| 32 | total_tokens = tokens_per_frame * num_frames |
| 33 | evs_num_tokens = int(total_tokens * (1 - q)) |
| 34 | min_num_tokens = tokens_per_frame |
| 35 | return max(min_num_tokens, evs_num_tokens) |
| 36 | |
| 37 | |
| 38 | def compute_retention_mask( |
| 39 | video_embeds: torch.Tensor, |
| 40 | video_size_thw: torch.LongTensor | tuple[int, int, int], |
| 41 | spatial_merge_size: int, |
| 42 | q: float, |
| 43 | ) -> torch.Tensor: |
| 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 | """ |
| 59 | T, H, W = map(int, video_size_thw) |
| 60 | |
| 61 | # Use reshape instead of einops to avoid graph breaks |
| 62 | video_embeds = video_embeds.reshape( |
| 63 | T, |
| 64 | H // spatial_merge_size, |
| 65 | W // spatial_merge_size, |
| 66 | video_embeds.size(-1), |
| 67 | ) |
| 68 | tokens_per_frame = (H // spatial_merge_size) * (W // spatial_merge_size) |
| 69 | # Core EVS |
| 70 | similarity = torch.nn.functional.cosine_similarity( |
| 71 | video_embeds[1:, ...], video_embeds[:-1, ...], dim=-1 |
| 72 | ) |
| 73 | dissimilarity = 1 - similarity |
| 74 | |
| 75 | # Always ensure we include all tokens from the first frame |
| 76 | dissimilarity = torch.cat( |
| 77 | [255 * torch.ones_like(video_embeds[:1, :, :, 0]), dissimilarity], dim=0 |
| 78 | ) |
| 79 | |
| 80 | dissimilarity_flat = dissimilarity.view(-1) |
| 81 | order = torch.argsort(dissimilarity_flat, dim=-1, descending=True, stable=True) |
| 82 | retain_num_tokens = compute_retained_tokens_count( |
| 83 | tokens_per_frame=tokens_per_frame, num_frames=T, q=q |
| 84 | ) |
| 85 | topk_indices = order[:retain_num_tokens] |
| 86 | |
| 87 | retention_mask = torch.zeros_like(dissimilarity_flat, dtype=torch.bool) |
| 88 | retention_mask.index_fill_(0, topk_indices, True) |
| 89 | retention_mask = retention_mask.reshape(dissimilarity.size()) |
| 90 | |
| 91 | mask = retention_mask.view(-1) # "T H W -> (T H W)" |
| 92 | return mask |
| 93 | |
| 94 | |
| 95 | def compute_mrope_for_media( |
| 96 | video_size_thw: torch.LongTensor, |
| 97 | spatial_merge_size: int, |
| 98 | tokens_per_second: float = 1.0, |
| 99 | video_second_per_grid: float = 1.0, |
| 100 | ) -> torch.Tensor: |
| 101 | """Computes the mrope for video embeddings based on the grid dimensions. |
| 102 | Computed mrope positions match original qwen 2.5 implementation, |
| 103 | but positions are built for media being the first element in sequence. |
| 104 | |
| 105 | Args: |
| 106 | video_size_thw: Media size (num frames, rows, cols) |
| 107 | spatial_merge_size: Size reduction for rows & cols dimensions. |
| 108 | tokens_per_second: Number of tokens per second. |
| 109 | video_second_per_grid: Number of seconds per video. |
| 110 | |
| 111 | Returns: |
| 112 | Tensor of shape `(T * H * W, 4)` where last dimension |
| 113 | represents mrope positions [0:3), while the last channel |
| 114 | contains value of llm_grid_w repeated for all positions. |
| 115 | |
| 116 | """ |
| 117 | llm_grid_t = video_size_thw[0] |
| 118 | llm_grid_h = video_size_thw[1] // spatial_merge_size |
| 119 | llm_grid_w = video_size_thw[2] // spatial_merge_size |
| 120 | |
| 121 | t_index = ( |
| 122 | ( |
| 123 | torch.arange(llm_grid_t) |
| 124 | .view(-1, 1) |
| 125 | .expand(-1, llm_grid_h * llm_grid_w) |
| 126 | .mul(tokens_per_second * video_second_per_grid) |
| 127 | ) |
| 128 | .long() |
| 129 | .flatten() |
| 130 | ) |
| 131 | h_index = ( |
| 132 | torch.arange(llm_grid_h) |
| 133 | .view(1, -1, 1) |
| 134 | .expand(llm_grid_t, -1, llm_grid_w) |
| 135 | .flatten() |
| 136 | ) |
| 137 | w_index = ( |
| 138 | torch.arange(llm_grid_w) |
| 139 | .view(1, 1, -1) |
| 140 | .expand(llm_grid_t, llm_grid_h, -1) |
| 141 | .flatten() |
| 142 | ) |
| 143 | llm_grid_w = ( |
| 144 | torch.tensor([llm_grid_w]) |
| 145 | .view(1, 1, 1) |
| 146 | .expand(llm_grid_t, llm_grid_h, llm_grid_w) |
| 147 | .flatten() |
| 148 | ) |
| 149 | |
| 150 | positions = torch.stack([t_index, h_index, w_index, llm_grid_w], dim=1) |
| 151 | return positions |
| 152 | |
| 153 | |
| 154 | def recompute_mrope_positions( |
| 155 | input_ids: torch.LongTensor, |
| 156 | multimodal_positions: list[torch.Tensor], |
| 157 | mrope_positions: torch.LongTensor, |
| 158 | num_computed_tokens: int, |
| 159 | vision_start_token_id: int, |
| 160 | image_token_id: int, |
| 161 | video_token_id: int, |
| 162 | ) -> tuple[torch.LongTensor, int]: |
| 163 | """Update part of input mrope positions. |
| 164 | Original mrope_positions are computed incorrectly, so once we prune media |
| 165 | tokens we should reflect this in the mrope positions for the LLM. |
| 166 | |
| 167 | This method supports chunked prefill approach where |
| 168 | multimodal_embeddings are passed to LLM in chunks, so input |
| 169 | multimodal_embeddings may contain zero, some or even some part of all |
| 170 | multimodal_embeddings for a given prompt. |
| 171 | |
| 172 | Each multimodal_positions has 4 or 5 extra channels |
| 173 | (first 3 channels correspond to the original 3 mrope positions; |
| 174 | remaining channels vary by model — see below). Provided multimodal_positions |
| 175 | do not reflect location of media position in sequence - they are computed |
| 176 | like the media is in the 0-th position in the sequence. |
| 177 | |
| 178 | Method works as follows: it recomputes mrope_positions starting from the |
| 179 | `num_computed_tokens` for `total_len_of_multimodal_embeddings` and then |
| 180 | shifts all text tokens that goes after total_len_of_multimodal_embeddings. |
| 181 | |
| 182 | It also handles case when multimodal_embeddings is partial |
| 183 | (e.g. one media is split into two prefill stages) |
| 184 | |
| 185 | Args: |
| 186 | input_ids: (N,) All input tokens of the prompt (entire sequence). |
| 187 | multimodal_positions: List of mrope positions for each media. |
| 188 | If a given element is of shape (4, N), it is assumed to only describe |
| 189 | positions for video / image embeddings. This is the case of e.g. Qwen2.5 VL, |
| 190 | where each multimodal input is a contiguous chunk of embeddings. |
| 191 | The expected channels are [t, h, w, max_width]. |
| 192 | If it is of shape (5, N), it is assumed to possibly describe positions for |
| 193 | both video / image embeddings, as well as text embeddings. This is the case |
| 194 | of e.g. Qwen3 VL, where each video inputs are comprised of individual |
| 195 | frames' embeddings, interleaved with embeddings for timestamp tokens, |
| 196 | and vision start / end tokens. The expected channels are |
| 197 | [t, h, w, is_vision_start, is_vision]. |
| 198 | mrope_positions: Existing mrope positions (4, N) for entire sequence. |
| 199 | num_computed_tokens: A number of computed tokens so far. |
| 200 | vision_start_token_id: Token indicating start of vision media. |
| 201 | image_token_id: Image token id |
| 202 | video_token_id: Video token id |
| 203 | |
| 204 | Returns: |
| 205 | Tuple of (mrope_positions, mrope_position_delta). |
| 206 | |
| 207 | """ |
| 208 | # Tensors |
| 209 | positions: torch.LongTensor = typing.cast( |
| 210 | torch.LongTensor, mrope_positions.clone() |
| 211 | ) # (3, N) |
| 212 | N = input_ids.numel() |
| 213 | |
| 214 | image_mask = input_ids.eq(image_token_id) |
| 215 | video_mask = input_ids.eq(video_token_id) |
| 216 | media_mask = image_mask | video_mask |
| 217 | text_mask = ~media_mask |
| 218 | |
| 219 | # Early exit: no media in this chunk |
| 220 | if len(multimodal_positions) == 0: |
| 221 | delta = int((positions.max().item() + 1) - N) if positions.numel() else -N |
| 222 | return positions, delta |
| 223 | |
| 224 | total_mm_tokens = torch.count_nonzero(media_mask) |
| 225 | seen_mm_tokens = torch.count_nonzero(media_mask[:num_computed_tokens]) |
| 226 | |
| 227 | # Early exit: we've updated positions for all media tokens |
| 228 | # (and consequently - for all remaining text tokens) |
| 229 | if seen_mm_tokens == total_mm_tokens: |
| 230 | delta = int((positions.max().item() + 1) - N) if positions.numel() else -N |
| 231 | return positions, delta |
| 232 | |
| 233 | vision_start_indices = (input_ids == vision_start_token_id).nonzero(as_tuple=True)[ |
| 234 | 0 |
| 235 | ] |
| 236 | |
| 237 | for mm_pos in multimodal_positions: |
| 238 | # Each mm_pos can be a complete embedding for single media |
| 239 | # or it can be a part of a single media (due to chunked prefill) |
| 240 | |
| 241 | # Cases to cover |
| 242 | # - Current prefill chunk has no vision start indexes at all |
| 243 | # - Vision start token appeared in previous prefill round |
| 244 | # - Regular case |
| 245 | has_video_tokens = False |
| 246 | num_timestamp_tokens = 0 |
| 247 | if mm_pos.shape[0] == 5 and mm_pos.shape[1] > 0: |
| 248 | # mm_pos[4, :] indicates which positions are for video embeddings. |
| 249 | # If there are no video embeddings, skip timestamp adjustment. |
| 250 | has_video_tokens = torch.any(mm_pos[4, :]).item() |
| 251 | if has_video_tokens: |
| 252 | # Channel 3 flags VISION_START tokens. Timestamp tokens |
| 253 | # precede the first VISION_START, so its index gives us the |
| 254 | # exact timestamp count. This is robust even when early |
| 255 | # frames have all their video tokens pruned (which would |
| 256 | # push argmax(channel 4) far into a later frame). |
| 257 | first_vs = (mm_pos[3, :] == 1).nonzero(as_tuple=True)[0] |
| 258 | num_timestamp_tokens = first_vs[0].item() if len(first_vs) > 0 else 0 |
| 259 | |
| 260 | seen_vision_start_indices = vision_start_indices[ |
| 261 | vision_start_indices < num_computed_tokens |
| 262 | ] |
| 263 | |
| 264 | if len(seen_vision_start_indices): |
| 265 | # If we have encountered some vision start indexes, |
| 266 | # then we should check the condition: |
| 267 | # | --- prefill 1 ------| ---- prefill 2 ----- | |
| 268 | # | TTTTTTTTTSVVVVVVVVVV|VVVVVVTTTTTTTTTTTTTTTT| |
| 269 | last_vision_start_token = seen_vision_start_indices[-1] |
| 270 | seem_mm_tokens_before_last_vision_start = torch.count_nonzero( |
| 271 | media_mask[:last_vision_start_token] |
| 272 | ) |
| 273 | in_the_middle_of_media = ( |
| 274 | seen_mm_tokens > seem_mm_tokens_before_last_vision_start |
| 275 | ) |
| 276 | # For Qwen3 VL, we can be inside a media segment even before any |
| 277 | # video tokens appear (timestamp tokens are text). If we've passed |
| 278 | # the last vision_start token but haven't reached the first video |
| 279 | # embedding, treat this as "in the middle of media". |
| 280 | if ( |
| 281 | not in_the_middle_of_media |
| 282 | and has_video_tokens |
| 283 | and num_computed_tokens > last_vision_start_token |
| 284 | and num_computed_tokens |
| 285 | <= last_vision_start_token + num_timestamp_tokens + 1 |
| 286 | ): |
| 287 | in_the_middle_of_media = True |
| 288 | |
| 289 | if in_the_middle_of_media: |
| 290 | mm_embeddings_seen = ( |
| 291 | seen_mm_tokens - seem_mm_tokens_before_last_vision_start |
| 292 | ) |
| 293 | global_mm_start = last_vision_start_token |
| 294 | else: |
| 295 | # We have completed previous mm_embedding part and |
| 296 | # ready to start a new one |
| 297 | next_vision_start_token = vision_start_indices[ |
| 298 | vision_start_indices >= num_computed_tokens |
| 299 | ][0] |
| 300 | mm_embeddings_seen = 0 |
| 301 | global_mm_start = next_vision_start_token |
| 302 | |
| 303 | else: |
| 304 | # If there were no vision start indexes so far, |
| 305 | # let's find first vision start index |
| 306 | next_vision_start_token = vision_start_indices[ |
| 307 | vision_start_indices >= num_computed_tokens |
| 308 | ][0] |
| 309 | |
| 310 | mm_embeddings_seen = 0 |
| 311 | global_mm_start = next_vision_start_token |
| 312 | |
| 313 | # For Qwen3 VL, mm_pos includes timestamp tokens before vision_start |
| 314 | # when starting a new media. Adjust global_mm_start to point to where |
| 315 | # the sequence actually begins (before timestamp tokens). |
| 316 | adjusted_for_timestamps = False |
| 317 | if mm_pos.shape[0] == 5 and mm_embeddings_seen == 0 and has_video_tokens: |
| 318 | # NOTE: -1 is because there is a vision start token right after |
| 319 | # timestamp tokens before any video embeddings appear. |
| 320 | |
| 321 | # Adjust global_mm_start to point to the first timestamp token |
| 322 | # instead of the vision_start token. |
| 323 | global_mm_start -= num_timestamp_tokens |
| 324 | adjusted_for_timestamps = True |
| 325 | |
| 326 | # Offset calculation depends on whether we adjusted for timestamp tokens |
| 327 | if adjusted_for_timestamps: |
| 328 | # Start from position before the first timestamp token |
| 329 | base = positions[-1, global_mm_start - 1] + 1 |
| 330 | local_start = global_mm_start + mm_embeddings_seen |
| 331 | else: |
| 332 | # Original logic: start after vision_start_token |
| 333 | base = positions[-1, global_mm_start] + 1 |
| 334 | local_start = global_mm_start + 1 + mm_embeddings_seen |
| 335 | |
| 336 | local_end = local_start + mm_pos.shape[1] |
| 337 | positions[:, local_start:local_end] = mm_pos[0:3] + base |
| 338 | |
| 339 | # For Qwen3 VL (5-channel), use the maximum position reached across |
| 340 | # all tokens (both video and text) in all dimensions (t, h, w). |
| 341 | # For Qwen2.5 VL (4-channel), mm_pos[3, 0] is the max width. |
| 342 | if mm_pos.shape[0] == 5: |
| 343 | offset = mm_pos[0:3, :].max() + base + 1 |
| 344 | else: |
| 345 | offset = mm_pos[3, 0] + base |
| 346 | |
| 347 | text_pos_sum = torch.cumsum(text_mask[local_end:].long(), dim=0) |
| 348 | |
| 349 | positions[:, local_end:N] = text_pos_sum + offset - 1 |
| 350 | |
| 351 | # Include distance to the next vision start token |
| 352 | num_computed_tokens += mm_pos.shape[1] |
| 353 | |
| 354 | mrope_positions_delta = (positions.max() + 1 - N).item() |
| 355 | return positions, mrope_positions_delta |