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 12, 17:15 UTC
[LoRA] Cleanup FusedMoEWithLoRA (#29187)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Nov 22, 2025 at 12:01 UTC
1489902b531bb649f8110c94572b2d8b753a72cc
3 files changed
+94
-99
vllm/lora/layers/fused_moe.py
+90
-95
@@ -42,6 +42,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
42
self.tp_size = get_tensor_model_parallel_world_size()
43
self.tp_rank = get_tensor_model_parallel_rank()
44
self.device = base_layer.w2_weight.device
45
+ self.w13_slices = 2
46
self._inject_lora_into_fused_moe()
47
48
def _normalize_keys(self, config: dict[str, int | None]) -> dict[str, int | None]:
@@ -60,8 +61,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
61
def _get_lora_moe_configs(
62
self,
63
op_prefix: str,
63
- lora_a_stacked: torch.Tensor,
64
- lora_b_stacked: torch.Tensor,
64
+ num_loras: int,
65
+ rank: int,
66
num_slices: int,
67
M: int,
68
layer: FusedMoE,
@@ -69,23 +70,25 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
70
config_dtype: str,
71
):
72
if envs.VLLM_TUNED_CONFIG_FOLDER:
73
+ hidden_size = layer.hidden_size
74
+ intermediate_size = layer.intermediate_size_per_partition
75
shrink_config = get_lora_op_configs(
76
op_type=f"fused_moe_lora_{op_prefix}_shrink",
74
- max_loras=lora_a_stacked.shape[0],
77
+ max_loras=num_loras,
78
batch=M,
76
- hidden_size=lora_a_stacked.shape[-1],
77
- rank=lora_a_stacked.shape[-2],
79
+ hidden_size=hidden_size,
80
+ rank=rank,
81
num_slices=num_slices,
79
- moe_intermediate_size=lora_b_stacked.shape[-2],
82
+ moe_intermediate_size=intermediate_size,
83
)
84
expand_config = get_lora_op_configs(
85
op_type=f"fused_moe_lora_{op_prefix}_expand",
83
- max_loras=lora_a_stacked.shape[0],
86
+ max_loras=num_loras,
87
batch=M,
85
- hidden_size=lora_a_stacked.shape[-1],
86
- rank=lora_a_stacked.shape[-2],
88
+ hidden_size=hidden_size, # lora_a_stacked.shape[-1],
89
+ rank=rank,
90
num_slices=num_slices,
88
- moe_intermediate_size=lora_b_stacked.shape[-2],
91
+ moe_intermediate_size=intermediate_size, # lora_b_stacked.shape[-2],
92
)
93
else: # fall back to the default config
94
get_config_func = functools.partial(
@@ -152,12 +155,12 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
155
CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
156
num_tokens = hidden_states.size(0)
157
M = min(num_tokens, CHUNK_SIZE)
155
-
158
+ max_lora_rank = self.w13_lora_a_stacked[0].shape[-2]
159
shrink_config, expand_config = self._get_lora_moe_configs(
160
op_prefix="w13",
158
- lora_a_stacked=self.w1_lora_a_stacked,
159
- lora_b_stacked=self.w1_lora_b_stacked,
160
- num_slices=2,
161
+ num_loras=self.max_loras,
162
+ rank=max_lora_rank,
163
+ num_slices=self.w13_slices,
164
M=M,
165
layer=layer,
166
top_k=top_k,
@@ -165,7 +168,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
168
)
169
170
# get the block size of m from customized config or default config
168
- max_loras = self.w1_lora_a_stacked.shape[0]
171
(
172
sorted_token_ids_lora,
173
expert_ids_lora,
@@ -175,7 +177,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
177
num_tokens,
178
shrink_config["BLOCK_SIZE_M"],
179
self.base_layer.local_num_experts,
178
- max_loras,
180
+ self.max_loras,
181
self.adapter_enabled,
182
expert_map,
183
)
@@ -186,17 +188,15 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
188
num_tokens_post_padded_lora
189
)
190
189
- w13_lora_a_stacked = [self.w1_lora_a_stacked, self.w3_lora_a_stacked]
190
- w13_lora_b_stacked = [self.w1_lora_b_stacked, self.w3_lora_b_stacked]
191
- max_lora_rank = self.w1_lora_a_stacked.shape[-2]
192
- expert_ids_lora = expert_ids_lora.view(max_loras, -1)
193
- sorted_token_ids_lora = sorted_token_ids_lora.view(max_loras, -1)
191
+ expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
192
+ sorted_token_ids_lora = sorted_token_ids_lora.view(self.max_loras, -1)
193
+ #
194
195
self.punica_wrapper.add_lora_fused_moe(
196
input.view(-1, top_k, input.shape[-1]),
197
hidden_states,
198
- w13_lora_a_stacked,
199
- w13_lora_b_stacked,
198
+ self.w13_lora_a_stacked,
199
+ self.w13_lora_b_stacked,
200
topk_weights,
201
sorted_token_ids_lora,
202
expert_ids_lora,
@@ -230,11 +230,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
230
CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
231
num_tokens = hidden_states.size(0)
232
M = min(num_tokens, CHUNK_SIZE)
233
-
233
+ max_lora_rank = self.w2_lora_a_stacked.shape[-2]
234
shrink_config, expand_config = self._get_lora_moe_configs(
235
op_prefix="w2",
236
- lora_a_stacked=self.w2_lora_a_stacked,
237
- lora_b_stacked=self.w2_lora_b_stacked,
236
+ num_loras=self.max_loras,
237
+ rank=max_lora_rank,
238
num_slices=1,
239
M=M,
240
layer=layer,
@@ -247,20 +247,19 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
247
num_tokens_post_padded_lora = moe_state_dict[
248
"num_tokens_post_padded_lora"
249
]
250
- max_loras = self.w1_lora_a_stacked.shape[0]
251
- expert_ids_lora = expert_ids_lora.view(max_loras, -1)
252
- sorted_token_ids_lora = sorted_token_ids_lora.view(max_loras, -1)
250
+
251
+ expert_ids_lora = expert_ids_lora.view(self.max_loras, -1)
252
+ sorted_token_ids_lora = sorted_token_ids_lora.view(self.max_loras, -1)
253
intermediate_cache2 = moe_state_dict["intermediate_cache2"]
254
intermediate_cache3 = args[0]
255
- max_lora_rank = self.w2_lora_a_stacked.shape[-2]
255
256
shard_size_w2 = divide(self.base_layer.hidden_size, self.tp_size)
257
258
self.punica_wrapper.add_lora_fused_moe(
259
intermediate_cache3,
260
intermediate_cache2,
262
- [self.w2_lora_a_stacked],
263
- [self.w2_lora_b_stacked],
261
+ (self.w2_lora_a_stacked,),
262
+ (self.w2_lora_b_stacked,),
263
topk_weights,
264
sorted_token_ids_lora,
265
expert_ids_lora,
@@ -289,7 +288,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
288
fused_experts.moe_sum = moe_sum_decorator(
289
self.base_layer, fused_experts.moe_sum
290
)
292
-
291
self.base_layer.quant_method = FusedMoEModularMethod(
292
self.base_layer.quant_method, m_fused_moe_fn
293
)
@@ -301,33 +299,42 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
299
model_config: PretrainedConfig | None = None,
300
) -> None:
301
"""Initializes lora matrices."""
302
+ assert self.w13_slices == 2
303
+ self.max_loras = lora_config.max_loras
304
self.fully_sharded = lora_config.fully_sharded_loras
305
306
self.adapter_enabled = torch.tensor(
307
[0] * (max_loras + 1), dtype=torch.int, device=self.device
308
)
309
310
- self.w1_lora_a_stacked = torch.zeros(
311
- (
312
- max_loras,
313
- self.base_layer.local_num_experts,
314
- lora_config.max_lora_rank
315
- if not self.fully_sharded
316
- else divide(lora_config.max_lora_rank, self.tp_size),
317
- self.base_layer.hidden_size,
318
- ),
319
- dtype=lora_config.lora_dtype,
320
- device=self.device,
310
+ self.w13_lora_a_stacked = tuple(
311
+ torch.zeros(
312
+ (
313
+ max_loras,
314
+ self.base_layer.local_num_experts,
315
+ lora_config.max_lora_rank
316
+ if not self.fully_sharded
317
+ else divide(lora_config.max_lora_rank, self.tp_size),
318
+ self.base_layer.hidden_size,
319
+ ),
320
+ dtype=lora_config.lora_dtype,
321
+ device=self.device,
322
+ )
323
+ for _ in range(self.w13_slices)
324
)
322
- self.w1_lora_b_stacked = torch.zeros(
323
- (
324
- max_loras,
325
- self.base_layer.local_num_experts,
326
- self.base_layer.intermediate_size_per_partition,
327
- lora_config.max_lora_rank,
328
- ),
329
- dtype=lora_config.lora_dtype,
330
- device=self.device,
325
+
326
+ self.w13_lora_b_stacked = tuple(
327
+ torch.zeros(
328
+ (
329
+ max_loras,
330
+ self.base_layer.local_num_experts,
331
+ self.base_layer.intermediate_size_per_partition,
332
+ lora_config.max_lora_rank,
333
+ ),
334
+ dtype=lora_config.lora_dtype,
335
+ device=self.device,
336
+ )
337
+ for _ in range(self.w13_slices)
338
)
339
340
self.w2_lora_a_stacked = torch.zeros(
@@ -353,29 +360,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
360
device=self.device,
361
)
362
356
- self.w3_lora_a_stacked = torch.zeros(
357
- (
358
- max_loras,
359
- self.base_layer.local_num_experts,
360
- lora_config.max_lora_rank
361
- if not self.fully_sharded
362
- else divide(lora_config.max_lora_rank, self.tp_size),
363
- self.base_layer.hidden_size,
364
- ),
365
- dtype=lora_config.lora_dtype,
366
- device=self.device,
367
- )
368
- self.w3_lora_b_stacked = torch.zeros(
369
- (
370
- max_loras,
371
- self.base_layer.local_num_experts,
372
- self.base_layer.intermediate_size_per_partition,
373
- lora_config.max_lora_rank,
374
- ),
375
- dtype=lora_config.lora_dtype,
376
- device=self.device,
377
- )
378
-
363
# They will be used by 'LoRALayerWeights.create_dummy_lora_weights'
364
# to create a dummy LoRA weights.
365
self.lora_a_stacked = []
@@ -383,20 +367,28 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
367
for lora_id in range(max_loras):
368
for experts_id in range(self.base_layer.local_num_experts):
369
# gate_proj,down_proj,up_proj
386
- self.lora_a_stacked.append(self.w1_lora_a_stacked[lora_id][experts_id])
370
+ self.lora_a_stacked.append(
371
+ self.w13_lora_a_stacked[0][lora_id][experts_id]
372
+ )
373
self.lora_a_stacked.append(self.w2_lora_a_stacked[lora_id][experts_id])
388
- self.lora_a_stacked.append(self.w3_lora_a_stacked[lora_id][experts_id])
374
+ self.lora_a_stacked.append(
375
+ self.w13_lora_a_stacked[1][lora_id][experts_id]
376
+ )
377
390
- self.lora_b_stacked.append(self.w1_lora_b_stacked[lora_id][experts_id])
378
+ self.lora_b_stacked.append(
379
+ self.w13_lora_b_stacked[0][lora_id][experts_id]
380
+ )
381
self.lora_b_stacked.append(self.w2_lora_b_stacked[lora_id][experts_id])
392
- self.lora_b_stacked.append(self.w3_lora_b_stacked[lora_id][experts_id])
382
+ self.lora_b_stacked.append(
383
+ self.w13_lora_b_stacked[1][lora_id][experts_id]
384
+ )
385
386
def reset_lora(self, index: int):
387
"""Resets the lora weights at index back to 0."""
396
- self.w1_lora_a_stacked[index] = 0
397
- self.w1_lora_b_stacked[index] = 0
398
- self.w3_lora_a_stacked[index] = 0
399
- self.w3_lora_b_stacked[index] = 0
388
+ for pos in range(self.w13_slices):
389
+ self.w13_lora_a_stacked[pos][index] = 0
390
+ self.w13_lora_b_stacked[pos][index] = 0
391
+
392
self.w2_lora_a_stacked[index] = 0
393
self.w2_lora_b_stacked[index] = 0
394
self.adapter_enabled[index] = 0
@@ -434,7 +426,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
426
if self.fully_sharded:
427
# Based on S-LoRA, we slice W1 and W3 A along the rank dim,
428
# and W2 B along the hidden_size dim.
437
- w13_shard_size = self.w1_lora_a_stacked[index, eid].shape[0]
429
+ w13_shard_size = self.w13_lora_a_stacked[0][index, eid].shape[0]
430
w13_start_idx = self.tp_rank * w13_shard_size
431
w13_end_idx = (self.tp_rank + 1) * w13_shard_size
432
w1_lora_a = w1_lora_a[w13_start_idx:w13_end_idx, :]
@@ -444,29 +436,32 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
436
w2_start_idx = self.tp_rank * w2_shard_size
437
w2_end_idx = (self.tp_rank + 1) * w2_shard_size
438
w2_lora_b = w2_lora_b[w2_start_idx:w2_end_idx, :]
447
-
448
- self.w1_lora_a_stacked[
439
+ # w1 lora_a
440
+ self.w13_lora_a_stacked[0][
441
index, eid, : w1_lora_a.shape[0], : w1_lora_a.shape[1]
442
].copy_(w1_lora_a, non_blocking=True)
451
-
452
- self.w3_lora_a_stacked[
443
+ # w3 lora_a
444
+ self.w13_lora_a_stacked[1][
445
index, eid, : w3_lora_a.shape[0], : w3_lora_a.shape[1]
446
].copy_(w3_lora_a, non_blocking=True)
447
456
- self.w2_lora_b_stacked[
457
- index, eid, : w2_lora_b.shape[0], : w2_lora_b.shape[1]
458
- ].copy_(w2_lora_b, non_blocking=True)
459
-
460
- self.w1_lora_b_stacked[
448
+ # w1 lora_b
449
+ self.w13_lora_b_stacked[0][
450
index, eid, : w1_lora_b.shape[0], : w1_lora_b.shape[1]
451
].copy_(w1_lora_b, non_blocking=True)
463
- self.w3_lora_b_stacked[
452
+ # w3 lora_b
453
+ self.w13_lora_b_stacked[1][
454
index, eid, : w3_lora_b.shape[0], : w3_lora_b.shape[1]
455
].copy_(w3_lora_b, non_blocking=True)
456
+
457
self.w2_lora_a_stacked[
458
index, eid, : w2_lora_a.shape[0], : w2_lora_a.shape[1]
459
].copy_(w2_lora_a, non_blocking=True)
460
461
+ self.w2_lora_b_stacked[
462
+ index, eid, : w2_lora_b.shape[0], : w2_lora_b.shape[1]
463
+ ].copy_(w2_lora_b, non_blocking=True)
464
+
465
@classmethod
466
def can_replace_layer(
467
cls,
vllm/lora/punica_wrapper/punica_base.py
+2
-2
@@ -470,8 +470,8 @@ class PunicaWrapperBase(PunicaWrapperABC):
470
self,
471
y: torch.Tensor,
472
x: torch.Tensor,
473
- lora_a_stacked: list[torch.Tensor],
474
- lora_b_stacked: list[torch.Tensor],
473
+ lora_a_stacked: tuple[torch.Tensor, ...],
474
+ lora_b_stacked: tuple[torch.Tensor, ...],
475
topk_weights: torch.Tensor,
476
sorted_token_ids: torch.Tensor,
477
expert_ids: torch.Tensor,
vllm/lora/punica_wrapper/punica_gpu.py
+2
-2
@@ -360,8 +360,8 @@ class PunicaWrapperGPU(PunicaWrapperBase):
360
self,
361
y: torch.Tensor,
362
x: torch.Tensor,
363
- lora_a_stacked: list[torch.Tensor],
364
- lora_b_stacked: list[torch.Tensor],
363
+ lora_a_stacked: tuple[torch.Tensor, ...],
364
+ lora_b_stacked: tuple[torch.Tensor, ...],
365
topk_weights: torch.Tensor,
366
sorted_token_ids: torch.Tensor,
367
expert_ids: torch.Tensor,