[FlyDSL] Raw-dialect cleanup (others): gdr_decode / chunk_gated_delta_h / pa_mqa_logits_fp4

aiter/ops/flydsl/__init__.py

@@ -39,19 +39,19 @@
 
     from .fmha_kernels import flydsl_flash_attn_func
     from .gemm_kernels import flydsl_hgemm, flydsl_preshuffle_gemm_a8
-    from .kernels.fp8_mqa_logits import (
+    from .kernels.mqa_logits.fp8_mqa_logits import (
         DEFAULT_VARIANT as FP8_MQA_LOGITS_DEFAULT_VARIANT,
     )
-    from .kernels.fp8_mqa_logits import (
+    from .kernels.mqa_logits.fp8_mqa_logits import (
         KERNEL_VARIANTS as FP8_MQA_LOGITS_VARIANTS,
     )
-    from .kernels.fp8_mqa_logits import (
+    from .kernels.mqa_logits.fp8_mqa_logits import (
         flydsl_fp8_mqa_logits,
     )
-    from .kernels.pa_mqa_logits_fp4 import (
+    from .kernels.mqa_logits.pa_mqa_logits_fp4 import (
         flydsl_pa_mqa_logits_fp4,
     )
-    from .kernels.pa_mqa_logits_fp4_prefill import (
+    from .kernels.mqa_logits.pa_mqa_logits_fp4_prefill import (
         compute_varqlen_windows,
         flydsl_pa_mqa_logits_fp4_prefill,
         flydsl_pa_mqa_logits_fp4_varqlen,

aiter/ops/flydsl/kernels/chunk_gated_delta_h.py

@@ -16,21 +16,15 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl._mlir import ir
-from flydsl._mlir.dialects import llvm as _llvm
 from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
 from flydsl.expr.typing import T
 
-from .tensor_shim import GTensor, _to_raw
+from .tensor_shim import GTensor
 
 _LOG2E = math.log2(math.e)  # 1.4426950408889634
 _LLVM_GEP_DYNAMIC = -2147483648
 
 
-def _llvm_lds_ptr_ty():
-    return ir.Type.parse("!llvm.ptr<3>")
-
-
 def _make_fast_exp(g_is_log2_scaled: bool):
     """Return the ``exp`` helper for this kernel compile.
 
@@ -292,12 +286,13 @@ def _lds_vec_read_k_bf16x8(elem_idx):
         # -- ds_read_b64_tr_b16 helper (gfx950) --
         v4bf16_type = T.vec(4, T.bf16)
 
+        lds_ptr_ty = lds_h_ptr.type
+
         def _ds_read_tr_bf16x4(lds_byte_offset):
-            byte_i64 = fx.Int64(lds_byte_offset).ir_value()
-            ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result
+            # Raw int LDS byte offset -> bf16 shared ptr; fx resolves the
+            # backend address space (-> !llvm.ptr<3>) for the ds_read intrinsic.
+            ptr = fx.to_llvm_ptr(fx.inttoptr(lds_ptr_ty, fx.Int64(lds_byte_offset)))
             raw = rocdl.ds_read_tr16_b64(v4bf16_type, ptr).result
-            # Wrap as Vector so call-sites can use .shuffle()/.bitcast()
-            # method-style API instead of the bare vector.shuffle wrapper.
             return fx.Vector(raw, (4,), fx.BFloat16)
 
         # ds_read_b64_tr_b16 lane decomposition
@@ -404,9 +399,11 @@ def _ds_read_tr_bf16x4(lds_byte_offset):
                 g_off = w_base + safe_row * stride_w + fx.Int32(kb * 64) + load_col_base
                 w_prefetch_init.append(w_.vec_load((fx.Int64(g_off),), LOAD_VEC_WIDTH))
 
-        init_state = [_to_raw(v) for v in h_accs] + [
-            _to_raw(v) for v in w_prefetch_init
-        ]
+        # h_accs are fx.Vector, w_prefetch_init are raw buffer_load results;
+        # ``range(init=)`` runs each element through ``as_ir_value`` and the
+        # ``yield`` below through the scf-yield rewriter, so both fx values and
+        # raw ir.Values are accepted directly -- no manual _to_raw needed.
+        init_state = list(h_accs) + list(w_prefetch_init)
         c_zero = fx.Int64(0)
         c_one = fx.Int64(1)
         nt_idx = fx.Int64(NT)
@@ -926,9 +923,10 @@ def _emit_vn_store(off, value):
                         (fx.Int64(w_next_prefetch_off[i_wp]),), LOAD_VEC_WIDTH
                     )
 
-            results = yield [_to_raw(v) for v in h_accs_in] + [
-                _to_raw(v) for v in w_next_prefetch
-            ]
+            # h_accs_in are fx.Vector, w_next_prefetch are raw buffer_load
+            # results; the scf-yield rewriter converts fx values via ir_value()
+            # and passes raw ir.Values through unchanged.
+            results = yield list(h_accs_in) + list(w_next_prefetch)
 
         h_accs_final = list(results[:NUM_H_ACCS])
 

aiter/ops/flydsl/kernels/gdr_decode.py

@@ -8,20 +8,16 @@
 from flydsl._mlir.dialects import (
     gpu as mlir_gpu,
 )
-from flydsl.expr import arith, const_expr, range_constexpr, rocdl
+from flydsl.expr import const_expr, range_constexpr, rocdl
 from flydsl.expr.typing import T
 
-from aiter.ops.flydsl.kernels import vector
-
 from .tensor_shim import (
     GTensor,
     _to_raw,
     get_dtype_bytes,
     get_dtype_in_kernel,
 )
 
-fm_fast = arith.FastMathFlags.fast
-
 
 @functools.lru_cache(maxsize=1024)
 def create_vk_gdr_decode_kernel(
@@ -98,18 +94,17 @@ def gdr_decode_kernel(
         out: fx.Tensor,
         batch_size: fx.Int32,
     ):
-        scale = arith.constant(SCALE_VALUE, type=T.f32)
-        softplus_beta_ = arith.constant(softplus_beta, type=T.f32)
-        softplus_threshold_ = arith.constant(softplus_threshold, type=T.f32)
+        scale = fx.Float32(SCALE_VALUE)
+        softplus_beta_ = fx.Float32(softplus_beta)
+        softplus_threshold_ = fx.Float32(softplus_threshold)
 
         dtype_ = get_dtype_in_kernel(dtype)
         fx_dtype_ = fx.BFloat16 if dtype == "bf16" else fx.Float16
         A_log_dtype_ = get_dtype_in_kernel(A_log_dtype)
         state_dtype_ = get_dtype_in_kernel(state_dtype)
-        # i32_0 = arith.constant(0, type=T.i32)
-        f32_0 = arith.constant(0.0, type=T.f32)
-        f32_1 = arith.constant(1.0, type=T.f32)
-        width_i32 = arith.constant(WARP_SIZE, type=T.i32)
+        f32_0 = fx.Float32(0.0)
+        f32_1 = fx.Float32(1.0)
+        width_i32 = _to_raw(fx.Int32(WARP_SIZE))
         vec_t = T.vec(VALUES_PER_THREAD_K, dtype_)
         acc_vec_t = T.vec(VALUES_PER_THREAD_K, T.f32)
 
@@ -170,12 +165,11 @@ def gdr_decode_kernel(
         def fast_exp(x, use_exp2=True):
             if const_expr(use_exp2):
                 log2e = 1.4426950408889634
-                out = rocdl.exp2(T.f32, x * log2e)
-                return out
-            return fx.math.exp(x, fastmath=fm_fast)
+                return rocdl.exp2(T.f32, _to_raw(fx.Float32(x) * log2e))
+            return fx.math.exp(x, fastmath=fx.FastMathFlags.fast)
 
         def fast_log1p(x):
-            return fx.math.log1p(x, fastmath=fm_fast)
+            return fx.math.log1p(x, fastmath=fx.FastMathFlags.fast)
 
         # Skip CG-pad slots (indices sentinel < 0). The guarded body is a
         # closure so the runtime `if` sees an opaque call (no GTensor "state"
@@ -220,12 +214,16 @@ def _do_decode():
                 r_beta = f32_1 / (f32_1 + fast_exp(-r_b))
                 r_g = fast_exp(r_g_value)
 
-                r_g_vec = vector.BroadcastOp(acc_vec_t, r_g).vector
+                r_g_vec = fx.Vector.filled(
+                    VALUES_PER_THREAD_K, fx.Float32(r_g), fx.Float32
+                )
 
                 sq_vecs = [0] * WARP_TILE_K_ITERS
                 sk_vecs = [0] * WARP_TILE_K_ITERS
 
-                scale_vec = vector.BroadcastOp(acc_vec_t, scale).vector
+                scale_vec = fx.Vector.filled(
+                    VALUES_PER_THREAD_K, fx.Float32(scale), fx.Float32
+                )
 
                 for ki in range_constexpr(WARP_TILE_K_ITERS):
                     warp_k_vec_i = warp_k_vec_start + ki * WARP_TILE_K
@@ -239,11 +237,13 @@ def _do_decode():
                     sk_vecs[ki] = k_vec.extf(acc_vec_t)
 
                 if const_expr(use_qk_l2norm):
-                    sum_q_partial_vec = vector.from_elements(
-                        acc_vec_t, [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)]
+                    sum_q_partial_vec = fx.Vector.from_elements(
+                        [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)],
+                        fx.Float32,
                     )
-                    sum_k_partial_vec = vector.from_elements(
-                        acc_vec_t, [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)]
+                    sum_k_partial_vec = fx.Vector.from_elements(
+                        [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)],
+                        fx.Float32,
                     )
                     for ki in range_constexpr(WARP_TILE_K_ITERS):
                         sum_q_partial_vec = (
@@ -279,49 +279,49 @@ def _do_decode():
                     ).shuffleResult
                     inv_norm_q = fx.math.rsqrt(local_sum_q + 1e-6)
                     inv_norm_k = fx.math.rsqrt(local_sum_k + 1e-6)
-                    inv_norm_q_vec = vector.BroadcastOp(
-                        acc_vec_t, _to_raw(inv_norm_q)
-                    ).vector
-                    inv_norm_k_vec = vector.BroadcastOp(
-                        acc_vec_t, _to_raw(inv_norm_k)
-                    ).vector
+                    inv_norm_q_vec = fx.Vector.filled(
+                        VALUES_PER_THREAD_K, fx.Float32(inv_norm_q), fx.Float32
+                    )
+                    inv_norm_k_vec = fx.Vector.filled(
+                        VALUES_PER_THREAD_K, fx.Float32(inv_norm_k), fx.Float32
+                    )
                     for ki in range_constexpr(WARP_TILE_K_ITERS):
                         sq_vecs[ki] = sq_vecs[ki] * inv_norm_q_vec * scale_vec
                         sk_vecs[ki] = sk_vecs[ki] * inv_norm_k_vec
                 else:
                     for ki in range_constexpr(WARP_TILE_K_ITERS):
                         sq_vecs[ki] = sq_vecs[ki] * scale_vec
 
-                dot_kq_vec = vector.from_elements(
-                    acc_vec_t, [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)]
+                dot_kq_vec = fx.Vector.from_elements(
+                    [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)], fx.Float32
                 )
                 for ki in range_constexpr(WARP_TILE_K_ITERS):
-                    dot_kq_vec = vector.FMAOp(
-                        sk_vecs[ki], sq_vecs[ki], dot_kq_vec
-                    ).result
-                dot_kq = fx.Vector(dot_kq_vec).reduce(fx.ReductionOp.ADD)
+                    dot_kq_vec = fx.math.fma(sk_vecs[ki], sq_vecs[ki], dot_kq_vec)
+                dot_kq = dot_kq_vec.reduce(fx.ReductionOp.ADD)
                 for offset in WARP_THREADS_K_SHFL_OFFSETS:
                     dot_kq = dot_kq + dot_kq.shuffle_xor(offset, WARP_SIZE)
 
                 for vi in range_constexpr(WARP_TILE_V_ITERS):
                     global_v_i = global_v_start + vi * WARP_GROUP_TILE_V
                     r_v = v_tensor[b_i, sq_i, hv_i, global_v_i].extf(T.f32)
 
-                    sum_hk = vector.from_elements(
-                        acc_vec_t, [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)]
+                    sum_hk = fx.Vector.from_elements(
+                        [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)],
+                        fx.Float32,
                     )
-                    sum_hq_old = vector.from_elements(
-                        acc_vec_t, [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)]
+                    sum_hq_old = fx.Vector.from_elements(
+                        [f32_0 for i in range_constexpr(VALUES_PER_THREAD_K)],
+                        fx.Float32,
                     )
 
                     for ki in range_constexpr(WARP_TILE_K_ITERS):
                         state_vecs[vi * WARP_TILE_K_ITERS + ki] *= r_g_vec
                         h_cur = state_vecs[vi * WARP_TILE_K_ITERS + ki]
-                        sum_hk = vector.FMAOp(h_cur, sk_vecs[ki], sum_hk).result
-                        sum_hq_old = vector.FMAOp(h_cur, sq_vecs[ki], sum_hq_old).result
+                        sum_hk = fx.math.fma(h_cur, sk_vecs[ki], sum_hk)
+                        sum_hq_old = fx.math.fma(h_cur, sq_vecs[ki], sum_hq_old)
 
-                    sum_hk = fx.Vector(sum_hk).reduce(fx.ReductionOp.ADD)
-                    sum_hq_old = fx.Vector(sum_hq_old).reduce(fx.ReductionOp.ADD)
+                    sum_hk = sum_hk.reduce(fx.ReductionOp.ADD)
+                    sum_hq_old = sum_hq_old.reduce(fx.ReductionOp.ADD)
 
                     for offset in WARP_THREADS_K_SHFL_OFFSETS:
                         sum_hk = sum_hk + sum_hk.shuffle_xor(offset, WARP_SIZE)
@@ -337,14 +337,16 @@ def _do_decode():
                         mode="idx",
                     ).shuffleResult
                     sum_hq = sum_hq_old + v_new * dot_kq
-                    v_new_bcast = vector.BroadcastOp(acc_vec_t, v_new)
+                    v_new_bcast = fx.Vector.filled(
+                        VALUES_PER_THREAD_K, fx.Float32(v_new), fx.Float32
+                    )
 
                     for ki in range_constexpr(WARP_TILE_K_ITERS):
-                        h_new = vector.FMAOp(
+                        h_new = fx.math.fma(
                             sk_vecs[ki],
                             v_new_bcast,
                             state_vecs[vi * WARP_TILE_K_ITERS + ki],
-                        ).result
+                        )
                         state_vecs[vi * WARP_TILE_K_ITERS + ki] = h_new
 
                     sum_hq = sum_hq.to(fx_dtype_)

aiter/ops/flydsl/kernels/mqa_logits/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

aiter/ops/flydsl/kernels/fp8_mqa_logits.py → aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py

@@ -26,7 +26,7 @@
 
 from aiter.jit.utils.chip_info import get_gfx
 
-from .tensor_shim import GTensor, _run_compiled
+from ..tensor_shim import GTensor, _run_compiled
 
 Vec = fx.Vector