diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py
index 4caf085e3..5e4e6c218 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
@@ -119,18 +118,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:
@@ -141,10 +128,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)
 
@@ -349,11 +332,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)
@@ -404,11 +383,13 @@ 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):
-    return _lane_pair_reduce(_reduce_score_pair(v_s, zero_f, _fadd, fm_fast), _fadd, fm_fast)
+    reducer = lambda a, b, _fm: a + b  # noqa: E731
+    return _lane_pair_reduce(_reduce_score_pair(v_s, zero_f, reducer, fm_fast), reducer, fm_fast)
 
 
 def _sub_score_pair(v_s, row_max, fm_fast):
@@ -416,9 +397,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()
 
 
@@ -432,9 +413,9 @@ 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)
-    lo = [fmath.fma(s_lo[r], scale, neg_scaled_max, fastmath=fm_fast) for r in range_constexpr(16)]
-    hi = [fmath.fma(s_hi[r], scale, neg_scaled_max, fastmath=fm_fast) for r in range_constexpr(16)]
+    neg_scaled_max = zero_f - scale * row_max_raw
+    lo = [fx.fma(s_lo[r], scale, neg_scaled_max, fastmath=fm_fast) for r in range_constexpr(16)]
+    hi = [fx.fma(s_hi[r], scale, neg_scaled_max, fastmath=fm_fast) for r in range_constexpr(16)]
     return Vec.from_elements(lo, fx.Float32).ir_value(), Vec.from_elements(hi, fx.Float32).ir_value()
 
 
@@ -472,15 +453,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):
@@ -926,10 +907,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)
@@ -2869,51 +2850,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):
@@ -2921,10 +2900,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):
@@ -2991,7 +2970,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):
@@ -3001,13 +2979,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
@@ -3018,8 +2996,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
 
 
@@ -3049,11 +3027,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.
@@ -3222,13 +3196,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:
@@ -3632,11 +3600,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
@@ -3682,19 +3646,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)
@@ -3713,8 +3677,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(
@@ -3724,11 +3688,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)]
@@ -3740,7 +3704,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
 
@@ -3752,7 +3716,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())
@@ -4216,11 +4180,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)
@@ -4434,25 +4394,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
@@ -5012,10 +4960,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)
@@ -5027,7 +4975,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.
@@ -5060,13 +5008,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)
@@ -5087,8 +5035,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())
@@ -5106,7 +5054,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 = (
@@ -5116,7 +5064,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))
 
@@ -5125,8 +5073,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())
@@ -5143,7 +5091,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 = (
@@ -5152,7 +5100,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)
 
@@ -5372,7 +5320,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):
@@ -5385,9 +5333,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)),
@@ -5397,7 +5345,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)
@@ -5412,7 +5360,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])
@@ -5431,11 +5379,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)))
 
@@ -5492,13 +5436,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<fast>.
+    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_swa.py b/kernels/attention/pa_decode_swa.py
index 74a092c5e..829a322a7 100644
--- a/kernels/attention/pa_decode_swa.py
+++ b/kernels/attention/pa_decode_swa.py
@@ -8,6 +8,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 kernels.attention.pa_common import _compute_block_base_dw_i64, _prefetch_q_chunks
@@ -230,10 +231,10 @@ def _finish_q_fragments(
         q_abs_i32 = q_i32 & abs_mask
         q_abs = q_abs_i32.bitcast(fx.Float32)
         chunk_max = q_abs.reduce("max")
-        local_max = local_max.maximumf(chunk_max)
+        local_max = fx.maxnumf(local_max, chunk_max)
 
     for sh in [8, 4, 2, 1]:
-        local_max = local_max.maximumf(dpp_utils.dpp_xor_f32(local_max, sh))
+        local_max = fx.maxnumf(local_max, dpp_utils.dpp_xor_f32(local_max, sh))
     query_scale_lane = fx.Float32(
         arith.select(
             local_max > c_zero_f,
@@ -264,14 +265,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
@@ -474,7 +475,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(
@@ -494,9 +495,9 @@ def _store_vmax_warp(partition_start, *, seq_end=None):
                         vs_i = vector.extract(as_ir_value(vs), static_position=[i], dynamic_position=[])
                         vs_i = arith.select(kv_tok < seq_end, vs_i, zero_f)
                         vs = vector.insert(vs_i, vs, static_position=[i], dynamic_position=[])
-                v_max_warp = v_max_warp.maximumf(fx.Vector(vs).reduce("max"))
+                v_max_warp = fx.maxnumf(v_max_warp, fx.Vector(vs).reduce("max"))
             for sh in [32, 16]:
-                v_max_warp = v_max_warp.maximumf(v_max_warp.shuffle_xor(arith.constant(sh, type=T.i32), c_w))
+                v_max_warp = fx.maxnumf(v_max_warp, v_max_warp.shuffle_xor(arith.constant(sh, type=T.i32), c_w))
             fx.ptr_store(
                 fx.Vector.from_elements([v_max_warp], dtype=fx.Float32),
                 softmax_base + sm_vmax_wr_off,
@@ -559,9 +560,9 @@ def _qk_and_intra_softmax(
             if const_expr(kv_tok_base is not None):
                 logits_vec = _apply_token_mask_vec(logits_vec, td, kv_tok_base, causal_bound, seq_start, neg_inf)
                 d_out[td] = logits_vec
-            qk_max = qk_max.maximumf(fx.Vector(logits_vec).reduce("max"))
+            qk_max = fx.maxnumf(qk_max, fx.Vector(logits_vec).reduce("max"))
         for sh in [32, 16]:
-            qk_max = qk_max.maximumf(qk_max.shuffle_xor(arith.constant(sh, type=T.i32), c_w))
+            qk_max = fx.maxnumf(qk_max, qk_max.shuffle_xor(arith.constant(sh, type=T.i32), c_w))
         fx.ptr_store(
             fx.Vector.from_elements([qk_max], dtype=fx.Float32),
             softmax_base + sm_max_off,
@@ -587,20 +588,20 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs):
         partition_max = neg_inf
         partition_sum = zero_f
         warp_rescale_factors = []
-        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):
             w_max = max_vec[w]
-            partition_max = partition_max.maximumf(w_max)
+            partition_max = fx.maxnumf(partition_max, w_max)
             warp_rescale_factors.append(w_max)
-        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):
             diff_w = warp_rescale_factors[w] - partition_max
             if const_expr(needs_mask):
                 diff_w = arith.select(partition_max > neg_inf, diff_w, zero_f)
             wf = exp2_f32_fast(diff_w * fx.Float32(LOG2E).ir_value())
             w_sum = sum_vec[w]
-            wf_sum = arith.mulf(arith.unwrap(w_sum), arith.unwrap(wf), fastmath=arith.FastMathFlags.contract)
-            partition_sum = arith.addf(arith.unwrap(partition_sum), wf_sum, fastmath=arith.FastMathFlags.contract)
+            wf_sum = w_sum * wf
+            partition_sum = partition_sum + wf_sum
             warp_rescale_factors[w] = wf
 
         my_warp_rescale = warp_rescale_factors[0]
@@ -611,7 +612,7 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs):
                 my_warp_rescale,
             )
 
-        new_rmax = rmax.maximumf(partition_max)
+        new_rmax = fx.maxnumf(rmax, partition_max)
         if const_expr(needs_mask):
             accum_scale = arith.select(
                 rmax > neg_inf,
@@ -627,13 +628,9 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs):
             accum_scale = exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value())
             part_to_new = exp2_f32_fast((partition_max - new_rmax) * fx.Float32(LOG2E).ir_value())
 
-        accum_sum = arith.mulf(arith.unwrap(accum_scale), arith.unwrap(rsum), fastmath=arith.FastMathFlags.contract)
-        partition_sum_scaled = arith.mulf(
-            arith.unwrap(partition_sum),
-            arith.unwrap(part_to_new),
-            fastmath=arith.FastMathFlags.contract,
-        )
-        rsum = arith.addf(accum_sum, partition_sum_scaled, fastmath=arith.FastMathFlags.contract)
+        accum_sum = accum_scale * rsum
+        partition_sum_scaled = partition_sum * part_to_new
+        rsum = accum_sum + partition_sum_scaled
         rmax = new_rmax
         accum_scale_vec = vector.broadcast(T.f32x4, arith.unwrap(accum_scale))
         for vhe in range_constexpr(vhe_loop):
@@ -641,10 +638,10 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs):
 
         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 = v_max_global.maximumf(w_vmax)
+                v_max_global = fx.maxnumf(v_max_global, w_vmax)
             v_max_scaled = v_max_global * fx.Float32(1.0 / FP8_MAX).ir_value()
             v_max_safe_scaled = v_max_scaled + fx.Float32(1e-8 / FP8_MAX).ir_value()
             norm_factor = rcp_f32(v_max_safe_scaled)
@@ -671,7 +668,6 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs):
 
     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)
         for vhe in range_constexpr(vhe_loop):
             tmp_out = arith.constant_vector(0.0, T.f32x4)
@@ -680,7 +676,7 @@ def _pv_mfma(v_ops, outs, v_correction):
                 for j in range_constexpr(2):
                     p_elem_off = pv_prob_i64_elem_offs[vt * 2 + j]
                     p_i64 = fx.ptr_load(
-                        fx.recast_iter(fx.Int64, logits_base) + (p_elem_off),
+                        fx.recast_iter(fx.Int64, logits_base) + p_elem_off,
                         result_type=fx.Vector.make_type(1, fx.Int64),
                     )[0]
                     tmp_out = rocdl.mfma_f32_16x16x32_fp8_fp8(
@@ -694,11 +690,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 (
@@ -804,7 +796,7 @@ def pa_decode_sw_reduce_kernel(
         def _wave_reduce_max_full(val):
             red = val
             for sh in [32, 16, 8, 4, 2, 1]:
-                red = red.maximumf(red.shuffle_xor(fx.Int32(sh), c_w))
+                red = fx.maxnumf(red, red.shuffle_xor(fx.Int32(sh), c_w))
             return red
 
         def _wave_reduce_sum_full(val):
@@ -850,7 +842,7 @@ def _block_reduce(val, mode):
             def _wave_reduce_max(val):
                 red = val
                 for sh in reduce_shuffle_offsets:
-                    red = red.maximumf(red.shuffle_xor(fx.Int32(sh), c_w))
+                    red = fx.maxnumf(red, red.shuffle_xor(fx.Int32(sh), c_w))
                 return red
 
             def _wave_reduce_sum(val):
@@ -894,7 +886,7 @@ def _wave_reduce_sum(val):
             )
             inv_global_exp_sum = rcp_f32(safe_global_exp_sum)
             weight_local = scaled_sum * inv_global_exp_sum
-            weight_local_i32 = arith.bitcast(T.i32, arith.unwrap(weight_local))
+            weight_local_i32 = weight_local.bitcast(fx.Int32)
 
             acc = c_zero_f
             for part_idx in range_constexpr(max_context_partition_num):
@@ -930,7 +922,7 @@ def _wave_reduce_sum(val):
                 part_max_raw = buffer_ops.buffer_load(ml_rsrc, es_off, vec_width=1, dtype=fx.Float32)
                 part_max = arith.select(in_chunk, part_max_raw, c_neg_inf)
                 chunk_max = _block_reduce(part_max, "max")
-                global_max = global_max.maximumf(chunk_max)
+                global_max = fx.maxnumf(global_max, chunk_max)
 
             global_exp_sum = c_zero_f
             for chunk_base in range(0, max_context_partition_num, block_threads):
@@ -1042,27 +1034,28 @@ def launch_pa_decode_sw_reduce(
         num_kv_heads,
         stream: fx.Stream = fx.Stream(None),
     ):
-        pa_decode_sw_reduce_kernel(
-            output,
-            exp_sums,
-            max_logits,
-            logits,
-            stride_output_bs,
-            stride_output_len,
-            stride_output_kv_head,
-            stride_output_group_size,
-            stride_exp_sums_seq,
-            stride_exp_sums_head,
-            stride_exp_sums_part,
-            stride_logits_seq,
-            stride_logits_head,
-            stride_logits_part,
-            stride_logits_group,
-        ).launch(
-            grid=(batch_size, num_kv_heads, query_seq_len * query_group_size),
-            block=(block_threads, 1, 1),
-            stream=stream,
-        )
+        with CompilationContext.compile_hints({"fastmath": arith.FastMathFlags.contract}):
+            pa_decode_sw_reduce_kernel(
+                output,
+                exp_sums,
+                max_logits,
+                logits,
+                stride_output_bs,
+                stride_output_len,
+                stride_output_kv_head,
+                stride_output_group_size,
+                stride_exp_sums_seq,
+                stride_exp_sums_head,
+                stride_exp_sums_part,
+                stride_logits_seq,
+                stride_logits_head,
+                stride_logits_part,
+                stride_logits_group,
+            ).launch(
+                grid=(batch_size, num_kv_heads, query_seq_len * query_group_size),
+                block=(block_threads, 1, 1),
+                stream=stream,
+            )
 
     return {
         "launch": launch_pa_decode_sw_reduce,
@@ -1570,39 +1563,40 @@ def launch_pa_decode_sw(
         gz: Int32,
         stream: fx.Stream = fx.Stream(None),
     ):
-        pa_decode_sw_kernel(
-            es,
-            ml,
-            to,
-            out,
-            q,
-            kc,
-            vc,
-            bt,
-            cl,
-            ks,
-            vs,
-            s_q_seq,
-            s_q_head,
-            s_k_block,
-            s_k_head,
-            s_v_block,
-            s_v_head,
-            s_es_seq,
-            s_es_head,
-            s_es_part,
-            s_to_seq,
-            s_to_head,
-            s_to_part,
-            s_to_group,
-            s_out_bs,
-            s_out_len,
-            s_out_kv_head,
-            s_out_group_size,
-            s_bt_seq,
-            s_ks_block,
-            s_ks_head,
-        ).launch(grid=(gx, gy, gz), block=(BLOCK_THREADS, 1, 1), stream=stream)
+        with CompilationContext.compile_hints({"fastmath": arith.FastMathFlags.contract}):
+            pa_decode_sw_kernel(
+                es,
+                ml,
+                to,
+                out,
+                q,
+                kc,
+                vc,
+                bt,
+                cl,
+                ks,
+                vs,
+                s_q_seq,
+                s_q_head,
+                s_k_block,
+                s_k_head,
+                s_v_block,
+                s_v_head,
+                s_es_seq,
+                s_es_head,
+                s_es_part,
+                s_to_seq,
+                s_to_head,
+                s_to_part,
+                s_to_group,
+                s_out_bs,
+                s_out_len,
+                s_out_kv_head,
+                s_out_group_size,
+                s_bt_seq,
+                s_ks_block,
+                s_ks_head,
+            ).launch(grid=(gx, gy, gz), block=(BLOCK_THREADS, 1, 1), stream=stream)
 
     return {
         "launch": launch_pa_decode_sw,
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 916339187..dfdac3a22 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 c8f8f89bc..12db70492 100644
--- a/kernels/comm/flydsl_dispatch_combine_intranode_kernel.py
+++ b/kernels/comm/flydsl_dispatch_combine_intranode_kernel.py
@@ -307,10 +307,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 c4b00db97..e2b174fdb 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
 
@@ -63,7 +62,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..db0918052 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)
@@ -268,7 +266,7 @@ def _i8x4_in_i32_to_bf16x4_i64(val_i32, arith, vector, scale_val=None):
 
     c16 = fx.Int32(16)
     c_ffff0000 = fx.Int32(0xFFFF0000)
-    bits = [arith.bitcast(T.i32, f) for f in f32_vals]
+    bits = [f.bitcast(fx.Int32) for f in f32_vals]
     i32_lo = (bits[0] >> c16) | (bits[1] & c_ffff0000)
     i32_hi = (bits[2] >> c16) | (bits[3] & c_ffff0000)
     return _pack_i32_pair_to_i64(i32_lo, i32_hi, vector)
@@ -379,7 +377,7 @@ def _int4_to_bf16x4_i64_gfx950(packed32, nibble_offsets, arith, vector, scale_va
         # Truncate f32→bf16 via bit-shift (exact for scaled int values).
         c16_shift = fx.Int32(16)
         c_ffff0000 = fx.Int32(0xFFFF0000)
-        bf16_vals = [arith.bitcast(T.i32, _av(v)) for v in f32_vals]
+        bf16_vals = [v.bitcast(fx.Int32) for v in f32_vals]
         i32_lo = (bf16_vals[0] >> c16_shift) | (bf16_vals[1] & c_ffff0000)
         i32_hi = (bf16_vals[2] >> c16_shift) | (bf16_vals[3] & c_ffff0000)
 
@@ -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,14 +794,14 @@ 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
 
@@ -816,10 +814,10 @@ def extract_bf16_scale(arith, scale_raw_i32, ku: int):
     """
     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 (scale_raw_i32 << 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 (scale_raw_i32 & 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/gemm/hgemm_splitk.py b/kernels/gemm/hgemm_splitk.py
index 0138431cc..3da3717cd 100644
--- a/kernels/gemm/hgemm_splitk.py
+++ b/kernels/gemm/hgemm_splitk.py
@@ -288,8 +288,8 @@ def __barrier(vmcnt=0, use_s_barrier=True):
 
         def zero_c():
             # zero c if current block is the first block
-            is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0))
-            cond_ks0 = arith.cmpi(arith.CmpIPredicate.eq, ks_idx, fx.Index(0))
+            is_t0_cond = as_ir_value(fx.Index(tid) == fx.Index(0))
+            cond_ks0 = as_ir_value(ks_idx == fx.Index(0))
             cond_ks0_if = scf.IfOp(cond_ks0, results_=[], has_else=False)
             with ir.InsertionPoint(cond_ks0_if.then_block):
                 zero_vec = vector.broadcast(T.vec(LDG_VEC_SIZE, dtype_), c_zero_d)
@@ -301,11 +301,11 @@ def zero_c():
                     init_vec = zero_vec
                     if const_expr(HAS_BIAS):
                         init_vec = BIAS_.vec_load((n_offset + n_local_idx,), LDG_VEC_SIZE)
-                    cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m))
+                    cond_boundary = as_ir_value(row_idx < fx.Index(m))
                     cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                     with ir.InsertionPoint(cond_boundary_if.then_block):
                         bytes_offset = C_.linear_offset((row_idx, n_offset + n_local_idx))
-                        bytes_offset_i32 = arith.index_cast(T.i32, bytes_offset)
+                        bytes_offset_i32 = as_ir_value(fx.Int32(bytes_offset))
                         c_ptr = get_llvm_ptr(C, bytes_offset_i32, DTYPE_BYTES)
                         llvm.InlineAsmOp(
                             None,
@@ -333,7 +333,7 @@ def zero_c():
 
         def split_k_barrier():
             # spin-wait until signal triggered
-            is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0))
+            is_t0_cond = as_ir_value(fx.Index(tid) == fx.Index(0))
             is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False)
             with ir.InsertionPoint(is_t0_cond_if.then_block):
                 init_cur = arith.constant(0, type=T.i32)
@@ -362,7 +362,7 @@ def split_k_barrier():
             is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False)
             with ir.InsertionPoint(is_t0_cond_if.then_block):
                 arrive_idx = atomic_add(semaphore, signal_idx, arith.constant(1, type=T.i32), dtype_bytes=4)
-                cond_ksl = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(arrive_idx), fx.Index(SPLIT_K - 1))
+                cond_ksl = as_ir_value(fx.Index(arrive_idx) == fx.Index(SPLIT_K - 1))
                 cond_ksl_if = scf.IfOp(cond_ksl, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_ksl_if.then_block):
                     semaphore_[signal_idx] = arith.constant(0, type=T.i32)
@@ -379,7 +379,7 @@ def ldg_a(k_offset):
                 k_local_idx = global_tid % LDG_A_X_THREADS * LDG_VEC_SIZE
                 row_idx = m_offset + fx.Index(m_local_idx)
                 safe_row_idx = arith.select(
-                    arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m)),
+                    row_idx < fx.Index(m),
                     row_idx,
                     fx.Index(0),
                 )
@@ -401,7 +401,7 @@ def sts_a(vecs, lds_stage):
         def get_dma_copy_warp_offset():
             warp_offset = rocdl.readfirstlane(
                 T.i64,
-                arith.index_cast(T.i64, fx.Index(wid) * arith.constant(WARP_SIZE * DMA_BYTES, index=True)),
+                as_ir_value(fx.Int64(wid) * fx.Int64(WARP_SIZE * DMA_BYTES)),
             )
             return warp_offset
 
@@ -425,14 +425,14 @@ def ldg_sts_a_async(k_offset, lds_stage):
                 col_in_bytes = swizzle_xor16(m_local_idx, col_in_bytes, k_blocks16)
                 row_idx = m_offset + fx.Index(m_local_idx)
                 safe_row_idx = arith.select(
-                    arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m)),
+                    row_idx < fx.Index(m),
                     row_idx,
                     fx.Index(0),
                 )
                 col_idx = fx.Index(k_offset + col_in_bytes // DTYPE_BYTES)
                 # get offset
                 global_offset = A_.linear_offset((safe_row_idx, col_idx)) * DTYPE_BYTES
-                global_offset = arith.index_cast(T.i32, global_offset)
+                global_offset = as_ir_value(fx.Int32(global_offset))
                 # get lds ptr
                 if const_expr(i == 0):
                     lds_byte_offset = as_off(lds_stage, 0, 0) * DTYPE_BYTES
@@ -456,14 +456,14 @@ def ldg_sts_b_async(k_offset, lds_stage):
                 col_in_bytes = swizzle_xor16(n_local_idx, col_in_bytes, k_blocks16)
                 row_idx = n_offset + fx.Index(n_local_idx)
                 safe_row_idx = arith.select(
-                    arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(n)),
+                    row_idx < fx.Index(n),
                     row_idx,
                     fx.Index(0),
                 )
                 col_idx = fx.Index(k_offset + col_in_bytes // DTYPE_BYTES)
                 # get offset
                 global_offset = B_.linear_offset((safe_row_idx, col_idx)) * DTYPE_BYTES
-                global_offset = arith.index_cast(T.i32, global_offset)
+                global_offset = as_ir_value(fx.Int32(global_offset))
                 # get lds ptr
                 if const_expr(i == 0):
                     lds_byte_offset = bs_off(lds_stage, 0, 0) * DTYPE_BYTES
@@ -586,7 +586,7 @@ def hot_loop_scheduler():
                 # ================ Reordered ================
                 rocdl.sched_barrier(0)
 
-            init_state = [ks_begin, arith.constant(0, index=True)] + c_frags
+            init_state = [ks_begin, fx.Index(0)] + c_frags
             for bki, state in range(0, BLOCK_K_LOOPS - (STAGES - 1), 1, init=init_state):
                 k_offset = state[0]
                 current_stage = fx.Index(state[1])
@@ -629,7 +629,7 @@ def hot_loop_scheduler():
                 # ================ Reordered ================
                 rocdl.sched_barrier(0)
 
-            init_state = [ks_begin, arith.constant(0, index=True)] + c_frags + b_frags_next
+            init_state = [ks_begin, fx.Index(0)] + c_frags + b_frags_next
             for bki, state in range(1, BLOCK_K_LOOPS, init=init_state):
                 k_offset = state[0]
                 current_stage = fx.Index(state[1])
@@ -682,7 +682,7 @@ def hot_loop_scheduler():
                 n_local_idx = fx.Int64(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE)
                 m_global_idx = m_offset + m_local_idx
                 n_global_idx = n_offset + n_local_idx
-                cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, m_global_idx, fx.Index(m))
+                cond_boundary = as_ir_value(m_global_idx < fx.Index(m))
                 cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_boundary_if.then_block):
                     pk_val = fx.ptr_load(
@@ -719,7 +719,7 @@ def hot_loop_scheduler():
                 m_local_idx = fx.Int64(global_tid // LDG_C_X_THREADS)
                 n_local_idx = fx.Int64(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE)
                 m_global_idx = m_offset + m_local_idx
-                cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, m_global_idx, fx.Index(m))
+                cond_boundary = as_ir_value(m_global_idx < fx.Index(m))
                 cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_boundary_if.then_block):
                     vec = fx.ptr_load(
diff --git a/kernels/gemm/small_m_hgemm.py b/kernels/gemm/small_m_hgemm.py
index 32b685be9..cb882e6a0 100644
--- a/kernels/gemm/small_m_hgemm.py
+++ b/kernels/gemm/small_m_hgemm.py
@@ -630,12 +630,7 @@ def _store(offset, value, vec_size=1):
         ]
         tile_n_offsets = [tile_block_n_idx * fx.Index(BLOCK_N) for tile_block_n_idx in tile_block_n_indices]
         tile_actives = [
-            arith.cmpi(
-                arith.CmpIPredicate.ult,
-                tile_block_n_idx,
-                fx.Index(block_n_tiles),
-            )
-            for tile_block_n_idx in tile_block_n_indices
+            as_ir_value(tile_block_n_idx < fx.Index(block_n_tiles)) for tile_block_n_idx in tile_block_n_indices
         ]
         tile_signal_indices = [
             fx.block_idx.x * fx.Int32(block_n_tiles) + arith.index_cast(T.i32, tile_block_n_idx)
@@ -667,14 +662,14 @@ def zero_c_tile(c_g, bias_g, tile_n_offset):
                 init_vec = zero_vec
                 if const_expr(HAS_BIAS):
                     init_vec = bias_g.vec_load((tile_n_offset + n_local_idx,), LDG_VEC_SIZE)
-                cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m))
+                cond_boundary = as_ir_value(row_idx < fx.Index(m))
                 cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_boundary_if.then_block):
                     c_g.vec_store((row_idx, tile_n_offset + n_local_idx), init_vec, LDG_VEC_SIZE)
                     scf.YieldOp([])
 
         def prepare_split_k_tile(c_g, bias_g, tile_n_offset, tile_signal_idx):
-            is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0))
+            is_t0_cond = as_ir_value(fx.Index(tid) == fx.Index(0))
             is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False)
             with ir.InsertionPoint(is_t0_cond_if.then_block):
                 semaphore_ptr = get_llvm_ptr(semaphore, tile_signal_idx, 4)
@@ -691,7 +686,7 @@ def prepare_split_k_tile(c_g, bias_g, tile_n_offset, tile_signal_idx):
             gpu.barrier()
             arrive_idx = fx.Index(bc_[0])
 
-            first_arrival = arith.cmpi(arith.CmpIPredicate.eq, arrive_idx, fx.Index(0))
+            first_arrival = as_ir_value(arrive_idx == fx.Index(0))
             first_arrival_if = scf.IfOp(first_arrival, results_=[], has_else=False)
             with ir.InsertionPoint(first_arrival_if.then_block):
                 zero_c_tile(c_g, bias_g, tile_n_offset)
@@ -740,7 +735,7 @@ def split_k_barrier(tile_signal_idx):
             rocdl.sched_barrier(0)
             gpu.barrier()
 
-            is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0))
+            is_t0_cond = as_ir_value(fx.Index(tid) == fx.Index(0))
             is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[T.i32], has_else=True)
             with ir.InsertionPoint(is_t0_cond_if.then_block):
                 semaphore_ptr = get_llvm_ptr(semaphore, tile_signal_idx, 4)
@@ -776,13 +771,9 @@ def ldg_a(k_offset):
                 k_local_idx = global_tid % LDG_A_X_THREADS * LDG_VEC_SIZE
                 row_idx = m_offset + fx.Index(m_local_idx)
                 col_idx = fx.Index(k_offset + k_local_idx)
-                slot_valid = arith.cmpi(
-                    arith.CmpIPredicate.ult,
-                    fx.Index(global_tid),
-                    fx.Index(LDG_A_TOTAL_VECS),
-                )
-                valid_row = arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m))
-                can_load = arith.andi(slot_valid, valid_row)
+                slot_valid = fx.Index(global_tid) < fx.Index(LDG_A_TOTAL_VECS)
+                valid_row = row_idx < fx.Index(m)
+                can_load = as_ir_value(slot_valid & valid_row)
                 load_if = scf.IfOp(
                     can_load,
                     results_=[T.vec(LDG_VEC_SIZE, dtype_)],
@@ -802,11 +793,7 @@ def sts_a(vecs, lds_stage):
                 k_local_idx = global_tid % LDG_A_X_THREADS * LDG_VEC_SIZE
                 col_in_bytes = k_local_idx * DTYPE_BYTES
                 col_in_bytes = swizzle_xor16(m_local_idx, col_in_bytes, k_blocks16)
-                slot_valid = arith.cmpi(
-                    arith.CmpIPredicate.ult,
-                    fx.Index(global_tid),
-                    fx.Index(LDG_A_TOTAL_VECS),
-                )
+                slot_valid = as_ir_value(fx.Index(global_tid) < fx.Index(LDG_A_TOTAL_VECS))
                 store_if = scf.IfOp(slot_valid, results_=[], has_else=False)
                 with ir.InsertionPoint(store_if.then_block):
                     as_.vec_store(
@@ -825,14 +812,10 @@ def ldg_sts_a_async(k_offset, lds_stage):
                 col_in_bytes = swizzle_xor16(m_local_idx, col_in_bytes, k_blocks16)
                 row_idx = m_offset + fx.Index(m_local_idx)
                 col_idx = fx.Index(k_offset + col_in_bytes // DTYPE_BYTES)
-                slot_valid = arith.cmpi(
-                    arith.CmpIPredicate.ult,
-                    fx.Index(global_tid),
-                    fx.Index(LDG_A_TOTAL_VECS_AS),
-                )
+                slot_valid = as_ir_value(fx.Index(global_tid) < fx.Index(LDG_A_TOTAL_VECS_AS))
                 slot_if = scf.IfOp(slot_valid, results_=[], has_else=False)
                 with ir.InsertionPoint(slot_if.then_block):
-                    valid_row = arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m))
+                    valid_row = as_ir_value(row_idx < fx.Index(m))
                     cond_if = scf.IfOp(valid_row, results_=[], has_else=True)
                     with ir.InsertionPoint(cond_if.then_block):
                         global_offset = A_.linear_offset((row_idx, col_idx)) * DTYPE_BYTES
@@ -924,7 +907,7 @@ def store_split_k_tile(c_tensor, c_g, c_s, tile_n_offset):
                 n_local_idx = fx.Index(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE)
                 m_global_idx = m_offset + m_local_idx
                 n_global_idx = tile_n_offset + n_local_idx
-                cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, m_global_idx, fx.Index(m))
+                cond_boundary = as_ir_value(m_global_idx < fx.Index(m))
                 cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_boundary_if.then_block):
                     pk_val = c_s.vec_load((m_local_idx, n_local_idx), LDG_VEC_SIZE)
@@ -970,7 +953,7 @@ def store_c_tile(bias_g, c_g, c_s, tile_n_offset):
                 m_local_idx = fx.Index(global_tid // LDG_C_X_THREADS)
                 n_local_idx = fx.Index(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE)
                 m_global_idx = m_offset + m_local_idx
-                cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, m_global_idx, fx.Index(m))
+                cond_boundary = as_ir_value(m_global_idx < fx.Index(m))
                 cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_boundary_if.then_block):
                     vec = c_s.vec_load((m_local_idx, n_local_idx), LDG_VEC_SIZE)
@@ -1020,11 +1003,7 @@ def ldg_sts_b_async(bs_s, k_offset, lds_stage, tile_n_offset):
                     col_in_bytes = k_local_idx * DTYPE_BYTES
                     col_in_bytes = swizzle_xor16(n_local_idx, col_in_bytes, k_blocks16)
                     col_idx = fx.Index(k_offset + col_in_bytes // DTYPE_BYTES)
-                    slot_valid = arith.cmpi(
-                        arith.CmpIPredicate.ult,
-                        fx.Index(global_tid),
-                        fx.Index(LDG_B_TOTAL_VECS_AS),
-                    )
+                    slot_valid = as_ir_value(fx.Index(global_tid) < fx.Index(LDG_B_TOTAL_VECS_AS))
                     slot_if = scf.IfOp(slot_valid, results_=[], has_else=False)
                     with ir.InsertionPoint(slot_if.then_block):
                         global_offset = B_.linear_offset((tile_n_offset + fx.Index(n_local_idx), col_idx))
@@ -1101,17 +1080,13 @@ def hot_loop_scheduler():
                     rocdl.sched_barrier(0)
 
                 UNROLL = EFFECTIVE_B_TO_LDS_UNROLL
-                init_state = [ks_begin, arith.constant(0, index=True)] + c_frags_local
+                init_state = [ks_begin, fx.Index(0)] + c_frags_local
                 for bki, state in range(0, BLOCK_K_LOOPS - 1, UNROLL, init=init_state):
                     k_offset = state[0]
                     current_stage = fx.Index(state[1])
                     c_frags_local = state[2 : 2 + C_FRAGS_LEN]
                     for unroll_i in range_constexpr(UNROLL):
-                        cond = arith.cmpi(
-                            arith.CmpIPredicate.ult,
-                            fx.Index(bki + unroll_i),
-                            fx.Index(BLOCK_K_LOOPS - 1),
-                        )
+                        cond = as_ir_value(fx.Index(bki + unroll_i) < fx.Index(BLOCK_K_LOOPS - 1))
                         cond_if = scf.IfOp(
                             cond,
                             results_=[T.vec(WMMA_C_FRAG_VALUES, T.f32)] * C_FRAGS_LEN + [T.index, T.i32],
@@ -1183,7 +1158,7 @@ def hot_loop_scheduler():
 
             TOTAL_C_FRAGS_LEN = C_FRAGS_LEN * N_TILE_REPEAT
             TOTAL_B_FRAGS_LEN = B_FRAGS_LEN * N_TILE_REPEAT
-            init_state = [ks_begin, arith.constant(0, index=True)] + c_frags + a_frags + b_frags
+            init_state = [ks_begin, fx.Index(0)] + c_frags + a_frags + b_frags
             for _, state in range(1, BLOCK_K_LOOPS, init=init_state):
                 k_offset = state[0]
                 current_stage = fx.Index(state[1])
diff --git a/kernels/gemm/splitk_hgemm.py b/kernels/gemm/splitk_hgemm.py
index a19cb8282..93bd1be73 100644
--- a/kernels/gemm/splitk_hgemm.py
+++ b/kernels/gemm/splitk_hgemm.py
@@ -305,7 +305,7 @@ def swizzle_for_cache_reuse(pid):
 
         def zero_c():
             # get arrive index within split-k group
-            is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0))
+            is_t0_cond = as_ir_value(fx.Index(tid) == fx.Index(0))
             is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False)
             with ir.InsertionPoint(is_t0_cond_if.then_block):
                 semaphore_ptr = get_llvm_ptr(semaphore, signal_idx, 4)
@@ -322,7 +322,7 @@ def zero_c():
             gpu.barrier()
             arrive_idx = fx.Index(fx.memref_load(bc_view, 0))
             # zero c if current block is the first arrived block
-            cond_ks0 = arith.cmpi(arith.CmpIPredicate.eq, arrive_idx, fx.Index(0))
+            cond_ks0 = as_ir_value(arrive_idx == fx.Index(0))
             cond_ks0_if = scf.IfOp(cond_ks0, results_=[], has_else=False)
             with ir.InsertionPoint(cond_ks0_if.then_block):
                 zero_vec = vector.broadcast(T.vec(LDG_VEC_SIZE, dtype_), c_zero_d)
@@ -334,11 +334,11 @@ def zero_c():
                     init_vec = zero_vec
                     if const_expr(HAS_BIAS):
                         init_vec = BIAS_.vec_load((n_offset + n_local_idx,), LDG_VEC_SIZE)
-                    cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m))
+                    cond_boundary = as_ir_value(row_idx < fx.Index(m))
                     cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                     with ir.InsertionPoint(cond_boundary_if.then_block):
                         bytes_offset = C_.linear_offset((row_idx, n_offset + n_local_idx))
-                        bytes_offset_i32 = arith.index_cast(T.i32, bytes_offset)
+                        bytes_offset_i32 = as_ir_value(fx.Int32(bytes_offset))
                         c_ptr = get_llvm_ptr(C, bytes_offset_i32, DTYPE_BYTES)
                         llvm.InlineAsmOp(
                             None,
@@ -373,7 +373,7 @@ def zero_c():
 
         def split_k_barrier():
             # spin-wait until signal triggered
-            is_t0_cond = arith.cmpi(arith.CmpIPredicate.eq, fx.Index(tid), fx.Index(0))
+            is_t0_cond = as_ir_value(fx.Index(tid) == fx.Index(0))
             is_t0_cond_if = scf.IfOp(is_t0_cond, results_=[], has_else=False)
             with ir.InsertionPoint(is_t0_cond_if.then_block):
                 init_cur = arith.constant(0, type=T.i32)
@@ -410,11 +410,7 @@ def split_k_barrier():
                     syncscope="agent",
                     alignment=4,
                 ).result
-                cond_ksl = arith.cmpi(
-                    arith.CmpIPredicate.eq,
-                    fx.Index(arrive_idx),
-                    fx.Index(2 * SPLIT_K - 1),
-                )
+                cond_ksl = as_ir_value(fx.Index(arrive_idx) == fx.Index(2 * SPLIT_K - 1))
                 cond_ksl_if = scf.IfOp(cond_ksl, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_ksl_if.then_block):
                     semaphore_[signal_idx] = arith.constant(0, type=T.i32)
@@ -431,7 +427,7 @@ def ldg_a(k_offset):
                 k_local_idx = global_tid % LDG_A_X_THREADS * LDG_VEC_SIZE
                 row_idx = m_offset + fx.Index(m_local_idx)
                 safe_row_idx = arith.select(
-                    arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m)),
+                    row_idx < fx.Index(m),
                     row_idx,
                     fx.Index(0),
                 )
@@ -457,7 +453,7 @@ def ldg_b(k_offset):
                 k_local_idx = global_tid % LDG_B_X_THREADS * LDG_VEC_SIZE
                 row_idx = n_offset + fx.Index(n_local_idx)
                 safe_row_idx = arith.select(
-                    arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(n)),
+                    row_idx < fx.Index(n),
                     row_idx,
                     fx.Index(0),
                 )
@@ -478,10 +474,7 @@ def sts_b(vecs, lds_stage):
         def get_dma_copy_warp_offset():
             warp_offset = rocdl.readfirstlane(
                 T.i64,
-                arith.index_cast(
-                    T.i64,
-                    fx.Index(wid) * arith.constant(WARP_SIZE * DMA_BYTES, index=True),
-                ),
+                as_ir_value(fx.Int64(wid) * fx.Int64(WARP_SIZE * DMA_BYTES)),
             )
             return warp_offset
 
@@ -494,14 +487,14 @@ def ldg_sts_a_async(k_offset, lds_stage):
                 col_in_bytes = swizzle_xor16(m_local_idx, col_in_bytes, k_blocks16)
                 row_idx = m_offset + fx.Index(m_local_idx)
                 safe_row_idx = arith.select(
-                    arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(m)),
+                    row_idx < fx.Index(m),
                     row_idx,
                     fx.Index(0),
                 )
                 col_idx = fx.Index(k_offset + col_in_bytes // DTYPE_BYTES)
                 # get offset
                 global_offset = A_.linear_offset((safe_row_idx, col_idx)) * DTYPE_BYTES
-                global_offset = arith.index_cast(T.i32, global_offset)
+                global_offset = as_ir_value(fx.Int32(global_offset))
                 # get lds ptr
                 if const_expr(i == 0):
                     lds_ptr_base = _lds_a3_ptr(a_lds_i64, fx.Index(lds_stage) * (BLOCK_M * BLOCK_K))
@@ -531,14 +524,14 @@ def ldg_sts_b_async(k_offset, lds_stage):
                 col_in_bytes = swizzle_xor16(n_local_idx, col_in_bytes, k_blocks16)
                 row_idx = n_offset + fx.Index(n_local_idx)
                 safe_row_idx = arith.select(
-                    arith.cmpi(arith.CmpIPredicate.ult, row_idx, fx.Index(n)),
+                    row_idx < fx.Index(n),
                     row_idx,
                     fx.Index(0),
                 )
                 col_idx = fx.Index(k_offset + col_in_bytes // DTYPE_BYTES)
                 # get offset
                 global_offset = B_.linear_offset((safe_row_idx, col_idx)) * DTYPE_BYTES
-                global_offset = arith.index_cast(T.i32, global_offset)
+                global_offset = as_ir_value(fx.Int32(global_offset))
                 # get lds ptr
                 if const_expr(i == 0):
                     lds_ptr_base = _lds_a3_ptr(b_lds_i64, fx.Index(lds_stage) * (BLOCK_N * BLOCK_K))
@@ -672,7 +665,7 @@ def hot_loop_scheduler():
                 # ================ Reordered ================
                 rocdl.sched_barrier(0)
 
-            init_state = [ks_begin, arith.constant(0, index=True)] + c_frags + b_frags_next
+            init_state = [ks_begin, fx.Index(0)] + c_frags + b_frags_next
             for bki, state in range(0, BLOCK_K_LOOPS - 1, 1, init=init_state):
                 k_offset = state[0]
                 current_stage = fx.Index(state[1])
@@ -733,7 +726,7 @@ def hot_loop_scheduler():
                         rocdl.sched_mfma(mfma_.consume(AVG_MFMA_COUNT))
                 rocdl.sched_barrier(0)
 
-            init_state = [ks_begin, arith.constant(0, index=True)] + c_frags + a_frags + b_frags
+            init_state = [ks_begin, fx.Index(0)] + c_frags + a_frags + b_frags
             for bki, state in range(1, BLOCK_K_LOOPS, init=init_state):
                 k_offset = state[0]
                 current_stage = fx.Index(state[1])
@@ -790,7 +783,7 @@ def hot_loop_scheduler():
                 n_local_idx = fx.Index(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE)
                 m_global_idx = m_offset + m_local_idx
                 n_global_idx = n_offset + n_local_idx
-                cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, m_global_idx, fx.Index(m))
+                cond_boundary = as_ir_value(m_global_idx < fx.Index(m))
                 cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_boundary_if.then_block):
                     pk_val = cs_load_vec(m_local_idx, n_local_idx, LDG_VEC_SIZE)
@@ -829,7 +822,7 @@ def hot_loop_scheduler():
                 m_local_idx = fx.Index(global_tid // LDG_C_X_THREADS)
                 n_local_idx = fx.Index(global_tid % LDG_C_X_THREADS * LDG_VEC_SIZE)
                 m_global_idx = m_offset + m_local_idx
-                cond_boundary = arith.cmpi(arith.CmpIPredicate.ult, m_global_idx, fx.Index(m))
+                cond_boundary = as_ir_value(m_global_idx < fx.Index(m))
                 cond_boundary_if = scf.IfOp(cond_boundary, results_=[], has_else=False)
                 with ir.InsertionPoint(cond_boundary_if.then_block):
                     vec = cs_load_vec(m_local_idx, n_local_idx, LDG_VEC_SIZE)
diff --git a/kernels/moe/moe_gemm_2stage/gemm1.py b/kernels/moe/moe_gemm_2stage/gemm1.py
index 2c90d2bee..cb3b11713 100644
--- a/kernels/moe/moe_gemm_2stage/gemm1.py
+++ b/kernels/moe/moe_gemm_2stage/gemm1.py
@@ -331,7 +331,7 @@ def silu(x):
             fx.make_layout((tokens_i32_v, k_i32_v), stride=(k_i32_v, 1))
 
             # B preshuffle layout: match GEMM test helper exactly.
-            c_n_total = arith.index(experts * (2 * inter_dim))
+            c_n_total = fx.Index(experts * (2 * inter_dim))
             # For packed int4 (W4A8/W4A16/W4A_FP8), kpack_bytes=8.
             kpack_bytes = 8 if w_is_int4 else 16
             w_elem_bytes = 1 if w_is_int4 else elem_bytes
@@ -343,7 +343,7 @@ def silu(x):
                 elem_bytes=w_elem_bytes,
             )
             layout_b = b_layout.layout_b
-            (k_in * arith.index(int(elem_bytes))) // fx.Index(64)
+            (k_in * fx.Index(int(elem_bytes))) // fx.Index(64)
 
             shape_lds = fx.make_shape(tile_m, tile_k)
             stride_lds = fx.make_stride(lds_stride, 1)
@@ -358,9 +358,9 @@ def silu(x):
 
             if const_expr(_is_splitk):
                 bz = gpu.block_id("z")  # K-batch id
-                k_base_idx = bz * arith.index(_k_per_batch)
+                k_base_idx = bz * fx.Index(_k_per_batch)
             else:
-                k_base_idx = arith.index(0)
+                k_base_idx = fx.Index(0)
 
             # Block validity: compute as early as possible so invalid blocks skip all buffer-resource
             # setup, LDS pointer math, and gmem prefetch work.
@@ -371,17 +371,17 @@ def silu(x):
                 num_records_bytes=fx.Index(4),
             )
             max_token_id_i32 = buffer_ops.buffer_load(maxids_rsrc, fx.Index(0), vec_width=1, dtype=T.i32)
-            bx_m_i32 = arith.index_cast(T.i32, bx_m)
-            blk_valid = arith.cmpi(arith.CmpIPredicate.ult, bx_m_i32, max_token_id_i32)
+            bx_m_i32 = fx.Int32(bx_m)
+            blk_valid = bx_m_i32 < max_token_id_i32
             # Common constants/atoms (hoisted): keep IR small like GEMM.
             # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs).
-            k_blocks16 = arith.index(tile_k_bytes // 16)
+            k_blocks16 = fx.Index(tile_k_bytes // 16)
             layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1))
             layout_lane16 = fx.make_layout((4, 16), stride=(16, 1))
 
             # Everything below is gated by `blk_valid` to avoid doing buffer-resource setup and
             # gmem work for padding blocks.
-            _if_blk = scf.IfOp(blk_valid)
+            _if_blk = scf.IfOp(as_ir_value(blk_valid))
             with _if_then(_if_blk):
                 base_ptr = allocator.get_base()
                 lds_x_ptr = SmemPtr(
@@ -411,7 +411,7 @@ def silu(x):
 
                 # X: [tokens, k] bytes = tokens*k*elem_bytes
                 x_rows = tokens_in * (c_topk if x_is_token_slot else fx.Index(1))
-                x_nbytes_idx = x_rows * k_in * arith.index(int(elem_bytes))
+                x_nbytes_idx = x_rows * k_in * fx.Index(int(elem_bytes))
                 x_rsrc = buffer_ops.create_buffer_resource(arg_x, max_size=False, num_records_bytes=x_nbytes_idx)
 
                 w_rsrc = buffer_ops.create_buffer_resource(arg_w, max_size=False)
@@ -453,7 +453,7 @@ def silu(x):
                 # Expert id for this M tile (keep address math in `index`)
                 expert_i32 = buffer_ops.buffer_load(expert_rsrc, bx, vec_width=1, dtype=T.i32)
                 expert_idx = arith.index_cast(T.index, expert_i32)
-                inter2_idx = arith.index(2 * inter_dim)
+                inter2_idx = fx.Index(2 * inter_dim)
                 expert_off_idx = expert_idx * inter2_idx  # index
 
                 # ---- X gmem->reg prefetch (match preshuffle GEMM mapping) ----
@@ -478,15 +478,15 @@ def silu(x):
                 num_x_loads = bytes_per_thread_x // x_load_bytes
                 chunk_i32 = x_load_bytes // 4  # dwords per chunk (1/2/4)
 
-                c_k_div4 = (k_in * arith.index(int(elem_bytes))) // fx.Index(4)
-                c_k_div4_i32 = arith.index_cast(T.i32, c_k_div4)
+                c_k_div4 = (k_in * fx.Index(int(elem_bytes))) // fx.Index(4)
+                c_k_div4_i32 = fx.Int32(c_k_div4)
                 fx.make_layout((tokens_i32_v, c_k_div4_i32), stride=(c_k_div4_i32, 1))
                 tile_k_dwords = (int(tile_k) * int(elem_bytes)) // 4
                 layout_x_tile_div4 = fx.make_layout((tile_m, tile_k_dwords), stride=(tile_k_dwords, 1))
                 c_chunk_i32 = fx.Index(chunk_i32)
                 tx_i32_base = tx * c_chunk_i32
                 mask24 = fx.Int32(0xFFFFFF)
-                tokens_i32 = arith.index_cast(T.i32, tokens_in)
+                tokens_i32 = fx.Int32(tokens_in)
                 topk_i32 = fx.Int32(topk)
 
                 def x_tile_chunk_coord_i32(i: int):
@@ -515,7 +515,7 @@ def x_tile_chunk_coord_i32(i: int):
                     t_raw = fused_i & mask24
                     # NOTE: aiter moe_sorting uses sentinel token_id == tokens for padding.
                     # Do NOT rely on buffer OOB semantics for X loads; explicitly mask to a safe row.
-                    t_valid_i32 = arith.cmpi(arith.CmpIPredicate.ult, t_raw, tokens_i32)
+                    t_valid_i32 = t_raw < tokens_i32
                     if const_expr(x_is_token_slot):
                         s_raw = fused_i >> 24
                         # X is indexed by token-slot in **slot-major** order:
@@ -557,7 +557,7 @@ def load_x(idx_i32):
 
                 def load_x_tile(base_k):
                     """Prefetch the per-thread X tile portion (gmem -> regs) for a given K base (in elements)."""
-                    base_k_div4 = (base_k * arith.index(int(elem_bytes))) // fx.Index(4)
+                    base_k_div4 = (base_k * fx.Index(int(elem_bytes))) // fx.Index(4)
                     parts = []
                     for i in range_constexpr(num_x_loads):
                         idx_i32 = x_row_base_div4[i] + base_k_div4 + x_col_local_i32[i]
@@ -583,9 +583,9 @@ def load_x_tile(base_k):
                 # may differ (e.g. 8 for int4 weights), but that only affects B preshuffle.
                 row_a_lds = lane_mod_16
                 a_kpack_elems = 16 // elem_bytes
-                col_offset_base = lane_div_16 * arith.index(int(a_kpack_elems))
+                col_offset_base = lane_div_16 * fx.Index(int(a_kpack_elems))
                 col_offset_base_bytes = (
-                    col_offset_base if elem_bytes == 1 else (col_offset_base * arith.index(int(elem_bytes)))
+                    col_offset_base if elem_bytes == 1 else (col_offset_base * fx.Index(int(elem_bytes)))
                 )
 
                 # Dynamic N tiling within block (same as existing kernels)
@@ -608,7 +608,7 @@ def load_x_tile(base_k):
                 c_n0_static = experts * (2 * inter_dim) // 16
                 layout_n_blk_intra = fx.make_layout((c_n0_static, 16), stride=(16, 1))
                 for ni in range_constexpr(num_acc_n):
-                    offset = arith.index(ni * 16)
+                    offset = fx.Index(ni * 16)
                     col_g = by_n + n_tile_base
                     col_g = col_g + offset
                     col_g = col_g + lane_mod_16
@@ -782,7 +782,7 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_base):
                 def lds_load_packs_k64(curr_row_a_lds, col_base_bytes, lds_base):
                     col_base_swz_bytes = swizzle_xor16(curr_row_a_lds, col_base_bytes, k_blocks16)
                     col_base_swz = (
-                        col_base_swz_bytes if elem_bytes == 1 else (col_base_swz_bytes // arith.index(int(elem_bytes)))
+                        col_base_swz_bytes if elem_bytes == 1 else (col_base_swz_bytes // fx.Index(int(elem_bytes)))
                     )
                     idx_a16 = preshuffle_crd2idx((fx.Int32(curr_row_a_lds), fx.Int32(col_base_swz)), layout_lds)
                     idx_a16 = idx_a16 + lds_base
@@ -883,11 +883,11 @@ def _acc_scaled_f32(f32_acc_vec, f32_partial_vec, scale_val):
                         for ku in range_constexpr(k_unroll):
                             b_gate_raw = b_gate_tile_in[ku]
                             b_up_raw = b_up_tile_in[ku]
-                            ki64 = arith.index(ku * 64)
+                            ki64 = fx.Index(ku * 64)
                             col_base = col_offset_base_bytes + ki64
 
                             for mi in range_constexpr(m_repeat):
-                                mi_val = arith.index(mi * 16)
+                                mi_val = fx.Index(mi * 16)
                                 curr_row_a_lds = row_a_lds + mi_val
 
                                 if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)):
@@ -967,11 +967,11 @@ def _acc_scaled_f32(f32_acc_vec, f32_partial_vec, scale_val):
                         for ku in range_constexpr(k_unroll):
                             b_gate_packs0, b_gate_packs1 = b_gate_tile_in[ku]
                             b_up_packs0, b_up_packs1 = b_up_tile_in[ku]
-                            ki64 = arith.index(ku * 64)
+                            ki64 = fx.Index(ku * 64)
                             col_base = col_offset_base_bytes + ki64
 
                             for mi in range_constexpr(m_repeat):
-                                mi_val = arith.index(mi * 16)
+                                mi_val = fx.Index(mi * 16)
                                 curr_row_a_lds = row_a_lds + mi_val
 
                                 if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)):
@@ -998,7 +998,7 @@ def _acc_scaled_f32(f32_acc_vec, f32_partial_vec, scale_val):
                     return gate_list, up_list, epilogue_pf
 
                 # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ----------------
-                lds_tile_elems = arith.index(tile_m * lds_stride)
+                lds_tile_elems = fx.Index(tile_m * lds_stride)
                 lds_base_cur = fx.Index(0)
                 lds_base_nxt = lds_tile_elems
 
@@ -1053,8 +1053,8 @@ def hot_loop_scheduler():
 
                 # Ping-pong main loop (2 tiles per iteration), leaving 2 tail tiles.
                 # Uses scf.for with loop-carried accumulators, B-tile prefetch, and A0 LDS prefetch.
-                arith.index(tile_k * 2)
-                c_tile_k = arith.index(tile_k)
+                fx.Index(tile_k * 2)
+                c_tile_k = fx.Index(tile_k)
                 total_tiles = int(_k_per_batch) // int(tile_k)
                 pair_iters = max((total_tiles - 2) // 2, 0)
 
@@ -1184,7 +1184,7 @@ def _unflatten_b_tile(vals):
                     b_gate_cur = _unflatten_b_tile(list(loop_results[_p_bg:_p_bu]))
                     b_up_cur = _unflatten_b_tile(list(loop_results[_p_bu:_p_a0]))
                     a0_prefetch_pong = (loop_results[_p_a0], loop_results[_p_a0 + 1])
-                k_tail1 = k_base_idx + arith.index(_k_per_batch - tile_k)
+                k_tail1 = k_base_idx + fx.Index(_k_per_batch - tile_k)
                 x_regs_ping = load_x_tile(k_tail1)
                 b_gate_ping = load_b_tile(k_tail1, n_blk_gate, n_intra_gate)
                 b_up_ping = load_b_tile(k_tail1, n_blk_up, n_intra_up)
@@ -1256,7 +1256,7 @@ def _unflatten_b_tile(vals):
                 # Epilogue hoists to keep IR + Python build time small:
                 col_i32_list = []
                 for ni in range_constexpr(num_acc_n):
-                    col_i32_list.append(arith.index_cast(T.i32, col_g_list[ni]))
+                    col_i32_list.append(fx.Int32(col_g_list[ni]))
 
                 lane_div_16 * fx.Index(4)
                 inter_i32_local = inter_i32_v
@@ -1304,7 +1304,7 @@ def write_row_to_lds_splitk(
                         # Load per-row scale_x (sx) — same logic as normal epilogue.
                         fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32)
                         t2 = fused2 & mask24_i32
-                        t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v)
+                        t_valid = as_ir_value(t2 < tokens_i32_v)
                         if const_expr(x_is_token_slot):
                             s2 = fused2 >> 24
                             ts2 = s2 * tokens_i32_v + t2
@@ -1349,25 +1349,25 @@ def precompute_row_splitk(*, row_local, row):
                         fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32)
                         t2 = fused2 & mask24_i32
                         s2 = fused2 >> 24
-                        t_ok = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v)
+                        t_ok = as_ir_value(t2 < tokens_i32_v)
                         t_idx = arith.index_cast(T.index, t2)
                         s_idx = arith.index_cast(T.index, s2)
-                        ts_idx = t_idx * arith.index(topk) + s_idx
+                        ts_idx = t_idx * fx.Index(topk) + s_idx
                         if const_expr(_splitk_use_bf16 and not _needs_global_atomic_bf16_s1):
                             # For buffer atomics: compute relative byte offset from buffer base
-                            row_byte_off = ts_idx * arith.index(_split_k_out_row_stride)
+                            row_byte_off = ts_idx * fx.Index(_split_k_out_row_stride)
                             return (row_byte_off, t_ok)
                         else:
                             # For global atomics: compute absolute address
-                            row_byte_base = out_base_idx + ts_idx * arith.index(_split_k_out_row_stride)
+                            row_byte_base = out_base_idx + ts_idx * fx.Index(_split_k_out_row_stride)
                             return (row_byte_base, t_ok)
 
                     _splitk_zero_i32 = [fx.Int32(0) if _splitk_use_bf16 else None]
 
                     def store_pair_splitk(*, row_local, row, row_ctx, col_pair0, col_g0, frag):
                         row_byte_ctx = row_ctx
-                        col_idx = col_g0 + arith.index(_split_k_n_offset[0])
-                        byte_off_col = col_idx * arith.index(out_elem_bytes)
+                        col_idx = col_g0 + fx.Index(_split_k_n_offset[0])
+                        byte_off_col = col_idx * fx.Index(out_elem_bytes)
                         if const_expr(_splitk_use_bf16):
                             _z = _splitk_zero_i32[0]
                             if const_expr(_needs_global_atomic_bf16_s1):
@@ -1386,7 +1386,7 @@ def store_pair_splitk(*, row_local, row, row_ctx, col_pair0, col_g0, frag):
                                 )
                             else:
                                 # gfx950+: buffer_atomic_pk_add_bf16
-                                byte_off_i32 = arith.index_cast(T.i32, row_byte_ctx + byte_off_col)
+                                byte_off_i32 = fx.Int32(row_byte_ctx + byte_off_col)
                                 buffer_atomic_add(frag, out_rsrc, byte_off_i32, _z, _z)
                         else:
                             # f32 atomic: global atomicrmw fadd
@@ -1490,7 +1490,7 @@ def write_row_to_lds(
                         s2 = fused2 >> 24
                         # aiter moe_sorting uses sentinel token_id == tokens for padding.
                         # Do NOT rely on buffer OOB semantics for scale loads; explicitly mask.
-                        t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v)
+                        t_valid = as_ir_value(t2 < tokens_i32_v)
                         if const_expr(x_is_token_slot):
                             # slot-major: slot*tokens + token
                             ts2 = s2 * tokens_i32_v + t2
@@ -1566,11 +1566,11 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag):
                         # OOB buffer stores are not guaranteed to be safe on all paths, so predicate explicitly.
                         fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32)
                         t2 = fused2 & mask24_i32
-                        t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v)
+                        t_valid = as_ir_value(t2 < tokens_i32_v)
                         _if_valid = scf.IfOp(t_valid)
                         with _if_then(_if_valid):
                             idx0 = row_ctx
-                            col_i32 = arith.index_cast(T.i32, col_g0)
+                            col_i32 = fx.Int32(col_g0)
                             idx_out = idx0 + col_i32
                             # Vectorized fp16 store (EVec=4).
                             buffer_ops.buffer_store(frag, out_rsrc, idx_out)
@@ -1609,7 +1609,7 @@ def _stage1_store_row(*, mi: int, ii: int, row_in_tile, row):
                     s2_raw = fused2 >> 24
                     t2 = t2_raw
                     s2 = s2_raw
-                    t_valid = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32_v)
+                    t_valid = as_ir_value(t2 < tokens_i32_v)
 
                     # Do NOT rely on buffer OOB semantics for scale loads; explicitly mask.
                     if const_expr(x_is_token_slot):
diff --git a/kernels/moe/moe_gemm_2stage/gemm2.py b/kernels/moe/moe_gemm_2stage/gemm2.py
index 9203b0a67..61a8ed305 100644
--- a/kernels/moe/moe_gemm_2stage/gemm2.py
+++ b/kernels/moe/moe_gemm_2stage/gemm2.py
@@ -330,11 +330,11 @@ def moe_gemm2(
             # A2 layout (flatten token-slot -> M; use i32 for fly.make_shape).
             topk_idx = fx.Index(topk)
             m_in = tokens_in * topk_idx
-            m_i32_v = arith.index_cast(T.i32, m_in)
+            m_i32_v = fx.Int32(m_in)
             fx.make_layout((m_i32_v, k_i32_v), stride=(k_i32_v, 1))
 
             # B preshuffle layout: [experts*model_dim, inter_dim]
-            c_n_total = arith.index(experts * model_dim)
+            c_n_total = fx.Index(experts * model_dim)
             # For packed int4 (W4A8/W4A16/W4A_FP8), kpack_bytes=8.
             kpack_bytes = 8 if w_is_int4 else 16
             w_elem_bytes = 1 if w_is_int4 else elem_bytes
@@ -346,7 +346,7 @@ def moe_gemm2(
                 elem_bytes=w_elem_bytes,
             )
             layout_b = b_layout.layout_b
-            (k_in * arith.index(int(elem_bytes))) // fx.Index(64)
+            (k_in * fx.Index(int(elem_bytes))) // fx.Index(64)
 
             shape_lds = fx.make_shape(tile_m, tile_k)
             stride_lds = fx.make_stride(lds_stride, 1)
@@ -360,7 +360,7 @@ def moe_gemm2(
             bx = gpu.block_id("y")  # tile along sorted M
 
             # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs).
-            k_blocks16 = arith.index(tile_k_bytes // 16)
+            k_blocks16 = fx.Index(tile_k_bytes // 16)
             layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1))
             layout_lane16 = fx.make_layout((4, 16), stride=(16, 1))
             fx.make_layout((tile_m, tile_k), stride=(tile_k, 1))
@@ -391,7 +391,7 @@ def moe_gemm2(
             c_topk = fx.Index(topk)
 
             # X(A2): [tokens*topk, inter_dim] bytes = tokens*topk*k*elem_bytes
-            x_nbytes_idx = (tokens_in * c_topk) * k_in * arith.index(int(elem_bytes))
+            x_nbytes_idx = (tokens_in * c_topk) * k_in * fx.Index(int(elem_bytes))
             x_rsrc = buffer_ops.create_buffer_resource(arg_x, max_size=False, num_records_bytes=x_nbytes_idx)
 
             w_rsrc = buffer_ops.create_buffer_resource(arg_w, max_size=False)
@@ -444,8 +444,8 @@ def moe_gemm2(
                 num_records_bytes=fx.Index(4),
             )
             num_valid_i32 = buffer_ops.buffer_load(numids_rsrc, fx.Index(0), vec_width=1, dtype=T.i32)
-            bx_m_i32 = arith.index_cast(T.i32, bx_m)
-            blk_valid = arith.cmpi(arith.CmpIPredicate.ult, bx_m_i32, num_valid_i32)
+            bx_m_i32 = fx.Int32(bx_m)
+            blk_valid = bx_m_i32 < num_valid_i32
 
             def _moe_gemm2_then_body():
                 # Expert id for this M tile.
@@ -476,8 +476,8 @@ def _moe_gemm2_then_body():
                 num_x_loads = bytes_per_thread_x // x_load_bytes
                 chunk_i32 = x_load_bytes // 4  # dwords per chunk (1/2/4)
 
-                c_k_div4 = (k_in * arith.index(int(elem_bytes))) // fx.Index(4)
-                c_k_div4_i32 = arith.index_cast(T.i32, c_k_div4)
+                c_k_div4 = (k_in * fx.Index(int(elem_bytes))) // fx.Index(4)
+                c_k_div4_i32 = fx.Int32(c_k_div4)
                 fx.make_layout((m_i32_v, c_k_div4_i32), stride=(c_k_div4_i32, 1))
                 tile_k_dwords = (int(tile_k) * int(elem_bytes)) // 4
                 layout_x_tile_div4 = fx.make_layout((tile_m, tile_k_dwords), stride=(tile_k_dwords, 1))
@@ -487,7 +487,7 @@ def _moe_gemm2_then_body():
                 topk_i32 = fx.Int32(topk)
                 mask24 = fx.Int32(0xFFFFFF)
                 # Sentinel clamp uses `tokens` as the upper bound: t_valid = (t < tokens).
-                tokens_i32 = arith.index_cast(T.i32, tokens_in)
+                tokens_i32 = fx.Int32(tokens_in)
 
                 def x_tile_chunk_coord_i32(i: int):
                     return tile_chunk_coord_i32(
@@ -532,8 +532,8 @@ def load_x(idx_i32):
                     s_i32 = fused_i >> 24
                     # aiter moe_sorting uses sentinel token_id == tokens for padding.
                     # Do NOT rely on buffer OOB semantics for A2/scale loads; explicitly mask.
-                    t_valid = arith.cmpi(arith.CmpIPredicate.ult, t_i32, tokens_i32)
-                    s_valid = arith.cmpi(arith.CmpIPredicate.ult, s_i32, topk_i32)
+                    t_valid = t_i32 < tokens_i32
+                    s_valid = s_i32 < topk_i32
                     ts_valid = t_valid & s_valid
                     t_safe = ts_valid.select(t_i32, fx.Int32(0))
                     s_safe = ts_valid.select(s_i32, fx.Int32(0))
@@ -543,7 +543,7 @@ def load_x(idx_i32):
                     x_row_base_div4.append(row_ts_idx * c_k_div4)
 
                 def load_x_tile(base_k):
-                    base_k_div4 = (base_k * arith.index(int(elem_bytes))) // fx.Index(4)
+                    base_k_div4 = (base_k * fx.Index(int(elem_bytes))) // fx.Index(4)
                     parts = []
                     for i in range_constexpr(num_x_loads):
                         idx_i32 = x_row_base_div4[i] + base_k_div4 + x_col_local_i32[i]
@@ -567,9 +567,9 @@ def load_x_tile(base_k):
                 row_a_lds = lane_mod_16
                 # A-side kpack is always 16 bytes; kpack_bytes is B-side (may be 8 for int4).
                 a_kpack_elems = 16 // elem_bytes
-                col_offset_base = lane_div_16 * arith.index(int(a_kpack_elems))
+                col_offset_base = lane_div_16 * fx.Index(int(a_kpack_elems))
                 col_offset_base_bytes = (
-                    col_offset_base if elem_bytes == 1 else (col_offset_base * arith.index(int(elem_bytes)))
+                    col_offset_base if elem_bytes == 1 else (col_offset_base * fx.Index(int(elem_bytes)))
                 )
 
                 # Dynamic N tiling within block.
@@ -589,7 +589,7 @@ def load_x_tile(base_k):
                 c_n0_static = experts * model_dim // 16
                 layout_n_blk_intra = fx.make_layout((c_n0_static, 16), stride=(16, 1))
                 for ni in range_constexpr(num_acc_n):
-                    offset = arith.index(ni * 16)
+                    offset = fx.Index(ni * 16)
                     col_g = by_n + n_tile_base + offset + lane_mod_16
                     col_g_list.append(col_g)
 
@@ -752,7 +752,7 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_base):
                 def lds_load_packs_k64(curr_row_a_lds, col_base_bytes, lds_base):
                     col_base_swz_bytes = swizzle_xor16(curr_row_a_lds, col_base_bytes, k_blocks16)
                     col_base_swz = (
-                        col_base_swz_bytes if elem_bytes == 1 else (col_base_swz_bytes // arith.index(int(elem_bytes)))
+                        col_base_swz_bytes if elem_bytes == 1 else (col_base_swz_bytes // fx.Index(int(elem_bytes)))
                     )
                     idx_a16 = preshuffle_crd2idx((fx.Int32(curr_row_a_lds), fx.Int32(col_base_swz)), layout_lds)
                     idx_a16 = idx_a16 + lds_base
@@ -804,7 +804,7 @@ def compute_tile(
                             lane_div_16_mul4_pf = lane_div_16 * fx.Index(4)
                             ii_idx_list_pf = [fx.Index(ii) for ii in range(4)]
                             for mi in range_constexpr(m_repeat):
-                                mi_base_pf = arith.index(mi * 16)
+                                mi_base_pf = fx.Index(mi * 16)
                                 for ii in range_constexpr(4):
                                     row_off_pf = lane_div_16_mul4_pf + ii_idx_list_pf[ii]
                                     row_in_tile_pf = mi_base_pf + row_off_pf
@@ -860,11 +860,11 @@ def _acc_scaled_f32(f32_acc_vec, f32_partial_vec, scale_val):
                         _pending_acc = None
                         for ku in range_constexpr(k_unroll):
                             b_raw = b_tile_in[ku]
-                            ki64 = arith.index(ku * 64)
+                            ki64 = fx.Index(ku * 64)
                             col_base = col_offset_base_bytes + ki64
 
                             for mi in range_constexpr(m_repeat):
-                                mi_val = arith.index(mi * 16)
+                                mi_val = fx.Index(mi * 16)
                                 curr_row_a_lds = row_a_lds + mi_val
 
                                 if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)):
@@ -911,11 +911,11 @@ def _acc_scaled_f32(f32_acc_vec, f32_partial_vec, scale_val):
                     else:
                         for ku in range_constexpr(k_unroll):
                             b_packs0, b_packs1 = b_tile_in[ku]
-                            ki64 = arith.index(ku * 64)
+                            ki64 = fx.Index(ku * 64)
                             col_base = col_offset_base_bytes + ki64
 
                             for mi in range_constexpr(m_repeat):
-                                mi_val = arith.index(mi * 16)
+                                mi_val = fx.Index(mi * 16)
                                 curr_row_a_lds = row_a_lds + mi_val
 
                                 if const_expr((a0_prefetch is not None) and (ku == 0) and (mi == 0)):
@@ -935,7 +935,7 @@ def _acc_scaled_f32(f32_acc_vec, f32_partial_vec, scale_val):
                     return acc_list, epilogue_pf
 
                 # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ----------------
-                lds_tile_elems = arith.index(tile_m * lds_stride)
+                lds_tile_elems = fx.Index(tile_m * lds_stride)
                 lds_base_cur = fx.Index(0)
                 lds_base_nxt = lds_tile_elems
 
@@ -1049,8 +1049,8 @@ def hot_loop_scheduler():
                 if const_expr(k_main2_py < 0):
                     k_main2_py = 0
 
-                arith.index(tile_k * 2)
-                c_tile_k_s2 = arith.index(tile_k)
+                fx.Index(tile_k * 2)
+                c_tile_k_s2 = fx.Index(tile_k)
                 pair_iters = k_main2_py // (int(tile_k) * 2)
 
                 # B-tile data layout per k_unroll entry (3 variants):
@@ -1222,8 +1222,8 @@ def _stage2_row_atomic(*, mi: int, ii: int, row_in_tile, row):
 
                         # Mask sentinel (token_id==tokens, slot==topk) to avoid OOB scale_x loads.
                         # For invalid rows, force sx=0 so they contribute exactly 0 to output.
-                        t_ok = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32)
-                        s_ok = arith.cmpi(arith.CmpIPredicate.ult, s2, topk_i32_v)
+                        t_ok = t2 < tokens_i32
+                        s_ok = s2 < topk_i32_v
                         ts_ok = t_ok & s_ok
                         t2_safe = ts_ok.select(t2, fx.Int32(0))
                         s2_safe = ts_ok.select(s2, fx.Int32(0))
@@ -1261,7 +1261,7 @@ def _stage2_row_atomic(*, mi: int, ii: int, row_in_tile, row):
                             v = v * sx * sw
                             if const_expr(doweight_stage2):
                                 v = v * tw
-                            col_i32 = arith.index_cast(T.i32, col_g)
+                            col_i32 = fx.Int32(col_g)
                             idx_elem = idx0 + col_i32
                             byte_off = idx_elem * c4_i32
                             atomic_add_f32(v, byte_off)
@@ -1299,8 +1299,8 @@ def write_row_to_lds(
                         t2 = fused2 & mask24_i32
                         s2 = fused2 >> 24
                         # Explicitly mask sentinel token/slot to avoid OOB scale_x loads.
-                        t_ok = arith.cmpi(arith.CmpIPredicate.ult, t2, tokens_i32)
-                        s_ok = arith.cmpi(arith.CmpIPredicate.ult, s2, topk_i32_v)
+                        t_ok = t2 < tokens_i32
+                        s_ok = s2 < topk_i32_v
                         ts_ok = t_ok & s_ok
                         t2_safe = ts_ok.select(t2, fx.Int32(0))
                         s2_safe = ts_ok.select(s2, fx.Int32(0))
@@ -1349,14 +1349,14 @@ def precompute_row(*, row_local, row):
                         # Return (fused_i32, row_valid_i1) so the epilogue can skip the entire row
                         # for invalid tail rows (CK-style), avoiding per-store branching.
                         fused2 = buffer_ops.buffer_load(sorted_rsrc, row, vec_width=1, dtype=T.i32)
-                        row_i32 = arith.index_cast(T.i32, row)
-                        row_valid0 = arith.cmpi(arith.CmpIPredicate.ult, row_i32, num_valid_i32)
+                        row_i32 = fx.Int32(row)
+                        row_valid0 = row_i32 < num_valid_i32
                         t = fused2 & mask24_i32
                         s = fused2 >> 24
-                        t_ok = arith.cmpi(arith.CmpIPredicate.ult, t, tokens_i32)
-                        s_ok = arith.cmpi(arith.CmpIPredicate.ult, s, topk_i32_v)
+                        t_ok = t < tokens_i32
+                        s_ok = s < topk_i32_v
                         row_valid = row_valid0 & t_ok & s_ok
-                        return (fused2, row_valid)
+                        return (fused2, as_ir_value(row_valid))
 
                     def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag):
                         fused = row_ctx
@@ -1366,7 +1366,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag):
                         if const_expr(not bool(accumulate)):
                             ts = t * topk_i32_v + s
                             idx0 = ts * model_i32
-                        col_i32 = arith.index_cast(T.i32, col_g0)
+                        col_i32 = fx.Int32(col_g0)
                         idx_elem = idx0 + col_i32
                         idx_elem_even = idx_elem & mask_even_i32
                         if const_expr(_needs_global_atomic_bf16):
@@ -1420,7 +1420,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag):
                         store_pair=store_pair,
                     )
 
-            _if_blk = scf.IfOp(blk_valid)
+            _if_blk = scf.IfOp(as_ir_value(blk_valid))
             with _if_then(_if_blk):
                 _moe_gemm2_then_body()
 
diff --git a/kernels/moe/mxfp_moe/gemm1.py b/kernels/moe/mxfp_moe/gemm1.py
index 559c09bdf..dc5e22b48 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 6168cce8a..f51f85a2e 100644
--- a/kernels/moe/mxfp_moe/gemm2.py
+++ b/kernels/moe/mxfp_moe/gemm2.py
@@ -12,7 +12,6 @@
     _a_lds_swz_block_idx,
     _a_lds_swz_block_layout,
     _e8m0_from_amax,
-    _fabs_f32,
     _gep1,
     _gep3,
     _global_base_ptr1,
@@ -769,14 +768,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 071e36e38..848c85c90 100644
--- a/kernels/moe/mxfp_moe/mxfp4_gemm_common.py
+++ b/kernels/moe/mxfp_moe/mxfp4_gemm_common.py
@@ -158,14 +158,10 @@ def _a_lds_swz_block_idx(swz_layout, row, block_col):
     return fx.Int32(crd2idx([fx.Int64(row), fx.Int64(block_col)], swz_layout))
 
 
-def _fabs_f32(x):
-    return fx.Float32(llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(x)], [], []))
-
-
 def _e8m0_roundup(amax_f32):
     wi = fx.Int32(_raw(amax_f32 * fx.Float32(1.0 / 6.0)).bitcast(T.i32))
     bexp = (wi + fx.Int32(0x7FFFFF)).shrui(fx.Int32(23)) & fx.Int32(0xFF)
-    lt = arith.cmpi(arith.CmpIPredicate.ult, _raw(bexp), _raw(fx.Int32(254)))
+    lt = fx.as_ir_value((bexp) < (fx.Int32(254)))
     return fx.Int32(arith.select(lt, _raw(bexp), _raw(fx.Int32(254))))
 
 
@@ -196,7 +192,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)))
 
 
diff --git a/kernels/moe/silu_and_mul_fq.py b/kernels/moe/silu_and_mul_fq.py
index b3cd08ff2..2c2677c8d 100644
--- a/kernels/moe/silu_and_mul_fq.py
+++ b/kernels/moe/silu_and_mul_fq.py
@@ -26,7 +26,7 @@
 from flydsl._mlir.dialects import llvm, scf, vector
 from flydsl.compiler.kernel_function import CompilationContext
 from flydsl.expr import arith, as_ir_value, const_expr, range_constexpr
-from flydsl.expr.arith import ArithValue, CmpIPredicate
+from flydsl.expr.arith import ArithValue
 from flydsl.expr.typing import Int32, T
 from kernels.common import buffer_ops
 
@@ -157,14 +157,14 @@ def _load_bias_scalar(offset):
         token_num_i32 = ArithValue(token_num)
         bid_i32 = ArithValue(bid)
 
-        row_in_range = arith.cmpi(CmpIPredicate.ult, bid_i32, num_valid)
+        row_in_range = bid_i32 < num_valid
         fused_tid_val = buffer_ops.buffer_load(tid_rsrc, bid_i32, vec_width=1, dtype=i32)
         mask24 = arith.constant(0xFFFFFF, type=i32)
         token_id = fused_tid_val & mask24
         slot_id = ArithValue(fused_tid_val) >> arith.constant(24, type=i32)
-        t_ok = arith.cmpi(CmpIPredicate.ult, token_id, token_num_i32)
-        s_ok = arith.cmpi(CmpIPredicate.ult, slot_id, topk_i32)
-        is_valid = arith.andi(row_in_range, arith.andi(t_ok, s_ok))
+        t_ok = token_id < token_num_i32
+        s_ok = slot_id < topk_i32
+        is_valid = as_ir_value(row_in_range & t_ok & s_ok)
 
         if const_expr(_need_fp4):
 
@@ -174,11 +174,8 @@ def _f32_to_e2m1(qx_f32):
                 qx = qx_f32.bitcast(i32)
                 s = qx & c0x80000000_i32
                 qx_abs = qx & c0x7FFFFFFF_i32
-                denormal_mask = arith.cmpi(CmpIPredicate.ult, qx_abs, c0x3F800000_i32)
-                normal_mask = arith.andi(
-                    arith.cmpi(CmpIPredicate.ult, qx_abs, c0x40C00000_i32),
-                    arith.cmpi(CmpIPredicate.uge, qx_abs, c0x3F800000_i32),
-                )
+                denormal_mask = qx_abs < c0x3F800000_i32
+                normal_mask = (qx_abs < c0x40C00000_i32) & (qx_abs >= c0x3F800000_i32)
 
                 denorm_f32 = qx_abs.bitcast(f32) + c0x4A800000_i32.bitcast(f32)
                 denormal_x = denorm_f32.bitcast(i32) - c0x4A800000_i32
@@ -197,7 +194,7 @@ def _f32_to_e2m1(qx_f32):
         for iter_idx in range_constexpr((inter_dim + COLS_PER_ITER - 1) // COLS_PER_ITER):
             col0 = thread_id * arith.constant(VEC, type=i32) + arith.constant(iter_idx * COLS_PER_ITER, type=i32)
 
-            col_valid = arith.cmpi(CmpIPredicate.ult, col0, inter_dim_i32)
+            col_valid = as_ir_value(col0 < inter_dim_i32)
             _if_col = scf.IfOp(col_valid)
             with ir.InsertionPoint(_if_col.then_block):
 
@@ -298,8 +295,8 @@ def _f32_to_e2m1(qx_f32):
                     if const_expr(_need_quant):
                         local_max = c0_f32
                         for vi in range_constexpr(VEC):
-                            abs_v = llvm.call_intrinsic(f32, "llvm.fabs.f32", [act_vals[vi]], [], [])
-                            local_max = arith.maximumf(local_max, abs_v)
+                            abs_v = fx.absf(act_vals[vi])
+                            local_max = fx.maxnumf(local_max, abs_v)
 
                         for sh_dist in SHUFFLE_DISTS:
                             off = arith.constant(sh_dist, type=i32)
@@ -413,7 +410,7 @@ def _f32_to_e2m1(qx_f32):
                                     )
 
                         lane_in_blk = col0 & c31_i32
-                        _if_sw = scf.IfOp(arith.cmpi(CmpIPredicate.eq, lane_in_blk, c0_i32))
+                        _if_sw = scf.IfOp(as_ir_value(lane_in_blk == c0_i32))
                         with ir.InsertionPoint(_if_sw.then_block):
                             row_s = bid_i32
                             col_s = col0 >> c5_i32
@@ -456,7 +453,7 @@ def _f32_to_e2m1(qx_f32):
                 with ir.InsertionPoint(_if_valid.else_block):
                     if const_expr(_need_quant):
                         lane_in_blk_p = col0 & c31_i32
-                        _if_sw_p = scf.IfOp(arith.cmpi(CmpIPredicate.eq, lane_in_blk_p, c0_i32))
+                        _if_sw_p = scf.IfOp(as_ir_value(lane_in_blk_p == c0_i32))
                         with ir.InsertionPoint(_if_sw_p.then_block):
                             row_s_p = bid_i32
                             col_s_p = col0 >> c5_i32