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 13, 23:15 UTC
[Kernel] Enable FusedMoEModularKernel support bias (#27754)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Nov 1, 2025 at 10:05 UTC
bc4486d60962a0e251fe8d98fc26cfca206fd54c
2 files changed
+15
-30
vllm/lora/layers/fused_moe.py
+15
-28
@@ -15,9 +15,7 @@ from vllm.distributed.parallel_state import (
15
from vllm.lora.layers.base import BaseLayerWithLoRA
16
from vllm.model_executor.layers.fused_moe import FusedMoE
17
from vllm.model_executor.layers.fused_moe.config import (
18
- FUSED_MOE_UNQUANTIZED_CONFIG,
18
_get_config_dtype_str,
20
- mxfp4_w4a16_moe_quant_config,
19
)
20
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
21
modular_marlin_fused_moe,
@@ -26,13 +24,16 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
24
modular_triton_fused_moe,
25
try_get_optimal_moe_config,
26
)
29
-from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4Config
27
28
29
class FusedMoEWithLoRA(BaseLayerWithLoRA):
30
def __init__(self, base_layer: FusedMoE) -> None:
31
super().__init__()
32
self.base_layer = base_layer
33
+
34
+ assert not self.base_layer.use_ep, (
35
+ "EP support for Fused MoE LoRA is not implemented yet."
36
+ )
37
self.tp_size = get_tensor_model_parallel_world_size()
38
self.tp_rank = get_tensor_model_parallel_rank()
39
self.device = base_layer.w2_weight.device
@@ -42,17 +43,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
43
moe_state_dict = {}
44
top_k = self.base_layer.top_k
45
45
- if self.base_layer.quant_config is None:
46
- quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
47
- elif not isinstance(self.base_layer.quant_config, Mxfp4Config):
48
- quant_config = self.base_layer.quant_config
49
- else:
50
- quant_config = mxfp4_w4a16_moe_quant_config(
51
- w1_bias=self.base_layer.w13_bias,
52
- w2_bias=self.base_layer.w2_bias,
53
- w1_scale=self.base_layer.w13_weight_scale,
54
- w2_scale=self.base_layer.w2_weight_scale,
55
- )
46
+ self.base_layer.ensure_moe_quant_config_init()
47
+ quant_config = self.base_layer.quant_method.moe_quant_config
48
49
m_fused_moe_fn = (
50
modular_triton_fused_moe(
@@ -69,7 +61,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
61
moe_state_dict["hidden_states"] = kwargs["hidden_states"]
62
moe_state_dict["topk_ids"] = kwargs["topk_ids"]
63
moe_state_dict["topk_weights"] = kwargs["topk_weights"]
72
- moe_state_dict["global_num_experts"] = kwargs["global_num_experts"]
64
moe_state_dict["expert_map"] = kwargs["expert_map"]
65
moe_state_dict["apply_router_weight_on_input"] = kwargs[
66
"apply_router_weight_on_input"
@@ -86,7 +77,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
77
hidden_states = moe_state_dict["hidden_states"]
78
topk_weights = moe_state_dict["topk_weights"]
79
curr_topk_ids = moe_state_dict["topk_ids"]
89
- global_num_experts = moe_state_dict["global_num_experts"]
80
+
81
expert_map = moe_state_dict["expert_map"]
82
83
config_dtype = _get_config_dtype_str(
@@ -118,7 +109,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
109
curr_topk_ids,
110
num_tokens,
111
config["BLOCK_SIZE_M"],
121
- global_num_experts,
112
+ self.base_layer.local_num_experts,
113
max_loras,
114
expert_map,
115
)
@@ -236,14 +227,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
227
) -> None:
228
"""Initializes lora matrices."""
229
239
- assert not self.base_layer.use_ep, (
240
- "EP support for Fused MoE LoRA is not implemented yet."
241
- )
242
-
230
self.w1_lora_a_stacked = torch.zeros(
231
(
232
max_loras,
246
- self.base_layer.global_num_experts,
233
+ self.base_layer.local_num_experts,
234
lora_config.max_lora_rank,
235
self.base_layer.hidden_size,
236
),
@@ -253,7 +240,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
240
self.w1_lora_b_stacked = torch.zeros(
241
(
242
max_loras,
256
- self.base_layer.global_num_experts,
243
+ self.base_layer.local_num_experts,
244
self.base_layer.intermediate_size_per_partition,
245
lora_config.max_lora_rank,
246
),
@@ -264,7 +251,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
251
self.w2_lora_a_stacked = torch.zeros(
252
(
253
max_loras,
267
- self.base_layer.global_num_experts,
254
+ self.base_layer.local_num_experts,
255
lora_config.max_lora_rank,
256
self.base_layer.intermediate_size_per_partition,
257
),
@@ -274,7 +261,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
261
self.w2_lora_b_stacked = torch.zeros(
262
(
263
max_loras,
277
- self.base_layer.global_num_experts,
264
+ self.base_layer.local_num_experts,
265
self.base_layer.hidden_size,
266
lora_config.max_lora_rank,
267
),
@@ -285,7 +272,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
272
self.w3_lora_a_stacked = torch.zeros(
273
(
274
max_loras,
288
- self.base_layer.global_num_experts,
275
+ self.base_layer.local_num_experts,
276
lora_config.max_lora_rank,
277
self.base_layer.hidden_size,
278
),
@@ -295,7 +282,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
282
self.w3_lora_b_stacked = torch.zeros(
283
(
284
max_loras,
298
- self.base_layer.global_num_experts,
285
+ self.base_layer.local_num_experts,
286
self.base_layer.intermediate_size_per_partition,
287
lora_config.max_lora_rank,
288
),
@@ -308,7 +295,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
295
self.lora_a_stacked = []
296
self.lora_b_stacked = []
297
for lora_id in range(max_loras):
311
- for experts_id in range(self.base_layer.global_num_experts):
298
+ for experts_id in range(self.base_layer.local_num_experts):
299
# gate_proj,down_proj,up_proj
300
self.lora_a_stacked.append(self.w1_lora_a_stacked[lora_id][experts_id])
301
self.lora_a_stacked.append(self.w2_lora_a_stacked[lora_id][experts_id])
vllm/model_executor/layers/fused_moe/layer.py
-2
@@ -672,8 +672,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
672
apply_router_weight_on_input=apply_router_weight_on_input,
673
)
674
elif self.fused_experts is not None:
675
- if self.moe.has_bias:
676
- raise ValueError("FusedMoEModularKernel does not support bias.")
675
result = self.fused_experts(
676
hidden_states=x,
677
w1=layer.w13_weight,