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 14, 05:30 UTC
[Core] Handle MoE LoRA edge cases (#27335)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Oct 22, 2025 at 21:14 UTC
abf3db40ef0af73f53f00780371b1742b7cb6f1d
2 files changed
+7
-11
vllm/lora/layers/fused_moe.py
+7
-10
@@ -74,7 +74,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
74
moe_state_dict["apply_router_weight_on_input"] = kwargs[
75
"apply_router_weight_on_input"
76
]
77
- moe_state_dict["max_loras"] = layer.w1_lora_a_stacked.shape[0]
77
result = func(*args, **kwargs)
78
return result
79
@@ -89,7 +88,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
88
curr_topk_ids = moe_state_dict["topk_ids"]
89
global_num_experts = moe_state_dict["global_num_experts"]
90
expert_map = moe_state_dict["expert_map"]
92
- max_loras = moe_state_dict["max_loras"]
91
92
config_dtype = _get_config_dtype_str(
93
dtype=hidden_states.dtype,
@@ -110,6 +108,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
108
block_shape=layer.quant_method.moe_quant_config.block_shape,
109
)
110
111
+ max_loras = self.w1_lora_a_stacked.shape[0]
112
config = get_config_func(M)
113
(
114
sorted_token_ids_lora,
@@ -161,7 +160,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
160
def wrapper(*args, **kwargs):
161
hidden_states = moe_state_dict["hidden_states"]
162
topk_weights = moe_state_dict["topk_weights"]
164
- max_loras = moe_state_dict["max_loras"]
163
164
config_dtype = _get_config_dtype_str(
165
dtype=hidden_states.dtype,
@@ -189,7 +187,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
187
num_tokens_post_padded_lora = moe_state_dict[
188
"num_tokens_post_padded_lora"
189
]
192
-
190
+ max_loras = self.w1_lora_a_stacked.shape[0]
191
expert_ids_lora = expert_ids_lora.view(max_loras, -1)
192
sorted_token_ids_lora = sorted_token_ids_lora.view(max_loras, -1)
193
intermediate_cache2 = moe_state_dict["intermediate_cache2"]
@@ -305,12 +303,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
303
device=self.device,
304
)
305
308
- self.base_layer.w1_lora_a_stacked = self.w1_lora_a_stacked
309
- self.base_layer.w1_lora_b_stacked = self.w1_lora_b_stacked
310
- self.base_layer.w2_lora_a_stacked = self.w2_lora_a_stacked
311
- self.base_layer.w2_lora_b_stacked = self.w2_lora_b_stacked
312
- self.base_layer.w3_lora_a_stacked = self.w3_lora_a_stacked
313
- self.base_layer.w3_lora_b_stacked = self.w3_lora_b_stacked
306
# They will be used by 'LoRALayerWeights.create_dummy_lora_weights'
307
# to create a dummy LoRA weights.
308
self.lora_a_stacked = []
@@ -343,6 +335,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
335
embeddings_tensor: torch.Tensor | None,
336
bias: torch.Tensor | None = None,
337
):
338
+ self.reset_lora(index)
339
"""Overwrites lora tensors at index."""
340
for eid in range(len(lora_a) // 3):
341
w1_lora_a = lora_a[eid * 3]
@@ -352,6 +345,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
345
w2_lora_b = lora_b[eid * 3 + 1]
346
w3_lora_b = lora_b[eid * 3 + 2]
347
348
+ # Handle the case of adding LoRA to only a subset of experts
349
+ if w1_lora_a is None or w2_lora_a is None or w3_lora_a is None:
350
+ continue
351
+
352
if self.tp_size > 1:
353
shard_size = self.base_layer.intermediate_size_per_partition
354
start_idx = self.tp_rank * shard_size
vllm/lora/models.py
-1
@@ -426,7 +426,6 @@ class LoRAModelManager:
426
for module_name, module in self.modules.items():
427
module_lora = self._get_lora_layer_weights(lora_model, module_name)
428
if module_lora:
429
- module_lora.optimize()
429
# Note (gnovack) - If MOE lora weights are not split into
430
# num_experts chunks, we split them here
431
if isinstance(module, FusedMoEWithLoRA) and torch.is_tensor(