diff --git a/aiter/aot/flydsl/gemm.py b/aiter/aot/flydsl/gemm.py index 56668e1bf5..630257b931 100644 --- a/aiter/aot/flydsl/gemm.py +++ b/aiter/aot/flydsl/gemm.py @@ -52,11 +52,12 @@ parse_wmma_kernel_name as parse_ptpc_wmma_kernel_name, ) from aiter.ops.flydsl.gemm_kernels import ( - SPLIT_K_SEMAPHORE_MAX_LEN, + _split_k_workspace_slots, get_flydsl_splitk_hgemm_kernel_params, ) from aiter.ops.flydsl.kernels.hgemm_dispatch import compile_flydsl_hgemm_kernel from aiter.ops.flydsl.kernels.preshuffle_gemm import compile_preshuffle_gemm +from aiter.ops.flydsl.kernels.splitk_reduce import compile_splitk_reduce_kernel from aiter.ops.flydsl.mxfp8_128_bpreshuffle_gemm_gfx1250 import ( parse_wmma_kernel_name as parse_mxfp8_128_wmma_kernel_name, ) @@ -260,16 +261,9 @@ def _compile_hgemm_to_cache( a = torch.empty((m, k), device=dev, dtype=torch_dtype) b = torch.empty((n, k), device=dev, dtype=torch_dtype) bias = torch.empty((n,), device=dev, dtype=torch_dtype) - semaphore = torch.zeros( - (SPLIT_K_SEMAPHORE_MAX_LEN,), - device=dev, - dtype=torch.int32, - ) - signal = torch.zeros( - (SPLIT_K_SEMAPHORE_MAX_LEN,), - device=dev, - dtype=torch.int32, - ) + # Split-K fp32 workspace slot. Only the pointer's presence matters at + # compile time, never its size, so a 1-element stand-in is enough. + workspace = torch.empty((1,), device=dev, dtype=torch.float32) stream = fx.Stream(0) exe = compile_flydsl_hgemm_kernel( @@ -295,8 +289,8 @@ def _compile_hgemm_to_cache( c_to_lds=c_to_lds, has_bias=has_bias, ) - # FlyDSL JIT does not accept None for tensor slots; pass real buffers for - # optional bias and split-K sync tensors. + # FlyDSL JIT does not accept None for tensor slots; pass a real buffer for + # the optional bias. launch_bias = bias if has_bias else b _compile_executable_to_cache( exe, @@ -305,11 +299,30 @@ def _compile_hgemm_to_cache( _ptr_view_safe(b), _ptr_view_safe(launch_bias), m, - _ptr_view_safe(semaphore), - _ptr_view_safe(signal), + _ptr_view_safe(workspace), stream, ) + # A split-K config is two launches, not one: the main kernel above leaves + # fp32 partials in the workspace and the reduce kernel sums them into C. + # Precompiling only the first would still leave every split-K config JIT-ing + # on first use, which is the cost AOT exists to remove. + if split_k > 1: + reduce_exe = compile_splitk_reduce_kernel( + dtype, + n, + _split_k_workspace_slots(split_k, block_k_warps, kernel_family), + HAS_BIAS=has_bias, + ) + _compile_executable_to_cache( + reduce_exe, + _ptr_view_safe(out), + _ptr_view_safe(workspace), + _ptr_view_safe(launch_bias), + m, + stream, + ) + def _compile_preshuffle_to_cache( *, diff --git a/aiter/ops/flydsl/gemm_kernels.py b/aiter/ops/flydsl/gemm_kernels.py index 6a81d1da6f..95456df502 100644 --- a/aiter/ops/flydsl/gemm_kernels.py +++ b/aiter/ops/flydsl/gemm_kernels.py @@ -22,11 +22,13 @@ from .kernels.hgemm_dispatch import compile_flydsl_hgemm_kernel # from .kernels.small_m_hgemm import iter_small_m_registry_configs +from .kernels.splitk_reduce import compile_splitk_reduce_kernel from .kernels.tensor_shim import _run_compiled from .utils import get_shared_memory_per_block, is_flydsl_available __all__ = [ "flydsl_hgemm", + "flydsl_splitk_prewarm_capture_workspace", ] @@ -36,7 +38,13 @@ def _get_dtypes(): return dtypes -SPLIT_K_SEMAPHORE_MAX_LEN = 256 +# Global accesses in these kernels go through an AMD buffer descriptor, so the +# workspace is bounded by that descriptor rather than by any pointer arithmetic: +# `num_records` is a 32-bit BYTE count (clamped to 0xFFFFFFFF in +# `buffer_ops.create_buffer_resource*`) and the per-lane voffset is a 32-bit +# element offset scaled to bytes. Addressing therefore wraps at 4GiB; cap the +# workspace at half of that so the largest slot offset keeps a 2x margin. +SPLIT_K_WORKSPACE_MAX_BYTES = 1 << 31 FIXED_STAGE = 2 FIXED_C_TO_LDS = False KERNEL_ASYNC_COPY = get_rocm_arch() != "gfx942" @@ -64,11 +72,6 @@ def _get_dtypes(): r"_(?Pgfx[0-9a-z]+)$" ) -SplitKStreamKey = tuple[int, int] -SPLIT_K_GLOBAL_SEMAPHORE: dict[SplitKStreamKey, torch.Tensor] = {} -SPLIT_K_GLOBAL_SIGNAL: dict[SplitKStreamKey, torch.Tensor] = {} - - # Keep the generic auto-generated catalog aligned with the upstream FlyDSL # reference tuning space. The wider local one-off search space introduced # gfx950-faulting candidates (for example tile_k=160 and tile_n=160/192), @@ -175,13 +178,6 @@ def flydsl_kernel_name( return name -def _stream_cache_key(stream: torch.cuda.Stream) -> SplitKStreamKey: - device_index = stream.device.index - if device_index is None: - raise ValueError(f"Unable to determine device index for stream {stream!r}") - return (device_index, int(stream.cuda_stream)) - - def _normalize_launch_stream( device: torch.device, stream: torch.cuda.Stream | None, @@ -697,31 +693,157 @@ def _register_all_configs(): _register_all_configs() -@functools.lru_cache(maxsize=128) -def _get_split_k_tensors( +# --------------------------------------------------------------------------- +# Split-K fp32 workspace (workspace + reduce combine) +# --------------------------------------------------------------------------- +# +# One growable fp32 buffer per device. Growth is monotonic and superseded +# buffers are RETAINED: unlike opus (which dereferences a device-resident +# `ws_handle->ptr` and therefore survives a post-capture grow), FlyDSL bakes the +# raw pointer into the launch args at capture time, so an already-captured graph +# must keep the exact buffer it captured. Keeping the old allocation alive makes +# that pointer valid forever; growth is at least 2x, so the retained total is +# bounded by roughly 2x the live size. +_SPLIT_K_WS: dict[int, torch.Tensor] = {} +_SPLIT_K_WS_RETIRED: list[torch.Tensor] = [] + + +def _split_k_workspace_elems(m: int, n: int, slots: int) -> int: + """Element count of the unpadded `[slots, m, n]` fp32 workspace.""" + return slots * m * n + + +def _split_k_workspace_slots( + split_k: int, block_k_warps: int, kernel_family: str +) -> int: + """Slot count of the `[slots, m, n]` workspace for one kernel config. + + Each slice-K warp group gets its own slot, so its partial is reduced in fp32 + by the reduce kernel instead of through a bf16 LDS combine; small_m has no + slice-K. This is part of the layout contract between the main kernel, the + reduce kernel and the AOT precompiler, so all three read it from here rather + than restating the arithmetic. + """ + slice_k_slots = block_k_warps if kernel_family == KERNEL_FAMILY_HGEMM else 1 + return split_k * slice_k_slots + + +def _get_split_k_workspace( device: torch.device, - stream: torch.cuda.Stream, -) -> tuple[torch.Tensor, torch.Tensor]: - semaphore = torch.zeros( - (SPLIT_K_SEMAPHORE_MAX_LEN,), dtype=torch.int32, device=device - ) - signal = torch.zeros((SPLIT_K_SEMAPHORE_MAX_LEN,), dtype=torch.int32, device=device) - return semaphore, signal + elems: int, + stream: torch.cuda.Stream | None = None, +) -> torch.Tensor: + """Return an fp32 workspace of at least `elems` elements on `device`. + + Never grows during CUDA-graph capture: allocating inside capture would put + the buffer in the graph's private pool and is exactly the class of + capture-time state the redesign removes. Callers that capture should size + the workspace first (see `flydsl_splitk_prewarm_capture_workspace`). + """ + if device.type != "cuda": + raise ValueError(f"split-K workspace requires a CUDA device, got {device}") + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + device = torch.device("cuda", device_index) + + nbytes = elems * 4 + if nbytes > SPLIT_K_WORKSPACE_MAX_BYTES: + raise ValueError( + f"FlyDSL split-K workspace would need {nbytes} bytes, above the " + f"{SPLIT_K_WORKSPACE_MAX_BYTES}-byte limit imposed by 32-bit buffer " + "descriptor addressing; use a smaller split_k for this shape" + ) + + ws = _SPLIT_K_WS.get(device_index) + if ws is not None and ws.numel() >= elems: + return ws + + # Only the grow path below reaches here, so the context managers are off the + # hot path. Both the capture check and the allocation run under them: capture + # state is per (device, current stream), so checking the *current* stream + # while allocating on an explicitly passed `stream` can disagree -- it would + # miss a capturing launch stream and allocate into the graph's private pool, + # which is precisely what this guard exists to prevent. `torch.cuda.stream` + # accepts None as a no-op. + with torch.cuda.device(device), torch.cuda.stream(stream): + if torch.cuda.is_current_stream_capturing(): + have = 0 if ws is None else ws.numel() + raise RuntimeError( + "FlyDSL split-K workspace must be sized before CUDA graph capture " + f"(need {elems} fp32 elements, have {have}). Run this shape eagerly " + "once, or call " + "aiter.ops.flydsl.gemm_kernels.flydsl_splitk_prewarm_capture_workspace(...) " + "on the capture stream, before capturing." + ) + + grow_to = max(elems, 0 if ws is None else 2 * ws.numel()) + if grow_to * 4 > SPLIT_K_WORKSPACE_MAX_BYTES: + grow_to = elems + new_ws = torch.empty(grow_to, dtype=torch.float32, device=device) + if ws is not None: + # Retained, not freed: a captured graph may still hold this pointer. + _SPLIT_K_WS_RETIRED.append(ws) + _SPLIT_K_WS[device_index] = new_ws + return new_ws + +def _graph_capture_stream() -> torch.cuda.Stream: + """The stream `torch.cuda.graph` captures on when called without `stream=`. -def _check_split_k_semaphore_capacity( - m: int, n: int, tile_m: int, tile_n: int, split_k: int + Mirrors torch's own lazy init (and `aiter/tuned_gemm.py:: + _opus_graph_capture_stream`) so the workspace is registered on the exact + stream a later `with torch.cuda.graph(g):` will use. + """ + g = torch.cuda.graphs.graph + if getattr(g, "default_capture_stream", None) is None: + g.default_capture_stream = torch.cuda.Stream() + return g.default_capture_stream + + +def flydsl_splitk_prewarm_capture_workspace( + m: int, + n: int, + *, + split_k: int, + block_k_warps: int = 1, + device: torch.device | None = None, + stream: torch.cuda.Stream | None = None, ) -> None: + """Size the split-K fp32 workspace on the graph capture stream, before capture. + + No-op when already capturing (too late to allocate), for `split_k <= 1` + (that path never touches the workspace), or when the buffer is already big + enough -- the steady state, so this is cheap to call per GEMM from dispatch. + Follows the opus precedent in + `aiter/tuned_gemm.py::_opus_prewarm_capture_workspace`. + """ if split_k <= 1: return - bm = (m + tile_m - 1) // tile_m - bn = n // tile_n - required = bm * bn - if required > SPLIT_K_SEMAPHORE_MAX_LEN: - raise ValueError( - "Split-K semaphore capacity exceeded: " - f"requires {required} counters, max supported is {SPLIT_K_SEMAPHORE_MAX_LEN}" - ) + # Resolve the index too: an index-less `cuda` device misses the cache lookup + # below and would sync on every call. + if device is None or device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + # Capture state is per device and `torch.cuda.Stream()` binds to whichever + # device is current, so resolve both under `device`: prewarming for a + # non-current device would otherwise test, and register the workspace on, + # some other device's stream. + with torch.cuda.device(device): + if torch.cuda.is_current_stream_capturing(): + return + capture_stream = _graph_capture_stream() if stream is None else stream + # Deliberately an upper bound rather than `_split_k_workspace_slots`: + # a prewarm that guesses the family wrong must over-allocate, never + # under-allocate, since under-allocating resurfaces as a hard error at + # capture time. small_m simply leaves the extra slots unused. + elems = _split_k_workspace_elems(m, n, split_k * block_k_warps) + before = _SPLIT_K_WS.get(device.index) + ws = _get_split_k_workspace(device, elems, capture_stream) + # Sync only when this call actually allocated: on the warm path there is + # nothing outstanding, and a host-side sync per GEMM would make this + # unusable from dispatch. + if ws is not before: + capture_stream.synchronize() @functools.lru_cache(maxsize=16384) @@ -820,6 +942,26 @@ def _compile_flydsl_hgemm( has_bias=has_bias, ) + # Split-K combine: the main kernel writes fp32 partials into + # `[slots, m, n]` and a second launch reduces them into C. The stream + # supplies the ordering (a dependency edge between two nodes inside a + # captured graph); nothing is shared between blocks. + is_split_k = split_k > 1 + ws_slots = _split_k_workspace_slots(split_k, block_k_warps, kernel_family) + # `_split_k_workspace_elems` is `slots * m * n`; only `m` varies per call, so + # precompute the constant factor and keep the launcher to one multiply. + ws_slots_n = ws_slots * n + reduce_kernel = ( + compile_splitk_reduce_kernel( + dtype, + n, + ws_slots, + HAS_BIAS=has_bias, + ) + if is_split_k + else None + ) + def launcher( out: torch.Tensor, a: torch.Tensor, @@ -839,18 +981,48 @@ def launcher( launch_bias = b if bias is None else bias runtime_m = int(a.shape[0]) launch_stream = _normalize_launch_stream(a.device, stream) - _check_split_k_semaphore_capacity(runtime_m, n, tile_m, tile_n, split_k) - semaphore, signal = _get_split_k_tensors(a.device, launch_stream) - return _run_compiled( + if not is_split_k: + # The workspace slot is never dereferenced without split-K; the + # kernel ABI just needs a valid pointer there. + return _run_compiled( + kernel, + ptr_arg(out), + ptr_arg(a), + ptr_arg(b), + ptr_arg(launch_bias), + runtime_m, + ptr_arg(out), + fx.Stream(launch_stream), + ) + workspace = _get_split_k_workspace( + a.device, + runtime_m * ws_slots_n, + launch_stream, + ) + # The split-K path issues two launches, so build the argument wrappers + # that both share exactly once. `ptr_arg` and `fx.Stream` each allocate + # a python object per call, and this launcher is on the eager hot path. + out_ptr = ptr_arg(out) + ws_ptr = ptr_arg(workspace) + bias_ptr = ptr_arg(launch_bias) + fx_stream = fx.Stream(launch_stream) + _run_compiled( kernel, - ptr_arg(out), + out_ptr, ptr_arg(a), ptr_arg(b), - ptr_arg(launch_bias), + bias_ptr, + runtime_m, + ws_ptr, + fx_stream, + ) + return _run_compiled( + reduce_kernel, + out_ptr, + ws_ptr, + bias_ptr, runtime_m, - ptr_arg(semaphore), - ptr_arg(signal), - fx.Stream(launch_stream), + fx_stream, ) return launcher diff --git a/aiter/ops/flydsl/kernels/small_m_hgemm.py b/aiter/ops/flydsl/kernels/small_m_hgemm.py index 042a791beb..8f99ed8f68 100644 --- a/aiter/ops/flydsl/kernels/small_m_hgemm.py +++ b/aiter/ops/flydsl/kernels/small_m_hgemm.py @@ -500,8 +500,8 @@ def compile_small_m_hgemm_kernel( BLOCK_K_BYTES = BLOCK_K * DTYPE_BYTES - # LDS layout: C output (and the split-K arrival counter) alias the A tile - # region; B has its own field only on the B_TO_LDS path. + # LDS layout: the C output field aliases the A tile region; B has its own + # field only on the B_TO_LDS path. A_FIELD_ELEMS = max(STAGES * BLOCK_M * BLOCK_K, BLOCK_M * BLOCK_N) B_FIELD_ELEMS = STAGES * BLOCK_N * BLOCK_K if B_TO_LDS else 0 assert (A_FIELD_ELEMS + B_FIELD_ELEMS) * DTYPE_BYTES <= MAX_LDS_BYTES @@ -550,12 +550,9 @@ def small_m_hgemm_kernel( B: fx.Pointer, BIAS: fx.Pointer, m: fx.Int32, - semaphore: fx.Pointer, - signal: fx.Pointer, + WS: fx.Pointer, ): dtype_ = get_dtype_in_kernel(dtype) - _ptr_type = ir.Type.parse("!llvm.ptr<1>") - _i64_type = T.i64 c_zero_d = arith.constant(0.0, type=dtype_) acc_init = arith.constant_vector(0.0, T.vec(WMMA_C_FRAG_VALUES, T.f32)) zero_a_vec = vector.broadcast(T.vec(LDG_VEC_SIZE, dtype_), c_zero_d) @@ -575,8 +572,7 @@ def small_m_hgemm_kernel( # LDS accessors: linear element offsets mirroring the old STensor shapes. # as_/bs_ = (stage, row, col) over (STAGES, BLOCK*, BLOCK_K); cs_ = - # (row, col) over (BLOCK_M, BLOCK_N) aliasing the A field; the split-K - # arrival counter reinterprets the A field as i32. + # (row, col) over (BLOCK_M, BLOCK_N) aliasing the A field. def as_store(stage, row, col, value): elem_off = ( fx.Int64(stage) * (BLOCK_M * BLOCK_K) @@ -619,9 +615,11 @@ def cs_load_vec(row, col, vec_size): ) if const_expr(IS_SPLIT_K): - bc_i32_ptr = fx.recast_iter(fx.Int32, a_lds_ptr) - semaphore_ = GTensor(semaphore, dtype=T.i32, shape=(-1,)) - signal_ = GTensor(signal, dtype=T.i32, shape=(-1,)) + # fp32 split-K workspace, logically [SPLIT_K, m, N] flattened to + # [SPLIT_K * m, N] (no slice-K here, so one slot per split). + # Unpadded: stores are masked to row < m. + # Bias is folded by the reduce kernel. + WS_ = GTensor(WS, dtype=T.f32, shape=(-1, n)) tid = fx.Int32(fx.thread_idx.x) wid = tid // WARP_SIZE @@ -650,11 +648,6 @@ def cs_load_vec(row, col, vec_size): ) for tile_block_n_idx in tile_block_n_indices ] - tile_signal_indices = [ - fx.block_idx.x * fx.Int32(block_n_tiles) - + arith.index_cast(T.i32, tile_block_n_idx) - for tile_block_n_idx in tile_block_n_indices - ] k_blocks16 = fx.Int32(BLOCK_K_BYTES // 16) warp_m_idx = fx.Int32(0) @@ -671,136 +664,6 @@ def cs_load_vec(row, col, vec_size): zero_b_frag = vector.broadcast(B_FRAG_T, c_zero_d) c_frags = [acc_init] * (C_FRAGS_LEN * N_TILE_REPEAT) - def zero_c_tile(c_g, bias_g, tile_n_offset): - zero_vec = vector.broadcast(T.vec(LDG_VEC_SIZE, dtype_), c_zero_d) - for i in range_constexpr(LDG_REG_C_COUNT): - global_tid = BLOCK_THREADS * i + tid - m_local_idx = global_tid // LDG_C_X_THREADS - n_local_idx = global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE - row_idx = m_offset + fx.Index(m_local_idx) - init_vec = zero_vec - if const_expr(HAS_BIAS): - init_vec = bias_g.vec_load( - (tile_n_offset + n_local_idx,), LDG_VEC_SIZE - ) - cond_boundary = arith.cmpi( - arith.CmpIPredicate.ult, row_idx, fx.Index(m) - ) - cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False) - with ir.InsertionPoint(cond_boundary_if.then_block): - c_g.vec_store( - (row_idx, tile_n_offset + n_local_idx), init_vec, LDG_VEC_SIZE - ) - scf.YieldOp([]) - - def get_llvm_ptr(ptr, offset, dtype_bytes): - base_ptr = arith.index_cast(_i64_type, fx.ptrtoint(ptr)) - byte_offset = arith.index_cast( - T.i64, fx.Index(offset) * fx.Index(dtype_bytes) - ) - llvm_ptr = llvm.AddOp( - base_ptr, byte_offset, llvm.IntegerOverflowFlags(0) - ).result - llvm_ptr = llvm.IntToPtrOp(_ptr_type, llvm_ptr).result - return llvm_ptr._value if hasattr(llvm_ptr, "_value") else llvm_ptr - - def prepare_split_k_tile(c_g, bias_g, tile_n_offset, tile_signal_idx): - is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0)) - is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False) - with ir.InsertionPoint(is_t0_cond_if.then_block): - semaphore_ptr = get_llvm_ptr(semaphore, tile_signal_idx, 4) - prev = llvm.AtomicRMWOp( - llvm.AtomicBinOp.add, - semaphore_ptr, - arith.constant(1, type=T.i32), - llvm.AtomicOrdering.monotonic, - syncscope="agent", - alignment=4, - ).result - fx.ptr_store(prev, bc_i32_ptr) - scf.YieldOp([]) - gpu.barrier() - arrive_idx = fx.Index(fx.ptr_load(bc_i32_ptr)) - - first_arrival = arith.cmpi(arith.CmpIPredicate.eq, arrive_idx, fx.Index(0)) - first_arrival_if = scf.IfOp(first_arrival, results_=[], has_else=False) - with ir.InsertionPoint(first_arrival_if.then_block): - zero_c_tile(c_g, bias_g, tile_n_offset) - llvm.InlineAsmOp( - None, - [], - "s_waitcnt vmcnt(0)", - "", - has_side_effects=True, - ) - gpu.barrier() - is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False) - with ir.InsertionPoint(is_t0_cond_if.then_block): - signal_ptr = get_llvm_ptr(signal, tile_signal_idx, 4) - llvm.InlineAsmOp( - None, - [signal_ptr, arith.constant(1, type=T.i32)], - "global_store_dword $0, $1, off sc0 sc1", - "v,v", - has_side_effects=True, - ) - scf.YieldOp([]) - gpu.barrier() - scf.YieldOp([]) - - def split_k_barrier(tile_signal_idx): - init_cur = arith.constant(0, type=T.i32) - w = scf.WhileOp([T.i32], [init_cur]) - before = ir.Block.create_at_start(w.before, [T.i32]) - after = ir.Block.create_at_start(w.after, [T.i32]) - with ir.InsertionPoint(before): - cur = before.arguments[0] - need_wait = arith.CmpIOp( - arith.CmpIPredicate.eq, cur, arith.constant(0, type=T.i32) - ).result - scf.ConditionOp(need_wait, [cur]) - with ir.InsertionPoint(after): - signal_ptr = get_llvm_ptr(signal, tile_signal_idx, 4) - data = llvm.InlineAsmOp( - T.i32, - [signal_ptr], - "global_load_dword $0, $1, off sc1", - "=v,v", - has_side_effects=True, - ).result - rocdl.s_waitcnt(0) - scf.YieldOp([data]) - rocdl.sched_barrier(0) - gpu.barrier() - - is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0)) - is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[T.i32], has_else=True) - with ir.InsertionPoint(is_t0_cond_if.then_block): - semaphore_ptr = get_llvm_ptr(semaphore, tile_signal_idx, 4) - arrive_idx = llvm.AtomicRMWOp( - llvm.AtomicBinOp.add, - semaphore_ptr, - arith.constant(1, type=T.i32), - llvm.AtomicOrdering.monotonic, - syncscope="agent", - alignment=4, - ).result - scf.YieldOp([arrive_idx]) - with ir.InsertionPoint(is_t0_cond_if.else_block): - scf.YieldOp([arith.constant(0, type=T.i32)]) - - last_departure = arith.cmpi( - arith.CmpIPredicate.eq, - is_t0_cond_if.results[0], - arith.constant(2 * SPLIT_K - 1, type=T.i32), - ) - last_departure_if = scf.IfOp(last_departure, results_=[], has_else=False) - with ir.InsertionPoint(last_departure_if.then_block): - semaphore_[tile_signal_idx] = arith.constant(0, type=T.i32) - signal_[tile_signal_idx] = arith.constant(0, type=T.i32) - scf.YieldOp([]) - gpu.barrier() - def ldg_a(k_offset): vecs = [] for i in range_constexpr(LDG_REG_A_COUNT): @@ -962,61 +825,6 @@ def block_mma_sync(a_frags, b_frags, c_frags): ) return c_frags_new - def store_split_k_tile(c_tensor, c_g, tile_n_offset): - out_raw = c_tensor - out_base_int = arith.index_cast(_i64_type, fx.ptrtoint(out_raw)) - for i in range_constexpr(LDG_REG_C_COUNT): - global_tid = BLOCK_THREADS * i + tid - m_local_idx = fx.Index(global_tid // LDG_C_X_THREADS) - n_local_idx = fx.Index(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE) - m_global_idx = m_offset + m_local_idx - n_global_idx = tile_n_offset + n_local_idx - cond_boundary = arith.cmpi( - arith.CmpIPredicate.ult, m_global_idx, fx.Index(m) - ) - cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False) - with ir.InsertionPoint(cond_boundary_if.then_block): - pk_val = cs_load_vec(m_local_idx, n_local_idx, LDG_VEC_SIZE) - linear_bytes_offset = ( - c_g.linear_offset((m_global_idx, n_global_idx)) * DTYPE_BYTES - ) - vec2_ty = T.vec(2, dtype_) - for vec_idx in range_constexpr(LDG_VEC_SIZE // 2): - e0 = vector.extract( - pk_val, - static_position=[vec_idx * 2], - dynamic_position=[], - ) - e1 = vector.extract( - pk_val, - static_position=[vec_idx * 2 + 1], - dynamic_position=[], - ) - pair = vector.from_elements(vec2_ty, [e0, e1]) - pair_byte_offset = arith.index_cast( - T.i64, - linear_bytes_offset + fx.Index(vec_idx * 2 * DTYPE_BYTES), - ) - pair_addr_i64 = llvm.AddOp( - out_base_int, - pair_byte_offset, - llvm.IntegerOverflowFlags(0), - ).result - pair_ptr = llvm.IntToPtrOp(_ptr_type, pair_addr_i64).result - pair_ptr_v = ( - pair_ptr._value if hasattr(pair_ptr, "_value") else pair_ptr - ) - pair_v = pair._value if hasattr(pair, "_value") else pair - llvm.AtomicRMWOp( - llvm.AtomicBinOp.fadd, - pair_ptr_v, - pair_v, - llvm.AtomicOrdering.monotonic, - syncscope="agent", - alignment=4, - ) - scf.YieldOp([]) - def store_c_tile(bias_g, c_g, tile_n_offset): for i in range_constexpr(LDG_REG_C_COUNT): global_tid = BLOCK_THREADS * i + tid @@ -1059,19 +867,43 @@ def write_c_frags_to_lds(tile_c_frags_): ) cs_store_scalar(lds_m_idx, lds_n_idx, val.truncf(dtype_)) - if const_expr(IS_SPLIT_K and not B_TO_LDS): - for tile_i in range_constexpr(N_TILE_REPEAT): - tile_init_if = scf.IfOp( - tile_actives[tile_i], results_=[], has_else=False - ) - with ir.InsertionPoint(tile_init_if.then_block): - prepare_split_k_tile( - C_, - BIAS_, - tile_n_offsets[tile_i], - tile_signal_indices[tile_i], + def store_split_k_tile_ws(tile_c_frags_, tile_n_offset): + """fp32 workspace epilogue: MFMA accumulators straight to slot ks_idx. + + No LDS staging, no barrier, no atomics -- this split owns its slice + of the workspace, so the stores cannot race anything. + + The workspace is [SPLIT_K, m, N]: the slot stride is the real m and + the store is masked to row < m, because the reduce kernel only reads + rows [0, m) of each slot. Writing the whole TILE_M tile unmasked + would be pure write amplification on the skinny-M shapes this family + exists for. + """ + ws_row_base = ks_idx * fx.Index(m) + for ii in range_constexpr(WARP_M_STEPS): + warp_atom_m_idx = warp_m_idx + ii * WARP_ATOM_M + for kk in range_constexpr(WMMA_C_FRAG_VALUES): + m_global = m_offset + fx.Index( + warp_atom_m_idx + stmatrix_c_m_vec_idx + kk ) - scf.YieldOp([]) + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, m_global, fx.Index(m) + ) + row_valid_if = scf.IfOp(row_valid, results_=[], has_else=False) + with ir.InsertionPoint(row_valid_if.then_block): + ws_row = ws_row_base + m_global + for jj in range_constexpr(WARP_N_STEPS): + warp_atom_n_idx = warp_n_idx + jj * WARP_ATOM_N + ws_col = tile_n_offset + fx.Index( + warp_atom_n_idx + stmatrix_c_n_idx + ) + val = vector.extract( + tile_c_frags_[ii * WARP_N_STEPS + jj], + static_position=[kk], + dynamic_position=[], + ) + WS_[(ws_row, ws_col)] = val + scf.YieldOp([]) if const_expr(B_TO_LDS): @@ -1141,10 +973,8 @@ def lds_matrix_b(lds_stage): b_frags[kk * WARP_N_STEPS + ii] = vec return b_frags - def run_b_to_lds_tile(tile_n_offset, tile_signal_idx): + def run_b_to_lds_tile(tile_n_offset): c_frags_local = [acc_init] * C_FRAGS_LEN - if const_expr(IS_SPLIT_K): - prepare_split_k_tile(C_, BIAS_, tile_n_offset, tile_signal_idx) ldg_sts_a_async(ks_begin, 0) ldg_sts_b_async(ks_begin, 0, tile_n_offset) @@ -1233,23 +1063,23 @@ def hot_loop_scheduler(): b_frags = lds_matrix_b(current_stage) c_frags_local = block_mma_sync(a_frags, b_frags, c_frags_local) - write_c_frags_to_lds(c_frags_local) - gpu.barrier() if const_expr(IS_SPLIT_K): - split_k_barrier(tile_signal_idx) - store_split_k_tile(C, C_, tile_n_offset) + store_split_k_tile_ws(c_frags_local, tile_n_offset) + # Separates this tile's LDS reads from the next tile's + # async LDS writes; the C path itself never touches LDS. + gpu.barrier() else: + write_c_frags_to_lds(c_frags_local) + gpu.barrier() store_c_tile(BIAS_, C_, tile_n_offset) - gpu.barrier() + gpu.barrier() for tile_i in range_constexpr(tile_group): tile_exec_if = scf.IfOp( tile_actives[tile_i], results_=[], has_else=False ) with ir.InsertionPoint(tile_exec_if.then_block): - run_b_to_lds_tile( - tile_n_offsets[tile_i], tile_signal_indices[tile_i] - ) + run_b_to_lds_tile(tile_n_offsets[tile_i]) scf.YieldOp([]) else: sts_a(ldg_a(ks_begin), 0) @@ -1370,14 +1200,15 @@ def hot_loop_scheduler(): tile_actives[tile_i], results_=[], has_else=False ) with ir.InsertionPoint(tile_store_if.then_block): - write_c_frags_to_lds(tile_c_frags[tile_i]) - gpu.barrier() if const_expr(IS_SPLIT_K): - split_k_barrier(tile_signal_indices[tile_i]) - store_split_k_tile(C, C_, tile_n_offsets[tile_i]) + store_split_k_tile_ws( + tile_c_frags[tile_i], tile_n_offsets[tile_i] + ) else: + write_c_frags_to_lds(tile_c_frags[tile_i]) + gpu.barrier() store_c_tile(BIAS_, C_, tile_n_offsets[tile_i]) - gpu.barrier() + gpu.barrier() scf.YieldOp([]) @flyc.jit @@ -1387,8 +1218,7 @@ def launch_small_m_hgemm_kernel( B: fx.Pointer, BIAS: fx.Pointer, m: fx.Int32, - semaphore: fx.Pointer, - signal: fx.Pointer, + WS: fx.Pointer, stream: fx.Stream, ): ctx = CompilationContext.get_current() @@ -1403,7 +1233,7 @@ def launch_small_m_hgemm_kernel( tile_group = PERSISTENT_N_TILES if const_expr(PERSISTENT_N) else N_TILE_REPEAT bn = (n // BLOCK_N + tile_group - 1) // tile_group small_m_hgemm_kernel._func.__name__ = KERNEL_NAME - small_m_hgemm_kernel(C, A, B, BIAS, m, semaphore, signal).launch( + small_m_hgemm_kernel(C, A, B, BIAS, m, WS).launch( grid=(bm, bn, SPLIT_K), block=(BLOCK_THREADS, 1, 1), stream=stream, diff --git a/aiter/ops/flydsl/kernels/splitk_hgemm.py b/aiter/ops/flydsl/kernels/splitk_hgemm.py index bdcbb0a228..8e028a451f 100644 --- a/aiter/ops/flydsl/kernels/splitk_hgemm.py +++ b/aiter/ops/flydsl/kernels/splitk_hgemm.py @@ -16,8 +16,6 @@ from .tensor_shim import GTensor, get_dtype_in_kernel -SPLIT_K_SEMAPHORE_MAX_LEN = 256 - def swizzle_xor16(row, col_in_bytes, k_blocks16): return col_in_bytes ^ ((row % k_blocks16) * 16) @@ -250,11 +248,9 @@ def hgemm_kernel( B: fx.Pointer, BIAS: fx.Pointer, m: fx.Int32, - semaphore: fx.Pointer, - signal: fx.Pointer, + WS: fx.Pointer, ): dtype_ = get_dtype_in_kernel(dtype) - c_zero_d = arith.constant(0.0, type=dtype_) acc_init = arith.constant_vector(0.0, T.vec(WMMA_C_FRAG_VALUES, T.f32)) A_ = GTensor(A, dtype=dtype_, shape=(-1, k)) @@ -262,6 +258,11 @@ def hgemm_kernel( C_ = GTensor(C, dtype=dtype_, shape=(-1, n)) if const_expr(HAS_BIAS): BIAS_ = GTensor(BIAS, dtype=dtype_, shape=(n,)) + if const_expr(IS_SPLIT_K): + # fp32 split-K workspace, logically [SLOTS, m, N] flattened to + # [SLOTS * m, N], SLOTS = SPLIT_K * BLOCK_K_WARPS. Unpadded: stores + # are masked to row < m. Bias is folded by the reduce kernel. + WS_ = GTensor(WS, dtype=T.f32, shape=(-1, n)) lds = fx.SharedAllocator().allocate(SharedStorage) a_lds_ptr = lds.pipeline.a_lds.peek().ptr c_lds_ptr = lds.c_lds.peek().ptr @@ -339,11 +340,6 @@ def cs_load_vec(k_slice, row, col, vec_size): result_type=fx.Vector.make_type(vec_size, fx_dtype), ) - if const_expr(IS_SPLIT_K): - semaphore_ = GTensor(semaphore, dtype=T.i32, shape=(-1,)) - signal_ = GTensor(signal, dtype=T.i32, shape=(-1,)) - signal_idx = fx.Int32(fx.block_idx.x) - tid = fx.thread_idx.x wid = tid // WARP_SIZE wid_mn = wid % BLOCK_MN_WARPS @@ -379,130 +375,6 @@ def __barrier(vmcnt=0, use_s_barrier=True): asm = f"s_waitcnt vmcnt({vmcnt})" llvm.InlineAsmOp(None, [], asm, "", has_side_effects=True) - def get_llvm_ptr( - ptr, - offset, - dtype_bytes, - ptr_type=ir.Type.parse("!llvm.ptr<1>"), # noqa: B008 - ): - base_ptr = arith.index_cast(T.i64, fx.ptrtoint(ptr)) - byte_offset = arith.index_cast( - T.i64, fx.Index(offset) * fx.Index(dtype_bytes) - ) - llvm_ptr = llvm.AddOp( - base_ptr, byte_offset, llvm.IntegerOverflowFlags(0) - ).result - llvm_ptr = llvm.IntToPtrOp(ptr_type, llvm_ptr).result - ptr_v = ( - llvm_ptr._value if const_expr(hasattr(llvm_ptr, "_value")) else llvm_ptr - ) - return ptr_v - - def zero_c(): - # zero c if current block is the first block - is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0)) - cond_ks0 = arith.cmpi(arith.CmpIPredicate.eq, ks_idx, fx.Index(0)) - cond_ks0_if = scf.IfOp(cond_ks0, results_=[], has_else=False) - with ir.InsertionPoint(cond_ks0_if.then_block): - zero_vec = vector.broadcast(T.vec(LDG_VEC_SIZE, dtype_), c_zero_d) - for i in range_constexpr(LDG_REG_C_COUNT): - global_tid = BLOCK_THREADS * i + tid - m_local_idx = global_tid // LDG_C_X_THREADS - n_local_idx = global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE - row_idx = m_offset + fx.Index(m_local_idx) - init_vec = zero_vec - if const_expr(HAS_BIAS): - init_vec = BIAS_.vec_load( - (n_offset + n_local_idx,), LDG_VEC_SIZE - ) - cond_boundary = arith.cmpi( - arith.CmpIPredicate.ult, row_idx, fx.Index(m) - ) - cond_boundary_if = scf.IfOp( - cond_boundary, results_=[], has_else=False - ) - with ir.InsertionPoint(cond_boundary_if.then_block): - bytes_offset = C_.linear_offset( - (row_idx, n_offset + n_local_idx) - ) - bytes_offset_i32 = arith.index_cast(T.i32, bytes_offset) - c_ptr = get_llvm_ptr(C, bytes_offset_i32, DTYPE_BYTES) - llvm.InlineAsmOp( - None, - [c_ptr, init_vec], - "global_store_dwordx4 $0, $1, off sc0 sc1", - "v,v", - has_side_effects=True, - ) - scf.YieldOp([]) - gpu.barrier() - # trigger signal when zeroc is done by the first arrived block - is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False) - with ir.InsertionPoint(is_t0_cond_if.then_block): - signal_ptr = get_llvm_ptr(signal, signal_idx, 4) - llvm.InlineAsmOp( - None, - [signal_ptr, arith.constant(1, type=T.i32)], - "global_store_dword $0, $1, off sc0 sc1", - "v,v", - has_side_effects=True, - ) - scf.YieldOp([]) - gpu.barrier() - scf.YieldOp([]) - - def split_k_barrier(): - # spin-wait until signal triggered - is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0)) - is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False) - with ir.InsertionPoint(is_t0_cond_if.then_block): - init_cur = arith.constant(0, type=T.i32) - w = scf.WhileOp([T.i32], [init_cur]) - before = ir.Block.create_at_start(w.before, [T.i32]) - after = ir.Block.create_at_start(w.after, [T.i32]) - with ir.InsertionPoint(before): - cur = before.arguments[0] - need_wait = arith.CmpIOp( - arith.CmpIPredicate.eq, cur, arith.constant(0, type=T.i32) - ).result - scf.ConditionOp(need_wait, [cur]) - with ir.InsertionPoint(after): - signal_ptr = get_llvm_ptr(signal, signal_idx, 4) - data = llvm.InlineAsmOp( - T.i32, - [signal_ptr], - "global_load_dword $0, $1, off sc1", - "=v,v", - has_side_effects=True, - ).result - rocdl.s_waitcnt(0) - scf.YieldOp([data]) - scf.YieldOp([]) - rocdl.sched_barrier(0) - gpu.barrier() - # clean semaphore and signal if this is the last block within split-k group - is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False) - with ir.InsertionPoint(is_t0_cond_if.then_block): - semaphore_ptr = get_llvm_ptr(semaphore, signal_idx, 4) - arrive_idx = llvm.AtomicRMWOp( - llvm.AtomicBinOp.add, - semaphore_ptr, - arith.constant(1, type=T.i32), - llvm.AtomicOrdering.monotonic, - syncscope="agent", - alignment=4, - ).result - cond_ksl = arith.cmpi( - arith.CmpIPredicate.eq, fx.Index(arrive_idx), fx.Index(SPLIT_K - 1) - ) - cond_ksl_if = scf.IfOp(cond_ksl, results_=[], has_else=False) - with ir.InsertionPoint(cond_ksl_if.then_block): - semaphore_[signal_idx] = arith.constant(0, type=T.i32) - signal_[signal_idx] = arith.constant(0, type=T.i32) - scf.YieldOp([]) - scf.YieldOp([]) - gpu.barrier() - def ldg_a(k_offset): vecs = [] for i in range_constexpr(LDG_REG_A_COUNT): @@ -866,9 +738,6 @@ def load_a_frag(m_step, warp_atom_k_idx): warp_offset = get_dma_copy_warp_offset() - if const_expr(IS_SPLIT_K): - zero_c() - if const_expr(B_TO_LDS): for s in range_constexpr(STAGES - 1): @@ -996,9 +865,54 @@ def hot_loop_scheduler(): b_frags = results[2 + C_FRAGS_LEN :] c_frags = ldmatrix_compute_tile_streaming(current_stage, c_frags, b_frags) - # write to lds stmatrix_c_m_vec_idx = w_tid // WMMA_N * WMMA_C_FRAG_VALUES stmatrix_c_n_idx = w_tid % WMMA_N + + if const_expr(IS_SPLIT_K): + # fp32 workspace epilogue: straight from the MFMA accumulators to a + # slot nobody else writes. No LDS staging, no barrier, no atomics. + # Slot = ks_idx * BLOCK_K_WARPS + wid_k, so a slice-K (BLOCK_K_WARPS + # > 1) partial is reduced in fp32 by the reduce kernel too instead of + # via a bf16 LDS combine. + # + # The workspace is [SLOTS, m, N] -- the slot stride is the real m, + # not the grid-aligned M_PAD, and the store is masked to row < m. + # The reduce kernel only ever reads rows [0, m) of each slot + # (row = block_idx.y over a grid of m rows), so an unmasked write + # across the whole BLOCK_M tile is pure write amplification: at + # m=1 with BLOCK_M=32 it moved 32x the bytes the reduce reads. + # The compare costs a few instructions per store; the padded + # traffic costs orders of magnitude more on skinny-M shapes. + ws_slot = ks_idx * fx.Index(BLOCK_K_WARPS) + fx.Index(wid_k) + ws_row_base = ws_slot * fx.Index(m) + for ii in range_constexpr(WARP_M_STEPS): + warp_atom_m_idx = warp_m_idx + ii * WARP_ATOM_M + for kk in range_constexpr(WMMA_C_FRAG_VALUES): + m_global = m_offset + fx.Index( + warp_atom_m_idx + stmatrix_c_m_vec_idx + kk + ) + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, m_global, fx.Index(m) + ) + row_valid_if = scf.IfOp(row_valid, results_=[], has_else=False) + with ir.InsertionPoint(row_valid_if.then_block): + ws_row = ws_row_base + m_global + for jj in range_constexpr(WARP_N_STEPS): + warp_atom_n_idx = warp_n_idx + jj * WARP_ATOM_N + ws_col = n_offset + fx.Index( + warp_atom_n_idx + stmatrix_c_n_idx + ) + val = vector.extract( + c_frags[ii * WARP_N_STEPS + jj], + static_position=[kk], + dynamic_position=[], + ) + WS_[(ws_row, ws_col)] = val + scf.YieldOp([]) + # Traced at compile time: everything below is the split_k == 1 path. + return + + # write to lds gpu.barrier() for ii in range_constexpr(WARP_M_STEPS): warp_atom_m_idx = warp_m_idx + ii * WARP_ATOM_M @@ -1019,76 +933,25 @@ def hot_loop_scheduler(): cs_store_scalar(0, lds_m_idx, lds_n_idx, val) # write back to global - if const_expr(IS_SPLIT_K): - split_k_barrier() - for i in range_constexpr(LDG_REG_C_COUNT): - global_tid = BLOCK_THREADS * i + tid - m_local_idx = fx.Index(global_tid // LDG_C_X_THREADS) - n_local_idx = fx.Index(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE) - m_global_idx = m_offset + m_local_idx - n_global_idx = n_offset + n_local_idx - cond_boundary = arith.cmpi( - arith.CmpIPredicate.ult, m_global_idx, fx.Index(m) - ) - cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False) - with ir.InsertionPoint(cond_boundary_if.then_block): - pk_val = cs_load_vec(0, m_local_idx, n_local_idx, LDG_VEC_SIZE) - for ksi in range_constexpr(1, BLOCK_K_WARPS): - pk_val += cs_load_vec( - ksi, m_local_idx, n_local_idx, LDG_VEC_SIZE - ) - linear_offset_c = C_.linear_offset((m_global_idx, n_global_idx)) - # split to vec2s - vec2_ty = T.vec(2, dtype_) - for vec_idx in range_constexpr(LDG_VEC_SIZE // 2): - e0 = vector.extract( - pk_val, static_position=[vec_idx * 2], dynamic_position=[] - ) - e1 = vector.extract( - pk_val, - static_position=[vec_idx * 2 + 1], - dynamic_position=[], - ) - pair = vector.from_elements(vec2_ty, [e0, e1]) - pair_v = ( - pair._value if const_expr(hasattr(pair, "_value")) else pair - ) - pair_ptr_v = get_llvm_ptr( - C, fx.Int32(linear_offset_c + vec_idx * 2), DTYPE_BYTES - ) - llvm.AtomicRMWOp( - llvm.AtomicBinOp.fadd, - pair_ptr_v, - pair_v, - llvm.AtomicOrdering.monotonic, - syncscope="agent", - alignment=4, - ) - scf.YieldOp([]) - else: - gpu.barrier() - for i in range_constexpr(LDG_REG_C_COUNT): - global_tid = BLOCK_THREADS * i + tid - m_local_idx = fx.Index(global_tid // LDG_C_X_THREADS) - n_local_idx = fx.Index(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE) - m_global_idx = m_offset + m_local_idx - cond_boundary = arith.cmpi( - arith.CmpIPredicate.ult, m_global_idx, fx.Index(m) - ) - cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False) - with ir.InsertionPoint(cond_boundary_if.then_block): - vec = cs_load_vec(0, m_local_idx, n_local_idx, LDG_VEC_SIZE) - for ksi in range_constexpr(1, BLOCK_K_WARPS): - vec += cs_load_vec(ksi, m_local_idx, n_local_idx, LDG_VEC_SIZE) - if const_expr(HAS_BIAS): - bias_vec = BIAS_.vec_load( - (n_offset + n_local_idx,), LDG_VEC_SIZE - ) - vec = vec + bias_vec - C_.vec_store( - (m_global_idx, n_offset + n_local_idx), vec, LDG_VEC_SIZE - ) - scf.YieldOp([]) + gpu.barrier() + for i in range_constexpr(LDG_REG_C_COUNT): + global_tid = BLOCK_THREADS * i + tid + m_local_idx = fx.Index(global_tid // LDG_C_X_THREADS) + n_local_idx = fx.Index(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE) + m_global_idx = m_offset + m_local_idx + cond_boundary = arith.cmpi( + arith.CmpIPredicate.ult, m_global_idx, fx.Index(m) + ) + cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False) + with ir.InsertionPoint(cond_boundary_if.then_block): + vec = cs_load_vec(0, m_local_idx, n_local_idx, LDG_VEC_SIZE) + for ksi in range_constexpr(1, BLOCK_K_WARPS): + vec += cs_load_vec(ksi, m_local_idx, n_local_idx, LDG_VEC_SIZE) + if const_expr(HAS_BIAS): + bias_vec = BIAS_.vec_load((n_offset + n_local_idx,), LDG_VEC_SIZE) + vec = vec + bias_vec + C_.vec_store((m_global_idx, n_offset + n_local_idx), vec, LDG_VEC_SIZE) + scf.YieldOp([]) @flyc.jit def launch_hgemm_kernel( @@ -1097,8 +960,7 @@ def launch_hgemm_kernel( B: fx.Pointer, BIAS: fx.Pointer, m: fx.Int32, - semaphore: fx.Pointer, - signal: fx.Pointer, + WS: fx.Pointer, stream: fx.Stream, ): bm = (m + BLOCK_M - 1) // BLOCK_M @@ -1111,7 +973,7 @@ def launch_hgemm_kernel( if USE_8WAVE_PIPE else None ) - hgemm_kernel(C, A, B, BIAS, m, semaphore, signal).launch( + hgemm_kernel(C, A, B, BIAS, m, WS).launch( grid=(bm * N_BLOCKS, SPLIT_K, 1), block=(BLOCK_THREADS, 1, 1), stream=stream, diff --git a/aiter/ops/flydsl/kernels/splitk_reduce.py b/aiter/ops/flydsl/kernels/splitk_reduce.py new file mode 100644 index 0000000000..2d77d20833 --- /dev/null +++ b/aiter/ops/flydsl/kernels/splitk_reduce.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Split-K reduce kernel for the FlyDSL HGEMM family. + +Tile-agnostic companion to the split-K path in `splitk_hgemm.py` / +`small_m_hgemm.py`: the main kernel writes each split's **fp32** accumulator +into its own disjoint slice of a workspace laid out ``[SLOTS, m, N]``; +this kernel sums the slices, folds bias once, casts to the output dtype and +writes C exactly once. + +Ordering between the two launches comes from the stream (a dependency edge +between two nodes inside a captured graph). Nothing is shared between blocks, +so there is no coordination state a CUDA-graph replay can poison. + +Design mirrors `csrc/opus_gemm/include/gfx950/splitk_reduce_gfx950.cuh` (that +one is C++ with no python binding, so it is a reference, not a callable). + +Layout contract (must match the main kernel): + * The workspace is unpadded ``[SLOTS, m, N]``: slice ``s`` starts at element + ``s * m * N``. The main kernel masks its stores to ``row < m``, because + this kernel only ever reads rows ``[0, m)`` of a slot (grid.y covers m + real rows), so writing a grid-aligned M_PAD tile would be pure write + amplification -- 32x at m=1 with BLOCK_M=32. + * ``SLOTS = SPLIT_K * BLOCK_K_WARPS`` for the generic hgemm family (each + slice-K warp group gets its own slot so the slice combine is also fp32), + ``SLOTS = SPLIT_K`` for small_m (no slice-K there). + +Grid: ``(ceil(N / (VEC * THREADS)), m, 1)``; each thread owns ``VEC`` fp32 lanes +along N (one ``buffer_load_dwordx4`` per slot for ``VEC=4``). +""" + +from __future__ import annotations + +import functools + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import scf +from flydsl.expr import arith, const_expr, range_constexpr +from flydsl.expr.typing import T + +from .tensor_shim import GTensor, get_dtype_in_kernel + +__all__ = [ + "REDUCE_VEC", + "compile_splitk_reduce_kernel", + "reduce_block_threads", +] + +# fp32 lanes per thread. 4 == one buffer_load_dwordx4 per split slot, and the +# bf16 store is a single dwordx2. Every supported N is a multiple of 16 (tile_n +# is a multiple of 16 and N % tile_n == 0), so VEC=4 never produces a partial +# vector tail. +REDUCE_VEC = 4 +REDUCE_MIN_THREADS = 64 +REDUCE_MAX_THREADS = 256 + + +def reduce_block_threads(n: int, vec: int = REDUCE_VEC) -> int: + """Block size for the reduce grid: enough waves to cover N, capped.""" + lanes = max(1, (n + vec - 1) // vec) + threads = ((lanes + 63) // 64) * 64 + return max(REDUCE_MIN_THREADS, min(REDUCE_MAX_THREADS, threads)) + + +@functools.lru_cache(maxsize=4096) +def compile_splitk_reduce_kernel( + dtype: str, + n: int, + SLOTS: int, + HAS_BIAS: bool = False, + VEC: int = REDUCE_VEC, + THREADS: int | None = None, +): + if SLOTS < 2: + raise ValueError(f"split-K reduce needs SLOTS >= 2, got {SLOTS}") + if n % VEC != 0: + raise ValueError(f"split-K reduce needs n % {VEC} == 0, got n={n}") + if THREADS is None: + THREADS = reduce_block_threads(n, VEC) + + TILE_N = VEC * THREADS + N_TILES = (n + TILE_N - 1) // TILE_N + + KERNEL_NAME = ( + f"splitk_reduce_{dtype}_n{n}_s{SLOTS}_v{VEC}x{THREADS}" + f"{'_BIAS' if HAS_BIAS else ''}" + ) + + @flyc.kernel(known_block_size=[THREADS, 1, 1]) + def splitk_reduce_kernel( + C: fx.Pointer, + WS: fx.Pointer, + BIAS: fx.Pointer, + m: fx.Int32, + ): + dtype_ = get_dtype_in_kernel(dtype) + f32_vec_t = T.vec(VEC, T.f32) + out_vec_t = T.vec(VEC, dtype_) + + C_ = GTensor(C, dtype=dtype_, shape=(-1, n)) + WS_ = GTensor(WS, dtype=T.f32, shape=(-1, n)) + if const_expr(HAS_BIAS): + BIAS_ = GTensor(BIAS, dtype=dtype_, shape=(n,)) + + tid = fx.Index(fx.thread_idx.x) + row = fx.Index(fx.block_idx.y) + n_base = fx.Index(fx.block_idx.x) * TILE_N + tid * VEC + + # Slot stride: the main kernel writes an unpadded [SLOTS, m, N] + # workspace, masking its stores to row < m. + m_rows = fx.Index(m) + + in_range = arith.cmpi(arith.CmpIPredicate.ult, n_base, fx.Index(n)) + in_range_if = scf.IfOp(in_range, results_=[], has_else=False) + with ir.InsertionPoint(in_range_if.then_block): + acc = WS_.vec_load((row, n_base), VEC) + # Fixed summation order -> bit-reproducible across runs. + for s in range_constexpr(1, SLOTS): + part = WS_.vec_load((row + m_rows * s, n_base), VEC) + acc = arith.addf(acc, part) + if const_expr(HAS_BIAS): + bias_vec = BIAS_.vec_load((n_base,), VEC) + acc = arith.addf(acc, arith.extf(f32_vec_t, bias_vec)) + C_.vec_store((row, n_base), arith.truncf(out_vec_t, acc), VEC) + scf.YieldOp([]) + + @flyc.jit + def launch_splitk_reduce_kernel( + C: fx.Pointer, + WS: fx.Pointer, + BIAS: fx.Pointer, + m: fx.Int32, + stream: fx.Stream, + ): + splitk_reduce_kernel._func.__name__ = KERNEL_NAME + splitk_reduce_kernel(C, WS, BIAS, m).launch( + grid=(N_TILES, m, 1), + block=(THREADS, 1, 1), + stream=stream, + ) + + return launch_splitk_reduce_kernel diff --git a/aiter/ops/flydsl/test_flydsl_splitk_hgemm.py b/aiter/ops/flydsl/test_flydsl_splitk_hgemm.py index 2f7201fad0..ee2a1ada14 100644 --- a/aiter/ops/flydsl/test_flydsl_splitk_hgemm.py +++ b/aiter/ops/flydsl/test_flydsl_splitk_hgemm.py @@ -39,6 +39,12 @@ DEFAULT_PASS_PCT = 99.9 DEFAULT_INPUT_SEED = 20260401 +# `max_delta_limit` is set to ONE bf16 ULP at each case's output magnitude +# (ulp(x) = 2^(floor(log2 x) - 7)): the split-K combine sums fp32 partials and +# rounds once, so the result is correctly rounded and any multi-ULP drift is a +# real regression. Measured max_delta is 0.0 on all four cases -- the earlier +# 4-ULP limits and 99% pass rates existed to tolerate the bf16 atomic combine +# that has since been removed. SPLITK_PRECISION_CASES = [ { "name": "splitk8_tile32_m104_n384_k7168", @@ -51,6 +57,8 @@ "pack_n": 1, "split_k": 8, "b_preshuffle": False, + "pass_pct": 100.0, + "max_delta_limit": 8.0, # max|ref| ~1856 }, { "name": "splitk4_tile16_m1_n7168_k512", @@ -63,6 +71,8 @@ "pack_n": 1, "split_k": 2, "b_preshuffle": False, + "pass_pct": 100.0, + "max_delta_limit": 1.0, # max|ref| ~142 }, { "name": "splitk16_tile32_m1_n2112_k7168_warp2x2_blds", @@ -78,8 +88,8 @@ "block_n_warps": 2, "b_to_lds": True, "b_preshuffle": False, - "pass_pct": 99.0, - "max_delta_limit": 32.0, + "pass_pct": 100.0, + "max_delta_limit": 8.0, # max|ref| ~1808 }, { "name": "splitk8_tile32_m1_n3072_k1536_warp2x2_blds", @@ -95,8 +105,8 @@ "block_n_warps": 2, "b_to_lds": True, "b_preshuffle": False, - "pass_pct": 99.0, - "max_delta_limit": 8.0, + "pass_pct": 100.0, + "max_delta_limit": 2.0, # max|ref| ~390 }, ] diff --git a/aiter/tuned_gemm.py b/aiter/tuned_gemm.py index 74d7ff2850..237b3142ba 100644 --- a/aiter/tuned_gemm.py +++ b/aiter/tuned_gemm.py @@ -137,6 +137,28 @@ def _opus_prewarm_capture_workspace(inp, weights, solidx, splitK, bias, otype): ) +def _flydsl_prewarm_capture_workspace(inp, weights, flydsl_config) -> None: + """The opus analogue above, for the FlyDSL split-K workspace: size it before + capture, since it cannot grow inside one. Cheap once warm; soft-fails so a + prewarm that cannot run never breaks eager callers. + """ + try: + aiter.ops.flydsl.gemm_kernels.flydsl_splitk_prewarm_capture_workspace( + inp.shape[0], + weights.shape[0], + split_k=int(flydsl_config.get("split_k", 1)), + block_k_warps=int(flydsl_config.get("block_k_warps", 1)), + device=inp.device, + ) + except Exception as e: # noqa: BLE001 + logger.warning( + f"FlyDSL split-K workspace prewarm on the graph capture stream failed " + f"({type(e).__name__}: {e}); HIP graph capture of this shape may abort. " + f"Run it eagerly once before capturing, or call " + f"flydsl_splitk_prewarm_capture_workspace(...) on your capture stream." + ) + + this_dir = os.path.dirname(os.path.abspath(__file__)) @@ -554,6 +576,10 @@ def flydsl_gemm( config["kernelName"] ) stages = flydsl_config.get("stages", flydsl_config.get("stage", 2)) + # Size the split-K workspace on the capture stream now: it cannot grow inside + # a capture, and an unwarmed shape raises there. Same reason and placement as + # _opus_prewarm_capture_workspace in gemm() below. + _flydsl_prewarm_capture_workspace(inp, weights, flydsl_config) fused_bias = None if ( bias is not None diff --git a/op_tests/test_flydsl_splitk_workspace.py b/op_tests/test_flydsl_splitk_workspace.py new file mode 100644 index 0000000000..26ea31a3f9 --- /dev/null +++ b/op_tests/test_flydsl_splitk_workspace.py @@ -0,0 +1,500 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FlyDSL split-K HGEMM: workspace + reduce combine. + +Each split writes an fp32 partial into its own slot of a `[slots, m, n]` +workspace and a separate reduce kernel sums them, so the combine carries no +cross-block coordination state. + +Two independent things live here: + +* **Correctness** (pytest): the result matches an fp32 torch reference, it is + bit-reproducible run to run, and a CUDA-graph replay matches the eager + result. `CASES` covers every structurally distinct split-K epilogue in the + two kernel families -- generic hgemm (plain and slice-K) and small_m (plain, + B_TO_LDS, wide-N repeat, persistent-N). +* **Performance** (`main()`): a `@benchmark` sweep over shape x split_k x + tiling that reports us / TFLOPS / TB per s per config and prints a markdown + summary table. That bandwidth figure deliberately counts the fp32 workspace + round trip, which is this design's main cost (see `test_flydsl_splitk_gemm`). + +Usage: + python op_tests/test_flydsl_splitk_workspace.py # perf sweep + table + pytest -q op_tests/test_flydsl_splitk_workspace.py # correctness only +""" + +from __future__ import annotations + +import argparse +import itertools +from collections.abc import Mapping +from types import MappingProxyType +from typing import Any, NamedTuple + +import pandas as pd +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("ROCm not available. Skipping GPU tests.", allow_module_level=True) + +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx +from aiter.ops.flydsl.utils import is_flydsl_available +from aiter.test_common import benchmark, checkAllclose, run_perftest + +if not is_flydsl_available(): + pytest.skip("flydsl is not installed.", allow_module_level=True) + +# Reuse the library's own tiling validators rather than restating their algebra +# here: a tiling that they reject would abort the launch, and the sweep must +# filter those out before launching (see `_supported`). +from aiter.ops.flydsl.gemm_kernels import ( + _get_split_k_workspace, + _validate_hgemm_tiling, + flydsl_hgemm, + flydsl_splitk_prewarm_capture_workspace, +) +from aiter.ops.flydsl.kernels.small_m_hgemm import _validate_small_m_registry_config + +# Archs the FlyDSL split-K HGEMM family is built and validated for. +SUPPORTED_GFX = ("gfx942", "gfx950") +# `compile_small_m_hgemm_kernel` rejects gfx942 (it targets the async-copy bf16 +# path only), so the small_m cases are gfx950-only. +SMALL_M_GFX = ("gfx950",) + +# Whole-module arch gate for pytest. `main()` gates separately with the same +# allow-list (the skill's rule: gate in main(), not by returning from inside the +# @benchmark fn, which would still emit an args-only row). +pytestmark = pytest.mark.skipif( + get_gfx() not in SUPPORTED_GFX, + reason=f"FlyDSL split-K HGEMM unsupported on {get_gfx()}", +) + +DEFAULT_SEED = 20260401 +# Workspace state each replay starts from: a plain replay first, then values +# that would corrupt the result if the combine read anything it had not +# written itself in the same launch. +REPLAY_POISONS = (None, 1e30, float("nan"), 1e30) +# Fraction of elements allowed to fail isclose(atol=rtol=1e-2) vs the fp32 ref. +MAX_MISMATCH_RATIO = 0.001 + + +class Case(NamedTuple): + """One shape + tiling. `kernel` is passed straight to `flydsl_hgemm`.""" + + name: str + mnk: tuple[int, int, int] + tiles: tuple[int, int, int] # tile_m, tile_n, tile_k + split_k: int + # A NamedTuple default is shared by every instance that omits it, so make it + # actually immutable rather than relying on nobody mutating it. + overrides: Mapping[str, Any] = MappingProxyType({}) + + @property + def kernel(self) -> dict[str, Any]: + # Warp defaults match the runner in test_flydsl_splitk_hgemm.py. + tile_m, tile_n, tile_k = self.tiles + return { + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "split_k": self.split_k, + "block_m_warps": 1, + "block_n_warps": 4, + **self.overrides, + } + + @property + def is_small_m(self) -> bool: + return self.overrides.get("kernel_family") == "small_m" + + +# Shared by four cases below, so immutable for the same reason as the default. +# `|` on a mappingproxy returns a plain dict (PEP 584), which is what the +# per-case overrides want anyway. +SMALL_M = MappingProxyType({"kernel_family": "small_m"}) + +CASES = [ + Case("hgemm_m104_n384_k7168_spk8", (104, 384, 7168), (32, 64, 128), 8), + Case("hgemm_m1_n7168_k512_spk2", (1, 7168, 512), (16, 128, 128), 2), + # Slice-K: BLOCK_K_WARPS=2 gives every k-slice warp group its own workspace + # slot, so the slice combine is fp32 as well. + Case( + "hgemm_m104_n384_k7168_spk8_slicek2", + (104, 384, 7168), + (32, 64, 128), + 8, + {"block_n_warps": 2, "block_k_warps": 2}, + ), + Case("small_m_m1_n7168_k512_spk2", (1, 7168, 512), (16, 128, 128), 2, SMALL_M), + # small_m B_TO_LDS path (`run_b_to_lds_tile`), higher split_k. + Case( + "small_m_m1_n7168_k512_spk8_blds", + (1, 7168, 512), + (16, 128, 64), + 8, + SMALL_M | {"b_to_lds": True}, + ), + # small_m wide-N repeat path (N_TILE_REPEAT > 1, non-B_TO_LDS). + Case( + "small_m_m1_n7168_k512_spk4_nr2", + (1, 7168, 512), + (16, 64, 128), + 4, + SMALL_M | {"block_n_warps": 1, "n_tile_repeat": 2}, + ), + # small_m persistent-N path (PERSISTENT_N_TILES > 1, B_TO_LDS). + Case( + "small_m_m1_n7168_k512_spk4_pn4", + (1, 7168, 512), + (16, 128, 128), + 4, + SMALL_M | {"block_n_warps": 2, "b_to_lds": True, "persistent_n_tiles": 4}, + ), +] + + +def skip_if_unsupported(case: Case) -> None: + if case.is_small_m and get_gfx() not in SMALL_M_GFX: + pytest.skip(f"small_m kernel unsupported on {get_gfx()}") + + +def make_inputs(case: Case, seed: int = DEFAULT_SEED): + gen = torch.Generator(device="cuda") + gen.manual_seed(seed) + m, n, k = case.mnk + kw = {"generator": gen, "device": "cuda", "dtype": torch.bfloat16} + return torch.rand((m, k), **kw), torch.rand((n, k), **kw) + + +def run_torch(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """fp32 reference. Compared against, never timed.""" + return a.float() @ b.float().t() + + +def run_case(case: Case, a: torch.Tensor, b: torch.Tensor, out=None) -> torch.Tensor: + m, n, _ = case.mnk + if out is None: + out = torch.empty((m, n), dtype=torch.bfloat16, device=a.device) + flydsl_hgemm(a, b, out, **case.kernel) + return out + + +@pytest.mark.parametrize("case", CASES, ids=[c.name for c in CASES]) +def test_workspace_matches_reference(case: Case): + skip_if_unsupported(case) + a, b = make_inputs(case) + out = run_case(case, a, b) + torch.cuda.synchronize() + close = torch.isclose(out.float(), run_torch(a, b), atol=1e-2, rtol=1e-2) + mismatch = int((~close).sum().item()) + assert ( + mismatch / close.numel() < MAX_MISMATCH_RATIO + ), f"{case.name}: {mismatch}/{close.numel()} elements outside tolerance" + + +@pytest.mark.parametrize("case", CASES, ids=[c.name for c in CASES]) +def test_workspace_is_deterministic(case: Case): + skip_if_unsupported(case) + a, b = make_inputs(case) + first = run_case(case, a, b).clone() + torch.cuda.synchronize() + second = run_case(case, a, b).clone() + torch.cuda.synchronize() + assert torch.equal(first, second), f"{case.name}: not bit-reproducible" + + +@pytest.mark.parametrize("case", CASES, ids=[c.name for c in CASES]) +def test_graph_replay_matches_eager(case: Case): + """A captured replay must reproduce the eager result exactly, even after + the workspace has been deliberately corrupted. + + The corruption is the point. A clean capture-and-replay passes on the old + cross-block-counter combine too, so on its own it proves nothing; the + counter only wedged the kernel once it was left dirty. Here the equivalent + adversarial state is the workspace itself, so it is filled with a huge + value and with NaN between replays. Both survive because the combine reads + only slots it wrote in the same launch -- there is no carried state. + """ + skip_if_unsupported(case) + a, b = make_inputs(case) + m, n, _ = case.mnk + out = torch.empty((m, n), dtype=torch.bfloat16, device=a.device) + + # Eager pass: compiles the kernels and sizes the workspace. + run_case(case, a, b, out) + torch.cuda.synchronize() + eager = out.clone() + + # The workspace cannot grow during capture, so size it on the stream + # torch.cuda.graph() will capture on first. + flydsl_splitk_prewarm_capture_workspace( + m, + n, + split_k=case.split_k, + block_k_warps=case.overrides.get("block_k_warps", 1), + device=a.device, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_case(case, a, b, out) + + slots = case.split_k * case.overrides.get("block_k_warps", 1) + workspace = _get_split_k_workspace(a.device, slots * m * n) + for replay, poison in enumerate(REPLAY_POISONS): + if poison is not None: + workspace.fill_(poison) + out.zero_() + torch.cuda.synchronize() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager), ( + f"{case.name}: replay {replay} diverged from eager " + f"(workspace pre-filled with {poison})" + ) + + +# --------------------------------------------------------------------------- +# Performance sweep +# --------------------------------------------------------------------------- + +# Tilings the sweep tries per (shape, split_k). Anything the library's +# validators reject for a given shape is filtered out in `main()` before launch. +SWEEP_TILINGS = [ + # tile_m, tile_n, tile_k, block_n_warps, block_k_warps + (16, 128, 128, 4, 1), + (32, 64, 128, 4, 1), + (32, 64, 128, 2, 2), # slice-K: one extra workspace slot per k-slice group + (32, 128, 64, 4, 1), + (64, 128, 128, 4, 1), +] + +# STAGES is fixed at 2 by the kernel, and the pipeline needs at least that many +# K iterations per split (`assert BLOCK_K_LOOPS >= STAGES` in splitk_hgemm.py). +FIXED_STAGES = 2 + + +def _supported( + m, n, k, split_k, tile_m, tile_n, tile_k, block_n_warps, block_k_warps, family +): + """True when this (shape, split_k, tiling) is a launchable config. + + Filters before launch instead of letting the kernel raise, so an + unsupported combination leaves no row in the table rather than a failed one. + """ + if family == "small_m": + # small_m hard-wires tile_m=16 / block_m_warps=1 and rejects gfx942. + if tile_m != 16 or block_k_warps != 1 or get_gfx() not in SMALL_M_GFX: + return False + try: + _validate_small_m_registry_config( + m, + n, + k, + tile_n=tile_n, + tile_k=tile_k, + split_k=split_k, + block_n_warps=block_n_warps, + n_tile_repeat=1, + persistent_n_tiles=1, + waves_per_eu=0, + b_to_lds_unroll=0, + b_to_lds=False, + ) + except ValueError: + return False + return (k // split_k) // tile_k >= FIXED_STAGES + try: + _validate_hgemm_tiling( + m, + n, + k, + dtype="bf16", + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + pack_n=1, + split_k=split_k, + stages=FIXED_STAGES, + block_m_warps=1, + block_n_warps=block_n_warps, + block_k_warps=block_k_warps, + b_to_lds=True, + ) + except ValueError: + return False + return (k // split_k) // tile_k >= FIXED_STAGES + + +@benchmark() +def test_flydsl_splitk_gemm( + m, n, k, split_k, family, tile_m, tile_n, tile_k, block_n_warps, block_k_warps +): + kernel = { + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "split_k": split_k, + "block_m_warps": 1, + "block_n_warps": block_n_warps, + "block_k_warps": block_k_warps, + } + if family == "small_m": + kernel["kernel_family"] = "small_m" + + gen = torch.Generator(device="cuda") + gen.manual_seed(DEFAULT_SEED) + a = torch.rand((m, k), generator=gen, device="cuda", dtype=dtypes.bf16) + b = torch.rand((n, k), generator=gen, device="cuda", dtype=dtypes.bf16) + out = torch.empty((m, n), dtype=dtypes.bf16, device="cuda") + ref = run_torch(a, b) + + # A[m,k] @ B[n,k]^T -> C[m,n]: + # FLOPs = 2 * m * n * k (multiply-add). + # Bytes must include the fp32 split-K workspace, because trading the + # in-place atomic combine for a workspace is exactly what this design does: + # every split writes its partial once and the reduce kernel reads it once, + # so the workspace contributes 2x its size. Size it the way the kernel + # actually allocates it: SLOTS = split_k * block_k_warps (each k-slice warp + # group owns a slot) over the real m rows -- the layout is unpadded and the + # stores are masked to row < m. + slots = split_k * (block_k_warps if family != "small_m" else 1) + ws_bytes = 2 * slots * m * n * 4 if split_k > 1 else 0 + flops = 2 * m * n * k + nbytes = (m * k + n * k) * a.element_size() + m * n * out.element_size() + ws_bytes + + ret = {"gfx": get_gfx()} + res, us = run_perftest(lambda: flydsl_hgemm(a, b, out, **kernel)) + err = checkAllclose( + ref.to(dtypes.fp32), + res.to(dtypes.fp32), + rtol=1e-2, + atol=1e-2, + msg=f"{family} split-K hgemm m={m} n={n} k={k} split_k={split_k}", + ) + ret["flydsl us"] = us + ret["flydsl TFLOPS"] = flops / us / 1e6 + ret["flydsl TB/s"] = nbytes / us / 1e6 + ret["flydsl err"] = err + return ret + + +# The skill names the swept function `test_*`, but it takes shape arguments and +# is driven by `main()`, not by pytest. Opt it out of collection so `pytest -q` +# on this file runs only the correctness cases above. +test_flydsl_splitk_gemm.__test__ = False + + +def main(): + # Positive allow-list: an unknown new card must not run an unbuilt kernel. + if get_gfx() not in SUPPORTED_GFX: + aiter.logger.warning( + "FlyDSL split-K HGEMM unsupported on %s; skipping", get_gfx() + ) + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="config input of test", + ) + parser.add_argument( + "-d", + "--dtype", + type=dtypes.str2Dtype, + choices=[dtypes.d_dtypes["bf16"]], + nargs="*", + default="bf16,", + metavar="{bf16}", + help="""Data type. The FlyDSL split-K path is bf16-only. + e.g.: -d bf16""", + ) + parser.add_argument( + "-s", + "--mnk", + type=dtypes.str2tuple, + nargs="*", + default=[ + # The shapes the split-K path is actually selected for: skinny-M + # decode GEMMs where K is long enough to be worth splitting. + (1, 7168, 512), + (1, 7168, 7168), + (1, 3072, 1536), + (16, 7168, 512), + (104, 384, 7168), + (128, 2112, 7168), + ], + help="""Shape (m, n, k). + e.g.: -s 104,384,7168""", + ) + parser.add_argument( + "--split-k", + type=int, + nargs="*", + default=[2, 4, 8], + help="""Split-K factors to sweep. + e.g.: --split-k 4 8""", + ) + parser.add_argument( + "--family", + type=str, + nargs="*", + choices=["hgemm", "small_m"], + default=["hgemm", "small_m"], + help="""Kernel family. + e.g.: --family hgemm""", + ) + args = parser.parse_args() + + for dtype in args.dtype: + del dtype # bf16-only; kept so the axis stays a swept list + df = [] + for (m, n, k), split_k, family, tiling in itertools.product( + args.mnk, args.split_k, args.family, SWEEP_TILINGS + ): + tile_m, tile_n, tile_k, block_n_warps, block_k_warps = tiling + if not _supported( + m, + n, + k, + split_k, + tile_m, + tile_n, + tile_k, + block_n_warps, + block_k_warps, + family, + ): + continue + df.append( + test_flydsl_splitk_gemm( + m, + n, + k, + split_k, + family, + tile_m, + tile_n, + tile_k, + block_n_warps, + block_k_warps, + ) + ) + if not df: + aiter.logger.warning("no supported split-K configs for this sweep") + continue + df = pd.DataFrame(df) + try: + table = df.to_markdown(index=False) + except ImportError: + # to_markdown needs the optional `tabulate` package; plain fallback, + # mirroring op_tests/test_flydsl_grouped_gemm_gfx1250.py. + table = df.to_string(index=False) + aiter.logger.info("flydsl split-K hgemm summary (markdown):\n%s", table) + + +if __name__ == "__main__": + main()