diff --git a/kernels/attention/pa_decode_fp8.py b/kernels/attention/pa_decode_fp8.py index cdd36a7a..b1867529 100644 --- a/kernels/attention/pa_decode_fp8.py +++ b/kernels/attention/pa_decode_fp8.py @@ -12,67 +12,24 @@ Requires: aiter's get_pa_metadata_v1 (module_pa_metadata.so) """ -import functools -import math +from __future__ import annotations import torch -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import vector -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 from kernels.attention.pa_decode_swa import compile_pa_decode_sw, compile_pa_decode_sw_reduce +from kernels.attention.pa_decode_tile import pa_decode_tile from kernels.attention.pa_metadata import compile_pa_decode_metadata -from kernels.common import buffer_ops, dpp_utils from kernels.common.tensor_shim import _run_compiled -from kernels.common.utils import ( - cdiv, - exp2_f32_fast, - global_load_i32, - global_load_i64x2, - global_ptr_from_addr, - rcp_f32, - udiv_const, - unflatten_k, - urem_const, -) +from kernels.common.utils import cdiv # ── Kernel geometry constants ──────────────────────────────────────── -KV_BLOCK_SIZE = 1024 # physical page size (matches SP3 kBlockSize) KV_COMPUTE_BLOCK = 256 # tile size (matches SP3 kTileKV) # Persistent-grid oversubscription for the metadata decode path: launch # CU_count * this many workgroups so the HW keeps multiple workgroups resident # per CU (memory-latency hiding). 1 = original (1 wg/CU). _PA_METADATA_GRID_OVERSUB = 3 -NUM_WARPS = 4 -WARP_SIZE = 64 -BLOCK_THREADS = NUM_WARPS * WARP_SIZE # 256 MFMA_N = 16 -TOKENS_PER_WARP = KV_COMPUTE_BLOCK // NUM_WARPS # 64 -TLOOP = TOKENS_PER_WARP // MFMA_N # 4 -ROWS_PER_WARP = WARP_SIZE // MFMA_N # 4 -FP8_ELEMS_16B = 16 # 16 FP8 per 16-byte load -QKHE_PER_FETCH = FP8_ELEMS_16B * ROWS_PER_WARP # 64 - -VTLOOP = NUM_WARPS # 4 -Q_ELEMS_PER_LANE = 8 -Q_CHUNKS_PER_LANE = Q_ELEMS_PER_LANE // 4 - -# LDS sizes -PROB_ROW_STRIDE_BYTES = 40 # 32 data + 8 padding -> 0 bank conflict -LDS_LOGITS_BYTES = NUM_WARPS * 4 * MFMA_N * PROB_ROW_STRIDE_BYTES # 10240 -LDS_SOFTMAX_BYTES = 2 * NUM_WARPS * MFMA_N * 4 # 512 -LDS_SCALE_V_PADDING = 4 # break K/V same-bank paired writes -LDS_SCALE_V_OFFSET = KV_COMPUTE_BLOCK + LDS_SCALE_V_PADDING -LDS_SCALE_BYTES = (LDS_SCALE_V_OFFSET + KV_COMPUTE_BLOCK) * 4 # K/V per-token scale staging - -FP8_MAX = 240.0 -LOG2E = 1.4426950408889634 - _PACKED_FP8_QUERY_DTYPES = tuple( dtype for dtype in ( @@ -84,565 +41,6 @@ ) -def _flatten_v_results(v_results, vhe_loop: int = 2): - """v_results[vt][vhe] = i64x2 → flat list of 2 * VTLOOP * vhe_loop scalar i64 - values, in the same order ``_unflatten_v_results`` expects. Used to carry - V data through scf.for state (which only accepts scalar values).""" - flat = [] - for vt in range(VTLOOP): - for vhe in range(vhe_loop): - v_i64x2 = fx.Vector(v_results[vt][vhe]) - flat.append(v_i64x2[0]) - flat.append(v_i64x2[1]) - return flat - - -def _unflatten_v_results(v_flat, vhe_loop: int = 2): - """Inverse of ``_flatten_v_results``: rebuild v_results[vt][vhe] = i64x2.""" - v_results = [] - idx = 0 - for vt in range(VTLOOP): - vhe_data = [] - for vhe in range(vhe_loop): - v_i64x2 = vector.from_elements(T.vec(2, T.i64), [as_ir_value(v_flat[idx]), as_ir_value(v_flat[idx + 1])]) - vhe_data.append(v_i64x2) - idx += 2 - v_results.append(vhe_data) - return v_results - - -def _build_pa_thread_invariants( - warp_id, - lane16id, - rowid, - *, - per_token_kv, -): - c_tokens_per_warp = fx.Int32(TOKENS_PER_WARP) - kv_tok_thread_base = warp_id * c_tokens_per_warp + rowid * 4 - rowid_8x8 = rowid >> fx.Int32(1) - offset_in_slot = rowid & fx.Int32(1) - prob_wr_thread_base = ( - warp_id * fx.Int32(4 * MFMA_N * PROB_ROW_STRIDE_BYTES) - + lane16id * fx.Int32(PROB_ROW_STRIDE_BYTES) - + rowid_8x8 * fx.Int32(8) - + offset_in_slot * 4 - ) - pv_prob_read_base = rowid * fx.Int32(MFMA_N * PROB_ROW_STRIDE_BYTES) + lane16id * fx.Int32(PROB_ROW_STRIDE_BYTES) - - sm_lane_wave_base = lane16id * fx.Int32(NUM_WARPS) - sm_max_off = sm_lane_wave_base + warp_id - sm_sum_off = fx.Int32(NUM_WARPS * MFMA_N) + sm_lane_wave_base + warp_id - sm_rd_max_offs = [sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] - sm_rd_sum_offs = [fx.Int32(NUM_WARPS * MFMA_N) + sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] - - sm_vmax_wr_off = None - sm_vmax_rd_offs = None - if const_expr(per_token_kv): - sm_vmax_wr_off = fx.Int32(2 * NUM_WARPS * MFMA_N) + sm_lane_wave_base + warp_id - sm_vmax_rd_offs = [fx.Int32(2 * NUM_WARPS * MFMA_N) + sm_lane_wave_base + fx.Int32(w) for w in range(NUM_WARPS)] - - return ( - kv_tok_thread_base, - prob_wr_thread_base, - pv_prob_read_base, - sm_max_off, - sm_sum_off, - sm_rd_max_offs, - sm_rd_sum_offs, - sm_vmax_wr_off, - sm_vmax_rd_offs, - ) - - -def _compute_mtp_group_state( - lane16id, - local_qhead_idx, - *, - mtp_group_idx, - query_length, - query_group_size, -): - g_off = mtp_group_idx * 16 - lane_pair_raw = lane16id + fx.Int32(g_off) - c_total_pairs = fx.Int32(query_length * query_group_size) - c_pair_max = fx.Int32(query_length * query_group_size - 1) - c_ql_m1 = fx.Int32(query_length - 1) - - if const_expr((query_length * query_group_size) % MFMA_N == 0): - lane_pair = lane_pair_raw - else: - lane_pair = arith.select(lane_pair_raw < c_total_pairs, lane_pair_raw, c_pair_max) - qi_raw = udiv_const(lane_pair, query_group_size) - if const_expr((query_length * query_group_size) % MFMA_N == 0): - qi_val = qi_raw - else: - qi_val = arith.select(qi_raw < c_ql_m1, qi_raw, c_ql_m1) - qhi_pos = urem_const(lane_pair, query_group_size) - - lqh_pair_raw = local_qhead_idx + fx.Int32(g_off) - if const_expr((query_length * query_group_size) % MFMA_N == 0): - lqh_pair = lqh_pair_raw - else: - lqh_pair = arith.select(lqh_pair_raw < c_total_pairs, lqh_pair_raw, c_pair_max) - lqi_raw = udiv_const(lqh_pair, query_group_size) - if const_expr((query_length * query_group_size) % MFMA_N == 0): - qi_for_q = lqi_raw - else: - qi_for_q = arith.select(lqi_raw < c_ql_m1, lqi_raw, c_ql_m1) - local_qhead_idx_for_q = urem_const(lqh_pair, query_group_size) - return qi_val, qhi_pos, qi_for_q, local_qhead_idx_for_q - - -@flyc.jit -def _finish_q_fragments( - logits_lds_i32, - logits_lds_i64, - softmax_lds_f32, - q_chunks, - lane16id, - rowid, - local_qhead_idx, - *, - head_size: int, - qkhe_loop: int, - q_lanes_per_head: int, -): - # LDS Q layout (compact, per-qhead contiguous): - # Q[head=h][hd=d] at byte offset h * HEAD_SIZE + d (FP8 after conversion) - # Total Q footprint = 16 qheads * HEAD_SIZE bytes, aliased with the later P - # writes via `logits_lds_i32 / logits_lds_i64` (same base). For HEAD_SIZE=64, - # only the first 8 lanes write Q for each qhead. - # - # Writer: thread (warp_id W, rowid R', lane16id L') owns qhead = W*4 + R' = - # `local_qhead_idx`, and within that qhead owns the 8 FP8 elements at - # head_dim [L'*8 .. L'*8+7]. We therefore write 2 i32 words (= 1 i64 = 8 B) - # at `local_qhead_idx * HEAD_SIZE + lane16id * 8`. - # - # Reader: MFMA lane layout for mfma_f32_16x16x32_fp8_fp8 (B = Q^T, N = qhead, - # K = head_dim) — reverse-engineered from `_load_k_flat`: thread (rowid R, - # lane16id L) consumes, for k_step = qkhe*2 + qkr, - # Q[head = L][hd = (qkhe*4 + R) * 16 + qkr * 8 + 0..7] - # i.e. the read byte offset is `L * HEAD_SIZE + qkhe*64 + R*16 + qkr*8`. - c_head_size = fx.Int32(head_size) - lds_q_base = local_qhead_idx * c_head_size + lane16id * 8 - abs_mask = fx.Vector.filled(4, 0x7FFFFFFF, fx.Int32) - c_zero_f = fx.Float32(0.0) - c_one_f = fx.Float32(1.0) - - q_f32_chunks = [] - local_max = c_zero_f - for q_src in q_chunks: - q_f32 = fx.Vector(q_src).to(fx.Float32) - q_f32_chunks.append(q_f32) - q_i32 = q_f32.bitcast(fx.Int32) - q_abs_i32 = q_i32 & abs_mask - q_abs = q_abs_i32.bitcast(fx.Float32) - chunk_max = q_abs.reduce("max") - local_max = fx.maxnumf(local_max, chunk_max) - - for sh in [8, 4, 2, 1]: - 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, - local_max * fx.Float32(1.0 / FP8_MAX).ir_value(), - c_one_f, - ) - ) - inv_query_scale = rcp_f32(query_scale_lane) - q_words = [] - for q_f32 in q_f32_chunks: - p = q_f32 * inv_query_scale - lo = rocdl.cvt_pk_fp8_f32(T.i32, p[0], p[1], fx.Int32(0), False) - q_words.append(rocdl.cvt_pk_fp8_f32(T.i32, p[2], p[3], lo, True)) - q_w0, q_w1 = q_words - - if lane16id == fx.Int32(0): - fx.ptr_store( - fx.Vector.from_elements([query_scale_lane], dtype=fx.Float32), - softmax_lds_f32 + local_qhead_idx, - ) - - v01 = fx.Vector.from_elements([q_w0, q_w1], dtype=fx.Int32) - lds_q_i32 = lds_q_base >> fx.Int32(2) - if const_expr(q_lanes_per_head < MFMA_N): - if lane16id < fx.Int32(q_lanes_per_head): - fx.ptr_store(v01, logits_lds_i32 + lds_q_i32) - else: - fx.ptr_store(v01, logits_lds_i32 + lds_q_i32) - - q_frags = [] - gpu.barrier() - query_scale_lane = fx.ptr_load(softmax_lds_f32 + (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): - # See layout comment above. Byte offset: - # lane16id * HEAD_SIZE + qkhe*64 + rowid*16 + qkr*8 - lds_rd_byte = lane16id * c_head_size + fx.Int32(qkhe << 6) + (rowid << fx.Int32(4)) + fx.Int32(qkr << 3) - lds_rd_base = lds_rd_byte >> fx.Int32(3) - q_v1 = fx.ptr_load(logits_lds_i64 + (lds_rd_base), result_type=fx.Vector.make_type(1, fx.Int64)) - q_frags.append(q_v1[0]) - return q_frags, query_scale_lane - - -def _prefetch_mtp_group_query( - q_rsrc, - batch_idx, - kv_h, - stride_q_seq, - stride_q_head, - lane16id, - local_qhead_idx, - *, - mtp_group_idx, - query_length, - query_group_size, - query_load_is_bf16, - q_lanes_per_head, -): - qi_val, qhi_pos, qi_for_q, local_qhead_idx_for_q = _compute_mtp_group_state( - lane16id, - local_qhead_idx, - mtp_group_idx=mtp_group_idx, - query_length=query_length, - query_group_size=query_group_size, - ) - q_row = batch_idx * arith.constant(query_length, type=T.i32) + qi_for_q - q_base = ( - q_row * stride_q_seq - + (kv_h * arith.constant(query_group_size, type=T.i32) + local_qhead_idx_for_q) * stride_q_head - ) - q_chunks = _prefetch_q_chunks( - q_rsrc, - q_base, - lane16id, - query_load_is_bf16=query_load_is_bf16, - q_lanes_per_head=q_lanes_per_head, - ) - return qi_val, qhi_pos, q_chunks - - -def _finish_mtp_group_q_fragments( - logits_lds_i32, - logits_lds_i64, - softmax_lds_f32, - mtp_prefetch, - lane16id, - rowid, - local_qhead_idx, - *, - head_size: int, - qkhe_loop: int, - q_lanes_per_head: int, -): - qi_val, qhi_pos, q_chunks = mtp_prefetch - q_frags, query_scale_lane = _finish_q_fragments( - logits_lds_i32, - logits_lds_i64, - softmax_lds_f32, - q_chunks, - lane16id, - rowid, - local_qhead_idx, - head_size=head_size, - qkhe_loop=qkhe_loop, - q_lanes_per_head=q_lanes_per_head, - ) - return qi_val, qhi_pos, q_frags, query_scale_lane - - -def _normalize_pa_output(running_sum, outs, zero_f): - one_f = fx.Float32(1.0).ir_value() - safe_sum = arith.select(running_sum > zero_f, running_sum, one_f) - inv_sum = rcp_f32(safe_sum) - inv_sum_vec = vector.broadcast(T.f32x4, inv_sum) - return [out * inv_sum_vec for out in outs] - - -@flyc.jit -def _make_pa_phase_helpers( - *, - per_token_q, - per_token_kv, - needs_mask, - query_length, - logits_lds_i32, - logits_lds_i64, - softmax_lds_f32, - scale_lds_f32, - softmax_scale_base, - softmax_q_scale, - k_scale_val, - scale, - v_scale_val, - warp_id, - lane16id, - rowid, - kv_tok_thread_base, - prob_wr_thread_base, - pv_prob_read_base, - sm_max_off, - sm_sum_off, - sm_rd_max_offs, - sm_rd_sum_offs, - sm_vmax_wr_off, - sm_vmax_rd_offs, - c_w, - neg_inf, - zero_f, - qkhe_loop: int = 2, - vhe_loop: int = 2, -): - apply_causal_mask = needs_mask or query_length > 1 - pv_prob_i64_indices = [] - for vt in range_constexpr(VTLOOP): - for j in range_constexpr(2): - p_byte = ( - arith.constant(vt * 4 * MFMA_N * PROB_ROW_STRIDE_BYTES, type=T.i32) - + pv_prob_read_base - + arith.constant(j * 8, type=T.i32) - ) - pv_prob_i64_indices.append(p_byte >> fx.Int32(3)) - - 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_lds_f32 + (_scale_row_base(td)), result_type=fx.Vector.make_type(4, fx.Float32) - ).ir_value() - - def _load_v_scale_vec(td: int): - return fx.ptr_load( - scale_lds_f32 + (fx.Int32(LDS_SCALE_V_OFFSET) + _scale_row_base(td)), - result_type=fx.Vector.make_type(4, fx.Float32), - ).ir_value() - - def _get_k_scale_vec(td: int, k_scale_vecs=None): - if const_expr(per_token_kv): - return k_scale_vecs[td] - return _load_k_scale_vec(td) - - def _get_v_scale_vec(td: int, v_scale_vecs=None): - if const_expr(per_token_kv): - return v_scale_vecs[td] - return _load_v_scale_vec(td) - - def _store_vmax_warp(partition_start, *, seq_end=None, v_scale_vecs=None): - if const_expr(per_token_kv): - kv_tok_base = partition_start + kv_tok_thread_base if const_expr(seq_end is not None) else None - v_max_warp = zero_f - for td in range_constexpr(TLOOP): - vs = _get_v_scale_vec(td, v_scale_vecs) - for i in range_constexpr(4): - if const_expr(kv_tok_base is not None): - kv_tok = kv_tok_base + arith.constant(td * MFMA_N + i, type=T.i32) - 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 = fx.maxnumf(v_max_warp, fx.Vector(vs).reduce("max")) - for sh in [32, 16]: - 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_lds_f32 + sm_vmax_wr_off, - ) - - def _token_vec_i32(kv_tok_base, td: int): - kv_tok_td_base = kv_tok_base + arith.constant(td * MFMA_N, type=T.i32) - return fx.Vector.from_elements( - [kv_tok_td_base + arith.constant(i, type=T.i32) for i in range_constexpr(4)], - dtype=fx.Int32, - ) - - def _apply_token_mask_vec(logit_vec, td: int, kv_tok_base, causal_bound, false_value): - tok_vec = _token_vec_i32(kv_tok_base, td) - if const_expr(apply_causal_mask): - in_range = tok_vec < causal_bound - return arith.select(in_range, logit_vec, vector.broadcast(T.f32x4, arith.unwrap(false_value))) - return logit_vec - - def _qk_and_intra_softmax( - k_ops, - partition_start, - q_frags, - causal_bound, - query_scale_lane=None, - *, - preloaded_scales=None, - ): - if const_expr(preloaded_scales is not None): - k_scale_vecs, v_scale_vecs = preloaded_scales - - query_scale_vec = None - if const_expr(per_token_q): - query_scale_vec = vector.broadcast(T.f32x4, query_scale_lane * softmax_scale_base) - d_out = [] - for td in range_constexpr(TLOOP): - acc = arith.constant_vector(0.0, T.f32x4) - for k_step in range_constexpr(qkhe_loop * 2): - acc = rocdl.mfma_f32_16x16x32_fp8_fp8(T.f32x4, [k_ops[td][k_step], q_frags[k_step], acc, 0, 0, 0]) - if const_expr(per_token_kv): - k_scale_vec = _get_k_scale_vec(td, k_scale_vecs) - scale_vec = ( - k_scale_vec * query_scale_vec - if const_expr(per_token_q) - else k_scale_vec * vector.broadcast(T.f32x4, softmax_q_scale) - ) - d_out.append(acc * scale_vec) - else: - if const_expr(per_token_q): - d_out.append(acc * (query_scale_vec * vector.broadcast(T.f32x4, k_scale_val))) - else: - d_out.append(acc * vector.broadcast(T.f32x4, scale)) - - kv_tok_base = partition_start + kv_tok_thread_base if const_expr(apply_causal_mask) else None - qk_max = neg_inf - for td in range_constexpr(TLOOP): - logits_vec = d_out[td] - if const_expr(kv_tok_base is not None): - logits_vec = _apply_token_mask_vec(logits_vec, td, kv_tok_base, causal_bound, neg_inf) - d_out[td] = logits_vec - qk_max = fx.maxnumf(qk_max, fx.Vector(logits_vec).reduce("max")) - for sh in [32, 16]: - 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_lds_f32 + sm_max_off, - ) - - if const_expr(per_token_kv): - return d_out, v_scale_vecs - return d_out - - 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_lds_f32 + (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]) - - new_rmax = fx.maxnumf(rmax, partition_max) - safe_eff_max = arith.select(partition_max > neg_inf, new_rmax, zero_f) if const_expr(needs_mask) else new_rmax - local_exp_sum = zero_f - for td in range_constexpr(TLOOP): - diff_vec = fx.Vector(d_out[td]) - vector.broadcast(T.f32x4, arith.unwrap(safe_eff_max)) - p_vec = exp2_f32_fast(diff_vec * vector.broadcast(T.f32x4, arith.unwrap(fx.Float32(LOG2E)))) - local_exp_sum = local_exp_sum + fx.Vector(p_vec).reduce("add") - d_out[td] = p_vec - for sh in [32, 16]: - local_exp_sum = local_exp_sum + local_exp_sum.shuffle_xor(arith.constant(sh, type=T.i32), c_w) - fx.ptr_store( - fx.Vector.from_elements([local_exp_sum], dtype=fx.Float32), - softmax_lds_f32 + sm_sum_off, - ) - if const_expr(needs_mask): - accum_scale = arith.select( - rmax > neg_inf, - exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()), - zero_f, - ) - else: - accum_scale = exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()) - - gpu.barrier() - sum_vec = fx.ptr_load(softmax_lds_f32 + (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 - ) - - 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) - rmax = new_rmax - accum_scale_vec = vector.broadcast(T.f32x4, arith.unwrap(accum_scale)) - for vhe in range_constexpr(vhe_loop): - outs[vhe] = outs[vhe] * accum_scale_vec - - if const_expr(per_token_kv): - v_max_global = zero_f - vmax_vec = fx.ptr_load( - softmax_lds_f32 + (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) - 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) - v_correction = v_max_scaled - _vec_norm_p = arith.unwrap(norm_factor) - for td in range_constexpr(TLOOP): - d_out[td] = d_out[td] * (_get_v_scale_vec(td, v_scale_vecs) * vector.broadcast(T.f32x4, _vec_norm_p)) - else: - v_correction = v_scale_val - - for td in range_constexpr(TLOOP): - p0 = vector.extract(as_ir_value(d_out[td]), static_position=[0], dynamic_position=[]) - p1 = vector.extract(as_ir_value(d_out[td]), static_position=[1], dynamic_position=[]) - p2 = vector.extract(as_ir_value(d_out[td]), static_position=[2], dynamic_position=[]) - p3 = vector.extract(as_ir_value(d_out[td]), static_position=[3], dynamic_position=[]) - lo = rocdl.cvt_pk_fp8_f32(T.i32, p0, p1, arith.constant(0, type=T.i32), False) - pk = rocdl.cvt_pk_fp8_f32(T.i32, p2, p3, lo, True) - byte_base = prob_wr_thread_base + arith.constant(td * MFMA_N * PROB_ROW_STRIDE_BYTES, type=T.i32) - i32_off = byte_base >> fx.Int32(2) - pk_vec = vector.from_elements(T.vec(1, T.i32), [as_ir_value(pk)]) - fx.ptr_store(pk_vec, logits_lds_i32 + i32_off) - return rmax, rsum, outs, v_correction - - 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 ── - # `p_i64` depends only on (vt, j), NOT on vhe, so the previous - # per-vhe inner LDS load was redundant: VHELOOP × VTLOOP*2 reads - # of the same VTLOOP*2 LDS slots. Issue all VTLOOP*2 ds_read_b64 - # ops once at the start so the compiler pipelines them — lgkmcnt - # drains during the address arithmetic before the MFMA chain. - p_i64_all = [] - for vt in range_constexpr(VTLOOP): - for j in range_constexpr(2): - p_i64_idx = pv_prob_i64_indices[vt * 2 + j] - p_i64_all.append( - fx.ptr_load(logits_lds_i64 + (p_i64_idx), result_type=fx.Vector.make_type(1, fx.Int64))[0] - ) - - for vhe in range_constexpr(vhe_loop): - tmp_out = arith.constant_vector(0.0, T.f32x4) - for vt in range_constexpr(VTLOOP): - v_i64x2 = fx.Vector(v_ops[vt][vhe]) - for j in range_constexpr(2): - tmp_out = rocdl.mfma_f32_16x16x32_fp8_fp8( - T.f32x4, - [ - v_i64x2[j], - p_i64_all[vt * 2 + j], - tmp_out, - 0, - 0, - 0, - ], - ) - outs[vhe] = arith.addf( - arith.mulf(tmp_out, v_correction_vec, fastmath=fm_contract), - outs[vhe], - fastmath=fm_contract, - ) - return outs - - return ( - _store_vmax_warp, - _qk_and_intra_softmax, - _cross_warp_softmax_and_prob_pack, - _pv_mfma, - ) - - # ===================================================================== # Launch API — Persistent Scheduling mode # ===================================================================== @@ -877,825 +275,6 @@ def get_recommended_splits( _PA_DECODE_PS_SMALL_BLOCK_SIZES = (16, 64) -@flyc.jit -def _pa_small_block_load_k_flat( - k_global_ptr, - kv_h_i32, - stride_k_block_i32, - stride_k_head_i32, - lane16id_i32, - rowid_i32, - *, - block_size: int, - phys_blocks, - qkhe_loop: int = 2, -): - """Load K data for one warp's 64-token slice of a 256-token partition. - - Returns ``k_flat`` (a list of ``TLOOP * qkhe_loop * 2`` i64 scalars) compatible - with ``unflatten_k`` and downstream MFMA invocations. - """ - c_he_stride_dw = fx.Int32(block_size * FP8_ELEMS_16B // 4) - c_tok_stride_dw = fx.Int32(FP8_ELEMS_16B // 4) - k_he_off_dw = [rowid_i32 * c_he_stride_dw + fx.Int32(qkhe * 4) * c_he_stride_dw for qkhe in range(qkhe_loop)] - k_head_off = kv_h_i32 * stride_k_head_i32 - - k_flat = [] - if const_expr(block_size == 64): - # Each warp owns exactly one physical block (64 tokens). - phys_block = phys_blocks - k_block_base_dw = _compute_block_base_dw_i64(phys_block, stride_k_block_i32, k_head_off) - for td in range_constexpr(TLOOP): - within_block_token = fx.Int32(td * MFMA_N) + lane16id_i32 - kbo_dw = within_block_token * c_tok_stride_dw - for qkhe in range_constexpr(qkhe_loop): - ka_dw = k_block_base_dw + fx.Int64(kbo_dw + k_he_off_dw[qkhe]) - k2 = global_load_i64x2(k_global_ptr, ka_dw * fx.Int64(4)) - k2_words = fx.Vector(k2) - k_flat.append(k2_words[0]) - k_flat.append(k2_words[1]) - else: - # block_size == 16: each warp spans 4 blocks (one MFMA tile per block). - within_block_token = lane16id_i32 - kbo_dw = within_block_token * c_tok_stride_dw - for td in range_constexpr(TLOOP): - phys_block = phys_blocks[td] - k_block_base_dw = _compute_block_base_dw_i64(phys_block, stride_k_block_i32, k_head_off) - for qkhe in range_constexpr(qkhe_loop): - ka_dw = k_block_base_dw + fx.Int64(kbo_dw + k_he_off_dw[qkhe]) - k2 = global_load_i64x2(k_global_ptr, ka_dw * fx.Int64(4)) - rocdl.sched_barrier(rocdl.mask_vmem_rd) - k2_words = fx.Vector(k2) - k_flat.append(k2_words[0]) - k_flat.append(k2_words[1]) - return k_flat - - -@flyc.jit -def _pa_small_block_load_v_trans( - v_global_ptr, - kv_h_i32, - stride_v_block_i32, - stride_v_head_i32, - warp_id_i32, - lane16id_i32, - rowid_i32, - v_phys_blocks, - *, - block_size: int, - head_size: int = 128, - vhe_loop: int = 2, -): - """Load V tiles for one CTA's 256-token partition (``trans_v=True``). - - Returns ``v_results[vt][vhe]`` (i64x2) indexed exactly as the reference - ``_load_v_and_scales`` so it can be passed as ``preloaded_v_and_scales``. - """ - v_head_off = kv_h_i32 * stride_v_head_i32 - vhead_elems = [ - fx.Int32(vhe * NUM_WARPS * MFMA_N) + warp_id_i32 * fx.Int32(MFMA_N) + lane16id_i32 for vhe in range(vhe_loop) - ] - vhead_elem_dw = [vhead_elems[vhe] * fx.Int32(FP8_ELEMS_16B // 4) for vhe in range(vhe_loop)] - c_subblock_dw = fx.Int32(head_size * FP8_ELEMS_16B // 4) - - v_results = [] - for vt in range_constexpr(VTLOOP): - phys_block = v_phys_blocks[vt] - if const_expr(block_size == 64): - # vt selects the physical block (4 blocks per partition); rowid - # selects the 16-token sub-block within that physical block. - sub_block_idx = rowid_i32 - else: - # block_size == 16: (vt * 4 + rowid) selects the block; only one - # 16-token sub-block per physical block, so sub_block_idx == 0. - sub_block_idx = fx.Int32(0) - v_block_base_dw = _compute_block_base_dw_i64(phys_block, stride_v_block_i32, v_head_off) - vhe_data = [] - for vhe in range_constexpr(vhe_loop): - va_dw_delta = sub_block_idx * c_subblock_dw + vhead_elem_dw[vhe] - va_byte = (v_block_base_dw + fx.Int64(va_dw_delta)) * fx.Int64(4) - v_i64x2 = global_load_i64x2(v_global_ptr, va_byte) - vhe_data.append(v_i64x2) - v_results.append(vhe_data) - return v_results - - -@functools.lru_cache(maxsize=256) -def compile_pa_decode_ps( - *, - block_size: int, - max_context_partition_num: int, - softmax_scale: float = None, - trans_v: bool = True, - query_group_size: int = 16, - per_token_kv: bool = False, - query_length: int = 1, - query_input_dtype: str = "bf16", - head_dim: int = 128, -): - """Compile the small-block partition kernel. See module-level comment.""" - if block_size not in _PA_DECODE_PS_SMALL_BLOCK_SIZES: - raise ValueError( - f"compile_pa_decode_ps: unsupported block_size={block_size}; " - f"expected one of {_PA_DECODE_PS_SMALL_BLOCK_SIZES}." - ) - if query_input_dtype not in ("bf16", "f16"): - raise ValueError("compile_pa_decode_ps currently expects bf16/f16 query inputs.") - if not trans_v: - raise NotImplementedError("compile_pa_decode_ps: trans_v=False not yet supported.") - if head_dim % QKHE_PER_FETCH != 0 or head_dim % (MFMA_N * NUM_WARPS) != 0 or head_dim % Q_ELEMS_PER_LANE != 0: - raise ValueError(f"Unsupported head_dim={head_dim}; must be a multiple of {MFMA_N * NUM_WARPS}.") - _HEAD = head_dim - _QKHELOOP = head_dim // QKHE_PER_FETCH - _VHELOOP = head_dim // MFMA_N // NUM_WARPS - _Q_LANES_PER_HEAD = head_dim // Q_ELEMS_PER_LANE - _N_K_h = TLOOP * _QKHELOOP * 2 - _N_V_FLAT_h = 2 * VTLOOP * _VHELOOP - - query_load_is_bf16 = query_input_dtype == "bf16" - if softmax_scale is None: - softmax_scale = 1.0 / (head_dim**0.5) - _softmax_scale = float(softmax_scale) - _block_size = block_size - _blocks_per_partition = KV_COMPUTE_BLOCK // _block_size - - _mtp_groups = max(1, math.ceil(query_length * query_group_size / 16)) - - # LDS allocation — same layout as compile_pa_decode_metadata's small-block - # path. per_token_kv adds a cross-warp v_scale_max region (appended to the - # softmax block) and a K/V per-token scale staging region. - LDS_VMAX_BYTES = NUM_WARPS * MFMA_N * 4 if const_expr(per_token_kv) else 0 - LDS_SOFTMAX_TOTAL = LDS_SOFTMAX_BYTES + LDS_VMAX_BYTES - - _LOGITS_I32 = LDS_LOGITS_BYTES // 4 - _SOFTMAX_F32 = LDS_SOFTMAX_TOTAL // 4 - _SCALE_F32 = LDS_SCALE_BYTES // 4 - _BT_I32 = NUM_WARPS * TLOOP - - if per_token_kv: - - @fx.struct - class SharedStorage: - logits: fx.Array[fx.Int32, _LOGITS_I32, 16] - softmax: fx.Array[fx.Float32, _SOFTMAX_F32, 16] - scale: fx.Array[fx.Float32, _SCALE_F32, 16] - bt: fx.Array[fx.Int32, _BT_I32, 16] - - else: - - @fx.struct - class SharedStorage: - logits: fx.Array[fx.Int32, _LOGITS_I32, 16] - softmax: fx.Array[fx.Float32, _SOFTMAX_F32, 16] - bt: fx.Array[fx.Int32, _BT_I32, 16] - - @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) - def pa_decode_ps_kernel( - # Raw-pointer kernargs: bare i64 data_ptr() (compact s_load prologue); - # shapes/strides come from the Int32 stride args below. - exp_sums_ptr: fx.Int64, - max_logits_ptr: fx.Int64, - tmp_out_ptr: fx.Int64, - query_ptr: fx.Int64, - key_cache_ptr: fx.Int64, - value_cache_ptr: fx.Int64, - block_tables_ptr: fx.Int64, - context_lengths_ptr: fx.Int64, - key_scale_ptr: fx.Int64, - value_scale_ptr: fx.Int64, - stride_q_seq: Int32, - stride_q_head: Int32, - stride_k_block: Int32, - stride_k_head: Int32, - stride_v_block: Int32, - stride_v_head: Int32, - stride_es_seq: Int32, - stride_es_head: Int32, - stride_es_part: Int32, - stride_to_seq: Int32, - stride_to_head: Int32, - stride_to_part: Int32, - stride_to_group: Int32, - stride_bt_seq: Int32, - # Per-token K/V scale strides (per_token_kv only), metadata layout - # `[num_blocks, num_kv_heads, block_size]`: - # stride_ks_block = num_kv_heads * block_size - # stride_ks_head = block_size - # Both 0 for per-tensor. - stride_ks_block: Int32, - stride_ks_head: Int32, - ): - tid = fx.Int32(gpu.thread_id("x")) - batch_idx = fx.Int32(gpu.block_id("x")) - kv_h = fx.Int32(gpu.block_id("y")) - partition_idx = fx.Int32(gpu.block_id("z")) - - cl_global_ptr = global_ptr_from_addr(context_lengths_ptr) - context_len = global_load_i32(cl_global_ptr, batch_idx) - - lane16id = tid & fx.Int32(15) - rowid = (tid >> fx.Int32(4)) & fx.Int32(3) - warp_id = tid >> fx.Int32(6) - - q_rsrc = buffer_ops.create_buffer_resource_from_addr(query_ptr) - k_global_ptr = global_ptr_from_addr(key_cache_ptr) - v_global_ptr = global_ptr_from_addr(value_cache_ptr) - # block_tables needs a real OOB bound (HW bounds-check returns 0 for - # empty-slot reads past the table); raw pointers carry no size, so pass - # the exact byte size: grid_dim.x (num_seqs) rows * stride_bt_seq i32. - _bt_records_bytes = fx.Int32(gpu.grid_dim.x) * stride_bt_seq * fx.Int32(4) - bt_rsrc = buffer_ops.create_buffer_resource_from_addr(block_tables_ptr, num_records_bytes=_bt_records_bytes) - es_rsrc = buffer_ops.create_buffer_resource_from_addr(exp_sums_ptr) - ml_rsrc = buffer_ops.create_buffer_resource_from_addr(max_logits_ptr) - to_rsrc = buffer_ops.create_buffer_resource_from_addr(tmp_out_ptr) - ks_rsrc = buffer_ops.create_buffer_resource_from_addr(key_scale_ptr) - vs_rsrc = buffer_ops.create_buffer_resource_from_addr(value_scale_ptr) - - q_scale_val = arith.constant(1.0, type=T.f32) - # Per-tensor K/V scales are loaded from index 0; per_token_kv uses - # per-token scales cross-iter-prefetched into VGPR and staged to LDS - # (see _load_my_kv_scale_from_vgpr / _stage_kv_scale_to_lds). - if const_expr(per_token_kv): - k_scale_val = arith.constant(1.0, type=T.f32) - v_scale_val = arith.constant(1.0, type=T.f32) - else: - k_scale_val = buffer_ops.buffer_load(ks_rsrc, arith.constant(0, type=T.i32), vec_width=1) - v_scale_val = buffer_ops.buffer_load(vs_rsrc, arith.constant(0, type=T.i32), vec_width=1) - - lds = fx.SharedAllocator().allocate(SharedStorage).peek() - logits_lds_i32 = lds.logits.ptr - logits_lds_i64 = fx.recast_iter(fx.Int64, lds.logits.ptr) - softmax_lds_f32 = lds.softmax.ptr - bt_lds_i32 = lds.bt.ptr - if const_expr(per_token_kv): - scale_lds_f32 = lds.scale.ptr - else: - scale_lds_f32 = None - - _softmax_scale_const = arith.constant(_softmax_scale, type=T.f32) - _softmax_q_scale = _softmax_scale_const * q_scale_val - _scale = _softmax_q_scale * k_scale_val - c_w = arith.constant(WARP_SIZE, type=T.i32) - NEG_INF = arith.constant(float("-inf"), type=T.f32) - ZERO_F = arith.constant(0.0, type=T.f32) - c_cps = arith.constant(KV_COMPUTE_BLOCK, type=T.i32) - c_query_group_size = arith.constant(query_group_size, type=T.i32) - - local_qhead_idx = warp_id * arith.constant(4, type=T.i32) + rowid - - ( - _kv_tok_thread_base, - _prob_wr_thread_base, - _pv_prob_read_base, - _sm_max_off, - _sm_sum_off, - _sm_rd_max_offs, - _sm_rd_sum_offs, - _sm_vmax_wr_off, - _sm_vmax_rd_offs, - ) = _build_pa_thread_invariants( - warp_id, - lane16id, - rowid, - per_token_kv=per_token_kv, - ) - - ( - _store_vmax_warp, - _qk_and_intra_softmax, - _cross_warp_softmax_and_prob_pack, - _pv_mfma, - ) = _make_pa_phase_helpers( - per_token_q=True, - per_token_kv=per_token_kv, - needs_mask=True, - query_length=query_length, - logits_lds_i32=logits_lds_i32, - logits_lds_i64=logits_lds_i64, - softmax_lds_f32=softmax_lds_f32, - scale_lds_f32=scale_lds_f32, - softmax_scale_base=_softmax_scale_const, - softmax_q_scale=_softmax_q_scale, - k_scale_val=k_scale_val, - scale=_scale, - v_scale_val=v_scale_val, - warp_id=warp_id, - lane16id=lane16id, - rowid=rowid, - kv_tok_thread_base=_kv_tok_thread_base, - prob_wr_thread_base=_prob_wr_thread_base, - pv_prob_read_base=_pv_prob_read_base, - sm_max_off=_sm_max_off, - sm_sum_off=_sm_sum_off, - sm_rd_max_offs=_sm_rd_max_offs, - sm_rd_sum_offs=_sm_rd_sum_offs, - sm_vmax_wr_off=_sm_vmax_wr_off, - sm_vmax_rd_offs=_sm_vmax_rd_offs, - c_w=c_w, - neg_inf=NEG_INF, - zero_f=ZERO_F, - qkhe_loop=_QKHELOOP, - vhe_loop=_VHELOOP, - ) - - def _store_partition_results(eqgs_lane, running_sum, running_max, outs_norm): - for vhe in range_constexpr(_VHELOOP): - hs_base = fx.Int32(vhe * NUM_WARPS * MFMA_N) + warp_id * fx.Int32(MFMA_N) + rowid * fx.Int32(4) - to_off = ( - batch_idx * stride_to_seq - + kv_h * stride_to_head - + partition_idx * stride_to_part - + eqgs_lane * stride_to_group - + hs_base - ) - out_bf16 = fx.Vector(outs_norm[vhe]).to(fx.BFloat16) - buffer_ops.buffer_store(out_bf16, to_rsrc, to_off) - es_off = batch_idx * stride_es_seq + kv_h * stride_es_head + partition_idx * stride_es_part + eqgs_lane - buffer_ops.buffer_store(fx.Float32(running_sum), es_rsrc, es_off) - buffer_ops.buffer_store(fx.Float32(running_max), ml_rsrc, es_off) - - # Slot covers one or more contiguous 256-token sub-partitions. The - # inner scf.for loop walks those sub-partitions with online-softmax - # loop-carried state, mirroring the Gluon `for sequence_partition_idx` - # loop in `paged_attention_decode_ps`. - c_max_parts = arith.constant(max_context_partition_num, type=T.i32) - num_total_partitions = (context_len + c_cps - fx.Int32(1)) >> fx.Int32(8) - page_size_partitions = (num_total_partitions + c_max_parts - fx.Int32(1)) // c_max_parts - local_partition_start = partition_idx * page_size_partitions - local_partition_end_raw = (partition_idx + fx.Int32(1)) * page_size_partitions - local_partition_end = arith.select( - local_partition_end_raw < num_total_partitions, - local_partition_end_raw, - num_total_partitions, - ) - - def _unwrap(v): - return v.ir_value() if hasattr(v, "ir_value") else v - - # Pack/unpack loop state. State is `_mtp_groups` accumulators, each a - # tuple of (rmax, rsum, outs...), plus the current sub-partition's K - # and V tiles (k_flat: _N_K i64 scalars, v_flat: 2 * _N_V i64 scalars - # — each V element is i64x2, flattened to two scalars). Both K and V - # are loop-carried so the body can use them while we prefetch the - # NEXT iteration's K and V (ping-pong). - state_width = 2 + _VHELOOP - - def _pack_states(states, k_flat, v_flat, k_scale_scalar=None, v_scale_scalar=None): - flat = [] - for st in states: - rmax, rsum = st[0], st[1] - outs = [st[2 + vhe] for vhe in range_constexpr(_VHELOOP)] - flat.extend([_unwrap(rmax), _unwrap(rsum)]) - flat.extend(_unwrap(out) for out in outs) - flat.extend(_unwrap(v) for v in k_flat) - flat.extend(_unwrap(v) for v in v_flat) - # Per-token K and V scale scalars are cross-iter-prefetched into VGPR - # via the loop's init/yield path (per_token_kv only). - if const_expr(per_token_kv): - flat.append(_unwrap(k_scale_scalar)) - flat.append(_unwrap(v_scale_scalar)) - return flat - - def _unpack_states(flat): - base = state_width * _mtp_groups - states = [ - tuple(flat[state_width * i + j] for j in range_constexpr(state_width)) - for i in range_constexpr(_mtp_groups) - ] - k_flat = list(flat[base : base + _N_K_h]) - v_flat = list(flat[base + _N_K_h : base + _N_K_h + _N_V_FLAT_h]) - if const_expr(per_token_kv): - _scale_base = base + _N_K_h + _N_V_FLAT_h - k_scale_scalar = flat[_scale_base] - v_scale_scalar = flat[_scale_base + 1] - else: - k_scale_scalar = None - v_scale_scalar = None - return states, k_flat, v_flat, k_scale_scalar, v_scale_scalar - - init_states = [ - tuple([NEG_INF, ZERO_F] + [arith.constant_vector(0.0, T.f32x4) for _ in range_constexpr(_VHELOOP)]) - for _ in range(_mtp_groups) - ] - - loop_start = fx.Index(arith.unwrap(local_partition_start)) - loop_end = fx.Index(arith.unwrap(local_partition_end)) - loop_step = arith.index(1) - last_partition_idx = local_partition_end - fx.Int32(1) - - def _pa_small_block_stage_phys_blocks(partition_block_base): - # bt offset is wave-uniform (batch_idx and warp_id are constant - # per wave, partition_block_base is workgroup-uniform). Use - # s_buffer_load to route through SMEM cache and land the result - # in SGPRs directly — eliminates the vmcnt(0) drain (was 25% of - # all kernel stalls) and the downstream readfirstlane. - if const_expr(block_size == 64): - bt_elem_off = batch_idx * stride_bt_seq + partition_block_base + warp_id - phys_blocks = buffer_ops.buffer_load(bt_rsrc, bt_elem_off, vec_width=1, is_scalar=True) - else: - bt_elem_off = batch_idx * stride_bt_seq + partition_block_base + warp_id * fx.Int32(TLOOP) - phys_blocks = buffer_ops.buffer_load(bt_rsrc, bt_elem_off, vec_width=TLOOP, is_scalar=True) - return phys_blocks - - def _pa_small_block_store_phys_blocks_to_lds(phys_block_vec): - if (lane16id | rowid) == fx.Int32(0): - if const_expr(block_size == 64): - # block_size=64: `_stage_phys_blocks` returned vec_width=1 - # → scalar i32, not a Vector. Wrap in a 1-element - # Each warp writes 1 i32 to bt_lds_i32[warp_id]; - # `_load_v_phys_blocks_from_lds` reads back the 4-elem - # vec starting at offset 0. - fx.ptr_store( - fx.Vector.from_elements([phys_block_vec], dtype=fx.Int32), - bt_lds_i32 + warp_id, - ) - else: - fx.ptr_store(phys_block_vec, bt_lds_i32 + warp_id * fx.Int32(TLOOP)) - - def _pa_small_block_load_v_phys_blocks_from_lds(): - v_phys_blocks = [] - if const_expr(block_size == 64): - phys_block_vec = fx.ptr_load( - bt_lds_i32 + (fx.Int32(0)), result_type=fx.Vector.make_type(VTLOOP, fx.Int32) - ) - for vt in range_constexpr(VTLOOP): - v_phys_blocks.append(phys_block_vec[vt]) - else: - for vt in range_constexpr(VTLOOP): - bt_lds_off = fx.Int32(vt * TLOOP) + rowid - phys_block = fx.ptr_load(bt_lds_i32 + (bt_lds_off), result_type=fx.Vector.make_type(1, fx.Int32))[0] - v_phys_blocks.append(phys_block) - return v_phys_blocks - - # Pre-load the FIRST (== reverse-order start = last partition) sub- - # partition's block-table entries before Q setup so the dependent K - # prefetch below does not also pay the table latency. - # Empty-slot guard: when num_total_partitions < max_context_partition_num, - # CTAs with partition_idx >= num_total_partitions get - # local_partition_start >= num_total_partitions and the inner loop runs - # 0 iters. But the prologue still issues block-table + K reads using - # `last_partition_idx`; clamp to 0 so all reads stay in-bounds (the - # results are unused since the loop never executes). - _safe_init_partition = arith.select( - local_partition_start < num_total_partitions, - last_partition_idx, - arith.constant(0, type=T.i32), - ) - first_block_base = _safe_init_partition * fx.Int32(_blocks_per_partition) - first_phys_blocks = _pa_small_block_stage_phys_blocks(first_block_base) - - # Pre-load Q for every MTP group ONCE before the KV loop. Each group's - # q_frags / qi / qhi / qscale are kept in registers across the entire - # KV loop, so we pay the Q-load cost (global → LDS → registers) exactly - # once per CTA regardless of how many sub-partitions the slot covers. - - q_frags_per_mtp = [] - qi_per_mtp = [] - qhi_per_mtp = [] - qscale_per_mtp = [] - for _mtp_g in range_constexpr(_mtp_groups): - mtp_prefetch = _prefetch_mtp_group_query( - q_rsrc, - batch_idx, - kv_h, - stride_q_seq, - stride_q_head, - lane16id, - local_qhead_idx, - mtp_group_idx=_mtp_g, - query_length=query_length, - query_group_size=query_group_size, - query_load_is_bf16=query_load_is_bf16, - q_lanes_per_head=_Q_LANES_PER_HEAD, - ) - qi_val, qhi_pos, q_frags, query_scale_lane = _finish_mtp_group_q_fragments( - logits_lds_i32, - logits_lds_i64, - softmax_lds_f32, - mtp_prefetch, - lane16id, - rowid, - local_qhead_idx, - head_size=_HEAD, - qkhe_loop=_QKHELOOP, - q_lanes_per_head=_Q_LANES_PER_HEAD, - ) - q_frags_per_mtp.append(q_frags) - qi_per_mtp.append(qi_val) - qhi_per_mtp.append(qhi_pos) - qscale_per_mtp.append(query_scale_lane) - - _pa_small_block_store_phys_blocks_to_lds(first_phys_blocks) - - # Per-token K/V scale: cross-iter VGPR prefetch (per_token_kv only). - # Each thread loads its own (k, v) scale scalar one iteration ahead - # (loop-carried) and stages it warp-locally to scale_lds, avoiding the - # in-loop bt_lds round-trip + cross-warp barrier. scale_idx uses the - # fp32 layout `[num_blocks, num_kv_heads, block_size]`, shared by K/V. - def _load_my_kv_scale_from_vgpr(phys_blocks_): - # phys_blocks_ from _pa_small_block_stage_phys_blocks (scalar bs64 / - # TLOOP-vec bs16) — reused from the K prefetch, so no bt_lds read. - if const_expr(_block_size == 64): - my_phys = fx.Int32(phys_blocks_) - tok_in_page = rowid * fx.Int32(MFMA_N) + lane16id - else: - # block_size==16: the warp owns TLOOP pages; rowid selects which. - my_phys = fx.Int32( - vector.extract( - as_ir_value(arith.unwrap(phys_blocks_)), - static_position=[ir.ShapedType.get_dynamic_size()], - dynamic_position=[as_ir_value(fx.Index(rowid))], - ) - ) - tok_in_page = lane16id - scale_idx = my_phys * stride_ks_block + kv_h * stride_ks_head + tok_in_page - k_scale_scalar = buffer_ops.buffer_load(ks_rsrc, scale_idx, vec_width=1, dtype=fx.Float32) - v_scale_scalar = buffer_ops.buffer_load(vs_rsrc, scale_idx, vec_width=1, dtype=fx.Float32) - return k_scale_scalar, v_scale_scalar - - def _stage_kv_scale_to_lds(k_scale_scalar, v_scale_scalar): - # Warp-local: warp w writes slots [w*64, w*64+64) and reads only - # those back, so the RAW is intra-wave (compiler lgkmcnt) — no barrier. - t = warp_id * fx.Int32(WARP_SIZE) + rowid * fx.Int32(MFMA_N) + lane16id - fx.ptr_store(fx.Vector.from_elements([k_scale_scalar], dtype=fx.Float32), scale_lds_f32 + t) - fx.ptr_store( - fx.Vector.from_elements([v_scale_scalar], dtype=fx.Float32), - scale_lds_f32 + (fx.Int32(LDS_SCALE_V_OFFSET) + t), - ) - - def _load_small_block_scale_vecs(): - k_scale_vecs = [] - v_scale_vecs = [] - for td in range_constexpr(TLOOP): - row = _kv_tok_thread_base + arith.constant(td * MFMA_N, type=T.i32) - k_scale_vecs.append( - fx.ptr_load(scale_lds_f32 + (row), result_type=fx.Vector.make_type(4, fx.Float32)).ir_value() - ) - v_scale_vecs.append( - fx.ptr_load( - scale_lds_f32 + (fx.Int32(LDS_SCALE_V_OFFSET) + row), - result_type=fx.Vector.make_type(4, fx.Float32), - ).ir_value() - ) - return k_scale_vecs, v_scale_vecs - - # Pre-load the FIRST sub-partition's K so the loop body can issue the - # next sub-partition's K prefetch in parallel with the current K's QK - # MFMA. For empty slots (loop_start == loop_end), this k_flat0 is - # computed using local_partition_start but never used because the - # loop runs 0 iterations — the block_table buffer_load is bounded so - # any OOB lookup safely returns 0 (→ block 0, then masked out by the - # softmax causal bound). - k_flat0 = _pa_small_block_load_k_flat( - k_global_ptr, - kv_h, - stride_k_block, - stride_k_head, - lane16id, - rowid, - block_size=_block_size, - phys_blocks=first_phys_blocks, - qkhe_loop=_QKHELOOP, - ) - gpu.barrier() - # ── Prologue V load ── - # V is cross-iter prefetched (ping-pong with K). Issue iter 0's V - # load here so the loop body can issue iter N+1's V at the END of - # iter N alongside K, both hidden behind the next iter's QK MFMA. - # `_pa_small_block_load_v_phys_blocks_from_lds` reads the LDS-staged - # first_phys_blocks written above; the barrier guarantees visibility. - _v_phys_blocks0 = _pa_small_block_load_v_phys_blocks_from_lds() - _v_results0 = _pa_small_block_load_v_trans( - v_global_ptr, - kv_h, - stride_v_block, - stride_v_head, - warp_id, - lane16id, - rowid, - _v_phys_blocks0, - block_size=_block_size, - head_size=_HEAD, - vhe_loop=_VHELOOP, - ) - v_flat0 = _flatten_v_results(_v_results0, vhe_loop=_VHELOOP) - # Prefetch iter 0's K/V scale (loop-carried via `init`). - if const_expr(per_token_kv): - k_scale_init, v_scale_init = _load_my_kv_scale_from_vgpr(first_phys_blocks) - else: - k_scale_init = None - v_scale_init = None - # No runtime `if _is_valid:` around this loop: the if-rewriter turns a - # body containing `ast.Yield` into a generator (empty then-region). Run - # unconditionally — empty slots iterate 0x and yield the init state. - for sub_part_ib, state in range( - loop_start, - loop_end, - loop_step, - init=_pack_states(init_states, k_flat0, v_flat0, k_scale_init, v_scale_init), - ): - cur_states, k_flat, v_flat, k_scale_cur, v_scale_cur = _unpack_states(state) - # Reverse iteration: scf.for walks sub_part_ib forward over - # [local_partition_start, local_partition_end); remap to walk - # sub_part_i32 from last_partition_idx down to local_partition_start - # so the sink-prone partition 0 is processed last. - _sub_raw_i32 = arith.index_cast(T.i32, sub_part_ib) - sub_part_i32 = last_partition_idx - (_sub_raw_i32 - local_partition_start) - sub_token_start = sub_part_i32 * c_cps - - # Both K and V come from the loop-carried state (prefetched at the - # END of the previous iteration). K's VMEM latency overlaps prev - # iter's PV MFMA; V's latency overlaps the entire next iter QK + - # softmax compute before PV consumes it. - k_ops = unflatten_k(k_flat, qkhe_loop=_QKHELOOP) - v_results = _unflatten_v_results(v_flat, vhe_loop=_VHELOOP) - - # Stage this iter's K/V scale (prefetched last iter, latency hidden). - if const_expr(per_token_kv): - _stage_kv_scale_to_lds(k_scale_cur, v_scale_cur) - k_scale_vecs, v_scale_vecs = _load_small_block_scale_vecs() - - # Compute the NEXT sub-partition's K base address (clamped to - # local_partition_start so the prefetch on the final loop - # iteration doesn't walk before the block_table window — - # k_next_flat is yielded out but never consumed since the loop - # terminates). Reverse iteration: next == sub_part_i32 - 1. - next_part_i32 = sub_part_i32 - fx.Int32(1) - next_safe_part = arith.select(next_part_i32 >= local_partition_start, next_part_i32, local_partition_start) - next_block_base = next_safe_part * fx.Int32(_blocks_per_partition) - - new_states = [] - k_next_flat = None - k_scale_next = None - v_scale_next = None - for _mtp_g in range_constexpr(_mtp_groups): - state = cur_states[_mtp_g] - rmax, rsum = state[0], state[1] - outs = [state[2 + vhe] for vhe in range_constexpr(_VHELOOP)] - causal_bound = context_len + arith.constant(1 - query_length, type=T.i32) + qi_per_mtp[_mtp_g] - - if const_expr(per_token_kv): - d_out, v_scales = _qk_and_intra_softmax( - k_ops, - sub_token_start, - q_frags_per_mtp[_mtp_g], - causal_bound, - query_scale_lane=qscale_per_mtp[_mtp_g], - preloaded_scales=(k_scale_vecs, v_scale_vecs), - ) - else: - d_out = _qk_and_intra_softmax( - k_ops, - sub_token_start, - q_frags_per_mtp[_mtp_g], - causal_bound, - query_scale_lane=qscale_per_mtp[_mtp_g], - ) - v_scales = None - - if const_expr(_mtp_g == _mtp_groups - 1): - next_phys_blocks = _pa_small_block_stage_phys_blocks(next_block_base) - # Prefetch next iter's K/V scale from the VGPR phys blocks. - if const_expr(per_token_kv): - k_scale_next, v_scale_next = _load_my_kv_scale_from_vgpr(next_phys_blocks) - - # per_token_kv needs the cross-warp v_scale_max staged to LDS so - # _cross_warp_softmax_and_prob_pack can read it for norm_factor. - if const_expr(per_token_kv): - _store_vmax_warp(sub_token_start, seq_end=context_len, v_scale_vecs=v_scales) - - gpu.barrier() - - rmax, rsum, outs, v_correction = _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs, v_scales) - - # Issue the next sub-partition's K prefetch on the LAST MTP - # iter, after cross_warp_softmax_and_prob_pack but BEFORE - # _pv_mfma — same hoist as in pa_decode_metadata_kenrel's - # _process_block_split. This lets the K VMEM load latency - # overlap with the upcoming PV MFMA compute. - if const_expr(_mtp_g == _mtp_groups - 1): - _pa_small_block_store_phys_blocks_to_lds(next_phys_blocks) - k_next_flat = _pa_small_block_load_k_flat( - k_global_ptr, - kv_h, - stride_k_block, - stride_k_head, - lane16id, - rowid, - block_size=_block_size, - phys_blocks=next_phys_blocks, - qkhe_loop=_QKHELOOP, - ) - gpu.barrier() - outs = _pv_mfma(v_results, outs, v_correction) - new_states.append(tuple([rmax, rsum] + outs)) - - # ── Cross-iter V prefetch (ping-pong) ── - # Issue NEXT iter's V load AFTER PV MFMA: the current iter's V - # vgprs are now consumed and can be reused. V phys_blocks come - # from the LDS-staged `next_phys_blocks` written above (the - # barrier after K prefetch ensures cross-warp visibility). The - # V VMEM latency is hidden behind next iter's QK MFMA + softmax. - _v_phys_blocks_next = _pa_small_block_load_v_phys_blocks_from_lds() - _v_next_results = _pa_small_block_load_v_trans( - v_global_ptr, - kv_h, - stride_v_block, - stride_v_head, - warp_id, - lane16id, - rowid, - _v_phys_blocks_next, - block_size=_block_size, - head_size=_HEAD, - vhe_loop=_VHELOOP, - ) - v_next_flat = _flatten_v_results(_v_next_results, vhe_loop=_VHELOOP) - - results = yield _pack_states(new_states, k_next_flat, v_next_flat, k_scale_next, v_scale_next) - - # Normalize and store one output slot per MTP group. - final_states, _final_k_flat, _final_v_flat, _final_k_scale, _final_v_scale = _unpack_states(results) - for _mtp_g in range_constexpr(_mtp_groups): - final_state = final_states[_mtp_g] - rmax_raw, rsum_raw = final_state[0], final_state[1] - outs_raw = [final_state[2 + vhe] for vhe in range_constexpr(_VHELOOP)] - running_max = fx.Float32(rmax_raw) - running_sum = fx.Float32(rsum_raw) - outs = [fx.Vector(out_raw) for out_raw in outs_raw] - outs_norm = _normalize_pa_output(running_sum, outs, ZERO_F) - eqgs_lane = qi_per_mtp[_mtp_g] * c_query_group_size + qhi_per_mtp[_mtp_g] - _store_partition_results(eqgs_lane, running_sum, running_max, outs_norm) - - @flyc.jit - def launch_pa_decode_ps_small_block( - exp_sums: fx.Int64, - max_logits: fx.Int64, - tmp_out: fx.Int64, - query: fx.Int64, - key_cache: fx.Int64, - value_cache: fx.Int64, - block_tables: fx.Int64, - context_lengths: fx.Int64, - key_scale: fx.Int64, - value_scale: fx.Int64, - s_q_seq: Int32, - s_q_head: Int32, - s_k_block: Int32, - s_k_head: Int32, - s_v_block: Int32, - s_v_head: Int32, - s_es_seq: Int32, - s_es_head: Int32, - s_es_part: Int32, - s_to_seq: Int32, - s_to_head: Int32, - s_to_part: Int32, - s_to_group: Int32, - s_bt_seq: Int32, - s_ks_block: Int32, - s_ks_head: Int32, - gx: Int32, - gy: Int32, - gz: Int32, - stream: fx.Stream = fx.Stream(None), - ): - pa_decode_ps_kernel( - exp_sums, - max_logits, - tmp_out, - query, - key_cache, - value_cache, - block_tables, - context_lengths, - key_scale, - value_scale, - 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_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_ps_small_block, - "kernel": pa_decode_ps_kernel, - "mtp_groups": _mtp_groups, - } - - def pa_decode_ps_launch( output: torch.Tensor, query: torch.Tensor, @@ -1896,7 +475,7 @@ def pa_decode_ps_launch( ) return "ps_sw_partitioned" - # ── small-block (block_size 16/64) → grid partition kernel + reduce ── + # ── small-block (block_size 16/64) → tile kernel ── # Key cache shape is [num_blocks, num_kv_heads, head_size // 16, block_size, 16]. block_size = key_cache.shape[-2] if block_size in _PA_DECODE_PS_SMALL_BLOCK_SIZES: @@ -1906,108 +485,39 @@ def pa_decode_ps_launch( "(per-sequence physical block index table)." ) batch_size = context_lengths.shape[0] - head_size = query.shape[-1] - eqgs = query_length * query_group_size - context_partition_size = KV_COMPUTE_BLOCK - blocks_per_partition = context_partition_size // block_size - if max_context_partition_num == 0: - max_context_partition_num = get_recommended_splits( - batch_size, - num_kv_heads, - split_kv_blocks=blocks_per_partition, - ) - if is_graph_capturing and (exp_sums is None or max_logits is None or temporary_output is None): - raise ValueError( - "CUDA graph capture requires preallocated `exp_sums`, `max_logits`, " - "and `temporary_output` for the small-block PS path." - ) - if exp_sums is None: - exp_sums = torch.zeros( - batch_size, num_kv_heads, max_context_partition_num, eqgs, device=dev, dtype=torch.float32 - ) - if max_logits is None: - max_logits = torch.full( - (batch_size, num_kv_heads, max_context_partition_num, eqgs), - float("-inf"), - device=dev, - dtype=torch.float32, - ) - if temporary_output is None: - temporary_output = torch.zeros( - batch_size, num_kv_heads, max_context_partition_num, eqgs, head_size, device=dev, dtype=torch.bfloat16 - ) - compiled_small = compile_pa_decode_ps( - block_size=block_size, - max_context_partition_num=max_context_partition_num, + if is_graph_capturing: + # Buffer sizes must be fixed ahead of capture and stay identical + # across every replay, so require the caller to have preallocated + # exp_sums/max_logits/temporary_output, exactly as the other PS + # paths already require. + if exp_sums is None or max_logits is None or temporary_output is None: + raise ValueError( + "CUDA graph capture requires preallocated `exp_sums`, `max_logits`, " + "and `temporary_output` for the tile-backed small-block PS path." + ) + # pa_decode_tile requires an exact [num_blocks, num_kv_heads, + # block_size] per-token scale shape; callers here may pass an extra + # trailing singleton dim (e.g. from a pertoken-quant helper), which + # reshape away without changing the strides. + if per_token_kv: + num_blocks = key_cache.shape[0] + key_scale = key_scale.reshape(num_blocks, num_kv_heads, block_size) + value_scale = value_scale.reshape(num_blocks, num_kv_heads, block_size) + pa_decode_tile( + output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, softmax_scale=softmax_scale, - trans_v=trans_v, - query_group_size=query_group_size, - per_token_kv=per_token_kv, - query_length=query_length, - query_input_dtype=query_input_dtype, - head_dim=int(head_size), - ) - output_5d = output.reshape(batch_size, query_length, num_kv_heads, query_group_size, head_size) - _run_compiled( - compiled_small["launch"], - exp_sums.data_ptr(), - max_logits.data_ptr(), - temporary_output.data_ptr(), - query.data_ptr(), - key_cache.data_ptr(), - value_cache.data_ptr(), - block_tables.data_ptr(), - context_lengths.data_ptr(), - key_scale.data_ptr(), - value_scale.data_ptr(), - query.stride(0), - query.stride(1), - key_cache.stride(0), - key_cache.stride(1), - value_cache.stride(0), - value_cache.stride(1), - exp_sums.stride(0), - exp_sums.stride(1), - exp_sums.stride(2), - temporary_output.stride(0), - temporary_output.stride(1), - temporary_output.stride(2), - temporary_output.stride(3), - block_tables.stride(0), - stride_ks_block, - stride_ks_head, - batch_size, - num_kv_heads, - max_context_partition_num, - s, - ) - compiled_sw_reduce = compile_pa_decode_sw_reduce( - max_context_partition_num=max_context_partition_num, - query_seq_len=query_length, - query_group_size=query_group_size, - head_size=head_size, - output_dtype_str=_get_output_dtype_str(output), - ) - _run_compiled( - compiled_sw_reduce["launch"], - output_5d.data_ptr(), - exp_sums.data_ptr(), - max_logits.data_ptr(), - temporary_output.data_ptr(), - output_5d.stride(0), - output_5d.stride(1), - output_5d.stride(2), - output_5d.stride(3), - exp_sums.stride(0), - exp_sums.stride(1), - exp_sums.stride(2), - temporary_output.stride(0), - temporary_output.stride(1), - temporary_output.stride(2), - temporary_output.stride(3), - batch_size, - num_kv_heads, - s, + stream=s, + num_partitions=max_context_partition_num, + pmax=max_logits, + psum=exp_sums, + pout=temporary_output, ) return "ps_small_block" @@ -2078,11 +588,11 @@ def pa_decode_ps_launch( s, ) - from kernels.attention.pa_metadata import pa_ps_reduce + from kernels.attention.pa_metadata import pa_metadata_reduce # Deterministic FlyDSL reduce replaces the racy aiter pa_reduce_v1/mla_reduce_v1 # (root cause of the flaky test_pa NaN). Same partial layout / reduce maps. - pa_ps_reduce( + pa_metadata_reduce( partial_output=partial_output[query_length:], partial_lse=partial_lse[query_length:], reduce_indptr=metadata["reduce_indptr"], diff --git a/kernels/attention/pa_decode_swa.py b/kernels/attention/pa_decode_swa.py index 44e26e60..74a092c5 100644 --- a/kernels/attention/pa_decode_swa.py +++ b/kernels/attention/pa_decode_swa.py @@ -9,12 +9,12 @@ import flydsl.expr as fx from flydsl._mlir.dialects import vector from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, rocdl -from flydsl.expr import math as fly_math from flydsl.expr.typing import Int32, T from kernels.attention.pa_common import _compute_block_base_dw_i64, _prefetch_q_chunks from kernels.common import buffer_ops, dpp_utils from kernels.common.utils import ( cdiv, + exp2_f32_fast, global_load_i32, global_load_i64x2, global_ptr_from_addr, @@ -66,10 +66,6 @@ def _get_sw_mtp_pair_offset(mtp_group_idx: int, mtp_subgroup_idx: int = 0) -> in return mtp_group_idx * MFMA_N + mtp_subgroup_idx * MFMA_N -def _exp2_f32_fast(value): - return fly_math.exp2(value, fastmath=arith.FastMathFlags.fast) - - def _load_k_flat( k_global_ptr, k_block_base_dw_i64, @@ -575,7 +571,7 @@ def _qk_and_intra_softmax( safe_qk_max = arith.select(qk_max > neg_inf, qk_max, zero_f) if const_expr(kv_tok_base is not None) else qk_max for td in range_constexpr(TLOOP): diff_vec = fx.Vector(d_out[td]) - vector.broadcast(T.f32x4, arith.unwrap(safe_qk_max)) - p_vec = _exp2_f32_fast(diff_vec * vector.broadcast(T.f32x4, arith.unwrap(fx.Float32(LOG2E)))) + p_vec = exp2_f32_fast(diff_vec * vector.broadcast(T.f32x4, arith.unwrap(fx.Float32(LOG2E)))) exp_sum = exp_sum + fx.Vector(p_vec).reduce("add") d_out[td] = p_vec for sh in [32, 16]: @@ -601,7 +597,7 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs): 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()) + 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) @@ -619,17 +615,17 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs): if const_expr(needs_mask): accum_scale = arith.select( rmax > neg_inf, - _exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()), + exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()), zero_f, ) part_to_new = arith.select( partition_max > neg_inf, - _exp2_f32_fast((partition_max - new_rmax) * fx.Float32(LOG2E).ir_value()), + exp2_f32_fast((partition_max - new_rmax) * fx.Float32(LOG2E).ir_value()), zero_f, ) else: - 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_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( @@ -665,12 +661,9 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs): d_out[td] = d_out[td] * vector.broadcast(T.f32x4, arith.unwrap(prob_scale)) for td in range_constexpr(TLOOP): - p0 = vector.extract(as_ir_value(d_out[td]), static_position=[0], dynamic_position=[]) - p1 = vector.extract(as_ir_value(d_out[td]), static_position=[1], dynamic_position=[]) - p2 = vector.extract(as_ir_value(d_out[td]), static_position=[2], dynamic_position=[]) - p3 = vector.extract(as_ir_value(d_out[td]), static_position=[3], dynamic_position=[]) - lo = rocdl.cvt_pk_fp8_f32(T.i32, p0, p1, arith.constant(0, type=T.i32), False) - pk = rocdl.cvt_pk_fp8_f32(T.i32, p2, p3, lo, True) + pv = fx.Vector(d_out[td]) + lo = rocdl.cvt_pk_fp8_f32(T.i32, pv[0], pv[1], arith.constant(0, type=T.i32), False) + pk = rocdl.cvt_pk_fp8_f32(T.i32, pv[2], pv[3], lo, True) elem_base = prob_wr_thread_base + arith.constant(td * MFMA_N * (PROB_ROW_STRIDE_BYTES // 4), type=T.i32) pk_vec = fx.Vector.from_elements([pk], dtype=fx.Int32) fx.ptr_store(pk_vec, logits_base + elem_base) @@ -737,7 +730,19 @@ def compile_pa_decode_sw_reduce( query_group_size: int, head_size: int, output_dtype_str: str, + logits_dtype_str: str = "bf16", ): + # Partition partials (`logits`) are read at this dtype; defaults to bf16 + # (this kernel's original, still-default behavior for existing callers). + # Callers with an fp8 query should still keep this at bf16 (fp8 has too + # little precision for a re-accumulated intermediate) -- only real f16/f32 + # queries should pick a matching non-bf16 value here. + if logits_dtype_str == "f32": + LOGITS_DTYPE = fx.Float32 + elif logits_dtype_str == "f16": + LOGITS_DTYPE = fx.Float16 + else: + LOGITS_DTYPE = fx.BFloat16 block_threads = head_size assert block_threads > 0, "head_size must be positive" assert block_threads <= 1024, "head_size must fit in one workgroup" @@ -877,7 +882,7 @@ def _wave_reduce_sum(val): global_max = _wave_reduce_max(part_max) part_scale = arith.select( lane_in_range, - _exp2_f32_fast((part_max - global_max) * c_log2e), + exp2_f32_fast((part_max - global_max) * c_log2e), c_zero_f, ) scaled_sum = part_sum * part_scale @@ -904,8 +909,8 @@ def _wave_reduce_sum(val): + eqgs_idx * stride_logits_group + tid ) - part_logits_bf16 = buffer_ops.buffer_load(logits_rsrc, logits_off, vec_width=1, dtype=fx.BFloat16) - part_logits = fx.Float32(part_logits_bf16) + part_logits_raw = buffer_ops.buffer_load(logits_rsrc, logits_off, vec_width=1, dtype=LOGITS_DTYPE) + part_logits = fx.Float32(part_logits_raw) acc = acc + part_logits * weight else: # Fallback for unusually large sliding-window partition counts. @@ -946,7 +951,7 @@ def _wave_reduce_sum(val): part_max = arith.select(in_chunk, part_max_raw, c_neg_inf) part_scale = arith.select( in_chunk, - _exp2_f32_fast((part_max - global_max) * c_log2e), + exp2_f32_fast((part_max - global_max) * c_log2e), c_zero_f, ) chunk_sum = _block_reduce(part_sum * part_scale, "sum") @@ -976,7 +981,7 @@ def _wave_reduce_sum(val): if in_chunk: part_sum = part_sum_raw part_max = part_max_raw - part_scale = _exp2_f32_fast((part_max - global_max) * c_log2e) + part_scale = exp2_f32_fast((part_max - global_max) * c_log2e) weight = part_sum * part_scale * inv_global_exp_sum part_idx_idx = fx.Int32(part_i32) fx.memref_store(weight, part_weights_lds, part_idx_idx) @@ -995,8 +1000,8 @@ def _wave_reduce_sum(val): + eqgs_idx * stride_logits_group + tid ) - part_logits_bf16 = buffer_ops.buffer_load(logits_rsrc, logits_off, vec_width=1, dtype=fx.BFloat16) - part_logits = fx.Float32(part_logits_bf16) + part_logits_raw = buffer_ops.buffer_load(logits_rsrc, logits_off, vec_width=1, dtype=LOGITS_DTYPE) + part_logits = fx.Float32(part_logits_raw) acc = acc + part_logits * weight query_idx = udiv_const(eqgs_idx, query_group_size) diff --git a/kernels/attention/pa_decode_tile.py b/kernels/attention/pa_decode_tile.py new file mode 100644 index 00000000..c89e7409 --- /dev/null +++ b/kernels/attention/pa_decode_tile.py @@ -0,0 +1,1206 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Readable tile-programming reference for paged-attention fp8 decode. + +K/V are fp8 e4m3 (FNUZ on gfx942, OCP on gfx950) fed straight into +``mfma_f32_16x16x32_fp8_fp8``; Q (bf16/f16) and the softmax probabilities P are +quantized to fp8 too. Scales fold out of the matmuls (q/key scale into the QK +score, value scale + 1/FP8_MAX into the epilogue); softmax max/sum stay f32. +``key_scale``/``value_scale`` are either a ``[1]`` per-tensor scalar or a +``[num_blocks, num_kv_heads, block_size]`` per-token tensor (chosen by rank). + +``block_size`` (16/64) and ``head_dim`` (multiple of 64) are compile-time +constants. Layouts are logical, not production's preshuffle. + +* ``query`` [num_seqs, num_q_heads, head_dim] f16/bf16 (head_dim contiguous) +* ``key_cache`` [num_blocks, num_kv_heads, head_dim//16, block_size, 16] fp8 +* ``value_cache`` [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (trans_v) + or [num_blocks, num_kv_heads, head_dim, block_size] (plain), by rank +* ``block_tables`` [num_seqs, max_blocks_per_seq] int32 +* ``context_lengths`` [num_seqs] int32 +* ``output`` [num_seqs, num_q_heads, head_dim] same dtype as query + +One CTA (4 waves) per (seq, kv_head) runs a flash-style online softmax over +256-token blocks; the 4 waves split tokens for Q·Kᵀ and head-dim for P·V, with +an LDS round-trip on P transposing ownership between the two MMAs. +""" + +import functools + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx +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 +from kernels.common.tensor_shim import _run_compiled +from kernels.common.utils import ( + cdiv, + exp2_amdgcn_scalar, + exp2_f32_fast, + rcp_f32, +) + +MFMA_MNK = 16 # M = N = 16 for the MMA atom; also query rows handled per CTA (padded to 16) +MFMA_K = 32 # fp8 MFMA contracts K = 32 per instruction (mfma_f32_16x16x32_fp8_fp8) +WAVE = 64 +LOG2E = 1.4426950408889634 + + +@functools.lru_cache(maxsize=None) +def compile_pa_decode_tile( + *, + head_dim: int, + query_group_size: int, + block_size: int, + num_partitions: int = 1, + softmax_scale: float | None = None, + query_dtype: str = "f16", + per_token_kv: bool = False, + query_length: int = 1, + trans_v: bool = True, +): + """Build the tile-programming PA-decode kernel + launch wrapper. + + ``block_size``, ``head_dim``, and ``query_dtype`` are compile-time + constants (lru_cache keys). ``query_length`` (MTP) and ``query_group_size`` + flatten into ``TOTAL_ROWS = query_length * query_group_size``, tiled into + ``M_TILES = ceil(TOTAL_ROWS / 16)`` independent 16-row MFMA tiles; each + extra M-tile duplicates loop-carried state, so VGPR/LDS/occupancy scale + roughly linearly with ``M_TILES``. + """ + is_gfx950 = "gfx95" in get_rocm_arch() + FP8 = fx.Float8E4M3FN if is_gfx950 else fx.Float8E4M3FNUZ + FP8_MAX = 448.0 if is_gfx950 else 240.0 # max representable magnitude of the format above + + assert head_dim % MFMA_MNK == 0, f"head_dim {head_dim} must be a multiple of {MFMA_MNK}" + assert block_size in (16, 64), f"pa_decode_tile only supports block_size in (16, 64), got {block_size}" + assert query_dtype in ( + "f16", + "bf16", + ), f"pa_decode_tile only supports query_dtype in ('f16', 'bf16'), got {query_dtype}" + Q_DTYPE = fx.BFloat16 if query_dtype == "bf16" else fx.Float16 + + assert head_dim % 64 == 0, f"pa_decode_tile only supports head_dim that's a multiple of 64, got {head_dim}" + assert query_length >= 1, f"query_length must be >= 1, got {query_length}" + # Flattened query-row axis (MTP outer, GQA head inner), tiled into 16-row M-tiles. + TOTAL_ROWS = query_length * query_group_size + M_TILES = cdiv(TOTAL_ROWS, MFMA_MNK) + ROWS_PADDED = M_TILES * MFMA_MNK + # PV layout: V=A, P=B -> output [head-dim (row), query-row (col=lane16)], + # generalized over head_dim via the VHE_CHUNKS loop. + NWARP = 4 # 4 waves / CTA + TOK_PER_WARP = 64 # tokens each warp owns per compute block (matches production KV_COMPUTE_BLOCK) + TILE_TOK = NWARP * TOK_PER_WARP # 256 tokens / compute block + PAGES_PER_CHUNK = TOK_PER_WARP // block_size # pages spanned by one 64-token warp-chunk: 1 (bs=64) or 4 (bs=16) + assert head_dim % (NWARP * MFMA_MNK) == 0, "head_dim must split across the 4 warps for PV" + + # head_dim splits into 16-element chunks (QK_CHUNK_ELEMS, one dwordx4 load); + # RGROUP_QUARTERS of them make a 64-element fetch group, QKHE_LOOP groups total. + RGROUP_QUARTERS = 4 + QK_CHUNK_ELEMS = 16 + QKHE_LOOP = head_dim // (RGROUP_QUARTERS * QK_CHUNK_ELEMS) + assert QKHE_LOOP >= 1, f"head_dim {head_dim} must be at least {RGROUP_QUARTERS * QK_CHUNK_ELEMS}" + N_SUBCHUNKS = QKHE_LOOP * (QK_CHUNK_ELEMS // 8) + + # Q-quant chunk width: NQCHUNK stays fixed at 16 (tied to `lane16`'s role + # as the absmax butterfly width); QCHUNK scales with head_dim instead. + NQCHUNK = 16 + QCHUNK = head_dim // NQCHUNK # f16 elements per lane's load chunk (8 for head_dim=128, 4 for head_dim=64) + + VHE_CHUNKS = head_dim // (NWARP * MFMA_MNK) # 2 for head_dim=128, 1 for head_dim=64 + + if softmax_scale is None: + softmax_scale = 1.0 / (head_dim**0.5) + NP = int(num_partitions) # context partitions (grid.z); compile-time constant + + BLOCK_THREADS = NWARP * WAVE # 256 + + # ── LDS layout (shared across the 4 warps) ── + # sQ: fp8[ROWS_PADDED,head_dim] staged+quantized query. sP: fp8[16,TILE_TOK] + # quantized probs. sQscale: f32[ROWS_PADDED]. sLmax/sLsum: cross-warp scratch. + # sVPage: V page-index broadcast. sKScale/sVScale/sVScaleMax: per-token K/V + # scale staging. No sO/sM/sL/sCorr: PV output is register-resident/loop-carried + # and stored straight to global (V=A/P=B swap). + f32 = 4 + sQ_bytes = ROWS_PADDED * head_dim * 1 # fp8 + sP_off = sQ_bytes + # +16B row padding breaks a 32-bank LDS conflict on the P-pack writes while + # keeping the row 16B-aligned for PV's ds_read_b128. + SP_ROW_BYTES = TILE_TOK + 16 + sP_bytes = MFMA_MNK * SP_ROW_BYTES # fp8, padded rows + sQscale_off = sP_off + sP_bytes + NWARP_PAD = NWARP + 1 # coprime with 32 banks -> no LDS bank conflict on the cross-warp scratch + # Phase-split slices sLmax per M-tile so all pass-1 writes share one barrier. + sLmax_off = sQscale_off + ROWS_PADDED * f32 + sLsum_off = sLmax_off + M_TILES * MFMA_MNK * NWARP_PAD * f32 + # V page-index broadcast (V's page depends on `rgroup`, shared across warps). + sVPage_off = sLsum_off + MFMA_MNK * NWARP_PAD * f32 + sVPage_bytes = NWARP * PAGES_PER_CHUNK * 4 # i32 + # Per-token KV-scale staging (per_token only), double-buffered so the tt+1 + # prefetch (into the other buffer) doesn't clobber the current tile's scales. + KV_BUFS = 2 if per_token_kv else 1 + KV_BUF_STRIDE = 2 * NWARP * TOK_PER_WARP * f32 # k-region + v-region, one buffer + sKScale_off = sVPage_off + sVPage_bytes + sVScale_off = sKScale_off + NWARP * TOK_PER_WARP * f32 + sKVScale_bytes = KV_BUFS * KV_BUF_STRIDE if per_token_kv else 0 + sVScaleMax_off = sKScale_off + sKVScale_bytes + sVScaleMax_bytes = NWARP_PAD * f32 if per_token_kv else 0 # m-independent: one cross-warp slot + total_bytes = sVScaleMax_off + sVScaleMax_bytes + + # LDS blob: one i32 array carved into typed byte-offset views (see the + # `_lds_*` helpers in the kernel). Every region size above is 4-byte + # aligned, so `total_bytes // 4` covers the blob exactly. + @fx.struct + class SharedStorage: + buf: fx.Array[fx.Int32, total_bytes // 4, 16] + + @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) + def pa_decode_tile_kernel( + output_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] (written directly when NP==1) + # per-partition partial outputs (combined by the reduce kernel when NP>1): + pmax_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row max + psum_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size] row sum + pout_ptr: fx.Tensor, # [num_seqs*query_length, num_kv_heads, num_partitions, query_group_size, head_dim] Q_DTYPE, normalized O_p/l_p + query_ptr: fx.Tensor, # [num_seqs*query_length, num_q_heads, head_dim] -- row = seq*query_length + qi (MTP position) + key_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, head_dim//16, block_size, 16] (blocked, see module docstring) + value_cache_ptr: fx.Tensor, # [num_blocks, num_kv_heads, block_size//16, head_dim, 16] (blocked, see module docstring) + block_tables_ptr: fx.Tensor, # [num_seqs, max_blocks_per_seq] + context_lengths_ptr: fx.Tensor, # [num_seqs] + key_scale_ptr: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token + value_scale_ptr: fx.Tensor, # same shape as key_scale_ptr + max_blocks_per_seq: fx.Int32, + stride_ks_block: fx.Int32, + stride_ks_head: fx.Int32, + stride_q_row: fx.Int32, + stride_q_head: fx.Int32, + ): + tid = fx.Int32(gpu.thread_id("x")) + warp = tid // WAVE # 0..NWARP-1 + lane = tid - warp * WAVE # 0..63 + seq = fx.Int32(gpu.block_id("x")) + kv_h = fx.Int32(gpu.block_id("y")) + part = fx.Int32(gpu.block_id("z")) # context partition handled by this CTA + n_kv = fx.Int32(gpu.grid_dim.y) # num_kv_heads == gridDim.y + + # fx.copy-based flat loaders: K/V/context_len over a raw pointer, Q over + # a buffer resource. + def _make_flat_loader(tensor_ptr, elem_ty, reg_width, copy_op): + use_buffer_resource = isinstance(copy_op, fx.rocdl.CopyOpCDNA3BufferCopyType) + copy_atom = fx.make_copy_atom(copy_op, elem_ty) + reg = fx.make_rmem_tensor(fx.make_layout(reg_width, 1), elem_ty) + base_iter = ( + fx.get_iter(fx.rocdl.make_buffer_tensor(tensor_ptr, max_size=True)) + if use_buffer_resource + else fx.get_iter(tensor_ptr) + ) + flat = fx.Tensor(fx.make_view(base_iter, fx.make_layout(1 << 30, 1))) + div = fx.logical_divide(flat, fx.make_layout(1, 1)) + + def _load(elem_idx): + fx.copy(copy_atom, fx.slice(div, (None, elem_idx)), reg) + return fx.Vector(fx.memref_load_vec(reg)) + + return _load + + _k_load_fp8x16 = _make_flat_loader(key_cache_ptr, FP8, 16, fx.UniversalCopy128b()) + _v_load_fp8x16 = _make_flat_loader(value_cache_ptr, FP8, 16, fx.UniversalCopy128b()) + # Per-lane Q chunk (QCHUNK 16-bit elems) fetched in QLOAD_UNIT-wide + # pieces (128b max per buffer load): head_dim=256 needs 2 pieces. + QLOAD_UNIT = QCHUNK if QCHUNK < 8 else 8 + N_QLOADS = QCHUNK // QLOAD_UNIT + _q_copy_op = fx.rocdl.BufferCopy128b() if QLOAD_UNIT == 8 else fx.rocdl.BufferCopy64b() + _q_load_chunk = _make_flat_loader(query_ptr, Q_DTYPE, QLOAD_UNIT, _q_copy_op) + _ctxlen_load = _make_flat_loader(context_lengths_ptr, fx.Int32, 1, fx.rocdl.BufferCopy32b()) + + def _k_load16(byte_off): + return _k_load_fp8x16(byte_off).bitcast(fx.Int64) + + def _v_load16(byte_off): + return _v_load_fp8x16(byte_off).bitcast(fx.Int64) + + context_len = fx.Int32(_ctxlen_load(seq)[0]) + # Bound block_tables to its real extent: the last (partial) 256-token tile + # can index a page past ceil(context/block_size); the bounded resource + # returns page 0 for that out-of-range read instead of faulting (those + # tail tokens are masked out anyway). + bt_num_records_bytes = fx.Index(gpu.grid_dim.x) * fx.Index(max_blocks_per_seq) * 4 # int32 entries + bt_rsrc = buffer_ops.create_buffer_resource( + block_tables_ptr, max_size=False, num_records_bytes=bt_num_records_bytes + ) + ks_rsrc = buffer_ops.create_buffer_resource(key_scale_ptr, max_size=True) + vs_rsrc = buffer_ops.create_buffer_resource(value_scale_ptr, max_size=True) + # Per-tensor: a single global scale, read once. Per-token: read + # per-token instead (see _kv_scale_ops/_stage_kv_scale_to_lds below). + if const_expr(not per_token_kv): + key_scale = fx.Int32( + buffer_ops.buffer_load(ks_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True) + ).bitcast(fx.Float32) + value_scale = fx.Int32( + buffer_ops.buffer_load(vs_rsrc, arith.constant(0, type=T.i32), vec_width=1, is_scalar=True) + ).bitcast(fx.Float32) + + num_tiles = cdiv(context_len, TILE_TOK) + tiles_per_part = cdiv(num_tiles, NP) + part_start = part * tiles_per_part + part_end_raw = part_start + tiles_per_part + part_end = arith.select(part_end_raw < num_tiles, part_end_raw, num_tiles) + + # One i8 blob carved into typed byte-offset pointers. `lds_base` is an + # ir.Value pointer (safe inside scf control flow); the Python `lds` + # handle is not. + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_base = fx.recast_iter(fx.Uint8, lds.buf.ptr) # byte-addressed base + + def _lds_ptr(byte_off, elem_ty): + # Pin the element-size alignment explicitly (a bare recast off the + # byte base would gcd the alignment down to 1). + p = fx.add_offset(lds_base, fx.make_int_tuple(byte_off)) + ptr_ty = fx.PointerType.get(elem_ty.ir_type, fx.AddressSpace.Shared, dsl_size_of(elem_ty)) + return fx.recast_iter(ptr_ty, p) + + def _lds_load(byte_off, elem_ty, n): + return fx.ptr_load(_lds_ptr(byte_off, elem_ty), result_type=fx.Vector.make_type(n, elem_ty)) + + def _lds_store(byte_off, elem_ty, vec): + fx.ptr_store(vec, _lds_ptr(byte_off, elem_ty)) + + c16 = 16 + rgroup = lane // c16 # 0..3: which quarter-wave (paired with warp -> query row) + lane16 = lane - rgroup * c16 # 0..15: this row's head-dim chunk index + + TOK_CHUNK = NWARP * MFMA_MNK # 64 + NCHUNK = TILE_TOK // TOK_CHUNK # 4 + + def _load_phys_scalar(page, vec_width=1): + result = buffer_ops.buffer_load( + bt_rsrc, seq * max_blocks_per_seq + page, vec_width=vec_width, is_scalar=True + ) + return fx.Int32(result) if vec_width == 1 else result + + def _v_page_fetch_and_stage(tt_i32): + # V's page depends on `rgroup` (shared across warps): warp w fetches + # its rgroup row and broadcasts via LDS (read back by _v_page_read_row). + base_page = tt_i32 * TILE_TOK // block_size # tile start is page-aligned + fetched = _load_phys_scalar(base_page + warp * PAGES_PER_CHUNK, PAGES_PER_CHUNK) + if lane == 0: + fetched_vec = ( + fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) + if const_expr(PAGES_PER_CHUNK == 1) + else fx.Vector(fetched) + ) + _lds_store(sVPage_off + warp * (PAGES_PER_CHUNK * 4), fx.Int32, fetched_vec) + + def _v_page_read_row(): + off = sVPage_off + rgroup * (PAGES_PER_CHUNK * 4) + return _lds_load(off, fx.Int32, PAGES_PER_CHUNK) + + def _kv_buf_off(tt_val): + # ping-pong buffer byte offset (0 when single-buffered). + if const_expr(per_token_kv): + return (tt_val & fx.Int32(1)) * KV_BUF_STRIDE + return 0 + + def _stage_kv_scale_to_lds(phys_vec, buf_off=0): + if const_expr(block_size == 64): + phys = fx.Int32(phys_vec[0]) + base_tok = lane16 * NCHUNK + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + base_tok + k_scale_vec = fx.Vector(buffer_ops.buffer_load(ks_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32)) + v_scale_vec = fx.Vector(buffer_ops.buffer_load(vs_rsrc, scale_idx, vec_width=NCHUNK, dtype=fx.Float32)) + slot = (warp * TOK_PER_WARP + base_tok) * f32 + _lds_store(sKScale_off + buf_off + slot, fx.Float32, k_scale_vec) + _lds_store(sVScale_off + buf_off + slot, fx.Float32, v_scale_vec) + else: + # block_size==16: each lane stages its own rgroup's page/token, + # so the 4 sub-blocks stage in parallel across rgroup-groups. + phys = fx.Int32(phys_vec[rgroup]) + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + lane16 + k_scale_scalar = fx.Float32(buffer_ops.buffer_load(ks_rsrc, scale_idx, vec_width=1, dtype=fx.Float32)) + v_scale_scalar = fx.Float32(buffer_ops.buffer_load(vs_rsrc, scale_idx, vec_width=1, dtype=fx.Float32)) + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + slot = (warp * TOK_PER_WARP + rgroup * c16 + lane16) * f32 + _lds_store( + sKScale_off + buf_off + slot, + fx.Float32, + fx.Vector.from_elements([k_scale_scalar], dtype=fx.Float32), + ) + _lds_store( + sVScale_off + buf_off + slot, + fx.Float32, + fx.Vector.from_elements([v_scale_scalar], dtype=fx.Float32), + ) + + def _load_scale_vec(base_off, a, buf_off=0): + # This lane's 4 per-token scales for chunk `a` from an LDS scale region. + slot = (warp * TOK_CHUNK + a * c16 + rgroup * 4) * f32 + return _lds_load(base_off + buf_off + slot, fx.Float32, 4) + + def _load_kv_scale_vecs(a, buf_off=0): + return _load_scale_vec(sKScale_off, a, buf_off), _load_scale_vec(sVScale_off, a, buf_off) + + # ── raw dwordx4 K load (A operand) ── + # token = warp*TOK_CHUNK + a*c16 + lane16 (the softmax mask and P-pack + # write position below must encode this same formula). + def _k_ops(phys, a): + within_page_tok = (a * c16 + lane16) % block_size + ops = [] + for qkhe in range_constexpr(QKHE_LOOP): + he_idx = qkhe * RGROUP_QUARTERS + rgroup + base = (((phys * n_kv + kv_h) * QCHUNK + he_idx) * block_size + within_page_tok) * QK_CHUNK_ELEMS + w = _k_load16(base) # head[he_idx*16 : +16] -> k_step 2*qkhe, 2*qkhe+1 + if const_expr(block_size == 16): + # help the scheduler overlap the PAGES_PER_CHUNK gathered loads + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + ops.extend([w[0], w[1]]) + return ops # N_SUBCHUNKS i64 operands + + def _k_ops_flat(tt_i32): + base_page = tt_i32 * TILE_TOK // block_size # tile start is page-aligned + fetched = _load_phys_scalar(base_page + warp * PAGES_PER_CHUNK, PAGES_PER_CHUNK) + phys_vec = ( + fx.Vector.from_elements([fx.Int32(fetched)], dtype=fx.Int32) + if const_expr(PAGES_PER_CHUNK == 1) + else fx.Vector(fetched) + ) + flat = [] + for a in range_constexpr(NCHUNK): + phys = fx.Int32(phys_vec[(a * c16) // block_size]) + flat.extend(_k_ops(phys, a)) + if const_expr(head_dim == 64): + fx.rocdl.sched_vmem(len(flat) // 2) + + return fx.Vector.from_elements(flat, dtype=fx.Int64), phys_vec + + # ── prologue: prefetch the first tile's K ── + num_tiles_m1 = num_tiles - 1 + start_safe = arith.select(part_start < num_tiles, part_start, num_tiles_m1) + k_pf0, phys_vec0 = _k_ops_flat(start_safe) + # V page-index prefetch, issued here too for the same overlap; the + # LDS write is visible after the barrier below. + _v_page_fetch_and_stage(start_safe) + if const_expr(per_token_kv): + _stage_kv_scale_to_lds(phys_vec0, _kv_buf_off(fx.Int32(start_safe))) + + # per_token_kv has no single global key_scale/value_scale: scale_qk + # drops the key_scale factor (folded in per-token, see masked_chunks + # below) and v_scale_f is unused (replaced by v_max_scaled). + if const_expr(per_token_kv): + scale_qk = fx.Float32(softmax_scale * LOG2E) + else: + scale_qk = fx.Float32(softmax_scale * LOG2E) * fx.Float32(key_scale) + 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.) + fm_nnan = arith.FastMathFlags.nnan + + def _row_off(byte_off, m_idx, width, elem_ty): + return byte_off + m_idx * (width * dsl_size_of(elem_ty)) + + def _ld1(byte_off, m_idx): + return _lds_load(_row_off(byte_off, m_idx, 1, fx.Float32), fx.Float32, 1)[0] + + def _st1(byte_off, m_idx, val): + _lds_store( + _row_off(byte_off, m_idx, 1, fx.Float32), fx.Float32, fx.Vector.from_elements([val], dtype=fx.Float32) + ) + + # f32[16, NWARP] cross-warp scratch: scalar write at (row, warp), vec read of a row's NWARP valid slots. + def _st_lw(base_off, row, w, val): + off = base_off + (row * NWARP_PAD + w) * 4 + _lds_store(off, fx.Float32, fx.Vector.from_elements([val], dtype=fx.Float32)) + + def _ld_lw_row(base_off, row): + off = base_off + row * (NWARP_PAD * 4) + return _lds_load(off, fx.Float32, NWARP) + + def _f32_to_fp8_words(vf32): + # f32 -> fp8 must use the HW cvt (arith.truncf to fp8 doesn't lower); + # pack 4 f32 -> 1 i32 (4 fp8) via two cvt_pk_fp8_f32 calls. + n = vf32.shape[0] + words = [] + for i in range_constexpr(n // 4): + b = i * 4 + lo = fx.rocdl.cvt_pk_fp8_f32(T.i32, vf32[b], vf32[b + 1], 0, False) + words.append(fx.rocdl.cvt_pk_fp8_f32(T.i32, vf32[b + 2], vf32[b + 3], lo, True)) + return fx.Vector.from_elements(words, dtype=fx.Int32) + + def _st_words(byte_off, words): + _lds_store(byte_off, fx.Int32, words) + + qh_local = warp * 4 + rgroup # 0..15: this thread's query row within an M-tile + + # Each M-tile quantizes 16 rows of the flattened (MTP position, GQA + # head) axis: `flat_idx = m*16 + qh_local`, decomposed as + # `qi = flat_idx // query_group_size`, `gs_head = flat_idx % + # query_group_size` (same convention as `_mtp_groups`). No + # cross-M-tile dependency, so no barriers needed between iterations. + def _quant_q_row(m, qi, gs_head, q_row_off): + qh0 = kv_h * query_group_size + gs_head + row_byte0 = ((seq * query_length + qi) * stride_q_row + qh0 * stride_q_head) * 2 # 16-bit float = 2B/elem + base_elem = (row_byte0 + lane16 * (QCHUNK * 2)) // 2 # byte offset -> element index + # QCHUNK elements, loaded as N_QLOADS contiguous QLOAD_UNIT pieces + # (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) + for u in range_constexpr(1, N_QLOADS): + absmax = fx.maxnumf(absmax, fmath.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)) + + q_scale = absmax * fx.Float32(1.0 / FP8_MAX) + inv = fx.Float32(rcp_f32(fx.maxnumf(q_scale, fx.Float32(1e-20)))) + inv_b = fx.Vector.from_elements([inv], dtype=fx.Float32).broadcast_to(QLOAD_UNIT) + + for u in range_constexpr(N_QLOADS): + q_scaled_unit = q_units[u].to(fx.Float32) * inv_b + _st_words( + q_row_off + qh_local * head_dim + lane16 * QCHUNK + u * QLOAD_UNIT, + _f32_to_fp8_words(q_scaled_unit), + ) + if lane16 == 0: + # Transposed [qh][m] (not [m][qh]) so the whole M_TILES-wide + # row for a fixed qh is contiguous, letting the KV-loop read + # it back in one wide load instead of M_TILES separate + # narrow ones -- see the read site below. + _st1(sQscale_off, qh_local * M_TILES + m, q_scale) + + for m in range_constexpr(M_TILES): + flat_idx = m * MFMA_MNK + qh_local + qi = flat_idx // query_group_size + gs_head = flat_idx - qi * query_group_size + q_row_off = m * MFMA_MNK * head_dim + # `flat_idx < TOTAL_ROWS` is statically true for every lane except + # possibly on the last M-tile (only it can be a partial tile), so + # skip the runtime branch (and its EXEC-mask overhead) entirely + # for every other M-tile. + if const_expr((m + 1) * MFMA_MNK <= TOTAL_ROWS): + _quant_q_row(m, qi, gs_head, q_row_off) + elif flat_idx < TOTAL_ROWS: + _quant_q_row(m, qi, gs_head, q_row_off) + else: + _st_words( + q_row_off + qh_local * head_dim + lane16 * QCHUNK, + fx.Vector.filled(QCHUNK // 4, 0, fx.Int32), + ) + if lane16 == 0: + _st1(sQscale_off, qh_local * M_TILES + m, ZERO_F) + + gpu.barrier() + + # First tile's V page-index row, now visible after the barrier above + # (the fetch+LDS-store was issued earlier, alongside k_pf0). + v_page_pf0 = _v_page_read_row() + + # Q is the B operand, read raw from sQ once per M-tile and held in + # registers. MUST use the exact same (qkhe, rgroup, qkr) -> head_dim + # permutation as K's `_k_ops`. + q_ops_all = [] + for m in range_constexpr(M_TILES): + q_row_off = m * MFMA_MNK * head_dim + for qkhe in range_constexpr(QKHE_LOOP): + he_idx = qkhe * RGROUP_QUARTERS + rgroup + chunk = _lds_load(q_row_off + lane16 * head_dim + he_idx * QK_CHUNK_ELEMS, fx.Int64, 2) + q_ops_all.extend([chunk[0], chunk[1]]) + # q_ops_all[m*N_SUBCHUNKS+s] for s=0..N_SUBCHUNKS-1, s = qkhe*2+qkr, + # = M-tile m's head[he_idx*16+qkr*8 : +8] of qhead=lane16 + + # QK in NCHUNK chunks of 4 tokens: each chunk yields a f32x4 + # C-fragment, so softmax processes 4 scores at a time (low VGPR peak). + _ct = [ + fx.Vector.from_elements([float(a * c16 + r) for r in range_constexpr(4)]) for a in range_constexpr(NCHUNK) + ] + # P·V is loop-tiled over head-dim (like production's VHELOOP): each + # step computes O[:, vh*VHE_SIZE:+VHE_SIZE] instead of materializing + # the full [16, head_dim] at once. + VHE_SIZE = head_dim // VHE_CHUNKS + OP_ELEMS = MFMA_MNK * VHE_SIZE // (NWARP * WAVE) # PV C-fragment elements/lane/chunk (probed = 4) + + # ── raw dwordx4 V load (B operand) ── + # One dwordx4 load per (16-token sub-block, head_elem); trans_v only + # changes the offset formula. `sub`/`step` walk pages/16-token sub-blocks. + NVOPS = TILE_TOK // MFMA_K # 8 PV k_steps (256 tokens / K=32) + STEPS_PER_PAGE = block_size // 16 + + def _v_ops(phys_row, vh): + head_group = ((vh * VHE_SIZE) // 16) + warp + head_element = head_group * 16 + lane16 + ops = [] + for sub in range_constexpr(PAGES_PER_CHUNK): + for step in range_constexpr(STEPS_PER_PAGE): + if const_expr(trans_v): + base = (((phys_row[sub] * n_kv + kv_h) * STEPS_PER_PAGE + step) * head_dim + head_element) * 16 + else: + base = ((phys_row[sub] * n_kv + kv_h) * head_dim + head_element) * block_size + step * 16 + w = _v_load16(base) + if const_expr(block_size == 16): + # help the scheduler overlap the per-page gathered loads (see _k_ops) + fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) + ops.extend([w[0], w[1]]) + if const_expr(head_dim == 64): + fx.rocdl.sched_vmem(len(ops) // 2) + return ops # NVOPS i64, the 64-token contiguous run for this head + + # query_length==1: plain wave-uniform `context_len` (folding in the + # per-lane `lane16` would cost a live VGPR for no behavioral change). + if const_expr(query_length == 1): + causal_bound = [context_len for _m in range_constexpr(M_TILES)] + else: + causal_bound = [ + context_len - (query_length - 1) + (m * MFMA_MNK + lane16) // query_group_size + for m in range_constexpr(M_TILES) + ] + + K_SLOT, V_SLOT = 0, 1 + # Per-M-tile loop-carried state after the K/V slots: one output chunk + # per VHE_CHUNKS (2 for head_dim=128, 4 for head_dim=256, 1 for 64), + # then running-max + running-denom. + STATE_PER_M = VHE_CHUNKS + 2 + + def _o_slot(m, vh): + return 2 + STATE_PER_M * m + vh + + def _m_slot(m): + return 2 + STATE_PER_M * m + VHE_CHUNKS + + def _l_slot(m): + return 2 + STATE_PER_M * m + VHE_CHUNKS + 1 + + o_zero = fx.Vector.filled(OP_ELEMS, 0.0, fx.Float32) + init_state = [k_pf0, v_page_pf0] + for _m in range_constexpr(M_TILES): + init_state.extend([o_zero] * VHE_CHUNKS + [NEG_INF, ZERO_F]) + for tt, ostate in range(part_start, part_end, 1, init=init_state): + k_cur = ostate[K_SLOT] # this tile's prefetched K, as one (NCHUNK*N_SUBCHUNKS,) i64 vector + v_page_cur = ostate[V_SLOT] # this tile's V pages, as one PAGES_PER_CHUNK-wide i32 vector + tt = fx.Int32(tt) + tok0 = tt * TILE_TOK + # per_tensor phase-split: let IGLP interleave MFMA with softmax + # VALU/LDS to hide the MFMA-hazard s_nop (per_token is VGPR-cliffed). + if const_expr(not per_token_kv and M_TILES > 1): + fx.rocdl.iglp_opt(0) + + tt1 = tt + 1 + + next_state = [None, None] # slots 0/1 (K_SLOT/V_SLOT) filled in at m==0 below + + # This tile's ping-pong scale buffer; the tt+1 prefetch stages the other. + cur_kv_buf = _kv_buf_off(tt) + + # V data doesn't depend on `m`; hoist once for the phase-split + # instead of reloading per M-tile. + v_vh_shared = None + if const_expr(M_TILES > 1): + v_vh_shared = [_v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS)] + + # q_scale doesn't depend on `m`; read the whole M_TILES-wide row once + # (contiguous via the transposed [qh][m] sQscale layout). + q_scale_vec = None + if const_expr(M_TILES > 1): + q_scale_vec = _lds_load(sQscale_off + lane16 * (M_TILES * f32), fx.Float32, M_TILES) + + def _lmax_off_m(m): + return sLmax_off + (m * MFMA_MNK * NWARP_PAD * f32 if const_expr(M_TILES > 1) else 0) + + # Phase-split (M_TILES>1): pass-1 (QK+mask+max) loops all M-tiles into + # per-M-tile LDS slices so they share ONE barrier. M_TILES==1: else below. + if const_expr(M_TILES > 1): + masked_chunks_saved = [None] * M_TILES + scale_saved = [None] * M_TILES + + # per_token: compute the m-independent pv_max once, hold only + # k_scale across Phase A, re-read v_scale for Phase B after the + # barrier (peak scale liveness 16, not 32). + k_scale_shared = None + if const_expr(per_token_kv): + v_scale_A = [_load_scale_vec(sVScale_off, a, cur_kv_buf) for a in range_constexpr(NCHUNK)] + pv_max = fx.Float32(0.0) + for a in range_constexpr(NCHUNK): + pv_max = fx.maxnumf(pv_max, v_scale_A[a].reduce(ReductionOp.MAX)) + for sh in (16, 32): + pv_max = fx.maxnumf(pv_max, pv_max.shuffle_xor(sh, WAVE)) + _st_lw(sVScaleMax_off, 0, warp, pv_max) + k_scale_shared = [_load_scale_vec(sKScale_off, a, cur_kv_buf) for a in range_constexpr(NCHUNK)] + + for m in range_constexpr(M_TILES): + frag_Ss = [] + for a in range_constexpr(NCHUNK): + acc = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(N_SUBCHUNKS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, [k_cur[a * N_SUBCHUNKS + s], q_ops_all[m * N_SUBCHUNKS + s], acc, 0, 0, 0] + ) + frag_Ss.append(fx.Vector(acc)) + + scale = scale_qk * fx.Float32(q_scale_vec[m]) + n_valid_tile = (causal_bound[m] - tok0).to(fx.Float32) + base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) + thr = fx.Vector.from_elements([n_valid_tile - base_tok_f], dtype=fx.Float32).broadcast_to(4) + neg4 = fx.Vector.filled(4, float("-inf"), fx.Float32) + + if const_expr(per_token_kv): + scaled_frags = [frag_Ss[a] * k_scale_shared[a] for a in range_constexpr(NCHUNK)] + else: + scaled_frags = frag_Ss + + masked_chunks = [(_ct[a] < thr).select(scaled_frags[a], neg4) for a in range_constexpr(NCHUNK)] + + pm = fx.Float32(float("-inf")) + for a in range_constexpr(NCHUNK): + pm = fx.maxnumf( + pm, masked_chunks[a].reduce(ReductionOp.MAX, fastmath=fm_nnan), fastmath=fm_nnan + ) + for sh in (16, 32): + pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE), fastmath=fm_nnan) + _st_lw(_lmax_off_m(m), lane16, warp, pm * scale) + + masked_chunks_saved[m] = masked_chunks + scale_saved[m] = scale + + # tt+1 K/V/scale prefetch, issued once here so the V-page read + # reuses this barrier. + k_next = k_cur + if tt1 < part_end: + k_next, phys_vec1 = _k_ops_flat(tt1) + _v_page_fetch_and_stage(tt1) + if const_expr(per_token_kv): + # stage tt+1 into the OTHER buffer (current v_scale re-read below) + _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) + next_state[K_SLOT] = k_next + + gpu.barrier() + + v_page_next = v_page_cur + if tt1 < part_end: + v_page_next = _v_page_read_row() + next_state[V_SLOT] = v_page_next + + # Phase B v_scale: M_TILES>=4 re-reads per chunk (VGPR cliff); + # smaller M_TILES hold all NCHUNK. Both read the current buffer. + v_scale_shared = None + if const_expr(per_token_kv and M_TILES < 4): + v_scale_shared = [_load_scale_vec(sVScale_off, a, cur_kv_buf) for a in range_constexpr(NCHUNK)] + + for m in range_constexpr(M_TILES): + o_acc = [ostate[_o_slot(m, vh)] for vh in range_constexpr(VHE_CHUNKS)] + m_prev = ostate[_m_slot(m)] # this thread's own running max, carried from last tile + l_prev = ostate[_l_slot(m)] # this thread's own running denom, carried from last tile + + masked_chunks = masked_chunks_saved[m] + scale = scale_saved[m] + + v_max_scaled = None + norm_factor_b = None + if const_expr(per_token_kv): + v_max_global = _ld_lw_row(sVScaleMax_off, 0).reduce(ReductionOp.MAX) + v_max_scaled = v_max_global * fx.Float32(1.0 / FP8_MAX) + v_max_safe = v_max_scaled + fx.Float32(1e-8 / FP8_MAX) + norm_factor = fx.Float32(rcp_f32(v_max_safe)) + norm_factor_b = fx.Vector.from_elements([norm_factor], dtype=fx.Float32).broadcast_to(4) + + m_new = fx.maxnumf( + m_prev, + _ld_lw_row(_lmax_off_m(m), lane16).reduce(ReductionOp.MAX, fastmath=fm_nnan), + fastmath=fm_nnan, + ) + # Fully-invalid row: use 0 as the effective max so masked lanes + # give exp2(-inf-0)==0 (avoids the -inf-(-inf) cancellation). + safe_max = arith.select(m_new > NEG_INF, m_new, ZERO_F) + m_new_b = fx.Vector.from_elements([safe_max], dtype=fx.Float32).broadcast_to(4) + ls = fx.Float32(0.0) + words = [] + for a in range_constexpr(NCHUNK): + Pa = fx.Vector(exp2_f32_fast(masked_chunks[a] * scale - m_new_b)) + ls = ls + Pa.reduce(ReductionOp.ADD) + if const_expr(per_token_kv): + v_sc = ( + _load_scale_vec(sVScale_off, a, cur_kv_buf) + if const_expr(M_TILES >= 4) + else v_scale_shared[a] + ) + p_scaled = Pa * v_sc * norm_factor_b + else: + p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) + words.append(_f32_to_fp8_words(p_scaled)[0]) + + p_off0 = sP_off + lane16 * SP_ROW_BYTES + warp * TOK_CHUNK + rgroup * 4 + # The NCHUNK P words scatter across the row at stride c16//4 + # i32 (the token->fp8-lane interleave the PV ds_read_b128 + # expects); one strided store per chunk. + for a in range_constexpr(NCHUNK): + _lds_store( + 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) + # 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) + corr_reg = fx.Float32(exp2_amdgcn_scalar(safe_prev - safe_max)) + 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) + + 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) + for vh in range_constexpr(VHE_CHUNKS): + v_vh = v_vh_shared[vh] + acc = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(NVOPS): + # SWAPPED operands (V=A, P=B): output row = + # head-dim, output col = query-row=lane16. + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8(T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0]) + op = fx.Vector(acc) + if const_expr(per_token_kv): + op = op * fx.Vector.from_elements([v_max_scaled], dtype=fx.Float32).broadcast_to(OP_ELEMS) + o_acc[vh] = o_acc[vh] * corr_b + op + next_state.extend([*o_acc, m_new, l_new]) + # sP and sLsum are reused by the next M-tile. Synchronize + # all waves after their LDS reads before any wave overwrites + # those regions, then retire this tile's dependency chain. + if const_expr(m < M_TILES - 1): + gpu.barrier() + fx.rocdl.sched_barrier(0) + else: + # M_TILES==1 single tile (m==0 for the _o_slot/_m_slot/_l_slot helpers). + o_acc = [ostate[_o_slot(0, vh)] for vh in range_constexpr(VHE_CHUNKS)] + m_prev = ostate[_m_slot(0)] # running max, carried from last tile + l_prev = ostate[_l_slot(0)] # running denom, carried from last tile + # QK: each NCHUNK chunk accumulates N_SUBCHUNKS k_steps into an f32x4. + frag_Ss = [] + for a in range_constexpr(NCHUNK): + acc = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(N_SUBCHUNKS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, [k_cur[a * N_SUBCHUNKS + s], q_ops_all[s], acc, 0, 0, 0] + ) + frag_Ss.append(fx.Vector(acc)) + # tt+1 K/V/scale prefetch, issued here to reuse the pass-1 barrier. + k_next = k_cur + if tt1 < part_end: + k_next, phys_vec1 = _k_ops_flat(tt1) + _v_page_fetch_and_stage(tt1) + if const_expr(per_token_kv): + _stage_kv_scale_to_lds(phys_vec1, _kv_buf_off(tt1)) + next_state[K_SLOT] = k_next + # Softmax: each lane owns one qhead (lane%16); register reduce + shuffle_xor. + scale = scale_qk * _ld1(sQscale_off, lane16) # per-qhead positive score scale + n_valid_tile = (causal_bound[0] - tok0).to(fx.Float32) + base_tok_f = fx.Int32(warp * TOK_CHUNK + rgroup * 4).to(fx.Float32) + thr = fx.Vector.from_elements([n_valid_tile - base_tok_f], dtype=fx.Float32).broadcast_to(4) + neg4 = fx.Vector.filled(4, -1e30, fx.Float32) + # per_token_kv: K-scale varies per token, so fold it in BEFORE the max-reduce. + v_scale_vecs = None + if const_expr(per_token_kv): + v_scale_vecs = [] + scaled_frags = [] + for a in range_constexpr(NCHUNK): + k_scale_vec, v_scale_vec = _load_kv_scale_vecs(a, cur_kv_buf) + v_scale_vecs.append(v_scale_vec) + scaled_frags.append(frag_Ss[a] * k_scale_vec) + else: + scaled_frags = frag_Ss + # Reused in pass 2 below, halving the mask instruction count. + masked_chunks = [(_ct[a] < thr).select(scaled_frags[a], neg4) for a in range_constexpr(NCHUNK)] + # pass 1: per-warp max for this qhead + pm = fx.Float32(float("-inf")) + for a in range_constexpr(NCHUNK): + pm = fx.maxnumf(pm, masked_chunks[a].reduce(ReductionOp.MAX)) + for sh in (16, 32): + pm = fx.maxnumf(pm, pm.shuffle_xor(sh, WAVE)) + _st_lw(sLmax_off, lane16, warp, pm * scale) # redundant across the 4 lanes sharing this qhead + # per_token_kv: max V-scale for the per-tile fp8 normalization + # (any positive value is correct, so skip the causal mask). + if const_expr(per_token_kv): + pv_max = fx.Float32(0.0) + for a in range_constexpr(NCHUNK): + pv_max = fx.maxnumf(pv_max, v_scale_vecs[a].reduce(ReductionOp.MAX)) + for sh in (16, 32): + pv_max = fx.maxnumf(pv_max, pv_max.shuffle_xor(sh, WAVE)) + _st_lw(sVScaleMax_off, 0, warp, pv_max) + gpu.barrier() + # V page-index row for next tile, now visible after the barrier. + v_page_next = v_page_cur + if tt1 < part_end: + v_page_next = _v_page_read_row() + next_state[V_SLOT] = v_page_next + # per_token_kv: combine warps' max V-scale into the tile's + # normalization factor (also the PV correction below). + v_max_scaled = None + norm_factor_b = None + if const_expr(per_token_kv): + v_max_global = _ld_lw_row(sVScaleMax_off, 0).reduce(ReductionOp.MAX) + v_max_scaled = v_max_global * fx.Float32(1.0 / FP8_MAX) + v_max_safe = v_max_scaled + fx.Float32(1e-8 / FP8_MAX) + norm_factor = fx.Float32(rcp_f32(v_max_safe)) + norm_factor_b = fx.Vector.from_elements([norm_factor], dtype=fx.Float32).broadcast_to(4) + # pass 2: global max over warps -> exp -> fp8 P pack (-> sP) -> sum + m_new = fx.maxnumf(m_prev, _ld_lw_row(sLmax_off, lane16).reduce(ReductionOp.MAX)) + m_new_b = fx.Vector.from_elements([m_new], dtype=fx.Float32).broadcast_to(4) + ls = fx.Float32(0.0) + words = [] + zero4_p = fx.Vector.filled(4, 0.0, fx.Float32) + for a in range_constexpr(NCHUNK): + # re-mask Pa so a fully-masked chunk contributes exactly 0 + valid_a = masked_chunks[a] > fx.Vector.filled(4, -1e29, fx.Float32) + Pa = valid_a.select(fx.Vector(exp2_f32_fast(masked_chunks[a] * scale - m_new_b)), zero4_p) + ls = ls + Pa.reduce(ReductionOp.ADD) + if const_expr(per_token_kv): + v_scale_this = ( + _load_scale_vec(sVScale_off, a, cur_kv_buf) + if const_expr(head_dim == 64) + else v_scale_vecs[a] + ) + p_scaled = Pa * v_scale_this * norm_factor_b + else: + p_scaled = Pa * fx.Vector.filled(4, FP8_MAX, fx.Float32) + words.append(_f32_to_fp8_words(p_scaled)[0]) + p_off0 = sP_off + lane16 * SP_ROW_BYTES + warp * TOK_CHUNK + rgroup * 4 + # NCHUNK P words scatter at stride c16//4 i32 (see phase-split). + for a in range_constexpr(NCHUNK): + _lds_store( + p_off0 + a * (c16 // 4) * f32, fx.Int32, fx.Vector.from_elements([words[a]], dtype=fx.Int32) + ) + 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) + # 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 + ) + 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 + # to hide the latency behind). + v_vh_batch = [_v_ops(v_page_cur, vh) for vh in range_constexpr(VHE_CHUNKS)] + for vh in range_constexpr(VHE_CHUNKS): + v_vh = v_vh_batch[vh] + acc = arith.constant_vector(0.0, T.f32x4) + for s in range_constexpr(NVOPS): + acc = fx.rocdl.mfma_f32_16x16x32_fp8_fp8(T.f32x4, [v_vh[s], p_ops[s], acc, 0, 0, 0]) + op = fx.Vector(acc) + if const_expr(per_token_kv): + op = op * fx.Vector.from_elements([v_max_scaled], dtype=fx.Float32).broadcast_to(OP_ELEMS) + o_acc[vh] = o_acc[vh] * corr_b + op + next_state.extend([*o_acc, m_new, l_new]) + results = yield next_state + o_final = results + + # Direct-store epilogue: after the PV swap each lane holds its 4 head-dim + # values for one query-row and writes them straight to global. + inv_fp8 = fx.Float32(1.0 / FP8_MAX) + for m in range_constexpr(M_TILES): + row = m * MFMA_MNK + lane16 # flat (mtp, gqa) query-row for this lane + l_row = o_final[_l_slot(m)] + safe_l = arith.select(l_row > ZERO_F, l_row, fx.Float32(1.0)) + inv_l = fx.Float32(rcp_f32(safe_l)) + 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_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 + qh = kv_h * query_group_size + gs_head_e + + def _emit(o_norm, sub): + if const_expr(NP == 1): + out_row = output_ptr[seq * query_length + qi_e, qh, None] + out_chunk = fx.slice(fx.logical_divide(out_row, fx.make_layout(OP_ELEMS, 1)), (None, sub)) + out_chunk.store(o_norm) + else: + base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row + pout_div = fx.logical_divide(pout_ptr, fx.make_layout(OP_ELEMS, 1)) + pout_chunk = fx.slice(pout_div, (None, base * (head_dim // OP_ELEMS) + sub)) + pout_chunk.store(o_norm) + + for vh in range_constexpr(VHE_CHUNKS): + o_slot = _o_slot(m, vh) + o_norm = (o_final[o_slot] * o_scale_b).to(Q_DTYPE) + head_base = vh * (NWARP * MFMA_MNK) + warp * MFMA_MNK + rgroup * OP_ELEMS + sub = head_base // OP_ELEMS + # Guard the partial last tile's out-of-range rows (folded away for full tiles). + if row < TOTAL_ROWS: + _emit(o_norm, sub) + + if const_expr(NP > 1): + if warp == 0 and rgroup == 0: + base = ((seq * n_kv + kv_h) * NP + part) * TOTAL_ROWS + row + if row < TOTAL_ROWS: + # Convert the running max from log2 units (scale_qk folds + # in LOG2E) to natural-log units: the shared reduce + # re-applies LOG2E itself when combining partitions. + pmax_ptr[base] = o_final[_m_slot(m)] * fx.Float32(1.0 / LOG2E) + psum_ptr[base] = l_row + + @flyc.jit + def pa_decode_tile_launch( + output: fx.Tensor, + pmax: fx.Tensor, + psum: fx.Tensor, + pout: fx.Tensor, + query: fx.Tensor, + key_cache: fx.Tensor, + value_cache: fx.Tensor, + block_tables: fx.Tensor, + context_lengths: fx.Tensor, + key_scale: fx.Tensor, # [1] per-tensor OR [num_blocks, num_kv_heads, block_size] per-token + value_scale: fx.Tensor, # same shape as key_scale + max_blocks_per_seq: fx.Int32, + num_seqs: fx.Int32, + num_kv_heads: fx.Int32, + stride_ks_block: fx.Int32, + stride_ks_head: fx.Int32, + stride_q_row: fx.Int32, + 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) + + return {"launch": pa_decode_tile_launch, "kernel": pa_decode_tile_kernel} + + +def pa_decode_tile( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + key_scale: float | torch.Tensor, + value_scale: float | torch.Tensor, + softmax_scale: float | None = None, + stream=None, + *, + num_partitions: int | None = None, + pmax: torch.Tensor | None = None, + psum: torch.Tensor | None = None, + pout: torch.Tensor | None = None, +) -> None: + """Host entry point. See module docstring for tensor layouts. + + ``num_partitions``/``pmax``/``psum``/``pout`` are optional caller overrides + (e.g. for CUDA-graph capture, where nothing may be allocated on-the-fly); + when omitted they are picked/allocated here. + """ + num_seqs = context_lengths.shape[0] + total_q_rows, num_q_heads, head_dim = query.shape + assert ( + total_q_rows % num_seqs == 0 + ), f"query.shape[0] ({total_q_rows}) must be a multiple of context_lengths.shape[0] ({num_seqs})" + query_length = total_q_rows // num_seqs + _, num_kv_heads, num_hgroups, block_size, hgroup_width = key_cache.shape + + assert num_hgroups == head_dim // 16 and hgroup_width == 16 + assert block_size in (16, 64), f"pa_decode_tile only supports block_size in (16, 64), got {block_size}" + + trans_v = value_cache.dim() == 5 + if trans_v: + _, v_num_kv_heads, v_subblocks, v_head_dim, v_width = value_cache.shape + assert ( + v_head_dim == head_dim and v_width == 16 and v_subblocks == block_size // 16 + ), f"value_cache shape {tuple(value_cache.shape)} doesn't match block_size={block_size}, head_dim={head_dim}" + else: + _, v_num_kv_heads, v_head_dim, v_block_size = value_cache.shape + assert v_head_dim == head_dim and v_block_size == block_size, ( + f"value_cache shape {tuple(value_cache.shape)} doesn't match " + f"block_size={block_size}, head_dim={head_dim}" + ) + assert v_num_kv_heads == num_kv_heads + assert block_tables.dtype == torch.int32, f"block_tables must be int32, got {block_tables.dtype}" + assert context_lengths.dtype == torch.int32, f"context_lengths must be int32, got {context_lengths.dtype}" + query_group_size = num_q_heads // num_kv_heads + max_blocks_per_seq = block_tables.shape[1] + if query.dtype == torch.bfloat16: + query_dtype = "bf16" + elif query.dtype == torch.float16: + query_dtype = "f16" + else: + raise ValueError(f"pa_decode_tile only supports f16/bf16 query, got {query.dtype}") + assert ( + output.dtype == query.dtype + ), f"pa_decode_tile requires output.dtype == query.dtype, got {output.dtype} vs {query.dtype}" + + assert query.stride(2) == 1, f"pa_decode_tile requires a contiguous head_dim axis, got strides {query.stride()}" + + dev = query.device + key_scale_t = key_scale if isinstance(key_scale, torch.Tensor) else torch.tensor([float(key_scale)], device=dev) + value_scale_t = ( + value_scale if isinstance(value_scale, torch.Tensor) else torch.tensor([float(value_scale)], device=dev) + ) + per_token_kv = key_scale_t.dim() > 1 + if per_token_kv: + assert value_scale_t.dim() > 1, "value_scale must also be per-token (dim>1) when key_scale is per-token" + assert ( + key_scale_t.shape == value_scale_t.shape + ), f"key_scale/value_scale shape mismatch: {tuple(key_scale_t.shape)} vs {tuple(value_scale_t.shape)}" + assert key_scale_t.shape == (key_cache.shape[0], num_kv_heads, block_size), ( + "per-token key_scale/value_scale must be [num_blocks, num_kv_heads, block_size] " + f"matching the KV cache, got {tuple(key_scale_t.shape)}" + ) + stride_ks_block = int(key_scale_t.stride(0)) + stride_ks_head = int(key_scale_t.stride(1)) + else: + stride_ks_block = 0 + stride_ks_head = 0 + assert ( + key_scale_t.dtype == torch.float32 and key_scale_t.device == dev + ), f"key_scale tensor must be float32 on {dev}, got {key_scale_t.dtype} on {key_scale_t.device}" + assert ( + value_scale_t.dtype == torch.float32 and value_scale_t.device == dev + ), f"value_scale tensor must be float32 on {dev}, got {value_scale_t.dtype} on {value_scale_t.device}" + + if not num_partitions: + from kernels.attention.pa_decode_fp8 import KV_COMPUTE_BLOCK, get_recommended_splits + + blocks_per_partition = KV_COMPUTE_BLOCK // block_size + num_partitions = get_recommended_splits( + num_seqs, + num_kv_heads, + split_kv_blocks=blocks_per_partition, + ) + + compiled = compile_pa_decode_tile( + head_dim=head_dim, + query_group_size=query_group_size, + block_size=int(block_size), + num_partitions=num_partitions, + softmax_scale=softmax_scale, + query_dtype=query_dtype, + per_token_kv=per_token_kv, + query_length=query_length, + trans_v=trans_v, + ) + from kernels.attention.pa_decode_fp8 import _is_current_stream_capturing + + is_graph_capturing = _is_current_stream_capturing() + if num_partitions == 1: + # NP==1 writes output directly; partials unused (caller buffers ignored). + if pmax is None: + if is_graph_capturing: + raise ValueError( + "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " + "even when num_partitions==1 (nothing may be allocated mid-capture)." + ) + pmax = psum = pout = torch.empty(1, dtype=torch.float32, device=dev) + else: + total_rows = query_length * query_group_size + expected_scalar_shape = (num_seqs, num_kv_heads, num_partitions, total_rows) + if pmax is None or psum is None or pout is None: + if is_graph_capturing: + raise ValueError( + "CUDA graph capture requires preallocated `pmax`/`psum`/`pout` " + "for num_partitions>1 (nothing may be allocated mid-capture)." + ) + pmax = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) + psum = torch.empty(*expected_scalar_shape, dtype=torch.float32, device=dev) + pout = torch.empty(*expected_scalar_shape, head_dim, dtype=output.dtype, device=dev) + else: + assert pmax.shape == expected_scalar_shape, f"pmax shape {tuple(pmax.shape)} != {expected_scalar_shape}" + assert psum.shape == expected_scalar_shape, f"psum shape {tuple(psum.shape)} != {expected_scalar_shape}" + assert pout.shape == ( + *expected_scalar_shape, + head_dim, + ), f"pout shape {tuple(pout.shape)} != {(*expected_scalar_shape, head_dim)}" + s = stream or torch.cuda.current_stream() + + _run_compiled( + compiled["launch"], + output, + pmax.view(-1), + psum.view(-1), + pout.view(-1), + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale_t, + value_scale_t, + int(max_blocks_per_seq), + int(num_seqs), + int(num_kv_heads), + stride_ks_block, + stride_ks_head, + int(query.stride(0)), + int(query.stride(1)), + s, + ) + if num_partitions > 1: + from kernels.attention.pa_decode_fp8 import _get_output_dtype_str + from kernels.attention.pa_decode_swa import compile_pa_decode_sw_reduce + + reduce_compiled = compile_pa_decode_sw_reduce( + max_context_partition_num=num_partitions, + query_seq_len=query_length, + query_group_size=query_group_size, + head_size=head_dim, + output_dtype_str=_get_output_dtype_str(output), + logits_dtype_str=_get_output_dtype_str(pout), + ) + _run_compiled( + reduce_compiled["launch"], + output.data_ptr(), + psum.data_ptr(), # exp_sums + pmax.data_ptr(), # max_logits + pout.data_ptr(), # logits (already-normalized query_dtype partials) + query_length * output.stride(0), # stride_output_bs: per TRUE seq, spans all query_length rows + output.stride(0), # stride_output_len: per MTP position (query_length==1: unused, multiplied by 0) + query_group_size * output.stride(1), + output.stride(1), + pmax.stride(0), + pmax.stride(1), + pmax.stride(2), + pout.stride(0), + pout.stride(1), + pout.stride(2), + pout.stride(3), + int(num_seqs), + int(num_kv_heads), + s, + ) diff --git a/kernels/attention/pa_metadata.py b/kernels/attention/pa_metadata.py index 3fbb9d20..7a4b7616 100644 --- a/kernels/attention/pa_metadata.py +++ b/kernels/attention/pa_metadata.py @@ -726,12 +726,9 @@ def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs, v_scale_vecs): v_correction = v_scale_val for td in range_constexpr(TLOOP): - p0 = vector.extract(as_ir_value(d_out[td]), static_position=[0], dynamic_position=[]) - p1 = vector.extract(as_ir_value(d_out[td]), static_position=[1], dynamic_position=[]) - p2 = vector.extract(as_ir_value(d_out[td]), static_position=[2], dynamic_position=[]) - p3 = vector.extract(as_ir_value(d_out[td]), static_position=[3], dynamic_position=[]) - lo = rocdl.cvt_pk_fp8_f32(T.i32, p0, p1, arith.constant(0, type=T.i32), False) - pk = rocdl.cvt_pk_fp8_f32(T.i32, p2, p3, lo, True) + pv = fx.Vector(d_out[td]) + lo = rocdl.cvt_pk_fp8_f32(T.i32, pv[0], pv[1], arith.constant(0, type=T.i32), False) + pk = rocdl.cvt_pk_fp8_f32(T.i32, pv[2], pv[3], lo, True) elem_base = prob_wr_thread_base + arith.constant(td * MFMA_N * (PROB_ROW_STRIDE_BYTES // 4), type=T.i32) pk_vec = fx.Vector.from_elements([pk], dtype=fx.Int32) fx.ptr_store(pk_vec, logits_base + elem_base) @@ -1500,8 +1497,8 @@ def pa_decode_metadata_kenrel( # ── Work loop bounds ── # wi[cu_id] and wi[cu_id+1] are adjacent int32; load both in one vec2 load. work_bounds = buffer_ops.buffer_load(wi_rsrc, cu_id, vec_width=2, dtype=T.i32) - work_start = vector.extract(as_ir_value(work_bounds), static_position=[0], dynamic_position=[]) - work_end = vector.extract(as_ir_value(work_bounds), static_position=[1], dynamic_position=[]) + work_start = fx.Vector(work_bounds)[0] + work_end = fx.Vector(work_bounds)[1] # Outer work loop — each work item = one (batch, kv_head_range, kv_page_range) _work_start_idx = fx.Index(arith.unwrap(work_start)) @@ -1520,12 +1517,14 @@ def pa_decode_metadata_kenrel( wi_hi = buffer_ops.buffer_load( winfo_rsrc, info_base + arith.constant(4, type=T.i32), vec_width=4, dtype=T.i32 ) - batch_idx = vector.extract(as_ir_value(wi_lo), static_position=[0], dynamic_position=[]) - partial_idx = vector.extract(as_ir_value(wi_lo), static_position=[1], dynamic_position=[]) - qo_start = vector.extract(as_ir_value(wi_lo), static_position=[2], dynamic_position=[]) - kv_start = vector.extract(as_ir_value(wi_hi), static_position=[0], dynamic_position=[]) - kv_end = vector.extract(as_ir_value(wi_hi), static_position=[1], dynamic_position=[]) - q_head_range = vector.extract(as_ir_value(wi_hi), static_position=[3], dynamic_position=[]) + wi_lo_v = fx.Vector(wi_lo) + batch_idx = wi_lo_v[0] + partial_idx = wi_lo_v[1] + qo_start = wi_lo_v[2] + wi_hi_v = fx.Vector(wi_hi) + kv_start = wi_hi_v[0] + kv_end = wi_hi_v[1] + q_head_range = wi_hi_v[3] # work_info.kv_start/kv_end are cumulative partition indices (256-token # units, summed across batches). partition_indptr[batch] gives the @@ -1537,8 +1536,9 @@ def pa_decode_metadata_kenrel( # array. kv_page_end clamps small-block page-gather reads so the last # (partial) partition never reads past the sequence. _kvind2 = buffer_ops.buffer_load(kvindptr_rsrc, batch_idx, vec_width=2, dtype=T.i32) - kv_page_base = vector.extract(as_ir_value(_kvind2), static_position=[0], dynamic_position=[]) - kv_page_end = vector.extract(as_ir_value(_kvind2), static_position=[1], dynamic_position=[]) + _kvind2_v = fx.Vector(_kvind2) + kv_page_base = _kvind2_v[0] + kv_page_end = _kvind2_v[1] local_part_start = kv_start - kv_part_base # Derive kv_head from q_head_range @@ -2078,7 +2078,7 @@ def launch_pa_decode_metadata( # reduce_partial_map: slot → partial row base (in the sliced buffer). # Direct (non-split) outputs have empty groups (indptr delta 0) and are skipped. @functools.lru_cache(maxsize=64) -def compile_pa_ps_reduce( +def compile_pa_metadata_reduce( *, query_length: int, num_query_heads: int, @@ -2089,7 +2089,7 @@ def compile_pa_ps_reduce( assert 0 < block_threads <= 1024, "head_size must fit in one workgroup" @flyc.kernel(known_block_size=(block_threads, 1, 1)) - def pa_ps_reduce_kernel( + def pa_metadata_reduce_kernel( final_output_ptr: fx.Tensor, partial_output_ptr: fx.Tensor, partial_lse_ptr: fx.Tensor, @@ -2166,7 +2166,7 @@ def pa_ps_reduce_kernel( buffer_ops.buffer_store(out_val, out_rsrc, out_off) @flyc.jit - def launch_pa_ps_reduce( + def launch_pa_metadata_reduce( final_output, partial_output, partial_lse, @@ -2180,7 +2180,7 @@ def launch_pa_ps_reduce( num_groups, stream: fx.Stream = fx.Stream(None), ): - pa_ps_reduce_kernel( + pa_metadata_reduce_kernel( final_output, partial_output, partial_lse, @@ -2197,7 +2197,7 @@ def launch_pa_ps_reduce( stream=stream, ) - return {"launch": launch_pa_ps_reduce, "kernel": pa_ps_reduce_kernel} + return {"launch": launch_pa_metadata_reduce, "kernel": pa_metadata_reduce_kernel} _PA_PS_REDUCE_DTYPE_STR = { @@ -2207,7 +2207,7 @@ def launch_pa_ps_reduce( } -def pa_ps_reduce( +def pa_metadata_reduce( *, partial_output: torch.Tensor, partial_lse: torch.Tensor, @@ -2231,7 +2231,7 @@ def pa_ps_reduce( stride_po_row = num_query_heads * head_size stride_pl_row = num_query_heads out_dtype_str = _PA_PS_REDUCE_DTYPE_STR[final_output.dtype] - compiled = compile_pa_ps_reduce( + compiled = compile_pa_metadata_reduce( query_length=int(max_seqlen_q), num_query_heads=int(num_query_heads), head_size=int(head_size), diff --git a/kernels/common/utils.py b/kernels/common/utils.py index 4dc3b336..e1613810 100644 --- a/kernels/common/utils.py +++ b/kernels/common/utils.py @@ -35,11 +35,8 @@ def exp2_f32_fast(value): raw = arith.unwrap(value) if hasattr(value, "ir_value") or hasattr(value, "type") else value ty = raw.type if isinstance(ty, ir.VectorType): - n = ty.shape[0] - elems = [] - for i in range(n): - scalar = _vector_dialect.extract(raw, static_position=[i], dynamic_position=[]) - elems.append(exp2_amdgcn_scalar(scalar)) + vec = fx.Vector(raw) + elems = [exp2_amdgcn_scalar(vec[i]) for i in range(ty.shape[0])] return _vector_dialect.from_elements(ty, elems) return exp2_amdgcn_scalar(raw) diff --git a/tests/kernels/test_pa.py b/tests/kernels/test_pa.py index c81da790..da100c3e 100644 --- a/tests/kernels/test_pa.py +++ b/tests/kernels/test_pa.py @@ -11,7 +11,6 @@ from pathlib import Path from typing import Dict, List, Optional, Tuple, Union -import numpy as np import pandas as pd import pytest import torch @@ -21,7 +20,6 @@ import aiter from aiter import dtypes, per_tensor_quant, pertoken_quant from aiter.ops.triton.gluon.pa_decode_gluon import get_recommended_splits - from aiter.test_common import checkAllclose except Exception as exc: pytest.skip(f"aiter is not available: {exc}", allow_module_level=True) @@ -69,11 +67,6 @@ "fp8": torch.uint8, } -CASE_SET_NAME_OPTIONS = [ - "normal_accuracy", - "sliding_window_accuracy", -] - COMPUTE_TYPE_OPTIONS = ["fp8"] KV_VARLEN_OPTIONS = [False, True] TRANS_V_OPTIONS = [True] @@ -95,64 +88,6 @@ def setup_seed(seed: int) -> None: torch.backends.cudnn.deterministic = True -def compare_arrays( - arr1: np.ndarray, - arr2: np.ndarray, - k: int = 5, - thresholds: List[float] = [0, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1e0, 1e1], -) -> Dict[str, object]: - if arr1.shape != arr2.shape: - raise ValueError("Input arrays must have the same shape") - arr1 = arr1.astype(np.float32) - arr2 = arr2.astype(np.float32) - diff = np.abs(arr1 - arr2) - total_elements = arr1.size - result: Dict[str, object] = { - "top_k_diff": [], - "threshold_stats": [], - "max_diff": float(diff.max()), - "max_diff_thr": float((diff / (1.0 + np.abs(arr2))).max()), - } - flat_diff = diff.flatten() - top_k_indices = np.argpartition(flat_diff, -k)[-k:] - top_k_indices = top_k_indices[np.argsort(-flat_diff[top_k_indices])] - orig_indices = np.unravel_index(top_k_indices, diff.shape) - for i in range(k): - idx = tuple(dim[i] for dim in orig_indices) - result["top_k_diff"].append( - { - "value": float(diff[idx]), - "position": idx, - "arr1_value": float(arr1[idx]), - "arr2_value": float(arr2[idx]), - } - ) - for i in range(len(thresholds) - 1): - lower = thresholds[i] - upper = thresholds[i + 1] - mask = (diff >= lower) & (diff < upper) - count = int(np.sum(mask)) - result["threshold_stats"].append( - { - "range": f"[{lower:.1e}, {upper:.1e})", - "count": count, - "percentage": 100.0 * count / total_elements, - } - ) - mask = diff >= thresholds[-1] - count = int(np.sum(mask)) - result["threshold_stats"].append( - { - "range": f">={thresholds[-1]:.1e}", - "count": count, - "percentage": 100.0 * count / total_elements, - } - ) - print(f"diff.abs.max={result['max_diff']}") - print(f"max_diff_thr={result['max_diff_thr']}") - return result - - def get_kv_cache_torch_dtype( cache_dtype: Optional[Union[str, torch.dtype]], model_dtype: Optional[Union[str, torch.dtype]] = None, @@ -601,7 +536,7 @@ def run_flydsl_ps( def get_tolerance(*, kv_varlen: bool, sliding_window: int) -> float: - diff_tolerance = 5e-3 + diff_tolerance = 8e-3 if kv_varlen: diff_tolerance = 5e-2 if sliding_window > 0: @@ -611,11 +546,6 @@ def get_tolerance(*, kv_varlen: bool, sliding_window: int) -> float: return diff_tolerance -def get_ps_vs_gluon_tolerance(ps_tolerance: float, gluon_tolerance: float) -> float: - """Cross-check tolerance should not be stricter than the Gluon reference itself.""" - return max(ps_tolerance, gluon_tolerance) - - def dtype_to_name(dtype: torch.dtype) -> str: for name, candidate in dtypes.d_dtypes.items(): if candidate == dtype: @@ -623,24 +553,6 @@ def dtype_to_name(dtype: torch.dtype) -> str: return str(dtype) -def summarize_comparison( - name: str, - actual: torch.Tensor, - expected: torch.Tensor, - *, - atol: float, - rtol: float, -) -> Tuple[int, Dict[str, object]]: - err = checkAllclose(expected, actual, atol=atol, rtol=rtol, msg=f"[{name}]") - err = 1 if err > 0 else 0 - diff_result = compare_arrays( - actual.to(torch.float32).detach().cpu().numpy(), - expected.to(torch.float32).detach().cpu().numpy(), - ) - print(f"{name} {'PASSED' if err == 0 else 'FAILED'}") - return err, diff_result - - def run_pa_decode_ps_test( context_length: int, batch_size: int, @@ -820,13 +732,8 @@ def gluon_call() -> None: gluon_time = measure_us(gluon_call) gluon_tol = get_tolerance(kv_varlen=kv_varlen, sliding_window=sliding_window) print("\nGluon vs Torch:") - err_gluon, gluon_diff = summarize_comparison( - "Gluon vs Torch", - gluon_output, - reference_output, - atol=gluon_tol, - rtol=gluon_tol, - ) + torch.testing.assert_close(gluon_output, reference_output, atol=gluon_tol, rtol=gluon_tol) + print("Gluon vs Torch PASSED") kv_page_indices, kv_indptr = build_ps_page_data( block_tables_list, @@ -849,10 +756,16 @@ def gluon_call() -> None: ps_value_scale: torch.Tensor = value_scale_original flydsl_ps_output = torch.empty_like(reference_output) + # Match pa_decode_ps_kernel: each split unit is one 256-token partition, + # containing context_partition_size // block_size physical KV blocks. + blocks_per_partition = context_partition_size // block_size max_context_partition_num = get_recommended_splits( - sliding_window, - context_partition_size, - query_length, + batch_size, + num_kv_heads, + blocks_per_partition, + sliding_window=sliding_window, + context_partition_size=context_partition_size, + query_length=query_length, ) # Preallocate the FlyDSL intermediate buffers (partial exp-sums / max-logits / # output) unconditionally so CUDA-graph capture works for every path, not just @@ -897,24 +810,13 @@ def flydsl_ps_call() -> None: flydsl_ps_time = measure_us(flydsl_ps_call) ps_tol = get_tolerance(kv_varlen=kv_varlen, sliding_window=sliding_window) print("\nFlyDSL PS vs Torch:") - err_flydsl_ps, flydsl_ps_diff = summarize_comparison( - "FlyDSL PS vs Torch", - flydsl_ps_output, - reference_output, - atol=ps_tol, - rtol=ps_tol, - ) + torch.testing.assert_close(flydsl_ps_output, reference_output, atol=ps_tol, rtol=ps_tol) + print("FlyDSL PS vs Torch PASSED") if HAS_GLUON: results["us_gluon"] = gluon_time - results["err_gluon"] = err_gluon - results["gluon_max_diff"] = float(gluon_diff["max_diff"]) - results["gluon_max_diff_thr"] = float(gluon_diff["max_diff_thr"]) results["us_flydsl_ps"] = flydsl_ps_time - results["err_flydsl_ps"] = err_flydsl_ps - results["flydsl_ps_max_diff"] = float(flydsl_ps_diff["max_diff"]) - results["flydsl_ps_max_diff_thr"] = float(flydsl_ps_diff["max_diff_thr"]) return results @@ -1169,79 +1071,92 @@ def parse_arg_and_run_test(sample_rate0: float = None, *, output_tag: str = TEST results_df.to_csv(output_file, index=False) print(f"\nResults saved to {output_file}") print(f"\nSummary:\n{results_df}") - flydsl_errors = int(results_df["err_flydsl_ps"].sum()) + print("\nAll PS-only tests passed!") - if flydsl_errors: - raise AssertionError(f"{flydsl_errors} FlyDSL PS case(s) exceeded the Torch-reference tolerance") - print("\nAll PS-only tests passed!") +@pytest.mark.parametrize("compute_type", ["fp8"]) +@pytest.mark.parametrize("context_partition_size", [256]) +@pytest.mark.parametrize("head_size", [128, 256]) +@pytest.mark.parametrize("num_heads", [(8, 1), (16, 1), (4, 1)]) +@pytest.mark.parametrize("query_length", [1, 2, 3, 4]) +@pytest.mark.parametrize("quant_mode", ["per_token", "per_tensor"]) +@pytest.mark.parametrize("context_length", [1027, 8192]) +@pytest.mark.parametrize("batch_size", [3, 81, 128]) +@pytest.mark.parametrize("trans_v", [True, False]) +@pytest.mark.parametrize("kv_varlen", [False, True]) +@pytest.mark.parametrize("block_size", [16, 64]) +@pytest.mark.parametrize("sliding_window", [0]) +def test_normal_accuracy( + compute_type: str, + context_partition_size: int, + head_size: int, + num_heads: Tuple[int, int], + query_length: int, + quant_mode: str, + context_length: int, + batch_size: int, + trans_v: bool, + kv_varlen: bool, + block_size: int, + sliding_window: int, +) -> None: + run_pa_decode_ps_test( + context_length=context_length, + batch_size=batch_size, + num_heads=num_heads, + head_size=head_size, + block_size=block_size, + compute_type=dtypes.d_dtypes[compute_type], + query_length=query_length, + quant_mode=quant_mode, + context_partition_size=context_partition_size, + trans_v=trans_v, + kv_varlen=kv_varlen, + sliding_window=sliding_window, + ) -def normal_accuracy_test() -> None: - global BLOCK_SIZE_OPTIONS - global QUERY_LENGTH_OPTIONS - global BATCH_SIZE_OPTIONS - global HEAD_CONFIGURATIONS - global CONTEXT_LENGTH_OPTIONS - global COMPUTE_TYPE_OPTIONS - global QUANT_MODE_OPTIONS - global HEAD_DIMENSION_OPTIONS - global TRANS_V_OPTIONS - global KV_VARLEN_OPTIONS - global CONTEXT_PARTITION_SIZE_OPTIONS - global SLIDING_WINDOW_OPTIONS - COMPUTE_TYPE_OPTIONS = ["fp8"] - CONTEXT_PARTITION_SIZE_OPTIONS = [256] - HEAD_DIMENSION_OPTIONS = [128] - HEAD_CONFIGURATIONS = [(8, 1), (16, 1)] - QUERY_LENGTH_OPTIONS = [1, 2, 3, 4] - QUANT_MODE_OPTIONS = ["per_token", "per_tensor"] - CONTEXT_LENGTH_OPTIONS = [1027] - BATCH_SIZE_OPTIONS = [3, 81] - TRANS_V_OPTIONS = [True] - KV_VARLEN_OPTIONS = [False, True] - BLOCK_SIZE_OPTIONS = [1024] - SLIDING_WINDOW_OPTIONS = [0] - parse_arg_and_run_test(output_tag="ps_normal_accuracy") - - -def sliding_window_accuracy_test() -> None: - global BLOCK_SIZE_OPTIONS - global QUERY_LENGTH_OPTIONS - global BATCH_SIZE_OPTIONS - global HEAD_CONFIGURATIONS - global CONTEXT_LENGTH_OPTIONS - global COMPUTE_TYPE_OPTIONS - global QUANT_MODE_OPTIONS - global HEAD_DIMENSION_OPTIONS - global TRANS_V_OPTIONS - global KV_VARLEN_OPTIONS - global CONTEXT_PARTITION_SIZE_OPTIONS - global SLIDING_WINDOW_OPTIONS - COMPUTE_TYPE_OPTIONS = ["fp8"] - CONTEXT_PARTITION_SIZE_OPTIONS = [256] - HEAD_DIMENSION_OPTIONS = [128] - HEAD_CONFIGURATIONS = [(8, 1), (16, 1)] - QUERY_LENGTH_OPTIONS = [1, 2, 3, 4] - QUANT_MODE_OPTIONS = ["per_token"] - CONTEXT_LENGTH_OPTIONS = [8192] - BATCH_SIZE_OPTIONS = [128] - TRANS_V_OPTIONS = [True] - KV_VARLEN_OPTIONS = [True] - BLOCK_SIZE_OPTIONS = [16, 1024] - SLIDING_WINDOW_OPTIONS = [0] - parse_arg_and_run_test(output_tag="ps_sliding_window_accuracy") - - -@pytest.mark.parametrize("case_set_name", CASE_SET_NAME_OPTIONS) -def test_multi_case_set(case_set_name: str) -> None: - if case_set_name == "normal_accuracy": - normal_accuracy_test() - elif case_set_name == "sliding_window_accuracy": - sliding_window_accuracy_test() - else: - raise ValueError(f"Unsupported case set: {case_set_name}") +@pytest.mark.parametrize("compute_type", ["fp8"]) +@pytest.mark.parametrize("context_partition_size", [256]) +@pytest.mark.parametrize("head_size", [128]) +@pytest.mark.parametrize("num_heads", [(8, 1), (16, 1)]) +@pytest.mark.parametrize("query_length", [1, 2, 3, 4]) +@pytest.mark.parametrize("quant_mode", ["per_token"]) +@pytest.mark.parametrize("context_length", [8192]) +@pytest.mark.parametrize("batch_size", [128]) +@pytest.mark.parametrize("trans_v", [True]) +@pytest.mark.parametrize("kv_varlen", [True]) +@pytest.mark.parametrize("block_size", [1024]) +@pytest.mark.parametrize("sliding_window", [1023]) +def test_sliding_window_accuracy( + compute_type: str, + context_partition_size: int, + head_size: int, + num_heads: Tuple[int, int], + query_length: int, + quant_mode: str, + context_length: int, + batch_size: int, + trans_v: bool, + kv_varlen: bool, + block_size: int, + sliding_window: int, +) -> None: + run_pa_decode_ps_test( + context_length=context_length, + batch_size=batch_size, + num_heads=num_heads, + head_size=head_size, + block_size=block_size, + compute_type=dtypes.d_dtypes[compute_type], + query_length=query_length, + quant_mode=quant_mode, + context_partition_size=context_partition_size, + trans_v=trans_v, + kv_varlen=kv_varlen, + sliding_window=sliding_window, + ) if __name__ == "__main__": - sliding_window_accuracy_test() + parse_arg_and_run_test()