@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ce1eafd1a5
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 22:45 UTC

[Core] Initialize LoRA support for tower and connector in multi-modal models (#26674)

Signed-off-by: bk-201 <joy25810@foxmail.com> Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: prashanth058 <prashanth.dannamaneni@uipath.com> Co-authored-by: bk-201 <joy25810@foxmail.com> Co-authored-by: prashanth058 <prashanth.dannamaneni@uipath.com> Co-authored-by: Anexdeus <5142168@mail.ru>

Jee Jee Li committed Dec 26, 2025 at 20:48 UTC ce1eafd1a56bd5a4eb1ee743933a5e96e63d2116
20 files changed +634 -79
docs/features/lora.md
+4
@@ -275,6 +275,10 @@ The new format of `--lora-modules` is mainly to support the display of parent mo
275 }
276 ```
277
278 +## LoRA Support for Tower and Connector of Multi-Modal Model
279 +
280 +Currently, vLLM experimentally supports LoRA for the Tower and Connector components of multi-modal models. To enable this feature, you need to implement the corresponding token helper functions for the tower and connector. For more details on the rationale behind this approach, please refer to [PR 26674](https://github.com/vllm-project/vllm/pull/26674). We welcome contributions to extend LoRA support to additional models' tower and connector.
281 +
282 ## Default LoRA Models For Multimodal Models
283
284 Some models, e.g., [Granite Speech](https://huggingface.co/ibm-granite/granite-speech-3.3-8b) and [Phi-4-multimodal-instruct](https://huggingface.co/microsoft/Phi-4-multimodal-instruct) multimodal, contain LoRA adapter(s) that are expected to always be applied when a given modality is present. This can be a bit tedious to manage with the above approaches, as it requires the user to send the `LoRARequest` (offline) or to filter requests between the base model and LoRA model (server) depending on the content of the request's multimodal data.
tests/lora/conftest.py
+25
@@ -208,6 +208,31 @@ def qwen25vl_lora_files():
208 return snapshot_download(repo_id="jeeejeee/qwen25-vl-lora-pokemon")
209
210
211 +@pytest.fixture(scope="session")
212 +def qwen2vl_language_lora_files():
213 + return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-language")
214 +
215 +
216 +@pytest.fixture(scope="session")
217 +def qwen2vl_vision_tower_connector_lora_files():
218 + return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower-connector")
219 +
220 +
221 +@pytest.fixture(scope="session")
222 +def qwen2vl_vision_tower_lora_files():
223 + return snapshot_download(repo_id="prashanth058/qwen2vl-flickr-lora-tower")
224 +
225 +
226 +@pytest.fixture(scope="session")
227 +def qwen25vl_vision_lora_files():
228 + return snapshot_download(repo_id="EpochEcho/qwen2.5-3b-vl-lora-vision-connector")
229 +
230 +
231 +@pytest.fixture(scope="session")
232 +def qwen3vl_vision_lora_files():
233 + return snapshot_download(repo_id="EpochEcho/qwen3-4b-vl-lora-vision-connector")
234 +
235 +
236 @pytest.fixture(scope="session")
237 def tinyllama_lora_files():
238 return snapshot_download(repo_id="jashing/tinyllama-colorist-lora")
tests/lora/test_lora_manager.py
+21 -7
@@ -18,6 +18,7 @@ from vllm.lora.layers import (
18 from vllm.lora.lora_model import LoRAModel
19 from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
20 from vllm.lora.model_manager import (
21 + DEFAULT_LANGUAGE_WRAPPER_KEY,
22 LoRAMapping,
23 LoRAModelManager,
24 LRUCacheLoRAModelManager,
@@ -183,9 +184,11 @@ def test_lora_model_manager(dist_init, dummy_model, device):
184 assert manager.activate_adapter(2)
185 assert manager.lora_index_to_id[0] == 3
186 assert manager.lora_index_to_id[1] == 2
186 -
187 assert manager.device == device
188 - assert manager.punica_wrapper.device == device
188 + assert (
189 + manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
190 + == device
191 + )
192 assert hasattr(manager, "supported_lora_modules")
193 assert sorted(manager.supported_lora_modules) == [
194 "dense1",
@@ -278,8 +281,10 @@ def test_lora_lru_cache_model_manager(dist_init, dummy_model, device):
281 assert manager.remove_adapter(3)
282 with pytest.raises(ValueError):
283 assert manager.pin_adapter(3)
281 -
282 - assert manager.punica_wrapper.device == device
284 + assert (
285 + manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
286 + == device
287 + )
288 assert manager.device == device
289
290
@@ -402,7 +407,10 @@ def test_lru_lora_model_manager(dist_init, dummy_model, device):
407 assert manager.remove_oldest_adapter()
408
409 assert set(manager.list_adapters()) == {1}
405 - assert manager.punica_wrapper.device == device
410 + assert (
411 + manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
412 + == device
413 + )
414 assert manager.device == device
415
416
@@ -514,7 +522,10 @@ def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_pa
522 )
523
524 assert worker_adapter_manager.device == device
517 - assert worker_adapter_manager._adapter_manager.punica_wrapper.device == device
525 + punica_wrapper = worker_adapter_manager._adapter_manager.punica_wrapper_mapping.get(
526 + DEFAULT_LANGUAGE_WRAPPER_KEY
527 + )
528 + assert punica_wrapper.device == device
529
530
531 @pytest.mark.parametrize("device", DEVICES)
@@ -618,7 +629,10 @@ def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path
629 )
630
631 assert worker_adapter_manager.device == device
621 - assert worker_adapter_manager._adapter_manager.punica_wrapper.device == device
632 + punica_wrapper = worker_adapter_manager._adapter_manager.punica_wrapper_mapping.get(
633 + DEFAULT_LANGUAGE_WRAPPER_KEY
634 + )
635 + assert punica_wrapper.device == device
636
637
638 @pytest.mark.parametrize("device", DEVICES)
tests/lora/test_qwenvl.py renamed
+125 -4
@@ -14,9 +14,12 @@ class TestConfig:
14 lora_path: str
15 max_num_seqs: int = 2
16 max_loras: int = 2
17 - max_lora_rank: int = 16
18 - max_model_len: int = 4096
17 + max_lora_rank: int = 32
18 + enable_tower_connector_lora: bool = False
19 + max_model_len: int = 8192
20 + gpu_memory_utilization: float = 0.85
21 mm_processor_kwargs: dict[str, int] | None = None
22 + mm_processor_cache_gb: float = 4
23
24 def __post_init__(self):
25 if self.mm_processor_kwargs is None:
@@ -48,8 +51,11 @@ class Qwen2VLTester:
51 enable_lora=True,
52 max_loras=self.config.max_loras,
53 max_lora_rank=self.config.max_lora_rank,
54 + enable_tower_connector_lora=self.config.enable_tower_connector_lora,
55 trust_remote_code=True,
56 + gpu_memory_utilization=self.config.gpu_memory_utilization,
57 mm_processor_kwargs=self.config.mm_processor_kwargs,
58 + mm_processor_cache_gb=self.config.mm_processor_cache_gb,
59 max_model_len=self.config.max_model_len,
60 )
61
@@ -58,6 +64,7 @@ class Qwen2VLTester:
64 images: list[ImageAsset],
65 expected_outputs: list[str],
66 lora_id: int | None = None,
67 + lora_name: str | None = None,
68 temperature: float = 0,
69 max_tokens: int = 5,
70 ):
@@ -73,10 +80,11 @@ class Qwen2VLTester:
80 for asset in images
81 ]
82
76 - lora_request = LoRARequest(str(lora_id), lora_id, self.config.lora_path)
83 + lora_request = LoRARequest(
84 + lora_name if lora_name else str(lora_id), lora_id, self.config.lora_path
85 + )
86 outputs = self.llm.generate(inputs, sampling_params, lora_request=lora_request)
87 generated_texts = [output.outputs[0].text.strip() for output in outputs]
79 -
88 # Validate outputs
89 for generated, expected in zip(generated_texts, expected_outputs):
90 assert expected.startswith(generated), (
@@ -127,6 +135,22 @@ EXPECTED_OUTPUTS = [
135 "A majestic skyscraper stands tall, partially obscured by a vibrant canopy of cherry blossoms, against a clear blue sky.", # noqa: E501
136 ]
137
138 +EXPECTED_OUTPUTS_LANGUAGE = [
139 + "A stop sign is shown in an Asian city, with buildings and a car in the "
140 + "background.",
141 + "The Tokyo Skytree can be seen behind the pink blossoms of the cherry trees.",
142 +]
143 +
144 +EXPECTED_OUTPUTS_VISION = [
145 + "A stop sign in front of oriental buildings.",
146 + "A tree with pink flowers in front of it and a blue sky behind the flowers.",
147 +]
148 +
149 +EXPECTED_OUTPUTS_VISION_NO_CONNECTOR = [
150 + "A stop sign is located on the street of a Chinese neighborhood.",
151 + "A closeup shot of the Tokyo Skytree with pink flowers in the foreground.",
152 +]
153 +
154 # NOTE - beam search .text contains the whole text
155 EXPECTED_BEAM_SEARCH_OUTPUTS = [
156 [
@@ -137,6 +161,7 @@ EXPECTED_BEAM_SEARCH_OUTPUTS = [
161
162 QWEN2VL_MODEL_PATH = "Qwen/Qwen2-VL-2B-Instruct"
163 QWEN25VL_MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct"
164 +QWEN3VL_MODEL_PATH = "Qwen/Qwen3-VL-4B-Instruct"
165
166
167 def test_qwen2vl_lora(qwen2vl_lora_files):
@@ -175,3 +200,99 @@ def test_qwen25vl_lora(qwen25vl_lora_files):
200 # Test with different LoRA IDs
201 for lora_id in [1, 2]:
202 tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
203 +
204 +
205 +def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files):
206 + config = TestConfig(
207 + model_path=QWEN25VL_MODEL_PATH,
208 + lora_path=qwen25vl_vision_lora_files,
209 + # Currently, tower_connector_lora is incompatible with
210 + # the multi-modal processor cache.
211 + # TODO: Remove this restriction
212 + mm_processor_cache_gb=0,
213 + enable_tower_connector_lora=True,
214 + )
215 + tester = Qwen2VLTester(config)
216 + for lora_id in [1, 2]:
217 + tester.run_test(
218 + TEST_IMAGES,
219 + expected_outputs=EXPECTED_OUTPUTS,
220 + lora_id=lora_id,
221 + )
222 +
223 +
224 +def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files):
225 + config = TestConfig(
226 + model_path=QWEN3VL_MODEL_PATH,
227 + lora_path=qwen3vl_vision_lora_files,
228 + # Currently, tower_connector_lora is incompatible with
229 + # the multi-modal processor cache.
230 + # TODO: Remove this restriction
231 + mm_processor_cache_gb=0,
232 + enable_tower_connector_lora=True,
233 + )
234 + tester = Qwen2VLTester(config)
235 + for lora_id in [1, 2]:
236 + tester.run_test(
237 + TEST_IMAGES,
238 + expected_outputs=EXPECTED_OUTPUTS,
239 + lora_id=lora_id,
240 + )
241 +
242 +
243 +def test_qwen2vl_multiple_lora_types(
244 + qwen2vl_language_lora_files,
245 + qwen2vl_vision_tower_connector_lora_files,
246 + qwen2vl_vision_tower_lora_files,
247 +):
248 + """
249 + Test multiple LoRA adapter types (language, vision tower + connector,
250 + vision tower only) using the same LLM instance to verify mm_encoder_cache
251 + behavior with different LoRA requests.
252 +
253 + By reusing the same LLM instance across different LoRA requests, we ensure that
254 + the multimodal encoder cache correctly manages state transitions between
255 + language-only and vision-enabled LoRA adapters.
256 + """
257 + config = TestConfig(
258 + model_path=QWEN2VL_MODEL_PATH,
259 + # We'll override the lora_path for each specific test, but need to provide
260 + # an initial path for initialization
261 + lora_path=qwen2vl_language_lora_files,
262 + # Currently, tower_connector_lora is incompatible with
263 + # the multi-modal processor cache.
264 + # TODO: Remove this restriction
265 + mm_processor_cache_gb=0,
266 + enable_tower_connector_lora=True,
267 + )
268 + tester = Qwen2VLTester(config)
269 +
270 + # Test 1: Language-only LoRA adapter
271 + tester.config.lora_path = qwen2vl_language_lora_files
272 + for lora_id in [1, 2]:
273 + tester.run_test(
274 + TEST_IMAGES,
275 + expected_outputs=EXPECTED_OUTPUTS_LANGUAGE,
276 + lora_id=lora_id,
277 + lora_name="language_only",
278 + )
279 +
280 + # Test 2: Vision tower + connector LoRA adapter
281 + tester.config.lora_path = qwen2vl_vision_tower_connector_lora_files
282 + for lora_id in [3, 4]:
283 + tester.run_test(
284 + TEST_IMAGES,
285 + expected_outputs=EXPECTED_OUTPUTS_VISION,
286 + lora_id=lora_id,
287 + lora_name="vision_tower_connector",
288 + )
289 +
290 + # Test 3: Vision tower only LoRA adapter (no connector)
291 + tester.config.lora_path = qwen2vl_vision_tower_lora_files
292 + for lora_id in [5, 6]:
293 + tester.run_test(
294 + TEST_IMAGES,
295 + expected_outputs=EXPECTED_OUTPUTS_VISION_NO_CONNECTOR,
296 + lora_id=lora_id,
297 + lora_name="vision_tower",
298 + )
vllm/config/lora.py
+6
@@ -55,6 +55,11 @@ class LoRAConfig:
55 per prompt. When run in offline mode, the lora IDs for n modalities
56 will be automatically assigned to 1-n with the names of the modalities
57 in alphabetic order."""
58 + enable_tower_connector_lora: bool = False
59 + """If `True`, LoRA support for the tower (vision encoder) and connector
60 + of multimodal models will be enabled. This is an experimental feature and
61 + currently only supports some MM models such as the Qwen VL series. The default
62 + is False."""
63
64 def compute_hash(self) -> str:
65 """
@@ -73,6 +78,7 @@ class LoRAConfig:
78 factors.append(self.max_loras)
79 factors.append(self.fully_sharded_loras)
80 factors.append(self.lora_dtype)
81 + factors.append(self.enable_tower_connector_lora)
82
83 hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
84 return hash_str
vllm/engine/arg_utils.py
+19
@@ -483,6 +483,7 @@ class EngineArgs:
483 fully_sharded_loras: bool = LoRAConfig.fully_sharded_loras
484 max_cpu_loras: int | None = LoRAConfig.max_cpu_loras
485 lora_dtype: str | torch.dtype | None = LoRAConfig.lora_dtype
486 + enable_tower_connector_lora: bool = LoRAConfig.enable_tower_connector_lora
487
488 ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight
489 num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override
@@ -996,6 +997,10 @@ class EngineArgs:
997 "--lora-dtype",
998 **lora_kwargs["lora_dtype"],
999 )
1000 + lora_group.add_argument(
1001 + "--enable-tower-connector-lora",
1002 + **lora_kwargs["enable_tower_connector_lora"],
1003 + )
1004 lora_group.add_argument("--max-cpu-loras", **lora_kwargs["max_cpu_loras"])
1005 lora_group.add_argument(
1006 "--fully-sharded-loras", **lora_kwargs["fully_sharded_loras"]
@@ -1631,6 +1636,7 @@ class EngineArgs:
1636 default_mm_loras=self.default_mm_loras,
1637 fully_sharded_loras=self.fully_sharded_loras,
1638 lora_dtype=self.lora_dtype,
1639 + enable_tower_connector_lora=self.enable_tower_connector_lora,
1640 max_cpu_loras=self.max_cpu_loras
1641 if self.max_cpu_loras and self.max_cpu_loras > 0
1642 else None,
@@ -1639,6 +1645,19 @@ class EngineArgs:
1645 else None
1646 )
1647
1648 + if (
1649 + lora_config is not None
1650 + and lora_config.enable_tower_connector_lora
1651 + and self.mm_processor_cache_gb != 0
1652 + ):
1653 + raise ValueError(
1654 + "Currently, enable_tower_connector_lora is "
1655 + "incompatible with the multi-modal processor cache. "
1656 + "When enable_tower_connector_lora is set, "
1657 + "mm_processor_cache_gb must be 0, got %s",
1658 + self.mm_processor_cache_gb,
1659 + )
1660 +
1661 if (
1662 lora_config is not None
1663 and speculative_config is not None
vllm/lora/layers/__init__.py
+2 -1
@@ -18,7 +18,7 @@ from vllm.lora.layers.row_parallel_linear import (
18 RowParallelLinearWithLoRA,
19 RowParallelLinearWithShardedLoRA,
20 )
21 -from vllm.lora.layers.utils import LoRAMapping
21 +from vllm.lora.layers.utils import LoRAMapping, LoRAMappingType
22 from vllm.lora.layers.vocal_parallel_embedding import VocabParallelEmbeddingWithLoRA
23
24 __all__ = [
@@ -37,6 +37,7 @@ __all__ = [
37 "RowParallelLinearWithShardedLoRA",
38 "ReplicatedLinearWithLoRA",
39 "LoRAMapping",
40 + "LoRAMappingType",
41 "FusedMoEWithLoRA",
42 "FusedMoE3DWithLoRA",
43 ]
vllm/lora/layers/base_linear.py
+8 -1
@@ -122,7 +122,9 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
122 def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
123 output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
124
125 - # In Transformers modeling backend, x and output have extra batch dimension like
125 + original_shape = output.shape if output.ndim == 3 else None
126 +
127 + # In transformers backend, x and output have extra batch dimension like
128 # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim),
129 # therefore we need to flatten the batch dimensions.
130 if x.ndim == 3 and output.ndim == 3:
@@ -135,6 +137,11 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
137 if not current_platform.can_update_inplace():
138 output = lora_output
139
140 + # Reshape the flattened output back to its original shape,
141 + # as some MM encoders cannot handle flattened inputs.
142 + if original_shape is not None:
143 + output = output.reshape(original_shape)
144 +
145 return output
146
147 @property
vllm/lora/layers/utils.py
+8
@@ -2,16 +2,24 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 +from enum import Enum
6
7 import torch
8 import torch.nn as nn
9
10
11 +class LoRAMappingType(Enum):
12 + LANGUAGE = 1
13 + TOWER = 2
14 + CONNECTOR = 3
15 +
16 +
17 @dataclass
18 class LoRAMapping:
19 index_mapping: tuple[int, ...]
20 prompt_mapping: tuple[int, ...]
21 is_prefill: bool = False
22 + type: LoRAMappingType = LoRAMappingType.LANGUAGE
23
24 def __post_init__(self):
25 self.index_mapping = tuple(self.index_mapping)
vllm/lora/model_manager.py
+174 -41
@@ -9,12 +9,18 @@ import regex as re
9 import torch
10 from torch import nn
11
12 -from vllm.config.lora import LoRAConfig
12 +from vllm.config import VllmConfig
13 +from vllm.config.lora import LoRAConfig, ModelConfig
14 from vllm.logger import init_logger
14 -from vllm.lora.layers import BaseLayerWithLoRA, FusedMoE3DWithLoRA, LoRAMapping
15 +from vllm.lora.layers import (
16 + BaseLayerWithLoRA,
17 + FusedMoE3DWithLoRA,
18 + LoRAMapping,
19 + LoRAMappingType,
20 +)
21 from vllm.lora.lora_model import LoRAModel
22 from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
17 -from vllm.lora.punica_wrapper import get_punica_wrapper
23 +from vllm.lora.punica_wrapper import PunicaWrapperBase, get_punica_wrapper
24 from vllm.lora.utils import (
25 from_layer,
26 from_layer_logits_processor,
@@ -28,12 +34,15 @@ from vllm.model_executor.models import SupportsLoRA, supports_multimodal
34 from vllm.model_executor.models.interfaces import is_pooling_model
35 from vllm.model_executor.models.module_mapping import MultiModelKeys
36 from vllm.model_executor.models.utils import PPMissingLayer
37 +from vllm.multimodal import MULTIMODAL_REGISTRY
38 from vllm.utils.cache import LRUCache
39 from vllm.utils.platform_utils import is_pin_memory_available
40 +from vllm.v1.worker.utils import MultiModalBudget
41
42 logger = init_logger(__name__)
43
44 T = TypeVar("T")
45 +DEFAULT_LANGUAGE_WRAPPER_KEY = "language_model"
46
47
48 class AdapterLRUCache(LRUCache[int, T]):
@@ -58,6 +67,7 @@ class LoRAModelManager:
67 vocab_size: int,
68 lora_config: LoRAConfig,
69 device: torch.device,
70 + vllm_config: VllmConfig | None = None,
71 ):
72 """Create a LoRAModelManager and adapter for a given model.
73
@@ -71,6 +81,11 @@ class LoRAModelManager:
81 lora_config: the LoRA configuration.
82 """
83 self.model: SupportsLoRA = model
84 + self.supported_lora_modules = get_supported_lora_modules(self.model)
85 + assert self.supported_lora_modules, (
86 + f"No supported LoRA modules found in {self.model.__class__.__name__}."
87 + )
88 +
89 self._registered_adapters: dict[int, LoRAModel] = {}
90 # Dict instead of a set for compatibility with LRUCache.
91 self._active_adapters: dict[int, None] = {}
@@ -82,35 +97,121 @@ class LoRAModelManager:
97 self.max_num_batched_tokens = math.ceil(max_num_batched_tokens / 8) * 8
98 self.lora_index_to_id: list[int | None] = [None] * self.lora_slots
99 self.vocab_size = vocab_size
85 - self.punica_wrapper = get_punica_wrapper(
86 - max_num_batched_tokens,
87 - max_batches=self.max_num_seqs,
88 - device=self.device,
89 - max_loras=self.lora_config.max_loras,
90 - )
91 -
92 - self.supported_lora_modules = get_supported_lora_modules(self.model)
93 - assert self.supported_lora_modules, "No supported LoRA modules found in"
94 - f" {self.model.__class__.__name__}."
95 -
100 self.packed_modules_mapping = process_packed_modules_mapping(self.model)
97 - # Used to indicate whether the model is a multimodal model
98 - self.supports_mm: bool = (
99 - supports_multimodal(self.model)
100 - # In case the model only supports LoRA for
101 - # text modules (e.g. ChatGLM)
102 - and hasattr(self.model, "get_mm_mapping")
103 - )
101 +
102 self.is_pooling_model = is_pooling_model(self.model)
103 self.packed_modules: dict[str, list[str]] = {}
104 self.modules: dict[str, BaseLayerWithLoRA] = {}
105 # Dict instead of a set for compatibility with LRUCache.
106 self._last_mapping: LoRAMapping | None = None
107 self._is_3d_moe_model = is_moe_model(self.model) and self.model.is_3d_moe_weight
108 + self._init_punica_wrapper(max_num_batched_tokens, vllm_config)
109 self._create_lora_modules()
110
111 self.model.lora_manager = self
112
113 + def _init_punica_wrapper(
114 + self, max_num_batched_tokens: int, vllm_config: VllmConfig
115 + ) -> None:
116 + # Used to indicate whether the model is a multimodal model
117 + self.supports_mm: bool = (
118 + supports_multimodal(self.model)
119 + # In case the model only supports LoRA for
120 + # text modules (e.g. ChatGLM)
121 + and hasattr(self.model, "get_mm_mapping")
122 + )
123 + self.punica_wrapper_mapping: dict[str, PunicaWrapperBase] = {}
124 + if self.supports_mm:
125 + self._maybe_init_mm(vllm_config, max_num_batched_tokens)
126 + else:
127 + llm_punica_wrapper = get_punica_wrapper(
128 + max_num_batched_tokens,
129 + max_batches=self.max_num_seqs,
130 + device=self.device,
131 + max_loras=self.lora_config.max_loras,
132 + )
133 +
134 + self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY] = (
135 + llm_punica_wrapper
136 + )
137 +
138 + def _maybe_init_mm(self, vllm_config: VllmConfig, max_num_batched_tokens) -> None:
139 + self.supports_tower_connector_lora = False
140 + model_config: ModelConfig = vllm_config.model_config
141 + self.mm_mapping: MultiModelKeys = self.model.get_mm_mapping()
142 +
143 + # Only one language model can be included in the model.
144 + assert len(self.mm_mapping.language_model) == 1
145 +
146 + # Language model punica wrapper
147 + llm_punica_wrapper = get_punica_wrapper(
148 + max_num_batched_tokens,
149 + max_batches=self.max_num_seqs,
150 + device=self.device,
151 + max_loras=self.lora_config.max_loras,
152 + )
153 + lm_prefix = self.mm_mapping.language_model[0]
154 + self.punica_wrapper_mapping[lm_prefix] = llm_punica_wrapper
155 +
156 + if self.lora_config.enable_tower_connector_lora:
157 + self.mm_processor_info = MULTIMODAL_REGISTRY.create_processor(
158 + model_config
159 + ).info
160 + self.supports_tower_connector_lora = self.supports_mm and hasattr(
161 + self.model, "get_num_mm_encoder_tokens"
162 + )
163 + if not self.supports_tower_connector_lora:
164 + return
165 +
166 + logger.warning(
167 + "LoRA for the tower and connector of multimodal models is "
168 + "experimental and may contain bugs. Please report any related issues on "
169 + "GitHub if you encounter them."
170 + )
171 +
172 + mm_budget = MultiModalBudget(
173 + model_config,
174 + vllm_config.scheduler_config,
175 + MULTIMODAL_REGISTRY,
176 + )
177 + limit_per_prompt: int = max(
178 + self.mm_processor_info.get_allowed_mm_limits().values()
179 + )
180 + num_encoder_tokens = self.model.get_num_mm_encoder_tokens(
181 + mm_budget.get_encoder_budget()
182 + )
183 +
184 + # Tower wrappers
185 + tower_punica_wrapper = get_punica_wrapper(
186 + num_encoder_tokens,
187 + max_batches=self.max_num_seqs * limit_per_prompt,
188 + device=self.device,
189 + max_loras=self.lora_config.max_loras,
190 + )
191 + for prefix in self.mm_mapping.tower_model:
192 + self.punica_wrapper_mapping[prefix] = tower_punica_wrapper
193 +
194 + # Use wrapper for connector if present.
195 + if self.mm_mapping.connector:
196 + if hasattr(self.model, "get_num_mm_connector_tokens"):
197 + connector_tokens = self.model.get_num_mm_connector_tokens(
198 + num_encoder_tokens
199 + )
200 + connector_punica_wrapper = get_punica_wrapper(
201 + connector_tokens,
202 + max_batches=self.max_num_seqs * limit_per_prompt,
203 + device=self.device,
204 + max_loras=self.lora_config.max_loras,
205 + )
206 + for prefix in self.mm_mapping.connector:
207 + self.punica_wrapper_mapping[prefix] = connector_punica_wrapper
208 + else:
209 + logger.warning_once(
210 + "Connector LoRA support disabled: model does not implement "
211 + "get_num_mm_connector_tokens(). This method is required to "
212 + "determine the connector's token budget for LoRA operations."
213 + )
214 +
215 def __len__(self) -> int:
216 return len(self._registered_adapters)
217
@@ -237,8 +338,24 @@ class LoRAModelManager:
338 ) # type: ignore
339
340 def _set_adapter_mapping(self, mapping: LoRAMapping) -> None:
240 - # update lora states
241 - self.punica_wrapper.update_metadata(
341 + # Default to the main language model wrapper
342 + if not (self.supports_mm and self.supports_tower_connector_lora):
343 + target_prefix = (
344 + self.mm_mapping.language_model[0]
345 + if self.supports_mm
346 + else DEFAULT_LANGUAGE_WRAPPER_KEY
347 + )
348 + elif mapping.type == LoRAMappingType.TOWER and self.mm_mapping.tower_model:
349 + target_prefix = self.mm_mapping.tower_model[0]
350 + elif mapping.type == LoRAMappingType.CONNECTOR and self.mm_mapping.connector:
351 + target_prefix = self.mm_mapping.connector[0]
352 + else:
353 + target_prefix = self.mm_mapping.language_model[0]
354 +
355 + punica_wrapper = self._get_punica_wrapper(target_prefix)
356 + assert punica_wrapper is not None
357 +
358 + punica_wrapper.update_metadata(
359 mapping,
360 self.lora_index_to_id,
361 self.lora_slots + 1,
@@ -265,15 +382,17 @@ class LoRAModelManager:
382
383 if not self._match_target_modules(module_name):
384 continue
268 - # A temporary approach for multimodal models to support LoRA
269 - # TODO: Remove this restriction
270 - if self._filter_unsupported_mm_module(module_name):
385 +
386 + punica_wrapper = self._get_punica_wrapper(module_name)
387 + if punica_wrapper is None:
388 logger.warning(
272 - "Regarding multimodal models, vLLM currently only supports "
273 - "adding LoRA to language model, %s will be ignored.",
389 + "Regarding %s, vLLM currently only supports adding LoRA to"
390 + " language model, %s will be ignored.",
391 + self.model.__class__.__name__,
392 module_name,
393 )
394 continue
395 +
396 parts = module_name.split(".")[-1]
397 packed_moduled_lst = self.packed_modules_mapping.get(parts, [])
398 if isinstance(module, FusedMoE):
@@ -328,10 +447,10 @@ class LoRAModelManager:
447 if self.supports_mm and not isinstance(new_module, BaseLayerWithLoRA):
448 continue
449 self.register_module(module_name, new_module)
450 +
451 self._register_packed_modules(module_name)
452 # All lora layers share the same punica_wrapper based on reference.
333 - new_module.set_mapping(self.punica_wrapper)
334 - pass
453 + new_module.set_mapping(punica_wrapper)
454
455 def register_module(self, module_name: str, module: "BaseLayerWithLoRA"):
456 assert isinstance(module, BaseLayerWithLoRA), (
@@ -352,7 +471,7 @@ class LoRAModelManager:
471 if (
472 not self._match_target_modules(module_name)
473 or not isinstance(module, BaseLayerWithLoRA)
355 - or self._filter_unsupported_mm_module(module_name)
474 + or self._get_punica_wrapper(module_name) is None
475 ):
476 continue
477 parts = module_name.split(".")
@@ -441,17 +560,22 @@ class LoRAModelManager:
560 for target_module in self.supported_lora_modules
561 )
562
444 - def _filter_unsupported_mm_module(self, module_name: str) -> bool:
563 + def _get_punica_wrapper(self, module_name: str) -> PunicaWrapperBase | None:
564 """
446 - Regarding multimodal models, vLLM currently only supports adding LoRA to
447 - language model. LoRA for other modules, such as the vision tower, will
448 - be filtered out.
565 + Determine whether this module supports LoRA and which wrapper to use.
566 """
450 - if self.supports_mm:
451 - module_mapping: MultiModelKeys = self.model.get_mm_mapping()
452 - prefix_lst = module_mapping.connector + module_mapping.tower_model
453 - return any([module_name.startswith(prefix) for prefix in prefix_lst])
454 - return False
567 + # For language model (early return)
568 + if not self.supports_mm:
569 + return self.punica_wrapper_mapping[DEFAULT_LANGUAGE_WRAPPER_KEY]
570 +
571 + # For multimodal model
572 + # NOTE Sort by prefix length (descending) to match the longest prefix first
573 + # e.g., 'visual.merger' should match 'visual.merger' instead of 'visual.'
574 + for prefix in sorted(self.punica_wrapper_mapping.keys(), key=len, reverse=True):
575 + if module_name.startswith(prefix):
576 + return self.punica_wrapper_mapping[prefix]
577 +
578 + return None
579
580 def _register_packed_modules(self, module_full_name: str) -> None:
581 parts = module_full_name.split(".")
@@ -596,9 +720,16 @@ class LRUCacheLoRAModelManager(LoRAModelManager):
720 vocab_size: int,
721 lora_config: LoRAConfig,
722 device: torch.device,
723 + vllm_config: VllmConfig | None = None,
724 ):
725 super().__init__(
601 - model, max_num_seqs, max_num_batched_tokens, vocab_size, lora_config, device
726 + model,
727 + max_num_seqs,
728 + max_num_batched_tokens,
729 + vocab_size,
730 + lora_config,
731 + device,
732 + vllm_config,
733 )
734 self._registered_adapters: LoRALRUCache = LoRALRUCache(
735 self.capacity, self.deactivate_adapter
@@ -671,6 +802,7 @@ def create_lora_manager(
802 max_num_batched_tokens: int,
803 vocab_size: int,
804 lora_config: LoRAConfig,
805 + vllm_config: VllmConfig,
806 device: torch.device,
807 lora_manager_cls: type[LoRAModelManager] = LoRAModelManager,
808 **kwargs,
@@ -684,6 +816,7 @@ def create_lora_manager(
816 max_num_batched_tokens=max_num_batched_tokens,
817 vocab_size=vocab_size,
818 lora_config=lora_config,
819 + vllm_config=vllm_config,
820 device=device,
821 **kwargs,
822 )
vllm/lora/worker_manager.py
+10
@@ -69,6 +69,7 @@ class WorkerLoRAManager:
69 def create_lora_manager(
70 self,
71 model: torch.nn.Module,
72 + vllm_config: VllmConfig | None = None,
73 ) -> Any:
74 lora_manager = create_lora_manager(
75 model,
@@ -78,6 +79,7 @@ class WorkerLoRAManager:
79 lora_config=self.lora_config,
80 device=self.device,
81 lora_manager_cls=self._manager_cls,
82 + vllm_config=vllm_config,
83 )
84 self._adapter_manager = lora_manager
85 return lora_manager.model
@@ -161,6 +163,12 @@ class WorkerLoRAManager:
163 if mapping is not None:
164 self._adapter_manager.set_adapter_mapping(mapping)
165
166 + def supports_tower_connector_lora(self) -> bool:
167 + return (
168 + self._adapter_manager.supports_mm
169 + and self._adapter_manager.supports_tower_connector_lora
170 + )
171 +
172 def _apply_adapters(self, adapter_requests: set[Any]) -> None:
173 existing_adapters = self.list_adapters()
174 models_map = {
@@ -210,6 +218,7 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager):
218 def create_lora_manager(
219 self,
220 model: torch.nn.Module,
221 + vllm_config: VllmConfig | None = None,
222 ) -> Any:
223 lora_manager = create_lora_manager(
224 model,
@@ -219,6 +228,7 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager):
228 lora_config=self.lora_config,
229 device=self.device,
230 max_num_batched_tokens=self.max_num_batched_tokens,
231 + vllm_config=vllm_config,
232 )
233 self._adapter_manager = lora_manager
234 return lora_manager.model
vllm/model_executor/models/idefics3.py
+18
@@ -714,3 +714,21 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo
714 connector="model.connector",
715 tower_model="model.vision_model",
716 )
717 +
718 + def get_num_mm_encoder_tokens(
719 + self,
720 + num_image_tokens: int,
721 + ) -> int:
722 + hf_config = self.config
723 + scale_factor = hf_config.scale_factor
724 +
725 + return num_image_tokens * scale_factor**2
726 +
727 + def get_num_mm_connector_tokens(
728 + self,
729 + num_vision_tokens: int,
730 + ) -> int:
731 + hf_config = self.config
732 + scale_factor = hf_config.scale_factor
733 +
734 + return num_vision_tokens // scale_factor**2
vllm/model_executor/models/interfaces.py
+18
@@ -136,6 +136,24 @@ class SupportsMultiModal(Protocol):
136 """
137 ...
138
139 + def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
140 + """
141 + Implement this function to enable LoRA support
142 + for the tower module of the multi-modal model.
143 + Given the number of image tokens, output the number of
144 + multi-modal encoder tokens.
145 + """
146 + ...
147 +
148 + def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int:
149 + """
150 + Implement this function to enable LoRA support
151 + for the connector module of the multi-modal model.
152 + Given the number of vision tokens, output the number of
153 + multi-modal connector tokens.
154 + """
155 + ...
156 +
157 @classmethod
158 def get_language_model_spec(cls) -> tuple[nn.Module | None, str | None]:
159 """
vllm/model_executor/models/qwen2_5_vl.py
+20
@@ -1026,6 +1026,7 @@ class Qwen2_5_VLForConditionalGeneration(
1026 packed_modules_mapping = {
1027 "qkv_proj": ["q_proj", "k_proj", "v_proj"],
1028 "gate_up_proj": ["gate_proj", "up_proj"],
1029 + "qkv": ["qkv"], # For vision tower's already-packed QKV
1030 }
1031
1032 # To ensure correct weight loading and mapping.
@@ -1568,6 +1569,25 @@ class Qwen2_5_VLForConditionalGeneration(
1569 tower_model="visual.",
1570 )
1571
1572 + def get_num_mm_encoder_tokens(
1573 + self,
1574 + num_image_tokens: int,
1575 + ) -> int:
1576 + hf_config = self.config
1577 + vision_config = hf_config.vision_config
1578 + merge_size = vision_config.spatial_merge_size
1579 +
1580 + return num_image_tokens * merge_size**2
1581 +
1582 + def get_num_mm_connector_tokens(
1583 + self,
1584 + num_vision_tokens: int,
1585 + ) -> int:
1586 + hf_config = self.config
1587 + vision_config = hf_config.vision_config
1588 + merge_size = vision_config.spatial_merge_size
1589 + return num_vision_tokens // merge_size**2
1590 +
1591 @classmethod
1592 def get_language_model_spec(cls) -> tuple[nn.Module | None, str | None]:
1593 """
vllm/model_executor/models/qwen2_vl.py
+19
@@ -1491,6 +1491,25 @@ class Qwen2VLForConditionalGeneration(
1491 tower_model="visual.",
1492 )
1493
1494 + def get_num_mm_encoder_tokens(
1495 + self,
1496 + num_image_tokens: int,
1497 + ) -> int:
1498 + hf_config = self.config
1499 + vision_config = hf_config.vision_config
1500 + merge_size = vision_config.spatial_merge_size
1501 +
1502 + return num_image_tokens * merge_size**2
1503 +
1504 + def get_num_mm_connector_tokens(
1505 + self,
1506 + num_vision_tokens: int,
1507 + ) -> int:
1508 + hf_config = self.config
1509 + vision_config = hf_config.vision_config
1510 + merge_size = vision_config.spatial_merge_size
1511 + return num_vision_tokens // merge_size**2
1512 +
1513
1514 class Tarsier2MultiModalProcessor(Qwen2VLMultiModalProcessor):
1515 pass
vllm/model_executor/models/qwen3_vl.py
+21 -1
@@ -1240,6 +1240,7 @@ class Qwen3VLForConditionalGeneration(
1240 "gate_proj",
1241 "up_proj",
1242 ],
1243 + "qkv": ["qkv"], # For vision tower's already-packed QKV
1244 }
1245
1246 supports_encoder_tp_data = True
@@ -2087,10 +2088,29 @@ class Qwen3VLForConditionalGeneration(
2088 """
2089 return MultiModelKeys.from_string_field(
2090 language_model="language_model",
2090 - connector="visual.merger",
2091 + connector=["visual.merger", "visual.deepstack_merger_list"],
2092 tower_model="visual.",
2093 )
2094
2095 + def get_num_mm_encoder_tokens(
2096 + self,
2097 + num_image_tokens: int,
2098 + ) -> int:
2099 + hf_config = self.config
2100 + vision_config = hf_config.vision_config
2101 + merge_size = vision_config.spatial_merge_size
2102 +
2103 + return num_image_tokens * merge_size**2
2104 +
2105 + def get_num_mm_connector_tokens(
2106 + self,
2107 + num_vision_tokens: int,
2108 + ) -> int:
2109 + hf_config = self.config
2110 + vision_config = hf_config.vision_config
2111 + merge_size = vision_config.spatial_merge_size
2112 + return num_vision_tokens // merge_size**2
2113 +
2114 @classmethod
2115 def get_language_model_spec(cls) -> tuple[nn.Module | None, str | None]:
2116 """
vllm/v1/engine/input_processor.py
+22 -1
@@ -406,6 +406,24 @@ class InputProcessor:
406 mm_uuids[modality] = [f"{request_id}-{modality}-{i}" for i in range(n)]
407 return mm_uuids
408
409 + def _get_mm_identifier(
410 + self,
411 + mm_hash: str,
412 + lora_request: LoRARequest | None,
413 + ) -> str:
414 + """
415 + When enable_tower_connector_lora is True, multi-modal embeddings
416 + vary depending on the LoRA request. Therefore, the mm_hash must be
417 + generated based on the LoRA request to prevent incorrect cache hits.
418 + """
419 + if (
420 + lora_request is None
421 + or self.lora_config is None
422 + or not self.lora_config.enable_tower_connector_lora
423 + ):
424 + return mm_hash
425 + return f"{lora_request.lora_name}:{mm_hash}"
426 +
427 @staticmethod
428 def assign_request_id(request: EngineCoreRequest):
429 """Replace the externally supplied request ID with an internal request ID
@@ -539,7 +557,10 @@ class InputProcessor:
557 MultiModalFeatureSpec(
558 data=decoder_mm_inputs[modality][idx],
559 modality=modality,
542 - identifier=decoder_mm_hashes[modality][idx],
560 + identifier=self._get_mm_identifier(
561 + decoder_mm_hashes[modality][idx],
562 + lora_request,
563 + ),
564 mm_position=decoder_mm_positions[modality][idx],
565 )
566 )
vllm/v1/worker/gpu_model_runner.py
+79 -6
@@ -54,6 +54,7 @@ from vllm.forward_context import (
54 set_forward_context,
55 )
56 from vllm.logger import init_logger
57 +from vllm.lora.layers import LoRAMapping, LoRAMappingType
58 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
59 from vllm.model_executor.layers.rotary_embedding import (
60 MRotaryEmbedding,
@@ -79,7 +80,11 @@ from vllm.model_executor.models.interfaces_base import (
80 is_text_generation_model,
81 )
82 from vllm.multimodal import MULTIMODAL_REGISTRY
82 -from vllm.multimodal.inputs import BatchedTensorInputs, MultiModalKwargsItem
83 +from vllm.multimodal.inputs import (
84 + BatchedTensorInputs,
85 + MultiModalKwargsItem,
86 + PlaceholderRange,
87 +)
88 from vllm.multimodal.utils import group_mm_kwargs_by_modality
89 from vllm.pooling_params import PoolingParams
90 from vllm.sampling_params import SamplingType
@@ -2095,7 +2100,11 @@ class GPUModelRunner(
2100 def _batch_mm_inputs_from_scheduler(
2101 self,
2102 scheduler_output: "SchedulerOutput",
2098 - ) -> tuple[list[str], list[MultiModalKwargsItem]]:
2103 + ) -> tuple[
2104 + list[str],
2105 + list[MultiModalKwargsItem],
2106 + list[tuple[str, PlaceholderRange]],
2107 + ]:
2108 """Batch multimodal inputs from scheduled encoder inputs.
2109
2110 Args:
@@ -2103,16 +2112,20 @@ class GPUModelRunner(
2112 inputs.
2113
2114 Returns:
2106 - A tuple of (mm_hashes, mm_kwargs) where:
2115 + A tuple of (mm_hashes, mm_kwargs, mm_lora_refs) where:
2116 - mm_hashes: List of multimodal hashes for each item
2117 - mm_kwargs: List of multimodal kwargs for each item
2118 + - mm_lora_refs: List of (req_id, placeholder_range) for each item
2119 """
2120 scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
2121 if not scheduled_encoder_inputs:
2112 - return [], []
2122 + return [], [], []
2123
2124 mm_hashes = list[str]()
2125 mm_kwargs = list[MultiModalKwargsItem]()
2126 + # Multimodal LoRA reference info to map each multimodal item
2127 + # back to its request & position
2128 + mm_lora_refs = list[tuple[str, PlaceholderRange]]()
2129 for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
2130 req_state = self.requests[req_id]
2131
@@ -2123,13 +2136,16 @@ class GPUModelRunner(
2136
2137 mm_hashes.append(mm_feature.identifier)
2138 mm_kwargs.append(mm_feature.data)
2139 + mm_lora_refs.append((req_id, mm_feature.mm_position))
2140
2127 - return mm_hashes, mm_kwargs
2141 + return mm_hashes, mm_kwargs, mm_lora_refs
2142
2143 def _execute_mm_encoder(
2144 self, scheduler_output: "SchedulerOutput"
2145 ) -> list[torch.Tensor]:
2132 - mm_hashes, mm_kwargs = self._batch_mm_inputs_from_scheduler(scheduler_output)
2146 + mm_hashes, mm_kwargs, mm_lora_refs = self._batch_mm_inputs_from_scheduler(
2147 + scheduler_output
2148 + )
2149
2150 if not mm_kwargs:
2151 return []
@@ -2142,6 +2158,63 @@ class GPUModelRunner(
2158 # multimodal inputs. The proper solution should be reordering the
2159 # encoder outputs.
2160 model = cast(SupportsMultiModal, self.model)
2161 +
2162 + if self.lora_config and self.lora_manager.supports_tower_connector_lora():
2163 + # Build LoRA mappings independently for encoder inputs
2164 + # (encoder batch structure is different from main batch)
2165 + prompt_lora_mapping = []
2166 + token_lora_mapping = []
2167 + lora_requests = set()
2168 + encoder_token_counts = []
2169 +
2170 + for req_id, pos_info in mm_lora_refs:
2171 + req_idx = self.input_batch.req_id_to_index[req_id]
2172 + lora_id = int(self.input_batch.request_lora_mapping[req_idx])
2173 +
2174 + # Prefer pos_info.get_num_embeds to count precise MM embedding tokens.
2175 + num_tokens = self.model.get_num_mm_encoder_tokens( # type: ignore[attr-defined]
2176 + pos_info.get_num_embeds
2177 + )
2178 + prompt_lora_mapping.append(lora_id)
2179 + token_lora_mapping.extend([lora_id] * num_tokens)
2180 + encoder_token_counts.append(num_tokens)
2181 +
2182 + if lora_id > 0:
2183 + lora_request = self.input_batch.lora_id_to_lora_request.get(lora_id)
2184 + if lora_request is not None:
2185 + lora_requests.add(lora_request)
2186 +
2187 + # Set tower adapter mapping
2188 + tower_mapping = LoRAMapping(
2189 + tuple(token_lora_mapping),
2190 + tuple(prompt_lora_mapping),
2191 + is_prefill=True,
2192 + type=LoRAMappingType.TOWER,
2193 + )
2194 + self.lora_manager.set_active_adapters(lora_requests, tower_mapping)
2195 +
2196 + if hasattr(self.model, "get_num_mm_connector_tokens"):
2197 + post_op_counts = [
2198 + self.model.get_num_mm_connector_tokens(num_tokens) # type: ignore[attr-defined]
2199 + for num_tokens in encoder_token_counts
2200 + ]
2201 +
2202 + connector_token_mapping = np.repeat(
2203 + np.array(prompt_lora_mapping, dtype=np.int32),
2204 + np.array(post_op_counts, dtype=np.int32),
2205 + )
2206 + connector_mapping = LoRAMapping(
2207 + index_mapping=tuple(connector_token_mapping.tolist()),
2208 + prompt_mapping=tuple(prompt_lora_mapping),
2209 + is_prefill=True,
2210 + type=LoRAMappingType.CONNECTOR,
2211 + )
2212 +
2213 + self.lora_manager.set_active_adapters(
2214 + lora_requests,
2215 + connector_mapping,
2216 + )
2217 +
2218 encoder_outputs: list[torch.Tensor] = []
2219 for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
2220 mm_kwargs,
vllm/v1/worker/lora_model_runner_mixin.py
+27 -14
@@ -5,6 +5,7 @@ Define LoRA functionality mixin for model runners.
5 """
6
7 from contextlib import contextmanager
8 +from typing import TypeAlias
9
10 import numpy as np
11 import torch
@@ -13,14 +14,14 @@ import torch.nn as nn
14 from vllm.config import VllmConfig
15 from vllm.config.lora import LoRAConfig
16 from vllm.logger import init_logger
16 -from vllm.lora.layers import LoRAMapping
17 +from vllm.lora.layers import LoRAMapping, LoRAMappingType
18 from vllm.lora.request import LoRARequest
19 from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
19 -from vllm.model_executor.models import supports_lora, supports_multimodal
20 +from vllm.model_executor.models import supports_lora
21 from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch
22 from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch
23
23 -InputBatch = TPUInputBatch | GPUInputBatch
24 +InputBatch: TypeAlias = TPUInputBatch | GPUInputBatch
25
26 logger = init_logger(__name__)
27
@@ -28,29 +29,28 @@ logger = init_logger(__name__)
29 # Defined as a mixin for GPUModelRunner
30 class LoRAModelRunnerMixin:
31 def load_lora_model(
31 - self, model: nn.Module, vllm_config: VllmConfig, device: torch.device
32 + self,
33 + model: nn.Module,
34 + vllm_config: VllmConfig,
35 + device: torch.device,
36 ) -> nn.Module:
37 if not supports_lora(model):
38 raise ValueError(f"{model.__class__.__name__} does not support LoRA yet.")
39
36 - if supports_multimodal(model):
37 - logger.warning(
38 - "Regarding multimodal models, vLLM currently "
39 - "only supports adding LoRA to language model."
40 - )
40 # Add LoRA Manager to the Model Runner
41 self.lora_manager = LRUCacheWorkerLoRAManager(
42 vllm_config,
43 device,
44 model.embedding_modules,
45 )
47 - return self.lora_manager.create_lora_manager(model)
46 + return self.lora_manager.create_lora_manager(model, vllm_config)
47
48 def _set_active_loras(
49 self,
50 prompt_lora_mapping: tuple[int, ...],
51 token_lora_mapping: tuple[int, ...],
52 lora_requests: set[LoRARequest],
53 + mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
54 ) -> None:
55 self._ensure_lora_enabled()
56
@@ -59,7 +59,10 @@ class LoRAModelRunnerMixin:
59 # On cuda platforms we use the same kernels for prefill and
60 # decode and this flag is generally ignored.
61 lora_mapping = LoRAMapping(
62 - token_lora_mapping, prompt_lora_mapping, is_prefill=True
62 + token_lora_mapping,
63 + prompt_lora_mapping,
64 + is_prefill=True,
65 + type=mapping_type,
66 )
67 self.lora_manager.set_active_adapters(lora_requests, lora_mapping)
68
@@ -72,6 +75,7 @@ class LoRAModelRunnerMixin:
75 input_batch: InputBatch,
76 num_scheduled_tokens: np.ndarray,
77 num_sampled_tokens: np.ndarray | None = None,
78 + mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
79 ) -> None:
80 if num_sampled_tokens is None:
81 num_sampled_tokens = np.ones_like(num_scheduled_tokens, dtype=np.int32)
@@ -83,7 +87,7 @@ class LoRAModelRunnerMixin:
87 input_batch.make_lora_inputs(num_scheduled_tokens, num_sampled_tokens)
88 )
89 return self._set_active_loras(
86 - prompt_lora_mapping, token_lora_mapping, lora_requests
90 + prompt_lora_mapping, token_lora_mapping, lora_requests, mapping_type
91 )
92
93 @contextmanager
@@ -127,6 +131,7 @@ class LoRAModelRunnerMixin:
131 self,
132 lora_config: LoRAConfig | None,
133 num_scheduled_tokens: np.ndarray,
134 + mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
135 num_sampled_tokens: np.ndarray | None = None,
136 activate_lora: bool = True,
137 ):
@@ -168,7 +173,10 @@ class LoRAModelRunnerMixin:
173 }
174
175 self._set_active_loras(
171 - tuple(sample_lora_mapping), tuple(token_lora_mapping), lora_requests
176 + tuple(sample_lora_mapping),
177 + tuple(token_lora_mapping),
178 + lora_requests,
179 + mapping_type,
180 )
181
182 yield
@@ -181,11 +189,16 @@ class LoRAModelRunnerMixin:
189 num_sampled_tokens: np.ndarray,
190 activate_lora: bool = True,
191 remove_lora: bool = True,
192 + mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
193 ):
194 with (
195 self.maybe_setup_dummy_loras(lora_config, remove_lora),
196 self.maybe_select_dummy_loras(
188 - lora_config, num_scheduled_tokens, num_sampled_tokens, activate_lora
197 + lora_config,
198 + num_scheduled_tokens,
199 + mapping_type,
200 + num_sampled_tokens,
201 + activate_lora,
202 ),
203 ):
204 yield
vllm/v1/worker/tpu_model_runner.py
+8 -3
@@ -31,7 +31,8 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_
31 from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
32 from vllm.forward_context import set_forward_context
33 from vllm.logger import init_logger
34 -from vllm.lora.layers import BaseLayerWithLoRA
34 +from vllm.lora.layers import BaseLayerWithLoRA, LoRAMappingType
35 +from vllm.lora.request import LoRARequest
36 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
37 from vllm.model_executor.model_loader import get_model_loader
38 from vllm.model_executor.model_loader.tpu import TPUModelLoader
@@ -1468,11 +1469,15 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1469 self._hidden_states_dtype = out.dtype
1470
1471 def _set_active_loras(
1471 - self, prompt_lora_mapping, token_lora_mapping, lora_requests
1472 + self,
1473 + prompt_lora_mapping: tuple[int, ...],
1474 + token_lora_mapping: tuple[int, ...],
1475 + lora_requests: set[LoRARequest],
1476 + mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
1477 ) -> None:
1478 torch_xla.sync(wait=False) # Captures input updates
1479 super()._set_active_loras(
1475 - prompt_lora_mapping, token_lora_mapping, lora_requests
1480 + prompt_lora_mapping, token_lora_mapping, lora_requests, mapping_type
1481 )
1482 torch_xla.sync(wait=False) # Captures metadata updates
1483