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 3, 11:45 UTC
[LoRA] Support 2D and 3D MoE LoRA adapter at the same time (#42242)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Jee Jee Li committed
May 18, 2026 at 15:22 UTC
7d5b033782681acee274f4f379c9fadc557fd7e8
16 files changed
+391
-9
docs/features/lora.md
+51
@@ -248,6 +248,57 @@ Now, you can specify a base_model_name alongside the name and path using JSON fo
248
249
To provide the backward compatibility support, you can still use the old key-value format (name=path), but the `base_model_name` will remain unspecified in that case.
250
251
+## Mixing 2D and 3D MoE LoRA Adapters
252
+
253
+To serve 2D-format(based on `megatron`) and 3D-format (based on `peft`) adapters from the same engine instance, start the server with `--enable-mixed-moe-lora-format`
254
+and declare the layout of each adapter explicitly via the `is_3d_lora_weight` field.
255
+
256
+Server startup (static modules):
257
+
258
+```bash
259
+vllm serve Qwen/Qwen3.6-35B-A3B \
260
+ --enable-lora \
261
+ --enable-mixed-moe-lora-format \
262
+ --tensor-parallel-size 4 \
263
+ --enable-expert-parallel \
264
+ --lora-modules \
265
+ '{"name": "lora-2d", "path": "jeeejeee/qwen36-35ba3b-2d-weights-poken-lora", "is_3d_lora_weight": false}' \
266
+ '{"name": "lora-3d", "path": "jeeejeee/qwen36-35ba3b-moe-all-linear-poken-lora", "is_3d_lora_weight": true}'
267
+```
268
+
269
+Dynamic load via `/v1/load_lora_adapter`:
270
+
271
+```bash
272
+curl -X POST http://localhost:8000/v1/load_lora_adapter \
273
+-H "Content-Type: application/json" \
274
+-d '{
275
+ "lora_name": "lora-3d",
276
+ "lora_path": "/path/to/3d-format-lora",
277
+ "is_3d_lora_weight": true
278
+}'
279
+```
280
+
281
+!!! warning "You must know your adapter's layout"
282
+ Under `--enable-mixed-moe-lora-format`, vLLM trusts whatever
283
+ `is_3d_lora_weight` the caller declares — it does **not** inspect the
284
+ checkpoint to verify. A wrong declaration will load weights into the
285
+ wrong stacked buffers and silently produce garbage outputs, with no
286
+ error at load time. Confirm the layout before serving:
287
+
288
+ - **2D (per-expert, megatron-style)** → set `is_3d_lora_weight: false`.
289
+ Adapter keys look like `...experts.{idx}.gate_proj.lora_A.weight`,
290
+ `...experts.{idx}.up_proj.lora_A.weight`,
291
+ `...experts.{idx}.down_proj.lora_A.weight` — one set per expert.
292
+ - **3D (fused, peft-style)** → set `is_3d_lora_weight: true`.
293
+ Adapter keys look like `...experts.gate_up_proj.lora_A.weight`,
294
+ `...experts.down_proj.lora_A.weight` — a single tensor that stacks
295
+ all experts on the leading dim.
296
+
297
+When `--enable-mixed-moe-lora-format` is **not** set, `is_3d_lora_weight`
298
+is ignored: vLLM picks the wrapper from the base model's
299
+`is_3d_moe_weight` and the adapter is required to match. The field is
300
+also ignored for non-MoE models.
301
+
302
## LoRA model lineage in model card
303
304
The new format of `--lora-modules` is mainly to support the display of parent model information in the model card. Here's an explanation of how your current response supports this:
tests/lora/conftest.py
+10
@@ -313,6 +313,16 @@ def qwen35_vl_lora_files():
313
return snapshot_download(repo_id="jeeejeee/qwen35-4b-all-linear-pokemon-lora")
314
315
316
+@pytest.fixture(scope="session")
317
+def qwen36_moe_2d_lora_files():
318
+ return snapshot_download(repo_id="jeeejeee/qwen36-35ba3b-2d-weights-poken-lora")
319
+
320
+
321
+@pytest.fixture(scope="session")
322
+def qwen36_moe_3d_lora_files():
323
+ return snapshot_download(repo_id="jeeejeee/qwen36-35ba3b-moe-all-linear-poken-lora")
324
+
325
+
326
@pytest.fixture
327
def reset_default_device():
328
"""
tests/lora/test_qwen36_moe_lora.py
new
+156
@@ -0,0 +1,156 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+import vllm
7
+import vllm.config
8
+from vllm.assets.image import ImageAsset
9
+from vllm.lora.request import LoRARequest
10
+
11
+from ..utils import multi_gpu_test
12
+
13
+MODEL_PATH = "Qwen/Qwen3.6-35B-A3B"
14
+
15
+LORA_2D_ID = 1
16
+LORA_3D_ID = 2
17
+
18
+PROMPT_TEMPLATE = """<|im_start|>user
19
+<|vision_start|><|image_pad|><|vision_end|>What is in the image?<|im_end|>
20
+<|im_start|>assistant
21
+<think>
22
+
23
+</think>
24
+"""
25
+
26
+# Visual captioning prompts: each image will be paired with one LoRA in the
27
+# mixed-batch case so we can check per-prompt routing.
28
+VL_TEST_IMAGES = [
29
+ ImageAsset("stop_sign"),
30
+ ImageAsset("cherry_blossom"),
31
+]
32
+
33
+
34
+def _build_prompts() -> list[dict]:
35
+ return [
36
+ {
37
+ "prompt": PROMPT_TEMPLATE,
38
+ "multi_modal_data": {"image": asset.pil_image},
39
+ }
40
+ for asset in VL_TEST_IMAGES
41
+ ]
42
+
43
+
44
+def _generate(llm: vllm.LLM, lora_request) -> list[str]:
45
+ outputs = llm.generate(
46
+ _build_prompts(),
47
+ vllm.SamplingParams(temperature=0, max_tokens=128),
48
+ lora_request=lora_request,
49
+ )
50
+ return [out.outputs[0].text.strip() for out in outputs]
51
+
52
+
53
+def _run_mixed_2d_3d_lora_test(
54
+ lora_2d_files: str,
55
+ lora_3d_files: str,
56
+ tensor_parallel_size: int,
57
+ fully_sharded_loras: bool,
58
+) -> None:
59
+ llm = vllm.LLM(
60
+ model=MODEL_PATH,
61
+ max_model_len=4096,
62
+ enable_lora=True,
63
+ enable_mixed_moe_lora_format=True,
64
+ max_loras=2,
65
+ max_lora_rank=8,
66
+ max_num_seqs=4,
67
+ enforce_eager=True,
68
+ tensor_parallel_size=tensor_parallel_size,
69
+ enable_expert_parallel=not fully_sharded_loras,
70
+ fully_sharded_loras=fully_sharded_loras,
71
+ trust_remote_code=True,
72
+ enable_tower_connector_lora=True,
73
+ mm_processor_cache_gb=0,
74
+ limit_mm_per_prompt={"image": 1},
75
+ compilation_config=vllm.config.CompilationConfig(
76
+ cudagraph_specialize_lora=False,
77
+ ),
78
+ )
79
+
80
+ lora_2d = LoRARequest(
81
+ "lora_2d",
82
+ LORA_2D_ID,
83
+ lora_2d_files,
84
+ is_3d_lora_weight=False,
85
+ )
86
+ lora_3d = LoRARequest(
87
+ "lora_3d",
88
+ LORA_3D_ID,
89
+ lora_3d_files,
90
+ is_3d_lora_weight=True,
91
+ )
92
+
93
+ # Reference: each adapter alone over both prompts.
94
+ outputs_2d_alone = _generate(llm, lora_2d)
95
+ outputs_3d_alone = _generate(llm, lora_3d)
96
+
97
+ assert len(outputs_2d_alone) == len(VL_TEST_IMAGES)
98
+ assert len(outputs_3d_alone) == len(VL_TEST_IMAGES)
99
+ for text in outputs_2d_alone + outputs_3d_alone:
100
+ assert text, "Empty output from single-adapter LoRA generation"
101
+
102
+ # Mixed batch: prompt 0 uses the 2D adapter, prompt 1 uses the 3D
103
+ # adapter. Per-prompt outputs must match the standalone runs.
104
+ mixed_outputs = _generate(llm, [lora_2d, lora_3d])
105
+
106
+ assert mixed_outputs[0] == outputs_2d_alone[0], (
107
+ f"Mixed-batch 2D output {mixed_outputs[0]!r} does not match "
108
+ f"standalone 2D output {outputs_2d_alone[0]!r}"
109
+ )
110
+ assert mixed_outputs[1] == outputs_3d_alone[1], (
111
+ f"Mixed-batch 3D output {mixed_outputs[1]!r} does not match "
112
+ f"standalone 3D output {outputs_3d_alone[1]!r}"
113
+ )
114
+
115
+ # Reverse assignment: neither adapter should be silently aliased.
116
+ swapped_outputs = _generate(llm, [lora_3d, lora_2d])
117
+ assert swapped_outputs[0] == outputs_3d_alone[0], (
118
+ f"Swapped-batch 3D output {swapped_outputs[0]!r} does not match "
119
+ f"standalone 3D output {outputs_3d_alone[0]!r}"
120
+ )
121
+ assert swapped_outputs[1] == outputs_2d_alone[1], (
122
+ f"Swapped-batch 2D output {swapped_outputs[1]!r} does not match "
123
+ f"standalone 2D output {outputs_2d_alone[1]!r}"
124
+ )
125
+
126
+
127
+@pytest.mark.skip(reason="This model is too big, so skip this test temporarily.")
128
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
129
+@multi_gpu_test(num_gpus=2)
130
+def test_qwen36_moe_mixed_2d_3d_lora_tp2(
131
+ qwen36_moe_2d_lora_files,
132
+ qwen36_moe_3d_lora_files,
133
+ fully_sharded_loras,
134
+):
135
+ _run_mixed_2d_3d_lora_test(
136
+ lora_2d_files=qwen36_moe_2d_lora_files,
137
+ lora_3d_files=qwen36_moe_3d_lora_files,
138
+ tensor_parallel_size=2,
139
+ fully_sharded_loras=fully_sharded_loras,
140
+ )
141
+
142
+
143
+@pytest.mark.skip(reason="This model is too big, so skip this test temporarily.")
144
+@pytest.mark.parametrize("fully_sharded_loras", [False, True])
145
+@multi_gpu_test(num_gpus=4)
146
+def test_qwen36_moe_mixed_2d_3d_lora_tp4(
147
+ qwen36_moe_2d_lora_files,
148
+ qwen36_moe_3d_lora_files,
149
+ fully_sharded_loras,
150
+):
151
+ _run_mixed_2d_3d_lora_test(
152
+ lora_2d_files=qwen36_moe_2d_lora_files,
153
+ lora_3d_files=qwen36_moe_3d_lora_files,
154
+ tensor_parallel_size=4,
155
+ fully_sharded_loras=fully_sharded_loras,
156
+ )
tests/lora/test_qwen3moe_tp.py
+7
-2
@@ -77,6 +77,7 @@ def test_qwen3moe_lora(qwen3moe_lora_files):
77
enforce_eager=True,
78
trust_remote_code=True,
79
enable_chunked_prefill=True,
80
+ enable_mixed_moe_lora_format=True,
81
)
82
83
generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
@@ -85,7 +86,8 @@ def test_qwen3moe_lora(qwen3moe_lora_files):
86
87
@multi_gpu_test(num_gpus=2)
88
@pytest.mark.parametrize("ep", [False, True])
88
-def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files):
89
+@pytest.mark.parametrize("enable_mixed_moe_lora_format", [False, True])
90
+def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files, enable_mixed_moe_lora_format):
91
llm = vllm.LLM(
92
MODEL_PATH,
93
max_model_len=1024,
@@ -94,6 +96,7 @@ def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files):
96
trust_remote_code=True,
97
enable_expert_parallel=ep,
98
tensor_parallel_size=2,
99
+ enable_mixed_moe_lora_format=enable_mixed_moe_lora_format,
100
)
101
102
generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
@@ -102,7 +105,8 @@ def test_qwen3moe_lora_tp2(ep, qwen3moe_lora_files):
105
106
@multi_gpu_test(num_gpus=4)
107
@pytest.mark.parametrize("ep", [False, True])
105
-def test_qwen3moe_lora_tp4(ep, qwen3moe_lora_files):
108
+@pytest.mark.parametrize("enable_mixed_moe_lora_format", [False, True])
109
+def test_qwen3moe_lora_tp4(ep, qwen3moe_lora_files, enable_mixed_moe_lora_format):
110
llm = vllm.LLM(
111
MODEL_PATH,
112
max_model_len=1024,
@@ -111,6 +115,7 @@ def test_qwen3moe_lora_tp4(ep, qwen3moe_lora_files):
115
trust_remote_code=True,
116
enable_expert_parallel=ep,
117
tensor_parallel_size=4,
118
+ enable_mixed_moe_lora_format=enable_mixed_moe_lora_format,
119
)
120
121
generate_and_test(llm, qwen3moe_lora_files, lora_id=1)
vllm/config/lora.py
+7
@@ -71,6 +71,12 @@ class LoRAConfig:
71
for variable LoRA usage patterns at the cost of increased startup time and
72
memory usage. Only takes effect when cudagraph_specialize_lora is True.
73
"""
74
+ enable_mixed_moe_lora_format: bool = False
75
+ """If True, force the engine to use the universal 2D MoE LoRA wrapper
76
+ (`FusedMoEWithLoRA`) regardless of the model's `is_3d_moe_weight` flag, so
77
+ that 2D-format and 3D-format MoE LoRA adapters can be served in the same
78
+ deployment. Only meaningful forMoE models; ignored otherwise. Default False
79
+ keeps the existing model-driven behavior."""
80
81
def compute_hash(self) -> str:
82
"""
@@ -90,6 +96,7 @@ class LoRAConfig:
96
factors.append(self.fully_sharded_loras)
97
factors.append(self.lora_dtype)
98
factors.append(self.enable_tower_connector_lora)
99
+ factors.append(self.enable_mixed_moe_lora_format)
100
# target_modules affects which modules get LoRA applied
101
factors.append(
102
tuple(sorted(self.target_modules)) if self.target_modules else None
vllm/engine/arg_utils.py
+6
@@ -584,6 +584,7 @@ class EngineArgs:
584
lora_target_modules: list[str] | None = LoRAConfig.target_modules
585
enable_tower_connector_lora: bool = LoRAConfig.enable_tower_connector_lora
586
specialize_active_lora: bool = LoRAConfig.specialize_active_lora
587
+ enable_mixed_moe_lora_format: bool = LoRAConfig.enable_mixed_moe_lora_format
588
589
ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight
590
num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override
@@ -1272,6 +1273,10 @@ class EngineArgs:
1273
lora_group.add_argument(
1274
"--specialize-active-lora", **lora_kwargs["specialize_active_lora"]
1275
)
1276
+ lora_group.add_argument(
1277
+ "--enable-mixed-moe-lora-format",
1278
+ **lora_kwargs["enable_mixed_moe_lora_format"],
1279
+ )
1280
1281
# Observability arguments
1282
observability_kwargs = get_kwargs(ObservabilityConfig)
@@ -2009,6 +2014,7 @@ class EngineArgs:
2014
target_modules=self.lora_target_modules,
2015
enable_tower_connector_lora=self.enable_tower_connector_lora,
2016
specialize_active_lora=self.specialize_active_lora,
2017
+ enable_mixed_moe_lora_format=self.enable_mixed_moe_lora_format,
2018
max_cpu_loras=self.max_cpu_loras
2019
if self.max_cpu_loras and self.max_cpu_loras > 0
2020
else None,
vllm/entrypoints/openai/models/protocol.py
+1
@@ -16,3 +16,4 @@ class LoRAModulePath:
16
name: str
17
path: str
18
base_model_name: str | None = None
19
+ is_3d_lora_weight: bool = False
vllm/entrypoints/openai/models/serving.py
+4
-1
@@ -121,7 +121,9 @@ class OpenAIServingModels:
121
return
122
for lora in self.static_lora_modules:
123
load_request = LoadLoRAAdapterRequest(
124
- lora_path=lora.path, lora_name=lora.name
124
+ lora_path=lora.path,
125
+ lora_name=lora.name,
126
+ is_3d_lora_weight=lora.is_3d_lora_weight,
127
)
128
load_result = await self.load_lora_adapter(
129
request=load_request, base_model_name=lora.base_model_name
@@ -177,6 +179,7 @@ class OpenAIServingModels:
179
lora_int_id=lora_int_id,
180
lora_path=lora_path,
181
load_inplace=request.load_inplace,
182
+ is_3d_lora_weight=request.is_3d_lora_weight,
183
)
184
if base_model_name is not None and self.is_base_model(base_model_name):
185
lora_request.base_model_name = base_model_name
vllm/entrypoints/serve/lora/api_router.py
+1
@@ -37,6 +37,7 @@ def attach_router(app: FastAPI):
37
"lora_name": "body.name",
38
"lora_path": "body.src",
39
"load_inplace": "body.load_inplace || `false`",
40
+ "is_3d_lora_weight": "body.is_3d_lora_weight || `false`",
41
},
42
)
43
@router.post("/v1/load_lora_adapter", dependencies=[Depends(validate_json_request)])
vllm/entrypoints/serve/lora/protocol.py
+1
@@ -8,6 +8,7 @@ class LoadLoRAAdapterRequest(BaseModel):
8
lora_name: str
9
lora_path: str
10
load_inplace: bool = False
11
+ is_3d_lora_weight: bool = False
12
13
14
class UnloadLoRAAdapterRequest(BaseModel):
vllm/lora/layers/fused_moe.py
+4
@@ -37,6 +37,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
37
# For non-gated MoE (is_act_and_mul=False), only 1 slice is needed
38
# since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3)
39
self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1
40
+ # Mirrors per-(lora_id) layout of `self.lora_a_stacked` (built in
41
+ # `create_lora_weights`) so `create_dummy_lora`'s n_slices fallback
42
+ # matches `lora_a_stacked` length under EP.
43
+ self.n_slices = base_layer.local_num_experts * (self._w13_slices + 1)
44
45
self.base_layer.ensure_moe_quant_config_init()
46
if getattr(self.base_layer.quant_method, "supports_internal_mk", False):
vllm/lora/lora_model.py
+7
@@ -29,12 +29,17 @@ class LoRAModel:
29
lora_model_id: int,
30
rank: int,
31
loras: dict[str, LoRALayerWeights],
32
+ is_3d_lora_weight: bool = False,
33
) -> None:
34
"""
35
Args:
36
lora_model_id: The integer id for the lora model.
37
rank: lora rank.
38
loras: module name -> weights for lora-replaced layers.
39
+ is_3d_lora_weight: Whether the on-disk MoE adapter is in the 3D
40
+ fused (gate_up_proj / down_proj) layout. Propagated from the
41
+ originating LoRARequest. Only consulted by the LoRA model
42
+ manager when enable_mixed_moe_lora_format is on.
43
44
"""
45
self.id = lora_model_id
@@ -44,6 +49,7 @@ class LoRAModel:
49
)
50
self.rank = rank
51
self.loras: dict[str, LoRALayerWeights] = loras
52
+ self.is_3d_lora_weight = is_3d_lora_weight
53
54
def clone(self, lora_model_id: int) -> "LoRAModel":
55
"""Return a copy of the object with different ids.
@@ -53,6 +59,7 @@ class LoRAModel:
59
lora_model_id,
60
rank=self.rank,
61
loras=self.loras.copy(),
62
+ is_3d_lora_weight=self.is_3d_lora_weight,
63
)
64
65
def get_lora(self, module_name: str) -> LoRALayerWeights | None:
vllm/lora/model_manager.py
+117
-3
@@ -102,7 +102,6 @@ class LoRAModelManager:
102
self.max_num_batched_tokens = math.ceil(max_num_batched_tokens / 8) * 8
103
self.lora_index_to_id: list[int | None] = [None] * self.lora_slots
104
self.vocab_size = vocab_size
105
- self.packed_modules_mapping = process_packed_modules_mapping(self.model)
105
106
self.is_pooling_model = is_pooling_model(self.model)
107
self.packed_modules: dict[str, list[str]] = {}
@@ -110,7 +109,20 @@ class LoRAModelManager:
109
# Dict instead of a set for compatibility with LRUCache.
110
self._last_mapping: LoRAMapping | None = None
111
is_moe = is_moe_model(self.model)
113
- self._is_3d_moe_model = is_moe and self.model.is_3d_moe_weight
112
+ # Whether the underlying model class declares 3D fused MoE weights.
113
+ self._model_is_3d_moe = is_moe and self.model.is_3d_moe_weight
114
+ # When the engine is started with enable_mixed_moe_lora_format=True
115
+ # we force the universal 2D wrapper (FusedMoEWithLoRA) regardless of
116
+ # the model's 3D flag, so 2D and 3D adapters can coexist.
117
+ self._enable_mixed_moe_lora_format = (
118
+ is_moe and lora_config.enable_mixed_moe_lora_format
119
+ )
120
+ self._is_3d_moe_model = (
121
+ self._model_is_3d_moe and not self._enable_mixed_moe_lora_format
122
+ )
123
+ self.packed_modules_mapping = process_packed_modules_mapping(
124
+ self.model, force_2d_moe=self._enable_mixed_moe_lora_format
125
+ )
126
self._is_non_gated_moe = is_moe and self.model.is_non_gated_moe
127
self._init_punica_wrapper(max_num_batched_tokens, vllm_config)
128
self._create_lora_modules()
@@ -748,7 +760,16 @@ class LoRAModelManager:
760
if isinstance(module, FusedMoE3DWithLoRA):
761
self._stack_moe_lora_weights(lora_model, module, module_name)
762
elif isinstance(module, FusedMoEWithLoRA):
751
- self._slice_moe_lora_ep(lora_model, module, module_name)
763
+ # When mixed mode is enabled the universal 2D wrapper has to
764
+ # absorb both 2D and 3D-format adapters. 3D-format adapters
765
+ # need to be split into per-(w1, w2, w3) tensors before the
766
+ # 2D set_lora can copy them into the stacked buffers.
767
+ if self._enable_mixed_moe_lora_format and getattr(
768
+ lora_model, "is_3d_lora_weight", False
769
+ ):
770
+ self._convert_3d_to_2d_moe_lora(lora_model, module, module_name)
771
+ else:
772
+ self._slice_moe_lora_ep(lora_model, module, module_name)
773
774
first_lora: LoRALayerWeights = next(iter(lora_model.loras.values()))
775
assert first_lora.lora_a is not None
@@ -871,6 +892,99 @@ class LoRAModelManager:
892
module_lora.lora_a = lora_a
893
module_lora.lora_b = lora_b
894
895
+ def _convert_3d_to_2d_moe_lora(
896
+ self,
897
+ lora_model: LoRAModel,
898
+ module: FusedMoEWithLoRA,
899
+ module_name: str,
900
+ ) -> None:
901
+ """Convert a 3D-format MoE LoRA checkpoint into the 2D pack layout
902
+ that `FusedMoEWithLoRA.set_lora` expects.
903
+
904
+ On disk the 3D PEFT layout stores two flat tensor pairs per layer:
905
+ - `{module}.base_layer.lora_{A,B}`: gate_up_proj, with shapes
906
+ `(rank * num_experts, hidden)` / `(intermediate * 2,
907
+ rank * num_experts)`
908
+ - `{module}.lora_{A,B}`: down_proj, with shapes
909
+ `(rank * num_experts, intermediate)` / `(hidden,
910
+ rank * num_experts)`
911
+ The 2D wrapper expects three stacked per-expert tensors,
912
+ `[w1, w2, w3]`, with `(num_experts, rank, in)` for lora_a and
913
+ `(num_experts, out, rank)` for lora_b. In the 3D layout w1
914
+ (gate_proj) and w3 (up_proj) share the rank-r intermediate
915
+ representation, so both halves use the same lora_a tensor.
916
+
917
+ Only invoked when `enable_mixed_moe_lora_format=True` and the
918
+ source LoRARequest declares `is_3d_lora_weight=True`.
919
+ """
920
+ gate_up_proj_lora = self._get_lora_layer_weights(
921
+ lora_model, module_name + ".base_layer"
922
+ )
923
+ down_proj_lora = self._get_lora_layer_weights(lora_model, module_name)
924
+ if gate_up_proj_lora is None or down_proj_lora is None:
925
+ # Either the adapter omits the experts entirely or the file
926
+ # layout differs from what this path supports; leave the entry
927
+ # untouched so set_lora can raise a clear error if needed.
928
+ return
929
+
930
+ local_num_experts = module.base_layer.local_num_experts
931
+ global_num_experts = module.base_layer.global_num_experts
932
+ ep_rank = module.base_layer.ep_rank
933
+ expert_start = ep_rank * local_num_experts
934
+ expert_end = expert_start + local_num_experts
935
+
936
+ # Reshape and EP-slice into per-expert 3D tensors. This mirrors
937
+ # `_stack_moe_lora_weights`; for non-EP runs the slice is a no-op.
938
+ gate_up_a = gate_up_proj_lora.lora_a.reshape(
939
+ global_num_experts, -1, gate_up_proj_lora.lora_a.shape[-1]
940
+ )[expert_start:expert_end].contiguous()
941
+ gate_up_b = (
942
+ gate_up_proj_lora.lora_b.reshape(
943
+ gate_up_proj_lora.lora_b.shape[0], -1, global_num_experts
944
+ )[..., expert_start:expert_end]
945
+ .permute(2, 0, 1)
946
+ .contiguous()
947
+ )
948
+ down_a = down_proj_lora.lora_a.reshape(
949
+ global_num_experts, -1, down_proj_lora.lora_a.shape[-1]
950
+ )[expert_start:expert_end].contiguous()
951
+ down_b = (
952
+ down_proj_lora.lora_b.reshape(
953
+ down_proj_lora.lora_b.shape[0], -1, global_num_experts
954
+ )[..., expert_start:expert_end]
955
+ .permute(2, 0, 1)
956
+ .contiguous()
957
+ )
958
+
959
+ # Split the fused gate_up_proj output dim into separate w1 / w3
960
+ # halves. GPT-OSS interleaves them along the output dim, all other
961
+ # 3D MoE checkpoints we know about concatenate them.
962
+ intermediate_x2 = gate_up_b.shape[1]
963
+ if intermediate_x2 % 2 != 0:
964
+ raise ValueError(
965
+ "Expected gate_up_proj LoRA-B output dim to be 2 * intermediate, "
966
+ f"got {intermediate_x2}."
967
+ )
968
+ intermediate = intermediate_x2 // 2
969
+ base_arch = self.model.config.architectures[0]
970
+ if base_arch == "GptOssForCausalLM":
971
+ w1_b = gate_up_b[:, ::2, :].contiguous()
972
+ w3_b = gate_up_b[:, 1::2, :].contiguous()
973
+ else:
974
+ w1_b = gate_up_b[:, :intermediate, :].contiguous()
975
+ w3_b = gate_up_b[:, intermediate:, :].contiguous()
976
+
977
+ # In the 3D layout w1 and w3 share the same rank-r mid
978
+ # representation, so they reuse the same lora_a tensor. The 2D
979
+ # wrapper's set_lora copies whatever it gets here into independent
980
+ # per-slice buffers, so the sharing is purely a CPU-side memory
981
+ # optimization and does not affect numerics.
982
+ down_proj_lora.lora_a = [gate_up_a, down_a, gate_up_a]
983
+ down_proj_lora.lora_b = [w1_b, down_b, w3_b]
984
+ # Drop the redundant base_layer entry to avoid double pin_memory
985
+ # and to keep the activation path looking up only the wrapper key.
986
+ lora_model.loras.pop(module_name + ".base_layer", None)
987
+
988
def _slice_moe_lora_ep(
989
self,
990
lora_model: LoRAModel,
vllm/lora/request.py
+7
@@ -28,6 +28,13 @@ class LoRARequest(
28
base_model_name: str | None = msgspec.field(default=None)
29
tensorizer_config_dict: dict | None = None
30
load_inplace: bool = False
31
+ is_3d_lora_weight: bool = False
32
+ """Whether this adapter's MoE weights are stored in the 3D fused
33
+ `gate_up_proj` / `down_proj` layout (one fused tensor per layer) or the
34
+ 2D per-expert split layout (separate `gate_proj` / `up_proj` / `down_proj`
35
+ tensors per expert). Only consulted when the engine is started with
36
+ `enable_mixed_moe_lora_format=True`; otherwise it is ignored and the
37
+ on-disk format is inferred from the base model."""
38
39
def __post_init__(self):
40
if self.lora_int_id < 1:
vllm/lora/utils.py
+8
-3
@@ -355,7 +355,9 @@ def get_adapter_absolute_path(lora_path: str) -> str:
355
return local_snapshot_path
356
357
358
-def process_packed_modules_mapping(model: nn.Module) -> dict[str, list[str]]:
358
+def process_packed_modules_mapping(
359
+ model: nn.Module, force_2d_moe: bool = False
360
+) -> dict[str, list[str]]:
361
if is_moe_model(model):
362
if moe_packed_mapping := get_moe_expert_mapping(model):
363
# This method generates and returns a dictionary mapping packed module
@@ -364,8 +366,11 @@ def process_packed_modules_mapping(model: nn.Module) -> dict[str, list[str]]:
366
# the expert indices are expanded based on the configured number
367
# of routed experts.
368
packed_modules_mapping = get_packed_modules_mapping(model)
367
- if not model.is_3d_moe_weight:
368
- # 3D MoE LoRA does not need `packed_modules_mapping`
369
+ # The 2D mapping is needed when the model itself is 2D, or when
370
+ # the engine forces the universal 2D wrapper via
371
+ # enable_mixed_moe_lora_format (so 3D models can also load 2D
372
+ # adapters through FusedMoEWithLoRA).
373
+ if (not model.is_3d_moe_weight) or force_2d_moe:
374
# Filter out malformed entries: non-gated MoE has empty
375
# ckpt_up_proj_name which results in weight_name containing ".."
376
# (e.g., "experts.0.." instead of "experts.0.layer_name.")
vllm/lora/worker_manager.py
+4
@@ -141,6 +141,10 @@ class WorkerLoRAManager:
141
weights_mapper=hf_to_vllm_mapper,
142
skip_prefixes=lora_skip_prefixes,
143
)
144
+ # Stamp the on-disk MoE layout onto the loaded model so the
145
+ # adapter manager can route 3D-format checkpoints through the
146
+ # 3D->2D conversion when running under the universal 2D wrapper.
147
+ lora.is_3d_lora_weight = lora_request.is_3d_lora_weight
148
149
except FileNotFoundError as e:
150
# FileNotFoundError should be raised if both