DSL-ify raw arith float ops in kernels (flash/mla/pa/moe), fastmath v…

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,
+        )

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,17 +876,15 @@ 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))
                 if rgroup == 0:
                     _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}