diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index c8b5a7b47..7218b6997 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -22,7 +22,6 @@ from flydsl._mlir.dialects.fly_rocdl import TargetAddressSpace as _TargetAddressSpace from flydsl.compiler.ast_rewriter import ReplaceIfWithDispatch from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl -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 _to_raw as as_mlir_value @@ -122,18 +121,6 @@ def _ds_read_tr8_b64_imm(result_type, addr_i32, imm_offset=0): # Arithmetic and inline-asm primitives -def _fadd(a, b, fm_fast): - return arith.addf(as_mlir_value(a), as_mlir_value(b), fastmath=fm_fast) - - -def _fsub(a, b, fm_fast): - return arith.subf(as_mlir_value(a), as_mlir_value(b), fastmath=fm_fast) - - -def _fmul(a, b, fm_fast): - return arith.mulf(as_mlir_value(a), as_mlir_value(b), fastmath=fm_fast) - - def _tree_reduce(vals, binop): items = list(vals) while len(items) > 1: @@ -144,10 +131,6 @@ def _tree_reduce(vals, binop): return items[0] -def _fmax(a, b, fm_fast): - return arith.MaxNumFOp(as_mlir_value(a), as_mlir_value(b), fastmath=fm_fast).result - - def _mfma_acc(a, b, c, _mma_atom, mfma_acc_vec_type): return fly.mma_atom_call_ssa([mfma_acc_vec_type], _mma_atom, a, b, c) @@ -352,11 +335,7 @@ def _scale_v_p(traits, v_p, scale_scalar, elem_dtype, fm_fast): p_all_f32_op = llvm.FPExtOp(v32f32_type, as_mlir_value(p_all)) p_all_f32_op.operation.attributes["fastmathFlags"] = fm_fast_attr scale_vec = Vec.from_elements([scale_scalar], fx.Float32).broadcast_to(traits.PV_K_STEPS * 2 * 8) - p_scaled_f32 = arith.mulf( - as_mlir_value(scale_vec), - as_mlir_value(p_all_f32_op.result), - fastmath=fm_fast, - ) + p_scaled_f32 = as_mlir_value(scale_vec * Vec(p_all_f32_op.result)) p_scaled_bf16_op = llvm.FPTruncOp(v32bf16_type, p_scaled_f32) p_scaled_bf16_op.operation.attributes["fastmathFlags"] = fm_fast_attr return _v_vec32_to_p(traits, p_scaled_bf16_op.result, elem_dtype=elem_dtype) @@ -407,13 +386,15 @@ def _lane_pair_reduce(v, reducer, fm_fast): def _score_pair_max(v_s, neg_inf, fm_fast): - return _lane_pair_reduce(_reduce_score_pair(v_s, neg_inf, _fmax, fm_fast), _fmax, fm_fast) + reducer = lambda a, b, _fm: fx.maxnumf(a, b) # noqa: E731 + return _lane_pair_reduce(_reduce_score_pair(v_s, neg_inf, reducer, fm_fast), reducer, fm_fast) def _score_pair_sum(v_s, zero_f, fm_fast): s_lo, s_hi = _score_lists_to_vecs(v_s) tile = Vec(s_lo) + Vec(s_hi) - return _lane_pair_reduce(tile.reduce("add", init_val=zero_f, fastmath=fm_fast), _fadd, fm_fast) + reducer = lambda a, b, _fm: a + b # noqa: E731 + return _lane_pair_reduce(tile.reduce("add", init_val=zero_f, fastmath=fm_fast), reducer, fm_fast) def _sub_score_pair(v_s, row_max, fm_fast): @@ -421,9 +402,9 @@ def _sub_score_pair(v_s, row_max, fm_fast): lo_sub = [] hi_sub = [] for r in range_constexpr(16): - lo_sub.append(_fsub(s_lo[r], row_max, fm_fast)) + lo_sub.append(s_lo[r] - row_max) for r in range_constexpr(16): - hi_sub.append(_fsub(s_hi[r], row_max, fm_fast)) + hi_sub.append(s_hi[r] - row_max) return Vec.from_elements(lo_sub, fx.Float32).ir_value(), Vec.from_elements(hi_sub, fx.Float32).ir_value() @@ -437,11 +418,11 @@ def _scale_sub_score_pair(v_s, row_max_raw, scale, zero_f, fm_fast): ``-inf`` masked lanes stay ``-inf`` (scale > 0), matching the un-fused path. """ s_lo, s_hi = v_s - neg_scaled_max = _fsub(zero_f, _fmul(scale, row_max_raw, fm_fast), fm_fast) + neg_scaled_max = zero_f - scale * row_max_raw scale_v = Vec.from_elements([scale], fx.Float32).broadcast_to(16) nsm_v = Vec.from_elements([neg_scaled_max], fx.Float32).broadcast_to(16) - lo = fmath.fma(Vec(s_lo), scale_v, nsm_v, fastmath=fm_fast) - hi = fmath.fma(Vec(s_hi), scale_v, nsm_v, fastmath=fm_fast) + lo = fx.fma(Vec(s_lo), scale_v, nsm_v, fastmath=fm_fast) + hi = fx.fma(Vec(s_hi), scale_v, nsm_v, fastmath=fm_fast) return as_mlir_value(lo), as_mlir_value(hi) @@ -479,15 +460,15 @@ def _safe_l_inv(l_row, zero_f): def _rescale_from_tile_max(m_row, m_tile_max, fm_fast): - row_max = _fmax(m_row, m_tile_max, fm_fast) - rescale = rocdl.exp2(T.f32, as_mlir_value(_fsub(m_row, row_max, fm_fast))) + row_max = fx.maxnumf(m_row, m_tile_max) + rescale = rocdl.exp2(T.f32, as_mlir_value(m_row - row_max)) return row_max, rescale def _scale_o_accs(v_o, scale_scalar, traits, fm_fast): scale_vec = Vec.from_elements([scale_scalar], fx.Float32).broadcast_to(16) for dc in range_constexpr(traits.D_CHUNKS): - v_o[dc] = _fmul(Vec(v_o[dc]), scale_vec, fm_fast) + v_o[dc] = Vec(v_o[dc]) * scale_vec def _causal_pair_thresholds(kv_vectorized): @@ -947,10 +928,10 @@ def _init_dualwave_thread_mapping(ctx): ctx.lane_mod_32 = ctx.lane % 32 ctx.lane_div_32 = ctx.lane // 32 - _tid_i32 = as_mlir_value(fx.Int32(ctx.tid)) + _tid_i32 = fx.Int32(ctx.tid) _wave_id_uni_i32 = rocdl.readfirstlane( T.i32, - arith.divsi(_tid_i32, as_mlir_value(fx.Int32(traits.WARP_SIZE))), + (_tid_i32 // fx.Int32(traits.WARP_SIZE)).ir_value(), ) ctx.stagger_i32 = arith.divsi(_wave_id_uni_i32, as_mlir_value(fx.Int32(4))) ctx.wave_id_uni = fx.Index(_wave_id_uni_i32) @@ -2898,51 +2879,49 @@ def _exp2(self, x): def online_softmax_stats(self, m_running, s_raw_lo, s_raw_hi): ctx = self.ctx traits = ctx.traits - fm_fast = ctx.fm_fast if const_expr(os.getenv("FLYDSL_FLASH_ATTN_FUNC_TREE_REDUCE", "0") == "1"): def _max_pair(a, b): - return _fmax(a, b, fm_fast) + return fx.maxnumf(a, b) local_max = _tree_reduce(list(s_raw_lo) + list(s_raw_hi), _max_pair) else: local_max = s_raw_lo[0] for r in range_constexpr(15): - local_max = _fmax(local_max, s_raw_lo[r + 1], fm_fast) + local_max = fx.maxnumf(local_max, s_raw_lo[r + 1]) for r in range_constexpr(16): - local_max = _fmax(local_max, s_raw_hi[r], fm_fast) - row_max = _fmax(local_max, self.reduction_peer(local_max), fm_fast) - m_new_raw = _fmax(m_running, row_max, fm_fast) + local_max = fx.maxnumf(local_max, s_raw_hi[r]) + row_max = fx.maxnumf(local_max, self.reduction_peer(local_max)) + m_new_raw = fx.maxnumf(m_running, row_max) if const_expr(traits.CAUSAL): - m_new_raw = _fmax(m_new_raw, ctx.c_neg_floor, fm_fast) + m_new_raw = fx.maxnumf(m_new_raw, ctx.c_neg_floor) - diff_m_scaled = _fmul(_fsub(m_running, m_new_raw, fm_fast), ctx.c_sm_scale_log2e, fm_fast) + diff_m_scaled = (m_running - m_new_raw) * ctx.c_sm_scale_log2e corr = self._exp2(diff_m_scaled) - neg_scaled_max = _fsub(ctx.c_zero_f, _fmul(ctx.c_sm_scale_log2e, m_new_raw, fm_fast), fm_fast) + neg_scaled_max = ctx.c_zero_f - ctx.c_sm_scale_log2e * m_new_raw return m_new_raw, corr, neg_scaled_max def online_softmax(self, m_running, l_running, s_raw_lo, s_raw_hi): ctx = self.ctx - fm_fast = ctx.fm_fast m_new_raw, corr, neg_scaled_max = self.online_softmax_stats(m_running, s_raw_lo, s_raw_hi) p_vals_lo = [] p_vals_hi = [] local_sum = ctx.c_zero_f for r in range_constexpr(16): - diff_lo = fmath.fma(s_raw_lo[r], ctx.c_sm_scale_log2e, neg_scaled_max, fastmath=ctx.fm_fast) + diff_lo = fx.fma(s_raw_lo[r], ctx.c_sm_scale_log2e, neg_scaled_max, fastmath=ctx.fm_fast) p_lo = self._exp2(diff_lo) p_vals_lo.append(p_lo) - local_sum = _fadd(local_sum, p_lo, fm_fast) + local_sum = local_sum + p_lo for r in range_constexpr(16): - diff_hi = fmath.fma(s_raw_hi[r], ctx.c_sm_scale_log2e, neg_scaled_max, fastmath=ctx.fm_fast) + diff_hi = fx.fma(s_raw_hi[r], ctx.c_sm_scale_log2e, neg_scaled_max, fastmath=ctx.fm_fast) p_hi = self._exp2(diff_hi) p_vals_hi.append(p_hi) - local_sum = _fadd(local_sum, p_hi, fm_fast) + local_sum = local_sum + p_hi - tile_sum = _fadd(local_sum, self.reduction_peer(local_sum), fm_fast) - l_new = _fadd(_fmul(corr, l_running, fm_fast), tile_sum, fm_fast) + tile_sum = local_sum + self.reduction_peer(local_sum) + l_new = corr * l_running + tile_sum return m_new_raw, l_new, corr, p_vals_lo, p_vals_hi def rescale_o_accs(self, o_accs, corr): @@ -2950,10 +2929,10 @@ def rescale_o_accs(self, o_accs, corr): traits = ctx.traits corr_vec = Vec.from_elements([corr], fx.Float32).broadcast_to(16) if const_expr(not traits.USE_HW_TR): - o_accs[0] = _fmul(Vec(o_accs[0]), corr_vec, ctx.fm_fast) + o_accs[0] = Vec(o_accs[0]) * corr_vec else: for dc in range_constexpr(traits.D_CHUNKS): - o_accs[dc] = _fmul(Vec(o_accs[dc]), corr_vec, ctx.fm_fast) + o_accs[dc] = Vec(o_accs[dc]) * corr_vec return o_accs, corr_vec def build_p_packs(self, p_vals): @@ -3020,7 +2999,6 @@ def gemm2_gpfetch_fused( ): ctx = self.ctx traits = ctx.traits - fm_fast = ctx.fm_fast local_sum = ctx.c_zero_f if const_expr(not traits.USE_HW_TR): for dc in range_constexpr(1, traits.D_CHUNKS): @@ -3030,13 +3008,13 @@ def gemm2_gpfetch_fused( p_exp_lo = [] p_exp_hi = [] for j in range_constexpr(traits.MFMA_LANE_K): - diff_lo = fmath.fma(s_raw_lo[p_base + j], ctx.c_sm_scale_log2e, neg_scaled_max, fastmath=ctx.fm_fast) + diff_lo = fx.fma(s_raw_lo[p_base + j], ctx.c_sm_scale_log2e, neg_scaled_max, fastmath=ctx.fm_fast) p_exp_lo.append(self._exp2(diff_lo)) - diff_hi = fmath.fma(s_raw_hi[p_base + j], ctx.c_sm_scale_log2e, neg_scaled_max, fastmath=ctx.fm_fast) + diff_hi = fx.fma(s_raw_hi[p_base + j], ctx.c_sm_scale_log2e, neg_scaled_max, fastmath=ctx.fm_fast) p_exp_hi.append(self._exp2(diff_hi)) for j in range_constexpr(traits.MFMA_LANE_K): - local_sum = _fadd(local_sum, p_exp_lo[j], fm_fast) - local_sum = _fadd(local_sum, p_exp_hi[j], fm_fast) + local_sum = local_sum + p_exp_lo[j] + local_sum = local_sum + p_exp_hi[j] p_lo = self.bf16_trunc_pack_v4(p_exp_lo) p_hi = self.bf16_trunc_pack_v4(p_exp_hi) v_lo = [None] * traits.D_CHUNKS @@ -3047,8 +3025,8 @@ def gemm2_gpfetch_fused( o_accs[dc] = gemm_helper.mfma_acc(v_lo[dc], p_lo, o_accs[dc]) for dc in range_constexpr(traits.D_CHUNKS): o_accs[dc] = gemm_helper.mfma_acc(v_hi[dc], p_hi, o_accs[dc]) - tile_sum = _fadd(local_sum, self.reduction_peer(local_sum), fm_fast) - l_new = _fadd(_fmul(corr, l_running, fm_fast), tile_sum, fm_fast) + tile_sum = local_sum + self.reduction_peer(local_sum) + l_new = corr * l_running + tile_sum return o_accs, l_new @@ -3078,11 +3056,7 @@ def store_lse(self, m_final, l_final, q_row): # LSE = sm_scale * m_raw + ln(l); natural log, softmax scale folded in # (fully-masked row has l == 0 -> -inf). ctx = self.ctx - lse_val = _fadd( - _fmul(m_final, ctx.c_sm_scale, ctx.fm_fast), - fmath.log(as_mlir_value(l_final), fastmath=ctx.fm_fast), - ctx.fm_fast, - ) + lse_val = m_final * ctx.c_sm_scale + fx.log(l_final, fastmath=ctx.fm_fast) lse_local = ctx.q_head_idx * ctx.seq_len_v + q_row # One writer per row: low half-wave + in-bounds q_row; else redirect to the # dropped OOB sentinel. @@ -3251,13 +3225,7 @@ def init_types_and_constants(self, head_dim_runtime=None): c_log2e_f = fx.Float32(_LOG2E) # LSE store folds the log2->ln conversion (m_row is sm_scale*log2e-scaled). self.c_ln2_f = fx.Float32(1.0 / _LOG2E) - self.c_sm_scale_log2e = fx.Float32( - arith.mulf( - as_mlir_value(fmath.rsqrt(head_dim_f32, fastmath=self.fm_fast)), - as_mlir_value(c_log2e_f), - fastmath=self.fm_fast, - ) - ) + self.c_sm_scale_log2e = fx.rsqrt(head_dim_f32, fastmath=self.fm_fast) * c_log2e_f def init_runtime_indices(self, seq_len=None, seq_len_kv=None, stride_q_n=None, stride_kv_n=None): if seq_len is None: @@ -3661,11 +3629,7 @@ def scale_all(self, q_all_bf16): scale_vec = Vec.from_elements([self.c_sm_scale_log2e], fx.Float32).broadcast_to( traits.K_STEPS_QK * traits.MFMA_LANE_K ) - q_all_scaled_f32 = arith.mulf( - as_mlir_value(scale_vec), - as_mlir_value(q_all_f32), - fastmath=self.fm_fast, - ) + q_all_scaled_f32 = as_mlir_value(scale_vec * Vec(q_all_f32)) q_all_scaled_bf16_op = llvm.FPTruncOp(v64bf16_type, q_all_scaled_f32) q_all_scaled_bf16_op.operation.attributes["fastmathFlags"] = fm_fast_attr q_all_scaled_bf16 = q_all_scaled_bf16_op.result @@ -3711,19 +3675,19 @@ def reduce_max(self, v_s): return _score_pair_max(v_s, self.c_neg_inf, self.fm_fast) def floor_masked_max(self, row_max): - return _fmax(row_max, self.c_neg_floor, self.fm_fast) + return fx.maxnumf(row_max, self.c_neg_floor) def rescale_from_tile_max(self, m_row, m_tile_max): return _rescale_from_tile_max(m_row, m_tile_max, self.fm_fast) def apply_l_rescale(self, l_row, rescale): - return _fmul(l_row, rescale, self.fm_fast) + return l_row * rescale def exp2(self, v_s, start, length): return _exp2_score_slice(v_s, start, length) def reduce_sum(self, l_row, v_p): - return _fadd(l_row, _score_pair_sum(v_p, self.c_zero_f, self.fm_fast), self.fm_fast) + return l_row + _score_pair_sum(v_p, self.c_zero_f, self.fm_fast) def sub_m(self, v_s, row_max): return _sub_score_pair(v_s, row_max, self.fm_fast) @@ -3742,8 +3706,8 @@ def scale_o(self, v_o, scale_scalar): _scale_o_accs(v_o, scale_scalar, self.traits, self.fm_fast) def rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): - m_new = _fmax(m_row, m_tile_max, self.fm_fast) - corr = rocdl.exp2(T.f32, as_mlir_value(_fsub(m_row, m_new, self.fm_fast))) + m_new = fx.maxnumf(m_row, m_tile_max) + corr = rocdl.exp2(T.f32, as_mlir_value(m_row - m_new)) self.scale_o(v_o, corr) v_o = _anchor_v_o(self.traits, v_o) v_p = _scale_v_p( @@ -3753,11 +3717,11 @@ def rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): elem_dtype=self.elem_dtype, fm_fast=self.fm_fast, ) - l_row = _fmul(l_row, corr, self.fm_fast) + l_row = l_row * corr return v_o, m_new, l_row, v_p def _lazy_rescale_o_rescale(self, _n, *_st, v_o, m_row, l_row, m_tile_max, v_p): - corr = rocdl.exp2(T.f32, as_mlir_value(_fsub(m_row, m_tile_max, self.fm_fast))) + corr = rocdl.exp2(T.f32, as_mlir_value(m_row - m_tile_max)) scaled_accs = list(v_o) self.scale_o(scaled_accs, corr) out = [as_mlir_value(scaled_accs[dc]) for dc in range(self.traits.D_CHUNKS)] @@ -3769,7 +3733,7 @@ def _lazy_rescale_o_rescale(self, _n, *_st, v_o, m_row, l_row, m_tile_max, v_p): fm_fast=self.fm_fast, ) out.append(_v_p_to_vec32(scaled_p)) - out.append(as_mlir_value(_fmul(l_row, corr, self.fm_fast))) + out.append(as_mlir_value(l_row * corr)) out.append(_anchor_scalar_f32(m_tile_max)) return out @@ -3781,7 +3745,7 @@ def lazy_rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): @flyc.jit def _lazy_rescale_o(v_o, m_row, l_row, m_tile_max, v_p): c_eight_f = fx.Float32(traits.DUALWAVE_SWP_RESCALE_THRESHOLD) - m_diff = _fsub(m_tile_max, m_row, self.fm_fast) + m_diff = m_tile_max - m_row below = fx.Float32(m_diff) <= c_eight_f ballot = rocdl.ballot(T.i64, as_mlir_value(below)) all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64()) @@ -4245,11 +4209,7 @@ def _store_lse_row(self, m_row, l_row, q_row): lse_per_batch_elems = fx.Index(traits.NUM_HEADS_Q) * self.seq_len_v lse_per_batch_bytes = lse_per_batch_elems * fx.Index(4) lse_rsrc = _make_ws_rsrc(lse_base_i64, self.batch_idx * lse_per_batch_bytes, lse_per_batch_bytes) - lse_val = _fadd( - _fmul(m_row, self.c_ln2_f, self.fm_fast), - fmath.log(as_mlir_value(l_row), fastmath=self.fm_fast), - self.fm_fast, - ) + lse_val = m_row * self.c_ln2_f + fx.log(l_row, fastmath=self.fm_fast) lse_local = self.q_head_idx * self.seq_len_v + q_row # One writer per row: low half-wave + in-bounds q_row; else the dropped OOB sentinel. lse_off_row = (q_row < self.seqlen_q_v).select(lse_local, lse_per_batch_elems) @@ -4463,25 +4423,13 @@ def _load_scale_scalar(tensor): head_dim_f32 = fx.Float32(fx.Int32(self.head_dim_runtime)) c_log2e_f = fx.Float32(_LOG2E) - c_sm_scale_log2e = fx.Float32( - arith.mulf( - as_mlir_value(fmath.rsqrt(head_dim_f32, fastmath=self.fm_fast)), - as_mlir_value(c_log2e_f), - fastmath=self.fm_fast, - ) - ) + c_sm_scale_log2e = fx.rsqrt(head_dim_f32, fastmath=self.fm_fast) * c_log2e_f _qd = _load_scale_scalar(self.QDescale) _kd = _load_scale_scalar(self.KDescale) self.vd_fp8 = _load_scale_scalar(self.VDescale) # fp8 feeds raw Q/K into the MFMA, so q/k descale * softmax scale multiplies # the fp32 logits after QK. - self.c_logit_scale = fx.Float32( - arith.mulf( - as_mlir_value(c_sm_scale_log2e), - as_mlir_value(arith.mulf(as_mlir_value(_qd), as_mlir_value(_kd), fastmath=self.fm_fast)), - fastmath=self.fm_fast, - ) - ) + self.c_logit_scale = c_sm_scale_log2e * (_qd * _kd) def init_tile_bounds(self): traits = self.traits @@ -5043,10 +4991,10 @@ def reduce_max(self, v_s): return _score_pair_max(v_s, self.c_neg_inf, self.fm_fast) def max2(self, a, b): - return _fmax(a, b, self.fm_fast) + return fx.maxnumf(a, b) def floor_masked_max(self, row_max): - return _fmax(row_max, self.c_neg_floor, self.fm_fast) + return fx.maxnumf(row_max, self.c_neg_floor) def sub_m(self, v_s, row_max): return _scale_sub_score_pair(v_s, row_max, self.c_logit_scale, self.c_zero_f, self.fm_fast) @@ -5058,7 +5006,7 @@ def tile_sum(self, v_p): return _score_pair_sum(v_p, self.c_zero_f, self.fm_fast) def reduce_sum(self, l_row, v_p): - return _fadd(l_row, self.tile_sum(v_p), self.fm_fast) + return l_row + self.tile_sum(v_p) def cast_p(self, v_p): # Pack the finished softmax probabilities into v8 bf16 P packs for PV. @@ -5091,13 +5039,13 @@ def safe_l_inv(self, l_row): return _safe_l_inv(l_row, self.c_zero_f) def rescale_from_tile_max(self, m_row, m_tile_max): - row_max = _fmax(m_row, m_tile_max, self.fm_fast) - diff_scaled = _fmul(_fsub(m_row, row_max, self.fm_fast), self.c_logit_scale, self.fm_fast) + row_max = fx.maxnumf(m_row, m_tile_max) + diff_scaled = (m_row - row_max) * self.c_logit_scale rescale = rocdl.exp2(T.f32, as_mlir_value(diff_scaled)) return row_max, rescale def apply_l_rescale(self, l_row, rescale): - return _fmul(l_row, rescale, self.fm_fast) + return l_row * rescale def rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): m_new, corr = self.rescale_from_tile_max(m_row, m_tile_max) @@ -5118,8 +5066,8 @@ def v_vec32_to_p(self, v_p_all): def lazy_rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): @flyc.jit def _run(v_o, m_row, l_row, m_tile_max, v_p): - m_diff = _fsub(m_tile_max, m_row, self.fm_fast) - m_diff_scaled = _fmul(m_diff, self.c_logit_scale, self.fm_fast) + m_diff = m_tile_max - m_row + m_diff_scaled = m_diff * self.c_logit_scale below = fx.Float32(m_diff_scaled) <= self.c_eight_f ballot = rocdl.ballot(T.i64, as_mlir_value(below)) all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64()) @@ -5137,7 +5085,7 @@ def _run(v_o, m_row, l_row, m_tile_max, v_p): if fx.Boolean(all_below): pass else: - corr = rocdl.exp2(T.f32, as_mlir_value(_fsub(self.c_zero_f, m_diff_scaled, self.fm_fast))) + corr = rocdl.exp2(T.f32, as_mlir_value(self.c_zero_f - m_diff_scaled)) scaled_accs = list(v_o) self.scale_o(scaled_accs, corr) o0, o1, o2, o3 = ( @@ -5147,7 +5095,7 @@ def _run(v_o, m_row, l_row, m_tile_max, v_p): as_mlir_value(scaled_accs[3]), ) vp_out = self.v_p_to_vec32(self.scale_v_p(v_p, corr)) - l_out = as_mlir_value(_fmul(l_row, corr, self.fm_fast)) + l_out = as_mlir_value(l_row * corr) m_out = self.anchor_scalar_f32(m_tile_max) return ([o0, o1, o2, o3], m_out, l_out, self.v_vec32_to_p(vp_out)) @@ -5156,8 +5104,8 @@ def _run(v_o, m_row, l_row, m_tile_max, v_p): def lazy_correct_o(self, v_o, m_row, l_row, m_tile_max): @flyc.jit def _run(v_o, m_row, l_row, m_tile_max): - m_diff = _fsub(m_tile_max, m_row, self.fm_fast) - m_diff_scaled = _fmul(m_diff, self.c_logit_scale, self.fm_fast) + m_diff = m_tile_max - m_row + m_diff_scaled = m_diff * self.c_logit_scale below = fx.Float32(m_diff_scaled) <= self.c_eight_f ballot = rocdl.ballot(T.i64, as_mlir_value(below)) all_below = arith.cmpi(arith.CmpIPredicate.eq, as_mlir_value(ballot), _read_exec_i64()) @@ -5174,7 +5122,7 @@ def _run(v_o, m_row, l_row, m_tile_max): if fx.Boolean(all_below): pass else: - corr = rocdl.exp2(T.f32, as_mlir_value(_fsub(self.c_zero_f, m_diff_scaled, self.fm_fast))) + corr = rocdl.exp2(T.f32, as_mlir_value(self.c_zero_f - m_diff_scaled)) scaled_accs = list(v_o) self.scale_o(scaled_accs, corr) o0, o1, o2, o3 = ( @@ -5183,7 +5131,7 @@ def _run(v_o, m_row, l_row, m_tile_max): as_mlir_value(scaled_accs[2]), as_mlir_value(scaled_accs[3]), ) - l_out = as_mlir_value(_fmul(l_row, corr, self.fm_fast)) + l_out = as_mlir_value(l_row * corr) m_out = self.anchor_scalar_f32(m_tile_max) return ([o0, o1, o2, o3], m_out, l_out) @@ -5403,7 +5351,7 @@ def load_ml_rows(self): def reduce_m_max(self, m_s): m_max = m_s[0] for i in range_constexpr(self.traits.NUM_KV_SPLITS - 1): - m_max = _fmax(m_max, m_s[i + 1], self.fm_fast) + m_max = fx.maxnumf(m_max, m_s[i + 1]) return m_max def init_accumulators(self): @@ -5416,9 +5364,9 @@ def accumulate_split(self, acc, den, split_i, m_i, l_i, m_max): @flyc.jit def _accum_split(acc, den): if fx.Float32(l_i) > fx.Float32(0.0): - w = rocdl.exp2(T.f32, as_mlir_value(_fsub(m_i, m_max, self.fm_fast))) - wl = _fmul(w, l_i, self.fm_fast) - den = _fadd(den, wl, self.fm_fast) + w = rocdl.exp2(T.f32, as_mlir_value(m_i - m_max)) + wl = w * l_i + den = den + wl o2_raw = buffer_ops.buffer_load( orsrc_i, as_mlir_value(fx.Int32(local_o_idx_i)), @@ -5428,7 +5376,7 @@ def _accum_split(acc, den): o2_i32 = ir.Value(o2_raw) o4 = Vec(o2_i32, (2,), fx.Int32).bitcast(self.elem_dtype).to(fx.Float32) w4 = Vec.from_elements([fx.Float32(wl)], fx.Float32).broadcast_to(4) - acc = _fadd(acc, _fmul(w4, o4, self.fm_fast), self.fm_fast) + acc = acc + w4 * o4 return acc, den return _accum_split(acc, den) @@ -5443,7 +5391,7 @@ def pack_output(self, acc, den): inv_rcp = rocdl.rcp(T.f32, den) inv = (fx.Float32(den) > self.c_zero_f).select(inv_rcp, self.c_zero_f) inv4 = Vec.from_elements([fx.Float32(inv)], fx.Float32).broadcast_to(4) - out4 = Vec(_fmul(acc, inv4, self.fm_fast), (4,), fx.Float32) + out4 = Vec(acc * inv4, (4,), fx.Float32) if const_expr(self.traits.DTYPE_STR == "bf16"): lo = rocdl.cvt_pk_bf16_f32(out4[0], out4[1]) hi = rocdl.cvt_pk_bf16_f32(out4[2], out4[3]) @@ -5462,11 +5410,7 @@ def store_lse(self, m_max, den): lse_per_batch_elems = fx.Index(self.traits.NUM_HEADS_Q) * self.seq_len_v lse_per_batch_bytes = lse_per_batch_elems * fx.Index(4) lse_rsrc = _make_ws_rsrc(lse_base_i64, self.batch_idx * lse_per_batch_bytes, lse_per_batch_bytes) - lse_val = _fadd( - _fmul(m_max, self.c_ln2_f, self.fm_fast), - fmath.log(as_mlir_value(den), fastmath=self.fm_fast), - self.fm_fast, - ) + lse_val = m_max * self.c_ln2_f + fx.log(den, fastmath=self.fm_fast) lse_off = fx.Index((self.col == fx.Index(0)).select(self.local_ml_idx, lse_per_batch_elems)) buffer_ops.buffer_store(as_mlir_value(fx.Float32(lse_val)), lse_rsrc, as_mlir_value(fx.Int32(lse_off))) @@ -5523,13 +5467,7 @@ def _stagger_extra_barrier_if_one(stagger_i32): def _debug_atomic_inc_lazy_count(byte_offset, debug_counts_rsrc): - rocdl.raw_buffer_atomic_fadd( - as_mlir_value(fx.Float32(1.0)), - debug_counts_rsrc, - as_mlir_value(fx.Int32(byte_offset)), - as_mlir_value(fx.Int32(0)), - as_mlir_value(fx.Int32(0)), - ) + rocdl.raw_buffer_atomic(as_mlir_value(fx.Float32(1.0))) + debug_counts_rsrc @flyc.jit diff --git a/kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py b/kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py index e641f3dd8..a2f32bfab 100644 --- a/kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py +++ b/kernels/attention/mla_fwd_decode_m16x8_fp8_fp8.py @@ -19,8 +19,8 @@ import flydsl.expr as fx from flydsl._mlir import ir from flydsl._mlir.dialects import llvm +from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl -from flydsl.expr import math as fmath from flydsl.expr.arith import _to_raw as _raw from flydsl.expr.typing import T from flydsl.expr.typing import Vector as Vec @@ -329,32 +329,10 @@ def kn_mla_fwd_decode_m16x8_fp8_fp8( # ---- Types ---- fm_fast = arith.FastMathFlags.fast - # fastmath without ninf: safe for operations that may encounter -inf - # (boundary masking sets OOB attention scores to -inf) - fm_no_inf = ( - arith.FastMathFlags.nnan - | arith.FastMathFlags.nsz - | arith.FastMathFlags.arcp - | arith.FastMathFlags.contract - | arith.FastMathFlags.afn - | arith.FastMathFlags.reassoc - ) def _mfma_fp8(result_type, operands, **kw): return rocdl.mfma_f32_16x16x32_fp8_fp8(result_type, operands, **kw) - def _fadd(a, b, fastmath=fm_no_inf): - return arith.addf(_raw(a), _raw(b), fastmath=fastmath) - - def _fsub(a, b, fastmath=fm_no_inf): - return arith.subf(_raw(a), _raw(b), fastmath=fastmath) - - def _fmul(a, b, fastmath=fm_no_inf): - return arith.mulf(_raw(a), _raw(b), fastmath=fastmath) - - def _fmax(a, b, fastmath=fm_no_inf): - return arith.maximumf(_raw(a), _raw(b), fastmath=fastmath) - # ---- LDS setup ---- lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_base_idx = ArithValue(_raw(fx.ptrtoint(lds.storage.ptr))).index_cast(T.index) @@ -757,7 +735,7 @@ def _warp_reduce_max_16(val): """Butterfly max reduce across MFMA column groups (strides 32, 16).""" w = _f32(val) for sh in [32, 16]: - w = _fmax(w, _shfl_xor_f32(w, sh), fm_no_inf) + w = fx.maxnumf(w, _shfl_xor_f32(w, sh)) return w def _warp_reduce_add_16(val): @@ -917,7 +895,7 @@ def _softmax( # Local max local_max = scaled[0] for i in range_constexpr(1, P_VALS_PER_THR): - local_max = _fmax(local_max, scaled[i], fm_no_inf) + local_max = fx.maxnumf(local_max, scaled[i]) # Warp reduce max (within 16-lane groups) local_max = _warp_reduce_max_16(local_max) @@ -927,18 +905,18 @@ def _softmax( new_row_max = local_max rescale = c_one_f32 else: - new_row_max = _fmax(local_max, row_max_old, fm_no_inf) + new_row_max = fx.maxnumf(local_max, row_max_old) # rescale = exp2((old_max - new_max) * log2e) - diff = _fsub(row_max_old, new_row_max, fm_no_inf) - rescale = _fast_exp2(_fmul(diff, c_log2e, fm_no_inf)) + diff = row_max_old - new_row_max + rescale = _fast_exp2(diff * c_log2e) # exp(p - max) for each value, and sum p_exp_vals = [None] * P_VALS_PER_THR local_sum = c_zero_f32 for i in range_constexpr(P_VALS_PER_THR): - exp_arg = _fmul(_fsub(scaled[i], new_row_max, fm_no_inf), c_log2e, fm_no_inf) + exp_arg = (scaled[i] - new_row_max) * c_log2e p_exp_vals[i] = _fast_exp2(exp_arg) - local_sum = _fadd(local_sum, p_exp_vals[i], fm_no_inf) + local_sum = local_sum + p_exp_vals[i] # Warp reduce sum local_sum = _warp_reduce_add_16(local_sum) @@ -947,7 +925,7 @@ def _softmax( if const_expr(is_first_iter): row_sum_e_new = local_sum else: - row_sum_e_new = _fadd(_f32(rescale) * row_sum_e_old, local_sum, fm_no_inf) + row_sum_e_new = _f32(rescale) * row_sum_e_old + local_sum return p_exp_vals, new_row_max, row_sum_e_new, rescale @@ -1824,8 +1802,8 @@ def _v_base_i32(p_lds_kv_base): def _write_lse(pqo_loc_i32, rm, rse): """Write LSE for split output (first 16 lanes per warp).""" if ArithValue(lane_idx) < 16: - log2_sum = fmath.log2(rse, fastmath=fm_fast) - lse = fmath.fma(log2_sum, c_inv_log2e, rm, fastmath=fm_fast) + log2_sum = fx.log2(rse, fastmath=fm_fast) + lse = fx.fma(log2_sum, c_inv_log2e, rm, fastmath=fm_fast) row_idx = _raw(ArithValue(lane_idx) + warp_idx * 16 + _idx(pqo_loc_i32) * NUM_QO_HEADS) buffer_ops.buffer_store(lse, split_lse_rsrc, row_idx) @@ -2078,19 +2056,22 @@ def launch_mla_fwd_decode_m16x8_fp8_fp8( ): """JIT host function: configures grid/block and launches the kernel.""" assert TOTAL_LDS_BYTES <= lds_size, f"Kernel requires {TOTAL_LDS_BYTES} bytes LDS but CU budget is {lds_size}" - kn_mla_fwd_decode_m16x8_fp8_fp8( - query, - kv_buffer, - kv_page_indices, - work_indptr, - work_info_set, - final_output, - split_output, - split_lse, - softmax_scale, - ).launch( - grid=(num_cus, 1, 1), - block=(NUM_THREADS, 1, 1), - smem=0, - stream=stream, - ) + # DSL arithmetic (+ - * .maximumf) picks up fastmath from the ambient hint; + # enable it for the whole traced body so ops emit fastmath. + with CompilationContext.compile_hints({"fast_fp_math": True}): + kn_mla_fwd_decode_m16x8_fp8_fp8( + query, + kv_buffer, + kv_page_indices, + work_indptr, + work_info_set, + final_output, + split_output, + split_lse, + softmax_scale, + ).launch( + grid=(num_cus, 1, 1), + block=(NUM_THREADS, 1, 1), + smem=0, + stream=stream, + ) diff --git a/kernels/attention/pa_decode_tile.py b/kernels/attention/pa_decode_tile.py index c89e74095..c02b476db 100644 --- a/kernels/attention/pa_decode_tile.py +++ b/kernels/attention/pa_decode_tile.py @@ -32,9 +32,9 @@ import flydsl.compiler as flyc import flydsl.expr as fx +from flydsl.compiler.kernel_function import CompilationContext from flydsl.compiler.protocol import dsl_size_of from flydsl.expr import arith, const_expr, gpu, range_constexpr -from flydsl.expr import math as fmath from flydsl.expr.typing import ReductionOp, T from flydsl.runtime.device import get_rocm_arch from kernels.common import buffer_ops, dpp_utils @@ -397,7 +397,6 @@ def _k_ops_flat(tt_i32): v_scale_f = fx.Float32(value_scale) NEG_INF = fx.Float32(float("-inf")) ZERO_F = fx.Float32(0.0) - fm_contract = arith.FastMathFlags.contract # Softmax scores are finite or the -inf mask sentinel -- never NaN -- so # nnan lets maxnum lower to a bare v_max (no v_cmp_u NaN check + its s_nop # hazard) and fuse to v_max3. (ninf must NOT be set: -inf is load-bearing.) @@ -452,9 +451,9 @@ def _quant_q_row(m, qi, gs_head, q_row_off): # (a buffer load is 128b max); head_dim=256 splits into 2 pieces. q_units = [_q_load_chunk(base_elem + u * QLOAD_UNIT) for u in range_constexpr(N_QLOADS)] - absmax = fmath.absf(q_units[0]).reduce(ReductionOp.MAX).to(fx.Float32) + absmax = fx.absf(q_units[0]).reduce(ReductionOp.MAX).to(fx.Float32) for u in range_constexpr(1, N_QLOADS): - absmax = fx.maxnumf(absmax, fmath.absf(q_units[u]).reduce(ReductionOp.MAX).to(fx.Float32)) + absmax = fx.maxnumf(absmax, fx.absf(q_units[u]).reduce(ReductionOp.MAX).to(fx.Float32)) for sh in (8, 4, 2, 1): absmax = fx.maxnumf(absmax, dpp_utils.dpp_xor_f32(absmax, sh)) @@ -742,7 +741,7 @@ def _lmax_off_m(m): p_off0 + a * (c16 // 4) * f32, fx.Int32, fx.Vector.from_elements([words[a]], dtype=fx.Int32) ) for sh in (16, 32): - ls = ls.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) + ls = ls + ls.shuffle_xor(sh, WAVE) # PV output is [head-dim, query-row=lane16] after the operand # swap, so correction/denominator are per-lane scalars (no sCorr). safe_prev = arith.select(m_prev > NEG_INF, m_prev, ZERO_F) @@ -751,9 +750,7 @@ def _lmax_off_m(m): _st_lw(sLsum_off, lane16, warp, ls) gpu.barrier() gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) - l_new = fx.Float32( - arith.mulf(arith.unwrap(l_prev), arith.unwrap(corr_reg), fastmath=fm_contract) - ).addf(gsum, fastmath=fm_contract) + l_new = l_prev * corr_reg + gsum p_ops = _lds_load(sP_off + lane16 * SP_ROW_BYTES + rgroup * 64, fx.Int64, NVOPS) @@ -879,7 +876,7 @@ def _lmax_off_m(m): if const_expr(head_dim == 64): fx.rocdl.sched_dswr(NCHUNK) for sh in (16, 32): - ls = ls.addf(ls.shuffle_xor(sh, WAVE), fastmath=fm_contract) + ls = ls + ls.shuffle_xor(sh, WAVE) # PV (V=A, P=B) -> output [head-dim, query-row=lane16]; same as # the phase-split path. corr_reg = fx.Float32(exp2_amdgcn_scalar(m_prev - m_new)) @@ -887,9 +884,7 @@ def _lmax_off_m(m): _st_lw(sLsum_off, lane16, warp, ls) gpu.barrier() gsum = _ld_lw_row(sLsum_off, lane16).reduce(ReductionOp.ADD) - l_new = fx.Float32(arith.mulf(arith.unwrap(l_prev), arith.unwrap(corr_reg), fastmath=fm_contract)).addf( - gsum, fastmath=fm_contract - ) + l_new = l_prev * corr_reg + gsum p_ops = _lds_load(sP_off + lane16 * SP_ROW_BYTES + rgroup * 64, fx.Int64, NVOPS) corr_b = fx.Vector.from_elements([corr_reg], dtype=fx.Float32).broadcast_to(OP_ELEMS) # Single tile: batch both vh's V loads upfront (no sibling chain @@ -919,13 +914,7 @@ def _lmax_off_m(m): if const_expr(per_token_kv): o_scale = inv_l else: - o_scale = fx.Float32( - arith.mulf( - arith.unwrap(inv_l), - arith.unwrap(v_scale_f * inv_fp8), - fastmath=fm_contract, - ) - ) + o_scale = inv_l * (v_scale_f * inv_fp8) o_scale_b = fx.Vector.from_elements([o_scale], dtype=fx.Float32).broadcast_to(OP_ELEMS) qi_e = row // query_group_size gs_head_e = row - qi_e * query_group_size @@ -983,24 +972,25 @@ def pa_decode_tile_launch( stride_q_head: fx.Int32, stream: fx.Stream = fx.Stream(None), ): - pa_decode_tile_kernel( - output, - pmax, - psum, - pout, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - max_blocks_per_seq, - stride_ks_block, - stride_ks_head, - stride_q_row, - stride_q_head, - ).launch(grid=(num_seqs, num_kv_heads, NP), block=(BLOCK_THREADS, 1, 1), stream=stream) + with CompilationContext.compile_hints({"fastmath": arith.FastMathFlags.contract}): + pa_decode_tile_kernel( + output, + pmax, + psum, + pout, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + max_blocks_per_seq, + stride_ks_block, + stride_ks_head, + stride_q_row, + stride_q_head, + ).launch(grid=(num_seqs, num_kv_heads, NP), block=(BLOCK_THREADS, 1, 1), stream=stream) return {"launch": pa_decode_tile_launch, "kernel": pa_decode_tile_kernel} diff --git a/kernels/attention/pa_metadata.py b/kernels/attention/pa_metadata.py index 7a4b76162..86800a76f 100644 --- a/kernels/attention/pa_metadata.py +++ b/kernels/attention/pa_metadata.py @@ -42,6 +42,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx from flydsl._mlir.dialects import vector +from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl from flydsl.expr.typing import Int32, T from flydsl.runtime.device import get_rocm_arch @@ -332,14 +333,14 @@ def _finish_q_fragments( q_frags = [] gpu.barrier() - query_scale_lane = fx.ptr_load(softmax_base + (lane16id), result_type=fx.Vector.make_type(1, fx.Float32))[ + query_scale_lane = fx.ptr_load(softmax_base + lane16id, result_type=fx.Vector.make_type(1, fx.Float32))[ 0 ].ir_value() for qkhe in range_constexpr(qkhe_loop): for qkr in range_constexpr(2): lds_rd = lane16id * fx.Int32(head_size // 8) + fx.Int32(qkhe * 8) + rowid * fx.Int32(2) + fx.Int32(qkr) q_v1 = fx.ptr_load( - fx.recast_iter(fx.Int64, logits_base) + (lds_rd), result_type=fx.Vector.make_type(1, fx.Int64) + fx.recast_iter(fx.Int64, logits_base) + lds_rd, result_type=fx.Vector.make_type(1, fx.Int64) ) q_frags.append(q_v1[0]) return q_frags, query_scale_lane @@ -540,7 +541,7 @@ def _load_v_and_scales( for td in range_constexpr(TLOOP): scale_row_base = kv_tok_thread_base + fx.Int32(td * MFMA_N) k_scale_vecs.append( - fx.ptr_load(scale_base + (scale_row_base), result_type=fx.Vector.make_type(4, fx.Float32)) + fx.ptr_load(scale_base + scale_row_base, result_type=fx.Vector.make_type(4, fx.Float32)) ) v_scale_vecs.append( fx.ptr_load( @@ -556,7 +557,7 @@ def _scale_row_base(td: int): return kv_tok_thread_base + fx.Int32(td * MFMA_N) def _load_k_scale_vec(td: int): - return fx.ptr_load(scale_base + (_scale_row_base(td)), result_type=fx.Vector.make_type(4, fx.Float32)) + return fx.ptr_load(scale_base + _scale_row_base(td), result_type=fx.Vector.make_type(4, fx.Float32)) def _load_v_scale_vec(td: int): return fx.ptr_load( @@ -668,7 +669,7 @@ def _qk_and_intra_softmax( def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs, v_scale_vecs): partition_max = neg_inf partition_sum = zero_f - max_vec = fx.ptr_load(softmax_base + (sm_rd_max_offs[0]), result_type=fx.Vector.make_type(4, fx.Float32)) + max_vec = fx.ptr_load(softmax_base + sm_rd_max_offs[0], result_type=fx.Vector.make_type(4, fx.Float32)) for w in range_constexpr(NUM_WARPS): partition_max = fx.maxnumf(partition_max, max_vec[w]) @@ -696,14 +697,12 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs, v_scale_vecs): accum_scale = exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()) gpu.barrier() - sum_vec = fx.ptr_load(softmax_base + (sm_rd_sum_offs[0]), result_type=fx.Vector.make_type(4, fx.Float32)) + sum_vec = fx.ptr_load(softmax_base + sm_rd_sum_offs[0], result_type=fx.Vector.make_type(4, fx.Float32)) for w in range_constexpr(NUM_WARPS): - partition_sum = arith.addf( - arith.unwrap(partition_sum), arith.unwrap(sum_vec[w]), fastmath=arith.FastMathFlags.contract - ) + partition_sum = partition_sum + sum_vec[w] - accum_sum = arith.mulf(arith.unwrap(accum_scale), arith.unwrap(rsum), fastmath=arith.FastMathFlags.contract) - rsum = arith.addf(accum_sum, arith.unwrap(partition_sum), fastmath=arith.FastMathFlags.contract) + accum_sum = accum_scale * rsum + rsum = accum_sum + partition_sum rmax = new_rmax accum_scale_vec = vector.broadcast(T.f32x4, arith.unwrap(accum_scale)) for vhe in range_constexpr(vhe_loop): @@ -711,7 +710,7 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs, v_scale_vecs): if const_expr(per_token_kv): v_max_global = zero_f - vmax_vec = fx.ptr_load(softmax_base + (sm_vmax_rd_offs[0]), result_type=fx.Vector.make_type(4, fx.Float32)) + vmax_vec = fx.ptr_load(softmax_base + sm_vmax_rd_offs[0], result_type=fx.Vector.make_type(4, fx.Float32)) for w in range_constexpr(NUM_WARPS): w_vmax = vmax_vec[w] v_max_global = fx.maxnumf(v_max_global, w_vmax) @@ -736,7 +735,6 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs, v_scale_vecs): def _pv_mfma(v_ops, outs, v_correction): v_correction = fx.Float32(v_correction).ir_value() - fm_contract = arith.FastMathFlags.contract v_correction_vec = vector.broadcast(T.f32x4, v_correction) # ── Batch-load all P_i64 from LDS upfront ── @@ -751,7 +749,7 @@ def _pv_mfma(v_ops, outs, v_correction): p_i64_off = pv_prob_i64_elems[vt * 2 + j] p_i64_all.append( fx.ptr_load( - fx.recast_iter(fx.Int64, logits_base) + (p_i64_off), + fx.recast_iter(fx.Int64, logits_base) + p_i64_off, result_type=fx.Vector.make_type(1, fx.Int64), )[0] ) @@ -772,11 +770,7 @@ def _pv_mfma(v_ops, outs, v_correction): 0, ], ) - outs[vhe] = arith.addf( - arith.mulf(tmp_out, v_correction_vec, fastmath=fm_contract), - outs[vhe], - fastmath=fm_contract, - ) + outs[vhe] = fx.Vector(tmp_out) * v_correction_vec + outs[vhe] return outs return ( @@ -1000,12 +994,12 @@ def _sel(cond_b, a, b): for sh in _shuffle_offsets: sum_blocks = sum_blocks + sum_blocks.shuffle_xor(arith.constant(sh, type=i32), c_ws.ir_value()) - average = fx.Int32(arith.divui(sum_blocks.ir_value(), c_nspk.ir_value())) - reminder = fx.Int32(arith.remui(sum_blocks.ir_value(), c_nspk.ir_value())) + average = sum_blocks // c_nspk + reminder = sum_blocks % c_nspk def _remain_for_cid(cid_val): # remain = average + (1 if (cid % num_splits_per_khead) < reminder else 0) - mod = fx.Int32(arith.remui(cid_val.ir_value(), c_nspk.ir_value())) + mod = cid_val % c_nspk return average + _sel(mod < reminder, 1, 0) # ---- Phase 2: per khead, flattened CU x batch scheduler ---- @@ -1192,17 +1186,18 @@ def launch_pa_metadata_v1( num_batches: Int32, stream: fx.Stream = fx.Stream(None), ): - pa_metadata_v1_kernel( - seqlens_qo_indptr, - pages_kv_indptr, - context_lens, - work_indptr, - work_info, - reduce_indptr, - reduce_final_map, - reduce_partial_map, - num_batches, - ).launch(grid=(1, 1, 1), block=(warp_size, 1, 1), stream=stream) + with CompilationContext.compile_hints({"fastmath": arith.FastMathFlags.contract}): + pa_metadata_v1_kernel( + seqlens_qo_indptr, + pages_kv_indptr, + context_lens, + work_indptr, + work_info, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + num_batches, + ).launch(grid=(1, 1, 1), block=(warp_size, 1, 1), stream=stream) return {"kernel": pa_metadata_v1_kernel, "launch": launch_pa_metadata_v1} @@ -1503,10 +1498,10 @@ def pa_decode_metadata_kenrel( # Outer work loop — each work item = one (batch, kv_head_range, kv_page_range) _work_start_idx = fx.Index(arith.unwrap(work_start)) _work_end_idx = fx.Index(arith.unwrap(work_end)) - _work_step = arith.index(1) + _work_step = fx.Index(1) for _wi in range(_work_start_idx, _work_end_idx, _work_step): - work_idx = arith.index_cast(T.i32, _wi) + work_idx = fx.Int32(_wi) # ── Load work_info[work_idx] — 8 × int32, as 2 × vec4 loads ── # info_base is a multiple of 8, so both dwordx4 loads are naturally @@ -1627,9 +1622,9 @@ def _unwrap(v): # (sink-prone partition 0 processed last for online-softmax stability). num_parts_in_work = kv_end - kv_start last_part_idx_val = num_parts_in_work - c_one - _loop_start_g = arith.index(0) + _loop_start_g = fx.Index(0) _loop_stop_g = fx.Index(arith.unwrap(num_parts_in_work)) - _loop_step_g = arith.index(1) + _loop_step_g = fx.Index(1) _mtp_groups = math.ceil(query_length * query_group_size / 16) @@ -1677,7 +1672,7 @@ def _meta_load_v_phys_from_lds(): for vt in range_constexpr(VTLOOP): bt_lds_off = arith.constant(vt * TLOOP, type=T.i32) + rowid v_phys_blocks.append( - fx.ptr_load(bt_base + (bt_lds_off), result_type=fx.Vector.make_type(1, fx.Int32))[0] + fx.ptr_load(bt_base + bt_lds_off, result_type=fx.Vector.make_type(1, fx.Int32))[0] ) return v_phys_blocks @@ -1805,7 +1800,7 @@ def _unpack_states_kv(flat): # Reverse iteration: scf.for walks ib forward (0..N-1); remap to # the local partition index lp = N-1..0 so the sink-prone first # partition is processed last. - rel_part = last_part_idx_val - arith.index_cast(T.i32, ib) + rel_part = last_part_idx_val - as_ir_value(fx.Int32(ib)) lp = local_part_start + rel_part next_rel = rel_part - c_one next_rel_clamped = arith.select(next_rel >= c_zero_i32, next_rel, c_zero_i32) @@ -2144,7 +2139,7 @@ def pa_metadata_reduce_kernel( v = buffer_ops.buffer_load( po_rsrc, prow * stride_po_row + qhead * c_head + tid, vec_width=1, dtype=T.f32 ) - m_new = m.maximumf(lse) + m_new = fx.maxnumf(m, lse) scale_old = exp2_f32_fast((m - m_new) * c_log2e) w = exp2_f32_fast((lse - m_new) * c_log2e) denom_new = denom * scale_old + w @@ -2180,22 +2175,23 @@ def launch_pa_metadata_reduce( num_groups, stream: fx.Stream = fx.Stream(None), ): - pa_metadata_reduce_kernel( - final_output, - partial_output, - partial_lse, - reduce_indptr, - reduce_final_map, - reduce_partial_map, - stride_out_seq, - stride_out_head, - stride_po_row, - stride_pl_row, - ).launch( - grid=(num_groups, num_query_heads, query_length), - block=(block_threads, 1, 1), - stream=stream, - ) + with CompilationContext.compile_hints({"fastmath": arith.FastMathFlags.contract}): + pa_metadata_reduce_kernel( + final_output, + partial_output, + partial_lse, + reduce_indptr, + reduce_final_map, + reduce_partial_map, + stride_out_seq, + stride_out_head, + stride_po_row, + stride_pl_row, + ).launch( + grid=(num_groups, num_query_heads, query_length), + block=(block_threads, 1, 1), + stream=stream, + ) return {"launch": launch_pa_metadata_reduce, "kernel": pa_metadata_reduce_kernel} diff --git a/kernels/attention/qk_norm_rope_quant.py b/kernels/attention/qk_norm_rope_quant.py index 1403c80e4..a4ac4237c 100644 --- a/kernels/attention/qk_norm_rope_quant.py +++ b/kernels/attention/qk_norm_rope_quant.py @@ -290,7 +290,7 @@ def load_vec(div_tensor, idx, *, layout=full_lay, atom=full_atom, dt=elem_dtype) def _ptr_buffer_resource(ptr, num_records_bytes=None): addr = fx.ptrtoint(ptr) - addr_i64 = arith.index_cast(T.i64, addr) + addr_i64 = as_ir_value(fx.Int64(addr)) if num_records_bytes is None: return buffer_ops.create_buffer_resource_from_addr(addr_i64) return buffer_ops.create_buffer_resource_from_addr(addr_i64, num_records_bytes=num_records_bytes) diff --git a/kernels/comm/flydsl_dispatch_combine_intranode_kernel.py b/kernels/comm/flydsl_dispatch_combine_intranode_kernel.py index d4f26025e..cde6d8273 100644 --- a/kernels/comm/flydsl_dispatch_combine_intranode_kernel.py +++ b/kernels/comm/flydsl_dispatch_combine_intranode_kernel.py @@ -309,10 +309,12 @@ def ep_dispatch_intranode( expert_id = buffer_load(_r_out_idx_local, smoe_idx, vec_width=1, dtype=T.i32()) local_expert_id = expert_id - rank * experts_per_rank - # MUST be unsigned ``ult``: signed ``slt`` would mis-classify - # negative ``local_expert_id`` (non-local experts) as local - # and trigger illegal global access in WarpCopy below. - is_local = arith.cmpi(arith.CmpIPredicate.ult, local_expert_id, fx.Int32(experts_per_rank)) + # MUST stay unsigned (``ult``): a signed compare would + # mis-classify negative ``local_expert_id`` (non-local experts) + # as local and trigger illegal global access in WarpCopy below. + # fx.Uint32 reinterprets the same bits as unsigned, so ``<`` + # emits ``ult``. + is_local = fx.Uint32(local_expert_id) < fx.Uint32(experts_per_rank) packed_slot_lane0 = fx.Int32(0) if lane == 0: diff --git a/kernels/common/mem_ops.py b/kernels/common/mem_ops.py index 4efaa5b10..8c6553b7e 100644 --- a/kernels/common/mem_ops.py +++ b/kernels/common/mem_ops.py @@ -18,8 +18,7 @@ from flydsl._mlir.dialects import arith as _std_arith from flydsl._mlir.dialects import fly as _fly from flydsl._mlir.dialects import llvm as _llvm -from flydsl.expr import arith as _expr_arith -from flydsl.expr import const_expr, rocdl +from flydsl.expr import as_ir_value, const_expr, rocdl from flydsl.expr.typing import T from kernels.common import buffer_ops @@ -62,7 +61,7 @@ def get_llvm_ptr(ptr, offset, dtype_bytes, ptr_type=None): ptr_type = ir.Type.parse("!llvm.ptr<1>") base_ptr = _fly.extract_aligned_pointer_as_index(ptr_type, ptr) base_ptr = _llvm.PtrToIntOp(T.i64, base_ptr).result - byte_offset = _expr_arith.index_cast(T.i64, fx.Index(offset) * fx.Index(dtype_bytes)) + byte_offset = as_ir_value(fx.Int64(fx.Index(offset) * fx.Index(dtype_bytes))) llvm_ptr = _llvm.AddOp(base_ptr, byte_offset, _llvm.IntegerOverflowFlags(0)).result llvm_ptr = _llvm.IntToPtrOp(ptr_type, llvm_ptr).result return llvm_ptr._value if const_expr(hasattr(llvm_ptr, "_value")) else llvm_ptr diff --git a/kernels/common/mma/mfma_epilogues.py b/kernels/common/mma/mfma_epilogues.py index ae1c03ccd..75ef0d0bf 100644 --- a/kernels/common/mma/mfma_epilogues.py +++ b/kernels/common/mma/mfma_epilogues.py @@ -37,7 +37,6 @@ import flydsl.expr as fx from flydsl._mlir import ir -from flydsl._mlir.dialects.arith import CmpIPredicate from flydsl.expr.typing import T from kernels.common.kernels_common import _if_then @@ -69,7 +68,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) for ii in range_constexpr(4): row_off = lane_div_16_mul4 + ii_idx_list[ii] row_in_tile = mi_base + row_off @@ -157,11 +156,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) + _half_thr_idx = fx.Index(_half_threads) + _zero_idx = fx.Index(0) - _is_group_b = arith.cmpi(CmpIPredicate.uge, tx, _half_thr_idx) + _is_group_b = fx.as_ir_value((tx) >= (_half_thr_idx)) # -- write phase (all waves, each to its group's LDS buffer) -- n_tile_base_v = n_tile_base @@ -209,10 +208,10 @@ def _write_row_split(mi: int, ii: int, row_in_tile, row): # -- 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) + c_nlane_s = fx.Index(CShuffleNLane_s) 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) if frag_elem_type is None: frag_elem_type = T.f16 @@ -222,7 +221,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) row_local = row_base_m + m_lane_s row = bx_m_v + row_local row_ctx_raw = precompute_row(row_local=row_local, row=row) if precompute_row is not None else None @@ -238,7 +237,7 @@ def _write_row_split(mi: int, ii: int, row_in_tile, row): def _do_store_row_split(): 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)) col_pair0_local = col_base_nr + (n_lane_s * c_evec) lds_idx = row_base_lds + col_pair0_local @@ -273,7 +272,7 @@ def _do_store_row_split(): # ===================== 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)) n_tile_base_v = n_tile_base col_base_local = n_tile_base_v + lane_mod_16 # index within [0,tile_n) @@ -332,7 +331,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) row_local = row_base_m + m_lane row = bx_m_v + row_local @@ -356,7 +355,7 @@ def _do_store_row(): 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)) col_pair0 = col_base_nr + (n_lane * c_evec) # even col within tile lds_idx_pair = row_base_lds + col_pair0 diff --git a/kernels/common/mma/mfma_preshuffle_pipeline.py b/kernels/common/mma/mfma_preshuffle_pipeline.py index 7da3e2020..732d81594 100644 --- a/kernels/common/mma/mfma_preshuffle_pipeline.py +++ b/kernels/common/mma/mfma_preshuffle_pipeline.py @@ -35,10 +35,8 @@ def swizzle_xor16(row, col, k_blocks16): k_blocks16 is always a power of 2 (tile_k_bytes / 16), so use bitwise AND instead of remui to save ~10 VALU cycles on CDNA. """ - from flydsl.expr import arith as _swz_arith - - mask = k_blocks16 - _swz_arith.index(1) - rem = _swz_arith.andi(row, mask) + mask = k_blocks16 - fx.Index(1) + rem = (row) & (mask) return col ^ (rem * 16) @@ -137,11 +135,11 @@ def make_preshuffle_scale_layout( stride_k0 = c4 * stride_klane stride_n0 = c_k1 * stride_k0 - c_mn1_i32 = arith.index_cast(T.i32, c_mn1) - c_k1_i32 = arith.index_cast(T.i32, c_k1) - stride_n0_i32 = arith.index_cast(T.i32, stride_n0) - stride_k0_i32 = arith.index_cast(T.i32, stride_k0) - stride_klane_i32 = arith.index_cast(T.i32, stride_klane) + c_mn1_i32 = fx.Int32(c_mn1) + c_k1_i32 = fx.Int32(c_k1) + stride_n0_i32 = fx.Int32(stride_n0) + stride_k0_i32 = fx.Int32(stride_k0) + stride_klane_i32 = fx.Int32(stride_klane) layout_scale = fx.make_layout( (c_mn1_i32, c_k1_i32, 4, 16), @@ -187,10 +185,10 @@ def make_preshuffle_b_layout( if elem_bytes not in (1, 2): raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") - c_k_bytes = c_k * arith.constant(int(elem_bytes), index=True) + c_k_bytes = c_k * fx.Index(int(elem_bytes)) n0 = c_n // c16 - c_kpack_elems = c_kpack if elem_bytes == 1 else (c_kpack // arith.constant(int(elem_bytes), index=True)) + c_kpack_elems = c_kpack if elem_bytes == 1 else (c_kpack // fx.Index(int(elem_bytes))) stride_nlane = c_kpack_elems @@ -212,12 +210,12 @@ def make_preshuffle_b_layout( stride_n0 = c_k0 * stride_k0 kpack_elems_static = kpack_bytes if elem_bytes == 1 else kpack_bytes // elem_bytes - n0_i32 = arith.index_cast(T.i32, n0) - c_k0_i32 = arith.index_cast(T.i32, c_k0) - stride_n0_i32 = arith.index_cast(T.i32, stride_n0) - stride_k0_i32 = arith.index_cast(T.i32, stride_k0) - stride_klane_i32 = arith.index_cast(T.i32, stride_klane) - stride_nlane_i32 = arith.index_cast(T.i32, stride_nlane) + n0_i32 = fx.Int32(n0) + c_k0_i32 = fx.Int32(c_k0) + stride_n0_i32 = fx.Int32(stride_n0) + stride_k0_i32 = fx.Int32(stride_k0) + stride_klane_i32 = fx.Int32(stride_klane) + stride_nlane_i32 = fx.Int32(stride_nlane) stride_b = (stride_n0_i32, stride_k0_i32, stride_klane_i32, stride_nlane_i32, 1) layout_b = fx.make_layout((n0_i32, c_k0_i32, klane_dim, 16, kpack_elems_static), stride_b) @@ -439,12 +437,12 @@ def load_b_pack_k32( raise ValueError(f"elem_bytes must be 1 or 2, got {elem_bytes!r}") c64 = fx.Index(64) - base_k_bytes = base_k * arith.constant(int(elem_bytes), index=True) + base_k_bytes = base_k * fx.Index(int(elem_bytes)) k0_base = base_k_bytes // c64 - k0 = k0_base + arith.constant(ki_step // 2, index=True) + k0 = k0_base + fx.Index(ki_step // 2) k1 = lane_div_16 half_bytes = kpack_bytes // 2 - k2_base = arith.constant((ki_step % 2) * half_bytes, index=True) + k2_base = fx.Index((ki_step % 2) * half_bytes) coord_pack = (n_blk, k0, k1, n_intra, fx.Index(0)) idx_pack = preshuffle_crd2idx(tuple(fx.Int32(c) for c in coord_pack), layout_b) @@ -501,7 +499,7 @@ def tile_chunk_coord_i32( """Map (thread, chunk_id) -> (row_local, col_local_i32) for X/A loads.""" if chunk_i32 not in (1, 2, 4): raise ValueError(f"chunk_i32 must be one of (1,2,4), got {chunk_i32!r}") - chunk_off_i32 = arith.constant(i * total_threads * chunk_i32, index=True) + chunk_off_i32 = fx.Index(i * total_threads * chunk_i32) tile_idx_i32 = tx_i32_base + chunk_off_i32 coord_local = fx.idx2crd(fx.Int32(tile_idx_i32), layout_tile_div4) row_local = fx.get(coord_local, 0) @@ -796,30 +794,31 @@ def _load_groupwise_scale( c_npm1 = fx.Index(num_pairs - 1) dword_base = expert_offset * c_npm1 + n_global dword_elem = dword_base + pair_idx * c_npe - dword_idx = arith.index_cast(T.i32, dword_elem) + dword_idx = fx.Int32(dword_elem) scale_val = buffer_ops.buffer_load(scale_rsrc, dword_idx, vec_width=1, dtype=T.i32) else: # (E, G, N) layout with f32 dtype c_gm1 = fx.Index(num_groups - 1) base_scale = expert_offset * c_gm1 + n_global elem_idx = base_scale + group_idx * c_npe - scale_idx_i32 = arith.index_cast(T.i32, elem_idx) + scale_idx_i32 = fx.Int32(elem_idx) scale_val = buffer_ops.buffer_load(scale_rsrc, scale_idx_i32, vec_width=1, dtype=T.f32) return scale_val -def extract_bf16_scale(arith, scale_raw_i32, ku: int): +def extract_bf16_scale(scale_raw_i32, ku: int): """Extract f32 scale from raw i32 dword loaded by bf16 groupwise path. In the ``(E, G//2, N, 2)`` layout two adjacent groups share one dword. ``ku`` determines which half: even ku → low bf16, odd ku → high bf16. """ + v = fx.Uint32(scale_raw_i32) if ku % 2 == 0: # Low bf16: shift left by 16 to place in upper 16 bits → f32 - return arith.bitcast(T.f32, scale_raw_i32 << fx.Int32(16)) + return (v << fx.Int32(16)).bitcast(fx.Float32) else: # High bf16: mask upper 16 bits → f32 - return arith.bitcast(T.f32, scale_raw_i32 & fx.Int32(0xFFFF0000)) + return (v & fx.Int32(0xFFFF0000)).bitcast(fx.Float32) # --------------------------------------------------------------------------- diff --git a/kernels/conv/conv3d_implicit.py b/kernels/conv/conv3d_implicit.py index 6b8f3c01a..94072c288 100644 --- a/kernels/conv/conv3d_implicit.py +++ b/kernels/conv/conv3d_implicit.py @@ -565,12 +565,12 @@ def _big_store(off_nk_i64, value): def _valid_raw(row, col): if const_expr(_row_chk and n_tail): - return arith.andi(row < fx.Index(npq), col < fx.Index(k)) + return fx.as_ir_value((row < fx.Index(npq)) & (col < fx.Index(k))) if const_expr(_row_chk): v = row < fx.Index(npq) - return arith.andi(v, v) + return fx.as_ir_value(v & v) v = col < fx.Index(k) - return arith.andi(v, v) + return fx.as_ir_value(v & v) def store_acc(): for mi in range_constexpr(MI_M): diff --git a/kernels/moe/mxfp_moe/gemm1.py b/kernels/moe/mxfp_moe/gemm1.py index f147cdfcc..687e27caa 100644 --- a/kernels/moe/mxfp_moe/gemm1.py +++ b/kernels/moe/mxfp_moe/gemm1.py @@ -11,7 +11,6 @@ from . import dpp_utils from .mxfp4_gemm_common import ( _e8m0_from_amax, - _fabs_f32, _global_i32_at, _global_i32_buffer_tiles, _global_i32_buffer_view, @@ -604,9 +603,9 @@ def acc_load(idx): up_vs[ee] = acc_load(acc_idx(row_local, up_col)) result = _silu_mul_batch(gate_vs, up_vs) - local_max = _fabs_f32(result[0]) + local_max = fx.absf(result[0]) for ee in range_constexpr(1, 8): - local_max = local_max.maximumf(_fabs_f32(result[ee])) + local_max = fx.maxnumf(local_max, fx.absf(result[ee])) lm_i = _inline_dpp_quad_amax(fx.Int32(_raw(local_max).bitcast(T.i32))) local_max = fx.Float32(_raw(lm_i).bitcast(T.f32)) diff --git a/kernels/moe/mxfp_moe/gemm2.py b/kernels/moe/mxfp_moe/gemm2.py index f25c6d401..3c4250ad1 100644 --- a/kernels/moe/mxfp_moe/gemm2.py +++ b/kernels/moe/mxfp_moe/gemm2.py @@ -923,14 +923,13 @@ def _issue_load(mr, half): if _bi + 1 < len(_blocks): _r_next, _grp_next, _col0_next = _issue_load(*_blocks[_bi + 1]) if True: - amax_f = _raw(_fabs_f32(r[0])) + amax_f = fx.absf(r[0]) for e in range_constexpr(1, 8): - abs_e = _raw(_fabs_f32(r[e])) - amax_f = arith.maxnumf(amax_f, abs_e) - amax = arith.shrui(arith.bitcast(T.i32, amax_f), _raw(fx.Int32(16))) - amax_dpp = _raw(_inline_dpp_quad_amax(amax)) - f32b = arith.shli(amax_dpp, _raw(fx.Int32(16))) - e8m0, qscale_f = _e8m0_from_amax(fx.Float32(arith.bitcast(T.f32, f32b))) + amax_f = fx.maxnumf(amax_f, fx.absf(r[e])) + amax = amax_f.bitcast(fx.Uint32) >> fx.Int32(16) + amax_dpp = _inline_dpp_quad_amax(amax) + f32b = amax_dpp << fx.Int32(16) + e8m0, qscale_f = _e8m0_from_amax(f32b.bitcast(fx.Float32)) e8 = _raw(e8m0) qscale = _raw(qscale_f) packed = _raw(fx.Int32(0)) diff --git a/kernels/moe/mxfp_moe/mxfp4_gemm_common.py b/kernels/moe/mxfp_moe/mxfp4_gemm_common.py index 1de85420c..ffb06af23 100644 --- a/kernels/moe/mxfp_moe/mxfp4_gemm_common.py +++ b/kernels/moe/mxfp_moe/mxfp4_gemm_common.py @@ -257,7 +257,7 @@ def _silu_mul_batch(gs, us): def _umax_i32(a, b): - is_gt = arith.cmpi(arith.CmpIPredicate.ugt, _raw(a), _raw(b)) + is_gt = fx.as_ir_value((a) > (b)) return fx.Int32(arith.select(is_gt, _raw(a), _raw(b)))