| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | import torch.nn as nn |
| 5 | |
| 6 | from vllm.config import ModelConfig |
| 7 | from vllm.model_executor.models.interfaces import supports_multimodal_pruning |
| 8 | from vllm.multimodal.utils import ( |
| 9 | copy_mm_embedding_modality, |
| 10 | get_mm_features_in_window, |
| 11 | ) |
| 12 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 13 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 14 | from vllm.v1.worker.gpu.mm.rope import RopeState |
| 15 | from vllm.v1.worker.gpu.states import RequestState |
| 16 | |
| 17 | |
| 18 | class MultiModalPruner: |
| 19 | """Recomputes M-RoPE positions for multimodal models that prune embeddings |
| 20 | (e.g. Qwen2.5-VL / Qwen3-VL / Nemotron-Nano-VL Efficient Video Sampling). |
| 21 | |
| 22 | Pruning models append their mrope-position channels to the (variable-count) |
| 23 | media embeddings from `embed_multimodal`. Those channels must be split off and |
| 24 | used to recompute mrope positions before the embeddings are merged. |
| 25 | """ |
| 26 | |
| 27 | def __init__( |
| 28 | self, |
| 29 | model: nn.Module, |
| 30 | rope_state: RopeState, |
| 31 | encoder_cache: EncoderCache, |
| 32 | inputs_embeds_size: int, |
| 33 | ) -> None: |
| 34 | self.model = model |
| 35 | self.rope_state = rope_state |
| 36 | self.encoder_cache = encoder_cache |
| 37 | # The cleaned embedding width: pruning models append their mrope-position |
| 38 | # channels as trailing columns, so embeds[:, :inputs_embeds_size] strips them. |
| 39 | self.inputs_embeds_size = inputs_embeds_size |
| 40 | |
| 41 | def strip(self, mm_embeds: list[torch.Tensor]) -> list[torch.Tensor]: |
| 42 | """Draft forward: strip the appended position channels only. |
| 43 | |
| 44 | Stripping is per-embedding, so no per-request segmentation is needed. The |
| 45 | speculator reuses the target's already-recomputed positions, hence there is |
| 46 | no position write-back here. |
| 47 | """ |
| 48 | stripped: list[torch.Tensor] = [] |
| 49 | for mm in mm_embeds: |
| 50 | out = mm[:, : self.inputs_embeds_size] |
| 51 | stripped.append(copy_mm_embedding_modality(mm, out)) |
| 52 | return stripped |
| 53 | |
| 54 | def recompute( |
| 55 | self, |
| 56 | mm_embeds: list[torch.Tensor], |
| 57 | input_batch: InputBatch, |
| 58 | req_states: RequestState, |
| 59 | ) -> list[torch.Tensor]: |
| 60 | """Target forward: split the appended mrope-position channels off each |
| 61 | request's media embeddings, recompute the corrected mrope positions, and |
| 62 | stage them back into RopeState. Returns the cleaned, flattened embeddings. |
| 63 | """ |
| 64 | cleaned: list[torch.Tensor] = [] |
| 65 | pos = 0 |
| 66 | req_idx_list = input_batch.idx_mapping_np.tolist() |
| 67 | prefill_lens_list = input_batch.prefill_len_np.tolist() |
| 68 | num_computed_list = input_batch.num_computed_prefill_tokens_np.tolist() |
| 69 | num_scheduled_list = input_batch.num_scheduled_tokens.tolist() |
| 70 | for batch_idx, req_id in enumerate(input_batch.req_ids): |
| 71 | num_computed = num_computed_list[batch_idx] |
| 72 | query_end = num_computed + num_scheduled_list[batch_idx] |
| 73 | num_req_embeds = self._num_window_embeds(req_id, num_computed, query_end) |
| 74 | if num_req_embeds == 0: |
| 75 | continue |
| 76 | req_embeds = mm_embeds[pos : pos + num_req_embeds] |
| 77 | pos += num_req_embeds |
| 78 | |
| 79 | req_idx = req_idx_list[batch_idx] |
| 80 | prefill_len = prefill_lens_list[batch_idx] |
| 81 | input_ids = req_states.all_token_ids.gpu[req_idx, :prefill_len] |
| 82 | mrope_positions = self.rope_state.read_prefill_positions( |
| 83 | req_idx, prefill_len |
| 84 | ).long() |
| 85 | req_cleaned, new_positions, delta = self.model.recompute_mrope_positions( |
| 86 | input_ids=input_ids, |
| 87 | multimodal_embeddings=req_embeds, |
| 88 | mrope_positions=mrope_positions, |
| 89 | num_computed_tokens=num_computed, |
| 90 | ) |
| 91 | self.rope_state.update_prefill_positions(req_idx, new_positions, delta) |
| 92 | cleaned.extend( |
| 93 | copy_mm_embedding_modality(src, dst) |
| 94 | for src, dst in zip(req_embeds, req_cleaned) |
| 95 | ) |
| 96 | |
| 97 | assert pos == len(mm_embeds) |
| 98 | return cleaned |
| 99 | |
| 100 | def _num_window_embeds(self, req_id: str, query_start: int, query_end: int) -> int: |
| 101 | """Count the media items contributing embeddings to [query_start, |
| 102 | query_end), mirroring EncoderRunner.gather_mm_embeddings' per-request |
| 103 | windowing so the flat mm_embeds list can be re-segmented per request. |
| 104 | |
| 105 | Note: This logic is intentionally duplicated here rather than being emitted |
| 106 | from gather_mm_embeddings, to keep the main path cleaner, since this is a niche |
| 107 | feature. |
| 108 | """ |
| 109 | mm_features = self.encoder_cache.mm_features[req_id] |
| 110 | lo, hi = get_mm_features_in_window( |
| 111 | mm_features, start=query_start, end=query_end |
| 112 | ) |
| 113 | count = 0 |
| 114 | for mm_feature in mm_features[lo:hi]: |
| 115 | pos_info = mm_feature.mm_position |
| 116 | start_idx = max(query_start - pos_info.offset, 0) |
| 117 | end_idx = min(query_end - pos_info.offset, pos_info.length) |
| 118 | embeds_start, embeds_end = pos_info.get_embeds_indices_in_range( |
| 119 | start_idx, end_idx |
| 120 | ) |
| 121 | if embeds_start != embeds_end: |
| 122 | count += 1 |
| 123 | return count |
| 124 | |
| 125 | |
| 126 | def maybe_create_mm_pruner( |
| 127 | model_config: ModelConfig, |
| 128 | model: nn.Module, |
| 129 | rope_state: RopeState | None, |
| 130 | encoder_cache: EncoderCache | None, |
| 131 | ) -> MultiModalPruner | None: |
| 132 | """Create a MultiModalPruner if the model prunes embeddings and uses M-RoPE.""" |
| 133 | if ( |
| 134 | rope_state is None |
| 135 | or encoder_cache is None |
| 136 | or model_config.multimodal_config is None |
| 137 | or not model_config.multimodal_config.is_multimodal_pruning_enabled() |
| 138 | or not supports_multimodal_pruning(model) |
| 139 | ): |
| 140 | return None |
| 141 | |
| 142 | return MultiModalPruner( |
| 143 | model, rope_state, encoder_cache, model_config.get_inputs_embeds_size() |
| 144 | ) |