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, 01:30 UTC
fix(rocm): Enable non-gated MoE (is_act_and_mul=False) support on ROCm (#32244)
Signed-off-by: rabi <ramishra@redhat.com>
Rabi Mishra committed
Jan 16, 2026 at 13:01 UTC
b66b0d6abb955f9209a0d88b1dc245f4c1c9ff98
2 files changed
+15
-4
vllm/model_executor/layers/fused_moe/layer.py
+8
-4
@@ -448,9 +448,13 @@ class FusedMoE(CustomOp):
448
)
449
450
# ROCm aiter shared experts fusion
451
- self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
451
+ # AITER only supports gated activations (silu/gelu), so disable it
452
+ # for non-gated MoE (is_act_and_mul=False)
453
+ self.rocm_aiter_fmoe_enabled = (
454
+ rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul
455
+ )
456
self.aiter_fmoe_shared_expert_enabled = (
453
- rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
457
+ rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and is_act_and_mul
458
)
459
460
self.num_fused_shared_experts = (
@@ -619,9 +623,9 @@ class FusedMoE(CustomOp):
623
# for heuristic purposes, so it must be initialized first.
624
self.quant_method: FusedMoEMethodBase = _get_quant_method()
625
622
- if not self.moe_config.is_act_and_mul and not current_platform.is_cuda():
626
+ if not self.moe_config.is_act_and_mul and not current_platform.is_cuda_alike():
627
raise NotImplementedError(
624
- "is_act_and_mul=False is supported only for CUDA for now"
628
+ "is_act_and_mul=False is supported only for CUDA and ROCm for now"
629
)
630
631
if self.enable_eplb and not self.quant_method.supports_eplb:
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+7
@@ -8,6 +8,7 @@ from torch.nn import Module
8
9
import vllm.envs as envs
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.logger import init_logger
13
from vllm.model_executor.custom_op import CustomOp
14
from vllm.model_executor.layers.fused_moe.config import (
@@ -56,6 +57,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
57
use_ep=self.moe.moe_parallel_config.use_ep,
58
use_dp=self.moe.moe_parallel_config.dp_size > 1,
59
)
60
+
61
+ # AITER only supports gated activations (silu/gelu), so disable it
62
+ # for non-gated MoE (is_act_and_mul=False)
63
+ self.rocm_aiter_moe_enabled = (
64
+ rocm_aiter_ops.is_fused_moe_enabled() and moe.is_act_and_mul
65
+ )
66
self.kernel: mk.FusedMoEModularKernel | None = None
67
68
@property