[Triton] DS FP4 fusions redo

aiter/fused_moe.py

@@ -25,6 +25,7 @@
 from aiter.jit.utils.torch_guard import torch_compile_guard
 from aiter.utility import fp4_utils
 from aiter.utility.fp4_utils import moe_mxfp4_sort
+from aiter.ops.triton.fused_mxfp4_quant import fused_dynamic_mxfp4_quant_moe_sort
 
 BLOCK_SIZE_M = 32
 
@@ -790,7 +791,7 @@ def fused_moe_2stages(
     bias2=None,
 ):
     quant_func = get_quant(quant_type)
-
+    token_num_quant_moe_sort_switch = 1024
     token_num, _ = hidden_states.shape
     E, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape)
     dtype = moe_out.dtype
@@ -822,19 +823,29 @@ def fused_moe_2stages(
         a1 = hidden_states.to(dtype)
         a1_scale = None
     elif quant_type == QuantType.per_1x32:
-        a1, a1_scale = quant_func(
-            hidden_states,
-            scale=a1_scale,
-            quant_dtype=q_dtype_a,
-            num_rows=num_local_tokens,
-        )
-        a1_scale = moe_mxfp4_sort(
-            a1_scale,
-            sorted_ids=sorted_ids,
-            num_valid_ids=num_valid_ids,
-            token_num=token_num,
-            block_size=block_size_M,
-        )
+        if token_num <= token_num_quant_moe_sort_switch:
+            a1, a1_scale = fused_dynamic_mxfp4_quant_moe_sort(
+                hidden_states,
+                sorted_ids=sorted_ids,
+                num_valid_ids=num_valid_ids,
+                token_num=token_num,
+                topk=1,
+                block_size=block_size_M,
+            )
+        else:
+            a1, a1_scale = quant_func(
+                hidden_states,
+                scale=a1_scale,
+                quant_dtype=q_dtype_a,
+                num_rows=num_local_tokens,
+            )
+            a1_scale = moe_mxfp4_sort(
+                a1_scale,
+                sorted_ids=sorted_ids,
+                num_valid_ids=num_valid_ids,
+                token_num=token_num,
+                block_size=block_size_M,
+            )
     elif hidden_states.dtype != q_dtype_a:
         if quant_type == QuantType.per_1x128 and metadata.stage1.func is asm_stage1:
             quant_func = functools.partial(quant_func, transpose_scale=True)
@@ -887,22 +898,31 @@ def fused_moe_2stages(
         a2_scale = None
     elif quant_type == QuantType.per_1x32:
         a2 = a2.view(-1, inter_dim)
-        a2, a2_scale = quant_func(
-            a2,
-            scale=a2_scale,
-            quant_dtype=q_dtype_a,
-            num_rows=num_local_tokens,
-            num_rows_factor=topk,
-        )
+        if token_num <= token_num_quant_moe_sort_switch:
+            a2, a2_scale = fused_dynamic_mxfp4_quant_moe_sort(
+                a2,
+                sorted_ids=sorted_ids,
+                num_valid_ids=num_valid_ids,
+                token_num=token_num,
+                topk=topk,
+                block_size=block_size_M,
+            )
+        else:
+            a2, a2_scale = quant_func(
+                a2,
+                scale=a2_scale,
+                quant_dtype=q_dtype_a,
+                num_rows=num_local_tokens,
+                num_rows_factor=topk,
+            )
+            a2_scale = moe_mxfp4_sort(
+                a2_scale[: token_num * topk, :].view(token_num, topk, -1),
+                sorted_ids=sorted_ids,
+                num_valid_ids=num_valid_ids,
+                token_num=token_num,
+                block_size=block_size_M,
+            )
         a2 = a2.view(token_num, topk, -1)
-        a2_scale = moe_mxfp4_sort(
-            a2_scale[: token_num * topk, :].view(token_num, topk, -1),
-            sorted_ids=sorted_ids,
-            num_valid_ids=num_valid_ids,
-            token_num=token_num,
-            block_size=block_size_M,
-        )
-
     elif quant_type == QuantType.per_1x128 and metadata.stage1.func is asm_stage1:
         a2_v = a2[:token_num, :, :]
         a2_scale = (

aiter/ops/triton/_triton_kernels/batched_gemm_afp4wfp4_pre_quant.py → aiter/ops/triton/_triton_kernels/batched_gemm_a16wfp4.py

@@ -14,8 +14,8 @@
 from ..utils.core import AITER_TRITON_CONFIGS_PATH
 from .quant import _mxfp4_quant_op
 
-_batched_gemm_afp4_wfp4_pre_quant_repr = make_kernel_repr(
-    "_batched_gemm_afp4_wfp4_pre_quant_kernel",
+_batched_gemm_a16wfp4_repr = make_kernel_repr(
+    "_batched_gemm_a16wfp4_kernel",
     [
         "BLOCK_SIZE_M",
         "BLOCK_SIZE_N",
@@ -25,13 +25,15 @@
         "SPLITK_BLOCK_SIZE",
         "EVEN_K",
         "GRID_MN",
+        "PRE_QUANT",
+        "HAVE_Y_SCALE",
         "cache_modifier",
     ],
 )
 
 
-_batched_gemm_afp4_wfp4_pre_quant_reduce_repr = make_kernel_repr(
-    "_batched_gemm_afp4_wfp4_pre_quant_reduce_kernel",
+_batched_gemm_a16wfp4_reduce_repr = make_kernel_repr(
+    "_batched_gemm_a16wfp4_reduce_kernel",
     [
         "BLOCK_SIZE_M",
         "BLOCK_SIZE_N",
@@ -50,12 +52,13 @@
         * triton.cdiv(args["N"], args["BLOCK_SIZE_N"]),
     }
 )
-@triton.jit(repr=_batched_gemm_afp4_wfp4_pre_quant_repr)
-def _batched_gemm_afp4_wfp4_pre_quant_kernel(
+@triton.jit(repr=_batched_gemm_a16wfp4_repr)
+def _batched_gemm_a16wfp4_kernel(
     a_ptr,
     b_ptr,
     c_ptr,
     b_scales_ptr,
+    c_scale_ptr,
     M,
     N,
     K,
@@ -81,6 +84,8 @@ def _batched_gemm_afp4_wfp4_pre_quant_kernel(
     SPLITK_BLOCK_SIZE: tl.constexpr,
     EVEN_K: tl.constexpr,
     GRID_MN: tl.constexpr,
+    PRE_QUANT: tl.constexpr,
+    HAVE_Y_SCALE: tl.constexpr,
     cache_modifier: tl.constexpr,
 ):
     """
@@ -121,6 +126,12 @@ def _batched_gemm_afp4_wfp4_pre_quant_kernel(
     stride_cb = tl.cast(stride_cb, tl.int64)
     pid_batch = tl.cast(pid_batch, tl.int64)
 
+    if HAVE_Y_SCALE:
+        c_scale = tl.load(c_scale_ptr)
+    else:
+        c_scale = 1
+    c_scale_rcprl = (1 / c_scale).to(tl.float32)
+
     if NUM_KSPLIT == 1:
         remap_xcd(pid, GRID_MN)
 
@@ -183,13 +194,16 @@ def _batched_gemm_afp4_wfp4_pre_quant_kernel(
                 b = tl.load(b_ptrs, cache_modifier=cache_modifier)
             else:
                 a_bf16 = tl.load(
-                    a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0
+                    a_ptrs, mask=offs_k_bf16[None, :] < K - k * BLOCK_SIZE_K, other=0
                 )
                 b = tl.load(
                     b_ptrs, mask=offs_k[:, None] < K - k * (BLOCK_SIZE_K // 2), other=0
                 )
 
-            a, a_scales = _mxfp4_quant_op(a_bf16, BLOCK_SIZE_K, BLOCK_SIZE_M, 32)
+            if PRE_QUANT:  # TODO add PRE_QUANT = False
+                a, a_scales = _mxfp4_quant_op(
+                    a_bf16, BLOCK_SIZE_K, BLOCK_SIZE_M, SCALE_GROUP_SIZE
+                )
 
             accumulator += tl.dot_scaled(a, a_scales, "e2m1", b, b_scales, "e2m1")
 
@@ -198,6 +212,9 @@ def _batched_gemm_afp4_wfp4_pre_quant_kernel(
             b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk
             b_scale_ptrs += (BLOCK_SIZE_K // SCALE_GROUP_SIZE) * stride_bsk
 
+        if HAVE_Y_SCALE:
+            accumulator = accumulator * c_scale_rcprl
+
         c = accumulator.to(c_ptr.type.element_ty)
 
         # Write back the block of the output matrix C with masks.
@@ -214,8 +231,8 @@ def _batched_gemm_afp4_wfp4_pre_quant_kernel(
         tl.store(c_ptrs, c, mask=c_mask)
 
 
-@triton.jit(repr=_batched_gemm_afp4_wfp4_pre_quant_reduce_repr)
-def _batched_gemm_afp4_wfp4_pre_quant_reduce_kernel(
+@triton.jit(repr=_batched_gemm_a16wfp4_reduce_repr)
+def _batched_gemm_a16wfp4_reduce_kernel(
     c_in_ptr,
     c_out_ptr,
     M,