[gfx1250][FlyDSL] Unify&Rename GEMM kernels and refactor LDS load

aiter/aot/flydsl/gemm.py

@@ -14,7 +14,7 @@
   - ``flydsl_gemm2_*``                        split-K HGEMM kernels
   - ``flydsl_bpreshuflle_*``                  a8w8 preshuffle GEMM kernels
   - ``flydsl_bpreshuffle_wmma_*``             gfx1250 a8w8 ptpc GEMM kernels
-  - ``flydsl_blockscale_bpreshuffle_wmma_*``  gfx1250 a8w8 blockscale GEMM kernels
+  - ``flydsl_mxfp8_128_bpreshuffle_wmma_*``   gfx1250 mxfp8_128 GEMM kernels
 
 Usage:
     # Compile all unique FlyDSL GEMM kernels from default CSVs
@@ -48,9 +48,6 @@
     run_jobs_parallel,
 )
 from aiter.jit.core import AITER_CONFIGS
-from aiter.ops.flydsl.blockscale_bpreshuffle_gemm_gfx1250 import (
-    parse_wmma_kernel_name as parse_blockscale_wmma_kernel_name,
-)
 from aiter.ops.flydsl.bpreshuffle_gemm_gfx1250 import (
     parse_wmma_kernel_name as parse_ptpc_wmma_kernel_name,
 )
@@ -60,6 +57,9 @@
 )
 from aiter.ops.flydsl.kernels.hgemm_dispatch import compile_flydsl_hgemm_kernel
 from aiter.ops.flydsl.kernels.preshuffle_gemm import compile_preshuffle_gemm
+from aiter.ops.flydsl.mxfp8_128_bpreshuffle_gemm_gfx1250 import (
+    parse_wmma_kernel_name as parse_mxfp8_128_wmma_kernel_name,
+)
 
 # Keep the default AOT coverage aligned with runtime config resolution.
 DEFAULT_CSVS = [
@@ -153,11 +153,11 @@ def parse_csv(csv_path: str):
 
             if kernel_name.startswith("flydsl_bpreshuflle_"):
                 params = _parse_preshuffle_kernel_name(kernel_name)
-            elif kernel_name.startswith("flydsl_blockscale_bpreshuffle_wmma_"):
-                params = parse_blockscale_wmma_kernel_name(kernel_name)
+            elif kernel_name.startswith("flydsl_mxfp8_128_bpreshuffle_wmma_"):
+                params = parse_mxfp8_128_wmma_kernel_name(kernel_name)
                 if params is not None:
                     params = dict(params)
-                    params["kind"] = "blockscale_wmma"
+                    params["kind"] = "mxfp8_128_wmma"
             elif kernel_name.startswith("flydsl_bpreshuffle_wmma_"):
                 params = parse_ptpc_wmma_kernel_name(kernel_name)
                 if params is not None:
@@ -376,7 +376,7 @@ def _compile_preshuffle_to_cache(
     )
 
 
-def _compile_blockscale_wmma_to_cache(
+def _compile_mxfp8_128_wmma_to_cache(
     *,
     m: int,
     n: int,
@@ -396,9 +396,7 @@ def _compile_blockscale_wmma_to_cache(
 
     import torch
 
-    from aiter.ops.flydsl.kernels.gemm_a8w8_blockscale_gfx1250 import (
-        launch_gemm_a8w8_bsc_col,
-    )
+    from aiter.ops.flydsl.kernels.gemm_a8w8_gfx1250 import launch_gemm_a8w8
 
     dev = torch.device("cpu")
     k_blocks = (k + 127) // 128
@@ -410,7 +408,7 @@ def _compile_blockscale_wmma_to_cache(
     stream = fx.Stream(0)
 
     with compile_only_env():
-        launch_gemm_a8w8_bsc_col(
+        launch_gemm_a8w8(
             _ptr_view_safe(out),
             _ptr_view_safe(xq),
             _ptr_view_safe(wq),
@@ -432,6 +430,7 @@ def _compile_blockscale_wmma_to_cache(
             num_buffers,
             cluster_m,
             cluster_n,
+            True,
         )
 
 
@@ -455,7 +454,7 @@ def _compile_ptpc_wmma_to_cache(
 
     import torch
 
-    from aiter.ops.flydsl.kernels.gemm_a8w8_ptpc_gfx1250 import launch_gemm_a8w8_ptpc
+    from aiter.ops.flydsl.kernels.gemm_a8w8_gfx1250 import launch_gemm_a8w8
 
     dev = torch.device("cpu")
     xq = torch.empty((m, k), device=dev, dtype=torch.uint8)
@@ -466,7 +465,7 @@ def _compile_ptpc_wmma_to_cache(
     stream = fx.Stream(0)
 
     with compile_only_env():
-        launch_gemm_a8w8_ptpc(
+        launch_gemm_a8w8(
             _ptr_view_safe(out),
             _ptr_view_safe(xq),
             _ptr_view_safe(wq),
@@ -476,6 +475,7 @@ def _compile_ptpc_wmma_to_cache(
             stream,
             n,
             k,
+            0,
             xq.stride(0),
             out.stride(0),
             tile_m,
@@ -487,6 +487,7 @@ def _compile_ptpc_wmma_to_cache(
             num_buffers,
             cluster_m,
             cluster_n,
+            False,
         )
 
 
@@ -530,8 +531,8 @@ def compile_one_config(
                 _compile_hgemm_to_cache(m=m, n=n, k=k, **hgemm_kwargs)
             elif kind == "preshuffle":
                 _compile_preshuffle_to_cache(m=m, n=n, k=k, **kwargs)
-            elif kind == "blockscale_wmma":
-                _compile_blockscale_wmma_to_cache(m=m, n=n, k=k, **kwargs)
+            elif kind == "mxfp8_128_wmma":
+                _compile_mxfp8_128_wmma_to_cache(m=m, n=n, k=k, **kwargs)
             elif kind == "ptpc_wmma":
                 _compile_ptpc_wmma_to_cache(m=m, n=n, k=k, **kwargs)
             else:
@@ -583,7 +584,7 @@ def main():
 
     hgemm_jobs = [j for j in all_jobs if j["kind"] == "hgemm"]
     preshuffle_jobs = [j for j in all_jobs if j["kind"] == "preshuffle"]
-    blockscale_wmma_jobs = [j for j in all_jobs if j["kind"] == "blockscale_wmma"]
+    mxfp8_128_wmma_jobs = [j for j in all_jobs if j["kind"] == "mxfp8_128_wmma"]
     ptpc_wmma_jobs = [j for j in all_jobs if j["kind"] == "ptpc_wmma"]
 
     print("=" * 72)
@@ -593,7 +594,7 @@ def main():
         print(f"  CSV:              {csv_path}")
     print(f"  HGEMM jobs:       {len(hgemm_jobs)}")
     print(f"  Preshuffle jobs:  {len(preshuffle_jobs)}")
-    print(f"  Blockscale wmma jobs: {len(blockscale_wmma_jobs)}")
+    print(f"  MXFP8_128 wmma jobs: {len(mxfp8_128_wmma_jobs)}")
     print(f"  PTPC wmma jobs:   {len(ptpc_wmma_jobs)}")
     print(f"  Total jobs:       {len(all_jobs)}")
     print(f"  Cache dir:        {cache_dir}")
@@ -607,7 +608,7 @@ def main():
     print(f"\n--- Compiling {len(all_jobs)} kernels ---")
     results = run_jobs_parallel(
         compile_one_config,
-        hgemm_jobs + preshuffle_jobs + blockscale_wmma_jobs + ptpc_wmma_jobs,
+        hgemm_jobs + preshuffle_jobs + mxfp8_128_wmma_jobs + ptpc_wmma_jobs,
     )
 
     total_elapsed = time.time() - total_t0