[HIP] [FlyDSL] add mxfp8 w_scale 1*32 (scalea/b nonpreshuffle)

aiter/aot/flydsl/gemm.py

@@ -20,6 +20,7 @@
   - ``flydsl_mxfp8_128_bpreshuffle_compute_wmma_*`` gfx1250 compute-bound mxfp8_128 kernels
   - ``flydsl_mxfp8_32_bpreshuffle_{,compute_}wmma_*`` the same kernels on 1x32 scales
   - ``flydsl_mxpsh_*``                        gfx950 MX-microscale preshuffle GEMM
+  - ``flydsl_mxfp8_8w_*``                     gfx950 8-wave mxfp8 (1x32 row scales) GEMM
   - ``flydsl_bmm_mxfp8_mfma_*``               gfx950 mxscale batched GEMM kernels
 
 Usage:
@@ -69,6 +70,11 @@
     compile_8wave_gemm,
     parse_8wave_kernel_name,
 )
+from aiter.ops.flydsl.gemm_a8w8_mxfp8_8wave import NAME_PREFIX as MXFP8_8WAVE_PREFIX
+from aiter.ops.flydsl.gemm_a8w8_mxfp8_8wave import (
+    compile_mxfp8_8wave_gemm,
+    parse_mxfp8_8wave_kernel_name,
+)
 from aiter.ops.flydsl.gemm_kernels import (
     SPLIT_K_SEMAPHORE_MAX_LEN,
     get_flydsl_hgemm_kernel_params,
@@ -112,6 +118,7 @@
 from aiter.ops.flydsl.mxfp8_bpreshuffle_gemm_gfx1250 import (
     parse_wmma_kernel_name as parse_mxfp8_128_wmma_kernel_name,
 )
+from aiter.ops.flydsl.mxscale_preshuffle_kernels import MXFP8_ROW_SCALES
 
 # Keep the default AOT coverage aligned with runtime config resolution.
 DEFAULT_CSVS = [
@@ -228,6 +235,9 @@ def parse_csv(csv_path: str):
     jobs = []
     seen = set()
     batched_skipped = set()
+    # The mxfp8 1x32 table's gfx950 rows name mxpsh kernels that read the
+    # unshuffled per-1x32 scales; every other table's mxpsh rows are blockscale.
+    mxfp8_table = "a8w8_mxfp8_bpreshuffle" in os.path.basename(csv_path)
 
     with open(csv_path, newline="") as f:
         reader = csv.DictReader(f)
@@ -278,12 +288,20 @@ def parse_csv(csv_path: str):
                 if params is not None:
                     params = dict(params)
                     params["kind"] = "mxscale_preshuffle"
-                    params["blockscale"] = True
+                    if mxfp8_table:
+                        params.update(dict(MXFP8_ROW_SCALES))
+                    else:  # the a8w8 blockscale contract
+                        params.update(a_scale_block="1x128", b_scale_block="128x128")
             elif kernel_name.startswith("flydsl_bpreshuffle_8w_"):
                 params = parse_8wave_kernel_name(kernel_name)
                 if params is not None:
                     params = dict(params)
                     params["kind"] = "8wave"
+            elif kernel_name.startswith(f"{MXFP8_8WAVE_PREFIX}_"):
+                params = parse_mxfp8_8wave_kernel_name(kernel_name)
+                if params is not None:
+                    params = dict(params)
+                    params["kind"] = "mxfp8_8wave"
             elif kernel_name.startswith(
                 (
                     f"{MXFP8_128_WMMA_PREFIX}_",
@@ -686,14 +704,20 @@ def _compile_mxscale_preshuffle_to_cache(
     waves_per_eu: int,
     xcd_swizzle: int = 0,
     split_k: int = 1,
-    blockscale: bool = True,
+    a_scale_block: str = "1x128",
+    b_scale_block: str = "128x128",
+    a_scale_layout: str = "mfma_tile",
+    b_scale_layout: str = "mfma_tile",
+    fused_split_k: bool = False,
     **kwargs,
 ):
     """gfx950 MX-microscale preshuffle GEMM (``flydsl_mxpsh_*``).
 
-    ``blockscale`` selects the scale format; the two modes are distinct
-    Constexprs and therefore distinct binaries, so each is its own job. Keep
-    the CSV's logical M so the M<=16 fused split-K specializations are also AOT'd.
+    Each operand's scale block and layout are distinct Constexprs and therefore
+    distinct binaries, so each combination is its own job: the a8w8 blockscale
+    tables' 1x128 / 128x128 mfma_tile, and gemm_a8w8_mxfp8_bpreshuffle's 1x32 row.
+    Keep the CSV's logical M so the M<=16 fused split-K specializations are also
+    AOT'd.
     """
     del kwargs
 
@@ -712,16 +736,28 @@ def _compile_mxscale_preshuffle_to_cache(
     a = torch.empty((m, a_bytes), device=dev, dtype=torch.uint8)
     b = torch.empty((n, b_bytes), device=dev, dtype=torch.uint8)
     out = torch.empty((m, n), device=dev, dtype=out_torch_dtype)
-    if not blockscale:
-        raise NotImplementedError(
-            "mxscale preshuffle AOT covers the a8w8 blockscale mode only. "
-            "Enabling the per-1x32 MX scale layout needs matching tuner and "
-            "runner support (a4w4 / a6w4), not just dummy operands here."
-        )
     rows = (m + 31) // 32 * 32
     K1 = (k + 255) // 256
-    a_scale = torch.empty(rows * 2 * K1, device=dev, dtype=torch.uint8)
-    b_scale = torch.empty((n // 128) * 4 * K1, device=dev, dtype=torch.uint8)
+    # Dummy scale buffers of each operand's (block, layout) size.
+    shapes = {
+        ("1x128", "mfma_tile"): (rows * 2 * K1,),
+        ("128x128", "mfma_tile"): ((n // 128) * 4 * K1,),
+        ("1x32", "row"): None,  # per operand below: (rows, K // 32)
+    }
+    scales = []
+    for block, layout, nrows in (
+        (a_scale_block, a_scale_layout, m),
+        (b_scale_block, b_scale_layout, n),
+    ):
+        if (block, layout) not in shapes:
+            raise NotImplementedError(
+                f"mxscale preshuffle AOT covers the a8w8 blockscale and 1x32 row "
+                f"scales; got {block} {layout}. A 1x32 mfma_tile scale (a4w4 / "
+                f"a6w4) needs matching tuner and runner support first."
+            )
+        shape = shapes[(block, layout)] or (nrows, k // 32)
+        scales.append(torch.empty(shape, device=dev, dtype=torch.uint8))
+    a_scale, b_scale = scales
 
     with compile_only_env():
         flydsl_mxscale_preshuffle_gemm(
@@ -738,8 +774,12 @@ def _compile_mxscale_preshuffle_to_cache(
             waves_per_eu=waves_per_eu,
             xcd_swizzle=xcd_swizzle,
             split_k=split_k,
-            blockscale=blockscale,
             stream=fx.Stream(0),
+            a_scale_block=a_scale_block,
+            b_scale_block=b_scale_block,
+            a_scale_layout=a_scale_layout,
+            b_scale_layout=b_scale_layout,
+            fused_split_k=fused_split_k,
         )
 
 
@@ -777,6 +817,38 @@ def _compile_8wave_to_cache(
     _compile_executable_to_cache(exe, a, b, out, scale_a, scale_b, m, n, fx.Stream(0))
 
 
+def _compile_mxfp8_8wave_to_cache(
+    *,
+    m: int,
+    n: int,
+    k: int,
+    waves_per_eu: int,
+    xcd_swizzle: int,
+    **kwargs,
+):
+    """gfx950 8-wave mxfp8 GEMM (``flydsl_mxfp8_8w_*``) on 1x32 row scales,
+    with the dtypes flydsl_mxfp8_8wave_gemm launches it with."""
+    del kwargs
+
+    import torch
+
+    dev = torch.device("cpu")
+    a = torch.empty((m * k,), device=dev, dtype=torch.uint8)
+    b = torch.empty((n * k,), device=dev, dtype=torch.uint8)
+    out = torch.empty((m * n,), device=dev, dtype=torch.bfloat16)
+    scale_a = torch.empty((m * k // 32,), device=dev, dtype=torch.uint8)
+    scale_b = torch.empty((n * k // 32,), device=dev, dtype=torch.uint8)
+
+    exe = compile_mxfp8_8wave_gemm(
+        K=k,
+        waves_per_eu=waves_per_eu,
+        xcd_swizzle=int(xcd_swizzle),
+        scale_preshuffled=False,
+    )
+    # The 8-wave launcher takes (A, B, C, ...), as _compile_8wave_to_cache.
+    _compile_executable_to_cache(exe, a, b, out, scale_a, scale_b, m, n, fx.Stream(0))
+
+
 def _compile_bmm_mfma_to_cache(
     *,
     kernel_name: str,
@@ -813,6 +885,7 @@ def _compile_mxfp8_wmma_to_cache(
     cluster_m: int,
     cluster_n: int,
     a_preshuffle: bool = False,
+    a_scale_preshuffled: bool = False,
     persistent_n_tiles: int = 1,
     fused_splitk: bool = True,
     scale_block: int = SCALE_BLOCK_SIZE,
@@ -832,8 +905,8 @@ def _compile_mxfp8_wmma_to_cache(
 
     dev = torch.device("cpu")
     k_blocks = (k + scale_block - 1) // scale_block
-    if scale_block == MX32_BLOCK_K:  # m32k4 / n32k4
-        a_rows, b_rows = (m + 31) // 32 * 32, n
+    if scale_block == MX32_BLOCK_K:  # row-major or shuffled A; shuffled B
+        a_rows, b_rows = (((m + 31) // 32) * 32 if a_scale_preshuffled else m), n
     else:
         a_rows, b_rows = m, (n + 127) // 128
     xq = torch.empty((m, k), device=dev, dtype=torch.uint8)
@@ -894,6 +967,7 @@ def _compile_mxfp8_wmma_to_cache(
                         persistent_n_tiles,
                         fused,
                         bounded_m,
+                        a_scale_preshuffled,
                     )
             else:
                 launch(
@@ -904,6 +978,7 @@ def _compile_mxfp8_wmma_to_cache(
                     0,
                     1,
                     a_preshuffle,
+                    a_scale_preshuffled,
                 )
         if split_k > 1:
             compile_gemm_a8w8_splitk_reduce(split_k=split_k, out_dtype_str="bf16")(
@@ -1040,6 +1115,8 @@ def compile_one_config(
                 _compile_mxscale_preshuffle_to_cache(m=m, n=n, k=k, **kwargs)
             elif kind == "8wave":
                 _compile_8wave_to_cache(m=m, n=n, k=k, **kwargs)
+            elif kind == "mxfp8_8wave":
+                _compile_mxfp8_8wave_to_cache(m=m, n=n, k=k, **kwargs)
             elif kind == "mxfp8_wmma":
                 _compile_mxfp8_wmma_to_cache(
                     kernel_name=kernel_name,