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 15, 03:30 UTC
[V1] [ROCm] Enable EP with AITER Fused MoE (#20270)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Jul 1, 2025 at 09:48 UTC
02cabff207ca68094a73ba21296c82cdbcb1d1a5
4 files changed
+15
-5
vllm/model_executor/layers/fused_moe/layer.py
+1
-1
@@ -646,13 +646,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
646
indices_type=self.topk_indices_dtype)
647
648
if self.rocm_aiter_moe_enabled:
649
- assert expert_map is None
649
return self.rocm_aiter_fused_experts(
650
hidden_states=x,
651
w1=layer.w13_weight,
652
w2=layer.w2_weight,
653
topk_weights=topk_weights,
654
topk_ids=topk_ids,
655
+ expert_map=expert_map,
656
activation=activation,
657
apply_router_weight_on_input=apply_router_weight_on_input)
658
else:
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+9
-2
@@ -315,7 +315,8 @@ def rocm_aiter_fused_experts(
315
w2_scale: Optional[torch.Tensor] = None,
316
a1_scale: Optional[torch.Tensor] = None,
317
a2_scale: Optional[torch.Tensor] = None,
318
- block_shape: Optional[list[int]] = None) -> torch.Tensor:
318
+ block_shape: Optional[list[int]] = None,
319
+ expert_map: Optional[torch.Tensor] = None) -> torch.Tensor:
320
321
activation_method = (ActivationMethod.SILU
322
if activation == "silu" else ActivationMethod.GELU)
@@ -323,6 +324,11 @@ def rocm_aiter_fused_experts(
324
topk_weights = topk_weights.to(torch.float32)
325
topk_ids = topk_ids.to(torch.int32)
326
327
+ if expert_map is not None:
328
+ expert_mask = (expert_map > -1).to(torch.int32)
329
+ else:
330
+ expert_mask = None
331
+
332
# w8a8 per-channel quantization
333
if per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:
334
# AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input`
@@ -346,7 +352,7 @@ def rocm_aiter_fused_experts(
352
fc2_smooth_scale=None,
353
a16=False,
354
per_tensor_quant_scale=None,
349
- expert_mask=None,
355
+ expert_mask=expert_mask,
356
activation_method=activation_method)
357
358
else:
@@ -378,6 +384,7 @@ def rocm_aiter_fused_experts(
384
w2,
385
topk_weights,
386
topk_ids,
387
+ expert_mask=expert_mask,
388
quant_method=quant_method,
389
activation_method=activation_method,
390
w1_scale=w1_scale,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+2
-1
@@ -633,7 +633,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
633
w1_scale=layer.w13_weight_scale,
634
w2_scale=layer.w2_weight_scale,
635
a1_scale=layer.w13_input_scale,
636
- a2_scale=layer.w2_input_scale)
636
+ a2_scale=layer.w2_input_scale,
637
+ expert_map=expert_map)
638
if self.use_marlin:
639
assert activation == "silu", (
640
f"{activation} not supported for Marlin MoE.")
vllm/model_executor/layers/quantization/fp8.py
+3
-1
@@ -442,6 +442,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
442
"""
443
444
def __init__(self, quant_config: Fp8Config):
445
+
446
from vllm.model_executor.layers.fused_moe import fused_experts
447
self.quant_config = quant_config
448
self.block_quant = self.quant_config.weight_block_size is not None
@@ -879,7 +880,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
880
if self.block_quant else layer.w2_weight_scale),
881
a1_scale=layer.w13_input_scale,
882
a2_scale=layer.w2_input_scale,
882
- block_shape=self.quant_config.weight_block_size)
883
+ block_shape=self.quant_config.weight_block_size,
884
+ expert_map=expert_map)
885
elif self.use_marlin:
886
assert activation == "silu", (
887
f"{activation} not supported for Marlin MoE.")