diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 9642bfd79f..80dd8ee306 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -24,8 +24,11 @@ from vllm.distributed.parallel_state import ( initialize_model_parallel, ) from vllm.distributed.weight_transfer.base import SparseWeightPatch +from vllm.lora.layers import LoRAMappingType +from vllm.lora.request import LoRARequest from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 +from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams from vllm.utils.mem_constants import GiB_bytes @@ -44,6 +47,9 @@ from vllm.v1.kv_cache_interface import ( from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.metadata import SpecDecodeMetadata +from vllm.v1.worker.gpu.lora_utils import LoraState +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras from vllm.v1.worker.gpu_input_batch import InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.utils import select_common_block_size @@ -314,6 +320,79 @@ def test_select_common_block_size_no_valid_option(): select_common_block_size(48, [backend_a, backend_b]) +def test_set_active_mm_loras_builds_tower_and_connector_mappings(): + model = Mock() + model.get_num_mm_encoder_tokens.side_effect = lambda num_embeds: num_embeds + 1 + model.get_mm_mapping.return_value = SimpleNamespace(connector=True) + model.get_num_mm_connector_tokens.side_effect = lambda num_tokens: num_tokens + 10 + + lora_manager = Mock() + lora_manager.supports_tower_connector_lora.return_value = True + + encoder_cache = EncoderCache() + encoder_cache.mm_features["req-with-lora"] = [ + MultiModalFeatureSpec( + data=None, + modality="image", + identifier="img-0", + mm_position=PlaceholderRange(offset=0, length=2), + ), + MultiModalFeatureSpec( + data=None, + modality="image", + identifier="img-1", + mm_position=PlaceholderRange(offset=2, length=3), + ), + ] + encoder_cache.mm_features["req-no-lora"] = [ + MultiModalFeatureSpec( + data=None, + modality="image", + identifier="img-2", + mm_position=PlaceholderRange(offset=0, length=1), + ) + ] + + lora_state = LoraState(max_num_reqs=4) + lora_request = LoRARequest("vision-lora", 7, "/tmp/vision-lora") + lora_state.add_request("req-with-lora", 0, lora_request) + lora_state.add_request("req-no-lora", 1, None) + + set_active_mm_loras( + model=model, + lora_manager=lora_manager, + encoder_cache=encoder_cache, + req_id_to_index={ + "req-with-lora": 0, + "req-no-lora": 1, + }, + lora_state=lora_state, + scheduled_encoder_inputs={ + "req-with-lora": [1, 0], + "req-no-lora": [0], + "missing-req": [0], + }, + ) + + assert lora_manager.set_active_adapters.call_count == 2 + + tower_requests, tower_mapping = lora_manager.set_active_adapters.call_args_list[ + 0 + ].args + assert tower_requests == {lora_request} + assert tower_mapping.type is LoRAMappingType.TOWER + assert tower_mapping.prompt_mapping == (7, 7, 0) + assert tower_mapping.index_mapping == (7, 7, 7, 7, 7, 7, 7, 0, 0) + + connector_requests, connector_mapping = ( + lora_manager.set_active_adapters.call_args_list[1].args + ) + assert connector_requests == {lora_request} + assert connector_mapping.type is LoRAMappingType.CONNECTOR + assert connector_mapping.prompt_mapping == (7, 7, 0) + assert connector_mapping.index_mapping == ((7,) * 14 + (7,) * 13 + (0,) * 12) + + def test_update_states_new_request(model_runner, dist_init): req_id = "req_0" diff --git a/vllm/v1/worker/gpu/mm/lora.py b/vllm/v1/worker/gpu/mm/lora.py new file mode 100644 index 0000000000..492914a74b --- /dev/null +++ b/vllm/v1/worker/gpu/mm/lora.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import numpy as np + +from vllm.lora.layers import LoRAMapping, LoRAMappingType +from vllm.lora.worker_manager import WorkerLoRAManager +from vllm.v1.worker.gpu.lora_utils import LoraState +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache + + +def set_active_mm_loras( + model: Any, + lora_manager: WorkerLoRAManager, + encoder_cache: EncoderCache | None, + req_id_to_index: dict[str, int], + lora_state: LoraState, + scheduled_encoder_inputs: dict[str, list[int]], +) -> None: + if ( + not scheduled_encoder_inputs + or encoder_cache is None + or not lora_manager.supports_tower_connector_lora() + ): + return + + prompt_lora_mapping: list[int] = [] + token_lora_mapping: list[int] = [] + lora_requests = set() + encoder_token_counts: list[int] = [] + + # iterate through images + for req_id, encoder_input_ids in scheduled_encoder_inputs.items(): + req_idx = req_id_to_index.get(req_id) + if req_idx is None: + continue + + lora_id = int(lora_state.lora_ids[req_idx]) + mm_features = encoder_cache.mm_features[req_id] + + # iterate through visual tokens + for mm_input_id in encoder_input_ids: + pos_info = mm_features[mm_input_id].mm_position + num_tokens = model.get_num_mm_encoder_tokens(pos_info.get_num_embeds()) + prompt_lora_mapping.append(lora_id) + token_lora_mapping.extend([lora_id] * num_tokens) + encoder_token_counts.append(num_tokens) + + if lora_id > 0: + lora_request = lora_state.lora_requests.get(req_id) + if lora_request is not None: + lora_requests.add(lora_request) + + if not prompt_lora_mapping: + return + + lora_manager.set_active_adapters( + lora_requests, + LoRAMapping( + tuple(token_lora_mapping), + tuple(prompt_lora_mapping), + is_prefill=True, + type=LoRAMappingType.TOWER, + ), + ) + + mm_mapping = model.get_mm_mapping() if hasattr(model, "get_mm_mapping") else None + if ( + mm_mapping is None + or not mm_mapping.connector + or not hasattr(model, "get_num_mm_connector_tokens") + ): + return + + connector_token_mapping = np.repeat( + np.array(prompt_lora_mapping, dtype=np.int32), + np.array( + [ + model.get_num_mm_connector_tokens(num_tokens) + for num_tokens in encoder_token_counts + ], + dtype=np.int32, + ), + ) + lora_manager.set_active_adapters( + lora_requests, + LoRAMapping( + index_mapping=tuple(connector_token_mapping.tolist()), + prompt_mapping=tuple(prompt_lora_mapping), + is_prefill=True, + type=LoRAMappingType.CONNECTOR, + ), + ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 367147b0b4..50b78474bb 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -91,6 +91,7 @@ from vllm.v1.worker.gpu.kv_connector import ( ) from vllm.v1.worker.gpu.lora_utils import LoraState from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras from vllm.v1.worker.gpu.model_states import init_model_state from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner from vllm.v1.worker.gpu.pp_utils import PPHandler @@ -1207,6 +1208,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Only first PP rank prepares multimodal embeddings. # NOTE(woosuk): We must call get_mm_embeddings even during dummy runs # to obtain inputs_embeds, because the compiled model expects this input. + if self.lora_config is not None: + set_active_mm_loras( + model=self.model, + lora_manager=self.lora_manager, + encoder_cache=self.encoder_cache, + req_id_to_index=self.req_states.req_id_to_index, + lora_state=self.lora_state, + scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs, + ) inputs_embeds = self.model_state.get_mm_embeddings( scheduler_output.scheduled_encoder_inputs, input_batch )