Skip to content

[FlyDSL] clean attn-aux kernels: fx.* modernization cleanup #5113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
coderfeli wants to merge 9 commits into
base: main
Choose a base branch
from
Draft

[FlyDSL] clean attn-aux kernels: fx.* modernization cleanup #5113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
c8ec17b
[FLYDSL] attn-aux kernels: fx.* modernization cleanup
coderfeli Aug 29, 2026
7d51960
[FLYDSL] attn-aux: deep fx collapse under fast_fp_math + gfx1201 fx.I…
coderfeli Aug 29, 2026
d6c63b4
[FLYDSL MOE] mfma_epilogues: drop dialect-module params, scf.IfOp→@fl…
coderfeli Aug 29, 2026
c9d4336
[FLYDSL] attn-aux + mfma_epilogues: drive raw IR to the proven floor
coderfeli Aug 29, 2026
bdce863
[FLYDSL] migrate buffer_ops -> make_buffer_tensor (attn + moe_sorting)
coderfeli Aug 30, 2026
36ecbbc
[FLYDSL] moe_sorting: last 4 buffer_ops -> guarded direct ptr store
coderfeli Aug 30, 2026
d79c55c
[FLYDSL] scatter emitter: buffer_ops -> direct global ptr stores
coderfeli Aug 30, 2026
614f7cd
[FLYDSL] dedup buf_tensor/buf_load/buf_store into _common
coderfeli Aug 30, 2026
85974a8
[FLYDSL] drop dead vecVf32 kwarg from the scatter emitter
coderfeli Aug 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 39 additions & 44 deletions aiter/ops/flydsl/kernels/flash_attn_func_gfx1201.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,6 @@
from flydsl.expr import math as fmath
from flydsl.expr.typing import T
from flydsl.expr.typing import Vector as Vec
from flydsl.expr.utils.arith import ArithValue
from flydsl.expr.utils.arith import _to_raw as _raw

from aiter.ops.flydsl.kernels import buffer_ops
Expand Down Expand Up @@ -282,7 +281,7 @@ def load_global_v8f16(base_ptr, base_idx):
return _load_global_half_vec(base_ptr, base_idx, v8f16_type)

def _bitcast_i32(value):
return fx.Int32(ArithValue(value).bitcast(fx.Int32.ir_type))
return fx.Float32(value).bitcast(fx.Int32)

def _pack_bf16_pair(lo, hi, shift, mask):
lo_i32 = _bitcast_i32(lo)
Expand All @@ -298,12 +297,12 @@ def bf16_trunc_pack_v8(f32_vals):
pairs.append(
_pack_bf16_pair(f32_vals[j * 2], f32_vals[j * 2 + 1], _c16, _cmask)
)
return Vec.from_elements(pairs, fx.Int32).bitcast(elem_dtype).ir_value()
return Vec.from_elements(pairs, fx.Int32).bitcast(elem_dtype)

def k_buf_base(buf_id):
if const_expr(isinstance(buf_id, int)):
return fx.Index(buf_id * LDS_K_TILE_SIZE)
return buf_id * fx.Index(LDS_K_TILE_SIZE)
return fx.Int64(buf_id * LDS_K_TILE_SIZE)
return buf_id * fx.Int64(LDS_K_TILE_SIZE)

def v_buf_base(buf_id):
return fx.Index(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE)
Expand All @@ -314,7 +313,7 @@ def coop_load_k(tile_start, buf_id=0):
row_offset = batch * ROWS_PER_BATCH_LOAD
row_idx = tile_start + load_row_in_batch + row_offset
if const_expr(KV_NEEDS_GUARD):
row_valid = load_row_in_batch < fx.Index(BLOCK_N)
row_valid = load_row_in_batch < fx.Int64(BLOCK_N)
if row_valid:
g_idx = global_idx(row_idx, load_col_base)
lds_row = load_row_in_batch + row_offset
Expand Down Expand Up @@ -346,7 +345,7 @@ def coop_store_v_lds(vecs, buf_id=0):
for batch in range_constexpr(NUM_BATCHES_KV):
row_offset = batch * ROWS_PER_BATCH_LOAD
if const_expr(KV_NEEDS_GUARD):
row_valid = load_row_in_batch < fx.Index(BLOCK_N)
row_valid = load_row_in_batch < fx.Int64(BLOCK_N)
if row_valid:
lds_row = load_row_in_batch + row_offset
_v_store_row_major(v_base, lds_row, vecs[batch])
Expand All @@ -362,14 +361,14 @@ def coop_store_v_lds(vecs, buf_id=0):
# to `v_cmp_gt_u64_e64` and cause an ISA hash drift even though both
# variants are semantically equivalent for non-negative offsets.
q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, _raw(q_row), _raw(seq_len_v))
q_row_safe = fx.Index(ArithValue(q_in_bounds).select(q_row, fx.Index(0)))
c_zero_v8f16 = Vec.filled(8, 0.0, elem_dtype).ir_value()
q_row_safe = fx.Int64(q_in_bounds.select(q_row, fx.Index(0)))
c_zero_v8f16 = Vec.filled(8, 0.0, elem_dtype)
q_b_packs = []
for ks in range_constexpr(K_STEPS_QK):
q_col = fx.Index(ks * K_STEP_QK) + klane * WMMA_LANE_K
g_idx = global_idx(q_row_safe, q_col)
raw = load_global_v8f16(q_ptr, g_idx)
q_b_packs.append(ArithValue(q_in_bounds).select(raw, c_zero_v8f16))
q_b_packs.append(q_in_bounds.select(raw, c_zero_v8f16))

# ---- Constants ----
c_neg_inf = fx.Float32(float("-inf"))
Expand All @@ -385,14 +384,12 @@ def reduction_peer(v_f32):

_q_end = q_start + BLOCK_M
if const_expr(CAUSAL):
kv_upper = fx.Index(
ArithValue(_q_end < seq_len_v).select(_q_end, seq_len_v)
)
kv_upper = fx.Int64((_q_end < seq_len_v).select(_q_end, seq_len_v))
else:
kv_upper = seq_len_v

# ---- Opt4: Pre-issue first V global load before loop ----
_v_vecs_init = coop_load_v_global(fx.Index(0))
_v_vecs_init = coop_load_v_global(fx.Int64(0))

init_args = [_raw(c_neg_inf), _raw(c_zero_f)]
for _ in range_constexpr(D_CHUNKS):
Expand Down Expand Up @@ -421,18 +418,18 @@ def reduction_peer(v_f32):
s_accs = [_raw(c_zero_v8f32) for _ in range(NUM_S_ACCS)]

for ks in range_constexpr(K_STEPS_QK):
k_col = fx.Index(ks * K_STEP_QK) + klane * WMMA_LANE_K
k_col = fx.Int64(ks * K_STEP_QK) + klane * WMMA_LANE_K

for st_idx in range_constexpr(N_SUB_TILES):
st_base_row = st_idx * K_SUB_N

k_row_a = lane16 + fx.Index(st_base_row)
k_row_a = lane16 + fx.Int64(st_base_row)
k_lds_a = k_base + k_row_a * K_STRIDE + k_col
k_pack_a = fx.ptr_load(
lds_kv + fx.Int32(k_lds_a), result_type=v8f16_type
)

k_row_b = lane16 + fx.Index(st_base_row + 16)
k_row_b = lane16 + fx.Int64(st_base_row + 16)
k_lds_b = k_base + k_row_b * K_STRIDE + k_col
k_pack_b = fx.ptr_load(
lds_kv + fx.Int32(k_lds_b), result_type=v8f16_type
Expand Down Expand Up @@ -486,38 +483,38 @@ def reduction_peer(v_f32):
klane_off_i32 = klane_i32 * fx.Int32(8)
# st=0
_b0 = kv_start_i32 + fx.Int32(0) + klane_off_i32
s_v0 = ArithValue(_b0 > q_row_i32).select(c_neg_inf, s_v0)
s_v0 = (_b0 > q_row_i32).select(c_neg_inf, s_v0)
_b1 = kv_start_i32 + fx.Int32(1) + klane_off_i32
s_v1 = ArithValue(_b1 > q_row_i32).select(c_neg_inf, s_v1)
s_v1 = (_b1 > q_row_i32).select(c_neg_inf, s_v1)
_b2 = kv_start_i32 + fx.Int32(2) + klane_off_i32
s_v2 = ArithValue(_b2 > q_row_i32).select(c_neg_inf, s_v2)
s_v2 = (_b2 > q_row_i32).select(c_neg_inf, s_v2)
_b3 = kv_start_i32 + fx.Int32(3) + klane_off_i32
s_v3 = ArithValue(_b3 > q_row_i32).select(c_neg_inf, s_v3)
s_v3 = (_b3 > q_row_i32).select(c_neg_inf, s_v3)
_b4 = kv_start_i32 + fx.Int32(4) + klane_off_i32
s_v4 = ArithValue(_b4 > q_row_i32).select(c_neg_inf, s_v4)
s_v4 = (_b4 > q_row_i32).select(c_neg_inf, s_v4)
_b5 = kv_start_i32 + fx.Int32(5) + klane_off_i32
s_v5 = ArithValue(_b5 > q_row_i32).select(c_neg_inf, s_v5)
s_v5 = (_b5 > q_row_i32).select(c_neg_inf, s_v5)
_b6 = kv_start_i32 + fx.Int32(6) + klane_off_i32
s_v6 = ArithValue(_b6 > q_row_i32).select(c_neg_inf, s_v6)
s_v6 = (_b6 > q_row_i32).select(c_neg_inf, s_v6)
_b7 = kv_start_i32 + fx.Int32(7) + klane_off_i32
s_v7 = ArithValue(_b7 > q_row_i32).select(c_neg_inf, s_v7)
s_v7 = (_b7 > q_row_i32).select(c_neg_inf, s_v7)
# st=1 (st_base=16)
_b8 = kv_start_i32 + fx.Int32(16) + klane_off_i32
s_v8 = ArithValue(_b8 > q_row_i32).select(c_neg_inf, s_v8)
s_v8 = (_b8 > q_row_i32).select(c_neg_inf, s_v8)
_b9 = kv_start_i32 + fx.Int32(17) + klane_off_i32
s_v9 = ArithValue(_b9 > q_row_i32).select(c_neg_inf, s_v9)
s_v9 = (_b9 > q_row_i32).select(c_neg_inf, s_v9)
_b10 = kv_start_i32 + fx.Int32(18) + klane_off_i32
s_v10 = ArithValue(_b10 > q_row_i32).select(c_neg_inf, s_v10)
s_v10 = (_b10 > q_row_i32).select(c_neg_inf, s_v10)
_b11 = kv_start_i32 + fx.Int32(19) + klane_off_i32
s_v11 = ArithValue(_b11 > q_row_i32).select(c_neg_inf, s_v11)
s_v11 = (_b11 > q_row_i32).select(c_neg_inf, s_v11)
_b12 = kv_start_i32 + fx.Int32(20) + klane_off_i32
s_v12 = ArithValue(_b12 > q_row_i32).select(c_neg_inf, s_v12)
s_v12 = (_b12 > q_row_i32).select(c_neg_inf, s_v12)
_b13 = kv_start_i32 + fx.Int32(21) + klane_off_i32
s_v13 = ArithValue(_b13 > q_row_i32).select(c_neg_inf, s_v13)
s_v13 = (_b13 > q_row_i32).select(c_neg_inf, s_v13)
_b14 = kv_start_i32 + fx.Int32(22) + klane_off_i32
s_v14 = ArithValue(_b14 > q_row_i32).select(c_neg_inf, s_v14)
s_v14 = (_b14 > q_row_i32).select(c_neg_inf, s_v14)
_b15 = kv_start_i32 + fx.Int32(23) + klane_off_i32
s_v15 = ArithValue(_b15 > q_row_i32).select(c_neg_inf, s_v15)
s_v15 = (_b15 > q_row_i32).select(c_neg_inf, s_v15)
s_raw = [
s_v0,
s_v1,
Expand Down Expand Up @@ -547,7 +544,7 @@ def reduction_peer(v_f32):
# ---- Opt2: rocdl.exp2 ----
diff_m_raw = _fsub(m_running, m_new_raw)
diff_m_scaled = _fmul(diff_m_raw, c_sm_scale_log2e)
corr = rocdl.exp2(ir.F32Type.get(), _raw(diff_m_scaled))
corr = fx.rocdl.exp2(ir.F32Type.get(), _raw(diff_m_scaled))

scaled_max = _fmul(c_sm_scale_log2e, m_new_raw)
neg_scaled_max = _fsub(c_zero_f, scaled_max)
Expand All @@ -556,7 +553,7 @@ def reduction_peer(v_f32):
local_sum = _raw(c_zero_f)
for r in range_constexpr(NUM_S_VALS):
diff = fmath.fma(s_raw[r], _raw(c_sm_scale_log2e), neg_scaled_max)
p = rocdl.exp2(ir.F32Type.get(), _raw(diff))
p = fx.rocdl.exp2(ir.F32Type.get(), _raw(diff))
p_vals.append(p)
local_sum = _fadd(local_sum, p)

Expand All @@ -565,7 +562,7 @@ def reduction_peer(v_f32):
l_corr = _fmul(corr, l_running)
l_new = _fadd(l_corr, tile_sum)

corr_vec = Vec.from_elements([corr], fx.Float32).broadcast_to(8).ir_value()
corr_vec = Vec.from_elements([corr], fx.Float32).broadcast_to(8)
for dc in range_constexpr(D_CHUNKS):
o_accs[dc] = _fmul(o_accs[dc], corr_vec)

Expand All @@ -588,27 +585,25 @@ def reduction_peer(v_f32):
elem_list = []
for j in range_constexpr(8):
elem_list.append(fx.Float32(p_slice[j]).to(elem_dtype))
p_packs_st.append(
Vec.from_elements(elem_list, elem_dtype).ir_value()
)
p_packs_st.append(Vec.from_elements(elem_list, elem_dtype))
p_packs_all.append(p_packs_st)

# ==== GEMM2: O += V^T @ P (software pipelined, row-major V) ====
# Opt3: Prefetch next V pack while current WMMA executes
v_base = v_buf_base(0)

def _load_v_rowmajor(st_kv_base_val, pks_val, dc_val, v_base=v_base):
d_pos = fx.Index(dc_val * D_CHUNK) + lane16
d_pos = fx.Int64(dc_val * D_CHUNK) + lane16
v_elems = []
for k_sub in range_constexpr(8):
kv_row = (
fx.Index(st_kv_base_val + pks_val * PV_K_STEP)
fx.Int64(st_kv_base_val + pks_val * PV_K_STEP)
+ klane * WMMA_LANE_K
+ fx.Index(k_sub)
)
v_lds_idx = v_base + kv_row * V_STRIDE + d_pos
v_elems.append(fx.ptr_load(lds_kv + fx.Int32(v_lds_idx)))
return Vec.from_elements(v_elems, elem_dtype).ir_value()
return Vec.from_elements(v_elems, elem_dtype)

# Software pipeline: preload first V pack
cur_v_packs = []
Expand Down Expand Up @@ -657,12 +652,12 @@ def _load_v_rowmajor(st_kv_base_val, pks_val, dc_val, v_base=v_base):
o_finals = [loop_results[2 + dc] for dc in range_constexpr(D_CHUNKS)]

inv_l = arith.divf(_raw(c_one_f), _raw(l_final), fastmath=fm_fast)
inv_l_vec = Vec.from_elements([inv_l], fx.Float32).broadcast_to(8).ir_value()
inv_l_vec = Vec.from_elements([inv_l], fx.Float32).broadcast_to(8)

if q_in_bounds:
for dc in range_constexpr(D_CHUNKS):
o_norm_vec = _fmul(o_finals[dc], inv_l_vec)
o_trunc = Vec(o_norm_vec).to(elem_dtype).ir_value()
o_trunc = Vec(o_norm_vec).to(elem_dtype)
d_col = fx.Index(dc * D_CHUNK) + klane * 8
o_global = global_idx(q_row, d_col)
_store_global_half(o_ptr, o_global, o_trunc)
Expand Down
Loading
Loading