Skip to content

vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe

MXFP8 (1x32 block, E8M0) MoE via AITER's FlyDSL two-stage grouped GEMM (gfx950); alternative to Mxfp8NativeTritonExperts. Routes through aiter.fused_moe (per_1x32, gate_mode=INTERLEAVE); weights are preshuffled in convert_to_fp8_moe_kernel_format.

Classes:

  • AiterMxfp8Experts –

    MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950).

AiterMxfp8Experts

Bases: Mxfp8TritonExpertsBase

MXFP8 MoE through AITER's FlyDSL two-stage grouped GEMM (gfx950).

Source code in vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
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:
        return current_platform.supports_mx() and rocm_aiter_ops.is_fused_moe_enabled()

    @staticmethod
    def _supports_parallel_config(moe_parallel_config) -> bool:
        return True

    @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 != MoEActivation.SWIGLUOAI_UNINTERLEAVE
        ):
            return False, (
                "kernel hardcodes SwiGLU-OAI activation and requires "
                f"activation={MoEActivation.SWIGLUOAI_UNINTERLEAVE.value}; "
                f"got activation={moe_config.activation.value}"
            )
        if is_supported 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,
        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)

        # 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=ActivationType.Swiglu.value,
            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,
            gate_mode=GateMode.INTERLEAVE.value,
            swiglu_limit=swiglu_limit,
            output_dtype=output.dtype,
        )
        output.copy_(out.to(output.dtype))