[FlyDSL] One-stage split-K for the a8w8 preshuffle GEMM

aiter/aot/flydsl/gemm.py

@@ -97,7 +97,10 @@
     r"(?P<tile_m>\d+)x(?P<tile_n>\d+)x(?P<tile_k>\d+)_"
     r"(?P<qa>[A-Z0-9]+)_(?P<qw>[A-Z0-9]+)_(?P<out>[A-Z0-9]+)_"
     r"(?P<async_copy>\d+)x(?P<waves_per_eu>\d+)(?:x(?P<xcd_swizzle>\d+))?(?:x(?P<lds_stage>\d+))?_"
-    r"(?P<scheduler>[A-Za-z][A-Za-z0-9]*)$"
+    r"(?!ks\d+$)(?P<scheduler>[A-Za-z][A-Za-z0-9]*)"
+    # Trailing _ksN, emitted only for k_split > 1, so pre-split-K names still
+    # match. Without it they fail fullmatch and drop out of the AOT build.
+    r"(?:_ks(?P<k_split>\d+))?$"
 )
 _SHORT_DTYPE = {
     "F8": "fp8",
@@ -147,6 +150,7 @@ def _parse_preshuffle_kernel_name(name: str) -> dict | None:
         "xcd_swizzle": int(m.group("xcd_swizzle")) if m.group("xcd_swizzle") else 0,
         "lds_stage": int(m.group("lds_stage")) if m.group("lds_stage") else 2,
         "scheduler": m.group("scheduler"),
+        "k_split": int(m.group("k_split")) if m.group("k_split") else 1,
     }
 
 
@@ -354,10 +358,12 @@ def _compile_preshuffle_to_cache(
     xcd_swizzle: int = 0,
     lds_stage: int = 2,
     scheduler: str = "Default",
+    k_split: int = 1,
     **kwargs,
 ):
     del kwargs
     enable_scheduler = str(scheduler).lower() != "off"
+    k_split = int(k_split)
 
     import torch
 
@@ -367,7 +373,23 @@ def _compile_preshuffle_to_cache(
     # FlyDSL preshuffle kernels consume raw quantized bytes for fp8/int8 paths.
     a = torch.empty((m * k,), device=dev, dtype=torch.int8)
     b = torch.empty((n * k,), device=dev, dtype=torch.int8)
+    from aiter.ops.flydsl.gemm_kernels import (
+        PRESHUFFLE_SPLIT_K_MAX_TILES,
+        PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS,
+    )
+
+    # Sized to the same bounds the runtime uses, so the signatures match.
     out = torch.empty((m * n,), device=dev, dtype=out_torch_dtype)
+    workspace = (
+        torch.empty(PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS, device=dev, dtype=torch.float32)
+        if k_split > 1
+        else out
+    )
+    semaphore = torch.zeros(
+        PRESHUFFLE_SPLIT_K_MAX_TILES if k_split > 1 else 0,
+        device=dev,
+        dtype=torch.int32,
+    )
     scale_a = torch.empty((max(m, 1),), device=dev, dtype=torch.float32)
     scale_b = torch.empty((max(n, 1),), device=dev, dtype=torch.float32)
     bias = torch.empty(0, device=dev, dtype=out_torch_dtype)
@@ -386,13 +408,16 @@ def _compile_preshuffle_to_cache(
         enable_scheduler=enable_scheduler,
         xcd_swizzle=xcd_swizzle,
         lds_stage=lds_stage,
+        split_k=k_split,
     )
     # The layout-API launcher uses fx.Tensor args (it builds views via
     # fx.get_iter/make_view), so pass flat torch tensors directly rather
     # than raw pointers (pointer args would fail GetIterOp type checks).
     _compile_executable_to_cache(
         exe,
+        workspace,
         out,
+        semaphore,
         a,
         b,
         scale_a,

aiter/ops/flydsl/gemm_kernels.py

@@ -971,6 +971,54 @@ def _get_compile_fn():
     return _flydsl_compile_fn
 
 
+# Fixed size rather than one buffer per shape: a shape-keyed cache grows without
+# limit and can evict a buffer a captured CUDA graph still points at. The bounds
+# come from k_split_candidates, which keeps tile_count under CU_NUM and
+# k_split * tile_count at four per CU.
+# Mirrors preshuffle_gemm.PRESHUFFLE_M_MAX; duplicated so this module imports
+# without FlyDSL present.
+PRESHUFFLE_M_MAX = 65536
+
+PRESHUFFLE_SPLIT_K_MAX_TILES = 256
+PRESHUFFLE_SPLIT_K_MAX_TILE_ELEMS = 32 * 128
+PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS = (
+    4 * PRESHUFFLE_SPLIT_K_MAX_TILES * PRESHUFFLE_SPLIT_K_MAX_TILE_ELEMS
+)
+
+
+@functools.lru_cache(maxsize=128)
+def _get_preshuffle_split_buffers(
+    device: torch.device,
+    stream: torch.cuda.Stream,
+) -> tuple[Tensor, Tensor]:
+    # Safe to reuse: launches on a stream are ordered and the reduction hands
+    # the semaphore back zeroed.
+    workspace = torch.empty(
+        PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS, dtype=torch.float32, device=device
+    )
+    semaphore = torch.zeros(
+        PRESHUFFLE_SPLIT_K_MAX_TILES, dtype=torch.int32, device=device
+    )
+    return workspace, semaphore
+
+
+def _check_preshuffle_split_capacity(
+    m: int, n: int, tile_m: int, tile_n: int, split_k: int
+) -> None:
+    tiles = ((m + tile_m - 1) // tile_m) * (n // tile_n)
+    if tiles > PRESHUFFLE_SPLIT_K_MAX_TILES:
+        raise RuntimeError(
+            f"[FlyDSL] split_k needs {tiles} tile semaphores, "
+            f"more than {PRESHUFFLE_SPLIT_K_MAX_TILES}"
+        )
+    elems = split_k * m * n
+    if elems > PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS:
+        raise RuntimeError(
+            f"[FlyDSL] split_k needs a {elems}-element fp32 workspace, "
+            f"more than {PRESHUFFLE_SPLIT_K_WORKSPACE_ELEMS}"
+        )
+
+
 def flydsl_preshuffle_gemm_a8(
     XQ: Tensor,
     WQ: Tensor,
@@ -985,8 +1033,9 @@ def flydsl_preshuffle_gemm_a8(
     xcd_swizzle: int = 0,
     lds_stage: int = 2,
     enable_scheduler: bool = True,
+    split_k: int = 1,
 ) -> Tensor:
-    """Compile (cached via lru_cache) and run a FlyDSL preshuffle GEMM kernel."""
+    """Compile and run FlyDSL preshuffle GEMM, optionally with fp32 split-K."""
     compile_fn = _get_compile_fn()
     if compile_fn is None:
         raise RuntimeError("[FlyDSL] compile function not available")
@@ -995,14 +1044,24 @@ def flydsl_preshuffle_gemm_a8(
     m, k = XQ.shape[0], XQ.shape[-1]
     n = WQ.shape[0]
 
+    if m > PRESHUFFLE_M_MAX:
+        raise RuntimeError(
+            f"[FlyDSL] M ({m}) exceeds {PRESHUFFLE_M_MAX}; the preshuffle kernel "
+            f"views A and C through a layout bounded by that many rows."
+        )
     if n % tile_n != 0:
         raise RuntimeError(
             f"[FlyDSL] N ({n}) is not a multiple of tile_n ({tile_n}). "
             f"Arguments not supported! Skipping gemm!"
         )
-    if k % tile_k != 0:
+    if split_k < 1 or k % split_k != 0:
+        raise RuntimeError(
+            f"[FlyDSL] K ({k}) must be divisible by split_k ({split_k})."
+        )
+    if (k // split_k) % tile_k != 0:
         raise RuntimeError(
-            f"[FlyDSL] K ({k}) is not a multiple of tile_k ({tile_k}). "
+            f"[FlyDSL] K/split_k ({k // split_k}) is not a multiple of "
+            f"tile_k ({tile_k}). "
             f"Arguments not supported! Skipping gemm!"
         )
 
@@ -1037,6 +1096,7 @@ def flydsl_preshuffle_gemm_a8(
         enable_scheduler=bool(enable_scheduler),
         xcd_swizzle=int(xcd_swizzle),
         lds_stage=int(lds_stage),
+        split_k=int(split_k),
     )
 
     def _as_i8(t):
@@ -1047,12 +1107,24 @@ def _as_i8(t):
     # epilogue != "none"). Pass an empty tensor as a placeholder for the
     # default epilogue="none" path.
     _dummy_bias = torch.empty(0, dtype=Out.dtype, device=Out.device)
+    if split_k > 1:
+        _check_preshuffle_split_capacity(m, n, tile_m, tile_n, split_k)
+        workspace, semaphore = _get_preshuffle_split_buffers(
+            Out.device, torch.cuda.current_stream(device=Out.device)
+        )
+    else:
+        workspace = out_contig
+        # dtype is part of the executable's cache signature, so this must match
+        # what the AOT pre-compile passes or every non-split-K kernel misses it.
+        semaphore = torch.empty(0, dtype=torch.int32, device=Out.device)
     # The layout-API launcher (PR #754) takes fx.Tensor args (it builds views via
     # fx.get_iter/make_view), so pass flat torch tensors directly rather than raw
     # pointers.
     _run_compiled(
         exe,
+        workspace.view(-1),
         out_contig.view(-1),
+        semaphore,
         _as_i8(XQ.contiguous()).view(-1),
         _as_i8(WQ.contiguous()).view(-1),
         x_scale.contiguous().view(-1),

aiter/ops/flydsl/gemm_tune/flydsl_gemm_a8w8_bpreshuffle_common.py

@@ -67,6 +67,7 @@ class kernelInstance:
     xcd_swizzle: int  # 0=off, >0=group size for XCD remap
     lds_stage: int = 2  # 2=double-buffer ping-pong, 1=single A-LDS buffer (half LDS)
     sScheduler: str = "Default"  # scheduler hints on; "Off" = compiler default
+    k_split: int = 1  # >1 splits the K loop over gridDim.z (fp32 workspace + reduce)
 
     @property
     def enable_scheduler(self) -> bool:
@@ -99,6 +100,7 @@ def name(self) -> str:
                 ),
                 self.sScheduler.lower(),
             ]
+            + ([f"ks{self.k_split}"] if self.k_split > 1 else [])
         )
 
 
@@ -114,6 +116,7 @@ def _ki(
     q_dtype_w="fp8",
     dtype="bf16",
     scheduler="Default",
+    k_split=1,
 ):
     return kernelInstance(
         tile_m,
@@ -127,6 +130,7 @@ def _ki(
         xcd_swizzle,
         lds_stage,
         scheduler,
+        k_split,
     )
 
 
@@ -227,6 +231,8 @@ def kernel_fits_shape(ki: kernelInstance, M: int, N: int, K: int) -> bool:
         return False
     if N % ki.tile_n != 0 or K % ki.tile_k != 0:
         return False
+    if ki.k_split > 1 and (K // ki.tile_k) % ki.k_split != 0:
+        return False
     if _padded_m(M) % ki.tile_m != 0:
         return False
     num_ctas = ((M + ki.tile_m - 1) // ki.tile_m) * (N // ki.tile_n)
@@ -329,6 +335,32 @@ def _estimate_max_wpe(tile_m: int, tile_n: int, total_vgpr: int = 512) -> int:
     return int(total_vgpr / max(est_per_wave, 1))
 
 
+# Legal values are the divisors of K//tile_k, which is shape-dependent, so they
+# are enumerated rather than hardcoded.
+K_SPLIT_MIN_TILES_PER_SLICE = 2  # keep the ping-pong loop fed
+K_SPLIT_MAX_CTA_OVERSUBSCRIBE = 4  # no point going far past one CU each
+
+
+def k_split_candidates(ki, M: int, N: int, K: int, cu_num: int = 256) -> list[int]:
+    """Split-K values worth benchmarking; 1 is excluded, the caller has it.
+
+    Empty once the tile grid already fills the GPU -- splitting would only add
+    the reduce pass. That bound also caps the fp32 workspace, since for a given
+    tile grid it caps M.
+    """
+    if ki.k_split != 1 or K % ki.tile_k:
+        return []
+    base_ctas = ((M + ki.tile_m - 1) // ki.tile_m) * (N // ki.tile_n)
+    if base_ctas >= cu_num:
+        return []
+    n_tiles = K // ki.tile_k
+    max_split = min(
+        n_tiles // K_SPLIT_MIN_TILES_PER_SLICE,
+        max(2, cu_num * K_SPLIT_MAX_CTA_OVERSUBSCRIBE // base_ctas),
+    )
+    return [d for d in range(2, max_split + 1) if n_tiles % d == 0]
+
+
 def _build_kernels_list(tiles, total_vgpr=512):
     kl = {}
     idx = 0