[ROCm][MiniMax-M3.1] MXFP8 linear on AITER's tuned FlyDSL GEMM

vllm/_aiter_ops.py

@@ -4427,6 +4427,23 @@ def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k_bytes: int) -> bool:
         except (AssertionError, ImportError):
             return False
 
+    @staticmethod
+    @functools.cache
+    def is_mxfp8_bpreshuffle_gemm_tuned(n: int, k: int) -> bool:
+        """Whether gemm_a8w8_blockscale_bpreshuffle has tuned 1x32 E8M0 rows
+        for (N, K) on the current arch."""
+        if not current_platform.is_rocm():
+            return False
+        try:
+            from aiter.ops.gemm_op_a8w8 import get_mxscale_bpreshuffle_config
+
+            return all(
+                get_mxscale_bpreshuffle_config(m, n, k, "1x32") is not None
+                for m in (1, 8192)
+            )
+        except ImportError:
+            return False
+
     @staticmethod
     @functools.cache
     def is_asm_fp4_gemm_tuned(n: int, k_bytes: int) -> bool:

vllm/model_executor/kernels/linear/__init__.py

@@ -109,6 +109,7 @@
     Mxfp8LinearKernel,
     Mxfp8LinearLayerConfig,
 )
+from vllm.model_executor.kernels.linear.mxfp8.aiter import AiterMxfp8LinearKernel
 from vllm.model_executor.kernels.linear.mxfp8.b12x import (
     B12xMxfp8LinearKernel,
 )
@@ -360,6 +361,7 @@ def _batch_invariant_fp8_kernel(kernel, config, quantization: str):
         AiterFp8BlockScaledMMKernel,
         AiterPerTokenFp8ScaledMMLinearKernel,
         AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
+        AiterMxfp8LinearKernel,
         AiterMxfp4LinearKernel,
     },
     "machete": {
@@ -577,6 +579,9 @@ def _resolve_backend_kernels(
         TorchMxFp8LinearKernel,
     ],
     PlatformEnum.ROCM: [
+        # AITER's FlyDSL GEMM; can_implement() claims only the (N, K) shapes
+        # AITER has tuned, the rest fall through to the next kernel.
+        AiterMxfp8LinearKernel,
         # Native CDNA4 (gfx950) MX linear; is_supported() gates to gfx95x and
         # falls through to BF16 emulation (hipBLASLt) elsewhere / on regression.
         RocmDotScaledMxfp8LinearKernel,

vllm/model_executor/kernels/linear/mxfp8/aiter.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import torch
+
+import vllm.envs as envs
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    mxfp8_e4m3_quantize,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+
+class AiterMxfp8LinearKernel(Mxfp8LinearKernel):
+    """MXFP8 linear on AITER's tuned FlyDSL GEMM (gfx950).
+
+    Weights are (16, 16)-preshuffled once at load; activation and weight
+    scales stay row-major E8M0, [M, K / 32] and [N, K / 32].
+    """
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "AITER is only supported on ROCm"
+        if not current_platform.supports_mx():
+            return False, "native MX requires CDNA4 (gfx95x)"
+        if envs.VLLM_BATCH_INVARIANT:
+            return False, "tuned buckets use split K; not batch invariant"
+        if not rocm_aiter_ops.is_linear_enabled():
+            return False, "AITER linear is disabled or not supported"
+        # MiniMax-M3 mono decode re-tiles the MXFP8 QKV, output and separate
+        # shared-expert weights from the row-major layout; only the shared
+        # experts check for preshuffled weights.
+        if os.getenv("VLLM_ROCM_MINIMAX_M3_MONO", "0") in ("1", "true", "True"):
+            return False, "VLLM_ROCM_MINIMAX_M3_MONO needs row-major MXFP8 weights"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, config: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        N, K = config.weight_shape
+        if config.bmm_batch_size is not None:
+            return False, "batched MXFP8 GEMM is not supported"
+        # The FlyDSL GEMM writes bf16 only.
+        if torch.get_default_dtype() != torch.bfloat16:
+            return False, "requires a bf16 model"
+        if not rocm_aiter_ops.is_mxfp8_bpreshuffle_gemm_tuned(N, K):
+            return False, f"AITER has no tuned MXFP8 config for N={N}, K={K}"
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data
+        N, K = weight.shape
+        weight_scale = layer.weight_scale.data[:N, : K // MXFP8_BLOCK_SIZE]
+        replace_parameter(
+            layer, "weight", rocm_aiter_ops.shuffle_weight(weight.contiguous())
+        )
+        replace_parameter(
+            layer,
+            "weight_scale",
+            weight_scale.contiguous().view(torch.float8_e8m0fnu),
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        out_shape = (*x.shape[:-1], layer.weight.shape[0])
+        x_q, x_scale = mxfp8_e4m3_quantize(x.reshape(-1, x.shape[-1]))
+        out = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
+            x_q,
+            layer.weight,
+            x_scale.view(torch.float8_e8m0fnu),
+            layer.weight_scale,
+            x.dtype,
+        )
+        out = out.reshape(out_shape)
+        if bias is not None:
+            out = out + bias
+        return out