[Triton/Gluon] Merge MoE MXFP8 GEMM into existing _moe_gemm_a8w8 kernel

aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a8w8.py

@@ -120,6 +120,7 @@ def unswizzle_mx_scale_cdna4(
         "N_EXPTS_ACT",
         "APPLY_SWIGLU",
         "SWIGLU_ADD_RESIDUAL",
+        "USE_FNUZ",
     ],
 )
 
@@ -182,6 +183,9 @@ def _moe_gemm_a8w8(
     SPLIT_K: tl.constexpr,
     W_CACHE_MODIFIER: tl.constexpr,
     UPCAST_INDICES: tl.constexpr = False,
+    # Use fnuz FP8 (float8_e4m3fnuz / gfx942) instead of OCP FP8 (float8e4nv / gfx950).
+    # When True, replaces tl.dot_scaled with manual E8M0→fp32 dequant + tl.dot.
+    USE_FNUZ: tl.constexpr = False,
 ):
     tl.assume(stride_y_k >= 0)
     tl.assume(stride_y_m >= 0)
@@ -211,7 +215,10 @@ def _moe_gemm_a8w8(
     MX_PACK_DIVISOR: tl.constexpr = 32
     w_type: tl.constexpr = W.dtype.element_ty
     if is_w_microscaled:
-        tl.static_assert(w_type == tl.float8e4nv, "mx_weight_ptr must be float8e4nv")
+        tl.static_assert(
+            w_type == tl.float8e4nv or (USE_FNUZ and w_type == tl.float8e4b8),
+            "mx_weight_ptr must be float8e4nv (OCP) or float8e4b8/fnuz when USE_FNUZ=True",
+        )
         tl.static_assert(
             WMxScale.dtype.element_ty == tl.uint8, "mx_scale_ptr must be uint8"
         )
@@ -221,7 +228,10 @@ def _moe_gemm_a8w8(
         )
     x_type: tl.constexpr = X.dtype.element_ty
     if is_x_microscaled:
-        tl.static_assert(x_type == tl.float8e4nv, "mx_act_ptr must be float8e4nv")
+        tl.static_assert(
+            x_type == tl.float8e4nv or (USE_FNUZ and x_type == tl.float8e4b8),
+            "mx_act_ptr must be float8e4nv (OCP) or float8e4b8/fnuz when USE_FNUZ=True",
+        )
         tl.static_assert(
             XMxScale.dtype.element_ty == tl.uint8, "mx_scale_ptr must be uint8"
         )
@@ -358,9 +368,39 @@ def _moe_gemm_a8w8(
         else:
             w_scales = tl.full((BLOCK_N, MX_SCALE_BLOCK_K), 127, dtype=tl.uint8)
 
-        acc = tl.dot_scaled(
-            x, x_scales, "e4m3", w, w_scales, "e4m3", acc=acc, fast_math=True
-        )
+        if USE_FNUZ:
+            # fnuz path (gfx942): manual E8M0→fp32, then tl.dot.
+            # x_scales: [BLOCK_M, MX_SCALE_BLOCK_K], w_scales: [BLOCK_N, MX_SCALE_BLOCK_K]
+            a_f32 = x.to(tl.float32)
+            b_f32 = w.to(tl.float32)
+            a_sc = (x_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True)
+            b_sc = (w_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True)
+            # broadcast each block-scale across its MX_PACK_DIVISOR elements.
+            # a_sc: [BLOCK_M, MX_SCALE_BLOCK_K] → a_sc_full: [BLOCK_M, BLOCK_K]
+            # b_sc: [BLOCK_N, MX_SCALE_BLOCK_K], w layout is [K, N] so we need
+            # b_sc_full transposed: [BLOCK_K, BLOCK_N]
+            a_sc_full = tl.reshape(
+                tl.broadcast_to(
+                    a_sc[:, :, None], (BLOCK_M, MX_SCALE_BLOCK_K, MX_PACK_DIVISOR)
+                ),
+                (BLOCK_M, BLOCK_K),
+            )
+            b_sc_t = tl.reshape(
+                tl.broadcast_to(
+                    b_sc[:, :, None], (BLOCK_N, MX_SCALE_BLOCK_K, MX_PACK_DIVISOR)
+                ),
+                (BLOCK_N, BLOCK_K),
+            )
+            b_sc_full = tl.trans(b_sc_t)  # [BLOCK_K, BLOCK_N]
+            acc += tl.dot(
+                a_f32 * a_sc_full,
+                b_f32 * b_sc_full,
+                input_precision="ieee",
+            )
+        else:
+            acc = tl.dot_scaled(
+                x, x_scales, "e4m3", w, w_scales, "e4m3", acc=acc, fast_math=True
+            )
 
         if is_w_microscaled:
             WMxScalePtrs += (PACKED_MX_BLOCK * SPLIT_K) * stride_w_mx_k
@@ -399,9 +439,32 @@ def _moe_gemm_a8w8(
         else:
             w_scales = tl.full((BLOCK_N, MX_SCALE_BLOCK_K), 127, dtype=tl.uint8)
 
-        acc = tl.dot_scaled(
-            x, x_scales, "e4m3", w, w_scales, "e4m3", acc=acc, fast_math=True
-        )
+        if USE_FNUZ:
+            a_f32 = x.to(tl.float32)
+            b_f32 = w.to(tl.float32)
+            a_sc = (x_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True)
+            b_sc = (w_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True)
+            a_sc_full = tl.reshape(
+                tl.broadcast_to(
+                    a_sc[:, :, None], (BLOCK_M, MX_SCALE_BLOCK_K, MX_PACK_DIVISOR)
+                ),
+                (BLOCK_M, BLOCK_K),
+            )
+            b_sc_full = tl.reshape(
+                tl.broadcast_to(
+                    b_sc[:, None, :], (BLOCK_N, MX_PACK_DIVISOR, MX_SCALE_BLOCK_K)
+                ),
+                (BLOCK_N, BLOCK_K),
+            )
+            acc += tl.dot(
+                a_f32 * a_sc_full,
+                tl.trans(b_f32 * b_sc_full),
+                input_precision="ieee",
+            )
+        else:
+            acc = tl.dot_scaled(
+                x, x_scales, "e4m3", w, w_scales, "e4m3", acc=acc, fast_math=True
+            )
 
     # scalar fp8 scale
     if X_static_scale is not None:
@@ -418,8 +481,8 @@ def _moe_gemm_a8w8(
         if pid_k == 0:
             bias = tl.load(BPtrs, mask=mask_n, other=0, cache_modifier=W_CACHE_MODIFIER)
         else:
-            bias = tl.full([BLOCK_N], 0, dtype=tl.float32)
-        acc = acc + bias[None, :]
+            bias = tl.zeros([BLOCK_N], dtype=B.dtype.element_ty)
+        acc = acc + bias.to(tl.float32)[None, :]
     if APPLY_SWIGLU and SPLIT_K == 1:
         out = _swiglu(acc, alpha, limit, ADD_RESIDUAL=SWIGLU_ADD_RESIDUAL)
         tl.static_assert(

aiter/ops/triton/configs/gfx942/triton/moe/mxfp8_fnuz/DEFAULT.json

@@ -0,0 +1,8 @@
+{
+    "moe_gemm_mxfp8": {
+        "BLOCK_M": 64,
+        "BLOCK_N": 128,
+        "num_warps": 4,
+        "num_stages": 1
+    }
+}

aiter/ops/triton/configs/gfx950/triton/moe/mxfp8_fnuz/DEFAULT.json

@@ -0,0 +1,8 @@
+{
+    "moe_gemm_mxfp8": {
+        "BLOCK_M": 64,
+        "BLOCK_N": 128,
+        "num_warps": 4,
+        "num_stages": 1
+    }
+}