[Triton/Gluon] [GFX12] mxfp8 gemm

aiter/configs/model_configs/dsv4_a8w8_blockscale_bpreshuffle_tuned_gemm.csv

@@ -905,9 +905,9 @@ gfx1250,256,4096,4096,4096,flydsl,528,1,32.6565,flydsl_mxfp8_128_bpreshuffle_wmm
 gfx1250,256,8192,4096,4096,flydsl,528,1,58.3829,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,4708.19,2011.56,0.0
 gfx1250,256,16384,4096,4096,flydsl,528,1,101.7874,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,5401.02,2142.74,0.0
 gfx1250,256,32768,4096,4096,flydsl,528,1,199.2404,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,5518.52,2105.15,0.0
-gfx1250,256,16,7168,768,flydsl,4,1,3.2568,flydsl_mxfp8_128_bpreshuffle_wmma_t16x32x256_mw1_nw2_nb3_sk1_cm1_cn1,54.09,1764.52,0.0
+gfx1250,256,16,7168,768,triton,4,1,6.05,gluon,54.09,1764.52,0.0
 gfx1250,256,32,7168,768,flydsl,22,1,3.4209,flydsl_mxfp8_128_bpreshuffle_wmma_t16x64x256_mw1_nw4_nb3_sk1_cm1_cn1,102.99,1750.52,0.0
-gfx1250,256,64,7168,768,flydsl,131,1,3.6056,flydsl_mxfp8_128_bpreshuffle_wmma_t32x64x256_mw2_nw2_nb3_sk1_cm1_cn1,195.43,1794.9,0.0
+gfx1250,256,64,7168,768,triton,131,1,6.09,gluon,195.43,1794.9,0.0
 gfx1250,256,128,7168,768,flydsl,164,1,3.9758,flydsl_mxfp8_128_bpreshuffle_wmma_t32x128x256_mw2_nw2_nb3_sk1_cm1_cn1,354.47,1870.9,0.0
 gfx1250,256,256,7168,768,flydsl,407,1,4.4411,flydsl_mxfp8_128_bpreshuffle_wmma_t128x64x256_mw2_nw2_nb3_sk1_cm1_cn1,634.66,2110.21,0.0
 gfx1250,256,512,7168,768,flydsl,209,1,5.4108,flydsl_mxfp8_128_bpreshuffle_wmma_t128x128x256_mw2_nw2_nb3_sk1_cm1_cn1,1041.83,2446.64,0.0
@@ -917,9 +917,9 @@ gfx1250,256,4096,7168,768,flydsl,223,1,20.749,flydsl_mxfp8_128_bpreshuffle_wmma_
 gfx1250,256,8192,7168,768,flydsl,263,1,35.6411,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x256_mw2_nw2_nb2_sk1_cm1_cn1,2530.63,3626.07,0.0
 gfx1250,256,16384,7168,768,flydsl,529,1,65.7262,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x256_mw2_nw2_nb2_sk1_cm1_cn1,2744.55,3848.83,0.0
 gfx1250,256,32768,7168,768,flydsl,529,1,126.7629,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x256_mw2_nw2_nb2_sk1_cm1_cn1,2846.08,3947.79,0.0
-gfx1250,256,16,2048,7168,flydsl,29,1,6.3618,flydsl_mxfp8_128_bpreshuffle_wmma_t16x64x512_mw1_nw4_nb4_sk1_cm1_cn1,73.84,2335.86,0.0
+gfx1250,256,16,2048,7168,triton,29,1,9.72,gluon,73.84,2335.86,0.0
 gfx1250,256,32,2048,7168,flydsl,111,1,6.3816,flydsl_mxfp8_128_bpreshuffle_wmma_t32x32x512_mw2_nw2_nb4_sk1_cm1_cn1,147.22,2356.86,0.0
-gfx1250,256,64,2048,7168,flydsl,111,1,6.6749,flydsl_mxfp8_128_bpreshuffle_wmma_t32x32x512_mw2_nw2_nb4_sk1_cm1_cn1,281.51,2307.29,0.0
+gfx1250,256,64,2048,7168,triton,111,1,9.88,gluon,281.51,2307.29,0.0
 gfx1250,256,128,2048,7168,flydsl,111,1,7.4047,flydsl_mxfp8_128_bpreshuffle_wmma_t32x32x512_mw2_nw2_nb4_sk1_cm1_cn1,507.53,2177.25,0.0
 gfx1250,256,256,2048,7168,flydsl,140,1,9.4924,flydsl_mxfp8_128_bpreshuffle_wmma_t32x64x512_mw2_nw2_nb3_sk1_cm1_cn1,791.81,1850.29,0.0
 gfx1250,256,512,2048,7168,flydsl,2165,4,11.289,flydsl_mxfp8_128_bpreshuffle_wmma_t128x128x256_mw2_nw2_nb3_sk4_cm1_cn1,1331.6,1811.25,0.0
@@ -929,7 +929,7 @@ gfx1250,256,4096,2048,7168,flydsl,3172,1,26.7639,flydsl_mxfp8_128_bpreshuffle_co
 gfx1250,256,8192,2048,7168,flydsl,3157,1,40.2099,flydsl_mxfp8_128_bpreshuffle_compute_wmma_t256x256x128_mw2_nw2_nb2_sk1_cm1_cn4,5981.57,2659.91,0.0
 gfx1250,256,16384,2048,7168,flydsl,3157,1,76.353,flydsl_mxfp8_128_bpreshuffle_compute_wmma_t256x256x128_mw2_nw2_nb2_sk1_cm1_cn4,6300.16,2609.32,0.0
 gfx1250,256,32768,2048,7168,flydsl,3157,1,140.89,flydsl_mxfp8_128_bpreshuffle_compute_wmma_t256x256x128_mw2_nw2_nb2_sk1_cm1_cn4,6828.54,2723.96,0.0
-gfx1250,256,16,7168,4096,flydsl,33,1,5.0839,flydsl_mxfp8_128_bpreshuffle_wmma_t16x64x1024_mw1_nw4_nb3_sk1_cm1_cn1,184.8,5833.13,0.0
+gfx1250,256,16,7168,4096,triton,33,1,7.27,gluon,184.8,5833.13,0.0
 gfx1250,256,32,7168,4096,flydsl,113,1,5.3092,flydsl_mxfp8_128_bpreshuffle_wmma_t32x32x1024_mw2_nw2_nb3_sk1_cm1_cn1,353.92,5641.14,0.0
 gfx1250,256,64,7168,4096,flydsl,268,1,6.846,flydsl_mxfp8_128_bpreshuffle_wmma_t64x32x512_mw4_nw1_nb4_sk1_cm1_cn1,548.95,4460.97,0.0
 gfx1250,256,128,7168,4096,flydsl,283,1,7.8422,flydsl_mxfp8_128_bpreshuffle_wmma_t64x64x512_mw2_nw2_nb4_sk1_cm1_cn1,958.43,4044.71,0.0
@@ -941,9 +941,9 @@ gfx1250,256,4096,7168,4096,flydsl,262,1,56.9068,flydsl_mxfp8_128_bpreshuffle_wmm
 gfx1250,256,8192,7168,4096,flydsl,262,1,100.3682,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,4792.72,1796.93,0.0
 gfx1250,256,16384,7168,4096,flydsl,528,1,172.8295,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,5566.6,1917.21,0.0
 gfx1250,256,32768,7168,4096,flydsl,528,1,340.6371,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,5648.67,1859.28,0.0
-gfx1250,256,16,1536,7168,flydsl,29,1,6.296,flydsl_mxfp8_128_bpreshuffle_wmma_t16x64x512_mw1_nw4_nb4_sk1_cm1_cn1,55.96,1774.76,0.0
+gfx1250,256,16,1536,7168,triton,29,1,8.84,gluon,55.96,1774.76,0.0
 gfx1250,256,32,1536,7168,flydsl,111,1,6.2954,flydsl_mxfp8_128_bpreshuffle_wmma_t32x32x512_mw2_nw2_nb4_sk1_cm1_cn1,111.93,1800.95,0.0
-gfx1250,256,64,1536,7168,flydsl,111,1,6.5214,flydsl_mxfp8_128_bpreshuffle_wmma_t32x32x512_mw2_nw2_nb4_sk1_cm1_cn1,216.1,1788.79,0.0
+gfx1250,256,64,1536,7168,triton,111,1,9.12,gluon,216.1,1788.79,0.0
 gfx1250,256,128,1536,7168,flydsl,111,1,7.0521,flydsl_mxfp8_128_bpreshuffle_wmma_t32x32x512_mw2_nw2_nb4_sk1_cm1_cn1,399.68,1747.11,0.0
 gfx1250,256,256,1536,7168,flydsl,138,1,8.4731,flydsl_mxfp8_128_bpreshuffle_wmma_t32x64x512_mw1_nw4_nb4_sk1_cm1_cn1,665.3,1608.8,0.0
 gfx1250,256,512,1536,7168,flydsl,136,1,10.5785,flydsl_mxfp8_128_bpreshuffle_wmma_t64x64x512_mw2_nw2_nb4_sk1_cm1_cn1,1065.77,1536.41,0.0
@@ -953,9 +953,9 @@ gfx1250,256,4096,1536,7168,flydsl,217,1,26.7965,flydsl_mxfp8_128_bpreshuffle_wmm
 gfx1250,256,8192,1536,7168,flydsl,262,1,39.8716,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,4524.24,2380.04,0.0
 gfx1250,256,16384,1536,7168,flydsl,528,1,79.0006,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,4566.77,2263.05,0.0
 gfx1250,256,32768,1536,7168,flydsl,528,1,134.0169,flydsl_mxfp8_128_bpreshuffle_wmma_t256x256x128_mw2_nw2_nb4_sk1_cm1_cn1,5384.06,2585.9,0.0
-gfx1250,256,16,8192,1536,flydsl,7,1,4.2533,flydsl_mxfp8_128_bpreshuffle_wmma_t16x32x512_mw1_nw2_nb3_sk1_cm1_cn1,94.67,3025.8,0.0
+gfx1250,256,16,8192,1536,triton,7,1,6.75,gluon,94.67,3025.8,0.0
 gfx1250,256,32,8192,1536,flydsl,110,1,4.7064,flydsl_mxfp8_128_bpreshuffle_wmma_t32x32x512_mw2_nw2_nb3_sk1_cm1_cn1,171.11,2795.42,0.0
-gfx1250,256,64,8192,1536,flydsl,132,1,5.0964,flydsl_mxfp8_128_bpreshuffle_wmma_t32x64x256_mw2_nw2_nb4_sk1_cm1_cn1,316.03,2694.02,0.0
+gfx1250,256,64,8192,1536,triton,132,1,7.00,gluon,316.03,2694.02,0.0
 gfx1250,256,128,8192,1536,flydsl,280,1,5.9456,flydsl_mxfp8_128_bpreshuffle_wmma_t64x64x256_mw2_nw2_nb4_sk1_cm1_cn1,541.78,2502.13,0.0
 gfx1250,256,256,8192,1536,flydsl,408,1,5.7752,flydsl_mxfp8_128_bpreshuffle_wmma_t128x64x256_mw2_nw2_nb4_sk1_cm1_cn1,1115.54,2973.13,0.0
 gfx1250,256,512,8192,1536,flydsl,3162,1,6.9948,flydsl_mxfp8_128_bpreshuffle_compute_wmma_t128x128x128_mw2_nw2_nb4_sk1_cm1_cn4,1842.07,3110.59,0.0
@@ -965,9 +965,9 @@ gfx1250,256,4096,8192,1536,flydsl,3157,1,24.9559,flydsl_mxfp8_128_bpreshuffle_co
 gfx1250,256,8192,8192,1536,flydsl,3157,1,41.2433,flydsl_mxfp8_128_bpreshuffle_compute_wmma_t256x256x128_mw2_nw2_nb2_sk1_cm1_cn4,4998.59,3864.47,0.0
 gfx1250,256,16384,8192,1536,flydsl,3157,1,81.8328,flydsl_mxfp8_128_bpreshuffle_compute_wmma_t256x256x128_mw2_nw2_nb2_sk1_cm1_cn4,5038.53,3741.58,0.0
 gfx1250,256,32768,8192,1536,flydsl,3157,1,157.7343,flydsl_mxfp8_128_bpreshuffle_compute_wmma_t256x256x128_mw2_nw2_nb2_sk1_cm1_cn4,5227.99,3802.51,0.0
-gfx1250,256,16,16384,1536,flydsl,45,1,4.8973,flydsl_mxfp8_128_bpreshuffle_wmma_t16x128x256_mw1_nw4_nb4_sk1_cm1_cn1,164.44,5250.79,0.0
+gfx1250,256,16,16384,1536,triton,45,1,7.01,gluon,164.44,5250.79,0.0
 gfx1250,256,32,16384,1536,flydsl,165,1,5.5571,flydsl_mxfp8_128_bpreshuffle_wmma_t32x128x256_mw2_nw2_nb4_sk1_cm1_cn1,289.83,4726.13,0.0
-gfx1250,256,64,16384,1536,flydsl,299,1,6.0058,flydsl_mxfp8_128_bpreshuffle_wmma_t64x128x256_mw1_nw4_nb4_sk1_cm1_cn1,536.35,4555.81,0.0
+gfx1250,256,64,16384,1536,triton,299,1,7.31,gluon,536.35,4555.81,0.0
 gfx1250,256,128,16384,1536,flydsl,424,1,7.0387,flydsl_mxfp8_128_bpreshuffle_wmma_t128x128x512_mw2_nw2_nb2_sk1_cm1_cn1,915.29,4199.18,0.0
 gfx1250,256,256,16384,1536,flydsl,422,1,8.0812,flydsl_mxfp8_128_bpreshuffle_wmma_t128x128x256_mw2_nw2_nb3_sk1_cm1_cn1,1594.43,4200.82,0.0
 gfx1250,256,512,16384,1536,flydsl,259,1,11.5139,flydsl_mxfp8_128_bpreshuffle_wmma_t256x128x256_mw4_nw1_nb3_sk1_cm1_cn1,2238.15,3711.12,0.0

aiter/ops/gemm_op_a8w8.py

@@ -936,22 +936,59 @@ def gemm_a8w8_blockscale_bpreshuffle(
     else:
         Y = torch.empty(m, n, dtype=dtype, device=XQ.device)
 
-    use_gfx1250_flydsl_mxfp8_128 = (
+    use_gfx1250_flydsl_or_triton_mxfp8_128 = (
         get_gfx() == "gfx1250"
         and x_scale.dtype == dtypes.fp8_e8m0
         and w_scale.dtype == dtypes.fp8_e8m0
     )
-    if use_gfx1250_flydsl_mxfp8_128:
-        if not is_flydsl_available():
-            raise RuntimeError(
-                "gfx1250 mxfp8_128 bpreshuffle (fp8_e8m0 scales) requires FlyDSL"
-            )
+    if use_gfx1250_flydsl_or_triton_mxfp8_128:
         config = get_CKGEMM_config(
             m,
             n,
             k,
             AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE,
         )
+        # A tuned triton/gluon row wins over flydsl on the SAME mxfp8_128
+        # operands: gemm_afp8wfp8_preshuffle consumes the e8m0 scales natively
+        # (wmma_scaled on gluon / dot_scaled on triton), so unlike the fp32
+        # blockscale path further down there is no scale widening, and the (N, K)
+        # shuffled weight + column-major x_scale layout is a direct fit.
+        # Deliberately checked BEFORE the FlyDSL availability gate so a
+        # triton-tuned shape does not require FlyDSL to be installed.
+        if config is not None and config["libtype"] == "triton":
+            from aiter.ops.triton.gemm.basic.gemm_afp8wfp8 import (
+                gemm_afp8wfp8_preshuffle as _gemm_afp8wfp8_preshuffle_triton,
+            )
+
+            # Same convention as the fp32 blockscale triton branch below:
+            # kernelName optionally carries the backend hint ("triton"/"gluon"),
+            # anything else -> None (auto gluon->triton detection).
+            kernelName = str(config.get("kernelName", ""))
+            backend = kernelName if kernelName in ("triton", "gluon") else None
+            return _gemm_afp8wfp8_preshuffle_triton(
+                XQ,
+                WQ,
+                # wmma_scaled/dot_scaled take uint8-typed scale operands; e8m0 is
+                # bit-identical, so a view is the whole conversion.
+                x_scale.view(torch.uint8),
+                w_scale.view(torch.uint8),
+                dtype=dtype,
+                y=Y,  # honor caller out= (zero-copy); Y = out or fresh empty
+                x_scale_group_size=128,
+                # The mxfp8_128 contract, not a guess: x_scale bytes are
+                # column-major (K // 128, M) -- what per_group_quant_hip(
+                # transpose_scale=True) emits, and what the flydsl runner
+                # hard-codes as x_scale_transposed=True. It is NOT inferable from
+                # strides here: that buffer is a contiguous (M, K // 128) tensor
+                # whose *bytes* are transposed, so the stride(0) != 1 probe the
+                # fp32 branches use would read it as row-major.
+                is_x_scale_transposed=True,
+                backend=backend,
+            )
+        if not is_flydsl_available():
+            raise RuntimeError(
+                "gfx1250 mxfp8_128 bpreshuffle (fp8_e8m0 scales) requires FlyDSL"
+            )
         if config is not None and config["libtype"] == "flydsl":
             return gemm_a8w8_mxfp8_128_bpreshuffle_flydsl(
                 XQ, WQ, x_scale, w_scale, Y, config