From 727dd89dec224b004fcdd814a527488269de1b1e Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Mon, 13 Apr 2026 21:08:19 +0000 Subject: [PATCH 01/34] flydsl implementation of a16w16 gemm, WIP --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 857 ++++++++++++++++++ .../gemm/basic/test_gemm_a16w16.py | 129 ++- 2 files changed, 951 insertions(+), 35 deletions(-) create mode 100644 aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py new file mode 100644 index 00000000000..62132375d6b --- /dev/null +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -0,0 +1,857 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +s +import torch +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, tdm_ops, vector +from flydsl.expr.arith import _to_raw as _raw +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr, get_op_result_or_value +from flydsl.expr import idx2crd +from typing import Optional + +# WMMA 16×16×32 +WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 +WAVE_SIZE = 32 + +# LDS padding per row (8 elements -> 16 bytes) +LDS_PAD_A = 8 +LDS_PAD_B = 8 + +_STAGE_NAMES = tuple(f"stage{i}" for i in range(16)) + +def _apply_activation_scalar(val, activation: str): + """Apply activation to a single f32 scalar.""" + from flydsl._mlir.dialects import math as _math + + if activation == "relu": + zero = arith.constant(0.0, type=T.f32) + return arith.select(val > zero, val, zero) + + elif activation in ("silu", "silu_exp2"): + neg = arith.constant(0.0, type=T.f32) - val + exp_neg = _math.exp(neg) + one = arith.constant(1.0, type=T.f32) + denom = one + exp_neg + return val / denom + + elif activation == "gelu": + import math + inv_sqrt2 = arith.constant(1.0 / math.sqrt(2.0), type=T.f32) + scaled = val * inv_sqrt2 + erf_val = _math.erf(scaled) + one = arith.constant(1.0, type=T.f32) + half = arith.constant(0.5, type=T.f32) + return half * val * (one + erf_val) + + elif activation == "gelu_tanh": + import math + sqrt_2_over_pi = arith.constant(math.sqrt(2.0 / math.pi), type=T.f32) + coeff = arith.constant(0.044715, type=T.f32) + one = arith.constant(1.0, type=T.f32) + two = arith.constant(2.0, type=T.f32) + half = arith.constant(0.5, type=T.f32) + x3 = val * val * val + inner = sqrt_2_over_pi * (val + coeff * x3) + # tanh(z) = 1 - 2/(1 + exp(2z)) + exp2x = _math.exp(two * inner) + tanh_val = one - two / (one + exp2x) + return half * val * (one + tanh_val) + + else: + return val + +def compile_gemm_a16w16( + *, + M: int = 0, + N: int = 0, + K: int, + tile_m: int = 128, + tile_n: int = 128, + tile_k: int = 32, + m_warp: int = 2, + n_warp: int = 4, + in_dtype: str = "fp16", + out_dtype: str = None, + num_buffers: int = 2, + waves_per_eu: int = None, + l2_prefetch_distance: int = 2, + activation: Optional[str] = None, + add_bias: bool = False, + physical_mk: bool = True, # True=M-major (row-major X), False=K-major (col-major X) + physical_kn: bool = False, # False=N-major (row-major W), True=K-major (col-major/transposed W) +): + """Compile and return a launch function for the A16W16 GEMM kernel. + Returns a callable: launch_fn(y, x, w, bias, M, N, stream=stream) + """ + _ = (M, N) + + # ── Input validation ── + if num_buffers < 2: + raise ValueError(f"num_buffers must be >= 2, got {num_buffers}") + if in_dtype not in ("fp16", "bf16"): + raise ValueError(f"in_dtype must be 'fp16' or 'bf16', got {in_dtype!r}") + + effective_waves_per_eu = waves_per_eu + is_f16 = in_dtype == "fp16" + + if out_dtype is None: + out_dtype = "f16" if is_f16 else "bf16" + if out_dtype not in ("f32", "f16", "bf16"): + raise ValueError(f"out_dtype must be 'f32', 'f16', or 'bf16', got {out_dtype!r}") + + elem_bytes = 2 + elem_bytes_d = 2 if out_dtype in ("f16", "bf16") else 4 + num_warps = m_warp * n_warp + block_threads = num_warps * WAVE_SIZE + + # ── Tile dimension validation ── + if K % tile_k != 0: + raise ValueError(f"K must be divisible by tile_k={tile_k}, got K={K}") + if tile_k % WMMA_K != 0: + raise ValueError(f"tile_k must be a multiple of {WMMA_K}, got {tile_k}") + if tile_m % WMMA_M != 0: + raise ValueError(f"tile_m must be a multiple of {WMMA_M}, got {tile_m}") + if tile_n % WMMA_N != 0: + raise ValueError(f"tile_n must be a multiple of {WMMA_N}, got {tile_n}") + if (tile_k & (tile_k - 1)) != 0: + raise ValueError(f"tile_k must be a power of 2 for TDM, got {tile_k}") + + # ── Physical layout validation ── + if physical_kn: + if N == 0: + raise ValueError( + "N must be specified (> 0) at compile time when physical_kn=True, " + "because it is used as the TDM stride for the (K, N) weight layout") + if (tile_n & (tile_n - 1)) != 0: + raise ValueError( + f"tile_n must be a power of 2 when physical_kn=True " + f"(TDM pad_interval requirement), got {tile_n}") + if not physical_mk: + if M == 0: + raise ValueError( + "M must be specified (> 0) at compile time when physical_mk=False, " + "because it is used as the TDM stride for the (K, M) activation layout") + if (tile_m & (tile_m - 1)) != 0: + raise ValueError( + f"tile_m must be a power of 2 when physical_mk=False " + f"(TDM pad_interval requirement), got {tile_m}") + + # ── Warp tile dimensions ── + warp_tile_m = tile_m // m_warp + warp_tile_n = tile_n // n_warp + if warp_tile_m % WMMA_M != 0: + raise ValueError(f"warp_tile_m={warp_tile_m} must be a multiple of {WMMA_M}") + if warp_tile_n % WMMA_N != 0: + raise ValueError(f"warp_tile_n={warp_tile_n} must be a multiple of {WMMA_N}") + + # ── K-dimension tiling ── + num_k_tiles = K // tile_k + if num_k_tiles < num_buffers - 1: + raise ValueError( + f"{num_buffers}-stage buffering requires num_k_tiles >= {num_buffers - 1}, " + f"got {num_k_tiles} (K={K}, tile_k={tile_k})") + + # ── Architecture check ── + gpu_arch = str(get_hip_arch(timeout_s=300)) + assert gpu_arch.startswith("gfx1250"), f"Expected gfx1250, got {gpu_arch}" + + wmma_op = rocdl.wmma_f32_16x16x32_f16 if is_f16 else rocdl.wmma_f32_16x16x32_bf16 + + # ── Compute repetition counts ── + k_wmma_steps = tile_k // WMMA_K + + def _elem_type(): + return T.f16 if is_f16 else T.bf16 + + wmma_m_rep = warp_tile_m // WMMA_M + wmma_n_rep = warp_tile_n // WMMA_N + n_accs = wmma_m_rep * wmma_n_rep + n_frags_per_tile = k_wmma_steps * (wmma_n_rep + wmma_m_rep) + + # ── LDS layout ── + # physical_mk=True: A in LDS as [tile_m, tile_k + pad] + # physical_mk=False: A in LDS as [tile_k, tile_m + pad] (K-major) + if physical_mk: + lds_a_stride = tile_k + LDS_PAD_A + lds_a_elems = tile_m * lds_a_stride + LDS_PAD_A + else: + lds_a_stride = tile_m + LDS_PAD_A + lds_a_elems = tile_k * lds_a_stride + LDS_PAD_A + + if physical_kn: + lds_b_stride = tile_n + LDS_PAD_B + lds_b_elems = tile_k * lds_b_stride + LDS_PAD_B + else: + lds_b_stride = tile_k + LDS_PAD_B + lds_b_elems = tile_n * lds_b_stride + LDS_PAD_B + + buf_size_elems = lds_a_elems + lds_b_elems + + # ── LDS allocation per pipeline stage ── + stage_allocators = [] + stage_a_offsets = [] + stage_b_offsets = [] + for i in range(num_buffers): + name = _STAGE_NAMES[i] + alloc = SmemAllocator(None, arch=gpu_arch, global_sym_name=f"a16w16_{name}") + off = alloc._align(alloc.ptr, 16) + alloc.ptr = off + buf_size_elems * elem_bytes + stage_allocators.append(alloc) + stage_a_offsets.append(off) + stage_b_offsets.append(off + lds_a_elems * elem_bytes) + + # ── Pipeline iteration counts ── + pre_loaded = num_buffers - 1 + main_loop_iters = (num_k_tiles - pre_loaded) // num_buffers + tail_tiles = (num_k_tiles - pre_loaded) % num_buffers + pre_loaded + + @flyc.kernel + def kernel_gemm_a16w16( + arg_y: fx.Tensor, arg_x: fx.Tensor, arg_w: fx.Tensor, + arg_bias: fx.Tensor, i32_m: fx.Int32, i32_n: fx.Int32, + ): + rocdl.disable_xdl_arb_stall() + + # ── Thread/block indexing ── + tx = gpu.thread_id("x") + bx = gpu.block_id("x") + by = gpu.block_id("y") + blk_m = bx * arith.index(tile_m) + blk_n = by * arith.index(tile_n) + + # ── Thread-to-warp decomposition ── + layout_thr = fx.make_layout( + (m_warp, n_warp, 2, 16), + (n_warp * WAVE_SIZE, WAVE_SIZE, 16, 1)) + thr_coord = idx2crd(tx, layout_thr) + wave_m_idx, wave_n_idx, lane_kgrp, lane16 = ( + fx.get(thr_coord, 0), fx.get(thr_coord, 1), + fx.get(thr_coord, 2), fx.get(thr_coord, 3)) + + warp_m_base = wave_m_idx * arith.index(warp_tile_m) + warp_n_base = wave_n_idx * arith.index(warp_tile_n) + elem_ty = _elem_type() + + m_idx = arith.index_cast(T.index, i32_m.ir_value()) + n_stride = arith.index_cast(T.index, i32_n.ir_value()) + + # Buffer resource descriptors + y_nrec = m_idx * n_stride * arith.index(elem_bytes_d) + y_rsrc = buffer_ops.create_buffer_resource(arg_y, num_records_bytes=y_nrec) + if add_bias: + bias_rsrc = buffer_ops.create_buffer_resource(arg_bias, max_size=True) + + def make_a_desc(k_base, lds_a_mem_ref): + """TDM descriptor for A tile. Swaps dims when physical_mk=False (K-major).""" + if physical_mk: + return tdm_ops.make_tensor_descriptor_2d( + global_ptr=arg_x, lds_memref=lds_a_mem_ref, + global_offset=(blk_m, k_base), + tensor_shape=(tile_m, tile_k), strides=(K, 1), + tile_shape=(tile_m, tile_k), elem_bytes=elem_bytes, + pad_interval=tile_k, pad_amount=LDS_PAD_A, + num_warps=num_warps) + else: + return tdm_ops.make_tensor_descriptor_2d( + global_ptr=arg_x, lds_memref=lds_a_mem_ref, + global_offset=(k_base, blk_m), + tensor_shape=(tile_k, tile_m), strides=(M, 1), + tile_shape=(tile_k, tile_m), elem_bytes=elem_bytes, + pad_interval=tile_m, pad_amount=LDS_PAD_A, + num_warps=num_warps) + + def make_b_desc(k_base, lds_b_mem_ref): + """TDM descriptor for B tile. Swaps dims when physical_kn=True (K-major).""" + if physical_kn: + return tdm_ops.make_tensor_descriptor_2d( + global_ptr=arg_w, lds_memref=lds_b_mem_ref, + global_offset=(k_base, blk_n), + tensor_shape=(tile_k, tile_n), strides=(N, 1), + tile_shape=(tile_k, tile_n), elem_bytes=elem_bytes, + pad_interval=tile_n, pad_amount=LDS_PAD_B, + num_warps=num_warps) + return tdm_ops.make_tensor_descriptor_2d( + global_ptr=arg_w, lds_memref=lds_b_mem_ref, + global_offset=(blk_n, k_base), + tensor_shape=(tile_n, tile_k), strides=(K, 1), + tile_shape=(tile_n, tile_k), elem_bytes=elem_bytes, + pad_interval=tile_k, pad_amount=LDS_PAD_B, + num_warps=num_warps) + + def issue_tdm_load(desc): + tdm_ops.tensor_load_2d(desc) + + def copy_a_to_lds(k_base, lds_a_mem_ref): + issue_tdm_load(make_a_desc(k_base, lds_a_mem_ref)) + + def copy_b_to_lds(k_base, lds_b_mem_ref): + issue_tdm_load(make_b_desc(k_base, lds_b_mem_ref)) + + def _get_lds_memref(lds_ptr): + if isinstance(lds_ptr, SmemPtr): + return get_op_result_or_value(lds_ptr.get()) + return get_op_result_or_value(lds_ptr) + + def _precompute_lane_bases(lds_ptr, warp_base, reps, lds_stride): + """Pre-compute LDS base addresses for each WMMA rep (M-major layout).""" + lds_buffer = _get_lds_memref(lds_ptr) + row_stride_off = (warp_base + lane16) * arith.index(lds_stride) + k_lane_off = lane_kgrp * arith.index(8) + bases = [] + for rep in range_constexpr(reps): + base = row_stride_off + arith.index(rep * WMMA_M * lds_stride) + k_lane_off + bases.append(base) + return lds_buffer, bases + + def _precompute_a_lane_bases(lds_ptr): + """A fragment lane bases. Uses transpose-load addressing when K-major.""" + if physical_mk: + return _precompute_lane_bases( + lds_ptr, warp_m_base, wmma_m_rep, lds_a_stride) + + lds_buffer = _get_lds_memref(lds_ptr) + lane8 = lane16 % arith.index(8) + lane_mgrp = lane16 / arith.index(8) + k_lane_off = (lane_kgrp * arith.index(8) + lane8) \ + * arith.index(lds_a_stride) + m_lane_off = lane_mgrp * arith.index(8) + bases = [] + for wm in range_constexpr(wmma_m_rep): + m_col = warp_m_base + arith.index(wm * WMMA_M) + m_lane_off + bases.append(k_lane_off + m_col) + return lds_buffer, bases + + def _precompute_b_lane_bases(lds_ptr): + """B fragment lane bases. Uses transpose-load addressing when K-major.""" + if not physical_kn: + return _precompute_lane_bases( + lds_ptr, warp_n_base, wmma_n_rep, lds_b_stride) + + lds_buffer = _get_lds_memref(lds_ptr) + lane8 = lane16 % arith.index(8) + lane_ngrp = lane16 / arith.index(8) + k_lane_off = (lane_kgrp * arith.index(8) + lane8) \ + * arith.index(lds_b_stride) + n_lane_off = lane_ngrp * arith.index(8) + bases = [] + for wn in range_constexpr(wmma_n_rep): + n_col = warp_n_base + arith.index(wn * WMMA_N) + n_lane_off + bases.append(k_lane_off + n_col) + return lds_buffer, bases + + def load_wmma_frag_tr(lds_buffer, b_lane_base, ks): + """Load WMMA B fragment via ds_load_tr16_b128 (K-major LDS).""" + vec8_ty = ir.VectorType.get([8], elem_ty) + results = [] + for k_half in range_constexpr(2): + k_row_off = (ks * WMMA_K + k_half * 16) * lds_b_stride + elem_off = b_lane_base + arith.index(k_row_off) + v = rocdl.lds_transpose_load( + vec8_ty, lds_buffer, elem_off, elem_bytes) + results.append(v) + return vector.shuffle(results[0], results[1], list(range(16))) + + def load_wmma_frag(lds_buffer, lane_base, ks): + """Load WMMA fragment via ds_read_b128 (M/N-major LDS).""" + vec8_ty = ir.VectorType.get([8], elem_ty) + off0 = lane_base + arith.index(ks * WMMA_K) + off1 = lane_base + arith.index(ks * WMMA_K + 16) + v0 = vector.load_op(vec8_ty, lds_buffer, [off0]) + v1 = vector.load_op(vec8_ty, lds_buffer, [off1]) + return vector.shuffle(v0, v1, list(range(16))) + + _load_b_frag = load_wmma_frag_tr if physical_kn else load_wmma_frag + + def load_wmma_frag_tr_a(lds_buffer, a_lane_base, ks): + """Load WMMA A fragment via ds_load_tr16_b128 (K-major LDS).""" + vec8_ty = ir.VectorType.get([8], elem_ty) + results = [] + for k_half in range_constexpr(2): + k_row_off = (ks * WMMA_K + k_half * 16) * lds_a_stride + elem_off = a_lane_base + arith.index(k_row_off) + v = rocdl.lds_transpose_load( + vec8_ty, lds_buffer, elem_off, elem_bytes) + results.append(v) + return vector.shuffle(results[0], results[1], list(range(16))) + + _load_a_frag = load_wmma_frag_tr_a if not physical_mk else load_wmma_frag + + def load_tile_frags(lds_a_ptr, lds_b_ptr): + """Bulk-load all k-slice fragments for one tile (no overlap).""" + a_buf, a_bases = _precompute_a_lane_bases(lds_a_ptr) + b_buf, b_bases = _precompute_b_lane_bases(lds_b_ptr) + frags = [] + for ks in range_constexpr(k_wmma_steps): + for wn in range_constexpr(wmma_n_rep): + frags.append(_load_b_frag(b_buf, b_bases[wn], ks)) + for wm in range_constexpr(wmma_m_rep): + frags.append(_load_a_frag(a_buf, a_bases[wm], ks)) + return frags + + def wmma_tile(accs_in, tile_frags): + """Execute all WMMAs for one tile using pre-loaded fragments.""" + current_accs = list(accs_in) + for ks in range_constexpr(k_wmma_steps): + base = ks * (wmma_n_rep + wmma_m_rep) + b_frags = [tile_frags[base + wn] + for wn in range_constexpr(wmma_n_rep)] + a_frags = [tile_frags[base + wmma_n_rep + wm] + for wm in range_constexpr(wmma_m_rep)] + for wm in range_constexpr(wmma_m_rep): + for wn in range_constexpr(wmma_n_rep): + idx = wm * wmma_n_rep + wn + # ISA operand order: (B, A, C) — reversed from math + current_accs[idx] = wmma_op( + T.vec(8, T.f32), + b_frags[wn], a_frags[wm], current_accs[idx], + signA=False, signB=False, modC=0, + reuseA=False, reuseB=(wn > 0), + ).result + return current_accs + + def compute_and_prefetch(accs_in, cur_frags, next_lds_a_ptr, next_lds_b_ptr): + """Interleave WMMA on current frags with LDS loads for next tile.""" + current_accs = list(accs_in) + next_a_buf, next_a_bases = _precompute_a_lane_bases(next_lds_a_ptr) + next_b_buf, next_b_bases = _precompute_b_lane_bases(next_lds_b_ptr) + next_frags = [] + + for ks in range_constexpr(k_wmma_steps): + # Phase 1: ds_read for next tile + for wn in range_constexpr(wmma_n_rep): + next_frags.append(_load_b_frag(next_b_buf, next_b_bases[wn], ks)) + for wm in range_constexpr(wmma_m_rep): + next_frags.append(_load_a_frag(next_a_buf, next_a_bases[wm], ks)) + + # Phase 2: WMMA for curr tile + base = ks * (wmma_n_rep + wmma_m_rep) + b_frags = [cur_frags[base + wn] + for wn in range_constexpr(wmma_n_rep)] + a_frags = [cur_frags[base + wmma_n_rep + wm] + for wm in range_constexpr(wmma_m_rep)] + for wm in range_constexpr(wmma_m_rep): + for wn in range_constexpr(wmma_n_rep): + idx = wm * wmma_n_rep + wn + current_accs[idx] = wmma_op( + T.vec(8, T.f32), + b_frags[wn], a_frags[wm], current_accs[idx], + signA=False, signB=False, modC=0, + reuseA=False, reuseB=(wn > 0), + ).result + + return current_accs, next_frags + + def _l2_prefetch(k_base): + """Issue L2 cache prefetch hints for a future K-tile.""" + if l2_prefetch_distance <= 0: + return + pf_k = k_base + arith.index(l2_prefetch_distance * tile_k) + + if physical_mk: + tdm_ops.l2_prefetch_tile( + arg_x, (blk_m, pf_k), (tile_m, tile_k), (K, 1), + elem_bytes=elem_bytes, thread_id=tx, block_threads=block_threads) + else: + tdm_ops.l2_prefetch_tile( + arg_x, (pf_k, blk_m), (tile_k, tile_m), (M, 1), + elem_bytes=elem_bytes, thread_id=tx, block_threads=block_threads) + if physical_kn: + tdm_ops.l2_prefetch_tile( + arg_w, (pf_k, blk_n), (tile_k, tile_n), (N, 1), + elem_bytes=elem_bytes, thread_id=tx, + block_threads=block_threads) + else: + tdm_ops.l2_prefetch_tile( + arg_w, (blk_n, pf_k), (tile_n, tile_k), (K, 1), + elem_bytes=elem_bytes, thread_id=tx, + block_threads=block_threads) + + _half_out = out_dtype in ("f16", "bf16") + _out_elem = T.f16 if out_dtype == "f16" else (T.bf16 if out_dtype == "bf16" else None) + _bias_elem = T.f16 if is_f16 else T.bf16 + + def _widen_to_f32(val): + from flydsl._mlir.dialects import arith as _std_arith + return _std_arith.ExtFOp(T.f32, _raw(val)).result + + def _apply_bias_and_activation(accs): + """Add bias and/or apply activation to accumulators.""" + if not add_bias and activation is None: + return accs + + for wm in range_constexpr(wmma_m_rep): + for wn in range_constexpr(wmma_n_rep): + idx = wm * wmma_n_rep + wn + acc = accs[idx] + + if add_bias: + col_base = (blk_n + warp_n_base + + arith.index(wn * WMMA_N) + + lane_kgrp * arith.index(8)) + col_base_i32 = arith.index_cast(T.i32, col_base) + col_base_hi = col_base_i32 + arith.constant(4, type=T.i32) + + bv_lo = buffer_ops.buffer_load( + bias_rsrc, col_base_i32, + vec_width=4, dtype=_bias_elem) + bv_hi = buffer_ops.buffer_load( + bias_rsrc, col_base_hi, + vec_width=4, dtype=_bias_elem) + + bias_elems = [] + for i in range_constexpr(4): + b_val = vector.extract( + bv_lo, static_position=[i], dynamic_position=[]) + bias_elems.append(_widen_to_f32(b_val)) + for i in range_constexpr(4): + b_val = vector.extract( + bv_hi, static_position=[i], dynamic_position=[]) + bias_elems.append(_widen_to_f32(b_val)) + + bias_vec = vector.from_elements(T.vec(8, T.f32), bias_elems) + acc = acc + bias_vec + + if activation is not None: + new_elems = [] + for i in range_constexpr(8): + val = vector.extract( + acc, static_position=[i], dynamic_position=[]) + val = _apply_activation_scalar(val, activation) + new_elems.append(val) + acc = vector.from_elements(T.vec(8, T.f32), new_elems) + + accs[idx] = acc + return accs + + def epilogue_stores(final_accs): + """Write accumulators to global output Y.""" + for wm in range_constexpr(wmma_m_rep): + for wn in range_constexpr(wmma_n_rep): + idx = wm * wmma_n_rep + wn + row = blk_m + warp_m_base + arith.index(wm * WMMA_M) + lane16 + col_base = (blk_n + warp_n_base + arith.index(wn * WMMA_N) + + lane_kgrp * arith.index(8)) + + if _half_out: + h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) + i32_vec = vector.bitcast(T.vec(4, T.i32), h_vec) + c_off_bytes = (row * n_stride + col_base) * arith.index(elem_bytes_d) + buffer_ops.buffer_store( + i32_vec, y_rsrc, c_off_bytes, offset_is_bytes=True) + else: + for half in range_constexpr(2): + vals = [vector.extract( + final_accs[idx], + static_position=[half * 4 + vi], + dynamic_position=[]) + for vi in range_constexpr(4)] + vec4 = vector.from_elements(T.vec(4, T.f32), vals) + col = col_base + arith.index(half * 4) + c_off = row * n_stride + col + buffer_ops.buffer_store(vec4, y_rsrc, c_off) + + # Step 1: Initialize accumulators + acc_zero = arith.constant_vector(0.0, T.vec(8, T.f32)) + accs = [acc_zero] * n_accs + + # Step 2: Set up LDS buffer pointers + base_ptrs = [sa.get_base() for sa in stage_allocators] + stages_a = [ + SmemPtr(base_ptrs[i], stage_a_offsets[i], elem_ty, shape=(lds_a_elems,)) + for i in range_constexpr(num_buffers) + ] + stages_b = [ + SmemPtr(base_ptrs[i], stage_b_offsets[i], elem_ty, shape=(lds_b_elems,)) + for i in range_constexpr(num_buffers) + ] + stages_a_mem = [stages_a[i].get() for i in range_constexpr(num_buffers)] + stages_b_mem = [stages_b[i].get() for i in range_constexpr(num_buffers)] + + # Step 3: Prologue — pre-load tiles into LDS + for i in range_constexpr(pre_loaded): + copy_a_to_lds(arith.index(i * tile_k), stages_a_mem[i]) + copy_b_to_lds(arith.index(i * tile_k), stages_b_mem[i]) + + # L2 prefetch prologue + if l2_prefetch_distance > 0: + for pf_i in range_constexpr(min(l2_prefetch_distance, num_k_tiles)): + pf_k = arith.index(pf_i * tile_k) + if physical_mk: + tdm_ops.l2_prefetch_tile( + arg_x, (blk_m, pf_k), (tile_m, tile_k), (K, 1), + elem_bytes=elem_bytes, thread_id=tx, + block_threads=block_threads) + else: + tdm_ops.l2_prefetch_tile( + arg_x, (pf_k, blk_m), (tile_k, tile_m), (M, 1), + elem_bytes=elem_bytes, thread_id=tx, + block_threads=block_threads) + if physical_kn: + tdm_ops.l2_prefetch_tile( + arg_w, (pf_k, blk_n), (tile_k, tile_n), (N, 1), + elem_bytes=elem_bytes, thread_id=tx, + block_threads=block_threads) + else: + tdm_ops.l2_prefetch_tile( + arg_w, (blk_n, pf_k), (tile_n, tile_k), (K, 1), + elem_bytes=elem_bytes, thread_id=tx, + block_threads=block_threads) + + # Prepare first TDM descriptors for main loop + if main_loop_iters > 0: + _first_load_stage_0 = pre_loaded % num_buffers + _first_load_k_0 = pre_loaded * tile_k + pending_desc_a = make_a_desc(arith.index(_first_load_k_0), + stages_a_mem[_first_load_stage_0]) + pending_desc_b = make_b_desc(arith.index(_first_load_k_0), + stages_b_mem[_first_load_stage_0]) + + main_end = main_loop_iters * num_buffers * tile_k + + # Wait for pre-loaded tiles + tdm_ops.tensor_wait(2 * (num_buffers - 2)) + gpu.barrier() + + # Pre-load fragments for first compute tile + cur_frags = load_tile_frags(stages_a[0], stages_b[0]) + + # Step 4: Main loop + if main_loop_iters > 0: + init_descs = [pending_desc_a.dgroup0, pending_desc_a.dgroup1, + pending_desc_b.dgroup0, pending_desc_b.dgroup1] + + for iv, state in range(0, main_end, num_buffers * tile_k, + init=list(accs) + init_descs + cur_frags): + accs_in = list(state[:n_accs]) + cur_desc_a = tdm_ops.TDMDescriptor2D(state[n_accs], state[n_accs + 1]) + cur_desc_b = tdm_ops.TDMDescriptor2D(state[n_accs + 2], state[n_accs + 3]) + cur_frags = list(state[n_accs + 4:n_accs + 4 + n_frags_per_tile]) + + for s in range_constexpr(num_buffers): + issue_tdm_load(cur_desc_a) + issue_tdm_load(cur_desc_b) + + if s == num_buffers - 1: + _l2_prefetch(iv + arith.index(s * tile_k)) + + if s < num_buffers - 1: + _next_load_stage = (s + 1 + pre_loaded) % num_buffers + _next_load_k_off = (s + 1 + pre_loaded) * tile_k + cur_desc_a = make_a_desc( + iv + arith.index(_next_load_k_off), + stages_a_mem[_next_load_stage]) + cur_desc_b = make_b_desc( + iv + arith.index(_next_load_k_off), + stages_b_mem[_next_load_stage]) + else: + _next_load_stage = pre_loaded % num_buffers + _next_load_k_off = pre_loaded * tile_k + _next_step = num_buffers * tile_k + cur_desc_a = make_a_desc( + iv + arith.index(_next_step + _next_load_k_off), + stages_a_mem[_next_load_stage]) + cur_desc_b = make_b_desc( + iv + arith.index(_next_step + _next_load_k_off), + stages_b_mem[_next_load_stage]) + + tdm_ops.tensor_wait(2 * (num_buffers - 2)) + gpu.barrier() + + _next_compute = (s + 1) % num_buffers + accs_in, cur_frags = compute_and_prefetch( + accs_in, cur_frags, + stages_a[_next_compute], stages_b[_next_compute]) + + out_descs = [cur_desc_a.dgroup0, cur_desc_a.dgroup1, + cur_desc_b.dgroup0, cur_desc_b.dgroup1] + results = yield list(accs_in) + out_descs + cur_frags + + accs = list(results[:n_accs]) + cur_frags = list(results[n_accs + 4:n_accs + 4 + n_frags_per_tile]) + + # Step 5: drain remaining tiles + _tail_base_k = main_loop_iters * num_buffers + _extra_loads = tail_tiles - pre_loaded + + for t in range_constexpr(tail_tiles - 1): + if t < _extra_loads: + load_tile_idx = _tail_base_k + pre_loaded + t + load_stage = (pre_loaded + t) % num_buffers + copy_a_to_lds(arith.index(load_tile_idx * tile_k), + stages_a_mem[load_stage]) + copy_b_to_lds(arith.index(load_tile_idx * tile_k), + stages_b_mem[load_stage]) + + _epi_outstanding = 2 * (pre_loaded + min(t + 1, _extra_loads) - t - 1) + tdm_ops.tensor_wait(_epi_outstanding) + gpu.barrier() + + _next_epi_stage = (_tail_base_k + t + 1) % num_buffers + accs, cur_frags = compute_and_prefetch( + accs, cur_frags, + stages_a[_next_epi_stage], stages_b[_next_epi_stage]) + + # Step 6: Final WMMA + accs = wmma_tile(accs, cur_frags) + + # Step 7: Bias, activation, and store + accs = _apply_bias_and_activation(accs) + epilogue_stores(accs) + + cache_tag = (in_dtype, out_dtype, K, tile_m, tile_n, tile_k, m_warp, n_warp, + num_buffers, effective_waves_per_eu, l2_prefetch_distance, + activation, add_bias, physical_mk, physical_kn, M, N) + + @flyc.jit + def launch_gemm_a16w16( + arg_y: fx.Tensor, arg_x: fx.Tensor, arg_w: fx.Tensor, + arg_bias: fx.Tensor, i32_m: fx.Int32, i32_n: fx.Int32, + stream: fx.Stream, + ): + _ = cache_tag + ctx = CompilationContext.get_current() + + # Finalize LDS allocations + with ir.InsertionPoint(ctx.gpu_module_body): + for alloc in stage_allocators: + alloc.finalized = False + for alloc in stage_allocators: + alloc.finalize() + + # Grid dimensions + idx_m = arith.index_cast(T.index, i32_m.ir_value()) + idx_n = arith.index_cast(T.index, i32_n.ir_value()) + gx = _raw((idx_m + arith.index(tile_m - 1)) / arith.index(tile_m)) + gy = _raw((idx_n + arith.index(tile_n - 1)) / arith.index(tile_n)) + + # Emit kernel + launcher = kernel_gemm_a16w16(arg_y, arg_x, arg_w, arg_bias, i32_m, i32_n) + + # Set waves_per_eu + for op in ctx.gpu_module_body.operations: + if hasattr(op, 'attributes') and op.OPERATION_NAME == "gpu.func": + if effective_waves_per_eu is not None: + _wpe = int(effective_waves_per_eu) + if _wpe >= 1: + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + ir.IntegerType.get_signless(32), _wpe) + + launcher.launch( + grid=(gx, gy, 1), + block=(block_threads, 1, 1), + stream=stream, + ) + + return launch_gemm_a16w16 + +def gemm_a16w16( + x: torch.Tensor, + w: torch.Tensor, + bias: Optional[torch.Tensor] = None, + dtype: torch.dtype = torch.float16, + y: Optional[torch.Tensor] = None, + activation: Optional[str] = None, + tile_m: int = 128, + tile_n: int = 128, + tile_k: int = 32, + m_warp: int = 2, + n_warp: int = 4, + num_buffers: int = 2, + waves_per_eu: int = None, + l2_prefetch_distance: int = 2, +): + """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" + assert x.dtype in (torch.float16, torch.bfloat16), f"x must be fp16/bf16, got {x.dtype}" + assert w.dtype in (torch.float16, torch.bfloat16), f"w must be fp16/bf16, got {w.dtype}" + assert x.shape[1] == w.shape[1], "Incompatible K dimensions" + + M, K = x.shape + N = w.shape[0] + + # ── Layout ── + if x.stride(1) == 1: + physical_mk = True + elif x.stride(0) == 1: + physical_mk = False + + if w.stride(1) == 1: + physical_kn = False + elif w.stride(0) == 1: + physical_kn = True + + # ── K-pad ── + K_padded = ((K + tile_k - 1) // tile_k) * tile_k + if K_padded != K: + pad_size = K_padded - K + if physical_mk: + x = torch.nn.functional.pad(x, (0, pad_size)) + else: + x = torch.nn.functional.pad(x.T, (0, 0, 0, pad_size)).T + + if physical_kn: + if w.stride(1) == 1: + w = torch.nn.functional.pad(w, (0, 0, 0, pad_size)) + else: + w = torch.nn.functional.pad(w.T, (0, 0, 0, pad_size)).T + else: + w = torch.nn.functional.pad(w, (0, pad_size)) + K = K_padded + + # ── N-padding ── + N_stride = ((N + tile_n - 1) // tile_n) * tile_n + + # ── Output allocation ── + if y is not None: + y_buf = y if N_stride == N else torch.empty((M, N_stride), device=x.device, dtype=dtype) + else: + y_buf = torch.empty((M, N_stride), device=x.device, dtype=dtype) if N_stride != N \ + else torch.empty((M, N), device=x.device, dtype=dtype) + + if bias is None: + bias = torch.empty(0, device=x.device, dtype=dtype) + + in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" + if dtype == torch.float16: + out_dtype_str = "f16" + elif dtype == torch.bfloat16: + out_dtype_str = "bf16" + else: + out_dtype_str = "f32" + + launch_fn = compile_gemm_a16w16( + M=M if not physical_mk else 0, + N=N, + K=K, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + m_warp=m_warp, + n_warp=n_warp, + in_dtype=in_dtype_str, + out_dtype=out_dtype_str, + num_buffers=num_buffers, + waves_per_eu=waves_per_eu, + l2_prefetch_distance=l2_prefetch_distance, + activation=activation, + add_bias=(bias.numel() > 0), + physical_mk=physical_mk, + physical_kn=physical_kn, + ) + + launch_fn(y_buf, x, w, bias, M, N_stride, stream=torch.cuda.current_stream()) + + if N_stride != N: + result = y_buf[:, :N] + if y is not None: + y.copy_(result) + return y + return result + return y_buf + + +__all__ = ["compile_gemm_a16w16", "gemm_a16w16"] diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py index ebb76a2b23f..216d8be539f 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py +++ b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py @@ -6,9 +6,26 @@ import pytest from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 from aiter.ops.triton.gemm.basic.gemm_a16w16_atomic import gemm_a16w16_atomic +import importlib.util from op_tests.triton_tests.utils.types import str_to_torch_dtype +def get_gpu_arch(): + """Get the GPU architecture name (e.g., 'gfx1250', 'gfx942').""" + if not torch.cuda.is_available(): + return None + props = torch.cuda.get_device_properties(0) + return getattr(props, "gcnArchName", None) + + +def is_flydsl_supported(): + """Check if flydsl kernels are supported on the current GPU.""" + if importlib.util.find_spec("flydsl") is None: + return False + arch = get_gpu_arch() + return arch is not None and arch.startswith("gfx1250") + + def generate_gemm_a16w16_inputs(M, N, K, dtype, layout="TN", output=True, bias=False): if isinstance(dtype, str): dtype = str_to_torch_dtype[dtype] @@ -38,22 +55,48 @@ def generate_gemm_a16w16_inputs(M, N, K, dtype, layout="TN", output=True, bias=F return x, weight, bias_tensor, out_dtype, y +def _flydsl_gemm_a16w16(x, w, bias=None, dtype=torch.bfloat16, y=None, activation=None): + """Lazy import and call the flydsl gemm_a16w16 kernel.""" + from aiter.ops.flydsl.kernels.gemm_a16w16_gfx1250 import ( + gemm_a16w16 as flydsl_gemm, + ) + return flydsl_gemm(x, w, bias=bias, dtype=dtype, y=y, activation=activation) + + def get_x_vals(): - x_vals = [(1, 1, 1)] # minimal case - x_vals += [(3, 5, 2)] # irregular shape - x_vals += [(1024 * v, 1024 * v, 1024 * v) for v in (1, 2, 4, 5, 8)] - x_vals += [(2**i, 256, 7168) for i in range(5, 9)] # DSR1 router GEMM - # GPT-OSS-120B attention projections - x_vals += [(2**i, 5120, 2880) for i in range(5, 9)] # GPTOSS QKV input projection - x_vals += [(2**i, 2880, 4096) for i in range(5, 9)] # output projection - x_vals += [(2**i, 128, 2880) for i in range(5, 9)] # Router GEMM - x_vals += [(v, 106496, 16384) for v in (256, 4096)] # LL3 405B FC1 + x_vals = [ + (1, 1, 1), + (1, 16, 16), + (16, 1, 16), + (16, 16, 1), + # Irregular shapes (masking & OOB) + (3, 5, 7), + (17, 33, 65), + (63, 127, 255), + (65, 129, 257), + # + (64, 64, 64), + (128, 128, 128), + # Multiple blocks + (128, 256, 512), + (256, 512, 256), + # Asymmetric shapes + (32, 256, 128), + (256, 32, 128), + (128, 128, 1024), + (1024, 128, 128), + (1536, 512, 768), + ] return x_vals # Test plain BF16 GEMMs - the most common types. @pytest.mark.parametrize("M, N, K", get_x_vals()) -def test_gemm_a16_w16(M: int, N: int, K: int): +@pytest.mark.parametrize("backend", ["triton", "flydsl"]) +def test_gemm_a16_w16(M: int, N: int, K: int, backend): + if backend == "flydsl" and not is_flydsl_supported(): + pytest.skip("FlyDSL not supported (requires gfx1250 and flydsl package)") + x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( M, N, @@ -64,31 +107,41 @@ def test_gemm_a16_w16(M: int, N: int, K: int): torch_out = F.linear(x, w, bias=None) - triton_out = gemm_a16w16( - x, - w, - ) + if backend == "flydsl": + kernel_out = _flydsl_gemm_a16w16(x, w, dtype=torch.bfloat16) + else: + kernel_out = gemm_a16w16(x, w) - torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) # Smaller set for testing activations, setting the output tensor and dtype def get_fewer_x_vals(): - x_vals = [(16, 1024, 1024)] - x_vals += [(128, 8192, 512)] - x_vals += [(256, 512, 8192)] - x_vals += [(1024 * v, 1024 * v, 1024 * v) for v in (1, 5, 8)] + x_vals = [ + (64, 64, 64), + (128, 256, 512), + (256, 512, 256), + (128, 128, 1024), + (1024, 128, 128), + (1536, 512, 768), + ] return x_vals # A smaller set of shapes that tests fused activations, different dtypes # and output tensor arg. We don't want the larger set above to test # all these combinations. -@pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu"]) +@pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu", "silu_exp2"]) @pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16_activation(M: int, N: int, K: int, dtype, output, activation): +@pytest.mark.parametrize("backend", ["triton", "flydsl"]) +def test_gemm_a16_w16_activation( + M: int, N: int, K: int, dtype, output, activation, backend +): + if backend == "flydsl" and not is_flydsl_supported(): + pytest.skip("FlyDSL not supported (requires gfx1250 and flydsl package)") + x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( M, N, @@ -102,24 +155,27 @@ def test_gemm_a16_w16_activation(M: int, N: int, K: int, dtype, output, activati torch_out = F.gelu(torch_out) elif activation == "gelu_tanh": torch_out = F.gelu(torch_out, approximate="tanh") - elif activation == "silu": + elif activation in ("silu", "silu_exp2"): torch_out = F.silu(torch_out) - triton_out = gemm_a16w16( - x, - w, - None, - out_dtype, - y, - activation=activation, - ) + if backend == "flydsl": + kernel_out = _flydsl_gemm_a16w16( + x, w, dtype=out_dtype if not isinstance(out_dtype, tuple) else dtype, + y=y, activation=activation, + ) + else: + kernel_out = gemm_a16w16(x, w, None, out_dtype, y, activation=activation) - torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) @pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("layout", ["TT", "NN", "NT"]) -def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout): +@pytest.mark.parametrize("layout", ["TN", "TT", "NN", "NT"]) +@pytest.mark.parametrize("backend", ["triton", "flydsl"]) +def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout, backend): + if backend == "flydsl" and not is_flydsl_supported(): + pytest.skip("FlyDSL not supported (requires gfx1250 and flydsl package)") + torch.cuda.empty_cache() # Helps avoid hangs in large tests x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( @@ -128,9 +184,12 @@ def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout): torch_out = F.linear(x, w, bias=None) - triton_out = gemm_a16w16(x, w, None, out_dtype, y) + if backend == "flydsl": + kernel_out = _flydsl_gemm_a16w16(x, w, dtype=torch.bfloat16) + else: + kernel_out = gemm_a16w16(x, w, None, out_dtype, y) - torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-1) @pytest.mark.parametrize("M, N, K", get_x_vals()) From aa50b0e966a6eae0d2f56423deb1059cbabd3cc2 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Mon, 20 Apr 2026 19:00:53 +0000 Subject: [PATCH 02/34] Code styling and unit test fixes --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 471 ++++++++++++------ op_tests/flydsl_tests/test_gemm_a16w16.py | 161 ++++++ .../gemm/basic/test_gemm_a16w16.py | 69 +-- 3 files changed, 493 insertions(+), 208 deletions(-) create mode 100644 op_tests/flydsl_tests/test_gemm_a16w16.py diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index 62132375d6b..50c79cb68e3 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -s + import torch import flydsl.compiler as flyc import flydsl.expr as fx @@ -14,16 +14,17 @@ from flydsl.expr import idx2crd from typing import Optional -# WMMA 16×16×32 +# WMMA 16×16×32 WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 WAVE_SIZE = 32 -# LDS padding per row (8 elements -> 16 bytes) +# LDS padding per row (8 elements -> 16 bytes) LDS_PAD_A = 8 LDS_PAD_B = 8 _STAGE_NAMES = tuple(f"stage{i}" for i in range(16)) + def _apply_activation_scalar(val, activation: str): """Apply activation to a single f32 scalar.""" from flydsl._mlir.dialects import math as _math @@ -41,6 +42,7 @@ def _apply_activation_scalar(val, activation: str): elif activation == "gelu": import math + inv_sqrt2 = arith.constant(1.0 / math.sqrt(2.0), type=T.f32) scaled = val * inv_sqrt2 erf_val = _math.erf(scaled) @@ -50,6 +52,7 @@ def _apply_activation_scalar(val, activation: str): elif activation == "gelu_tanh": import math + sqrt_2_over_pi = arith.constant(math.sqrt(2.0 / math.pi), type=T.f32) coeff = arith.constant(0.044715, type=T.f32) one = arith.constant(1.0, type=T.f32) @@ -65,6 +68,7 @@ def _apply_activation_scalar(val, activation: str): else: return val + def compile_gemm_a16w16( *, M: int = 0, @@ -82,8 +86,8 @@ def compile_gemm_a16w16( l2_prefetch_distance: int = 2, activation: Optional[str] = None, add_bias: bool = False, - physical_mk: bool = True, # True=M-major (row-major X), False=K-major (col-major X) - physical_kn: bool = False, # False=N-major (row-major W), True=K-major (col-major/transposed W) + physical_mk: bool = True, # True=M-major (row-major X), False=K-major (col-major X) + physical_kn: bool = False, # False=N-major (row-major W), True=K-major (col-major/transposed W) ): """Compile and return a launch function for the A16W16 GEMM kernel. Returns a callable: launch_fn(y, x, w, bias, M, N, stream=stream) @@ -102,7 +106,9 @@ def compile_gemm_a16w16( if out_dtype is None: out_dtype = "f16" if is_f16 else "bf16" if out_dtype not in ("f32", "f16", "bf16"): - raise ValueError(f"out_dtype must be 'f32', 'f16', or 'bf16', got {out_dtype!r}") + raise ValueError( + f"out_dtype must be 'f32', 'f16', or 'bf16', got {out_dtype!r}" + ) elem_bytes = 2 elem_bytes_d = 2 if out_dtype in ("f16", "bf16") else 4 @@ -126,20 +132,24 @@ def compile_gemm_a16w16( if N == 0: raise ValueError( "N must be specified (> 0) at compile time when physical_kn=True, " - "because it is used as the TDM stride for the (K, N) weight layout") + "because it is used as the TDM stride for the (K, N) weight layout" + ) if (tile_n & (tile_n - 1)) != 0: raise ValueError( f"tile_n must be a power of 2 when physical_kn=True " - f"(TDM pad_interval requirement), got {tile_n}") + f"(TDM pad_interval requirement), got {tile_n}" + ) if not physical_mk: if M == 0: raise ValueError( "M must be specified (> 0) at compile time when physical_mk=False, " - "because it is used as the TDM stride for the (K, M) activation layout") + "because it is used as the TDM stride for the (K, M) activation layout" + ) if (tile_m & (tile_m - 1)) != 0: raise ValueError( f"tile_m must be a power of 2 when physical_mk=False " - f"(TDM pad_interval requirement), got {tile_m}") + f"(TDM pad_interval requirement), got {tile_m}" + ) # ── Warp tile dimensions ── warp_tile_m = tile_m // m_warp @@ -154,10 +164,11 @@ def compile_gemm_a16w16( if num_k_tiles < num_buffers - 1: raise ValueError( f"{num_buffers}-stage buffering requires num_k_tiles >= {num_buffers - 1}, " - f"got {num_k_tiles} (K={K}, tile_k={tile_k})") + f"got {num_k_tiles} (K={K}, tile_k={tile_k})" + ) # ── Architecture check ── - gpu_arch = str(get_hip_arch(timeout_s=300)) + gpu_arch = str(get_hip_arch()) assert gpu_arch.startswith("gfx1250"), f"Expected gfx1250, got {gpu_arch}" wmma_op = rocdl.wmma_f32_16x16x32_f16 if is_f16 else rocdl.wmma_f32_16x16x32_bf16 @@ -212,8 +223,12 @@ def _elem_type(): @flyc.kernel def kernel_gemm_a16w16( - arg_y: fx.Tensor, arg_x: fx.Tensor, arg_w: fx.Tensor, - arg_bias: fx.Tensor, i32_m: fx.Int32, i32_n: fx.Int32, + arg_y: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_bias: fx.Tensor, + i32_m: fx.Int32, + i32_n: fx.Int32, ): rocdl.disable_xdl_arb_stall() @@ -226,12 +241,15 @@ def kernel_gemm_a16w16( # ── Thread-to-warp decomposition ── layout_thr = fx.make_layout( - (m_warp, n_warp, 2, 16), - (n_warp * WAVE_SIZE, WAVE_SIZE, 16, 1)) + (m_warp, n_warp, 2, 16), (n_warp * WAVE_SIZE, WAVE_SIZE, 16, 1) + ) thr_coord = idx2crd(tx, layout_thr) wave_m_idx, wave_n_idx, lane_kgrp, lane16 = ( - fx.get(thr_coord, 0), fx.get(thr_coord, 1), - fx.get(thr_coord, 2), fx.get(thr_coord, 3)) + fx.get(thr_coord, 0), + fx.get(thr_coord, 1), + fx.get(thr_coord, 2), + fx.get(thr_coord, 3), + ) warp_m_base = wave_m_idx * arith.index(warp_tile_m) warp_n_base = wave_n_idx * arith.index(warp_tile_n) @@ -240,7 +258,7 @@ def kernel_gemm_a16w16( m_idx = arith.index_cast(T.index, i32_m.ir_value()) n_stride = arith.index_cast(T.index, i32_n.ir_value()) - # Buffer resource descriptors + # Buffer resource descriptors y_nrec = m_idx * n_stride * arith.index(elem_bytes_d) y_rsrc = buffer_ops.create_buffer_resource(arg_y, num_records_bytes=y_nrec) if add_bias: @@ -250,38 +268,58 @@ def make_a_desc(k_base, lds_a_mem_ref): """TDM descriptor for A tile. Swaps dims when physical_mk=False (K-major).""" if physical_mk: return tdm_ops.make_tensor_descriptor_2d( - global_ptr=arg_x, lds_memref=lds_a_mem_ref, + global_ptr=arg_x, + lds_memref=lds_a_mem_ref, global_offset=(blk_m, k_base), - tensor_shape=(tile_m, tile_k), strides=(K, 1), - tile_shape=(tile_m, tile_k), elem_bytes=elem_bytes, - pad_interval=tile_k, pad_amount=LDS_PAD_A, - num_warps=num_warps) + tensor_shape=(tile_m, tile_k), + strides=(K, 1), + tile_shape=(tile_m, tile_k), + elem_bytes=elem_bytes, + pad_interval=tile_k, + pad_amount=LDS_PAD_A, + num_warps=num_warps, + ) else: return tdm_ops.make_tensor_descriptor_2d( - global_ptr=arg_x, lds_memref=lds_a_mem_ref, + global_ptr=arg_x, + lds_memref=lds_a_mem_ref, global_offset=(k_base, blk_m), - tensor_shape=(tile_k, tile_m), strides=(M, 1), - tile_shape=(tile_k, tile_m), elem_bytes=elem_bytes, - pad_interval=tile_m, pad_amount=LDS_PAD_A, - num_warps=num_warps) + tensor_shape=(tile_k, tile_m), + strides=(M, 1), + tile_shape=(tile_k, tile_m), + elem_bytes=elem_bytes, + pad_interval=tile_m, + pad_amount=LDS_PAD_A, + num_warps=num_warps, + ) def make_b_desc(k_base, lds_b_mem_ref): """TDM descriptor for B tile. Swaps dims when physical_kn=True (K-major).""" if physical_kn: return tdm_ops.make_tensor_descriptor_2d( - global_ptr=arg_w, lds_memref=lds_b_mem_ref, + global_ptr=arg_w, + lds_memref=lds_b_mem_ref, global_offset=(k_base, blk_n), - tensor_shape=(tile_k, tile_n), strides=(N, 1), - tile_shape=(tile_k, tile_n), elem_bytes=elem_bytes, - pad_interval=tile_n, pad_amount=LDS_PAD_B, - num_warps=num_warps) + tensor_shape=(tile_k, tile_n), + strides=(N, 1), + tile_shape=(tile_k, tile_n), + elem_bytes=elem_bytes, + pad_interval=tile_n, + pad_amount=LDS_PAD_B, + num_warps=num_warps, + ) return tdm_ops.make_tensor_descriptor_2d( - global_ptr=arg_w, lds_memref=lds_b_mem_ref, + global_ptr=arg_w, + lds_memref=lds_b_mem_ref, global_offset=(blk_n, k_base), - tensor_shape=(tile_n, tile_k), strides=(K, 1), - tile_shape=(tile_n, tile_k), elem_bytes=elem_bytes, - pad_interval=tile_k, pad_amount=LDS_PAD_B, - num_warps=num_warps) + tensor_shape=(tile_n, tile_k), + strides=(K, 1), + tile_shape=(tile_n, tile_k), + elem_bytes=elem_bytes, + pad_interval=tile_k, + pad_amount=LDS_PAD_B, + num_warps=num_warps, + ) def issue_tdm_load(desc): tdm_ops.tensor_load_2d(desc) @@ -304,7 +342,9 @@ def _precompute_lane_bases(lds_ptr, warp_base, reps, lds_stride): k_lane_off = lane_kgrp * arith.index(8) bases = [] for rep in range_constexpr(reps): - base = row_stride_off + arith.index(rep * WMMA_M * lds_stride) + k_lane_off + base = ( + row_stride_off + arith.index(rep * WMMA_M * lds_stride) + k_lane_off + ) bases.append(base) return lds_buffer, bases @@ -312,13 +352,15 @@ def _precompute_a_lane_bases(lds_ptr): """A fragment lane bases. Uses transpose-load addressing when K-major.""" if physical_mk: return _precompute_lane_bases( - lds_ptr, warp_m_base, wmma_m_rep, lds_a_stride) + lds_ptr, warp_m_base, wmma_m_rep, lds_a_stride + ) lds_buffer = _get_lds_memref(lds_ptr) lane8 = lane16 % arith.index(8) lane_mgrp = lane16 / arith.index(8) - k_lane_off = (lane_kgrp * arith.index(8) + lane8) \ - * arith.index(lds_a_stride) + k_lane_off = (lane_kgrp * arith.index(8) + lane8) * arith.index( + lds_a_stride + ) m_lane_off = lane_mgrp * arith.index(8) bases = [] for wm in range_constexpr(wmma_m_rep): @@ -330,13 +372,15 @@ def _precompute_b_lane_bases(lds_ptr): """B fragment lane bases. Uses transpose-load addressing when K-major.""" if not physical_kn: return _precompute_lane_bases( - lds_ptr, warp_n_base, wmma_n_rep, lds_b_stride) + lds_ptr, warp_n_base, wmma_n_rep, lds_b_stride + ) lds_buffer = _get_lds_memref(lds_ptr) lane8 = lane16 % arith.index(8) lane_ngrp = lane16 / arith.index(8) - k_lane_off = (lane_kgrp * arith.index(8) + lane8) \ - * arith.index(lds_b_stride) + k_lane_off = (lane_kgrp * arith.index(8) + lane8) * arith.index( + lds_b_stride + ) n_lane_off = lane_ngrp * arith.index(8) bases = [] for wn in range_constexpr(wmma_n_rep): @@ -351,8 +395,7 @@ def load_wmma_frag_tr(lds_buffer, b_lane_base, ks): for k_half in range_constexpr(2): k_row_off = (ks * WMMA_K + k_half * 16) * lds_b_stride elem_off = b_lane_base + arith.index(k_row_off) - v = rocdl.lds_transpose_load( - vec8_ty, lds_buffer, elem_off, elem_bytes) + v = rocdl.lds_transpose_load(vec8_ty, lds_buffer, elem_off, elem_bytes) results.append(v) return vector.shuffle(results[0], results[1], list(range(16))) @@ -374,8 +417,7 @@ def load_wmma_frag_tr_a(lds_buffer, a_lane_base, ks): for k_half in range_constexpr(2): k_row_off = (ks * WMMA_K + k_half * 16) * lds_a_stride elem_off = a_lane_base + arith.index(k_row_off) - v = rocdl.lds_transpose_load( - vec8_ty, lds_buffer, elem_off, elem_bytes) + v = rocdl.lds_transpose_load(vec8_ty, lds_buffer, elem_off, elem_bytes) results.append(v) return vector.shuffle(results[0], results[1], list(range(16))) @@ -398,19 +440,25 @@ def wmma_tile(accs_in, tile_frags): current_accs = list(accs_in) for ks in range_constexpr(k_wmma_steps): base = ks * (wmma_n_rep + wmma_m_rep) - b_frags = [tile_frags[base + wn] - for wn in range_constexpr(wmma_n_rep)] - a_frags = [tile_frags[base + wmma_n_rep + wm] - for wm in range_constexpr(wmma_m_rep)] + b_frags = [tile_frags[base + wn] for wn in range_constexpr(wmma_n_rep)] + a_frags = [ + tile_frags[base + wmma_n_rep + wm] + for wm in range_constexpr(wmma_m_rep) + ] for wm in range_constexpr(wmma_m_rep): for wn in range_constexpr(wmma_n_rep): idx = wm * wmma_n_rep + wn # ISA operand order: (B, A, C) — reversed from math current_accs[idx] = wmma_op( T.vec(8, T.f32), - b_frags[wn], a_frags[wm], current_accs[idx], - signA=False, signB=False, modC=0, - reuseA=False, reuseB=(wn > 0), + b_frags[wn], + a_frags[wm], + current_accs[idx], + signA=False, + signB=False, + modC=0, + reuseA=False, + reuseB=(wn > 0), ).result return current_accs @@ -422,26 +470,32 @@ def compute_and_prefetch(accs_in, cur_frags, next_lds_a_ptr, next_lds_b_ptr): next_frags = [] for ks in range_constexpr(k_wmma_steps): - # Phase 1: ds_read for next tile + # Phase 1: ds_read for next tile for wn in range_constexpr(wmma_n_rep): next_frags.append(_load_b_frag(next_b_buf, next_b_bases[wn], ks)) for wm in range_constexpr(wmma_m_rep): next_frags.append(_load_a_frag(next_a_buf, next_a_bases[wm], ks)) - # Phase 2: WMMA for curr tile + # Phase 2: WMMA for curr tile base = ks * (wmma_n_rep + wmma_m_rep) - b_frags = [cur_frags[base + wn] - for wn in range_constexpr(wmma_n_rep)] - a_frags = [cur_frags[base + wmma_n_rep + wm] - for wm in range_constexpr(wmma_m_rep)] + b_frags = [cur_frags[base + wn] for wn in range_constexpr(wmma_n_rep)] + a_frags = [ + cur_frags[base + wmma_n_rep + wm] + for wm in range_constexpr(wmma_m_rep) + ] for wm in range_constexpr(wmma_m_rep): for wn in range_constexpr(wmma_n_rep): idx = wm * wmma_n_rep + wn current_accs[idx] = wmma_op( T.vec(8, T.f32), - b_frags[wn], a_frags[wm], current_accs[idx], - signA=False, signB=False, modC=0, - reuseA=False, reuseB=(wn > 0), + b_frags[wn], + a_frags[wm], + current_accs[idx], + signA=False, + signB=False, + modC=0, + reuseA=False, + reuseB=(wn > 0), ).result return current_accs, next_frags @@ -454,29 +508,54 @@ def _l2_prefetch(k_base): if physical_mk: tdm_ops.l2_prefetch_tile( - arg_x, (blk_m, pf_k), (tile_m, tile_k), (K, 1), - elem_bytes=elem_bytes, thread_id=tx, block_threads=block_threads) + arg_x, + (blk_m, pf_k), + (tile_m, tile_k), + (K, 1), + elem_bytes=elem_bytes, + thread_id=tx, + block_threads=block_threads, + ) else: tdm_ops.l2_prefetch_tile( - arg_x, (pf_k, blk_m), (tile_k, tile_m), (M, 1), - elem_bytes=elem_bytes, thread_id=tx, block_threads=block_threads) + arg_x, + (pf_k, blk_m), + (tile_k, tile_m), + (M, 1), + elem_bytes=elem_bytes, + thread_id=tx, + block_threads=block_threads, + ) if physical_kn: tdm_ops.l2_prefetch_tile( - arg_w, (pf_k, blk_n), (tile_k, tile_n), (N, 1), - elem_bytes=elem_bytes, thread_id=tx, - block_threads=block_threads) + arg_w, + (pf_k, blk_n), + (tile_k, tile_n), + (N, 1), + elem_bytes=elem_bytes, + thread_id=tx, + block_threads=block_threads, + ) else: tdm_ops.l2_prefetch_tile( - arg_w, (blk_n, pf_k), (tile_n, tile_k), (K, 1), - elem_bytes=elem_bytes, thread_id=tx, - block_threads=block_threads) + arg_w, + (blk_n, pf_k), + (tile_n, tile_k), + (K, 1), + elem_bytes=elem_bytes, + thread_id=tx, + block_threads=block_threads, + ) _half_out = out_dtype in ("f16", "bf16") - _out_elem = T.f16 if out_dtype == "f16" else (T.bf16 if out_dtype == "bf16" else None) + _out_elem = ( + T.f16 if out_dtype == "f16" else (T.bf16 if out_dtype == "bf16" else None) + ) _bias_elem = T.f16 if is_f16 else T.bf16 def _widen_to_f32(val): from flydsl._mlir.dialects import arith as _std_arith + return _std_arith.ExtFOp(T.f32, _raw(val)).result def _apply_bias_and_activation(accs): @@ -490,27 +569,32 @@ def _apply_bias_and_activation(accs): acc = accs[idx] if add_bias: - col_base = (blk_n + warp_n_base - + arith.index(wn * WMMA_N) - + lane_kgrp * arith.index(8)) + col_base = ( + blk_n + + warp_n_base + + arith.index(wn * WMMA_N) + + lane_kgrp * arith.index(8) + ) col_base_i32 = arith.index_cast(T.i32, col_base) col_base_hi = col_base_i32 + arith.constant(4, type=T.i32) bv_lo = buffer_ops.buffer_load( - bias_rsrc, col_base_i32, - vec_width=4, dtype=_bias_elem) + bias_rsrc, col_base_i32, vec_width=4, dtype=_bias_elem + ) bv_hi = buffer_ops.buffer_load( - bias_rsrc, col_base_hi, - vec_width=4, dtype=_bias_elem) + bias_rsrc, col_base_hi, vec_width=4, dtype=_bias_elem + ) bias_elems = [] for i in range_constexpr(4): b_val = vector.extract( - bv_lo, static_position=[i], dynamic_position=[]) + bv_lo, static_position=[i], dynamic_position=[] + ) bias_elems.append(_widen_to_f32(b_val)) for i in range_constexpr(4): b_val = vector.extract( - bv_hi, static_position=[i], dynamic_position=[]) + bv_hi, static_position=[i], dynamic_position=[] + ) bias_elems.append(_widen_to_f32(b_val)) bias_vec = vector.from_elements(T.vec(8, T.f32), bias_elems) @@ -520,7 +604,8 @@ def _apply_bias_and_activation(accs): new_elems = [] for i in range_constexpr(8): val = vector.extract( - acc, static_position=[i], dynamic_position=[]) + acc, static_position=[i], dynamic_position=[] + ) val = _apply_activation_scalar(val, activation) new_elems.append(val) acc = vector.from_elements(T.vec(8, T.f32), new_elems) @@ -534,22 +619,32 @@ def epilogue_stores(final_accs): for wn in range_constexpr(wmma_n_rep): idx = wm * wmma_n_rep + wn row = blk_m + warp_m_base + arith.index(wm * WMMA_M) + lane16 - col_base = (blk_n + warp_n_base + arith.index(wn * WMMA_N) - + lane_kgrp * arith.index(8)) + col_base = ( + blk_n + + warp_n_base + + arith.index(wn * WMMA_N) + + lane_kgrp * arith.index(8) + ) if _half_out: h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) i32_vec = vector.bitcast(T.vec(4, T.i32), h_vec) - c_off_bytes = (row * n_stride + col_base) * arith.index(elem_bytes_d) + c_off_bytes = (row * n_stride + col_base) * arith.index( + elem_bytes_d + ) buffer_ops.buffer_store( - i32_vec, y_rsrc, c_off_bytes, offset_is_bytes=True) + i32_vec, y_rsrc, c_off_bytes, offset_is_bytes=True + ) else: for half in range_constexpr(2): - vals = [vector.extract( - final_accs[idx], - static_position=[half * 4 + vi], - dynamic_position=[]) - for vi in range_constexpr(4)] + vals = [ + vector.extract( + final_accs[idx], + static_position=[half * 4 + vi], + dynamic_position=[], + ) + for vi in range_constexpr(4) + ] vec4 = vector.from_elements(T.vec(4, T.f32), vals) col = col_base + arith.index(half * 4) c_off = row * n_stride + col @@ -583,33 +678,55 @@ def epilogue_stores(final_accs): pf_k = arith.index(pf_i * tile_k) if physical_mk: tdm_ops.l2_prefetch_tile( - arg_x, (blk_m, pf_k), (tile_m, tile_k), (K, 1), - elem_bytes=elem_bytes, thread_id=tx, - block_threads=block_threads) + arg_x, + (blk_m, pf_k), + (tile_m, tile_k), + (K, 1), + elem_bytes=elem_bytes, + thread_id=tx, + block_threads=block_threads, + ) else: tdm_ops.l2_prefetch_tile( - arg_x, (pf_k, blk_m), (tile_k, tile_m), (M, 1), - elem_bytes=elem_bytes, thread_id=tx, - block_threads=block_threads) + arg_x, + (pf_k, blk_m), + (tile_k, tile_m), + (M, 1), + elem_bytes=elem_bytes, + thread_id=tx, + block_threads=block_threads, + ) if physical_kn: tdm_ops.l2_prefetch_tile( - arg_w, (pf_k, blk_n), (tile_k, tile_n), (N, 1), - elem_bytes=elem_bytes, thread_id=tx, - block_threads=block_threads) + arg_w, + (pf_k, blk_n), + (tile_k, tile_n), + (N, 1), + elem_bytes=elem_bytes, + thread_id=tx, + block_threads=block_threads, + ) else: tdm_ops.l2_prefetch_tile( - arg_w, (blk_n, pf_k), (tile_n, tile_k), (K, 1), - elem_bytes=elem_bytes, thread_id=tx, - block_threads=block_threads) + arg_w, + (blk_n, pf_k), + (tile_n, tile_k), + (K, 1), + elem_bytes=elem_bytes, + thread_id=tx, + block_threads=block_threads, + ) # Prepare first TDM descriptors for main loop if main_loop_iters > 0: _first_load_stage_0 = pre_loaded % num_buffers _first_load_k_0 = pre_loaded * tile_k - pending_desc_a = make_a_desc(arith.index(_first_load_k_0), - stages_a_mem[_first_load_stage_0]) - pending_desc_b = make_b_desc(arith.index(_first_load_k_0), - stages_b_mem[_first_load_stage_0]) + pending_desc_a = make_a_desc( + arith.index(_first_load_k_0), stages_a_mem[_first_load_stage_0] + ) + pending_desc_b = make_b_desc( + arith.index(_first_load_k_0), stages_b_mem[_first_load_stage_0] + ) main_end = main_loop_iters * num_buffers * tile_k @@ -620,17 +737,27 @@ def epilogue_stores(final_accs): # Pre-load fragments for first compute tile cur_frags = load_tile_frags(stages_a[0], stages_b[0]) - # Step 4: Main loop + # Step 4: Main loop if main_loop_iters > 0: - init_descs = [pending_desc_a.dgroup0, pending_desc_a.dgroup1, - pending_desc_b.dgroup0, pending_desc_b.dgroup1] - - for iv, state in range(0, main_end, num_buffers * tile_k, - init=list(accs) + init_descs + cur_frags): + init_descs = [ + pending_desc_a.dgroup0, + pending_desc_a.dgroup1, + pending_desc_b.dgroup0, + pending_desc_b.dgroup1, + ] + + for iv, state in range( + 0, + main_end, + num_buffers * tile_k, + init=list(accs) + init_descs + cur_frags, + ): accs_in = list(state[:n_accs]) cur_desc_a = tdm_ops.TDMDescriptor2D(state[n_accs], state[n_accs + 1]) - cur_desc_b = tdm_ops.TDMDescriptor2D(state[n_accs + 2], state[n_accs + 3]) - cur_frags = list(state[n_accs + 4:n_accs + 4 + n_frags_per_tile]) + cur_desc_b = tdm_ops.TDMDescriptor2D( + state[n_accs + 2], state[n_accs + 3] + ) + cur_frags = list(state[n_accs + 4 : n_accs + 4 + n_frags_per_tile]) for s in range_constexpr(num_buffers): issue_tdm_load(cur_desc_a) @@ -644,35 +771,46 @@ def epilogue_stores(final_accs): _next_load_k_off = (s + 1 + pre_loaded) * tile_k cur_desc_a = make_a_desc( iv + arith.index(_next_load_k_off), - stages_a_mem[_next_load_stage]) + stages_a_mem[_next_load_stage], + ) cur_desc_b = make_b_desc( iv + arith.index(_next_load_k_off), - stages_b_mem[_next_load_stage]) + stages_b_mem[_next_load_stage], + ) else: _next_load_stage = pre_loaded % num_buffers _next_load_k_off = pre_loaded * tile_k _next_step = num_buffers * tile_k cur_desc_a = make_a_desc( iv + arith.index(_next_step + _next_load_k_off), - stages_a_mem[_next_load_stage]) + stages_a_mem[_next_load_stage], + ) cur_desc_b = make_b_desc( iv + arith.index(_next_step + _next_load_k_off), - stages_b_mem[_next_load_stage]) + stages_b_mem[_next_load_stage], + ) tdm_ops.tensor_wait(2 * (num_buffers - 2)) gpu.barrier() _next_compute = (s + 1) % num_buffers accs_in, cur_frags = compute_and_prefetch( - accs_in, cur_frags, - stages_a[_next_compute], stages_b[_next_compute]) - - out_descs = [cur_desc_a.dgroup0, cur_desc_a.dgroup1, - cur_desc_b.dgroup0, cur_desc_b.dgroup1] + accs_in, + cur_frags, + stages_a[_next_compute], + stages_b[_next_compute], + ) + + out_descs = [ + cur_desc_a.dgroup0, + cur_desc_a.dgroup1, + cur_desc_b.dgroup0, + cur_desc_b.dgroup1, + ] results = yield list(accs_in) + out_descs + cur_frags accs = list(results[:n_accs]) - cur_frags = list(results[n_accs + 4:n_accs + 4 + n_frags_per_tile]) + cur_frags = list(results[n_accs + 4 : n_accs + 4 + n_frags_per_tile]) # Step 5: drain remaining tiles _tail_base_k = main_loop_iters * num_buffers @@ -682,10 +820,12 @@ def epilogue_stores(final_accs): if t < _extra_loads: load_tile_idx = _tail_base_k + pre_loaded + t load_stage = (pre_loaded + t) % num_buffers - copy_a_to_lds(arith.index(load_tile_idx * tile_k), - stages_a_mem[load_stage]) - copy_b_to_lds(arith.index(load_tile_idx * tile_k), - stages_b_mem[load_stage]) + copy_a_to_lds( + arith.index(load_tile_idx * tile_k), stages_a_mem[load_stage] + ) + copy_b_to_lds( + arith.index(load_tile_idx * tile_k), stages_b_mem[load_stage] + ) _epi_outstanding = 2 * (pre_loaded + min(t + 1, _extra_loads) - t - 1) tdm_ops.tensor_wait(_epi_outstanding) @@ -693,8 +833,8 @@ def epilogue_stores(final_accs): _next_epi_stage = (_tail_base_k + t + 1) % num_buffers accs, cur_frags = compute_and_prefetch( - accs, cur_frags, - stages_a[_next_epi_stage], stages_b[_next_epi_stage]) + accs, cur_frags, stages_a[_next_epi_stage], stages_b[_next_epi_stage] + ) # Step 6: Final WMMA accs = wmma_tile(accs, cur_frags) @@ -703,14 +843,34 @@ def epilogue_stores(final_accs): accs = _apply_bias_and_activation(accs) epilogue_stores(accs) - cache_tag = (in_dtype, out_dtype, K, tile_m, tile_n, tile_k, m_warp, n_warp, - num_buffers, effective_waves_per_eu, l2_prefetch_distance, - activation, add_bias, physical_mk, physical_kn, M, N) + cache_tag = ( + in_dtype, + out_dtype, + K, + tile_m, + tile_n, + tile_k, + m_warp, + n_warp, + num_buffers, + effective_waves_per_eu, + l2_prefetch_distance, + activation, + add_bias, + physical_mk, + physical_kn, + M, + N, + ) @flyc.jit def launch_gemm_a16w16( - arg_y: fx.Tensor, arg_x: fx.Tensor, arg_w: fx.Tensor, - arg_bias: fx.Tensor, i32_m: fx.Int32, i32_n: fx.Int32, + arg_y: fx.Tensor, + arg_x: fx.Tensor, + arg_w: fx.Tensor, + arg_bias: fx.Tensor, + i32_m: fx.Int32, + i32_n: fx.Int32, stream: fx.Stream, ): _ = cache_tag @@ -734,12 +894,13 @@ def launch_gemm_a16w16( # Set waves_per_eu for op in ctx.gpu_module_body.operations: - if hasattr(op, 'attributes') and op.OPERATION_NAME == "gpu.func": + if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": if effective_waves_per_eu is not None: _wpe = int(effective_waves_per_eu) if _wpe >= 1: op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( - ir.IntegerType.get_signless(32), _wpe) + ir.IntegerType.get_signless(32), _wpe + ) launcher.launch( grid=(gx, gy, 1), @@ -749,6 +910,7 @@ def launch_gemm_a16w16( return launch_gemm_a16w16 + def gemm_a16w16( x: torch.Tensor, w: torch.Tensor, @@ -766,8 +928,14 @@ def gemm_a16w16( l2_prefetch_distance: int = 2, ): """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" - assert x.dtype in (torch.float16, torch.bfloat16), f"x must be fp16/bf16, got {x.dtype}" - assert w.dtype in (torch.float16, torch.bfloat16), f"w must be fp16/bf16, got {w.dtype}" + assert x.dtype in ( + torch.float16, + torch.bfloat16, + ), f"x must be fp16/bf16, got {x.dtype}" + assert w.dtype in ( + torch.float16, + torch.bfloat16, + ), f"w must be fp16/bf16, got {w.dtype}" assert x.shape[1] == w.shape[1], "Incompatible K dimensions" M, K = x.shape @@ -807,10 +975,17 @@ def gemm_a16w16( # ── Output allocation ── if y is not None: - y_buf = y if N_stride == N else torch.empty((M, N_stride), device=x.device, dtype=dtype) + y_buf = ( + y + if N_stride == N + else torch.empty((M, N_stride), device=x.device, dtype=dtype) + ) else: - y_buf = torch.empty((M, N_stride), device=x.device, dtype=dtype) if N_stride != N \ - else torch.empty((M, N), device=x.device, dtype=dtype) + y_buf = ( + torch.empty((M, N_stride), device=x.device, dtype=dtype) + if N_stride != N + else torch.empty((M, N), device=x.device, dtype=dtype) + ) if bias is None: bias = torch.empty(0, device=x.device, dtype=dtype) diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py new file mode 100644 index 00000000000..c8cb5e59499 --- /dev/null +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""FlyDSL-only unit test for gemm_a16w16 on gfx1250. + +Intentionally isolated from the triton test harness: does NOT import +`aiter.ops.flydsl` or `aiter.ops.triton` at top level, because those +trigger a build of `module_aiter_core` (via `aiter.ops.enum`) which is +orthogonal to the flydsl kernel under test. The kernel file is loaded +directly by file path. +""" + +import importlib.util +import os + +import pytest +import torch +import torch.nn.functional as F + +_KERNEL_PATH = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "aiter", + "ops", + "flydsl", + "kernels", + "gemm_a16w16_gfx1250.py", + ) +) + + +def _get_gpu_arch(): + if not torch.cuda.is_available(): + return None + return getattr(torch.cuda.get_device_properties(0), "gcnArchName", None) + + +def _flydsl_available(): + if importlib.util.find_spec("flydsl") is None: + return False + arch = _get_gpu_arch() + return arch is not None and arch.startswith("gfx1250") + + +if not _flydsl_available(): + pytest.skip( + "FlyDSL a16w16 tests require gfx1250 and the flydsl package.", + allow_module_level=True, + ) + + +def _load_kernel(): + """Load the kernel module by file path to skip aiter.ops.flydsl.__init__.""" + spec = importlib.util.spec_from_file_location("_flydsl_a16w16_kernel", _KERNEL_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.gemm_a16w16 + + +gemm_a16w16 = _load_kernel() + + +def _generate_inputs(M, N, K, dtype, layout="TN", output=False, bias=False): + if layout[0] == "T": + x = torch.randn((M, K), dtype=dtype, device="cuda") + else: + x = torch.randn((K, M), dtype=dtype, device="cuda").T + + if layout[1] == "T": + w = torch.randn((K, N), dtype=dtype, device="cuda").T + else: + w = torch.randn((N, K), dtype=dtype, device="cuda") + + bias_tensor = None + if bias: + bias_tensor = torch.randn((N,), dtype=dtype, device="cuda") + + y = torch.empty((M, N), dtype=dtype, device="cuda") if output else None + return x, w, bias_tensor, y + + +def get_x_vals(): + return [ + (1, 1, 1), + (1, 16, 16), + (16, 1, 16), + (16, 16, 1), + # Irregular + (3, 5, 7), + (17, 33, 65), + (63, 127, 255), + (65, 129, 257), + # Aligned + (64, 64, 64), + (128, 128, 128), + # Multi-block + (128, 256, 512), + (256, 512, 256), + # Asymmetric + (32, 256, 128), + (256, 32, 128), + (128, 128, 1024), + (1024, 128, 128), + (1536, 512, 768), + ] + + +def get_fewer_x_vals(): + return [ + (64, 64, 64), + (128, 256, 512), + (256, 512, 256), + (128, 128, 1024), + (1024, 128, 128), + (1536, 512, 768), + ] + + +@pytest.mark.parametrize("M, N, K", get_x_vals()) +def test_gemm_a16_w16(M, N, K): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu", "silu_exp2"]) +@pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("output", [True, False]) +def test_gemm_a16_w16_activation(M, N, K, dtype, output, activation): + torch.cuda.empty_cache() + x, w, _, y = _generate_inputs(M, N, K, dtype, output=output) + + torch_out = F.linear(x, w, bias=None) + if activation == "gelu": + torch_out = F.gelu(torch_out) + elif activation == "gelu_tanh": + torch_out = F.gelu(torch_out, approximate="tanh") + elif activation in ("silu", "silu_exp2"): + torch_out = F.silu(torch_out) + + kernel_out = gemm_a16w16(x, w, dtype=dtype, y=y, activation=activation) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize("M, N, K", get_x_vals()) +@pytest.mark.parametrize("layout", ["TN", "TT", "NN", "NT"]) +def test_gemm_a16_w16_layout(M, N, K, layout): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16, layout=layout) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-1) diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py index 216d8be539f..191d50cecf3 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py +++ b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py @@ -6,26 +6,9 @@ import pytest from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 from aiter.ops.triton.gemm.basic.gemm_a16w16_atomic import gemm_a16w16_atomic -import importlib.util from op_tests.triton_tests.utils.types import str_to_torch_dtype -def get_gpu_arch(): - """Get the GPU architecture name (e.g., 'gfx1250', 'gfx942').""" - if not torch.cuda.is_available(): - return None - props = torch.cuda.get_device_properties(0) - return getattr(props, "gcnArchName", None) - - -def is_flydsl_supported(): - """Check if flydsl kernels are supported on the current GPU.""" - if importlib.util.find_spec("flydsl") is None: - return False - arch = get_gpu_arch() - return arch is not None and arch.startswith("gfx1250") - - def generate_gemm_a16w16_inputs(M, N, K, dtype, layout="TN", output=True, bias=False): if isinstance(dtype, str): dtype = str_to_torch_dtype[dtype] @@ -55,14 +38,6 @@ def generate_gemm_a16w16_inputs(M, N, K, dtype, layout="TN", output=True, bias=F return x, weight, bias_tensor, out_dtype, y -def _flydsl_gemm_a16w16(x, w, bias=None, dtype=torch.bfloat16, y=None, activation=None): - """Lazy import and call the flydsl gemm_a16w16 kernel.""" - from aiter.ops.flydsl.kernels.gemm_a16w16_gfx1250 import ( - gemm_a16w16 as flydsl_gemm, - ) - return flydsl_gemm(x, w, bias=bias, dtype=dtype, y=y, activation=activation) - - def get_x_vals(): x_vals = [ (1, 1, 1), @@ -92,11 +67,7 @@ def get_x_vals(): # Test plain BF16 GEMMs - the most common types. @pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("backend", ["triton", "flydsl"]) -def test_gemm_a16_w16(M: int, N: int, K: int, backend): - if backend == "flydsl" and not is_flydsl_supported(): - pytest.skip("FlyDSL not supported (requires gfx1250 and flydsl package)") - +def test_gemm_a16_w16(M: int, N: int, K: int): x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( M, N, @@ -107,12 +78,9 @@ def test_gemm_a16_w16(M: int, N: int, K: int, backend): torch_out = F.linear(x, w, bias=None) - if backend == "flydsl": - kernel_out = _flydsl_gemm_a16w16(x, w, dtype=torch.bfloat16) - else: - kernel_out = gemm_a16w16(x, w) + triton_out = gemm_a16w16(x, w) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) # Smaller set for testing activations, setting the output tensor and dtype @@ -135,13 +103,7 @@ def get_fewer_x_vals(): @pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("output", [True, False]) -@pytest.mark.parametrize("backend", ["triton", "flydsl"]) -def test_gemm_a16_w16_activation( - M: int, N: int, K: int, dtype, output, activation, backend -): - if backend == "flydsl" and not is_flydsl_supported(): - pytest.skip("FlyDSL not supported (requires gfx1250 and flydsl package)") - +def test_gemm_a16_w16_activation(M: int, N: int, K: int, dtype, output, activation): x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( M, N, @@ -158,24 +120,14 @@ def test_gemm_a16_w16_activation( elif activation in ("silu", "silu_exp2"): torch_out = F.silu(torch_out) - if backend == "flydsl": - kernel_out = _flydsl_gemm_a16w16( - x, w, dtype=out_dtype if not isinstance(out_dtype, tuple) else dtype, - y=y, activation=activation, - ) - else: - kernel_out = gemm_a16w16(x, w, None, out_dtype, y, activation=activation) + triton_out = gemm_a16w16(x, w, None, out_dtype, y, activation=activation) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) @pytest.mark.parametrize("M, N, K", get_x_vals()) @pytest.mark.parametrize("layout", ["TN", "TT", "NN", "NT"]) -@pytest.mark.parametrize("backend", ["triton", "flydsl"]) -def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout, backend): - if backend == "flydsl" and not is_flydsl_supported(): - pytest.skip("FlyDSL not supported (requires gfx1250 and flydsl package)") - +def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout): torch.cuda.empty_cache() # Helps avoid hangs in large tests x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( @@ -184,12 +136,9 @@ def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout, backend): torch_out = F.linear(x, w, bias=None) - if backend == "flydsl": - kernel_out = _flydsl_gemm_a16w16(x, w, dtype=torch.bfloat16) - else: - kernel_out = gemm_a16w16(x, w, None, out_dtype, y) + triton_out = gemm_a16w16(x, w, None, out_dtype, y) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-1) + torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) @pytest.mark.parametrize("M, N, K", get_x_vals()) From eb82c24533714b473cb31130779b1751ba0468b8 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Mon, 20 Apr 2026 19:02:54 +0000 Subject: [PATCH 03/34] UT update --- op_tests/flydsl_tests/test_gemm_a16w16.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index c8cb5e59499..519f5278e13 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -1,13 +1,5 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""FlyDSL-only unit test for gemm_a16w16 on gfx1250. - -Intentionally isolated from the triton test harness: does NOT import -`aiter.ops.flydsl` or `aiter.ops.triton` at top level, because those -trigger a build of `module_aiter_core` (via `aiter.ops.enum`) which is -orthogonal to the flydsl kernel under test. The kernel file is loaded -directly by file path. -""" import importlib.util import os From 6e04e37a4dde3f6755a424f3d034d90f28495427 Mon Sep 17 00:00:00 2001 From: Satya Nikhil Date: Sat, 11 Jul 2026 02:37:50 +0000 Subject: [PATCH 04/34] code updates, best runs --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 807 ++++++++++-------- 1 file changed, 438 insertions(+), 369 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index 50c79cb68e3..b7aa30d0127 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -6,11 +6,20 @@ import flydsl.expr as fx from flydsl._mlir import ir from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, tdm_ops, vector +from flydsl.expr import ( + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + tdm_ops, + vector, +) from flydsl.expr.arith import _to_raw as _raw from flydsl.expr.typing import T from flydsl.runtime.device import get_rocm_arch as get_hip_arch -from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr, get_op_result_or_value +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr from flydsl.expr import idx2crd from typing import Optional @@ -18,11 +27,23 @@ WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 WAVE_SIZE = 32 -# LDS padding per row (8 elements -> 16 bytes) +# LDS row padding (8 elems = 16 bytes) LDS_PAD_A = 8 LDS_PAD_B = 8 -_STAGE_NAMES = tuple(f"stage{i}" for i in range(16)) +_STAGE_NAMES = ("ping", "pong", "pang", "pung") + + +def _disable_unroll_on_enclosing_loop(): + """Tag the enclosing scf.for with llvm.loop metadata so it stays rolled at ASM level.""" + block = ir.InsertionPoint.current.block + op = block.owner + if op.name != "scf.for": + return + anno = ir.Attribute.parse( + "#llvm.loop_annotation, disableNonforced = true>" + ) + op.attributes["loop_annotation"] = anno def _apply_activation_scalar(val, activation: str): @@ -87,16 +108,16 @@ def compile_gemm_a16w16( activation: Optional[str] = None, add_bias: bool = False, physical_mk: bool = True, # True=M-major (row-major X), False=K-major (col-major X) - physical_kn: bool = False, # False=N-major (row-major W), True=K-major (col-major/transposed W) + physical_kn: bool = False, # False=N-major (row-major W), True=K-major (transposed W) + loop_carried_load_percent: Optional[int] = None, + kernarg_preload: bool = False, + use_manual_barrier: bool = False, ): - """Compile and return a launch function for the A16W16 GEMM kernel. - Returns a callable: launch_fn(y, x, w, bias, M, N, stream=stream) - """ + """Compile the A16W16 GEMM kernel; returns launch_fn(y, x, w, bias, M, N, stream=stream).""" _ = (M, N) - # ── Input validation ── - if num_buffers < 2: - raise ValueError(f"num_buffers must be >= 2, got {num_buffers}") + if not (2 <= num_buffers <= 8): + raise ValueError(f"num_buffers must be between 2 and 8, got {num_buffers}") if in_dtype not in ("fp16", "bf16"): raise ValueError(f"in_dtype must be 'fp16' or 'bf16', got {in_dtype!r}") @@ -115,7 +136,6 @@ def compile_gemm_a16w16( num_warps = m_warp * n_warp block_threads = num_warps * WAVE_SIZE - # ── Tile dimension validation ── if K % tile_k != 0: raise ValueError(f"K must be divisible by tile_k={tile_k}, got K={K}") if tile_k % WMMA_K != 0: @@ -127,7 +147,6 @@ def compile_gemm_a16w16( if (tile_k & (tile_k - 1)) != 0: raise ValueError(f"tile_k must be a power of 2 for TDM, got {tile_k}") - # ── Physical layout validation ── if physical_kn: if N == 0: raise ValueError( @@ -151,7 +170,6 @@ def compile_gemm_a16w16( f"(TDM pad_interval requirement), got {tile_m}" ) - # ── Warp tile dimensions ── warp_tile_m = tile_m // m_warp warp_tile_n = tile_n // n_warp if warp_tile_m % WMMA_M != 0: @@ -159,7 +177,6 @@ def compile_gemm_a16w16( if warp_tile_n % WMMA_N != 0: raise ValueError(f"warp_tile_n={warp_tile_n} must be a multiple of {WMMA_N}") - # ── K-dimension tiling ── num_k_tiles = K // tile_k if num_k_tiles < num_buffers - 1: raise ValueError( @@ -167,13 +184,11 @@ def compile_gemm_a16w16( f"got {num_k_tiles} (K={K}, tile_k={tile_k})" ) - # ── Architecture check ── gpu_arch = str(get_hip_arch()) assert gpu_arch.startswith("gfx1250"), f"Expected gfx1250, got {gpu_arch}" wmma_op = rocdl.wmma_f32_16x16x32_f16 if is_f16 else rocdl.wmma_f32_16x16x32_bf16 - # ── Compute repetition counts ── k_wmma_steps = tile_k // WMMA_K def _elem_type(): @@ -182,11 +197,10 @@ def _elem_type(): wmma_m_rep = warp_tile_m // WMMA_M wmma_n_rep = warp_tile_n // WMMA_N n_accs = wmma_m_rep * wmma_n_rep - n_frags_per_tile = k_wmma_steps * (wmma_n_rep + wmma_m_rep) + N_A_FRAGS = wmma_m_rep * k_wmma_steps + N_B_FRAGS = wmma_n_rep * k_wmma_steps - # ── LDS layout ── - # physical_mk=True: A in LDS as [tile_m, tile_k + pad] - # physical_mk=False: A in LDS as [tile_k, tile_m + pad] (K-major) + # LDS layout: A is [tile_m, tile_k+pad] (M-major) or [tile_k, tile_m+pad] (K-major) if physical_mk: lds_a_stride = tile_k + LDS_PAD_A lds_a_elems = tile_m * lds_a_stride + LDS_PAD_A @@ -201,25 +215,28 @@ def _elem_type(): lds_b_stride = tile_k + LDS_PAD_B lds_b_elems = tile_n * lds_b_stride + LDS_PAD_B - buf_size_elems = lds_a_elems + lds_b_elems - - # ── LDS allocation per pipeline stage ── - stage_allocators = [] - stage_a_offsets = [] - stage_b_offsets = [] - for i in range(num_buffers): - name = _STAGE_NAMES[i] - alloc = SmemAllocator(None, arch=gpu_arch, global_sym_name=f"a16w16_{name}") - off = alloc._align(alloc.ptr, 16) - alloc.ptr = off + buf_size_elems * elem_bytes - stage_allocators.append(alloc) - stage_a_offsets.append(off) - stage_b_offsets.append(off + lds_a_elems * elem_bytes) - - # ── Pipeline iteration counts ── - pre_loaded = num_buffers - 1 - main_loop_iters = (num_k_tiles - pre_loaded) // num_buffers - tail_tiles = (num_k_tiles - pre_loaded) % num_buffers + pre_loaded + lds_a_data_bytes = lds_a_elems * elem_bytes + lds_b_data_bytes = lds_b_elems * elem_bytes + + # Unified LDS allocator: contiguous [A0..A_nb-1 | B0..B_nb-1] ring slots + unified_alloc = SmemAllocator( + None, arch=gpu_arch, global_sym_name="a16w16_unified" + ) + unified_a_off = unified_alloc._align(unified_alloc.ptr, 16) + unified_alloc.ptr = unified_a_off + num_buffers * lds_a_data_bytes + unified_b_off = unified_alloc._align(unified_alloc.ptr, 16) + unified_alloc.ptr = unified_b_off + num_buffers * lds_b_data_bytes + + stage_a_data_off = [ + unified_a_off + i * lds_a_data_bytes for i in range(num_buffers) + ] + stage_b_data_off = [ + unified_b_off + i * lds_b_data_bytes for i in range(num_buffers) + ] + stage_allocators = [unified_alloc] * num_buffers + + # TDMs per K-tile: A + B + _TDMS_PER_TILE = 2 @flyc.kernel def kernel_gemm_a16w16( @@ -232,14 +249,13 @@ def kernel_gemm_a16w16( ): rocdl.disable_xdl_arb_stall() - # ── Thread/block indexing ── tx = gpu.thread_id("x") bx = gpu.block_id("x") by = gpu.block_id("y") blk_m = bx * arith.index(tile_m) blk_n = by * arith.index(tile_n) - # ── Thread-to-warp decomposition ── + # Thread -> warp decomposition layout_thr = fx.make_layout( (m_warp, n_warp, 2, 16), (n_warp * WAVE_SIZE, WAVE_SIZE, 16, 1) ) @@ -261,12 +277,12 @@ def kernel_gemm_a16w16( # Buffer resource descriptors y_nrec = m_idx * n_stride * arith.index(elem_bytes_d) y_rsrc = buffer_ops.create_buffer_resource(arg_y, num_records_bytes=y_nrec) - if add_bias: + if const_expr(add_bias): bias_rsrc = buffer_ops.create_buffer_resource(arg_bias, max_size=True) def make_a_desc(k_base, lds_a_mem_ref): """TDM descriptor for A tile. Swaps dims when physical_mk=False (K-major).""" - if physical_mk: + if const_expr(physical_mk): return tdm_ops.make_tensor_descriptor_2d( global_ptr=arg_x, lds_memref=lds_a_mem_ref, @@ -295,7 +311,7 @@ def make_a_desc(k_base, lds_a_mem_ref): def make_b_desc(k_base, lds_b_mem_ref): """TDM descriptor for B tile. Swaps dims when physical_kn=True (K-major).""" - if physical_kn: + if const_expr(physical_kn): return tdm_ops.make_tensor_descriptor_2d( global_ptr=arg_w, lds_memref=lds_b_mem_ref, @@ -321,23 +337,102 @@ def make_b_desc(k_base, lds_b_mem_ref): num_warps=num_warps, ) - def issue_tdm_load(desc): - tdm_ops.tensor_load_2d(desc) - - def copy_a_to_lds(k_base, lds_a_mem_ref): - issue_tdm_load(make_a_desc(k_base, lds_a_mem_ref)) - - def copy_b_to_lds(k_base, lds_b_mem_ref): - issue_tdm_load(make_b_desc(k_base, lds_b_mem_ref)) - - def _get_lds_memref(lds_ptr): - if isinstance(lds_ptr, SmemPtr): - return get_op_result_or_value(lds_ptr.get()) - return get_op_result_or_value(lds_ptr) + # Unified LDS pointers (all NB slots) + per-stage aliases + big_a = SmemPtr( + unified_alloc.get_base(), + unified_a_off, + elem_ty, + shape=(num_buffers * lds_a_elems,), + ) + big_b = SmemPtr( + unified_alloc.get_base(), + unified_b_off, + elem_ty, + shape=(num_buffers * lds_b_elems,), + ) + big_a_mem = big_a.get() + big_b_mem = big_b.get() + stages_a = [ + SmemPtr( + unified_alloc.get_base(), + stage_a_data_off[i], + elem_ty, + shape=(lds_a_elems,), + ) + for i in range(num_buffers) + ] + stages_b = [ + SmemPtr( + unified_alloc.get_base(), + stage_b_data_off[i], + elem_ty, + shape=(lds_b_elems,), + ) + for i in range(num_buffers) + ] + stages_a_mem = [p.get() for p in stages_a] + stages_b_mem = [p.get() for p in stages_b] + + # TDM descriptors built once at entry; lo32 advances per K-tile, LDS base per ring slot. + _desc_a_init = make_a_desc(arith.index(0), stages_a_mem[0]) + _desc_b_init = make_b_desc(arith.index(0), stages_b_mem[0]) + dgroup1_a = _desc_a_init.dgroup1 + dgroup1_w = _desc_b_init.dgroup1 + addr_hi_a = vector.extract( + _desc_a_init.dgroup0, static_position=[3], dynamic_position=[] + ) + addr_hi_w = vector.extract( + _desc_b_init.dgroup0, static_position=[3], dynamic_position=[] + ) + addr_lo_init_a = vector.extract( + _desc_a_init.dgroup0, static_position=[2], dynamic_position=[] + ) + addr_lo_init_w = vector.extract( + _desc_b_init.dgroup0, static_position=[2], dynamic_position=[] + ) + a_lds_base_i32 = vector.extract( + _desc_a_init.dgroup0, static_position=[1], dynamic_position=[] + ) + b_lds_base_i32 = vector.extract( + _desc_b_init.dgroup0, static_position=[1], dynamic_position=[] + ) + slot_stride_a_i32 = arith.constant(lds_a_data_bytes, type=T.i32) + slot_stride_b_i32 = arith.constant(lds_b_data_bytes, type=T.i32) + slot_stride_a_elems_i32 = arith.constant(lds_a_elems, type=T.i32) + slot_stride_b_elems_i32 = arith.constant(lds_b_elems, type=T.i32) + + # Per-K-tile lo32 byte advance (tile_k*bytes row-major; outer-stride*bytes if K-major). + adv_a_i32 = arith.constant( + (tile_k if physical_mk else tile_k * M) * elem_bytes, type=T.i32 + ) + adv_w_i32 = arith.constant( + (tile_k if not physical_kn else tile_k * N) * elem_bytes, type=T.i32 + ) + pred_const = arith.constant(1, type=T.i32) + + def _buf_idx_to_i32(buf_idx): + """Accept either Python int (prologue/drain) or i32 SSA (main loop).""" + if const_expr(isinstance(buf_idx, int)): + return arith.constant(buf_idx, type=T.i32) + return buf_idx + + def issue_tdm_loads(buf_idx, lo_a, lo_w): + """Issue A+B TDMs for one K-tile into LDS slot buf_idx; returns advanced (lo_a, lo_w).""" + buf_i32 = _buf_idx_to_i32(buf_idx) + a_addr = arith.addi(a_lds_base_i32, arith.muli(buf_i32, slot_stride_a_i32)) + b_addr = arith.addi(b_lds_base_i32, arith.muli(buf_i32, slot_stride_b_i32)) + dg0_a = vector.from_elements( + T.vec(4, T.i32), [pred_const, a_addr, lo_a, addr_hi_a] + ) + tdm_ops.tensor_load_2d(tdm_ops.TDMDescriptor2D(dg0_a, dgroup1_a)) + dg0_w = vector.from_elements( + T.vec(4, T.i32), [pred_const, b_addr, lo_w, addr_hi_w] + ) + tdm_ops.tensor_load_2d(tdm_ops.TDMDescriptor2D(dg0_w, dgroup1_w)) + return arith.addi(lo_a, adv_a_i32), arith.addi(lo_w, adv_w_i32) - def _precompute_lane_bases(lds_ptr, warp_base, reps, lds_stride): - """Pre-compute LDS base addresses for each WMMA rep (M-major layout).""" - lds_buffer = _get_lds_memref(lds_ptr) + def _precompute_lane_bases(warp_base, reps, lds_stride): + """Per-lane LDS element offsets for `reps` WMMA frags (M/N-major).""" row_stride_off = (warp_base + lane16) * arith.index(lds_stride) k_lane_off = lane_kgrp * arith.index(8) bases = [] @@ -346,16 +441,13 @@ def _precompute_lane_bases(lds_ptr, warp_base, reps, lds_stride): row_stride_off + arith.index(rep * WMMA_M * lds_stride) + k_lane_off ) bases.append(base) - return lds_buffer, bases + return bases - def _precompute_a_lane_bases(lds_ptr): + def _precompute_a_lane_bases(): """A fragment lane bases. Uses transpose-load addressing when K-major.""" - if physical_mk: - return _precompute_lane_bases( - lds_ptr, warp_m_base, wmma_m_rep, lds_a_stride - ) + if const_expr(physical_mk): + return _precompute_lane_bases(warp_m_base, wmma_m_rep, lds_a_stride) - lds_buffer = _get_lds_memref(lds_ptr) lane8 = lane16 % arith.index(8) lane_mgrp = lane16 / arith.index(8) k_lane_off = (lane_kgrp * arith.index(8) + lane8) * arith.index( @@ -366,16 +458,13 @@ def _precompute_a_lane_bases(lds_ptr): for wm in range_constexpr(wmma_m_rep): m_col = warp_m_base + arith.index(wm * WMMA_M) + m_lane_off bases.append(k_lane_off + m_col) - return lds_buffer, bases + return bases - def _precompute_b_lane_bases(lds_ptr): + def _precompute_b_lane_bases(): """B fragment lane bases. Uses transpose-load addressing when K-major.""" - if not physical_kn: - return _precompute_lane_bases( - lds_ptr, warp_n_base, wmma_n_rep, lds_b_stride - ) + if const_expr(not physical_kn): + return _precompute_lane_bases(warp_n_base, wmma_n_rep, lds_b_stride) - lds_buffer = _get_lds_memref(lds_ptr) lane8 = lane16 % arith.index(8) lane_ngrp = lane16 / arith.index(8) k_lane_off = (lane_kgrp * arith.index(8) + lane8) * arith.index( @@ -386,7 +475,10 @@ def _precompute_b_lane_bases(lds_ptr): for wn in range_constexpr(wmma_n_rep): n_col = warp_n_base + arith.index(wn * WMMA_N) + n_lane_off bases.append(k_lane_off + n_col) - return lds_buffer, bases + return bases + + a_lane_bases = _precompute_a_lane_bases() + b_lane_bases = _precompute_b_lane_bases() def load_wmma_frag_tr(lds_buffer, b_lane_base, ks): """Load WMMA B fragment via ds_load_tr16_b128 (K-major LDS).""" @@ -423,129 +515,62 @@ def load_wmma_frag_tr_a(lds_buffer, a_lane_base, ks): _load_a_frag = load_wmma_frag_tr_a if not physical_mk else load_wmma_frag - def load_tile_frags(lds_a_ptr, lds_b_ptr): - """Bulk-load all k-slice fragments for one tile (no overlap).""" - a_buf, a_bases = _precompute_a_lane_bases(lds_a_ptr) - b_buf, b_bases = _precompute_b_lane_bases(lds_b_ptr) - frags = [] + def load_a_frags(buf_idx): + """Load all A frags for one K-tile from LDS slot buf_idx; index a_frags[ks*wmma_m_rep+wm].""" + if const_expr(isinstance(buf_idx, int)): + slot_off_a = arith.index(buf_idx * lds_a_elems) + else: + slot_off_a = arith.index_cast( + T.index, arith.muli(buf_idx, slot_stride_a_elems_i32) + ) + a_frags = [] for ks in range_constexpr(k_wmma_steps): - for wn in range_constexpr(wmma_n_rep): - frags.append(_load_b_frag(b_buf, b_bases[wn], ks)) for wm in range_constexpr(wmma_m_rep): - frags.append(_load_a_frag(a_buf, a_bases[wm], ks)) - return frags + a_frags.append( + _load_a_frag(big_a_mem, a_lane_bases[wm] + slot_off_a, ks) + ) + return a_frags - def wmma_tile(accs_in, tile_frags): - """Execute all WMMAs for one tile using pre-loaded fragments.""" - current_accs = list(accs_in) + def load_b_frags(buf_idx): + """Load all B frags for one K-tile from LDS slot buf_idx; index b_frags[ks*wmma_n_rep+wn].""" + if const_expr(isinstance(buf_idx, int)): + slot_off_b = arith.index(buf_idx * lds_b_elems) + else: + slot_off_b = arith.index_cast( + T.index, arith.muli(buf_idx, slot_stride_b_elems_i32) + ) + b_frags = [] for ks in range_constexpr(k_wmma_steps): - base = ks * (wmma_n_rep + wmma_m_rep) - b_frags = [tile_frags[base + wn] for wn in range_constexpr(wmma_n_rep)] - a_frags = [ - tile_frags[base + wmma_n_rep + wm] - for wm in range_constexpr(wmma_m_rep) - ] - for wm in range_constexpr(wmma_m_rep): - for wn in range_constexpr(wmma_n_rep): - idx = wm * wmma_n_rep + wn - # ISA operand order: (B, A, C) — reversed from math - current_accs[idx] = wmma_op( - T.vec(8, T.f32), - b_frags[wn], - a_frags[wm], - current_accs[idx], - signA=False, - signB=False, - modC=0, - reuseA=False, - reuseB=(wn > 0), - ).result - return current_accs + for wn in range_constexpr(wmma_n_rep): + b_frags.append( + _load_b_frag(big_b_mem, b_lane_bases[wn] + slot_off_b, ks) + ) + return b_frags - def compute_and_prefetch(accs_in, cur_frags, next_lds_a_ptr, next_lds_b_ptr): - """Interleave WMMA on current frags with LDS loads for next tile.""" + def wmma_tile(accs_in, a_frags, b_frags): + """Execute all WMMAs for one tile (ks-outer accumulate); reuseB=(wn>0) hint kept.""" current_accs = list(accs_in) - next_a_buf, next_a_bases = _precompute_a_lane_bases(next_lds_a_ptr) - next_b_buf, next_b_bases = _precompute_b_lane_bases(next_lds_b_ptr) - next_frags = [] - for ks in range_constexpr(k_wmma_steps): - # Phase 1: ds_read for next tile - for wn in range_constexpr(wmma_n_rep): - next_frags.append(_load_b_frag(next_b_buf, next_b_bases[wn], ks)) - for wm in range_constexpr(wmma_m_rep): - next_frags.append(_load_a_frag(next_a_buf, next_a_bases[wm], ks)) - - # Phase 2: WMMA for curr tile - base = ks * (wmma_n_rep + wmma_m_rep) - b_frags = [cur_frags[base + wn] for wn in range_constexpr(wmma_n_rep)] - a_frags = [ - cur_frags[base + wmma_n_rep + wm] - for wm in range_constexpr(wmma_m_rep) + a_f = [ + a_frags[ks * wmma_m_rep + wm] for wm in range_constexpr(wmma_m_rep) + ] + b_f = [ + b_frags[ks * wmma_n_rep + wn] for wn in range_constexpr(wmma_n_rep) ] for wm in range_constexpr(wmma_m_rep): for wn in range_constexpr(wmma_n_rep): idx = wm * wmma_n_rep + wn + # ISA operand order: (B, A, C) — reversed from math current_accs[idx] = wmma_op( T.vec(8, T.f32), - b_frags[wn], - a_frags[wm], + b_f[wn], + a_f[wm], current_accs[idx], - signA=False, - signB=False, modC=0, reuseA=False, reuseB=(wn > 0), ).result - - return current_accs, next_frags - - def _l2_prefetch(k_base): - """Issue L2 cache prefetch hints for a future K-tile.""" - if l2_prefetch_distance <= 0: - return - pf_k = k_base + arith.index(l2_prefetch_distance * tile_k) - - if physical_mk: - tdm_ops.l2_prefetch_tile( - arg_x, - (blk_m, pf_k), - (tile_m, tile_k), - (K, 1), - elem_bytes=elem_bytes, - thread_id=tx, - block_threads=block_threads, - ) - else: - tdm_ops.l2_prefetch_tile( - arg_x, - (pf_k, blk_m), - (tile_k, tile_m), - (M, 1), - elem_bytes=elem_bytes, - thread_id=tx, - block_threads=block_threads, - ) - if physical_kn: - tdm_ops.l2_prefetch_tile( - arg_w, - (pf_k, blk_n), - (tile_k, tile_n), - (N, 1), - elem_bytes=elem_bytes, - thread_id=tx, - block_threads=block_threads, - ) - else: - tdm_ops.l2_prefetch_tile( - arg_w, - (blk_n, pf_k), - (tile_n, tile_k), - (K, 1), - elem_bytes=elem_bytes, - thread_id=tx, - block_threads=block_threads, - ) + return current_accs _half_out = out_dtype in ("f16", "bf16") _out_elem = ( @@ -560,7 +585,7 @@ def _widen_to_f32(val): def _apply_bias_and_activation(accs): """Add bias and/or apply activation to accumulators.""" - if not add_bias and activation is None: + if const_expr(not add_bias and activation is None): return accs for wm in range_constexpr(wmma_m_rep): @@ -568,7 +593,7 @@ def _apply_bias_and_activation(accs): idx = wm * wmma_n_rep + wn acc = accs[idx] - if add_bias: + if const_expr(add_bias): col_base = ( blk_n + warp_n_base @@ -600,7 +625,7 @@ def _apply_bias_and_activation(accs): bias_vec = vector.from_elements(T.vec(8, T.f32), bias_elems) acc = acc + bias_vec - if activation is not None: + if const_expr(activation is not None): new_elems = [] for i in range_constexpr(8): val = vector.extract( @@ -626,7 +651,7 @@ def epilogue_stores(final_accs): + lane_kgrp * arith.index(8) ) - if _half_out: + if const_expr(_half_out): h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) i32_vec = vector.bitcast(T.vec(4, T.i32), h_vec) c_off_bytes = (row * n_stride + col_base) * arith.index( @@ -650,196 +675,188 @@ def epilogue_stores(final_accs): c_off = row * n_stride + col buffer_ops.buffer_store(vec4, y_rsrc, c_off) - # Step 1: Initialize accumulators + # Accumulators acc_zero = arith.constant_vector(0.0, T.vec(8, T.f32)) accs = [acc_zero] * n_accs - # Step 2: Set up LDS buffer pointers - base_ptrs = [sa.get_base() for sa in stage_allocators] - stages_a = [ - SmemPtr(base_ptrs[i], stage_a_offsets[i], elem_ty, shape=(lds_a_elems,)) - for i in range_constexpr(num_buffers) - ] - stages_b = [ - SmemPtr(base_ptrs[i], stage_b_offsets[i], elem_ty, shape=(lds_b_elems,)) - for i in range_constexpr(num_buffers) - ] - stages_a_mem = [stages_a[i].get() for i in range_constexpr(num_buffers)] - stages_b_mem = [stages_b[i].get() for i in range_constexpr(num_buffers)] - - # Step 3: Prologue — pre-load tiles into LDS - for i in range_constexpr(pre_loaded): - copy_a_to_lds(arith.index(i * tile_k), stages_a_mem[i]) - copy_b_to_lds(arith.index(i * tile_k), stages_b_mem[i]) - - # L2 prefetch prologue - if l2_prefetch_distance > 0: - for pf_i in range_constexpr(min(l2_prefetch_distance, num_k_tiles)): - pf_k = arith.index(pf_i * tile_k) - if physical_mk: - tdm_ops.l2_prefetch_tile( - arg_x, - (blk_m, pf_k), - (tile_m, tile_k), - (K, 1), - elem_bytes=elem_bytes, - thread_id=tx, - block_threads=block_threads, - ) - else: - tdm_ops.l2_prefetch_tile( - arg_x, - (pf_k, blk_m), - (tile_k, tile_m), - (M, 1), - elem_bytes=elem_bytes, - thread_id=tx, - block_threads=block_threads, - ) - if physical_kn: - tdm_ops.l2_prefetch_tile( - arg_w, - (pf_k, blk_n), - (tile_k, tile_n), - (N, 1), - elem_bytes=elem_bytes, - thread_id=tx, - block_threads=block_threads, - ) - else: - tdm_ops.l2_prefetch_tile( - arg_w, - (blk_n, pf_k), - (tile_n, tile_k), - (K, 1), - elem_bytes=elem_bytes, - thread_id=tx, - block_threads=block_threads, - ) - - # Prepare first TDM descriptors for main loop - if main_loop_iters > 0: - _first_load_stage_0 = pre_loaded % num_buffers - _first_load_k_0 = pre_loaded * tile_k - pending_desc_a = make_a_desc( - arith.index(_first_load_k_0), stages_a_mem[_first_load_stage_0] - ) - pending_desc_b = make_b_desc( - arith.index(_first_load_k_0), stages_b_mem[_first_load_stage_0] - ) - - main_end = main_loop_iters * num_buffers * tile_k - - # Wait for pre-loaded tiles - tdm_ops.tensor_wait(2 * (num_buffers - 2)) + # State carry: [accs (N_ACCS) | cur_a frags (N_A_FRAGS) | cur_b frags (N_B_FRAGS)] + N_ACCS = n_accs + + def _pack_state(accs_, a_, b_): + return list(accs_) + list(a_) + list(b_) + + def _unpack_state(state): + i = 0 + accs_ = list(state[i : i + N_ACCS]) + i += N_ACCS + a_ = list(state[i : i + N_A_FRAGS]) + i += N_A_FRAGS + b_ = list(state[i : i + N_B_FRAGS]) + i += N_B_FRAGS + return accs_, a_, b_ + + # ── Prologue (memory-bound): fill NB-1 LDS buffers, then load tile 0 ── + lo_a = addr_lo_init_a + lo_w = addr_lo_init_w + for i in range_constexpr(num_buffers - 1): + lo_a, lo_w = issue_tdm_loads(i, lo_a, lo_w) + + # Wait retires tile 0 (NB-2 in flight); the 2-iter WAR gap lets the main loop use one barrier. + tdm_ops.tensor_wait((num_buffers - 2) * _TDMS_PER_TILE) gpu.barrier() - # Pre-load fragments for first compute tile - cur_frags = load_tile_frags(stages_a[0], stages_b[0]) - - # Step 4: Main loop - if main_loop_iters > 0: - init_descs = [ - pending_desc_a.dgroup0, - pending_desc_a.dgroup1, - pending_desc_b.dgroup0, - pending_desc_b.dgroup1, - ] - - for iv, state in range( - 0, - main_end, - num_buffers * tile_k, - init=list(accs) + init_descs + cur_frags, - ): - accs_in = list(state[:n_accs]) - cur_desc_a = tdm_ops.TDMDescriptor2D(state[n_accs], state[n_accs + 1]) - cur_desc_b = tdm_ops.TDMDescriptor2D( - state[n_accs + 2], state[n_accs + 3] - ) - cur_frags = list(state[n_accs + 4 : n_accs + 4 + n_frags_per_tile]) - - for s in range_constexpr(num_buffers): - issue_tdm_load(cur_desc_a) - issue_tdm_load(cur_desc_b) - - if s == num_buffers - 1: - _l2_prefetch(iv + arith.index(s * tile_k)) - - if s < num_buffers - 1: - _next_load_stage = (s + 1 + pre_loaded) % num_buffers - _next_load_k_off = (s + 1 + pre_loaded) * tile_k - cur_desc_a = make_a_desc( - iv + arith.index(_next_load_k_off), - stages_a_mem[_next_load_stage], - ) - cur_desc_b = make_b_desc( - iv + arith.index(_next_load_k_off), - stages_b_mem[_next_load_stage], - ) - else: - _next_load_stage = pre_loaded % num_buffers - _next_load_k_off = pre_loaded * tile_k - _next_step = num_buffers * tile_k - cur_desc_a = make_a_desc( - iv + arith.index(_next_step + _next_load_k_off), - stages_a_mem[_next_load_stage], - ) - cur_desc_b = make_b_desc( - iv + arith.index(_next_step + _next_load_k_off), - stages_b_mem[_next_load_stage], - ) - - tdm_ops.tensor_wait(2 * (num_buffers - 2)) - gpu.barrier() - - _next_compute = (s + 1) % num_buffers - accs_in, cur_frags = compute_and_prefetch( - accs_in, - cur_frags, - stages_a[_next_compute], - stages_b[_next_compute], - ) - - out_descs = [ - cur_desc_a.dgroup0, - cur_desc_a.dgroup1, - cur_desc_b.dgroup0, - cur_desc_b.dgroup1, + cur_a = load_a_frags(0) + cur_b = load_b_frags(0) + + # ── Main loop (wait-first, unrolled by 2): ping-pong two register banks, prefetch T+1/T+2 ── + main_loop_iters = num_k_tiles - (num_buffers - 1) + nb_const_i32 = arith.constant(num_buffers, type=T.i32) + one_i32 = arith.constant(1, type=T.i32) + load_idx_init = arith.constant(num_buffers - 1, type=T.i32) + compute_idx_init = arith.constant(0, type=T.i32) + + def _wg_barrier(): + # gpu.barrier() drains ds_loads via fences; manual variant is a bare s_barrier (no fences, faster but racy). + if const_expr(use_manual_barrier): + rocdl.s_barrier_signal(-1) + rocdl.s_barrier_wait(-1) + else: + gpu.barrier() + + def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): + """Compute K-tile cidx from carried frags; if do_prefetch, load cidx+1 into a fresh bank.""" + load_buf_i32 = arith.remui(lidx, nb_const_i32) + + # Wait at the TOP (above the TDM -> does not count this iter's refill). + tdm_ops.tensor_wait(max(0, num_buffers - 3) * _TDMS_PER_TILE) + # Refill LDS slot load_buf with the next-to-issue K-tile. + lo_a_, lo_w_ = issue_tdm_loads(load_buf_i32, lo_a_, lo_w_) + if const_expr(num_buffers == 2): + # NB=2: drain the just-issued TDM to close the RAW hazard on the slot we refilled. + tdm_ops.tensor_wait(0) + # Single barrier: RAW wall for this tile's T+1 ds_reads and (across the backedge) WAR wall for the next refill. + _wg_barrier() + + if const_expr(do_prefetch): + next_buf_i32 = arith.remui(arith.addi(cidx, one_i32), nb_const_i32) + # Preload next tile's A+B, then WMMA (loads co-execute with it); raises peak liveness to ~2 banks. + na = load_a_frags(next_buf_i32) + nb = load_b_frags(next_buf_i32) + accs_ = wmma_tile(accs_, ca, cb) + nbank = (na, nb) + else: + accs_ = wmma_tile(accs_, ca, cb) + nbank = None + return accs_, nbank, lo_a_, lo_w_ + + if const_expr(main_loop_iters > 0): + lo_a_s = lo_a + lo_w_s = lo_w + load_idx_s = load_idx_init + compute_idx_s = compute_idx_init + num_pairs = main_loop_iters // 2 + + if const_expr(num_pairs > 0): + init_state = _pack_state(accs, cur_a, cur_b) + [ + lo_a_s, + lo_w_s, + load_idx_s, + compute_idx_s, ] - results = yield list(accs_in) + out_descs + cur_frags - - accs = list(results[:n_accs]) - cur_frags = list(results[n_accs + 4 : n_accs + 4 + n_frags_per_tile]) - - # Step 5: drain remaining tiles - _tail_base_k = main_loop_iters * num_buffers - _extra_loads = tail_tiles - pre_loaded + results = init_state + for pair_step, state in range(0, num_pairs, 1, init=init_state): + _disable_unroll_on_enclosing_loop() + cidx = state[-1] + lidx = state[-2] + lw = state[-3] + lx = state[-4] + p_accs, b0a, b0b = _unpack_state(state[:-4]) + + # sub-0: tile T from bank0, prefetch T+1 -> bank1. + p_accs, bank1, lx, lw = _run_tile( + p_accs, b0a, b0b, lx, lw, cidx, lidx, do_prefetch=True + ) + b1a, b1b = bank1 + cidx1 = arith.addi(cidx, one_i32) + lidx1 = arith.addi(lidx, one_i32) - for t in range_constexpr(tail_tiles - 1): - if t < _extra_loads: - load_tile_idx = _tail_base_k + pre_loaded + t - load_stage = (pre_loaded + t) % num_buffers - copy_a_to_lds( - arith.index(load_tile_idx * tile_k), stages_a_mem[load_stage] - ) - copy_b_to_lds( - arith.index(load_tile_idx * tile_k), stages_b_mem[load_stage] + # sub-1: tile T+1 from bank1, prefetch T+2 -> bank0 (yielded). + p_accs, bank0, lx, lw = _run_tile( + p_accs, b1a, b1b, lx, lw, cidx1, lidx1, do_prefetch=True + ) + nb0a, nb0b = bank0 + cidx2 = arith.addi(cidx1, one_i32) + lidx2 = arith.addi(lidx1, one_i32) + + new_state = _pack_state(p_accs, nb0a, nb0b) + [ + lx, + lw, + lidx2, + cidx2, + ] + results = yield new_state + + accs, cur_a, cur_b = _unpack_state(results[:-4]) + lo_a_s = results[-4] + lo_w_s = results[-3] + load_idx_s = results[-2] + compute_idx_s = results[-1] + else: + accs = list(accs) + + # Leftover odd tile: compute the carried bank and prefetch the first drain tile into the carry. + if const_expr(main_loop_iters % 2 == 1): + accs, leftover_bank, lo_a_s, lo_w_s = _run_tile( + accs, + cur_a, + cur_b, + lo_a_s, + lo_w_s, + compute_idx_s, + load_idx_s, + do_prefetch=True, ) + cur_a, cur_b = leftover_bank + compute_idx_s = arith.addi(compute_idx_s, one_i32) - _epi_outstanding = 2 * (pre_loaded + min(t + 1, _extra_loads) - t - 1) - tdm_ops.tensor_wait(_epi_outstanding) - gpu.barrier() - - _next_epi_stage = (_tail_base_k + t + 1) % num_buffers - accs, cur_frags = compute_and_prefetch( - accs, cur_frags, stages_a[_next_epi_stage], stages_b[_next_epi_stage] - ) - - # Step 6: Final WMMA - accs = wmma_tile(accs, cur_frags) + final_compute_idx = compute_idx_s + else: + accs = list(accs) + # No main loop ran — drain starts at compute_idx = 0. + final_compute_idx = arith.constant(0, type=T.i32) + + # ── Drain (fully unrolled): consume carried frags, prefetch next bank per tile; final tile does no wait/barrier/ds_load ── + drain_count_d = ( + num_buffers - 1 + if main_loop_iters > 0 + else min(num_buffers - 1, num_k_tiles) + ) + drain_base = main_loop_iters if main_loop_iters > 0 else 0 + + def _drain_wait_for(tile_idx): + # Wait for compile-time drain tile tile_idx to land. + tdm_ops.tensor_wait(max(0, num_k_tiles - 1 - tile_idx) * _TDMS_PER_TILE) + + accs = list(accs) + drain_a, drain_b = cur_a, cur_b + for j in range_constexpr(drain_count_d): + tile_idx = drain_base + j + if const_expr(j < drain_count_d - 1): + # Not the last tile: wait + barrier + prefetch tile_idx+1, interleaved with the current WMMA. + next_tile = tile_idx + 1 + _drain_wait_for(next_tile) + gpu.barrier() + next_a = load_a_frags(next_tile % num_buffers) + accs = wmma_tile(accs, drain_a, drain_b) + next_b = load_b_frags(next_tile % num_buffers) + drain_a, drain_b = next_a, next_b + else: + # Final K-tile: compute the carried frags — no wait, barrier, or ds_loads. + accs = wmma_tile(accs, drain_a, drain_b) - # Step 7: Bias, activation, and store + # ── Epilogue: bias, activation, and store ── + if const_expr(num_buffers > 2): + rocdl.sched_barrier(0) accs = _apply_bias_and_activation(accs) epilogue_stores(accs) @@ -861,6 +878,9 @@ def epilogue_stores(final_accs): physical_kn, M, N, + loop_carried_load_percent, + kernarg_preload, + use_manual_barrier, ) @flyc.jit @@ -902,12 +922,55 @@ def launch_gemm_a16w16( ir.IntegerType.get_signless(32), _wpe ) + # Flat work-group size hint. + flat_wg_attr = ir.StringAttr.get(f"{block_threads},{block_threads}") + for op in ctx.gpu_module_body.operations: + if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": + op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr + + # Experimental: amdgpu-loop-carried-load-percent passthrough. + if loop_carried_load_percent is not None: + lcv = ir.ArrayAttr.get( + [ + ir.ArrayAttr.get( + [ + ir.StringAttr.get("amdgpu-loop-carried-load-percent"), + ir.StringAttr.get(str(int(loop_carried_load_percent))), + ] + ) + ] + ) + for op in ctx.gpu_module_body.operations: + if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": + op.attributes["passthrough"] = lcv + + # Mark kernel args inreg so AMDGPU preloads them into user SGPRs. + if kernarg_preload: + inreg_attr = ir.UnitAttr.get() + for op in ctx.gpu_module_body.operations: + if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": + num_args = len(op.regions[0].blocks[0].arguments) + per_arg = [ + ir.DictAttr.get({"llvm.inreg": inreg_attr}) + for _ in range(num_args) + ] + op.attributes["arg_attrs"] = ir.ArrayAttr.get(per_arg) + launcher.launch( grid=(gx, gy, 1), block=(block_threads, 1, 1), stream=stream, ) + # Backend scheduling hints: only cl::opts present in THIS LLVM are enabled; others raise "Unknown LLVM option" and stay commented. + launch_gemm_a16w16.compile_hints["llvm_options"] = { + "amdgpu-expert-scheduling-mode": True, # valid bool cl::opt (GFX12+ only) + # Not present in this LLVM; re-enable when built against a downstream LLVM that defines them: + # "amdgpu-anti-hints-for-va-vdst", "amdgpu-enable-static-simulator", + # "amdgpu-static-sim-inline", "amdgpu-block-carried-latency", + # "amdgpu-sched-strategy" ("coexec" is also unrecognized upstream) + } + return launch_gemm_a16w16 @@ -926,6 +989,9 @@ def gemm_a16w16( num_buffers: int = 2, waves_per_eu: int = None, l2_prefetch_distance: int = 2, + kernarg_preload: bool = False, + loop_carried_load_percent: Optional[int] = None, + use_manual_barrier: bool = False, ): """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" assert x.dtype in ( @@ -941,7 +1007,7 @@ def gemm_a16w16( M, K = x.shape N = w.shape[0] - # ── Layout ── + # Layout from strides if x.stride(1) == 1: physical_mk = True elif x.stride(0) == 1: @@ -952,7 +1018,7 @@ def gemm_a16w16( elif w.stride(0) == 1: physical_kn = True - # ── K-pad ── + # K-pad up to a tile_k multiple K_padded = ((K + tile_k - 1) // tile_k) * tile_k if K_padded != K: pad_size = K_padded - K @@ -970,10 +1036,10 @@ def gemm_a16w16( w = torch.nn.functional.pad(w, (0, pad_size)) K = K_padded - # ── N-padding ── + # N-pad up to a tile_n multiple N_stride = ((N + tile_n - 1) // tile_n) * tile_n - # ── Output allocation ── + # Output allocation if y is not None: y_buf = ( y @@ -1016,6 +1082,9 @@ def gemm_a16w16( add_bias=(bias.numel() > 0), physical_mk=physical_mk, physical_kn=physical_kn, + kernarg_preload=kernarg_preload, + loop_carried_load_percent=loop_carried_load_percent, + use_manual_barrier=use_manual_barrier, ) launch_fn(y_buf, x, w, bias, M, N_stride, stream=torch.cuda.current_stream()) From a3563205dc1cd742425ed950ef2ec8913cfa3165 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 16 Jul 2026 20:35:26 +0000 Subject: [PATCH 05/34] Kernel updates,split k for smaller configs and unroll variable --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 347 +++++++++++++----- aiter/ops/shuffle.py | 38 ++ op_tests/flydsl_tests/test_gemm_a16w16.py | 128 ++++++- 3 files changed, 397 insertions(+), 116 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index b7aa30d0127..8caf045a3ee 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -112,6 +112,10 @@ def compile_gemm_a16w16( loop_carried_load_percent: Optional[int] = None, kernarg_preload: bool = False, use_manual_barrier: bool = False, + split_k: int = 1, + sched_strategy: Optional[str] = None, + barrier_signal_wait_latency: Optional[int] = None, + main_loop_unroll: bool = False, ): """Compile the A16W16 GEMM kernel; returns launch_fn(y, x, w, bias, M, N, stream=stream).""" _ = (M, N) @@ -120,6 +124,15 @@ def compile_gemm_a16w16( raise ValueError(f"num_buffers must be between 2 and 8, got {num_buffers}") if in_dtype not in ("fp16", "bf16"): raise ValueError(f"in_dtype must be 'fp16' or 'bf16', got {in_dtype!r}") + # Experimental LLVM scheduling levers (injected via llvm_options below). + if sched_strategy is not None and sched_strategy not in ( + "max-ilp", + "max-memory-clause", + ): + raise ValueError( + "sched_strategy must be None, 'max-ilp', or 'max-memory-clause', " + f"got {sched_strategy!r}" + ) effective_waves_per_eu = waves_per_eu is_f16 = in_dtype == "fp16" @@ -138,6 +151,10 @@ def compile_gemm_a16w16( if K % tile_k != 0: raise ValueError(f"K must be divisible by tile_k={tile_k}, got K={K}") + if split_k < 1: + raise ValueError(f"split_k must be >= 1, got {split_k}") + if K % split_k != 0: + raise ValueError(f"K must be divisible by split_k={split_k}, got K={K}") if tile_k % WMMA_K != 0: raise ValueError(f"tile_k must be a multiple of {WMMA_K}, got {tile_k}") if tile_m % WMMA_M != 0: @@ -177,7 +194,13 @@ def compile_gemm_a16w16( if warp_tile_n % WMMA_N != 0: raise ValueError(f"warp_tile_n={warp_tile_n} must be a multiple of {WMMA_N}") - num_k_tiles = K // tile_k + # split_k>1 splits K into split_k grid-z chunks; num_k_tiles is per-split. + split_k_chunk = K // split_k + if split_k_chunk % tile_k != 0: + raise ValueError( + f"K/split_k must be divisible by tile_k={tile_k}, got {split_k_chunk}" + ) + num_k_tiles = split_k_chunk // tile_k if num_k_tiles < num_buffers - 1: raise ValueError( f"{num_buffers}-stage buffering requires num_k_tiles >= {num_buffers - 1}, " @@ -219,9 +242,7 @@ def _elem_type(): lds_b_data_bytes = lds_b_elems * elem_bytes # Unified LDS allocator: contiguous [A0..A_nb-1 | B0..B_nb-1] ring slots - unified_alloc = SmemAllocator( - None, arch=gpu_arch, global_sym_name="a16w16_unified" - ) + unified_alloc = SmemAllocator(None, arch=gpu_arch, global_sym_name="a16w16_unified") unified_a_off = unified_alloc._align(unified_alloc.ptr, 16) unified_alloc.ptr = unified_a_off + num_buffers * lds_a_data_bytes unified_b_off = unified_alloc._align(unified_alloc.ptr, 16) @@ -254,6 +275,10 @@ def kernel_gemm_a16w16( by = gpu.block_id("y") blk_m = bx * arith.index(tile_m) blk_n = by * arith.index(tile_n) + # split_k>1: grid-z block picks the K-chunk origin (compile-time gated). + if const_expr(split_k > 1): + bz = gpu.block_id("z") + split_k_base = bz * arith.index(split_k_chunk) # Thread -> warp decomposition layout_thr = fx.make_layout( @@ -374,8 +399,12 @@ def make_b_desc(k_base, lds_b_mem_ref): stages_b_mem = [p.get() for p in stages_b] # TDM descriptors built once at entry; lo32 advances per K-tile, LDS base per ring slot. - _desc_a_init = make_a_desc(arith.index(0), stages_a_mem[0]) - _desc_b_init = make_b_desc(arith.index(0), stages_b_mem[0]) + if const_expr(split_k > 1): + _desc_a_init = make_a_desc(split_k_base, stages_a_mem[0]) + _desc_b_init = make_b_desc(split_k_base, stages_b_mem[0]) + else: + _desc_a_init = make_a_desc(arith.index(0), stages_a_mem[0]) + _desc_b_init = make_b_desc(arith.index(0), stages_b_mem[0]) dgroup1_a = _desc_a_init.dgroup1 dgroup1_w = _desc_b_init.dgroup1 addr_hi_a = vector.extract( @@ -639,7 +668,9 @@ def _apply_bias_and_activation(accs): return accs def epilogue_stores(final_accs): - """Write accumulators to global output Y.""" + """Write accumulators to Y (buffer_store, or atomic-fadd if split_k>1).""" + if const_expr(split_k > 1): + zero_i32 = fx.Int32(0) for wm in range_constexpr(wmma_m_rep): for wn in range_constexpr(wmma_n_rep): idx = wm * wmma_n_rep + wn @@ -652,28 +683,70 @@ def epilogue_stores(final_accs): ) if const_expr(_half_out): - h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) - i32_vec = vector.bitcast(T.vec(4, T.i32), h_vec) - c_off_bytes = (row * n_stride + col_base) * arith.index( - elem_bytes_d - ) - buffer_ops.buffer_store( - i32_vec, y_rsrc, c_off_bytes, offset_is_bytes=True - ) - else: - for half in range_constexpr(2): - vals = [ - vector.extract( - final_accs[idx], - static_position=[half * 4 + vi], + if const_expr(split_k > 1): + h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) + c_off_bytes = (row * n_stride + col_base) * arith.index( + elem_bytes_d + ) + pair_ty = T.vec(2, _out_elem) + for pair in range_constexpr(4): + e0 = vector.extract( + h_vec, + static_position=[pair * 2], + dynamic_position=[], + ) + e1 = vector.extract( + h_vec, + static_position=[pair * 2 + 1], dynamic_position=[], ) - for vi in range_constexpr(4) - ] - vec4 = vector.from_elements(T.vec(4, T.f32), vals) - col = col_base + arith.index(half * 4) - c_off = row * n_stride + col - buffer_ops.buffer_store(vec4, y_rsrc, c_off) + pair_vec = vector.from_elements(pair_ty, [e0, e1]) + byte_off = arith.index_cast( + T.i32, c_off_bytes + arith.index(pair * 4) + ) + rocdl.raw_ptr_buffer_atomic_fadd( + pair_vec, y_rsrc, byte_off, zero_i32, zero_i32 + ) + else: + h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) + i32_vec = vector.bitcast(T.vec(4, T.i32), h_vec) + c_off_bytes = (row * n_stride + col_base) * arith.index( + elem_bytes_d + ) + buffer_ops.buffer_store( + i32_vec, y_rsrc, c_off_bytes, offset_is_bytes=True + ) + else: + if const_expr(split_k > 1): + for half in range_constexpr(2): + base = row * n_stride + col_base + arith.index(half * 4) + for vi in range_constexpr(4): + val = vector.extract( + final_accs[idx], + static_position=[half * 4 + vi], + dynamic_position=[], + ) + byte_off = arith.index_cast( + T.i32, + (base + arith.index(vi)) * arith.index(4), + ) + rocdl.raw_ptr_buffer_atomic_fadd( + val, y_rsrc, byte_off, zero_i32, zero_i32 + ) + else: + for half in range_constexpr(2): + vals = [ + vector.extract( + final_accs[idx], + static_position=[half * 4 + vi], + dynamic_position=[], + ) + for vi in range_constexpr(4) + ] + vec4 = vector.from_elements(T.vec(4, T.f32), vals) + col = col_base + arith.index(half * 4) + c_off = row * n_stride + col + buffer_ops.buffer_store(vec4, y_rsrc, c_off) # Accumulators acc_zero = arith.constant_vector(0.0, T.vec(8, T.f32)) @@ -727,13 +800,14 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): """Compute K-tile cidx from carried frags; if do_prefetch, load cidx+1 into a fresh bank.""" load_buf_i32 = arith.remui(lidx, nb_const_i32) - # Wait at the TOP (above the TDM -> does not count this iter's refill). - tdm_ops.tensor_wait(max(0, num_buffers - 3) * _TDMS_PER_TILE) - # Refill LDS slot load_buf with the next-to-issue K-tile. + # Refill first, then wait: keeps NB-2 tiles' TDM in flight (depth-2). lo_a_, lo_w_ = issue_tdm_loads(load_buf_i32, lo_a_, lo_w_) if const_expr(num_buffers == 2): - # NB=2: drain the just-issued TDM to close the RAW hazard on the slot we refilled. + # NB=2: only 2 slots -> drain the just-issued TDM (RAW hazard). tdm_ops.tensor_wait(0) + else: + tdm_ops.tensor_wait((num_buffers - 2) * _TDMS_PER_TILE) + # Single barrier: RAW wall for this tile's T+1 ds_reads and (across the backedge) WAR wall for the next refill. _wg_barrier() @@ -750,80 +824,113 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): return accs_, nbank, lo_a_, lo_w_ if const_expr(main_loop_iters > 0): - lo_a_s = lo_a - lo_w_s = lo_w - load_idx_s = load_idx_init - compute_idx_s = compute_idx_init - num_pairs = main_loop_iters // 2 + if const_expr(main_loop_unroll): + # Unroll-by-2: 2 tiles/trip, ping-pong 2 register banks + odd leftover. + lo_a_s = lo_a + lo_w_s = lo_w + load_idx_s = load_idx_init + compute_idx_s = compute_idx_init + num_pairs = main_loop_iters // 2 + + if const_expr(num_pairs > 0): + init_state = _pack_state(accs, cur_a, cur_b) + [ + lo_a_s, + lo_w_s, + load_idx_s, + compute_idx_s, + ] + results = init_state + for pair_step, state in range(0, num_pairs, 1, init=init_state): + _disable_unroll_on_enclosing_loop() + cidx = state[-1] + lidx = state[-2] + lw = state[-3] + lx = state[-4] + p_accs, b0a, b0b = _unpack_state(state[:-4]) + + # sub-0: tile T from bank0, prefetch T+1 -> bank1. + p_accs, bank1, lx, lw = _run_tile( + p_accs, b0a, b0b, lx, lw, cidx, lidx, do_prefetch=True + ) + b1a, b1b = bank1 + cidx1 = arith.addi(cidx, one_i32) + lidx1 = arith.addi(lidx, one_i32) - if const_expr(num_pairs > 0): + # sub-1: tile T+1 from bank1, prefetch T+2 -> bank0 (yielded). + p_accs, bank0, lx, lw = _run_tile( + p_accs, b1a, b1b, lx, lw, cidx1, lidx1, do_prefetch=True + ) + nb0a, nb0b = bank0 + cidx2 = arith.addi(cidx1, one_i32) + lidx2 = arith.addi(lidx1, one_i32) + + new_state = _pack_state(p_accs, nb0a, nb0b) + [ + lx, + lw, + lidx2, + cidx2, + ] + results = yield new_state + + accs, cur_a, cur_b = _unpack_state(results[:-4]) + lo_a_s = results[-4] + lo_w_s = results[-3] + load_idx_s = results[-2] + compute_idx_s = results[-1] + else: + accs = list(accs) + + # Leftover odd tile: compute the carried bank, prefetch first drain tile. + if const_expr(main_loop_iters % 2 == 1): + accs, leftover_bank, lo_a_s, lo_w_s = _run_tile( + accs, + cur_a, + cur_b, + lo_a_s, + lo_w_s, + compute_idx_s, + load_idx_s, + do_prefetch=True, + ) + cur_a, cur_b = leftover_bank + compute_idx_s = arith.addi(compute_idx_s, one_i32) + else: + # Single tile/trip: compute T from carried bank, prefetch T+1. init_state = _pack_state(accs, cur_a, cur_b) + [ - lo_a_s, - lo_w_s, - load_idx_s, - compute_idx_s, + lo_a, + lo_w, + load_idx_init, + compute_idx_init, ] results = init_state - for pair_step, state in range(0, num_pairs, 1, init=init_state): + for step, state in range(0, main_loop_iters, 1, init=init_state): _disable_unroll_on_enclosing_loop() cidx = state[-1] lidx = state[-2] lw = state[-3] lx = state[-4] - p_accs, b0a, b0b = _unpack_state(state[:-4]) + p_accs, ca, cb = _unpack_state(state[:-4]) - # sub-0: tile T from bank0, prefetch T+1 -> bank1. - p_accs, bank1, lx, lw = _run_tile( - p_accs, b0a, b0b, lx, lw, cidx, lidx, do_prefetch=True + # tile T from the carried bank, prefetch T+1 -> next bank (yielded). + p_accs, bank, lx, lw = _run_tile( + p_accs, ca, cb, lx, lw, cidx, lidx, do_prefetch=True ) - b1a, b1b = bank1 + next_a, next_b = bank cidx1 = arith.addi(cidx, one_i32) lidx1 = arith.addi(lidx, one_i32) - # sub-1: tile T+1 from bank1, prefetch T+2 -> bank0 (yielded). - p_accs, bank0, lx, lw = _run_tile( - p_accs, b1a, b1b, lx, lw, cidx1, lidx1, do_prefetch=True - ) - nb0a, nb0b = bank0 - cidx2 = arith.addi(cidx1, one_i32) - lidx2 = arith.addi(lidx1, one_i32) - - new_state = _pack_state(p_accs, nb0a, nb0b) + [ + new_state = _pack_state(p_accs, next_a, next_b) + [ lx, lw, - lidx2, - cidx2, + lidx1, + cidx1, ] results = yield new_state accs, cur_a, cur_b = _unpack_state(results[:-4]) - lo_a_s = results[-4] - lo_w_s = results[-3] - load_idx_s = results[-2] - compute_idx_s = results[-1] - else: - accs = list(accs) - - # Leftover odd tile: compute the carried bank and prefetch the first drain tile into the carry. - if const_expr(main_loop_iters % 2 == 1): - accs, leftover_bank, lo_a_s, lo_w_s = _run_tile( - accs, - cur_a, - cur_b, - lo_a_s, - lo_w_s, - compute_idx_s, - load_idx_s, - do_prefetch=True, - ) - cur_a, cur_b = leftover_bank - compute_idx_s = arith.addi(compute_idx_s, one_i32) - - final_compute_idx = compute_idx_s else: accs = list(accs) # No main loop ran — drain starts at compute_idx = 0. - final_compute_idx = arith.constant(0, type=T.i32) # ── Drain (fully unrolled): consume carried frags, prefetch next bank per tile; final tile does no wait/barrier/ds_load ── drain_count_d = ( @@ -881,6 +988,10 @@ def _drain_wait_for(tile_idx): loop_carried_load_percent, kernarg_preload, use_manual_barrier, + split_k, + sched_strategy, + barrier_signal_wait_latency, + main_loop_unroll, ) @flyc.jit @@ -908,6 +1019,7 @@ def launch_gemm_a16w16( idx_n = arith.index_cast(T.index, i32_n.ir_value()) gx = _raw((idx_m + arith.index(tile_m - 1)) / arith.index(tile_m)) gy = _raw((idx_n + arith.index(tile_n - 1)) / arith.index(tile_n)) + gz = split_k # 1 at default -> grid tuple identical to original # Emit kernel launcher = kernel_gemm_a16w16(arg_y, arg_x, arg_w, arg_bias, i32_m, i32_n) @@ -957,19 +1069,23 @@ def launch_gemm_a16w16( op.attributes["arg_attrs"] = ir.ArrayAttr.get(per_arg) launcher.launch( - grid=(gx, gy, 1), + grid=(gx, gy, gz), block=(block_threads, 1, 1), stream=stream, ) - # Backend scheduling hints: only cl::opts present in THIS LLVM are enabled; others raise "Unknown LLVM option" and stay commented. - launch_gemm_a16w16.compile_hints["llvm_options"] = { + # Backend cl::opt hints; only options present in THIS build are set. + _llvm_opts = { "amdgpu-expert-scheduling-mode": True, # valid bool cl::opt (GFX12+ only) - # Not present in this LLVM; re-enable when built against a downstream LLVM that defines them: - # "amdgpu-anti-hints-for-va-vdst", "amdgpu-enable-static-simulator", - # "amdgpu-static-sim-inline", "amdgpu-block-carried-latency", - # "amdgpu-sched-strategy" ("coexec" is also unrecognized upstream) } + # Experimental scheduling levers (in-process LLVM); None => cl::init default. + if sched_strategy is not None: + # cl::opt: max-ilp (GCNMaxILP) or max-memory-clause. + _llvm_opts["amdgpu-sched-strategy"] = sched_strategy + if barrier_signal_wait_latency is not None: + # cl::opt init(35): synthetic barrier signal->wait latency. + _llvm_opts["amdgpu-barrier-signal-wait-latency"] = barrier_signal_wait_latency + launch_gemm_a16w16.compile_hints["llvm_options"] = _llvm_opts return launch_gemm_a16w16 @@ -992,6 +1108,10 @@ def gemm_a16w16( kernarg_preload: bool = False, loop_carried_load_percent: Optional[int] = None, use_manual_barrier: bool = False, + split_k: int = 1, + sched_strategy: Optional[str] = None, + barrier_signal_wait_latency: Optional[int] = None, + main_loop_unroll: bool = False, ): """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" assert x.dtype in ( @@ -1039,31 +1159,42 @@ def gemm_a16w16( # N-pad up to a tile_n multiple N_stride = ((N + tile_n - 1) // tile_n) * tile_n - # Output allocation - if y is not None: + # split_k>1 half-out: accumulate in fp32 workspace, cast at end (bf16 atomic lossy). + _half_dtype = dtype in (torch.float16, torch.bfloat16) + _splitk_f32_accum = split_k > 1 and _half_dtype + + in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" + if dtype == torch.float16: + out_dtype_str = "f16" + elif dtype == torch.bfloat16: + out_dtype_str = "bf16" + else: + out_dtype_str = "f32" + buf_dtype = dtype + if _splitk_f32_accum: + out_dtype_str = "f32" # kernel accumulates fp32 partials + buf_dtype = torch.float32 # fp32 workspace, cast to `dtype` after launch + + # split_k>1 atomic-fadd needs a zeroed buffer; fp32-accum uses a fresh workspace. + _alloc = torch.zeros if split_k > 1 else torch.empty + if y is not None and not _splitk_f32_accum: y_buf = ( y if N_stride == N - else torch.empty((M, N_stride), device=x.device, dtype=dtype) + else _alloc((M, N_stride), device=x.device, dtype=buf_dtype) ) + if split_k > 1 and y_buf is y: + y_buf.zero_() else: y_buf = ( - torch.empty((M, N_stride), device=x.device, dtype=dtype) + _alloc((M, N_stride), device=x.device, dtype=buf_dtype) if N_stride != N - else torch.empty((M, N), device=x.device, dtype=dtype) + else _alloc((M, N), device=x.device, dtype=buf_dtype) ) if bias is None: bias = torch.empty(0, device=x.device, dtype=dtype) - in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" - if dtype == torch.float16: - out_dtype_str = "f16" - elif dtype == torch.bfloat16: - out_dtype_str = "bf16" - else: - out_dtype_str = "f32" - launch_fn = compile_gemm_a16w16( M=M if not physical_mk else 0, N=N, @@ -1085,10 +1216,22 @@ def gemm_a16w16( kernarg_preload=kernarg_preload, loop_carried_load_percent=loop_carried_load_percent, use_manual_barrier=use_manual_barrier, + split_k=split_k, + sched_strategy=sched_strategy, + barrier_signal_wait_latency=barrier_signal_wait_latency, + main_loop_unroll=main_loop_unroll, ) launch_fn(y_buf, x, w, bias, M, N_stride, stream=torch.cuda.current_stream()) + if _splitk_f32_accum: + # Cast the fp32 accumulation workspace back to the requested half dtype. + result = (y_buf[:, :N] if N_stride != N else y_buf).to(dtype) + if y is not None: + y.copy_(result) + return y + return result + if N_stride != N: result = y_buf[:, :N] if y is not None: diff --git a/aiter/ops/shuffle.py b/aiter/ops/shuffle.py index a442a16e683..ba81cf38f6d 100644 --- a/aiter/ops/shuffle.py +++ b/aiter/ops/shuffle.py @@ -111,3 +111,41 @@ def shuffle_scale_a16w4( shfl_scale = shfl_scale.permute(0, 1, 4, 6, 3, 5, 2).contiguous() # print("shf_scale shape:", shfl_scale.shape) return shfl_scale.view(*src.shape).contiguous() + + +def shuffle_weight_gfx1250(w: torch.Tensor) -> torch.Tensor: + """ + Preshuffle weights for gfx1250 WMMA. + + For 2D input (N, K): view as (N//16, 16, K//32, 2, 16) -> + permute(0, 2, 3, 1, 4) -> reshape (N//16, K*16). + For 3D input (E, N, K) or (E, K, N): transpose to (E, N, K) first, + then apply the same pattern per-expert. + + The result is reshaped to (N//16, K*16) for TDM-optimal loading. + """ + x_type = w.dtype + if hasattr(torch, "float4_e2m1fn_x2") and x_type == torch.float4_e2m1fn_x2: + w = w.view(torch.uint8) + + if w.ndim == 2: + N, K = w.shape + assert N % 16 == 0, f"N={N} must be divisible by 16" + assert K % 32 == 0, f"K={K} must be divisible by 32" + w = w.view(N // 16, 16, K // 32, 2, 16) + w = w.permute(0, 2, 3, 1, 4).contiguous() + w = w.view(N // 16, K * 16) + elif w.ndim == 3: + E, K, N = w.shape + assert K % 32 == 0, f"K={K} must be divisible by 32" + assert N % 16 == 0, f"N={N} must be divisible by 16" + w = w.transpose(-1, -2) # (E, N, K) + w = w.view(E, N // 16, 16, K // 32, 2, 16) + w = w.permute(0, 1, 3, 4, 2, 5).contiguous() + w = w.view(E, N // 16, K * 16) + w = w.transpose(-1, -2) # (E, K*16, N//16) + else: + raise ValueError(f"Expected 2D or 3D tensor, got {w.ndim}D") + + w = w.view(x_type) + return w diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index 519f5278e13..c0a9fcbc92c 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -74,27 +74,19 @@ def _generate_inputs(M, N, K, dtype, layout="TN", output=False, bias=False): def get_x_vals(): return [ - (1, 1, 1), - (1, 16, 16), - (16, 1, 16), - (16, 16, 1), - # Irregular - (3, 5, 7), - (17, 33, 65), - (63, 127, 255), - (65, 129, 257), # Aligned (64, 64, 64), (128, 128, 128), + (256, 256, 256), # Multi-block (128, 256, 512), (256, 512, 256), # Asymmetric (32, 256, 128), (256, 32, 128), + # Large K (drives the flat main loop + drain) (128, 128, 1024), (1024, 128, 128), - (1536, 512, 768), ] @@ -105,7 +97,6 @@ def get_fewer_x_vals(): (256, 512, 256), (128, 128, 1024), (1024, 128, 128), - (1536, 512, 768), ] @@ -115,7 +106,7 @@ def test_gemm_a16_w16(M, N, K): x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16) + kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, num_buffers=3) torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) @@ -136,7 +127,9 @@ def test_gemm_a16_w16_activation(M, N, K, dtype, output, activation): elif activation in ("silu", "silu_exp2"): torch_out = F.silu(torch_out) - kernel_out = gemm_a16w16(x, w, dtype=dtype, y=y, activation=activation) + kernel_out = gemm_a16w16( + x, w, dtype=dtype, y=y, activation=activation, num_buffers=3 + ) torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) @@ -148,6 +141,113 @@ def test_gemm_a16_w16_layout(M, N, K, layout): x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16, layout=layout) torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16) + kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, num_buffers=3) torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-1) + + +# 32x32 output tile, deep K (tile_k=128); tile_n=32 needs n_warp=2, K multiple of 128. +@pytest.mark.parametrize( + "M, N, K", [(32, 32, 256), (128, 128, 256), (256, 256, 512), (64, 64, 1024)] +) +def test_gemm_a16_w16_tile_32x32x128(M, N, K): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16( + x, + w, + dtype=torch.bfloat16, + num_buffers=3, + tile_m=32, + tile_n=32, + tile_k=128, + m_warp=2, + n_warp=2, + ) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +# Split-K: partition K across split_k grid-z workgroups, atomic-add into a pre-zeroed C. +def _splitk_configs(): + # (M, N, K, tile_m, tile_n, tile_k, m_warp, n_warp, num_buffers, split_k) — all valid. + return [ + (128, 128, 512, 64, 64, 128, 2, 2, 2, 2), # 4 tiles, 2/split + (128, 128, 512, 64, 64, 128, 2, 2, 2, 4), # 4 tiles, 1/split + (64, 64, 768, 64, 64, 128, 2, 2, 2, 3), # 6 tiles, 2/split + (64, 64, 768, 64, 64, 128, 2, 2, 2, 6), # 6 tiles, 1/split + (128, 256, 1024, 64, 64, 128, 2, 2, 2, 4), # 8 tiles, 2/split + (256, 128, 1024, 32, 32, 128, 2, 2, 2, 8), # 8 tiles, 1/split (max split) + (128, 128, 1024, 64, 64, 128, 2, 2, 3, 4), # nb=3, 2/split + ] + + +@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb,sk", _splitk_configs()) +def test_gemm_a16_w16_split_k(M, N, K, tm, tn, tk, mw, nw, nb, sk): + """Primary split-K correctness: fp32 output (lossless) -> split_k=N == split_k=1.""" + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + common = dict( + dtype=torch.float32, + tile_m=tm, + tile_n=tn, + tile_k=tk, + m_warp=mw, + n_warp=nw, + num_buffers=nb, + ) + out_single = gemm_a16w16(x, w, split_k=1, **common).clone() + out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() + + torch.testing.assert_close(out_split, out_single, atol=1e-2, rtol=1e-3) + + +@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb,sk", _splitk_configs()) +def test_gemm_a16_w16_split_k_bf16(M, N, K, tm, tn, tk, mw, nw, nb, sk): + """bf16-output split-K at O(1) magnitudes (scaled inputs) must match split_k=1.""" + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + x = (x.float() * 0.1).to(torch.bfloat16) + w = (w.float() * 0.1).to(torch.bfloat16) + + common = dict( + dtype=torch.bfloat16, + tile_m=tm, + tile_n=tn, + tile_k=tk, + m_warp=mw, + n_warp=nw, + num_buffers=nb, + ) + out_single = gemm_a16w16(x, w, split_k=1, **common).clone() + out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() + + torch.testing.assert_close(out_split, out_single, atol=1e-1, rtol=2e-2) + + +@pytest.mark.parametrize("sk", [2, 4]) +def test_gemm_a16_w16_split_k_output_buffer_zeroed(sk): + """split_k>1 atomic-ADDS into C; the wrapper must zero it each call (no doubling).""" + M, N, K = 128, 128, 512 + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + y = torch.empty((M, N), dtype=torch.float32, device="cuda") + + common = dict( + dtype=torch.float32, + y=y, + tile_m=64, + tile_n=64, + tile_k=128, + m_warp=2, + n_warp=2, + num_buffers=2, + split_k=sk, + ) + out1 = gemm_a16w16(x, w, **common).clone() + out2 = gemm_a16w16(x, w, **common).clone() # reuse y -> must re-zero, not 2x + + torch.testing.assert_close(out2, out1, atol=1e-2, rtol=1e-3) From c61cee85e497aff7c1dc7411f3450e9d4ff9c698 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Mon, 20 Jul 2026 10:52:22 -0400 Subject: [PATCH 06/34] Update shuffle.py --- aiter/ops/shuffle.py | 38 -------------------------------------- 1 file changed, 38 deletions(-) diff --git a/aiter/ops/shuffle.py b/aiter/ops/shuffle.py index ba81cf38f6d..a442a16e683 100644 --- a/aiter/ops/shuffle.py +++ b/aiter/ops/shuffle.py @@ -111,41 +111,3 @@ def shuffle_scale_a16w4( shfl_scale = shfl_scale.permute(0, 1, 4, 6, 3, 5, 2).contiguous() # print("shf_scale shape:", shfl_scale.shape) return shfl_scale.view(*src.shape).contiguous() - - -def shuffle_weight_gfx1250(w: torch.Tensor) -> torch.Tensor: - """ - Preshuffle weights for gfx1250 WMMA. - - For 2D input (N, K): view as (N//16, 16, K//32, 2, 16) -> - permute(0, 2, 3, 1, 4) -> reshape (N//16, K*16). - For 3D input (E, N, K) or (E, K, N): transpose to (E, N, K) first, - then apply the same pattern per-expert. - - The result is reshaped to (N//16, K*16) for TDM-optimal loading. - """ - x_type = w.dtype - if hasattr(torch, "float4_e2m1fn_x2") and x_type == torch.float4_e2m1fn_x2: - w = w.view(torch.uint8) - - if w.ndim == 2: - N, K = w.shape - assert N % 16 == 0, f"N={N} must be divisible by 16" - assert K % 32 == 0, f"K={K} must be divisible by 32" - w = w.view(N // 16, 16, K // 32, 2, 16) - w = w.permute(0, 2, 3, 1, 4).contiguous() - w = w.view(N // 16, K * 16) - elif w.ndim == 3: - E, K, N = w.shape - assert K % 32 == 0, f"K={K} must be divisible by 32" - assert N % 16 == 0, f"N={N} must be divisible by 16" - w = w.transpose(-1, -2) # (E, N, K) - w = w.view(E, N // 16, 16, K // 32, 2, 16) - w = w.permute(0, 1, 3, 4, 2, 5).contiguous() - w = w.view(E, N // 16, K * 16) - w = w.transpose(-1, -2) # (E, K*16, N//16) - else: - raise ValueError(f"Expected 2D or 3D tensor, got {w.ndim}D") - - w = w.view(x_type) - return w From 3ae9e68226532c5444cecbe59adc6b89255c6fd9 Mon Sep 17 00:00:00 2001 From: Satya Nikhil Date: Mon, 20 Jul 2026 17:43:48 +0000 Subject: [PATCH 07/34] Added compute bound variant + black/ruff --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 110 +++++++++++++++++- op_tests/flydsl_tests/test_gemm_a16w16.py | 14 +++ 2 files changed, 122 insertions(+), 2 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index 8caf045a3ee..b79ff44f816 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -33,6 +33,8 @@ _STAGE_NAMES = ("ping", "pong", "pang", "pung") +_SCHED_ALLOW_SALU = 1 << 2 + def _disable_unroll_on_enclosing_loop(): """Tag the enclosing scf.for with llvm.loop metadata so it stays rolled at ASM level.""" @@ -116,12 +118,25 @@ def compile_gemm_a16w16( sched_strategy: Optional[str] = None, barrier_signal_wait_latency: Optional[int] = None, main_loop_unroll: bool = False, + variant: str = "bandwidth_bound", ): - """Compile the A16W16 GEMM kernel; returns launch_fn(y, x, w, bias, M, N, stream=stream).""" + """Compile the A16W16 GEMM kernel; returns launch_fn(y, x, w, bias, M, N, stream=stream). + + variant: + "bandwidth_bound" (default) — steady-state loads the whole next fragment + bank up front, then runs all WMMAs (burst load, ~2 banks live). + "compute_bound" — in-place fragment rotation: each fragment's next-tile + ds_load is hoisted to right after its last WMMA use so the load + co-executes with the remaining WMMAs (~1 bank live + trickle). + """ _ = (M, N) if not (2 <= num_buffers <= 8): raise ValueError(f"num_buffers must be between 2 and 8, got {num_buffers}") + if variant not in ("bandwidth_bound", "compute_bound"): + raise ValueError( + "variant must be 'bandwidth_bound' or 'compute_bound', " f"got {variant!r}" + ) if in_dtype not in ("fp16", "bf16"): raise ValueError(f"in_dtype must be 'fp16' or 'bf16', got {in_dtype!r}") # Experimental LLVM scheduling levers (injected via llvm_options below). @@ -601,6 +616,53 @@ def wmma_tile(accs_in, a_frags, b_frags): ).result return current_accs + def _load_one_a_frag(buf_i32, ks, wm): + """ds_load a single A fragment (ks, wm) from LDS ring slot buf_i32 (i32 SSA).""" + slot_off_a = arith.index_cast( + T.index, arith.muli(buf_i32, slot_stride_a_elems_i32) + ) + return _load_a_frag(big_a_mem, a_lane_bases[wm] + slot_off_a, ks) + + def _load_one_b_frag(buf_i32, ks, wn): + """ds_load a single B fragment (ks, wn) from LDS ring slot buf_i32 (i32 SSA).""" + slot_off_b = arith.index_cast( + T.index, arith.muli(buf_i32, slot_stride_b_elems_i32) + ) + return _load_b_frag(big_b_mem, b_lane_bases[wn] + slot_off_b, ks) + + def _wmma_rotate(accs_in, a_frags, b_frags, next_buf_i32): + """compute_bound steady-state: run this tile's WMMAs, and the instant a + fragment reaches its last use, issue its next-tile ds_load so the load + co-executes with the remaining WMMAs. Returns (accs, next_a, next_b) where + next_* are the freshly rotated-in fragments for the next K-tile.""" + current_accs = list(accs_in) + next_a = [None] * N_A_FRAGS + next_b = [None] * N_B_FRAGS + for ks in range_constexpr(k_wmma_steps): + for wm in range_constexpr(wmma_m_rep): + a_f = a_frags[ks * wmma_m_rep + wm] + for wn in range_constexpr(wmma_n_rep): + idx = wm * wmma_n_rep + wn + current_accs[idx] = wmma_op( + T.vec(8, T.f32), + b_frags[ks * wmma_n_rep + wn], + a_f, + current_accs[idx], + modC=0, + reuseA=False, + reuseB=(wn > 0), + ).result + rocdl.sched_barrier(_SCHED_ALLOW_SALU) + next_a[ks * wmma_m_rep + wm] = _load_one_a_frag( + next_buf_i32, ks, wm + ) + rocdl.sched_barrier(_SCHED_ALLOW_SALU) + for wn in range_constexpr(wmma_n_rep): + next_b[ks * wmma_n_rep + wn] = _load_one_b_frag( + next_buf_i32, ks, wn + ) + return current_accs, next_a, next_b + _half_out = out_dtype in ("f16", "bf16") _out_elem = ( T.f16 if out_dtype == "f16" else (T.bf16 if out_dtype == "bf16" else None) @@ -823,7 +885,48 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): nbank = None return accs_, nbank, lo_a_, lo_w_ - if const_expr(main_loop_iters > 0): + if const_expr(variant == "compute_bound"): + if const_expr(main_loop_iters > 0): + init_state = _pack_state(accs, cur_a, cur_b) + [ + lo_a, + lo_w, + load_idx_init, + compute_idx_init, + ] + results = init_state + for step, state in range(0, main_loop_iters, 1, init=init_state): + _disable_unroll_on_enclosing_loop() + cidx = state[-1] + lidx = state[-2] + lw = state[-3] + lx = state[-4] + p_accs, ca, cb = _unpack_state(state[:-4]) + + load_buf_i32 = arith.remui(lidx, nb_const_i32) + lx, lw = issue_tdm_loads(load_buf_i32, lx, lw) + if const_expr(num_buffers == 2): + tdm_ops.tensor_wait(0) + else: + tdm_ops.tensor_wait((num_buffers - 2) * _TDMS_PER_TILE) + _wg_barrier() + + next_buf_i32 = arith.remui(arith.addi(cidx, one_i32), nb_const_i32) + p_accs, next_a, next_b = _wmma_rotate(p_accs, ca, cb, next_buf_i32) + cidx1 = arith.addi(cidx, one_i32) + lidx1 = arith.addi(lidx, one_i32) + + new_state = _pack_state(p_accs, next_a, next_b) + [ + lx, + lw, + lidx1, + cidx1, + ] + results = yield new_state + + accs, cur_a, cur_b = _unpack_state(results[:-4]) + else: + accs = list(accs) + elif const_expr(main_loop_iters > 0): if const_expr(main_loop_unroll): # Unroll-by-2: 2 tiles/trip, ping-pong 2 register banks + odd leftover. lo_a_s = lo_a @@ -992,6 +1095,7 @@ def _drain_wait_for(tile_idx): sched_strategy, barrier_signal_wait_latency, main_loop_unroll, + variant, ) @flyc.jit @@ -1112,6 +1216,7 @@ def gemm_a16w16( sched_strategy: Optional[str] = None, barrier_signal_wait_latency: Optional[int] = None, main_loop_unroll: bool = False, + variant: str = "bandwidth_bound", ): """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" assert x.dtype in ( @@ -1220,6 +1325,7 @@ def gemm_a16w16( sched_strategy=sched_strategy, barrier_signal_wait_latency=barrier_signal_wait_latency, main_loop_unroll=main_loop_unroll, + variant=variant, ) launch_fn(y_buf, x, w, bias, M, N_stride, stream=torch.cuda.current_stream()) diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index c0a9fcbc92c..acf88520149 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -111,6 +111,20 @@ def test_gemm_a16_w16(M, N, K): torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) +@pytest.mark.parametrize("M, N, K", get_x_vals()) +@pytest.mark.parametrize("num_buffers", [2, 3]) +def test_gemm_a16_w16_compute_bound(M, N, K, num_buffers): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16( + x, w, dtype=torch.bfloat16, num_buffers=num_buffers, variant="compute_bound" + ) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + @pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu", "silu_exp2"]) @pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) From 56f7c1741e97e6a219bd20dc0d0e8acb6598c2ff Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 24 Jul 2026 17:41:18 +0000 Subject: [PATCH 08/34] black/ruff --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 34 +++++------ op_tests/flydsl_tests/test_gemm_a16w16.py | 58 +++++++++---------- .../gemm/basic/test_gemm_a16w16.py | 12 ++-- 3 files changed, 52 insertions(+), 52 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index b79ff44f816..265148e7f4b 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -import torch + import flydsl.compiler as flyc import flydsl.expr as fx +import torch from flydsl._mlir import ir from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import ( @@ -11,6 +12,7 @@ buffer_ops, const_expr, gpu, + idx2crd, range_constexpr, rocdl, tdm_ops, @@ -20,8 +22,6 @@ from flydsl.expr.typing import T from flydsl.runtime.device import get_rocm_arch as get_hip_arch from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr -from flydsl.expr import idx2crd -from typing import Optional # WMMA 16×16×32 WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 @@ -103,20 +103,20 @@ def compile_gemm_a16w16( m_warp: int = 2, n_warp: int = 4, in_dtype: str = "fp16", - out_dtype: str = None, + out_dtype: str | None = None, num_buffers: int = 2, - waves_per_eu: int = None, + waves_per_eu: int | None = None, l2_prefetch_distance: int = 2, - activation: Optional[str] = None, + activation: str | None = None, add_bias: bool = False, physical_mk: bool = True, # True=M-major (row-major X), False=K-major (col-major X) physical_kn: bool = False, # False=N-major (row-major W), True=K-major (transposed W) - loop_carried_load_percent: Optional[int] = None, + loop_carried_load_percent: int | None = None, kernarg_preload: bool = False, use_manual_barrier: bool = False, split_k: int = 1, - sched_strategy: Optional[str] = None, - barrier_signal_wait_latency: Optional[int] = None, + sched_strategy: str | None = None, + barrier_signal_wait_latency: int | None = None, main_loop_unroll: bool = False, variant: str = "bandwidth_bound", ): @@ -1041,7 +1041,7 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): if main_loop_iters > 0 else min(num_buffers - 1, num_k_tiles) ) - drain_base = main_loop_iters if main_loop_iters > 0 else 0 + drain_base = max(0, main_loop_iters) def _drain_wait_for(tile_idx): # Wait for compile-time drain tile tile_idx to land. @@ -1197,24 +1197,24 @@ def launch_gemm_a16w16( def gemm_a16w16( x: torch.Tensor, w: torch.Tensor, - bias: Optional[torch.Tensor] = None, + bias: torch.Tensor | None = None, dtype: torch.dtype = torch.float16, - y: Optional[torch.Tensor] = None, - activation: Optional[str] = None, + y: torch.Tensor | None = None, + activation: str | None = None, tile_m: int = 128, tile_n: int = 128, tile_k: int = 32, m_warp: int = 2, n_warp: int = 4, num_buffers: int = 2, - waves_per_eu: int = None, + waves_per_eu: int | None = None, l2_prefetch_distance: int = 2, kernarg_preload: bool = False, - loop_carried_load_percent: Optional[int] = None, + loop_carried_load_percent: int | None = None, use_manual_barrier: bool = False, split_k: int = 1, - sched_strategy: Optional[str] = None, - barrier_signal_wait_latency: Optional[int] = None, + sched_strategy: str | None = None, + barrier_signal_wait_latency: int | None = None, main_loop_unroll: bool = False, variant: str = "bandwidth_bound", ): diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index acf88520149..a67409acbdc 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -204,15 +204,15 @@ def test_gemm_a16_w16_split_k(M, N, K, tm, tn, tk, mw, nw, nb, sk): torch.cuda.empty_cache() x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - common = dict( - dtype=torch.float32, - tile_m=tm, - tile_n=tn, - tile_k=tk, - m_warp=mw, - n_warp=nw, - num_buffers=nb, - ) + common = { + "dtype": torch.float32, + "tile_m": tm, + "tile_n": tn, + "tile_k": tk, + "m_warp": mw, + "n_warp": nw, + "num_buffers": nb, + } out_single = gemm_a16w16(x, w, split_k=1, **common).clone() out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() @@ -227,15 +227,15 @@ def test_gemm_a16_w16_split_k_bf16(M, N, K, tm, tn, tk, mw, nw, nb, sk): x = (x.float() * 0.1).to(torch.bfloat16) w = (w.float() * 0.1).to(torch.bfloat16) - common = dict( - dtype=torch.bfloat16, - tile_m=tm, - tile_n=tn, - tile_k=tk, - m_warp=mw, - n_warp=nw, - num_buffers=nb, - ) + common = { + "dtype": torch.bfloat16, + "tile_m": tm, + "tile_n": tn, + "tile_k": tk, + "m_warp": mw, + "n_warp": nw, + "num_buffers": nb, + } out_single = gemm_a16w16(x, w, split_k=1, **common).clone() out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() @@ -250,17 +250,17 @@ def test_gemm_a16_w16_split_k_output_buffer_zeroed(sk): x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) y = torch.empty((M, N), dtype=torch.float32, device="cuda") - common = dict( - dtype=torch.float32, - y=y, - tile_m=64, - tile_n=64, - tile_k=128, - m_warp=2, - n_warp=2, - num_buffers=2, - split_k=sk, - ) + common = { + "dtype": torch.float32, + "y": y, + "tile_m": 64, + "tile_n": 64, + "tile_k": 128, + "m_warp": 2, + "n_warp": 2, + "num_buffers": 2, + "split_k": sk, + } out1 = gemm_a16w16(x, w, **common).clone() out2 = gemm_a16w16(x, w, **common).clone() # reuse y -> must re-zero, not 2x diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py index 44ddad08200..12c67f0312c 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py +++ b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import pytest import torch import torch.nn.functional as F -import pytest -from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16, _is_gluon_available + +from aiter.ops.triton.gemm.basic.gemm_a16w16 import _is_gluon_available, gemm_a16w16 from aiter.ops.triton.gemm.basic.gemm_a16w16_atomic import gemm_a16w16_atomic from op_tests.triton_tests.utils.types import str_to_torch_dtype @@ -66,7 +67,6 @@ def get_x_vals(): (17, 33, 65), (63, 127, 255), (65, 129, 257), - # (64, 64, 64), (128, 128, 128), # Multiple blocks @@ -93,7 +93,7 @@ def test_gemm_a16_w16(M: int, N: int, K: int, backend, kernel_type): pytest.skip("Gluon not supported on this architecture") _skip_if_triton_on_gfx1250(backend) - x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( + x, w, _, _out_dtype, _y = generate_gemm_a16w16_inputs( M, N, K, @@ -200,7 +200,7 @@ def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout, backend, kernel_typ def test_gemm_a16_w16_atomic(M: int, N: int, K: int, output): torch.cuda.empty_cache() # Helps avoid hangs in large tests - x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( + x, w, _, _out_dtype, y = generate_gemm_a16w16_inputs( M, N, K, torch.bfloat16, output=output ) @@ -221,7 +221,7 @@ def test_gemm_a16_w16_atomic(M: int, N: int, K: int, output): def test_gemm_a16_w16_atomic_layout(M: int, N: int, K: int, layout): torch.cuda.empty_cache() # Helps avoid hangs in large tests - x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( + x, w, _, _out_dtype, y = generate_gemm_a16w16_inputs( M, N, K, torch.bfloat16, layout=layout, output=True ) From 3244365b8dd01e860b3d2704ac74fa54debc436e Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 24 Jul 2026 17:44:57 +0000 Subject: [PATCH 09/34] black/ruff --- aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index 265148e7f4b..71f108aebec 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -1130,13 +1130,16 @@ def launch_gemm_a16w16( # Set waves_per_eu for op in ctx.gpu_module_body.operations: - if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": - if effective_waves_per_eu is not None: - _wpe = int(effective_waves_per_eu) - if _wpe >= 1: - op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( - ir.IntegerType.get_signless(32), _wpe - ) + if ( + hasattr(op, "attributes") + and op.OPERATION_NAME == "gpu.func" + and effective_waves_per_eu is not None + ): + _wpe = int(effective_waves_per_eu) + if _wpe >= 1: + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + ir.IntegerType.get_signless(32), _wpe + ) # Flat work-group size hint. flat_wg_attr = ir.StringAttr.get(f"{block_threads},{block_threads}") From 3dfa6ee085cb6d7302dee9a1d997734046e43328 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Mon, 27 Jul 2026 23:03:01 +0000 Subject: [PATCH 10/34] new api --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 571 +++++++++--------- 1 file changed, 287 insertions(+), 284 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index 71f108aebec..1f7d91394bc 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -6,22 +6,22 @@ import flydsl.expr as fx import torch from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm as llvm_dialect +from flydsl._mlir.dialects import rocdl as raw_rocdl from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import ( arith, - buffer_ops, const_expr, gpu, idx2crd, range_constexpr, rocdl, tdm_ops, - vector, ) from flydsl.expr.arith import _to_raw as _raw from flydsl.expr.typing import T from flydsl.runtime.device import get_rocm_arch as get_hip_arch -from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from flydsl.utils.smem_allocator import check_smem_capacity # WMMA 16×16×32 WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 @@ -36,6 +36,53 @@ _SCHED_ALLOW_SALU = 1 << 2 +def _align_up(n, a): + return ((n + a - 1) // a) * a + + +def _i32_const(v): + return _raw(arith.constant(int(v), type=T.i32)) + + +def _to_i32_offset(offset): + if isinstance(offset, int): + return _i32_const(offset) + offset = arith.unwrap(offset) + if not isinstance(offset.type, ir.IntegerType) or offset.type.width != 32: + return _raw(arith.index_cast(T.i32, offset)) + return offset + + +def _buffer_load(rsrc, offset, vec_width=1, dtype=None): + if dtype is None: + dtype = T.f32 + off = _to_i32_offset(offset) + off = _raw(arith.muli(off, _i32_const(dtype.width // 8))) + result_type = dtype if vec_width == 1 else T.vec(vec_width, dtype) + return rocdl.RawPtrBufferLoadOp( + result_type, rsrc, off, _i32_const(0), _i32_const(0) + ).result + + +def _buffer_store(data, rsrc, offset, offset_is_bytes=False): + data = arith.unwrap(data) + off = _to_i32_offset(offset) + if not offset_is_bytes: + data_type = data.type + element_type = ( + data_type.element_type if hasattr(data_type, "element_type") else data_type + ) + off = _raw(arith.muli(off, _i32_const(element_type.width // 8))) + rocdl.RawPtrBufferStoreOp(data, rsrc, off, _i32_const(0), _i32_const(0)) + + +def _make_buffer_rsrc(tensor, num_records_bytes=None): + buf_tensor = fx.rocdl.make_buffer_tensor( + tensor, max_size=True, num_records_bytes=num_records_bytes + ) + return fx.rocdl.get_buffer_rsrc(fx.get_iter(buf_tensor)) + + def _disable_unroll_on_enclosing_loop(): """Tag the enclosing scf.for with llvm.loop metadata so it stays rolled at ASM level.""" block = ir.InsertionPoint.current.block @@ -225,13 +272,13 @@ def compile_gemm_a16w16( gpu_arch = str(get_hip_arch()) assert gpu_arch.startswith("gfx1250"), f"Expected gfx1250, got {gpu_arch}" - wmma_op = rocdl.wmma_f32_16x16x32_f16 if is_f16 else rocdl.wmma_f32_16x16x32_bf16 - k_wmma_steps = tile_k // WMMA_K def _elem_type(): return T.f16 if is_f16 else T.bf16 + _fx_elem = fx.Float16 if is_f16 else fx.BFloat16 + wmma_m_rep = warp_tile_m // WMMA_M wmma_n_rep = warp_tile_n // WMMA_N n_accs = wmma_m_rep * wmma_n_rep @@ -257,19 +304,32 @@ def _elem_type(): lds_b_data_bytes = lds_b_elems * elem_bytes # Unified LDS allocator: contiguous [A0..A_nb-1 | B0..B_nb-1] ring slots - unified_alloc = SmemAllocator(None, arch=gpu_arch, global_sym_name="a16w16_unified") - unified_a_off = unified_alloc._align(unified_alloc.ptr, 16) - unified_alloc.ptr = unified_a_off + num_buffers * lds_a_data_bytes - unified_b_off = unified_alloc._align(unified_alloc.ptr, 16) - unified_alloc.ptr = unified_b_off + num_buffers * lds_b_data_bytes - - stage_a_data_off = [ - unified_a_off + i * lds_a_data_bytes for i in range(num_buffers) - ] - stage_b_data_off = [ - unified_b_off + i * lds_b_data_bytes for i in range(num_buffers) - ] - stage_allocators = [unified_alloc] * num_buffers + unified_a_off = 0 + unified_b_off = _align_up(unified_a_off + num_buffers * lds_a_data_bytes, 16) + lds_arena_bytes = _align_up(unified_b_off + num_buffers * lds_b_data_bytes, 16) + check_smem_capacity(lds_arena_bytes, gpu_arch) + + if physical_mk: + a_tile_shape = (tile_m, tile_k) + a_outer_stride = K + a_pad_interval = tile_k + a_imm_bytes = tile_k * elem_bytes + else: + a_tile_shape = (tile_k, tile_m) + a_outer_stride = M + a_pad_interval = tile_m + a_imm_bytes = tile_k * M * elem_bytes + + if physical_kn: + b_tile_shape = (tile_k, tile_n) + b_outer_stride = N + b_pad_interval = tile_n + b_imm_bytes = tile_k * N * elem_bytes + else: + b_tile_shape = (tile_n, tile_k) + b_outer_stride = K + b_pad_interval = tile_k + b_imm_bytes = tile_k * elem_bytes # TDMs per K-tile: A + B _TDMS_PER_TILE = 2 @@ -294,6 +354,8 @@ def kernel_gemm_a16w16( if const_expr(split_k > 1): bz = gpu.block_id("z") split_k_base = bz * arith.index(split_k_chunk) + else: + split_k_base = arith.index(0) # Thread -> warp decomposition layout_thr = fx.make_layout( @@ -316,164 +378,105 @@ def kernel_gemm_a16w16( # Buffer resource descriptors y_nrec = m_idx * n_stride * arith.index(elem_bytes_d) - y_rsrc = buffer_ops.create_buffer_resource(arg_y, num_records_bytes=y_nrec) + y_rsrc = _make_buffer_rsrc(arg_y, num_records_bytes=y_nrec) if const_expr(add_bias): - bias_rsrc = buffer_ops.create_buffer_resource(arg_bias, max_size=True) + bias_rsrc = _make_buffer_rsrc(arg_bias) - def make_a_desc(k_base, lds_a_mem_ref): - """TDM descriptor for A tile. Swaps dims when physical_mk=False (K-major).""" - if const_expr(physical_mk): - return tdm_ops.make_tensor_descriptor_2d( - global_ptr=arg_x, - lds_memref=lds_a_mem_ref, - global_offset=(blk_m, k_base), - tensor_shape=(tile_m, tile_k), - strides=(K, 1), - tile_shape=(tile_m, tile_k), - elem_bytes=elem_bytes, - pad_interval=tile_k, - pad_amount=LDS_PAD_A, - num_warps=num_warps, - ) - else: - return tdm_ops.make_tensor_descriptor_2d( - global_ptr=arg_x, - lds_memref=lds_a_mem_ref, - global_offset=(k_base, blk_m), - tensor_shape=(tile_k, tile_m), - strides=(M, 1), - tile_shape=(tile_k, tile_m), - elem_bytes=elem_bytes, - pad_interval=tile_m, - pad_amount=LDS_PAD_A, - num_warps=num_warps, - ) + lds_base_ptr = fx.SharedAllocator(static=True).allocate(lds_arena_bytes)._ptr + big_a_mem = fx.recast_iter(_fx_elem, fx.add_offset(lds_base_ptr, unified_a_off)) + big_b_mem = fx.recast_iter(_fx_elem, fx.add_offset(lds_base_ptr, unified_b_off)) + big_a_base_idx = arith.index_cast(T.index, arith.unwrap(fx.ptrtoint(big_a_mem))) + big_b_base_idx = arith.index_cast(T.index, arith.unwrap(fx.ptrtoint(big_b_mem))) + slot_stride_a_elems_i32 = arith.constant(lds_a_elems, type=T.i32) + slot_stride_b_elems_i32 = arith.constant(lds_b_elems, type=T.i32) - def make_b_desc(k_base, lds_b_mem_ref): - """TDM descriptor for B tile. Swaps dims when physical_kn=True (K-major).""" - if const_expr(physical_kn): - return tdm_ops.make_tensor_descriptor_2d( - global_ptr=arg_w, - lds_memref=lds_b_mem_ref, - global_offset=(k_base, blk_n), - tensor_shape=(tile_k, tile_n), - strides=(N, 1), - tile_shape=(tile_k, tile_n), - elem_bytes=elem_bytes, - pad_interval=tile_n, - pad_amount=LDS_PAD_B, - num_warps=num_warps, - ) - return tdm_ops.make_tensor_descriptor_2d( - global_ptr=arg_w, - lds_memref=lds_b_mem_ref, - global_offset=(blk_n, k_base), - tensor_shape=(tile_n, tile_k), - strides=(K, 1), - tile_shape=(tile_n, tile_k), - elem_bytes=elem_bytes, - pad_interval=tile_k, - pad_amount=LDS_PAD_B, - num_warps=num_warps, - ) + blk_m64 = fx.Int64(blk_m) + blk_n64 = fx.Int64(blk_n) + split_k_base64 = fx.Int64(split_k_base) - # Unified LDS pointers (all NB slots) + per-stage aliases - big_a = SmemPtr( - unified_alloc.get_base(), - unified_a_off, - elem_ty, - shape=(num_buffers * lds_a_elems,), - ) - big_b = SmemPtr( - unified_alloc.get_base(), - unified_b_off, - elem_ty, - shape=(num_buffers * lds_b_elems,), - ) - big_a_mem = big_a.get() - big_b_mem = big_b.get() - stages_a = [ - SmemPtr( - unified_alloc.get_base(), - stage_a_data_off[i], - elem_ty, - shape=(lds_a_elems,), - ) - for i in range(num_buffers) - ] - stages_b = [ - SmemPtr( - unified_alloc.get_base(), - stage_b_data_off[i], - elem_ty, - shape=(lds_b_elems,), + def _global_view(tensor, elem_off, shape, stride): + return fx.Tensor( + fx.make_view( + fx.add_offset(fx.get_iter(tensor), elem_off), + fx.make_layout(shape, stride), + ) ) - for i in range(num_buffers) - ] - stages_a_mem = [p.get() for p in stages_a] - stages_b_mem = [p.get() for p in stages_b] - # TDM descriptors built once at entry; lo32 advances per K-tile, LDS base per ring slot. - if const_expr(split_k > 1): - _desc_a_init = make_a_desc(split_k_base, stages_a_mem[0]) - _desc_b_init = make_b_desc(split_k_base, stages_b_mem[0]) + if const_expr(physical_mk): + a_base_off = blk_m64 * fx.Int64(K) + split_k_base64 else: - _desc_a_init = make_a_desc(arith.index(0), stages_a_mem[0]) - _desc_b_init = make_b_desc(arith.index(0), stages_b_mem[0]) - dgroup1_a = _desc_a_init.dgroup1 - dgroup1_w = _desc_b_init.dgroup1 - addr_hi_a = vector.extract( - _desc_a_init.dgroup0, static_position=[3], dynamic_position=[] + a_base_off = split_k_base64 * fx.Int64(M) + blk_m64 + gA = _global_view(arg_x, a_base_off, a_tile_shape, (a_outer_stride, 1)) + atom_a = fx.rocdl.make_tdm_atom( + gA, + [None, None], + strides=[fx.Int64(a_outer_stride), None], + num_warps=num_warps, + pad_interval=a_pad_interval, + pad_amount=LDS_PAD_A, ) - addr_hi_w = vector.extract( - _desc_b_init.dgroup0, static_position=[3], dynamic_position=[] - ) - addr_lo_init_a = vector.extract( - _desc_a_init.dgroup0, static_position=[2], dynamic_position=[] - ) - addr_lo_init_w = vector.extract( - _desc_b_init.dgroup0, static_position=[2], dynamic_position=[] - ) - a_lds_base_i32 = vector.extract( - _desc_a_init.dgroup0, static_position=[1], dynamic_position=[] - ) - b_lds_base_i32 = vector.extract( - _desc_b_init.dgroup0, static_position=[1], dynamic_position=[] - ) - slot_stride_a_i32 = arith.constant(lds_a_data_bytes, type=T.i32) - slot_stride_b_i32 = arith.constant(lds_b_data_bytes, type=T.i32) - slot_stride_a_elems_i32 = arith.constant(lds_a_elems, type=T.i32) - slot_stride_b_elems_i32 = arith.constant(lds_b_elems, type=T.i32) - # Per-K-tile lo32 byte advance (tile_k*bytes row-major; outer-stride*bytes if K-major). - adv_a_i32 = arith.constant( - (tile_k if physical_mk else tile_k * M) * elem_bytes, type=T.i32 - ) - adv_w_i32 = arith.constant( - (tile_k if not physical_kn else tile_k * N) * elem_bytes, type=T.i32 + if const_expr(physical_kn): + b_base_off = split_k_base64 * fx.Int64(N) + blk_n64 + else: + b_base_off = blk_n64 * fx.Int64(K) + split_k_base64 + gB = _global_view(arg_w, b_base_off, b_tile_shape, (b_outer_stride, 1)) + atom_b = fx.rocdl.make_tdm_atom( + gB, + [None, None], + strides=[fx.Int64(b_outer_stride), None], + num_warps=num_warps, + pad_interval=b_pad_interval, + pad_amount=LDS_PAD_B, ) - pred_const = arith.constant(1, type=T.i32) - def _buf_idx_to_i32(buf_idx): - """Accept either Python int (prologue/drain) or i32 SSA (main loop).""" + def _slot_elem_off(buf_idx, slot_elems, slot_stride_i32): if const_expr(isinstance(buf_idx, int)): - return arith.constant(buf_idx, type=T.i32) - return buf_idx - - def issue_tdm_loads(buf_idx, lo_a, lo_w): - """Issue A+B TDMs for one K-tile into LDS slot buf_idx; returns advanced (lo_a, lo_w).""" - buf_i32 = _buf_idx_to_i32(buf_idx) - a_addr = arith.addi(a_lds_base_i32, arith.muli(buf_i32, slot_stride_a_i32)) - b_addr = arith.addi(b_lds_base_i32, arith.muli(buf_i32, slot_stride_b_i32)) - dg0_a = vector.from_elements( - T.vec(4, T.i32), [pred_const, a_addr, lo_a, addr_hi_a] + return arith.index(buf_idx * slot_elems) + return arith.index_cast(T.index, arith.muli(buf_idx, slot_stride_i32)) + + def _imm64(k_tile, mul): + if const_expr(isinstance(k_tile, int)): + return fx.Int64(k_tile * mul) + as_index = arith.index_cast( + T.index, arith.muli(k_tile, arith.constant(mul, type=T.i32)) ) - tdm_ops.tensor_load_2d(tdm_ops.TDMDescriptor2D(dg0_a, dgroup1_a)) - dg0_w = vector.from_elements( - T.vec(4, T.i32), [pred_const, b_addr, lo_w, addr_hi_w] + return fx.Int64(arith.index_cast(T.i64, as_index)) + + def issue_tdm_loads(buf_idx, k_tile): + """Issue A+B TDMs for K-tile k_tile into LDS ring slot buf_idx.""" + fx.copy( + atom_a, + gA, + fx.Tensor( + fx.make_view( + fx.add_offset( + big_a_mem, + _slot_elem_off( + buf_idx, lds_a_elems, slot_stride_a_elems_i32 + ), + ), + fx.make_layout(a_tile_shape, (lds_a_stride, 1)), + ) + ), + imm_offset=_imm64(k_tile, a_imm_bytes), + ) + fx.copy( + atom_b, + gB, + fx.Tensor( + fx.make_view( + fx.add_offset( + big_b_mem, + _slot_elem_off( + buf_idx, lds_b_elems, slot_stride_b_elems_i32 + ), + ), + fx.make_layout(b_tile_shape, (lds_b_stride, 1)), + ) + ), + imm_offset=_imm64(k_tile, b_imm_bytes), ) - tdm_ops.tensor_load_2d(tdm_ops.TDMDescriptor2D(dg0_w, dgroup1_w)) - return arith.addi(lo_a, adv_a_i32), arith.addi(lo_w, adv_w_i32) def _precompute_lane_bases(warp_base, reps, lds_stride): """Per-lane LDS element offsets for `reps` WMMA frags (M/N-major).""" @@ -524,38 +527,61 @@ def _precompute_b_lane_bases(): a_lane_bases = _precompute_a_lane_bases() b_lane_bases = _precompute_b_lane_bases() + def _lds_transpose_load(result_type, lds_base_idx, elem_off): + addr_i32 = arith.index_cast( + T.i32, lds_base_idx + elem_off * arith.index(elem_bytes) + ) + ptr_val = llvm_dialect.inttoptr( + ir.Type.parse("!llvm.ptr<3>"), arith.unwrap(addr_i32) + ) + return raw_rocdl.ds_load_tr16_b128(result_type, ptr_val) + def load_wmma_frag_tr(lds_buffer, b_lane_base, ks): """Load WMMA B fragment via ds_load_tr16_b128 (K-major LDS).""" vec8_ty = ir.VectorType.get([8], elem_ty) + lds_base_idx = big_b_base_idx results = [] for k_half in range_constexpr(2): k_row_off = (ks * WMMA_K + k_half * 16) * lds_b_stride elem_off = b_lane_base + arith.index(k_row_off) - v = rocdl.lds_transpose_load(vec8_ty, lds_buffer, elem_off, elem_bytes) + v = _lds_transpose_load(vec8_ty, lds_base_idx, elem_off) results.append(v) - return vector.shuffle(results[0], results[1], list(range(16))) + return ( + fx.Vector(results[0]) + .shuffle(fx.Vector(results[1]), list(range(16))) + .ir_value() + ) def load_wmma_frag(lds_buffer, lane_base, ks): """Load WMMA fragment via ds_read_b128 (M/N-major LDS).""" vec8_ty = ir.VectorType.get([8], elem_ty) off0 = lane_base + arith.index(ks * WMMA_K) off1 = lane_base + arith.index(ks * WMMA_K + 16) - v0 = vector.load_op(vec8_ty, lds_buffer, [off0]) - v1 = vector.load_op(vec8_ty, lds_buffer, [off1]) - return vector.shuffle(v0, v1, list(range(16))) + v0 = arith.unwrap( + fx.ptr_load(fx.add_offset(lds_buffer, off0), result_type=vec8_ty) + ) + v1 = arith.unwrap( + fx.ptr_load(fx.add_offset(lds_buffer, off1), result_type=vec8_ty) + ) + return fx.Vector(v0).shuffle(fx.Vector(v1), list(range(16))).ir_value() _load_b_frag = load_wmma_frag_tr if physical_kn else load_wmma_frag def load_wmma_frag_tr_a(lds_buffer, a_lane_base, ks): """Load WMMA A fragment via ds_load_tr16_b128 (K-major LDS).""" vec8_ty = ir.VectorType.get([8], elem_ty) + lds_base_idx = big_a_base_idx results = [] for k_half in range_constexpr(2): k_row_off = (ks * WMMA_K + k_half * 16) * lds_a_stride elem_off = a_lane_base + arith.index(k_row_off) - v = rocdl.lds_transpose_load(vec8_ty, lds_buffer, elem_off, elem_bytes) + v = _lds_transpose_load(vec8_ty, lds_base_idx, elem_off) results.append(v) - return vector.shuffle(results[0], results[1], list(range(16))) + return ( + fx.Vector(results[0]) + .shuffle(fx.Vector(results[1]), list(range(16))) + .ir_value() + ) _load_a_frag = load_wmma_frag_tr_a if not physical_mk else load_wmma_frag @@ -591,30 +617,40 @@ def load_b_frags(buf_idx): ) return b_frags + wmma_atom = fx.make_mma_atom( + fx.rocdl.WMMA(WMMA_M, WMMA_N, WMMA_K, _fx_elem, fx.Float32) + ) + + def _rmem_vec(v, n, dtype): + t = fx.make_rmem_tensor(n, dtype) + t.store(fx.Vector(arith.unwrap(v))) + return t + + def _rmem_frag_lists(a_frags, b_frags): + return ( + [_rmem_vec(f, 16, _fx_elem) for f in a_frags], + [_rmem_vec(f, 16, _fx_elem) for f in b_frags], + ) + def wmma_tile(accs_in, a_frags, b_frags): - """Execute all WMMAs for one tile (ks-outer accumulate); reuseB=(wn>0) hint kept.""" - current_accs = list(accs_in) + """Execute all WMMAs for one tile (ks-outer accumulate).""" + a_frag_ts, b_frag_ts = _rmem_frag_lists(a_frags, b_frags) + acc_ts = [_rmem_vec(a, 8, fx.Float32) for a in accs_in] for ks in range_constexpr(k_wmma_steps): a_f = [ - a_frags[ks * wmma_m_rep + wm] for wm in range_constexpr(wmma_m_rep) + a_frag_ts[ks * wmma_m_rep + wm] + for wm in range_constexpr(wmma_m_rep) ] b_f = [ - b_frags[ks * wmma_n_rep + wn] for wn in range_constexpr(wmma_n_rep) + b_frag_ts[ks * wmma_n_rep + wn] + for wn in range_constexpr(wmma_n_rep) ] for wm in range_constexpr(wmma_m_rep): for wn in range_constexpr(wmma_n_rep): idx = wm * wmma_n_rep + wn # ISA operand order: (B, A, C) — reversed from math - current_accs[idx] = wmma_op( - T.vec(8, T.f32), - b_f[wn], - a_f[wm], - current_accs[idx], - modC=0, - reuseA=False, - reuseB=(wn > 0), - ).result - return current_accs + fx.gemm(wmma_atom, acc_ts[idx], b_f[wn], a_f[wm], acc_ts[idx]) + return [arith.unwrap(t.load()) for t in acc_ts] def _load_one_a_frag(buf_i32, ks, wm): """ds_load a single A fragment (ks, wm) from LDS ring slot buf_i32 (i32 SSA).""" @@ -635,23 +671,22 @@ def _wmma_rotate(accs_in, a_frags, b_frags, next_buf_i32): fragment reaches its last use, issue its next-tile ds_load so the load co-executes with the remaining WMMAs. Returns (accs, next_a, next_b) where next_* are the freshly rotated-in fragments for the next K-tile.""" - current_accs = list(accs_in) + a_frag_ts, b_frag_ts = _rmem_frag_lists(a_frags, b_frags) + acc_ts = [_rmem_vec(a, 8, fx.Float32) for a in accs_in] next_a = [None] * N_A_FRAGS next_b = [None] * N_B_FRAGS for ks in range_constexpr(k_wmma_steps): for wm in range_constexpr(wmma_m_rep): - a_f = a_frags[ks * wmma_m_rep + wm] + a_f = a_frag_ts[ks * wmma_m_rep + wm] for wn in range_constexpr(wmma_n_rep): idx = wm * wmma_n_rep + wn - current_accs[idx] = wmma_op( - T.vec(8, T.f32), - b_frags[ks * wmma_n_rep + wn], + fx.gemm( + wmma_atom, + acc_ts[idx], + b_frag_ts[ks * wmma_n_rep + wn], a_f, - current_accs[idx], - modC=0, - reuseA=False, - reuseB=(wn > 0), - ).result + acc_ts[idx], + ) rocdl.sched_barrier(_SCHED_ALLOW_SALU) next_a[ks * wmma_m_rep + wm] = _load_one_a_frag( next_buf_i32, ks, wm @@ -661,12 +696,17 @@ def _wmma_rotate(accs_in, a_frags, b_frags, next_buf_i32): next_b[ks * wmma_n_rep + wn] = _load_one_b_frag( next_buf_i32, ks, wn ) - return current_accs, next_a, next_b + return [arith.unwrap(t.load()) for t in acc_ts], next_a, next_b _half_out = out_dtype in ("f16", "bf16") _out_elem = ( T.f16 if out_dtype == "f16" else (T.bf16 if out_dtype == "bf16" else None) ) + _out_num = ( + fx.Float16 + if out_dtype == "f16" + else (fx.BFloat16 if out_dtype == "bf16" else None) + ) _bias_elem = T.f16 if is_f16 else T.bf16 def _widen_to_f32(val): @@ -694,37 +734,33 @@ def _apply_bias_and_activation(accs): col_base_i32 = arith.index_cast(T.i32, col_base) col_base_hi = col_base_i32 + arith.constant(4, type=T.i32) - bv_lo = buffer_ops.buffer_load( + bv_lo = _buffer_load( bias_rsrc, col_base_i32, vec_width=4, dtype=_bias_elem ) - bv_hi = buffer_ops.buffer_load( + bv_hi = _buffer_load( bias_rsrc, col_base_hi, vec_width=4, dtype=_bias_elem ) bias_elems = [] for i in range_constexpr(4): - b_val = vector.extract( - bv_lo, static_position=[i], dynamic_position=[] - ) + b_val = fx.Vector(bv_lo)[i].ir_value() bias_elems.append(_widen_to_f32(b_val)) for i in range_constexpr(4): - b_val = vector.extract( - bv_hi, static_position=[i], dynamic_position=[] - ) + b_val = fx.Vector(bv_hi)[i].ir_value() bias_elems.append(_widen_to_f32(b_val)) - bias_vec = vector.from_elements(T.vec(8, T.f32), bias_elems) + bias_vec = fx.Vector.from_elements( + bias_elems, fx.Float32 + ).ir_value() acc = acc + bias_vec if const_expr(activation is not None): new_elems = [] for i in range_constexpr(8): - val = vector.extract( - acc, static_position=[i], dynamic_position=[] - ) + val = fx.Vector(acc)[i].ir_value() val = _apply_activation_scalar(val, activation) new_elems.append(val) - acc = vector.from_elements(T.vec(8, T.f32), new_elems) + acc = fx.Vector.from_elements(new_elems, fx.Float32).ir_value() accs[idx] = acc return accs @@ -750,19 +786,12 @@ def epilogue_stores(final_accs): c_off_bytes = (row * n_stride + col_base) * arith.index( elem_bytes_d ) - pair_ty = T.vec(2, _out_elem) for pair in range_constexpr(4): - e0 = vector.extract( - h_vec, - static_position=[pair * 2], - dynamic_position=[], - ) - e1 = vector.extract( - h_vec, - static_position=[pair * 2 + 1], - dynamic_position=[], - ) - pair_vec = vector.from_elements(pair_ty, [e0, e1]) + e0 = fx.Vector(h_vec)[pair * 2].ir_value() + e1 = fx.Vector(h_vec)[pair * 2 + 1].ir_value() + pair_vec = fx.Vector.from_elements( + [e0, e1], _out_num + ).ir_value() byte_off = arith.index_cast( T.i32, c_off_bytes + arith.index(pair * 4) ) @@ -771,11 +800,11 @@ def epilogue_stores(final_accs): ) else: h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) - i32_vec = vector.bitcast(T.vec(4, T.i32), h_vec) + i32_vec = fx.Vector(h_vec).bitcast(fx.Int32).ir_value() c_off_bytes = (row * n_stride + col_base) * arith.index( elem_bytes_d ) - buffer_ops.buffer_store( + _buffer_store( i32_vec, y_rsrc, c_off_bytes, offset_is_bytes=True ) else: @@ -783,11 +812,9 @@ def epilogue_stores(final_accs): for half in range_constexpr(2): base = row * n_stride + col_base + arith.index(half * 4) for vi in range_constexpr(4): - val = vector.extract( - final_accs[idx], - static_position=[half * 4 + vi], - dynamic_position=[], - ) + val = fx.Vector(final_accs[idx])[ + half * 4 + vi + ].ir_value() byte_off = arith.index_cast( T.i32, (base + arith.index(vi)) * arith.index(4), @@ -798,17 +825,15 @@ def epilogue_stores(final_accs): else: for half in range_constexpr(2): vals = [ - vector.extract( - final_accs[idx], - static_position=[half * 4 + vi], - dynamic_position=[], - ) + fx.Vector(final_accs[idx])[half * 4 + vi].ir_value() for vi in range_constexpr(4) ] - vec4 = vector.from_elements(T.vec(4, T.f32), vals) + vec4 = fx.Vector.from_elements( + vals, fx.Float32 + ).ir_value() col = col_base + arith.index(half * 4) c_off = row * n_stride + col - buffer_ops.buffer_store(vec4, y_rsrc, c_off) + _buffer_store(vec4, y_rsrc, c_off) # Accumulators acc_zero = arith.constant_vector(0.0, T.vec(8, T.f32)) @@ -831,10 +856,8 @@ def _unpack_state(state): return accs_, a_, b_ # ── Prologue (memory-bound): fill NB-1 LDS buffers, then load tile 0 ── - lo_a = addr_lo_init_a - lo_w = addr_lo_init_w for i in range_constexpr(num_buffers - 1): - lo_a, lo_w = issue_tdm_loads(i, lo_a, lo_w) + issue_tdm_loads(i, i) # Wait retires tile 0 (NB-2 in flight); the 2-iter WAR gap lets the main loop use one barrier. tdm_ops.tensor_wait((num_buffers - 2) * _TDMS_PER_TILE) @@ -858,12 +881,12 @@ def _wg_barrier(): else: gpu.barrier() - def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): + def _run_tile(accs_, ca, cb, cidx, lidx, do_prefetch): """Compute K-tile cidx from carried frags; if do_prefetch, load cidx+1 into a fresh bank.""" load_buf_i32 = arith.remui(lidx, nb_const_i32) # Refill first, then wait: keeps NB-2 tiles' TDM in flight (depth-2). - lo_a_, lo_w_ = issue_tdm_loads(load_buf_i32, lo_a_, lo_w_) + issue_tdm_loads(load_buf_i32, lidx) if const_expr(num_buffers == 2): # NB=2: only 2 slots -> drain the just-issued TDM (RAW hazard). tdm_ops.tensor_wait(0) @@ -883,13 +906,11 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): else: accs_ = wmma_tile(accs_, ca, cb) nbank = None - return accs_, nbank, lo_a_, lo_w_ + return accs_, nbank if const_expr(variant == "compute_bound"): if const_expr(main_loop_iters > 0): init_state = _pack_state(accs, cur_a, cur_b) + [ - lo_a, - lo_w, load_idx_init, compute_idx_init, ] @@ -898,12 +919,10 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): _disable_unroll_on_enclosing_loop() cidx = state[-1] lidx = state[-2] - lw = state[-3] - lx = state[-4] - p_accs, ca, cb = _unpack_state(state[:-4]) + p_accs, ca, cb = _unpack_state(state[:-2]) load_buf_i32 = arith.remui(lidx, nb_const_i32) - lx, lw = issue_tdm_loads(load_buf_i32, lx, lw) + issue_tdm_loads(load_buf_i32, lidx) if const_expr(num_buffers == 2): tdm_ops.tensor_wait(0) else: @@ -916,29 +935,23 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): lidx1 = arith.addi(lidx, one_i32) new_state = _pack_state(p_accs, next_a, next_b) + [ - lx, - lw, lidx1, cidx1, ] results = yield new_state - accs, cur_a, cur_b = _unpack_state(results[:-4]) + accs, cur_a, cur_b = _unpack_state(results[:-2]) else: accs = list(accs) elif const_expr(main_loop_iters > 0): if const_expr(main_loop_unroll): # Unroll-by-2: 2 tiles/trip, ping-pong 2 register banks + odd leftover. - lo_a_s = lo_a - lo_w_s = lo_w load_idx_s = load_idx_init compute_idx_s = compute_idx_init num_pairs = main_loop_iters // 2 if const_expr(num_pairs > 0): init_state = _pack_state(accs, cur_a, cur_b) + [ - lo_a_s, - lo_w_s, load_idx_s, compute_idx_s, ] @@ -947,37 +960,31 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): _disable_unroll_on_enclosing_loop() cidx = state[-1] lidx = state[-2] - lw = state[-3] - lx = state[-4] - p_accs, b0a, b0b = _unpack_state(state[:-4]) + p_accs, b0a, b0b = _unpack_state(state[:-2]) # sub-0: tile T from bank0, prefetch T+1 -> bank1. - p_accs, bank1, lx, lw = _run_tile( - p_accs, b0a, b0b, lx, lw, cidx, lidx, do_prefetch=True + p_accs, bank1 = _run_tile( + p_accs, b0a, b0b, cidx, lidx, do_prefetch=True ) b1a, b1b = bank1 cidx1 = arith.addi(cidx, one_i32) lidx1 = arith.addi(lidx, one_i32) # sub-1: tile T+1 from bank1, prefetch T+2 -> bank0 (yielded). - p_accs, bank0, lx, lw = _run_tile( - p_accs, b1a, b1b, lx, lw, cidx1, lidx1, do_prefetch=True + p_accs, bank0 = _run_tile( + p_accs, b1a, b1b, cidx1, lidx1, do_prefetch=True ) nb0a, nb0b = bank0 cidx2 = arith.addi(cidx1, one_i32) lidx2 = arith.addi(lidx1, one_i32) new_state = _pack_state(p_accs, nb0a, nb0b) + [ - lx, - lw, lidx2, cidx2, ] results = yield new_state - accs, cur_a, cur_b = _unpack_state(results[:-4]) - lo_a_s = results[-4] - lo_w_s = results[-3] + accs, cur_a, cur_b = _unpack_state(results[:-2]) load_idx_s = results[-2] compute_idx_s = results[-1] else: @@ -985,12 +992,10 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): # Leftover odd tile: compute the carried bank, prefetch first drain tile. if const_expr(main_loop_iters % 2 == 1): - accs, leftover_bank, lo_a_s, lo_w_s = _run_tile( + accs, leftover_bank = _run_tile( accs, cur_a, cur_b, - lo_a_s, - lo_w_s, compute_idx_s, load_idx_s, do_prefetch=True, @@ -1000,8 +1005,6 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): else: # Single tile/trip: compute T from carried bank, prefetch T+1. init_state = _pack_state(accs, cur_a, cur_b) + [ - lo_a, - lo_w, load_idx_init, compute_idx_init, ] @@ -1010,27 +1013,23 @@ def _run_tile(accs_, ca, cb, lo_a_, lo_w_, cidx, lidx, do_prefetch): _disable_unroll_on_enclosing_loop() cidx = state[-1] lidx = state[-2] - lw = state[-3] - lx = state[-4] - p_accs, ca, cb = _unpack_state(state[:-4]) + p_accs, ca, cb = _unpack_state(state[:-2]) # tile T from the carried bank, prefetch T+1 -> next bank (yielded). - p_accs, bank, lx, lw = _run_tile( - p_accs, ca, cb, lx, lw, cidx, lidx, do_prefetch=True + p_accs, bank = _run_tile( + p_accs, ca, cb, cidx, lidx, do_prefetch=True ) next_a, next_b = bank cidx1 = arith.addi(cidx, one_i32) lidx1 = arith.addi(lidx, one_i32) new_state = _pack_state(p_accs, next_a, next_b) + [ - lx, - lw, lidx1, cidx1, ] results = yield new_state - accs, cur_a, cur_b = _unpack_state(results[:-4]) + accs, cur_a, cur_b = _unpack_state(results[:-2]) else: accs = list(accs) # No main loop ran — drain starts at compute_idx = 0. @@ -1111,13 +1110,6 @@ def launch_gemm_a16w16( _ = cache_tag ctx = CompilationContext.get_current() - # Finalize LDS allocations - with ir.InsertionPoint(ctx.gpu_module_body): - for alloc in stage_allocators: - alloc.finalized = False - for alloc in stage_allocators: - alloc.finalize() - # Grid dimensions idx_m = arith.index_cast(T.index, i32_m.ir_value()) idx_n = arith.index_cast(T.index, i32_n.ir_value()) @@ -1181,9 +1173,20 @@ def launch_gemm_a16w16( stream=stream, ) + # Commented out until coexec branch is merged. + # launch_gemm_a8w8_blockscale.compile_hints["llvm_options"] = { + # "amdgpu-expert-scheduling-mode": True, + # "amdgpu-anti-hints-for-va-vdst": True, + # "amdgpu-enable-static-simulator": True, + # "amdgpu-static-sim-inline": True, + # "amdgpu-sched-strategy": "coexec", + # # "amdgpu-block-carried-latency": EnumOpt("all"), # enable per the note above + # } + # Backend cl::opt hints; only options present in THIS build are set. _llvm_opts = { "amdgpu-expert-scheduling-mode": True, # valid bool cl::opt (GFX12+ only) + "unroll-threshold": 0, } # Experimental scheduling levers (in-process LLVM); None => cl::init default. if sched_strategy is not None: From 48b36d0643a3c70a6bb75e10268db8f44415fd5e Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Wed, 29 Jul 2026 21:58:29 +0000 Subject: [PATCH 11/34] Cleanup pt2 --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 1289 +++++------------ 1 file changed, 384 insertions(+), 905 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index 1f7d91394bc..eda173563d5 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - import flydsl.compiler as flyc import flydsl.expr as fx import torch @@ -23,16 +22,8 @@ from flydsl.runtime.device import get_rocm_arch as get_hip_arch from flydsl.utils.smem_allocator import check_smem_capacity -# WMMA 16×16×32 WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 -WAVE_SIZE = 32 - -# LDS row padding (8 elems = 16 bytes) -LDS_PAD_A = 8 -LDS_PAD_B = 8 - -_STAGE_NAMES = ("ping", "pong", "pang", "pung") - +WAVE_SIZE, LDS_PAD_A, LDS_PAD_B = 32, 8, 8 _SCHED_ALLOW_SALU = 1 << 2 @@ -83,60 +74,29 @@ def _make_buffer_rsrc(tensor, num_records_bytes=None): return fx.rocdl.get_buffer_rsrc(fx.get_iter(buf_tensor)) -def _disable_unroll_on_enclosing_loop(): - """Tag the enclosing scf.for with llvm.loop metadata so it stays rolled at ASM level.""" - block = ir.InsertionPoint.current.block - op = block.owner - if op.name != "scf.for": - return - anno = ir.Attribute.parse( - "#llvm.loop_annotation, disableNonforced = true>" - ) - op.attributes["loop_annotation"] = anno - +def apply_activation_scalar(val, activation: str): + import math -def _apply_activation_scalar(val, activation: str): - """Apply activation to a single f32 scalar.""" from flydsl._mlir.dialects import math as _math if activation == "relu": zero = arith.constant(0.0, type=T.f32) - return arith.select(val > zero, val, zero) - - elif activation in ("silu", "silu_exp2"): - neg = arith.constant(0.0, type=T.f32) - val - exp_neg = _math.exp(neg) - one = arith.constant(1.0, type=T.f32) - denom = one + exp_neg - return val / denom - - elif activation == "gelu": - import math - - inv_sqrt2 = arith.constant(1.0 / math.sqrt(2.0), type=T.f32) - scaled = val * inv_sqrt2 - erf_val = _math.erf(scaled) - one = arith.constant(1.0, type=T.f32) - half = arith.constant(0.5, type=T.f32) - return half * val * (one + erf_val) - - elif activation == "gelu_tanh": - import math - - sqrt_2_over_pi = arith.constant(math.sqrt(2.0 / math.pi), type=T.f32) - coeff = arith.constant(0.044715, type=T.f32) - one = arith.constant(1.0, type=T.f32) + return (val > zero).select(val, zero) + one, half = arith.constant(1.0, type=T.f32), arith.constant(0.5, type=T.f32) + if activation in ("silu", "silu_exp2"): + return val / (one + _math.exp(arith.constant(0.0, type=T.f32) - val)) + if activation == "gelu": + scaled = val * arith.constant(1.0 / math.sqrt(2.0), type=T.f32) + return half * val * (one + _math.erf(scaled)) + if activation == "gelu_tanh": two = arith.constant(2.0, type=T.f32) - half = arith.constant(0.5, type=T.f32) - x3 = val * val * val - inner = sqrt_2_over_pi * (val + coeff * x3) + inner = arith.constant(math.sqrt(2.0 / math.pi), type=T.f32) * ( + val + arith.constant(0.044715, type=T.f32) * val * val * val + ) # tanh(z) = 1 - 2/(1 + exp(2z)) - exp2x = _math.exp(two * inner) - tanh_val = one - two / (one + exp2x) + tanh_val = one - two / (one + _math.exp(two * inner)) return half * val * (one + tanh_val) - - else: - return val + return val def compile_gemm_a16w16( @@ -158,181 +118,105 @@ def compile_gemm_a16w16( add_bias: bool = False, physical_mk: bool = True, # True=M-major (row-major X), False=K-major (col-major X) physical_kn: bool = False, # False=N-major (row-major W), True=K-major (transposed W) - loop_carried_load_percent: int | None = None, kernarg_preload: bool = False, - use_manual_barrier: bool = False, split_k: int = 1, sched_strategy: str | None = None, - barrier_signal_wait_latency: int | None = None, main_loop_unroll: bool = False, variant: str = "bandwidth_bound", ): - """Compile the A16W16 GEMM kernel; returns launch_fn(y, x, w, bias, M, N, stream=stream). - - variant: - "bandwidth_bound" (default) — steady-state loads the whole next fragment - bank up front, then runs all WMMAs (burst load, ~2 banks live). - "compute_bound" — in-place fragment rotation: each fragment's next-tile - ds_load is hoisted to right after its last WMMA use so the load - co-executes with the remaining WMMAs (~1 bank live + trickle). - """ _ = (M, N) - if not (2 <= num_buffers <= 8): - raise ValueError(f"num_buffers must be between 2 and 8, got {num_buffers}") - if variant not in ("bandwidth_bound", "compute_bound"): - raise ValueError( - "variant must be 'bandwidth_bound' or 'compute_bound', " f"got {variant!r}" - ) - if in_dtype not in ("fp16", "bf16"): - raise ValueError(f"in_dtype must be 'fp16' or 'bf16', got {in_dtype!r}") - # Experimental LLVM scheduling levers (injected via llvm_options below). - if sched_strategy is not None and sched_strategy not in ( - "max-ilp", - "max-memory-clause", - ): - raise ValueError( - "sched_strategy must be None, 'max-ilp', or 'max-memory-clause', " - f"got {sched_strategy!r}" - ) + def _req(cond, msg): + if not cond: + raise ValueError(msg) - effective_waves_per_eu = waves_per_eu + if out_dtype is None: + out_dtype = "f16" if in_dtype == "fp16" else "bf16" is_f16 = in_dtype == "fp16" + effective_waves_per_eu = waves_per_eu - if out_dtype is None: - out_dtype = "f16" if is_f16 else "bf16" - if out_dtype not in ("f32", "f16", "bf16"): - raise ValueError( - f"out_dtype must be 'f32', 'f16', or 'bf16', got {out_dtype!r}" - ) + _req(2 <= num_buffers <= 8, f"num_buffers must be in [2, 8], got {num_buffers}") + _req(variant in ("bandwidth_bound", "compute_bound"), f"bad variant {variant!r}") + _req(in_dtype in ("fp16", "bf16"), f"bad in_dtype {in_dtype!r}") + _req(out_dtype in ("f32", "f16", "bf16"), f"bad out_dtype {out_dtype!r}") + _req( + sched_strategy in (None, "max-ilp", "max-memory-clause"), + f"bad sched_strategy {sched_strategy!r}", + ) + _req(split_k >= 1, f"split_k must be >= 1, got {split_k}") elem_bytes = 2 elem_bytes_d = 2 if out_dtype in ("f16", "bf16") else 4 - num_warps = m_warp * n_warp - block_threads = num_warps * WAVE_SIZE - - if K % tile_k != 0: - raise ValueError(f"K must be divisible by tile_k={tile_k}, got K={K}") - if split_k < 1: - raise ValueError(f"split_k must be >= 1, got {split_k}") - if K % split_k != 0: - raise ValueError(f"K must be divisible by split_k={split_k}, got K={K}") - if tile_k % WMMA_K != 0: - raise ValueError(f"tile_k must be a multiple of {WMMA_K}, got {tile_k}") - if tile_m % WMMA_M != 0: - raise ValueError(f"tile_m must be a multiple of {WMMA_M}, got {tile_m}") - if tile_n % WMMA_N != 0: - raise ValueError(f"tile_n must be a multiple of {WMMA_N}, got {tile_n}") - if (tile_k & (tile_k - 1)) != 0: - raise ValueError(f"tile_k must be a power of 2 for TDM, got {tile_k}") + num_warps, block_threads = m_warp * n_warp, m_warp * n_warp * WAVE_SIZE + warp_tile_m, warp_tile_n = tile_m // m_warp, tile_n // n_warp + split_k_chunk = K // split_k + num_k_tiles = split_k_chunk // tile_k - if physical_kn: - if N == 0: - raise ValueError( - "N must be specified (> 0) at compile time when physical_kn=True, " - "because it is used as the TDM stride for the (K, N) weight layout" - ) - if (tile_n & (tile_n - 1)) != 0: - raise ValueError( - f"tile_n must be a power of 2 when physical_kn=True " - f"(TDM pad_interval requirement), got {tile_n}" - ) - if not physical_mk: - if M == 0: - raise ValueError( - "M must be specified (> 0) at compile time when physical_mk=False, " - "because it is used as the TDM stride for the (K, M) activation layout" - ) - if (tile_m & (tile_m - 1)) != 0: - raise ValueError( - f"tile_m must be a power of 2 when physical_mk=False " - f"(TDM pad_interval requirement), got {tile_m}" - ) + for _nm, _v, _d in ( + ("K", K, tile_k), + ("K", K, split_k), + ("K/split_k", split_k_chunk, tile_k), + ("tile_k", tile_k, WMMA_K), + ("tile_m", tile_m, WMMA_M), + ("tile_n", tile_n, WMMA_N), + ("warp_tile_m", warp_tile_m, WMMA_M), + ("warp_tile_n", warp_tile_n, WMMA_N), + ): + _req(_v % _d == 0, f"{_nm} ({_v}) must be a multiple of {_d}") - warp_tile_m = tile_m // m_warp - warp_tile_n = tile_n // n_warp - if warp_tile_m % WMMA_M != 0: - raise ValueError(f"warp_tile_m={warp_tile_m} must be a multiple of {WMMA_M}") - if warp_tile_n % WMMA_N != 0: - raise ValueError(f"warp_tile_n={warp_tile_n} must be a multiple of {WMMA_N}") + _req( + tile_k & (tile_k - 1) == 0, f"tile_k must be a power of 2 for TDM, got {tile_k}" + ) - # split_k>1 splits K into split_k grid-z chunks; num_k_tiles is per-split. - split_k_chunk = K // split_k - if split_k_chunk % tile_k != 0: - raise ValueError( - f"K/split_k must be divisible by tile_k={tile_k}, got {split_k_chunk}" - ) - num_k_tiles = split_k_chunk // tile_k - if num_k_tiles < num_buffers - 1: - raise ValueError( - f"{num_buffers}-stage buffering requires num_k_tiles >= {num_buffers - 1}, " - f"got {num_k_tiles} (K={K}, tile_k={tile_k})" - ) + if physical_kn: + _req(N > 0, "N must be > 0 at compile time when physical_kn=True") + _req(tile_n & (tile_n - 1) == 0, f"tile_n must be a power of 2, got {tile_n}") + if not physical_mk: + _req(M > 0, "M must be > 0 at compile time when physical_mk=False") + _req(tile_m & (tile_m - 1) == 0, f"tile_m must be a power of 2, got {tile_m}") + _req( + num_k_tiles >= num_buffers - 1, + f"num_k_tiles {num_k_tiles} < {num_buffers - 1} (K={K}, tile_k={tile_k})", + ) gpu_arch = str(get_hip_arch()) assert gpu_arch.startswith("gfx1250"), f"Expected gfx1250, got {gpu_arch}" k_wmma_steps = tile_k // WMMA_K - - def _elem_type(): - return T.f16 if is_f16 else T.bf16 - _fx_elem = fx.Float16 if is_f16 else fx.BFloat16 - - wmma_m_rep = warp_tile_m // WMMA_M - wmma_n_rep = warp_tile_n // WMMA_N + wmma_m_rep, wmma_n_rep = warp_tile_m // WMMA_M, warp_tile_n // WMMA_N n_accs = wmma_m_rep * wmma_n_rep - N_A_FRAGS = wmma_m_rep * k_wmma_steps - N_B_FRAGS = wmma_n_rep * k_wmma_steps + N_ACCS = n_accs + N_A_FRAGS, N_B_FRAGS = wmma_m_rep * k_wmma_steps, wmma_n_rep * k_wmma_steps - # LDS layout: A is [tile_m, tile_k+pad] (M-major) or [tile_k, tile_m+pad] (K-major) if physical_mk: - lds_a_stride = tile_k + LDS_PAD_A + a_tile_shape, a_outer_stride, a_pad_interval = (tile_m, tile_k), K, tile_k + a_imm_bytes, lds_a_stride = tile_k * elem_bytes, tile_k + LDS_PAD_A lds_a_elems = tile_m * lds_a_stride + LDS_PAD_A else: - lds_a_stride = tile_m + LDS_PAD_A + a_tile_shape, a_outer_stride, a_pad_interval = (tile_k, tile_m), M, tile_m + a_imm_bytes, lds_a_stride = tile_k * M * elem_bytes, tile_m + LDS_PAD_A lds_a_elems = tile_k * lds_a_stride + LDS_PAD_A if physical_kn: - lds_b_stride = tile_n + LDS_PAD_B + b_tile_shape, b_outer_stride, b_pad_interval = (tile_k, tile_n), N, tile_n + b_imm_bytes, lds_b_stride = tile_k * N * elem_bytes, tile_n + LDS_PAD_B lds_b_elems = tile_k * lds_b_stride + LDS_PAD_B else: - lds_b_stride = tile_k + LDS_PAD_B + b_tile_shape, b_outer_stride, b_pad_interval = (tile_n, tile_k), K, tile_k + b_imm_bytes, lds_b_stride = tile_k * elem_bytes, tile_k + LDS_PAD_B lds_b_elems = tile_n * lds_b_stride + LDS_PAD_B - lds_a_data_bytes = lds_a_elems * elem_bytes - lds_b_data_bytes = lds_b_elems * elem_bytes - - # Unified LDS allocator: contiguous [A0..A_nb-1 | B0..B_nb-1] ring slots + lds_a_data_bytes, lds_b_data_bytes = ( + lds_a_elems * elem_bytes, + lds_b_elems * elem_bytes, + ) unified_a_off = 0 unified_b_off = _align_up(unified_a_off + num_buffers * lds_a_data_bytes, 16) lds_arena_bytes = _align_up(unified_b_off + num_buffers * lds_b_data_bytes, 16) check_smem_capacity(lds_arena_bytes, gpu_arch) - if physical_mk: - a_tile_shape = (tile_m, tile_k) - a_outer_stride = K - a_pad_interval = tile_k - a_imm_bytes = tile_k * elem_bytes - else: - a_tile_shape = (tile_k, tile_m) - a_outer_stride = M - a_pad_interval = tile_m - a_imm_bytes = tile_k * M * elem_bytes - - if physical_kn: - b_tile_shape = (tile_k, tile_n) - b_outer_stride = N - b_pad_interval = tile_n - b_imm_bytes = tile_k * N * elem_bytes - else: - b_tile_shape = (tile_n, tile_k) - b_outer_stride = K - b_pad_interval = tile_k - b_imm_bytes = tile_k * elem_bytes - - # TDMs per K-tile: A + B - _TDMS_PER_TILE = 2 + _TDMS_PER_TILE = 2 # A + B @flyc.kernel def kernel_gemm_a16w16( @@ -344,20 +228,16 @@ def kernel_gemm_a16w16( i32_n: fx.Int32, ): rocdl.disable_xdl_arb_stall() + elem_ty = T.f16 if is_f16 else T.bf16 - tx = gpu.thread_id("x") - bx = gpu.block_id("x") - by = gpu.block_id("y") - blk_m = bx * arith.index(tile_m) - blk_n = by * arith.index(tile_n) - # split_k>1: grid-z block picks the K-chunk origin (compile-time gated). + tx, bx, by = gpu.thread_id("x"), gpu.block_id("x"), gpu.block_id("y") + blk_m, blk_n = bx * arith.index(tile_m), by * arith.index(tile_n) if const_expr(split_k > 1): bz = gpu.block_id("z") split_k_base = bz * arith.index(split_k_chunk) else: split_k_base = arith.index(0) - # Thread -> warp decomposition layout_thr = fx.make_layout( (m_warp, n_warp, 2, 16), (n_warp * WAVE_SIZE, WAVE_SIZE, 16, 1) ) @@ -371,12 +251,8 @@ def kernel_gemm_a16w16( warp_m_base = wave_m_idx * arith.index(warp_tile_m) warp_n_base = wave_n_idx * arith.index(warp_tile_n) - elem_ty = _elem_type() - m_idx = arith.index_cast(T.index, i32_m.ir_value()) n_stride = arith.index_cast(T.index, i32_n.ir_value()) - - # Buffer resource descriptors y_nrec = m_idx * n_stride * arith.index(elem_bytes_d) y_rsrc = _make_buffer_rsrc(arg_y, num_records_bytes=y_nrec) if const_expr(add_bias): @@ -389,52 +265,41 @@ def kernel_gemm_a16w16( big_b_base_idx = arith.index_cast(T.index, arith.unwrap(fx.ptrtoint(big_b_mem))) slot_stride_a_elems_i32 = arith.constant(lds_a_elems, type=T.i32) slot_stride_b_elems_i32 = arith.constant(lds_b_elems, type=T.i32) - - blk_m64 = fx.Int64(blk_m) - blk_n64 = fx.Int64(blk_n) + blk_m64, blk_n64 = fx.Int64(blk_m), fx.Int64(blk_n) split_k_base64 = fx.Int64(split_k_base) - def _global_view(tensor, elem_off, shape, stride): - return fx.Tensor( + def _mk_side(tensor, base_off, shape, outer_stride, pad_interval, pad_amount): + """Global tile view + its TDM atom (identical shape for A and B).""" + gt = fx.Tensor( fx.make_view( - fx.add_offset(fx.get_iter(tensor), elem_off), - fx.make_layout(shape, stride), + fx.add_offset(fx.get_iter(tensor), base_off), + fx.make_layout(shape, (outer_stride, 1)), ) ) + return gt, fx.rocdl.make_tdm_atom( + gt, + [None, None], + strides=[fx.Int64(outer_stride), None], + num_warps=num_warps, + pad_interval=pad_interval, + pad_amount=pad_amount, + ) if const_expr(physical_mk): a_base_off = blk_m64 * fx.Int64(K) + split_k_base64 else: a_base_off = split_k_base64 * fx.Int64(M) + blk_m64 - gA = _global_view(arg_x, a_base_off, a_tile_shape, (a_outer_stride, 1)) - atom_a = fx.rocdl.make_tdm_atom( - gA, - [None, None], - strides=[fx.Int64(a_outer_stride), None], - num_warps=num_warps, - pad_interval=a_pad_interval, - pad_amount=LDS_PAD_A, + gA, atom_a = _mk_side( + arg_x, a_base_off, a_tile_shape, a_outer_stride, a_pad_interval, LDS_PAD_A ) - if const_expr(physical_kn): b_base_off = split_k_base64 * fx.Int64(N) + blk_n64 else: b_base_off = blk_n64 * fx.Int64(K) + split_k_base64 - gB = _global_view(arg_w, b_base_off, b_tile_shape, (b_outer_stride, 1)) - atom_b = fx.rocdl.make_tdm_atom( - gB, - [None, None], - strides=[fx.Int64(b_outer_stride), None], - num_warps=num_warps, - pad_interval=b_pad_interval, - pad_amount=LDS_PAD_B, + gB, atom_b = _mk_side( + arg_w, b_base_off, b_tile_shape, b_outer_stride, b_pad_interval, LDS_PAD_B ) - def _slot_elem_off(buf_idx, slot_elems, slot_stride_i32): - if const_expr(isinstance(buf_idx, int)): - return arith.index(buf_idx * slot_elems) - return arith.index_cast(T.index, arith.muli(buf_idx, slot_stride_i32)) - def _imm64(k_tile, mul): if const_expr(isinstance(k_tile, int)): return fx.Int64(k_tile * mul) @@ -443,179 +308,122 @@ def _imm64(k_tile, mul): ) return fx.Int64(arith.index_cast(T.i64, as_index)) + def _slot_off(buf_idx, slot_elems, stride_i32): + if const_expr(isinstance(buf_idx, int)): + return arith.index(buf_idx * slot_elems) + return arith.index_cast(T.index, arith.muli(buf_idx, stride_i32)) + + def _lane_bases(warp_base, reps, lds_stride, transpose): + if const_expr(not transpose): + row_off = (warp_base + lane16) * arith.index(lds_stride) + k_off = lane_kgrp * arith.index(8) + return [ + row_off + arith.index(rep * WMMA_M * lds_stride) + k_off + for rep in range_constexpr(reps) + ] + k_off = ( + lane_kgrp * arith.index(8) + lane16 % arith.index(8) + ) * arith.index(lds_stride) + grp_off = (lane16 / arith.index(8)) * arith.index(8) + return [ + k_off + warp_base + arith.index(rep * WMMA_M) + grp_off + for rep in range_constexpr(reps) + ] + + # One spec per operand side; A and B differ only in these fields. + MEM, BASE, BASES, STRIDE, TR, ELEMS, SSTRIDE, REPS, ATOM, GT, SHAPE, IMM = ( + range(12) + ) + A_SIDE = ( + big_a_mem, + big_a_base_idx, + _lane_bases(warp_m_base, wmma_m_rep, lds_a_stride, not physical_mk), + lds_a_stride, + not physical_mk, + lds_a_elems, + slot_stride_a_elems_i32, + wmma_m_rep, + atom_a, + gA, + a_tile_shape, + a_imm_bytes, + ) + B_SIDE = ( + big_b_mem, + big_b_base_idx, + _lane_bases(warp_n_base, wmma_n_rep, lds_b_stride, physical_kn), + lds_b_stride, + physical_kn, + lds_b_elems, + slot_stride_b_elems_i32, + wmma_n_rep, + atom_b, + gB, + b_tile_shape, + b_imm_bytes, + ) + def issue_tdm_loads(buf_idx, k_tile): - """Issue A+B TDMs for K-tile k_tile into LDS ring slot buf_idx.""" - fx.copy( - atom_a, - gA, - fx.Tensor( - fx.make_view( - fx.add_offset( - big_a_mem, - _slot_elem_off( - buf_idx, lds_a_elems, slot_stride_a_elems_i32 + """A then B TDM for K-tile k_tile into LDS ring slot buf_idx.""" + for sd in (A_SIDE, B_SIDE): + fx.copy( + sd[ATOM], + sd[GT], + fx.Tensor( + fx.make_view( + fx.add_offset( + sd[MEM], + _slot_off(buf_idx, sd[ELEMS], sd[SSTRIDE]), ), - ), - fx.make_layout(a_tile_shape, (lds_a_stride, 1)), - ) - ), - imm_offset=_imm64(k_tile, a_imm_bytes), - ) - fx.copy( - atom_b, - gB, - fx.Tensor( - fx.make_view( - fx.add_offset( - big_b_mem, - _slot_elem_off( - buf_idx, lds_b_elems, slot_stride_b_elems_i32 - ), - ), - fx.make_layout(b_tile_shape, (lds_b_stride, 1)), - ) - ), - imm_offset=_imm64(k_tile, b_imm_bytes), - ) - - def _precompute_lane_bases(warp_base, reps, lds_stride): - """Per-lane LDS element offsets for `reps` WMMA frags (M/N-major).""" - row_stride_off = (warp_base + lane16) * arith.index(lds_stride) - k_lane_off = lane_kgrp * arith.index(8) - bases = [] - for rep in range_constexpr(reps): - base = ( - row_stride_off + arith.index(rep * WMMA_M * lds_stride) + k_lane_off + fx.make_layout(sd[SHAPE], (sd[STRIDE], 1)), + ) + ), + imm_offset=_imm64(k_tile, sd[IMM]), ) - bases.append(base) - return bases - - def _precompute_a_lane_bases(): - """A fragment lane bases. Uses transpose-load addressing when K-major.""" - if const_expr(physical_mk): - return _precompute_lane_bases(warp_m_base, wmma_m_rep, lds_a_stride) - - lane8 = lane16 % arith.index(8) - lane_mgrp = lane16 / arith.index(8) - k_lane_off = (lane_kgrp * arith.index(8) + lane8) * arith.index( - lds_a_stride - ) - m_lane_off = lane_mgrp * arith.index(8) - bases = [] - for wm in range_constexpr(wmma_m_rep): - m_col = warp_m_base + arith.index(wm * WMMA_M) + m_lane_off - bases.append(k_lane_off + m_col) - return bases - - def _precompute_b_lane_bases(): - """B fragment lane bases. Uses transpose-load addressing when K-major.""" - if const_expr(not physical_kn): - return _precompute_lane_bases(warp_n_base, wmma_n_rep, lds_b_stride) - - lane8 = lane16 % arith.index(8) - lane_ngrp = lane16 / arith.index(8) - k_lane_off = (lane_kgrp * arith.index(8) + lane8) * arith.index( - lds_b_stride - ) - n_lane_off = lane_ngrp * arith.index(8) - bases = [] - for wn in range_constexpr(wmma_n_rep): - n_col = warp_n_base + arith.index(wn * WMMA_N) + n_lane_off - bases.append(k_lane_off + n_col) - return bases - - a_lane_bases = _precompute_a_lane_bases() - b_lane_bases = _precompute_b_lane_bases() - - def _lds_transpose_load(result_type, lds_base_idx, elem_off): - addr_i32 = arith.index_cast( - T.i32, lds_base_idx + elem_off * arith.index(elem_bytes) - ) - ptr_val = llvm_dialect.inttoptr( - ir.Type.parse("!llvm.ptr<3>"), arith.unwrap(addr_i32) - ) - return raw_rocdl.ds_load_tr16_b128(result_type, ptr_val) - def load_wmma_frag_tr(lds_buffer, b_lane_base, ks): - """Load WMMA B fragment via ds_load_tr16_b128 (K-major LDS).""" - vec8_ty = ir.VectorType.get([8], elem_ty) - lds_base_idx = big_b_base_idx - results = [] + def _frag(sd, slot_off, ks, rep): + vec8_ty = T.vec(8, elem_ty) + halves = [] for k_half in range_constexpr(2): - k_row_off = (ks * WMMA_K + k_half * 16) * lds_b_stride - elem_off = b_lane_base + arith.index(k_row_off) - v = _lds_transpose_load(vec8_ty, lds_base_idx, elem_off) - results.append(v) - return ( - fx.Vector(results[0]) - .shuffle(fx.Vector(results[1]), list(range(16))) - .ir_value() - ) - - def load_wmma_frag(lds_buffer, lane_base, ks): - """Load WMMA fragment via ds_read_b128 (M/N-major LDS).""" - vec8_ty = ir.VectorType.get([8], elem_ty) - off0 = lane_base + arith.index(ks * WMMA_K) - off1 = lane_base + arith.index(ks * WMMA_K + 16) - v0 = arith.unwrap( - fx.ptr_load(fx.add_offset(lds_buffer, off0), result_type=vec8_ty) - ) - v1 = arith.unwrap( - fx.ptr_load(fx.add_offset(lds_buffer, off1), result_type=vec8_ty) - ) - return fx.Vector(v0).shuffle(fx.Vector(v1), list(range(16))).ir_value() - - _load_b_frag = load_wmma_frag_tr if physical_kn else load_wmma_frag - - def load_wmma_frag_tr_a(lds_buffer, a_lane_base, ks): - """Load WMMA A fragment via ds_load_tr16_b128 (K-major LDS).""" - vec8_ty = ir.VectorType.get([8], elem_ty) - lds_base_idx = big_a_base_idx - results = [] - for k_half in range_constexpr(2): - k_row_off = (ks * WMMA_K + k_half * 16) * lds_a_stride - elem_off = a_lane_base + arith.index(k_row_off) - v = _lds_transpose_load(vec8_ty, lds_base_idx, elem_off) - results.append(v) + if const_expr(sd[TR]): + off = ( + sd[BASES][rep] + + slot_off + + arith.index((ks * WMMA_K + k_half * 16) * sd[STRIDE]) + ) + addr = arith.index_cast( + T.i32, sd[BASE] + off * arith.index(elem_bytes) + ) + ptr = llvm_dialect.inttoptr( + ir.Type.parse("!llvm.ptr<3>"), arith.unwrap(addr) + ) + halves.append(raw_rocdl.ds_load_tr16_b128(vec8_ty, ptr)) + else: + off = ( + sd[BASES][rep] + + slot_off + + arith.index(ks * WMMA_K + k_half * 16) + ) + halves.append( + arith.unwrap( + fx.ptr_load( + fx.add_offset(sd[MEM], off), result_type=vec8_ty + ) + ) + ) return ( - fx.Vector(results[0]) - .shuffle(fx.Vector(results[1]), list(range(16))) + fx.Vector(halves[0]) + .shuffle(fx.Vector(halves[1]), list(range(16))) .ir_value() ) - _load_a_frag = load_wmma_frag_tr_a if not physical_mk else load_wmma_frag - - def load_a_frags(buf_idx): - """Load all A frags for one K-tile from LDS slot buf_idx; index a_frags[ks*wmma_m_rep+wm].""" - if const_expr(isinstance(buf_idx, int)): - slot_off_a = arith.index(buf_idx * lds_a_elems) - else: - slot_off_a = arith.index_cast( - T.index, arith.muli(buf_idx, slot_stride_a_elems_i32) - ) - a_frags = [] - for ks in range_constexpr(k_wmma_steps): - for wm in range_constexpr(wmma_m_rep): - a_frags.append( - _load_a_frag(big_a_mem, a_lane_bases[wm] + slot_off_a, ks) - ) - return a_frags - - def load_b_frags(buf_idx): - """Load all B frags for one K-tile from LDS slot buf_idx; index b_frags[ks*wmma_n_rep+wn].""" - if const_expr(isinstance(buf_idx, int)): - slot_off_b = arith.index(buf_idx * lds_b_elems) - else: - slot_off_b = arith.index_cast( - T.index, arith.muli(buf_idx, slot_stride_b_elems_i32) - ) - b_frags = [] - for ks in range_constexpr(k_wmma_steps): - for wn in range_constexpr(wmma_n_rep): - b_frags.append( - _load_b_frag(big_b_mem, b_lane_bases[wn] + slot_off_b, ks) - ) - return b_frags + def load_frags(sd, buf_idx): + off = _slot_off(buf_idx, sd[ELEMS], sd[SSTRIDE]) + return [ + _frag(sd, off, ks, rep) + for ks in range_constexpr(k_wmma_steps) + for rep in range_constexpr(sd[REPS]) + ] wmma_atom = fx.make_mma_atom( fx.rocdl.WMMA(WMMA_M, WMMA_N, WMMA_K, _fx_elem, fx.Float32) @@ -626,152 +434,73 @@ def _rmem_vec(v, n, dtype): t.store(fx.Vector(arith.unwrap(v))) return t - def _rmem_frag_lists(a_frags, b_frags): - return ( - [_rmem_vec(f, 16, _fx_elem) for f in a_frags], - [_rmem_vec(f, 16, _fx_elem) for f in b_frags], - ) - - def wmma_tile(accs_in, a_frags, b_frags): - """Execute all WMMAs for one tile (ks-outer accumulate).""" - a_frag_ts, b_frag_ts = _rmem_frag_lists(a_frags, b_frags) + def wmma_tile(accs_in, a_frags, b_frags, rotate_buf=None): + a_ts = [_rmem_vec(f, 16, _fx_elem) for f in a_frags] + b_ts = [_rmem_vec(f, 16, _fx_elem) for f in b_frags] acc_ts = [_rmem_vec(a, 8, fx.Float32) for a in accs_in] - for ks in range_constexpr(k_wmma_steps): - a_f = [ - a_frag_ts[ks * wmma_m_rep + wm] - for wm in range_constexpr(wmma_m_rep) - ] - b_f = [ - b_frag_ts[ks * wmma_n_rep + wn] - for wn in range_constexpr(wmma_n_rep) - ] - for wm in range_constexpr(wmma_m_rep): - for wn in range_constexpr(wmma_n_rep): - idx = wm * wmma_n_rep + wn - # ISA operand order: (B, A, C) — reversed from math - fx.gemm(wmma_atom, acc_ts[idx], b_f[wn], a_f[wm], acc_ts[idx]) - return [arith.unwrap(t.load()) for t in acc_ts] - - def _load_one_a_frag(buf_i32, ks, wm): - """ds_load a single A fragment (ks, wm) from LDS ring slot buf_i32 (i32 SSA).""" - slot_off_a = arith.index_cast( - T.index, arith.muli(buf_i32, slot_stride_a_elems_i32) - ) - return _load_a_frag(big_a_mem, a_lane_bases[wm] + slot_off_a, ks) - - def _load_one_b_frag(buf_i32, ks, wn): - """ds_load a single B fragment (ks, wn) from LDS ring slot buf_i32 (i32 SSA).""" - slot_off_b = arith.index_cast( - T.index, arith.muli(buf_i32, slot_stride_b_elems_i32) - ) - return _load_b_frag(big_b_mem, b_lane_bases[wn] + slot_off_b, ks) - - def _wmma_rotate(accs_in, a_frags, b_frags, next_buf_i32): - """compute_bound steady-state: run this tile's WMMAs, and the instant a - fragment reaches its last use, issue its next-tile ds_load so the load - co-executes with the remaining WMMAs. Returns (accs, next_a, next_b) where - next_* are the freshly rotated-in fragments for the next K-tile.""" - a_frag_ts, b_frag_ts = _rmem_frag_lists(a_frags, b_frags) - acc_ts = [_rmem_vec(a, 8, fx.Float32) for a in accs_in] - next_a = [None] * N_A_FRAGS - next_b = [None] * N_B_FRAGS + rotate = rotate_buf is not None + next_a = [None] * N_A_FRAGS if rotate else None + next_b = [None] * N_B_FRAGS if rotate else None + if const_expr(rotate): + a_off = _slot_off(rotate_buf, lds_a_elems, slot_stride_a_elems_i32) + b_off = _slot_off(rotate_buf, lds_b_elems, slot_stride_b_elems_i32) for ks in range_constexpr(k_wmma_steps): for wm in range_constexpr(wmma_m_rep): - a_f = a_frag_ts[ks * wmma_m_rep + wm] + a_f = a_ts[ks * wmma_m_rep + wm] for wn in range_constexpr(wmma_n_rep): idx = wm * wmma_n_rep + wn fx.gemm( wmma_atom, acc_ts[idx], - b_frag_ts[ks * wmma_n_rep + wn], + b_ts[ks * wmma_n_rep + wn], a_f, acc_ts[idx], ) + if const_expr(rotate): + rocdl.sched_barrier(_SCHED_ALLOW_SALU) + next_a[ks * wmma_m_rep + wm] = _frag(A_SIDE, a_off, ks, wm) + if const_expr(rotate): rocdl.sched_barrier(_SCHED_ALLOW_SALU) - next_a[ks * wmma_m_rep + wm] = _load_one_a_frag( - next_buf_i32, ks, wm - ) - rocdl.sched_barrier(_SCHED_ALLOW_SALU) - for wn in range_constexpr(wmma_n_rep): - next_b[ks * wmma_n_rep + wn] = _load_one_b_frag( - next_buf_i32, ks, wn - ) + for wn in range_constexpr(wmma_n_rep): + next_b[ks * wmma_n_rep + wn] = _frag(B_SIDE, b_off, ks, wn) return [arith.unwrap(t.load()) for t in acc_ts], next_a, next_b _half_out = out_dtype in ("f16", "bf16") - _out_elem = ( - T.f16 if out_dtype == "f16" else (T.bf16 if out_dtype == "bf16" else None) - ) _out_num = ( - fx.Float16 - if out_dtype == "f16" - else (fx.BFloat16 if out_dtype == "bf16" else None) + fx.Float16 if out_dtype == "f16" else fx.BFloat16 if _half_out else None ) - _bias_elem = T.f16 if is_f16 else T.bf16 - - def _widen_to_f32(val): - from flydsl._mlir.dialects import arith as _std_arith - - return _std_arith.ExtFOp(T.f32, _raw(val)).result - - def _apply_bias_and_activation(accs): - """Add bias and/or apply activation to accumulators.""" - if const_expr(not add_bias and activation is None): - return accs + _bias_elem = elem_ty - for wm in range_constexpr(wmma_m_rep): + def epilogue_stores(final_accs): + if const_expr(split_k > 1): + zero_i32 = fx.Int32(0) + if const_expr(add_bias): + bias_vecs = [] for wn in range_constexpr(wmma_n_rep): - idx = wm * wmma_n_rep + wn - acc = accs[idx] - - if const_expr(add_bias): - col_base = ( - blk_n - + warp_n_base - + arith.index(wn * WMMA_N) - + lane_kgrp * arith.index(8) - ) - col_base_i32 = arith.index_cast(T.i32, col_base) - col_base_hi = col_base_i32 + arith.constant(4, type=T.i32) - - bv_lo = _buffer_load( - bias_rsrc, col_base_i32, vec_width=4, dtype=_bias_elem - ) - bv_hi = _buffer_load( - bias_rsrc, col_base_hi, vec_width=4, dtype=_bias_elem + col_i32 = arith.index_cast( + T.i32, + blk_n + + warp_n_base + + arith.index(wn * WMMA_N) + + lane_kgrp * arith.index(8), + ) + elems = [] + for half in range_constexpr(2): + bv = _buffer_load( + bias_rsrc, + col_i32 + arith.constant(half * 4, type=T.i32), + vec_width=4, + dtype=_bias_elem, ) - - bias_elems = [] for i in range_constexpr(4): - b_val = fx.Vector(bv_lo)[i].ir_value() - bias_elems.append(_widen_to_f32(b_val)) - for i in range_constexpr(4): - b_val = fx.Vector(bv_hi)[i].ir_value() - bias_elems.append(_widen_to_f32(b_val)) - - bias_vec = fx.Vector.from_elements( - bias_elems, fx.Float32 - ).ir_value() - acc = acc + bias_vec - - if const_expr(activation is not None): - new_elems = [] - for i in range_constexpr(8): - val = fx.Vector(acc)[i].ir_value() - val = _apply_activation_scalar(val, activation) - new_elems.append(val) - acc = fx.Vector.from_elements(new_elems, fx.Float32).ir_value() - - accs[idx] = acc - return accs + elems.append(fx.Vector(bv)[i].to(fx.Float32).ir_value()) + bias_vecs.append( + fx.Vector.from_elements(elems, fx.Float32).ir_value() + ) - def epilogue_stores(final_accs): - """Write accumulators to Y (buffer_store, or atomic-fadd if split_k>1).""" - if const_expr(split_k > 1): - zero_i32 = fx.Int32(0) for wm in range_constexpr(wmma_m_rep): for wn in range_constexpr(wmma_n_rep): - idx = wm * wmma_n_rep + wn + acc = final_accs[wm * wmma_n_rep + wn] row = blk_m + warp_m_base + arith.index(wm * WMMA_M) + lane16 col_base = ( blk_n @@ -779,260 +508,143 @@ def epilogue_stores(final_accs): + arith.index(wn * WMMA_N) + lane_kgrp * arith.index(8) ) + if const_expr(add_bias): + acc = acc + bias_vecs[wn] + if const_expr(activation is not None): + acc = fx.Vector.from_elements( + [ + apply_activation_scalar( + fx.Vector(acc)[i].ir_value(), activation + ) + for i in range_constexpr(8) + ], + fx.Float32, + ).ir_value() if const_expr(_half_out): + h_vec = fx.Vector(acc).to(_out_num) + c_off_bytes = (row * n_stride + col_base) * arith.index( + elem_bytes_d + ) if const_expr(split_k > 1): - h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) - c_off_bytes = (row * n_stride + col_base) * arith.index( - elem_bytes_d - ) for pair in range_constexpr(4): - e0 = fx.Vector(h_vec)[pair * 2].ir_value() - e1 = fx.Vector(h_vec)[pair * 2 + 1].ir_value() pair_vec = fx.Vector.from_elements( - [e0, e1], _out_num + [ + h_vec[pair * 2].ir_value(), + h_vec[pair * 2 + 1].ir_value(), + ], + _out_num, ).ir_value() - byte_off = arith.index_cast( - T.i32, c_off_bytes + arith.index(pair * 4) - ) rocdl.raw_ptr_buffer_atomic_fadd( - pair_vec, y_rsrc, byte_off, zero_i32, zero_i32 + pair_vec, + y_rsrc, + arith.index_cast( + T.i32, c_off_bytes + arith.index(pair * 4) + ), + zero_i32, + zero_i32, ) else: - h_vec = arith.trunc_f(T.vec(8, _out_elem), final_accs[idx]) - i32_vec = fx.Vector(h_vec).bitcast(fx.Int32).ir_value() - c_off_bytes = (row * n_stride + col_base) * arith.index( - elem_bytes_d - ) _buffer_store( - i32_vec, y_rsrc, c_off_bytes, offset_is_bytes=True + h_vec.bitcast(fx.Int32).ir_value(), + y_rsrc, + c_off_bytes, + offset_is_bytes=True, + ) + elif const_expr(split_k > 1): + for e in range_constexpr(8): + rocdl.raw_ptr_buffer_atomic_fadd( + fx.Vector(acc)[e].ir_value(), + y_rsrc, + arith.index_cast( + T.i32, + (row * n_stride + col_base + arith.index(e)) + * arith.index(4), + ), + zero_i32, + zero_i32, ) else: - if const_expr(split_k > 1): - for half in range_constexpr(2): - base = row * n_stride + col_base + arith.index(half * 4) - for vi in range_constexpr(4): - val = fx.Vector(final_accs[idx])[ - half * 4 + vi - ].ir_value() - byte_off = arith.index_cast( - T.i32, - (base + arith.index(vi)) * arith.index(4), - ) - rocdl.raw_ptr_buffer_atomic_fadd( - val, y_rsrc, byte_off, zero_i32, zero_i32 - ) - else: - for half in range_constexpr(2): - vals = [ - fx.Vector(final_accs[idx])[half * 4 + vi].ir_value() + for half in range_constexpr(2): + vec4 = fx.Vector.from_elements( + [ + fx.Vector(acc)[half * 4 + vi].ir_value() for vi in range_constexpr(4) - ] - vec4 = fx.Vector.from_elements( - vals, fx.Float32 - ).ir_value() - col = col_base + arith.index(half * 4) - c_off = row * n_stride + col - _buffer_store(vec4, y_rsrc, c_off) - - # Accumulators - acc_zero = arith.constant_vector(0.0, T.vec(8, T.f32)) - accs = [acc_zero] * n_accs - - # State carry: [accs (N_ACCS) | cur_a frags (N_A_FRAGS) | cur_b frags (N_B_FRAGS)] - N_ACCS = n_accs + ], + fx.Float32, + ).ir_value() + col = col_base + arith.index(half * 4) + _buffer_store(vec4, y_rsrc, row * n_stride + col) def _pack_state(accs_, a_, b_): return list(accs_) + list(a_) + list(b_) def _unpack_state(state): - i = 0 - accs_ = list(state[i : i + N_ACCS]) - i += N_ACCS - a_ = list(state[i : i + N_A_FRAGS]) - i += N_A_FRAGS - b_ = list(state[i : i + N_B_FRAGS]) - i += N_B_FRAGS - return accs_, a_, b_ - - # ── Prologue (memory-bound): fill NB-1 LDS buffers, then load tile 0 ── + out, i = [], 0 + for n in (N_ACCS, N_A_FRAGS, N_B_FRAGS): + out.append(list(state[i : i + n])) + i += n + return out + + def _run_tile(accs_, ca, cb, cidx, lidx): + issue_tdm_loads(arith.remui(lidx, nb_const_i32), lidx) + tdm_ops.tensor_wait( + 0 if num_buffers == 2 else (num_buffers - 2) * _TDMS_PER_TILE + ) + gpu.barrier() + next_buf_i32 = arith.remui(arith.addi(cidx, one_i32), nb_const_i32) + if const_expr(variant == "compute_bound"): + return wmma_tile(accs_, ca, cb, rotate_buf=next_buf_i32) + na = load_frags(A_SIDE, next_buf_i32) + nb = load_frags(B_SIDE, next_buf_i32) + accs_, _, _ = wmma_tile(accs_, ca, cb) + return accs_, na, nb + + # Accumulators + acc_zero = arith.constant_vector(0.0, T.vec(8, T.f32)) + accs = [acc_zero] * n_accs + + # Prologue: for i in range_constexpr(num_buffers - 1): issue_tdm_loads(i, i) - - # Wait retires tile 0 (NB-2 in flight); the 2-iter WAR gap lets the main loop use one barrier. tdm_ops.tensor_wait((num_buffers - 2) * _TDMS_PER_TILE) gpu.barrier() + cur_a, cur_b = load_frags(A_SIDE, 0), load_frags(B_SIDE, 0) - cur_a = load_a_frags(0) - cur_b = load_b_frags(0) - - # ── Main loop (wait-first, unrolled by 2): ping-pong two register banks, prefetch T+1/T+2 ── main_loop_iters = num_k_tiles - (num_buffers - 1) nb_const_i32 = arith.constant(num_buffers, type=T.i32) one_i32 = arith.constant(1, type=T.i32) load_idx_init = arith.constant(num_buffers - 1, type=T.i32) compute_idx_init = arith.constant(0, type=T.i32) - def _wg_barrier(): - # gpu.barrier() drains ds_loads via fences; manual variant is a bare s_barrier (no fences, faster but racy). - if const_expr(use_manual_barrier): - rocdl.s_barrier_signal(-1) - rocdl.s_barrier_wait(-1) - else: - gpu.barrier() - - def _run_tile(accs_, ca, cb, cidx, lidx, do_prefetch): - """Compute K-tile cidx from carried frags; if do_prefetch, load cidx+1 into a fresh bank.""" - load_buf_i32 = arith.remui(lidx, nb_const_i32) - - # Refill first, then wait: keeps NB-2 tiles' TDM in flight (depth-2). - issue_tdm_loads(load_buf_i32, lidx) - if const_expr(num_buffers == 2): - # NB=2: only 2 slots -> drain the just-issued TDM (RAW hazard). - tdm_ops.tensor_wait(0) - else: - tdm_ops.tensor_wait((num_buffers - 2) * _TDMS_PER_TILE) - - # Single barrier: RAW wall for this tile's T+1 ds_reads and (across the backedge) WAR wall for the next refill. - _wg_barrier() - - if const_expr(do_prefetch): - next_buf_i32 = arith.remui(arith.addi(cidx, one_i32), nb_const_i32) - # Preload next tile's A+B, then WMMA (loads co-execute with it); raises peak liveness to ~2 banks. - na = load_a_frags(next_buf_i32) - nb = load_b_frags(next_buf_i32) - accs_ = wmma_tile(accs_, ca, cb) - nbank = (na, nb) - else: - accs_ = wmma_tile(accs_, ca, cb) - nbank = None - return accs_, nbank - - if const_expr(variant == "compute_bound"): - if const_expr(main_loop_iters > 0): - init_state = _pack_state(accs, cur_a, cur_b) + [ - load_idx_init, - compute_idx_init, - ] - results = init_state - for step, state in range(0, main_loop_iters, 1, init=init_state): - _disable_unroll_on_enclosing_loop() - cidx = state[-1] - lidx = state[-2] - p_accs, ca, cb = _unpack_state(state[:-2]) - - load_buf_i32 = arith.remui(lidx, nb_const_i32) - issue_tdm_loads(load_buf_i32, lidx) - if const_expr(num_buffers == 2): - tdm_ops.tensor_wait(0) - else: - tdm_ops.tensor_wait((num_buffers - 2) * _TDMS_PER_TILE) - _wg_barrier() - - next_buf_i32 = arith.remui(arith.addi(cidx, one_i32), nb_const_i32) - p_accs, next_a, next_b = _wmma_rotate(p_accs, ca, cb, next_buf_i32) - cidx1 = arith.addi(cidx, one_i32) - lidx1 = arith.addi(lidx, one_i32) - - new_state = _pack_state(p_accs, next_a, next_b) + [ - lidx1, - cidx1, - ] - results = yield new_state - - accs, cur_a, cur_b = _unpack_state(results[:-2]) - else: - accs = list(accs) - elif const_expr(main_loop_iters > 0): - if const_expr(main_loop_unroll): - # Unroll-by-2: 2 tiles/trip, ping-pong 2 register banks + odd leftover. - load_idx_s = load_idx_init - compute_idx_s = compute_idx_init - num_pairs = main_loop_iters // 2 - - if const_expr(num_pairs > 0): - init_state = _pack_state(accs, cur_a, cur_b) + [ - load_idx_s, - compute_idx_s, - ] - results = init_state - for pair_step, state in range(0, num_pairs, 1, init=init_state): - _disable_unroll_on_enclosing_loop() - cidx = state[-1] - lidx = state[-2] - p_accs, b0a, b0b = _unpack_state(state[:-2]) - - # sub-0: tile T from bank0, prefetch T+1 -> bank1. - p_accs, bank1 = _run_tile( - p_accs, b0a, b0b, cidx, lidx, do_prefetch=True - ) - b1a, b1b = bank1 - cidx1 = arith.addi(cidx, one_i32) - lidx1 = arith.addi(lidx, one_i32) - - # sub-1: tile T+1 from bank1, prefetch T+2 -> bank0 (yielded). - p_accs, bank0 = _run_tile( - p_accs, b1a, b1b, cidx1, lidx1, do_prefetch=True - ) - nb0a, nb0b = bank0 - cidx2 = arith.addi(cidx1, one_i32) - lidx2 = arith.addi(lidx1, one_i32) - - new_state = _pack_state(p_accs, nb0a, nb0b) + [ - lidx2, - cidx2, - ] - results = yield new_state - - accs, cur_a, cur_b = _unpack_state(results[:-2]) - load_idx_s = results[-2] - compute_idx_s = results[-1] - else: - accs = list(accs) - - # Leftover odd tile: compute the carried bank, prefetch first drain tile. - if const_expr(main_loop_iters % 2 == 1): - accs, leftover_bank = _run_tile( - accs, - cur_a, - cur_b, - compute_idx_s, - load_idx_s, - do_prefetch=True, - ) - cur_a, cur_b = leftover_bank - compute_idx_s = arith.addi(compute_idx_s, one_i32) - else: - # Single tile/trip: compute T from carried bank, prefetch T+1. - init_state = _pack_state(accs, cur_a, cur_b) + [ - load_idx_init, - compute_idx_init, - ] - results = init_state - for step, state in range(0, main_loop_iters, 1, init=init_state): - _disable_unroll_on_enclosing_loop() - cidx = state[-1] - lidx = state[-2] - p_accs, ca, cb = _unpack_state(state[:-2]) - - # tile T from the carried bank, prefetch T+1 -> next bank (yielded). - p_accs, bank = _run_tile( - p_accs, ca, cb, cidx, lidx, do_prefetch=True - ) - next_a, next_b = bank - cidx1 = arith.addi(cidx, one_i32) - lidx1 = arith.addi(lidx, one_i32) - - new_state = _pack_state(p_accs, next_a, next_b) + [ - lidx1, - cidx1, - ] - results = yield new_state - - accs, cur_a, cur_b = _unpack_state(results[:-2]) + TILES_PER_TRIP = 2 if variant == "bandwidth_bound" and main_loop_unroll else 1 + num_trips = main_loop_iters // TILES_PER_TRIP + load_idx_s, compute_idx_s = load_idx_init, compute_idx_init + + if const_expr(num_trips > 0): + init_state = _pack_state(accs, cur_a, cur_b) + [ + load_idx_init, + compute_idx_init, + ] + results = init_state + for trip, state in range(0, num_trips, 1, init=init_state): + cidx, lidx = state[-1], state[-2] + p_accs, ca, cb = _unpack_state(state[:-2]) + for _sub in range_constexpr(TILES_PER_TRIP): + p_accs, ca, cb = _run_tile(p_accs, ca, cb, cidx, lidx) + cidx = arith.addi(cidx, one_i32) + lidx = arith.addi(lidx, one_i32) + results = yield _pack_state(p_accs, ca, cb) + [lidx, cidx] + + accs, cur_a, cur_b = _unpack_state(results[:-2]) + load_idx_s, compute_idx_s = results[-2], results[-1] else: accs = list(accs) - # No main loop ran — drain starts at compute_idx = 0. + + if const_expr(main_loop_iters % TILES_PER_TRIP == 1): + accs, cur_a, cur_b = _run_tile( + accs, cur_a, cur_b, compute_idx_s, load_idx_s + ) # ── Drain (fully unrolled): consume carried frags, prefetch next bank per tile; final tile does no wait/barrier/ds_load ── drain_count_d = ( @@ -1042,61 +654,27 @@ def _run_tile(accs_, ca, cb, cidx, lidx, do_prefetch): ) drain_base = max(0, main_loop_iters) - def _drain_wait_for(tile_idx): - # Wait for compile-time drain tile tile_idx to land. - tdm_ops.tensor_wait(max(0, num_k_tiles - 1 - tile_idx) * _TDMS_PER_TILE) - accs = list(accs) drain_a, drain_b = cur_a, cur_b for j in range_constexpr(drain_count_d): tile_idx = drain_base + j if const_expr(j < drain_count_d - 1): - # Not the last tile: wait + barrier + prefetch tile_idx+1, interleaved with the current WMMA. next_tile = tile_idx + 1 - _drain_wait_for(next_tile) + tdm_ops.tensor_wait( + max(0, num_k_tiles - 1 - next_tile) * _TDMS_PER_TILE + ) gpu.barrier() - next_a = load_a_frags(next_tile % num_buffers) - accs = wmma_tile(accs, drain_a, drain_b) - next_b = load_b_frags(next_tile % num_buffers) + next_a = load_frags(A_SIDE, next_tile % num_buffers) + accs, _, _ = wmma_tile(accs, drain_a, drain_b) + next_b = load_frags(B_SIDE, next_tile % num_buffers) drain_a, drain_b = next_a, next_b else: - # Final K-tile: compute the carried frags — no wait, barrier, or ds_loads. - accs = wmma_tile(accs, drain_a, drain_b) + accs, _, _ = wmma_tile(accs, drain_a, drain_b) - # ── Epilogue: bias, activation, and store ── if const_expr(num_buffers > 2): rocdl.sched_barrier(0) - accs = _apply_bias_and_activation(accs) epilogue_stores(accs) - cache_tag = ( - in_dtype, - out_dtype, - K, - tile_m, - tile_n, - tile_k, - m_warp, - n_warp, - num_buffers, - effective_waves_per_eu, - l2_prefetch_distance, - activation, - add_bias, - physical_mk, - physical_kn, - M, - N, - loop_carried_load_percent, - kernarg_preload, - use_manual_barrier, - split_k, - sched_strategy, - barrier_signal_wait_latency, - main_loop_unroll, - variant, - ) - @flyc.jit def launch_gemm_a16w16( arg_y: fx.Tensor, @@ -1107,94 +685,37 @@ def launch_gemm_a16w16( i32_n: fx.Int32, stream: fx.Stream, ): - _ = cache_tag ctx = CompilationContext.get_current() - # Grid dimensions idx_m = arith.index_cast(T.index, i32_m.ir_value()) idx_n = arith.index_cast(T.index, i32_n.ir_value()) gx = _raw((idx_m + arith.index(tile_m - 1)) / arith.index(tile_m)) gy = _raw((idx_n + arith.index(tile_n - 1)) / arith.index(tile_n)) - gz = split_k # 1 at default -> grid tuple identical to original - # Emit kernel launcher = kernel_gemm_a16w16(arg_y, arg_x, arg_w, arg_bias, i32_m, i32_n) - # Set waves_per_eu - for op in ctx.gpu_module_body.operations: - if ( - hasattr(op, "attributes") - and op.OPERATION_NAME == "gpu.func" - and effective_waves_per_eu is not None - ): - _wpe = int(effective_waves_per_eu) - if _wpe >= 1: - op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( - ir.IntegerType.get_signless(32), _wpe - ) - - # Flat work-group size hint. flat_wg_attr = ir.StringAttr.get(f"{block_threads},{block_threads}") + wpe = int(effective_waves_per_eu or 0) + for op in ctx.gpu_module_body.operations: if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr - - # Experimental: amdgpu-loop-carried-load-percent passthrough. - if loop_carried_load_percent is not None: - lcv = ir.ArrayAttr.get( - [ - ir.ArrayAttr.get( - [ - ir.StringAttr.get("amdgpu-loop-carried-load-percent"), - ir.StringAttr.get(str(int(loop_carried_load_percent))), - ] + if wpe >= 1: + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + ir.IntegerType.get_signless(32), wpe ) - ] - ) - for op in ctx.gpu_module_body.operations: - if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": - op.attributes["passthrough"] = lcv - - # Mark kernel args inreg so AMDGPU preloads them into user SGPRs. - if kernarg_preload: - inreg_attr = ir.UnitAttr.get() - for op in ctx.gpu_module_body.operations: - if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": - num_args = len(op.regions[0].blocks[0].arguments) - per_arg = [ - ir.DictAttr.get({"llvm.inreg": inreg_attr}) - for _ in range(num_args) - ] - op.attributes["arg_attrs"] = ir.ArrayAttr.get(per_arg) + if kernarg_preload: + inreg = ir.DictAttr.get({"llvm.inreg": ir.UnitAttr.get()}) + n_args = len(op.regions[0].blocks[0].arguments) + op.attributes["arg_attrs"] = ir.ArrayAttr.get([inreg] * n_args) launcher.launch( - grid=(gx, gy, gz), - block=(block_threads, 1, 1), - stream=stream, + grid=(gx, gy, split_k), block=(block_threads, 1, 1), stream=stream ) - # Commented out until coexec branch is merged. - # launch_gemm_a8w8_blockscale.compile_hints["llvm_options"] = { - # "amdgpu-expert-scheduling-mode": True, - # "amdgpu-anti-hints-for-va-vdst": True, - # "amdgpu-enable-static-simulator": True, - # "amdgpu-static-sim-inline": True, - # "amdgpu-sched-strategy": "coexec", - # # "amdgpu-block-carried-latency": EnumOpt("all"), # enable per the note above - # } - - # Backend cl::opt hints; only options present in THIS build are set. - _llvm_opts = { - "amdgpu-expert-scheduling-mode": True, # valid bool cl::opt (GFX12+ only) - "unroll-threshold": 0, - } - # Experimental scheduling levers (in-process LLVM); None => cl::init default. + _llvm_opts = {"amdgpu-expert-scheduling-mode": True, "unroll-threshold": 0} if sched_strategy is not None: - # cl::opt: max-ilp (GCNMaxILP) or max-memory-clause. _llvm_opts["amdgpu-sched-strategy"] = sched_strategy - if barrier_signal_wait_latency is not None: - # cl::opt init(35): synthetic barrier signal->wait latency. - _llvm_opts["amdgpu-barrier-signal-wait-latency"] = barrier_signal_wait_latency launch_gemm_a16w16.compile_hints["llvm_options"] = _llvm_opts return launch_gemm_a16w16 @@ -1216,40 +737,21 @@ def gemm_a16w16( waves_per_eu: int | None = None, l2_prefetch_distance: int = 2, kernarg_preload: bool = False, - loop_carried_load_percent: int | None = None, - use_manual_barrier: bool = False, split_k: int = 1, sched_strategy: str | None = None, - barrier_signal_wait_latency: int | None = None, main_loop_unroll: bool = False, variant: str = "bandwidth_bound", ): """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" - assert x.dtype in ( - torch.float16, - torch.bfloat16, - ), f"x must be fp16/bf16, got {x.dtype}" - assert w.dtype in ( - torch.float16, - torch.bfloat16, - ), f"w must be fp16/bf16, got {w.dtype}" + _half = (torch.float16, torch.bfloat16) + assert x.dtype in _half, f"x must be fp16/bf16, got {x.dtype}" + assert w.dtype in _half, f"w must be fp16/bf16, got {w.dtype}" assert x.shape[1] == w.shape[1], "Incompatible K dimensions" - M, K = x.shape - N = w.shape[0] - - # Layout from strides - if x.stride(1) == 1: - physical_mk = True - elif x.stride(0) == 1: - physical_mk = False + M, K, N = x.shape[0], x.shape[1], w.shape[0] + physical_mk = x.stride(1) == 1 + physical_kn = w.stride(1) != 1 - if w.stride(1) == 1: - physical_kn = False - elif w.stride(0) == 1: - physical_kn = True - - # K-pad up to a tile_k multiple K_padded = ((K + tile_k - 1) // tile_k) * tile_k if K_padded != K: pad_size = K_padded - K @@ -1267,26 +769,14 @@ def gemm_a16w16( w = torch.nn.functional.pad(w, (0, pad_size)) K = K_padded - # N-pad up to a tile_n multiple N_stride = ((N + tile_n - 1) // tile_n) * tile_n - - # split_k>1 half-out: accumulate in fp32 workspace, cast at end (bf16 atomic lossy). - _half_dtype = dtype in (torch.float16, torch.bfloat16) - _splitk_f32_accum = split_k > 1 and _half_dtype - + _splitk_f32_accum = split_k > 1 and dtype in _half in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" - if dtype == torch.float16: - out_dtype_str = "f16" - elif dtype == torch.bfloat16: - out_dtype_str = "bf16" - else: - out_dtype_str = "f32" + out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get(dtype, "f32") buf_dtype = dtype if _splitk_f32_accum: - out_dtype_str = "f32" # kernel accumulates fp32 partials - buf_dtype = torch.float32 # fp32 workspace, cast to `dtype` after launch + out_dtype_str, buf_dtype = "f32", torch.float32 - # split_k>1 atomic-fadd needs a zeroed buffer; fp32-accum uses a fresh workspace. _alloc = torch.zeros if split_k > 1 else torch.empty if y is not None and not _splitk_f32_accum: y_buf = ( @@ -1325,32 +815,21 @@ def gemm_a16w16( physical_mk=physical_mk, physical_kn=physical_kn, kernarg_preload=kernarg_preload, - loop_carried_load_percent=loop_carried_load_percent, - use_manual_barrier=use_manual_barrier, split_k=split_k, sched_strategy=sched_strategy, - barrier_signal_wait_latency=barrier_signal_wait_latency, main_loop_unroll=main_loop_unroll, variant=variant, ) launch_fn(y_buf, x, w, bias, M, N_stride, stream=torch.cuda.current_stream()) + result = y_buf[:, :N] if N_stride != N else y_buf if _splitk_f32_accum: - # Cast the fp32 accumulation workspace back to the requested half dtype. - result = (y_buf[:, :N] if N_stride != N else y_buf).to(dtype) - if y is not None: - y.copy_(result) - return y - return result - - if N_stride != N: - result = y_buf[:, :N] - if y is not None: - y.copy_(result) - return y + result = result.to(dtype) + if y is None or result is y: return result - return y_buf + y.copy_(result) + return y __all__ = ["compile_gemm_a16w16", "gemm_a16w16"] From 3fa14fe41ad8d5a8024a3c8740f9971777944458 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 30 Jul 2026 16:39:11 +0000 Subject: [PATCH 12/34] Reverting UT changes that existed for AM --- .../gemm/basic/test_gemm_a16w16.py | 232 ++++++++---------- 1 file changed, 101 insertions(+), 131 deletions(-) diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py index 12c67f0312c..b6a7df7215b 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py +++ b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py @@ -1,33 +1,15 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -import pytest import torch import torch.nn.functional as F - -from aiter.ops.triton.gemm.basic.gemm_a16w16 import _is_gluon_available, gemm_a16w16 +import pytest +from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 from aiter.ops.triton.gemm.basic.gemm_a16w16_atomic import gemm_a16w16_atomic from op_tests.triton_tests.utils.types import str_to_torch_dtype -def is_gluon_supported(): - """gluon a16w16 kernels are only available on supported archs (gfx1250).""" - return _is_gluon_available() - - -def _skip_if_triton_on_gfx1250(backend): - """gfx1250 only ships gluon-format a16w16 configs, so the triton backend has - no usable config there; skip the triton backend on gfx1250.""" - if backend != "triton": - return - from aiter.ops.triton.utils._triton.arch_info import get_arch - - if "gfx1250" in (get_arch() or ""): - pytest.skip("triton backend has no gfx1250 a16w16 config (gluon-only arch)") - - def generate_gemm_a16w16_inputs(M, N, K, dtype, layout="TN", output=True, bias=False): - torch.manual_seed(0) if isinstance(dtype, str): dtype = str_to_torch_dtype[dtype] @@ -57,88 +39,47 @@ def generate_gemm_a16w16_inputs(M, N, K, dtype, layout="TN", output=True, bias=F def get_x_vals(): - x_vals = [ - (1, 1, 1), - (1, 16, 16), - (16, 1, 16), - (16, 16, 1), - # Irregular shapes (masking & OOB) - (3, 5, 7), - (17, 33, 65), - (63, 127, 255), - (65, 129, 257), - (64, 64, 64), - (128, 128, 128), - # Multiple blocks - (128, 256, 512), - (256, 512, 256), - # Asymmetric shapes - (32, 256, 128), - (256, 32, 128), - (128, 128, 1024), - (1024, 128, 128), - (1536, 512, 768), - ] - return x_vals - - -# Test plain BF16 GEMMs - the most common types. -@pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("backend", ["triton", "gluon"]) -@pytest.mark.parametrize("kernel_type", ["bandwidth_bound", "compute_bound"]) -def test_gemm_a16_w16(M: int, N: int, K: int, backend, kernel_type): - if backend == "triton" and kernel_type != "bandwidth_bound": - pytest.skip("kernel_type only applies to the gluon backend") - if backend == "gluon" and not is_gluon_supported(): - pytest.skip("Gluon not supported on this architecture") - _skip_if_triton_on_gfx1250(backend) - - x, w, _, _out_dtype, _y = generate_gemm_a16w16_inputs( - M, - N, - K, - dtype=torch.bfloat16, - output=False, - ) - - torch_out = F.linear(x, w, bias=None) - - triton_out = gemm_a16w16(x, w, backend=backend, kernel_type=kernel_type) - - torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) - - -# Smaller set for testing activations, setting the output tensor and dtype -def get_fewer_x_vals(): - x_vals = [ - (64, 64, 64), - (128, 256, 512), - (256, 512, 256), - (128, 128, 1024), - (1024, 128, 128), - (1536, 512, 768), + x_vals = [(1, 1, 1)] # minimal case + x_vals += [(3, 5, 2)] # irregular shape + x_vals += [(1024 * v, 1024 * v, 1024 * v) for v in range(1, 9)] + x_vals += [(4864, 4096, 8192), (9728, 8192, 65536), (4864, 8192, 4160)] + x_vals += [(2**i, 256, 7168) for i in range(5, 9)] + x_vals += [ + (1, 1280, 8192), + (32, 1280, 8192), + (64, 1280, 8192), + (128, 1280, 8192), + (192, 1280, 8192), + (256, 1280, 8192), + (320, 1280, 8192), + (512, 1280, 8192), + (1024, 1280, 8192), + (2048, 1280, 8192), + (4096, 1280, 8192), + (8192, 1280, 8192), + (16384, 1280, 8192), + (1, 8192, 1024), + (32, 8192, 1024), + (64, 8192, 1024), + (128, 8192, 1024), + (192, 8192, 1024), + (256, 8192, 1024), + (320, 8192, 1024), + (512, 8192, 1024), + (1024, 8192, 1024), + (2048, 8192, 1024), + (4096, 8192, 1024), + (8192, 8192, 1024), + (16384, 8192, 1024), ] return x_vals -# A smaller set of shapes that tests fused activations, different dtypes -# and output tensor arg. We don't want the larger set above to test -# all these combinations. @pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu", "silu_exp2"]) -@pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) +@pytest.mark.parametrize("M, N, K", get_x_vals()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("output", [True, False]) -@pytest.mark.parametrize("backend", ["triton", "gluon"]) -@pytest.mark.parametrize("kernel_type", ["bandwidth_bound", "compute_bound"]) -def test_gemm_a16_w16_activation( - M: int, N: int, K: int, dtype, output, activation, backend, kernel_type -): - if backend == "triton" and kernel_type != "bandwidth_bound": - pytest.skip("kernel_type only applies to the gluon backend") - if backend == "gluon" and not is_gluon_supported(): - pytest.skip("Gluon not supported on this architecture") - _skip_if_triton_on_gfx1250(backend) - +def test_gemm_a16_w16_activation(M: int, N: int, K: int, dtype, output, activation): x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( M, N, @@ -152,82 +93,111 @@ def test_gemm_a16_w16_activation( torch_out = F.gelu(torch_out) elif activation == "gelu_tanh": torch_out = F.gelu(torch_out, approximate="tanh") - elif activation in ("silu", "silu_exp2"): + elif activation == "silu": + torch_out = F.silu(torch_out) + elif activation == "silu_exp2": torch_out = F.silu(torch_out) - triton_out = gemm_a16w16( - x, - w, - None, - out_dtype, - y, - activation=activation, - backend=backend, - kernel_type=kernel_type, - ) + if output: + triton_out = gemm_a16w16( + x, + w, + None, + out_dtype, + y, + activation=activation, + ) + else: + triton_out = gemm_a16w16( + x, + w, + None, + out_dtype, + activation=activation, + ) torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) @pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("layout", ["TT", "NN", "NT"]) -@pytest.mark.parametrize("backend", ["triton", "gluon"]) -@pytest.mark.parametrize("kernel_type", ["bandwidth_bound", "compute_bound"]) -def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout, backend, kernel_type): - if backend == "triton" and kernel_type != "bandwidth_bound": - pytest.skip("kernel_type only applies to the gluon backend") - if backend == "gluon" and not is_gluon_supported(): - pytest.skip("Gluon not supported on this architecture") - _skip_if_triton_on_gfx1250(backend) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("output", [True, False]) +def test_gemm_a16_w16(M: int, N: int, K: int, dtype, output): + torch.cuda.empty_cache() # Helps avoid hangs in large tests + x, w, bias, out_dtype, y = generate_gemm_a16w16_inputs( + M, N, K, dtype, output=output, bias=True + ) + + torch_out = F.linear(x, w, bias=bias) + + if output: + triton_out = gemm_a16w16(x, w, bias, out_dtype, y) + else: + triton_out = gemm_a16w16(x, w, bias, out_dtype) + + torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) + + +@pytest.mark.parametrize("M, N, K", get_x_vals()) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("layout", ["TT", "NN", "NT"]) +@pytest.mark.parametrize("output", [True, False]) +def test_gemm_a16_w16_layout(M: int, N: int, K: int, dtype, layout, output): torch.cuda.empty_cache() # Helps avoid hangs in large tests x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( - M, N, K, torch.bfloat16, layout=layout, output=False + M, N, K, dtype, layout=layout, output=output ) torch_out = F.linear(x, w, bias=None) - triton_out = gemm_a16w16( - x, w, None, out_dtype, y, backend=backend, kernel_type=kernel_type - ) + if output: + triton_out = gemm_a16w16(x, w, None, out_dtype, y) + else: + triton_out = gemm_a16w16(x, w, None, out_dtype) torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) @pytest.mark.parametrize("M, N, K", get_x_vals()) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16_atomic(M: int, N: int, K: int, output): +def test_gemm_a16_w16_atomic(M: int, N: int, K: int, dtype, output): torch.cuda.empty_cache() # Helps avoid hangs in large tests - x, w, _, _out_dtype, y = generate_gemm_a16w16_inputs( - M, N, K, torch.bfloat16, output=output - ) + x, w, _, out_dtype, y = generate_gemm_a16w16_inputs(M, N, K, dtype, output=output) torch_out = F.linear(x, w, bias=None) # Accumulation in bf16/fp16 leads to precision loss, cast y to fp32 to prevent that if output: y = y.to(torch.float32).zero_() - triton_out = gemm_a16w16_atomic(x, w, torch.float32, y).to(torch.bfloat16) + triton_out = gemm_a16w16_atomic(x, w, torch.float32, y).to(dtype) else: - triton_out = gemm_a16w16_atomic(x, w, dtype=torch.float32).to(torch.bfloat16) + triton_out = gemm_a16w16_atomic(x, w, dtype=torch.float32).to(dtype) torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) -@pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) +@pytest.mark.parametrize("M, N, K", get_x_vals()) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("layout", ["TT", "NN", "NT"]) -def test_gemm_a16_w16_atomic_layout(M: int, N: int, K: int, layout): +@pytest.mark.parametrize("output", [True, False]) +def test_gemm_a16_w16_atomic_layout(M: int, N: int, K: int, dtype, layout, output): torch.cuda.empty_cache() # Helps avoid hangs in large tests - x, w, _, _out_dtype, y = generate_gemm_a16w16_inputs( - M, N, K, torch.bfloat16, layout=layout, output=True + x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( + M, N, K, dtype, layout=layout, output=output ) torch_out = F.linear(x, w, bias=None) - y = y.to(torch.float32).zero_() - triton_out = gemm_a16w16_atomic(x, w, torch.float32, y).to(torch.bfloat16) + # Accumulation in bf16/fp16 leads to precision loss, cast y to fp32 to prevent that + if output: + y = y.to(torch.float32).zero_() + triton_out = gemm_a16w16_atomic(x, w, torch.float32, y).to(dtype) + else: + triton_out = gemm_a16w16_atomic(x, w, dtype=torch.float32).to(dtype) torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) From 5f4d6decc94f5155014860729e7a5f04357212a5 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 30 Jul 2026 16:43:44 +0000 Subject: [PATCH 13/34] Revert prev commit --- .../gemm/basic/test_gemm_a16w16.py | 211 ++++++++++-------- 1 file changed, 112 insertions(+), 99 deletions(-) diff --git a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py index b6a7df7215b..d98487772ec 100644 --- a/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py +++ b/op_tests/triton_tests/gemm/basic/test_gemm_a16w16.py @@ -1,15 +1,33 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import pytest import torch import torch.nn.functional as F -import pytest -from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 + +from aiter.ops.triton.gemm.basic.gemm_a16w16 import _is_gluon_available, gemm_a16w16 from aiter.ops.triton.gemm.basic.gemm_a16w16_atomic import gemm_a16w16_atomic from op_tests.triton_tests.utils.types import str_to_torch_dtype +def is_gluon_supported(): + """gluon a16w16 kernels are only available on supported archs (gfx1250).""" + return _is_gluon_available() + + +def _skip_if_triton_on_gfx1250(backend): + """gfx1250 only ships gluon-format a16w16 configs, so the triton backend has + no usable config there; skip the triton backend on gfx1250.""" + if backend != "triton": + return + from aiter.ops.triton.utils._triton.arch_info import get_arch + + if "gfx1250" in (get_arch() or ""): + pytest.skip("triton backend has no gfx1250 a16w16 config (gluon-only arch)") + + def generate_gemm_a16w16_inputs(M, N, K, dtype, layout="TN", output=True, bias=False): + torch.manual_seed(0) if isinstance(dtype, str): dtype = str_to_torch_dtype[dtype] @@ -41,45 +59,69 @@ def generate_gemm_a16w16_inputs(M, N, K, dtype, layout="TN", output=True, bias=F def get_x_vals(): x_vals = [(1, 1, 1)] # minimal case x_vals += [(3, 5, 2)] # irregular shape - x_vals += [(1024 * v, 1024 * v, 1024 * v) for v in range(1, 9)] - x_vals += [(4864, 4096, 8192), (9728, 8192, 65536), (4864, 8192, 4160)] - x_vals += [(2**i, 256, 7168) for i in range(5, 9)] - x_vals += [ - (1, 1280, 8192), - (32, 1280, 8192), - (64, 1280, 8192), - (128, 1280, 8192), - (192, 1280, 8192), - (256, 1280, 8192), - (320, 1280, 8192), - (512, 1280, 8192), - (1024, 1280, 8192), - (2048, 1280, 8192), - (4096, 1280, 8192), - (8192, 1280, 8192), - (16384, 1280, 8192), - (1, 8192, 1024), - (32, 8192, 1024), - (64, 8192, 1024), - (128, 8192, 1024), - (192, 8192, 1024), - (256, 8192, 1024), - (320, 8192, 1024), - (512, 8192, 1024), - (1024, 8192, 1024), - (2048, 8192, 1024), - (4096, 8192, 1024), - (8192, 8192, 1024), - (16384, 8192, 1024), - ] + x_vals += [(1024 * v, 1024 * v, 1024 * v) for v in (1, 2, 4, 5, 8)] + x_vals += [(2**i, 256, 7168) for i in range(5, 9)] # DSR1 router GEMM + # GPT-OSS-120B attention projections + x_vals += [(2**i, 5120, 2880) for i in range(5, 9)] # GPTOSS QKV input projection + x_vals += [(2**i, 2880, 4096) for i in range(5, 9)] # output projection + x_vals += [(2**i, 128, 2880) for i in range(5, 9)] # Router GEMM + x_vals += [(v, 106496, 16384) for v in (256, 4096)] # LL3 405B FC1 return x_vals -@pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu", "silu_exp2"]) +# Test plain BF16 GEMMs - the most common types. @pytest.mark.parametrize("M, N, K", get_x_vals()) +@pytest.mark.parametrize("backend", ["triton", "gluon"]) +@pytest.mark.parametrize("kernel_type", ["bandwidth_bound", "compute_bound"]) +def test_gemm_a16_w16(M: int, N: int, K: int, backend, kernel_type): + if backend == "triton" and kernel_type != "bandwidth_bound": + pytest.skip("kernel_type only applies to the gluon backend") + if backend == "gluon" and not is_gluon_supported(): + pytest.skip("Gluon not supported on this architecture") + _skip_if_triton_on_gfx1250(backend) + + x, w, _, _out_dtype, _y = generate_gemm_a16w16_inputs( + M, + N, + K, + dtype=torch.bfloat16, + output=False, + ) + + torch_out = F.linear(x, w, bias=None) + + triton_out = gemm_a16w16(x, w, backend=backend, kernel_type=kernel_type) + + torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) + + +# Smaller set for testing activations, setting the output tensor and dtype +def get_fewer_x_vals(): + x_vals = [(16, 1024, 1024)] + x_vals += [(128, 8192, 512)] + x_vals += [(256, 512, 8192)] + x_vals += [(1024 * v, 1024 * v, 1024 * v) for v in (1, 5, 8)] + return x_vals + + +# A smaller set of shapes that tests fused activations, different dtypes +# and output tensor arg. We don't want the larger set above to test +# all these combinations. +@pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu"]) +@pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16_activation(M: int, N: int, K: int, dtype, output, activation): +@pytest.mark.parametrize("backend", ["triton", "gluon"]) +@pytest.mark.parametrize("kernel_type", ["bandwidth_bound", "compute_bound"]) +def test_gemm_a16_w16_activation( + M: int, N: int, K: int, dtype, output, activation, backend, kernel_type +): + if backend == "triton" and kernel_type != "bandwidth_bound": + pytest.skip("kernel_type only applies to the gluon backend") + if backend == "gluon" and not is_gluon_supported(): + pytest.skip("Gluon not supported on this architecture") + _skip_if_triton_on_gfx1250(backend) + x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( M, N, @@ -95,109 +137,80 @@ def test_gemm_a16_w16_activation(M: int, N: int, K: int, dtype, output, activati torch_out = F.gelu(torch_out, approximate="tanh") elif activation == "silu": torch_out = F.silu(torch_out) - elif activation == "silu_exp2": - torch_out = F.silu(torch_out) - - if output: - triton_out = gemm_a16w16( - x, - w, - None, - out_dtype, - y, - activation=activation, - ) - else: - triton_out = gemm_a16w16( - x, - w, - None, - out_dtype, - activation=activation, - ) - - torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) - -@pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16(M: int, N: int, K: int, dtype, output): - torch.cuda.empty_cache() # Helps avoid hangs in large tests - - x, w, bias, out_dtype, y = generate_gemm_a16w16_inputs( - M, N, K, dtype, output=output, bias=True + triton_out = gemm_a16w16( + x, + w, + None, + out_dtype, + y, + activation=activation, + backend=backend, + kernel_type=kernel_type, ) - torch_out = F.linear(x, w, bias=bias) - - if output: - triton_out = gemm_a16w16(x, w, bias, out_dtype, y) - else: - triton_out = gemm_a16w16(x, w, bias, out_dtype) - - torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) + torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-2) @pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("layout", ["TT", "NN", "NT"]) -@pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16_layout(M: int, N: int, K: int, dtype, layout, output): +@pytest.mark.parametrize("backend", ["triton", "gluon"]) +@pytest.mark.parametrize("kernel_type", ["bandwidth_bound", "compute_bound"]) +def test_gemm_a16_w16_layout(M: int, N: int, K: int, layout, backend, kernel_type): + if backend == "triton" and kernel_type != "bandwidth_bound": + pytest.skip("kernel_type only applies to the gluon backend") + if backend == "gluon" and not is_gluon_supported(): + pytest.skip("Gluon not supported on this architecture") + _skip_if_triton_on_gfx1250(backend) + torch.cuda.empty_cache() # Helps avoid hangs in large tests x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( - M, N, K, dtype, layout=layout, output=output + M, N, K, torch.bfloat16, layout=layout, output=False ) torch_out = F.linear(x, w, bias=None) - if output: - triton_out = gemm_a16w16(x, w, None, out_dtype, y) - else: - triton_out = gemm_a16w16(x, w, None, out_dtype) + triton_out = gemm_a16w16( + x, w, None, out_dtype, y, backend=backend, kernel_type=kernel_type + ) torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) @pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16_atomic(M: int, N: int, K: int, dtype, output): +def test_gemm_a16_w16_atomic(M: int, N: int, K: int, output): torch.cuda.empty_cache() # Helps avoid hangs in large tests - x, w, _, out_dtype, y = generate_gemm_a16w16_inputs(M, N, K, dtype, output=output) + x, w, _, _out_dtype, y = generate_gemm_a16w16_inputs( + M, N, K, torch.bfloat16, output=output + ) torch_out = F.linear(x, w, bias=None) # Accumulation in bf16/fp16 leads to precision loss, cast y to fp32 to prevent that if output: y = y.to(torch.float32).zero_() - triton_out = gemm_a16w16_atomic(x, w, torch.float32, y).to(dtype) + triton_out = gemm_a16w16_atomic(x, w, torch.float32, y).to(torch.bfloat16) else: - triton_out = gemm_a16w16_atomic(x, w, dtype=torch.float32).to(dtype) + triton_out = gemm_a16w16_atomic(x, w, dtype=torch.float32).to(torch.bfloat16) torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) -@pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) @pytest.mark.parametrize("layout", ["TT", "NN", "NT"]) -@pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16_atomic_layout(M: int, N: int, K: int, dtype, layout, output): +def test_gemm_a16_w16_atomic_layout(M: int, N: int, K: int, layout): torch.cuda.empty_cache() # Helps avoid hangs in large tests - x, w, _, out_dtype, y = generate_gemm_a16w16_inputs( - M, N, K, dtype, layout=layout, output=output + x, w, _, _out_dtype, y = generate_gemm_a16w16_inputs( + M, N, K, torch.bfloat16, layout=layout, output=True ) torch_out = F.linear(x, w, bias=None) - # Accumulation in bf16/fp16 leads to precision loss, cast y to fp32 to prevent that - if output: - y = y.to(torch.float32).zero_() - triton_out = gemm_a16w16_atomic(x, w, torch.float32, y).to(dtype) - else: - triton_out = gemm_a16w16_atomic(x, w, dtype=torch.float32).to(dtype) + y = y.to(torch.float32).zero_() + triton_out = gemm_a16w16_atomic(x, w, torch.float32, y).to(torch.bfloat16) torch.testing.assert_close(triton_out, torch_out, atol=1e-1, rtol=1e-1) From 0cdcc43382f8b883c70ad99d8d23b212a26ea63b Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 30 Jul 2026 18:40:24 +0000 Subject: [PATCH 14/34] small refactor --- .../ops/flydsl/kernels/gemm_a16w16_gfx1250.py | 59 ++++--------------- 1 file changed, 11 insertions(+), 48 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py index eda173563d5..327c67b8100 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py @@ -22,6 +22,8 @@ from flydsl.runtime.device import get_rocm_arch as get_hip_arch from flydsl.utils.smem_allocator import check_smem_capacity +from aiter.ops.flydsl.kernels import buffer_ops + WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 WAVE_SIZE, LDS_PAD_A, LDS_PAD_B = 32, 8, 8 _SCHED_ALLOW_SALU = 1 << 2 @@ -31,49 +33,6 @@ def _align_up(n, a): return ((n + a - 1) // a) * a -def _i32_const(v): - return _raw(arith.constant(int(v), type=T.i32)) - - -def _to_i32_offset(offset): - if isinstance(offset, int): - return _i32_const(offset) - offset = arith.unwrap(offset) - if not isinstance(offset.type, ir.IntegerType) or offset.type.width != 32: - return _raw(arith.index_cast(T.i32, offset)) - return offset - - -def _buffer_load(rsrc, offset, vec_width=1, dtype=None): - if dtype is None: - dtype = T.f32 - off = _to_i32_offset(offset) - off = _raw(arith.muli(off, _i32_const(dtype.width // 8))) - result_type = dtype if vec_width == 1 else T.vec(vec_width, dtype) - return rocdl.RawPtrBufferLoadOp( - result_type, rsrc, off, _i32_const(0), _i32_const(0) - ).result - - -def _buffer_store(data, rsrc, offset, offset_is_bytes=False): - data = arith.unwrap(data) - off = _to_i32_offset(offset) - if not offset_is_bytes: - data_type = data.type - element_type = ( - data_type.element_type if hasattr(data_type, "element_type") else data_type - ) - off = _raw(arith.muli(off, _i32_const(element_type.width // 8))) - rocdl.RawPtrBufferStoreOp(data, rsrc, off, _i32_const(0), _i32_const(0)) - - -def _make_buffer_rsrc(tensor, num_records_bytes=None): - buf_tensor = fx.rocdl.make_buffer_tensor( - tensor, max_size=True, num_records_bytes=num_records_bytes - ) - return fx.rocdl.get_buffer_rsrc(fx.get_iter(buf_tensor)) - - def apply_activation_scalar(val, activation: str): import math @@ -254,9 +213,13 @@ def kernel_gemm_a16w16( m_idx = arith.index_cast(T.index, i32_m.ir_value()) n_stride = arith.index_cast(T.index, i32_n.ir_value()) y_nrec = m_idx * n_stride * arith.index(elem_bytes_d) - y_rsrc = _make_buffer_rsrc(arg_y, num_records_bytes=y_nrec) + y_rsrc = fx.rocdl.get_buffer_rsrc( + fx.rocdl.make_buffer_ptr(fx.get_iter(arg_y), y_nrec) + ) if const_expr(add_bias): - bias_rsrc = _make_buffer_rsrc(arg_bias) + bias_rsrc = fx.rocdl.get_buffer_rsrc( + fx.rocdl.make_buffer_ptr(fx.get_iter(arg_bias)) + ) lds_base_ptr = fx.SharedAllocator(static=True).allocate(lds_arena_bytes)._ptr big_a_mem = fx.recast_iter(_fx_elem, fx.add_offset(lds_base_ptr, unified_a_off)) @@ -486,7 +449,7 @@ def epilogue_stores(final_accs): ) elems = [] for half in range_constexpr(2): - bv = _buffer_load( + bv = buffer_ops.buffer_load( bias_rsrc, col_i32 + arith.constant(half * 4, type=T.i32), vec_width=4, @@ -545,7 +508,7 @@ def epilogue_stores(final_accs): zero_i32, ) else: - _buffer_store( + buffer_ops.buffer_store( h_vec.bitcast(fx.Int32).ir_value(), y_rsrc, c_off_bytes, @@ -574,7 +537,7 @@ def epilogue_stores(final_accs): fx.Float32, ).ir_value() col = col_base + arith.index(half * 4) - _buffer_store(vec4, y_rsrc, row * n_stride + col) + buffer_ops.buffer_store(vec4, y_rsrc, row * n_stride + col) def _pack_state(accs_, a_, b_): return list(accs_) + list(a_) + list(b_) From 7ab8dc3991cc32e8d594d9435183b64ad51b91c1 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 30 Jul 2026 19:26:32 +0000 Subject: [PATCH 15/34] Splitting up kernel / wrapper --- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 142 ++++++++++++++++++ ...x1250.py => gemm_a16w16_kernel_gfx1250.py} | 114 +------------- op_tests/flydsl_tests/test_gemm_a16w16.py | 3 +- 3 files changed, 144 insertions(+), 115 deletions(-) create mode 100644 aiter/ops/flydsl/gemm_a16w16_gfx1250.py rename aiter/ops/flydsl/kernels/{gemm_a16w16_gfx1250.py => gemm_a16w16_kernel_gfx1250.py} (88%) diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py new file mode 100644 index 00000000000..f4608b0b57b --- /dev/null +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""gfx1250 FlyDSL backend for the A16W16 GEMM: torch-facing wrapper. + +The compiled kernel lives in ``kernels/gemm_a16w16_kernel_gfx1250.py``; this +module owns layout detection, padding, output allocation and dtype mapping. +""" + +from __future__ import annotations + +import torch + +_compile_gemm_a16w16 = None + + +def _lazy_import(): + """Defer the kernel import so ``import aiter`` does not pull in flydsl.""" + global _compile_gemm_a16w16 + if _compile_gemm_a16w16 is not None: + return + # Absolute (not relative) so this module also works when loaded by file path. + from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( + compile_gemm_a16w16, + ) + + _compile_gemm_a16w16 = compile_gemm_a16w16 + + +def gemm_a16w16( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + dtype: torch.dtype = torch.float16, + y: torch.Tensor | None = None, + activation: str | None = None, + tile_m: int = 128, + tile_n: int = 128, + tile_k: int = 32, + m_warp: int = 2, + n_warp: int = 4, + num_buffers: int = 2, + waves_per_eu: int | None = None, + l2_prefetch_distance: int = 2, + kernarg_preload: bool = False, + split_k: int = 1, + sched_strategy: str | None = None, + main_loop_unroll: bool = False, + variant: str = "bandwidth_bound", +): + """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" + _lazy_import() + _half = (torch.float16, torch.bfloat16) + assert x.dtype in _half, f"x must be fp16/bf16, got {x.dtype}" + assert w.dtype in _half, f"w must be fp16/bf16, got {w.dtype}" + assert x.shape[1] == w.shape[1], "Incompatible K dimensions" + + M, K, N = x.shape[0], x.shape[1], w.shape[0] + physical_mk = x.stride(1) == 1 + physical_kn = w.stride(1) != 1 + + K_padded = ((K + tile_k - 1) // tile_k) * tile_k + if K_padded != K: + pad_size = K_padded - K + if physical_mk: + x = torch.nn.functional.pad(x, (0, pad_size)) + else: + x = torch.nn.functional.pad(x.T, (0, 0, 0, pad_size)).T + + if physical_kn: + if w.stride(1) == 1: + w = torch.nn.functional.pad(w, (0, 0, 0, pad_size)) + else: + w = torch.nn.functional.pad(w.T, (0, 0, 0, pad_size)).T + else: + w = torch.nn.functional.pad(w, (0, pad_size)) + K = K_padded + + N_stride = ((N + tile_n - 1) // tile_n) * tile_n + _splitk_f32_accum = split_k > 1 and dtype in _half + in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" + out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get(dtype, "f32") + buf_dtype = dtype + if _splitk_f32_accum: + out_dtype_str, buf_dtype = "f32", torch.float32 + + _alloc = torch.zeros if split_k > 1 else torch.empty + if y is not None and not _splitk_f32_accum: + y_buf = ( + y + if N_stride == N + else _alloc((M, N_stride), device=x.device, dtype=buf_dtype) + ) + if split_k > 1 and y_buf is y: + y_buf.zero_() + else: + y_buf = ( + _alloc((M, N_stride), device=x.device, dtype=buf_dtype) + if N_stride != N + else _alloc((M, N), device=x.device, dtype=buf_dtype) + ) + + if bias is None: + bias = torch.empty(0, device=x.device, dtype=dtype) + + launch_fn = _compile_gemm_a16w16( + M=M if not physical_mk else 0, + N=N, + K=K, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + m_warp=m_warp, + n_warp=n_warp, + in_dtype=in_dtype_str, + out_dtype=out_dtype_str, + num_buffers=num_buffers, + waves_per_eu=waves_per_eu, + l2_prefetch_distance=l2_prefetch_distance, + activation=activation, + add_bias=(bias.numel() > 0), + physical_mk=physical_mk, + physical_kn=physical_kn, + kernarg_preload=kernarg_preload, + split_k=split_k, + sched_strategy=sched_strategy, + main_loop_unroll=main_loop_unroll, + variant=variant, + ) + + launch_fn(y_buf, x, w, bias, M, N_stride, stream=torch.cuda.current_stream()) + + result = y_buf[:, :N] if N_stride != N else y_buf + if _splitk_f32_accum: + result = result.to(dtype) + if y is None or result is y: + return result + y.copy_(result) + return y + + +__all__ = ["gemm_a16w16"] diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py similarity index 88% rename from aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py rename to aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index 327c67b8100..7f91e5c25a4 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -3,7 +3,6 @@ import flydsl.compiler as flyc import flydsl.expr as fx -import torch from flydsl._mlir import ir from flydsl._mlir.dialects import llvm as llvm_dialect from flydsl._mlir.dialects import rocdl as raw_rocdl @@ -684,115 +683,4 @@ def launch_gemm_a16w16( return launch_gemm_a16w16 -def gemm_a16w16( - x: torch.Tensor, - w: torch.Tensor, - bias: torch.Tensor | None = None, - dtype: torch.dtype = torch.float16, - y: torch.Tensor | None = None, - activation: str | None = None, - tile_m: int = 128, - tile_n: int = 128, - tile_k: int = 32, - m_warp: int = 2, - n_warp: int = 4, - num_buffers: int = 2, - waves_per_eu: int | None = None, - l2_prefetch_distance: int = 2, - kernarg_preload: bool = False, - split_k: int = 1, - sched_strategy: str | None = None, - main_loop_unroll: bool = False, - variant: str = "bandwidth_bound", -): - """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" - _half = (torch.float16, torch.bfloat16) - assert x.dtype in _half, f"x must be fp16/bf16, got {x.dtype}" - assert w.dtype in _half, f"w must be fp16/bf16, got {w.dtype}" - assert x.shape[1] == w.shape[1], "Incompatible K dimensions" - - M, K, N = x.shape[0], x.shape[1], w.shape[0] - physical_mk = x.stride(1) == 1 - physical_kn = w.stride(1) != 1 - - K_padded = ((K + tile_k - 1) // tile_k) * tile_k - if K_padded != K: - pad_size = K_padded - K - if physical_mk: - x = torch.nn.functional.pad(x, (0, pad_size)) - else: - x = torch.nn.functional.pad(x.T, (0, 0, 0, pad_size)).T - - if physical_kn: - if w.stride(1) == 1: - w = torch.nn.functional.pad(w, (0, 0, 0, pad_size)) - else: - w = torch.nn.functional.pad(w.T, (0, 0, 0, pad_size)).T - else: - w = torch.nn.functional.pad(w, (0, pad_size)) - K = K_padded - - N_stride = ((N + tile_n - 1) // tile_n) * tile_n - _splitk_f32_accum = split_k > 1 and dtype in _half - in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" - out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get(dtype, "f32") - buf_dtype = dtype - if _splitk_f32_accum: - out_dtype_str, buf_dtype = "f32", torch.float32 - - _alloc = torch.zeros if split_k > 1 else torch.empty - if y is not None and not _splitk_f32_accum: - y_buf = ( - y - if N_stride == N - else _alloc((M, N_stride), device=x.device, dtype=buf_dtype) - ) - if split_k > 1 and y_buf is y: - y_buf.zero_() - else: - y_buf = ( - _alloc((M, N_stride), device=x.device, dtype=buf_dtype) - if N_stride != N - else _alloc((M, N), device=x.device, dtype=buf_dtype) - ) - - if bias is None: - bias = torch.empty(0, device=x.device, dtype=dtype) - - launch_fn = compile_gemm_a16w16( - M=M if not physical_mk else 0, - N=N, - K=K, - tile_m=tile_m, - tile_n=tile_n, - tile_k=tile_k, - m_warp=m_warp, - n_warp=n_warp, - in_dtype=in_dtype_str, - out_dtype=out_dtype_str, - num_buffers=num_buffers, - waves_per_eu=waves_per_eu, - l2_prefetch_distance=l2_prefetch_distance, - activation=activation, - add_bias=(bias.numel() > 0), - physical_mk=physical_mk, - physical_kn=physical_kn, - kernarg_preload=kernarg_preload, - split_k=split_k, - sched_strategy=sched_strategy, - main_loop_unroll=main_loop_unroll, - variant=variant, - ) - - launch_fn(y_buf, x, w, bias, M, N_stride, stream=torch.cuda.current_stream()) - - result = y_buf[:, :N] if N_stride != N else y_buf - if _splitk_f32_accum: - result = result.to(dtype) - if y is None or result is y: - return result - y.copy_(result) - return y - - -__all__ = ["compile_gemm_a16w16", "gemm_a16w16"] +__all__ = ["compile_gemm_a16w16"] diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index a67409acbdc..3c3d9c3fd69 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -16,7 +16,6 @@ "aiter", "ops", "flydsl", - "kernels", "gemm_a16w16_gfx1250.py", ) ) @@ -44,7 +43,7 @@ def _flydsl_available(): def _load_kernel(): """Load the kernel module by file path to skip aiter.ops.flydsl.__init__.""" - spec = importlib.util.spec_from_file_location("_flydsl_a16w16_kernel", _KERNEL_PATH) + spec = importlib.util.spec_from_file_location("_flydsl_a16w16_wrapper", _KERNEL_PATH) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod.gemm_a16w16 From d8e2bdb1f82564e7a15ea3b862f5a606b9dfa9be Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 30 Jul 2026 19:33:13 +0000 Subject: [PATCH 16/34] black/ruff --- op_tests/flydsl_tests/test_gemm_a16w16.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index 3c3d9c3fd69..5a18a31f48e 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -43,7 +43,9 @@ def _flydsl_available(): def _load_kernel(): """Load the kernel module by file path to skip aiter.ops.flydsl.__init__.""" - spec = importlib.util.spec_from_file_location("_flydsl_a16w16_wrapper", _KERNEL_PATH) + spec = importlib.util.spec_from_file_location( + "_flydsl_a16w16_wrapper", _KERNEL_PATH + ) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod.gemm_a16w16 From 6c0fdf18e2ad2c7b70609e2afdfc72a323fcbc6c Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 31 Jul 2026 16:58:29 +0000 Subject: [PATCH 17/34] refactor with claude skill --- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 94 ++++-- .../kernels/gemm_a16w16_kernel_gfx1250.py | 286 +++++++----------- 2 files changed, 173 insertions(+), 207 deletions(-) diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index f4608b0b57b..cb89b2c3048 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -2,29 +2,64 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. """gfx1250 FlyDSL backend for the A16W16 GEMM: torch-facing wrapper. - -The compiled kernel lives in ``kernels/gemm_a16w16_kernel_gfx1250.py``; this -module owns layout detection, padding, output allocation and dtype mapping. """ from __future__ import annotations +import functools + import torch _compile_gemm_a16w16 = None +_run_compiled = None +_fx = None def _lazy_import(): - """Defer the kernel import so ``import aiter`` does not pull in flydsl.""" - global _compile_gemm_a16w16 + global _compile_gemm_a16w16, _run_compiled, _fx if _compile_gemm_a16w16 is not None: return # Absolute (not relative) so this module also works when loaded by file path. + import flydsl.expr as fx + from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( compile_gemm_a16w16, ) + from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled as run_compiled _compile_gemm_a16w16 = compile_gemm_a16w16 + _run_compiled = run_compiled + _fx = fx + + +_CFG_KEYS = ( + "M", + "N", + "K", + "tile_m", + "tile_n", + "tile_k", + "m_warp", + "n_warp", + "in_dtype", + "out_dtype", + "num_buffers", + "waves_per_eu", + "activation", + "add_bias", + "physical_mk", + "physical_kn", + "kernarg_preload", + "split_k", + "sched_strategy", + "main_loop_unroll", + "variant", +) + + +@functools.lru_cache(maxsize=1024) +def _cached_launcher(*cfg): + return _compile_gemm_a16w16(**dict(zip(_CFG_KEYS, cfg))) def gemm_a16w16( @@ -41,7 +76,6 @@ def gemm_a16w16( n_warp: int = 4, num_buffers: int = 2, waves_per_eu: int | None = None, - l2_prefetch_distance: int = 2, kernarg_preload: bool = False, split_k: int = 1, sched_strategy: str | None = None, @@ -103,32 +137,32 @@ def gemm_a16w16( if bias is None: bias = torch.empty(0, device=x.device, dtype=dtype) - launch_fn = _compile_gemm_a16w16( - M=M if not physical_mk else 0, - N=N, - K=K, - tile_m=tile_m, - tile_n=tile_n, - tile_k=tile_k, - m_warp=m_warp, - n_warp=n_warp, - in_dtype=in_dtype_str, - out_dtype=out_dtype_str, - num_buffers=num_buffers, - waves_per_eu=waves_per_eu, - l2_prefetch_distance=l2_prefetch_distance, - activation=activation, - add_bias=(bias.numel() > 0), - physical_mk=physical_mk, - physical_kn=physical_kn, - kernarg_preload=kernarg_preload, - split_k=split_k, - sched_strategy=sched_strategy, - main_loop_unroll=main_loop_unroll, - variant=variant, + launch_fn = _cached_launcher( + M if not physical_mk else 0, + N, + K, + tile_m, + tile_n, + tile_k, + m_warp, + n_warp, + in_dtype_str, + out_dtype_str, + num_buffers, + waves_per_eu, + activation, + bias.numel() > 0, + physical_mk, + physical_kn, + kernarg_preload, + split_k, + sched_strategy, + main_loop_unroll, + variant, ) - launch_fn(y_buf, x, w, bias, M, N_stride, stream=torch.cuda.current_stream()) + stream = torch.cuda.current_stream(device=x.device).cuda_stream + _run_compiled(launch_fn, y_buf, x, w, bias, M, N_stride, _fx.Stream(stream)) result = y_buf[:, :N] if N_stride != N else y_buf if _splitk_f32_accum: diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index 7f91e5c25a4..fc414930ed4 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -1,58 +1,45 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import math + import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir import ir from flydsl._mlir.dialects import llvm as llvm_dialect -from flydsl._mlir.dialects import rocdl as raw_rocdl from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import ( - arith, - const_expr, - gpu, - idx2crd, - range_constexpr, - rocdl, - tdm_ops, -) -from flydsl.expr.arith import _to_raw as _raw +from flydsl.compiler.protocol import dsl_size_of +from flydsl.expr import const_expr, gpu, idx2crd, range_constexpr, rocdl, tdm_ops from flydsl.expr.typing import T from flydsl.runtime.device import get_rocm_arch as get_hip_arch from flydsl.utils.smem_allocator import check_smem_capacity -from aiter.ops.flydsl.kernels import buffer_ops - WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 WAVE_SIZE, LDS_PAD_A, LDS_PAD_B = 32, 8, 8 _SCHED_ALLOW_SALU = 1 << 2 -def _align_up(n, a): - return ((n + a - 1) // a) * a +def _byte_off_i32(elem_off, elem_bytes): + return (fx.Int32(elem_off) * elem_bytes).ir_value() -def apply_activation_scalar(val, activation: str): - import math +def _f32(x): + return fx.Float32(fx.Float32(x).ir_value()) - from flydsl._mlir.dialects import math as _math +def apply_activation_scalar(val, activation: str): if activation == "relu": - zero = arith.constant(0.0, type=T.f32) + zero = _f32(0.0) return (val > zero).select(val, zero) - one, half = arith.constant(1.0, type=T.f32), arith.constant(0.5, type=T.f32) + one, half = _f32(1.0), _f32(0.5) if activation in ("silu", "silu_exp2"): - return val / (one + _math.exp(arith.constant(0.0, type=T.f32) - val)) + return val / (one + fx.exp(_f32(0.0) - val)) if activation == "gelu": - scaled = val * arith.constant(1.0 / math.sqrt(2.0), type=T.f32) - return half * val * (one + _math.erf(scaled)) + return half * val * (one + fx.erf(val * (1.0 / math.sqrt(2.0)))) if activation == "gelu_tanh": - two = arith.constant(2.0, type=T.f32) - inner = arith.constant(math.sqrt(2.0 / math.pi), type=T.f32) * ( - val + arith.constant(0.044715, type=T.f32) * val * val * val - ) - # tanh(z) = 1 - 2/(1 + exp(2z)) - tanh_val = one - two / (one + _math.exp(two * inner)) + two = _f32(2.0) + inner = math.sqrt(2.0 / math.pi) * (val + 0.044715 * val * val * val) + tanh_val = one - two / (one + fx.exp(two * inner)) return half * val * (one + tanh_val) return val @@ -71,19 +58,16 @@ def compile_gemm_a16w16( out_dtype: str | None = None, num_buffers: int = 2, waves_per_eu: int | None = None, - l2_prefetch_distance: int = 2, activation: str | None = None, add_bias: bool = False, physical_mk: bool = True, # True=M-major (row-major X), False=K-major (col-major X) - physical_kn: bool = False, # False=N-major (row-major W), True=K-major (transposed W) + physical_kn: bool = False, # False=N-major (row-major W), True=K-major (W^T) kernarg_preload: bool = False, split_k: int = 1, sched_strategy: str | None = None, main_loop_unroll: bool = False, variant: str = "bandwidth_bound", ): - _ = (M, N) - def _req(cond, msg): if not cond: raise ValueError(msg) @@ -91,7 +75,6 @@ def _req(cond, msg): if out_dtype is None: out_dtype = "f16" if in_dtype == "fp16" else "bf16" is_f16 = in_dtype == "fp16" - effective_waves_per_eu = waves_per_eu _req(2 <= num_buffers <= 8, f"num_buffers must be in [2, 8], got {num_buffers}") _req(variant in ("bandwidth_bound", "compute_bound"), f"bad variant {variant!r}") @@ -143,8 +126,7 @@ def _req(cond, msg): k_wmma_steps = tile_k // WMMA_K _fx_elem = fx.Float16 if is_f16 else fx.BFloat16 wmma_m_rep, wmma_n_rep = warp_tile_m // WMMA_M, warp_tile_n // WMMA_N - n_accs = wmma_m_rep * wmma_n_rep - N_ACCS = n_accs + N_ACCS = n_accs = wmma_m_rep * wmma_n_rep N_A_FRAGS, N_B_FRAGS = wmma_m_rep * k_wmma_steps, wmma_n_rep * k_wmma_steps if physical_mk: @@ -165,14 +147,12 @@ def _req(cond, msg): b_imm_bytes, lds_b_stride = tile_k * elem_bytes, tile_k + LDS_PAD_B lds_b_elems = tile_n * lds_b_stride + LDS_PAD_B - lds_a_data_bytes, lds_b_data_bytes = ( - lds_a_elems * elem_bytes, - lds_b_elems * elem_bytes, - ) - unified_a_off = 0 - unified_b_off = _align_up(unified_a_off + num_buffers * lds_a_data_bytes, 16) - lds_arena_bytes = _align_up(unified_b_off + num_buffers * lds_b_data_bytes, 16) - check_smem_capacity(lds_arena_bytes, gpu_arch) + @fx.struct + class SharedStorage: + a: fx.Array[_fx_elem, num_buffers * lds_a_elems, 16] + b: fx.Array[_fx_elem, num_buffers * lds_b_elems, 16] + + check_smem_capacity(dsl_size_of(SharedStorage), gpu_arch) _TDMS_PER_TILE = 2 # A + B @@ -186,52 +166,46 @@ def kernel_gemm_a16w16( i32_n: fx.Int32, ): rocdl.disable_xdl_arb_stall() - elem_ty = T.f16 if is_f16 else T.bf16 + elem_ty = _fx_elem.ir_type tx, bx, by = gpu.thread_id("x"), gpu.block_id("x"), gpu.block_id("y") - blk_m, blk_n = bx * arith.index(tile_m), by * arith.index(tile_n) + blk_m, blk_n = fx.Uint64(bx) * tile_m, fx.Uint64(by) * tile_n if const_expr(split_k > 1): - bz = gpu.block_id("z") - split_k_base = bz * arith.index(split_k_chunk) + split_k_base = fx.Uint64(gpu.block_id("z")) * split_k_chunk else: - split_k_base = arith.index(0) + split_k_base = fx.Uint64(0) layout_thr = fx.make_layout( (m_warp, n_warp, 2, 16), (n_warp * WAVE_SIZE, WAVE_SIZE, 16, 1) ) thr_coord = idx2crd(tx, layout_thr) wave_m_idx, wave_n_idx, lane_kgrp, lane16 = ( - fx.get(thr_coord, 0), - fx.get(thr_coord, 1), - fx.get(thr_coord, 2), - fx.get(thr_coord, 3), + fx.Uint64(fx.get(thr_coord, 0)), + fx.Uint64(fx.get(thr_coord, 1)), + fx.Uint64(fx.get(thr_coord, 2)), + fx.Uint64(fx.get(thr_coord, 3)), ) - warp_m_base = wave_m_idx * arith.index(warp_tile_m) - warp_n_base = wave_n_idx * arith.index(warp_tile_n) - m_idx = arith.index_cast(T.index, i32_m.ir_value()) - n_stride = arith.index_cast(T.index, i32_n.ir_value()) - y_nrec = m_idx * n_stride * arith.index(elem_bytes_d) + warp_m_base, warp_n_base = wave_m_idx * warp_tile_m, wave_n_idx * warp_tile_n + m_idx, n_stride = fx.Uint64(i32_m), fx.Uint64(i32_n) y_rsrc = fx.rocdl.get_buffer_rsrc( - fx.rocdl.make_buffer_ptr(fx.get_iter(arg_y), y_nrec) + fx.rocdl.make_buffer_ptr( + fx.get_iter(arg_y), m_idx * n_stride * elem_bytes_d + ) ) if const_expr(add_bias): bias_rsrc = fx.rocdl.get_buffer_rsrc( fx.rocdl.make_buffer_ptr(fx.get_iter(arg_bias)) ) - lds_base_ptr = fx.SharedAllocator(static=True).allocate(lds_arena_bytes)._ptr - big_a_mem = fx.recast_iter(_fx_elem, fx.add_offset(lds_base_ptr, unified_a_off)) - big_b_mem = fx.recast_iter(_fx_elem, fx.add_offset(lds_base_ptr, unified_b_off)) - big_a_base_idx = arith.index_cast(T.index, arith.unwrap(fx.ptrtoint(big_a_mem))) - big_b_base_idx = arith.index_cast(T.index, arith.unwrap(fx.ptrtoint(big_b_mem))) - slot_stride_a_elems_i32 = arith.constant(lds_a_elems, type=T.i32) - slot_stride_b_elems_i32 = arith.constant(lds_b_elems, type=T.i32) + lds = fx.SharedAllocator(static=True).allocate(SharedStorage).peek() + big_a_mem, big_b_mem = lds.a.ptr, lds.b.ptr + big_a_base_idx = fx.Uint64(fx.ptrtoint(big_a_mem)) + big_b_base_idx = fx.Uint64(fx.ptrtoint(big_b_mem)) blk_m64, blk_n64 = fx.Int64(blk_m), fx.Int64(blk_n) split_k_base64 = fx.Int64(split_k_base) def _mk_side(tensor, base_off, shape, outer_stride, pad_interval, pad_amount): - """Global tile view + its TDM atom (identical shape for A and B).""" gt = fx.Tensor( fx.make_view( fx.add_offset(fx.get_iter(tensor), base_off), @@ -241,23 +215,23 @@ def _mk_side(tensor, base_off, shape, outer_stride, pad_interval, pad_amount): return gt, fx.rocdl.make_tdm_atom( gt, [None, None], - strides=[fx.Int64(outer_stride), None], + strides=[outer_stride, None], num_warps=num_warps, pad_interval=pad_interval, pad_amount=pad_amount, ) if const_expr(physical_mk): - a_base_off = blk_m64 * fx.Int64(K) + split_k_base64 + a_base_off = blk_m64 * K + split_k_base64 else: - a_base_off = split_k_base64 * fx.Int64(M) + blk_m64 + a_base_off = split_k_base64 * M + blk_m64 gA, atom_a = _mk_side( arg_x, a_base_off, a_tile_shape, a_outer_stride, a_pad_interval, LDS_PAD_A ) if const_expr(physical_kn): - b_base_off = split_k_base64 * fx.Int64(N) + blk_n64 + b_base_off = split_k_base64 * N + blk_n64 else: - b_base_off = blk_n64 * fx.Int64(K) + split_k_base64 + b_base_off = blk_n64 * K + split_k_base64 gB, atom_b = _mk_side( arg_w, b_base_off, b_tile_shape, b_outer_stride, b_pad_interval, LDS_PAD_B ) @@ -265,37 +239,30 @@ def _mk_side(tensor, base_off, shape, outer_stride, pad_interval, pad_amount): def _imm64(k_tile, mul): if const_expr(isinstance(k_tile, int)): return fx.Int64(k_tile * mul) - as_index = arith.index_cast( - T.index, arith.muli(k_tile, arith.constant(mul, type=T.i32)) - ) - return fx.Int64(arith.index_cast(T.i64, as_index)) + return fx.Int64(fx.Int32(k_tile) * mul) # mul in i32, widen for TDM - def _slot_off(buf_idx, slot_elems, stride_i32): + def _slot_off(buf_idx, slot_elems): if const_expr(isinstance(buf_idx, int)): - return arith.index(buf_idx * slot_elems) - return arith.index_cast(T.index, arith.muli(buf_idx, stride_i32)) + return fx.Uint64(buf_idx * slot_elems) + return fx.Uint64(fx.Int32(buf_idx) * slot_elems) # mul in i32, widen def _lane_bases(warp_base, reps, lds_stride, transpose): if const_expr(not transpose): - row_off = (warp_base + lane16) * arith.index(lds_stride) - k_off = lane_kgrp * arith.index(8) + row_off = (warp_base + lane16) * lds_stride + k_off = lane_kgrp * 8 return [ - row_off + arith.index(rep * WMMA_M * lds_stride) + k_off + row_off + rep * WMMA_M * lds_stride + k_off for rep in range_constexpr(reps) ] - k_off = ( - lane_kgrp * arith.index(8) + lane16 % arith.index(8) - ) * arith.index(lds_stride) - grp_off = (lane16 / arith.index(8)) * arith.index(8) + k_off = (lane_kgrp * 8 + lane16 % 8) * lds_stride + grp_off = (lane16 // 8) * 8 return [ - k_off + warp_base + arith.index(rep * WMMA_M) + grp_off + k_off + warp_base + rep * WMMA_M + grp_off for rep in range_constexpr(reps) ] # One spec per operand side; A and B differ only in these fields. - MEM, BASE, BASES, STRIDE, TR, ELEMS, SSTRIDE, REPS, ATOM, GT, SHAPE, IMM = ( - range(12) - ) + MEM, BASE, BASES, STRIDE, TR, ELEMS, REPS, ATOM, GT, SHAPE, IMM = range(11) A_SIDE = ( big_a_mem, big_a_base_idx, @@ -303,7 +270,6 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): lds_a_stride, not physical_mk, lds_a_elems, - slot_stride_a_elems_i32, wmma_m_rep, atom_a, gA, @@ -317,7 +283,6 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): lds_b_stride, physical_kn, lds_b_elems, - slot_stride_b_elems_i32, wmma_n_rep, atom_b, gB, @@ -326,7 +291,6 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): ) def issue_tdm_loads(buf_idx, k_tile): - """A then B TDM for K-tile k_tile into LDS ring slot buf_idx.""" for sd in (A_SIDE, B_SIDE): fx.copy( sd[ATOM], @@ -335,7 +299,7 @@ def issue_tdm_loads(buf_idx, k_tile): fx.make_view( fx.add_offset( sd[MEM], - _slot_off(buf_idx, sd[ELEMS], sd[SSTRIDE]), + _slot_off(buf_idx, sd[ELEMS]), ), fx.make_layout(sd[SHAPE], (sd[STRIDE], 1)), ) @@ -351,36 +315,20 @@ def _frag(sd, slot_off, ks, rep): off = ( sd[BASES][rep] + slot_off - + arith.index((ks * WMMA_K + k_half * 16) * sd[STRIDE]) - ) - addr = arith.index_cast( - T.i32, sd[BASE] + off * arith.index(elem_bytes) - ) - ptr = llvm_dialect.inttoptr( - ir.Type.parse("!llvm.ptr<3>"), arith.unwrap(addr) + + (ks * WMMA_K + k_half * 16) * sd[STRIDE] ) - halves.append(raw_rocdl.ds_load_tr16_b128(vec8_ty, ptr)) + addr = fx.Int32(sd[BASE] + off * elem_bytes).ir_value() + ptr = llvm_dialect.inttoptr(ir.Type.parse("!llvm.ptr<3>"), addr) + halves.append(fx.Vector(rocdl.ds_load_tr16_b128(vec8_ty, ptr))) else: - off = ( - sd[BASES][rep] - + slot_off - + arith.index(ks * WMMA_K + k_half * 16) - ) + off = sd[BASES][rep] + slot_off + (ks * WMMA_K + k_half * 16) halves.append( - arith.unwrap( - fx.ptr_load( - fx.add_offset(sd[MEM], off), result_type=vec8_ty - ) - ) + fx.ptr_load(fx.add_offset(sd[MEM], off), result_type=vec8_ty) ) - return ( - fx.Vector(halves[0]) - .shuffle(fx.Vector(halves[1]), list(range(16))) - .ir_value() - ) + return halves[0].shuffle(halves[1], list(range(16))) def load_frags(sd, buf_idx): - off = _slot_off(buf_idx, sd[ELEMS], sd[SSTRIDE]) + off = _slot_off(buf_idx, sd[ELEMS]) return [ _frag(sd, off, ks, rep) for ks in range_constexpr(k_wmma_steps) @@ -393,7 +341,7 @@ def load_frags(sd, buf_idx): def _rmem_vec(v, n, dtype): t = fx.make_rmem_tensor(n, dtype) - t.store(fx.Vector(arith.unwrap(v))) + t.store(v if isinstance(v, fx.Vector) else fx.Vector(v)) return t def wmma_tile(accs_in, a_frags, b_frags, rotate_buf=None): @@ -404,8 +352,8 @@ def wmma_tile(accs_in, a_frags, b_frags, rotate_buf=None): next_a = [None] * N_A_FRAGS if rotate else None next_b = [None] * N_B_FRAGS if rotate else None if const_expr(rotate): - a_off = _slot_off(rotate_buf, lds_a_elems, slot_stride_a_elems_i32) - b_off = _slot_off(rotate_buf, lds_b_elems, slot_stride_b_elems_i32) + a_off = _slot_off(rotate_buf, lds_a_elems) + b_off = _slot_off(rotate_buf, lds_b_elems) for ks in range_constexpr(k_wmma_steps): for wm in range_constexpr(wmma_m_rep): a_f = a_ts[ks * wmma_m_rep + wm] @@ -425,37 +373,32 @@ def wmma_tile(accs_in, a_frags, b_frags, rotate_buf=None): rocdl.sched_barrier(_SCHED_ALLOW_SALU) for wn in range_constexpr(wmma_n_rep): next_b[ks * wmma_n_rep + wn] = _frag(B_SIDE, b_off, ks, wn) - return [arith.unwrap(t.load()) for t in acc_ts], next_a, next_b + return [t.load() for t in acc_ts], next_a, next_b _half_out = out_dtype in ("f16", "bf16") _out_num = ( fx.Float16 if out_dtype == "f16" else fx.BFloat16 if _half_out else None ) - _bias_elem = elem_ty def epilogue_stores(final_accs): + zero_off = fx.Int32(0).ir_value() if const_expr(split_k > 1): - zero_i32 = fx.Int32(0) + zero_i32 = zero_off if const_expr(add_bias): bias_vecs = [] for wn in range_constexpr(wmma_n_rep): - col_i32 = arith.index_cast( - T.i32, - blk_n - + warp_n_base - + arith.index(wn * WMMA_N) - + lane_kgrp * arith.index(8), - ) + col_i32 = blk_n + warp_n_base + wn * WMMA_N + lane_kgrp * 8 elems = [] for half in range_constexpr(2): - bv = buffer_ops.buffer_load( + bv = rocdl.raw_ptr_buffer_load( + T.vec(4, elem_ty), bias_rsrc, - col_i32 + arith.constant(half * 4, type=T.i32), - vec_width=4, - dtype=_bias_elem, + _byte_off_i32(col_i32 + half * 4, elem_bytes), + zero_off, + zero_off, ) for i in range_constexpr(4): - elems.append(fx.Vector(bv)[i].to(fx.Float32).ir_value()) + elems.append(fx.Vector(bv)[i].to(fx.Float32)) bias_vecs.append( fx.Vector.from_elements(elems, fx.Float32).ir_value() ) @@ -463,13 +406,8 @@ def epilogue_stores(final_accs): for wm in range_constexpr(wmma_m_rep): for wn in range_constexpr(wmma_n_rep): acc = final_accs[wm * wmma_n_rep + wn] - row = blk_m + warp_m_base + arith.index(wm * WMMA_M) + lane16 - col_base = ( - blk_n - + warp_n_base - + arith.index(wn * WMMA_N) - + lane_kgrp * arith.index(8) - ) + row = blk_m + warp_m_base + wm * WMMA_M + lane16 + col_base = blk_n + warp_n_base + wn * WMMA_N + lane_kgrp * 8 if const_expr(add_bias): acc = acc + bias_vecs[wn] if const_expr(activation is not None): @@ -485,9 +423,7 @@ def epilogue_stores(final_accs): if const_expr(_half_out): h_vec = fx.Vector(acc).to(_out_num) - c_off_bytes = (row * n_stride + col_base) * arith.index( - elem_bytes_d - ) + c_off_bytes = (row * n_stride + col_base) * elem_bytes_d if const_expr(split_k > 1): for pair in range_constexpr(4): pair_vec = fx.Vector.from_elements( @@ -500,29 +436,24 @@ def epilogue_stores(final_accs): rocdl.raw_ptr_buffer_atomic_fadd( pair_vec, y_rsrc, - arith.index_cast( - T.i32, c_off_bytes + arith.index(pair * 4) - ), + fx.Int32(c_off_bytes + pair * 4).ir_value(), zero_i32, zero_i32, ) else: - buffer_ops.buffer_store( + rocdl.raw_ptr_buffer_store( h_vec.bitcast(fx.Int32).ir_value(), y_rsrc, - c_off_bytes, - offset_is_bytes=True, + fx.Int32(c_off_bytes).ir_value(), + zero_off, + zero_off, ) elif const_expr(split_k > 1): for e in range_constexpr(8): rocdl.raw_ptr_buffer_atomic_fadd( fx.Vector(acc)[e].ir_value(), y_rsrc, - arith.index_cast( - T.i32, - (row * n_stride + col_base + arith.index(e)) - * arith.index(4), - ), + _byte_off_i32(row * n_stride + col_base + e, 4), zero_i32, zero_i32, ) @@ -535,8 +466,14 @@ def epilogue_stores(final_accs): ], fx.Float32, ).ir_value() - col = col_base + arith.index(half * 4) - buffer_ops.buffer_store(vec4, y_rsrc, row * n_stride + col) + col = col_base + half * 4 + rocdl.raw_ptr_buffer_store( + vec4, + y_rsrc, + _byte_off_i32(row * n_stride + col, 4), + zero_off, + zero_off, + ) def _pack_state(accs_, a_, b_): return list(accs_) + list(a_) + list(b_) @@ -549,12 +486,12 @@ def _unpack_state(state): return out def _run_tile(accs_, ca, cb, cidx, lidx): - issue_tdm_loads(arith.remui(lidx, nb_const_i32), lidx) + issue_tdm_loads(lidx % num_buffers, lidx) tdm_ops.tensor_wait( 0 if num_buffers == 2 else (num_buffers - 2) * _TDMS_PER_TILE ) gpu.barrier() - next_buf_i32 = arith.remui(arith.addi(cidx, one_i32), nb_const_i32) + next_buf_i32 = (cidx + 1) % num_buffers if const_expr(variant == "compute_bound"): return wmma_tile(accs_, ca, cb, rotate_buf=next_buf_i32) na = load_frags(A_SIDE, next_buf_i32) @@ -562,8 +499,7 @@ def _run_tile(accs_, ca, cb, cidx, lidx): accs_, _, _ = wmma_tile(accs_, ca, cb) return accs_, na, nb - # Accumulators - acc_zero = arith.constant_vector(0.0, T.vec(8, T.f32)) + acc_zero = fx.Vector.filled(8, fx.Float32(0.0), fx.Float32) accs = [acc_zero] * n_accs # Prologue: @@ -574,10 +510,7 @@ def _run_tile(accs_, ca, cb, cidx, lidx): cur_a, cur_b = load_frags(A_SIDE, 0), load_frags(B_SIDE, 0) main_loop_iters = num_k_tiles - (num_buffers - 1) - nb_const_i32 = arith.constant(num_buffers, type=T.i32) - one_i32 = arith.constant(1, type=T.i32) - load_idx_init = arith.constant(num_buffers - 1, type=T.i32) - compute_idx_init = arith.constant(0, type=T.i32) + load_idx_init, compute_idx_init = fx.Uint32(num_buffers - 1), fx.Uint32(0) TILES_PER_TRIP = 2 if variant == "bandwidth_bound" and main_loop_unroll else 1 num_trips = main_loop_iters // TILES_PER_TRIP @@ -590,16 +523,17 @@ def _run_tile(accs_, ca, cb, cidx, lidx): ] results = init_state for trip, state in range(0, num_trips, 1, init=init_state): - cidx, lidx = state[-1], state[-2] + cidx, lidx = fx.Uint32(state[-1]), fx.Uint32(state[-2]) p_accs, ca, cb = _unpack_state(state[:-2]) for _sub in range_constexpr(TILES_PER_TRIP): p_accs, ca, cb = _run_tile(p_accs, ca, cb, cidx, lidx) - cidx = arith.addi(cidx, one_i32) - lidx = arith.addi(lidx, one_i32) + cidx, lidx = cidx + 1, lidx + 1 results = yield _pack_state(p_accs, ca, cb) + [lidx, cidx] accs, cur_a, cur_b = _unpack_state(results[:-2]) - load_idx_s, compute_idx_s = results[-2], results[-1] + # loop results come back untyped; re-wrap so % stays unsigned (remui) + load_idx_s = fx.Uint32(results[-2]) + compute_idx_s = fx.Uint32(results[-1]) else: accs = list(accs) @@ -608,7 +542,7 @@ def _run_tile(accs_, ca, cb, cidx, lidx): accs, cur_a, cur_b, compute_idx_s, load_idx_s ) - # ── Drain (fully unrolled): consume carried frags, prefetch next bank per tile; final tile does no wait/barrier/ds_load ── + # Drain drain_count_d = ( num_buffers - 1 if main_loop_iters > 0 @@ -649,15 +583,13 @@ def launch_gemm_a16w16( ): ctx = CompilationContext.get_current() - idx_m = arith.index_cast(T.index, i32_m.ir_value()) - idx_n = arith.index_cast(T.index, i32_n.ir_value()) - gx = _raw((idx_m + arith.index(tile_m - 1)) / arith.index(tile_m)) - gy = _raw((idx_n + arith.index(tile_n - 1)) / arith.index(tile_n)) + gx = (fx.Uint64(i32_m) + (tile_m - 1)) // tile_m + gy = (fx.Uint64(i32_n) + (tile_n - 1)) // tile_n launcher = kernel_gemm_a16w16(arg_y, arg_x, arg_w, arg_bias, i32_m, i32_n) flat_wg_attr = ir.StringAttr.get(f"{block_threads},{block_threads}") - wpe = int(effective_waves_per_eu or 0) + wpe = int(waves_per_eu or 0) for op in ctx.gpu_module_body.operations: if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": From eb941936bc9e4c4d040915e1c606ce8f721c33f4 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 31 Jul 2026 17:01:13 +0000 Subject: [PATCH 18/34] black/ruff --- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index cb89b2c3048..c543331df28 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""gfx1250 FlyDSL backend for the A16W16 GEMM: torch-facing wrapper. -""" +"""gfx1250 FlyDSL backend for the A16W16 GEMM: torch-facing wrapper.""" from __future__ import annotations From d46e4036225e75e8a173e092beeb0009dd07b122 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Wed, 5 Aug 2026 01:00:12 +0000 Subject: [PATCH 19/34] kernel updates based off comments --- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 32 +++- .../kernels/gemm_a16w16_kernel_gfx1250.py | 178 ++++++++++++------ 2 files changed, 147 insertions(+), 63 deletions(-) diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index c543331df28..9927716df73 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -11,14 +11,15 @@ _compile_gemm_a16w16 = None _run_compiled = None +_flyc = None _fx = None def _lazy_import(): - global _compile_gemm_a16w16, _run_compiled, _fx + global _compile_gemm_a16w16, _run_compiled, _flyc, _fx if _compile_gemm_a16w16 is not None: return - # Absolute (not relative) so this module also works when loaded by file path. + import flydsl.compiler as flyc import flydsl.expr as fx from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( @@ -28,9 +29,25 @@ def _lazy_import(): _compile_gemm_a16w16 = compile_gemm_a16w16 _run_compiled = run_compiled + _flyc = flyc _fx = fx +_FX_DTYPE = {} + + +def _p(t): + if not _FX_DTYPE: + _FX_DTYPE.update( + { + torch.bfloat16: _fx.BFloat16, + torch.float16: _fx.Float16, + torch.float32: _fx.Float32, + } + ) + return _flyc.from_c_void_p(_FX_DTYPE[t.dtype], t.data_ptr()) + + _CFG_KEYS = ( "M", "N", @@ -161,7 +178,16 @@ def gemm_a16w16( ) stream = torch.cuda.current_stream(device=x.device).cuda_stream - _run_compiled(launch_fn, y_buf, x, w, bias, M, N_stride, _fx.Stream(stream)) + _run_compiled( + launch_fn, + _p(y_buf), + _p(x), + _p(w), + _p(bias), + M, + N_stride, + _fx.Stream(stream), + ) result = y_buf[:, :N] if N_stride != N else y_buf if _splitk_f32_accum: diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index fc414930ed4..ad9677fa25a 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -7,7 +7,6 @@ import flydsl.expr as fx from flydsl._mlir import ir from flydsl._mlir.dialects import llvm as llvm_dialect -from flydsl.compiler.kernel_function import CompilationContext from flydsl.compiler.protocol import dsl_size_of from flydsl.expr import const_expr, gpu, idx2crd, range_constexpr, rocdl, tdm_ops from flydsl.expr.typing import T @@ -17,6 +16,8 @@ WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 WAVE_SIZE, LDS_PAD_A, LDS_PAD_B = 32, 8, 8 _SCHED_ALLOW_SALU = 1 << 2 +KERNARG_PRELOAD_COUNT = 8 +SCOPE_DEV = 16 def _byte_off_i32(elem_off, elem_bytes): @@ -147,10 +148,22 @@ def _req(cond, msg): b_imm_bytes, lds_b_stride = tile_k * elem_bytes, tile_k + LDS_PAD_B lds_b_elems = tile_n * lds_b_stride + LDS_PAD_B - @fx.struct - class SharedStorage: - a: fx.Array[_fx_elem, num_buffers * lds_a_elems, 16] - b: fx.Array[_fx_elem, num_buffers * lds_b_elems, 16] + _USE_XT = split_k > 1 and out_dtype == "f32" + + if _USE_XT: + + @fx.struct + class SharedStorage: + a: fx.Array[_fx_elem, num_buffers * lds_a_elems, 16] + b: fx.Array[_fx_elem, num_buffers * lds_b_elems, 16] + xt: fx.Array[fx.Float32, num_warps * WMMA_M * WMMA_N, 16] + + else: + + @fx.struct + class SharedStorage: + a: fx.Array[_fx_elem, num_buffers * lds_a_elems, 16] + b: fx.Array[_fx_elem, num_buffers * lds_b_elems, 16] check_smem_capacity(dsl_size_of(SharedStorage), gpu_arch) @@ -158,10 +171,10 @@ class SharedStorage: @flyc.kernel def kernel_gemm_a16w16( - arg_y: fx.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_bias: fx.Tensor, + arg_y: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_bias: fx.Pointer, i32_m: fx.Int32, i32_n: fx.Int32, ): @@ -189,14 +202,8 @@ def kernel_gemm_a16w16( warp_m_base, warp_n_base = wave_m_idx * warp_tile_m, wave_n_idx * warp_tile_n m_idx, n_stride = fx.Uint64(i32_m), fx.Uint64(i32_n) y_rsrc = fx.rocdl.get_buffer_rsrc( - fx.rocdl.make_buffer_ptr( - fx.get_iter(arg_y), m_idx * n_stride * elem_bytes_d - ) + fx.rocdl.make_buffer_ptr(arg_y, m_idx * n_stride * elem_bytes_d) ) - if const_expr(add_bias): - bias_rsrc = fx.rocdl.get_buffer_rsrc( - fx.rocdl.make_buffer_ptr(fx.get_iter(arg_bias)) - ) lds = fx.SharedAllocator(static=True).allocate(SharedStorage).peek() big_a_mem, big_b_mem = lds.a.ptr, lds.b.ptr @@ -205,35 +212,54 @@ def kernel_gemm_a16w16( blk_m64, blk_n64 = fx.Int64(blk_m), fx.Int64(blk_n) split_k_base64 = fx.Int64(split_k_base) - def _mk_side(tensor, base_off, shape, outer_stride, pad_interval, pad_amount): + def _mk_side( + ptr, base_off, shape, outer_stride, pad_interval, pad_amount, extents + ): gt = fx.Tensor( fx.make_view( - fx.add_offset(fx.get_iter(tensor), base_off), + fx.add_offset(ptr, base_off), fx.make_layout(shape, (outer_stride, 1)), ) ) return gt, fx.rocdl.make_tdm_atom( gt, - [None, None], + extents, strides=[outer_stride, None], num_warps=num_warps, pad_interval=pad_interval, pad_amount=pad_amount, ) + m_oob = fx.Int32(m_idx - blk_m) + n_oob = fx.Int32(fx.Uint64(N) - blk_n) if const_expr(N > 0) else None + a_extents = [m_oob, None] if const_expr(physical_mk) else [None, m_oob] + b_extents = [None, n_oob] if const_expr(physical_kn) else [n_oob, None] + if const_expr(physical_mk): a_base_off = blk_m64 * K + split_k_base64 else: a_base_off = split_k_base64 * M + blk_m64 gA, atom_a = _mk_side( - arg_x, a_base_off, a_tile_shape, a_outer_stride, a_pad_interval, LDS_PAD_A + arg_x, + a_base_off, + a_tile_shape, + a_outer_stride, + a_pad_interval, + LDS_PAD_A, + a_extents, ) if const_expr(physical_kn): b_base_off = split_k_base64 * N + blk_n64 else: b_base_off = blk_n64 * K + split_k_base64 gB, atom_b = _mk_side( - arg_w, b_base_off, b_tile_shape, b_outer_stride, b_pad_interval, LDS_PAD_B + arg_w, + b_base_off, + b_tile_shape, + b_outer_stride, + b_pad_interval, + LDS_PAD_B, + b_extents, ) def _imm64(k_tile, mul): @@ -380,25 +406,42 @@ def wmma_tile(accs_in, a_frags, b_frags, rotate_buf=None): fx.Float16 if out_dtype == "f16" else fx.BFloat16 if _half_out else None ) + if const_expr(add_bias): + bias_lay = fx.make_layout(4, 1) + bias_tiles = fx.logical_divide( + fx.rocdl.make_buffer_tensor( + fx.make_view(arg_bias, fx.make_layout(N, 1)), + max_size=False, + num_records_bytes=N * elem_bytes, + ), + bias_lay, + ) + bias_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), _fx_elem) + + def load_bias4(tile_idx): + r = fx.make_rmem_tensor(bias_lay, _fx_elem) + fx.copy_atom_call(bias_atom, fx.slice(bias_tiles, (None, tile_idx)), r) + return r.load() + def epilogue_stores(final_accs): zero_off = fx.Int32(0).ir_value() if const_expr(split_k > 1): zero_i32 = zero_off + dev_scope = fx.Int32(SCOPE_DEV).ir_value() + if const_expr(_USE_XT): + xt = fx.add_offset( + lds.xt.ptr, + (wave_m_idx * n_warp + wave_n_idx) * (WMMA_M * WMMA_N), + ) if const_expr(add_bias): bias_vecs = [] for wn in range_constexpr(wmma_n_rep): - col_i32 = blk_n + warp_n_base + wn * WMMA_N + lane_kgrp * 8 + col_tile = (blk_n + warp_n_base + wn * WMMA_N + lane_kgrp * 8) // 4 elems = [] for half in range_constexpr(2): - bv = rocdl.raw_ptr_buffer_load( - T.vec(4, elem_ty), - bias_rsrc, - _byte_off_i32(col_i32 + half * 4, elem_bytes), - zero_off, - zero_off, - ) + bv = load_bias4(col_tile + half) for i in range_constexpr(4): - elems.append(fx.Vector(bv)[i].to(fx.Float32)) + elems.append(bv[i].to(fx.Float32)) bias_vecs.append( fx.Vector.from_elements(elems, fx.Float32).ir_value() ) @@ -438,7 +481,7 @@ def epilogue_stores(final_accs): y_rsrc, fx.Int32(c_off_bytes + pair * 4).ir_value(), zero_i32, - zero_i32, + dev_scope, ) else: rocdl.raw_ptr_buffer_store( @@ -450,13 +493,28 @@ def epilogue_stores(final_accs): ) elif const_expr(split_k > 1): for e in range_constexpr(8): - rocdl.raw_ptr_buffer_atomic_fadd( - fx.Vector(acc)[e].ir_value(), - y_rsrc, - _byte_off_i32(row * n_stride + col_base + e, 4), - zero_i32, - zero_i32, + fx.ptr_store( + fx.Vector(acc)[e], + fx.add_offset(xt, lane16 * WMMA_N + lane_kgrp * 8 + e), ) + t_base = blk_m + warp_m_base + wm * WMMA_M + t_col = blk_n + warp_n_base + wn * WMMA_N + lane16 + for e in range_constexpr(8): + if t_base + 2 * e < m_idx: + rocdl.raw_ptr_buffer_atomic_fadd( + fx.ptr_load( + fx.add_offset( + xt, (lane_kgrp + 2 * e) * WMMA_N + lane16 + ) + ).ir_value(), + y_rsrc, + _byte_off_i32( + (t_base + lane_kgrp + 2 * e) * n_stride + t_col, + 4, + ), + zero_i32, + dev_scope, + ) else: for half in range_constexpr(2): vec4 = fx.Vector.from_elements( @@ -573,41 +631,41 @@ def _run_tile(accs_, ca, cb, cidx, lidx): @flyc.jit def launch_gemm_a16w16( - arg_y: fx.Tensor, - arg_x: fx.Tensor, - arg_w: fx.Tensor, - arg_bias: fx.Tensor, + arg_y: fx.Pointer, + arg_x: fx.Pointer, + arg_w: fx.Pointer, + arg_bias: fx.Pointer, i32_m: fx.Int32, i32_n: fx.Int32, stream: fx.Stream, ): - ctx = CompilationContext.get_current() - gx = (fx.Uint64(i32_m) + (tile_m - 1)) // tile_m gy = (fx.Uint64(i32_n) + (tile_n - 1)) // tile_n - launcher = kernel_gemm_a16w16(arg_y, arg_x, arg_w, arg_bias, i32_m, i32_n) - - flat_wg_attr = ir.StringAttr.get(f"{block_threads},{block_threads}") - wpe = int(waves_per_eu or 0) - - for op in ctx.gpu_module_body.operations: - if hasattr(op, "attributes") and op.OPERATION_NAME == "gpu.func": - op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr - if wpe >= 1: - op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( - ir.IntegerType.get_signless(32), wpe - ) - if kernarg_preload: - inreg = ir.DictAttr.get({"llvm.inreg": ir.UnitAttr.get()}) - n_args = len(op.regions[0].blocks[0].arguments) - op.attributes["arg_attrs"] = ir.ArrayAttr.get([inreg] * n_args) + wpe = int(waves_per_eu) if waves_per_eu is not None else 0 + launcher = kernel_gemm_a16w16( + arg_y, + arg_x, + arg_w, + arg_bias, + i32_m, + i32_n, + value_attrs={ + "rocdl.flat_work_group_size": f"{block_threads},{block_threads}", + "rocdl.waves_per_eu": wpe if wpe >= 1 else None, + }, + ) launcher.launch( grid=(gx, gy, split_k), block=(block_threads, 1, 1), stream=stream ) - _llvm_opts = {"amdgpu-expert-scheduling-mode": True, "unroll-threshold": 0} + _llvm_opts = { + "amdgpu-expert-scheduling-mode": True, + "unroll-threshold": 0, + "amdgpu-kernarg-preload": kernarg_preload, + "amdgpu-kernarg-preload-count": KERNARG_PRELOAD_COUNT, + } if sched_strategy is not None: _llvm_opts["amdgpu-sched-strategy"] = sched_strategy launch_gemm_a16w16.compile_hints["llvm_options"] = _llvm_opts From 91ae387f604ffa6454f0bd78e4f8e741df525ffe Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Wed, 5 Aug 2026 14:09:29 +0000 Subject: [PATCH 20/34] black/ruff --- aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index ad9677fa25a..842decc1e2f 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -17,7 +17,7 @@ WAVE_SIZE, LDS_PAD_A, LDS_PAD_B = 32, 8, 8 _SCHED_ALLOW_SALU = 1 << 2 KERNARG_PRELOAD_COUNT = 8 -SCOPE_DEV = 16 +SCOPE_DEV = 16 def _byte_off_i32(elem_off, elem_bytes): From da29bb5f2add517530c0a2b92bc2e9be92f799a4 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 6 Aug 2026 21:53:21 -0400 Subject: [PATCH 21/34] Update gemm_a16w16_gfx1250.py --- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 494 +++++++++++++++--------- 1 file changed, 313 insertions(+), 181 deletions(-) diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index 9927716df73..12b8041c606 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -1,201 +1,333 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""gfx1250 FlyDSL backend for the A16W16 GEMM: torch-facing wrapper.""" +import importlib.util +import os -from __future__ import annotations +import pytest +import torch +import torch.nn.functional as F -import functools +_KERNEL_PATH = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "aiter", + "ops", + "flydsl", + "gemm_a16w16_gfx1250.py", + ) +) -import torch -_compile_gemm_a16w16 = None -_run_compiled = None -_flyc = None -_fx = None +def _get_gpu_arch(): + if not torch.cuda.is_available(): + return None + return getattr(torch.cuda.get_device_properties(0), "gcnArchName", None) + +def _flydsl_available(): + if importlib.util.find_spec("flydsl") is None: + return False + arch = _get_gpu_arch() + return arch is not None and arch.startswith("gfx1250") -def _lazy_import(): - global _compile_gemm_a16w16, _run_compiled, _flyc, _fx - if _compile_gemm_a16w16 is not None: - return - import flydsl.compiler as flyc - import flydsl.expr as fx - from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( - compile_gemm_a16w16, +if not _flydsl_available(): + pytest.skip( + "FlyDSL a16w16 tests require gfx1250 and the flydsl package.", + allow_module_level=True, ) - from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled as run_compiled - - _compile_gemm_a16w16 = compile_gemm_a16w16 - _run_compiled = run_compiled - _flyc = flyc - _fx = fx - - -_FX_DTYPE = {} - - -def _p(t): - if not _FX_DTYPE: - _FX_DTYPE.update( - { - torch.bfloat16: _fx.BFloat16, - torch.float16: _fx.Float16, - torch.float32: _fx.Float32, - } - ) - return _flyc.from_c_void_p(_FX_DTYPE[t.dtype], t.data_ptr()) - - -_CFG_KEYS = ( - "M", - "N", - "K", - "tile_m", - "tile_n", - "tile_k", - "m_warp", - "n_warp", - "in_dtype", - "out_dtype", - "num_buffers", - "waves_per_eu", - "activation", - "add_bias", - "physical_mk", - "physical_kn", - "kernarg_preload", - "split_k", - "sched_strategy", - "main_loop_unroll", - "variant", -) -@functools.lru_cache(maxsize=1024) -def _cached_launcher(*cfg): - return _compile_gemm_a16w16(**dict(zip(_CFG_KEYS, cfg))) - - -def gemm_a16w16( - x: torch.Tensor, - w: torch.Tensor, - bias: torch.Tensor | None = None, - dtype: torch.dtype = torch.float16, - y: torch.Tensor | None = None, - activation: str | None = None, - tile_m: int = 128, - tile_n: int = 128, - tile_k: int = 32, - m_warp: int = 2, - n_warp: int = 4, - num_buffers: int = 2, - waves_per_eu: int | None = None, - kernarg_preload: bool = False, - split_k: int = 1, - sched_strategy: str | None = None, - main_loop_unroll: bool = False, - variant: str = "bandwidth_bound", -): - """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" - _lazy_import() - _half = (torch.float16, torch.bfloat16) - assert x.dtype in _half, f"x must be fp16/bf16, got {x.dtype}" - assert w.dtype in _half, f"w must be fp16/bf16, got {w.dtype}" - assert x.shape[1] == w.shape[1], "Incompatible K dimensions" - - M, K, N = x.shape[0], x.shape[1], w.shape[0] - physical_mk = x.stride(1) == 1 - physical_kn = w.stride(1) != 1 - - K_padded = ((K + tile_k - 1) // tile_k) * tile_k - if K_padded != K: - pad_size = K_padded - K - if physical_mk: - x = torch.nn.functional.pad(x, (0, pad_size)) - else: - x = torch.nn.functional.pad(x.T, (0, 0, 0, pad_size)).T - - if physical_kn: - if w.stride(1) == 1: - w = torch.nn.functional.pad(w, (0, 0, 0, pad_size)) - else: - w = torch.nn.functional.pad(w.T, (0, 0, 0, pad_size)).T - else: - w = torch.nn.functional.pad(w, (0, pad_size)) - K = K_padded - - N_stride = ((N + tile_n - 1) // tile_n) * tile_n - _splitk_f32_accum = split_k > 1 and dtype in _half - in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" - out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get(dtype, "f32") - buf_dtype = dtype - if _splitk_f32_accum: - out_dtype_str, buf_dtype = "f32", torch.float32 - - _alloc = torch.zeros if split_k > 1 else torch.empty - if y is not None and not _splitk_f32_accum: - y_buf = ( - y - if N_stride == N - else _alloc((M, N_stride), device=x.device, dtype=buf_dtype) - ) - if split_k > 1 and y_buf is y: - y_buf.zero_() +def _load_kernel(): + """Load the kernel module by file path to skip aiter.ops.flydsl.__init__.""" + spec = importlib.util.spec_from_file_location( + "_flydsl_a16w16_wrapper", _KERNEL_PATH + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.gemm_a16w16 + + +gemm_a16w16 = _load_kernel() + + +def _generate_inputs(M, N, K, dtype, layout="TN", output=False, bias=False): + if layout[0] == "T": + x = torch.randn((M, K), dtype=dtype, device="cuda") + else: + x = torch.randn((K, M), dtype=dtype, device="cuda").T + + if layout[1] == "T": + w = torch.randn((K, N), dtype=dtype, device="cuda").T else: - y_buf = ( - _alloc((M, N_stride), device=x.device, dtype=buf_dtype) - if N_stride != N - else _alloc((M, N), device=x.device, dtype=buf_dtype) - ) - - if bias is None: - bias = torch.empty(0, device=x.device, dtype=dtype) - - launch_fn = _cached_launcher( - M if not physical_mk else 0, - N, - K, - tile_m, - tile_n, - tile_k, - m_warp, - n_warp, - in_dtype_str, - out_dtype_str, - num_buffers, - waves_per_eu, - activation, - bias.numel() > 0, - physical_mk, - physical_kn, - kernarg_preload, - split_k, - sched_strategy, - main_loop_unroll, - variant, + w = torch.randn((N, K), dtype=dtype, device="cuda") + + bias_tensor = None + if bias: + bias_tensor = torch.randn((N,), dtype=dtype, device="cuda") + + y = torch.empty((M, N), dtype=dtype, device="cuda") if output else None + return x, w, bias_tensor, y + + +def get_x_vals(): + return [ + # Aligned + (64, 64, 64), + (128, 128, 128), + (256, 256, 256), + # Multi-block + (128, 256, 512), + (256, 512, 256), + # Asymmetric + (32, 256, 128), + (256, 32, 128), + # Large K (drives the flat main loop + drain) + (128, 128, 1024), + (1024, 128, 128), + ] + + +def get_fewer_x_vals(): + return [ + (64, 64, 64), + (128, 256, 512), + (256, 512, 256), + (128, 128, 1024), + (1024, 128, 128), + ] + + +@pytest.mark.parametrize("M, N, K", get_x_vals()) +def test_gemm_a16_w16(M, N, K): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, num_buffers=3) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize("M, N, K", get_x_vals()) +@pytest.mark.parametrize("num_buffers", [2, 3]) +def test_gemm_a16_w16_compute_bound(M, N, K, num_buffers): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16( + x, w, dtype=torch.bfloat16, num_buffers=num_buffers, variant="compute_bound" ) - stream = torch.cuda.current_stream(device=x.device).cuda_stream - _run_compiled( - launch_fn, - _p(y_buf), - _p(x), - _p(w), - _p(bias), - M, - N_stride, - _fx.Stream(stream), + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu", "silu_exp2"]) +@pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("output", [True, False]) +def test_gemm_a16_w16_activation(M, N, K, dtype, output, activation): + torch.cuda.empty_cache() + x, w, _, y = _generate_inputs(M, N, K, dtype, output=output) + + torch_out = F.linear(x, w, bias=None) + if activation == "gelu": + torch_out = F.gelu(torch_out) + elif activation == "gelu_tanh": + torch_out = F.gelu(torch_out, approximate="tanh") + elif activation in ("silu", "silu_exp2"): + torch_out = F.silu(torch_out) + + kernel_out = gemm_a16w16( + x, w, dtype=dtype, y=y, activation=activation, num_buffers=3 ) - result = y_buf[:, :N] if N_stride != N else y_buf - if _splitk_f32_accum: - result = result.to(dtype) - if y is None or result is y: - return result - y.copy_(result) - return y + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize("M, N, K", get_x_vals()) +@pytest.mark.parametrize("layout", ["TN", "TT", "NN", "NT"]) +def test_gemm_a16_w16_layout(M, N, K, layout): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16, layout=layout) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, num_buffers=3) + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-1) + + +# 32x32 output tile, deep K (tile_k=128); tile_n=32 needs n_warp=2, K multiple of 128. +@pytest.mark.parametrize( + "M, N, K", [(32, 32, 256), (128, 128, 256), (256, 256, 512), (64, 64, 1024)] +) +def test_gemm_a16_w16_tile_32x32x128(M, N, K): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16( + x, + w, + dtype=torch.bfloat16, + num_buffers=3, + tile_m=32, + tile_n=32, + tile_k=128, + m_warp=2, + n_warp=2, + ) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +# Split-K: partition K across split_k grid-z workgroups, atomic-add into a pre-zeroed C. +def _splitk_configs(): + # (M, N, K, tile_m, tile_n, tile_k, m_warp, n_warp, num_buffers, split_k) — all valid. + return [ + (128, 128, 512, 64, 64, 128, 2, 2, 2, 2), # 4 tiles, 2/split + (128, 128, 512, 64, 64, 128, 2, 2, 2, 4), # 4 tiles, 1/split + (64, 64, 768, 64, 64, 128, 2, 2, 2, 3), # 6 tiles, 2/split + (64, 64, 768, 64, 64, 128, 2, 2, 2, 6), # 6 tiles, 1/split + (128, 256, 1024, 64, 64, 128, 2, 2, 2, 4), # 8 tiles, 2/split + (256, 128, 1024, 32, 32, 128, 2, 2, 2, 8), # 8 tiles, 1/split (max split) + (128, 128, 1024, 64, 64, 128, 2, 2, 3, 4), # nb=3, 2/split + ] + + +@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb,sk", _splitk_configs()) +def test_gemm_a16_w16_split_k(M, N, K, tm, tn, tk, mw, nw, nb, sk): + """Primary split-K correctness: fp32 output (lossless) -> split_k=N == split_k=1.""" + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + common = { + "dtype": torch.float32, + "tile_m": tm, + "tile_n": tn, + "tile_k": tk, + "m_warp": mw, + "n_warp": nw, + "num_buffers": nb, + } + out_single = gemm_a16w16(x, w, split_k=1, **common).clone() + out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() + + torch.testing.assert_close(out_split, out_single, atol=1e-2, rtol=1e-3) + + +@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb,sk", _splitk_configs()) +def test_gemm_a16_w16_split_k_bf16(M, N, K, tm, tn, tk, mw, nw, nb, sk): + """bf16-output split-K at O(1) magnitudes (scaled inputs) must match split_k=1.""" + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + x = (x.float() * 0.1).to(torch.bfloat16) + w = (w.float() * 0.1).to(torch.bfloat16) + + common = { + "dtype": torch.bfloat16, + "tile_m": tm, + "tile_n": tn, + "tile_k": tk, + "m_warp": mw, + "n_warp": nw, + "num_buffers": nb, + } + out_single = gemm_a16w16(x, w, split_k=1, **common).clone() + out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() + + torch.testing.assert_close(out_split, out_single, atol=1e-1, rtol=2e-2) + + +@pytest.mark.parametrize("sk", [2, 4]) +def test_gemm_a16_w16_split_k_output_buffer_zeroed(sk): + """split_k>1 atomic-ADDS into C; the wrapper must zero it each call (no doubling).""" + M, N, K = 128, 128, 512 + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + y = torch.empty((M, N), dtype=torch.float32, device="cuda") + + common = { + "dtype": torch.float32, + "y": y, + "tile_m": 64, + "tile_n": 64, + "tile_k": 128, + "m_warp": 2, + "n_warp": 2, + "num_buffers": 2, + "split_k": sk, + } + out1 = gemm_a16w16(x, w, **common).clone() + out2 = gemm_a16w16(x, w, **common).clone() # reuse y -> must re-zero, not 2x + + torch.testing.assert_close(out2, out1, atol=1e-2, rtol=1e-3) + + +def _oob_configs(): + return [ + (100, 100, 256, 128, 128, 32, 2, 4, 3), + (129, 257, 512, 128, 128, 32, 2, 4, 3), + (65, 190, 256, 128, 128, 32, 2, 4, 3), + (250, 120, 512, 64, 64, 128, 2, 2, 2), + (100, 128, 256, 128, 128, 32, 2, 4, 3), + (128, 100, 256, 128, 128, 32, 2, 4, 3), + (33, 65, 256, 32, 32, 128, 2, 2, 3), + ] + + +@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb", _oob_configs()) +def test_gemm_a16_w16_oob_tile_edges(M, N, K, tm, tn, tk, mw, nw, nb): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16( + x, + w, + dtype=torch.bfloat16, + tile_m=tm, + tile_n=tn, + tile_k=tk, + m_warp=mw, + n_warp=nw, + num_buffers=nb, + ) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize( + "M,N,K,tm,tn,tk,mw,nw,nb", + [ + (100, 128, 256, 128, 128, 32, 2, 4, 3), + (65, 64, 512, 64, 64, 128, 2, 2, 2), + ], +) +def test_gemm_a16_w16_oob_guard_region(M, N, K, tm, tn, tk, mw, nw, nb): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + sentinel = 777.0 + parent = torch.full((M + tm, N), sentinel, dtype=torch.bfloat16, device="cuda") + y = parent[:M] + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16( + x, + w, + dtype=torch.bfloat16, + y=y, + tile_m=tm, + tile_n=tn, + tile_k=tk, + m_warp=mw, + n_warp=nw, + num_buffers=nb, + ) -__all__ = ["gemm_a16w16"] + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + assert torch.all(parent[M:] == sentinel) From 7062dbfb4853985c9a6e296782fc2165a1d1e5d8 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 7 Aug 2026 01:49:25 +0000 Subject: [PATCH 22/34] oob edge cases --- op_tests/flydsl_tests/test_gemm_a16w16.py | 65 +++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index 5a18a31f48e..12b8041c606 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -266,3 +266,68 @@ def test_gemm_a16_w16_split_k_output_buffer_zeroed(sk): out2 = gemm_a16w16(x, w, **common).clone() # reuse y -> must re-zero, not 2x torch.testing.assert_close(out2, out1, atol=1e-2, rtol=1e-3) + + +def _oob_configs(): + return [ + (100, 100, 256, 128, 128, 32, 2, 4, 3), + (129, 257, 512, 128, 128, 32, 2, 4, 3), + (65, 190, 256, 128, 128, 32, 2, 4, 3), + (250, 120, 512, 64, 64, 128, 2, 2, 2), + (100, 128, 256, 128, 128, 32, 2, 4, 3), + (128, 100, 256, 128, 128, 32, 2, 4, 3), + (33, 65, 256, 32, 32, 128, 2, 2, 3), + ] + + +@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb", _oob_configs()) +def test_gemm_a16_w16_oob_tile_edges(M, N, K, tm, tn, tk, mw, nw, nb): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16( + x, + w, + dtype=torch.bfloat16, + tile_m=tm, + tile_n=tn, + tile_k=tk, + m_warp=mw, + n_warp=nw, + num_buffers=nb, + ) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize( + "M,N,K,tm,tn,tk,mw,nw,nb", + [ + (100, 128, 256, 128, 128, 32, 2, 4, 3), + (65, 64, 512, 64, 64, 128, 2, 2, 2), + ], +) +def test_gemm_a16_w16_oob_guard_region(M, N, K, tm, tn, tk, mw, nw, nb): + torch.cuda.empty_cache() + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + sentinel = 777.0 + parent = torch.full((M + tm, N), sentinel, dtype=torch.bfloat16, device="cuda") + y = parent[:M] + + torch_out = F.linear(x, w, bias=None) + kernel_out = gemm_a16w16( + x, + w, + dtype=torch.bfloat16, + y=y, + tile_m=tm, + tile_n=tn, + tile_k=tk, + m_warp=mw, + n_warp=nw, + num_buffers=nb, + ) + + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + assert torch.all(parent[M:] == sentinel) From f52a1aa6907b5588ee4ee5431b917c29204508c9 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 7 Aug 2026 01:59:14 +0000 Subject: [PATCH 23/34] Revert "Update gemm_a16w16_gfx1250.py" This reverts commit da29bb5f2add517530c0a2b92bc2e9be92f799a4. --- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 494 +++++++++--------------- 1 file changed, 181 insertions(+), 313 deletions(-) diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index 12b8041c606..9927716df73 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -1,333 +1,201 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -import importlib.util -import os - -import pytest -import torch -import torch.nn.functional as F - -_KERNEL_PATH = os.path.abspath( - os.path.join( - os.path.dirname(__file__), - "..", - "..", - "aiter", - "ops", - "flydsl", - "gemm_a16w16_gfx1250.py", - ) -) +"""gfx1250 FlyDSL backend for the A16W16 GEMM: torch-facing wrapper.""" +from __future__ import annotations -def _get_gpu_arch(): - if not torch.cuda.is_available(): - return None - return getattr(torch.cuda.get_device_properties(0), "gcnArchName", None) +import functools +import torch -def _flydsl_available(): - if importlib.util.find_spec("flydsl") is None: - return False - arch = _get_gpu_arch() - return arch is not None and arch.startswith("gfx1250") - +_compile_gemm_a16w16 = None +_run_compiled = None +_flyc = None +_fx = None -if not _flydsl_available(): - pytest.skip( - "FlyDSL a16w16 tests require gfx1250 and the flydsl package.", - allow_module_level=True, - ) +def _lazy_import(): + global _compile_gemm_a16w16, _run_compiled, _flyc, _fx + if _compile_gemm_a16w16 is not None: + return + import flydsl.compiler as flyc + import flydsl.expr as fx -def _load_kernel(): - """Load the kernel module by file path to skip aiter.ops.flydsl.__init__.""" - spec = importlib.util.spec_from_file_location( - "_flydsl_a16w16_wrapper", _KERNEL_PATH + from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( + compile_gemm_a16w16, ) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - return mod.gemm_a16w16 - - -gemm_a16w16 = _load_kernel() - + from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled as run_compiled + + _compile_gemm_a16w16 = compile_gemm_a16w16 + _run_compiled = run_compiled + _flyc = flyc + _fx = fx + + +_FX_DTYPE = {} + + +def _p(t): + if not _FX_DTYPE: + _FX_DTYPE.update( + { + torch.bfloat16: _fx.BFloat16, + torch.float16: _fx.Float16, + torch.float32: _fx.Float32, + } + ) + return _flyc.from_c_void_p(_FX_DTYPE[t.dtype], t.data_ptr()) + + +_CFG_KEYS = ( + "M", + "N", + "K", + "tile_m", + "tile_n", + "tile_k", + "m_warp", + "n_warp", + "in_dtype", + "out_dtype", + "num_buffers", + "waves_per_eu", + "activation", + "add_bias", + "physical_mk", + "physical_kn", + "kernarg_preload", + "split_k", + "sched_strategy", + "main_loop_unroll", + "variant", +) -def _generate_inputs(M, N, K, dtype, layout="TN", output=False, bias=False): - if layout[0] == "T": - x = torch.randn((M, K), dtype=dtype, device="cuda") - else: - x = torch.randn((K, M), dtype=dtype, device="cuda").T - if layout[1] == "T": - w = torch.randn((K, N), dtype=dtype, device="cuda").T +@functools.lru_cache(maxsize=1024) +def _cached_launcher(*cfg): + return _compile_gemm_a16w16(**dict(zip(_CFG_KEYS, cfg))) + + +def gemm_a16w16( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + dtype: torch.dtype = torch.float16, + y: torch.Tensor | None = None, + activation: str | None = None, + tile_m: int = 128, + tile_n: int = 128, + tile_k: int = 32, + m_warp: int = 2, + n_warp: int = 4, + num_buffers: int = 2, + waves_per_eu: int | None = None, + kernarg_preload: bool = False, + split_k: int = 1, + sched_strategy: str | None = None, + main_loop_unroll: bool = False, + variant: str = "bandwidth_bound", +): + """Compute Y = X @ W^T + bias. Auto-detects physical layout from strides.""" + _lazy_import() + _half = (torch.float16, torch.bfloat16) + assert x.dtype in _half, f"x must be fp16/bf16, got {x.dtype}" + assert w.dtype in _half, f"w must be fp16/bf16, got {w.dtype}" + assert x.shape[1] == w.shape[1], "Incompatible K dimensions" + + M, K, N = x.shape[0], x.shape[1], w.shape[0] + physical_mk = x.stride(1) == 1 + physical_kn = w.stride(1) != 1 + + K_padded = ((K + tile_k - 1) // tile_k) * tile_k + if K_padded != K: + pad_size = K_padded - K + if physical_mk: + x = torch.nn.functional.pad(x, (0, pad_size)) + else: + x = torch.nn.functional.pad(x.T, (0, 0, 0, pad_size)).T + + if physical_kn: + if w.stride(1) == 1: + w = torch.nn.functional.pad(w, (0, 0, 0, pad_size)) + else: + w = torch.nn.functional.pad(w.T, (0, 0, 0, pad_size)).T + else: + w = torch.nn.functional.pad(w, (0, pad_size)) + K = K_padded + + N_stride = ((N + tile_n - 1) // tile_n) * tile_n + _splitk_f32_accum = split_k > 1 and dtype in _half + in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" + out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get(dtype, "f32") + buf_dtype = dtype + if _splitk_f32_accum: + out_dtype_str, buf_dtype = "f32", torch.float32 + + _alloc = torch.zeros if split_k > 1 else torch.empty + if y is not None and not _splitk_f32_accum: + y_buf = ( + y + if N_stride == N + else _alloc((M, N_stride), device=x.device, dtype=buf_dtype) + ) + if split_k > 1 and y_buf is y: + y_buf.zero_() else: - w = torch.randn((N, K), dtype=dtype, device="cuda") - - bias_tensor = None - if bias: - bias_tensor = torch.randn((N,), dtype=dtype, device="cuda") - - y = torch.empty((M, N), dtype=dtype, device="cuda") if output else None - return x, w, bias_tensor, y - - -def get_x_vals(): - return [ - # Aligned - (64, 64, 64), - (128, 128, 128), - (256, 256, 256), - # Multi-block - (128, 256, 512), - (256, 512, 256), - # Asymmetric - (32, 256, 128), - (256, 32, 128), - # Large K (drives the flat main loop + drain) - (128, 128, 1024), - (1024, 128, 128), - ] - - -def get_fewer_x_vals(): - return [ - (64, 64, 64), - (128, 256, 512), - (256, 512, 256), - (128, 128, 1024), - (1024, 128, 128), - ] - - -@pytest.mark.parametrize("M, N, K", get_x_vals()) -def test_gemm_a16_w16(M, N, K): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, num_buffers=3) - - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("num_buffers", [2, 3]) -def test_gemm_a16_w16_compute_bound(M, N, K, num_buffers): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16( - x, w, dtype=torch.bfloat16, num_buffers=num_buffers, variant="compute_bound" + y_buf = ( + _alloc((M, N_stride), device=x.device, dtype=buf_dtype) + if N_stride != N + else _alloc((M, N), device=x.device, dtype=buf_dtype) + ) + + if bias is None: + bias = torch.empty(0, device=x.device, dtype=dtype) + + launch_fn = _cached_launcher( + M if not physical_mk else 0, + N, + K, + tile_m, + tile_n, + tile_k, + m_warp, + n_warp, + in_dtype_str, + out_dtype_str, + num_buffers, + waves_per_eu, + activation, + bias.numel() > 0, + physical_mk, + physical_kn, + kernarg_preload, + split_k, + sched_strategy, + main_loop_unroll, + variant, ) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu", "silu_exp2"]) -@pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16_activation(M, N, K, dtype, output, activation): - torch.cuda.empty_cache() - x, w, _, y = _generate_inputs(M, N, K, dtype, output=output) - - torch_out = F.linear(x, w, bias=None) - if activation == "gelu": - torch_out = F.gelu(torch_out) - elif activation == "gelu_tanh": - torch_out = F.gelu(torch_out, approximate="tanh") - elif activation in ("silu", "silu_exp2"): - torch_out = F.silu(torch_out) - - kernel_out = gemm_a16w16( - x, w, dtype=dtype, y=y, activation=activation, num_buffers=3 + stream = torch.cuda.current_stream(device=x.device).cuda_stream + _run_compiled( + launch_fn, + _p(y_buf), + _p(x), + _p(w), + _p(bias), + M, + N_stride, + _fx.Stream(stream), ) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("layout", ["TN", "TT", "NN", "NT"]) -def test_gemm_a16_w16_layout(M, N, K, layout): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16, layout=layout) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, num_buffers=3) + result = y_buf[:, :N] if N_stride != N else y_buf + if _splitk_f32_accum: + result = result.to(dtype) + if y is None or result is y: + return result + y.copy_(result) + return y - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-1) - - -# 32x32 output tile, deep K (tile_k=128); tile_n=32 needs n_warp=2, K multiple of 128. -@pytest.mark.parametrize( - "M, N, K", [(32, 32, 256), (128, 128, 256), (256, 256, 512), (64, 64, 1024)] -) -def test_gemm_a16_w16_tile_32x32x128(M, N, K): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16( - x, - w, - dtype=torch.bfloat16, - num_buffers=3, - tile_m=32, - tile_n=32, - tile_k=128, - m_warp=2, - n_warp=2, - ) - - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -# Split-K: partition K across split_k grid-z workgroups, atomic-add into a pre-zeroed C. -def _splitk_configs(): - # (M, N, K, tile_m, tile_n, tile_k, m_warp, n_warp, num_buffers, split_k) — all valid. - return [ - (128, 128, 512, 64, 64, 128, 2, 2, 2, 2), # 4 tiles, 2/split - (128, 128, 512, 64, 64, 128, 2, 2, 2, 4), # 4 tiles, 1/split - (64, 64, 768, 64, 64, 128, 2, 2, 2, 3), # 6 tiles, 2/split - (64, 64, 768, 64, 64, 128, 2, 2, 2, 6), # 6 tiles, 1/split - (128, 256, 1024, 64, 64, 128, 2, 2, 2, 4), # 8 tiles, 2/split - (256, 128, 1024, 32, 32, 128, 2, 2, 2, 8), # 8 tiles, 1/split (max split) - (128, 128, 1024, 64, 64, 128, 2, 2, 3, 4), # nb=3, 2/split - ] - - -@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb,sk", _splitk_configs()) -def test_gemm_a16_w16_split_k(M, N, K, tm, tn, tk, mw, nw, nb, sk): - """Primary split-K correctness: fp32 output (lossless) -> split_k=N == split_k=1.""" - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - common = { - "dtype": torch.float32, - "tile_m": tm, - "tile_n": tn, - "tile_k": tk, - "m_warp": mw, - "n_warp": nw, - "num_buffers": nb, - } - out_single = gemm_a16w16(x, w, split_k=1, **common).clone() - out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() - - torch.testing.assert_close(out_split, out_single, atol=1e-2, rtol=1e-3) - - -@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb,sk", _splitk_configs()) -def test_gemm_a16_w16_split_k_bf16(M, N, K, tm, tn, tk, mw, nw, nb, sk): - """bf16-output split-K at O(1) magnitudes (scaled inputs) must match split_k=1.""" - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - x = (x.float() * 0.1).to(torch.bfloat16) - w = (w.float() * 0.1).to(torch.bfloat16) - - common = { - "dtype": torch.bfloat16, - "tile_m": tm, - "tile_n": tn, - "tile_k": tk, - "m_warp": mw, - "n_warp": nw, - "num_buffers": nb, - } - out_single = gemm_a16w16(x, w, split_k=1, **common).clone() - out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() - - torch.testing.assert_close(out_split, out_single, atol=1e-1, rtol=2e-2) - - -@pytest.mark.parametrize("sk", [2, 4]) -def test_gemm_a16_w16_split_k_output_buffer_zeroed(sk): - """split_k>1 atomic-ADDS into C; the wrapper must zero it each call (no doubling).""" - M, N, K = 128, 128, 512 - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - y = torch.empty((M, N), dtype=torch.float32, device="cuda") - - common = { - "dtype": torch.float32, - "y": y, - "tile_m": 64, - "tile_n": 64, - "tile_k": 128, - "m_warp": 2, - "n_warp": 2, - "num_buffers": 2, - "split_k": sk, - } - out1 = gemm_a16w16(x, w, **common).clone() - out2 = gemm_a16w16(x, w, **common).clone() # reuse y -> must re-zero, not 2x - - torch.testing.assert_close(out2, out1, atol=1e-2, rtol=1e-3) - - -def _oob_configs(): - return [ - (100, 100, 256, 128, 128, 32, 2, 4, 3), - (129, 257, 512, 128, 128, 32, 2, 4, 3), - (65, 190, 256, 128, 128, 32, 2, 4, 3), - (250, 120, 512, 64, 64, 128, 2, 2, 2), - (100, 128, 256, 128, 128, 32, 2, 4, 3), - (128, 100, 256, 128, 128, 32, 2, 4, 3), - (33, 65, 256, 32, 32, 128, 2, 2, 3), - ] - - -@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb", _oob_configs()) -def test_gemm_a16_w16_oob_tile_edges(M, N, K, tm, tn, tk, mw, nw, nb): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16( - x, - w, - dtype=torch.bfloat16, - tile_m=tm, - tile_n=tn, - tile_k=tk, - m_warp=mw, - n_warp=nw, - num_buffers=nb, - ) - - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize( - "M,N,K,tm,tn,tk,mw,nw,nb", - [ - (100, 128, 256, 128, 128, 32, 2, 4, 3), - (65, 64, 512, 64, 64, 128, 2, 2, 2), - ], -) -def test_gemm_a16_w16_oob_guard_region(M, N, K, tm, tn, tk, mw, nw, nb): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - sentinel = 777.0 - parent = torch.full((M + tm, N), sentinel, dtype=torch.bfloat16, device="cuda") - y = parent[:M] - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16( - x, - w, - dtype=torch.bfloat16, - y=y, - tile_m=tm, - tile_n=tn, - tile_k=tk, - m_warp=mw, - n_warp=nw, - num_buffers=nb, - ) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - assert torch.all(parent[M:] == sentinel) +__all__ = ["gemm_a16w16"] From 5b04d43c8e97eca123361005aedf248037c24191 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Wed, 19 Aug 2026 17:18:41 +0000 Subject: [PATCH 24/34] e2e testing and fixes --- .../model_configs/gptoss_bf16_tuned_gemm.csv | 6 +- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 22 +++++-- .../kernels/gemm_a16w16_kernel_gfx1250.py | 63 ++++++++++++++----- 3 files changed, 68 insertions(+), 23 deletions(-) diff --git a/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv b/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv index e3acd298cce..17ca8987492 100644 --- a/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv +++ b/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv @@ -138,7 +138,6 @@ gfx1250,256,8,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0 gfx1250,256,16,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,32,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,48,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 -gfx1250,256,64,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,128,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,256,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,512,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 @@ -154,7 +153,6 @@ gfx1250,256,8,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0, gfx1250,256,16,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,32,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,48,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 -gfx1250,256,64,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,128,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,256,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,512,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 @@ -170,7 +168,6 @@ gfx1250,256,8,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0, gfx1250,256,16,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,32,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,48,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 -gfx1250,256,64,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,128,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,256,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,256,512,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 @@ -259,3 +256,6 @@ gfx1250,96,2048,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,trit gfx1250,96,4096,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,96,8192,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,96,16384,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 +gfx1250,256,64,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl_gfx1250,0,2,6.82,flydsl_a16w16_gfx1250_t64x64x256_mw4_nw2_nb4_sk2_unroll1,0.0,276.75,4474.37 +gfx1250,256,64,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,flydsl_gfx1250,0,4,5.82,flydsl_a16w16_gfx1250_t64x64x128_mw4_nw2_nb6_sk4_unroll1,0.0,259.44,4207.20 +gfx1250,256,64,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl_gfx1250,0,6,3.0,flydsl_a16w16_gfx1250_t16x16x256_mw1_nw1_nb3_sk6_unroll1,0.0,15.73,374.10 diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index 9927716df73..52e31695878 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -49,7 +49,6 @@ def _p(t): _CFG_KEYS = ( - "M", "N", "K", "tile_m", @@ -106,8 +105,14 @@ def gemm_a16w16( assert x.shape[1] == w.shape[1], "Incompatible K dimensions" M, K, N = x.shape[0], x.shape[1], w.shape[0] - physical_mk = x.stride(1) == 1 - physical_kn = w.stride(1) != 1 + assert ( + x.stride(1) == 1 or x.stride(0) == 1 or 1 in (M, K) + ), f"gemm_a16w16: x needs a unit-stride dim for TDM, got strides {tuple(x.stride())}" + assert ( + w.stride(1) == 1 or w.stride(0) == 1 or 1 in (N, K) + ), f"gemm_a16w16: w needs a unit-stride dim for TDM, got strides {tuple(w.stride())}" + physical_mk = x.stride(1) == 1 or K == 1 + physical_kn = w.stride(1) != 1 and N > 1 K_padded = ((K + tile_k - 1) // tile_k) * tile_k if K_padded != K: @@ -154,7 +159,6 @@ def gemm_a16w16( bias = torch.empty(0, device=x.device, dtype=dtype) launch_fn = _cached_launcher( - M if not physical_mk else 0, N, K, tile_m, @@ -177,6 +181,14 @@ def gemm_a16w16( variant, ) + if physical_mk: + lda = x.stride(0) if M > 1 else K + else: + lda = x.stride(1) if K > 1 else M + if physical_kn: + ldb = w.stride(1) if K > 1 else N + else: + ldb = w.stride(0) if N > 1 else K stream = torch.cuda.current_stream(device=x.device).cuda_stream _run_compiled( launch_fn, @@ -186,6 +198,8 @@ def gemm_a16w16( _p(bias), M, N_stride, + lda, + ldb, _fx.Stream(stream), ) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index 842decc1e2f..fecb3393017 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -86,6 +86,11 @@ def _req(cond, msg): f"bad sched_strategy {sched_strategy!r}", ) _req(split_k >= 1, f"split_k must be >= 1, got {split_k}") + _req( + not (activation and split_k > 1), + "activation is applied per k-split partial; activation with split_k > 1 " + "would compute act(partial) summed instead of act(sum)", + ) elem_bytes = 2 elem_bytes_d = 2 if out_dtype in ("f16", "bf16") else 4 @@ -114,7 +119,6 @@ def _req(cond, msg): _req(N > 0, "N must be > 0 at compile time when physical_kn=True") _req(tile_n & (tile_n - 1) == 0, f"tile_n must be a power of 2, got {tile_n}") if not physical_mk: - _req(M > 0, "M must be > 0 at compile time when physical_mk=False") _req(tile_m & (tile_m - 1) == 0, f"tile_m must be a power of 2, got {tile_m}") _req( num_k_tiles >= num_buffers - 1, @@ -131,20 +135,20 @@ def _req(cond, msg): N_A_FRAGS, N_B_FRAGS = wmma_m_rep * k_wmma_steps, wmma_n_rep * k_wmma_steps if physical_mk: - a_tile_shape, a_outer_stride, a_pad_interval = (tile_m, tile_k), K, tile_k + a_tile_shape, a_pad_interval = (tile_m, tile_k), tile_k a_imm_bytes, lds_a_stride = tile_k * elem_bytes, tile_k + LDS_PAD_A lds_a_elems = tile_m * lds_a_stride + LDS_PAD_A else: - a_tile_shape, a_outer_stride, a_pad_interval = (tile_k, tile_m), M, tile_m - a_imm_bytes, lds_a_stride = tile_k * M * elem_bytes, tile_m + LDS_PAD_A + a_tile_shape, a_pad_interval = (tile_k, tile_m), tile_m + a_imm_bytes, lds_a_stride = None, tile_m + LDS_PAD_A lds_a_elems = tile_k * lds_a_stride + LDS_PAD_A if physical_kn: - b_tile_shape, b_outer_stride, b_pad_interval = (tile_k, tile_n), N, tile_n - b_imm_bytes, lds_b_stride = tile_k * N * elem_bytes, tile_n + LDS_PAD_B + b_tile_shape, b_pad_interval = (tile_k, tile_n), tile_n + b_imm_bytes, lds_b_stride = None, tile_n + LDS_PAD_B lds_b_elems = tile_k * lds_b_stride + LDS_PAD_B else: - b_tile_shape, b_outer_stride, b_pad_interval = (tile_n, tile_k), K, tile_k + b_tile_shape, b_pad_interval = (tile_n, tile_k), tile_k b_imm_bytes, lds_b_stride = tile_k * elem_bytes, tile_k + LDS_PAD_B lds_b_elems = tile_n * lds_b_stride + LDS_PAD_B @@ -177,6 +181,8 @@ def kernel_gemm_a16w16( arg_bias: fx.Pointer, i32_m: fx.Int32, i32_n: fx.Int32, + i32_lda: fx.Int32, + i32_ldb: fx.Int32, ): rocdl.disable_xdl_arb_stall() elem_ty = _fx_elem.ir_type @@ -228,6 +234,7 @@ def _mk_side( num_warps=num_warps, pad_interval=pad_interval, pad_amount=pad_amount, + early_timeout=True, ) m_oob = fx.Int32(m_idx - blk_m) @@ -235,34 +242,40 @@ def _mk_side( a_extents = [m_oob, None] if const_expr(physical_mk) else [None, m_oob] b_extents = [None, n_oob] if const_expr(physical_kn) else [n_oob, None] + lda64 = fx.Int64(fx.Uint64(i32_lda)) + ldb64 = fx.Int64(fx.Uint64(i32_ldb)) if const_expr(physical_mk): - a_base_off = blk_m64 * K + split_k_base64 + a_base_off = blk_m64 * lda64 + split_k_base64 else: - a_base_off = split_k_base64 * M + blk_m64 + a_base_off = split_k_base64 * lda64 + blk_m64 + a_stride_rt = fx.Int32(i32_lda) gA, atom_a = _mk_side( arg_x, a_base_off, a_tile_shape, - a_outer_stride, + a_stride_rt, a_pad_interval, LDS_PAD_A, a_extents, ) if const_expr(physical_kn): - b_base_off = split_k_base64 * N + blk_n64 + b_base_off = split_k_base64 * ldb64 + blk_n64 else: - b_base_off = blk_n64 * K + split_k_base64 + b_base_off = blk_n64 * ldb64 + split_k_base64 + b_stride_rt = fx.Int32(i32_ldb) gB, atom_b = _mk_side( arg_w, b_base_off, b_tile_shape, - b_outer_stride, + b_stride_rt, b_pad_interval, LDS_PAD_B, b_extents, ) def _imm64(k_tile, mul): + if const_expr(not isinstance(mul, int)): + return fx.Int64(fx.Int64(k_tile) * mul) if const_expr(isinstance(k_tile, int)): return fx.Int64(k_tile * mul) return fx.Int64(fx.Int32(k_tile) * mul) # mul in i32, widen for TDM @@ -287,6 +300,17 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): for rep in range_constexpr(reps) ] + a_imm_rt = ( + a_imm_bytes + if const_expr(a_imm_bytes is not None) + else lda64 * (tile_k * elem_bytes) + ) + b_imm_rt = ( + b_imm_bytes + if const_expr(b_imm_bytes is not None) + else ldb64 * (tile_k * elem_bytes) + ) + # One spec per operand side; A and B differ only in these fields. MEM, BASE, BASES, STRIDE, TR, ELEMS, REPS, ATOM, GT, SHAPE, IMM = range(11) A_SIDE = ( @@ -300,7 +324,7 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): atom_a, gA, a_tile_shape, - a_imm_bytes, + a_imm_rt, ) B_SIDE = ( big_b_mem, @@ -313,7 +337,7 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): atom_b, gB, b_tile_shape, - b_imm_bytes, + b_imm_rt, ) def issue_tdm_loads(buf_idx, k_tile): @@ -441,7 +465,10 @@ def epilogue_stores(final_accs): for half in range_constexpr(2): bv = load_bias4(col_tile + half) for i in range_constexpr(4): - elems.append(bv[i].to(fx.Float32)) + if const_expr(split_k > 1): + elems.append(bv[i].to(fx.Float32) * (1.0 / split_k)) + else: + elems.append(bv[i].to(fx.Float32)) bias_vecs.append( fx.Vector.from_elements(elems, fx.Float32).ir_value() ) @@ -637,6 +664,8 @@ def launch_gemm_a16w16( arg_bias: fx.Pointer, i32_m: fx.Int32, i32_n: fx.Int32, + i32_lda: fx.Int32, + i32_ldb: fx.Int32, stream: fx.Stream, ): gx = (fx.Uint64(i32_m) + (tile_m - 1)) // tile_m @@ -650,6 +679,8 @@ def launch_gemm_a16w16( arg_bias, i32_m, i32_n, + i32_lda, + i32_ldb, value_attrs={ "rocdl.flat_work_group_size": f"{block_threads},{block_threads}", "rocdl.waves_per_eu": wpe if wpe >= 1 else None, From 069f72de2b99d59cd48ac9eb8e608d22324acf53 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 20 Aug 2026 16:23:42 +0000 Subject: [PATCH 25/34] Wire flydsl_gfx1250 a16w16 dispatch into tuned_gemm solMap The tuned CSV rows added for gfx1250 (libtype=flydsl_gfx1250) had no matching solMap entry, so those bf16 shapes would KeyError at dispatch. Add the kernelName parser (tile/warp/buffer/split_k/unroll encoding) and the flydsl_gfx1250_gemm dispatch wrapper, and register it in solMap. Addresses PR review feedback. Co-Authored-By: Claude Fable 5 --- aiter/tuned_gemm.py | 58 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/aiter/tuned_gemm.py b/aiter/tuned_gemm.py index 3d0e919ccdf..b7529720ff7 100644 --- a/aiter/tuned_gemm.py +++ b/aiter/tuned_gemm.py @@ -458,6 +458,63 @@ def asm_gemm( return gemm_a16w16_asm(inp, weights, out_asm, bias, splitK, kernelName, bpreshuffle) +_FLY1250_NAME_RE = None + + +def _parse_flydsl_gfx1250_kname(name: str) -> dict: + """Parse flydsl_a16w16_gfx1250_t{tm}x{tn}x{tk}_mw{mw}_nw{nw}_nb{nb}_sk{sk}_unroll{0|1}.""" + global _FLY1250_NAME_RE + if _FLY1250_NAME_RE is None: + import re + + _FLY1250_NAME_RE = re.compile( + r"flydsl_a16w16_gfx1250_t(\d+)x(\d+)x(\d+)_mw(\d+)_nw(\d+)" + r"_nb(\d+)_sk(\d+)_unroll([01])" + ) + m = _FLY1250_NAME_RE.match(name) + assert m is not None, f"unrecognized flydsl_gfx1250 kernel name: {name!r}" + tm, tn, tk, mw, nw, nb, sk, un = (int(g) for g in m.groups()) + return dict( + tile_m=tm, + tile_n=tn, + tile_k=tk, + m_warp=mw, + n_warp=nw, + num_buffers=nb, + split_k=sk, + main_loop_unroll=bool(un), + ) + + +def flydsl_gfx1250_gemm( + inp: Tensor, + weights: Tensor, + solidx: int, + bias: Tensor | None = None, + otype: torch.dtype | None = None, + scale_a: Tensor | None = None, + scale_b: Tensor | None = None, + scale_c: Tensor | None = None, + bpreshuffle=False, + config: dict | None = None, +): + assert ( + scale_a is None and scale_b is None and scale_c is None + ), "flydsl_gfx1250 a16w16 does not support scaling." + assert not bpreshuffle, "flydsl_gfx1250 a16w16 expects unshuffled weights." + from aiter.ops.flydsl.gemm_a16w16_gfx1250 import gemm_a16w16 as _fly1250 + + cfg = _parse_flydsl_gfx1250_kname(config["kernelName"]) + return _fly1250( + inp, + weights, + bias=bias, + dtype=otype if otype is not None else inp.dtype, + waves_per_eu=None, + **cfg, + ) + + def flydsl_gemm( inp: Tensor, weights: Tensor, @@ -596,6 +653,7 @@ def triton_gemm( "asm": asm_gemm, "triton": triton_gemm, "flydsl": flydsl_gemm, + "flydsl_gfx1250": flydsl_gfx1250_gemm, "opus": opus_gemm, } From 71db74d130e8569fc84314775896b1f2095e89cc Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 20 Aug 2026 12:26:05 -0400 Subject: [PATCH 26/34] Update aiter/tuned_gemm.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- aiter/tuned_gemm.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/aiter/tuned_gemm.py b/aiter/tuned_gemm.py index b7529720ff7..65331e69f27 100644 --- a/aiter/tuned_gemm.py +++ b/aiter/tuned_gemm.py @@ -474,16 +474,16 @@ def _parse_flydsl_gfx1250_kname(name: str) -> dict: m = _FLY1250_NAME_RE.match(name) assert m is not None, f"unrecognized flydsl_gfx1250 kernel name: {name!r}" tm, tn, tk, mw, nw, nb, sk, un = (int(g) for g in m.groups()) - return dict( - tile_m=tm, - tile_n=tn, - tile_k=tk, - m_warp=mw, - n_warp=nw, - num_buffers=nb, - split_k=sk, - main_loop_unroll=bool(un), - ) + return { + "tile_m": tm, + "tile_n": tn, + "tile_k": tk, + "m_warp": mw, + "n_warp": nw, + "num_buffers": nb, + "split_k": sk, + "main_loop_unroll": bool(un), + } def flydsl_gfx1250_gemm( From 7338bfbc79f5cf005e8179b8d964c501ccd1063d Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 20 Aug 2026 18:26:58 +0000 Subject: [PATCH 27/34] Address review nits: fx-native LDS pointer, single typed casts - _frag transposed load: replace raw llvm_dialect.inttoptr + ir.Type.parse("!llvm.ptr<3>") with fx.PointerType/fx.inttoptr/ fx.to_llvm_ptr (matches chunk_gated_delta_h); drops the flydsl._mlir imports entirely. - Collapse redundant fx wraps in _f32/_imm64/_slot_off to single typed casts. No functional change: 172/172 op tests pass. Co-Authored-By: Claude Fable 5 --- .../kernels/gemm_a16w16_kernel_gfx1250.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index fecb3393017..934bc818c8c 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -5,8 +5,6 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import llvm as llvm_dialect from flydsl.compiler.protocol import dsl_size_of from flydsl.expr import const_expr, gpu, idx2crd, range_constexpr, rocdl, tdm_ops from flydsl.expr.typing import T @@ -25,7 +23,7 @@ def _byte_off_i32(elem_off, elem_bytes): def _f32(x): - return fx.Float32(fx.Float32(x).ir_value()) + return fx.Float32(x) def apply_activation_scalar(val, activation: str): @@ -275,15 +273,15 @@ def _mk_side( def _imm64(k_tile, mul): if const_expr(not isinstance(mul, int)): - return fx.Int64(fx.Int64(k_tile) * mul) + return fx.Int64(k_tile) * mul if const_expr(isinstance(k_tile, int)): return fx.Int64(k_tile * mul) - return fx.Int64(fx.Int32(k_tile) * mul) # mul in i32, widen for TDM + return fx.Int64(k_tile * mul) def _slot_off(buf_idx, slot_elems): if const_expr(isinstance(buf_idx, int)): return fx.Uint64(buf_idx * slot_elems) - return fx.Uint64(fx.Int32(buf_idx) * slot_elems) # mul in i32, widen + return fx.Uint64(buf_idx * slot_elems) def _lane_bases(warp_base, reps, lds_stride, transpose): if const_expr(not transpose): @@ -357,6 +355,12 @@ def issue_tdm_loads(buf_idx, k_tile): imm_offset=_imm64(k_tile, sd[IMM]), ) + lds_tr_ptr_ty = fx.PointerType.get( + elem_ty=_fx_elem.ir_type, + address_space=fx.AddressSpace.Shared, + alignment=elem_bytes, + ) + def _frag(sd, slot_off, ks, rep): vec8_ty = T.vec(8, elem_ty) halves = [] @@ -367,8 +371,8 @@ def _frag(sd, slot_off, ks, rep): + slot_off + (ks * WMMA_K + k_half * 16) * sd[STRIDE] ) - addr = fx.Int32(sd[BASE] + off * elem_bytes).ir_value() - ptr = llvm_dialect.inttoptr(ir.Type.parse("!llvm.ptr<3>"), addr) + addr = fx.Int32(sd[BASE] + off * elem_bytes) + ptr = fx.to_llvm_ptr(fx.inttoptr(lds_tr_ptr_ty, addr)) halves.append(fx.Vector(rocdl.ds_load_tr16_b128(vec8_ty, ptr))) else: off = sd[BASES][rep] + slot_off + (ks * WMMA_K + k_half * 16) From 9917e566b7b22bcd6c897633ffd71acf66e5ea91 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 20 Aug 2026 18:27:19 +0000 Subject: [PATCH 28/34] Support FlyDSL >=0.3.1: version-gate raw buffer op aux argument FlyDSL 0.3.1 changed raw_ptr_buffer_store/atomic_fadd's trailing aux from a required SSA operand (i32 Value) to an optional MLIR attribute; the old Value form raises on store and is silently dropped on atomic_fadd (losing the split-k device-scope bits). Gate the four call sites on the installed version at compile time: <=0.3.0 keeps the operand form, >=0.3.1 passes aux=SCOPE_DEV as an int and omits the zero aux on stores. Tested: 172/172 op tests on both flydsl 0.3.0 and 0.3.1. Co-Authored-By: Claude Fable 5 --- .../kernels/gemm_a16w16_kernel_gfx1250.py | 101 ++++++++++++------ 1 file changed, 68 insertions(+), 33 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index 934bc818c8c..3a30f2643c3 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -18,6 +18,15 @@ SCOPE_DEV = 16 +def _flydsl_aux_is_attr() -> bool: + from importlib.metadata import version + + return tuple(int(x) for x in version("flydsl").split(".")[:3]) >= (0, 3, 1) + + +_AUX_IS_ATTR = _flydsl_aux_is_attr() + + def _byte_off_i32(elem_off, elem_bytes): return (fx.Int32(elem_off) * elem_bytes).ir_value() @@ -507,21 +516,38 @@ def epilogue_stores(final_accs): ], _out_num, ).ir_value() - rocdl.raw_ptr_buffer_atomic_fadd( - pair_vec, + if const_expr(_AUX_IS_ATTR): + rocdl.raw_ptr_buffer_atomic_fadd( + pair_vec, + y_rsrc, + fx.Int32(c_off_bytes + pair * 4).ir_value(), + zero_i32, + aux=SCOPE_DEV, + ) + else: + rocdl.raw_ptr_buffer_atomic_fadd( + pair_vec, + y_rsrc, + fx.Int32(c_off_bytes + pair * 4).ir_value(), + zero_i32, + dev_scope, + ) + else: + if const_expr(_AUX_IS_ATTR): + rocdl.raw_ptr_buffer_store( + h_vec.bitcast(fx.Int32).ir_value(), y_rsrc, - fx.Int32(c_off_bytes + pair * 4).ir_value(), - zero_i32, - dev_scope, + fx.Int32(c_off_bytes).ir_value(), + zero_off, + ) + else: + rocdl.raw_ptr_buffer_store( + h_vec.bitcast(fx.Int32).ir_value(), + y_rsrc, + fx.Int32(c_off_bytes).ir_value(), + zero_off, + zero_off, ) - else: - rocdl.raw_ptr_buffer_store( - h_vec.bitcast(fx.Int32).ir_value(), - y_rsrc, - fx.Int32(c_off_bytes).ir_value(), - zero_off, - zero_off, - ) elif const_expr(split_k > 1): for e in range_constexpr(8): fx.ptr_store( @@ -532,20 +558,27 @@ def epilogue_stores(final_accs): t_col = blk_n + warp_n_base + wn * WMMA_N + lane16 for e in range_constexpr(8): if t_base + 2 * e < m_idx: - rocdl.raw_ptr_buffer_atomic_fadd( - fx.ptr_load( - fx.add_offset( - xt, (lane_kgrp + 2 * e) * WMMA_N + lane16 - ) - ).ir_value(), - y_rsrc, - _byte_off_i32( - (t_base + lane_kgrp + 2 * e) * n_stride + t_col, - 4, - ), - zero_i32, - dev_scope, + _xt_val = fx.ptr_load( + fx.add_offset( + xt, (lane_kgrp + 2 * e) * WMMA_N + lane16 + ) + ).ir_value() + _xt_off = _byte_off_i32( + (t_base + lane_kgrp + 2 * e) * n_stride + t_col, + 4, ) + if const_expr(_AUX_IS_ATTR): + rocdl.raw_ptr_buffer_atomic_fadd( + _xt_val, + y_rsrc, + _xt_off, + zero_i32, + aux=SCOPE_DEV, + ) + else: + rocdl.raw_ptr_buffer_atomic_fadd( + _xt_val, y_rsrc, _xt_off, zero_i32, dev_scope + ) else: for half in range_constexpr(2): vec4 = fx.Vector.from_elements( @@ -556,13 +589,15 @@ def epilogue_stores(final_accs): fx.Float32, ).ir_value() col = col_base + half * 4 - rocdl.raw_ptr_buffer_store( - vec4, - y_rsrc, - _byte_off_i32(row * n_stride + col, 4), - zero_off, - zero_off, - ) + _f32_off = _byte_off_i32(row * n_stride + col, 4) + if const_expr(_AUX_IS_ATTR): + rocdl.raw_ptr_buffer_store( + vec4, y_rsrc, _f32_off, zero_off + ) + else: + rocdl.raw_ptr_buffer_store( + vec4, y_rsrc, _f32_off, zero_off, zero_off + ) def _pack_state(accs_, a_, b_): return list(accs_) + list(a_) + list(b_) From 449ba4f6245d7642f655789ea222b6e950924802 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 21 Aug 2026 19:46:33 +0000 Subject: [PATCH 29/34] resolving comments --- .../kernels/gemm_a16w16_kernel_gfx1250.py | 124 +++++++----------- 1 file changed, 46 insertions(+), 78 deletions(-) diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index 3a30f2643c3..21efaf2c40e 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -15,16 +15,6 @@ WAVE_SIZE, LDS_PAD_A, LDS_PAD_B = 32, 8, 8 _SCHED_ALLOW_SALU = 1 << 2 KERNARG_PRELOAD_COUNT = 8 -SCOPE_DEV = 16 - - -def _flydsl_aux_is_attr() -> bool: - from importlib.metadata import version - - return tuple(int(x) for x in version("flydsl").split(".")[:3]) >= (0, 3, 1) - - -_AUX_IS_ATTR = _flydsl_aux_is_attr() def _byte_off_i32(elem_off, elem_bytes): @@ -214,8 +204,11 @@ def kernel_gemm_a16w16( warp_m_base, warp_n_base = wave_m_idx * warp_tile_m, wave_n_idx * warp_tile_n m_idx, n_stride = fx.Uint64(i32_m), fx.Uint64(i32_n) - y_rsrc = fx.rocdl.get_buffer_rsrc( - fx.rocdl.make_buffer_ptr(arg_y, m_idx * n_stride * elem_bytes_d) + gYp = fx.Tensor(fx.make_view(arg_y, fx.make_layout((1, 1), (1, 1)))) + gY = fx.rocdl.make_buffer_tensor( + fx.Tensor(fx.make_view(arg_y, fx.make_layout((1, 1), (1, 1)))), + max_size=False, + num_records_bytes=m_idx * n_stride * elem_bytes_d, ) lds = fx.SharedAllocator(static=True).allocate(SharedStorage).peek() @@ -461,10 +454,13 @@ def load_bias4(tile_idx): return r.load() def epilogue_stores(final_accs): - zero_off = fx.Int32(0).ir_value() + _acc_ty = _out_num if _half_out else fx.Float32 + st_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), _acc_ty) if const_expr(split_k > 1): - zero_i32 = zero_off - dev_scope = fx.Int32(SCOPE_DEV).ir_value() + add_atom = fx.make_copy_atom( + fx.UniversalAtomicAdd(_acc_ty, syncscope=fx.rocdl.SyncScope.Agent), + _acc_ty, + ) if const_expr(_USE_XT): xt = fx.add_offset( lds.xt.ptr, @@ -506,7 +502,7 @@ def epilogue_stores(final_accs): if const_expr(_half_out): h_vec = fx.Vector(acc).to(_out_num) - c_off_bytes = (row * n_stride + col_base) * elem_bytes_d + c_off = row * n_stride + col_base if const_expr(split_k > 1): for pair in range_constexpr(4): pair_vec = fx.Vector.from_elements( @@ -515,39 +511,19 @@ def epilogue_stores(final_accs): h_vec[pair * 2 + 1].ir_value(), ], _out_num, - ).ir_value() - if const_expr(_AUX_IS_ATTR): - rocdl.raw_ptr_buffer_atomic_fadd( - pair_vec, - y_rsrc, - fx.Int32(c_off_bytes + pair * 4).ir_value(), - zero_i32, - aux=SCOPE_DEV, - ) - else: - rocdl.raw_ptr_buffer_atomic_fadd( - pair_vec, - y_rsrc, - fx.Int32(c_off_bytes + pair * 4).ir_value(), - zero_i32, - dev_scope, + ) + if row < m_idx: + fx.copy( + add_atom, + _rmem_vec(pair_vec, 2, _out_num), + gYp[None, c_off + pair * 2], ) else: - if const_expr(_AUX_IS_ATTR): - rocdl.raw_ptr_buffer_store( - h_vec.bitcast(fx.Int32).ir_value(), - y_rsrc, - fx.Int32(c_off_bytes).ir_value(), - zero_off, - ) - else: - rocdl.raw_ptr_buffer_store( - h_vec.bitcast(fx.Int32).ir_value(), - y_rsrc, - fx.Int32(c_off_bytes).ir_value(), - zero_off, - zero_off, - ) + fx.copy( + st_atom, + _rmem_vec(h_vec, 8, _out_num), + gY[None, c_off], + ) elif const_expr(split_k > 1): for e in range_constexpr(8): fx.ptr_store( @@ -557,28 +533,24 @@ def epilogue_stores(final_accs): t_base = blk_m + warp_m_base + wm * WMMA_M t_col = blk_n + warp_n_base + wn * WMMA_N + lane16 for e in range_constexpr(8): - if t_base + 2 * e < m_idx: - _xt_val = fx.ptr_load( - fx.add_offset( - xt, (lane_kgrp + 2 * e) * WMMA_N + lane16 - ) - ).ir_value() - _xt_off = _byte_off_i32( - (t_base + lane_kgrp + 2 * e) * n_stride + t_col, - 4, + xt_row = t_base + lane_kgrp + 2 * e + if xt_row < m_idx: + xt_val = fx.Vector.from_elements( + [ + fx.ptr_load( + fx.add_offset( + xt, + (lane_kgrp + 2 * e) * WMMA_N + lane16, + ) + ).ir_value() + ], + fx.Float32, + ) + fx.copy( + add_atom, + _rmem_vec(xt_val, 1, fx.Float32), + gYp[None, xt_row * n_stride + t_col], ) - if const_expr(_AUX_IS_ATTR): - rocdl.raw_ptr_buffer_atomic_fadd( - _xt_val, - y_rsrc, - _xt_off, - zero_i32, - aux=SCOPE_DEV, - ) - else: - rocdl.raw_ptr_buffer_atomic_fadd( - _xt_val, y_rsrc, _xt_off, zero_i32, dev_scope - ) else: for half in range_constexpr(2): vec4 = fx.Vector.from_elements( @@ -587,17 +559,13 @@ def epilogue_stores(final_accs): for vi in range_constexpr(4) ], fx.Float32, - ).ir_value() + ) col = col_base + half * 4 - _f32_off = _byte_off_i32(row * n_stride + col, 4) - if const_expr(_AUX_IS_ATTR): - rocdl.raw_ptr_buffer_store( - vec4, y_rsrc, _f32_off, zero_off - ) - else: - rocdl.raw_ptr_buffer_store( - vec4, y_rsrc, _f32_off, zero_off, zero_off - ) + fx.copy( + st_atom, + _rmem_vec(vec4, 4, fx.Float32), + gY[None, row * n_stride + col], + ) def _pack_state(accs_, a_, b_): return list(accs_) + list(a_) + list(b_) From a94e8515fde9209c9b2727a21806083fdf3d887d Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Mon, 31 Aug 2026 18:22:46 +0000 Subject: [PATCH 30/34] removing bad copies, and unncessary padding --- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 66 ++---- .../kernels/gemm_a16w16_kernel_gfx1250.py | 208 +++++++++++++----- op_tests/flydsl_tests/test_gemm_a16w16.py | 95 ++++++++ 3 files changed, 267 insertions(+), 102 deletions(-) diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index 52e31695878..f71f6935615 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -11,25 +11,25 @@ _compile_gemm_a16w16 = None _run_compiled = None -_flyc = None +_ptr_arg = None _fx = None def _lazy_import(): - global _compile_gemm_a16w16, _run_compiled, _flyc, _fx + global _compile_gemm_a16w16, _run_compiled, _ptr_arg, _fx if _compile_gemm_a16w16 is not None: return - import flydsl.compiler as flyc import flydsl.expr as fx from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( compile_gemm_a16w16, ) from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled as run_compiled + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg _compile_gemm_a16w16 = compile_gemm_a16w16 _run_compiled = run_compiled - _flyc = flyc + _ptr_arg = ptr_arg _fx = fx @@ -45,7 +45,7 @@ def _p(t): torch.float32: _fx.Float32, } ) - return _flyc.from_c_void_p(_FX_DTYPE[t.dtype], t.data_ptr()) + return _ptr_arg(t, _FX_DTYPE[t.dtype]) _CFG_KEYS = ( @@ -114,24 +114,6 @@ def gemm_a16w16( physical_mk = x.stride(1) == 1 or K == 1 physical_kn = w.stride(1) != 1 and N > 1 - K_padded = ((K + tile_k - 1) // tile_k) * tile_k - if K_padded != K: - pad_size = K_padded - K - if physical_mk: - x = torch.nn.functional.pad(x, (0, pad_size)) - else: - x = torch.nn.functional.pad(x.T, (0, 0, 0, pad_size)).T - - if physical_kn: - if w.stride(1) == 1: - w = torch.nn.functional.pad(w, (0, 0, 0, pad_size)) - else: - w = torch.nn.functional.pad(w.T, (0, 0, 0, pad_size)).T - else: - w = torch.nn.functional.pad(w, (0, pad_size)) - K = K_padded - - N_stride = ((N + tile_n - 1) // tile_n) * tile_n _splitk_f32_accum = split_k > 1 and dtype in _half in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get(dtype, "f32") @@ -141,19 +123,20 @@ def gemm_a16w16( _alloc = torch.zeros if split_k > 1 else torch.empty if y is not None and not _splitk_f32_accum: - y_buf = ( - y - if N_stride == N - else _alloc((M, N_stride), device=x.device, dtype=buf_dtype) - ) - if split_k > 1 and y_buf is y: + assert ( + y.shape[0] == M and y.shape[1] == N + ), f"y must be ({M}, {N}), got {tuple(y.shape)}" + assert y.stride(1) == 1 or 1 in ( + M, + N, + ), f"gemm_a16w16: y needs unit column stride, got strides {tuple(y.stride())}" + y_buf = y + if split_k > 1: y_buf.zero_() else: - y_buf = ( - _alloc((M, N_stride), device=x.device, dtype=buf_dtype) - if N_stride != N - else _alloc((M, N), device=x.device, dtype=buf_dtype) - ) + y_buf = _alloc((M, N), device=x.device, dtype=buf_dtype) + ldy = y_buf.stride(0) if M > 1 else N + assert ldy >= N, f"gemm_a16w16: y row stride {ldy} < N {N}" if bias is None: bias = torch.empty(0, device=x.device, dtype=dtype) @@ -197,19 +180,18 @@ def gemm_a16w16( _p(w), _p(bias), M, - N_stride, + ldy, lda, ldb, _fx.Stream(stream), ) - result = y_buf[:, :N] if N_stride != N else y_buf - if _splitk_f32_accum: - result = result.to(dtype) - if y is None or result is y: - return result - y.copy_(result) - return y + if not _splitk_f32_accum: + return y_buf + if y is not None: + y.copy_(y_buf) + return y + return y_buf.to(dtype) __all__ = ["gemm_a16w16"] diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index 21efaf2c40e..29ca0adb70f 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -17,10 +17,6 @@ KERNARG_PRELOAD_COUNT = 8 -def _byte_off_i32(elem_off, elem_bytes): - return (fx.Int32(elem_off) * elem_bytes).ir_value() - - def _f32(x): return fx.Float32(x) @@ -94,12 +90,11 @@ def _req(cond, msg): num_warps, block_threads = m_warp * n_warp, m_warp * n_warp * WAVE_SIZE warp_tile_m, warp_tile_n = tile_m // m_warp, tile_n // n_warp split_k_chunk = K // split_k - num_k_tiles = split_k_chunk // tile_k + num_k_tiles = (split_k_chunk + tile_k - 1) // tile_k + k_rem = split_k_chunk % tile_k for _nm, _v, _d in ( - ("K", K, tile_k), ("K", K, split_k), - ("K/split_k", split_k_chunk, tile_k), ("tile_k", tile_k, WMMA_K), ("tile_m", tile_m, WMMA_M), ("tile_n", tile_n, WMMA_N), @@ -112,8 +107,8 @@ def _req(cond, msg): tile_k & (tile_k - 1) == 0, f"tile_k must be a power of 2 for TDM, got {tile_k}" ) + _req(N > 0, "N must be > 0 at compile time") if physical_kn: - _req(N > 0, "N must be > 0 at compile time when physical_kn=True") _req(tile_n & (tile_n - 1) == 0, f"tile_n must be a power of 2, got {tile_n}") if not physical_mk: _req(tile_m & (tile_m - 1) == 0, f"tile_m must be a power of 2, got {tile_m}") @@ -169,6 +164,7 @@ class SharedStorage: check_smem_capacity(dsl_size_of(SharedStorage), gpu_arch) _TDMS_PER_TILE = 2 # A + B + n_edge = N % tile_n @flyc.kernel def kernel_gemm_a16w16( @@ -177,7 +173,7 @@ def kernel_gemm_a16w16( arg_w: fx.Pointer, arg_bias: fx.Pointer, i32_m: fx.Int32, - i32_n: fx.Int32, + i32_ldy: fx.Int32, i32_lda: fx.Int32, i32_ldb: fx.Int32, ): @@ -203,12 +199,12 @@ def kernel_gemm_a16w16( ) warp_m_base, warp_n_base = wave_m_idx * warp_tile_m, wave_n_idx * warp_tile_n - m_idx, n_stride = fx.Uint64(i32_m), fx.Uint64(i32_n) + m_idx, ld_y = fx.Uint64(i32_m), fx.Uint64(i32_ldy) gYp = fx.Tensor(fx.make_view(arg_y, fx.make_layout((1, 1), (1, 1)))) gY = fx.rocdl.make_buffer_tensor( fx.Tensor(fx.make_view(arg_y, fx.make_layout((1, 1), (1, 1)))), max_size=False, - num_records_bytes=m_idx * n_stride * elem_bytes_d, + num_records_bytes=m_idx * ld_y * elem_bytes_d, ) lds = fx.SharedAllocator(static=True).allocate(SharedStorage).peek() @@ -276,13 +272,9 @@ def _mk_side( def _imm64(k_tile, mul): if const_expr(not isinstance(mul, int)): return fx.Int64(k_tile) * mul - if const_expr(isinstance(k_tile, int)): - return fx.Int64(k_tile * mul) return fx.Int64(k_tile * mul) def _slot_off(buf_idx, slot_elems): - if const_expr(isinstance(buf_idx, int)): - return fx.Uint64(buf_idx * slot_elems) return fx.Uint64(buf_idx * slot_elems) def _lane_bases(warp_base, reps, lds_stride, transpose): @@ -312,7 +304,9 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): ) # One spec per operand side; A and B differ only in these fields. - MEM, BASE, BASES, STRIDE, TR, ELEMS, REPS, ATOM, GT, SHAPE, IMM = range(11) + MEM, BASE, BASES, STRIDE, TR, ELEMS, REPS, ATOM, GT, SHAPE, IMM, KDIM = range( + 12 + ) A_SIDE = ( big_a_mem, big_a_base_idx, @@ -325,6 +319,7 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): gA, a_tile_shape, a_imm_rt, + 1 if physical_mk else 0, ) B_SIDE = ( big_b_mem, @@ -338,12 +333,18 @@ def _lane_bases(warp_base, reps, lds_stride, transpose): gB, b_tile_shape, b_imm_rt, + 0 if physical_kn else 1, ) def issue_tdm_loads(buf_idx, k_tile): + if const_expr(k_rem != 0): + k_left = fx.Int32(split_k_chunk) - fx.Int32(k_tile) * tile_k for sd in (A_SIDE, B_SIDE): + atom = sd[ATOM] + if const_expr(k_rem != 0): + atom = fx.atom_set_value(atom, f"extent_{sd[KDIM]}", k_left) fx.copy( - sd[ATOM], + atom, sd[GT], fx.Tensor( fx.make_view( @@ -438,24 +439,45 @@ def wmma_tile(accs_in, a_frags, b_frags, rotate_buf=None): if const_expr(add_bias): bias_lay = fx.make_layout(4, 1) - bias_tiles = fx.logical_divide( - fx.rocdl.make_buffer_tensor( - fx.make_view(arg_bias, fx.make_layout(N, 1)), - max_size=False, - num_records_bytes=N * elem_bytes, - ), - bias_lay, + bias_buf = fx.rocdl.make_buffer_tensor( + fx.make_view(arg_bias, fx.make_layout(N, 1)), + max_size=False, + num_records_bytes=N * elem_bytes, ) + bias_tiles = fx.logical_divide(bias_buf, bias_lay) bias_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), _fx_elem) + if const_expr(n_edge != 0): + bias_elems = fx.logical_divide(bias_buf, fx.make_layout(1, 1)) + bias1_atom = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), _fx_elem) def load_bias4(tile_idx): - r = fx.make_rmem_tensor(bias_lay, _fx_elem) - fx.copy_atom_call(bias_atom, fx.slice(bias_tiles, (None, tile_idx)), r) - return r.load() + if const_expr(n_edge == 0): + r = fx.make_rmem_tensor(bias_lay, _fx_elem) + fx.copy(bias_atom, fx.slice(bias_tiles, (None, tile_idx)), r) + return r.load() + elems = [] + for e in range_constexpr(4): + r1 = fx.make_rmem_tensor(1, _fx_elem) + fx.copy( + bias1_atom, + fx.slice(bias_elems, (None, tile_idx * 4 + e)), + r1, + ) + elems.append(fx.Vector(r1.load())[0]) + return fx.Vector.from_elements(elems, _fx_elem) def epilogue_stores(final_accs): _acc_ty = _out_num if _half_out else fx.Float32 st_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), _acc_ty) + if const_expr(n_edge != 0): + st1_atom = fx.make_copy_atom( + ( + fx.rocdl.BufferCopy16b() + if _half_out + else fx.rocdl.BufferCopy32b() + ), + _acc_ty, + ) if const_expr(split_k > 1): add_atom = fx.make_copy_atom( fx.UniversalAtomicAdd(_acc_ty, syncscope=fx.rocdl.SyncScope.Agent), @@ -492,9 +514,7 @@ def epilogue_stores(final_accs): if const_expr(activation is not None): acc = fx.Vector.from_elements( [ - apply_activation_scalar( - fx.Vector(acc)[i].ir_value(), activation - ) + apply_activation_scalar(fx.Vector(acc)[i], activation) for i in range_constexpr(8) ], fx.Float32, @@ -502,28 +522,68 @@ def epilogue_stores(final_accs): if const_expr(_half_out): h_vec = fx.Vector(acc).to(_out_num) - c_off = row * n_stride + col_base + c_off = row * ld_y + col_base if const_expr(split_k > 1): for pair in range_constexpr(4): pair_vec = fx.Vector.from_elements( - [ - h_vec[pair * 2].ir_value(), - h_vec[pair * 2 + 1].ir_value(), - ], + [h_vec[pair * 2], h_vec[pair * 2 + 1]], _out_num, ) - if row < m_idx: - fx.copy( - add_atom, - _rmem_vec(pair_vec, 2, _out_num), - gYp[None, c_off + pair * 2], + if const_expr(n_edge == 0): + if row < m_idx: + fx.copy( + add_atom, + _rmem_vec(pair_vec, 2, _out_num), + gYp[None, c_off + pair * 2], + ) + else: + n_left = ( + fx.Int32(N) + - fx.Int32(col_base) + - fx.Int32(pair * 2) ) + if row < m_idx: + if n_left >= fx.Int32(2): + fx.copy( + add_atom, + _rmem_vec(pair_vec, 2, _out_num), + gYp[None, c_off + pair * 2], + ) + if n_left == fx.Int32(1): + one_vec = fx.Vector.from_elements( + [h_vec[pair * 2]], _out_num + ) + fx.copy( + add_atom, + _rmem_vec(one_vec, 1, _out_num), + gYp[None, c_off + pair * 2], + ) else: - fx.copy( - st_atom, - _rmem_vec(h_vec, 8, _out_num), - gY[None, c_off], - ) + if const_expr(n_edge == 0): + fx.copy( + st_atom, + _rmem_vec(h_vec, 8, _out_num), + gY[None, c_off], + ) + else: + n_left = fx.Int32(N) - fx.Int32(col_base) + if n_left >= fx.Int32(8): + fx.copy( + st_atom, + _rmem_vec(h_vec, 8, _out_num), + gY[None, c_off], + ) + if n_left < fx.Int32(8): + for e in range_constexpr(8): + if fx.Int32(e) < n_left: + one_vec = fx.Vector.from_elements( + [h_vec[e]], _out_num + ) + fx.copy( + st1_atom, + _rmem_vec(one_vec, 1, _out_num), + gY[None, c_off + e], + ) elif const_expr(split_k > 1): for e in range_constexpr(8): fx.ptr_store( @@ -542,30 +602,58 @@ def epilogue_stores(final_accs): xt, (lane_kgrp + 2 * e) * WMMA_N + lane16, ) - ).ir_value() + ) ], fx.Float32, ) - fx.copy( - add_atom, - _rmem_vec(xt_val, 1, fx.Float32), - gYp[None, xt_row * n_stride + t_col], - ) + if const_expr(n_edge == 0): + fx.copy( + add_atom, + _rmem_vec(xt_val, 1, fx.Float32), + gYp[None, xt_row * ld_y + t_col], + ) + else: + if t_col < fx.Uint64(N): + fx.copy( + add_atom, + _rmem_vec(xt_val, 1, fx.Float32), + gYp[None, xt_row * ld_y + t_col], + ) else: for half in range_constexpr(2): vec4 = fx.Vector.from_elements( [ - fx.Vector(acc)[half * 4 + vi].ir_value() + fx.Vector(acc)[half * 4 + vi] for vi in range_constexpr(4) ], fx.Float32, ) col = col_base + half * 4 - fx.copy( - st_atom, - _rmem_vec(vec4, 4, fx.Float32), - gY[None, row * n_stride + col], - ) + if const_expr(n_edge == 0): + fx.copy( + st_atom, + _rmem_vec(vec4, 4, fx.Float32), + gY[None, row * ld_y + col], + ) + else: + n_left = fx.Int32(N) - fx.Int32(col) + if n_left >= fx.Int32(4): + fx.copy( + st_atom, + _rmem_vec(vec4, 4, fx.Float32), + gY[None, row * ld_y + col], + ) + if n_left < fx.Int32(4): + for e in range_constexpr(4): + if fx.Int32(e) < n_left: + one_vec = fx.Vector.from_elements( + [vec4[e]], fx.Float32 + ) + fx.copy( + st1_atom, + _rmem_vec(one_vec, 1, fx.Float32), + gY[None, row * ld_y + col + e], + ) def _pack_state(accs_, a_, b_): return list(accs_) + list(a_) + list(b_) @@ -670,13 +758,13 @@ def launch_gemm_a16w16( arg_w: fx.Pointer, arg_bias: fx.Pointer, i32_m: fx.Int32, - i32_n: fx.Int32, + i32_ldy: fx.Int32, i32_lda: fx.Int32, i32_ldb: fx.Int32, stream: fx.Stream, ): gx = (fx.Uint64(i32_m) + (tile_m - 1)) // tile_m - gy = (fx.Uint64(i32_n) + (tile_n - 1)) // tile_n + gy = (N + tile_n - 1) // tile_n wpe = int(waves_per_eu) if waves_per_eu is not None else 0 launcher = kernel_gemm_a16w16( @@ -685,7 +773,7 @@ def launch_gemm_a16w16( arg_w, arg_bias, i32_m, - i32_n, + i32_ldy, i32_lda, i32_ldb, value_attrs={ diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index 12b8041c606..cb402b1c685 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -331,3 +331,98 @@ def test_gemm_a16_w16_oob_guard_region(M, N, K, tm, tn, tk, mw, nw, nb): torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) assert torch.all(parent[M:] == sentinel) + + +@pytest.mark.parametrize("K, tile_k", [(48, 32), (100, 32), (130, 64), (1000, 64)]) +@pytest.mark.parametrize("layout", ["TN", "NN", "TT", "NT"]) +def test_gemm_a16_w16_ragged_k(K, tile_k, layout, monkeypatch): + torch.cuda.empty_cache() + M, N = 128, 192 + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16, layout=layout) + + def _no_pad(*args, **kwargs): + raise AssertionError("host-side pad/copy in gemm_a16w16 wrapper") + + monkeypatch.setattr(F, "pad", _no_pad) + torch_out = (x.float() @ w.float().T).to(torch.bfloat16) + kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, tile_k=tile_k) + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +def test_gemm_a16_w16_ragged_k_split_k(): + torch.cuda.empty_cache() + M, N, K, sk = 64, 64, 576, 2 + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + torch_out = x.float() @ w.float().T + kernel_out = gemm_a16w16( + x, + w, + dtype=torch.float32, + tile_m=64, + tile_n=64, + tile_k=64, + m_warp=2, + n_warp=2, + num_buffers=2, + split_k=sk, + ) + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + + +@pytest.mark.parametrize("N", [100, 101, 190, 257]) +@pytest.mark.parametrize("otype", [torch.bfloat16, torch.float32]) +def test_gemm_a16_w16_ragged_n_direct_store(N, otype): + torch.cuda.empty_cache() + M, K = 129, 256 + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + sentinel = 777.0 + parent = torch.full((M + 32, N + 64), sentinel, dtype=otype, device="cuda") + y = parent[:M, :N] + + torch_out = (x.float() @ w.float().T).to(otype) + out = gemm_a16w16(x, w, dtype=otype, y=y) + + assert out is y + torch.testing.assert_close(out, torch_out, atol=1e-1, rtol=1e-2) + assert torch.all(parent[:, N:] == sentinel) + assert torch.all(parent[M:] == sentinel) + + +@pytest.mark.parametrize("N", [100, 190]) +def test_gemm_a16_w16_ragged_n_split_k(N): + torch.cuda.empty_cache() + M, K, sk = 64, 512, 2 + x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) + sentinel = 777.0 + parent = torch.full((M + 8, N + 32), sentinel, dtype=torch.float32, device="cuda") + y = parent[:M, :N] + + torch_out = x.float() @ w.float().T + out = gemm_a16w16( + x, + w, + dtype=torch.float32, + y=y, + tile_m=64, + tile_n=64, + tile_k=64, + m_warp=2, + n_warp=2, + num_buffers=2, + split_k=sk, + ) + + assert out is y + torch.testing.assert_close(out, torch_out, atol=1e-1, rtol=1e-2) + assert torch.all(parent[:, N:] == sentinel) + assert torch.all(parent[M:] == sentinel) + + +@pytest.mark.parametrize("M, N, K", [(100, 190, 256), (65, 257, 512)]) +def test_gemm_a16_w16_ragged_n_bias_activation(M, N, K): + torch.cuda.empty_cache() + x, w, bias, _ = _generate_inputs(M, N, K, torch.bfloat16, bias=True) + + torch_out = F.relu(F.linear(x, w, bias=bias)) + kernel_out = gemm_a16w16(x, w, bias=bias, dtype=torch.bfloat16, activation="relu") + torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) From e659da702616838bdce34f1e256ac0df22282531 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Tue, 8 Sep 2026 22:22:55 +0000 Subject: [PATCH 31/34] config picker, refactor, kernel change to prevent copies with split k --- .../model_configs/gptoss_bf16_tuned_gemm.csv | 6 +- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 82 ++- .../kernels/gemm_a16w16_kernel_gfx1250.py | 334 +++++---- aiter/tuned_gemm.py | 78 +- op_tests/flydsl_tests/test_gemm_a16w16.py | 695 ++++++++---------- 5 files changed, 552 insertions(+), 643 deletions(-) diff --git a/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv b/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv index 17ca8987492..8570bea2ded 100644 --- a/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv +++ b/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv @@ -256,6 +256,6 @@ gfx1250,96,2048,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,trit gfx1250,96,4096,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,96,8192,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,96,16384,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 -gfx1250,256,64,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl_gfx1250,0,2,6.82,flydsl_a16w16_gfx1250_t64x64x256_mw4_nw2_nb4_sk2_unroll1,0.0,276.75,4474.37 -gfx1250,256,64,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,flydsl_gfx1250,0,4,5.82,flydsl_a16w16_gfx1250_t64x64x128_mw4_nw2_nb6_sk4_unroll1,0.0,259.44,4207.20 -gfx1250,256,64,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl_gfx1250,0,6,3.0,flydsl_a16w16_gfx1250_t16x16x256_mw1_nw1_nb3_sk6_unroll1,0.0,15.73,374.10 +gfx1250,256,64,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,2,6.82,flydsl_a16w16_gfx1250_t64x64x256_mw4_nw2_nb4_sk2_unroll1,0.0,276.75,4474.37 +gfx1250,256,64,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,4,5.82,flydsl_a16w16_gfx1250_t64x64x128_mw4_nw2_nb6_sk4_unroll1,0.0,259.44,4207.20 +gfx1250,256,64,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,6,3.0,flydsl_a16w16_gfx1250_t16x16x256_mw1_nw1_nb3_sk6_unroll1,0.0,15.73,374.10 diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index f71f6935615..7660dada613 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -6,6 +6,7 @@ from __future__ import annotations import functools +import re import torch @@ -43,6 +44,7 @@ def _p(t): torch.bfloat16: _fx.BFloat16, torch.float16: _fx.Float16, torch.float32: _fx.Float32, + torch.int32: _fx.Int32, } ) return _ptr_arg(t, _FX_DTYPE[t.dtype]) @@ -77,6 +79,47 @@ def _cached_launcher(*cfg): return _compile_gemm_a16w16(**dict(zip(_CFG_KEYS, cfg))) +# flydsl_a16w16_gfx1250_t{tm}x{tn}x{tk}_mw{mw}_nw{nw}_nb{nb}_sk{sk}_unroll{0|1} +_KERNEL_NAME_RE = re.compile( + r"^flydsl_a16w16_gfx1250_" + r"t(?P\d+)x(?P\d+)x(?P\d+)_" + r"mw(?P\d+)_nw(?P\d+)_" + r"nb(?P\d+)_sk(?P\d+)_unroll(?P[01])$" +) + + +def parse_kernel_name(name: str): + """Parse a flydsl_a16w16_gfx1250_ kernelName into its config dict, or None.""" + m = _KERNEL_NAME_RE.fullmatch(name) + if m is None: + return None + cfg = {k: int(v) for k, v in m.groupdict().items()} + cfg["main_loop_unroll"] = bool(cfg["main_loop_unroll"]) + return cfg + + +def run_gemm_a16w16_gfx1250( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None, + dtype: torch.dtype, + kernel_name: str, +) -> torch.Tensor: + """Dispatch entry: decode a tuned kernelName and run the kernel.""" + cfg = parse_kernel_name(kernel_name) + if cfg is None: + raise ValueError(f"[FlyDSL gfx1250] unrecognised kernelName: {kernel_name!r}") + return gemm_a16w16(x, w, bias=bias, dtype=dtype, **cfg) + + +_SPLIT_K_MAX_TILES = 1 << 16 + + +@functools.cache +def _split_k_counters(device, stream): + return torch.zeros(_SPLIT_K_MAX_TILES, dtype=torch.int32, device=device) + + def gemm_a16w16( x: torch.Tensor, w: torch.Tensor, @@ -114,15 +157,10 @@ def gemm_a16w16( physical_mk = x.stride(1) == 1 or K == 1 physical_kn = w.stride(1) != 1 and N > 1 - _splitk_f32_accum = split_k > 1 and dtype in _half in_dtype_str = "fp16" if x.dtype == torch.float16 else "bf16" out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get(dtype, "f32") - buf_dtype = dtype - if _splitk_f32_accum: - out_dtype_str, buf_dtype = "f32", torch.float32 - _alloc = torch.zeros if split_k > 1 else torch.empty - if y is not None and not _splitk_f32_accum: + if y is not None: assert ( y.shape[0] == M and y.shape[1] == N ), f"y must be ({M}, {N}), got {tuple(y.shape)}" @@ -131,16 +169,27 @@ def gemm_a16w16( N, ), f"gemm_a16w16: y needs unit column stride, got strides {tuple(y.stride())}" y_buf = y - if split_k > 1: - y_buf.zero_() else: - y_buf = _alloc((M, N), device=x.device, dtype=buf_dtype) + y_buf = torch.empty((M, N), device=x.device, dtype=dtype) ldy = y_buf.stride(0) if M > 1 else N assert ldy >= N, f"gemm_a16w16: y row stride {ldy} < N {N}" if bias is None: bias = torch.empty(0, device=x.device, dtype=dtype) + stream = torch.cuda.current_stream(device=x.device) + sem = _split_k_counters(x.device, stream) + if split_k > 1: + tiles = ((M + tile_m - 1) // tile_m) * ((N + tile_n - 1) // tile_n) + assert ( + tiles <= sem.numel() + ), f"gemm_a16w16: split_k needs {tiles} tile counters, max {sem.numel()}" + m_pad = ((M + tile_m - 1) // tile_m) * tile_m + n_pad = ((N + tile_n - 1) // tile_n) * tile_n + ws = torch.empty((split_k, m_pad, n_pad), device=x.device, dtype=torch.float32) + else: + ws = y_buf + launch_fn = _cached_launcher( N, K, @@ -172,26 +221,21 @@ def gemm_a16w16( ldb = w.stride(1) if K > 1 else N else: ldb = w.stride(0) if N > 1 else K - stream = torch.cuda.current_stream(device=x.device).cuda_stream _run_compiled( launch_fn, _p(y_buf), _p(x), _p(w), _p(bias), + _p(ws), + _p(sem), M, ldy, lda, ldb, - _fx.Stream(stream), + _fx.Stream(stream.cuda_stream), ) - - if not _splitk_f32_accum: - return y_buf - if y is not None: - y.copy_(y_buf) - return y - return y_buf.to(dtype) + return y_buf -__all__ = ["gemm_a16w16"] +__all__ = ["gemm_a16w16", "parse_kernel_name", "run_gemm_a16w16_gfx1250"] diff --git a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py index 29ca0adb70f..cfb9ab26e89 100644 --- a/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_a16w16_kernel_gfx1250.py @@ -11,9 +11,12 @@ from flydsl.runtime.device import get_rocm_arch as get_hip_arch from flydsl.utils.smem_allocator import check_smem_capacity +from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops + WMMA_M, WMMA_N, WMMA_K = 16, 16, 32 WAVE_SIZE, LDS_PAD_A, LDS_PAD_B = 32, 8, 8 _SCHED_ALLOW_SALU = 1 << 2 +CPOL_DEV = 0x10 KERNARG_PRELOAD_COUNT = 8 @@ -79,11 +82,6 @@ def _req(cond, msg): f"bad sched_strategy {sched_strategy!r}", ) _req(split_k >= 1, f"split_k must be >= 1, got {split_k}") - _req( - not (activation and split_k > 1), - "activation is applied per k-split partial; activation with split_k > 1 " - "would compute act(partial) summed instead of act(sum)", - ) elem_bytes = 2 elem_bytes_d = 2 if out_dtype in ("f16", "bf16") else 4 @@ -144,15 +142,13 @@ def _req(cond, msg): b_imm_bytes, lds_b_stride = tile_k * elem_bytes, tile_k + LDS_PAD_B lds_b_elems = tile_n * lds_b_stride + LDS_PAD_B - _USE_XT = split_k > 1 and out_dtype == "f32" - - if _USE_XT: + if split_k > 1: @fx.struct class SharedStorage: a: fx.Array[_fx_elem, num_buffers * lds_a_elems, 16] b: fx.Array[_fx_elem, num_buffers * lds_b_elems, 16] - xt: fx.Array[fx.Float32, num_warps * WMMA_M * WMMA_N, 16] + flag: fx.Array[fx.Int32, 1, 16] else: @@ -172,6 +168,8 @@ def kernel_gemm_a16w16( arg_x: fx.Pointer, arg_w: fx.Pointer, arg_bias: fx.Pointer, + arg_ws: fx.Pointer, + arg_sem: fx.Pointer, i32_m: fx.Int32, i32_ldy: fx.Int32, i32_lda: fx.Int32, @@ -200,12 +198,40 @@ def kernel_gemm_a16w16( warp_m_base, warp_n_base = wave_m_idx * warp_tile_m, wave_n_idx * warp_tile_n m_idx, ld_y = fx.Uint64(i32_m), fx.Uint64(i32_ldy) - gYp = fx.Tensor(fx.make_view(arg_y, fx.make_layout((1, 1), (1, 1)))) gY = fx.rocdl.make_buffer_tensor( fx.Tensor(fx.make_view(arg_y, fx.make_layout((1, 1), (1, 1)))), max_size=False, num_records_bytes=m_idx * ld_y * elem_bytes_d, ) + if const_expr(split_k > 1): + tiles_n = (N + tile_n - 1) // tile_n + n_pad = tiles_n * tile_n + m_pad = ((m_idx + (tile_m - 1)) // tile_m) * tile_m + plane = fx.Int64(m_pad) * n_pad + + def _ws_view(plane_idx): + return fx.rocdl.make_buffer_tensor( + fx.Tensor( + fx.make_view( + fx.add_offset(arg_ws, fx.Int64(plane_idx) * plane), + fx.make_layout((1, 1), (1, 1)), + ) + ), + max_size=False, + num_records_bytes=plane * 4, + ) + + gWs_mine = _ws_view(gpu.block_id("z")) + gWs = [_ws_view(s) for s in range_constexpr(split_k)] + tile_lin = fx.Uint64(bx) * tiles_n + fx.Uint64(by) + warp_lin = wave_m_idx * n_warp + wave_n_idx + ws_lane_base = (tile_lin * num_warps + warp_lin) * ( + n_accs * WAVE_SIZE * 8 + ) + (lane_kgrp * 16 + lane16) * 8 + sem_addr = ( + fx.Int64(fx.ptrtoint(arg_sem)) + + (fx.Int64(bx) * tiles_n + fx.Int64(by)) * 4 + ) lds = fx.SharedAllocator(static=True).allocate(SharedStorage).peek() big_a_mem, big_b_mem = lds.a.ptr, lds.b.ptr @@ -479,14 +505,8 @@ def epilogue_stores(final_accs): _acc_ty, ) if const_expr(split_k > 1): - add_atom = fx.make_copy_atom( - fx.UniversalAtomicAdd(_acc_ty, syncscope=fx.rocdl.SyncScope.Agent), - _acc_ty, - ) - if const_expr(_USE_XT): - xt = fx.add_offset( - lds.xt.ptr, - (wave_m_idx * n_warp + wave_n_idx) * (WMMA_M * WMMA_N), + ws_atom = fx.make_copy_atom( + fx.rocdl.BufferCopy128b(CPOL_DEV), fx.Float32 ) if const_expr(add_bias): bias_vecs = [] @@ -496,164 +516,148 @@ def epilogue_stores(final_accs): for half in range_constexpr(2): bv = load_bias4(col_tile + half) for i in range_constexpr(4): - if const_expr(split_k > 1): - elems.append(bv[i].to(fx.Float32) * (1.0 / split_k)) - else: - elems.append(bv[i].to(fx.Float32)) + elems.append(bv[i].to(fx.Float32)) bias_vecs.append( fx.Vector.from_elements(elems, fx.Float32).ir_value() ) - for wm in range_constexpr(wmma_m_rep): - for wn in range_constexpr(wmma_n_rep): - acc = final_accs[wm * wmma_n_rep + wn] - row = blk_m + warp_m_base + wm * WMMA_M + lane16 - col_base = blk_n + warp_n_base + wn * WMMA_N + lane_kgrp * 8 - if const_expr(add_bias): - acc = acc + bias_vecs[wn] - if const_expr(activation is not None): - acc = fx.Vector.from_elements( - [ - apply_activation_scalar(fx.Vector(acc)[i], activation) - for i in range_constexpr(8) - ], - fx.Float32, - ).ir_value() - - if const_expr(_half_out): - h_vec = fx.Vector(acc).to(_out_num) - c_off = row * ld_y + col_base - if const_expr(split_k > 1): - for pair in range_constexpr(4): - pair_vec = fx.Vector.from_elements( - [h_vec[pair * 2], h_vec[pair * 2 + 1]], - _out_num, - ) - if const_expr(n_edge == 0): - if row < m_idx: - fx.copy( - add_atom, - _rmem_vec(pair_vec, 2, _out_num), - gYp[None, c_off + pair * 2], - ) - else: - n_left = ( - fx.Int32(N) - - fx.Int32(col_base) - - fx.Int32(pair * 2) - ) - if row < m_idx: - if n_left >= fx.Int32(2): - fx.copy( - add_atom, - _rmem_vec(pair_vec, 2, _out_num), - gYp[None, c_off + pair * 2], - ) - if n_left == fx.Int32(1): - one_vec = fx.Vector.from_elements( - [h_vec[pair * 2]], _out_num - ) - fx.copy( - add_atom, - _rmem_vec(one_vec, 1, _out_num), - gYp[None, c_off + pair * 2], - ) - else: - if const_expr(n_edge == 0): - fx.copy( - st_atom, - _rmem_vec(h_vec, 8, _out_num), - gY[None, c_off], - ) - else: - n_left = fx.Int32(N) - fx.Int32(col_base) - if n_left >= fx.Int32(8): - fx.copy( - st_atom, - _rmem_vec(h_vec, 8, _out_num), - gY[None, c_off], - ) - if n_left < fx.Int32(8): - for e in range_constexpr(8): - if fx.Int32(e) < n_left: - one_vec = fx.Vector.from_elements( - [h_vec[e]], _out_num - ) - fx.copy( - st1_atom, - _rmem_vec(one_vec, 1, _out_num), - gY[None, c_off + e], - ) - elif const_expr(split_k > 1): - for e in range_constexpr(8): - fx.ptr_store( - fx.Vector(acc)[e], - fx.add_offset(xt, lane16 * WMMA_N + lane_kgrp * 8 + e), + def _vec4(acc, half): + return fx.Vector.from_elements( + [fx.Vector(acc)[half * 4 + vi] for vi in range_constexpr(4)], + fx.Float32, + ) + + def _ws_off(acc_idx): + return ws_lane_base + acc_idx * (WAVE_SIZE * 8) + + def store_partial(acc, acc_idx): + for half in range_constexpr(2): + fx.copy( + ws_atom, + _rmem_vec(_vec4(acc, half), 4, fx.Float32), + gWs_mine[None, _ws_off(acc_idx) + half * 4], + ) + + def issue_partial_loads(acc_idx): + """Issue this lane's loads of all split_k planes (no waits here).""" + frags = [] + for half in range_constexpr(2): + for sidx in range_constexpr(split_k): + r = fx.make_rmem_tensor(4, fx.Float32) + fx.copy( + ws_atom, gWs[sidx][None, _ws_off(acc_idx) + half * 4], r + ) + frags.append(r) + return frags + + def sum_partials(frags): + """Consume issue_partial_loads' fragments into one 8-wide fp32 vector.""" + elems = [] + for half in range_constexpr(2): + total = None + for sidx in range_constexpr(split_k): + v = fx.Vector(frags[half * split_k + sidx].load()) + total = v if sidx == 0 else total + v + for vi in range_constexpr(4): + elems.append(total[vi]) + return fx.Vector.from_elements(elems, fx.Float32) + + def store_final(acc, row, col_base, wn): + if const_expr(add_bias): + acc = acc + bias_vecs[wn] + if const_expr(activation is not None): + acc = fx.Vector.from_elements( + [ + apply_activation_scalar(fx.Vector(acc)[i], activation) + for i in range_constexpr(8) + ], + fx.Float32, + ).ir_value() + if const_expr(_half_out): + h_vec = fx.Vector(acc).to(_out_num) + c_off = row * ld_y + col_base + if const_expr(n_edge == 0): + fx.copy(st_atom, _rmem_vec(h_vec, 8, _out_num), gY[None, c_off]) + else: + n_left = fx.Int32(N) - fx.Int32(col_base) + if n_left >= fx.Int32(8): + fx.copy( + st_atom, _rmem_vec(h_vec, 8, _out_num), gY[None, c_off] ) - t_base = blk_m + warp_m_base + wm * WMMA_M - t_col = blk_n + warp_n_base + wn * WMMA_N + lane16 - for e in range_constexpr(8): - xt_row = t_base + lane_kgrp + 2 * e - if xt_row < m_idx: - xt_val = fx.Vector.from_elements( - [ - fx.ptr_load( - fx.add_offset( - xt, - (lane_kgrp + 2 * e) * WMMA_N + lane16, - ) - ) - ], - fx.Float32, - ) - if const_expr(n_edge == 0): + if n_left < fx.Int32(8): + for e in range_constexpr(8): + if fx.Int32(e) < n_left: + one_vec = fx.Vector.from_elements( + [h_vec[e]], _out_num + ) fx.copy( - add_atom, - _rmem_vec(xt_val, 1, fx.Float32), - gYp[None, xt_row * ld_y + t_col], + st1_atom, + _rmem_vec(one_vec, 1, _out_num), + gY[None, c_off + e], ) - else: - if t_col < fx.Uint64(N): - fx.copy( - add_atom, - _rmem_vec(xt_val, 1, fx.Float32), - gYp[None, xt_row * ld_y + t_col], - ) - else: - for half in range_constexpr(2): - vec4 = fx.Vector.from_elements( - [ - fx.Vector(acc)[half * 4 + vi] - for vi in range_constexpr(4) - ], - fx.Float32, + else: + for half in range_constexpr(2): + vec4 = _vec4(acc, half) + col = col_base + half * 4 + if const_expr(n_edge == 0): + fx.copy( + st_atom, + _rmem_vec(vec4, 4, fx.Float32), + gY[None, row * ld_y + col], ) - col = col_base + half * 4 - if const_expr(n_edge == 0): + else: + n_left = fx.Int32(N) - fx.Int32(col) + if n_left >= fx.Int32(4): fx.copy( st_atom, _rmem_vec(vec4, 4, fx.Float32), gY[None, row * ld_y + col], ) - else: - n_left = fx.Int32(N) - fx.Int32(col) - if n_left >= fx.Int32(4): - fx.copy( - st_atom, - _rmem_vec(vec4, 4, fx.Float32), - gY[None, row * ld_y + col], - ) - if n_left < fx.Int32(4): - for e in range_constexpr(4): - if fx.Int32(e) < n_left: - one_vec = fx.Vector.from_elements( - [vec4[e]], fx.Float32 - ) - fx.copy( - st1_atom, - _rmem_vec(one_vec, 1, fx.Float32), - gY[None, row * ld_y + col + e], - ) + if n_left < fx.Int32(4): + for e in range_constexpr(4): + if fx.Int32(e) < n_left: + one_vec = fx.Vector.from_elements( + [vec4[e]], fx.Float32 + ) + fx.copy( + st1_atom, + _rmem_vec(one_vec, 1, fx.Float32), + gY[None, row * ld_y + col + e], + ) + + def _coords(wm, wn): + row = blk_m + warp_m_base + wm * WMMA_M + lane16 + col_base = blk_n + warp_n_base + wn * WMMA_N + lane_kgrp * 8 + return row, col_base + + if const_expr(split_k == 1): + for wm in range_constexpr(wmma_m_rep): + for wn in range_constexpr(wmma_n_rep): + row, col_base = _coords(wm, wn) + store_final(final_accs[wm * wmma_n_rep + wn], row, col_base, wn) + else: + for acc_idx in range_constexpr(n_accs): + store_partial(final_accs[acc_idx], acc_idx) + comm_ops.waitcnt_all() + gpu.barrier() + if fx.Int32(tx) == fx.Int32(0): + arrival = fx.Int32(comm_ops.atomic_add_agent(sem_addr, fx.Int32(1))) + fx.ptr_store( + (arrival == fx.Int32(split_k - 1)).select( + fx.Int32(1), fx.Int32(0) + ), + lds.flag.ptr, + ) + gpu.barrier() + if fx.ptr_load(lds.flag.ptr) != fx.Int32(0): + pending = [issue_partial_loads(i) for i in range_constexpr(n_accs)] + for wm in range_constexpr(wmma_m_rep): + for wn in range_constexpr(wmma_n_rep): + row, col_base = _coords(wm, wn) + total = sum_partials(pending[wm * wmma_n_rep + wn]) + store_final(total, row, col_base, wn) + if fx.Int32(tx) == fx.Int32(0): + comm_ops.atomic_add_agent(sem_addr, fx.Int32(-split_k)) def _pack_state(accs_, a_, b_): return list(accs_) + list(a_) + list(b_) @@ -757,6 +761,8 @@ def launch_gemm_a16w16( arg_x: fx.Pointer, arg_w: fx.Pointer, arg_bias: fx.Pointer, + arg_ws: fx.Pointer, + arg_sem: fx.Pointer, i32_m: fx.Int32, i32_ldy: fx.Int32, i32_lda: fx.Int32, @@ -772,6 +778,8 @@ def launch_gemm_a16w16( arg_x, arg_w, arg_bias, + arg_ws, + arg_sem, i32_m, i32_ldy, i32_lda, diff --git a/aiter/tuned_gemm.py b/aiter/tuned_gemm.py index 58c4a377bc6..69f5ab56b41 100644 --- a/aiter/tuned_gemm.py +++ b/aiter/tuned_gemm.py @@ -162,9 +162,16 @@ def get_GEMM_A16W16_config( if config is not None: if config["libtype"] == "flydsl": if is_flydsl_available(): - flydsl_config = aiter.ops.flydsl.gemm_kernels.get_flydsl_splitk_hgemm_kernel_params( - config["kernelName"] - ) + if get_gfx() == "gfx1250": + from aiter.ops.flydsl.gemm_a16w16_gfx1250 import ( + parse_kernel_name, + ) + + flydsl_config = parse_kernel_name(config["kernelName"]) + else: + flydsl_config = aiter.ops.flydsl.gemm_kernels.get_flydsl_splitk_hgemm_kernel_params( + config["kernelName"] + ) if flydsl_config is None: logger.warning( f"FlyDSL kernel '{config['kernelName']}' from tuned config is not " @@ -471,63 +478,6 @@ def asm_gemm( return gemm_a16w16_asm(inp, weights, out_asm, bias, splitK, kernelName, bpreshuffle) -_FLY1250_NAME_RE = None - - -def _parse_flydsl_gfx1250_kname(name: str) -> dict: - """Parse flydsl_a16w16_gfx1250_t{tm}x{tn}x{tk}_mw{mw}_nw{nw}_nb{nb}_sk{sk}_unroll{0|1}.""" - global _FLY1250_NAME_RE - if _FLY1250_NAME_RE is None: - import re - - _FLY1250_NAME_RE = re.compile( - r"flydsl_a16w16_gfx1250_t(\d+)x(\d+)x(\d+)_mw(\d+)_nw(\d+)" - r"_nb(\d+)_sk(\d+)_unroll([01])" - ) - m = _FLY1250_NAME_RE.match(name) - assert m is not None, f"unrecognized flydsl_gfx1250 kernel name: {name!r}" - tm, tn, tk, mw, nw, nb, sk, un = (int(g) for g in m.groups()) - return { - "tile_m": tm, - "tile_n": tn, - "tile_k": tk, - "m_warp": mw, - "n_warp": nw, - "num_buffers": nb, - "split_k": sk, - "main_loop_unroll": bool(un), - } - - -def flydsl_gfx1250_gemm( - inp: Tensor, - weights: Tensor, - solidx: int, - bias: Tensor | None = None, - otype: torch.dtype | None = None, - scale_a: Tensor | None = None, - scale_b: Tensor | None = None, - scale_c: Tensor | None = None, - bpreshuffle=False, - config: dict | None = None, -): - assert ( - scale_a is None and scale_b is None and scale_c is None - ), "flydsl_gfx1250 a16w16 does not support scaling." - assert not bpreshuffle, "flydsl_gfx1250 a16w16 expects unshuffled weights." - from aiter.ops.flydsl.gemm_a16w16_gfx1250 import gemm_a16w16 as _fly1250 - - cfg = _parse_flydsl_gfx1250_kname(config["kernelName"]) - return _fly1250( - inp, - weights, - bias=bias, - dtype=otype if otype is not None else inp.dtype, - waves_per_eu=None, - **cfg, - ) - - def flydsl_gemm( inp: Tensor, weights: Tensor, @@ -543,6 +493,13 @@ def flydsl_gemm( assert ( scale_a is None and scale_b is None and scale_c is None ), "FlyDSL hgemm does not support scaling yet." + # gfx1250 runs the FlyDSL a16w16 WMMA kernel; other archs use the hgemm path. + if get_gfx() == "gfx1250": + from aiter.ops.flydsl.gemm_a16w16_gfx1250 import run_gemm_a16w16_gfx1250 + + return run_gemm_a16w16_gfx1250( + inp, weights, bias, otype or inp.dtype, config["kernelName"] + ) flydsl_config = aiter.ops.flydsl.gemm_kernels.get_flydsl_splitk_hgemm_kernel_params( config["kernelName"] ) @@ -666,7 +623,6 @@ def triton_gemm( "asm": asm_gemm, "triton": triton_gemm, "flydsl": flydsl_gemm, - "flydsl_gfx1250": flydsl_gfx1250_gemm, "opus": opus_gemm, } diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index cb402b1c685..ed9dc2ee676 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -1,428 +1,329 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -import importlib.util -import os +import argparse +import itertools +from functools import partial -import pytest +import pandas as pd import torch import torch.nn.functional as F -_KERNEL_PATH = os.path.abspath( - os.path.join( - os.path.dirname(__file__), - "..", - "..", - "aiter", - "ops", - "flydsl", - "gemm_a16w16_gfx1250.py", - ) -) - - -def _get_gpu_arch(): - if not torch.cuda.is_available(): - return None - return getattr(torch.cuda.get_device_properties(0), "gcnArchName", None) - - -def _flydsl_available(): - if importlib.util.find_spec("flydsl") is None: - return False - arch = _get_gpu_arch() - return arch is not None and arch.startswith("gfx1250") - - -if not _flydsl_available(): - pytest.skip( - "FlyDSL a16w16 tests require gfx1250 and the flydsl package.", - allow_module_level=True, - ) - - -def _load_kernel(): - """Load the kernel module by file path to skip aiter.ops.flydsl.__init__.""" - spec = importlib.util.spec_from_file_location( - "_flydsl_a16w16_wrapper", _KERNEL_PATH - ) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - return mod.gemm_a16w16 - - -gemm_a16w16 = _load_kernel() - - -def _generate_inputs(M, N, K, dtype, layout="TN", output=False, bias=False): +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl.gemm_a16w16_gfx1250 import gemm_a16w16 as flydsl_gemm_a16w16 +from aiter.ops.flydsl.utils import is_flydsl_available +from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 as triton_gemm_a16w16 +from aiter.test_common import benchmark, checkAllclose, run_perftest + +torch.set_default_device("cuda") + +SUPPORTED_GFX = ["gfx1250"] +LAYOUTS = ["TN", "TT", "NN", "NT"] +ACTIVATIONS = { + "gelu": F.gelu, + "gelu_tanh": partial(F.gelu, approximate="tanh"), + "silu": F.silu, + "silu_exp2": F.silu, + "relu": F.relu, +} +SENTINEL = 4096.0 + + +def generate_inputs(m, n, k, dtype, layout="TN", bias=False): if layout[0] == "T": - x = torch.randn((M, K), dtype=dtype, device="cuda") + x = torch.randn(m, k, dtype=dtype) else: - x = torch.randn((K, M), dtype=dtype, device="cuda").T - + x = torch.randn(k, m, dtype=dtype).T if layout[1] == "T": - w = torch.randn((K, N), dtype=dtype, device="cuda").T + w = torch.randn(k, n, dtype=dtype).T else: - w = torch.randn((N, K), dtype=dtype, device="cuda") - - bias_tensor = None - if bias: - bias_tensor = torch.randn((N,), dtype=dtype, device="cuda") - - y = torch.empty((M, N), dtype=dtype, device="cuda") if output else None - return x, w, bias_tensor, y - - -def get_x_vals(): - return [ - # Aligned - (64, 64, 64), - (128, 128, 128), - (256, 256, 256), - # Multi-block - (128, 256, 512), - (256, 512, 256), - # Asymmetric - (32, 256, 128), - (256, 32, 128), - # Large K (drives the flat main loop + drain) - (128, 128, 1024), - (1024, 128, 128), - ] - - -def get_fewer_x_vals(): - return [ - (64, 64, 64), - (128, 256, 512), - (256, 512, 256), - (128, 128, 1024), - (1024, 128, 128), - ] - - -@pytest.mark.parametrize("M, N, K", get_x_vals()) -def test_gemm_a16_w16(M, N, K): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, num_buffers=3) - - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("num_buffers", [2, 3]) -def test_gemm_a16_w16_compute_bound(M, N, K, num_buffers): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16( - x, w, dtype=torch.bfloat16, num_buffers=num_buffers, variant="compute_bound" - ) - - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize("activation", ["gelu", "gelu_tanh", "silu", "silu_exp2"]) -@pytest.mark.parametrize("M, N, K", get_fewer_x_vals()) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("output", [True, False]) -def test_gemm_a16_w16_activation(M, N, K, dtype, output, activation): - torch.cuda.empty_cache() - x, w, _, y = _generate_inputs(M, N, K, dtype, output=output) - - torch_out = F.linear(x, w, bias=None) - if activation == "gelu": - torch_out = F.gelu(torch_out) - elif activation == "gelu_tanh": - torch_out = F.gelu(torch_out, approximate="tanh") - elif activation in ("silu", "silu_exp2"): - torch_out = F.silu(torch_out) - - kernel_out = gemm_a16w16( - x, w, dtype=dtype, y=y, activation=activation, num_buffers=3 - ) - - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize("M, N, K", get_x_vals()) -@pytest.mark.parametrize("layout", ["TN", "TT", "NN", "NT"]) -def test_gemm_a16_w16_layout(M, N, K, layout): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16, layout=layout) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, num_buffers=3) - - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-1) - + w = torch.randn(n, k, dtype=dtype) + b = torch.randn(n, dtype=dtype) if bias else None + return x, w, b -# 32x32 output tile, deep K (tile_k=128); tile_n=32 needs n_warp=2, K multiple of 128. -@pytest.mark.parametrize( - "M, N, K", [(32, 32, 256), (128, 128, 256), (256, 256, 512), (64, 64, 1024)] -) -def test_gemm_a16_w16_tile_32x32x128(M, N, K): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16( - x, - w, - dtype=torch.bfloat16, - num_buffers=3, - tile_m=32, - tile_n=32, - tile_k=128, - m_warp=2, - n_warp=2, +def run_torch(x, w, bias=None, activation=None, dtype=dtypes.bf16): + out = F.linear( + x.to(dtypes.fp32), + w.to(dtypes.fp32), + None if bias is None else bias.to(dtypes.fp32), ) - - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -# Split-K: partition K across split_k grid-z workgroups, atomic-add into a pre-zeroed C. -def _splitk_configs(): - # (M, N, K, tile_m, tile_n, tile_k, m_warp, n_warp, num_buffers, split_k) — all valid. - return [ - (128, 128, 512, 64, 64, 128, 2, 2, 2, 2), # 4 tiles, 2/split - (128, 128, 512, 64, 64, 128, 2, 2, 2, 4), # 4 tiles, 1/split - (64, 64, 768, 64, 64, 128, 2, 2, 2, 3), # 6 tiles, 2/split - (64, 64, 768, 64, 64, 128, 2, 2, 2, 6), # 6 tiles, 1/split - (128, 256, 1024, 64, 64, 128, 2, 2, 2, 4), # 8 tiles, 2/split - (256, 128, 1024, 32, 32, 128, 2, 2, 2, 8), # 8 tiles, 1/split (max split) - (128, 128, 1024, 64, 64, 128, 2, 2, 3, 4), # nb=3, 2/split - ] - - -@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb,sk", _splitk_configs()) -def test_gemm_a16_w16_split_k(M, N, K, tm, tn, tk, mw, nw, nb, sk): - """Primary split-K correctness: fp32 output (lossless) -> split_k=N == split_k=1.""" - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - common = { - "dtype": torch.float32, - "tile_m": tm, - "tile_n": tn, - "tile_k": tk, - "m_warp": mw, - "n_warp": nw, - "num_buffers": nb, + if activation is not None: + out = ACTIVATIONS[activation](out) + return out.to(dtype) + + +def run_candidates(candidates, ref, m, n, k, in_bytes, out_bytes, msg): + # [m,k] @ [n,k]^T: FLOPs = 2*m*n*k; bytes = x + w read, y written. + flops = 2 * m * n * k + nbytes = (m * k + n * k) * in_bytes + m * n * out_bytes + ret = {"gfx": get_gfx()} + for name, fn in candidates.items(): + out, us = run_perftest(fn) + err = checkAllclose( + ref.to(dtypes.fp32), + out.to(dtypes.fp32), + rtol=1e-2, + atol=1e-2, + msg=f"{name}: {msg}", + ) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} err"] = err + return ret + + +@benchmark() +def test_gemm_a16w16(m, n, k, dtype, layout): + x, w, _ = generate_inputs(m, n, k, dtype, layout) + ref = run_torch(x, w, dtype=dtype) + candidates = { + "flydsl": lambda: flydsl_gemm_a16w16(x, w, dtype=dtype), + "flydsl_cb": lambda: flydsl_gemm_a16w16( + x, w, dtype=dtype, variant="compute_bound" + ), + "triton": lambda: triton_gemm_a16w16(x, w, dtype=dtype), } - out_single = gemm_a16w16(x, w, split_k=1, **common).clone() - out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() - - torch.testing.assert_close(out_split, out_single, atol=1e-2, rtol=1e-3) - + return run_candidates( + candidates, ref, m, n, k, x.element_size(), ref.element_size(), "gemm a16w16" + ) -@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb,sk", _splitk_configs()) -def test_gemm_a16_w16_split_k_bf16(M, N, K, tm, tn, tk, mw, nw, nb, sk): - """bf16-output split-K at O(1) magnitudes (scaled inputs) must match split_k=1.""" - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - x = (x.float() * 0.1).to(torch.bfloat16) - w = (w.float() * 0.1).to(torch.bfloat16) - common = { - "dtype": torch.bfloat16, - "tile_m": tm, - "tile_n": tn, - "tile_k": tk, - "m_warp": mw, - "n_warp": nw, - "num_buffers": nb, +@benchmark() +def test_gemm_a16w16_fused(m, n, k, dtype, activation, bias): + x, w, b = generate_inputs(m, n, k, dtype, bias=bias) + ref = run_torch(x, w, b, activation, dtype) + candidates = { + "flydsl": lambda: flydsl_gemm_a16w16( + x, w, bias=b, dtype=dtype, activation=activation + ), + "triton": lambda: triton_gemm_a16w16( + x, w, bias=b, dtype=dtype, activation=activation + ), } - out_single = gemm_a16w16(x, w, split_k=1, **common).clone() - out_split = gemm_a16w16(x, w, split_k=sk, **common).clone() - - torch.testing.assert_close(out_split, out_single, atol=1e-1, rtol=2e-2) - + return run_candidates( + candidates, + ref, + m, + n, + k, + x.element_size(), + ref.element_size(), + f"gemm a16w16 + {activation}", + ) -@pytest.mark.parametrize("sk", [2, 4]) -def test_gemm_a16_w16_split_k_output_buffer_zeroed(sk): - """split_k>1 atomic-ADDS into C; the wrapper must zero it each call (no doubling).""" - M, N, K = 128, 128, 512 - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - y = torch.empty((M, N), dtype=torch.float32, device="cuda") - common = { - "dtype": torch.float32, +@benchmark() +def test_gemm_a16w16_config( + m, + n, + k, + dtype, + otype, + tile_m, + tile_n, + tile_k, + m_warp, + n_warp, + num_buffers, + split_k, + unroll, +): + x, w, _ = generate_inputs(m, n, k, dtype) + ref = run_torch(x, w, dtype=otype) + parent = torch.full((m + tile_m, n + tile_n), SENTINEL, dtype=otype) + y = parent[:m, :n] + cfg = { + "dtype": otype, "y": y, - "tile_m": 64, - "tile_n": 64, - "tile_k": 128, - "m_warp": 2, - "n_warp": 2, - "num_buffers": 2, - "split_k": sk, + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "m_warp": m_warp, + "n_warp": n_warp, + "num_buffers": num_buffers, + "split_k": split_k, + "main_loop_unroll": bool(unroll), } - out1 = gemm_a16w16(x, w, **common).clone() - out2 = gemm_a16w16(x, w, **common).clone() # reuse y -> must re-zero, not 2x - - torch.testing.assert_close(out2, out1, atol=1e-2, rtol=1e-3) - - -def _oob_configs(): - return [ - (100, 100, 256, 128, 128, 32, 2, 4, 3), - (129, 257, 512, 128, 128, 32, 2, 4, 3), - (65, 190, 256, 128, 128, 32, 2, 4, 3), - (250, 120, 512, 64, 64, 128, 2, 2, 2), - (100, 128, 256, 128, 128, 32, 2, 4, 3), - (128, 100, 256, 128, 128, 32, 2, 4, 3), - (33, 65, 256, 32, 32, 128, 2, 2, 3), - ] - - -@pytest.mark.parametrize("M,N,K,tm,tn,tk,mw,nw,nb", _oob_configs()) -def test_gemm_a16_w16_oob_tile_edges(M, N, K, tm, tn, tk, mw, nw, nb): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16( - x, - w, - dtype=torch.bfloat16, - tile_m=tm, - tile_n=tn, - tile_k=tk, - m_warp=mw, - n_warp=nw, - num_buffers=nb, + assert flydsl_gemm_a16w16(x, w, **cfg) is y + candidates = { + "flydsl": lambda: flydsl_gemm_a16w16(x, w, **cfg), + "triton": lambda: triton_gemm_a16w16(x, w, dtype=otype), + } + ret = run_candidates( + candidates, + ref, + m, + n, + k, + x.element_size(), + ref.element_size(), + "gemm a16w16 config", ) + assert torch.all(parent[m:] == SENTINEL) and torch.all( + parent[:, n:] == SENTINEL + ), "flydsl wrote outside the [m, n] output view" + return ret - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize( - "M,N,K,tm,tn,tk,mw,nw,nb", - [ - (100, 128, 256, 128, 128, 32, 2, 4, 3), - (65, 64, 512, 64, 64, 128, 2, 2, 2), - ], -) -def test_gemm_a16_w16_oob_guard_region(M, N, K, tm, tn, tk, mw, nw, nb): - torch.cuda.empty_cache() - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - sentinel = 777.0 - parent = torch.full((M + tm, N), sentinel, dtype=torch.bfloat16, device="cuda") - y = parent[:M] - torch_out = F.linear(x, w, bias=None) - kernel_out = gemm_a16w16( - x, - w, - dtype=torch.bfloat16, - y=y, - tile_m=tm, - tile_n=tn, - tile_k=tk, - m_warp=mw, - n_warp=nw, - num_buffers=nb, +def summarize(name, rows): + aiter.logger.info( + "%s summary (markdown):\n%s", name, pd.DataFrame(rows).to_markdown(index=False) ) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - assert torch.all(parent[M:] == sentinel) +def main(): + if get_gfx() not in SUPPORTED_GFX or not is_flydsl_available(): + aiter.logger.warning( + "flydsl gemm_a16w16 needs gfx1250 + flydsl; skipping on %s", get_gfx() + ) + return -@pytest.mark.parametrize("K, tile_k", [(48, 32), (100, 32), (130, 64), (1000, 64)]) -@pytest.mark.parametrize("layout", ["TN", "NN", "TT", "NT"]) -def test_gemm_a16_w16_ragged_k(K, tile_k, layout, monkeypatch): - torch.cuda.empty_cache() - M, N = 128, 192 - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16, layout=layout) - - def _no_pad(*args, **kwargs): - raise AssertionError("host-side pad/copy in gemm_a16w16 wrapper") - - monkeypatch.setattr(F, "pad", _no_pad) - torch_out = (x.float() @ w.float().T).to(torch.bfloat16) - kernel_out = gemm_a16w16(x, w, dtype=torch.bfloat16, tile_k=tile_k) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -def test_gemm_a16_w16_ragged_k_split_k(): - torch.cuda.empty_cache() - M, N, K, sk = 64, 64, 576, 2 - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - torch_out = x.float() @ w.float().T - kernel_out = gemm_a16w16( - x, - w, - dtype=torch.float32, - tile_m=64, - tile_n=64, - tile_k=64, - m_warp=2, - n_warp=2, - num_buffers=2, - split_k=sk, + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="config input of test", ) - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) - - -@pytest.mark.parametrize("N", [100, 101, 190, 257]) -@pytest.mark.parametrize("otype", [torch.bfloat16, torch.float32]) -def test_gemm_a16_w16_ragged_n_direct_store(N, otype): - torch.cuda.empty_cache() - M, K = 129, 256 - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - sentinel = 777.0 - parent = torch.full((M + 32, N + 64), sentinel, dtype=otype, device="cuda") - y = parent[:M, :N] - - torch_out = (x.float() @ w.float().T).to(otype) - out = gemm_a16w16(x, w, dtype=otype, y=y) - - assert out is y - torch.testing.assert_close(out, torch_out, atol=1e-1, rtol=1e-2) - assert torch.all(parent[:, N:] == sentinel) - assert torch.all(parent[M:] == sentinel) - - -@pytest.mark.parametrize("N", [100, 190]) -def test_gemm_a16_w16_ragged_n_split_k(N): - torch.cuda.empty_cache() - M, K, sk = 64, 512, 2 - x, w, _, _ = _generate_inputs(M, N, K, torch.bfloat16) - sentinel = 777.0 - parent = torch.full((M + 8, N + 32), sentinel, dtype=torch.float32, device="cuda") - y = parent[:M, :N] - - torch_out = x.float() @ w.float().T - out = gemm_a16w16( - x, - w, - dtype=torch.float32, - y=y, - tile_m=64, - tile_n=64, - tile_k=64, - m_warp=2, - n_warp=2, - num_buffers=2, - split_k=sk, + parser.add_argument( + "-d", + "--dtype", + type=dtypes.str2Dtype, + choices=[dtypes.d_dtypes["bf16"], dtypes.d_dtypes["fp16"]], + nargs="*", + default="bf16,", + metavar="{bf16,fp16}", + help="""Input dtype of x and w. + e.g.: -d bf16""", ) - - assert out is y - torch.testing.assert_close(out, torch_out, atol=1e-1, rtol=1e-2) - assert torch.all(parent[:, N:] == sentinel) - assert torch.all(parent[M:] == sentinel) - - -@pytest.mark.parametrize("M, N, K", [(100, 190, 256), (65, 257, 512)]) -def test_gemm_a16_w16_ragged_n_bias_activation(M, N, K): - torch.cuda.empty_cache() - x, w, bias, _ = _generate_inputs(M, N, K, torch.bfloat16, bias=True) - - torch_out = F.relu(F.linear(x, w, bias=bias)) - kernel_out = gemm_a16w16(x, w, bias=bias, dtype=torch.bfloat16, activation="relu") - torch.testing.assert_close(kernel_out, torch_out, atol=1e-1, rtol=1e-2) + parser.add_argument( + "-o", + "--otype", + type=dtypes.str2Dtype, + choices=[dtypes.d_dtypes[t] for t in ("bf16", "fp16", "fp32")], + nargs="*", + default="bf16,fp32,", + metavar="{bf16,fp16,fp32}", + help="""Output dtype for the kernel-config sweep (-c): fp32 keeps split-K + lossless, bf16/fp16 take the f32-accumulate-then-cast path. + e.g.: -o fp32""", + ) + parser.add_argument( + "-s", + "--mnk", + type=dtypes.str2tuple, + nargs="*", + default=[ + (64, 64, 64), # single tile + (256, 256, 256), # aligned multi-tile + (32, 256, 128), # m < tile_m + (256, 32, 128), # n < tile_n + (128, 128, 1024), # long main loop + drain + (1024, 128, 128), + (100, 190, 256), # ragged m/n + (129, 257, 512), # one past a tile boundary + (128, 192, 48), # ragged k, K < 2 * tile_k + (65, 190, 1000), # ragged m/n/k + (64, 5120, 2880), + (64, 2880, 4096), + (64, 128, 2880), + ], + help="""Shape of mnk. + e.g.: -s 64,5120,2880""", + ) + parser.add_argument( + "-l", + "--layout", + type=str, + choices=LAYOUTS, + nargs="*", + default=LAYOUTS, + help="""(x, w) memory layout. x: T = row-major [m,k], N = transposed view; + w: N = row-major [n,k], T = transposed view. TN is the nn.Linear call. + e.g.: -l TN NT""", + ) + parser.add_argument( + "-a", + "--activation", + type=str, + choices=list(ACTIVATIONS), + nargs="*", + default=list(ACTIVATIONS), + help="""Fused epilogue activation. + e.g.: -a gelu silu""", + ) + parser.add_argument( + "--bias", + type=dtypes.str2bool, + nargs="*", + default=[False, True], + help="""Fuse a bias add into the epilogue. + e.g.: --bias 1""", + ) + parser.add_argument( + "-c", + "--config", + type=dtypes.str2tuple, + nargs="*", + default=[ + (64, 5120, 2880, 64, 64, 256, 4, 2, 4, 2, 1), + (64, 2880, 4096, 64, 64, 128, 4, 2, 6, 4, 1), + (64, 128, 2880, 16, 16, 256, 1, 1, 3, 6, 1), + # split-K: k-tiles per split from 1 up to 4, num_buffers 2 and 3 + (128, 128, 512, 64, 64, 128, 2, 2, 2, 2, 0), + (128, 128, 512, 64, 64, 128, 2, 2, 2, 4, 0), + (64, 64, 768, 64, 64, 128, 2, 2, 2, 3, 0), + (64, 64, 768, 64, 64, 128, 2, 2, 2, 6, 0), + (128, 256, 1024, 64, 64, 128, 2, 2, 2, 4, 0), + (256, 128, 1024, 32, 32, 128, 2, 2, 2, 8, 0), + (128, 128, 1024, 64, 64, 128, 2, 2, 3, 4, 0), + (64, 64, 576, 64, 64, 64, 2, 2, 2, 2, 0), # ragged k per split + (64, 100, 512, 64, 64, 64, 2, 2, 2, 2, 0), # ragged n + split-K + # ragged m/n edge tiles at each tile size + (100, 100, 256, 128, 128, 32, 2, 4, 3, 1, 0), + (129, 257, 512, 128, 128, 32, 2, 4, 3, 1, 0), + (65, 190, 256, 128, 128, 32, 2, 4, 3, 1, 0), + (250, 120, 512, 64, 64, 128, 2, 2, 2, 1, 0), + (33, 65, 256, 32, 32, 128, 2, 2, 3, 1, 0), + (64, 64, 1024, 32, 32, 128, 2, 2, 3, 1, 0), + # default tile with main-loop unroll + (128, 128, 1024, 128, 128, 32, 2, 4, 3, 1, 1), + ], + help="""Explicit kernel config as + m,n,k,tile_m,tile_n,tile_k,m_warp,n_warp,num_buffers,split_k,unroll + (the fields a tuned flydsl_a16w16_gfx1250_* kernel name encodes). + e.g.: -c 64,5120,2880,64,64,256,4,2,4,2,1""", + ) + args = parser.parse_args() + + for dtype in args.dtype: + summarize( + "gemm_a16w16", + [ + test_gemm_a16w16(m, n, k, dtype, layout) + for layout, (m, n, k) in itertools.product(args.layout, args.mnk) + ], + ) + summarize( + "gemm_a16w16 fused epilogue", + [ + test_gemm_a16w16_fused(m, n, k, dtype, act, bias) + for act, bias, (m, n, k) in itertools.product( + args.activation, args.bias, args.mnk + ) + ], + ) + summarize( + "gemm_a16w16 kernel config", + [ + test_gemm_a16w16_config(m, n, k, dtype, otype, *tile) + for otype, (m, n, k, *tile) in itertools.product( + args.otype, args.config + ) + ], + ) + + +if __name__ == "__main__": + main() From c56625eeecc502d66018bb1ca47f5d445cde91e5 Mon Sep 17 00:00:00 2001 From: Felix Li Date: Wed, 9 Sep 2026 06:04:25 +0000 Subject: [PATCH 32/34] Fix ruff I001 and drop deleted is_flydsl_available import Co-authored-by: Cursor --- op_tests/flydsl_tests/test_gemm_a16w16.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index ed9dc2ee676..0970b158154 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -13,7 +13,6 @@ from aiter import dtypes from aiter.jit.utils.chip_info import get_gfx from aiter.ops.flydsl.gemm_a16w16_gfx1250 import gemm_a16w16 as flydsl_gemm_a16w16 -from aiter.ops.flydsl.utils import is_flydsl_available from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 as triton_gemm_a16w16 from aiter.test_common import benchmark, checkAllclose, run_perftest @@ -176,9 +175,9 @@ def summarize(name, rows): def main(): - if get_gfx() not in SUPPORTED_GFX or not is_flydsl_available(): + if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning( - "flydsl gemm_a16w16 needs gfx1250 + flydsl; skipping on %s", get_gfx() + "flydsl gemm_a16w16 needs gfx1250; skipping on %s", get_gfx() ) return From 8d828dfd1e85062cfd399ce1abdd963f3592c7ef Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Thu, 10 Sep 2026 01:16:19 +0000 Subject: [PATCH 33/34] Addressing comments --- .../model_configs/gptoss_bf16_tuned_gemm.csv | 6 +- aiter/ops/flydsl/gemm_a16w16_gfx1250.py | 46 +---- aiter/ops/flydsl/gemm_kernels.py | 30 ++- aiter/tuned_gemm.py | 22 +- op_tests/flydsl_tests/test_gemm_a16w16.py | 195 ++++++++++-------- 5 files changed, 152 insertions(+), 147 deletions(-) diff --git a/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv b/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv index 6e75b69b074..744eaf4cd5d 100644 --- a/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv +++ b/aiter/configs/model_configs/gptoss_bf16_tuned_gemm.csv @@ -256,6 +256,6 @@ gfx1250,96,2048,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,trit gfx1250,96,4096,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,96,8192,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 gfx1250,96,16384,201088,2880,False,torch.bfloat16,torch.bfloat16,False,False,triton,0,0,0.0,auto,0.0,0.0,0.0 -gfx1250,256,64,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,2,6.82,flydsl_a16w16_gfx1250_t64x64x256_mw4_nw2_nb4_sk2_unroll1,0.0,276.75,4474.37 -gfx1250,256,64,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,4,5.82,flydsl_a16w16_gfx1250_t64x64x128_mw4_nw2_nb6_sk4_unroll1,0.0,259.44,4207.20 -gfx1250,256,64,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,6,3.0,flydsl_a16w16_gfx1250_t16x16x256_mw1_nw1_nb3_sk6_unroll1,0.0,15.73,374.10 +gfx1250,256,64,5120,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,2,6.82,flydsl_hgemm_abf16_wbf16_bf16_t64x64x256x4_ks2_w4x2x1_bias1_ktail0_gm0_pht_gfx1250,0.0,276.75,4474.37 +gfx1250,256,64,2880,4096,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,4,5.82,flydsl_hgemm_abf16_wbf16_bf16_t64x64x128x6_ks4_w4x2x1_bias1_ktail0_gm0_pht_gfx1250,0.0,259.44,4207.20 +gfx1250,256,64,128,2880,True,torch.bfloat16,torch.bfloat16,False,False,flydsl,0,6,3.0,flydsl_hgemm_abf16_wbf16_bf16_t16x16x256x3_ks6_w1x1x1_bias1_ktail0_gm0_pht_gfx1250,0.0,15.73,374.10 diff --git a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py index 7660dada613..8167e63e260 100644 --- a/aiter/ops/flydsl/gemm_a16w16_gfx1250.py +++ b/aiter/ops/flydsl/gemm_a16w16_gfx1250.py @@ -6,7 +6,6 @@ from __future__ import annotations import functools -import re import torch @@ -79,40 +78,7 @@ def _cached_launcher(*cfg): return _compile_gemm_a16w16(**dict(zip(_CFG_KEYS, cfg))) -# flydsl_a16w16_gfx1250_t{tm}x{tn}x{tk}_mw{mw}_nw{nw}_nb{nb}_sk{sk}_unroll{0|1} -_KERNEL_NAME_RE = re.compile( - r"^flydsl_a16w16_gfx1250_" - r"t(?P\d+)x(?P\d+)x(?P\d+)_" - r"mw(?P\d+)_nw(?P\d+)_" - r"nb(?P\d+)_sk(?P\d+)_unroll(?P[01])$" -) - - -def parse_kernel_name(name: str): - """Parse a flydsl_a16w16_gfx1250_ kernelName into its config dict, or None.""" - m = _KERNEL_NAME_RE.fullmatch(name) - if m is None: - return None - cfg = {k: int(v) for k, v in m.groupdict().items()} - cfg["main_loop_unroll"] = bool(cfg["main_loop_unroll"]) - return cfg - - -def run_gemm_a16w16_gfx1250( - x: torch.Tensor, - w: torch.Tensor, - bias: torch.Tensor | None, - dtype: torch.dtype, - kernel_name: str, -) -> torch.Tensor: - """Dispatch entry: decode a tuned kernelName and run the kernel.""" - cfg = parse_kernel_name(kernel_name) - if cfg is None: - raise ValueError(f"[FlyDSL gfx1250] unrecognised kernelName: {kernel_name!r}") - return gemm_a16w16(x, w, bias=bias, dtype=dtype, **cfg) - - -_SPLIT_K_MAX_TILES = 1 << 16 +_SPLIT_K_MAX_TILES = 4096 @functools.cache @@ -181,9 +147,11 @@ def gemm_a16w16( sem = _split_k_counters(x.device, stream) if split_k > 1: tiles = ((M + tile_m - 1) // tile_m) * ((N + tile_n - 1) // tile_n) - assert ( - tiles <= sem.numel() - ), f"gemm_a16w16: split_k needs {tiles} tile counters, max {sem.numel()}" + if tiles > _SPLIT_K_MAX_TILES: + raise RuntimeError( + f"[FlyDSL gfx1250] split_k needs {tiles} tile counters, " + f"more than {_SPLIT_K_MAX_TILES}" + ) m_pad = ((M + tile_m - 1) // tile_m) * tile_m n_pad = ((N + tile_n - 1) // tile_n) * tile_n ws = torch.empty((split_k, m_pad, n_pad), device=x.device, dtype=torch.float32) @@ -238,4 +206,4 @@ def gemm_a16w16( return y_buf -__all__ = ["gemm_a16w16", "parse_kernel_name", "run_gemm_a16w16_gfx1250"] +__all__ = ["gemm_a16w16"] diff --git a/aiter/ops/flydsl/gemm_kernels.py b/aiter/ops/flydsl/gemm_kernels.py index 9f10353d75a..97cc076443c 100644 --- a/aiter/ops/flydsl/gemm_kernels.py +++ b/aiter/ops/flydsl/gemm_kernels.py @@ -15,6 +15,7 @@ from aiter import logger from aiter.jit.utils.chip_info import get_gfx +from .gemm_a16w16_gfx1250 import gemm_a16w16 as gemm_a16w16_gfx1250 from .kernels.gemm_a16w16_gfx950 import ( SPLIT_K_SEMAPHORE_MAX_LEN, gemm_a16w16, @@ -138,10 +139,8 @@ def flydsl_hgemm( out_dtype: torch.dtype | None = None, stream: torch.cuda.Stream | None = None, ) -> torch.Tensor: - """Run the gfx950 A16W16 kernel with AITER's ``B[N, K]`` convention.""" + """Run the A16W16 kernel for this arch with AITER's ``B[N, K]`` convention.""" - if get_gfx() != "gfx950": - raise RuntimeError("The FlyDSL A16W16 kernel currently supports gfx950 only") if policy not in ("ft", "ht", "hti"): raise ValueError(f"Unsupported FlyDSL HGEMM policy: {policy!r}") launch_stream = ( @@ -150,6 +149,31 @@ def flydsl_hgemm( if launch_stream.device != a.device: raise ValueError(f"`stream` must be on {a.device}, got {launch_stream.device}") + gfx = get_gfx() + if gfx == "gfx1250": + if k_waves != 1: + raise ValueError("The gfx1250 FlyDSL A16W16 kernel supports k_waves=1 only") + with torch.cuda.stream(launch_stream): + return gemm_a16w16_gfx1250( + a, + b, + bias=bias, + dtype=out_dtype or a.dtype, + y=out, + tile_m=block_m, + tile_n=block_n, + tile_k=block_k, + m_warp=m_waves, + n_warp=n_waves, + num_buffers=stages, + split_k=split_k, + main_loop_unroll=policy in ("ht", "hti"), + ) + if gfx != "gfx950": + raise RuntimeError( + "The FlyDSL A16W16 kernel currently supports gfx950 and gfx1250 only" + ) + if not a.is_contiguous(): a = a.contiguous() if not b.is_contiguous(): diff --git a/aiter/tuned_gemm.py b/aiter/tuned_gemm.py index a2299ccc4ec..9f047bd876d 100644 --- a/aiter/tuned_gemm.py +++ b/aiter/tuned_gemm.py @@ -158,18 +158,11 @@ def get_GEMM_A16W16_config( ) if config is not None: if config["libtype"] == "flydsl": - if gfx == "gfx1250": - from aiter.ops.flydsl.gemm_a16w16_gfx1250 import ( - parse_kernel_name, - ) - - flydsl_config = parse_kernel_name(config["kernelName"]) - else: - flydsl_config = ( - _get_flydsl_gemm_kernels().get_flydsl_hgemm_kernel_params( - config["kernelName"] - ) + flydsl_config = ( + _get_flydsl_gemm_kernels().get_flydsl_hgemm_kernel_params( + config["kernelName"] ) + ) # None means the tuned CSV names a kernel absent from this # catalog version; it is unrelated to FlyDSL import availability. if flydsl_config is None: @@ -491,13 +484,6 @@ def flydsl_gemm( assert ( scale_a is None and scale_b is None and scale_c is None ), "FlyDSL hgemm does not support scaling yet." - # gfx1250 runs the FlyDSL a16w16 WMMA kernel; other archs use the hgemm path. - if get_gfx() == "gfx1250": - from aiter.ops.flydsl.gemm_a16w16_gfx1250 import run_gemm_a16w16_gfx1250 - - return run_gemm_a16w16_gfx1250( - inp, weights, bias, otype or inp.dtype, config["kernelName"] - ) flydsl_gemm_kernels = _get_flydsl_gemm_kernels() flydsl_config = flydsl_gemm_kernels.get_flydsl_hgemm_kernel_params( config["kernelName"] diff --git a/op_tests/flydsl_tests/test_gemm_a16w16.py b/op_tests/flydsl_tests/test_gemm_a16w16.py index 0970b158154..1f76199758f 100644 --- a/op_tests/flydsl_tests/test_gemm_a16w16.py +++ b/op_tests/flydsl_tests/test_gemm_a16w16.py @@ -12,13 +12,13 @@ import aiter from aiter import dtypes from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl.gemm_a16w16_gfx1250 import gemm_a16w16 as flydsl_gemm_a16w16 +from aiter.ops.flydsl.gemm_kernels import flydsl_hgemm from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16 as triton_gemm_a16w16 from aiter.test_common import benchmark, checkAllclose, run_perftest torch.set_default_device("cuda") -SUPPORTED_GFX = ["gfx1250"] +SUPPORTED_GFX = ["gfx950", "gfx1250"] LAYOUTS = ["TN", "TT", "NN", "NT"] ACTIVATIONS = { "gelu": F.gelu, @@ -29,6 +29,67 @@ } SENTINEL = 4096.0 +DEFAULT_MNK = { + "gfx1250": [ + (64, 64, 64), + (256, 256, 256), + (32, 256, 128), + (256, 32, 128), + (128, 128, 1024), + (1024, 128, 128), + (100, 190, 256), + (129, 257, 512), + (128, 192, 48), + (65, 190, 1000), + (64, 5120, 2880), + (64, 2880, 4096), + (64, 128, 2880), + ], + "gfx950": [ + (8, 4096, 4096), + (64, 4096, 4096), + (256, 4096, 4096), + (1024, 1024, 1024), + (32, 384, 7168), + (8, 7168, 2048), + ], +} +DEFAULT_CONFIGS = { + "gfx1250": [ + (64, 5120, 2880, 64, 64, 256, 4, 2, 4, 2, 1), + (64, 2880, 4096, 64, 64, 128, 4, 2, 6, 4, 1), + (64, 128, 2880, 16, 16, 256, 1, 1, 3, 6, 1), + (128, 128, 512, 64, 64, 128, 2, 2, 2, 2, 0), + (128, 128, 512, 64, 64, 128, 2, 2, 2, 4, 0), + (64, 64, 768, 64, 64, 128, 2, 2, 2, 3, 0), + (64, 64, 768, 64, 64, 128, 2, 2, 2, 6, 0), + (128, 256, 1024, 64, 64, 128, 2, 2, 2, 4, 0), + (256, 128, 1024, 32, 32, 128, 2, 2, 2, 8, 0), + (128, 128, 1024, 64, 64, 128, 2, 2, 3, 4, 0), + (64, 64, 576, 64, 64, 64, 2, 2, 2, 2, 0), + (64, 100, 512, 64, 64, 64, 2, 2, 2, 2, 0), + (100, 100, 256, 128, 128, 32, 2, 4, 3, 1, 0), + (129, 257, 512, 128, 128, 32, 2, 4, 3, 1, 0), + (65, 190, 256, 128, 128, 32, 2, 4, 3, 1, 0), + (250, 120, 512, 64, 64, 128, 2, 2, 2, 1, 0), + (33, 65, 256, 32, 32, 128, 2, 2, 3, 1, 0), + (64, 64, 1024, 32, 32, 128, 2, 2, 3, 1, 0), + (128, 128, 1024, 128, 128, 32, 2, 4, 3, 1, 1), + ], + "gfx950": [ + (64, 4096, 4096, 32, 32, 128, 2, 2, 6, 1, 0), + (128, 4096, 4096, 32, 64, 128, 1, 4, 4, 1, 0), + (256, 4096, 4096, 64, 64, 128, 4, 2, 4, 1, 0), + (512, 4096, 4096, 64, 128, 64, 2, 4, 5, 1, 0), + (1024, 4096, 4096, 128, 128, 64, 2, 4, 4, 1, 0), + (1024, 1024, 1024, 64, 64, 64, 1, 4, 4, 1, 0), + (2048, 2048, 2048, 128, 128, 64, 2, 4, 4, 1, 0), + (4096, 4096, 4096, 256, 256, 64, 2, 4, 2, 1, 1), + (8, 7168, 2048, 16, 16, 64, 1, 1, 8, 1, 0), + (32, 14336, 4096, 32, 64, 128, 2, 2, 5, 1, 0), + ], +} + def generate_inputs(m, n, k, dtype, layout="TN", bias=False): if layout[0] == "T": @@ -55,7 +116,6 @@ def run_torch(x, w, bias=None, activation=None, dtype=dtypes.bf16): def run_candidates(candidates, ref, m, n, k, in_bytes, out_bytes, msg): - # [m,k] @ [n,k]^T: FLOPs = 2*m*n*k; bytes = x + w read, y written. flops = 2 * m * n * k nbytes = (m * k + n * k) * in_bytes + m * n * out_bytes ret = {"gfx": get_gfx()} @@ -80,10 +140,7 @@ def test_gemm_a16w16(m, n, k, dtype, layout): x, w, _ = generate_inputs(m, n, k, dtype, layout) ref = run_torch(x, w, dtype=dtype) candidates = { - "flydsl": lambda: flydsl_gemm_a16w16(x, w, dtype=dtype), - "flydsl_cb": lambda: flydsl_gemm_a16w16( - x, w, dtype=dtype, variant="compute_bound" - ), + "flydsl": lambda: flydsl_hgemm(x, w, out_dtype=dtype), "triton": lambda: triton_gemm_a16w16(x, w, dtype=dtype), } return run_candidates( @@ -93,12 +150,12 @@ def test_gemm_a16w16(m, n, k, dtype, layout): @benchmark() def test_gemm_a16w16_fused(m, n, k, dtype, activation, bias): + from aiter.ops.flydsl.gemm_a16w16_gfx1250 import gemm_a16w16 + x, w, b = generate_inputs(m, n, k, dtype, bias=bias) ref = run_torch(x, w, b, activation, dtype) candidates = { - "flydsl": lambda: flydsl_gemm_a16w16( - x, w, bias=b, dtype=dtype, activation=activation - ), + "flydsl": lambda: gemm_a16w16(x, w, bias=b, dtype=dtype, activation=activation), "triton": lambda: triton_gemm_a16w16( x, w, bias=b, dtype=dtype, activation=activation ), @@ -122,34 +179,36 @@ def test_gemm_a16w16_config( k, dtype, otype, - tile_m, - tile_n, - tile_k, - m_warp, - n_warp, - num_buffers, + block_m, + block_n, + block_k, + m_waves, + n_waves, + stages, split_k, unroll, ): x, w, _ = generate_inputs(m, n, k, dtype) ref = run_torch(x, w, dtype=otype) - parent = torch.full((m + tile_m, n + tile_n), SENTINEL, dtype=otype) + parent = torch.full((m + block_m, n + block_n), SENTINEL, dtype=otype) y = parent[:m, :n] cfg = { - "dtype": otype, - "y": y, - "tile_m": tile_m, - "tile_n": tile_n, - "tile_k": tile_k, - "m_warp": m_warp, - "n_warp": n_warp, - "num_buffers": num_buffers, + "out": y, + "out_dtype": otype, + "block_m": block_m, + "block_n": block_n, + "block_k": block_k, + "stages": stages, "split_k": split_k, - "main_loop_unroll": bool(unroll), + "m_waves": m_waves, + "n_waves": n_waves, + "k_waves": 1, + "group_m": 0, + "policy": "ht" if unroll else "ft", } - assert flydsl_gemm_a16w16(x, w, **cfg) is y + assert flydsl_hgemm(x, w, **cfg) is y candidates = { - "flydsl": lambda: flydsl_gemm_a16w16(x, w, **cfg), + "flydsl": lambda: flydsl_hgemm(x, w, **cfg), "triton": lambda: triton_gemm_a16w16(x, w, dtype=otype), } ret = run_candidates( @@ -175,10 +234,9 @@ def summarize(name, rows): def main(): - if get_gfx() not in SUPPORTED_GFX: - aiter.logger.warning( - "flydsl gemm_a16w16 needs gfx1250; skipping on %s", get_gfx() - ) + gfx = get_gfx() + if gfx not in SUPPORTED_GFX: + aiter.logger.warning("flydsl a16w16 gemm unsupported on %s; skipping", gfx) return parser = argparse.ArgumentParser( @@ -213,21 +271,7 @@ def main(): "--mnk", type=dtypes.str2tuple, nargs="*", - default=[ - (64, 64, 64), # single tile - (256, 256, 256), # aligned multi-tile - (32, 256, 128), # m < tile_m - (256, 32, 128), # n < tile_n - (128, 128, 1024), # long main loop + drain - (1024, 128, 128), - (100, 190, 256), # ragged m/n - (129, 257, 512), # one past a tile boundary - (128, 192, 48), # ragged k, K < 2 * tile_k - (65, 190, 1000), # ragged m/n/k - (64, 5120, 2880), - (64, 2880, 4096), - (64, 128, 2880), - ], + default=DEFAULT_MNK[gfx], help="""Shape of mnk. e.g.: -s 64,5120,2880""", ) @@ -249,7 +293,7 @@ def main(): choices=list(ACTIVATIONS), nargs="*", default=list(ACTIVATIONS), - help="""Fused epilogue activation. + help="""Fused epilogue activation (gfx1250 kernel only). e.g.: -a gelu silu""", ) parser.add_argument( @@ -265,33 +309,10 @@ def main(): "--config", type=dtypes.str2tuple, nargs="*", - default=[ - (64, 5120, 2880, 64, 64, 256, 4, 2, 4, 2, 1), - (64, 2880, 4096, 64, 64, 128, 4, 2, 6, 4, 1), - (64, 128, 2880, 16, 16, 256, 1, 1, 3, 6, 1), - # split-K: k-tiles per split from 1 up to 4, num_buffers 2 and 3 - (128, 128, 512, 64, 64, 128, 2, 2, 2, 2, 0), - (128, 128, 512, 64, 64, 128, 2, 2, 2, 4, 0), - (64, 64, 768, 64, 64, 128, 2, 2, 2, 3, 0), - (64, 64, 768, 64, 64, 128, 2, 2, 2, 6, 0), - (128, 256, 1024, 64, 64, 128, 2, 2, 2, 4, 0), - (256, 128, 1024, 32, 32, 128, 2, 2, 2, 8, 0), - (128, 128, 1024, 64, 64, 128, 2, 2, 3, 4, 0), - (64, 64, 576, 64, 64, 64, 2, 2, 2, 2, 0), # ragged k per split - (64, 100, 512, 64, 64, 64, 2, 2, 2, 2, 0), # ragged n + split-K - # ragged m/n edge tiles at each tile size - (100, 100, 256, 128, 128, 32, 2, 4, 3, 1, 0), - (129, 257, 512, 128, 128, 32, 2, 4, 3, 1, 0), - (65, 190, 256, 128, 128, 32, 2, 4, 3, 1, 0), - (250, 120, 512, 64, 64, 128, 2, 2, 2, 1, 0), - (33, 65, 256, 32, 32, 128, 2, 2, 3, 1, 0), - (64, 64, 1024, 32, 32, 128, 2, 2, 3, 1, 0), - # default tile with main-loop unroll - (128, 128, 1024, 128, 128, 32, 2, 4, 3, 1, 1), - ], - help="""Explicit kernel config as - m,n,k,tile_m,tile_n,tile_k,m_warp,n_warp,num_buffers,split_k,unroll - (the fields a tuned flydsl_a16w16_gfx1250_* kernel name encodes). + default=DEFAULT_CONFIGS[gfx], + help="""Explicit flydsl_hgemm config as + m,n,k,block_m,block_n,block_k,m_waves,n_waves,stages,split_k,unroll + (unroll=1 -> policy ht; the fields of a tuned flydsl_hgemm_* csv name). e.g.: -c 64,5120,2880,64,64,256,4,2,4,2,1""", ) args = parser.parse_args() @@ -304,15 +325,21 @@ def main(): for layout, (m, n, k) in itertools.product(args.layout, args.mnk) ], ) - summarize( - "gemm_a16w16 fused epilogue", - [ - test_gemm_a16w16_fused(m, n, k, dtype, act, bias) - for act, bias, (m, n, k) in itertools.product( - args.activation, args.bias, args.mnk - ) - ], - ) + if gfx == "gfx1250": + summarize( + "gemm_a16w16 fused epilogue", + [ + test_gemm_a16w16_fused(m, n, k, dtype, act, bias) + for act, bias, (m, n, k) in itertools.product( + args.activation, args.bias, args.mnk + ) + ], + ) + else: + aiter.logger.warning( + "%s: skipping fused-epilogue table (activation fusion is gfx1250-only)", + gfx, + ) summarize( "gemm_a16w16 kernel config", [ From 2101f4b33a4a885a37c41ed89341700ac76f2525 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Fri, 11 Sep 2026 15:26:13 +0000 Subject: [PATCH 34/34] CI fix --- aiter/aot/flydsl/gemm.py | 117 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 115 insertions(+), 2 deletions(-) diff --git a/aiter/aot/flydsl/gemm.py b/aiter/aot/flydsl/gemm.py index c2db1829fbd..f429780cc1c 100755 --- a/aiter/aot/flydsl/gemm.py +++ b/aiter/aot/flydsl/gemm.py @@ -12,6 +12,7 @@ Supported kernel families: - ``flydsl_hgemm_*`` gfx950 A16W16 GEMM kernels + - ``flydsl_hgemm_*_gfx1250`` gfx1250 A16W16 GEMM kernels - ``flydsl_bpreshuflle_*`` a8w8 preshuffle GEMM kernels - ``flydsl_bpreshuffle_8w_*`` gfx950 8-wave a8w8 ptpc GEMM kernels - ``flydsl_bpreshuffle_wmma_*`` gfx1250 a8w8 ptpc GEMM kernels @@ -71,8 +72,12 @@ gemm_a16w16_gfx950, make_gemm_a16w16_param_and_validate, ) +from aiter.ops.flydsl.kernels.gemm_a16w16_kernel_gfx1250 import ( + compile_gemm_a16w16 as compile_gemm_a16w16_gfx1250, +) from aiter.ops.flydsl.kernels.kernels_common import run_cached from aiter.ops.flydsl.kernels.preshuffle_gemm import compile_preshuffle_gemm +from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg from aiter.ops.flydsl.mxfp8_128_bpreshuffle_gemm_gfx1250 import ( BLOCK_K as SCALE_BLOCK_SIZE, ) @@ -209,7 +214,11 @@ def parse_csv(csv_path: str): params = get_flydsl_hgemm_kernel_params(kernel_name) if params is not None: params = dict(params) - params["kind"] = "hgemm" + params["kind"] = ( + "a16w16_gfx1250" + if params["target_gfx"] == "gfx1250" + else "hgemm" + ) else: params = None @@ -376,6 +385,104 @@ def _compile_hgemm_to_cache( ) +def _compile_a16w16_gfx1250_to_cache( + *, + m: int, + n: int, + k: int, + dtype: str, + out_dtype: str, + block_m: int, + block_n: int, + block_k: int, + stages: int, + split_k: int, + m_waves: int, + n_waves: int, + k_waves: int, + group_m: int, + use_half_tile_interleaved: bool, + has_bias: bool, + target_gfx: str, + **kwargs, +): + del kwargs, group_m + + import torch + + if target_gfx != "gfx1250": + raise ValueError( + f"The FlyDSL gfx1250 A16W16 kernel only supports gfx1250, got {target_gfx}" + ) + if k_waves != 1: + raise ValueError( + f"The FlyDSL gfx1250 A16W16 kernel only supports k_waves=1, got {k_waves}" + ) + + torch_dtype = _torch_dtype_for_kernel(dtype) + torch_out_dtype = _torch_dtype_for_kernel(out_dtype) + if torch_out_dtype not in (torch_dtype, torch.float32): + raise ValueError( + f"Unsupported output dtype {out_dtype!r} for input dtype {dtype!r}" + ) + in_dtype_str = "fp16" if torch_dtype == torch.float16 else "bf16" + out_dtype_str = {torch.float16: "f16", torch.bfloat16: "bf16"}.get( + torch_out_dtype, "f32" + ) + fx_dtype = { + torch.bfloat16: fx.BFloat16, + torch.float16: fx.Float16, + torch.float32: fx.Float32, + torch.int32: fx.Int32, + } + launch = compile_gemm_a16w16_gfx1250( + N=n, + K=k, + tile_m=block_m, + tile_n=block_n, + tile_k=block_k, + m_warp=m_waves, + n_warp=n_waves, + in_dtype=in_dtype_str, + out_dtype=out_dtype_str, + num_buffers=stages, + add_bias=has_bias, + split_k=split_k, + main_loop_unroll=use_half_tile_interleaved, + ) + + with compile_only_env(): + dev = torch.device("cpu") + representative_extent = 8 + a = torch.empty((1, representative_extent), device=dev, dtype=torch_dtype) + b = torch.empty((1, representative_extent), device=dev, dtype=torch_dtype) + out = torch.empty((1, representative_extent), device=dev, dtype=torch_out_dtype) + bias = torch.empty( + (representative_extent if has_bias else 0,), + device=dev, + dtype=torch_dtype if has_bias else torch_out_dtype, + ) + workspace = torch.empty( + (representative_extent,), device=dev, dtype=torch.float32 + ) + semaphore = torch.zeros((representative_extent,), device=dev, dtype=torch.int32) + stream = fx.Stream(0) + dispatch_args = ( + ptr_arg(out, fx_dtype[out.dtype]), + ptr_arg(a, fx_dtype[a.dtype]), + ptr_arg(b, fx_dtype[b.dtype]), + ptr_arg(bias, fx_dtype[bias.dtype]), + ptr_arg(workspace, fx_dtype[workspace.dtype]), + ptr_arg(semaphore, fx_dtype[semaphore.dtype]), + m, + n, + k, + k, + stream, + ) + _compile_executable_to_cache(launch, *dispatch_args) + + def _compile_preshuffle_to_cache( *, m: int, @@ -679,7 +786,9 @@ def compile_one_config( t0 = time.time() try: - tensor_context = nullcontext() if kind == "hgemm" else FakeTensorMode() + tensor_context = ( + nullcontext() if kind in ("hgemm", "a16w16_gfx1250") else FakeTensorMode() + ) with ( override_env("FLYDSL_GPU_ARCH", aot_arch), tensor_context, @@ -688,6 +797,10 @@ def compile_one_config( hgemm_kwargs = dict(kwargs) hgemm_kwargs["target_gfx"] = aot_arch _compile_hgemm_to_cache(m=m, n=n, k=k, **hgemm_kwargs) + elif kind == "a16w16_gfx1250": + a16w16_kwargs = dict(kwargs) + a16w16_kwargs["target_gfx"] = aot_arch + _compile_a16w16_gfx1250_to_cache(m=m, n=n, k=k, **a16w16_kwargs) elif kind == "preshuffle": _compile_preshuffle_to_cache(m=m, n=n, k=k, **kwargs) elif kind == "8wave":