class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
"""MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950)."""
consumes_expert_mask = True
@property
def quant_dtype(self) -> torch.dtype | str | None:
return self.quant_config.quant_dtype
@property
def block_shape(self) -> list[int] | None:
return self.quant_config.block_shape
@property
def expects_unquantized_inputs(self) -> bool:
# aiter.fused_moe MXFP8-quantizes the activations internally.
return True
@staticmethod
def _supports_current_device() -> bool:
# ROCm ``supports_mx`` deliberately excludes gfx942, so gfx942 may
# import this module during backend discovery but never selects or
# executes the AITER MXFP8 experts path.
return (
current_platform.supports_mx()
and rocm_aiter_ops.is_fused_moe_enabled()
and rocm_aiter_ops.fused_moe_supports_gate_mode()
)
@staticmethod
def _supports_parallel_config(moe_parallel_config) -> bool:
return True
@staticmethod
def _supports_batch_invariance() -> bool:
return False
@staticmethod
def is_supported_config(
cls, moe_config, weight_key, activation_key, activation_format
):
is_supported, reason = super().is_supported_config(
cls, moe_config, weight_key, activation_key, activation_format
)
if is_supported and moe_config.activation not in (
MoEActivation.SILU,
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
):
return False, (
"kernel supports only SiLU or SwiGLU-OAI activation; got "
f"activation={moe_config.activation.value}"
)
if (
is_supported
and moe_config.activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE
and (
moe_config.swiglu_alpha is None
or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA)
or moe_config.swiglu_beta is None
or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA)
)
):
return False, (
"kernel hardcodes SwiGLU-OAI with "
f"alpha={_AITER_SWIGLU_ALPHA} and beta={_AITER_SWIGLU_BETA}; "
f"got swiglu_alpha={moe_config.swiglu_alpha} and "
f"swiglu_beta={moe_config.swiglu_beta}"
)
return is_supported, reason
def apply(
self,
output: torch.Tensor,
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
activation: MoEActivation,
global_num_experts: int,
expert_map: torch.Tensor | None,
a1q_scale: torch.Tensor | None,
a2_scale: torch.Tensor | None,
workspace13: torch.Tensor,
workspace2: torch.Tensor,
expert_tokens_meta: mk.ExpertTokensMetadata | None,
apply_router_weight_on_input: bool,
):
from aiter import ActivationType, QuantType
from aiter.ops.flydsl.moe_common import GateMode
from vllm._aiter_ops import rocm_aiter_ops
limit = self.quant_config.gemm1_clamp_limit
swiglu_limit = 0.0 if limit is None else float(limit)
if activation == MoEActivation.SILU:
activation_method = ActivationType.Silu.value
elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
activation_method = ActivationType.Swiglu.value
else:
activation_name = getattr(activation, "value", activation)
raise ValueError(
"AITER MXFP8 MoE supports only SiLU or packed SwiGLU-OAI; "
f"got activation={activation_name!r}"
)
# RoutedExperts.expert_map hands AITER experts the precomputed 0/1
# expert_mask (with trailing sentinel) instead of the vLLM expert_map.
expert_mask = expert_map
# Route through the graph-safe ``rocm_aiter_fused_moe`` custom op so the
# call is captured under HIP graphs / torch.compile (a direct
# ``aiter.fused_moe`` is opaque to the dispatcher). aiter requires FP32
# routing weights / INT32 ids.
out = rocm_aiter_ops.fused_moe(
hidden_states,
w1,
w2,
topk_weights.to(torch.float32),
topk_ids.to(torch.int32),
expert_mask=expert_mask,
activation_method=activation_method,
quant_method=QuantType.per_1x32.value,
doweight_stage1=apply_router_weight_on_input,
w1_scale=self.w1_scale_val,
w2_scale=self.w2_scale_val,
a1_scale=None,
a2_scale=None,
# MXFP8 w1 and its scales are preshuffled into AITER's GUGU layout
# during model loading. Declare that layout explicitly so AITER can
# dispatch the gfx950-optimized FlyDSL kernel instead of CK.
gate_mode=GateMode.INTERLEAVE.value,
swiglu_limit=swiglu_limit,
output_dtype=output.dtype,
)
output.copy_(out.to(output.dtype))