vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe

Classes:

• AiterExperts –
• AiterFp8BlockExperts –

  128x128 block-quantized weight scale, 128-grouped dynamic

• AiterFp8ChannelExperts –

  Per-output-channel static weight scale, per-token dynamic

• AiterFp8TensorExperts –

  Per-tensor static weight scale, either per-tensor static or

• AiterMxfp4Experts –

  MXFP4-quantized weights, with either unquantized (bf16/fp16) or

• AiterUnquantizedExperts –

  Unquantized (bf16/fp16) weights and activations.

Functions:

• aiter_moe_pads –

  Round the padding deltas to what AITER's MoE dispatch expects.

• derive_aiter_quant_method –

  AiterExperts._AITER_QUANT_METHOD for callers that bypass the leaves.

• inject_shared_expert_weights –

  Merge routed topk results with the shared expert buffer and inject

• mxfp4_w4a16_extra_kwargs –

  AiterMxfp4Experts._extra_fused_moe_kwargs for callers that bypass it.

• rocm_aiter_fused_experts –

  ROCm AITER fused MoE expert computation.

AiterExperts

Bases: FusedMoEExpertsModular

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

class AiterExperts(mk.FusedMoEExpertsModular):
    consumes_expert_mask = True

    _AITER_QUANT_METHOD: QuantMethod
    _SUPPORTS_ROUTER_WEIGHT_ON_INPUT = True

    def _extra_fused_moe_kwargs(self) -> dict:
        """Leaf-specific extra kwargs for `rocm_aiter_fused_experts`."""
        return {}

    @property
    def expects_unquantized_inputs(self) -> bool:
        # When paired with MoRI, the prepare/finalize handles FP8 or MXFP4
        # quantization during dispatch to reduce network traffic,
        # so we should not defer input quantization.
        # Otherwise, AITER fused MoE kernels handle input quantization
        # internally via a single fused kernel.
        return not self.moe_config.use_mori_epv1_kernels

    @staticmethod
    def activation_format() -> mk.FusedMoEActivationFormat:
        return mk.FusedMoEActivationFormat.Standard

    @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 not is_supported and not rocm_aiter_ops.is_fused_moe_enabled():
            reason = (
                f"{reason}. AITER MoE is not enabled — "
                "set VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_MOE=1 "
                "to enable it"
            )
        return is_supported, reason

    @staticmethod
    def _supports_current_device() -> bool:
        return rocm_aiter_ops.is_fused_moe_enabled()

    @staticmethod
    def _supports_no_act_and_mul() -> bool:
        return False

    @staticmethod
    def _supports_activation(activation: MoEActivation) -> bool:
        return activation in [
            MoEActivation.SILU,
            MoEActivation.GELU,
            MoEActivation.SITU,
            MoEActivation.SWIGLUOAI,
            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
        ]

    @staticmethod
    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
        return not (
            moe_parallel_config.use_fi_nvl_two_sided_kernels
            or moe_parallel_config.use_fi_nvl_one_sided_kernels
        )

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
        return TopKWeightAndReduceNoOP()

    def workspace_shapes(
        self,
        M: int,
        N: int,
        K: int,
        topk: int,
        global_num_experts: int,
        local_num_experts: int,
        expert_tokens_meta: mk.ExpertTokensMetadata | None,
        activation: MoEActivation,
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
        # Workspaces are managed internally by AITER.
        workspace1 = (0,)
        workspace2 = (0,)
        # MoRI dispatches MXFP4 activations packed two per byte, so K is
        # half the hidden size there.
        packed_fp4_input = (
            self.moe_config.use_mori_epv1_kernels and self.quant_config.use_mxfp4_w4a4
        )
        output = (M, K * 2 if packed_fp4_input else K)
        return (workspace1, workspace2, output)

    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,
    ):
        # TODO(rob): rocm_aiter_fused_experts uses self.quant_config's
        # a_scales for static quantization. Update this to fit better
        # with the interface once all quant integrations are complete.
        assert self._SUPPORTS_ROUTER_WEIGHT_ON_INPUT or not (
            apply_router_weight_on_input
        ), f"{type(self).__name__} does not support apply_router_weight_on_input"

        if expert_tokens_meta is not None:
            num_local_tokens = expert_tokens_meta.expert_num_tokens
        else:
            num_local_tokens = None

        result = rocm_aiter_fused_experts(
            hidden_states=hidden_states,
            w1=w1,
            w2=w2,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            activation=activation,
            apply_router_weight_on_input=apply_router_weight_on_input,
            expert_mask=expert_map,
            quant_config=self.quant_config,
            moe_config=self.moe_config,
            quant_method=self._AITER_QUANT_METHOD,
            a1q_scale=a1q_scale,
            num_local_tokens=num_local_tokens,
            output_dtype=output.dtype,
            moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),
            **self._extra_fused_moe_kwargs(),
        )
        _write_moe_output(output, result)

_extra_fused_moe_kwargs()

Leaf-specific extra kwargs for rocm_aiter_fused_experts.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

def _extra_fused_moe_kwargs(self) -> dict:
    """Leaf-specific extra kwargs for `rocm_aiter_fused_experts`."""
    return {}

AiterFp8BlockExperts

Bases: AiterExperts

128x128 block-quantized weight scale, 128-grouped dynamic activation scale.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

class AiterFp8BlockExperts(AiterExperts):
    """128x128 block-quantized weight scale, 128-grouped dynamic
    activation scale."""

    _AITER_QUANT_METHOD = QuantMethod.BLOCK_128x128
    _SUPPORTS_ROUTER_WEIGHT_ON_INPUT = False

    @staticmethod
    def _supports_quant_scheme(
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        SUPPORTED_W_A = [(kFp8Static128BlockSym, kFp8Dynamic128Sym)]
        return (weight_key, activation_key) in SUPPORTED_W_A

AiterFp8ChannelExperts

Bases: AiterExperts

Per-output-channel static weight scale, per-token dynamic activation scale.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

class AiterFp8ChannelExperts(AiterExperts):
    """Per-output-channel static weight scale, per-token dynamic
    activation scale."""

    _AITER_QUANT_METHOD = QuantMethod.PER_TOKEN

    @staticmethod
    def _supports_quant_scheme(
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        SUPPORTED_W_A = [(kFp8StaticChannelSym, kFp8DynamicTokenSym)]
        return (weight_key, activation_key) in SUPPORTED_W_A

    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,
    ):
        if not apply_router_weight_on_input:
            return super().apply(
                output=output,
                hidden_states=hidden_states,
                w1=w1,
                w2=w2,
                topk_weights=topk_weights,
                topk_ids=topk_ids,
                activation=activation,
                global_num_experts=global_num_experts,
                expert_map=expert_map,
                a1q_scale=a1q_scale,
                a2_scale=a2_scale,
                workspace13=workspace13,
                workspace2=workspace2,
                expert_tokens_meta=expert_tokens_meta,
                apply_router_weight_on_input=apply_router_weight_on_input,
            )

        # tkw1 applies topk_weights to FC1's output instead of FC2's.
        assert expert_tokens_meta is None, (
            "AITER tkw1 kernel does not support `num_local_tokens`"
        )
        _assert_topk1_for_router_weight_on_input(topk_weights)
        activation_method, _ = _resolve_aiter_activation(activation)
        result = rocm_aiter_ops.asm_moe_tkw1(
            hidden_states,
            w1,
            w2,
            topk_weights.to(torch.float32),
            topk_ids.to(torch.int32),
            fc1_scale=self.quant_config.w1_scale,
            fc2_scale=self.quant_config.w2_scale,
            fc1_smooth_scale=None,
            fc2_smooth_scale=None,
            a16=False,
            per_tensor_quant_scale=None,
            expert_mask=expert_map,
            activation_method=activation_method,
        )
        _write_moe_output(output, result)

AiterFp8TensorExperts

Bases: AiterExperts

Per-tensor static weight scale, either per-tensor static or per-tensor dynamic activation scale.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

class AiterFp8TensorExperts(AiterExperts):
    """Per-tensor static weight scale, either per-tensor static or
    per-tensor dynamic activation scale."""

    _AITER_QUANT_METHOD = QuantMethod.PER_TENSOR

    @staticmethod
    def _supports_quant_scheme(
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        SUPPORTED_W_A = [
            (kFp8StaticTensorSym, kFp8StaticTensorSym),
            (kFp8StaticTensorSym, kFp8DynamicTensorSym),
        ]
        return (weight_key, activation_key) in SUPPORTED_W_A

AiterMxfp4Experts

Bases: AiterExperts

MXFP4-quantized weights, with either unquantized (bf16/fp16) or MXFP4-quantized dynamic activations. gfx950-only (not gfx1250), except for Kimi-K3's MXFP4 experts requantized to packed int4 at load on gfx942.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

class AiterMxfp4Experts(AiterExperts):
    """MXFP4-quantized weights, with either unquantized (bf16/fp16) or
    MXFP4-quantized dynamic activations. gfx950-only (not gfx1250), except for
    Kimi-K3's MXFP4 experts requantized to packed int4 at load on gfx942."""

    _AITER_QUANT_METHOD = QuantMethod.BLOCK_1X32

    @staticmethod
    def _supports_quant_scheme(
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        from vllm.platforms.rocm import on_gfx942, on_gfx950, on_gfx1250

        if (weight_key, activation_key) == (kInt4Static32, None):
            return on_gfx942()
        SUPPORTED_W_A = [
            (kMxfp4Static, None),
            (kMxfp4Static, kMxfp4Dynamic),
        ]
        if (weight_key, activation_key) not in SUPPORTED_W_A:
            return False
        return on_gfx950() and not on_gfx1250()

    def _extra_fused_moe_kwargs(self) -> dict:
        return mxfp4_w4a16_extra_kwargs(self.quant_config, self.moe_config)

AiterUnquantizedExperts

Bases: AiterExperts

Unquantized (bf16/fp16) weights and activations.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

class AiterUnquantizedExperts(AiterExperts):
    """Unquantized (bf16/fp16) weights and activations."""

    _AITER_QUANT_METHOD = QuantMethod.NO

    @staticmethod
    def _supports_quant_scheme(
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        SUPPORTED_W_A = [(None, None)]
        return (weight_key, activation_key) in SUPPORTED_W_A

_resolve_aiter_activation(activation)

Map a vLLM activation to AITER's activation_method, plus a gate/up interleave hint (only SWIGLUOAI_UNINTERLEAVE sets it).

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

def _resolve_aiter_activation(
    activation: MoEActivation,
) -> tuple["ActivationMethod | ActivationType", bool | None]:
    """Map a vLLM activation to AITER's activation_method, plus a gate/up
    interleave hint (only SWIGLUOAI_UNINTERLEAVE sets it)."""
    activation_interleave = None
    activation_method: ActivationMethod | ActivationType
    if activation == MoEActivation.SILU:
        activation_method = ActivationMethod.SILU
    elif activation == MoEActivation.GELU:
        activation_method = ActivationMethod.GELU
    elif activation == MoEActivation.SWIGLUOAI:
        activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu")
    elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
        activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu")
        activation_interleave = False
    elif activation == MoEActivation.SITU:
        activation_method = rocm_aiter_ops.get_aiter_activation_type("situ")
    else:
        raise ValueError(f"Unsupported activation: {activation}")
    if activation_method is None:
        raise ValueError(f"AITER does not support activation: {activation}")
    return activation_method, activation_interleave

aiter_moe_pads(hidden_pad, intermediate_pad, activation, use_mxfp4_w4a4, tp_size)

Round the padding deltas to what AITER's MoE dispatch expects.

SITU's A16W4 FlyDSL kernel pads per gate/up half, so its pads pass through unrounded. Otherwise the pads are rounded to AITER's CK/FlyDSL tile granularity, and at tp_size == 1 the intermediate pad is doubled for the fused gate/up weight. The W4A4 SiLU path doubles it inside AITER, so it is passed undoubled; doubling it here as well masks live intermediate rows.

Where AITER doubles it (v0.1.24.post1): - CK-Tile stage 1, aiter/fused_moe.py: n_pad_zeros=intermediate_pad // 64 * 64 * (2 if use_g1u1 else 1). - FlyDSL stage 1 takes inter_dim_pad unchanged and doubles it in the kernel, aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py: inter_dim_pad_total = 2 * inter_dim_pad is subtracted from the fused gate/up width. Stage 2 masks the K tail with the undoubled value (base_k < inter_dim - inter_dim_pad). So both stage 1 kernels expect the undoubled pad. Recheck these lines on an AITER bump.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

def aiter_moe_pads(
    hidden_pad: int,
    intermediate_pad: int,
    activation: MoEActivation,
    use_mxfp4_w4a4: bool,
    tp_size: int,
) -> tuple[int, int]:
    """Round the padding deltas to what AITER's MoE dispatch expects.

    SITU's A16W4 FlyDSL kernel pads per gate/up half, so its pads pass through
    unrounded. Otherwise the pads are rounded to AITER's CK/FlyDSL tile
    granularity, and at tp_size == 1 the intermediate pad is doubled for the
    fused gate/up weight. The W4A4 SiLU path doubles it inside AITER, so it is
    passed undoubled; doubling it here as well masks live intermediate rows.

    Where AITER doubles it (v0.1.24.post1):
    - CK-Tile stage 1, aiter/fused_moe.py:
      ``n_pad_zeros=intermediate_pad // 64 * 64 * (2 if use_g1u1 else 1)``.
    - FlyDSL stage 1 takes ``inter_dim_pad`` unchanged and doubles it in the
      kernel, aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py:
      ``inter_dim_pad_total = 2 * inter_dim_pad`` is subtracted from the fused
      gate/up width. Stage 2 masks the K tail with the undoubled value
      (``base_k < inter_dim - inter_dim_pad``).
    So both stage 1 kernels expect the undoubled pad. Recheck these lines on an
    AITER bump.
    """
    if activation == MoEActivation.SITU:
        return hidden_pad, intermediate_pad
    double = tp_size == 1 and not (use_mxfp4_w4a4 and activation == MoEActivation.SILU)
    return (
        hidden_pad // 128 * 128,
        intermediate_pad // 64 * 64 * (2 if double else 1),
    )

derive_aiter_quant_method(quant_config)

AiterExperts._AITER_QUANT_METHOD for callers that bypass the leaves.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

def derive_aiter_quant_method(quant_config: FusedMoEQuantConfig) -> QuantMethod:
    """`AiterExperts._AITER_QUANT_METHOD` for callers that bypass the leaves."""
    if quant_config.use_mxfp4_w4a4 or quant_config.use_mxfp4_w4a16:
        return QuantMethod.BLOCK_1X32
    if quant_config.block_shape is not None and quant_config.use_fp8_w8a8:
        return QuantMethod.BLOCK_128x128
    if quant_config.use_fp8_w8a8 and quant_config.per_out_ch_quant:
        return QuantMethod.PER_TOKEN
    if quant_config.use_fp8_w8a8:
        return QuantMethod.PER_TENSOR
    return QuantMethod.NO

inject_shared_expert_weights(topk_weights, topk_ids, topk, num_fused_shared_experts, shared_expert_weights=None)

Merge routed topk results with the shared expert buffer and inject dynamic per-token shared expert gate values for AITER fusion.

For routers that already return the combined buffer (e.g. GroupedTopKRouter via rocm_aiter_grouped_topk), only the dynamic weight injection is needed. For routers that return only routed slots (e.g. FusedTopKRouter), this also copies the routed results into the pre-allocated combined buffer.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

def inject_shared_expert_weights(
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    topk: int,
    num_fused_shared_experts: int,
    shared_expert_weights: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Merge routed topk results with the shared expert buffer and inject
    dynamic per-token shared expert gate values for AITER fusion.

    For routers that already return the combined buffer (e.g. GroupedTopKRouter
    via rocm_aiter_grouped_topk), only the dynamic weight injection is needed.
    For routers that return only routed slots (e.g. FusedTopKRouter), this also
    copies the routed results into the pre-allocated combined buffer.
    """
    if num_fused_shared_experts == 0:
        return topk_weights, topk_ids

    assert aiter_topK_meta_data is not None, (
        "aiter_topK_meta_data is not initialized but "
        "num_fused_shared_experts > 0. Ensure init_aiter_topK_meta_data "
        "is called before routing."
    )

    total_topk_weights, total_topk_ids = aiter_topK_meta_data
    token = topk_weights.shape[0]

    assert total_topk_weights.shape[0] >= token, (
        f"AITER topK meta data supports {total_topk_weights.shape[0]} "
        f"tokens, but got {token} tokens."
    )

    total_topk_weights_slice = total_topk_weights[:token]
    total_topk_ids_slice = total_topk_ids[:token]

    if topk_weights.shape[1] == topk:
        total_topk_weights_slice[:, :topk] = topk_weights
        total_topk_ids_slice[:, :topk] = topk_ids
        topk_weights = total_topk_weights_slice
        topk_ids = total_topk_ids_slice

    if shared_expert_weights is not None:
        topk_weights[:, topk : topk + num_fused_shared_experts] = shared_expert_weights[
            :token
        ]

    return topk_weights, topk_ids

mxfp4_w4a16_extra_kwargs(quant_config, moe_config)

AiterMxfp4Experts._extra_fused_moe_kwargs for callers that bypass it.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

def mxfp4_w4a16_extra_kwargs(
    quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig
) -> dict:
    """`AiterMxfp4Experts._extra_fused_moe_kwargs` for callers that bypass it."""
    if not quant_config.use_mxfp4_w4a16:
        return {}

    from aiter.ops.flydsl.moe_common import GateMode

    # Opt-in DeepSeek V4.1 a4w4 weights are shuffled SEPARATED at load time.
    if moe_config.use_mxfp4_w4a4_dsv4:
        from aiter import dtypes

        return {
            "gate_mode_override": GateMode.SEPARATED.value,
            "q_dtype_a": dtypes.fp4x2,
            "bias1": quant_config.w1_bias,
            "bias2": quant_config.w2_bias,
        }
    return {
        "gate_mode_override": GateMode.INTERLEAVE.value,
        "bias1": quant_config.w1_bias,
        "bias2": quant_config.w2_bias,
    }

rocm_aiter_fused_experts(hidden_states, w1, w2, topk_weights, topk_ids, moe_config, quant_method, activation=MoEActivation.SILU, apply_router_weight_on_input=False, expert_mask=None, quant_config=None, a1q_scale=None, num_local_tokens=None, output_dtype=None, moe_sorting_dispatch_policy=0, gate_mode_override=None, bias1=None, bias2=None, q_dtype_a=None, shared_w1=None, shared_w2=None, shared_w1_scale=None, shared_w2_scale=None, shared_expert_id=-1)

ROCm AITER fused MoE expert computation.

Source code in vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

def rocm_aiter_fused_experts(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    moe_config: FusedMoEConfig,
    quant_method: QuantMethod,
    activation: MoEActivation = MoEActivation.SILU,
    apply_router_weight_on_input: bool = False,
    expert_mask: torch.Tensor | None = None,
    quant_config: FusedMoEQuantConfig | None = None,
    a1q_scale: torch.Tensor | None = None,
    num_local_tokens: torch.Tensor | None = None,
    output_dtype: torch.dtype | None = None,
    moe_sorting_dispatch_policy: int = 0,
    gate_mode_override: str | None = None,
    bias1: torch.Tensor | None = None,
    bias2: torch.Tensor | None = None,
    q_dtype_a: torch.dtype | None = None,
    shared_w1: torch.Tensor | None = None,
    shared_w2: torch.Tensor | None = None,
    shared_w1_scale: torch.Tensor | None = None,
    shared_w2_scale: torch.Tensor | None = None,
    shared_expert_id: int = -1,
) -> torch.Tensor:
    """ROCm AITER fused MoE expert computation."""
    if quant_config is None:
        quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

    activation_method, activation_interleave = _resolve_aiter_activation(activation)
    # All AITER Fused MoE kernels are expecting the following datatypes
    topk_weights = topk_weights.to(torch.float32)
    topk_ids = topk_ids.to(torch.int32)

    if apply_router_weight_on_input:
        _assert_topk1_for_router_weight_on_input(topk_weights)

    # Compute padding on-the-fly for CK MXFP4 kernels
    hidden_pad = 0
    intermediate_pad = 0
    assert moe_config.hidden_dim_unpadded is not None
    assert moe_config.intermediate_size_per_partition_unpadded is not None
    hidden_size = hidden_states.shape[1]
    if hidden_states.dtype == torch.float4_e2m1fn_x2:
        # Pre-quantized MXFP4 from MoRI dispatch, packed two per byte.
        hidden_size *= 2
    hidden_pad = hidden_size - moe_config.hidden_dim_unpadded
    intermediate_pad = (
        (
            moe_config.intermediate_size_per_partition
            - moe_config.intermediate_size_per_partition_unpadded
        )
        if moe_config.intermediate_pad is None
        else moe_config.intermediate_pad
    )

    hidden_pad, intermediate_pad = aiter_moe_pads(
        hidden_pad,
        intermediate_pad,
        activation,
        quant_config.use_mxfp4_w4a4,
        moe_config.tp_size,
    )

    # AITER's stage1 GEMM needs gate_mode to match how weights were shuffled
    # at load time (oracle/mxfp4.py).
    from aiter.ops.flydsl.moe_common import GateMode

    gate_mode = ""
    if activation == MoEActivation.SITU:
        gate_mode = (
            GateMode.INTERLEAVE.value
            if rocm_aiter_ops.is_fused_moe_situv2_gate_up_interleaved()
            else GateMode.SEPARATED.value
        )
    elif gate_mode_override is not None:
        gate_mode = gate_mode_override
    elif activation_interleave is not None:
        gate_mode = (
            GateMode.INTERLEAVE.value
            if activation_interleave
            else GateMode.SEPARATED.value
        )

    return rocm_aiter_ops.fused_moe(
        hidden_states,
        w1,
        w2,
        topk_weights,
        topk_ids,
        expert_mask=expert_mask,
        quant_method=quant_method.value,
        activation_method=activation_method,
        w1_scale=quant_config.w1_scale,
        w2_scale=quant_config.w2_scale,
        a1_scale=quant_config.a1_scale if a1q_scale is None else a1q_scale,
        a2_scale=quant_config.a2_scale,
        doweight_stage1=apply_router_weight_on_input,
        num_local_tokens=num_local_tokens,
        output_dtype=output_dtype,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
        gate_mode=gate_mode,
        bias1=bias1,
        bias2=bias2,
        moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
        swiglu_limit=(
            0.0 if moe_config.swiglu_limit is None else float(moe_config.swiglu_limit)
        ),
        beta=moe_config.activation_situ_beta,
        linear_beta=moe_config.activation_situ_linear_beta,
        shared_w1=shared_w1,
        shared_w2=shared_w2,
        shared_w1_scale=shared_w1_scale,
        shared_w2_scale=shared_w2_scale,
        shared_expert_id=shared_expert_id,
        q_dtype_a=q_dtype_a,
    )