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, 17:45 UTC
[MoE Refactor] Aiter Experts for BF16 MoE (#31542)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Yongye Zhu committed
Jan 5, 2026 at 14:52 UTC
776ca1e18784707a8e94f62086d2cbd7f42a20c5
1 file changed
+30
-39
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+30
-39
@@ -30,6 +30,9 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
30
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
31
MoEPrepareAndFinalizeNoEP,
32
)
33
+from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
34
+ AiterExperts,
35
+)
36
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
37
swap_w13_to_w31,
38
)
@@ -60,12 +63,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
63
super().__init__(moe)
64
65
self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
63
- if self.rocm_aiter_moe_enabled:
64
- from .rocm_aiter_fused_moe import rocm_aiter_fused_experts
65
-
66
- self.rocm_aiter_fused_experts = rocm_aiter_fused_experts
67
- else:
68
- self.rocm_aiter_fused_experts = None # type: ignore
66
67
# FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS
68
self.flashinfer_cutlass_moe_enabled = (
@@ -207,14 +204,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
204
layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
205
layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
206
210
- if self.rocm_aiter_moe_enabled:
211
- shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
212
- layer.w13_weight.data, layer.w2_weight.data
213
- )
214
-
215
- layer.w13_weight.data = shuffled_w13
216
- layer.w2_weight.data = shuffled_w2
217
-
207
if current_platform.is_xpu():
208
import intel_extension_for_pytorch as ipex
209
@@ -258,7 +247,21 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
247
layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
248
elif current_platform.is_cuda_alike():
249
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
261
- if self.flashinfer_cutlass_moe_enabled:
250
+ if self.rocm_aiter_moe_enabled:
251
+ shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
252
+ layer.w13_weight.data, layer.w2_weight.data
253
+ )
254
+ replace_parameter(layer, "w13_weight", shuffled_w13)
255
+ replace_parameter(layer, "w2_weight", shuffled_w2)
256
+
257
+ self.use_inplace = True
258
+ self.kernel = mk.FusedMoEModularKernel(
259
+ MoEPrepareAndFinalizeNoEP(),
260
+ AiterExperts(self.moe_quant_config),
261
+ shared_experts=None,
262
+ )
263
+
264
+ elif self.flashinfer_cutlass_moe_enabled:
265
self.use_inplace = False
266
# Swap halves to arrange as [w3; w1] (kernel expectation)
267
w13_weight = swap_w13_to_w31(layer.w13_weight.data)
@@ -315,30 +318,18 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
318
router_logits=router_logits,
319
)
320
318
- if self.rocm_aiter_moe_enabled:
319
- result = self.rocm_aiter_fused_experts(
320
- hidden_states=x,
321
- w1=layer.w13_weight,
322
- w2=layer.w2_weight,
323
- topk_weights=topk_weights,
324
- topk_ids=topk_ids,
325
- expert_map=layer.expert_map,
326
- activation=layer.activation,
327
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
328
- )
329
- else:
330
- result = self.kernel(
331
- hidden_states=x,
332
- w1=layer.w13_weight,
333
- w2=layer.w2_weight,
334
- topk_weights=topk_weights,
335
- topk_ids=topk_ids,
336
- inplace=self.use_inplace,
337
- activation=layer.activation,
338
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
339
- global_num_experts=layer.global_num_experts,
340
- expert_map=layer.expert_map,
341
- )
321
+ result = self.kernel(
322
+ hidden_states=x,
323
+ w1=layer.w13_weight,
324
+ w2=layer.w2_weight,
325
+ topk_weights=topk_weights,
326
+ topk_ids=topk_ids,
327
+ inplace=self.use_inplace,
328
+ activation=layer.activation,
329
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
330
+ global_num_experts=layer.global_num_experts,
331
+ expert_map=layer.expert_map,
332
+ )
333
334
return result
335