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 11, 21:15 UTC
[ROCm] [Fused Moe EP] Use binary expert mask for aiter fused moe kernel (#29773)
Signed-off-by: ZhiweiYan-96 <zhiwei.yan@amd.com>
Zhiwei committed
Dec 8, 2025 at 13:23 UTC
c6df05ebb49901510890aea7c7411a17a38c0a8f
2 files changed
+5
vllm/model_executor/layers/fused_moe/layer.py
+4
@@ -520,6 +520,10 @@ class FusedMoE(CustomOp):
520
self._init_aiter_shared_experts_topK_buffer(
521
vllm_config=vllm_config, dp_size=dp_size_
522
)
523
+ if self.use_ep and self.rocm_aiter_fmoe_enabled:
524
+ assert self.expert_mask is None or torch.all(
525
+ (expert_mask == 0) | (expert_mask == 1)
526
+ ), "Aiter Fused MoE kernel only supports expert_map with 0 and 1s."
527
528
assert intermediate_size % self.tp_size == 0
529
self.hidden_size = hidden_size
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1
@@ -633,6 +633,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
633
topk_ids=topk_ids,
634
activation=activation,
635
quant_config=self.moe_quant_config,
636
+ expert_map=expert_map,
637
)
638
else:
639
from vllm.model_executor.layers.fused_moe import fused_experts