vllm.model_executor.kernels.linear.mxfp4.aiter

Classes:

• AiterMxfp4LinearKernel –

  AITER-based native MXFP4 GEMM kernel for ROCm.

Functions:

• gemm_with_dynamic_quant –

  With hadamard_rotate, the HIP quant of the ASM path rotates each

AiterMxfp4LinearKernel

Bases: MxFp4LinearKernel

AITER-based native MXFP4 GEMM kernel for ROCm.

Source code in vllm/model_executor/kernels/linear/mxfp4/aiter.py

class AiterMxfp4LinearKernel(MxFp4LinearKernel):
    """AITER-based native MXFP4 GEMM kernel for ROCm."""

    def __init__(self, config: MxFp4LinearLayerConfig) -> None:
        super().__init__(config)
        self.asm_gemm_supported = rocm_aiter_ops.is_asm_fp4_gemm_supported()
        self.use_asm_gemm = False
        self.hadamard_rotate = False
        self.out_dtype = torch.get_default_dtype()

    @classmethod
    def is_supported(
        cls, compute_capability: int | None = None
    ) -> tuple[bool, str | None]:
        if not current_platform.supports_mx():
            return False, "current platform does not support native MXFP4 computation"

        from vllm._aiter_ops import is_aiter_found_and_supported
        from vllm.model_executor.kernels.linear import _get_linear_backend

        linear_backend = _get_linear_backend(quantization="mxfp4")

        if (
            current_platform.is_rocm()
            and current_platform.supports_mx()
            and "AiterMxfp4LinearKernel" not in envs.VLLM_DISABLED_KERNELS
            and linear_backend == "auto"
            and not is_aiter_found_and_supported()
        ):
            logger.warning_once(
                "This platform supports native MXFP4 W4A4 MOE "
                "computation via AITER MOE backend, but AITER is not "
                "found or not supported. Consider installing AITER: "
                "https://github.com/ROCm/aiter."
            )

        if is_aiter_found_and_supported():
            return True, None
        return False, "AITER not found or not supported on the current platform"

    @classmethod
    def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
        if config.activation_quant_key != kMxfp4Dynamic:
            return False, "only supports MXFP4 dynamic activation"
        return True, None

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        # Use preshuffled ASM (P) and Triton (D) GEMM when tuned in aiter
        if self.asm_gemm_supported:
            n, k_bytes = layer.weight.shape
            self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_tuned(
                n, k_bytes
            ) and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(n, k_bytes)

        if self.use_asm_gemm and not _asm_fp4_scale_swizzle_supported(
            layer.weight_scale.data
        ):
            logger.warning_once(
                "AITER ASM FP4 GEMM requires weight_scale dims divisible by "
                "(%d, %d), but this layer has weight_scale shape %s. Falling "
                "back to the AITER Triton FP4 GEMM for this layer.",
                _ASM_FP4_SCALE_ROW_MULTIPLE,
                _ASM_FP4_SCALE_COL_MULTIPLE,
                tuple(layer.weight_scale.data.shape),
            )
            self.use_asm_gemm = False

        if self.use_asm_gemm:
            from aiter.ops.shuffle import shuffle_weight

            weight_scale = layer.weight_scale.data
            sm, sn = weight_scale.shape
            weight_scale = weight_scale.view(sm // 32, 2, 16, sn // 8, 2, 4, 1)
            weight_scale = weight_scale.permute(0, 3, 5, 2, 4, 1, 6).contiguous()
            weight_scale = weight_scale.view(sm, sn)
            replace_parameter(layer, "weight_scale", weight_scale)

            replace_parameter(
                layer, "weight", shuffle_weight(layer.weight.data, layout=(16, 16))
            )
        else:
            replace_parameter(
                layer, "weight_scale", layer.weight_scale.data.T.contiguous()
            )

    def fuse_input_hadamard(self, size: int) -> bool:
        # Only the HIP activation quant of the ASM path rotates, in 32-blocks.
        self.hadamard_rotate = self.use_asm_gemm and size == 32
        return self.hadamard_rotate

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        y = torch.ops.vllm.gemm_with_dynamic_quant(
            x,
            layer.weight,
            layer.weight_scale,
            self.use_asm_gemm,
            self.out_dtype,
            hadamard_rotate=self.hadamard_rotate,
        )
        if bias is not None:
            y = y + bias
        return y

gemm_with_dynamic_quant(x, weight, weight_scale, rocm_use_aiter_fp4_asm_gemm=False, out_dtype=torch.bfloat16, x_scales=None, hadamard_rotate=False)

With hadamard_rotate, the HIP quant of the ASM path rotates each 32-element block of x by the orthonormal Sylvester Hadamard matrix (block_hadamard_rotate); weight must be rotated the same way.

Source code in vllm/model_executor/kernels/linear/mxfp4/aiter.py

def gemm_with_dynamic_quant(
    x: torch.Tensor,
    weight: torch.Tensor,
    weight_scale: torch.Tensor,
    rocm_use_aiter_fp4_asm_gemm: bool = False,
    out_dtype: torch.dtype | None = torch.bfloat16,
    x_scales: torch.Tensor | None = None,
    hadamard_rotate: bool = False,
) -> torch.Tensor:
    """With ``hadamard_rotate``, the HIP quant of the ASM path rotates each
    32-element block of ``x`` by the orthonormal Sylvester Hadamard matrix
    (`block_hadamard_rotate`); ``weight`` must be rotated the same way."""
    assert not (hadamard_rotate and x_scales is not None), (
        "hadamard_rotate needs an unquantized x"
    )
    # Only AITER releases with the fused rotation take the argument.
    quant_kwargs = {"hadamard_rotate": True} if hadamard_rotate else {}
    from aiter.ops.triton.gemm_afp4wfp4 import (
        gemm_afp4wfp4,
        gemm_afp4wfp4_preshuffle,
    )
    from aiter.ops.triton.quant import dynamic_mxfp4_quant

    if rocm_use_aiter_fp4_asm_gemm:
        from aiter import gemm_a4w4, per_1x32_f4_quant_hip

    M = x.shape[0]
    N = weight.shape[0]
    K = weight.shape[1]
    if rocm_use_aiter_fp4_asm_gemm:
        # The preshuffle kernel's tuned configs split K in some M buckets.
        if (
            M <= 64
            and not envs.VLLM_BATCH_INVARIANT
            and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K)
        ):
            if x_scales is None:
                # use hip quant kernel for performance
                x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=M >= 32, **quant_kwargs)
            else:
                x_q = x
                x_s = x_scales

            if M >= 32:
                x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1)
            else:
                x_s = x_s[:M, ...].view(torch.uint8)

            y = torch.empty(M, N, device=x_q.device, dtype=out_dtype)
            gemm_afp4wfp4_preshuffle(
                x_q.view(torch.uint8),
                weight.view(torch.uint8).view(weight.shape[0] // 16, -1),
                x_s,
                weight_scale.view(torch.uint8).view(
                    weight_scale.shape[0] // 32, -1
                ),
                out_dtype,
                y,
            )
        else:
            if x_scales is None:
                # use hip quant kernel for performance
                x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True, **quant_kwargs)
            else:
                x_q = x
                x_s = x_scales

            y = gemm_a4w4(
                x_q,
                weight.view(x_q.dtype),
                x_s,
                weight_scale.view(x_s.dtype),
                dtype=out_dtype,
                bpreshuffle=True,
            )
        return y[:M]
    else:
        assert not hadamard_rotate, "hadamard_rotate needs the ASM GEMM path"
        if x_scales is None:
            x_q, x_s = dynamic_mxfp4_quant(x)
        else:
            x_q = x
            x_s = x_scales
        y = torch.empty(
            x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype
        )

        config = None
        if envs.VLLM_BATCH_INVARIANT:
            # The configs split K in some M buckets (up to M <= 256 for
            # tuned shapes). Keep the tile, drop the split.
            from aiter.ops.triton.gemm_afp4wfp4 import _get_config

            config, _ = _get_config(M, N, K)
            config = dict(config, NUM_KSPLIT=1)

        gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y, config)
        return y