From 637f8cbbcb340645c4cd0544c0dca5df17afc763 Mon Sep 17 00:00:00 2001 From: Lingpeng Jin <103567126+valarLip@users.noreply.github.com> Date: Sat, 5 Sep 2026 05:33:48 +0000 Subject: [PATCH] perf: bound the FP4 MQA-logits store with a window-sized V# The per-token store was guarded in software -- `is_writer & (tok >= local_start) & (tok < local_end)` -- which lowers to 2 v_cmp + 2 s_and_b64 + s_or_b64 + s_and_saveexec_b64 + s_cbranch_execz per n-tile: 28 instructions per chunk, 11.6% of the steady-state loop body and second only to the reduction itself. Size a buffer descriptor to the window instead and num_records IS the test, in hardware. A token below the window underflows to a large unsigned offset and is dropped by the same bound, so one check covers both ends. Two things this gets wrong if done casually: - Lanes 16..63 hold redundant copies of the butterfly result and must not write. Parking them past num_records needs a CONSTANT offset, not a multiple of the window length: `token_base - local_start` is negative for a token below a non-zero window start, and adding the length to that lands back INSIDE the window, racing the lane that owns the column. - The window bounds must go through readfirstlane. They are wave-uniform but arrive from cta_info in VGPRs, and a divergent descriptor makes the compiler wrap the store in a waterfall loop -- 252 instructions, worse than the guard it replaces. Steady-state loop body 242 -> 203 instructions. MI355X, 2 interleaved reps, min over 15-20 iters: rows=8192 W~2064 124.4us -> 119.0us (-4.4%) rows=1024 W~32768 241.0us -> 225.4us (-6.5%) Same fold on the rectangular decode twin. Output is unchanged -- cells outside the window still hold the caller's -inf pre-fill, which test_flydsl_pa_mqa_logits_fp4{,_prefill}.py already assert on every case, including the one window set with a non-zero local_start. Co-Authored-By: Claude Opus 5 (1M context) --- .../kernels/mqa_logits/pa_mqa_logits_fp4.py | 55 ++++++++++++--- .../mqa_logits/pa_mqa_logits_fp4_common.py | 7 ++ .../mqa_logits/pa_mqa_logits_fp4_prefill.py | 67 ++++++++++++++----- 3 files changed, 103 insertions(+), 26 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py index 0a5ffabaf48..a2972fe6da5 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py @@ -16,6 +16,7 @@ from flydsl.expr.typing import Float4E2M1FN, Int32, T from .pa_mqa_logits_fp4_common import ( + _NON_WRITER_LANE_OFF, _i32_buffer, _load_vec4_i32, ) @@ -259,10 +260,16 @@ def pa_mqa_logits_fp4_kernel( ZERO_F = fx.Float32(0.0) c0_i32 = fx.Int32(0) - batch_packed = cta_info_vec[0] + # Wave-uniform by construction, but they arrive in VGPRs via the buffer + # load; the V# below needs num_records in an SGPR or the store is wrapped + # in a waterfall loop. + def _uniform(v): + return fx.Int32(fx.rocdl.readfirstlane(T.i32, v.ir_value())) + + batch_packed = _uniform(cta_info_vec[0]) chunk_start = cta_info_vec[1] chunk_count = cta_info_vec[2] - context_len = cta_info_vec[3] + context_len = _uniform(cta_info_vec[3]) # out row base folded into an f32 global pointer: sizeof(f32) = 4, so the # per-token store offset below stays small (no i32 overflow for large @@ -273,6 +280,33 @@ def pa_mqa_logits_fp4_kernel( pid_b = batch_packed // fx.Int32(next_n) pid_next_n = batch_packed % fx.Int32(next_n) + # A V# spanning exactly this row's visible range: num_records then IS + # the `out_token + mask_off < context_len` test, in hardware. Lanes + # 16..63 hold redundant copies of the butterfly result and must not + # write; `lane_div_16 * win_len` puts them past num_records. The clamp + # matters: a negative length wraps to a huge unsigned bound. + _mask_off = fx.Int32(next_n - 1) - pid_next_n + _win_raw = context_len - _mask_off + win_len = (_win_raw < fx.Int32(0)).select(fx.Int32(0), _win_raw) + out_win = fx.rocdl.make_buffer_tensor( + fx.make_view( + fx.recast_iter( + fx.PointerType.get(T.f32, out_base.memspace, 4), out_base + ), + fx.make_layout((win_len, 1), (1, 1)), + ), + max_size=False, + num_records_bytes=win_len * fx.Int32(4), + ) + out_lane_off = lane_mod_16 + (lane_div_16 > fx.Int32(0)).select( + fx.Int32(_NON_WRITER_LANE_OFF), fx.Int32(0) + ) + out_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 1) + out_reg_ty = fx.MemRefType.get( + T.f32, fx.LayoutType.get(1, 1), fx.AddressSpace.Register + ) + out_reg_lay = fx.make_layout(1, 1) + # Scaled FP4 16x16x128 MMA (opsel 0/0); one e8m0 word per operand. mfma_atom = fx.make_mma_atom( fx.rocdl.cdna4.MFMA_Scale( @@ -457,15 +491,14 @@ def _bperm_xor_add(val, sh): thread_sum = _bperm_xor_add(thread_sum, 32) thread_sum = thread_sum * weight_scale - is_writer = lane_div_16 < fx.Int32(1) - out_token = token_base + lane_mod_16 - mask_off = fx.Int32(next_n - 1) - pid_next_n - in_ctx = (out_token + mask_off) < context_len - # Sparse 1-writer scatter: guard the plain store instead of a V# - # OOB sentinel. Row base is folded into out_bt's pointer, so the - # store offset is just the (small) token index. - if is_writer & in_ctx: - fx.ptr_store(thread_sum, fx.add_offset(out_base, out_token)) + # Context bound and writer-lane guard both live in the V# built + # above; nothing is tested here. + off = token_base + out_lane_off + r_out = fx.memref_alloca(out_reg_ty, out_reg_lay) + fx.memref_store_vec( + fx.Vector.from_elements([thread_sum], dtype=fx.Float32), r_out + ) + fx.copy(out_atom, r_out, fx.slice(out_win, (off, None))) def _compute_chunk(kv_list_in, kvs_packed_list_in, c_i32_arg, nt0_accs_in=None): """Process chunk c using prefetched (kv, kvs_packed).""" diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_common.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_common.py index 4f2b01838bc..df9d1bef50c 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_common.py @@ -7,6 +7,13 @@ import flydsl.expr as fx from flydsl.expr.typing import T +# Offset that parks a non-writer lane outside any window's num_records. Must be +# a constant, not a multiple of the window length: a token below a non-zero +# `local_start` has a negative base offset, and adding the length to that lands +# back inside the window. Big enough that no in-range offset reaches it, small +# enough not to overflow i32 when added to one. +_NON_WRITER_LANE_OFF = 1 << 30 + def _i32_buffer(ptr, width=1): """OOB-checked global i32 buffer-tensor over ``ptr`` (mirrors a max_size V#). diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py index 530d0eb3019..de938996939 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py @@ -17,6 +17,7 @@ from flydsl.expr.typing import Float4E2M1FN, Int32, T from .pa_mqa_logits_fp4_common import ( + _NON_WRITER_LANE_OFF, _i32_buffer, _load_vec4_i32, ) @@ -472,8 +473,15 @@ def pa_mqa_logits_fp4_prefill_kernel( fx.make_view(cta_it, fx.make_layout((1 << 28, 4), (4, 1))) ) cta_info_vec = fx.Vector(_load_vec4_i32(cta_info_bt, fx.Int32(0))) - local_start = cta_info_bt[(fx.Int32(1), fx.Int32(0))] - local_end = cta_info_bt[(fx.Int32(1), fx.Int32(1))] + + # Wave-uniform by construction, but they arrive in VGPRs via the buffer + # load; the V# below needs num_records in an SGPR or the store is wrapped + # in a waterfall loop. + def _uniform(v): + return fx.Int32(fx.rocdl.readfirstlane(T.i32, v.ir_value())) + + local_start = _uniform(cta_info_bt[(fx.Int32(1), fx.Int32(0))]) + local_end = _uniform(cta_info_bt[(fx.Int32(1), fx.Int32(1))]) kv_bt = _i32_buffer(kv_cache_ptr, width=4) kvs_bt = _i32_buffer(kv_scale_ptr, width=1) @@ -487,10 +495,37 @@ def pa_mqa_logits_fp4_prefill_kernel( chunk_start = cta_info_vec[2] chunk_count = cta_info_vec[3] - # out row base folded into an f32 global pointer (sizeof(f32)=4); the - # per-token store offset below stays small (no i32 overflow). - _row_elems = fx.Int64(row_id) * fx.Int64(stride_out_row) - out_base = fx.add_offset(fx.get_iter(out_logits_ptr), _row_elems) + # A V# spanning exactly [local_start, local_end) of this row: num_records + # then IS the window test, in hardware. A token below the window + # underflows to a huge unsigned offset and is dropped by the same bound, + # so one check covers both ends. Lanes 16..63 hold redundant copies of + # the butterfly result and must not write; a large constant added to + # their offset puts them past num_records. It must be a CONSTANT, not a + # multiple of win_len: `token_base - local_start` is negative for a + # token below a non-zero window start, and adding win_len to that lands + # back INSIDE the window, racing the lane that owns the column. Both + # terms are chunk-invariant, so both live here. + win_len = local_end - local_start + _row_elems = fx.Int64(row_id) * fx.Int64(stride_out_row) + fx.Int64(local_start) + out_win = fx.rocdl.make_buffer_tensor( + fx.make_view( + fx.recast_iter( + fx.PointerType.get(T.f32, out_logits_ptr.memspace, 4), + fx.add_offset(fx.get_iter(out_logits_ptr), _row_elems), + ), + fx.make_layout((win_len, 1), (1, 1)), + ), + max_size=False, + num_records_bytes=win_len * fx.Int32(4), + ) + out_lane_off = lane_mod_16 + (lane_div_16 > fx.Int32(0)).select( + fx.Int32(_NON_WRITER_LANE_OFF), fx.Int32(0) + ) + out_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), 1) + out_reg_ty = fx.MemRefType.get( + T.f32, fx.LayoutType.get(1, 1), fx.AddressSpace.Register + ) + out_reg_lay = fx.make_layout(1, 1) # Q load (hoisted): per (k_tile, mi_idx) a thread loads its 16-byte FP4 # chunk for head row mi_idx*16+lane_mod_16. Q: [total_tokens, H, D/2] uint8. @@ -674,15 +709,17 @@ def _bperm_xor_add(val, sh): thread_sum = _bperm_xor_add(thread_sum, 32) # `weight_scale` already folded into `w_per_lane` (hoisted, once/wave). - # Only [local_start, local_end) is written (one writer lane per token); - # the rest stays at the caller's -inf pre-fill. Sparse 1-writer - # scatter: guard the plain store instead of a V# OOB sentinel. Row base - # is folded into out_base, so the store offset is the token index. - is_writer = lane_div_16 < fx.Int32(1) - out_token = token_base + lane_mod_16 - in_window = (out_token >= local_start) & (out_token < local_end) - if is_writer & in_window: - fx.ptr_store(thread_sum, fx.add_offset(out_base, out_token)) + # Window and writer-lane guards are both in the V#; nothing is + # tested here. Cells outside stay at the caller's -inf pre-fill. + r_out = fx.memref_alloca(out_reg_ty, out_reg_lay) + fx.memref_store_vec( + fx.Vector.from_elements([thread_sum], dtype=fx.Float32), r_out + ) + fx.copy( + out_atom, + r_out, + fx.slice(out_win, (token_base - local_start + out_lane_off, None)), + ) def _compute_chunk(kv_list_in, kvs_packed_list_in, c_i32_arg, nt0_accs_in=None): assert (