diff --git a/aiter/ops/flydsl/kernels/flash_attn_func_gfx1201.py b/aiter/ops/flydsl/kernels/flash_attn_func_gfx1201.py index c3990b5467..cbc8553e9f 100644 --- a/aiter/ops/flydsl/kernels/flash_attn_func_gfx1201.py +++ b/aiter/ops/flydsl/kernels/flash_attn_func_gfx1201.py @@ -44,7 +44,6 @@ from flydsl.expr import math as fmath from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec -from flydsl.expr.utils.arith import ArithValue from flydsl.expr.utils.arith import _to_raw as _raw from aiter.ops.flydsl.kernels import buffer_ops @@ -233,13 +232,13 @@ def wmma_acc(a_v8, b_v8, c_v8): ).result return rocdl.wmma_f32_16x16x16_f16(v8f32_type, a_v8, b_v8, c_v8).result - seq_len_v = fx.Index(seq_len) + seq_len_v = fx.Uint64(seq_len) lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_kv = lds.kv.ptr - block_id = fx.Index(gpu.block_idx.x) - tid = fx.Index(gpu.thread_idx.x) + block_id = fx.Uint64(gpu.block_idx.x) + tid = fx.Uint64(gpu.thread_idx.x) wave_id = tid // WARP_SIZE lane = tid % WARP_SIZE @@ -282,7 +281,7 @@ def load_global_v8f16(base_ptr, base_idx): return _load_global_half_vec(base_ptr, base_idx, v8f16_type) def _bitcast_i32(value): - return fx.Int32(ArithValue(value).bitcast(fx.Int32.ir_type)) + return fx.Float32(value).bitcast(fx.Int32) def _pack_bf16_pair(lo, hi, shift, mask): lo_i32 = _bitcast_i32(lo) @@ -298,23 +297,27 @@ def bf16_trunc_pack_v8(f32_vals): pairs.append( _pack_bf16_pair(f32_vals[j * 2], f32_vals[j * 2 + 1], _c16, _cmask) ) - return Vec.from_elements(pairs, fx.Int32).bitcast(elem_dtype).ir_value() + return Vec.from_elements(pairs, fx.Int32).bitcast(elem_dtype) def k_buf_base(buf_id): if const_expr(isinstance(buf_id, int)): - return fx.Index(buf_id * LDS_K_TILE_SIZE) - return buf_id * fx.Index(LDS_K_TILE_SIZE) + return fx.Int64(buf_id * LDS_K_TILE_SIZE) + return buf_id * fx.Int64(LDS_K_TILE_SIZE) def v_buf_base(buf_id): - return fx.Index(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) + return fx.Int64(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) def coop_load_k(tile_start, buf_id=0): + # `tile_start` may arrive as the loop induction variable (index type); + # normalize to i64 so it composes with the i64 thread offsets. index + # is i64 in the gfx1201 datalayout, so this cast is a no-op in the ISA. + tile_start = fx.Int64(tile_start) k_base = k_buf_base(buf_id) for batch in range_constexpr(NUM_BATCHES_KV): row_offset = batch * ROWS_PER_BATCH_LOAD row_idx = tile_start + load_row_in_batch + row_offset if const_expr(KV_NEEDS_GUARD): - row_valid = load_row_in_batch < fx.Index(BLOCK_N) + row_valid = load_row_in_batch < fx.Int64(BLOCK_N) if row_valid: g_idx = global_idx(row_idx, load_col_base) lds_row = load_row_in_batch + row_offset @@ -333,6 +336,7 @@ def _v_store_row_major(v_base, lds_row, vec): fx.ptr_store(Vec(vec), lds_kv + fx.Int32(lds_idx)) def coop_load_v_global(tile_start): + tile_start = fx.Int64(tile_start) vecs = [] for batch in range_constexpr(NUM_BATCHES_KV): row_offset = batch * ROWS_PER_BATCH_LOAD @@ -346,7 +350,7 @@ def coop_store_v_lds(vecs, buf_id=0): for batch in range_constexpr(NUM_BATCHES_KV): row_offset = batch * ROWS_PER_BATCH_LOAD if const_expr(KV_NEEDS_GUARD): - row_valid = load_row_in_batch < fx.Index(BLOCK_N) + row_valid = load_row_in_batch < fx.Int64(BLOCK_N) if row_valid: lds_row = load_row_in_batch + row_offset _v_store_row_major(v_base, lds_row, vecs[batch]) @@ -358,18 +362,18 @@ def coop_store_v_lds(vecs, buf_id=0): q_row = q_start + wave_q_offset + lane16 q_row_i32 = fx.Int32(q_row) # Use explicit signed-less-than predicate to match baseline ISA - # (`v_cmp_gt_i64_e64`). fx.Index defaults to unsigned which would lower + # (`v_cmp_gt_i64_e64`). An fx `<` on unsigned-typed operands would lower # to `v_cmp_gt_u64_e64` and cause an ISA hash drift even though both # variants are semantically equivalent for non-negative offsets. q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, _raw(q_row), _raw(seq_len_v)) - q_row_safe = fx.Index(ArithValue(q_in_bounds).select(q_row, fx.Index(0))) - c_zero_v8f16 = Vec.filled(8, 0.0, elem_dtype).ir_value() + q_row_safe = fx.Int64(q_in_bounds.select(q_row, fx.Int64(0))) + c_zero_v8f16 = Vec.filled(8, 0.0, elem_dtype) q_b_packs = [] for ks in range_constexpr(K_STEPS_QK): - q_col = fx.Index(ks * K_STEP_QK) + klane * WMMA_LANE_K + q_col = fx.Int64(ks * K_STEP_QK) + klane * WMMA_LANE_K g_idx = global_idx(q_row_safe, q_col) raw = load_global_v8f16(q_ptr, g_idx) - q_b_packs.append(ArithValue(q_in_bounds).select(raw, c_zero_v8f16)) + q_b_packs.append(q_in_bounds.select(raw, c_zero_v8f16)) # ---- Constants ---- c_neg_inf = fx.Float32(float("-inf")) @@ -385,14 +389,12 @@ def reduction_peer(v_f32): _q_end = q_start + BLOCK_M if const_expr(CAUSAL): - kv_upper = fx.Index( - ArithValue(_q_end < seq_len_v).select(_q_end, seq_len_v) - ) + kv_upper = fx.Int64((_q_end < seq_len_v).select(_q_end, seq_len_v)) else: kv_upper = seq_len_v # ---- Opt4: Pre-issue first V global load before loop ---- - _v_vecs_init = coop_load_v_global(fx.Index(0)) + _v_vecs_init = coop_load_v_global(fx.Int64(0)) init_args = [_raw(c_neg_inf), _raw(c_zero_f)] for _ in range_constexpr(D_CHUNKS): @@ -403,7 +405,7 @@ def reduction_peer(v_f32): loop_results = init_args for kv_block_start, inner_iter_args in range( - 0, kv_upper, BLOCK_N_OUT, init=init_args + fx.Int64(0), kv_upper, fx.Int64(BLOCK_N_OUT), init=init_args ): m_running = inner_iter_args[0] l_running = inner_iter_args[1] @@ -421,18 +423,18 @@ def reduction_peer(v_f32): s_accs = [_raw(c_zero_v8f32) for _ in range(NUM_S_ACCS)] for ks in range_constexpr(K_STEPS_QK): - k_col = fx.Index(ks * K_STEP_QK) + klane * WMMA_LANE_K + k_col = fx.Int64(ks * K_STEP_QK) + klane * WMMA_LANE_K for st_idx in range_constexpr(N_SUB_TILES): st_base_row = st_idx * K_SUB_N - k_row_a = lane16 + fx.Index(st_base_row) + k_row_a = lane16 + fx.Int64(st_base_row) k_lds_a = k_base + k_row_a * K_STRIDE + k_col k_pack_a = fx.ptr_load( lds_kv + fx.Int32(k_lds_a), result_type=v8f16_type ) - k_row_b = lane16 + fx.Index(st_base_row + 16) + k_row_b = lane16 + fx.Int64(st_base_row + 16) k_lds_b = k_base + k_row_b * K_STRIDE + k_col k_pack_b = fx.ptr_load( lds_kv + fx.Int32(k_lds_b), result_type=v8f16_type @@ -486,38 +488,38 @@ def reduction_peer(v_f32): klane_off_i32 = klane_i32 * fx.Int32(8) # st=0 _b0 = kv_start_i32 + fx.Int32(0) + klane_off_i32 - s_v0 = ArithValue(_b0 > q_row_i32).select(c_neg_inf, s_v0) + s_v0 = (_b0 > q_row_i32).select(c_neg_inf, s_v0) _b1 = kv_start_i32 + fx.Int32(1) + klane_off_i32 - s_v1 = ArithValue(_b1 > q_row_i32).select(c_neg_inf, s_v1) + s_v1 = (_b1 > q_row_i32).select(c_neg_inf, s_v1) _b2 = kv_start_i32 + fx.Int32(2) + klane_off_i32 - s_v2 = ArithValue(_b2 > q_row_i32).select(c_neg_inf, s_v2) + s_v2 = (_b2 > q_row_i32).select(c_neg_inf, s_v2) _b3 = kv_start_i32 + fx.Int32(3) + klane_off_i32 - s_v3 = ArithValue(_b3 > q_row_i32).select(c_neg_inf, s_v3) + s_v3 = (_b3 > q_row_i32).select(c_neg_inf, s_v3) _b4 = kv_start_i32 + fx.Int32(4) + klane_off_i32 - s_v4 = ArithValue(_b4 > q_row_i32).select(c_neg_inf, s_v4) + s_v4 = (_b4 > q_row_i32).select(c_neg_inf, s_v4) _b5 = kv_start_i32 + fx.Int32(5) + klane_off_i32 - s_v5 = ArithValue(_b5 > q_row_i32).select(c_neg_inf, s_v5) + s_v5 = (_b5 > q_row_i32).select(c_neg_inf, s_v5) _b6 = kv_start_i32 + fx.Int32(6) + klane_off_i32 - s_v6 = ArithValue(_b6 > q_row_i32).select(c_neg_inf, s_v6) + s_v6 = (_b6 > q_row_i32).select(c_neg_inf, s_v6) _b7 = kv_start_i32 + fx.Int32(7) + klane_off_i32 - s_v7 = ArithValue(_b7 > q_row_i32).select(c_neg_inf, s_v7) + s_v7 = (_b7 > q_row_i32).select(c_neg_inf, s_v7) # st=1 (st_base=16) _b8 = kv_start_i32 + fx.Int32(16) + klane_off_i32 - s_v8 = ArithValue(_b8 > q_row_i32).select(c_neg_inf, s_v8) + s_v8 = (_b8 > q_row_i32).select(c_neg_inf, s_v8) _b9 = kv_start_i32 + fx.Int32(17) + klane_off_i32 - s_v9 = ArithValue(_b9 > q_row_i32).select(c_neg_inf, s_v9) + s_v9 = (_b9 > q_row_i32).select(c_neg_inf, s_v9) _b10 = kv_start_i32 + fx.Int32(18) + klane_off_i32 - s_v10 = ArithValue(_b10 > q_row_i32).select(c_neg_inf, s_v10) + s_v10 = (_b10 > q_row_i32).select(c_neg_inf, s_v10) _b11 = kv_start_i32 + fx.Int32(19) + klane_off_i32 - s_v11 = ArithValue(_b11 > q_row_i32).select(c_neg_inf, s_v11) + s_v11 = (_b11 > q_row_i32).select(c_neg_inf, s_v11) _b12 = kv_start_i32 + fx.Int32(20) + klane_off_i32 - s_v12 = ArithValue(_b12 > q_row_i32).select(c_neg_inf, s_v12) + s_v12 = (_b12 > q_row_i32).select(c_neg_inf, s_v12) _b13 = kv_start_i32 + fx.Int32(21) + klane_off_i32 - s_v13 = ArithValue(_b13 > q_row_i32).select(c_neg_inf, s_v13) + s_v13 = (_b13 > q_row_i32).select(c_neg_inf, s_v13) _b14 = kv_start_i32 + fx.Int32(22) + klane_off_i32 - s_v14 = ArithValue(_b14 > q_row_i32).select(c_neg_inf, s_v14) + s_v14 = (_b14 > q_row_i32).select(c_neg_inf, s_v14) _b15 = kv_start_i32 + fx.Int32(23) + klane_off_i32 - s_v15 = ArithValue(_b15 > q_row_i32).select(c_neg_inf, s_v15) + s_v15 = (_b15 > q_row_i32).select(c_neg_inf, s_v15) s_raw = [ s_v0, s_v1, @@ -547,7 +549,7 @@ def reduction_peer(v_f32): # ---- Opt2: rocdl.exp2 ---- diff_m_raw = _fsub(m_running, m_new_raw) diff_m_scaled = _fmul(diff_m_raw, c_sm_scale_log2e) - corr = rocdl.exp2(ir.F32Type.get(), _raw(diff_m_scaled)) + corr = fx.rocdl.exp2(ir.F32Type.get(), _raw(diff_m_scaled)) scaled_max = _fmul(c_sm_scale_log2e, m_new_raw) neg_scaled_max = _fsub(c_zero_f, scaled_max) @@ -556,7 +558,7 @@ def reduction_peer(v_f32): local_sum = _raw(c_zero_f) for r in range_constexpr(NUM_S_VALS): diff = fmath.fma(s_raw[r], _raw(c_sm_scale_log2e), neg_scaled_max) - p = rocdl.exp2(ir.F32Type.get(), _raw(diff)) + p = fx.rocdl.exp2(ir.F32Type.get(), _raw(diff)) p_vals.append(p) local_sum = _fadd(local_sum, p) @@ -565,7 +567,7 @@ def reduction_peer(v_f32): l_corr = _fmul(corr, l_running) l_new = _fadd(l_corr, tile_sum) - corr_vec = Vec.from_elements([corr], fx.Float32).broadcast_to(8).ir_value() + corr_vec = Vec.from_elements([corr], fx.Float32).broadcast_to(8) for dc in range_constexpr(D_CHUNKS): o_accs[dc] = _fmul(o_accs[dc], corr_vec) @@ -588,9 +590,7 @@ def reduction_peer(v_f32): elem_list = [] for j in range_constexpr(8): elem_list.append(fx.Float32(p_slice[j]).to(elem_dtype)) - p_packs_st.append( - Vec.from_elements(elem_list, elem_dtype).ir_value() - ) + p_packs_st.append(Vec.from_elements(elem_list, elem_dtype)) p_packs_all.append(p_packs_st) # ==== GEMM2: O += V^T @ P (software pipelined, row-major V) ==== @@ -598,17 +598,17 @@ def reduction_peer(v_f32): v_base = v_buf_base(0) def _load_v_rowmajor(st_kv_base_val, pks_val, dc_val, v_base=v_base): - d_pos = fx.Index(dc_val * D_CHUNK) + lane16 + d_pos = fx.Int64(dc_val * D_CHUNK) + lane16 v_elems = [] for k_sub in range_constexpr(8): kv_row = ( - fx.Index(st_kv_base_val + pks_val * PV_K_STEP) + fx.Int64(st_kv_base_val + pks_val * PV_K_STEP) + klane * WMMA_LANE_K - + fx.Index(k_sub) + + fx.Int64(k_sub) ) v_lds_idx = v_base + kv_row * V_STRIDE + d_pos v_elems.append(fx.ptr_load(lds_kv + fx.Int32(v_lds_idx))) - return Vec.from_elements(v_elems, elem_dtype).ir_value() + return Vec.from_elements(v_elems, elem_dtype) # Software pipeline: preload first V pack cur_v_packs = [] @@ -644,7 +644,7 @@ def _load_v_rowmajor(st_kv_base_val, pks_val, dc_val, v_base=v_base): l_running = l_new # ---- Opt4: Issue NEXT iteration's V global load ---- - next_kv_start = kv_block_start + fx.Index(BLOCK_N_OUT) + next_kv_start = fx.Int64(kv_block_start) + fx.Int64(BLOCK_N_OUT) _v_vecs_next = coop_load_v_global(next_kv_start) _yield_args = [m_running, l_running] + o_accs @@ -657,13 +657,13 @@ def _load_v_rowmajor(st_kv_base_val, pks_val, dc_val, v_base=v_base): o_finals = [loop_results[2 + dc] for dc in range_constexpr(D_CHUNKS)] inv_l = arith.divf(_raw(c_one_f), _raw(l_final), fastmath=fm_fast) - inv_l_vec = Vec.from_elements([inv_l], fx.Float32).broadcast_to(8).ir_value() + inv_l_vec = Vec.from_elements([inv_l], fx.Float32).broadcast_to(8) if q_in_bounds: for dc in range_constexpr(D_CHUNKS): o_norm_vec = _fmul(o_finals[dc], inv_l_vec) - o_trunc = Vec(o_norm_vec).to(elem_dtype).ir_value() - d_col = fx.Index(dc * D_CHUNK) + klane * 8 + o_trunc = Vec(o_norm_vec).to(elem_dtype) + d_col = fx.Int64(dc * D_CHUNK) + klane * 8 o_global = global_idx(q_row, d_col) _store_global_half(o_ptr, o_global, o_trunc) @@ -681,8 +681,8 @@ def launch_flash_attn_func( ): ctx = CompilationContext.get_current() - bs_idx = fx.Index(batch_size) - sl_idx = fx.Index(seq_len) + bs_idx = fx.Uint64(batch_size) + sl_idx = fx.Uint64(seq_len) num_q_tiles = (sl_idx + BLOCK_M - 1) // BLOCK_M grid_x = bs_idx * num_q_tiles * NUM_HEADS diff --git a/aiter/ops/flydsl/kernels/fused_compress_attn.py b/aiter/ops/flydsl/kernels/fused_compress_attn.py index 21268ea3af..be7b4e3ca2 100644 --- a/aiter/ops/flydsl/kernels/fused_compress_attn.py +++ b/aiter/ops/flydsl/kernels/fused_compress_attn.py @@ -76,12 +76,10 @@ import flydsl.expr as fx import torch from flydsl._mlir.dialects import rocdl -from flydsl.expr import arith, const_expr, gpu, range_constexpr +from flydsl.expr import arith, const_expr, fastmath, gpu, range_constexpr from flydsl.expr import math as fmath -from flydsl.expr.arith import ArithValue, CmpFPredicate, CmpIPredicate from flydsl.expr.typing import Int32, Stream, T -from aiter.ops.flydsl.kernels import buffer_ops, vector from aiter.utility.mx_types import ( MxDtypeInt as _MxDtypeInt, ) @@ -97,11 +95,27 @@ state_slot_byte_offset, ) from .quant_utils import emit_f32_to_e2m1, emit_mx_e8m0_scale -from .tensor_shim import _run_compiled, _to_raw +from .tensor_shim import _run_compiled, _to_raw, ptr_buf_tensor # --- shape constants -------------------------------------------------------- BLOCK_THREADS = 64 # 1 wave64; D must be a multiple + +def _ptr_at_byte_off(tensor, base_i64): + """Global byte pointer at ``tensor``'s base + ``base_i64`` (64-bit byte offset). + + Used to fold a slot/block rebase into the pointer handed to + ``ptr_buf_tensor``, which re-derives the descriptor's element type -- so the + i8 carrier type here only carries the address and the ptrtoint/inttoptr + roundtrip folds away pre-ISA, keeping the emitted V# identical to the old + ``buf_tensor(base_i64=...)`` descriptor. + """ + pt = fx.PointerType.get(T.i8, address_space=fx.AddressSpace.Global, alignment=1) + return fx.inttoptr( + pt, fx.Int64(fx.ptrtoint(fx.get_iter(tensor))) + fx.Int64(base_i64) + ) + + # --- fp8 + e8m0 constants --------------------------------------------------- # Defer ``aiter.utility.dtypes`` import to first call (matches # qk_norm_rope_quant pattern). The aiter package is walked by setup.py's AOT @@ -136,11 +150,6 @@ def _fp8_const(): _FP4_K_TILE = 128 -# ============================================================================ -# scf helpers (copied verbatim from moe_gemm_2stage.py -- too small to share) -# ============================================================================ - - # ============================================================================ # Kernel builder # ============================================================================ @@ -312,7 +321,6 @@ def kernel( ): f32 = T.f32 i32 = T.i32 - vecVf32 = T.vec(VEC, T.f32) # --- thread / block ids --- pid = fx.block_idx.x # one program per plan row @@ -321,30 +329,33 @@ def kernel( # --- constants --- c_neg_inf = arith.constant(_NEG_INF, type=f32) c_zero_f32 = arith.constant(0.0, type=f32) - c_eps = arith.constant(rms_eps, type=f32) - c_inv_D = arith.constant(1.0 / D, type=f32) - c_log2e = arith.constant(_LOG2E, type=f32) + c_eps = fx.Float32(rms_eps) + c_inv_D = fx.Float32(1.0 / D) + c_log2e = fx.Float32(_LOG2E) def fexp_f32(x): - """exp(x) via exp2(x * log2e). Single v_exp_f32 on AMD.""" - return fx.rocdl.exp2(f32, x * c_log2e) + """exp(x) via exp2(x * log2e). Single v_exp_f32 on AMD. + + ``x`` is an fx.Float32; exp2 needs a raw operand, so wrap once here + (not at every call site).""" + return fx.rocdl.exp2(f32, _to_raw(x * c_log2e)) def wave_reduce_add(x): """Butterfly sum across wave64.""" - w = _to_raw(x) + w = fx.Float32(x) for sh_exp in range_constexpr(log2_block): off = BLOCK_THREADS // (2 << sh_exp) - peer = _to_raw(ArithValue(w).shuffle_xor(off, BLOCK_THREADS)) - w = arith.AddFOp(w, peer, fastmath=fm_fast).result + peer = w.shuffle_xor(off, BLOCK_THREADS) + w = w + peer return w def wave_reduce_max(x): """Butterfly max across wave64 (used by quant path).""" - w = _to_raw(x) + w = fx.Float32(x) for sh_exp in range_constexpr(log2_block): off = BLOCK_THREADS // (2 << sh_exp) - peer = _to_raw(ArithValue(w).shuffle_xor(off, BLOCK_THREADS)) - w = arith.maximumf(w, peer) + peer = w.shuffle_xor(off, BLOCK_THREADS) + w = fx.max(w, peer) return w # ---- Step 1: load plan row (single dwordx4) ---- @@ -352,13 +363,14 @@ def wave_reduce_max(x): # Fuse the 4 scalar loads into one buffer_load_dwordx4 + 4 extracts -- # saves 3 buffer-load instructions per program (visible at small N # where total program count is low). - plan_rsrc = buffer_ops.create_buffer_resource(plan, max_size=True) - plan_base = ArithValue(pid) * arith.constant(4, type=i32) - plan_vec = buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) - ragged_id = vector.extract(plan_vec, static_position=[0], dynamic_position=[]) - batch_id = vector.extract(plan_vec, static_position=[1], dynamic_position=[]) - position = vector.extract(plan_vec, static_position=[2], dynamic_position=[]) - window_len = vector.extract(plan_vec, static_position=[3], dynamic_position=[]) + plan_buf = ptr_buf_tensor(fx.get_iter(plan), fx.Int32) + plan_vec = fx.Vector( + fx.add_offset(fx.get_iter(plan_buf), fx.Int32(pid) * 4).load(T.vec(4, i32)) + ) + ragged_id = plan_vec[0] + batch_id = plan_vec[1] + position = plan_vec[2] + window_len = plan_vec[3] # ---- Step 2: sentinel-skip ---- # Sentinel-skip: run the whole body only for position >= 0. A bare @@ -368,16 +380,12 @@ def wave_reduce_max(x): # thread) and lowers it to scf.if, guarding every store inside. def _body(): # ---- Step 3: per-seq state slot ---- - slot_map_rsrc = buffer_ops.create_buffer_resource( - state_slot_mapping, max_size=True - ) - slot = buffer_ops.buffer_load( - slot_map_rsrc, batch_id, vec_width=1, dtype=i32 - ) + slot_map_buf = ptr_buf_tensor(fx.get_iter(state_slot_mapping), fx.Int32) + slot = fx.add_offset(fx.get_iter(slot_map_buf), batch_id).load(i32) # ---- Step 4: per-thread element-range bookkeeping ---- # This thread owns columns [tid*VEC, tid*VEC+VEC) of BLOCK_D. - tid_x_vec = ArithValue(tid) * arith.constant(VEC, type=i32) + tid_x_vec = fx.Int32(tid) * VEC # ---- Step 5: online-softmax accumulator init ---- # 3 * VEC fp32 scalars carried across K iters. @@ -418,104 +426,95 @@ def _online_softmax_update( w_old = w_lane[i] kv_old = kv_lane[i] - m_new = arith.maximumf(m_old, score) - is_first = arith.cmpf(CmpFPredicate.OEQ, m_old, c_neg_inf) - scale_active = fexp_f32(arith.subf(m_old, m_new)) - scale_v = arith.select(is_first, c_zero_f32, scale_active) - wk_active = fexp_f32(arith.subf(score, m_new)) + # -inf sentinel. maximumf keeps the ambient fast_fp_math + # (matches the old raw maximumf, which resolved + # fastmath). The OEQ compares against -inf MUST emit + # fastmath -- ambient fast lets `x == -inf` fold to + # false and corrupts the sentinel -- so scope only the `==` + # to fastmath(None) (== the old raw cmpf, no fastmath). + m_old_f = fx.Float32(m_old) + score_f = fx.Float32(score) + neg_inf_f = fx.Float32(c_neg_inf) + m_new = fx.max(m_old_f, score_f).ir_value() + with fastmath(None): + is_first = m_old_f == neg_inf_f + scale_active = fexp_f32(m_old_f - m_new) + scale_v = is_first.select(c_zero_f32, scale_active) + wk_active = fexp_f32(score_f - m_new) if const_expr(score_can_be_neg_inf): - is_pad_score = arith.cmpf(CmpFPredicate.OEQ, score, c_neg_inf) - w_k = arith.select(is_pad_score, c_zero_f32, wk_active) + with fastmath(None): + is_pad_score = score_f == neg_inf_f + w_k = is_pad_score.select(c_zero_f32, wk_active) else: w_k = wk_active new_m.append(m_new) new_kv.append( - arith.AddFOp( - arith.MulFOp(kv_old, scale_v, fastmath=fm_fast).result, - arith.MulFOp(w_k, kv_v, fastmath=fm_fast).result, - fastmath=fm_fast, - ).result + _to_raw( + fx.Float32(kv_old) * fx.Float32(scale_v) + + fx.Float32(w_k) * fx.Float32(kv_v) + ) ) new_w.append( - arith.AddFOp( - arith.MulFOp(w_old, scale_v, fastmath=fm_fast).result, - w_k, - fastmath=fm_fast, - ).result + _to_raw( + fx.Float32(w_old) * fx.Float32(scale_v) + fx.Float32(w_k) + ) ) return new_m, new_kv, new_w - def _load_bf16_vec_then_f32(rsrc, off_elems_i32): + def _load_bf16_vec_then_f32(buf, off_elems_i32): """Load VEC bf16 from byte-aligned dword stream -> fp32 VEC scalars. - Returns a list of VEC fp32 MLIR values. + ``buf`` is an i32 (dword) buffer-tensor; ``off_elems_i32`` is in + bf16 elements. Returns a list of VEC fp32 MLIR values. """ - off_dw = ArithValue(off_elems_i32) >> arith.constant(1, type=i32) + off_dw = fx.Int32(off_elems_i32) >> 1 # bf16 VEC = VEC * 2 bytes; for VEC ? {2, 4, 8} that's # {4, 8, 16} bytes = {1, 2, 4} dwords. dwords = (VEC + 1) // 2 # ceil(VEC*2 / 4) if const_expr(dwords == 1): - # buffer_load(vec_width=1) returns a scalar i32; wrap into - # vec<1xi32> before bitcasting to vec<2xbf16>. - raw_s = buffer_ops.buffer_load(rsrc, off_dw, vec_width=1, dtype=i32) - raw = vector.from_elements(T.vec(1, T.i32), [raw_s]) - else: - raw = buffer_ops.buffer_load( - rsrc, off_dw, vec_width=dwords, dtype=i32 + # width=1 returns a scalar i32; wrap into vec<1xi32> before + # bitcasting to vec<2xbf16>. + raw = fx.Vector.from_elements( + [fx.add_offset(fx.get_iter(buf), off_dw).load(i32)], + dtype=fx.Int32, ) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - out = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, - static_position=[i], - dynamic_position=[], + else: + raw = fx.Vector( + fx.add_offset(fx.get_iter(buf), off_dw).load(T.vec(dwords, i32)) ) - f32_v = arith.extf(f32, bf16_v) - out.append(f32_v) - return out + vec_bf16 = raw.bitcast(fx.BFloat16) + return [vec_bf16[i].to(fx.Float32) for i in range_constexpr(VEC)] - def _load_f32_vec(rsrc, off_elems_i32): + def _load_f32_vec(buf, off_elems_i32): """Load VEC fp32 from byte-aligned stream -> list of VEC fp32 scalars. For VEC=2 -> dwordx2; VEC=4 -> dwordx4; VEC=8 -> 2x dwordx4 (HW max). """ if const_expr(VEC <= 4): - vw = VEC - raw = buffer_ops.buffer_load( - rsrc, off_elems_i32, vec_width=vw, dtype=f32 + raw = fx.Vector( + fx.add_offset(fx.get_iter(buf), off_elems_i32).load( + T.vec(VEC, f32) + ) ) - return [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + return [raw[i] for i in range(VEC)] else: # VEC == 8 -> 2x dwordx4 assert VEC == 8 half = VEC // 2 - r0 = buffer_ops.buffer_load( - rsrc, off_elems_i32, vec_width=half, dtype=f32 - ) - r1 = buffer_ops.buffer_load( - rsrc, - ArithValue(off_elems_i32) + arith.constant(half, type=i32), - vec_width=half, - dtype=f32, + base = fx.Int32(off_elems_i32) + r0 = fx.Vector( + fx.add_offset(fx.get_iter(buf), base).load(T.vec(half, f32)) ) - out = [] - for i in range_constexpr(half): - out.append( - vector.extract(r0, static_position=[i], dynamic_position=[]) + r1 = fx.Vector( + fx.add_offset(fx.get_iter(buf), base + half).load( + T.vec(half, f32) ) - for i in range_constexpr(half): - out.append( - vector.extract(r1, static_position=[i], dynamic_position=[]) - ) - return out + ) + return [r0[i] for i in range(half)] + [r1[i] for i in range(half)] # Buffer resources reused across K iters. - kv_in_rsrc = buffer_ops.create_buffer_resource(kv_in, max_size=True) - score_in_rsrc = buffer_ops.create_buffer_resource(score_in, max_size=True) + kv_in_buf = ptr_buf_tensor(fx.get_iter(kv_in), fx.Int32) + score_in_buf = ptr_buf_tensor(fx.get_iter(score_in), fx.Int32) # State descriptors are rebased onto THIS program's slot. A buffer # offset is a 32-bit byte offset, so one descriptor can only reach # 4 GiB from its base; a state tensor whose slot stride is a @@ -523,17 +522,19 @@ def _load_f32_vec(rsrc, off_elems_i32): # far more than that across all slots. Folding the slot term into # the base — 64-bit pointer arithmetic, done once per program — # leaves the offset covering a single entry. - kv_state_rsrc = buffer_ops.create_buffer_resource( - kv_state, - max_size=True, - base_byte_offset=state_slot_byte_offset(slot, kv_state_slot_stride), + kv_state_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_state, state_slot_byte_offset(slot, kv_state_slot_stride) + ), + fx.Float32, ) - score_state_rsrc = buffer_ops.create_buffer_resource( - score_state, - max_size=True, - base_byte_offset=state_slot_byte_offset(slot, score_state_slot_stride), + score_state_buf = ptr_buf_tensor( + _ptr_at_byte_off( + score_state, state_slot_byte_offset(slot, score_state_slot_stride) + ), + fx.Float32, ) - ape_rsrc = buffer_ops.create_buffer_resource(ape, max_size=True) + ape_buf = ptr_buf_tensor(fx.get_iter(ape), fx.Float32) def _col_off_for_k(k_static_val): """Compute col_off ? {0, D} for OVERLAP (==head_dim when k >= RATIO), @@ -542,60 +543,39 @@ def _col_off_for_k(k_static_val): ``k_static_val`` may be a Python int (constexpr) or an MLIR i32 value. """ if const_expr(not overlap): - return arith.constant(0, type=i32) + return fx.Int32(0) if const_expr(isinstance(k_static_val, int)): - return arith.constant(D if k_static_val >= ratio else 0, type=i32) + return fx.Int32(D if k_static_val >= ratio else 0) # Dynamic: (k >= RATIO) ? D : 0 via select - is_b = arith.cmpi( - CmpIPredicate.sge, - k_static_val, - arith.constant(ratio, type=i32), - ) - return arith.select( - is_b, - arith.constant(D, type=i32), - arith.constant(0, type=i32), - ) + is_b = fx.Int32(k_static_val) >= fx.Int32(ratio) + return is_b.select(fx.Int32(D), fx.Int32(0)) # ---- Step 6: Phase 1 -- state cache loop (dynamic bound = window_len) ---- # window_len ? [0, K]. When 0, the loop is a no-op. - c_K_m1 = arith.constant(K - 1, type=i32) - c_state_size = arith.constant(state_size, type=i32) - - for k_static, state in range(0, _to_raw(window_len), 1, init=init_state): + for k_static, state in range(0, window_len, 1, init=init_state): m_lane, kv_lane, w_lane = _split_state(state) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) - s = arith.subi( - arith.addi( - arith.subi(_to_raw(position), c_K_m1), - k_i32, - ), - arith.constant(0, type=i32), - ) - is_pad = arith.cmpi(CmpIPredicate.slt, s, arith.constant(0, type=i32)) - s_safe = arith.select(is_pad, arith.constant(0, type=i32), s) - ring = arith.remui(s_safe, c_state_size) + k_i32 = fx.Int32(k_static) + s = fx.Int32(position) - (K - 1) + k_i32 + is_pad = s < fx.Int32(0) + s_safe = is_pad.select(fx.Int32(0), s) + ring = fx.Uint32(s_safe) % fx.Uint32(state_size) col_off = _col_off_for_k(k_i32) # Slot term already folded into the descriptor base. - base_kv_off = ( - ArithValue(ring) * ArithValue(kv_state_pos_stride) - + ArithValue(col_off) - + tid_x_vec - ) + base_kv_off = fx.Int32(ring) * kv_state_pos_stride + col_off + tid_x_vec base_sc_off = ( - ArithValue(ring) * ArithValue(score_state_pos_stride) - + ArithValue(col_off) - + tid_x_vec + fx.Int32(ring) * score_state_pos_stride + col_off + tid_x_vec ) - kv_v_lane = _load_f32_vec(kv_state_rsrc, base_kv_off) - sc_v_lane = _load_f32_vec(score_state_rsrc, base_sc_off) + kv_v_lane = _load_f32_vec(kv_state_buf, base_kv_off) + sc_v_lane = _load_f32_vec(score_state_buf, base_sc_off) sc_pad_lane = [] for i in range_constexpr(VEC): - sc_pad_lane.append(arith.select(is_pad, c_neg_inf, sc_v_lane[i])) + sc_pad_lane.append( + is_pad.select(c_neg_inf, fx.Float32(sc_v_lane[i])) + ) new_m, new_kv, new_w = _online_softmax_update( m_lane, kv_lane, w_lane, sc_pad_lane, kv_v_lane @@ -621,10 +601,10 @@ def _col_off_for_k(k_static_val): def _phase2_offsets(k_i32): """Compute (col_off, in_row, ape_row) for Phase 2 iter k.""" + k = fx.Int32(k_i32) col_off = _col_off_for_k(k_i32) - ape_row = arith.remui(k_i32, arith.constant(ratio, type=i32)) - tmp = arith.subi(c_K_m1, k_i32) - in_row = arith.subi(_to_raw(ragged_id), tmp) + ape_row = fx.Uint32(k) % fx.Uint32(ratio) + in_row = fx.Int32(ragged_id) - ((K - 1) - k) return col_off, in_row, ape_row def _phase2_issue_loads(k_i32): @@ -636,38 +616,21 @@ def _phase2_issue_loads(k_i32): k = K (one past the last legal iter) for prefetch tails. """ col_off, in_row, ape_row = _phase2_offsets(k_i32) - base_in_off = ( - ArithValue(in_row) * ArithValue(kv_in_row_stride) - + ArithValue(col_off) - + tid_x_vec - ) - base_sc_off = ( - ArithValue(in_row) * ArithValue(score_in_row_stride) - + ArithValue(col_off) - + tid_x_vec - ) - base_ape_off = ( - ArithValue(ape_row) * arith.constant(DIM_FULL, type=i32) - + ArithValue(col_off) - + tid_x_vec - ) - kv = _load_bf16_vec_then_f32(kv_in_rsrc, base_in_off) - sc = _load_bf16_vec_then_f32(score_in_rsrc, base_sc_off) - ape = _load_f32_vec(ape_rsrc, base_ape_off) + base_in_off = in_row * kv_in_row_stride + col_off + tid_x_vec + base_sc_off = in_row * score_in_row_stride + col_off + tid_x_vec + base_ape_off = fx.Int32(ape_row) * DIM_FULL + col_off + tid_x_vec + kv = _load_bf16_vec_then_f32(kv_in_buf, base_in_off) + sc = _load_bf16_vec_then_f32(score_in_buf, base_sc_off) + ape = _load_f32_vec(ape_buf, base_ape_off) return kv, sc, ape if const_expr(not enable_prefetch_input): - for k_static, state in range( - _to_raw(window_len), K, 1, init=phase1_state - ): + for k_static, state in range(window_len, K, 1, init=phase1_state): m_lane, kv_lane, w_lane = _split_state(state) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = fx.Int32(k_static) kv_a_lane, score_a_lane, ape_v_lane = _phase2_issue_loads(k_i32) score_k_lane = [ - arith.AddFOp( - score_a_lane[i], ape_v_lane[i], fastmath=fm_fast - ).result - for i in range(VEC) + _to_raw(score_a_lane[i] + ape_v_lane[i]) for i in range(VEC) ] new_m, new_kv, new_w = _online_softmax_update( m_lane, @@ -701,17 +664,15 @@ def _phase2_issue_loads(k_i32): # tail iter : k = K-1, consumes prefetched values, issues # no new prefetch. Gated by window_len < K so # that wl==K skips Phase 2 entirely. - c_K_m1_i32 = arith.constant(K - 1, type=i32) - k_prologue = arith.minsi(_to_raw(window_len), c_K_m1_i32) + # fx.min on signed fx.Int32 lowers to arith.minsi (byte-identical). + k_prologue = fx.min(fx.Int32(window_len), fx.Int32(K - 1)) pre_kv0, pre_sc0, pre_ape0 = _phase2_issue_loads(k_prologue) init_pf_state = ( list(phase1_state) + list(pre_kv0) + list(pre_sc0) + list(pre_ape0) ) loop_final = init_pf_state - for k_static, state in range( - _to_raw(window_len), K - 1, 1, init=init_pf_state - ): + for k_static, state in range(window_len, K - 1, 1, init=init_pf_state): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) @@ -719,13 +680,13 @@ def _phase2_issue_loads(k_i32): pre_sc = list(state[4 * VEC : 5 * VEC]) pre_ape = list(state[5 * VEC : 6 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = fx.Int32(k_static) # k+1 ? [window_len+1, K-1]: always in-bounds, no clamp. - k_next = arith.addi(k_i32, arith.constant(1, type=i32)) + k_next = k_i32 + fx.Int32(1) nxt_kv, nxt_sc, nxt_ape = _phase2_issue_loads(k_next) score_k_lane = [ - arith.AddFOp(pre_sc[i], pre_ape[i], fastmath=fm_fast).result + _to_raw(fx.Float32(pre_sc[i]) + fx.Float32(pre_ape[i])) for i in range(VEC) ] new_m, new_kv, new_w = _online_softmax_update( @@ -758,7 +719,7 @@ def _phase2_issue_loads(k_i32): pre_sc_t = list(loop_final[4 * VEC : 5 * VEC]) pre_ape_t = list(loop_final[5 * VEC : 6 * VEC]) score_k_lane_t = [ - arith.AddFOp(pre_sc_t[i], pre_ape_t[i], fastmath=fm_fast).result + _to_raw(fx.Float32(pre_sc_t[i]) + fx.Float32(pre_ape_t[i])) for i in range(VEC) ] _, new_kv_t, new_w_t = _online_softmax_update( @@ -784,48 +745,38 @@ def _phase2_issue_loads(k_i32): comp_lane = [] for i in range_constexpr(VEC): rcp_w = fx.rocdl.rcp(f32, w_final[i]) - comp_lane.append( - arith.MulFOp(kv_final[i], rcp_w, fastmath=fm_fast).result - ) + comp_lane.append(fx.Float32(kv_final[i]) * fx.Float32(rcp_w)) # ---- Step 9: RMSNorm (fp32) -- sum-of-squares across wave ---- - sq_local = arith.constant(0.0, type=f32) + sq_local = fx.Float32(0.0) for i in range_constexpr(VEC): - sq_local = arith.AddFOp( - sq_local, - arith.MulFOp(comp_lane[i], comp_lane[i], fastmath=fm_fast).result, - fastmath=fm_fast, - ).result + cl = comp_lane[i] + sq_local = sq_local + cl * cl sq_full = wave_reduce_add(sq_local) - var = arith.MulFOp(sq_full, c_inv_D, fastmath=fm_fast).result - rrms = fmath.rsqrt( - arith.AddFOp(var, c_eps, fastmath=fm_fast).result, fastmath=fm_fast - ) + var = sq_full * c_inv_D + rrms = fmath.rsqrt((var + c_eps).ir_value(), fastmath=fm_fast) # rms_weight: per-channel; this thread loads VEC values at tid*VEC. # Production atom passes bf16 (the param is cast at model load); # tests may pass fp32. Constexpr branch picks the right load. - rmsw_rsrc = buffer_ops.create_buffer_resource(rms_weight, max_size=True) if const_expr(rms_weight_is_bf16): - rmsw_lane = _load_bf16_vec_then_f32(rmsw_rsrc, tid_x_vec) + rmsw_buf = ptr_buf_tensor(fx.get_iter(rms_weight), fx.Int32) + rmsw_lane = _load_bf16_vec_then_f32(rmsw_buf, tid_x_vec) else: - rmsw_lane = _load_f32_vec(rmsw_rsrc, tid_x_vec) + rmsw_buf = ptr_buf_tensor(fx.get_iter(rms_weight), fx.Float32) + rmsw_lane = _load_f32_vec(rmsw_buf, tid_x_vec) normed_lane = [ - arith.MulFOp( - arith.MulFOp(comp_lane[i], rrms, fastmath=fm_fast).result, - rmsw_lane[i], - fastmath=fm_fast, - ).result + _to_raw(comp_lane[i] * fx.Float32(rrms) * fx.Float32(rmsw_lane[i])) for i in range(VEC) ] # ---- Step 10: GPT-J RoPE on RD tail ---- # is_rope = tid >= ROPE_THREAD_LO. RoPE applies only to those threads. - comp_pos_i32 = arith.muli( - arith.divsi(_to_raw(position), arith.constant(ratio, type=i32)), - arith.constant(ratio, type=i32), - ) + # position >= 0 here -> unsigned divide (byte-exact, cheaper ISA). + comp_pos_i32 = ( + fx.Uint32(position) // fx.Uint32(ratio) * fx.Uint32(ratio) + ).to(fx.Int32) # Always compute the rotated/passthrough values per-lane, then # store. ROPE-only threads load cos/sin; NOPE threads use the @@ -839,110 +790,73 @@ def _phase2_issue_loads(k_i32): # # cos/sin loads for NOPE threads are safe because we clamp the # row-relative index to 0 (a valid in-bounds position). - cos_rsrc = buffer_ops.create_buffer_resource(cos_cache, max_size=True) - sin_rsrc = buffer_ops.create_buffer_resource(sin_cache, max_size=True) - c_half_rd = arith.constant(RD // 2, type=i32) - cos_row_base = ArithValue(comp_pos_i32) * c_half_rd - - is_rope_t = arith.cmpi( - CmpIPredicate.sge, - _to_raw(tid), - arith.constant(ROPE_THREAD_LO, type=i32), - ) + cos_buf = ptr_buf_tensor(fx.get_iter(cos_cache), fx.BFloat16) + sin_buf = ptr_buf_tensor(fx.get_iter(sin_cache), fx.BFloat16) + cos_row_base = comp_pos_i32 * (RD // 2) + + is_rope_t = fx.Int32(tid) >= ROPE_THREAD_LO # rope_rel may be negative for NOPE threads; clamp to 0 so the # cos/sin load address is in-bounds (the loaded value is unused # because is_rope_t = false). - rope_rel_raw = ArithValue(tid) - arith.constant(ROPE_THREAD_LO, type=i32) - rope_rel = arith.maxsi(rope_rel_raw, arith.constant(0, type=i32)) - cs_lo = ArithValue(rope_rel) * arith.constant(PAIRS_PER_THREAD, type=i32) + rope_rel_raw = fx.Int32(tid) - ROPE_THREAD_LO + # fx.max on signed fx.Int32 lowers to arith.maxsi (byte-identical). + rope_rel = fx.max(rope_rel_raw, fx.Int32(0)) + cs_lo = rope_rel * PAIRS_PER_THREAD if const_expr(PAIRS_PER_THREAD == 1): - cos_b = buffer_ops.buffer_load( - cos_rsrc, - cos_row_base + cs_lo, - vec_width=1, - dtype=T.bf16, + cos_b = fx.add_offset(fx.get_iter(cos_buf), cos_row_base + cs_lo).load( + T.bf16 ) - sin_b = buffer_ops.buffer_load( - sin_rsrc, - cos_row_base + cs_lo, - vec_width=1, - dtype=T.bf16, + sin_b = fx.add_offset(fx.get_iter(sin_buf), cos_row_base + cs_lo).load( + T.bf16 ) - cos_vals = [arith.extf(f32, cos_b)] - sin_vals = [arith.extf(f32, sin_b)] + cos_vals = [fx.BFloat16(cos_b).to(fx.Float32)] + sin_vals = [fx.BFloat16(sin_b).to(fx.Float32)] else: - cos_vec = buffer_ops.buffer_load( - cos_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, - ) - sin_vec = buffer_ops.buffer_load( - sin_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, - ) - cos_vals = [ - arith.extf( - f32, - vector.extract( - cos_vec, static_position=[i], dynamic_position=[] - ), + cos_vec = fx.Vector( + fx.add_offset(fx.get_iter(cos_buf), cos_row_base + cs_lo).load( + T.vec(PAIRS_PER_THREAD, T.bf16) ) - for i in range(PAIRS_PER_THREAD) - ] - sin_vals = [ - arith.extf( - f32, - vector.extract( - sin_vec, static_position=[i], dynamic_position=[] - ), + ) + sin_vec = fx.Vector( + fx.add_offset(fx.get_iter(sin_buf), cos_row_base + cs_lo).load( + T.vec(PAIRS_PER_THREAD, T.bf16) ) - for i in range(PAIRS_PER_THREAD) - ] + ) + cos_vals = [cos_vec[i].to(fx.Float32) for i in range(PAIRS_PER_THREAD)] + sin_vals = [sin_vec[i].to(fx.Float32) for i in range(PAIRS_PER_THREAD)] # GPT-J pair rotation per VEC pair, then select rotated vs pass-through. + # ambient fast_fp_math -> `*`/`-`/`+` carry fastmath, same as the + # old explicit MulFOp/subf/AddFOp(fastmath=fast). rotated_lane = list(normed_lane) for k in range_constexpr(PAIRS_PER_THREAD): - e = normed_lane[2 * k] - o = normed_lane[2 * k + 1] + e = fx.Float32(normed_lane[2 * k]) + o = fx.Float32(normed_lane[2 * k + 1]) c = cos_vals[k] s = sin_vals[k] - new_e = arith.subf( - arith.MulFOp(e, c, fastmath=fm_fast).result, - arith.MulFOp(o, s, fastmath=fm_fast).result, - ) - new_o = arith.AddFOp( - arith.MulFOp(e, s, fastmath=fm_fast).result, - arith.MulFOp(o, c, fastmath=fm_fast).result, - fastmath=fm_fast, - ).result - rotated_lane[2 * k] = new_e - rotated_lane[2 * k + 1] = new_o + rotated_lane[2 * k] = e * c - o * s + rotated_lane[2 * k + 1] = e * s + o * c out_lane = [ - arith.select(is_rope_t, rotated_lane[i], normed_lane[i]) + _to_raw(is_rope_t.select(rotated_lane[i], normed_lane[i])) for i in range_constexpr(VEC) ] # ---- Step 11: Scatter (only when has_block_table) ---- if const_expr(has_block_table): # ci = position // ratio; block_in_seq = ci // k_per_block; - # slot_in_block = ci % k_per_block. - ci = arith.divsi(_to_raw(position), arith.constant(ratio, type=i32)) - block_in_seq = arith.divsi(ci, arith.constant(k_per_block, type=i32)) - slot_in_block = arith.remui(ci, arith.constant(k_per_block, type=i32)) + # slot_in_block = ci % k_per_block. All non-negative -> unsigned. + ci = fx.Uint32(position) // fx.Uint32(ratio) + block_in_seq = ci // fx.Uint32(k_per_block) + slot_in_block = (ci % fx.Uint32(k_per_block)).to(fx.Int32) # physical_block = block_table[batch_id, block_in_seq] - bt_rsrc = buffer_ops.create_buffer_resource(block_table, max_size=True) - bt_off = ArithValue(batch_id) * ArithValue( - block_table_seq_stride - ) + ArithValue(block_in_seq) - physical_block = buffer_ops.buffer_load( - bt_rsrc, bt_off, vec_width=1, dtype=i32 + bt_buf = ptr_buf_tensor(fx.get_iter(block_table), fx.Int32) + bt_off = fx.Int32(batch_id) * block_table_seq_stride + block_in_seq.to( + fx.Int32 ) + physical_block = fx.add_offset(fx.get_iter(bt_buf), bt_off).load(i32) if const_expr(not quant): # BF16 paged write. kv_cache layout: [NB, k_per_block, D]. @@ -950,65 +864,58 @@ def _phase2_issue_loads(k_i32): # (strides are in bf16 elements; caller passes elements.) # The block term rides on the descriptor's base, not on # the 32-bit offset -- see `block_base_bytes_i64`. - cache_off = ( - ArithValue(slot_in_block) * ArithValue(kv_cache_token_stride) - + tid_x_vec - ) + cache_off = slot_in_block * kv_cache_token_stride + tid_x_vec # Build a per-block GTensor and store VEC bf16 via dword path. # bf16 VEC ? {2, 4, 8} = {4, 8, 16} bytes = {1, 2, 4} dwords. out_vec_t = T.vec(VEC, T.bf16) - raw_vec = vector.from_elements(vecVf32, out_lane) + raw_vec = fx.Vector.from_elements(out_lane, dtype=fx.Float32) bf16_vec = raw_vec.truncf(out_vec_t) - out_rsrc = buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, kv_cache_block_stride, 2 + # bf16 kv_cache written as i32 dwords -> i32 buf; block base folded. + out_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_cache, + block_base_bytes_i64( + physical_block, kv_cache_block_stride, 2 + ), ), + fx.Int32, ) # cache_off is in bf16 elements; convert to dword for the i32-vec store. - cache_off_dw = ArithValue(cache_off) >> arith.constant(1, type=i32) + cache_off_dw = cache_off >> 1 dwords = (VEC + 1) // 2 - bf16_as_i32 = vector.bitcast(T.vec(dwords, T.i32), bf16_vec) + bf16_as_i32 = fx.Vector(bf16_vec).bitcast(fx.Int32) if const_expr(dwords == 1): # vec<1xi32> -> scalar i32 store - scalar_i32 = vector.extract( - bf16_as_i32, static_position=[0], dynamic_position=[] + fx.add_offset(fx.get_iter(out_buf), cache_off_dw).store( + bf16_as_i32[0] ) - buffer_ops.buffer_store(scalar_i32, out_rsrc, cache_off_dw) else: - buffer_ops.buffer_store(bf16_as_i32, out_rsrc, cache_off_dw) + fx.add_offset(fx.get_iter(out_buf), cache_off_dw).store( + bf16_as_i32 + ) elif const_expr(nm_asm): # -- group_fp8 (V4 nm-asm): nope fp8 + inline dup e8m0; rope bf16 # -> separate k_rope_buff. Shared emitter (byte-identical to HCA). -- # The block term rides on each descriptor's base, not on # the 32-bit offset -- see `block_base_bytes_i64`. - _nm_cache_base = ArithValue(slot_in_block) * ArithValue( - kv_cache_token_stride - ) - _nm_krope_base = ArithValue(slot_in_block) * ArithValue( - krope_token_stride - ) + _nm_cache_base = slot_in_block * kv_cache_token_stride + _nm_krope_base = slot_in_block * krope_token_stride emit_group_fp8_nm_asm_scatter( normed_lane=normed_lane, rotated_lane=rotated_lane, lane=tid, is_rope_t=is_rope_t, - cache_base=_to_raw(_nm_cache_base), - out_rsrc=buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( + cache_base=_nm_cache_base, + out_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(kv_cache))) + + fx.Int64( + block_base_bytes_i64( physical_block, kv_cache_block_stride, 1 - ), + ) ), - krope_base=_to_raw(_nm_krope_base), - krope_rsrc=buffer_ops.create_buffer_resource( - k_rope_buff, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, krope_block_stride, 2 - ), + krope_base=_nm_krope_base, + krope_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(k_rope_buff))) + + fx.Int64( + block_base_bytes_i64(physical_block, krope_block_stride, 2) ), VEC=VEC, NOPE=NOPE, @@ -1016,8 +923,6 @@ def _phase2_issue_loads(k_i32): log2_rts=log2_rts, ROPE_THREAD_LO=ROPE_THREAD_LO, wave_width=BLOCK_THREADS, - vecVf32=vecVf32, - fm_fast=fm_fast, ) elif const_expr(not quant_fp4): # ── QUANT=1: FP8 per-row scaled write + fp32 scale ── @@ -1041,32 +946,30 @@ def _phase2_issue_loads(k_i32): # since a single thread already has dword-aligned data. _, fp8_max = _fp8_const() - c_fp8_max = arith.constant(fp8_max, type=f32) - c_neg_fp8_max = arith.constant(-fp8_max, type=f32) - c_safety_floor = arith.constant(1e-4, type=f32) - c_inv_fp8_max = arith.constant(1.0 / fp8_max, type=f32) + c_fp8_max = fx.Float32(fp8_max) + c_neg_fp8_max = fx.Float32(-fp8_max) + c_safety_floor = fx.Float32(1e-4) + c_inv_fp8_max = fx.Float32(1.0 / fp8_max) # (a) per-lane amax - am_local = arith.constant(0.0, type=f32) + am_local = fx.Float32(0.0) for i in range_constexpr(VEC): - abs_v = fmath.absf(out_lane[i]) - am_local = arith.maximumf(am_local, abs_v) + am_local = fx.max(am_local, fx.Float32(fmath.absf(out_lane[i]))) amax = wave_reduce_max(am_local) - am_safe = arith.maximumf(amax, c_safety_floor) + am_safe = fx.max(amax, c_safety_floor) # (b) scale = am_safe / FP8_MAX, optionally ceil-pow2 - scale_raw = arith.MulFOp( - am_safe, c_inv_fp8_max, fastmath=fm_fast - ).result + # ambient fast_fp_math -> `*` == old MulFOp(fastmath=fast). + scale_raw = am_safe * c_inv_fp8_max if const_expr(use_ue8m0): # ceil-to-pow2 via bit trick: add 0x7FFFFF to mantissa, # mask off mantissa. If mantissa was 0, exp unchanged; # else exp += 1. - scale_i32 = scale_raw.bitcast(i32) - bits_up = ( - scale_i32 + arith.constant(0x7FFFFF, type=i32) - ) & arith.constant(0xFF800000, type=i32) - scale_v = bits_up.bitcast(f32) + scale_i32 = scale_raw.bitcast(fx.Int32) + bits_up = (scale_i32 + fx.Int32(0x7FFFFF)) & fx.Int32( + 0xFF800000 + ) + scale_v = bits_up.bitcast(fx.Float32) else: scale_v = scale_raw @@ -1078,28 +981,23 @@ def _phase2_issue_loads(k_i32): # for inputs that round to negative zero. Clamp small # negatives v ? (-2^-8, 0) to +0 first. Matches # _store_fp8_packed in qk_norm_rope_quant. - c_neg_uf = arith.constant(-(2.0**-8), type=f32) - c_zero = arith.constant(0.0, type=f32) + c_neg_uf = fx.Float32(-(2.0**-8)) + c_zero = fx.Float32(0.0) fp8_inputs = [] for i in range_constexpr(VEC): - v = arith.MulFOp( - out_lane[i], inv_scale, fastmath=fm_fast - ).result + v = fx.Float32(out_lane[i]) * fx.Float32(inv_scale) # clamp to [-FP8_MAX, +FP8_MAX] - v = arith.minimumf(arith.maximumf(v, c_neg_fp8_max), c_fp8_max) + v = fx.min(fx.max(v, c_neg_fp8_max), c_fp8_max) # NaN guard - is_tn = arith.andi( - arith.cmpf(CmpFPredicate.OLT, v, c_zero), - arith.cmpf(CmpFPredicate.OGT, v, c_neg_uf), - ) - v_safe = arith.select(is_tn, c_zero, v) + is_tn = (v < c_zero) & (v > c_neg_uf) + v_safe = _to_raw(is_tn.select(c_zero, v)) fp8_inputs.append(v_safe) # (e) pack VEC fp32 -> VEC fp8 bytes inside i32 seed # VEC=2: 1 cvt_pk_fp8_f32 call (places 2 bytes at index 0) # VEC=4: 2 calls (places 4 bytes at indices 0, 1) # VEC=8: 4 calls (places 8 bytes at indices 0..3 of 2 i32s) - c_p0 = arith.constant(0, type=i32) + c_p0 = fx.Int32(0).ir_value() if const_expr(VEC == 2): # Result in low 16 bits of i32 pk = rocdl.cvt_pk_fp8_f32( @@ -1107,10 +1005,9 @@ def _phase2_issue_loads(k_i32): ) # Pair cooperation: even tid stores dword with peer. # peer_pack (in low 16 bits) shifted to high 16 bits. - peer_pk = ArithValue(pk).shuffle_xor(1, BLOCK_THREADS) - dword = ArithValue(pk) | ( - ArithValue(peer_pk) << arith.constant(16, type=i32) - ) + pk = fx.Int32(pk) + peer_pk = pk.shuffle_xor(1, BLOCK_THREADS) + dword = pk | (peer_pk << fx.Int32(16)) elif const_expr(VEC == 4): # 4 bytes -> single i32, all in one thread. No coop. pk = rocdl.cvt_pk_fp8_f32( @@ -1141,12 +1038,17 @@ def _phase2_issue_loads(k_i32): # Both layouts share the same block base, which rides on # the descriptor rather than the 32-bit offset -- see # `block_base_bytes_i64`. Only the in-block offset differs. - out_rsrc = buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, kv_cache_block_stride, 1 + # fp8 packed dwords stored at BYTE offsets -> i8 buf so the + # element offset is a byte offset (matches offset_is_bytes); + # the i32 store value writes 4 bytes at that byte address. + out_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_cache, + block_base_bytes_i64( + physical_block, kv_cache_block_stride, 1 + ), ), + fx.Int8, ) if const_expr(preshuffle): @@ -1155,70 +1057,54 @@ def _phase2_issue_loads(k_i32): # + col_tile_id * (TILE * TILE) # + token_in_tile * TILE # + col_in_tile - c_TILE = arith.constant(_PRESHUFFLE_TILE, type=i32) - c_TILE_D = arith.constant(_PRESHUFFLE_TILE * D, type=i32) - c_TILE_TILE = arith.constant( - _PRESHUFFLE_TILE * _PRESHUFFLE_TILE, type=i32 - ) - token_tile_id = arith.divsi(slot_in_block, c_TILE) - token_in_tile = arith.remui(slot_in_block, c_TILE) + TILE = _PRESHUFFLE_TILE + token_tile_id = slot_in_block // TILE + token_in_tile = slot_in_block % TILE # d = tid * VEC; col_tile_id = d // TILE; col_in_tile = d % TILE - d_for_tid = ArithValue(tid) * arith.constant(VEC, type=i32) - col_tile_id = arith.divsi(d_for_tid, c_TILE) - col_in_tile = arith.remui(d_for_tid, c_TILE) + d_for_tid = fx.Int32(tid) * VEC + col_tile_id = d_for_tid // TILE + col_in_tile = d_for_tid % TILE in_block_off = ( - ArithValue(token_tile_id) * c_TILE_D - + ArithValue(col_tile_id) * c_TILE_TILE - + ArithValue(token_in_tile) * c_TILE - + ArithValue(col_in_tile) + token_tile_id * (TILE * D) + + col_tile_id * (TILE * TILE) + + token_in_tile * TILE + + col_in_tile ) else: # Linear layout: slot * D + tid * VEC (block base folded # into the descriptor above). - in_block_off = ArithValue(slot_in_block) * arith.constant( - D, type=i32 - ) + ArithValue(tid) * arith.constant(VEC, type=i32) + in_block_off = slot_in_block * D + fx.Int32(tid) * VEC + # in_block_off is a BYTE offset (i8 buf); the i32/vec store + # value writes its own width (4 / 8 bytes) at that address. byte_off = in_block_off if const_expr(VEC == 2): # Only even tid stores (its dword covers peer's bytes too). if (tid & 1) == 0: - buffer_ops.buffer_store( - dword, - out_rsrc, - byte_off, - offset_is_bytes=True, - ) + fx.add_offset(fx.get_iter(out_buf), byte_off).store(dword) elif const_expr(VEC == 4): - buffer_ops.buffer_store( - dword, out_rsrc, byte_off, offset_is_bytes=True - ) + fx.add_offset(fx.get_iter(out_buf), byte_off).store(dword) else: # VEC == 8: store 2 dwords (8 bytes) via vec<2xi32> - store_vec = vector.from_elements( - T.vec(2, i32), [dword[0], dword[1]] - ) - buffer_ops.buffer_store( - store_vec, - out_rsrc, - byte_off, - offset_is_bytes=True, + store_vec = fx.Vector.from_elements( + [dword[0], dword[1]], dtype=fx.Int32 ) + fx.add_offset(fx.get_iter(out_buf), byte_off).store(store_vec) # (f) lane-0 writes fp32 scale at cache_scale[phys, slot]. # Block term on the descriptor base, as above. if tid == 0: - cs_rsrc = buffer_ops.create_buffer_resource( - cache_scale, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, cache_scale_block_stride, 4 + cs_buf = ptr_buf_tensor( + _ptr_at_byte_off( + cache_scale, + block_base_bytes_i64( + physical_block, cache_scale_block_stride, 4 + ), ), + fx.Float32, ) - buffer_ops.buffer_store( - scale_v, cs_rsrc, fx.Int32(slot_in_block) - ) + fx.add_offset(fx.get_iter(cs_buf), slot_in_block).store(scale_v) else: # ── QUANT=1, FP4: per-group(32) e8m0 scale + E2M1 write ── # Mirrors dsv4_rotate_quant.cu's FP4 KV writer + the shared @@ -1236,30 +1122,18 @@ def _phase2_issue_loads(k_i32): PACKED_BYTES = VEC // 2 K_TILES = D // _FP4_K_TILE KVBS = k_per_block - c4_i32 = arith.constant(4, type=i32) - c23_i32 = arith.constant(23, type=i32) - c254_i32 = arith.constant(254, type=i32) - c16_i32 = arith.constant(16, type=i32) - c64_i32 = arith.constant(64, type=i32) - c32_i32 = arith.constant(_FP4_GROUP_SIZE, type=i32) # smallest-normal * fp4_max floor — guards all-zero groups, # matches dsv4_rotate_quant.cu eps_amax (bit-exact w/ ref). - c_eps_amax = arith.constant( - 6.0 * float.fromhex("0x1p-126"), type=f32 - ) + c_eps_amax = fx.Float32(6.0 * float.fromhex("0x1p-126")) # (a) per-lane amax, then butterfly group-reduce over NTG lanes. - am_local = arith.constant(0.0, type=f32) + am_grp = fx.Float32(0.0) for i in range_constexpr(VEC): - am_local = arith.maximumf(am_local, fmath.absf(out_lane[i])) - am_grp = _to_raw(am_local) + am_grp = fx.max(am_grp, fx.Float32(fmath.absf(out_lane[i]))) for sh_exp in range_constexpr(LOG2_NTG): off = NTG // (2 << sh_exp) - peer = _to_raw( - ArithValue(am_grp).shuffle_xor(off, BLOCK_THREADS) - ) - am_grp = arith.maximumf(am_grp, peer) - am_safe = arith.maximumf(am_grp, c_eps_amax) + am_grp = fx.max(am_grp, am_grp.shuffle_xor(off, BLOCK_THREADS)) + am_safe = _to_raw(fx.max(am_grp, c_eps_amax)) # (b) MX RoundUp e8m0 + multiplicative quant scale. e8m0 = emit_mx_e8m0_scale( @@ -1267,15 +1141,16 @@ def _phase2_issue_loads(k_i32): mode=_MxRoundInt.RoundUp, dtype=_MxDtypeInt.FP4_E2M1, ) - quant_exp = c254_i32 - e8m0 - quant_scale = (quant_exp << c23_i32).bitcast(f32) + quant_exp = fx.Int32(254) - fx.Int32(e8m0) + quant_scale = (quant_exp << fx.Int32(23)).bitcast(fx.Float32) # (c) per-element E2M1 nibble, pack VEC/2 bytes. + # emit_f32_to_e2m1 bitcasts its arg with a raw MLIR type, so + # pass a raw ir.Value; ambient fast_fp_math makes `*` carry + # fastmath=fast (byte-identical to the old MulFOp). nibs = [ emit_f32_to_e2m1( - arith.MulFOp( - out_lane[i], quant_scale, fastmath=fm_fast - ).result + (fx.Float32(out_lane[i]) * quant_scale).ir_value() ) for i in range_constexpr(VEC) ] @@ -1289,47 +1164,43 @@ def _phase2_issue_loads(k_i32): if preshuffle else k_per_block * (D // 2) ) - out_rsrc = buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, _fp4_block_bytes, 1 + out_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_cache, + block_base_bytes_i64(physical_block, _fp4_block_bytes, 1), ), + fx.Int8, ) # packed byte index within the row = (tid*VEC) / 2. - packed_start = ArithValue(tid_x_vec) >> arith.constant(1, type=i32) + packed_start = fx.Uint32(tid_x_vec) >> fx.Uint32(1) for b in range_constexpr(PACKED_BYTES): - byte_val = ArithValue(nibs[2 * b]) | ( - ArithValue(nibs[2 * b + 1]) << c4_i32 + byte_val = fx.Int32(nibs[2 * b]) | ( + fx.Int32(nibs[2 * b + 1]) << fx.Int32(4) ) - packed_idx = packed_start + arith.constant(b, type=i32) + packed_idx = fx.Int32(packed_start) + b if const_expr(preshuffle): # FP4 KV preshuffle [NB, k_tiles, 4, kvbs, 16] u8. - k_tile = arith.divsi(packed_idx, c64_i32) - rem = arith.remui(packed_idx, c64_i32) - group4 = arith.divsi(rem, c16_i32) - sub16 = arith.remui(rem, c16_i32) + # packed_idx / rem >= 0 -> unsigned divide/rem. + pk_u = fx.Uint32(packed_idx) + k_tile = fx.Int32(pk_u // fx.Uint32(64)) + rem_u = pk_u % fx.Uint32(64) + group4 = fx.Int32(rem_u // fx.Uint32(16)) + sub16 = fx.Int32(rem_u % fx.Uint32(16)) byte_off = ( - ArithValue(k_tile) - * arith.constant(4 * KVBS * 16, type=i32) - + ArithValue(group4) - * arith.constant(KVBS * 16, type=i32) - + ArithValue(slot_in_block) * c16_i32 - + ArithValue(sub16) + k_tile * (4 * KVBS * 16) + + group4 * (KVBS * 16) + + slot_in_block * 16 + + sub16 ) else: - byte_off = ArithValue(slot_in_block) * arith.constant( - D // 2, type=i32 - ) + ArithValue(packed_idx) - buffer_ops.buffer_store( - arith.trunci(T.i8, _to_raw(byte_val)), - out_rsrc, - _to_raw(byte_off), - offset_is_bytes=True, + byte_off = slot_in_block * (D // 2) + packed_idx + fx.add_offset(fx.get_iter(out_buf), byte_off).store( + fx.Int32(byte_val).to(fx.Int8) ) # (e) group-rep lane writes the e8m0 scale byte. - scale_group_idx = arith.divsi(tid_x_vec, c32_i32) + # tid*VEC >= 0 -> unsigned divide. + scale_group_idx = fx.Uint32(tid_x_vec) // fx.Uint32(_FP4_GROUP_SIZE) if tid % NTG == 0: # Block term on the descriptor base, as above; the u8 # scale plane packs a block into a constant byte count. @@ -1338,12 +1209,14 @@ def _phase2_issue_loads(k_i32): if preshuffle else k_per_block * (D // _FP4_GROUP_SIZE) ) - cs_rsrc = buffer_ops.create_buffer_resource( - cache_scale, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, _fp4_scale_block_bytes, 1 + cs_buf = ptr_buf_tensor( + _ptr_at_byte_off( + cache_scale, + block_base_bytes_i64( + physical_block, _fp4_scale_block_bytes, 1 + ), ), + fx.Int8, ) if const_expr(preshuffle): # scale [NB, k_tiles, 4, kvbs] u8, with the slot axis @@ -1353,27 +1226,20 @@ def _phase2_issue_loads(k_i32): # (KVS_NTPW == 4). Matches the op-test reference # writer `indexer_k_fp4_paged_preshuffle` and the # packed N_PHYS==1 readers in pa_mqa_logits_fp4*. - k_tile_s = arith.divsi(scale_group_idx, c4_i32) - group4_s = arith.remui(scale_group_idx, c4_i32) - sflat = ArithValue( - arith.remui(_to_raw(slot_in_block), c16_i32) - ) * c4_i32 + ArithValue( - arith.divsi(_to_raw(slot_in_block), c16_i32) - ) - cs_off = ( - ArithValue(k_tile_s) - * arith.constant(4 * KVBS, type=i32) - + ArithValue(group4_s) * arith.constant(KVBS, type=i32) - + sflat + sg_u = fx.Uint32(scale_group_idx) + k_tile_s = fx.Int32(sg_u // fx.Uint32(4)) + group4_s = fx.Int32(sg_u % fx.Uint32(4)) + slot_u = fx.Uint32(slot_in_block) + sflat = fx.Int32(slot_u % fx.Uint32(16)) * 4 + fx.Int32( + slot_u // fx.Uint32(16) ) + cs_off = k_tile_s * (4 * KVBS) + group4_s * KVBS + sflat else: - cs_off = ArithValue(slot_in_block) * arith.constant( - D // _FP4_GROUP_SIZE, type=i32 - ) + ArithValue(scale_group_idx) - buffer_ops.buffer_store( - arith.trunci(T.i8, _to_raw(e8m0)), - cs_rsrc, - _to_raw(cs_off), + cs_off = slot_in_block * (D // _FP4_GROUP_SIZE) + fx.Int32( + scale_group_idx + ) + fx.add_offset(fx.get_iter(cs_buf), cs_off).store( + fx.Int32(e8m0).to(fx.Int8) ) # e8m0 uint8 # else: warmup — no scatter, just consume compute. @@ -1626,126 +1492,102 @@ def kernel( ): f32 = T.f32 i32 = T.i32 - vecVf32 = T.vec(VEC, T.f32) pid = fx.block_idx.x tid = fx.thread_idx.x # 0 .. BLOCK_TH-1 c_neg_inf = arith.constant(_NEG_INF, type=f32) c_zero_f32 = arith.constant(0.0, type=f32) - c_zero_i32 = arith.constant(0, type=i32) - c_one_i32 = arith.constant(1, type=i32) - c_64 = arith.constant(BLOCK_THREADS, type=i32) - c_eps = arith.constant(rms_eps, type=f32) - c_inv_D = arith.constant(1.0 / D, type=f32) - c_log2e = arith.constant(_LOG2E, type=f32) - c_K_m1 = arith.constant(K - 1, type=i32) - c_K_per_wave = arith.constant(K_PER_WAVE, type=i32) - c_state_size = arith.constant(state_size, type=i32) - c_VEC = arith.constant(VEC, type=i32) - c_D = arith.constant(D, type=i32) + c_eps = fx.Float32(rms_eps) + c_inv_D = fx.Float32(1.0 / D) + c_log2e = fx.Float32(_LOG2E) def fexp_f32(x): - return fx.rocdl.exp2(f32, x * c_log2e) + # x is fx.Float32; exp2 needs a raw operand -> wrap once here. + return fx.rocdl.exp2(f32, _to_raw(x * c_log2e)) - wid = arith.divsi(_to_raw(tid), c_64) # ? [0, NW) - lid = arith.remui(_to_raw(tid), c_64) # ? [0, 64) + # tid >= 0 -> unsigned divide/rem. + tid_u = fx.Uint32(tid) + wid = (tid_u // fx.Uint32(BLOCK_THREADS)).to(fx.Int32) # ? [0, NW) + lid = (tid_u % fx.Uint32(BLOCK_THREADS)).to(fx.Int32) # ? [0, 64) # ---- plan row (single dwordx4) ---- - plan_rsrc = buffer_ops.create_buffer_resource(plan, max_size=True) - plan_base = ArithValue(pid) * arith.constant(4, type=i32) - plan_vec = buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) - ragged_id = vector.extract(plan_vec, static_position=[0], dynamic_position=[]) - batch_id = vector.extract(plan_vec, static_position=[1], dynamic_position=[]) - position = vector.extract(plan_vec, static_position=[2], dynamic_position=[]) - window_len = vector.extract(plan_vec, static_position=[3], dynamic_position=[]) + plan_buf = ptr_buf_tensor(fx.get_iter(plan), fx.Int32) + plan_vec = fx.Vector( + fx.add_offset(fx.get_iter(plan_buf), fx.Int32(pid) * 4).load(T.vec(4, i32)) + ) + ragged_id = plan_vec[0] + batch_id = plan_vec[1] + position = plan_vec[2] + window_len = plan_vec[3] # Sentinel-skip: whole body runs only for position >= 0, as a closure # under a runtime `if` (see the CSA kernel above for the rationale). def _body(): - slot_map_rsrc = buffer_ops.create_buffer_resource( - state_slot_mapping, max_size=True - ) - slot = buffer_ops.buffer_load( - slot_map_rsrc, batch_id, vec_width=1, dtype=i32 - ) + slot_map_buf = ptr_buf_tensor(fx.get_iter(state_slot_mapping), fx.Int32) + slot = fx.add_offset(fx.get_iter(slot_map_buf), batch_id).load(i32) # This lane owns columns [lid*VEC, lid*VEC+VEC) of head_dim. - lid_x_vec = ArithValue(lid) * c_VEC + lid_x_vec = lid * VEC - kv_in_rsrc = buffer_ops.create_buffer_resource(kv_in, max_size=True) - score_in_rsrc = buffer_ops.create_buffer_resource(score_in, max_size=True) + kv_in_buf = ptr_buf_tensor(fx.get_iter(kv_in), fx.Int32) + score_in_buf = ptr_buf_tensor(fx.get_iter(score_in), fx.Int32) # Rebased onto this program's slot — see `state_slot_byte_offset`. - kv_state_rsrc = buffer_ops.create_buffer_resource( - kv_state, - max_size=True, - base_byte_offset=state_slot_byte_offset(slot, kv_state_slot_stride), + kv_state_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_state, state_slot_byte_offset(slot, kv_state_slot_stride) + ), + fx.Float32, ) - score_state_rsrc = buffer_ops.create_buffer_resource( - score_state, - max_size=True, - base_byte_offset=state_slot_byte_offset(slot, score_state_slot_stride), + score_state_buf = ptr_buf_tensor( + _ptr_at_byte_off( + score_state, state_slot_byte_offset(slot, score_state_slot_stride) + ), + fx.Float32, ) - ape_rsrc = buffer_ops.create_buffer_resource(ape, max_size=True) + ape_buf = ptr_buf_tensor(fx.get_iter(ape), fx.Float32) def _col_off_for_k(k_i32): if const_expr(not overlap): - return c_zero_i32 - is_b = arith.cmpi( - CmpIPredicate.sge, k_i32, arith.constant(ratio, type=i32) - ) - return arith.select(is_b, c_D, c_zero_i32) + return fx.Int32(0) + return (fx.Int32(k_i32) >= ratio).select(fx.Int32(D), fx.Int32(0)) - def _load_f32_vec(rsrc, off_elems_i32): + def _load_f32_vec(buf, off_elems_i32): if const_expr(VEC <= 4): - raw = buffer_ops.buffer_load( - rsrc, off_elems_i32, vec_width=VEC, dtype=f32 + raw = fx.Vector( + fx.add_offset(fx.get_iter(buf), off_elems_i32).load( + T.vec(VEC, f32) + ) ) - return [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + return [raw[i] for i in range(VEC)] else: assert VEC == 8 half = VEC // 2 - r0 = buffer_ops.buffer_load( - rsrc, off_elems_i32, vec_width=half, dtype=f32 - ) - r1 = buffer_ops.buffer_load( - rsrc, - ArithValue(off_elems_i32) + arith.constant(half, type=i32), - vec_width=half, - dtype=f32, + base = fx.Int32(off_elems_i32) + r0 = fx.Vector( + fx.add_offset(fx.get_iter(buf), base).load(T.vec(half, f32)) ) - out = [] - for i in range_constexpr(half): - out.append( - vector.extract(r0, static_position=[i], dynamic_position=[]) + r1 = fx.Vector( + fx.add_offset(fx.get_iter(buf), base + half).load( + T.vec(half, f32) ) - for i in range_constexpr(half): - out.append( - vector.extract(r1, static_position=[i], dynamic_position=[]) - ) - return out + ) + return [r0[i] for i in range(half)] + [r1[i] for i in range(half)] - def _load_bf16_vec_then_f32(rsrc, off_elems_i32): - off_dw = ArithValue(off_elems_i32) >> c_one_i32 + def _load_bf16_vec_then_f32(buf, off_elems_i32): + off_dw = fx.Int32(off_elems_i32) >> 1 dwords = (VEC + 1) // 2 if const_expr(dwords == 1): - raw_s = buffer_ops.buffer_load(rsrc, off_dw, vec_width=1, dtype=i32) - raw = vector.from_elements(T.vec(1, T.i32), [raw_s]) - else: - raw = buffer_ops.buffer_load( - rsrc, off_dw, vec_width=dwords, dtype=i32 + raw = fx.Vector.from_elements( + [fx.add_offset(fx.get_iter(buf), off_dw).load(i32)], + dtype=fx.Int32, ) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - out = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, static_position=[i], dynamic_position=[] + else: + raw = fx.Vector( + fx.add_offset(fx.get_iter(buf), off_dw).load(T.vec(dwords, i32)) ) - out.append(arith.extf(f32, bf16_v)) - return out + vec_bf16 = raw.bitcast(fx.BFloat16) + return [vec_bf16[i].to(fx.Float32) for i in range_constexpr(VEC)] def _softmax_step(m_lane, kv_lane, w_lane, score_lane, kv_v_lane): """Padding-aware per-lane online-softmax update. Phase 2 scores @@ -1755,87 +1597,76 @@ def _softmax_step(m_lane, kv_lane, w_lane, score_lane, kv_v_lane): for i in range_constexpr(VEC): m_old = m_lane[i] score = score_lane[i] - m_new = arith.maximumf(m_old, score) - is_first = arith.cmpf(CmpFPredicate.OEQ, m_old, c_neg_inf) - scale_active = fexp_f32(arith.subf(m_old, m_new)) - scale_v = arith.select(is_first, c_zero_f32, scale_active) - wk_active = fexp_f32(arith.subf(score, m_new)) - is_pad = arith.cmpf(CmpFPredicate.OEQ, score, c_neg_inf) - w_k = arith.select(is_pad, c_zero_f32, wk_active) + # -inf sentinel. maximumf keeps the ambient fast_fp_math + # (matches the old raw maximumf, which resolved + # fastmath). The OEQ compares against -inf MUST emit + # fastmath -- ambient fast lets `x == -inf` fold to + # false and corrupts the sentinel -- so scope only the `==` + # to fastmath(None) (== the old raw cmpf, no fastmath). + m_old_f = fx.Float32(m_old) + score_f = fx.Float32(score) + neg_inf_f = fx.Float32(c_neg_inf) + m_new = fx.max(m_old_f, score_f).ir_value() + with fastmath(None): + is_first = m_old_f == neg_inf_f + scale_active = fexp_f32(m_old_f - m_new) + scale_v = is_first.select(c_zero_f32, scale_active) + wk_active = fexp_f32(score_f - m_new) + with fastmath(None): + is_pad = score_f == neg_inf_f + w_k = is_pad.select(c_zero_f32, wk_active) new_m.append(m_new) new_kv.append( - arith.AddFOp( - arith.MulFOp(kv_lane[i], scale_v, fastmath=fm_fast).result, - arith.MulFOp(w_k, kv_v_lane[i], fastmath=fm_fast).result, - fastmath=fm_fast, - ).result + _to_raw( + fx.Float32(kv_lane[i]) * fx.Float32(scale_v) + + fx.Float32(w_k) * fx.Float32(kv_v_lane[i]) + ) ) new_w.append( - arith.AddFOp( - arith.MulFOp(w_lane[i], scale_v, fastmath=fm_fast).result, - w_k, - fastmath=fm_fast, - ).result + _to_raw( + fx.Float32(w_lane[i]) * fx.Float32(scale_v) + + fx.Float32(w_k) + ) ) return new_m, new_kv, new_w def _phase1_loads(k_i32): - s = arith.addi(arith.subi(_to_raw(position), c_K_m1), k_i32) - is_pad = arith.cmpi(CmpIPredicate.slt, s, c_zero_i32) - s_safe = arith.select(is_pad, c_zero_i32, s) - ring = arith.remui(s_safe, c_state_size) + s = fx.Int32(position) - (K - 1) + fx.Int32(k_i32) + is_pad = s < fx.Int32(0) + s_safe = is_pad.select(fx.Int32(0), s) + ring = fx.Int32(s_safe) % state_size col_off = _col_off_for_k(k_i32) # Slot term already folded into the descriptor base. - base_kv = ( - ArithValue(ring) * ArithValue(kv_state_pos_stride) - + ArithValue(col_off) - + lid_x_vec - ) - base_sc = ( - ArithValue(ring) * ArithValue(score_state_pos_stride) - + ArithValue(col_off) - + lid_x_vec - ) - kv_v = _load_f32_vec(kv_state_rsrc, base_kv) - sc_v = _load_f32_vec(score_state_rsrc, base_sc) - sc_pad = [arith.select(is_pad, c_neg_inf, sc_v[i]) for i in range(VEC)] + base_kv = ring * kv_state_pos_stride + col_off + lid_x_vec + base_sc = ring * score_state_pos_stride + col_off + lid_x_vec + kv_v = _load_f32_vec(kv_state_buf, base_kv) + sc_v = _load_f32_vec(score_state_buf, base_sc) + sc_pad = [ + is_pad.select(c_neg_inf, fx.Float32(sc_v[i])) for i in range(VEC) + ] return kv_v, sc_pad def _phase2_loads(k_i32): col_off = _col_off_for_k(k_i32) - ape_row = arith.remui(k_i32, arith.constant(ratio, type=i32)) - in_row_raw = arith.subi(_to_raw(ragged_id), arith.subi(c_K_m1, k_i32)) - in_row = arith.maxsi(in_row_raw, c_zero_i32) - base_in = ( - ArithValue(in_row) * ArithValue(kv_in_row_stride) - + ArithValue(col_off) - + lid_x_vec - ) - base_sc = ( - ArithValue(in_row) * ArithValue(score_in_row_stride) - + ArithValue(col_off) - + lid_x_vec - ) - base_ape = ( - ArithValue(ape_row) * arith.constant(DIM_FULL, type=i32) - + ArithValue(col_off) - + lid_x_vec - ) - kv = _load_bf16_vec_then_f32(kv_in_rsrc, base_in) - sc = _load_bf16_vec_then_f32(score_in_rsrc, base_sc) - ape_v = _load_f32_vec(ape_rsrc, base_ape) - score = [ - arith.AddFOp(sc[i], ape_v[i], fastmath=fm_fast).result - for i in range(VEC) - ] + k = fx.Int32(k_i32) + ape_row = k % ratio + in_row_raw = fx.Int32(ragged_id) - ((K - 1) - k) + in_row = (in_row_raw > 0).select(in_row_raw, fx.Int32(0)) + base_in = in_row * kv_in_row_stride + col_off + lid_x_vec + base_sc = in_row * score_in_row_stride + col_off + lid_x_vec + base_ape = ape_row * DIM_FULL + col_off + lid_x_vec + kv = _load_bf16_vec_then_f32(kv_in_buf, base_in) + sc = _load_bf16_vec_then_f32(score_in_buf, base_sc) + ape_v = _load_f32_vec(ape_buf, base_ape) + score = [_to_raw(sc[i] + ape_v[i]) for i in range(VEC)] return kv, score # ---- this wave's K range [wid*KPW, (wid+1)*KPW), split at window_len ---- - k_start = ArithValue(wid) * c_K_per_wave - k_end = k_start + c_K_per_wave - wl = _to_raw(window_len) - split_lo = arith.maxsi(wl, _to_raw(k_start)) - split = arith.minsi(split_lo, _to_raw(k_end)) + k_start = wid * K_PER_WAVE + k_end = k_start + K_PER_WAVE + wl = fx.Int32(window_len) + split_lo = (wl > k_start).select(wl, k_start) + split = (split_lo < k_end).select(split_lo, k_end) init_m = [c_neg_inf for _ in range(VEC)] init_kv = [c_zero_f32 for _ in range(VEC)] @@ -1844,24 +1675,22 @@ def _phase2_loads(k_i32): # Phase 1 sub-loop [k_start, split): state cache. p1 = init_state - for k_static, state in range( - _to_raw(k_start), _to_raw(split), 1, init=init_state - ): + for k_static, state in range(k_start, split, 1, init=init_state): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = fx.Int32(k_static) kv_v, sc_v = _phase1_loads(k_i32) nm, nkv, nw = _softmax_step(m_lane, kv_lane, w_lane, sc_v, kv_v) p1 = yield list(nm) + list(nkv) + list(nw) # Phase 2 sub-loop [split, k_end): ragged input. final = p1 - for k_static, state in range(_to_raw(split), _to_raw(k_end), 1, init=p1): + for k_static, state in range(split, k_end, 1, init=p1): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = fx.Int32(k_static) kv_v, score = _phase2_loads(k_i32) nm, nkv, nw = _softmax_step(m_lane, kv_lane, w_lane, score, kv_v) final = yield list(nm) + list(nkv) + list(nw) @@ -1875,12 +1704,12 @@ def _phase2_loads(k_i32): lds_m_ptr = lds.lds_m.ptr lds_kv_ptr = lds.lds_kv.ptr lds_w_ptr = lds.lds_w.ptr - lds_thread_base = ArithValue(wid) * c_D + lid_x_vec + lds_thread_base = wid * D + lid_x_vec for i in range_constexpr(VEC): - idx_i = lds_thread_base + arith.constant(i, type=i32) - fx.ptr_store(m_local[i], lds_m_ptr + fx.Int32(idx_i)) - fx.ptr_store(kv_local[i], lds_kv_ptr + fx.Int32(idx_i)) - fx.ptr_store(w_local[i], lds_w_ptr + fx.Int32(idx_i)) + idx_i = lds_thread_base + i + fx.ptr_store(m_local[i], lds_m_ptr + idx_i) + fx.ptr_store(kv_local[i], lds_kv_ptr + idx_i) + fx.ptr_store(w_local[i], lds_w_ptr + idx_i) gpu.barrier() @@ -1888,224 +1717,175 @@ def _phase2_loads(k_i32): def _wave0(): comp_lane = [] for i in range_constexpr(VEC): - lane_off = lid_x_vec + arith.constant(i, type=i32) + lane_off = lid_x_vec + i m_g = fx.Float32(c_neg_inf) m_arr = [] for w in range_constexpr(NW): - idx_w = arith.constant(w * D, type=i32) + lane_off - m_w = fx.ptr_load(lds_m_ptr + fx.Int32(idx_w)) + idx_w = (w * D) + lane_off + m_w = fx.ptr_load(lds_m_ptr + idx_w) m_arr.append(m_w) m_g = m_g.maximumf(m_w) kv_sum = fx.Float32(0.0) w_sum = fx.Float32(0.0) for w in range_constexpr(NW): - idx_w = arith.constant(w * D, type=i32) + lane_off - kv_w = fx.ptr_load(lds_kv_ptr + fx.Int32(idx_w)) - w_w = fx.ptr_load(lds_w_ptr + fx.Int32(idx_w)) - scale_w = fx.Float32(fexp_f32(_to_raw(m_arr[w] - m_g))) + idx_w = (w * D) + lane_off + kv_w = fx.ptr_load(lds_kv_ptr + idx_w) + w_w = fx.ptr_load(lds_w_ptr + idx_w) + scale_w = fx.Float32(fexp_f32(m_arr[w] - m_g)) kv_sum = kv_sum + kv_w * scale_w w_sum = w_sum + w_w * scale_w rcp_w = fx.Float32(fx.rocdl.rcp(f32, _to_raw(w_sum))) - comp_lane.append(_to_raw(kv_sum * rcp_w)) + comp_lane.append(kv_sum * rcp_w) # ---- RMSNorm (wave-reduce sum-of-squares over wave 0) ---- def wave_reduce_add(x): - w = _to_raw(x) + w = fx.Float32(x) for sh_exp in range_constexpr(log2_block): off = BLOCK_THREADS // (2 << sh_exp) - peer = _to_raw(ArithValue(w).shuffle_xor(off, BLOCK_THREADS)) - w = arith.AddFOp(w, peer, fastmath=fm_fast).result + peer = w.shuffle_xor(off, BLOCK_THREADS) + w = w + peer return w - sq_local = arith.constant(0.0, type=f32) + sq_local = fx.Float32(0.0) for i in range_constexpr(VEC): - sq_local = arith.AddFOp( - sq_local, - arith.MulFOp( - comp_lane[i], comp_lane[i], fastmath=fm_fast - ).result, - fastmath=fm_fast, - ).result + cl = comp_lane[i] + sq_local = sq_local + cl * cl sq_full = wave_reduce_add(sq_local) - var = arith.MulFOp(sq_full, c_inv_D, fastmath=fm_fast).result - rrms = fmath.rsqrt( - arith.AddFOp(var, c_eps, fastmath=fm_fast).result, fastmath=fm_fast - ) + var = sq_full * c_inv_D + rrms = fmath.rsqrt((var + c_eps).ir_value(), fastmath=fm_fast) - rmsw_rsrc = buffer_ops.create_buffer_resource(rms_weight, max_size=True) if const_expr(rms_weight_is_bf16): - rmsw_lane = _load_bf16_vec_then_f32(rmsw_rsrc, lid_x_vec) + rmsw_buf = ptr_buf_tensor(fx.get_iter(rms_weight), fx.Int32) + rmsw_lane = _load_bf16_vec_then_f32(rmsw_buf, lid_x_vec) else: - rmsw_lane = _load_f32_vec(rmsw_rsrc, lid_x_vec) + rmsw_buf = ptr_buf_tensor(fx.get_iter(rms_weight), fx.Float32) + rmsw_lane = _load_f32_vec(rmsw_buf, lid_x_vec) normed_lane = [ - arith.MulFOp( - arith.MulFOp(comp_lane[i], rrms, fastmath=fm_fast).result, - rmsw_lane[i], - fastmath=fm_fast, - ).result + _to_raw(comp_lane[i] * fx.Float32(rrms) * fx.Float32(rmsw_lane[i])) for i in range(VEC) ] # ---- GPT-J RoPE on RD tail ---- - comp_pos_i32 = arith.muli( - arith.divsi(_to_raw(position), arith.constant(ratio, type=i32)), - arith.constant(ratio, type=i32), - ) - cos_rsrc = buffer_ops.create_buffer_resource(cos_cache, max_size=True) - sin_rsrc = buffer_ops.create_buffer_resource(sin_cache, max_size=True) - c_half_rd = arith.constant(RD // 2, type=i32) - cos_row_base = ArithValue(comp_pos_i32) * c_half_rd - - is_rope_t = arith.cmpi( - CmpIPredicate.sge, - _to_raw(lid), - arith.constant(ROPE_THREAD_LO, type=i32), - ) - rope_rel_raw = ArithValue(lid) - arith.constant( - ROPE_THREAD_LO, type=i32 - ) - rope_rel = arith.maxsi(rope_rel_raw, c_zero_i32) - cs_lo = ArithValue(rope_rel) * arith.constant( - PAIRS_PER_THREAD, type=i32 - ) + # position >= 0 -> unsigned divide (byte-exact, cheaper ISA). + comp_pos_i32 = ( + fx.Uint32(position) // fx.Uint32(ratio) * fx.Uint32(ratio) + ).to(fx.Int32) + cos_buf = ptr_buf_tensor(fx.get_iter(cos_cache), fx.BFloat16) + sin_buf = ptr_buf_tensor(fx.get_iter(sin_cache), fx.BFloat16) + cos_row_base = comp_pos_i32 * (RD // 2) + + is_rope_t = lid >= ROPE_THREAD_LO + rope_rel_raw = lid - ROPE_THREAD_LO + # fx.max on signed fx.Int32 lowers to arith.maxsi (byte-identical). + rope_rel = fx.max(rope_rel_raw, fx.Int32(0)) + cs_lo = rope_rel * PAIRS_PER_THREAD if const_expr(PAIRS_PER_THREAD == 1): - cos_b = buffer_ops.buffer_load( - cos_rsrc, cos_row_base + cs_lo, vec_width=1, dtype=T.bf16 - ) - sin_b = buffer_ops.buffer_load( - sin_rsrc, cos_row_base + cs_lo, vec_width=1, dtype=T.bf16 - ) - cos_vals = [arith.extf(f32, cos_b)] - sin_vals = [arith.extf(f32, sin_b)] + cos_b = fx.add_offset( + fx.get_iter(cos_buf), cos_row_base + cs_lo + ).load(T.bf16) + sin_b = fx.add_offset( + fx.get_iter(sin_buf), cos_row_base + cs_lo + ).load(T.bf16) + cos_vals = [fx.BFloat16(cos_b).to(fx.Float32)] + sin_vals = [fx.BFloat16(sin_b).to(fx.Float32)] else: - cos_vec = buffer_ops.buffer_load( - cos_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, + cos_vec = fx.Vector( + fx.add_offset(fx.get_iter(cos_buf), cos_row_base + cs_lo).load( + T.vec(PAIRS_PER_THREAD, T.bf16) + ) ) - sin_vec = buffer_ops.buffer_load( - sin_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, + sin_vec = fx.Vector( + fx.add_offset(fx.get_iter(sin_buf), cos_row_base + cs_lo).load( + T.vec(PAIRS_PER_THREAD, T.bf16) + ) ) cos_vals = [ - arith.extf( - f32, - vector.extract( - cos_vec, static_position=[i], dynamic_position=[] - ), - ) - for i in range(PAIRS_PER_THREAD) + cos_vec[i].to(fx.Float32) for i in range(PAIRS_PER_THREAD) ] sin_vals = [ - arith.extf( - f32, - vector.extract( - sin_vec, static_position=[i], dynamic_position=[] - ), - ) - for i in range(PAIRS_PER_THREAD) + sin_vec[i].to(fx.Float32) for i in range(PAIRS_PER_THREAD) ] rotated_lane = list(normed_lane) for kk in range_constexpr(PAIRS_PER_THREAD): - e = normed_lane[2 * kk] - o = normed_lane[2 * kk + 1] + e = fx.Float32(normed_lane[2 * kk]) + o = fx.Float32(normed_lane[2 * kk + 1]) cc = cos_vals[kk] ss = sin_vals[kk] - new_e = arith.subf( - arith.MulFOp(e, cc, fastmath=fm_fast).result, - arith.MulFOp(o, ss, fastmath=fm_fast).result, - ) - new_o = arith.AddFOp( - arith.MulFOp(e, ss, fastmath=fm_fast).result, - arith.MulFOp(o, cc, fastmath=fm_fast).result, - fastmath=fm_fast, - ).result - rotated_lane[2 * kk] = new_e - rotated_lane[2 * kk + 1] = new_o + rotated_lane[2 * kk] = e * cc - o * ss + rotated_lane[2 * kk + 1] = e * ss + o * cc out_lane = [ - arith.select(is_rope_t, rotated_lane[i], normed_lane[i]) + _to_raw(is_rope_t.select(rotated_lane[i], normed_lane[i])) for i in range_constexpr(VEC) ] # ---- paged scatter (BF16 or FP8). Emitted in wave 0; ``lid`` # (0..63) is the single-wave ``tid`` equivalent. ---- - ci = arith.divsi(_to_raw(position), arith.constant(ratio, type=i32)) - block_in_seq = arith.divsi(ci, arith.constant(k_per_block, type=i32)) - slot_in_block = arith.remui(ci, arith.constant(k_per_block, type=i32)) - bt_rsrc = buffer_ops.create_buffer_resource(block_table, max_size=True) - bt_off = ArithValue(batch_id) * ArithValue( - block_table_seq_stride - ) + ArithValue(block_in_seq) - physical_block = buffer_ops.buffer_load( - bt_rsrc, bt_off, vec_width=1, dtype=i32 + # position/ci >= 0 -> unsigned divide/rem. + ci = fx.Uint32(position) // fx.Uint32(ratio) + block_in_seq = ci // fx.Uint32(k_per_block) + slot_in_block = (ci % fx.Uint32(k_per_block)).to(fx.Int32) + bt_buf = ptr_buf_tensor(fx.get_iter(block_table), fx.Int32) + bt_off = fx.Int32(batch_id) * block_table_seq_stride + block_in_seq.to( + fx.Int32 ) + physical_block = fx.add_offset(fx.get_iter(bt_buf), bt_off).load(i32) if const_expr(not quant): # The block term rides on the descriptor's base, not on # the 32-bit offset -- see `block_base_bytes_i64`. - cache_off = ( - ArithValue(slot_in_block) * ArithValue(kv_cache_token_stride) - + lid_x_vec - ) + cache_off = slot_in_block * kv_cache_token_stride + lid_x_vec out_vec_t = T.vec(VEC, T.bf16) - raw_vec = vector.from_elements(vecVf32, out_lane) + raw_vec = fx.Vector.from_elements(out_lane, dtype=fx.Float32) bf16_vec = raw_vec.truncf(out_vec_t) - out_rsrc = buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, kv_cache_block_stride, 2 + # bf16 kv_cache written as i32 dwords -> i32 buf; block base folded. + out_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_cache, + block_base_bytes_i64( + physical_block, kv_cache_block_stride, 2 + ), ), + fx.Int32, ) - cache_off_dw = ArithValue(cache_off) >> c_one_i32 + cache_off_dw = cache_off >> 1 dwords = (VEC + 1) // 2 - bf16_as_i32 = vector.bitcast(T.vec(dwords, T.i32), bf16_vec) + bf16_as_i32 = fx.Vector(bf16_vec).bitcast(fx.Int32) if const_expr(dwords == 1): - scalar_i32 = vector.extract( - bf16_as_i32, static_position=[0], dynamic_position=[] + fx.add_offset(fx.get_iter(out_buf), cache_off_dw).store( + bf16_as_i32[0] ) - buffer_ops.buffer_store(scalar_i32, out_rsrc, cache_off_dw) else: - buffer_ops.buffer_store(bf16_as_i32, out_rsrc, cache_off_dw) + fx.add_offset(fx.get_iter(out_buf), cache_off_dw).store( + bf16_as_i32 + ) elif const_expr(nm_asm): # -- group_fp8 (V4 nm-asm): nope fp8 + inline dup e8m0; rope # bf16 -> separate k_rope_buff. Shared emitter, byte-identical # to the legacy single-wave / HCA paths (lane == wave-0 lid). -- # The block term rides on each descriptor's base, not on # the 32-bit offset -- see `block_base_bytes_i64`. - _nm_cache_base = ArithValue(slot_in_block) * ArithValue( - kv_cache_token_stride - ) - _nm_krope_base = ArithValue(slot_in_block) * ArithValue( - krope_token_stride - ) + _nm_cache_base = slot_in_block * kv_cache_token_stride + _nm_krope_base = slot_in_block * krope_token_stride emit_group_fp8_nm_asm_scatter( normed_lane=normed_lane, rotated_lane=rotated_lane, lane=lid, is_rope_t=is_rope_t, - cache_base=_to_raw(_nm_cache_base), - out_rsrc=buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( + cache_base=_nm_cache_base, + out_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(kv_cache))) + + fx.Int64( + block_base_bytes_i64( physical_block, kv_cache_block_stride, 1 - ), + ) ), - krope_base=_to_raw(_nm_krope_base), - krope_rsrc=buffer_ops.create_buffer_resource( - k_rope_buff, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, krope_block_stride, 2 - ), + krope_base=_nm_krope_base, + krope_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(k_rope_buff))) + + fx.Int64( + block_base_bytes_i64(physical_block, krope_block_stride, 2) ), VEC=VEC, NOPE=NOPE, @@ -2113,47 +1893,40 @@ def wave_reduce_add(x): log2_rts=log2_rts, ROPE_THREAD_LO=ROPE_THREAD_LO, wave_width=BLOCK_THREADS, - vecVf32=vecVf32, - fm_fast=fm_fast, ) elif const_expr(not quant_fp4): # ── FP8 per-row scaled write + fp32 scale (mirror legacy) ── # Wave-reduce-max over wave 0's 64 lanes; pair-coop dword # store via shuffle_xor(1) within the wave. def wave_reduce_max(x): - w = _to_raw(x) + w = fx.Float32(x) for sh_exp in range_constexpr(log2_block): off = BLOCK_THREADS // (2 << sh_exp) - peer = _to_raw( - ArithValue(w).shuffle_xor(off, BLOCK_THREADS) - ) - w = arith.maximumf(w, peer) + w = fx.max(w, w.shuffle_xor(off, BLOCK_THREADS)) return w _, fp8_max = _fp8_const() - c_fp8_max = arith.constant(fp8_max, type=f32) - c_neg_fp8_max = arith.constant(-fp8_max, type=f32) - c_safety_floor = arith.constant(1e-4, type=f32) - c_inv_fp8_max = arith.constant(1.0 / fp8_max, type=f32) + c_fp8_max = fx.Float32(fp8_max) + c_neg_fp8_max = fx.Float32(-fp8_max) + c_safety_floor = fx.Float32(1e-4) + c_inv_fp8_max = fx.Float32(1.0 / fp8_max) # (a) per-lane amax -> wave-reduce-max - am_local = arith.constant(0.0, type=f32) + am_local = fx.Float32(0.0) for i in range_constexpr(VEC): - abs_v = fmath.absf(out_lane[i]) - am_local = arith.maximumf(am_local, abs_v) + am_local = fx.max(am_local, fx.Float32(fmath.absf(out_lane[i]))) amax = wave_reduce_max(am_local) - am_safe = arith.maximumf(amax, c_safety_floor) + am_safe = fx.max(amax, c_safety_floor) # (b) scale = am_safe / FP8_MAX, optionally ceil-pow2 - scale_raw = arith.MulFOp( - am_safe, c_inv_fp8_max, fastmath=fm_fast - ).result + # ambient fast_fp_math -> `*` == old MulFOp(fastmath=fast). + scale_raw = am_safe * c_inv_fp8_max if const_expr(use_ue8m0): - scale_i32 = scale_raw.bitcast(i32) - bits_up = ( - scale_i32 + arith.constant(0x7FFFFF, type=i32) - ) & arith.constant(0xFF800000, type=i32) - scale_v = bits_up.bitcast(f32) + scale_i32 = scale_raw.bitcast(fx.Int32) + bits_up = (scale_i32 + fx.Int32(0x7FFFFF)) & fx.Int32( + 0xFF800000 + ) + scale_v = bits_up.bitcast(fx.Float32) else: scale_v = scale_raw @@ -2161,31 +1934,25 @@ def wave_reduce_max(x): inv_scale = fx.rocdl.rcp(f32, scale_v) # (d) per-lane fp8 cast: clamp + fnuz NaN guard - c_neg_uf = arith.constant(-(2.0**-8), type=f32) - c_zero = arith.constant(0.0, type=f32) + c_neg_uf = fx.Float32(-(2.0**-8)) + c_zero = fx.Float32(0.0) fp8_inputs = [] for i in range_constexpr(VEC): - v = arith.MulFOp( - out_lane[i], inv_scale, fastmath=fm_fast - ).result - v = arith.minimumf(arith.maximumf(v, c_neg_fp8_max), c_fp8_max) - is_tn = arith.andi( - arith.cmpf(CmpFPredicate.OLT, v, c_zero), - arith.cmpf(CmpFPredicate.OGT, v, c_neg_uf), - ) - v_safe = arith.select(is_tn, c_zero, v) + v = fx.Float32(out_lane[i]) * fx.Float32(inv_scale) + v = fx.min(fx.max(v, c_neg_fp8_max), c_fp8_max) + is_tn = (v < c_zero) & (v > c_neg_uf) + v_safe = _to_raw(is_tn.select(c_zero, v)) fp8_inputs.append(v_safe) # (e) pack VEC fp32 -> VEC fp8 bytes - c_p0 = arith.constant(0, type=i32) + c_p0 = fx.Int32(0).ir_value() if const_expr(VEC == 2): pk = rocdl.cvt_pk_fp8_f32( i32, fp8_inputs[0], fp8_inputs[1], c_p0, 0 ) - peer_pk = ArithValue(pk).shuffle_xor(1, BLOCK_THREADS) - dword = ArithValue(pk) | ( - ArithValue(peer_pk) << arith.constant(16, type=i32) - ) + pk = fx.Int32(pk) + peer_pk = pk.shuffle_xor(1, BLOCK_THREADS) + dword = pk | (peer_pk << fx.Int32(16)) elif const_expr(VEC == 4): pk = rocdl.cvt_pk_fp8_f32( i32, fp8_inputs[0], fp8_inputs[1], c_p0, 0 @@ -2212,68 +1979,60 @@ def wave_reduce_max(x): # Block base on the descriptor, not on the 32-bit offset # -- see `block_base_bytes_i64`. - out_rsrc = buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, kv_cache_block_stride, 1 + # fp8 packed dwords stored at BYTE offsets -> i8 buf (see + # _build_kernel); block base folded into the descriptor. + out_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_cache, + block_base_bytes_i64( + physical_block, kv_cache_block_stride, 1 + ), ), + fx.Int8, ) if const_expr(preshuffle): - c_TILE = arith.constant(_PRESHUFFLE_TILE, type=i32) - c_TILE_D = arith.constant(_PRESHUFFLE_TILE * D, type=i32) - c_TILE_TILE = arith.constant( - _PRESHUFFLE_TILE * _PRESHUFFLE_TILE, type=i32 - ) - token_tile_id = arith.divsi(slot_in_block, c_TILE) - token_in_tile = arith.remui(slot_in_block, c_TILE) - d_for_tid = ArithValue(lid) * arith.constant(VEC, type=i32) - col_tile_id = arith.divsi(d_for_tid, c_TILE) - col_in_tile = arith.remui(d_for_tid, c_TILE) + TILE = _PRESHUFFLE_TILE + token_tile_id = slot_in_block // TILE + token_in_tile = slot_in_block % TILE + d_for_tid = lid * VEC + col_tile_id = d_for_tid // TILE + col_in_tile = d_for_tid % TILE in_block_off = ( - ArithValue(token_tile_id) * c_TILE_D - + ArithValue(col_tile_id) * c_TILE_TILE - + ArithValue(token_in_tile) * c_TILE - + ArithValue(col_in_tile) + token_tile_id * (TILE * D) + + col_tile_id * (TILE * TILE) + + token_in_tile * TILE + + col_in_tile ) else: - in_block_off = ArithValue(slot_in_block) * arith.constant( - D, type=i32 - ) + ArithValue(lid) * arith.constant(VEC, type=i32) + in_block_off = slot_in_block * D + lid * VEC byte_off = in_block_off if const_expr(VEC == 2): if (lid & 1) == 0: - buffer_ops.buffer_store( - dword, out_rsrc, byte_off, offset_is_bytes=True - ) + fx.add_offset(fx.get_iter(out_buf), byte_off).store(dword) elif const_expr(VEC == 4): - buffer_ops.buffer_store( - dword, out_rsrc, byte_off, offset_is_bytes=True - ) + fx.add_offset(fx.get_iter(out_buf), byte_off).store(dword) else: - store_vec = vector.from_elements( - T.vec(2, i32), [dword[0], dword[1]] - ) - buffer_ops.buffer_store( - store_vec, out_rsrc, byte_off, offset_is_bytes=True + store_vec = fx.Vector.from_elements( + [dword[0], dword[1]], dtype=fx.Int32 ) + fx.add_offset(fx.get_iter(out_buf), byte_off).store(store_vec) # (f) lane-0 writes fp32 scale at cache_scale[phys, slot]. # Block term on the descriptor base, as above. if lid == 0: - cs_rsrc = buffer_ops.create_buffer_resource( - cache_scale, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, cache_scale_block_stride, 4 + cs_buf = ptr_buf_tensor( + _ptr_at_byte_off( + cache_scale, + block_base_bytes_i64( + physical_block, cache_scale_block_stride, 4 + ), ), + fx.Float32, ) - buffer_ops.buffer_store( - scale_v, cs_rsrc, fx.Int32(slot_in_block) - ) + fx.add_offset(fx.get_iter(cs_buf), slot_in_block).store(scale_v) else: # ── FP4: per-group(32) e8m0 scale + E2M1 write (mirror # legacy _build_kernel). Emitted in wave 0 where ``lid`` @@ -2287,28 +2046,16 @@ def wave_reduce_max(x): PACKED_BYTES = VEC // 2 K_TILES = D // _FP4_K_TILE KVBS = k_per_block - c4_i32 = arith.constant(4, type=i32) - c23_i32 = arith.constant(23, type=i32) - c254_i32 = arith.constant(254, type=i32) - c16_i32 = arith.constant(16, type=i32) - c64_i32 = arith.constant(64, type=i32) - c32_i32 = arith.constant(_FP4_GROUP_SIZE, type=i32) - c_eps_amax = arith.constant( - 6.0 * float.fromhex("0x1p-126"), type=f32 - ) + c_eps_amax = fx.Float32(6.0 * float.fromhex("0x1p-126")) # (a) per-lane amax, then butterfly group-reduce over NTG lanes. - am_local = arith.constant(0.0, type=f32) + am_grp = fx.Float32(0.0) for i in range_constexpr(VEC): - am_local = arith.maximumf(am_local, fmath.absf(out_lane[i])) - am_grp = _to_raw(am_local) + am_grp = fx.max(am_grp, fx.Float32(fmath.absf(out_lane[i]))) for sh_exp in range_constexpr(LOG2_NTG): off = NTG // (2 << sh_exp) - peer = _to_raw( - ArithValue(am_grp).shuffle_xor(off, BLOCK_THREADS) - ) - am_grp = arith.maximumf(am_grp, peer) - am_safe = arith.maximumf(am_grp, c_eps_amax) + am_grp = fx.max(am_grp, am_grp.shuffle_xor(off, BLOCK_THREADS)) + am_safe = _to_raw(fx.max(am_grp, c_eps_amax)) # (b) MX RoundUp e8m0 + multiplicative quant scale. e8m0 = emit_mx_e8m0_scale( @@ -2316,15 +2063,16 @@ def wave_reduce_max(x): mode=_MxRoundInt.RoundUp, dtype=_MxDtypeInt.FP4_E2M1, ) - quant_exp = c254_i32 - e8m0 - quant_scale = (quant_exp << c23_i32).bitcast(f32) + quant_exp = fx.Int32(254) - fx.Int32(e8m0) + quant_scale = (quant_exp << fx.Int32(23)).bitcast(fx.Float32) # (c) per-element E2M1 nibble, pack VEC/2 bytes. + # emit_f32_to_e2m1 bitcasts its arg with a raw MLIR type, so + # pass a raw ir.Value; ambient fast_fp_math makes `*` carry + # fastmath=fast (byte-identical to the old MulFOp). nibs = [ emit_f32_to_e2m1( - arith.MulFOp( - out_lane[i], quant_scale, fastmath=fm_fast - ).result + (fx.Float32(out_lane[i]) * quant_scale).ir_value() ) for i in range_constexpr(VEC) ] @@ -2336,47 +2084,43 @@ def wave_reduce_max(x): if preshuffle else k_per_block * (D // 2) ) - out_rsrc = buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, _fp4_block_bytes, 1 + out_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_cache, + block_base_bytes_i64(physical_block, _fp4_block_bytes, 1), ), + fx.Int8, ) - packed_start = ArithValue(lid_x_vec_i) >> arith.constant( - 1, type=i32 - ) + packed_start = fx.Uint32(lid_x_vec_i) >> fx.Uint32(1) for b in range_constexpr(PACKED_BYTES): - byte_val = ArithValue(nibs[2 * b]) | ( - ArithValue(nibs[2 * b + 1]) << c4_i32 + byte_val = fx.Int32(nibs[2 * b]) | ( + fx.Int32(nibs[2 * b + 1]) << fx.Int32(4) ) - packed_idx = packed_start + arith.constant(b, type=i32) + packed_idx = fx.Int32(packed_start) + b if const_expr(preshuffle): - k_tile = arith.divsi(packed_idx, c64_i32) - rem = arith.remui(packed_idx, c64_i32) - group4 = arith.divsi(rem, c16_i32) - sub16 = arith.remui(rem, c16_i32) + # packed_idx / rem >= 0 -> unsigned divide/rem. + pk_u = fx.Uint32(packed_idx) + k_tile = fx.Int32(pk_u // fx.Uint32(64)) + rem_u = pk_u % fx.Uint32(64) + group4 = fx.Int32(rem_u // fx.Uint32(16)) + sub16 = fx.Int32(rem_u % fx.Uint32(16)) byte_off = ( - ArithValue(k_tile) - * arith.constant(4 * KVBS * 16, type=i32) - + ArithValue(group4) - * arith.constant(KVBS * 16, type=i32) - + ArithValue(slot_in_block) * c16_i32 - + ArithValue(sub16) + k_tile * (4 * KVBS * 16) + + group4 * (KVBS * 16) + + slot_in_block * 16 + + sub16 ) else: - byte_off = ArithValue(slot_in_block) * arith.constant( - D // 2, type=i32 - ) + ArithValue(packed_idx) - buffer_ops.buffer_store( - arith.trunci(T.i8, _to_raw(byte_val)), - out_rsrc, - _to_raw(byte_off), - offset_is_bytes=True, + byte_off = slot_in_block * (D // 2) + packed_idx + fx.add_offset(fx.get_iter(out_buf), byte_off).store( + fx.Int32(byte_val).to(fx.Int8) ) # (e) group-rep lane writes the e8m0 scale byte. - scale_group_idx = arith.divsi(lid_x_vec_i, c32_i32) + # lid*VEC >= 0 -> unsigned divide. + scale_group_idx = fx.Uint32(lid_x_vec_i) // fx.Uint32( + _FP4_GROUP_SIZE + ) if lid % NTG == 0: # Block term on the descriptor base, as above. _fp4_scale_block_bytes = ( @@ -2384,12 +2128,14 @@ def wave_reduce_max(x): if preshuffle else k_per_block * (D // _FP4_GROUP_SIZE) ) - cs_rsrc = buffer_ops.create_buffer_resource( - cache_scale, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, _fp4_scale_block_bytes, 1 + cs_buf = ptr_buf_tensor( + _ptr_at_byte_off( + cache_scale, + block_base_bytes_i64( + physical_block, _fp4_scale_block_bytes, 1 + ), ), + fx.Int8, ) if const_expr(preshuffle): # scale [NB, k_tiles, 4, kvbs] u8, slot axis @@ -2397,27 +2143,20 @@ def wave_reduce_max(x): # (KVS_NTPW==4). Matches the legacy writer, the # op-test reference, and the packed N_PHYS==1 # readers in pa_mqa_logits_fp4*. - k_tile_s = arith.divsi(scale_group_idx, c4_i32) - group4_s = arith.remui(scale_group_idx, c4_i32) - sflat = ArithValue( - arith.remui(_to_raw(slot_in_block), c16_i32) - ) * c4_i32 + ArithValue( - arith.divsi(_to_raw(slot_in_block), c16_i32) - ) - cs_off = ( - ArithValue(k_tile_s) - * arith.constant(4 * KVBS, type=i32) - + ArithValue(group4_s) * arith.constant(KVBS, type=i32) - + sflat + sg_u = fx.Uint32(scale_group_idx) + k_tile_s = fx.Int32(sg_u // fx.Uint32(4)) + group4_s = fx.Int32(sg_u % fx.Uint32(4)) + slot_u = fx.Uint32(slot_in_block) + sflat = fx.Int32(slot_u % fx.Uint32(16)) * 4 + fx.Int32( + slot_u // fx.Uint32(16) ) + cs_off = k_tile_s * (4 * KVBS) + group4_s * KVBS + sflat else: - cs_off = ArithValue(slot_in_block) * arith.constant( - D // _FP4_GROUP_SIZE, type=i32 - ) + ArithValue(scale_group_idx) - buffer_ops.buffer_store( - arith.trunci(T.i8, _to_raw(e8m0)), - cs_rsrc, - _to_raw(cs_off), + cs_off = slot_in_block * (D // _FP4_GROUP_SIZE) + fx.Int32( + scale_group_idx + ) + fx.add_offset(fx.get_iter(cs_buf), cs_off).store( + fx.Int32(e8m0).to(fx.Int8) ) # e8m0 uint8 if wid == 0: diff --git a/aiter/ops/flydsl/kernels/fused_compress_attn_common.py b/aiter/ops/flydsl/kernels/fused_compress_attn_common.py index f59779d9f3..dc4382a081 100644 --- a/aiter/ops/flydsl/kernels/fused_compress_attn_common.py +++ b/aiter/ops/flydsl/kernels/fused_compress_attn_common.py @@ -10,17 +10,14 @@ byte-identical so the V4 nm-asm sparse-attn reader sees one layout). """ -from contextlib import contextmanager from functools import lru_cache -from flydsl._mlir import ir -from flydsl._mlir.dialects import rocdl, scf -from flydsl.expr import arith, range_constexpr -from flydsl.expr.arith import ArithValue, CmpFPredicate, CmpIPredicate +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import range_constexpr from flydsl.expr.typing import T from flydsl.runtime.device import get_rocm_arch -from aiter.ops.flydsl.kernels import buffer_ops, vector from aiter.utility.mx_types import ( MX_DEFAULT_ROUND_MODE as _MX_DEFAULT_MODE, ) @@ -29,19 +26,8 @@ ) from .quant_utils import emit_mx_e8m0_scale -from .tensor_shim import _to_raw - -@contextmanager -def _if_then(if_op): - """SCF IfOp then-region context manager. Auto-yields empty if missing.""" - with ir.InsertionPoint(if_op.then_block): - try: - yield if_op.then_block - finally: - blk = if_op.then_block - if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): - scf.YieldOp([]) +_AS_GLOBAL = fx.AddressSpace.Global @lru_cache(maxsize=1) @@ -65,12 +51,9 @@ def state_slot_byte_offset(slot, slot_stride_f32_elems): size. `get_element_ptr` adds this in 64-bit pointer arithmetic, so only the multiply needs widening, and the remaining offset covers one entry. """ - slot_i64 = arith.extsi(T.i64, buffer_ops._unwrap_value(slot)) - stride_i64 = arith.extsi(T.i64, buffer_ops._unwrap_value(slot_stride_f32_elems)) - return arith.muli( - arith.muli(slot_i64, stride_i64), - arith.constant(4, type=T.i64), # sizeof(f32) - ) + slot_i64 = fx.Int64(fx.Int32(slot)) + stride_i64 = fx.Int64(fx.Int32(slot_stride_f32_elems)) + return (slot_i64 * stride_i64 * fx.Int64(4)).ir_value() # x sizeof(f32) def block_base_bytes_i64(physical_block, block_stride, elem_bytes: int = 1): @@ -94,15 +77,34 @@ def block_base_bytes_i64(physical_block, block_stride, elem_bytes: int = 1): constants). ``elem_bytes`` converts: 1 for the fp8/uint8 entry caches, 2 for a bf16 one, 4 for the fp32 per-token scale. """ - blk_i64 = arith.extsi(T.i64, buffer_ops._unwrap_value(physical_block)) + blk_i64 = fx.Int64(fx.Int32(physical_block)) if isinstance(block_stride, int): - stride_i64 = arith.constant(block_stride, type=T.i64) + stride_i64 = fx.Int64(block_stride) else: - stride_i64 = arith.extsi(T.i64, buffer_ops._unwrap_value(block_stride)) - base = arith.muli(blk_i64, stride_i64) - if elem_bytes == 1: - return base - return arith.muli(base, arith.constant(elem_bytes, type=T.i64)) + stride_i64 = fx.Int64(fx.Int32(block_stride)) + base = blk_i64 * stride_i64 + if elem_bytes != 1: + base = base * fx.Int64(elem_bytes) + return base.ir_value() + + +def _global_ptr(base_i64, byte_off, elem_ir_type, align): + """Direct global pointer at ``base_i64 + byte_off``, element type ``elem_ir_type``. + + The block/slot term is already folded into ``base_i64`` (a 64-bit address), + so no 32-bit V# offset window and no 4 GiB wrap -- the reason the emitter + doesn't need a buffer resource. ``byte_off`` (i32, the lane's in-entry + position) is widened and folded into the base too, so the returned pointer + addresses the store site directly at element 0. The store is plan-bounded + and in-bounds (no OOB / sentinel), so a raw pointer is correct. + + ``elem_ir_type`` is a *scalar* type (i8 for the e8m0 byte, i32 for the fp8 / + rope dword vectors); a wider vector value is stored whole through it, with + ``align`` = the store's byte width so the wide store stays aligned. + """ + addr = base_i64 + fx.Int64(byte_off) + pt = fx.PointerType.get(elem_ir_type, address_space=_AS_GLOBAL, alignment=align) + return fx.inttoptr(pt, addr) def emit_group_fp8_nm_asm_scatter( @@ -111,119 +113,117 @@ def emit_group_fp8_nm_asm_scatter( rotated_lane, # list[VEC] f32: post-RoPE pe values (this lane's slice) lane, # i32: within-wave lane id (0..wave_width-1) is_rope_t, # i1: lane >= ROPE_THREAD_LO - cache_base, # i32: physical_block*kcache_block_stride + slot*kcache_token_stride - out_rsrc, # kv_cache buffer resource (fp8 entry [.., entry]) - krope_base, # i32: physical_block*krope_block_stride + slot*krope_token_stride - krope_rsrc, # k_rope_buff buffer resource (bf16 [.., RD]) + cache_base, # i32: byte offset of this token's fp8 entry within its block + out_base_i64, # i64: kv_cache base addr + physical_block*block_stride (bytes) + krope_base, # i32: byte offset of this token's rope row within its block + krope_base_i64, # i64: k_rope_buff base addr + physical_block*block_stride (bytes) VEC, # elems/lane (8 wave64, 16 wave32); must be a multiple of 4 NOPE, # nope_dim (head_dim - rope_head_dim) RTS, # threads per quant group (= group_size // VEC) log2_rts, ROPE_THREAD_LO, # first rope lane (= NOPE // VEC) wave_width, # 64 (wave64) or 32 (wave32) -- shuffle_xor width - vecVf32, # T.vec(VEC, f32) - fm_fast, # arith.FastMathFlags.fast ): """Emit the FP8 nope (1xG e8m0) + inline duplicated e8m0 scale + bf16 rope->separate buffer scatter (V4 nm-asm layout). Byte-identical across CSA / HCA / wave32. - Layout written into ``out_rsrc`` (fp8 entry, 1 byte/elem): + Stores go through direct global pointers built from ``out_base_i64`` / + ``krope_base_i64`` (the cache/rope base address with the per-block byte + offset already folded in). Layout written into kv_cache (fp8 entry, + 1 byte/elem): [0:NOPE) nope fp8 [NOPE:NOPE+2*nGroups) e8m0 group scale, each duplicated x2 - Rotated PE bf16 -> ``krope_rsrc`` at krope_base + (lane-ROPE_THREAD_LO)*VEC. + Rotated PE bf16 -> k_rope_buff at krope_base + (lane-ROPE_THREAD_LO)*VEC. """ - f32 = T.f32 i32 = T.i32 assert VEC % 4 == 0, f"group_fp8: VEC={VEC} must be a multiple of 4" - c0f = arith.constant(0.0, type=f32) - c_neg_uf = arith.constant(-(2.0**-8), type=f32) - c_zero_i32 = arith.constant(0, type=i32) - c_one_i32 = arith.constant(1, type=i32) + lane = fx.Int32(lane) + is_rope_t = fx.Boolean(is_rope_t) + cache_base = fx.Int32(cache_base) + krope_base = fx.Int32(krope_base) + normed = [fx.Float32(v) for v in normed_lane] + c0f = fx.Float32(0.0) + c_neg_uf = fx.Float32(-(2.0**-8)) # group-amax of |normed| over the RTS-thread group (shuffle_xor within wave) - amax_g = _to_raw(arith.constant(0.0, type=f32)) + amax_g = fx.Float32(0.0) for i in range_constexpr(VEC): - nv = arith.subf(c0f, normed_lane[i]) - av = arith.maximumf(normed_lane[i], nv) - amax_g = arith.maximumf(amax_g, av) + amax_g = fx.maximumf(amax_g, fx.maximumf(normed[i], c0f - normed[i])) for sh in range_constexpr(log2_rts): off = RTS >> (sh + 1) - peer = _to_raw(ArithValue(amax_g).shuffle_xor(off, wave_width)) - amax_g = arith.maximumf(amax_g, peer) - e8m0 = emit_mx_e8m0_scale(amax_g, mode=_MX_DEFAULT_MODE, dtype=group_fp8_mx_dtype()) - quant_exp = arith.constant(254, type=i32) - e8m0 - inv_scale = (quant_exp << arith.constant(23, type=i32)).bitcast(f32) + amax_g = fx.maximumf(amax_g, amax_g.shuffle_xor(off, wave_width)) + e8m0 = emit_mx_e8m0_scale( + amax_g.ir_value(), mode=_MX_DEFAULT_MODE, dtype=group_fp8_mx_dtype() + ) + quant_exp = fx.Int32(254) - fx.Int32(e8m0) + inv_scale = (quant_exp << fx.Int32(23)).bitcast(fx.Float32) # -- nope lanes: scaled fp8 + group-leader dup e8m0 byte -- - is_nope = arith.cmpi( - CmpIPredicate.slt, _to_raw(lane), arith.constant(ROPE_THREAD_LO, type=i32) - ) - _if_nope = scf.IfOp(is_nope) - with _if_then(_if_nope): - safe = [] - for i in range_constexpr(VEC): - sv = arith.MulFOp(normed_lane[i], inv_scale, fastmath=fm_fast).result - # e4m3fnuz -0->+0 clamp: small negatives -> +0 (cvt returns NaN otherwise) - is_tn = arith.andi( - arith.cmpf(CmpFPredicate.OLT, sv, c0f), - arith.cmpf(CmpFPredicate.OGT, sv, c_neg_uf), - ) - safe.append(arith.select(is_tn, c0f, sv)) - # pack VEC fp8 -> VEC/4 dwords (2 cvt_pk_fp8 per dword) - dwords = [] - for d in range_constexpr(VEC // 4): - pk = arith.constant(0, type=i32) - pk = rocdl.cvt_pk_fp8_f32(i32, safe[4 * d + 0], safe[4 * d + 1], pk, 0) - pk = rocdl.cvt_pk_fp8_f32(i32, safe[4 * d + 2], safe[4 * d + 3], pk, 1) - dwords.append(pk) - nope_off = ArithValue(cache_base) + ArithValue(lane) * arith.constant( - VEC, type=i32 - ) - store_vec = vector.from_elements(T.vec(VEC // 4, i32), dwords) - buffer_ops.buffer_store( - store_vec, out_rsrc, _to_raw(nope_off), offset_is_bytes=True - ) - group_id = ArithValue(lane) >> arith.constant(log2_rts, type=i32) - lane_in_group = ArithValue(lane) & arith.constant(RTS - 1, type=i32) - is_leader = arith.cmpi(CmpIPredicate.eq, _to_raw(lane_in_group), c_zero_i32) - _if_leader = scf.IfOp(is_leader) - with _if_then(_if_leader): - e8m0_i8 = arith.TruncIOp(T.i8, e8m0).result - sc_off = ( - ArithValue(cache_base) - + arith.constant(NOPE, type=i32) - + ArithValue(group_id) * arith.constant(2, type=i32) - ) - buffer_ops.buffer_store(e8m0_i8, out_rsrc, _to_raw(sc_off)) - buffer_ops.buffer_store( - e8m0_i8, out_rsrc, _to_raw(ArithValue(sc_off) + c_one_i32) - ) + # Guarded bodies live in local @flyc.jit helpers so a plain `if` lowers to + # scf.if: the DSL's if-rewrite fires on a decorated body, not on this + # imported emitter's own source (skill Sec.5). + @flyc.jit + def _nope_lanes(): + if lane < fx.Int32(ROPE_THREAD_LO): + safe = [] + for i in range_constexpr(VEC): + # inv_scale is fast-fp-math ambient -> `*` == the old + # MulFOp(fastmath=fast) (byte-identical under _DEFAULT_COMPILE_HINTS). + sv = normed[i] * inv_scale + # e4m3fnuz -0->+0 clamp: small negatives -> +0 (cvt returns NaN otherwise) + is_tn = (sv < c0f) & (sv > c_neg_uf) + safe.append(is_tn.select(c0f, sv)) + # pack VEC fp8 -> VEC/4 dwords (2 cvt_pk_fp8 per dword) + dwords = [] + for d in range_constexpr(VEC // 4): + pk = fx.Int32(0).ir_value() + pk = fx.rocdl.cvt_pk_fp8_f32( + i32, safe[4 * d + 0].ir_value(), safe[4 * d + 1].ir_value(), pk, 0 + ) + pk = fx.rocdl.cvt_pk_fp8_f32( + i32, safe[4 * d + 2].ir_value(), safe[4 * d + 3].ir_value(), pk, 1 + ) + dwords.append(fx.Int32(pk)) + nope_off = cache_base + lane * fx.Int32(VEC) + store_vec = fx.Vector.from_elements(dwords, fx.Int32) + # VEC fp8 bytes = VEC//4 i32 dwords at byte offset nope_off. + _global_ptr(out_base_i64, nope_off, i32, VEC).store(store_vec) + + if (lane & fx.Int32(RTS - 1)) == fx.Int32(0): + e8m0_i8 = fx.Int32(e8m0).to(fx.Int8) + group_id = lane >> fx.Int32(log2_rts) + sc_off = cache_base + fx.Int32(NOPE) + group_id * fx.Int32(2) + # e8m0 duplicated x2: one i8 byte at sc_off and sc_off+1. + sc_ptr = _global_ptr(out_base_i64, sc_off, T.i8, 1) + sc_ptr[0] = e8m0_i8 + sc_ptr[1] = e8m0_i8 + + _nope_lanes() # -- rope lanes: rotated bf16 -> separate k_rope_buff -- - _if_rope_q = scf.IfOp(is_rope_t) - with _if_then(_if_rope_q): - rope_rel = ArithValue(lane) - arith.constant(ROPE_THREAD_LO, type=i32) - krope_off = ArithValue(krope_base) + ArithValue(rope_rel) * arith.constant( - VEC, type=i32 - ) - rope_f32 = vector.from_elements(vecVf32, rotated_lane) - rope_bf16 = rope_f32.truncf(T.vec(VEC, T.bf16)) - dwr = (VEC + 1) // 2 - rope_i32 = vector.bitcast(T.vec(dwr, i32), rope_bf16) - krope_off_dw = ArithValue(krope_off) >> c_one_i32 - if dwr <= 4: - # VEC<=8 (wave64): single dwordx{dwr} store. - buffer_ops.buffer_store(rope_i32, krope_rsrc, _to_raw(krope_off_dw)) - else: - # VEC=16 (wave32) -> dwr=8: no dwordx8 store; split into 2x dwordx4. - c4_i32 = arith.constant(4, type=i32) - lo = vector.extract_strided_slice( - T.vec(4, i32), rope_i32, offsets=[0], sizes=[4], strides=[1] - ) - hi = vector.extract_strided_slice( - T.vec(4, i32), rope_i32, offsets=[4], sizes=[4], strides=[1] - ) - buffer_ops.buffer_store(lo, krope_rsrc, _to_raw(krope_off_dw)) - buffer_ops.buffer_store( - hi, krope_rsrc, _to_raw(ArithValue(krope_off_dw) + c4_i32) - ) + @flyc.jit + def _rope_lanes(): + if is_rope_t: + rope_rel = lane - fx.Int32(ROPE_THREAD_LO) + krope_off = krope_base + rope_rel * fx.Int32(VEC) # bf16 elements + rope_f32 = fx.Vector.from_elements(list(rotated_lane), fx.Float32) + rope_bf16 = rope_f32.truncf(T.vec(VEC, T.bf16)) + dwr = (VEC + 1) // 2 + rope_i32 = rope_bf16.bitcast(fx.Int32) # -> vec + # bf16 element offset -> byte offset (x2). dwr i32 dwords per rope row. + krope_byte = krope_off << fx.Int32(1) + if dwr <= 4: + # VEC<=8 (wave64): single dwordx{dwr} store. + _global_ptr(krope_base_i64, krope_byte, i32, 4 * dwr).store(rope_i32) + else: + # VEC=16 (wave32) -> dwr=8: no dwordx8 store; split into 2x dwordx4. + lo = fx.Vector.from_elements([rope_i32[k] for k in range(4)], fx.Int32) + hi = fx.Vector.from_elements( + [rope_i32[k + 4] for k in range(4)], fx.Int32 + ) + _global_ptr(krope_base_i64, krope_byte, i32, 16).store(lo) + _global_ptr(krope_base_i64, krope_byte + fx.Int32(16), i32, 16).store( + hi + ) + + _rope_lanes() diff --git a/aiter/ops/flydsl/kernels/fused_compress_attn_gfx1250.py b/aiter/ops/flydsl/kernels/fused_compress_attn_gfx1250.py index 7749456ff5..928581b599 100644 --- a/aiter/ops/flydsl/kernels/fused_compress_attn_gfx1250.py +++ b/aiter/ops/flydsl/kernels/fused_compress_attn_gfx1250.py @@ -25,7 +25,6 @@ # triggering a JIT recompile per dynamic-arg value). import math -from contextlib import contextmanager from functools import lru_cache import flydsl.compiler as flyc @@ -38,7 +37,7 @@ from flydsl.expr.arith import ArithValue, CmpFPredicate, CmpIPredicate from flydsl.expr.typing import Int32, Stream, T -from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels import buffer_ops from .fused_compress_attn_common import ( block_base_bytes_i64, @@ -77,23 +76,6 @@ def _fp8_const(): _PRESHUFFLE_TILE = 16 -# ============================================================================ -# scf helpers (copied verbatim from moe_gemm_2stage.py -- too small to share) -# ============================================================================ - - -@contextmanager -def _if_then(if_op): - """SCF IfOp then-region context manager. Auto-yields empty if missing.""" - with ir.InsertionPoint(if_op.then_block): - try: - yield if_op.then_block - finally: - blk = if_op.then_block - if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): - scf.YieldOp([]) - - # ============================================================================ # Kernel builder # ============================================================================ @@ -243,7 +225,6 @@ def kernel( ): f32 = T.f32 i32 = T.i32 - vecVf32 = T.vec(VEC, T.f32) # --- thread / block ids --- pid = fx.block_idx.x # one program per plan row @@ -286,22 +267,21 @@ def wave_reduce_max(x): # saves 3 buffer-load instructions per program (visible at small N # where total program count is low). plan_rsrc = buffer_ops.create_buffer_resource(plan, max_size=True) - plan_base = ArithValue(pid) * arith.constant(4, type=i32) - plan_vec = buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) - ragged_id = vector.extract(plan_vec, static_position=[0], dynamic_position=[]) - batch_id = vector.extract(plan_vec, static_position=[1], dynamic_position=[]) - position = vector.extract(plan_vec, static_position=[2], dynamic_position=[]) - window_len = vector.extract(plan_vec, static_position=[3], dynamic_position=[]) + plan_base = _to_raw(fx.Int32(pid) * 4) + plan_vec = fx.Vector( + buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) + ) + ragged_id = _to_raw(plan_vec[0]) + batch_id = _to_raw(plan_vec[1]) + position = _to_raw(plan_vec[2]) + window_len = _to_raw(plan_vec[3]) # ---- Step 2: sentinel-skip ---- - # Wrap the entire body in scf.IfOp(position >= 0). flydsl's - # `if cond: return` does NOT actually early-exit (tail kernel body - # still runs with stale values, OOB faults). The IfOp does. - is_active = arith.cmpi( - CmpIPredicate.sge, _to_raw(position), arith.constant(0, type=i32) - ) - _if_active = scf.IfOp(is_active) - with _if_then(_if_active): + # Guard the entire body on position >= 0. A bare `if cond: return` + # would NOT early-exit (the tail kernel body still runs with stale + # values -> OOB faults); an `if cond:` block scoping the whole body + # lowers to the same scf.if guard the raw IfOp built. + if fx.Int32(position) >= 0: # ---- Step 3: per-seq state slot ---- slot_map_rsrc = buffer_ops.create_buffer_resource( state_slot_mapping, max_size=True @@ -312,7 +292,7 @@ def wave_reduce_max(x): # ---- Step 4: per-thread element-range bookkeeping ---- # This thread owns columns [tid*VEC, tid*VEC+VEC) of BLOCK_D. - tid_x_vec = ArithValue(tid) * arith.constant(VEC, type=i32) + tid_x_vec = fx.Int32(tid) * VEC # ---- Step 5: online-softmax accumulator init ---- # 3 * VEC fp32 scalars carried across K iters. @@ -393,55 +373,37 @@ def _load_bf16_vec_then_f32(rsrc, off_elems_i32): # buffer_load(vec_width=1) returns a scalar i32; wrap into # vec<1xi32> before bitcasting to vec<2xbf16>. raw_s = buffer_ops.buffer_load(rsrc, off_dw, vec_width=1, dtype=i32) - raw = vector.from_elements(T.vec(1, T.i32), [raw_s]) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - out = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, - static_position=[i], - dynamic_position=[], - ) - f32_v = arith.extf(f32, bf16_v) - out.append(f32_v) - return out + raw = fx.Vector.from_elements([raw_s], dtype=fx.Int32) + vec_bf16 = raw.bitcast(fx.BFloat16) + # raw extf: .to() picks up the ambient fast_fp_math and would + # tag the widen fastmath (IR drift vs the plain extend). + return [arith.extf(f32, _to_raw(vec_bf16[i])) for i in range(VEC)] elif const_expr(dwords <= 4): - raw = buffer_ops.buffer_load( - rsrc, off_dw, vec_width=dwords, dtype=i32 - ) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - out = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, - static_position=[i], - dynamic_position=[], + raw = fx.Vector( + buffer_ops.buffer_load( + rsrc, off_dw, vec_width=dwords, dtype=i32 ) - f32_v = arith.extf(f32, bf16_v) - out.append(f32_v) - return out + ) + vec_bf16 = raw.bitcast(fx.BFloat16) + return [arith.extf(f32, _to_raw(vec_bf16[i])) for i in range(VEC)] else: # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4 half_dw = 4 half_bf16 = half_dw * 2 # 8 bf16 per chunk out = [] for chunk in range_constexpr(dwords // half_dw): - r = buffer_ops.buffer_load( - rsrc, - ArithValue(off_dw) - + arith.constant(chunk * half_dw, type=i32), - vec_width=half_dw, - dtype=i32, + r = fx.Vector( + buffer_ops.buffer_load( + rsrc, + ArithValue(off_dw) + + arith.constant(chunk * half_dw, type=i32), + vec_width=half_dw, + dtype=i32, + ) ) - vbf16 = vector.bitcast(T.vec(half_bf16, T.bf16), r) + vbf16 = r.bitcast(fx.BFloat16) for i in range_constexpr(half_bf16): - bf16_v = vector.extract( - vbf16, - static_position=[i], - dynamic_position=[], - ) - f32_v = arith.extf(f32, bf16_v) - out.append(f32_v) + out.append(arith.extf(f32, _to_raw(vbf16[i]))) return out def _load_f32_vec(rsrc, off_elems_i32): @@ -451,33 +413,29 @@ def _load_f32_vec(rsrc, off_elems_i32): VEC=16 -> 4x dwordx4 (HW max is dwordx4). """ if const_expr(VEC <= 4): - vw = VEC - raw = buffer_ops.buffer_load( - rsrc, off_elems_i32, vec_width=vw, dtype=f32 + raw = fx.Vector( + buffer_ops.buffer_load( + rsrc, off_elems_i32, vec_width=VEC, dtype=f32 + ) ) - return [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + return [_to_raw(raw[i]) for i in range(VEC)] else: # VEC in {8, 16} -> split into quarter=4 chunks quarter = 4 n_chunks = VEC // quarter out = [] for q in range_constexpr(n_chunks): - r = buffer_ops.buffer_load( - rsrc, - ArithValue(off_elems_i32) - + arith.constant(q * quarter, type=i32), - vec_width=quarter, - dtype=f32, + r = fx.Vector( + buffer_ops.buffer_load( + rsrc, + ArithValue(off_elems_i32) + + arith.constant(q * quarter, type=i32), + vec_width=quarter, + dtype=f32, + ) ) for i in range_constexpr(quarter): - out.append( - vector.extract( - r, static_position=[i], dynamic_position=[] - ) - ) + out.append(_to_raw(r[i])) return out # Buffer resources reused across K iters. @@ -507,26 +465,17 @@ def _col_off_for_k(k_static_val): if const_expr(isinstance(k_static_val, int)): return arith.constant(D if k_static_val >= ratio else 0, type=i32) # Dynamic: (k >= RATIO) ? D : 0 via select - is_b = arith.cmpi( - CmpIPredicate.sge, - k_static_val, - arith.constant(ratio, type=i32), - ) - return arith.select( - is_b, - arith.constant(D, type=i32), - arith.constant(0, type=i32), - ) + is_b = fx.Int32(k_static_val) >= ratio + return is_b.select(fx.Int32(D), fx.Int32(0)) # ---- Step 6: Phase 1 -- state cache loop (dynamic bound = window_len) ---- # window_len ? [0, K]. When 0, the loop is a no-op. c_K_m1 = arith.constant(K - 1, type=i32) - c_state_size = arith.constant(state_size, type=i32) for k_static, state in range(0, _to_raw(window_len), 1, init=init_state): m_lane, kv_lane, w_lane = _split_state(state) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = _to_raw(fx.Int32(k_static)) s = arith.subi( arith.addi( arith.subi(_to_raw(position), c_K_m1), @@ -534,20 +483,21 @@ def _col_off_for_k(k_static_val): ), arith.constant(0, type=i32), ) - is_pad = arith.cmpi(CmpIPredicate.slt, s, arith.constant(0, type=i32)) - s_safe = arith.select(is_pad, arith.constant(0, type=i32), s) - ring = arith.remui(s_safe, c_state_size) - col_off = _col_off_for_k(k_i32) - + s_fx = fx.Int32(s) + is_pad_b = s_fx < 0 + is_pad = is_pad_b.ir_value() + s_safe = is_pad_b.select(fx.Int32(0), s_fx) + ring = fx.Uint32(s_safe) % state_size # Slot term already folded into the descriptor base. + col_off_fx = fx.Int32(_col_off_for_k(k_i32)) base_kv_off = ( - ArithValue(ring) * ArithValue(kv_state_pos_stride) - + ArithValue(col_off) + fx.Int32(ring) * fx.Int32(kv_state_pos_stride) + + col_off_fx + tid_x_vec ) base_sc_off = ( - ArithValue(ring) * ArithValue(score_state_pos_stride) - + ArithValue(col_off) + fx.Int32(ring) * fx.Int32(score_state_pos_stride) + + col_off_fx + tid_x_vec ) @@ -583,9 +533,10 @@ def _col_off_for_k(k_static_val): def _phase2_offsets(k_i32): """Compute (col_off, in_row, ape_row) for Phase 2 iter k.""" col_off = _col_off_for_k(k_i32) - ape_row = arith.remui(k_i32, arith.constant(ratio, type=i32)) + # k_i32 >= 0 -> unsigned rem. + ape_row = fx.Uint32(k_i32) % ratio tmp = arith.subi(c_K_m1, k_i32) - in_row = arith.subi(_to_raw(ragged_id), tmp) + in_row = fx.Int32(ragged_id) - fx.Int32(tmp) return col_off, in_row, ape_row def _phase2_issue_loads(k_i32): @@ -597,21 +548,18 @@ def _phase2_issue_loads(k_i32): k = K (one past the last legal iter) for prefetch tails. """ col_off, in_row, ape_row = _phase2_offsets(k_i32) + col_off_fx = fx.Int32(col_off) base_in_off = ( - ArithValue(in_row) * ArithValue(kv_in_row_stride) - + ArithValue(col_off) + fx.Int32(in_row) * fx.Int32(kv_in_row_stride) + + col_off_fx + tid_x_vec ) base_sc_off = ( - ArithValue(in_row) * ArithValue(score_in_row_stride) - + ArithValue(col_off) - + tid_x_vec - ) - base_ape_off = ( - ArithValue(ape_row) * arith.constant(DIM_FULL, type=i32) - + ArithValue(col_off) + fx.Int32(in_row) * fx.Int32(score_in_row_stride) + + col_off_fx + tid_x_vec ) + base_ape_off = fx.Int32(ape_row) * DIM_FULL + col_off_fx + tid_x_vec kv = _load_bf16_vec_then_f32(kv_in_rsrc, base_in_off) sc = _load_bf16_vec_then_f32(score_in_rsrc, base_sc_off) ape = _load_f32_vec(ape_rsrc, base_ape_off) @@ -622,7 +570,7 @@ def _phase2_issue_loads(k_i32): _to_raw(window_len), K, 1, init=phase1_state ): m_lane, kv_lane, w_lane = _split_state(state) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = _to_raw(fx.Int32(k_static)) kv_a_lane, score_a_lane, ape_v_lane = _phase2_issue_loads(k_i32) score_k_lane = [ arith.AddFOp( @@ -680,7 +628,7 @@ def _phase2_issue_loads(k_i32): pre_sc = list(state[4 * VEC : 5 * VEC]) pre_ape = list(state[5 * VEC : 6 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = _to_raw(fx.Int32(k_static)) # k+1 ? [window_len+1, K-1]: always in-bounds, no clamp. k_next = arith.addi(k_i32, arith.constant(1, type=i32)) nxt_kv, nxt_sc, nxt_ape = _phase2_issue_loads(k_next) @@ -709,6 +657,14 @@ def _phase2_issue_loads(k_i32): # Tail iter at k=K-1. Gated by `window_len < K`: when wl==K # Phase 2 is empty and the IfOp returns phase1_state. + # + # KEPT as a raw value-yielding scf.IfOp (measured floor). Both + # arms produce the m/kv/w accumulator; a per-lane select is + # INVALID here (the then-arm's softmax update reads speculative + # pre_* prefetch state that is garbage when wl==K). A local + # @flyc.jit returning the branch tuple was TESTED and drifts the + # gfx1250 ISA hard (1599 -> 870 lines, wholesale reschedule) -- + # not byte-exact, so the raw IfOp stays. is_phase2_nonempty = arith.cmpi( CmpIPredicate.slt, _to_raw(window_len), @@ -793,10 +749,8 @@ def _phase2_issue_loads(k_i32): # ---- Step 10: GPT-J RoPE on RD tail ---- # is_rope = tid >= ROPE_THREAD_LO. RoPE applies only to those threads. - comp_pos_i32 = arith.muli( - arith.divsi(_to_raw(position), arith.constant(ratio, type=i32)), - arith.constant(ratio, type=i32), - ) + # position >= 0 (guarded by the sentinel-skip IfOp) -> unsigned div. + comp_pos_i32 = (fx.Uint32(position) // ratio) * ratio # Always compute the rotated/passthrough values per-lane, then # store. ROPE-only threads load cos/sin; NOPE threads use the @@ -812,20 +766,15 @@ def _phase2_issue_loads(k_i32): # row-relative index to 0 (a valid in-bounds position). cos_rsrc = buffer_ops.create_buffer_resource(cos_cache, max_size=True) sin_rsrc = buffer_ops.create_buffer_resource(sin_cache, max_size=True) - c_half_rd = arith.constant(RD // 2, type=i32) - cos_row_base = ArithValue(comp_pos_i32) * c_half_rd + cos_row_base = fx.Int32(comp_pos_i32) * (RD // 2) - is_rope_t = arith.cmpi( - CmpIPredicate.sge, - _to_raw(tid), - arith.constant(ROPE_THREAD_LO, type=i32), - ) + is_rope_t = (fx.Int32(tid) >= ROPE_THREAD_LO).ir_value() # rope_rel may be negative for NOPE threads; clamp to 0 so the # cos/sin load address is in-bounds (the loaded value is unused - # because is_rope_t = false). - rope_rel_raw = ArithValue(tid) - arith.constant(ROPE_THREAD_LO, type=i32) + # because is_rope_t = false). raw maxsi: no fx signed-int-max form. + rope_rel_raw = _to_raw(fx.Int32(tid) - ROPE_THREAD_LO) rope_rel = arith.maxsi(rope_rel_raw, arith.constant(0, type=i32)) - cs_lo = ArithValue(rope_rel) * arith.constant(PAIRS_PER_THREAD, type=i32) + cs_lo = fx.Int32(rope_rel) * PAIRS_PER_THREAD if const_expr(PAIRS_PER_THREAD == 1): cos_b = buffer_ops.buffer_load( @@ -843,34 +792,28 @@ def _phase2_issue_loads(k_i32): cos_vals = [arith.extf(f32, cos_b)] sin_vals = [arith.extf(f32, sin_b)] else: - cos_vec = buffer_ops.buffer_load( - cos_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, + cos_vec = fx.Vector( + buffer_ops.buffer_load( + cos_rsrc, + cos_row_base + cs_lo, + vec_width=PAIRS_PER_THREAD, + dtype=T.bf16, + ) ) - sin_vec = buffer_ops.buffer_load( - sin_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, + sin_vec = fx.Vector( + buffer_ops.buffer_load( + sin_rsrc, + cos_row_base + cs_lo, + vec_width=PAIRS_PER_THREAD, + dtype=T.bf16, + ) ) cos_vals = [ - arith.extf( - f32, - vector.extract( - cos_vec, static_position=[i], dynamic_position=[] - ), - ) + arith.extf(f32, _to_raw(cos_vec[i])) for i in range(PAIRS_PER_THREAD) ] sin_vals = [ - arith.extf( - f32, - vector.extract( - sin_vec, static_position=[i], dynamic_position=[] - ), - ) + arith.extf(f32, _to_raw(sin_vec[i])) for i in range(PAIRS_PER_THREAD) ] @@ -901,16 +844,18 @@ def _phase2_issue_loads(k_i32): # ---- Step 11: Scatter (only when has_block_table) ---- if const_expr(has_block_table): # ci = position // ratio; block_in_seq = ci // k_per_block; - # slot_in_block = ci % k_per_block. - ci = arith.divsi(_to_raw(position), arith.constant(ratio, type=i32)) - block_in_seq = arith.divsi(ci, arith.constant(k_per_block, type=i32)) - slot_in_block = arith.remui(ci, arith.constant(k_per_block, type=i32)) + # slot_in_block = ci % k_per_block. position >= 0 -> unsigned. + ci = fx.Uint32(position) // ratio + block_in_seq = _to_raw(ci // k_per_block) + slot_in_block = _to_raw(ci % k_per_block) + ci = _to_raw(ci) # physical_block = block_table[batch_id, block_in_seq] bt_rsrc = buffer_ops.create_buffer_resource(block_table, max_size=True) - bt_off = ArithValue(batch_id) * ArithValue( - block_table_seq_stride - ) + ArithValue(block_in_seq) + bt_off = _to_raw( + fx.Int32(batch_id) * fx.Int32(block_table_seq_stride) + + fx.Int32(block_in_seq) + ) physical_block = buffer_ops.buffer_load( bt_rsrc, bt_off, vec_width=1, dtype=i32 ) @@ -922,14 +867,15 @@ def _phase2_issue_loads(k_i32): # The block term rides on the descriptor's base, not on # the 32-bit offset -- see `block_base_bytes_i64`. cache_off = ( - ArithValue(slot_in_block) * ArithValue(kv_cache_token_stride) + fx.Int32(slot_in_block) * fx.Int32(kv_cache_token_stride) + tid_x_vec ) # Build a per-block GTensor and store VEC bf16 via dword path. # bf16 VEC ? {2, 4, 8, 16} = {4, 8, 16, 32} bytes = {1, 2, 4, 8} dwords. out_vec_t = T.vec(VEC, T.bf16) - raw_vec = vector.from_elements(vecVf32, out_lane) - bf16_vec = raw_vec.truncf(out_vec_t) + bf16_vec = fx.Vector.from_elements( + out_lane, dtype=fx.Float32 + ).truncf(out_vec_t) out_rsrc = buffer_ops.create_buffer_resource( kv_cache, max_size=True, @@ -940,31 +886,28 @@ def _phase2_issue_loads(k_i32): # cache_off is in bf16 elements; convert to dword for the i32-vec store. cache_off_dw = ArithValue(cache_off) >> arith.constant(1, type=i32) dwords = (VEC + 1) // 2 - bf16_as_i32 = vector.bitcast(T.vec(dwords, T.i32), bf16_vec) + bf16_as_i32 = fx.Vector(bf16_vec).bitcast(fx.Int32) if const_expr(dwords == 1): # vec<1xi32> -> scalar i32 store - scalar_i32 = vector.extract( - bf16_as_i32, static_position=[0], dynamic_position=[] + buffer_ops.buffer_store( + _to_raw(bf16_as_i32[0]), out_rsrc, cache_off_dw ) - buffer_ops.buffer_store(scalar_i32, out_rsrc, cache_off_dw) elif const_expr(dwords <= 4): - buffer_ops.buffer_store(bf16_as_i32, out_rsrc, cache_off_dw) + buffer_ops.buffer_store( + _to_raw(bf16_as_i32), out_rsrc, cache_off_dw + ) else: - # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4 + # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4. c4_i32 = arith.constant(4, type=i32) - lo = vector.extract_strided_slice( - T.vec(4, T.i32), - bf16_as_i32, - offsets=[0], - sizes=[4], - strides=[1], + lo = _to_raw( + fx.Vector.from_elements( + [bf16_as_i32[i] for i in range(4)], dtype=fx.Int32 + ) ) - hi = vector.extract_strided_slice( - T.vec(4, T.i32), - bf16_as_i32, - offsets=[4], - sizes=[4], - strides=[1], + hi = _to_raw( + fx.Vector.from_elements( + [bf16_as_i32[i] for i in range(4, 8)], dtype=fx.Int32 + ) ) buffer_ops.buffer_store(lo, out_rsrc, cache_off_dw) buffer_ops.buffer_store( @@ -975,10 +918,10 @@ def _phase2_issue_loads(k_i32): # as wave64 -- single source of truth). -- # The block term rides on each descriptor's base, not on # the 32-bit offset -- see `block_base_bytes_i64`. - _nm_cache_base = ArithValue(slot_in_block) * ArithValue( + _nm_cache_base = fx.Int32(slot_in_block) * fx.Int32( kv_cache_token_stride ) - _nm_krope_base = ArithValue(slot_in_block) * ArithValue( + _nm_krope_base = fx.Int32(slot_in_block) * fx.Int32( krope_token_stride ) emit_group_fp8_nm_asm_scatter( @@ -987,20 +930,16 @@ def _phase2_issue_loads(k_i32): lane=tid, is_rope_t=is_rope_t, cache_base=_to_raw(_nm_cache_base), - out_rsrc=buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( + out_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(kv_cache))) + + fx.Int64( + block_base_bytes_i64( physical_block, kv_cache_block_stride, 1 - ), + ) ), krope_base=_to_raw(_nm_krope_base), - krope_rsrc=buffer_ops.create_buffer_resource( - k_rope_buff, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, krope_block_stride, 2 - ), + krope_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(k_rope_buff))) + + fx.Int64( + block_base_bytes_i64(physical_block, krope_block_stride, 2) ), VEC=VEC, NOPE=NOPE, @@ -1008,8 +947,6 @@ def _phase2_issue_loads(k_i32): log2_rts=log2_rts, ROPE_THREAD_LO=ROPE_THREAD_LO, wave_width=BLOCK_THREADS, - vecVf32=vecVf32, - fm_fast=fm_fast, ) else: # -- QUANT=1: FP8 per-row scaled write + fp32 scale -- @@ -1080,7 +1017,7 @@ def _phase2_issue_loads(k_i32): out_lane[i], inv_scale, fastmath=fm_fast ).result # clamp to [-FP8_MAX, +FP8_MAX] - v = arith.minimumf(arith.maximumf(v, c_neg_fp8_max), c_fp8_max) + v = _to_raw(fx.min(fx.max(v, c_neg_fp8_max), c_fp8_max)) # NaN guard is_tn = arith.andi( arith.cmpf(CmpFPredicate.OLT, v, c_zero), @@ -1157,40 +1094,32 @@ def _phase2_issue_loads(k_i32): # + col_tile_id * (TILE * TILE) # + token_in_tile * TILE # + col_in_tile - c_TILE = arith.constant(_PRESHUFFLE_TILE, type=i32) - c_TILE_D = arith.constant(_PRESHUFFLE_TILE * D, type=i32) - c_TILE_TILE = arith.constant( - _PRESHUFFLE_TILE * _PRESHUFFLE_TILE, type=i32 - ) - token_tile_id = arith.divsi(slot_in_block, c_TILE) - token_in_tile = arith.remui(slot_in_block, c_TILE) + # slot_in_block, d are non-negative -> unsigned div/rem. + TILE = _PRESHUFFLE_TILE + slot_u = fx.Uint32(slot_in_block) + token_tile_id = slot_u // TILE + token_in_tile = slot_u % TILE # d = tid * VEC; col_tile_id = d // TILE; col_in_tile = d % TILE - d_for_tid = ArithValue(tid) * arith.constant(VEC, type=i32) - col_tile_id = arith.divsi(d_for_tid, c_TILE) - col_in_tile = arith.remui(d_for_tid, c_TILE) - in_block_off = ( - ArithValue(token_tile_id) * c_TILE_D - + ArithValue(col_tile_id) * c_TILE_TILE - + ArithValue(token_in_tile) * c_TILE - + ArithValue(col_in_tile) + d_for_tid = fx.Uint32(tid) * VEC + col_tile_id = d_for_tid // TILE + col_in_tile = d_for_tid % TILE + in_block_off = _to_raw( + token_tile_id * (TILE * D) + + col_tile_id * (TILE * TILE) + + token_in_tile * TILE + + col_in_tile ) else: # Linear layout: slot * D + tid * VEC - in_block_off = ArithValue(slot_in_block) * arith.constant( - D, type=i32 - ) + ArithValue(tid) * arith.constant(VEC, type=i32) + in_block_off = _to_raw( + fx.Int32(slot_in_block) * D + fx.Int32(tid) * VEC + ) byte_off = in_block_off if const_expr(VEC == 2): # Only even tid stores (its dword covers peer's bytes too). - is_even = arith.cmpi( - CmpIPredicate.eq, - arith.andi(_to_raw(tid), arith.constant(1, type=i32)), - arith.constant(0, type=i32), - ) - _if_even = scf.IfOp(is_even) - with _if_then(_if_even): + if (fx.Int32(tid) & 1) == 0: buffer_ops.buffer_store( dword, out_rsrc, @@ -1205,11 +1134,11 @@ def _phase2_issue_loads(k_i32): # VEC in {8, 16}: store n_dwords via dwordx4 chunks n_dw = VEC // 4 if const_expr(n_dw <= 4): - store_vec = vector.from_elements( - T.vec(n_dw, i32), list(dword) + store_vec = fx.Vector.from_elements( + list(dword), dtype=fx.Int32 ) buffer_ops.buffer_store( - store_vec, + _to_raw(store_vec), out_rsrc, byte_off, offset_is_bytes=True, @@ -1219,12 +1148,11 @@ def _phase2_issue_loads(k_i32): # kept for future-proofing) for chunk_start in range_constexpr(n_dw // 4): base = chunk_start * 4 - sv = vector.from_elements( - T.vec(4, i32), - list(dword[base : base + 4]), + sv = fx.Vector.from_elements( + list(dword[base : base + 4]), dtype=fx.Int32 ) buffer_ops.buffer_store( - sv, + _to_raw(sv), out_rsrc, ArithValue(byte_off) + arith.constant(base * 4, type=i32), @@ -1232,13 +1160,7 @@ def _phase2_issue_loads(k_i32): ) # (f) lane-0 writes fp32 scale at cache_scale[phys, slot] - is_lane0 = arith.cmpi( - CmpIPredicate.eq, - _to_raw(tid), - arith.constant(0, type=i32), - ) - _if_l0 = scf.IfOp(is_lane0) - with _if_then(_if_l0): + if fx.Int32(tid) == 0: cs_rsrc = buffer_ops.create_buffer_resource( cache_scale, max_size=True, @@ -1246,9 +1168,7 @@ def _phase2_issue_loads(k_i32): physical_block, cache_scale_block_stride, 4 ), ) - buffer_ops.buffer_store( - scale_v, cs_rsrc, ArithValue(slot_in_block) - ) + buffer_ops.buffer_store(scale_v, cs_rsrc, slot_in_block) # else: warmup -- no scatter, just consume compute. @flyc.jit @@ -1453,7 +1373,6 @@ def kernel( ): f32 = T.f32 i32 = T.i32 - vecVf32 = T.vec(VEC, T.f32) pid = fx.block_idx.x tid = fx.thread_idx.x # 0 .. BLOCK_TH-1 @@ -1462,36 +1381,32 @@ def kernel( c_zero_f32 = arith.constant(0.0, type=f32) c_zero_i32 = arith.constant(0, type=i32) c_one_i32 = arith.constant(1, type=i32) - c_WS = arith.constant(BLOCK_THREADS, type=i32) c_eps = arith.constant(rms_eps, type=f32) c_inv_D = arith.constant(1.0 / D, type=f32) c_log2e = arith.constant(_LOG2E, type=f32) c_K_m1 = arith.constant(K - 1, type=i32) - c_K_per_wave = arith.constant(K_PER_WAVE, type=i32) - c_state_size = arith.constant(state_size, type=i32) - c_VEC = arith.constant(VEC, type=i32) - c_D = arith.constant(D, type=i32) def fexp_f32(x): return llvm.call_intrinsic( f32, "llvm.amdgcn.exp2.f32", [x * c_log2e], [], [] ) - wid = arith.divsi(_to_raw(tid), c_WS) # ? [0, NW) - lid = arith.remui(_to_raw(tid), c_WS) # ? [0, 32) + wid = fx.Uint32(tid) // BLOCK_THREADS # -> [0, NW) + lid = fx.Uint32(tid) % BLOCK_THREADS # -> [0, 32) # ---- plan row (single dwordx4) ---- plan_rsrc = buffer_ops.create_buffer_resource(plan, max_size=True) - plan_base = ArithValue(pid) * arith.constant(4, type=i32) + plan_base = _to_raw(fx.Int32(pid) * 4) plan_vec = buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) - ragged_id = vector.extract(plan_vec, static_position=[0], dynamic_position=[]) - batch_id = vector.extract(plan_vec, static_position=[1], dynamic_position=[]) - position = vector.extract(plan_vec, static_position=[2], dynamic_position=[]) - window_len = vector.extract(plan_vec, static_position=[3], dynamic_position=[]) - - is_active = arith.cmpi(CmpIPredicate.sge, _to_raw(position), c_zero_i32) - _if_active = scf.IfOp(is_active) - with _if_then(_if_active): + plan_vec = fx.Vector(plan_vec) + ragged_id = _to_raw(plan_vec[0]) + batch_id = _to_raw(plan_vec[1]) + position = _to_raw(plan_vec[2]) + window_len = _to_raw(plan_vec[3]) + + # Sentinel-skip: guard the whole body on position >= 0 (same scf.if + # the raw IfOp built; a bare `if cond: return` would not early-exit). + if fx.Int32(position) >= 0: slot_map_rsrc = buffer_ops.create_buffer_resource( state_slot_mapping, max_size=True ) @@ -1500,7 +1415,7 @@ def fexp_f32(x): ) # This lane owns columns [lid*VEC, lid*VEC+VEC) of head_dim. - lid_x_vec = ArithValue(lid) * c_VEC + lid_x_vec = fx.Int32(lid) * VEC kv_in_rsrc = buffer_ops.create_buffer_resource(kv_in, max_size=True) score_in_rsrc = buffer_ops.create_buffer_resource(score_in, max_size=True) @@ -1520,86 +1435,72 @@ def fexp_f32(x): def _col_off_for_k(k_i32): if const_expr(not overlap): return c_zero_i32 - is_b = arith.cmpi( - CmpIPredicate.sge, k_i32, arith.constant(ratio, type=i32) - ) - return arith.select(is_b, c_D, c_zero_i32) + is_b = fx.Int32(k_i32) >= ratio + return is_b.select(fx.Int32(D), fx.Int32(0)) def _load_f32_vec(rsrc, off_elems_i32): if const_expr(VEC <= 4): - raw = buffer_ops.buffer_load( - rsrc, off_elems_i32, vec_width=VEC, dtype=f32 + raw = fx.Vector( + buffer_ops.buffer_load( + rsrc, off_elems_i32, vec_width=VEC, dtype=f32 + ) ) - return [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + return [_to_raw(raw[i]) for i in range(VEC)] else: # VEC in {8, 16} -> split into quarter=4 chunks quarter = 4 n_chunks = VEC // quarter out = [] for q in range_constexpr(n_chunks): - r = buffer_ops.buffer_load( - rsrc, - ArithValue(off_elems_i32) - + arith.constant(q * quarter, type=i32), - vec_width=quarter, - dtype=f32, + r = fx.Vector( + buffer_ops.buffer_load( + rsrc, + ArithValue(off_elems_i32) + + arith.constant(q * quarter, type=i32), + vec_width=quarter, + dtype=f32, + ) ) for i in range_constexpr(quarter): - out.append( - vector.extract( - r, static_position=[i], dynamic_position=[] - ) - ) + out.append(_to_raw(r[i])) return out def _load_bf16_vec_then_f32(rsrc, off_elems_i32): off_dw = ArithValue(off_elems_i32) >> c_one_i32 dwords = (VEC + 1) // 2 + # raw extf: .to() picks up the ambient fast_fp_math and would tag + # the bf16->f32 widen fastmath (IR drift vs the plain extend). if const_expr(dwords == 1): raw_s = buffer_ops.buffer_load(rsrc, off_dw, vec_width=1, dtype=i32) - raw = vector.from_elements(T.vec(1, T.i32), [raw_s]) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - out = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, static_position=[i], dynamic_position=[] - ) - out.append(arith.extf(f32, bf16_v)) - return out + raw = fx.Vector.from_elements([raw_s], dtype=fx.Int32) + vec_bf16 = raw.bitcast(fx.BFloat16) + return [arith.extf(f32, _to_raw(vec_bf16[i])) for i in range(VEC)] elif const_expr(dwords <= 4): - raw = buffer_ops.buffer_load( - rsrc, off_dw, vec_width=dwords, dtype=i32 - ) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - out = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, static_position=[i], dynamic_position=[] + raw = fx.Vector( + buffer_ops.buffer_load( + rsrc, off_dw, vec_width=dwords, dtype=i32 ) - out.append(arith.extf(f32, bf16_v)) - return out + ) + vec_bf16 = raw.bitcast(fx.BFloat16) + return [arith.extf(f32, _to_raw(vec_bf16[i])) for i in range(VEC)] else: # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4 half_dw = 4 half_bf16 = half_dw * 2 out = [] for chunk in range_constexpr(dwords // half_dw): - r = buffer_ops.buffer_load( - rsrc, - ArithValue(off_dw) - + arith.constant(chunk * half_dw, type=i32), - vec_width=half_dw, - dtype=i32, + r = fx.Vector( + buffer_ops.buffer_load( + rsrc, + ArithValue(off_dw) + + arith.constant(chunk * half_dw, type=i32), + vec_width=half_dw, + dtype=i32, + ) ) - vbf16 = vector.bitcast(T.vec(half_bf16, T.bf16), r) + vbf16 = r.bitcast(fx.BFloat16) for i in range_constexpr(half_bf16): - bf16_v = vector.extract( - vbf16, static_position=[i], dynamic_position=[] - ) - out.append(arith.extf(f32, bf16_v)) + out.append(arith.extf(f32, _to_raw(vbf16[i]))) return out def _softmax_step(m_lane, kv_lane, w_lane, score_lane, kv_v_lane): @@ -1636,19 +1537,21 @@ def _softmax_step(m_lane, kv_lane, w_lane, score_lane, kv_v_lane): def _phase1_loads(k_i32): s = arith.addi(arith.subi(_to_raw(position), c_K_m1), k_i32) - is_pad = arith.cmpi(CmpIPredicate.slt, s, c_zero_i32) - s_safe = arith.select(is_pad, c_zero_i32, s) - ring = arith.remui(s_safe, c_state_size) - col_off = _col_off_for_k(k_i32) + s_fx = fx.Int32(s) + is_pad_b = s_fx < 0 + is_pad = is_pad_b.ir_value() + s_safe = is_pad_b.select(fx.Int32(0), s_fx) + ring = fx.Uint32(s_safe) % state_size + col_off_fx = fx.Int32(_col_off_for_k(k_i32)) # Slot term already folded into the descriptor base. base_kv = ( - ArithValue(ring) * ArithValue(kv_state_pos_stride) - + ArithValue(col_off) + fx.Int32(ring) * fx.Int32(kv_state_pos_stride) + + col_off_fx + lid_x_vec ) base_sc = ( - ArithValue(ring) * ArithValue(score_state_pos_stride) - + ArithValue(col_off) + fx.Int32(ring) * fx.Int32(score_state_pos_stride) + + col_off_fx + lid_x_vec ) kv_v = _load_f32_vec(kv_state_rsrc, base_kv) @@ -1657,25 +1560,23 @@ def _phase1_loads(k_i32): return kv_v, sc_pad def _phase2_loads(k_i32): - col_off = _col_off_for_k(k_i32) - ape_row = arith.remui(k_i32, arith.constant(ratio, type=i32)) + col_off_fx = fx.Int32(_col_off_for_k(k_i32)) + # k_i32 >= 0 -> unsigned rem. + ape_row = fx.Uint32(k_i32) % ratio in_row_raw = arith.subi(_to_raw(ragged_id), arith.subi(c_K_m1, k_i32)) + # raw maxsi: no fx signed-int-max form. in_row = arith.maxsi(in_row_raw, c_zero_i32) base_in = ( - ArithValue(in_row) * ArithValue(kv_in_row_stride) - + ArithValue(col_off) + fx.Int32(in_row) * fx.Int32(kv_in_row_stride) + + col_off_fx + lid_x_vec ) base_sc = ( - ArithValue(in_row) * ArithValue(score_in_row_stride) - + ArithValue(col_off) - + lid_x_vec - ) - base_ape = ( - ArithValue(ape_row) * arith.constant(DIM_FULL, type=i32) - + ArithValue(col_off) + fx.Int32(in_row) * fx.Int32(score_in_row_stride) + + col_off_fx + lid_x_vec ) + base_ape = fx.Int32(ape_row) * DIM_FULL + col_off_fx + lid_x_vec kv = _load_bf16_vec_then_f32(kv_in_rsrc, base_in) sc = _load_bf16_vec_then_f32(score_in_rsrc, base_sc) ape_v = _load_f32_vec(ape_rsrc, base_ape) @@ -1686,8 +1587,8 @@ def _phase2_loads(k_i32): return kv, score # ---- this wave's K range [wid*KPW, (wid+1)*KPW), split at window_len ---- - k_start = ArithValue(wid) * c_K_per_wave - k_end = k_start + c_K_per_wave + k_start = fx.Int32(wid) * K_PER_WAVE + k_end = k_start + K_PER_WAVE wl = _to_raw(window_len) split_lo = arith.maxsi(wl, _to_raw(k_start)) split = arith.minsi(split_lo, _to_raw(k_end)) @@ -1705,7 +1606,7 @@ def _phase2_loads(k_i32): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = _to_raw(fx.Int32(k_static)) kv_v, sc_v = _phase1_loads(k_i32) nm, nkv, nw = _softmax_step(m_lane, kv_lane, w_lane, sc_v, kv_v) p1 = yield list(nm) + list(nkv) + list(nw) @@ -1716,7 +1617,7 @@ def _phase2_loads(k_i32): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = _to_raw(fx.Int32(k_static)) kv_v, score = _phase2_loads(k_i32) nm, nkv, nw = _softmax_step(m_lane, kv_lane, w_lane, score, kv_v) final = yield list(nm) + list(nkv) + list(nw) @@ -1730,35 +1631,33 @@ def _phase2_loads(k_i32): lds_m_ptr = lds.lds_m.ptr lds_kv_ptr = lds.lds_kv.ptr lds_w_ptr = lds.lds_w.ptr - lds_thread_base = ArithValue(wid) * c_D + lid_x_vec + lds_thread_base = fx.Int32(wid) * D + lid_x_vec for i in range_constexpr(VEC): - idx_i = lds_thread_base + arith.constant(i, type=i32) - fx.ptr_store(m_local[i], lds_m_ptr + fx.Int32(idx_i)) - fx.ptr_store(kv_local[i], lds_kv_ptr + fx.Int32(idx_i)) - fx.ptr_store(w_local[i], lds_w_ptr + fx.Int32(idx_i)) + idx_i = lds_thread_base + i + fx.ptr_store(m_local[i], lds_m_ptr + idx_i) + fx.ptr_store(kv_local[i], lds_kv_ptr + idx_i) + fx.ptr_store(w_local[i], lds_w_ptr + idx_i) gpu.barrier() # ---- wave 0: cross-wave reduce + norm + rope + scatter ---- - is_wave0 = arith.cmpi(CmpIPredicate.eq, wid, c_zero_i32) - _if_w0 = scf.IfOp(is_wave0) - with _if_then(_if_w0): + if fx.Int32(wid) == 0: comp_lane = [] for i in range_constexpr(VEC): - lane_off = lid_x_vec + arith.constant(i, type=i32) + lane_off = lid_x_vec + i m_g = fx.Float32(c_neg_inf) m_arr = [] for w in range_constexpr(NW): - idx_w = arith.constant(w * D, type=i32) + lane_off - m_w = fx.ptr_load(lds_m_ptr + fx.Int32(idx_w)) + idx_w = lane_off + (w * D) + m_w = fx.ptr_load(lds_m_ptr + idx_w) m_arr.append(m_w) m_g = m_g.maximumf(m_w) kv_sum = fx.Float32(0.0) w_sum = fx.Float32(0.0) for w in range_constexpr(NW): - idx_w = arith.constant(w * D, type=i32) + lane_off - kv_w = fx.ptr_load(lds_kv_ptr + fx.Int32(idx_w)) - w_w = fx.ptr_load(lds_w_ptr + fx.Int32(idx_w)) + idx_w = lane_off + (w * D) + kv_w = fx.ptr_load(lds_kv_ptr + idx_w) + w_w = fx.ptr_load(lds_w_ptr + idx_w) scale_w = fx.Float32(fexp_f32(_to_raw(m_arr[w] - m_g))) kv_sum = kv_sum + kv_w * scale_w w_sum = w_sum + w_w * scale_w @@ -1809,27 +1708,17 @@ def wave_reduce_add(x): ] # ---- GPT-J RoPE on RD tail ---- - comp_pos_i32 = arith.muli( - arith.divsi(_to_raw(position), arith.constant(ratio, type=i32)), - arith.constant(ratio, type=i32), - ) + # position >= 0 (sentinel-skip guard) -> unsigned div. + comp_pos_i32 = (fx.Uint32(position) // ratio) * ratio cos_rsrc = buffer_ops.create_buffer_resource(cos_cache, max_size=True) sin_rsrc = buffer_ops.create_buffer_resource(sin_cache, max_size=True) - c_half_rd = arith.constant(RD // 2, type=i32) - cos_row_base = ArithValue(comp_pos_i32) * c_half_rd + cos_row_base = fx.Int32(comp_pos_i32) * (RD // 2) - is_rope_t = arith.cmpi( - CmpIPredicate.sge, - _to_raw(lid), - arith.constant(ROPE_THREAD_LO, type=i32), - ) - rope_rel_raw = ArithValue(lid) - arith.constant( - ROPE_THREAD_LO, type=i32 - ) + is_rope_t = (fx.Int32(lid) >= ROPE_THREAD_LO).ir_value() + # raw maxsi: no fx signed-int-max form. + rope_rel_raw = _to_raw(fx.Int32(lid) - ROPE_THREAD_LO) rope_rel = arith.maxsi(rope_rel_raw, c_zero_i32) - cs_lo = ArithValue(rope_rel) * arith.constant( - PAIRS_PER_THREAD, type=i32 - ) + cs_lo = fx.Int32(rope_rel) * PAIRS_PER_THREAD if const_expr(PAIRS_PER_THREAD == 1): cos_b = buffer_ops.buffer_load( @@ -1841,34 +1730,28 @@ def wave_reduce_add(x): cos_vals = [arith.extf(f32, cos_b)] sin_vals = [arith.extf(f32, sin_b)] else: - cos_vec = buffer_ops.buffer_load( - cos_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, + cos_vec = fx.Vector( + buffer_ops.buffer_load( + cos_rsrc, + cos_row_base + cs_lo, + vec_width=PAIRS_PER_THREAD, + dtype=T.bf16, + ) ) - sin_vec = buffer_ops.buffer_load( - sin_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, + sin_vec = fx.Vector( + buffer_ops.buffer_load( + sin_rsrc, + cos_row_base + cs_lo, + vec_width=PAIRS_PER_THREAD, + dtype=T.bf16, + ) ) cos_vals = [ - arith.extf( - f32, - vector.extract( - cos_vec, static_position=[i], dynamic_position=[] - ), - ) + arith.extf(f32, _to_raw(cos_vec[i])) for i in range(PAIRS_PER_THREAD) ] sin_vals = [ - arith.extf( - f32, - vector.extract( - sin_vec, static_position=[i], dynamic_position=[] - ), - ) + arith.extf(f32, _to_raw(sin_vec[i])) for i in range(PAIRS_PER_THREAD) ] @@ -1897,13 +1780,16 @@ def wave_reduce_add(x): # ---- paged scatter (BF16 or FP8). Emitted in wave 0; ``lid`` # (0..63) is the single-wave ``tid`` equivalent. ---- - ci = arith.divsi(_to_raw(position), arith.constant(ratio, type=i32)) - block_in_seq = arith.divsi(ci, arith.constant(k_per_block, type=i32)) - slot_in_block = arith.remui(ci, arith.constant(k_per_block, type=i32)) + # position >= 0 (sentinel-skip guard) -> unsigned div/rem. + ci = fx.Uint32(position) // ratio + block_in_seq = _to_raw(ci // k_per_block) + slot_in_block = _to_raw(ci % k_per_block) + ci = _to_raw(ci) bt_rsrc = buffer_ops.create_buffer_resource(block_table, max_size=True) - bt_off = ArithValue(batch_id) * ArithValue( - block_table_seq_stride - ) + ArithValue(block_in_seq) + bt_off = _to_raw( + fx.Int32(batch_id) * fx.Int32(block_table_seq_stride) + + fx.Int32(block_in_seq) + ) physical_block = buffer_ops.buffer_load( bt_rsrc, bt_off, vec_width=1, dtype=i32 ) @@ -1912,12 +1798,13 @@ def wave_reduce_add(x): # The block term rides on the descriptor's base, not on # the 32-bit offset -- see `block_base_bytes_i64`. cache_off = ( - ArithValue(slot_in_block) * ArithValue(kv_cache_token_stride) + fx.Int32(slot_in_block) * fx.Int32(kv_cache_token_stride) + lid_x_vec ) out_vec_t = T.vec(VEC, T.bf16) - raw_vec = vector.from_elements(vecVf32, out_lane) - bf16_vec = raw_vec.truncf(out_vec_t) + bf16_vec = fx.Vector.from_elements( + out_lane, dtype=fx.Float32 + ).truncf(out_vec_t) out_rsrc = buffer_ops.create_buffer_resource( kv_cache, max_size=True, @@ -1927,30 +1814,27 @@ def wave_reduce_add(x): ) cache_off_dw = ArithValue(cache_off) >> c_one_i32 dwords = (VEC + 1) // 2 - bf16_as_i32 = vector.bitcast(T.vec(dwords, T.i32), bf16_vec) + bf16_as_i32 = fx.Vector(bf16_vec).bitcast(fx.Int32) if const_expr(dwords == 1): - scalar_i32 = vector.extract( - bf16_as_i32, static_position=[0], dynamic_position=[] + buffer_ops.buffer_store( + _to_raw(bf16_as_i32[0]), out_rsrc, cache_off_dw ) - buffer_ops.buffer_store(scalar_i32, out_rsrc, cache_off_dw) elif const_expr(dwords <= 4): - buffer_ops.buffer_store(bf16_as_i32, out_rsrc, cache_off_dw) + buffer_ops.buffer_store( + _to_raw(bf16_as_i32), out_rsrc, cache_off_dw + ) else: - # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4 + # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4. c4_i32 = arith.constant(4, type=i32) - lo = vector.extract_strided_slice( - T.vec(4, T.i32), - bf16_as_i32, - offsets=[0], - sizes=[4], - strides=[1], + lo = _to_raw( + fx.Vector.from_elements( + [bf16_as_i32[i] for i in range(4)], dtype=fx.Int32 + ) ) - hi = vector.extract_strided_slice( - T.vec(4, T.i32), - bf16_as_i32, - offsets=[4], - sizes=[4], - strides=[1], + hi = _to_raw( + fx.Vector.from_elements( + [bf16_as_i32[i] for i in range(4, 8)], dtype=fx.Int32 + ) ) buffer_ops.buffer_store(lo, out_rsrc, cache_off_dw) buffer_ops.buffer_store( @@ -2010,7 +1894,7 @@ def wave_reduce_max(x): v = arith.MulFOp( out_lane[i], inv_scale, fastmath=fm_fast ).result - v = arith.minimumf(arith.maximumf(v, c_neg_fp8_max), c_fp8_max) + v = _to_raw(fx.min(fx.max(v, c_neg_fp8_max), c_fp8_max)) is_tn = arith.andi( arith.cmpf(CmpFPredicate.OLT, v, c_zero), arith.cmpf(CmpFPredicate.OGT, v, c_neg_uf), @@ -2071,37 +1955,30 @@ def wave_reduce_max(x): ) if const_expr(preshuffle): - c_TILE = arith.constant(_PRESHUFFLE_TILE, type=i32) - c_TILE_D = arith.constant(_PRESHUFFLE_TILE * D, type=i32) - c_TILE_TILE = arith.constant( - _PRESHUFFLE_TILE * _PRESHUFFLE_TILE, type=i32 - ) - token_tile_id = arith.divsi(slot_in_block, c_TILE) - token_in_tile = arith.remui(slot_in_block, c_TILE) - d_for_tid = ArithValue(lid) * arith.constant(VEC, type=i32) - col_tile_id = arith.divsi(d_for_tid, c_TILE) - col_in_tile = arith.remui(d_for_tid, c_TILE) - in_block_off = ( - ArithValue(token_tile_id) * c_TILE_D - + ArithValue(col_tile_id) * c_TILE_TILE - + ArithValue(token_in_tile) * c_TILE - + ArithValue(col_in_tile) + # slot_in_block, d are non-negative -> unsigned div/rem. + TILE = _PRESHUFFLE_TILE + slot_u = fx.Uint32(slot_in_block) + token_tile_id = slot_u // TILE + token_in_tile = slot_u % TILE + d_for_tid = fx.Uint32(lid) * VEC + col_tile_id = d_for_tid // TILE + col_in_tile = d_for_tid % TILE + in_block_off = _to_raw( + token_tile_id * (TILE * D) + + col_tile_id * (TILE * TILE) + + token_in_tile * TILE + + col_in_tile ) else: - in_block_off = ArithValue(slot_in_block) * arith.constant( - D, type=i32 - ) + ArithValue(lid) * arith.constant(VEC, type=i32) + in_block_off = _to_raw( + fx.Int32(slot_in_block) * D + fx.Int32(lid) * VEC + ) byte_off = in_block_off if const_expr(VEC == 2): - is_even = arith.cmpi( - CmpIPredicate.eq, - arith.andi(_to_raw(lid), arith.constant(1, type=i32)), - arith.constant(0, type=i32), - ) - _if_even = scf.IfOp(is_even) - with _if_then(_if_even): + # Only even lid stores (its dword covers peer's bytes too). + if (fx.Int32(lid) & 1) == 0: buffer_ops.buffer_store( dword, out_rsrc, byte_off, offset_is_bytes=True ) @@ -2113,11 +1990,11 @@ def wave_reduce_max(x): # VEC in {8, 16}: store n_dwords via dwordx4 chunks n_dw = VEC // 4 if const_expr(n_dw <= 4): - store_vec = vector.from_elements( - T.vec(n_dw, i32), list(dword) + store_vec = fx.Vector.from_elements( + list(dword), dtype=fx.Int32 ) buffer_ops.buffer_store( - store_vec, + _to_raw(store_vec), out_rsrc, byte_off, offset_is_bytes=True, @@ -2125,12 +2002,11 @@ def wave_reduce_max(x): else: for chunk_start in range_constexpr(n_dw // 4): base = chunk_start * 4 - sv = vector.from_elements( - T.vec(4, i32), - list(dword[base : base + 4]), + sv = fx.Vector.from_elements( + list(dword[base : base + 4]), dtype=fx.Int32 ) buffer_ops.buffer_store( - sv, + _to_raw(sv), out_rsrc, ArithValue(byte_off) + arith.constant(base * 4, type=i32), @@ -2138,13 +2014,7 @@ def wave_reduce_max(x): ) # (f) lane-0 writes fp32 scale at cache_scale[phys, slot] - is_lane0 = arith.cmpi( - CmpIPredicate.eq, - _to_raw(lid), - arith.constant(0, type=i32), - ) - _if_l0 = scf.IfOp(is_lane0) - with _if_then(_if_l0): + if fx.Int32(lid) == 0: cs_rsrc = buffer_ops.create_buffer_resource( cache_scale, max_size=True, @@ -2152,9 +2022,7 @@ def wave_reduce_max(x): physical_block, cache_scale_block_stride, 4 ), ) - buffer_ops.buffer_store( - scale_v, cs_rsrc, ArithValue(slot_in_block) - ) + buffer_ops.buffer_store(scale_v, cs_rsrc, slot_in_block) @flyc.jit def launch_fused_compress_attn_ksplit( diff --git a/aiter/ops/flydsl/kernels/fused_compress_attn_hca.py b/aiter/ops/flydsl/kernels/fused_compress_attn_hca.py index dfc84e3d1b..d4bc4bb41b 100644 --- a/aiter/ops/flydsl/kernels/fused_compress_attn_hca.py +++ b/aiter/ops/flydsl/kernels/fused_compress_attn_hca.py @@ -48,17 +48,14 @@ import torch from flydsl.expr import arith, const_expr, gpu, range_constexpr from flydsl.expr import math as fmath -from flydsl.expr.arith import ArithValue, CmpFPredicate, CmpIPredicate from flydsl.expr.typing import Int32, Stream, T -from aiter.ops.flydsl.kernels import buffer_ops, vector - from .fused_compress_attn_common import ( block_base_bytes_i64, emit_group_fp8_nm_asm_scatter, state_slot_byte_offset, ) -from .tensor_shim import _run_compiled, _to_raw +from .tensor_shim import _run_compiled, _to_raw, ptr_buf_tensor BLOCK_THREADS = 64 # 1 wave64 SLICE = 64 # head_dim elements per block (grid-Y split) @@ -66,6 +63,21 @@ _LOG2E = math.log2(math.e) +def _ptr_at_byte_off(tensor, base_i64): + """Global byte pointer at ``tensor``'s base + ``base_i64`` (64-bit byte offset). + + Used to fold a slot/block rebase into the pointer handed to + ``ptr_buf_tensor``, which re-derives the descriptor's element type -- so the + i8 carrier type here only carries the address and the ptrtoint/inttoptr + roundtrip folds away pre-ISA, keeping the emitted V# identical to the old + ``buf_tensor(base_i64=...)`` descriptor. + """ + pt = fx.PointerType.get(T.i8, address_space=fx.AddressSpace.Global, alignment=1) + return fx.inttoptr( + pt, fx.Int64(fx.ptrtoint(fx.get_iter(tensor))) + fx.Int64(base_i64) + ) + + # ============================================================================ # Kernel A: compress_forward with multi-wave LDS K-split # ============================================================================ @@ -149,7 +161,6 @@ class SharedStorage: _kname = ( f"hca_compress_forward_D{D}_R{ratio}_NW{NW}_SL{SLICE_SZ}_S{state_size}_flydsl" ) - fm_fast = arith.FastMathFlags.fast @flyc.kernel(name=_kname, known_block_size=[BLOCK_TH, 1, 1]) def kernel( @@ -176,193 +187,163 @@ def kernel( sid = fx.block_idx.y tid = fx.thread_idx.x # 0..BLOCK_TH-1 - c_zero_i32 = arith.constant(0, type=i32) - c_one_i32 = arith.constant(1, type=i32) - c_64 = arith.constant(64, type=i32) + # -inf sentinel + its 0.0 partner feed the softmax maximumf / == compare, + # which must stay non-fast (a plain fx const would take ambient fast_fp_math + # and corrupt the -inf guard) -- keep both raw. c_neg_inf = arith.constant(_NEG_INF, type=f32) c_zero_f32 = arith.constant(0.0, type=f32) - c_log2e = arith.constant(_LOG2E, type=f32) - c_K_m1 = arith.constant(K - 1, type=i32) - c_K_per_wave = arith.constant(K_PER_WAVE, type=i32) - c_ratio = arith.constant(ratio, type=i32) - c_DIM_FULL = arith.constant(DIM_FULL, type=i32) - c_SLICE = arith.constant(SLICE_SZ, type=i32) - c_VEC = arith.constant(VEC, type=i32) - c_state_size = arith.constant(state_size, type=i32) + c_log2e = fx.Float32(_LOG2E) def fexp_f32(x): - return fx.rocdl.exp2(f32, x * c_log2e) + # x is fx.Float32; exp2 needs a raw operand -> wrap once here. + return fx.rocdl.exp2(f32, _to_raw(x * c_log2e)) # Per-thread wave / lane (block-local). - wid = arith.divsi(_to_raw(tid), c_64) # ? [0, NW) - lid = arith.remui(_to_raw(tid), c_64) # ? [0, 64) + wid = fx.Int32(tid) // 64 # -> [0, NW) + lid = fx.Int32(tid) % 64 # -> [0, 64) # -- Load plan row ---------------------------------------------- - plan_rsrc = buffer_ops.create_buffer_resource(plan, max_size=True) - plan_base = ArithValue(pid) * arith.constant(4, type=i32) - plan_vec = buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) - ragged_id = vector.extract(plan_vec, static_position=[0], dynamic_position=[]) - batch_id = vector.extract(plan_vec, static_position=[1], dynamic_position=[]) - position = vector.extract(plan_vec, static_position=[2], dynamic_position=[]) - window_len = vector.extract(plan_vec, static_position=[3], dynamic_position=[]) + plan_buf = ptr_buf_tensor(fx.get_iter(plan), fx.Int32) + plan_vec = fx.Vector( + fx.add_offset(fx.get_iter(plan_buf), fx.Int32(pid) * 4).load(T.vec(4, i32)) + ) + ragged_id = plan_vec[0] + batch_id = plan_vec[1] + position = plan_vec[2] + window_len = plan_vec[3] # Sentinel-skip: run the whole body only for position >= 0, as a closure # under a runtime `if` (rewriter sees an opaque call -> scf.if). def _body(): # Per-thread head_dim base: each thread owns VEC contiguous # elements starting at slice_base + lid * VEC. - slice_base_i32 = ArithValue(sid) * c_SLICE - col_off_base = slice_base_i32 + ArithValue(lid) * c_VEC - - slot_map_rsrc = buffer_ops.create_buffer_resource( - state_slot_mapping, max_size=True - ) - slot = buffer_ops.buffer_load( - slot_map_rsrc, batch_id, vec_width=1, dtype=i32 - ) - - kv_in_rsrc = buffer_ops.create_buffer_resource(kv_in, max_size=True) - score_in_rsrc = buffer_ops.create_buffer_resource(score_in, max_size=True) - # Rebased onto this program's slot — see `state_slot_byte_offset`. - kv_state_rsrc = buffer_ops.create_buffer_resource( - kv_state, - max_size=True, - base_byte_offset=state_slot_byte_offset(slot, kv_state_slot_stride), + col_off_base = fx.Int32(sid) * SLICE_SZ + fx.Int32(lid) * VEC + + slot_map_buf = ptr_buf_tensor(fx.get_iter(state_slot_mapping), fx.Int32) + slot = fx.add_offset(fx.get_iter(slot_map_buf), batch_id).load(i32) + + # bf16 inputs are read as i32 dwords (unaligned bit-extract) -> i32 buf. + kv_in_buf = ptr_buf_tensor(fx.get_iter(kv_in), fx.Int32) + score_in_buf = ptr_buf_tensor(fx.get_iter(score_in), fx.Int32) + # Rebased onto this program's slot — the 64-bit slot byte offset is + # folded into the descriptor base (see `state_slot_byte_offset`). + kv_state_buf = ptr_buf_tensor( + _ptr_at_byte_off( + kv_state, state_slot_byte_offset(slot, kv_state_slot_stride) + ), + fx.Float32, ) - score_state_rsrc = buffer_ops.create_buffer_resource( - score_state, - max_size=True, - base_byte_offset=state_slot_byte_offset(slot, score_state_slot_stride), + score_state_buf = ptr_buf_tensor( + _ptr_at_byte_off( + score_state, state_slot_byte_offset(slot, score_state_slot_stride) + ), + fx.Float32, ) - ape_rsrc = buffer_ops.create_buffer_resource(ape, max_size=True) + ape_buf = ptr_buf_tensor(fx.get_iter(ape), fx.Float32) - def _load_bf16_vec_to_f32(rsrc, base_off_elems_i32): + def _load_bf16_vec_to_f32(buf, base_off_elems_i32): """Load VEC contiguous bf16 elements starting at ``base_off_elems_i32`` -> list of VEC f32 values. VEC=1: unaligned-safe scalar via dword + bit-extract. - VEC>=2: vectorized i32 buffer_load + bitcast to bf16. + VEC>=2: vectorized i32 load + bitcast to bf16. """ + base_off = fx.Int32(base_off_elems_i32) if const_expr(VEC == 1): - off_dw = ArithValue(base_off_elems_i32) >> c_one_i32 - lane_in_dw = arith.andi(_to_raw(base_off_elems_i32), c_one_i32) - raw_s = buffer_ops.buffer_load(rsrc, off_dw, vec_width=1, dtype=i32) - hi = ArithValue(raw_s) >> arith.constant(16, type=i32) - lo_or_hi = arith.select( - arith.cmpi(CmpIPredicate.eq, lane_in_dw, c_zero_i32), - raw_s, - _to_raw(hi), - ) - lo16 = arith.andi(lo_or_hi, arith.constant(0xFFFF, type=i32)) - lo16_v = vector.from_elements(T.vec(1, T.i32), [lo16]) - bf16_pair = vector.bitcast(T.vec(2, T.bf16), lo16_v) - bf16_v = vector.extract( - bf16_pair, static_position=[0], dynamic_position=[] - ) - return [arith.extf(f32, bf16_v)] + off_dw = base_off >> 1 + lane_in_dw = base_off & 1 + raw_s = fx.Int32(fx.add_offset(fx.get_iter(buf), off_dw).load(i32)) + # logical (unsigned) shift for the hi-word extract: fx Int32 >> + # is arithmetic -> use Uint32 to keep v_lshrrev_b32. + hi = fx.Int32((fx.Uint32(raw_s) >> 16).ir_value()) + lo_or_hi = (lane_in_dw == fx.Int32(0)).select(raw_s, hi) + lo16 = lo_or_hi & 0xFFFF + lo16_v = fx.Vector.from_elements([lo16], dtype=fx.Int32) + bf16_pair = lo16_v.bitcast(fx.BFloat16) + return [bf16_pair[0].to(fx.Float32)] else: # base must be VEC-aligned (caller guarantees by # col_off_base = sid*SLICE + lid*VEC, both multiples of VEC). - off_dw = ArithValue(base_off_elems_i32) >> c_one_i32 + off_dw = base_off >> 1 dwords = VEC // 2 # VEC bf16 = VEC*2 bytes if const_expr(dwords == 1): - # buffer_load(vec_width=1) returns scalar i32; wrap - # into vec<1xi32> before bitcast to vec<2xbf16>. - raw_s = buffer_ops.buffer_load( - rsrc, off_dw, vec_width=1, dtype=i32 + # vec_width=1 returns scalar i32; wrap into vec<1xi32> + # before bitcast to vec<2xbf16>. + raw = fx.Vector.from_elements( + [fx.add_offset(fx.get_iter(buf), off_dw).load(i32)], + dtype=fx.Int32, ) - raw = vector.from_elements(T.vec(1, T.i32), [raw_s]) else: - raw = buffer_ops.buffer_load( - rsrc, off_dw, vec_width=dwords, dtype=i32 - ) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - out = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, static_position=[i], dynamic_position=[] + raw = fx.Vector( + fx.add_offset(fx.get_iter(buf), off_dw).load( + T.vec(dwords, i32) + ) ) - out.append(arith.extf(f32, bf16_v)) - return out + vec_bf16 = raw.bitcast(fx.BFloat16) + return [vec_bf16[i].to(fx.Float32) for i in range_constexpr(VEC)] - def _load_f32_vec(rsrc, base_off_elems_i32): + def _load_f32_vec(buf, base_off_elems_i32): """Load VEC f32 starting at base -> list of VEC f32 values.""" if const_expr(VEC <= 4): - raw = buffer_ops.buffer_load( - rsrc, base_off_elems_i32, vec_width=VEC, dtype=f32 - ) if const_expr(VEC == 1): - # vec_width=1 returns scalar, not 1-vec. - return [raw] - return [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + # width=1 returns scalar, not 1-vec. + return [ + fx.Float32( + fx.add_offset( + fx.get_iter(buf), base_off_elems_i32 + ).load(f32) + ) + ] + raw = fx.Vector( + fx.add_offset(fx.get_iter(buf), base_off_elems_i32).load( + T.vec(VEC, f32) + ) + ) + return [raw[i] for i in range(VEC)] else: # VEC == 8: AMD HW max is dwordx4 -> 2 loads. assert VEC == 8 half = VEC // 2 - r0 = buffer_ops.buffer_load( - rsrc, base_off_elems_i32, vec_width=half, dtype=f32 + base = fx.Int32(base_off_elems_i32) + r0 = fx.Vector( + fx.add_offset(fx.get_iter(buf), base).load(T.vec(half, f32)) ) - r1 = buffer_ops.buffer_load( - rsrc, - ArithValue(base_off_elems_i32) + arith.constant(half, type=i32), - vec_width=half, - dtype=f32, - ) - out = [] - for i in range_constexpr(half): - out.append( - vector.extract(r0, static_position=[i], dynamic_position=[]) - ) - for i in range_constexpr(half): - out.append( - vector.extract(r1, static_position=[i], dynamic_position=[]) + r1 = fx.Vector( + fx.add_offset(fx.get_iter(buf), base + half).load( + T.vec(half, f32) ) - return out + ) + return [r0[i] for i in range(half)] + [r1[i] for i in range(half)] def _issue_phase2_loads(k_i32): """Phase 2 (ragged input) loads. Returns (kv_list, sc_list, ape_list) each of length VEC.""" - ape_row = arith.remui(k_i32, c_ratio) - tmp = arith.subi(c_K_m1, k_i32) - in_row_raw = arith.subi(_to_raw(ragged_id), tmp) - in_row = arith.maxsi(in_row_raw, c_zero_i32) - base_in_off = ( - ArithValue(in_row) * ArithValue(kv_in_row_stride) + col_off_base - ) - base_sc_off = ( - ArithValue(in_row) * ArithValue(score_in_row_stride) + col_off_base - ) - base_ape_off = ArithValue(ape_row) * c_DIM_FULL + col_off_base - kv = _load_bf16_vec_to_f32(kv_in_rsrc, base_in_off) - sc = _load_bf16_vec_to_f32(score_in_rsrc, base_sc_off) - ape_v = _load_f32_vec(ape_rsrc, base_ape_off) + k = fx.Int32(k_i32) + ape_row = k % ratio + in_row_raw = fx.Int32(ragged_id) - (fx.Int32(K - 1) - k) + in_row = (in_row_raw > fx.Int32(0)).select(in_row_raw, fx.Int32(0)) + base_in_off = in_row * fx.Int32(kv_in_row_stride) + col_off_base + base_sc_off = in_row * fx.Int32(score_in_row_stride) + col_off_base + base_ape_off = ape_row * DIM_FULL + col_off_base + kv = _load_bf16_vec_to_f32(kv_in_buf, base_in_off) + sc = _load_bf16_vec_to_f32(score_in_buf, base_sc_off) + ape_v = _load_f32_vec(ape_buf, base_ape_off) return kv, sc, ape_v def _issue_phase1_loads(k_i32): """Phase 1 (state cache) loads. Returns (kv_list, sc_padded_list) each of length VEC. Score is -inf when s < 0.""" - s = arith.addi( - arith.subi(_to_raw(position), c_K_m1), - k_i32, - ) - is_pad = arith.cmpi(CmpIPredicate.slt, s, c_zero_i32) - s_safe = arith.select(is_pad, c_zero_i32, s) - ring = arith.remui(s_safe, c_state_size) + s = fx.Int32(position) - fx.Int32(K - 1) + fx.Int32(k_i32) + is_pad = s < fx.Int32(0) + s_safe = is_pad.select(fx.Int32(0), s) + ring = s_safe % state_size # Slot term already folded into the descriptor base. - base_kv_off = ( - ArithValue(ring) * ArithValue(kv_state_pos_stride) + col_off_base - ) - base_sc_off = ( - ArithValue(ring) * ArithValue(score_state_pos_stride) + col_off_base - ) - kv_list = _load_f32_vec(kv_state_rsrc, base_kv_off) - sc_list = _load_f32_vec(score_state_rsrc, base_sc_off) - sc_padded = [ - arith.select(is_pad, c_neg_inf, sc_list[i]) for i in range(VEC) - ] + base_kv_off = ring * fx.Int32(kv_state_pos_stride) + col_off_base + base_sc_off = ring * fx.Int32(score_state_pos_stride) + col_off_base + kv_list = _load_f32_vec(kv_state_buf, base_kv_off) + sc_list = _load_f32_vec(score_state_buf, base_sc_off) + neg_inf = fx.Float32(c_neg_inf) + sc_padded = [is_pad.select(neg_inf, sc_list[i]) for i in range(VEC)] return kv_list, sc_padded def _softmax_step_padded( @@ -373,40 +354,29 @@ def _softmax_step_padded( is also -inf). Safe in both Phase 1 (padding can occur) and Phase 2 (score finite -> pad-select branch is dead code). """ + neg_inf = fx.Float32(c_neg_inf) + zero = fx.Float32(c_zero_f32) new_m, new_kv, new_w = [], [], [] for i in range_constexpr(VEC): - m_old = m_old_list[i] - kv_old = kv_old_list[i] - w_old = w_old_list[i] - score_k = score_k_list[i] - kv_k = kv_k_list[i] - m_new = arith.maximumf(m_old, score_k) - is_first = arith.cmpf(CmpFPredicate.OEQ, m_old, c_neg_inf) - scale_active = fexp_f32(arith.subf(m_old, m_new)) - scale_v = arith.select(is_first, c_zero_f32, scale_active) - wk_active = fexp_f32(arith.subf(score_k, m_new)) - is_pad_score = arith.cmpf(CmpFPredicate.OEQ, score_k, c_neg_inf) - w_k = arith.select(is_pad_score, c_zero_f32, wk_active) - new_kv.append( - arith.AddFOp( - arith.MulFOp(kv_old, scale_v, fastmath=fm_fast).result, - arith.MulFOp(w_k, kv_k, fastmath=fm_fast).result, - fastmath=fm_fast, - ).result - ) - new_w.append( - arith.AddFOp( - arith.MulFOp(w_old, scale_v, fastmath=fm_fast).result, - w_k, - fastmath=fm_fast, - ).result - ) - new_m.append(m_new) + m_old = fx.Float32(m_old_list[i]) + kv_old = fx.Float32(kv_old_list[i]) + w_old = fx.Float32(w_old_list[i]) + score_k = fx.Float32(score_k_list[i]) + kv_k = fx.Float32(kv_k_list[i]) + m_new = m_old.maximumf(score_k) + is_first = m_old == neg_inf + scale_active = fx.Float32(fexp_f32(m_old - m_new)) + scale_v = is_first.select(zero, scale_active) + wk_active = fx.Float32(fexp_f32(score_k - m_new)) + w_k = (score_k == neg_inf).select(zero, wk_active) + new_kv.append((kv_old * scale_v + w_k * kv_k).ir_value()) + new_w.append((w_old * scale_v + w_k).ir_value()) + new_m.append(m_new.ir_value()) return new_m, new_kv, new_w # -- Wave's K range: [wid * K_PER_WAVE, (wid+1) * K_PER_WAVE) -- - k_start_i32 = ArithValue(wid) * c_K_per_wave - k_end_i32 = k_start_i32 + c_K_per_wave + k_start_i32 = fx.Int32(wid) * K_PER_WAVE + k_end_i32 = k_start_i32 + K_PER_WAVE # Split point inside this wave's K range. Each wave sees a # window_len-dependent slice of Phase 1 followed by Phase 2. @@ -417,9 +387,9 @@ def _softmax_step_padded( # ``split`` = clamp(wl, k_start, k_end) gives the boundary; # both sub-loops are empty when their bound collapses, so any # of the three cases naturally falls out. - wl_i32 = _to_raw(window_len) - split_lo = arith.maxsi(wl_i32, _to_raw(k_start_i32)) - split_i32 = arith.minsi(split_lo, _to_raw(k_end_i32)) + wl = fx.Int32(window_len) + split_lo = (wl > k_start_i32).select(wl, k_start_i32) + split_i32 = (split_lo < k_end_i32).select(split_lo, k_end_i32) # State is 3*VEC scalars: m_lane[VEC] + kv_lane[VEC] + w_lane[VEC]. init_m = [c_neg_inf for _ in range(VEC)] @@ -431,12 +401,12 @@ def _softmax_step_padded( # cache; padded softmax (score can be -inf). phase1_local = init_state for k_static, state in range( - _to_raw(k_start_i32), _to_raw(split_i32), 1, init=init_state + k_start_i32.ir_value(), split_i32.ir_value(), 1, init=init_state ): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = fx.Int32(k_static) kv_v, sc_v = _issue_phase1_loads(k_i32) new_m, new_kv, new_w = _softmax_step_padded( m_lane, kv_lane, w_lane, sc_v, kv_v @@ -449,16 +419,15 @@ def _softmax_step_padded( # Carry Phase 1's accumulator through as init. final = phase1_local for k_static, state in range( - _to_raw(split_i32), _to_raw(k_end_i32), 1, init=phase1_local + split_i32.ir_value(), k_end_i32.ir_value(), 1, init=phase1_local ): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = fx.Int32(k_static) p2_kv, p2_sc, p2_ape = _issue_phase2_loads(k_i32) p2_score = [ - arith.AddFOp(p2_sc[i], p2_ape[i], fastmath=fm_fast).result - for i in range(VEC) + (fx.Float32(p2_sc[i]) + p2_ape[i]).ir_value() for i in range(VEC) ] new_m, new_kv, new_w = _softmax_step_padded( m_lane, kv_lane, w_lane, p2_score, p2_kv @@ -477,12 +446,12 @@ def _softmax_step_padded( lds_m_ptr = lds.lds_m.ptr lds_kv_ptr = lds.lds_kv.ptr lds_w_ptr = lds.lds_w.ptr - lds_thread_base = ArithValue(wid) * c_SLICE + ArithValue(lid) * c_VEC + lds_thread_base = fx.Int32(wid) * SLICE_SZ + fx.Int32(lid) * VEC for i in range_constexpr(VEC): - idx_i = lds_thread_base + arith.constant(i, type=i32) - fx.ptr_store(m_local[i], lds_m_ptr + fx.Int32(idx_i)) - fx.ptr_store(kv_local[i], lds_kv_ptr + fx.Int32(idx_i)) - fx.ptr_store(w_local[i], lds_w_ptr + fx.Int32(idx_i)) + idx_i = lds_thread_base + i + fx.ptr_store(m_local[i], lds_m_ptr + idx_i) + fx.ptr_store(kv_local[i], lds_kv_ptr + idx_i) + fx.ptr_store(w_local[i], lds_w_ptr + idx_i) gpu.barrier() @@ -494,13 +463,13 @@ def _softmax_step_padded( def _wave0(): comp_list = [] for i in range_constexpr(VEC): - lane_off = ArithValue(lid) * c_VEC + arith.constant(i, type=i32) + lane_off = fx.Int32(lid) * VEC + i # Global max across NW waves for this element. m_g = fx.Float32(c_neg_inf) m_arr = [] for w in range_constexpr(NW): - idx_w = arith.constant(w * SLICE_SZ, type=i32) + lane_off - m_w = fx.ptr_load(lds_m_ptr + fx.Int32(idx_w)) + idx_w = w * SLICE_SZ + lane_off + m_w = fx.ptr_load(lds_m_ptr + idx_w) m_arr.append(m_w) m_g = m_g.maximumf(m_w) @@ -508,41 +477,34 @@ def _wave0(): kv_sum = fx.Float32(0.0) w_sum = fx.Float32(0.0) for w in range_constexpr(NW): - idx_w = arith.constant(w * SLICE_SZ, type=i32) + lane_off - kv_w = fx.ptr_load(lds_kv_ptr + fx.Int32(idx_w)) - w_w = fx.ptr_load(lds_w_ptr + fx.Int32(idx_w)) + idx_w = w * SLICE_SZ + lane_off + kv_w = fx.ptr_load(lds_kv_ptr + idx_w) + w_w = fx.ptr_load(lds_w_ptr + idx_w) m_w = m_arr[w] - scale_w = fx.Float32(fexp_f32(_to_raw(m_w - m_g))) + scale_w = fx.Float32(fexp_f32(m_w - m_g)) kv_sum = kv_sum + kv_w * scale_w w_sum = w_sum + w_w * scale_w - rcp_w = fx.Float32(fx.rocdl.rcp(f32, _to_raw(w_sum))) - comp_list.append(_to_raw(kv_sum * rcp_w)) + rcp_w = fx.Float32(fx.rocdl.rcp(f32, w_sum.ir_value())) + comp_list.append(kv_sum * rcp_w) # -- Vectorized write of VEC f32 comp values -- - out_rsrc = buffer_ops.create_buffer_resource( - kv_compressed, max_size=True - ) + out_buf = ptr_buf_tensor(fx.get_iter(kv_compressed), fx.Float32) out_off = ( - ArithValue(pid) * ArithValue(kv_compressed_row_stride) - + col_off_base + fx.Int32(pid) * fx.Int32(kv_compressed_row_stride) + col_off_base ) if const_expr(VEC == 1): - buffer_ops.buffer_store(comp_list[0], out_rsrc, out_off) + fx.add_offset(fx.get_iter(out_buf), out_off).store(comp_list[0]) elif const_expr(VEC <= 4): - out_vec = vector.from_elements(T.vec(VEC, T.f32), comp_list) - buffer_ops.buffer_store(out_vec, out_rsrc, out_off) + out_vec = fx.Vector.from_elements(comp_list, dtype=fx.Float32) + fx.add_offset(fx.get_iter(out_buf), out_off).store(out_vec) else: # VEC == 8: AMD HW max is dwordx4 -> 2 stores. assert VEC == 8 half = VEC // 2 - v0 = vector.from_elements(T.vec(half, T.f32), comp_list[0:half]) - v1 = vector.from_elements(T.vec(half, T.f32), comp_list[half:VEC]) - buffer_ops.buffer_store(v0, out_rsrc, out_off) - buffer_ops.buffer_store( - v1, - out_rsrc, - ArithValue(out_off) + arith.constant(half, type=i32), - ) + v0 = fx.Vector.from_elements(comp_list[0:half], dtype=fx.Float32) + v1 = fx.Vector.from_elements(comp_list[half:VEC], dtype=fx.Float32) + fx.add_offset(fx.get_iter(out_buf), out_off).store(v0) + fx.add_offset(fx.get_iter(out_buf), out_off + half).store(v1) if wid == 0: _wave0() @@ -678,36 +640,32 @@ def kernel( ): f32 = T.f32 i32 = T.i32 - vecVf32 = T.vec(VEC, T.f32) bid = fx.block_idx.x tid = fx.thread_idx.x # Pack KW plan rows per block: wave_id picks the row, lane indexes head_dim. - wave_id = ArithValue(tid) >> arith.constant(log2_wave, type=T.i32) - lane = ArithValue(tid) & arith.constant(WAVE - 1, type=T.i32) - pid = ArithValue(bid) * arith.constant(KW, type=T.i32) + wave_id + # logical shift (tid >= 0); fx Int32 >> is arithmetic -> use Uint32. + wave_id = fx.Int32((fx.Uint32(tid) >> log2_wave).ir_value()) + lane = fx.Int32(tid) & (WAVE - 1) + pid = fx.Int32(bid) * KW + wave_id - c_zero_i32 = arith.constant(0, type=i32) - c_one_i32 = arith.constant(1, type=i32) - c_eps = arith.constant(rms_eps, type=f32) - c_inv_D = arith.constant(1.0 / D, type=f32) - c_ratio = arith.constant(ratio, type=i32) - c_k_per_block = arith.constant(k_per_block, type=i32) + c_eps = fx.Float32(rms_eps) + c_inv_D = fx.Float32(1.0 / D) def wave_reduce_add(x): - w = _to_raw(x) + w = fx.Float32(x) for sh_exp in range_constexpr(log2_wave): off = WAVE // (2 << sh_exp) - peer = _to_raw(ArithValue(w).shuffle_xor(off, WAVE)) - w = arith.AddFOp(w, peer, fastmath=fm_fast).result + w = w + w.shuffle_xor(off, WAVE) return w # -- Load plan row -- - plan_rsrc = buffer_ops.create_buffer_resource(plan, max_size=True) - plan_base = ArithValue(pid) * arith.constant(4, type=i32) - plan_vec = buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) - batch_id = vector.extract(plan_vec, static_position=[1], dynamic_position=[]) - position = vector.extract(plan_vec, static_position=[2], dynamic_position=[]) + plan_buf = ptr_buf_tensor(fx.get_iter(plan), fx.Int32) + plan_vec = fx.Vector( + fx.add_offset(fx.get_iter(plan_buf), pid * 4).load(T.vec(4, i32)) + ) + batch_id = plan_vec[1] + position = plan_vec[2] # active = real plan row (position>=0 sentinel) AND within capacity (tail # waves of the last block have pid>=cap and must bail; their plan load is @@ -716,170 +674,117 @@ def wave_reduce_add(x): # Whole body as a closure under the runtime guard (opaque call -> scf.if). def _body(): - tid_x_vec = ArithValue(lane) * arith.constant(VEC, type=i32) + tid_x_vec = lane * VEC # -- Load kv_compressed[pid, tid*VEC : tid*VEC + VEC] -- - kvc_rsrc = buffer_ops.create_buffer_resource(kv_compressed, max_size=True) - base_off = ( - ArithValue(pid) * ArithValue(kv_compressed_row_stride) + tid_x_vec - ) + kvc_buf = ptr_buf_tensor(fx.get_iter(kv_compressed), fx.Float32) + base_off = pid * fx.Int32(kv_compressed_row_stride) + tid_x_vec # VEC ? {2, 4, 8}: VEC <= 4 -> single dwordx{VEC}; VEC=8 -> 2x dwordx4. if const_expr(VEC <= 4): - raw = buffer_ops.buffer_load( - kvc_rsrc, base_off, vec_width=VEC, dtype=f32 + raw = fx.Vector( + fx.add_offset(fx.get_iter(kvc_buf), base_off).load(T.vec(VEC, f32)) ) - comp_lane = [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + comp_lane = [raw[i] for i in range(VEC)] else: assert VEC == 8 half = 4 - r0 = buffer_ops.buffer_load( - kvc_rsrc, base_off, vec_width=half, dtype=f32 + r0 = fx.Vector( + fx.add_offset(fx.get_iter(kvc_buf), base_off).load(T.vec(half, f32)) ) - r1 = buffer_ops.buffer_load( - kvc_rsrc, - ArithValue(base_off) + arith.constant(half, type=i32), - vec_width=half, - dtype=f32, + r1 = fx.Vector( + fx.add_offset(fx.get_iter(kvc_buf), base_off + half).load( + T.vec(half, f32) + ) ) - comp_lane = [ - vector.extract(r0, static_position=[i], dynamic_position=[]) - for i in range(half) - ] + [ - vector.extract(r1, static_position=[i], dynamic_position=[]) - for i in range(half) - ] + comp_lane = [r0[i] for i in range(half)] + [r1[i] for i in range(half)] # -- RMSNorm (wave reduce-add of squares / D + eps; rsqrt) -- - sq_local = arith.constant(0.0, type=f32) + sq_local = fx.Float32(0.0) for i in range_constexpr(VEC): - sq_local = arith.AddFOp( - sq_local, - arith.MulFOp(comp_lane[i], comp_lane[i], fastmath=fm_fast).result, - fastmath=fm_fast, - ).result + sq_local = sq_local + comp_lane[i] * comp_lane[i] sq_full = wave_reduce_add(sq_local) - var = arith.MulFOp(sq_full, c_inv_D, fastmath=fm_fast).result - rrms = fmath.rsqrt( - arith.AddFOp(var, c_eps, fastmath=fm_fast).result, fastmath=fm_fast - ) + var = sq_full * c_inv_D + rrms = fmath.rsqrt((var + c_eps).ir_value(), fastmath=fm_fast) # rms_weight load - rmsw_rsrc = buffer_ops.create_buffer_resource(rms_weight, max_size=True) if const_expr(rms_weight_is_bf16): + # bf16 weights read as i32 dwords -> i32 buf. + rmsw_buf = ptr_buf_tensor(fx.get_iter(rms_weight), fx.Int32) dwords = (VEC + 1) // 2 - off_dw = ArithValue(tid_x_vec) >> c_one_i32 + off_dw = fx.Int32((fx.Uint32(tid_x_vec) >> 1).ir_value()) if const_expr(dwords == 1): - raw_s = buffer_ops.buffer_load( - rmsw_rsrc, off_dw, vec_width=1, dtype=i32 + raw = fx.Vector.from_elements( + [fx.add_offset(fx.get_iter(rmsw_buf), off_dw).load(i32)], + dtype=fx.Int32, ) - raw = vector.from_elements(T.vec(1, T.i32), [raw_s]) else: - raw = buffer_ops.buffer_load( - rmsw_rsrc, off_dw, vec_width=dwords, dtype=i32 - ) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - rmsw_lane = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, static_position=[i], dynamic_position=[] + raw = fx.Vector( + fx.add_offset(fx.get_iter(rmsw_buf), off_dw).load( + T.vec(dwords, i32) + ) ) - rmsw_lane.append(arith.extf(f32, bf16_v)) + vec_bf16 = raw.bitcast(fx.BFloat16) + rmsw_lane = [vec_bf16[i].to(fx.Float32) for i in range_constexpr(VEC)] else: + rmsw_buf = ptr_buf_tensor(fx.get_iter(rms_weight), fx.Float32) if const_expr(VEC <= 4): - raw = buffer_ops.buffer_load( - rmsw_rsrc, tid_x_vec, vec_width=VEC, dtype=f32 + raw = fx.Vector( + fx.add_offset(fx.get_iter(rmsw_buf), tid_x_vec).load( + T.vec(VEC, f32) + ) ) - rmsw_lane = [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + rmsw_lane = [raw[i] for i in range(VEC)] else: half = 4 - r0 = buffer_ops.buffer_load( - rmsw_rsrc, tid_x_vec, vec_width=half, dtype=f32 + r0 = fx.Vector( + fx.add_offset(fx.get_iter(rmsw_buf), tid_x_vec).load( + T.vec(half, f32) + ) ) - r1 = buffer_ops.buffer_load( - rmsw_rsrc, - ArithValue(tid_x_vec) + arith.constant(half, type=i32), - vec_width=half, - dtype=f32, + r1 = fx.Vector( + fx.add_offset(fx.get_iter(rmsw_buf), tid_x_vec + half).load( + T.vec(half, f32) + ) ) - rmsw_lane = [ - vector.extract(r0, static_position=[i], dynamic_position=[]) - for i in range(half) - ] + [ - vector.extract(r1, static_position=[i], dynamic_position=[]) - for i in range(half) + rmsw_lane = [r0[i] for i in range(half)] + [ + r1[i] for i in range(half) ] - normed_lane = [ - arith.MulFOp( - arith.MulFOp(comp_lane[i], rrms, fastmath=fm_fast).result, - rmsw_lane[i], - fastmath=fm_fast, - ).result - for i in range(VEC) - ] + normed_lane = [comp_lane[i] * rrms * rmsw_lane[i] for i in range(VEC)] # -- GPT-J RoPE on RD tail -- - comp_pos_i32 = arith.muli(arith.divsi(_to_raw(position), c_ratio), c_ratio) - cos_rsrc = buffer_ops.create_buffer_resource(cos_cache, max_size=True) - sin_rsrc = buffer_ops.create_buffer_resource(sin_cache, max_size=True) - c_half_rd = arith.constant(RD // 2, type=i32) - cos_row_base = ArithValue(comp_pos_i32) * c_half_rd - - is_rope_t = arith.cmpi( - CmpIPredicate.sge, - _to_raw(lane), - arith.constant(ROPE_THREAD_LO, type=i32), - ) - rope_rel_raw = ArithValue(lane) - arith.constant(ROPE_THREAD_LO, type=i32) - rope_rel = arith.maxsi(rope_rel_raw, c_zero_i32) - cs_lo = ArithValue(rope_rel) * arith.constant(PAIRS_PER_THREAD, type=i32) + comp_pos_i32 = (fx.Int32(position) // ratio) * ratio + cos_buf = ptr_buf_tensor(fx.get_iter(cos_cache), fx.BFloat16) + sin_buf = ptr_buf_tensor(fx.get_iter(sin_cache), fx.BFloat16) + cos_row_base = comp_pos_i32 * (RD // 2) + + is_rope_t = lane >= fx.Int32(ROPE_THREAD_LO) + rope_rel_raw = lane - ROPE_THREAD_LO + rope_rel = (rope_rel_raw > fx.Int32(0)).select(rope_rel_raw, fx.Int32(0)) + cs_lo = rope_rel * PAIRS_PER_THREAD if const_expr(PAIRS_PER_THREAD == 1): - cos_b = buffer_ops.buffer_load( - cos_rsrc, cos_row_base + cs_lo, vec_width=1, dtype=T.bf16 + cos_b = fx.add_offset(fx.get_iter(cos_buf), cos_row_base + cs_lo).load( + T.bf16 ) - sin_b = buffer_ops.buffer_load( - sin_rsrc, cos_row_base + cs_lo, vec_width=1, dtype=T.bf16 + sin_b = fx.add_offset(fx.get_iter(sin_buf), cos_row_base + cs_lo).load( + T.bf16 ) - cos_vals = [arith.extf(f32, cos_b)] - sin_vals = [arith.extf(f32, sin_b)] + cos_vals = [fx.BFloat16(cos_b).to(fx.Float32)] + sin_vals = [fx.BFloat16(sin_b).to(fx.Float32)] else: - cos_vec = buffer_ops.buffer_load( - cos_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, - ) - sin_vec = buffer_ops.buffer_load( - sin_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, - ) - cos_vals = [ - arith.extf( - f32, - vector.extract( - cos_vec, static_position=[i], dynamic_position=[] - ), + cos_vec = fx.Vector( + fx.add_offset(fx.get_iter(cos_buf), cos_row_base + cs_lo).load( + T.vec(PAIRS_PER_THREAD, T.bf16) ) - for i in range(PAIRS_PER_THREAD) - ] - sin_vals = [ - arith.extf( - f32, - vector.extract( - sin_vec, static_position=[i], dynamic_position=[] - ), + ) + sin_vec = fx.Vector( + fx.add_offset(fx.get_iter(sin_buf), cos_row_base + cs_lo).load( + T.vec(PAIRS_PER_THREAD, T.bf16) ) - for i in range(PAIRS_PER_THREAD) - ] + ) + cos_vals = [cos_vec[i].to(fx.Float32) for i in range(PAIRS_PER_THREAD)] + sin_vals = [sin_vec[i].to(fx.Float32) for i in range(PAIRS_PER_THREAD)] rotated_lane = list(normed_lane) for k in range_constexpr(PAIRS_PER_THREAD): @@ -887,59 +792,47 @@ def _body(): o = normed_lane[2 * k + 1] c = cos_vals[k] s = sin_vals[k] - new_e = arith.subf( - arith.MulFOp(e, c, fastmath=fm_fast).result, - arith.MulFOp(o, s, fastmath=fm_fast).result, - ) - new_o = arith.AddFOp( - arith.MulFOp(e, s, fastmath=fm_fast).result, - arith.MulFOp(o, c, fastmath=fm_fast).result, - fastmath=fm_fast, - ).result - rotated_lane[2 * k] = new_e - rotated_lane[2 * k + 1] = new_o + rotated_lane[2 * k] = e * c - o * s + rotated_lane[2 * k + 1] = e * s + o * c # -- Paged scatter dest (shared by bf16 / fp8) -- - ci = arith.divsi(_to_raw(position), c_ratio) - block_in_seq = arith.divsi(ci, c_k_per_block) - slot_in_block = arith.remui(ci, c_k_per_block) - bt_rsrc = buffer_ops.create_buffer_resource(block_table, max_size=True) - bt_off = ArithValue(batch_id) * ArithValue( - block_table_seq_stride - ) + ArithValue(block_in_seq) - physical_block = buffer_ops.buffer_load( - bt_rsrc, bt_off, vec_width=1, dtype=i32 + ci = fx.Int32(position) // ratio + block_in_seq = ci // k_per_block + slot_in_block = ci % k_per_block + bt_buf = ptr_buf_tensor(fx.get_iter(block_table), fx.Int32) + bt_off = ( + fx.Int32(batch_id) * fx.Int32(block_table_seq_stride) + block_in_seq ) + physical_block = fx.add_offset(fx.get_iter(bt_buf), bt_off).load(i32) # The block term rides on the descriptor's base, not on the # 32-bit offset -- see `block_base_bytes_i64`. What is left is one # block's worth, which fits by construction. - out_rsrc = buffer_ops.create_buffer_resource( - kv_cache, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, kv_cache_block_stride, 1 if quant else 2 - ), + # The block term rides on the descriptor's base (i64 fold), so the + # 32-bit voffset only spans one block. `block_base_bytes_i64` elem_bytes: + # fp8 kv_cache=1, bf16 kv_cache=2. + cache_block_base = block_base_bytes_i64( + physical_block, kv_cache_block_stride, 1 if quant else 2 ) - cache_base = ArithValue(slot_in_block) * ArithValue(kv_cache_token_stride) + cache_base = slot_in_block * fx.Int32(kv_cache_token_stride) if const_expr(quant): # -- group_fp8 (V4 nm-asm) via shared emitter (single source of truth - # shared with the CSA single-kernel; fp8 entry layout stays identical). -- - _krope_base = ArithValue(slot_in_block) * ArithValue(krope_token_stride) + # shared with the CSA single-kernel; fp8 entry layout stays identical). + # The emitter stores through direct global pointers built from the + # i64 block base (kv_cache / k_rope) we pass below. -- + _krope_base = slot_in_block * fx.Int32(krope_token_stride) emit_group_fp8_nm_asm_scatter( - normed_lane=normed_lane, - rotated_lane=rotated_lane, - lane=lane, - is_rope_t=is_rope_t, - cache_base=_to_raw(cache_base), - out_rsrc=out_rsrc, - krope_base=_to_raw(_krope_base), - krope_rsrc=buffer_ops.create_buffer_resource( - k_rope_buff, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, krope_block_stride, 2 - ), + normed_lane=[v.ir_value() for v in normed_lane], + rotated_lane=[v.ir_value() for v in rotated_lane], + lane=lane.ir_value(), + is_rope_t=is_rope_t.ir_value(), + cache_base=cache_base.ir_value(), + out_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(kv_cache))) + + fx.Int64(cache_block_base), + krope_base=_krope_base.ir_value(), + krope_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(k_rope_buff))) + + fx.Int64( + block_base_bytes_i64(physical_block, krope_block_stride, 2) ), VEC=VEC, NOPE=NOPE, @@ -947,29 +840,31 @@ def _body(): log2_rts=log2_rts, ROPE_THREAD_LO=ROPE_THREAD_LO, wave_width=WAVE, - vecVf32=vecVf32, - fm_fast=fm_fast, ) else: # ---- BF16 single-buffer scatter (nope + rope contiguous) ---- + # bf16 kv_cache written as i32 dwords -> i32 buf, block base folded. + out_buf = ptr_buf_tensor( + _ptr_at_byte_off(kv_cache, cache_block_base), + fx.Int32, + ) out_lane = [ - arith.select(is_rope_t, rotated_lane[i], normed_lane[i]) + is_rope_t.select(rotated_lane[i], normed_lane[i]) for i in range_constexpr(VEC) ] - cache_off = ArithValue(cache_base) + tid_x_vec - out_vec_t = T.vec(VEC, T.bf16) - raw_vec = vector.from_elements(vecVf32, out_lane) - bf16_vec = raw_vec.truncf(out_vec_t) - cache_off_dw = ArithValue(cache_off) >> c_one_i32 + cache_off = cache_base + tid_x_vec + bf16_vec = fx.Vector.from_elements(out_lane, dtype=fx.Float32).to( + fx.BFloat16 + ) + bf16_as_i32 = bf16_vec.bitcast(fx.Int32) + cache_off_dw = fx.Int32((fx.Uint32(cache_off) >> 1).ir_value()) dwords = (VEC + 1) // 2 - bf16_as_i32 = vector.bitcast(T.vec(dwords, T.i32), bf16_vec) if const_expr(dwords == 1): - scalar_i32 = vector.extract( - bf16_as_i32, static_position=[0], dynamic_position=[] + fx.add_offset(fx.get_iter(out_buf), cache_off_dw).store( + bf16_as_i32[0] ) - buffer_ops.buffer_store(scalar_i32, out_rsrc, cache_off_dw) else: - buffer_ops.buffer_store(bf16_as_i32, out_rsrc, cache_off_dw) + fx.add_offset(fx.get_iter(out_buf), cache_off_dw).store(bf16_as_i32) if is_active: _body() @@ -994,11 +889,7 @@ def launch_hca_norm_rope_scatter( stream: fx.Stream, ): # grid = ceil(cap / KW): KW plan rows packed per block. - cap_raw = _to_raw(plan_capacity) - nblocks = arith.divsi( - arith.addi(cap_raw, arith.constant(KW - 1, type=T.i32)), - arith.constant(KW, type=T.i32), - ) + nblocks = (fx.Int32(plan_capacity) + (KW - 1)) // KW idx_p = fx.Int64(nblocks) k = kernel( kv_compressed, diff --git a/aiter/ops/flydsl/kernels/fused_compress_attn_hca_gfx1250.py b/aiter/ops/flydsl/kernels/fused_compress_attn_hca_gfx1250.py index a6ef196e6a..18b07f0ca0 100644 --- a/aiter/ops/flydsl/kernels/fused_compress_attn_hca_gfx1250.py +++ b/aiter/ops/flydsl/kernels/fused_compress_attn_hca_gfx1250.py @@ -15,27 +15,25 @@ """ import math -from contextlib import contextmanager from functools import lru_cache import flydsl.compiler as flyc import flydsl.expr as fx import torch -from flydsl._mlir import ir -from flydsl._mlir.dialects import llvm, scf +from flydsl._mlir.dialects import llvm from flydsl.expr import arith, const_expr, gpu, range_constexpr from flydsl.expr import math as fmath -from flydsl.expr.arith import ArithValue, CmpFPredicate, CmpIPredicate +from flydsl.expr.arith import CmpFPredicate, CmpIPredicate from flydsl.expr.typing import Int32, Stream, T -from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels import buffer_ops from .fused_compress_attn_common import ( block_base_bytes_i64, emit_group_fp8_nm_asm_scatter, state_slot_byte_offset, ) -from .tensor_shim import _run_compiled, _to_raw +from .tensor_shim import _run_compiled BLOCK_THREADS = 32 # 1 wave32 (RDNA4 / gfx1250) SLICE = 32 # head_dim elements per block (grid-Y split) @@ -43,17 +41,6 @@ _LOG2E = math.log2(math.e) -@contextmanager -def _if_then(if_op): - with ir.InsertionPoint(if_op.then_block): - try: - yield if_op.then_block - finally: - blk = if_op.then_block - if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): - scf.YieldOp([]) - - # ============================================================================ # Kernel A: compress_forward with multi-wave LDS K-split # ============================================================================ @@ -164,44 +151,36 @@ def kernel( tid = fx.thread_idx.x # 0..BLOCK_TH-1 c_zero_i32 = arith.constant(0, type=i32) - c_one_i32 = arith.constant(1, type=i32) - c_WS = arith.constant(BLOCK_THREADS, type=i32) c_neg_inf = arith.constant(_NEG_INF, type=f32) c_zero_f32 = arith.constant(0.0, type=f32) c_log2e = arith.constant(_LOG2E, type=f32) - c_K_m1 = arith.constant(K - 1, type=i32) - c_K_per_wave = arith.constant(K_PER_WAVE, type=i32) - c_ratio = arith.constant(ratio, type=i32) - c_DIM_FULL = arith.constant(DIM_FULL, type=i32) - c_SLICE = arith.constant(SLICE_SZ, type=i32) - c_VEC = arith.constant(VEC, type=i32) - c_state_size = arith.constant(state_size, type=i32) def fexp_f32(x): return llvm.call_intrinsic( f32, "llvm.amdgcn.exp2.f32", [x * c_log2e], [], [] ) - # Per-thread wave / lane (block-local). - wid = arith.divsi(_to_raw(tid), c_WS) # ? [0, NW) - lid = arith.remui(_to_raw(tid), c_WS) # ? [0, 32) + # Per-thread wave / lane (block-local); tid >= 0 -> unsigned div/rem + # (divui/remui). Wrap back to Int32 for the signed i32 consumers. + wid = fx.Int32((fx.Uint32(tid) // BLOCK_THREADS).ir_value()) # -> [0, NW) + lid = fx.Int32((fx.Uint32(tid) % BLOCK_THREADS).ir_value()) # -> [0, 32) # -- Load plan row ---------------------------------------------- plan_rsrc = buffer_ops.create_buffer_resource(plan, max_size=True) - plan_base = ArithValue(pid) * arith.constant(4, type=i32) - plan_vec = buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) - ragged_id = vector.extract(plan_vec, static_position=[0], dynamic_position=[]) - batch_id = vector.extract(plan_vec, static_position=[1], dynamic_position=[]) - position = vector.extract(plan_vec, static_position=[2], dynamic_position=[]) - window_len = vector.extract(plan_vec, static_position=[3], dynamic_position=[]) - - is_active = arith.cmpi(CmpIPredicate.sge, _to_raw(position), c_zero_i32) - _if_active = scf.IfOp(is_active) - with _if_then(_if_active): + plan_vec = fx.Vector( + buffer_ops.buffer_load(plan_rsrc, fx.Int32(pid) * 4, vec_width=4, dtype=i32) + ) + ragged_id = plan_vec[0] + batch_id = plan_vec[1] + position = plan_vec[2] + window_len = plan_vec[3] + + # Sentinel-skip: run the whole body only for position >= 0, as a closure + # under a runtime `if` (rewriter sees an opaque call -> scf.if). + def _body(): # Per-thread head_dim base: each thread owns VEC contiguous # elements starting at slice_base + lid * VEC. - slice_base_i32 = ArithValue(sid) * c_SLICE - col_off_base = slice_base_i32 + ArithValue(lid) * c_VEC + col_off_base = fx.Int32(sid) * SLICE_SZ + lid * VEC slot_map_rsrc = buffer_ops.create_buffer_resource( state_slot_mapping, max_size=True @@ -232,27 +211,28 @@ def _load_bf16_vec_to_f32(rsrc, base_off_elems_i32): VEC=1: unaligned-safe scalar via dword + bit-extract. VEC>=2: vectorized i32 buffer_load + bitcast to bf16. """ + base_off = fx.Int32(base_off_elems_i32) + # logical (unsigned) >> for the dword offset (base_off >= 0): fx + # Int32 >> is arithmetic -> use Uint32 to keep shrui/v_lshrrev_b32. + off_dw = fx.Int32((fx.Uint32(base_off_elems_i32) >> 1).ir_value()) if const_expr(VEC == 1): - off_dw = ArithValue(base_off_elems_i32) >> c_one_i32 - lane_in_dw = arith.andi(_to_raw(base_off_elems_i32), c_one_i32) + lane_in_dw = base_off & 1 raw_s = buffer_ops.buffer_load(rsrc, off_dw, vec_width=1, dtype=i32) - hi = ArithValue(raw_s) >> arith.constant(16, type=i32) + # logical shift for the hi-word extract too. + hi = fx.Int32((fx.Uint32(raw_s) >> 16).ir_value()) lo_or_hi = arith.select( - arith.cmpi(CmpIPredicate.eq, lane_in_dw, c_zero_i32), + arith.cmpi(CmpIPredicate.eq, lane_in_dw.ir_value(), c_zero_i32), raw_s, - _to_raw(hi), + hi.ir_value(), ) lo16 = arith.andi(lo_or_hi, arith.constant(0xFFFF, type=i32)) - lo16_v = vector.from_elements(T.vec(1, T.i32), [lo16]) - bf16_pair = vector.bitcast(T.vec(2, T.bf16), lo16_v) - bf16_v = vector.extract( - bf16_pair, static_position=[0], dynamic_position=[] - ) - return [arith.extf(f32, bf16_v)] + lo16_v = fx.Vector.from_elements([lo16], dtype=fx.Int32) + bf16_pair = lo16_v.bitcast(fx.BFloat16) + # raw f32 for the explicit-fastmath float layer downstream. + return [bf16_pair[0].to(fx.Float32).ir_value()] else: # base must be VEC-aligned (caller guarantees by # col_off_base = sid*SLICE + lid*VEC, both multiples of VEC). - off_dw = ArithValue(base_off_elems_i32) >> c_one_i32 dwords = VEC // 2 # VEC bf16 = VEC*2 bytes if const_expr(dwords == 1): # buffer_load(vec_width=1) returns scalar i32; wrap @@ -260,22 +240,19 @@ def _load_bf16_vec_to_f32(rsrc, base_off_elems_i32): raw_s = buffer_ops.buffer_load( rsrc, off_dw, vec_width=1, dtype=i32 ) - raw = vector.from_elements(T.vec(1, T.i32), [raw_s]) + raw = fx.Vector.from_elements([raw_s], dtype=fx.Int32) else: - raw = buffer_ops.buffer_load( - rsrc, off_dw, vec_width=dwords, dtype=i32 - ) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - out = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, static_position=[i], dynamic_position=[] + raw = fx.Vector( + buffer_ops.buffer_load( + rsrc, off_dw, vec_width=dwords, dtype=i32 + ) ) - out.append(arith.extf(f32, bf16_v)) - return out + vec_bf16 = raw.bitcast(fx.BFloat16) + # raw f32 for the explicit-fastmath float layer downstream. + return [vec_bf16[i].to(fx.Float32).ir_value() for i in range(VEC)] def _load_f32_vec(rsrc, base_off_elems_i32): - """Load VEC f32 starting at base -> list of VEC f32 values.""" + """Load VEC f32 (raw ir.Values) starting at base -> list of VEC.""" if const_expr(VEC <= 4): raw = buffer_ops.buffer_load( rsrc, base_off_elems_i32, vec_width=VEC, dtype=f32 @@ -283,48 +260,38 @@ def _load_f32_vec(rsrc, base_off_elems_i32): if const_expr(VEC == 1): # vec_width=1 returns scalar, not 1-vec. return [raw] - return [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + return [fx.Vector(raw)[i].ir_value() for i in range(VEC)] else: # VEC == 8: AMD HW max is dwordx4 -> 2 loads. assert VEC == 8 half = VEC // 2 - r0 = buffer_ops.buffer_load( - rsrc, base_off_elems_i32, vec_width=half, dtype=f32 - ) - r1 = buffer_ops.buffer_load( - rsrc, - ArithValue(base_off_elems_i32) + arith.constant(half, type=i32), - vec_width=half, - dtype=f32, - ) - out = [] - for i in range_constexpr(half): - out.append( - vector.extract(r0, static_position=[i], dynamic_position=[]) + r0 = fx.Vector( + buffer_ops.buffer_load( + rsrc, base_off_elems_i32, vec_width=half, dtype=f32 ) - for i in range_constexpr(half): - out.append( - vector.extract(r1, static_position=[i], dynamic_position=[]) + ) + r1 = fx.Vector( + buffer_ops.buffer_load( + rsrc, + fx.Int32(base_off_elems_i32) + half, + vec_width=half, + dtype=f32, ) - return out + ) + return [r0[i].ir_value() for i in range(half)] + [ + r1[i].ir_value() for i in range(half) + ] def _issue_phase2_loads(k_i32): """Phase 2 (ragged input) loads. Returns (kv_list, sc_list, ape_list) each of length VEC.""" - ape_row = arith.remui(k_i32, c_ratio) - tmp = arith.subi(c_K_m1, k_i32) - in_row_raw = arith.subi(_to_raw(ragged_id), tmp) - in_row = arith.maxsi(in_row_raw, c_zero_i32) - base_in_off = ( - ArithValue(in_row) * ArithValue(kv_in_row_stride) + col_off_base - ) - base_sc_off = ( - ArithValue(in_row) * ArithValue(score_in_row_stride) + col_off_base - ) - base_ape_off = ArithValue(ape_row) * c_DIM_FULL + col_off_base + k = fx.Int32(k_i32) + ape_row = fx.Int32((fx.Uint32(k_i32) % ratio).ir_value()) + in_row_raw = fx.Int32(ragged_id) - (fx.Int32(K - 1) - k) + in_row = fx.max(in_row_raw, fx.Int32(0)) + base_in_off = in_row * fx.Int32(kv_in_row_stride) + col_off_base + base_sc_off = in_row * fx.Int32(score_in_row_stride) + col_off_base + base_ape_off = ape_row * DIM_FULL + col_off_base kv = _load_bf16_vec_to_f32(kv_in_rsrc, base_in_off) sc = _load_bf16_vec_to_f32(score_in_rsrc, base_sc_off) ape_v = _load_f32_vec(ape_rsrc, base_ape_off) @@ -333,20 +300,13 @@ def _issue_phase2_loads(k_i32): def _issue_phase1_loads(k_i32): """Phase 1 (state cache) loads. Returns (kv_list, sc_padded_list) each of length VEC. Score is -inf when s < 0.""" - s = arith.addi( - arith.subi(_to_raw(position), c_K_m1), - k_i32, - ) + s = (fx.Int32(position) - fx.Int32(K - 1) + fx.Int32(k_i32)).ir_value() is_pad = arith.cmpi(CmpIPredicate.slt, s, c_zero_i32) - s_safe = arith.select(is_pad, c_zero_i32, s) - ring = arith.remui(s_safe, c_state_size) + s_safe = fx.Int32(arith.select(is_pad, c_zero_i32, s)) + ring = fx.Int32((fx.Uint32(s_safe.ir_value()) % state_size).ir_value()) # Slot term already folded into the descriptor base. - base_kv_off = ( - ArithValue(ring) * ArithValue(kv_state_pos_stride) + col_off_base - ) - base_sc_off = ( - ArithValue(ring) * ArithValue(score_state_pos_stride) + col_off_base - ) + base_kv_off = ring * fx.Int32(kv_state_pos_stride) + col_off_base + base_sc_off = ring * fx.Int32(score_state_pos_stride) + col_off_base kv_list = _load_f32_vec(kv_state_rsrc, base_kv_off) sc_list = _load_f32_vec(score_state_rsrc, base_sc_off) sc_padded = [ @@ -369,13 +329,16 @@ def _softmax_step_padded( w_old = w_old_list[i] score_k = score_k_list[i] kv_k = kv_k_list[i] - m_new = arith.maximumf(m_old, score_k) + m_new = fx.max(fx.Float32(m_old), fx.Float32(score_k)).ir_value() is_first = arith.cmpf(CmpFPredicate.OEQ, m_old, c_neg_inf) scale_active = fexp_f32(arith.subf(m_old, m_new)) scale_v = arith.select(is_first, c_zero_f32, scale_active) wk_active = fexp_f32(arith.subf(score_k, m_new)) is_pad_score = arith.cmpf(CmpFPredicate.OEQ, score_k, c_neg_inf) w_k = arith.select(is_pad_score, c_zero_f32, wk_active) + # Explicit fastmath float layer: fx `+`/`*` drop fastmath + # here (the rocdl-fastmath pass does not re-add it on gfx1250) + # -> ISA drift (fmac vs split add/mul). Kept raw arith.*FOp. new_kv.append( arith.AddFOp( arith.MulFOp(kv_old, scale_v, fastmath=fm_fast).result, @@ -394,8 +357,8 @@ def _softmax_step_padded( return new_m, new_kv, new_w # -- Wave's K range: [wid * K_PER_WAVE, (wid+1) * K_PER_WAVE) -- - k_start_i32 = ArithValue(wid) * c_K_per_wave - k_end_i32 = k_start_i32 + c_K_per_wave + k_start_i32 = wid * K_PER_WAVE + k_end_i32 = k_start_i32 + K_PER_WAVE # Split point inside this wave's K range. Each wave sees a # window_len-dependent slice of Phase 1 followed by Phase 2. @@ -406,9 +369,8 @@ def _softmax_step_padded( # ``split`` = clamp(wl, k_start, k_end) gives the boundary; # both sub-loops are empty when their bound collapses, so any # of the three cases naturally falls out. - wl_i32 = _to_raw(window_len) - split_lo = arith.maxsi(wl_i32, _to_raw(k_start_i32)) - split_i32 = arith.minsi(split_lo, _to_raw(k_end_i32)) + # split = clamp(window_len, k_start, k_end) -> signed int max/min. + split_i32 = fx.min(fx.max(window_len, k_start_i32), k_end_i32) # State is 3*VEC scalars: m_lane[VEC] + kv_lane[VEC] + w_lane[VEC]. init_m = [c_neg_inf for _ in range(VEC)] @@ -420,12 +382,12 @@ def _softmax_step_padded( # cache; padded softmax (score can be -inf). phase1_local = init_state for k_static, state in range( - _to_raw(k_start_i32), _to_raw(split_i32), 1, init=init_state + k_start_i32.ir_value(), split_i32.ir_value(), 1, init=init_state ): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = fx.Int32(k_static) kv_v, sc_v = _issue_phase1_loads(k_i32) new_m, new_kv, new_w = _softmax_step_padded( m_lane, kv_lane, w_lane, sc_v, kv_v @@ -438,12 +400,12 @@ def _softmax_step_padded( # Carry Phase 1's accumulator through as init. final = phase1_local for k_static, state in range( - _to_raw(split_i32), _to_raw(k_end_i32), 1, init=phase1_local + split_i32.ir_value(), k_end_i32.ir_value(), 1, init=phase1_local ): m_lane = list(state[0:VEC]) kv_lane = list(state[VEC : 2 * VEC]) w_lane = list(state[2 * VEC : 3 * VEC]) - k_i32 = arith.index_cast(i32, _to_raw(k_static)) + k_i32 = fx.Int32(k_static) p2_kv, p2_sc, p2_ape = _issue_phase2_loads(k_i32) p2_score = [ arith.AddFOp(p2_sc[i], p2_ape[i], fastmath=fm_fast).result @@ -466,32 +428,28 @@ def _softmax_step_padded( lds_m_ptr = lds.lds_m.ptr lds_kv_ptr = lds.lds_kv.ptr lds_w_ptr = lds.lds_w.ptr - lds_thread_base = ArithValue(wid) * c_SLICE + ArithValue(lid) * c_VEC + lds_thread_base = wid * SLICE_SZ + lid * VEC for i in range_constexpr(VEC): - idx_i = lds_thread_base + arith.constant(i, type=i32) - fx.ptr_store(m_local[i], lds_m_ptr + fx.Int32(idx_i)) - fx.ptr_store(kv_local[i], lds_kv_ptr + fx.Int32(idx_i)) - fx.ptr_store(w_local[i], lds_w_ptr + fx.Int32(idx_i)) + idx_i = lds_thread_base + i + fx.ptr_store(m_local[i], lds_m_ptr + idx_i) + fx.ptr_store(kv_local[i], lds_kv_ptr + idx_i) + fx.ptr_store(w_local[i], lds_w_ptr + idx_i) gpu.barrier() # -- Cross-wave reduction: only wave 0 reads and reduces -- - # Wave 0's 64 threads cover SLICE_SZ = 64 * VEC head_dim elements - # (VEC elements per thread). For each owned element, the thread - # reads NW values from LDS (one per K-split wave) and computes - # the global online-softmax. - is_wave0 = arith.cmpi(CmpIPredicate.eq, wid, c_zero_i32) - _if_w0 = scf.IfOp(is_wave0) - with _if_then(_if_w0): + # Wave 0's 32 threads cover SLICE_SZ head_dim elements (VEC elements + # per thread). For each owned element, the thread reads NW values + # from LDS (one per K-split wave) and computes the global softmax. + def _wave0(): comp_list = [] for i in range_constexpr(VEC): - lane_off = ArithValue(lid) * c_VEC + arith.constant(i, type=i32) + lane_off = lid * VEC + i # Global max across NW waves for this element. m_g = fx.Float32(c_neg_inf) m_arr = [] for w in range_constexpr(NW): - idx_w = arith.constant(w * SLICE_SZ, type=i32) + lane_off - m_w = fx.ptr_load(lds_m_ptr + fx.Int32(idx_w)) + m_w = fx.ptr_load(lds_m_ptr + (lane_off + w * SLICE_SZ)) m_arr.append(m_w) m_g = m_g.maximumf(m_w) @@ -499,48 +457,48 @@ def _softmax_step_padded( kv_sum = fx.Float32(0.0) w_sum = fx.Float32(0.0) for w in range_constexpr(NW): - idx_w = arith.constant(w * SLICE_SZ, type=i32) + lane_off - kv_w = fx.ptr_load(lds_kv_ptr + fx.Int32(idx_w)) - w_w = fx.ptr_load(lds_w_ptr + fx.Int32(idx_w)) + idx_w = lane_off + w * SLICE_SZ + kv_w = fx.ptr_load(lds_kv_ptr + idx_w) + w_w = fx.ptr_load(lds_w_ptr + idx_w) m_w = m_arr[w] - scale_w = fx.Float32(fexp_f32(_to_raw(m_w - m_g))) + scale_w = fx.Float32(fexp_f32((m_w - m_g).ir_value())) kv_sum = kv_sum + kv_w * scale_w w_sum = w_sum + w_w * scale_w rcp_w = fx.Float32( llvm.call_intrinsic( - f32, "llvm.amdgcn.rcp.f32", [_to_raw(w_sum)], [], [] + f32, "llvm.amdgcn.rcp.f32", [w_sum.ir_value()], [], [] ) ) - comp_list.append(_to_raw(kv_sum * rcp_w)) + comp_list.append(kv_sum * rcp_w) # -- Vectorized write of VEC f32 comp values -- out_rsrc = buffer_ops.create_buffer_resource( kv_compressed, max_size=True ) out_off = ( - ArithValue(pid) * ArithValue(kv_compressed_row_stride) - + col_off_base + fx.Int32(pid) * fx.Int32(kv_compressed_row_stride) + col_off_base ) if const_expr(VEC == 1): - buffer_ops.buffer_store(comp_list[0], out_rsrc, out_off) + buffer_ops.buffer_store(comp_list[0].ir_value(), out_rsrc, out_off) elif const_expr(VEC <= 4): - out_vec = vector.from_elements(T.vec(VEC, T.f32), comp_list) - buffer_ops.buffer_store(out_vec, out_rsrc, out_off) + out_vec = fx.Vector.from_elements(comp_list, dtype=fx.Float32) + buffer_ops.buffer_store(out_vec.ir_value(), out_rsrc, out_off) else: # VEC > 4: AMD HW max is dwordx4 -> split into Nx dwordx4 stores. quarter = 4 n_chunks = VEC // quarter for q in range_constexpr(n_chunks): base = q * quarter - sv = vector.from_elements( - T.vec(quarter, T.f32), - comp_list[base : base + quarter], - ) - buffer_ops.buffer_store( - sv, - out_rsrc, - ArithValue(out_off) + arith.constant(base, type=i32), + sv = fx.Vector.from_elements( + comp_list[base : base + quarter], dtype=fx.Float32 ) + buffer_ops.buffer_store(sv.ir_value(), out_rsrc, out_off + base) + + if wid == 0: + _wave0() + + if fx.Int32(position) >= 0: + _body() @flyc.jit def launch_hca_compress_forward( @@ -562,8 +520,8 @@ def launch_hca_compress_forward( plan_capacity: fx.Int32, stream: fx.Stream, ): - idx_p = arith.index_cast(T.index, _to_raw(plan_capacity)) - idx_s = arith.index_cast(T.index, arith.constant(NUM_SPLIT, type=T.i32)) + idx_p = fx.Int64(plan_capacity) + idx_s = fx.Int64(NUM_SPLIT) k = kernel( kv_in, kv_in_row_stride, @@ -662,67 +620,59 @@ def kernel( ): f32 = T.f32 i32 = T.i32 - vecVf32 = T.vec(VEC, T.f32) pid = fx.block_idx.x tid = fx.thread_idx.x - c_zero_i32 = arith.constant(0, type=i32) - c_one_i32 = arith.constant(1, type=i32) c_eps = arith.constant(rms_eps, type=f32) c_inv_D = arith.constant(1.0 / D, type=f32) - c_ratio = arith.constant(ratio, type=i32) - c_k_per_block = arith.constant(k_per_block, type=i32) - def wave_reduce_add(x): - w = _to_raw(x) + def wave_reduce_add(w): + # w is a raw f32 ir.Value; keep the explicit-fastmath add (fx `+` + # drops fastmath -> ISA drift on gfx1250). for sh_exp in range_constexpr(log2_block): off = BLOCK_THREADS // (2 << sh_exp) - peer = _to_raw(ArithValue(w).shuffle_xor(off, BLOCK_THREADS)) + peer = fx.Float32(w).shuffle_xor(off, BLOCK_THREADS).ir_value() w = arith.AddFOp(w, peer, fastmath=fm_fast).result return w # -- Load plan row -- plan_rsrc = buffer_ops.create_buffer_resource(plan, max_size=True) - plan_base = ArithValue(pid) * arith.constant(4, type=i32) - plan_vec = buffer_ops.buffer_load(plan_rsrc, plan_base, vec_width=4, dtype=i32) - batch_id = vector.extract(plan_vec, static_position=[1], dynamic_position=[]) - position = vector.extract(plan_vec, static_position=[2], dynamic_position=[]) + plan_vec = fx.Vector( + buffer_ops.buffer_load(plan_rsrc, fx.Int32(pid) * 4, vec_width=4, dtype=i32) + ) + batch_id = plan_vec[1] + position = plan_vec[2] - is_active = arith.cmpi(CmpIPredicate.sge, _to_raw(position), c_zero_i32) - _if_active = scf.IfOp(is_active) - with _if_then(_if_active): - tid_x_vec = ArithValue(tid) * arith.constant(VEC, type=i32) + # Sentinel-skip: run the whole body only for position >= 0, as a closure + # under a runtime `if` (rewriter sees an opaque call -> scf.if). + def _body(): + tid_x_vec = fx.Int32(tid) * VEC # -- Load kv_compressed[pid, tid*VEC : tid*VEC + VEC] -- kvc_rsrc = buffer_ops.create_buffer_resource(kv_compressed, max_size=True) - base_off = ( - ArithValue(pid) * ArithValue(kv_compressed_row_stride) + tid_x_vec - ) + base_off = fx.Int32(pid) * fx.Int32(kv_compressed_row_stride) + tid_x_vec # VEC ? {2, 4, 8, 16}: VEC <= 4 -> single dwordx{VEC}; VEC>4 -> Nx dwordx4. + # comp_lane held as raw f32 ir.Values for the explicit-fastmath layer. if const_expr(VEC <= 4): - raw = buffer_ops.buffer_load( - kvc_rsrc, base_off, vec_width=VEC, dtype=f32 + raw = fx.Vector( + buffer_ops.buffer_load(kvc_rsrc, base_off, vec_width=VEC, dtype=f32) ) - comp_lane = [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + comp_lane = [raw[i].ir_value() for i in range(VEC)] else: quarter = 4 n_chunks = VEC // quarter comp_lane = [] for q in range_constexpr(n_chunks): - r = buffer_ops.buffer_load( - kvc_rsrc, - ArithValue(base_off) + arith.constant(q * quarter, type=i32), - vec_width=quarter, - dtype=f32, - ) - for i in range_constexpr(quarter): - comp_lane.append( - vector.extract(r, static_position=[i], dynamic_position=[]) + r = fx.Vector( + buffer_ops.buffer_load( + kvc_rsrc, + base_off + q * quarter, + vec_width=quarter, + dtype=f32, ) + ) + comp_lane += [r[i].ir_value() for i in range_constexpr(quarter)] # -- RMSNorm (wave reduce-add of squares / D + eps; rsqrt) -- sq_local = arith.constant(0.0, type=f32) @@ -742,30 +692,29 @@ def wave_reduce_add(x): rmsw_rsrc = buffer_ops.create_buffer_resource(rms_weight, max_size=True) if const_expr(rms_weight_is_bf16): dwords = (VEC + 1) // 2 - off_dw = ArithValue(tid_x_vec) >> c_one_i32 + # logical shift (tid_x_vec >= 0); fx Int32 >> is arithmetic. + off_dw = fx.Int32((fx.Uint32(tid_x_vec.ir_value()) >> 1).ir_value()) if const_expr(dwords == 1): raw_s = buffer_ops.buffer_load( rmsw_rsrc, off_dw, vec_width=1, dtype=i32 ) - raw = vector.from_elements(T.vec(1, T.i32), [raw_s]) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - rmsw_lane = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, static_position=[i], dynamic_position=[] - ) - rmsw_lane.append(arith.extf(f32, bf16_v)) + raw = fx.Vector.from_elements([raw_s], dtype=fx.Int32) + vec_bf16 = raw.bitcast(fx.BFloat16) + rmsw_lane = [ + vec_bf16[i].to(fx.Float32).ir_value() + for i in range_constexpr(VEC) + ] elif const_expr(dwords <= 4): - raw = buffer_ops.buffer_load( - rmsw_rsrc, off_dw, vec_width=dwords, dtype=i32 - ) - vec_bf16 = vector.bitcast(T.vec(VEC, T.bf16), raw) - rmsw_lane = [] - for i in range_constexpr(VEC): - bf16_v = vector.extract( - vec_bf16, static_position=[i], dynamic_position=[] + raw = fx.Vector( + buffer_ops.buffer_load( + rmsw_rsrc, off_dw, vec_width=dwords, dtype=i32 ) - rmsw_lane.append(arith.extf(f32, bf16_v)) + ) + vec_bf16 = raw.bitcast(fx.BFloat16) + rmsw_lane = [ + vec_bf16[i].to(fx.Float32).ir_value() + for i in range_constexpr(VEC) + ] else: # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4 half_dw = 4 @@ -774,44 +723,37 @@ def wave_reduce_add(x): for chunk in range_constexpr(dwords // half_dw): r = buffer_ops.buffer_load( rmsw_rsrc, - ArithValue(off_dw) - + arith.constant(chunk * half_dw, type=i32), + off_dw + chunk * half_dw, vec_width=half_dw, dtype=i32, ) - vbf16 = vector.bitcast(T.vec(half_bf16, T.bf16), r) - for i in range_constexpr(half_bf16): - bf16_v = vector.extract( - vbf16, static_position=[i], dynamic_position=[] - ) - rmsw_lane.append(arith.extf(f32, bf16_v)) + vbf16 = fx.Vector(r).bitcast(fx.BFloat16) + rmsw_lane += [ + vbf16[i].to(fx.Float32).ir_value() + for i in range_constexpr(half_bf16) + ] else: if const_expr(VEC <= 4): - raw = buffer_ops.buffer_load( - rmsw_rsrc, tid_x_vec, vec_width=VEC, dtype=f32 + raw = fx.Vector( + buffer_ops.buffer_load( + rmsw_rsrc, tid_x_vec, vec_width=VEC, dtype=f32 + ) ) - rmsw_lane = [ - vector.extract(raw, static_position=[i], dynamic_position=[]) - for i in range(VEC) - ] + rmsw_lane = [raw[i].ir_value() for i in range(VEC)] else: quarter = 4 n_chunks = VEC // quarter rmsw_lane = [] for q in range_constexpr(n_chunks): - r = buffer_ops.buffer_load( - rmsw_rsrc, - ArithValue(tid_x_vec) - + arith.constant(q * quarter, type=i32), - vec_width=quarter, - dtype=f32, - ) - for i in range_constexpr(quarter): - rmsw_lane.append( - vector.extract( - r, static_position=[i], dynamic_position=[] - ) + r = fx.Vector( + buffer_ops.buffer_load( + rmsw_rsrc, + tid_x_vec + q * quarter, + vec_width=quarter, + dtype=f32, ) + ) + rmsw_lane += [r[i].ir_value() for i in range_constexpr(quarter)] normed_lane = [ arith.MulFOp( @@ -822,21 +764,20 @@ def wave_reduce_add(x): for i in range(VEC) ] - # -- GPT-J RoPE on RD tail -- - comp_pos_i32 = arith.muli(arith.divsi(_to_raw(position), c_ratio), c_ratio) + # -- GPT-J RoPE on RD tail -- (position >= 0 -> unsigned div) + comp_pos_i32 = fx.Int32((fx.Uint32(position) // ratio).ir_value()) * ratio cos_rsrc = buffer_ops.create_buffer_resource(cos_cache, max_size=True) sin_rsrc = buffer_ops.create_buffer_resource(sin_cache, max_size=True) - c_half_rd = arith.constant(RD // 2, type=i32) - cos_row_base = ArithValue(comp_pos_i32) * c_half_rd + cos_row_base = comp_pos_i32 * (RD // 2) is_rope_t = arith.cmpi( CmpIPredicate.sge, - _to_raw(tid), + tid.ir_value(), arith.constant(ROPE_THREAD_LO, type=i32), ) - rope_rel_raw = ArithValue(tid) - arith.constant(ROPE_THREAD_LO, type=i32) - rope_rel = arith.maxsi(rope_rel_raw, c_zero_i32) - cs_lo = ArithValue(rope_rel) * arith.constant(PAIRS_PER_THREAD, type=i32) + rope_rel_raw = fx.Int32(tid) - ROPE_THREAD_LO + rope_rel = fx.max(rope_rel_raw, fx.Int32(0)) + cs_lo = rope_rel * PAIRS_PER_THREAD if const_expr(PAIRS_PER_THREAD == 1): cos_b = buffer_ops.buffer_load( @@ -845,37 +786,31 @@ def wave_reduce_add(x): sin_b = buffer_ops.buffer_load( sin_rsrc, cos_row_base + cs_lo, vec_width=1, dtype=T.bf16 ) - cos_vals = [arith.extf(f32, cos_b)] - sin_vals = [arith.extf(f32, sin_b)] + cos_vals = [fx.BFloat16(cos_b).to(fx.Float32).ir_value()] + sin_vals = [fx.BFloat16(sin_b).to(fx.Float32).ir_value()] else: - cos_vec = buffer_ops.buffer_load( - cos_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, + cos_vec = fx.Vector( + buffer_ops.buffer_load( + cos_rsrc, + cos_row_base + cs_lo, + vec_width=PAIRS_PER_THREAD, + dtype=T.bf16, + ) ) - sin_vec = buffer_ops.buffer_load( - sin_rsrc, - cos_row_base + cs_lo, - vec_width=PAIRS_PER_THREAD, - dtype=T.bf16, + sin_vec = fx.Vector( + buffer_ops.buffer_load( + sin_rsrc, + cos_row_base + cs_lo, + vec_width=PAIRS_PER_THREAD, + dtype=T.bf16, + ) ) cos_vals = [ - arith.extf( - f32, - vector.extract( - cos_vec, static_position=[i], dynamic_position=[] - ), - ) + cos_vec[i].to(fx.Float32).ir_value() for i in range(PAIRS_PER_THREAD) ] sin_vals = [ - arith.extf( - f32, - vector.extract( - sin_vec, static_position=[i], dynamic_position=[] - ), - ) + sin_vec[i].to(fx.Float32).ir_value() for i in range(PAIRS_PER_THREAD) ] @@ -885,6 +820,9 @@ def wave_reduce_add(x): o = normed_lane[2 * k + 1] c = cos_vals[k] s = sin_vals[k] + # NOTE: real part uses a non-fastmath subtract (default flags); + # explicit-fastmath MulFOp/AddFOp for the rest (fx ops drop + # fastmath on gfx1250 -> ISA drift). new_e = arith.subf( arith.MulFOp(e, c, fastmath=fm_fast).result, arith.MulFOp(o, s, fastmath=fm_fast).result, @@ -898,19 +836,24 @@ def wave_reduce_add(x): rotated_lane[2 * k + 1] = new_o # -- Paged scatter dest (shared by bf16 / fp8) -- - ci = arith.divsi(_to_raw(position), c_ratio) - block_in_seq = arith.divsi(ci, c_k_per_block) - slot_in_block = arith.remui(ci, c_k_per_block) + # position >= 0 (active guard) -> unsigned div/rem (divui/remui). + ci = fx.Int32((fx.Uint32(position) // ratio).ir_value()) + block_in_seq = fx.Int32( + (fx.Uint32(ci.ir_value()) // k_per_block).ir_value() + ) + slot_in_block = fx.Int32( + (fx.Uint32(ci.ir_value()) % k_per_block).ir_value() + ) bt_rsrc = buffer_ops.create_buffer_resource(block_table, max_size=True) - bt_off = ArithValue(batch_id) * ArithValue( - block_table_seq_stride - ) + ArithValue(block_in_seq) + bt_off = ( + fx.Int32(batch_id) * fx.Int32(block_table_seq_stride) + block_in_seq + ) physical_block = buffer_ops.buffer_load( bt_rsrc, bt_off, vec_width=1, dtype=i32 ) # The block term rides on the descriptor's base, not on the # 32-bit offset -- see `block_base_bytes_i64`. - cache_base = ArithValue(slot_in_block) * ArithValue(kv_cache_token_stride) + cache_base = slot_in_block * fx.Int32(kv_cache_token_stride) out_rsrc = buffer_ops.create_buffer_resource( kv_cache, max_size=True, @@ -921,22 +864,23 @@ def wave_reduce_add(x): if const_expr(quant): # -- group_fp8 (V4 nm-asm) via shared emitter (wave32; same layout - # as wave64 CSA/HCA -- single source of truth). -- - _krope_base = ArithValue(slot_in_block) * ArithValue(krope_token_stride) + # as wave64 CSA/HCA -- single source of truth). The emitter lives + # in _common and consumes raw ir.Values. -- + _krope_base = slot_in_block * fx.Int32(krope_token_stride) emit_group_fp8_nm_asm_scatter( normed_lane=normed_lane, rotated_lane=rotated_lane, lane=tid, is_rope_t=is_rope_t, - cache_base=_to_raw(cache_base), - out_rsrc=out_rsrc, - krope_base=_to_raw(_krope_base), - krope_rsrc=buffer_ops.create_buffer_resource( - k_rope_buff, - max_size=True, - base_byte_offset=block_base_bytes_i64( - physical_block, krope_block_stride, 2 - ), + cache_base=cache_base.ir_value(), + out_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(kv_cache))) + + fx.Int64( + block_base_bytes_i64(physical_block, kv_cache_block_stride, 1) + ), + krope_base=_krope_base.ir_value(), + krope_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(k_rope_buff))) + + fx.Int64( + block_base_bytes_i64(physical_block, krope_block_stride, 2) ), VEC=VEC, NOPE=NOPE, @@ -944,8 +888,6 @@ def wave_reduce_add(x): log2_rts=log2_rts, ROPE_THREAD_LO=ROPE_THREAD_LO, wave_width=BLOCK_THREADS, - vecVf32=vecVf32, - fm_fast=fm_fast, ) else: # ---- BF16 single-buffer scatter (nope + rope contiguous) ---- @@ -953,41 +895,37 @@ def wave_reduce_add(x): arith.select(is_rope_t, rotated_lane[i], normed_lane[i]) for i in range_constexpr(VEC) ] - cache_off = ArithValue(cache_base) + tid_x_vec + cache_off = cache_base + tid_x_vec out_vec_t = T.vec(VEC, T.bf16) - raw_vec = vector.from_elements(vecVf32, out_lane) + raw_vec = fx.Vector.from_elements(out_lane, dtype=fx.Float32) bf16_vec = raw_vec.truncf(out_vec_t) - cache_off_dw = ArithValue(cache_off) >> c_one_i32 + # logical shift (cache_off >= 0); fx Int32 >> is arithmetic. + cache_off_dw = fx.Int32( + (fx.Uint32(cache_off.ir_value()) >> 1).ir_value() + ) dwords = (VEC + 1) // 2 - bf16_as_i32 = vector.bitcast(T.vec(dwords, T.i32), bf16_vec) + bf16_as_i32 = bf16_vec.bitcast(fx.Int32) if const_expr(dwords == 1): - scalar_i32 = vector.extract( - bf16_as_i32, static_position=[0], dynamic_position=[] + buffer_ops.buffer_store( + bf16_as_i32[0].ir_value(), out_rsrc, cache_off_dw ) - buffer_ops.buffer_store(scalar_i32, out_rsrc, cache_off_dw) elif const_expr(dwords <= 4): - buffer_ops.buffer_store(bf16_as_i32, out_rsrc, cache_off_dw) - else: - # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4 - c4_i32 = arith.constant(4, type=i32) - lo = vector.extract_strided_slice( - T.vec(4, T.i32), - bf16_as_i32, - offsets=[0], - sizes=[4], - strides=[1], + buffer_ops.buffer_store( + bf16_as_i32.ir_value(), out_rsrc, cache_off_dw ) - hi = vector.extract_strided_slice( - T.vec(4, T.i32), - bf16_as_i32, - offsets=[4], - sizes=[4], - strides=[1], + else: + # dwords > 4 (VEC=16 -> dwords=8): split into 2x dwordx4. + lo = fx.Vector.from_elements( + [bf16_as_i32[i] for i in range(4)], dtype=fx.Int32 ) - buffer_ops.buffer_store(lo, out_rsrc, cache_off_dw) - buffer_ops.buffer_store( - hi, out_rsrc, ArithValue(cache_off_dw) + c4_i32 + hi = fx.Vector.from_elements( + [bf16_as_i32[i] for i in range(4, 8)], dtype=fx.Int32 ) + buffer_ops.buffer_store(lo.ir_value(), out_rsrc, cache_off_dw) + buffer_ops.buffer_store(hi.ir_value(), out_rsrc, cache_off_dw + 4) + + if fx.Int32(position) >= 0: + _body() @flyc.jit def launch_hca_norm_rope_scatter( @@ -1008,7 +946,7 @@ def launch_hca_norm_rope_scatter( plan_capacity: fx.Int32, stream: fx.Stream, ): - idx_p = arith.index_cast(T.index, _to_raw(plan_capacity)) + idx_p = fx.Int64(plan_capacity) k = kernel( kv_compressed, kv_compressed_row_stride, diff --git a/aiter/ops/flydsl/kernels/mfma_epilogues.py b/aiter/ops/flydsl/kernels/mfma_epilogues.py index 5a97f3664f..fe8a70c262 100644 --- a/aiter/ops/flydsl/kernels/mfma_epilogues.py +++ b/aiter/ops/flydsl/kernels/mfma_epilogues.py @@ -27,37 +27,24 @@ waves are partitioned into two groups (group A uses ``lds_out``, group B uses ``lds_out_split``), each handling half of the N dimension. -These helpers are intentionally *dialect-agnostic*: callers pass the dialect -modules (`arith`, `vector`, `gpu`) and the `range_constexpr` iterator. +The helpers import the `gpu` dialect and the `range_constexpr` iterator +directly, so callers no longer inject the dialect modules. """ from __future__ import annotations from collections.abc import Callable -from contextlib import contextmanager +import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir import ir -from flydsl._mlir.dialects.arith import CmpIPredicate +from flydsl.expr import arith, gpu, range_constexpr from flydsl.expr.typing import T - - -@contextmanager -def _if_then(if_op, scf): - """Compat helper for SCF IfOp then-region across old/new Python APIs.""" - with ir.InsertionPoint(if_op.then_block): - try: - yield if_op.then_block - finally: - blk = if_op.then_block - if (not blk.operations) or not isinstance(blk.operations[-1], scf.YieldOp): - scf.YieldOp([]) +from flydsl.expr.typing import Vector as Vec def default_epilog( *, - arith, - range_constexpr, m_repeat: int, lane_div_16, bx_m, @@ -68,8 +55,6 @@ def default_epilog( The mapping matches the common MFMA fragment layout used across kernels in this repo. Args: - arith: flydsl arith ext module. - range_constexpr: compile-time unrolled range helper. m_repeat: tile_m // 16 (python int). lane_div_16: index Value (0..3). bx_m: base row (index Value). For MoE, this is the base sorted-row for the tile. @@ -81,7 +66,7 @@ def default_epilog( ii_idx_list = [fx.Index(ii) for ii in range(4)] for mi in range_constexpr(m_repeat): - mi_base = arith.constant(mi * 16, index=True) + mi_base = fx.Index(mi * 16).ir_value() for ii in range_constexpr(4): row_off = lane_div_16_mul4 + ii_idx_list[ii] row_in_tile = mi_base + row_off @@ -91,11 +76,6 @@ def default_epilog( def c_shuffle_epilog( *, - arith, - vector, - gpu, - scf=None, - range_constexpr, # Tile params tile_m: int, tile_n: int, @@ -157,9 +137,6 @@ def c_shuffle_epilog( # Group B (waves N/2..N-1) uses lds_out_split, columns [tile_n/2, tile_n) # Each group writes/reads independently; same barriers synchronise all waves. if lds_out_split is not None: - if scf is None: - raise ValueError("scf module is required for split-LDS cshuffle") - _half_n = int(tile_n) // 2 _half_threads = int(block_size) // 2 EVec = int(e_vec) @@ -177,11 +154,11 @@ def c_shuffle_epilog( m_reps_s = int(tile_m) // CShuffleMLane_s n_reps_s = _half_n // (CShuffleNLane_s * EVec) - _half_n_idx = arith.constant(_half_n, index=True) - _half_thr_idx = arith.constant(_half_threads, index=True) - _zero_idx = arith.constant(0, index=True) + _half_n_idx = fx.Index(_half_n).ir_value() + _half_thr_idx = fx.Index(_half_threads).ir_value() + _zero_idx = fx.Index(0).ir_value() - _is_group_b = arith.cmpi(CmpIPredicate.uge, tx, _half_thr_idx) + _is_group_b = fx.Index(tx) >= fx.Index(_half_thr_idx) # -- write phase (all waves, each to its group's LDS buffer) -- n_tile_base_v = n_tile_base @@ -190,36 +167,36 @@ def c_shuffle_epilog( def _write_row_split(mi: int, ii: int, row_in_tile, row): row_base_lds = row_in_tile * _half_n_idx - _if_g = scf.IfOp(_is_group_b, has_else=True) - with ir.InsertionPoint(_if_g.then_block): - write_row_to_lds( - mi=mi, - ii=ii, - row_in_tile=row_in_tile, - row=row, - row_base_lds=row_base_lds, - col_base_local=col_base_local_b, - num_acc_n=num_acc_n, - lds_out=lds_out_split, - ) - scf.YieldOp([]) - with ir.InsertionPoint(_if_g.else_block): - write_row_to_lds( - mi=mi, - ii=ii, - row_in_tile=row_in_tile, - row=row, - row_base_lds=row_base_lds, - col_base_local=col_base_local_a, - num_acc_n=num_acc_n, - lds_out=lds_out, - ) - scf.YieldOp([]) + + @flyc.jit + def _write_group(): + if _is_group_b: + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_b, + num_acc_n=num_acc_n, + lds_out=lds_out_split, + ) + else: + write_row_to_lds( + mi=mi, + ii=ii, + row_in_tile=row_in_tile, + row=row, + row_base_lds=row_base_lds, + col_base_local=col_base_local_a, + num_acc_n=num_acc_n, + lds_out=lds_out, + ) + + _write_group() gpu.barrier() default_epilog( - arith=arith, - range_constexpr=range_constexpr, m_repeat=m_repeat, lane_div_16=lane_div_16, bx_m=bx_m, @@ -228,11 +205,11 @@ def _write_row_split(mi: int, ii: int, row_in_tile, row): gpu.barrier() # -- read phase (each group reads from its own LDS buffer) -- - tx_local = tx - arith.select(_is_group_b, _half_thr_idx, _zero_idx) - c_nlane_s = arith.constant(CShuffleNLane_s, index=True) + tx_local = tx - _is_group_b.select(_half_thr_idx, _zero_idx) + c_nlane_s = fx.Index(CShuffleNLane_s).ir_value() m_lane_s = tx_local // c_nlane_s n_lane_s = tx_local % c_nlane_s - c_evec = arith.constant(EVec, index=True) + c_evec = fx.Index(EVec).ir_value() if frag_elem_type is None: frag_elem_type = T.f16 @@ -242,7 +219,7 @@ def _write_row_split(mi: int, ii: int, row_in_tile, row): _precomputed_rows_s = [] for mr in range_constexpr(m_reps_s): - row_base_m = arith.constant(mr * CShuffleMLane_s, index=True) + row_base_m = fx.Index(mr * CShuffleMLane_s).ir_value() row_local = row_base_m + m_lane_s row = bx_m_v + row_local row_ctx_raw = ( @@ -252,12 +229,7 @@ def _write_row_split(mi: int, ii: int, row_in_tile, row): ) row_ctx = row_ctx_raw row_pred = None - if ( - scf is not None - and row_ctx_raw is not None - and isinstance(row_ctx_raw, tuple) - and len(row_ctx_raw) == 2 - ): + if isinstance(row_ctx_raw, tuple) and len(row_ctx_raw) == 2: row_ctx, row_pred = row_ctx_raw _precomputed_rows_s.append((row_local, row, row_ctx, row_pred)) @@ -267,23 +239,25 @@ def _write_row_split(mi: int, ii: int, row_in_tile, row): def _do_store_row_split(row=row, row_ctx=row_ctx, row_local=row_local): row_base_lds = row_local * _half_n_idx for nr in range_constexpr(n_reps_s): - col_base_nr = arith.constant( - nr * (CShuffleNLane_s * EVec), index=True - ) + col_base_nr = fx.Index(nr * (CShuffleNLane_s * EVec)).ir_value() col_pair0_local = col_base_nr + (n_lane_s * c_evec) lds_idx = row_base_lds + col_pair0_local - _if_ld = scf.IfOp(_is_group_b, [vec_frag], has_else=True) - with ir.InsertionPoint(_if_ld.then_block): - fb = vector.load_op(vec_frag, lds_out_split, [lds_idx]) - scf.YieldOp([fb]) - with ir.InsertionPoint(_if_ld.else_block): - fa = vector.load_op(vec_frag, lds_out, [lds_idx]) - scf.YieldOp([fa]) - frag = _if_ld.results[0] - - col_pair0 = col_pair0_local + arith.select( - _is_group_b, _half_n_idx, _zero_idx + # Select the per-group LDS source buffer, then issue a single + # load. Both buffers share the same `lds_idx`, so selecting the + # source (not the loaded value) keeps this to one ds_read; + # loading from both and selecting the result would double LDS + # traffic. A memref-typed ternary is a hard boundary with no + # numeric-`fx` wrapper, so the raw `arith` select is localized + # here (a value-returning `@flyc.jit` mis-merges the two + # branches and drifts numerically on multi-row tiles). + src = arith.ArithValue(_is_group_b.ir_value()).select( + lds_out_split, lds_out + ) + frag = Vec.load(vec_frag, src, [lds_idx]).ir_value() + + col_pair0 = col_pair0_local + _is_group_b.select( + _half_n_idx, _zero_idx ) store_pair( row_local=row_local, @@ -295,9 +269,15 @@ def _do_store_row_split(row=row, row_ctx=row_ctx, row_local=row_local): ) if row_pred is not None: - _if_row = scf.IfOp(row_pred) - with _if_then(_if_row, scf): - _do_store_row_split() + + # `_row_guard_split` is called immediately in this loop iteration, + # so capturing `row_pred`/`_do_store_row_split` by closure is safe. + @flyc.jit + def _row_guard_split(): + if row_pred: # noqa: B023 + _do_store_row_split() + + _row_guard_split() else: _do_store_row_split() @@ -306,7 +286,7 @@ def _do_store_row_split(row=row, row_ctx=row_ctx, row_local=row_local): # ===================== Standard (non-split) path below ===================== # ---------------- Step 1: write C tile to LDS (row-major, fp16) ---------------- - tile_n_idx = arith.constant(int(tile_n), index=True) + tile_n_idx = fx.Index(int(tile_n)).ir_value() n_tile_base_v = n_tile_base col_base_local = n_tile_base_v + lane_mod_16 # index within [0,tile_n) @@ -332,8 +312,6 @@ def _write_row(mi: int, ii: int, row_in_tile, row): # Ensure all LDS reads finished before the lds write. gpu.barrier() default_epilog( - arith=arith, - range_constexpr=range_constexpr, m_repeat=m_repeat, lane_div_16=lane_div_16, bx_m=bx_m, @@ -367,7 +345,7 @@ def _write_row(mi: int, ii: int, row_in_tile, row): # them instead of serializing each load with s_waitcnt vmcnt(0). _precomputed_rows = [] for mr in range_constexpr(m_reps_shuffle): - row_base_m = arith.constant(mr * CShuffleMLane, index=True) + row_base_m = fx.Index(mr * CShuffleMLane).ir_value() row_local = row_base_m + m_lane row = bx_m_v + row_local @@ -377,16 +355,11 @@ def _write_row(mi: int, ii: int, row_in_tile, row): else None ) - # Optional row-level predicate: if `precompute_row` returns `(ctx, pred_i1)` and `scf` - # is provided, we can skip the entire N-loop for invalid rows (cheaper than per-store checks). + # Optional row-level predicate: if `precompute_row` returns `(ctx, pred_i1)`, + # skip the entire N-loop for invalid rows (cheaper than per-store checks). row_ctx = row_ctx_raw row_pred = None - if ( - scf is not None - and row_ctx_raw is not None - and isinstance(row_ctx_raw, tuple) - and len(row_ctx_raw) == 2 - ): + if isinstance(row_ctx_raw, tuple) and len(row_ctx_raw) == 2: row_ctx, row_pred = row_ctx_raw _precomputed_rows.append((row_local, row, row_ctx, row_pred)) @@ -400,11 +373,11 @@ def _do_store_row(row=row, row_ctx=row_ctx, row_local=row_local): if _lds_row_base_offset is not None: row_base_lds = row_base_lds + _lds_row_base_offset for nr in range_constexpr(n_reps_shuffle): - col_base_nr = arith.constant(nr * (CShuffleNLane * EVec), index=True) + col_base_nr = fx.Index(nr * (CShuffleNLane * EVec)).ir_value() col_pair0 = col_base_nr + (n_lane * c_evec) # even col within tile lds_idx_pair = row_base_lds + col_pair0 - frag = vector.load_op(vec_frag, lds_out, [lds_idx_pair]) + frag = Vec.load(vec_frag, lds_out, [lds_idx_pair]).ir_value() store_pair( row_local=row_local, @@ -416,9 +389,15 @@ def _do_store_row(row=row, row_ctx=row_ctx, row_local=row_local): ) if row_pred is not None: - _if_row = scf.IfOp(row_pred) - with _if_then(_if_row, scf): - _do_store_row() + + # `_row_guard` is called immediately in this loop iteration, so + # capturing `row_pred`/`_do_store_row` by closure is safe. + @flyc.jit + def _row_guard(): + if row_pred: # noqa: B023 + _do_store_row() + + _row_guard() else: _do_store_row() @@ -427,17 +406,12 @@ def mfma_epilog( *, use_cshuffle: bool, # Common (always required) - arith, - range_constexpr, m_repeat: int, lane_div_16, bx_m, # Default epilog (required when use_cshuffle=False) body_row: Callable | None = None, # CShuffle epilog (required when use_cshuffle=True) - vector=None, - gpu=None, - scf=None, tile_m: int | None = None, tile_n: int | None = None, e_vec: int = 2, @@ -458,8 +432,6 @@ def mfma_epilog( if body_row is None: raise ValueError("mfma_epilog(use_cshuffle=False) requires `body_row`.") return default_epilog( - arith=arith, - range_constexpr=range_constexpr, m_repeat=m_repeat, lane_div_16=lane_div_16, bx_m=bx_m, @@ -467,11 +439,6 @@ def mfma_epilog( ) return c_shuffle_epilog( - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, tile_m=int(tile_m), tile_n=int(tile_n), e_vec=int(e_vec), diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py index bb0fae7c1e..c0654b5948 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py @@ -2824,8 +2824,6 @@ def fused_write(mi, ii, row_in_tile, row): ) default_epilog( - arith=arith, - range_constexpr=range_constexpr, m_repeat=m_repeat, lane_div_16=lane_div_16, bx_m=bx_m, @@ -2903,11 +2901,6 @@ def fused_read(_row_local=_row_local, row=row, rc=rc): gui_by_n = by_n // arith.constant(2, index=True) gui_n_tile_base = n_tile_base // arith.constant(2, index=True) c_shuffle_epilog( - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, tile_m=tile_m, tile_n=gui_tile_n, e_vec=e_vec, @@ -2931,11 +2924,6 @@ def fused_read(_row_local=_row_local, row=row, rc=rc): eff_e_vec = e_vec_sk acc = acc_gate c_shuffle_epilog( - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, tile_m=tile_m, tile_n=tile_n, e_vec=eff_e_vec, @@ -2962,11 +2950,6 @@ def fused_read(_row_local=_row_local, row=row, rc=rc): acc = acc_gate sk_n_offset[0] = 0 c_shuffle_epilog( - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, tile_m=tile_m, tile_n=tile_n, e_vec=eff_e_vec, @@ -2993,11 +2976,6 @@ def fused_read(_row_local=_row_local, row=row, rc=rc): acc = acc_up sk_n_offset[0] = inter_dim c_shuffle_epilog( - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, tile_m=tile_m, tile_n=tile_n, e_vec=eff_e_vec, @@ -3020,11 +2998,6 @@ def fused_read(_row_local=_row_local, row=row, rc=rc): ) else: c_shuffle_epilog( - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, tile_m=tile_m, tile_n=tile_n, e_vec=e_vec, @@ -5376,11 +5349,6 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): e_vec = 2 if accumulate else min(body_tile_n // 32, 8) rocdl.s_setprio(3) c_shuffle_epilog( - arith=arith, - vector=vector, - gpu=gpu, - scf=scf, - range_constexpr=range_constexpr, tile_m=tile_m, tile_n=body_tile_n, e_vec=e_vec, diff --git a/aiter/ops/flydsl/kernels/moe_sorting_kernel.py b/aiter/ops/flydsl/kernels/moe_sorting_kernel.py index 06c9e3073d..3147741e71 100644 --- a/aiter/ops/flydsl/kernels/moe_sorting_kernel.py +++ b/aiter/ops/flydsl/kernels/moe_sorting_kernel.py @@ -27,13 +27,10 @@ import torch from flydsl.expr import gpu, range_constexpr from flydsl.expr import rocdl as fly_rocdl -from flydsl.expr.arith import ArithValue from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec from flydsl.runtime.device import get_rocm_arch as get_hip_arch -from aiter.ops.flydsl.kernels import buffer_ops - from .kernels_common import get_warp_size from .tensor_shim import _run_compiled @@ -116,18 +113,17 @@ def _allwave_inclusive_prefix_sum(val, lane, wave, scratch_mr, NUM_WAVES, WARP_S @flyc.jit -def _zero_moe_buf_grid_stride(moe_buf_rsrc, gid_v4, stride_v4, total_v4, oob_idx): - """Grid-stride loop zeroing moe_buf via vectorized buffer_store.""" +def _zero_moe_buf_grid_stride(moe_buf_it, gid_v4, stride_v4, total_v4, oob_idx): + """Grid-stride loop zeroing moe_buf via vectorized (4xi32) stores.""" c_one = fx.Int32(1) niters = (total_v4 + stride_v4 - c_one) // stride_v4 c_zero_v4 = fx.Vector.filled(4, 0, fx.Int32) c4 = fx.Int32(4) - for _z in range(fx.Index(0), ArithValue(niters).index_cast(T.index), fx.Index(1)): + for _z in range(fx.Int32(0), niters, fx.Int32(1)): idx = gid_v4 + fx.Int32(_z) * stride_v4 valid = idx < total_v4 - buffer_ops.buffer_store( - c_zero_v4, moe_buf_rsrc, valid.select(idx * c4, oob_idx) - ) + # Store 4 i32 at element offset idx*4 (OOB sentinel drops the write). + fx.ptr_store(c_zero_v4, moe_buf_it + fx.Int64(valid.select(idx * c4, oob_idx))) def _extend_prefix_sum_serial(mr, start_block, E, load_fn, store_fn): @@ -146,41 +142,77 @@ def _extend_prefix_sum_serial(mr, start_block, E, load_fn, store_fn): @flyc.jit -def _write_expert_id_blocks(sorted_e_rsrc, local_eid, blk_start, n_blks): +def _write_expert_id_blocks(sorted_e_it, local_eid, blk_start, n_blks): """Write local_eid to sorted_expert_ids[blk_start .. blk_start+n_blks).""" - for _jb in range(fx.Index(0), ArithValue(n_blks).index_cast(T.index), fx.Index(1)): + for _jb in range(fx.Int32(0), n_blks, fx.Int32(1)): blk_idx = blk_start + fx.Int32(_jb) - buffer_ops.buffer_store(local_eid, sorted_e_rsrc, blk_idx) + _gst(sorted_e_it, local_eid, blk_idx) @flyc.jit def _fill_sentinel_slots( - sorted_ids_rsrc, sorted_w_rsrc, start, count, sentinel, block_size, tid, oob_idx + sorted_ids_it, sorted_w_it, start, count, sentinel, block_size, tid, oob_idx ): """Cooperative sentinel fill: threads fill [start, start+count) with sentinels.""" c_zero = fx.Int32(0) end = start + count niters = (count + fx.Int32(block_size) - fx.Int32(1)) // fx.Int32(block_size) - for _p in range(fx.Index(0), ArithValue(niters).index_cast(T.index), fx.Index(1)): + for _p in range(fx.Int32(0), niters, fx.Int32(1)): slot = start + fx.Int32(_p) * fx.Int32(block_size) + tid safe = (slot < end).select(slot, oob_idx) - buffer_ops.buffer_store(sentinel, sorted_ids_rsrc, safe) - buffer_ops.buffer_store(c_zero, sorted_w_rsrc, safe) + _gst(sorted_ids_it, sentinel, safe) + _gst(sorted_w_it, c_zero, safe) # --------------------------------------------------------------------------- # LDS helpers for multiphase kernels (module-level, used inside @flyc.kernel) # --------------------------------------------------------------------------- def _lds_load_raw(raw_ptr, idx): - """Load i32 from an LDS pointer at element offset `idx` (i32 or index).""" + """Load i32 from an LDS pointer at i32 element offset `idx`.""" return fx.ptr_load(raw_ptr + fx.Int64(idx)) def _lds_store_raw(raw_ptr, val, idx): - """Store i32 to an LDS pointer at element offset `idx` (i32 or index).""" + """Store i32 to an LDS pointer at i32 element offset `idx`.""" fx.ptr_store(val, raw_ptr + fx.Int64(idx)) +# --------------------------------------------------------------------------- +# Global buffer-tensor helpers. ``it`` is ``fx.get_iter(make_buffer_tensor(t))``: +# the iterator carries the OOB-checked V# descriptor, so a load/store at the +# sentinel element index 0x7FFFFFFF is dropped (store) / returns 0 (load), +# preserving the ``.select(idx, oob_idx)`` masking used throughout this file. +# Offsets are in ELEMENTS of the buffer's dtype. +# --------------------------------------------------------------------------- +def _buf_iter(tensor): + """OOB-checked V# iterator over a global tensor.""" + return fx.get_iter(fx.rocdl.make_buffer_tensor(tensor, max_size=True)) + + +def _i8_global_ptr(tensor): + """Raw i8 global pointer to `tensor` for byte-addressed stores. + + Used only for the uint8 mesh scatter, which is always issued under an + in-bounds guard (`if valid` / `if is_mine`) -- no OOB access occurs, so a + plain LLVM store suffices (no V# descriptor needed). i8 element == 1 byte, + so a byte offset is the element index directly (no *sizeof scaling). + """ + i8pt = fx.PointerType.get( + fx.Int8.ir_type, address_space=fx.AddressSpace.Global, alignment=1 + ) + return fx.inttoptr(i8pt, fx.Int64(fx.ptrtoint(fx.get_iter(tensor)))) + + +def _gld(it, idx): + """Scalar load at i32 element offset `idx` (OOB -> 0).""" + return fx.ptr_load(it + fx.Int64(idx)) + + +def _gst(it, val, idx): + """Scalar store at i32 element offset `idx` (OOB dropped).""" + fx.ptr_store(val, it + fx.Int64(idx)) + + _dummy_mask_cache = {} # device -> torch.Tensor(1, dtype=i32, value=1) _dummy_local_tokens_cache = {} # dummy placeholder when has_local_tokens=False @@ -284,36 +316,22 @@ def moe_sorting_oneshot_kernel( tokens = i32_tokens if has_local_tokens: - ltok_rsrc = buffer_ops.create_buffer_resource( - local_tokens_tensor, max_size=True - ) - tokens = buffer_ops.buffer_load( - ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 - ) + ltok_it = _buf_iter(local_tokens_tensor) + tokens = _gld(ltok_it, fx.Int32(0)) c_zero_i32 = fx.Int32(0) c_one_i32 = fx.Int32(1) c_oob_idx = fx.Int32(0x7FFFFFFF) c4_i32 = fx.Int32(4) - # Buffer resources (needed by both paths, defined at top level) - moe_buf_rsrc = buffer_ops.create_buffer_resource(moe_buf, max_size=True) - topk_ids_rsrc = buffer_ops.create_buffer_resource( - topk_ids_tensor, max_size=True - ) - weights_rsrc = buffer_ops.create_buffer_resource( - topk_weights_tensor, max_size=True - ) - sorted_ids_rsrc = buffer_ops.create_buffer_resource( - sorted_token_ids, max_size=True - ) - sorted_w_rsrc = buffer_ops.create_buffer_resource( - sorted_weights_out, max_size=True - ) - sorted_e_rsrc = buffer_ops.create_buffer_resource( - sorted_expert_ids, max_size=True - ) - nvalid_rsrc = buffer_ops.create_buffer_resource(num_valid_ids, max_size=True) - mask_rsrc = buffer_ops.create_buffer_resource(expert_mask_tensor, max_size=True) + # Buffer iterators (OOB-checked V# descriptors; needed by both paths). + moe_buf_it = _buf_iter(moe_buf) + topk_ids_it = _buf_iter(topk_ids_tensor) + weights_it = _buf_iter(topk_weights_tensor) + sorted_ids_it = _buf_iter(sorted_token_ids) + sorted_w_it = _buf_iter(sorted_weights_out) + sorted_e_it = _buf_iter(sorted_expert_ids) + nvalid_it = _buf_iter(num_valid_ids) + mask_it = _buf_iter(expert_mask_tensor) # LDS: capture field pointers ONCE — dominates all child scf.for/scf.if. lds = fx.SharedAllocator().allocate(SharedStorage).peek() @@ -335,7 +353,7 @@ def moe_sorting_oneshot_kernel( num_zero_blocks = gpu.grid_dim.x - c_one_i32 zero_stride_v4 = num_zero_blocks * fx.Int32(ONESHOT_BLOCK) _zero_moe_buf_grid_stride( - moe_buf_rsrc, + moe_buf_it, zero_gid_v4, zero_stride_v4, i32_moe_buf_elems >> fx.Int32(2), @@ -350,9 +368,8 @@ def moe_sorting_oneshot_kernel( idx = fx.Int32(i_clear) + tid is_valid = idx < fx.Int32(sub_tokens * smem_cols) safe_idx = is_valid.select(idx, c_zero_i32) - safe_idx_ix = ArithValue(safe_idx).index_cast(T.index) # Always store; out-of-bounds threads harmlessly write to index 0 - _lds_store_raw(mesh_mr, c_zero_i32, safe_idx_ix) + _lds_store_raw(mesh_mr, c_zero_i32, safe_idx) gpu.barrier() # Fill mesh: for each (token, topk_slot), write topk_slot+1 to mesh[token, expert_id] @@ -366,9 +383,7 @@ def moe_sorting_oneshot_kernel( topk_slot = safe_flat % c_topk global_idx = token_id * c_topk + topk_slot - eid = buffer_ops.buffer_load( - topk_ids_rsrc, global_idx, vec_width=1, dtype=T.i32 - ) + eid = _gld(topk_ids_it, global_idx) # mesh[token_id, eid] = topk_slot + 1 (valid threads only). # Invalid threads must NOT write to mesh[0] — that would race @@ -376,9 +391,8 @@ def moe_sorting_oneshot_kernel( mesh_addr = token_id * c_smem_cols + eid last_mesh_idx = fx.Int32(sub_tokens * smem_cols - 1) safe_mesh_addr = is_valid.select(mesh_addr, last_mesh_idx) - safe_mesh_ix = ArithValue(safe_mesh_addr).index_cast(T.index) val = is_valid.select(topk_slot + c_one_i32, c_zero_i32) - _lds_store_raw(mesh_mr, val, safe_mesh_ix) + _lds_store_raw(mesh_mr, val, safe_mesh_addr) gpu.barrier() # ===================== PHASE 2: Count + Prefix Sum ===================== @@ -407,8 +421,7 @@ def moe_sorting_oneshot_kernel( safe_sub = combined_valid.select(sub_idx, c_zero_i32) safe_eid = combined_valid.select(eid_local, c_zero_i32) mesh_rd_addr = safe_sub * c_smem_cols + safe_eid - mesh_rd_ix = ArithValue(mesh_rd_addr).index_cast(T.index) - mesh_val = _lds_load_raw(mesh_mr, mesh_rd_ix) + mesh_val = _lds_load_raw(mesh_mr, mesh_rd_addr) has_token = combined_valid.select( (mesh_val != c_zero_i32).select(c_one_i32, c_zero_i32), @@ -428,9 +441,8 @@ def moe_sorting_oneshot_kernel( # cumsum[0] which is harmless (cumsum[0] is always 0). write_valid = eid_valid & (lane_group_os == c_zero_i32) cs_idx = write_valid.select(eid_local + c_one_i32, c_zero_i32) - cs_ix = ArithValue(cs_idx).index_cast(T.index) cs_val = write_valid.select(cnt, c_zero_i32) - _lds_store_raw(cumsum_mr, cs_val, cs_ix) + _lds_store_raw(cumsum_mr, cs_val, cs_idx) gpu.barrier() # Phase 2b: Prefix sum over expert counts. @@ -440,15 +452,14 @@ def moe_sorting_oneshot_kernel( cvt_valid = cvt_eid < c_E # Safe index: valid → cumsum[eid+1], invalid → cumsum[0] (write 0, harmless) safe_cvt_idx = cvt_valid.select(cvt_eid + c_one_i32, c_zero_i32) - cvt_ix = ArithValue(safe_cvt_idx).index_cast(T.index) - raw_cnt_cvt = _lds_load_raw(cumsum_mr, cvt_ix) + raw_cnt_cvt = _lds_load_raw(cumsum_mr, safe_cvt_idx) blocks_cvt = (raw_cnt_cvt + c_unit - c_one_i32) // c_unit padded_cvt = (raw_cnt_cvt == c_zero_i32).select( c_zero_i32, blocks_cvt * c_unit ) # Valid threads write padded value; invalid threads write 0 to cumsum[0] _lds_store_raw( - cumsum_mr, cvt_valid.select(padded_cvt, c_zero_i32), cvt_ix + cumsum_mr, cvt_valid.select(padded_cvt, c_zero_i32), safe_cvt_idx ) gpu.barrier() @@ -460,19 +471,15 @@ def moe_sorting_oneshot_kernel( ep_eid = fx.Int32(i_ep) + tid ep_valid = ep_eid < c_E ep_safe_eid = ep_valid.select(ep_eid, c_zero_i32) - ep_m = buffer_ops.buffer_load( - mask_rsrc, ep_safe_eid, vec_width=1, dtype=T.i32 - ) + ep_m = _gld(mask_it, ep_safe_eid) should_zero = ep_valid & (ep_m == c_zero_i32) - ep_cs_ix = ArithValue( - ep_valid.select(ep_eid + c_one_i32, c_zero_i32) - ).index_cast(T.index) + ep_cs_idx = ep_valid.select(ep_eid + c_one_i32, c_zero_i32) _lds_store_raw( cumsum_mr, should_zero.select( - c_zero_i32, _lds_load_raw(cumsum_mr, ep_cs_ix) + c_zero_i32, _lds_load_raw(cumsum_mr, ep_cs_idx) ), - ep_cs_ix, + ep_cs_idx, ) gpu.barrier() @@ -481,13 +488,11 @@ def moe_sorting_oneshot_kernel( for _ps_chunk in range_constexpr(0, E, ONESHOT_BLOCK): ps_eid = fx.Int32(_ps_chunk) + tid ps_valid = ps_eid < c_E - ps_safe_ix = ArithValue( - ps_valid.select(ps_eid + c_one_i32, c_zero_i32) - ).index_cast(T.index) + ps_safe_idx = ps_valid.select(ps_eid + c_one_i32, c_zero_i32) ps_val = ps_valid.select( - _lds_load_raw(cumsum_mr, ps_safe_ix), c_zero_i32 + _lds_load_raw(cumsum_mr, ps_safe_idx), c_zero_i32 ) - _lds_store_raw(cumdup_mr, ps_val, ps_safe_ix) + _lds_store_raw(cumdup_mr, ps_val, ps_safe_idx) _lds_store_raw(cumdup_mr, c_zero_i32, c_zero_i32) gpu.barrier() @@ -502,9 +507,7 @@ def moe_sorting_oneshot_kernel( _lds_store_raw( cumdup_mr, ps_tid_valid.select(inclusive_ps, c_zero_i32), - ArithValue(ps_tid_valid.select(tid + c_one_i32, c_zero_i32)).index_cast( - T.index - ), + ps_tid_valid.select(tid + c_one_i32, c_zero_i32), ) gpu.barrier() @@ -521,10 +524,9 @@ def moe_sorting_oneshot_kernel( gpu.barrier() # Write num_valid_ids from cumdup[E] - cs_E_ix_ps = ArithValue(c_E).index_cast(T.index) - total_padded = _lds_load_raw(cumdup_mr, cs_E_ix_ps) - buffer_ops.buffer_store(total_padded, nvalid_rsrc, c_zero_i32) - buffer_ops.buffer_store(tokens, nvalid_rsrc, c_one_i32) + total_padded = _lds_load_raw(cumdup_mr, c_E) + _gst(nvalid_it, total_padded, c_zero_i32) + _gst(nvalid_it, tokens, c_one_i32) gpu.barrier() # Copy cumdup → cumsum (all threads, one expert per thread) @@ -532,9 +534,8 @@ def moe_sorting_oneshot_kernel( cp_idx = fx.Int32(i_cp) + tid cp_valid = cp_idx <= c_E safe_cp_idx = cp_valid.select(cp_idx, c_zero_i32) - cp_ix = ArithValue(safe_cp_idx).index_cast(T.index) - cp_val = _lds_load_raw(cumdup_mr, cp_ix) - _lds_store_raw(cumsum_mr, cp_val, cp_ix) + cp_val = _lds_load_raw(cumdup_mr, safe_cp_idx) + _lds_store_raw(cumsum_mr, cp_val, safe_cp_idx) gpu.barrier() if has_mask: @@ -544,14 +545,10 @@ def moe_sorting_oneshot_kernel( ml_eid = fx.Int32(i_ml) + tid ml_valid = ml_eid < c_E safe_ml_eid = ml_valid.select(ml_eid, c_zero_i32) - ml_mask = buffer_ops.buffer_load( - mask_rsrc, safe_ml_eid, vec_width=1, dtype=T.i32 - ) + ml_mask = _gld(mask_it, safe_ml_eid) ml_val = ml_valid.select(ml_mask, c_zero_i32) - ml_ix = ArithValue( - ml_valid.select(ml_eid + c_one_i32, c_zero_i32) - ).index_cast(T.index) - _lds_store_raw(cumdup_mr, ml_val, ml_ix) + ml_idx = ml_valid.select(ml_eid + c_one_i32, c_zero_i32) + _lds_store_raw(cumdup_mr, ml_val, ml_idx) _lds_store_raw(cumdup_mr, c_zero_i32, c_zero_i32) gpu.barrier() @@ -566,9 +563,7 @@ def moe_sorting_oneshot_kernel( _lds_store_raw( cumdup_mr, m_tid_valid.select(inclusive_m, c_zero_i32), - ArithValue( - m_tid_valid.select(tid + c_one_i32, c_zero_i32) - ).index_cast(T.index), + m_tid_valid.select(tid + c_one_i32, c_zero_i32), ) gpu.barrier() @@ -587,9 +582,8 @@ def moe_sorting_oneshot_kernel( ml_eid = fx.Int32(i_ml) + tid ml_valid = ml_eid < c_E safe_ml_eid = ml_valid.select(ml_eid, c_zero_i32) - ml_ix = ArithValue(safe_ml_eid).index_cast(T.index) _lds_store_raw( - cumdup_mr, ml_valid.select(safe_ml_eid, c_zero_i32), ml_ix + cumdup_mr, ml_valid.select(safe_ml_eid, c_zero_i32), safe_ml_eid ) gpu.barrier() @@ -600,38 +594,33 @@ def moe_sorting_oneshot_kernel( eid_wr_valid = eid_wr < c_E safe_eid_wr = eid_wr_valid.select(eid_wr, c_zero_i32) - cs_start_ix = ArithValue(safe_eid_wr).index_cast(T.index) - cs_end_ix = ArithValue(safe_eid_wr + c_one_i32).index_cast(T.index) - e_start = _lds_load_raw(cumsum_mr, cs_start_ix) + e_start = _lds_load_raw(cumsum_mr, safe_eid_wr) e_end = eid_wr_valid.select( - _lds_load_raw(cumsum_mr, cs_end_ix), e_start + _lds_load_raw(cumsum_mr, safe_eid_wr + c_one_i32), e_start ) - local_eid = _lds_load_raw(cumdup_mr, cs_start_ix) + local_eid = _lds_load_raw(cumdup_mr, safe_eid_wr) # Store cumdup: reuse cumdup for scatter phase position tracking. # Write e_start to cumdup[eid] (overwriting mask cumsum, no longer needed). - _lds_store_raw(cumdup_mr, e_start, cs_start_ix) + _lds_store_raw(cumdup_mr, e_start, safe_eid_wr) blk_start = e_start // c_unit blk_end = e_end // c_unit n_blks_wr = eid_wr_valid.select(blk_end - blk_start, c_zero_i32) - _write_expert_id_blocks(sorted_e_rsrc, local_eid, blk_start, n_blks_wr) + _write_expert_id_blocks(sorted_e_it, local_eid, blk_start, n_blks_wr) gpu.barrier() # Store cumdup[E] = cumsum[E]. # All threads write cumE to cumdup[E] (all write the same value, no race). - cs_E_ix = ArithValue(c_E).index_cast(T.index) - cumE = _lds_load_raw(cumsum_mr, cs_E_ix) - _lds_store_raw(cumdup_mr, cumE, cs_E_ix) + cumE = _lds_load_raw(cumsum_mr, c_E) + _lds_store_raw(cumdup_mr, cumE, c_E) gpu.barrier() # ====================== PRE-FILL: Sentinel fill (cooperative) =========== - total_padded_pre = _lds_load_raw( - cumdup_mr, ArithValue(c_E).index_cast(T.index) - ) + total_padded_pre = _lds_load_raw(cumdup_mr, c_E) _fill_sentinel_slots( - sorted_ids_rsrc, - sorted_w_rsrc, + sorted_ids_it, + sorted_w_it, c_zero_i32, total_padded_pre, c_sentinel | i32_tokens, @@ -652,16 +641,10 @@ def moe_sorting_oneshot_kernel( sc_expert_enabled = eid_sc_valid if has_mask: # EP: check if this expert is masked (skip scatter for masked experts) - sc_mask_val = buffer_ops.buffer_load( - mask_rsrc, - eid_sc_valid.select(eid_sc, c_zero_i32), - vec_width=1, - dtype=T.i32, - ) + sc_mask_val = _gld(mask_it, eid_sc_valid.select(eid_sc, c_zero_i32)) sc_expert_enabled = eid_sc_valid & (sc_mask_val != c_zero_i32) - cs_sc_ix = ArithValue(safe_eid_sc).index_cast(T.index) - position = _lds_load_raw(cumsum_mr, cs_sc_ix) + position = _lds_load_raw(cumsum_mr, safe_eid_sc) for i_sub2 in range_constexpr(0, sub_tokens, 8): # This lane handles sub_token (i_sub2 + lane_group_os). @@ -669,8 +652,7 @@ def moe_sorting_oneshot_kernel( my_sub_valid = sc_expert_enabled & (my_sub < c_sub_tokens) safe_my_sub = my_sub_valid.select(my_sub, c_zero_i32) my_mesh_addr = safe_my_sub * c_smem_cols + safe_eid_sc - my_mesh_ix = ArithValue(my_mesh_addr).index_cast(T.index) - my_x = _lds_load_raw(mesh_mr, my_mesh_ix) + my_x = _lds_load_raw(mesh_mr, my_mesh_addr) my_has_token = my_sub_valid & (my_x != c_zero_i32) local_cnt = my_has_token.select(c_one_i32, c_zero_i32) @@ -738,22 +720,20 @@ def moe_sorting_oneshot_kernel( topk_slot_sc = safe_x - c_one_i32 packed_id = (topk_slot_sc << fx.Int32(24)) | my_sub safe_slot = my_has_token.select(slot, c_oob_idx) - buffer_ops.buffer_store(packed_id, sorted_ids_rsrc, safe_slot) + _gst(sorted_ids_it, packed_id, safe_slot) w_addr = my_has_token.select( my_sub * c_topk + topk_slot_sc, c_zero_i32 ) - w_val_i32 = buffer_ops.buffer_load( - weights_rsrc, w_addr, vec_width=1, dtype=T.i32 - ) - buffer_ops.buffer_store(w_val_i32, sorted_w_rsrc, safe_slot) + w_val_i32 = _gld(weights_it, w_addr) + _gst(sorted_w_it, w_val_i32, safe_slot) # Advance position by batch total position = position + batch_total # Write back updated position (for padding phase). # Invalid lane groups write position (=0+0=0) to cumsum[0] which is harmless. - _lds_store_raw(cumsum_mr, position, cs_sc_ix) + _lds_store_raw(cumsum_mr, position, safe_eid_sc) gpu.barrier() # Padding already filled by PRE-FILL phase above (before scatter). @@ -848,20 +828,14 @@ def _compile_moe_sorting_multiphase( E = num_experts @flyc.jit - def _extend_local_idx_for_extra_experts( - cumsum_mr, mask_rsrc, K4_BLOCK, E, has_mask - ): + def _extend_local_idx_for_extra_experts(cumsum_mr, mask_it, K4_BLOCK, E, has_mask): """Thread-0: write local expert indices for experts >= K4_BLOCK to cumsum_mr.""" if has_mask: prev_local = _lds_load_raw(cumsum_mr, fx.Int32(K4_BLOCK - 1)) - prev_mask = buffer_ops.buffer_load( - mask_rsrc, fx.Int32(K4_BLOCK - 1), vec_width=1, dtype=T.i32 - ) + prev_mask = _gld(mask_it, fx.Int32(K4_BLOCK - 1)) prev_local = prev_local + prev_mask for _e3 in range_constexpr(K4_BLOCK, E): - e3_mask = buffer_ops.buffer_load( - mask_rsrc, fx.Int32(_e3), vec_width=1, dtype=T.i32 - ) + e3_mask = _gld(mask_it, fx.Int32(_e3)) _lds_store_raw(cumsum_mr, prev_local, fx.Int32(_e3)) prev_local = prev_local + e3_mask else: @@ -872,11 +846,11 @@ def _extend_local_idx_for_extra_experts( def _p23_scatter_mesh( tid, scatter_mr, - ws_rsrc, - weights_rsrc, - sorted_ids_rsrc, - sorted_w_rsrc, - mask_rsrc, + ws_it, + weights_it, + sorted_ids_it, + sorted_w_it, + mask_it, my_expert, my_start, my_end, @@ -894,9 +868,7 @@ def _p23_scatter_mesh( c_ff, c_oob_idx = fx.Int32(0xFF), fx.Int32(0x7FFFFFFF) p23_bid_enabled = c_one != c_zero if has_mask: - p23_bid_mask = buffer_ops.buffer_load( - mask_rsrc, my_expert, vec_width=1, dtype=T.i32 - ) + p23_bid_mask = _gld(mask_it, my_expert) p23_bid_enabled = p23_bid_mask != c_zero i32_words_per_row = i32_scan_words_per_row n_mesh_iters = (my_start != my_end).select( @@ -904,18 +876,16 @@ def _p23_scatter_mesh( ) mesh_row_i32_base = (my_expert * i32_mesh_stride) >> fx.Int32(2) for _si, state in range( - fx.Index(0), - ArithValue(n_mesh_iters).index_cast(T.index), - fx.Index(1), + fx.Int32(0), + n_mesh_iters, + fx.Int32(1), init=[my_start], ): position = state[0] word_idx = fx.Int32(_si) * fx.Int32(K4_BLOCK) + tid col_valid = p23_bid_enabled & (word_idx < i32_words_per_row) safe_word_idx = col_valid.select(word_idx, c_zero) - word = buffer_ops.buffer_load( - ws_rsrc, mesh_row_i32_base + safe_word_idx, vec_width=1, dtype=T.i32 - ) + word = _gld(ws_it, mesh_row_i32_base + safe_word_idx) x0 = word & c_ff x1 = (word >> fx.Int32(8)) & c_ff x2 = (word >> fx.Int32(16)) & c_ff @@ -957,46 +927,38 @@ def _p23_scatter_mesh( safe_slot_2 = h2.select(off2, c_oob_idx) off3 = off2 + h2.select(c_one, c_zero) safe_slot_3 = h3.select(off3, c_oob_idx) - w_val_0 = buffer_ops.buffer_load( - weights_rsrc, + w_val_0 = _gld( + weights_it, h0.select(base_col * c_topk + h0.select(x0 - c_one, c_zero), c_zero), - vec_width=1, - dtype=T.i32, ) - w_val_1 = buffer_ops.buffer_load( - weights_rsrc, + w_val_1 = _gld( + weights_it, h1.select( (base_col + c_one) * c_topk + h1.select(x1 - c_one, c_zero), c_zero ), - vec_width=1, - dtype=T.i32, ) - w_val_2 = buffer_ops.buffer_load( - weights_rsrc, + w_val_2 = _gld( + weights_it, h2.select( (base_col + fx.Int32(2)) * c_topk + h2.select(x2 - c_one, c_zero), c_zero, ), - vec_width=1, - dtype=T.i32, ) - w_val_3 = buffer_ops.buffer_load( - weights_rsrc, + w_val_3 = _gld( + weights_it, h3.select( (base_col + fx.Int32(3)) * c_topk + h3.select(x3 - c_one, c_zero), c_zero, ), - vec_width=1, - dtype=T.i32, ) - buffer_ops.buffer_store(pid_0, sorted_ids_rsrc, safe_slot_0) - buffer_ops.buffer_store(pid_1, sorted_ids_rsrc, safe_slot_1) - buffer_ops.buffer_store(pid_2, sorted_ids_rsrc, safe_slot_2) - buffer_ops.buffer_store(pid_3, sorted_ids_rsrc, safe_slot_3) - buffer_ops.buffer_store(w_val_0, sorted_w_rsrc, safe_slot_0) - buffer_ops.buffer_store(w_val_1, sorted_w_rsrc, safe_slot_1) - buffer_ops.buffer_store(w_val_2, sorted_w_rsrc, safe_slot_2) - buffer_ops.buffer_store(w_val_3, sorted_w_rsrc, safe_slot_3) + _gst(sorted_ids_it, pid_0, safe_slot_0) + _gst(sorted_ids_it, pid_1, safe_slot_1) + _gst(sorted_ids_it, pid_2, safe_slot_2) + _gst(sorted_ids_it, pid_3, safe_slot_3) + _gst(sorted_w_it, w_val_0, safe_slot_0) + _gst(sorted_w_it, w_val_1, safe_slot_1) + _gst(sorted_w_it, w_val_2, safe_slot_2) + _gst(sorted_w_it, w_val_3, safe_slot_3) pos_next = position + batch_total results = yield [pos_next] return results @@ -1012,12 +974,12 @@ def clear_workspace_kernel( i32_total_elems: fx.Int32, ): gid = gpu.block_idx.x * fx.Int32(K1_BLOCK) + gpu.thread_idx.x - ws_rsrc = buffer_ops.create_buffer_resource(workspace, max_size=True) + ws_it = _buf_iter(workspace) c_zero = fx.Int32(0) # Each thread stores exactly one element (no loop needed). valid = gid < i32_total_elems - buffer_ops.buffer_store(c_zero, ws_rsrc, valid.select(gid, c_zero)) + _gst(ws_it, c_zero, valid.select(gid, c_zero)) @flyc.jit def launch_clear_ws( @@ -1047,38 +1009,31 @@ def p0_scatter_kernel( ): gid = gpu.block_idx.x * fx.Int32(K2_BLOCK) + gpu.thread_idx.x stride = gpu.grid_dim.x * fx.Int32(K2_BLOCK) - topk_rsrc = buffer_ops.create_buffer_resource(topk_ids, max_size=True) - ws_rsrc = buffer_ops.create_buffer_resource(workspace, max_size=True) + topk_it = _buf_iter(topk_ids) + # uint8 byte-addressed mesh scatter into the i32 workspace, always under + # an in-bounds `if valid` guard -> a plain i8 global-pointer store. + ws_i8 = _i8_global_ptr(workspace) c_zero = fx.Int32(0) c_topk = fx.Int32(topk) c_one = fx.Int32(1) tokens_ = i32_tokens if has_local_tokens: - ltok_rsrc = buffer_ops.create_buffer_resource( - local_tokens_tensor, max_size=True - ) - tokens_ = buffer_ops.buffer_load( - ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 - ) + ltok_it = _buf_iter(local_tokens_tensor) + tokens_ = _gld(ltok_it, fx.Int32(0)) total = tokens_ * c_topk - _s = fx.Index(0) - _e = ArithValue(i32_niters).index_cast(T.index) - _one = fx.Index(1) - for _i in range(_s, _e, _one): + for _i in range(fx.Int32(0), i32_niters, fx.Int32(1)): flat = gid + fx.Int32(_i) * stride valid = flat < total safe_flat = valid.select(flat, c_zero) token_id = safe_flat // c_topk topk_slot = safe_flat % c_topk - eid = buffer_ops.buffer_load(topk_rsrc, safe_flat, vec_width=1, dtype=T.i32) + eid = _gld(topk_it, safe_flat) byte_offset = eid * i32_mesh_stride + token_id - val_i8 = ArithValue(topk_slot + c_one).trunci(T.i8) + val_i8 = fx.Int32(topk_slot + c_one).to(fx.Int8) if valid: - buffer_ops.buffer_store( - val_i8, ws_rsrc, byte_offset, offset_is_bytes=True - ) + ws_i8[byte_offset] = val_i8 @flyc.jit def launch_p0( @@ -1131,7 +1086,7 @@ def p1_count_kernel( lane = tid % WARP_SIZE wave = tid // WARP_SIZE - ws_rsrc = buffer_ops.create_buffer_resource(workspace, max_size=True) + ws_it = _buf_iter(workspace) c_zero = fx.Int32(0) c_one = fx.Int32(1) c_ff = fx.Int32(0xFF) @@ -1141,12 +1096,8 @@ def p1_count_kernel( # Data-dependent scan tokens_ = i32_tokens if has_local_tokens: - ltok_rsrc = buffer_ops.create_buffer_resource( - local_tokens_tensor, max_size=True - ) - tokens_ = buffer_ops.buffer_load( - ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 - ) + ltok_it = _buf_iter(local_tokens_tensor) + tokens_ = _gld(ltok_it, fx.Int32(0)) mesh_row_i32_base = (eid * i32_mesh_stride) >> fx.Int32(2) i32_scan_words_per_row = (tokens_ + fx.Int32(3)) >> fx.Int32(2) @@ -1155,23 +1106,21 @@ def p1_count_kernel( ) >> fx.Int32(K3_WORDS_PER_ITER_LOG2) if has_mask: - mask_rsrc = buffer_ops.create_buffer_resource( - expert_mask_tensor, max_size=True - ) - p1_mask = buffer_ops.buffer_load(mask_rsrc, eid, vec_width=1, dtype=T.i32) + mask_it = _buf_iter(expert_mask_tensor) + p1_mask = _gld(mask_it, eid) p1_is_local = p1_mask != c_zero p1_should_zero = (~p1_is_local) & (tid == c_zero) - buffer_ops.buffer_store( + _gst( + ws_it, c_zero, - ws_rsrc, p1_should_zero.select(i32_mesh_size + eid, fx.Int32(0x7FFFFFFF)), ) n_iters = p1_is_local.select(n_iters, c_zero) for _i, state in range( - fx.Index(0), - ArithValue(n_iters).index_cast(T.index), - fx.Index(1), + fx.Int32(0), + n_iters, + fx.Int32(1), init=[c_zero], ): cnt_so_far = state[0] @@ -1181,7 +1130,7 @@ def p1_count_kernel( ) valid = word_base < i32_scan_words_per_row safe_addr = mesh_row_i32_base + valid.select(word_base, c_zero) - vec4 = buffer_ops.buffer_load(ws_rsrc, safe_addr, vec_width=4, dtype=T.i32) + vec4 = fx.ptr_load(ws_it + fx.Int64(safe_addr), result_type=T.i32x4) iter_cnt = c_zero for _wi in range_constexpr(K3_VEC_WIDTH): @@ -1213,8 +1162,7 @@ def p1_count_kernel( # Cross-warp reduce via LDS: lane 0 of each warp writes partial sum is_lane0 = lane == c_zero if is_lane0: - wave_ix = ArithValue(wave).index_cast(T.index) - _lds_store_raw(reduce_mr, cnt, wave_ix) + _lds_store_raw(reduce_mr, cnt, wave) gpu.barrier() # Thread 0 sums all warp partials and writes to HBM @@ -1226,7 +1174,7 @@ def p1_count_kernel( cs_offset = i32_mesh_size + eid c_oob_idx = fx.Int32(0x7FFFFFFF) safe_cs = is_t0.select(cs_offset, c_oob_idx) - buffer_ops.buffer_store(total, ws_rsrc, safe_cs) + _gst(ws_it, total, safe_cs) @flyc.jit def launch_p1( @@ -1282,9 +1230,12 @@ def p0v2_kernel( lane = tid % WARP_SIZE wave = tid // WARP_SIZE - ws_rsrc = buffer_ops.create_buffer_resource(workspace, max_size=True) - mask_rsrc = buffer_ops.create_buffer_resource(expert_mask_tensor, max_size=True) - topk_rsrc = buffer_ops.create_buffer_resource(topk_ids, max_size=True) + ws_it = _buf_iter(workspace) + # uint8 byte-addressed scatter into the i32 workspace, always under an + # in-bounds `if is_mine` guard (see Phase 2) -> a plain i8 pointer store. + ws_i8 = _i8_global_ptr(workspace) + mask_it = _buf_iter(expert_mask_tensor) + topk_it = _buf_iter(topk_ids) c_zero = fx.Int32(0) c_oob = fx.Int32(0x7FFFFFFF) c_one = fx.Int32(1) @@ -1300,12 +1251,8 @@ def p0v2_kernel( tokens_ = i32_tokens if has_local_tokens: - ltok_rsrc = buffer_ops.create_buffer_resource( - local_tokens_tensor, max_size=True - ) - tokens_ = buffer_ops.buffer_load( - ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 - ) + ltok_it = _buf_iter(local_tokens_tensor) + tokens_ = _gld(ltok_it, fx.Int32(0)) # Phase 3 (count) only needs to scan words that can hold real mesh # bytes -- the first tokens_ columns (dynamic per-call count), never @@ -1324,31 +1271,25 @@ def p0v2_kernel( is_local_expert = c_one != c_zero # EP: load mask, write cumsum=0 for masked experts, set loop bounds to 0 if has_mask: - m_val = buffer_ops.buffer_load(mask_rsrc, eid, vec_width=1, dtype=T.i32) + m_val = _gld(mask_it, eid) is_local_expert = m_val != c_zero should_write_zero = (~is_local_expert) & (tid == c_zero) - buffer_ops.buffer_store( - c_zero, ws_rsrc, should_write_zero.select(i32_mesh_size + eid, c_oob) - ) + _gst(ws_it, c_zero, should_write_zero.select(i32_mesh_size + eid, c_oob)) clear_niters = is_local_expert.select(clear_niters, c_zero) scatter_niters = is_local_expert.select(scatter_niters, c_zero) count_niters = is_local_expert.select(count_niters, c_zero) # ---- Phase 1: Clear this expert's mesh row ---- - for _ci in range( - fx.Index(0), ArithValue(clear_niters).index_cast(T.index), fx.Index(1) - ): + for _ci in range(fx.Int32(0), clear_niters, fx.Int32(1)): word_idx = fx.Int32(_ci) * c_block + tid valid = word_idx < i32_words_per_row safe_idx = mesh_row_i32_base + valid.select(word_idx, c_zero) - buffer_ops.buffer_store(c_zero, ws_rsrc, valid.select(safe_idx, c_oob)) + _gst(ws_it, c_zero, valid.select(safe_idx, c_oob)) gpu.barrier() # ---- Phase 2: Scatter (scan all T*topk, filter by expert) ---- - for _si in range( - fx.Index(0), ArithValue(scatter_niters).index_cast(T.index), fx.Index(1) - ): + for _si in range(fx.Int32(0), scatter_niters, fx.Int32(1)): flat = fx.Int32(_si) * c_block + tid valid = flat < total_assignments safe_flat = valid.select(flat, c_zero) @@ -1364,27 +1305,23 @@ def p0v2_kernel( else safe_flat % c_topk ) - expert_id = buffer_ops.buffer_load( - topk_rsrc, safe_flat, vec_width=1, dtype=T.i32 - ) + expert_id = _gld(topk_it, safe_flat) is_mine = valid & (expert_id == eid) byte_offset = eid * i32_mesh_stride + token_id - val_i8 = ArithValue(is_mine.select(topk_slot + c_one, c_zero)).trunci(T.i8) - # Byte-mode buffer_store with OOB offset crashes on AMD GPUs. - # Use conditional branch to skip the store for non-matching threads. + val_i8 = fx.Int32(is_mine.select(topk_slot + c_one, c_zero)).to(fx.Int8) + # Guarded byte store: skip non-matching threads so no OOB byte offset + # is ever formed (a raw byte store would fault, not drop, at OOB). if is_mine: - buffer_ops.buffer_store( - val_i8, ws_rsrc, byte_offset, offset_is_bytes=True - ) + ws_i8[byte_offset] = val_i8 gpu.barrier() # ---- Phase 3: Count non-zero bytes + warp/cross-wave reduce ---- for _ki, state in range( - fx.Index(0), - ArithValue(count_niters).index_cast(T.index), - fx.Index(1), + fx.Int32(0), + count_niters, + fx.Int32(1), init=[c_zero], ): cnt_so_far = state[0] @@ -1392,7 +1329,7 @@ def p0v2_kernel( word_base = fx.Int32(_ki) * c_block + tid valid = word_base < i32_scan_words_per_row safe_addr = mesh_row_i32_base + valid.select(word_base, c_zero) - word = buffer_ops.buffer_load(ws_rsrc, safe_addr, vec_width=1, dtype=T.i32) + word = _gld(ws_it, safe_addr) b0 = word & c_ff b1 = (word >> fx.Int32(8)) & c_ff @@ -1418,8 +1355,7 @@ def p0v2_kernel( # Cross-warp reduce via LDS: lane 0 of each warp writes partial sum is_lane0 = lane == c_zero if is_lane0: - wave_ix = ArithValue(wave).index_cast(T.index) - _lds_store_raw(reduce_mr, cnt, wave_ix) + _lds_store_raw(reduce_mr, cnt, wave) gpu.barrier() # Thread 0 sums all warp partials and writes to HBM @@ -1431,7 +1367,7 @@ def p0v2_kernel( cs_offset = i32_mesh_size + eid c_oob_idx = fx.Int32(0x7FFFFFFF) safe_cs = is_t0.select(cs_offset, c_oob_idx) - buffer_ops.buffer_store(total, ws_rsrc, safe_cs) + _gst(ws_it, total, safe_cs) @flyc.jit def launch_p0v2( @@ -1500,18 +1436,12 @@ def p23_kernel( c_sentinel = fx.Int32(topk << 24) c_oob_idx = fx.Int32(0x7FFFFFFF) - # Buffer resources - ws_rsrc = buffer_ops.create_buffer_resource(workspace, max_size=True) - weights_rsrc = buffer_ops.create_buffer_resource( - topk_weights_tensor, max_size=True - ) - sorted_ids_rsrc = buffer_ops.create_buffer_resource( - sorted_token_ids, max_size=True - ) - sorted_w_rsrc = buffer_ops.create_buffer_resource( - sorted_weights_out, max_size=True - ) - mask_rsrc = buffer_ops.create_buffer_resource(expert_mask_tensor, max_size=True) + # Buffer iterators (OOB-checked V# descriptors). + ws_it = _buf_iter(workspace) + weights_it = _buf_iter(topk_weights_tensor) + sorted_ids_it = _buf_iter(sorted_token_ids) + sorted_w_it = _buf_iter(sorted_weights_out) + mask_it = _buf_iter(expert_mask_tensor) # LDS: cumsum[E+1] for prefix sums + cross-wave scratch lds = fx.SharedAllocator().allocate(K4SharedStorage).peek() @@ -1523,11 +1453,11 @@ def p23_kernel( # ================ MOE_BUF ZEROING (blocks >= E) ================== if is_zero_block: - moe_buf_rsrc = buffer_ops.create_buffer_resource(moe_buf, max_size=True) + moe_buf_it = _buf_iter(moe_buf) zero_gid_v4 = (bid - c_E) * fx.Int32(K4_BLOCK) + tid zero_stride_v4 = (gpu.grid_dim.x - c_E) * fx.Int32(K4_BLOCK) _zero_moe_buf_grid_stride( - moe_buf_rsrc, + moe_buf_it, zero_gid_v4, zero_stride_v4, i32_moe_buf_elems >> fx.Int32(2), @@ -1541,12 +1471,8 @@ def p23_kernel( tokens_ = i32_tokens if has_local_tokens: - ltok_rsrc = buffer_ops.create_buffer_resource( - local_tokens_tensor, max_size=True - ) - tokens_ = buffer_ops.buffer_load( - ltok_rsrc, fx.Int32(0), vec_width=1, dtype=T.i32 - ) + ltok_it = _buf_iter(local_tokens_tensor) + tokens_ = _gld(ltok_it, fx.Int32(0)) # Step 1: Load expert counts from workspace -> pad to unit_size -> LDS cumsum # Process E experts in chunks of K4_BLOCK (256). Most models have @@ -1561,30 +1487,20 @@ def p23_kernel( my_mask_val = c_one if has_mask: tid_has_expert = tid < c_E - my_mask_val = buffer_ops.buffer_load( - mask_rsrc, - tid_has_expert.select(tid, c_zero), - vec_width=1, - dtype=T.i32, - ) + my_mask_val = _gld(mask_it, tid_has_expert.select(tid, c_zero)) my_mask_val = tid_has_expert.select(my_mask_val, c_zero) for _chunk in range_constexpr(0, E, K4_BLOCK): expert_idx = fx.Int32(_chunk) + tid tid_valid_expert = expert_idx < c_E ws_cs_addr = i32_mesh_size + tid_valid_expert.select(expert_idx, c_zero) - raw_cnt = buffer_ops.buffer_load( - ws_rsrc, ws_cs_addr, vec_width=1, dtype=T.i32 - ) + raw_cnt = _gld(ws_it, ws_cs_addr) raw_cnt = tid_valid_expert.select(raw_cnt, c_zero) blocks = (raw_cnt + c_unit - c_one) // c_unit padded = (raw_cnt == c_zero).select(c_zero, blocks * c_unit) if has_mask: - chunk_mask = buffer_ops.buffer_load( - mask_rsrc, - tid_valid_expert.select(expert_idx, c_zero), - vec_width=1, - dtype=T.i32, + chunk_mask = _gld( + mask_it, tid_valid_expert.select(expert_idx, c_zero) ) chunk_mask = tid_valid_expert.select(chunk_mask, c_zero) padded = (chunk_mask == c_zero).select(c_zero, padded) @@ -1633,11 +1549,9 @@ def p23_kernel( # Block 0, thread 0 writes num_valid_ids if (bid == c_zero) & (tid == c_zero): - nvalid_rsrc = buffer_ops.create_buffer_resource( - num_valid_ids, max_size=True - ) - buffer_ops.buffer_store(total_padded, nvalid_rsrc, c_zero) - buffer_ops.buffer_store(tokens_, nvalid_rsrc, c_one) + nvalid_it = _buf_iter(num_valid_ids) + _gst(nvalid_it, total_padded, c_zero) + _gst(nvalid_it, tokens_, c_one) # Step 3: Write sorted_expert_ids for THIS expert (using local_idx_p23 for EP) # Store local_idx to LDS cumsum[tid], barrier, read cumsum[my_expert] @@ -1648,18 +1562,16 @@ def p23_kernel( gpu.barrier() if tid == c_zero: _extend_local_idx_for_extra_experts( - cumsum_mr, mask_rsrc, K4_BLOCK, E, has_mask + cumsum_mr, mask_it, K4_BLOCK, E, has_mask ) gpu.barrier() my_local_idx = _lds_load_raw(cumsum_mr, my_expert) - sorted_e_rsrc = buffer_ops.create_buffer_resource( - sorted_expert_ids, max_size=True - ) + sorted_e_it = _buf_iter(sorted_expert_ids) blk_start = my_start // c_unit blk_end = my_end // c_unit _write_expert_id_blocks( - sorted_e_rsrc, my_local_idx, blk_start, blk_end - blk_start + sorted_e_it, my_local_idx, blk_start, blk_end - blk_start ) # Step 4: Mesh-based scatter (EP mask + uint8 mesh read + DPP prefix sum + scatter) @@ -1668,11 +1580,11 @@ def p23_kernel( scatter_end_pos_t0 = _p23_scatter_mesh( tid, scatter_mr, - ws_rsrc, - weights_rsrc, - sorted_ids_rsrc, - sorted_w_rsrc, - mask_rsrc, + ws_it, + weights_it, + sorted_ids_it, + sorted_w_it, + mask_it, my_expert, my_start, my_end, @@ -1685,8 +1597,8 @@ def p23_kernel( # Step 5: Fill padding with sentinel for THIS expert (parallel) _fill_sentinel_slots( - sorted_ids_rsrc, - sorted_w_rsrc, + sorted_ids_it, + sorted_w_it, scatter_end_pos_t0, my_end - scatter_end_pos_t0, c_sentinel | i32_tokens, diff --git a/aiter/ops/flydsl/kernels/qk_norm_rope_quant.py b/aiter/ops/flydsl/kernels/qk_norm_rope_quant.py index c2befff9ec..fba15e52a5 100644 --- a/aiter/ops/flydsl/kernels/qk_norm_rope_quant.py +++ b/aiter/ops/flydsl/kernels/qk_norm_rope_quant.py @@ -18,7 +18,6 @@ import flydsl.compiler as flyc import flydsl.expr as fx import torch -from flydsl._mlir.dialects import llvm, rocdl from flydsl.expr import const_expr, ptrtoint, range_constexpr from flydsl.expr import math as fmath from flydsl.expr.arith import FastMathFlags @@ -50,12 +49,10 @@ def _imax(a, b): def _idiv(a, b): - """Truncating signed integer divide. ``//`` maps to arith.floordivsi, a - longer expansion; every dividend here is already clamped non-negative, so - the truncating op is both correct and cheaper.""" - from flydsl.expr import arith - - return fx.Int32(arith.divsi(a.ir_value(), b.ir_value())) + """Truncating integer divide. Signed ``//`` maps to arith.floordivsi, a + longer expansion; every dividend here is provably non-negative, so an + unsigned divide (divui) is both correct and cheaper.""" + return fx.Int32(fx.Uint32(a) // fx.Uint32(b)) _STATIC_ADAPTOR_CACHE = {} @@ -173,7 +170,7 @@ def _scalar_load(base_i64, idx, fx_dt, copy_bits): buf = _scalar_view(base_i64, fx_dt) atom = fx.make_copy_atom(fx.rocdl.BufferCopy(copy_bits), fx_dt) r = fx.make_rmem_tensor(fx.make_layout(1, 1), fx_dt) - fx.copy_atom_call(atom, fx.slice(buf, (idx, None)), r) + fx.copy(atom, fx.slice(buf, (idx, None)), r) return fx.memref_load_vec(r)[0] @@ -182,14 +179,13 @@ def _scalar_store(base_i64, idx, val, fx_dt, copy_bits): buf = _scalar_view(base_i64, fx_dt) atom = fx.make_copy_atom(fx.rocdl.BufferCopy(copy_bits), fx_dt) r = fx.make_rmem_tensor(fx.make_layout(1, 1), fx_dt) - fx.memref_store_vec(fx.Vector.from_elements([val.ir_value()], dtype=fx_dt), r) - fx.copy_atom_call(atom, r, fx.slice(buf, (idx, None))) + fx.memref_store_vec(fx.Vector.from_elements([val], dtype=fx_dt), r) + fx.copy(atom, r, fx.slice(buf, (idx, None))) def _store_bf16_tiled(vals_list, p_dst, copy, vec): """Convert VEC fp32 → bf16 and tiled-copy to ``p_dst``.""" - raw = [v.ir_value() if hasattr(v, "ir_value") else v for v in vals_list] - f32v = fx.Vector.from_elements(raw, dtype=fx.Float32) + f32v = fx.Vector.from_elements(vals_list, dtype=fx.Float32) bf16v = f32v.truncf(T.vec(vec, T.bf16)) frag = fx.make_fragment_like(p_dst) fx.memref_store_vec(bf16v, frag) @@ -288,8 +284,8 @@ def _store_fp8_packed_w32( for dw_idx in range_constexpr(n_dwords): base = dw_idx * 4 pk = fx.Int32(0).ir_value() - pk = rocdl.cvt_pk_fp8_f32(i32, safe[base + 0], safe[base + 1], pk, 0) - pk = rocdl.cvt_pk_fp8_f32(i32, safe[base + 2], safe[base + 3], pk, 1) + pk = fx.rocdl.cvt_pk_fp8_f32(i32, safe[base + 0], safe[base + 1], pk, 0) + pk = fx.rocdl.cvt_pk_fp8_f32(i32, safe[base + 2], safe[base + 3], pk, 1) dword_list.append(pk) off_bytes = row_base_bytes + idx * vec @@ -405,7 +401,7 @@ def load_vec( div_tensor, idx, *, layout=full_lay, atom=full_atom, dt=elem_dtype ): r = fx.make_rmem_tensor(layout, dt) - fx.copy_atom_call(atom, fx.slice(div_tensor, (None, idx)), r) + fx.copy(atom, fx.slice(div_tensor, (None, idx)), r) return fx.memref_load_vec(r) bid_x = fx.block_idx.x # 0..H-1 (Q head) or H (KV) @@ -534,8 +530,7 @@ def emit_body( scaled.append(xi * rstd) out_rmem = fx.make_rmem_tensor(full_lay, fx.Float32) - scaled_raw = [s.ir_value() for s in scaled] - scaled_vec = fx.Vector.from_elements(scaled_raw, dtype=fx.Float32) + scaled_vec = fx.Vector.from_elements(scaled, dtype=fx.Float32) fx.memref_store_vec(scaled_vec, out_rmem) is_rope = tid >= fx.Int32(ROPE_THREAD_LO) @@ -553,8 +548,8 @@ def emit_body( o = cur[2 * k + 1] c = cos_f32[k] s = sin_f32[k] - rope_elems.append((e * c - o * s).ir_value()) - rope_elems.append((e * s + o * c).ir_value()) + rope_elems.append(e * c - o * s) + rope_elems.append(e * s + o * c) rotated_vec = fx.Vector.from_elements(rope_elems, dtype=fx.Float32) fx.memref_store_vec(rotated_vec, out_rmem) @@ -938,10 +933,8 @@ def _load_weight_tensor(weight_tensor, tid_val): wdiv = fx.logical_divide(wbuf, half_lay) r0 = fx.make_rmem_tensor(half_lay, elem_dtype) r1 = fx.make_rmem_tensor(half_lay, elem_dtype) - fx.copy_atom_call(half_atom, fx.slice(wdiv, (None, tid_val * 2)), r0) - fx.copy_atom_call( - half_atom, fx.slice(wdiv, (None, tid_val * 2 + 1)), r1 - ) + fx.copy(half_atom, fx.slice(wdiv, (None, tid_val * 2)), r0) + fx.copy(half_atom, fx.slice(wdiv, (None, tid_val * 2 + 1)), r1) v0 = fx.memref_load_vec(r0).to(fx.Float32) v1 = fx.memref_load_vec(r1).to(fx.Float32) combined = v0.shuffle(v1, list(range(VEC))) @@ -982,7 +975,7 @@ def _load_bf16_raw(rsrc, off_dw): tok = _imin(bid_t * ROWS_PER_WG + tid_y, num_tokens - 1) bid_t = tok # all downstream token offsets use the clamped token - bid_t_idx = fx.Index(tok) + bid_t_idx = fx.Int64(tok) def _ptr_buffer_resource(ptr, num_records_bytes=None): addr = fx.ptrtoint(ptr) @@ -1034,8 +1027,8 @@ def emit_body( rope_rel = _imax(tid - fx.Int32(ROPE_THREAD_LO), fx.Int32(0)) cos_rmem = fx.make_rmem_tensor(rope_lay, elem_dtype) sin_rmem = fx.make_rmem_tensor(rope_lay, elem_dtype) - fx.copy_atom_call(rope_atom, fx.slice(cos_div, (None, rope_rel)), cos_rmem) - fx.copy_atom_call(rope_atom, fx.slice(sin_div, (None, rope_rel)), sin_rmem) + fx.copy(rope_atom, fx.slice(cos_div, (None, rope_rel)), cos_rmem) + fx.copy(rope_atom, fx.slice(sin_div, (None, rope_rel)), sin_rmem) cos_raw = fx.memref_load_vec(cos_rmem) sin_raw = fx.memref_load_vec(sin_rmem) @@ -1084,9 +1077,7 @@ def emit_body( factor = rstd * quant_scale scale_store = e8m0_biased.to(fx.Int8) else: - rcp_am = llvm.call_intrinsic( - f32, "llvm.amdgcn.rcp.f32", [am_safe.ir_value()], [], [] - ) + rcp_am = fx.Float32(fx.rocdl.rcp(f32, am_safe)) _fc = _fp8_const() factor = fx.Float32(_fc["max_over_sqrt2"]) * rcp_am scale_store = am_safe * rstd * fx.Float32(_fc["inv_max_sqrt2"]) @@ -1147,7 +1138,7 @@ def emit_body( VEC, ) - q_tok_off_bytes = bid_t_idx * fx.Index(H * D * 2) + q_tok_off_bytes = bid_t_idx * fx.Int64(H * D * 2) if bid_x < H: head_idx = bid_x @@ -1167,7 +1158,7 @@ def emit_body( qw_f32 = None if const_expr(quant): - q_tok_off_fp8 = bid_t_idx * fx.Index(H * D) + q_tok_off_fp8 = bid_t_idx * fx.Int64(H * D) qo_g_tmp = GTensor( q_out, dtype=T.i8, @@ -1222,7 +1213,7 @@ def emit_body( w_f32 = w_vec.to(fx.Float32) if const_expr(quant): - kv_tok_off_fp8 = bid_t_idx * fx.Index(D) + kv_tok_off_fp8 = bid_t_idx * fx.Int64(D) kvo_g_tmp = GTensor( kv_out, dtype=T.i8, @@ -1244,7 +1235,7 @@ def emit_body( scale_base_off=scale_base_off_kv, ) else: - kv_tok_off_bf16 = bid_t_idx * fx.Index(D * 2) + kv_tok_off_bf16 = bid_t_idx * fx.Int64(D * 2) kvo_g_tmp = GTensor( kv_out, dtype=T.bf16, @@ -1367,7 +1358,7 @@ def launch_qk_norm_rope_quant( num_tokens: fx.Int32, stream: fx.Stream, ): - grid_y = fx.Index(_idiv(num_tokens + ROWS_PER_WG - 1, fx.Int32(ROWS_PER_WG))) + grid_y = fx.Int64(_idiv(num_tokens + ROWS_PER_WG - 1, fx.Int32(ROWS_PER_WG))) k = kernel( q_in, kv_in, @@ -2010,9 +2001,7 @@ def norm_rope_store( scale_store = e8m0_biased.to(fx.Int8) elif const_expr(quant): am_safe = am.maximumf(fx.Float32(1e-12)) - rcp_am = llvm.call_intrinsic( - T.f32, "llvm.amdgcn.rcp.f32", [am_safe.ir_value()], [], [] - ) + rcp_am = fx.Float32(fx.rocdl.rcp(T.f32, am_safe)) fc = _fp8_const() factor = fx.Float32(fc["max_over_sqrt2"]) * rcp_am scaled = xw * factor @@ -2175,7 +2164,7 @@ def cs_of(tile_idx): # not landed; it stays latent only while the per-tile compute # happens to outlast the load. tdm_ops.tensor_wait(min(K - 1, CT - 1 - i)) - rocdl.s_wait_dscnt(0) + fx.rocdl.s_wait_dscnt(0) if const_expr(do_hoist): cosv, sinv = cs_cache[0], cs_cache[1] else: @@ -2332,7 +2321,7 @@ def launch( swa_cache_size, ) k.launch( - grid=(fx.Index(gx), 1, 1), + grid=(fx.Int64(gx), 1, 1), block=(BLOCK_THREADS, RT, 1), stream=stream, )