[Triton/Gluon] Add per token scale to a8w8 moe gemm

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

@@ -148,6 +148,8 @@ def _moe_gemm_a8w8(
     X_static_scale,
     W_static_scale,
     Quant_static_scale,
+    X_token_scale,
+    W_expt_scale,
     B,
     stride_b_e,  # Bias
     Gammas,
@@ -288,6 +290,11 @@ def _moe_gemm_a8w8(
         GatherIndx += start_m
         # no needs to bounds-check here because `offs_x_m` wraps around M dim
         offs_x_m = tl.load(GatherIndx + offs_x_m) // N_EXPTS_ACT
+    if X_token_scale is not None:
+        if GatherIndx is None:
+            offs_x_scale_m = start_m + offs_x_m
+        else:
+            offs_x_scale_m = offs_x_m
     offs_x_k = BLOCK_K * pid_k + tl.arange(0, BLOCK_K)
     XPtrs = (
         X
@@ -468,18 +475,16 @@ def _moe_gemm_a8w8(
                     ),
                     (BLOCK_M, BLOCK_K),
                 )
-                # Broadcasts on axis 1 where the EVEN_K branch uses axis 2;
-                # at most one can be right. Unreachable while the wrapper pins
-                # BLOCK_K == MX_PACK_DIVISOR, so left as-is.
-                b_sc_full = tl.reshape(
+                b_sc_t = tl.reshape(
                     tl.broadcast_to(
-                        b_sc[:, None, :], (BLOCK_N, MX_PACK_DIVISOR, MX_SCALE_BLOCK_K)
+                        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,
-                    tl.trans(b_f32 * b_sc_full),
+                    b_f32 * b_sc_full,
                     input_precision="ieee",
                 )
         else:
@@ -492,6 +497,12 @@ def _moe_gemm_a8w8(
         acc = acc * tl.load(X_static_scale)
     if W_static_scale is not None:
         acc = acc * tl.load(W_static_scale)
+    # per-token / per-expert fp8 scales
+    if X_token_scale is not None:
+        x_token_sc = tl.load(X_token_scale + offs_x_scale_m.to(index_type))
+        acc = acc * x_token_sc[:, None]
+    if W_expt_scale is not None:
+        acc = acc * tl.load(W_expt_scale + expt_id)
     # bias
     offs_m = BLOCK_M * block_id + tl.arange(0, BLOCK_M)
     offs_y_n = BLOCK_N * pid_n + tl.arange(0, BLOCK_N)

aiter/ops/triton/moe/moe_gemm_mxfp8.py

@@ -138,6 +138,8 @@ def moe_gemm_mxfp8(
         None,
         None,
         None,
+        None,
+        None,
         # bias
         bias,
         bias_stride,

aiter/ops/triton/moe/moe_op_gemm_a8w8.py

@@ -151,12 +151,36 @@ def moe_gemm_a8w8(
     swiglu_add_residual=True,
     unpadded_N=None,
     unpadded_K=None,
+    x_token_scale=None,
+    w_expt_scale=None,
 ):
     """
     Y[:, :] = 0.
     for e in num_experts:
         Y[idxs_y_m(e), :] += matmul(X[idxs_x_m(e), :], W[e, :, :])
+
+    x_token_scale: optional fp32 per-token activation scale, one entry per row
+        of x; applied to the accumulator before bias / activation.
+    w_expt_scale: optional fp32 per-expert weight scale [n_expts].
     """
+    if x_token_scale is not None:
+        assert (
+            x_token_scale.dtype == torch.float32
+        ), f"Expected fp32 x_token_scale, got {x_token_scale.dtype}"
+        x_token_scale = x_token_scale.reshape(-1).contiguous()
+        assert x_token_scale.numel() == x.shape[-2], (
+            f"x_token_scale must have one entry per row of x "
+            f"({x.shape[-2]}), got {x_token_scale.numel()}"
+        )
+    if w_expt_scale is not None:
+        assert (
+            w_expt_scale.dtype == torch.float32
+        ), f"Expected fp32 w_expt_scale, got {w_expt_scale.dtype}"
+        w_expt_scale = w_expt_scale.reshape(-1).contiguous()
+        assert w_expt_scale.numel() == w.shape[0], (
+            f"w_expt_scale must have one entry per expert "
+            f"({w.shape[0]}), got {w_expt_scale.numel()}"
+        )
     w_has_mx = w_scales is not None
     if w_has_mx:
         assert w.stride(-2) == 1, "`w` must be column-major when it has data-type mxfp"
@@ -250,6 +274,8 @@ def moe_gemm_a8w8(
         x_static_scale,
         w_static_scale,
         quant_static_scale,
+        x_token_scale,
+        w_expt_scale,
         bias,
         stride_bias,
         gammas,

op_tests/triton_tests/moe/test_moe_gemm_a8w8.py

@@ -221,6 +221,7 @@ class Case:
 @pytest.mark.parametrize("has_y_gammas", [False, True])
 @pytest.mark.parametrize("apply_swiglu", [False, True])
 @pytest.mark.parametrize("fused_quant", [False, True])
+@pytest.mark.parametrize("token_expt_scales", [False, True])
 def test_op(
     m,
     n,
@@ -230,6 +231,7 @@ def test_op(
     has_y_gammas,
     apply_swiglu,
     fused_quant,
+    token_expt_scales,
     n_expts_tot,
     n_expts_act,
     act_dtype_str,
@@ -263,6 +265,9 @@ def test_op(
     if act_mxfp8:
         act_dtype_str = act_dtype_str[2:]
 
+    if token_expt_scales and (act_mxfp8 or weight_mxfp8):
+        pytest.skip("token/expert scales are tested on the plain-fp8 path")
+
     weight_dtype = dtype_str_to_torch(weight_dtype_str)
     act_dtype = dtype_str_to_torch(act_dtype_str)
     m, rdata, gindx, sindx = init_routing_data(
@@ -282,6 +287,8 @@ def test_op(
         device=device,
     )
     x_ref, w_ref, bias_ref = x_tri.clone(), w_tri.clone(), bias_tri.clone()
+    x_token_scale = None
+    w_expt_scale = None
 
     if weight_mxfp8:
         w_tri, w_scale_tri = downcast_to_mxfp(w_tri, weight_dtype, axis=1)
@@ -300,9 +307,15 @@ def test_op(
     else:
         w_scale_tri = None
         swizzle_mx_scale = None
-        w_static_scale = w_tri.abs().max().float() / 448.0
-        w_tri = downcast_to_static_fp8_3d(w_tri, w_static_scale)
-        w_ref = (w_tri.float() * w_static_scale).to(torch.bfloat16)
+        if token_expt_scales:
+            w_static_scale = None
+            w_expt_scale = w_tri.abs().amax(dim=(1, 2)).float() / 448.0
+            w_tri = (w_tri / w_expt_scale[:, None, None]).to(weight_dtype)
+            w_ref = (w_tri.float() * w_expt_scale[:, None, None]).to(torch.bfloat16)
+        else:
+            w_static_scale = w_tri.abs().max().float() / 448.0
+            w_tri = downcast_to_static_fp8_3d(w_tri, w_static_scale)
+            w_ref = (w_tri.float() * w_static_scale).to(torch.bfloat16)
 
     if act_mxfp8:
         x_tri, x_mx_scales_tri = downcast_to_mxfp(x_tri, act_dtype, axis=-1)
@@ -311,9 +324,15 @@ def test_op(
         out_dtype = torch.bfloat16
     else:
         x_mx_scales_tri = None
-        x_static_scale = x_tri.abs().max().float() / 448.0
-        x_tri = downcast_to_static_fp8(x_tri, x_static_scale)
-        x_ref = (x_tri.float() * x_static_scale).to(torch.bfloat16)
+        if token_expt_scales:
+            x_static_scale = None
+            x_token_scale = x_tri.abs().amax(dim=-1).float() / 448.0
+            x_tri = (x_tri / x_token_scale[:, None]).to(act_dtype)
+            x_ref = (x_tri.float() * x_token_scale[:, None]).to(torch.bfloat16)
+        else:
+            x_static_scale = x_tri.abs().max().float() / 448.0
+            x_tri = downcast_to_static_fp8(x_tri, x_static_scale)
+            x_ref = (x_tri.float() * x_static_scale).to(torch.bfloat16)
         out_dtype = torch.float8_e4m3fn
 
     ref_y = moe_gemm_torch(
@@ -345,6 +364,8 @@ def test_op(
         swizzle_mx_scale,
         out_dtype,
         apply_swiglu,
+        x_token_scale=x_token_scale,
+        w_expt_scale=w_expt_scale,
     )
     if not act_mxfp8 and fused_quant:
         tri_y = (tri_y.float() * quant_static_scale).to(ref_y.dtype)