w8a8 matmul add splitk (#2700)

Author: Jennie1004

src/flag_gems/runtime/backend/_nvidia/hopper/ops/w8a8_block_fp8_matmul.py

@@ -345,6 +345,109 @@ def w8a8_block_fp8_matmul_kernel_host_tma(
     c_desc.store([offset_am, offset_bn], acc.to(c_desc.dtype))
 
 
+@libentry()
+@libtuner(
+    configs=runtime.ops_get_configs(
+        "w8a8_block_fp8_general_splitk",
+        yaml_path=EXPAND_CONFIG_FILENAME,
+    )
+    if os.environ.get("USE_FLAGTUNE") == "1"
+    else _get_placeholder_tuner_configs(pre_hook=None),
+    key=["M", "N", "K", "stride_am", "stride_bk"],
+    strategy=runtime.get_expand_config(
+        "w8a8_block_fp8_general_splitk", yaml_path=EXPAND_CONFIG_FILENAME
+    )["strategy"]
+    if os.environ.get("USE_FLAGTUNE") == "1"
+    else ["align32", "align32", "align32", "align32", "align32"],
+    warmup=5,
+    rep=5,
+)
+@triton.jit
+def w8a8_block_fp8_matmul_kernel_splitk(
+    A,
+    B,
+    C,
+    As,
+    Bs,
+    M,
+    N,
+    K,
+    group_n,
+    group_k,
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_As_m,
+    stride_As_k,
+    stride_Bs_k,
+    stride_Bs_n,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    pid_k = tl.program_id(1)
+
+    # grid_m = tl.cdiv(M, BLOCK_M)
+    grid_n = tl.cdiv(N, BLOCK_N)
+    pid_m = pid // grid_n
+    pid_n = pid % grid_n
+
+    offset_am = pid_m * BLOCK_M
+    offset_bn = pid_n * BLOCK_N
+    offs_am = offset_am + tl.arange(0, BLOCK_M)
+    offs_bn = offset_bn + tl.arange(0, BLOCK_N)
+
+    total_k_iters = tl.cdiv(K, BLOCK_K)
+    k_per_split = tl.cdiv(total_k_iters, SPLIT_K)
+    k_start = pid_k * k_per_split
+    k_end = min((pid_k + 1) * k_per_split, total_k_iters)
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(k_start, k_end):
+        offset_k = k * BLOCK_K
+        offs_k = offset_k + tl.arange(0, BLOCK_K)
+
+        a = tl.load(
+            A + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak,
+            mask=(offs_am[:, None] < M) & (offs_k[None, :] < K),
+            other=0.0,
+        )
+        b = tl.load(
+            B + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn,
+            mask=(offs_k[:, None] < K) & (offs_bn[None, :] < N),
+            other=0.0,
+        )
+
+        offs_ks = offset_k // group_k
+        a_s = tl.load(
+            As + offs_am * stride_As_m + offs_ks * stride_As_k,
+            mask=offs_am < M,
+            other=0.0,
+        )
+        b_s = tl.load(
+            Bs + offs_ks * stride_Bs_k + (offs_bn // group_n) * stride_Bs_n,
+            mask=offs_bn < N,
+            other=0.0,
+        )
+        acc += tl.dot(a, b, out_dtype=tl.float32) * a_s[:, None] * b_s[None, :]
+
+    offs_cm = offset_am + tl.arange(0, BLOCK_M)
+    offs_cn = offset_bn + tl.arange(0, BLOCK_N)
+    c_ptrs = C + offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn
+    mask = (offs_cm < M)[:, None] & (offs_cn < N)[None, :]
+    if C.dtype.element_ty == tl.bfloat16:
+        tl.atomic_add(c_ptrs, acc.to(tl.bfloat16), mask=mask)
+    elif C.dtype.element_ty == tl.float16:
+        tl.atomic_add(c_ptrs, acc.to(tl.float16), mask=mask)
+    else:
+        tl.atomic_add(c_ptrs, acc.to(tl.float32), mask=mask)
+
+
 def general_w8a8_block_fp8_matmul(a, b, c, a_s, b_s, M, N, K, group_n, group_k):
     logger.debug(
         "GEMS w8a8_block_fp8_matmul-hopper, [scenario]: general, [shape info]: [-, %s, %s, %s](batch, M, N, K), "
@@ -355,10 +458,88 @@ def general_w8a8_block_fp8_matmul(a, b, c, a_s, b_s, M, N, K, group_n, group_k):
         a.stride(0) == 1,
         b.stride(0) == 1,
     )
+
+    use_flagtune = os.environ.get("USE_FLAGTUNE") == "1"
+
+    # Split-K path for small-N, large-K shapes
+    if N <= 512 and K == 7168 and M < 8276:
+        if use_flagtune:
+            splitk_grid = lambda META: (
+                triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),
+                META["SPLIT_K"],
+            )
+            c.zero_()
+            with torch_device_fn.device(a.device):
+                w8a8_block_fp8_matmul_kernel_splitk[splitk_grid](
+                    a,
+                    b,
+                    c,
+                    a_s,
+                    b_s,
+                    M,
+                    N,
+                    K,
+                    group_n,
+                    group_k,
+                    a.stride(0),
+                    a.stride(1),
+                    b.stride(1),
+                    b.stride(0),
+                    c.stride(0),
+                    c.stride(1),
+                    a_s.stride(0),
+                    a_s.stride(1),
+                    b_s.stride(1),
+                    b_s.stride(0),
+                )
+        else:
+            SPLITK_BLOCK_K = 128
+            SPLITK_BLOCK_M = 16 if M <= 16 else 64
+            SPLITK_BLOCK_N = 64 if N > 256 else 32
+
+            grid_m = triton.cdiv(M, SPLITK_BLOCK_M)
+            grid_n = triton.cdiv(N, SPLITK_BLOCK_N)
+            grid_mn = grid_m * grid_n
+            total_k_iters = triton.cdiv(K, SPLITK_BLOCK_K)
+
+            SM_COUNT = torch.cuda.get_device_properties(a.device).multi_processor_count
+            split_k = min(total_k_iters, max(4, 2 * SM_COUNT // max(grid_mn, 1)))
+
+            c.zero_()
+            splitk_grid = (grid_mn, split_k)
+
+            with torch_device_fn.device(a.device):
+                w8a8_block_fp8_matmul_kernel_splitk.fn.fn[splitk_grid](
+                    a,
+                    b,
+                    c,
+                    a_s,
+                    b_s,
+                    M,
+                    N,
+                    K,
+                    group_n,
+                    group_k,
+                    a.stride(0),
+                    a.stride(1),
+                    b.stride(1),
+                    b.stride(0),
+                    c.stride(0),
+                    c.stride(1),
+                    a_s.stride(0),
+                    a_s.stride(1),
+                    b_s.stride(1),
+                    b_s.stride(0),
+                    BLOCK_M=SPLITK_BLOCK_M,
+                    BLOCK_N=SPLITK_BLOCK_N,
+                    BLOCK_K=SPLITK_BLOCK_K,
+                    SPLIT_K=split_k,
+                )
+        return c
+
     grid = lambda meta: (
         triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),
     )
-    use_flagtune = os.environ.get("USE_FLAGTUNE") == "1"
     fixed_meta = (
         None
         if use_flagtune

src/flag_gems/runtime/backend/_nvidia/hopper/w8a8_block_fp8_matmul_hopper_expand.yaml

@@ -35,6 +35,44 @@ w8a8_block_fp8_general:
      stride_am: default
      stride_bk: default
 
+w8a8_block_fp8_general_splitk:
+  - config:
+    param_map:
+      META:
+        BLOCK_M: block_m
+        BLOCK_N: block_n
+        BLOCK_K: block_k
+        SPLIT_K: split_k
+      num_stages: stages
+      num_warps: warps
+    block_m:
+      - 16
+      - 64
+      - 128
+    block_n:
+      - 32
+      - 64
+    block_k:
+      - 128
+    split_k:
+      - 4
+      - 8
+      - 16
+      - 28
+    stages:
+      - 1
+      - 2
+      - 3
+    warps:
+      - 4
+      - 8
+  - strategy:
+      M: default
+      N: default
+      K: default
+      stride_am: default
+      stride_bk: default
+
 w8a8_block_fp8_general_tma:
   - config:
     param_map:

src/flag_gems/runtime/configloader.py

@@ -30,6 +30,13 @@
         "align32",
         "align32",
     ],
+    "w8a8_block_fp8_general_splitk": [
+        "align32",
+        "align32",
+        "align32",
+        "align32",
+        "align32",
+    ],
     "w8a8_block_fp8_general_tma": [
         "align32",
         "align32",
@@ -59,6 +66,7 @@
     "baddbmm": ["M", "N", "K"],
     "mv": ["M", "N"],
     "w8a8_block_fp8_general": ["M", "N", "K", "stride_am", "stride_bk"],
+    "w8a8_block_fp8_general_splitk": ["M", "N", "K", "stride_am", "stride_bk"],
     "w8a8_block_fp8_general_tma": ["M", "N", "K", "stride_am", "stride_bk", "dtype"],
     "mm_general_tma": ["M", "N", "K", "stride_am", "stride_bk", "dtype"],
     "gemv": ["M", "K", "stride_am", "stride_bk"],
@@ -330,6 +338,27 @@ def _build_configs_by_op(self, op_name, ranges, pre_hook=None):
                 for w in ranges["w"]
             ]
 
+        if op_name == "w8a8_block_fp8_general_splitk":
+            return [
+                triton.Config(
+                    {
+                        "BLOCK_M": block_m,
+                        "BLOCK_N": block_n,
+                        "BLOCK_K": block_k,
+                        "SPLIT_K": split_k,
+                    },
+                    num_stages=s,
+                    num_warps=w,
+                    pre_hook=pre_hook,
+                )
+                for block_m in ranges["BLOCK_M"]
+                for block_n in ranges["BLOCK_N"]
+                for block_k in ranges["BLOCK_K"]
+                for split_k in ranges["SPLIT_K"]
+                for s in ranges["s"]
+                for w in ranges["w"]
+            ]
+
         return []
 
     def _build_single_expand_spec(
@@ -362,6 +391,9 @@ def _build_expand_registry(self):
             "w8a8_block_fp8_general": self._build_single_expand_spec(
                 "w8a8_block_fp8_general"
             ),
+            "w8a8_block_fp8_general_splitk": self._build_single_expand_spec(
+                "w8a8_block_fp8_general_splitk"
+            ),
             "w8a8_block_fp8_general_tma": self._build_single_expand_spec(
                 "w8a8_block_fp8_general_tma"
             ),