A high-throughput and memory-efficient inference and serving engine for LLMs
[Model Runner V2] Fix mrv2 mm lora issue (#44450)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 8, 2026 at 14:30 UTC
2c27c294c0a6f3c2e32a88d1bb460b0cd4b8e301
3 files changed
+183
tests/v1/worker/test_gpu_model_runner.py
+79
@@ -24,8 +24,11 @@ from vllm.distributed.parallel_state import (
24
initialize_model_parallel,
25
)
26
from vllm.distributed.weight_transfer.base import SparseWeightPatch
27
+from vllm.lora.layers import LoRAMappingType
28
+from vllm.lora.request import LoRARequest
29
from vllm.model_executor.layers.attention import Attention
30
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
31
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
32
from vllm.platforms import current_platform
33
from vllm.sampling_params import SamplingParams
34
from vllm.utils.mem_constants import GiB_bytes
@@ -44,6 +47,9 @@ from vllm.v1.kv_cache_interface import (
47
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
48
from vllm.v1.sample.metadata import SamplingMetadata
49
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
50
+from vllm.v1.worker.gpu.lora_utils import LoraState
51
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
52
+from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
53
from vllm.v1.worker.gpu_input_batch import InputBatch
54
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
55
from vllm.v1.worker.utils import select_common_block_size
@@ -314,6 +320,79 @@ def test_select_common_block_size_no_valid_option():
320
select_common_block_size(48, [backend_a, backend_b])
321
322
323
+def test_set_active_mm_loras_builds_tower_and_connector_mappings():
324
+ model = Mock()
325
+ model.get_num_mm_encoder_tokens.side_effect = lambda num_embeds: num_embeds + 1
326
+ model.get_mm_mapping.return_value = SimpleNamespace(connector=True)
327
+ model.get_num_mm_connector_tokens.side_effect = lambda num_tokens: num_tokens + 10
328
+
329
+ lora_manager = Mock()
330
+ lora_manager.supports_tower_connector_lora.return_value = True
331
+
332
+ encoder_cache = EncoderCache()
333
+ encoder_cache.mm_features["req-with-lora"] = [
334
+ MultiModalFeatureSpec(
335
+ data=None,
336
+ modality="image",
337
+ identifier="img-0",
338
+ mm_position=PlaceholderRange(offset=0, length=2),
339
+ ),
340
+ MultiModalFeatureSpec(
341
+ data=None,
342
+ modality="image",
343
+ identifier="img-1",
344
+ mm_position=PlaceholderRange(offset=2, length=3),
345
+ ),
346
+ ]
347
+ encoder_cache.mm_features["req-no-lora"] = [
348
+ MultiModalFeatureSpec(
349
+ data=None,
350
+ modality="image",
351
+ identifier="img-2",
352
+ mm_position=PlaceholderRange(offset=0, length=1),
353
+ )
354
+ ]
355
+
356
+ lora_state = LoraState(max_num_reqs=4)
357
+ lora_request = LoRARequest("vision-lora", 7, "/tmp/vision-lora")
358
+ lora_state.add_request("req-with-lora", 0, lora_request)
359
+ lora_state.add_request("req-no-lora", 1, None)
360
+
361
+ set_active_mm_loras(
362
+ model=model,
363
+ lora_manager=lora_manager,
364
+ encoder_cache=encoder_cache,
365
+ req_id_to_index={
366
+ "req-with-lora": 0,
367
+ "req-no-lora": 1,
368
+ },
369
+ lora_state=lora_state,
370
+ scheduled_encoder_inputs={
371
+ "req-with-lora": [1, 0],
372
+ "req-no-lora": [0],
373
+ "missing-req": [0],
374
+ },
375
+ )
376
+
377
+ assert lora_manager.set_active_adapters.call_count == 2
378
+
379
+ tower_requests, tower_mapping = lora_manager.set_active_adapters.call_args_list[
380
+ 0
381
+ ].args
382
+ assert tower_requests == {lora_request}
383
+ assert tower_mapping.type is LoRAMappingType.TOWER
384
+ assert tower_mapping.prompt_mapping == (7, 7, 0)
385
+ assert tower_mapping.index_mapping == (7, 7, 7, 7, 7, 7, 7, 0, 0)
386
+
387
+ connector_requests, connector_mapping = (
388
+ lora_manager.set_active_adapters.call_args_list[1].args
389
+ )
390
+ assert connector_requests == {lora_request}
391
+ assert connector_mapping.type is LoRAMappingType.CONNECTOR
392
+ assert connector_mapping.prompt_mapping == (7, 7, 0)
393
+ assert connector_mapping.index_mapping == ((7,) * 14 + (7,) * 13 + (0,) * 12)
394
+
395
+
396
def test_update_states_new_request(model_runner, dist_init):
397
req_id = "req_0"
398
vllm/v1/worker/gpu/mm/lora.py
new
+94
@@ -0,0 +1,94 @@
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
+
33
+ # iterate through images
34
+ for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
35
+ req_idx = req_id_to_index.get(req_id)
36
+ if req_idx is None:
37
+ continue
38
+
39
+ lora_id = int(lora_state.lora_ids[req_idx])
40
+ mm_features = encoder_cache.mm_features[req_id]
41
+
42
+ # iterate through visual tokens
43
+ for mm_input_id in encoder_input_ids:
44
+ pos_info = mm_features[mm_input_id].mm_position
45
+ num_tokens = model.get_num_mm_encoder_tokens(pos_info.get_num_embeds())
46
+ prompt_lora_mapping.append(lora_id)
47
+ token_lora_mapping.extend([lora_id] * num_tokens)
48
+ encoder_token_counts.append(num_tokens)
49
+
50
+ if lora_id > 0:
51
+ lora_request = lora_state.lora_requests.get(req_id)
52
+ if lora_request is not None:
53
+ lora_requests.add(lora_request)
54
+
55
+ if not prompt_lora_mapping:
56
+ return
57
+
58
+ lora_manager.set_active_adapters(
59
+ lora_requests,
60
+ LoRAMapping(
61
+ tuple(token_lora_mapping),
62
+ tuple(prompt_lora_mapping),
63
+ is_prefill=True,
64
+ type=LoRAMappingType.TOWER,
65
+ ),
66
+ )
67
+
68
+ mm_mapping = model.get_mm_mapping() if hasattr(model, "get_mm_mapping") else None
69
+ if (
70
+ mm_mapping is None
71
+ or not mm_mapping.connector
72
+ or not hasattr(model, "get_num_mm_connector_tokens")
73
+ ):
74
+ return
75
+
76
+ connector_token_mapping = np.repeat(
77
+ np.array(prompt_lora_mapping, dtype=np.int32),
78
+ np.array(
79
+ [
80
+ model.get_num_mm_connector_tokens(num_tokens)
81
+ for num_tokens in encoder_token_counts
82
+ ],
83
+ dtype=np.int32,
84
+ ),
85
+ )
86
+ lora_manager.set_active_adapters(
87
+ lora_requests,
88
+ LoRAMapping(
89
+ index_mapping=tuple(connector_token_mapping.tolist()),
90
+ prompt_mapping=tuple(prompt_lora_mapping),
91
+ is_prefill=True,
92
+ type=LoRAMappingType.CONNECTOR,
93
+ ),
94
+ )
vllm/v1/worker/gpu/model_runner.py
+10
@@ -91,6 +91,7 @@ from vllm.v1.worker.gpu.kv_connector import (
91
)
92
from vllm.v1.worker.gpu.lora_utils import LoraState
93
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
94
+from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
95
from vllm.v1.worker.gpu.model_states import init_model_state
96
from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
97
from vllm.v1.worker.gpu.pp_utils import PPHandler
@@ -1207,6 +1208,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1208
# Only first PP rank prepares multimodal embeddings.
1209
# NOTE(woosuk): We must call get_mm_embeddings even during dummy runs
1210
# to obtain inputs_embeds, because the compiled model expects this input.
1211
+ if self.lora_config is not None:
1212
+ set_active_mm_loras(
1213
+ model=self.model,
1214
+ lora_manager=self.lora_manager,
1215
+ encoder_cache=self.encoder_cache,
1216
+ req_id_to_index=self.req_states.req_id_to_index,
1217
+ lora_state=self.lora_state,
1218
+ scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,
1219
+ )
1220
inputs_embeds = self.model_state.get_mm_embeddings(
1221
scheduler_output.scheduled_encoder_inputs, input_batch
1222
)