[Model Runner V2] Fix mrv2 mm lora issue (#44450)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

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"
 

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,
+        ),
+    )

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
             )