| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any |
| 4 | |
| 5 | import numpy as np |
| 6 | |
| 7 | from vllm.lora.layers import LoRAMapping, LoRAMappingType |
| 8 | from vllm.lora.worker_manager import WorkerLoRAManager |
| 9 | from vllm.v1.worker.gpu.lora_utils import LoraState |
| 10 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 11 | |
| 12 | |
| 13 | def set_active_mm_loras( |
| 14 | model: Any, |
| 15 | lora_manager: WorkerLoRAManager, |
| 16 | encoder_cache: EncoderCache | None, |
| 17 | req_id_to_index: dict[str, int], |
| 18 | lora_state: LoraState, |
| 19 | scheduled_encoder_inputs: dict[str, list[int]], |
| 20 | ) -> None: |
| 21 | if ( |
| 22 | not scheduled_encoder_inputs |
| 23 | or encoder_cache is None |
| 24 | or not lora_manager.supports_tower_connector_lora() |
| 25 | ): |
| 26 | return |
| 27 | |
| 28 | prompt_lora_mapping: list[int] = [] |
| 29 | token_lora_mapping: list[int] = [] |
| 30 | lora_requests = set() |
| 31 | encoder_token_counts: list[int] = [] |
| 32 | connector_token_counts: list[int] = [] |
| 33 | |
| 34 | # iterate through images |
| 35 | for req_id, encoder_input_ids in scheduled_encoder_inputs.items(): |
| 36 | req_idx = req_id_to_index.get(req_id) |
| 37 | if req_idx is None: |
| 38 | continue |
| 39 | |
| 40 | lora_id = int(lora_state.lora_ids[req_idx]) |
| 41 | mm_features = encoder_cache.mm_features[req_id] |
| 42 | |
| 43 | # iterate through visual tokens |
| 44 | for mm_input_id in encoder_input_ids: |
| 45 | pos_info = mm_features[mm_input_id].mm_position |
| 46 | |
| 47 | tower_tokens, connector_tokens = model.get_mm_lora_token_counts( |
| 48 | modality=mm_features[mm_input_id].modality, |
| 49 | mm_kwargs=mm_features[mm_input_id].data, |
| 50 | num_mm_embeds=pos_info.get_num_embeds(), |
| 51 | ) |
| 52 | |
| 53 | prompt_lora_mapping.append(lora_id) |
| 54 | token_lora_mapping.extend([lora_id] * tower_tokens) |
| 55 | encoder_token_counts.append(tower_tokens) |
| 56 | connector_token_counts.append(connector_tokens) |
| 57 | |
| 58 | if lora_id > 0: |
| 59 | lora_request = lora_state.lora_requests.get(req_id) |
| 60 | if lora_request is not None: |
| 61 | lora_requests.add(lora_request) |
| 62 | |
| 63 | if not prompt_lora_mapping: |
| 64 | return |
| 65 | |
| 66 | lora_manager.set_active_adapters( |
| 67 | lora_requests, |
| 68 | LoRAMapping( |
| 69 | tuple(token_lora_mapping), |
| 70 | tuple(prompt_lora_mapping), |
| 71 | is_prefill=True, |
| 72 | type=LoRAMappingType.TOWER, |
| 73 | ), |
| 74 | ) |
| 75 | |
| 76 | mm_mapping = model.get_mm_mapping() if hasattr(model, "get_mm_mapping") else None |
| 77 | if ( |
| 78 | mm_mapping is None |
| 79 | or not mm_mapping.connector |
| 80 | or not all(count is not None for count in connector_token_counts) |
| 81 | ): |
| 82 | return |
| 83 | |
| 84 | connector_token_mapping = np.repeat( |
| 85 | np.array(prompt_lora_mapping, dtype=np.int32), |
| 86 | np.array(connector_token_counts, dtype=np.int32), |
| 87 | ) |
| 88 | lora_manager.set_active_adapters( |
| 89 | lora_requests, |
| 90 | LoRAMapping( |
| 91 | index_mapping=tuple(connector_token_mapping.tolist()), |
| 92 | prompt_mapping=tuple(prompt_lora_mapping), |
| 93 | is_prefill=True, |
| 94 | type=LoRAMappingType.CONNECTOR, |
| 95 | ), |
| 96 | ) |