From da930d50eddf70aa7fbbe7c9b6a4e0790b82e9b8 Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Mon, 3 Aug 2026 07:50:46 -0500 Subject: [PATCH 01/12] add bias, alibi bias, sink to flas attention gfx950 --- kernels/attention/flash_attn_gfx950.py | 258 +++++++- kernels/attention/flash_attn_interface.py | 154 ++++- tests/kernels/test_flash_attn_fwd.py | 685 ++++++++++++++++++++-- 3 files changed, 1041 insertions(+), 56 deletions(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index e72e7e163..3590cc68a 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -13,10 +13,14 @@ up to even, and a kv padding-mask on the non-causal path). """ +import math as host_math + import flydsl.compiler as flyc import flydsl.expr as fx from flydsl.compiler.kernel_function import CompilationContext from flydsl.expr import const_expr, range_constexpr +from flydsl.expr import math as fmath +from flydsl.expr.typing import Vector as Vec from flydsl.runtime.device import get_rocm_arch as get_hip_arch from kernels.attention.flash_attn_utils import ( DualwaveGemmHelper, @@ -32,7 +36,9 @@ _anchor_v_o, _anchor_v_p, _dualwave_sync_barrier, + _get_q_pack, _make_dualwave_swp_traits, + _mfma_acc, _s_barrier, _s_nop, _s_setprio, @@ -40,6 +46,8 @@ _sched_barrier, _sched_barrier_exp_pairs, _sched_barrier_pairs, + _seq_pad_col_base, + _seq_pad_score_threshold, _stagger_extra_barrier_if_one, _stagger_extra_barrier_if_zero, _v_pair_to_vec32, @@ -67,13 +75,22 @@ def build_flash_attn_dualwave_swp_module( paged=False, kv_cache_layout="linear", return_lse=False, + has_bias=False, + has_alibi=False, + has_sink=False, ): """Build an DUALWAVE_SWP flash_attn launcher for D=64/128 bf16/f16 on gfx950. Supports dense self-attention, varlen packed QKV, and paged-KV cache modes. Varlen uses cu_seqlens_q/kv with per-batch self-attention ranges. Paged mode keeps Q/O dense and maps KV tiles through BlockTable pages. - Varlen, paged, and split-K are mutually constrained by the caller.""" + Varlen, paged, and split-K are mutually constrained by the caller. + has_bias adds a dense elementwise attention bias and works on both the dense + and varlen paths; see the HAS_BIAS block below for the per-mode shapes. + has_alibi adds a per-head ALiBi positional bias computed analytically from a + slope table; see the HAS_ALIBI block below. + has_sink adds a per-head attention sink logit to the softmax denominator; see + the HAS_SINK block below. The three are independent and may be combined.""" gpu_arch = get_hip_arch() if not gpu_arch.startswith("gfx950"): @@ -98,6 +115,10 @@ def build_flash_attn_dualwave_swp_module( raise ValueError("vectorized layout requires HEAD_DIM and PageSize divisible by kVS") if VARLEN and SPLITK: raise ValueError("varlen is not supported together with num_kv_splits > 1") + HAS_BIAS = bool(has_bias) + HAS_ALIBI = bool(has_alibi) + HAS_SINK = bool(has_sink) + BIAS_LOG2E = host_math.log2(host_math.e) traits = _make_dualwave_swp_traits( num_heads, @@ -120,7 +141,8 @@ def build_flash_attn_dualwave_swp_module( return_lse=return_lse, ) traits.BLOCK_N_OUT // traits.BLOCK_N - _dualwave_swp_cache_tag = traits.cache_tag + + _dualwave_swp_cache_tag = (traits.cache_tag, HAS_BIAS, HAS_ALIBI, HAS_SINK) # Shared-memory layout: one 16B-aligned K/V region (K0/V0/K1/V1). _lds_elem_dtype = dtype_to_elem_type(traits.DTYPE_STR) @@ -149,12 +171,17 @@ def flash_attn_dualwave_swp_gfx950_kernel( CuSeqQ: fx.Tensor, CuSeqKv: fx.Tensor, BlockTable: fx.Tensor, + Bias: fx.Tensor, + AlibiSlopes: fx.Tensor, + Sink: fx.Tensor, seq_len: fx.Int32, seq_len_kv: fx.Int32, stride_q_n: fx.Int32, stride_kv_n: fx.Int32, head_dim_runtime: fx.Int32, block_table_stride: fx.Int32, + bias_stride0: fx.Int32, + alibi_stride_b: fx.Int32, ): ctx = DualwaveKernelContext( traits, @@ -203,6 +230,91 @@ def flash_attn_dualwave_swp_gfx950_kernel( gemm_helper = DualwaveGemmHelper(ctx) softmax_helper = DualwaveSoftmaxHelper(ctx) + # ---------------- attention bias ---------------- + if const_expr(HAS_BIAS): + _bias_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(Bias), fx.make_layout(1, 1)) + if const_expr(traits.KV_VECTORIZED): + _BIAS_VEC = 8 + _bias_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), ctx.elem_dtype) + else: + _BIAS_VEC = 4 + _bias_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), ctx.elem_dtype) + _BIAS_GROUPS = 16 // _BIAS_VEC + _bias_frags = [ + fx.make_rmem_tensor(fx.make_layout(_BIAS_VEC, 1), ctx.elem_dtype) for _ in range_constexpr(_BIAS_GROUPS) + ] + _bias_log2e = Vec.filled(_BIAS_VEC, BIAS_LOG2E, fx.Float32) + _bias_row_base_i32 = fx.Int32(ctx.q_tok_base) if const_expr(VARLEN) else None + + # ---------------- ALiBi ---------------- + if const_expr(HAS_ALIBI): + _alibi_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(AlibiSlopes), fx.make_layout(1, 1)) + _alibi_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + _alibi_frag = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + _alibi_idx_i32 = fx.Int32(ctx.batch_idx) * fx.Int32(alibi_stride_b) + fx.Int32(ctx.q_head_idx) + fx.copy(_alibi_atom, fx.slice(_alibi_div, (None, _alibi_idx_i32)), _alibi_frag) + + _alibi_neg_slope = Vec(_alibi_frag.load(), (1,), fx.Float32)[0] * fx.Float32(-BIAS_LOG2E) + + # ---------------- Attention sink ---------------- + # Split-K folds the sink in the combine pass instead, so skip the load here. + if const_expr(HAS_SINK and not traits.SPLITK): + + def _load_sink_log2(): + _sink_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(Sink), fx.make_layout(1, 1)) + _sink_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32) + _sink_frag = fx.make_rmem_tensor(fx.make_layout(1, 1), fx.Float32) + + fx.copy(_sink_atom, fx.slice(_sink_div, (None, fx.Int32(ctx.q_head_idx))), _sink_frag) + + return Vec(_sink_frag.load(), (1,), fx.Float32)[0] * fx.Float32(BIAS_LOG2E) + + def _score_bias_half(s_h, tile_idx, h): + col_base_h = _seq_pad_col_base(traits, tile_idx, lane_div_32=ctx.lane_div_32) + fx.Int32(32 * h) + src = Vec(s_h) + out = [src[r] for r in range_constexpr(16)] + + if const_expr(HAS_ALIBI): + rel0 = fx.Float32(ctx.q_row_i32 + ctx.delta_i32 - col_base_h) + for r in range_constexpr(16): + d = rel0 - fx.Float32(float(_seq_pad_score_threshold(traits, r))) + out[r] = fmath.absf(d) * _alibi_neg_slope + out[r] + + if const_expr(HAS_BIAS): + bias_row_i32 = ctx.q_row_i32 + if const_expr(VARLEN): + bias_row_i32 = bias_row_i32 + _bias_row_base_i32 + base = bias_row_i32 * fx.Int32(bias_stride0) + col_base_h + for g in range_constexpr(_BIAS_GROUPS): + col_off = _seq_pad_score_threshold(traits, g * _BIAS_VEC) + fx.copy(_bias_atom, fx.slice(_bias_div, (None, base + fx.Int32(col_off))), _bias_frags[g]) + for g in range_constexpr(_BIAS_GROUPS): + bv = Vec(Vec(_bias_frags[g].load(), (_BIAS_VEC,), ctx.elem_dtype).to(fx.Float32)) + r0 = g * _BIAS_VEC + acc = Vec.from_elements([out[r0 + e] for e in range_constexpr(_BIAS_VEC)], fx.Float32) + fused = bv * _bias_log2e + acc + for e in range_constexpr(_BIAS_VEC): + out[r0 + e] = fused[e] + + return Vec.from_elements(out, fx.Float32) + + def qk_scored(v_k, q_all_scaled_bf16, tile_idx): + if const_expr(not (HAS_BIAS or HAS_ALIBI)): + return gemm_helper.qk(v_k, q_all_scaled_bf16) + out_halves = [] + for h, k_h in ((0, v_k[0]), (1, v_k[1])): + acc = gemm_helper.c_zero_v16f32 + for ks in range_constexpr(traits.K_STEPS_QK): + acc = _mfma_acc( + k_h[ks], + _get_q_pack(traits, q_all_scaled_bf16, ks), + acc, + gemm_helper.mma_atom, + gemm_helper.mfma_acc_vec_type, + ) + out_halves.append(_score_bias_half(acc, tile_idx, h)) + return (out_halves[0], out_halves[1]) + def _main_body(): # Paged: stage the block-table row into LDS before any page-id ds_read. if const_expr(traits.PAGED): @@ -249,7 +361,8 @@ def _main_body(): # Prologue scores + first softmax pass for KV tile 0 if const_expr(traits.PAGED): pro_pageid_2_lds = page_ids.load_page_id_lds(page_ids.split_tile(2)) - v_s_0 = gemm_helper.qk(v_k, q_all_scaled_bf16) + + v_s_0 = qk_scored(v_k, q_all_scaled_bf16, ctx.split_tile(0)) _sched_barrier(0) if const_expr(traits.CAUSAL): @@ -326,7 +439,7 @@ def _main_body(): # Cluster 1 computes MMA0, finishes v_p_0 softmax, updates l_row, and casts P. if const_expr(traits.PAGED): c2_pageid_lds = page_ids.load_page_id_lds(j_idx) - v_s_1 = gemm_helper.qk(v_k, q_all_scaled_bf16) + v_s_1 = qk_scored(v_k, q_all_scaled_bf16, j_idx - 2) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) v_p_0 = softmax_helper.cast_p(v_p_0) @@ -401,7 +514,7 @@ def _main_body(): # Cluster 5 mirrors C1: MMA0, finish v_p_1 softmax, update l_row, and cast P. if const_expr(traits.PAGED): _c6_kpid_lds = page_ids.load_page_id_lds(j_idx + 1) - v_s_0 = gemm_helper.qk(v_k, q_all_scaled_bf16) + v_s_0 = qk_scored(v_k, q_all_scaled_bf16, j_idx - 1) v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_1) v_p_1 = softmax_helper.cast_p(v_p_1) @@ -492,7 +605,7 @@ def _main_body(): # Epilogue C1 (compute): MMA0 -> v_s_1; finish v_p_0 softmax (like C1). if const_expr(traits.PAGED): ec2_pageid_lds = page_ids.load_page_id_lds(max_m1) - v_s_1 = gemm_helper.qk(v_k, q_all_scaled_bf16) + v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m3) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) v_p_0 = softmax_helper.cast_p(v_p_0) @@ -559,7 +672,7 @@ def _main_body(): _dualwave_sync_barrier() # Epilogue C5 computes MMA0, folds rescale_e3 into l_row, and finishes v_p_1 softmax. - v_s_0 = gemm_helper.qk(v_k, q_all_scaled_bf16) + v_s_0 = qk_scored(v_k, q_all_scaled_bf16, max_m2) l_row = softmax_helper.apply_l_rescale(l_row, rescale_e3) v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_1) @@ -618,7 +731,7 @@ def _main_body(): _dualwave_sync_barrier() # Epilogue C9 computes the last-tile MMA0, folds rescale_e7 into l_row, and finishes v_p_0. - v_s_1 = gemm_helper.qk(v_k, q_all_scaled_bf16) + v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m1) l_row = softmax_helper.apply_l_rescale(l_row, rescale_e7) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) @@ -671,6 +784,9 @@ def _main_body(): # Epilogue C13 (compute): final P*V -> v_o holds the unnormalized output. v_o = gemm_helper.pv(v_p_1, v_packs_e13, v_o) + if const_expr(HAS_SINK and not traits.SPLITK): + m_row, l_row = softmax_helper.fold_sink(v_o, m_row, l_row, _load_sink_log2()) + # Normalize O; split-K stores normalized partials for later w_s * l_s reweighting. l_inv = softmax_helper.safe_l_inv(l_row) softmax_helper.scale_o(v_o, l_inv) @@ -715,11 +831,12 @@ def flash_attn_splitk_combine_kernel( O: fx.Tensor, # noqa: E741 WS: fx.Tensor, LSE: fx.Tensor, + Sink: fx.Tensor, batch_size: fx.Int32, seq_len: fx.Int32, stride_q_n: fx.Int32, ): - ctx = DualwaveSplitKCombineContext(traits, O, WS, batch_size, seq_len, stride_q_n, LSE=LSE) + ctx = DualwaveSplitKCombineContext(traits, O, WS, batch_size, seq_len, stride_q_n, LSE=LSE, Sink=Sink) ctx.init_types_and_constants() ctx.init_runtime_indices() ctx.init_thread_mapping(COMBINE_ROWS_PER_BLOCK, COMBINE_LANES_PER_ROW) @@ -729,7 +846,12 @@ def flash_attn_splitk_combine_kernel( combine = DualwaveSplitKCombineHelper(ctx) m_s, l_s = combine.load_ml_rows() m_max = combine.reduce_m_max(m_s) + + if const_expr(HAS_SINK): + m_max, sink_w = combine.fold_sink(m_max, BIAS_LOG2E) acc, den = combine.accumulate_splits(m_s, l_s, m_max) + if const_expr(HAS_SINK): + den = combine.add_sink_den(den, sink_w) if const_expr(traits.RETURN_LSE): combine.store_lse(m_max, den) o_pack = combine.pack_output(acc, den) @@ -746,6 +868,9 @@ def launch_flash_attn_dualwave_swp( CuSeqQ: fx.Tensor, CuSeqKv: fx.Tensor, BlockTable: fx.Tensor, + Bias: fx.Tensor, + AlibiSlopes: fx.Tensor, + Sink: fx.Tensor, batch_size: fx.Int32, seq_len: fx.Int32, seq_len_kv: fx.Int32, @@ -753,6 +878,8 @@ def launch_flash_attn_dualwave_swp( stride_kv_n: fx.Int32, head_dim_runtime: fx.Int32, block_table_stride: fx.Int32, + bias_stride0: fx.Int32, + alibi_stride_b: fx.Int32, stream: fx.Stream = fx.Stream(None), ): # Make shape/mode traits visible to the JIT cache key. @@ -784,12 +911,17 @@ def launch_flash_attn_dualwave_swp( CuSeqQ, CuSeqKv, BlockTable, + Bias, + AlibiSlopes, + Sink, seq_len, seq_len_kv, stride_q_n, stride_kv_n, head_dim_runtime, block_table_stride, + bias_stride0, + alibi_stride_b, value_attrs={ "rocdl.waves_per_eu": traits.WAVES_PER_EU, "rocdl.flat_work_group_size": f"{traits.BLOCK_SIZE},{traits.BLOCK_SIZE}", @@ -802,7 +934,7 @@ def launch_flash_attn_dualwave_swp( ) if const_expr(traits.SPLITK): combine_rows = bs_idx * traits.NUM_HEADS_Q * sl_idx - flash_attn_splitk_combine_kernel(O, DebugCounts, LSE, batch_size, seq_len, stride_q_n).launch( + flash_attn_splitk_combine_kernel(O, DebugCounts, LSE, Sink, batch_size, seq_len, stride_q_n).launch( grid=(combine_rows // COMBINE_ROWS_PER_BLOCK, 1, 1), block=(COMBINE_BLOCK, 1, 1), stream=stream, @@ -817,6 +949,56 @@ def launch_flash_attn_dualwave_swp( }, } + def _prep_alibi(alibi_slopes, placeholder): + if alibi_slopes is None: + if HAS_ALIBI: + raise ValueError( + "flash_attn_dualwave_swp was built with has_alibi=True but no `alibi_slopes` " + "tensor was provided; pass an fp32 (num_heads,) or (batch, num_heads) slope table." + ) + + return placeholder, 0 + if not HAS_ALIBI: + + raise ValueError( + "`alibi_slopes` was provided but flash_attn_dualwave_swp was built with " + "has_alibi=False; rebuild with has_alibi=True." + ) + if alibi_slopes.dim() not in (1, 2): + raise ValueError( + f"alibi_slopes must be 1D (num_heads,) or 2D (batch, num_heads), got shape {tuple(alibi_slopes.shape)}" + ) + if alibi_slopes.shape[-1] != num_heads: + raise ValueError( + f"alibi_slopes last dim must be num_heads={num_heads}, got shape {tuple(alibi_slopes.shape)}" + ) + + if not alibi_slopes.is_floating_point() or alibi_slopes.element_size() != 4: + raise ValueError(f"alibi_slopes must be fp32, got dtype {alibi_slopes.dtype}") + alibi_slopes = alibi_slopes.contiguous() + alibi_stride_b = alibi_slopes.stride(0) if alibi_slopes.dim() == 2 else 0 + return alibi_slopes.reshape(-1), alibi_stride_b + + def _prep_sink(sink, placeholder): + if sink is None: + if HAS_SINK: + raise ValueError( + "flash_attn_dualwave_swp was built with has_sink=True but no `sink` tensor " + "was provided; pass an fp32 (num_heads,) table." + ) + + return placeholder + if not HAS_SINK: + raise ValueError( + "`sink` was provided but flash_attn_dualwave_swp was built with " + "has_sink=False; rebuild with has_sink=True." + ) + if sink.dim() != 1 or sink.shape[0] != num_heads: + raise ValueError(f"sink must be 1D (num_heads={num_heads},), got shape {tuple(sink.shape)}") + if not sink.is_floating_point() or sink.element_size() != 4: + raise ValueError(f"sink must be fp32, got dtype {sink.dtype}") + return sink.contiguous() + def _launch( Q, K, @@ -835,6 +1017,10 @@ def _launch( cu_seqlens_kv=None, block_table=None, block_table_stride=None, + bias=None, + bias_stride0=None, + alibi_slopes=None, + sink=None, lse=None, stream=None, ): @@ -872,6 +1058,22 @@ def _launch( block_table = O if block_table_stride is None: block_table_stride = 0 + if bias is None: + if HAS_BIAS: + raise ValueError( + "flash_attn_dualwave_swp was built with has_bias=True but no `bias` tensor was " + "provided; pass a [total_q, max_seqlen_kv] bias (varlen) or a " + "[seq_len, seq_len_kv] bias (dense), with the same dtype as q." + ) + bias = O + bias_stride0 = 0 + else: + bias = bias.contiguous() + if bias_stride0 is None: + bias_stride0 = bias.stride(0) + bias = bias.view(-1) + alibi_slopes, alibi_stride_b = _prep_alibi(alibi_slopes, O) + sink = _prep_sink(sink, O) with CompilationContext.compile_hints(_dualwave_swp_compile_hints): if stream is None: return launch_flash_attn_dualwave_swp( @@ -884,6 +1086,9 @@ def _launch( cu_seqlens_q, cu_seqlens_kv, block_table, + bias, + alibi_slopes, + sink, batch_size, seq_len, seq_len_kv, @@ -891,6 +1096,8 @@ def _launch( stride_kv_n, head_dim_runtime, block_table_stride, + bias_stride0, + alibi_stride_b, ) return launch_flash_attn_dualwave_swp( Q, @@ -902,6 +1109,9 @@ def _launch( cu_seqlens_q, cu_seqlens_kv, block_table, + bias, + alibi_slopes, + sink, batch_size, seq_len, seq_len_kv, @@ -909,6 +1119,8 @@ def _launch( stride_kv_n, head_dim_runtime, block_table_stride, + bias_stride0, + alibi_stride_b, stream=stream, ) @@ -930,6 +1142,10 @@ def _compile( cu_seqlens_kv=None, block_table=None, block_table_stride=None, + bias=None, + bias_stride0=None, + alibi_slopes=None, + sink=None, lse=None, stream=None, ): @@ -961,6 +1177,23 @@ def _compile( block_table = O if block_table_stride is None: block_table_stride = 0 + + if bias is None: + if HAS_BIAS: + raise ValueError( + "flash_attn_dualwave_swp was built with has_bias=True but no `bias` tensor was " + "provided; pass a [total_q, max_seqlen_kv] bias (varlen) or a " + "[seq_len, seq_len_kv] bias (dense), with the same dtype as q." + ) + bias = O + bias_stride0 = 0 + else: + bias = bias.contiguous() + if bias_stride0 is None: + bias_stride0 = bias.stride(0) + bias = bias.view(-1) + alibi_slopes, alibi_stride_b = _prep_alibi(alibi_slopes, O) + sink = _prep_sink(sink, O) with CompilationContext.compile_hints(_dualwave_swp_compile_hints): return flyc.compile( launch_flash_attn_dualwave_swp, @@ -973,6 +1206,9 @@ def _compile( cu_seqlens_q, cu_seqlens_kv, block_table, + bias, + alibi_slopes, + sink, batch_size, seq_len, seq_len_kv, @@ -980,6 +1216,8 @@ def _compile( stride_kv_n, head_dim_runtime, block_table_stride, + bias_stride0, + alibi_stride_b, fx.Stream(stream), ) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 5ebbd8207..9d83cb092 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -133,6 +133,9 @@ def _build_dense_dualwave( debug_lazy_counts: bool, enable_stagger: bool, return_lse: bool = False, + has_bias: bool = False, + has_alibi: bool = False, + has_sink: bool = False, ): """Build (and cache) the dense gfx950 DUALWAVE_SWP launcher.""" from kernels.attention.flash_attn_gfx950 import build_flash_attn_dualwave_swp_module @@ -151,6 +154,9 @@ def _build_dense_dualwave( dualwave_swp_debug_lazy_counts=debug_lazy_counts, dualwave_swp_enable_stagger=enable_stagger, return_lse=return_lse, + has_bias=has_bias, + has_alibi=has_alibi, + has_sink=has_sink, ) @@ -197,6 +203,9 @@ def _build_varlen( debug_lazy_counts: bool, enable_stagger: bool, return_lse: bool = False, + has_bias: bool = False, + has_alibi: bool = False, + has_sink: bool = False, ): """Build (and cache) a varlen-mode launcher (gfx950 DUALWAVE_SWP, varlen=True).""" from kernels.attention.flash_attn_gfx950 import build_flash_attn_dualwave_swp_module @@ -216,6 +225,9 @@ def _build_varlen( dualwave_swp_debug_lazy_counts=debug_lazy_counts, dualwave_swp_enable_stagger=enable_stagger, return_lse=return_lse, + has_bias=has_bias, + has_alibi=has_alibi, + has_sink=has_sink, ) @@ -268,6 +280,9 @@ def _build_splitk( setprio: bool, enable_stagger: bool, return_lse: bool = False, + has_bias: bool = False, + has_alibi: bool = False, + has_sink: bool = False, ): """Build (and cache) a split-K launcher (gfx950 DUALWAVE_SWP, num_kv_splits>1).""" from kernels.attention.flash_attn_gfx950 import build_flash_attn_dualwave_swp_module @@ -285,6 +300,9 @@ def _build_splitk( dualwave_swp_setprio=setprio, dualwave_swp_enable_stagger=enable_stagger, return_lse=return_lse, + has_bias=has_bias, + has_alibi=has_alibi, + has_sink=has_sink, ) @@ -610,6 +628,15 @@ def flydsl_flash_attn_func( kv_cache_layout: str = "linear", # Split-K (gfx950 only, seq_len >= 384, D=64/128, bf16/f16). num_kv_splits: int = 1, + # Additive attention bias, folded into the scores after sm_scale and before + # masking. gfx950 DUALWAVE_SWP only (dense / varlen / split-K). + bias: Optional[torch.Tensor] = None, + # Per-head ALiBi slope table, computed analytically into the scores. Same + # path and restrictions as `bias`; the two may be combined. + alibi_slopes: Optional[torch.Tensor] = None, + # Per-head attention-sink logit: one extra softmax denominator term with no + # matching V row. Same path and restrictions as `bias`; freely combinable. + sink: Optional[torch.Tensor] = None, # fp8 dense ABI: per-tensor descales for pre-quantized e4m3fn Q/K/V. q_descale: Optional[torch.Tensor] = None, k_descale: Optional[torch.Tensor] = None, @@ -661,6 +688,32 @@ def flydsl_flash_attn_func( dense mode infers it from ``q.shape[1] != k.shape[1]``. block_table / seqlen_k: vLLM-style 2D block table metadata. num_kv_splits: Split-K factor (>1: gfx950 only, D=64/128, bf16/f16, seq>=384). + bias: Additive attention bias with the same dtype as q, folded in as + ``softmax(q @ k^T * sm_scale + bias)`` -- after the scale, before the + causal/padding mask. Dense: ``[Sq, Skv]``, broadcast over batch and + head. Varlen: ``[total_q, max_seqlen_kv]``, where the row is the + *global* packed q token index and the column is the *per-batch-local* + key index, broadcast over head. Routes to the gfx950 DUALWAVE_SWP + kernel; paged KV and fp8 raise NotImplementedError rather than + silently dropping the bias. + alibi_slopes: fp32 ALiBi slope table, ``[H]`` (broadcast over batch) or + ``[B, H]``, values positive. Adds + ``-slope * |i + seqlen_kv - seqlen_q - j|`` to the scores after the + 1/sqrt(D) scaling (the slope is not divided by it), bottom-right + aligned like the causal mask. Positions are measured *within* the + sequence, so varlen does not offset by the packed-token base. Same + kernel path and restrictions as ``bias``; the two may be combined. + sink: fp32 ``[H]`` per-head attention-sink logit -- one extra softmax + denominator term that has no matching V row:: + + O = sum_j exp(s_j - m) v_j / (exp(sink - m) + sum_j exp(s_j - m)) + + Consumed verbatim (no host-side scaling), so it lives in the same + post-sm_scale logit space as the scores. Applied in the epilogue, so + it touches no score element; under split-K the per-split partials + stay sink-free and the combine pass folds it in exactly once. Same + kernel path and restrictions as ``bias``; freely combinable with it + and with ``alibi_slopes``. q_descale / k_descale / v_descale: fp32 shape-[1] descales required for dense fp8 e4m3fn inputs. out: Optional pre-allocated output tensor. For fp8, output is bf16; @@ -697,6 +750,33 @@ def flydsl_flash_attn_func( raise NotImplementedError("flydsl_flash_attn_func: fp8 flash_attn does not support paged KV") if return_lse and paged_kv: raise NotImplementedError("flydsl_flash_attn_func: return_lse is not supported for paged KV") + has_bias = bias is not None + has_alibi = alibi_slopes is not None + has_sink = sink is not None + for _name, _t in (("bias", bias), ("alibi_slopes", alibi_slopes), ("sink", sink)): + if _t is None: + continue + if paged_kv: + raise NotImplementedError(f"flydsl_flash_attn_func: {_name} is not supported for paged KV") + if dtype_str == "fp8": + raise NotImplementedError(f"flydsl_flash_attn_func: {_name} is not supported for fp8") + if not _t.is_cuda or _t.device != q.device: + raise ValueError(f"flydsl_flash_attn_func: {_name} must be a CUDA tensor on {q.device}, got {_t.device}") + if has_bias: + if bias.dtype != q.dtype: + raise ValueError(f"flydsl_flash_attn_func: bias dtype must match q dtype {q.dtype}, got {bias.dtype}") + if bias.dim() != 2: + raise ValueError(f"flydsl_flash_attn_func: bias must be 2D, got {bias.dim()}D") + if has_alibi: + if alibi_slopes.dtype != torch.float32: + raise ValueError(f"flydsl_flash_attn_func: alibi_slopes must be float32, got {alibi_slopes.dtype}") + if alibi_slopes.dim() not in (1, 2): + raise ValueError(f"flydsl_flash_attn_func: alibi_slopes must be [H] or [B, H], got {alibi_slopes.dim()}D") + if has_sink: + if sink.dtype != torch.float32: + raise ValueError(f"flydsl_flash_attn_func: sink must be float32, got {sink.dtype}") + if sink.dim() != 1: + raise ValueError(f"flydsl_flash_attn_func: sink must be 1D [H], got {sink.dim()}D") if paged_kv: return _flydsl_flash_attn_paged( q, @@ -778,6 +858,37 @@ def flydsl_flash_attn_func( if D < 64 or D % 32 != 0: raise ValueError(f"flydsl_flash_attn_func: head_dim ({D}) must be >= 64 and a multiple of 32") + if has_bias: + # Bias rows are indexed by q token, columns by the per-batch-local key. + if varlen: + if bias.shape[0] != q.shape[0]: + raise ValueError( + f"flydsl_flash_attn_func: varlen bias must be [total_q, max_seqlen_kv] with " + f"total_q={q.shape[0]}, got {tuple(bias.shape)}" + ) + if max_seqlen_kv is not None and bias.shape[1] < int(max_seqlen_kv): + raise ValueError( + f"flydsl_flash_attn_func: varlen bias needs >= max_seqlen_kv={int(max_seqlen_kv)} " + f"columns, got {bias.shape[1]}" + ) + elif tuple(bias.shape) != (Sq, Skv): + raise ValueError(f"flydsl_flash_attn_func: dense bias must be [{Sq}, {Skv}], got {tuple(bias.shape)}") + + if has_alibi: + if alibi_slopes.shape[-1] != H: + raise ValueError( + f"flydsl_flash_attn_func: alibi_slopes last dim must be num_heads={H}, " + f"got {tuple(alibi_slopes.shape)}" + ) + if alibi_slopes.dim() == 2 and alibi_slopes.shape[0] != B: + raise ValueError( + f"flydsl_flash_attn_func: 2D alibi_slopes must be [batch={B}, num_heads={H}], " + f"got {tuple(alibi_slopes.shape)}" + ) + + if has_sink and sink.shape[0] != H: + raise ValueError(f"flydsl_flash_attn_func: sink must be [num_heads={H}], got {tuple(sink.shape)}") + splitk = num_kv_splits > 1 # ── split-K eligibility guard (SKIP analogous to run_splitk_config) ──── @@ -809,12 +920,19 @@ def flydsl_flash_attn_func( setprio=dualwave_swp_setprio, enable_stagger=dualwave_swp_enable_stagger, return_lse=return_lse, + has_bias=has_bias, + has_alibi=has_alibi, + has_sink=has_sink, ) elif varlen: # Short varlen attention uses generic light; long/debug stays on dualwave. _arch = _gpu_arch(q.device) _prefer_light = ( (not debug_lazy) + # The light (generic) kernel folds in neither bias nor ALiBi. + and (not has_bias) + and (not has_alibi) + and (not has_sink) and D in (64, 128) and dtype_str in ("bf16", "f16") and (not _arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) @@ -850,6 +968,9 @@ def flydsl_flash_attn_func( debug_lazy_counts=debug_lazy, enable_stagger=dualwave_swp_enable_stagger, return_lse=return_lse, + has_bias=has_bias, + has_alibi=has_alibi, + has_sink=has_sink, ) else: _arch = _gpu_arch(q.device) @@ -872,7 +993,20 @@ def flydsl_flash_attn_func( raise NotImplementedError( "flydsl_flash_attn_func: debug_counts requires the gfx950 DUALWAVE_SWP path" ) - if debug_lazy or (can_dualwave and _dense_routes_to_dualwave(B, Sq)): + if (has_bias or has_alibi or has_sink) and not can_dualwave: + _term = "bias" if has_bias else ("alibi_slopes" if has_alibi else "sink") + raise NotImplementedError( + f"flydsl_flash_attn_func: {_term} requires the gfx950 DUALWAVE_SWP path " + f"(D=64/128, bf16/f16, gfx950); got D={D}, dtype={dtype_str}, arch='{_arch or 'unknown'}'" + ) + # bias/ALiBi force dualwave: the generic dense kernel folds in neither. + if ( + debug_lazy + or has_bias + or has_alibi + or has_sink + or (can_dualwave and _dense_routes_to_dualwave(B, Sq)) + ): exe = _build_dense_dualwave( num_heads=H, num_kv_heads=num_kv_heads, @@ -887,6 +1021,9 @@ def flydsl_flash_attn_func( debug_lazy_counts=debug_lazy, enable_stagger=dualwave_swp_enable_stagger, return_lse=return_lse, + has_bias=has_bias, + has_alibi=has_alibi, + has_sink=has_sink, ) else: block_m, flat_work_group_size, path_tag = _dense_generic_tile(B, Sq, H, D, dtype_str, q.device) @@ -932,11 +1069,20 @@ def flydsl_flash_attn_func( lse = torch.empty((B, H, Sq), dtype=torch.float32, device=q.device) if return_lse else None # ── launch ────────────────────────────────────────────────────────── + _bias_kw = {} + if has_bias: + _bias_kw["bias"] = bias + if has_alibi: + _bias_kw["alibi_slopes"] = alibi_slopes + if has_sink: + _bias_kw["sink"] = sink if splitk: _ws = torch.empty(ws_elems, dtype=torch.float32, device=q.device) - exe(q_flat, k_flat, v_flat, o_flat, B, Sq, workspace=_ws, lse=lse, stream=launch_stream) + exe(q_flat, k_flat, v_flat, o_flat, B, Sq, workspace=_ws, lse=lse, stream=launch_stream, **_bias_kw) elif varlen: - kwargs = dict(cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, lse=lse, stream=launch_stream) + kwargs = dict( + cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, lse=lse, stream=launch_stream, **_bias_kw + ) if cross: kwargs["seq_len_kv"] = int(max_seqlen_kv) if debug_lazy: @@ -944,7 +1090,7 @@ def flydsl_flash_attn_func( else: exe(q_flat, k_flat, v_flat, o_flat, B, Sq, **kwargs) else: - kwargs: dict = dict(stream=launch_stream) + kwargs: dict = dict(stream=launch_stream, **_bias_kw) # fp8 has no LSE path (guarded above) and its launcher takes no `lse` arg. if dtype_str != "fp8": kwargs["lse"] = lse diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 982e8876d..30458a511 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -39,6 +39,11 @@ UNIFORM_RANGE = (-1, 1) DEFAULT_SEED = 123 PAGED_KV_MIN_CONTEXT_LENGTH = 16384 +# Target share of the softmax mass placed on the attention sink (see calibrate_sink). +DEFAULT_SINK_SHARE = 0.5 +# Attention-bias addressing limits (see bias_fits). +BIAS_BUFFER_MAX_BYTES = 0xFFFFFFFF # buffer descriptor num_records clamp +BIAS_I32_MAX_ELEMS = 2**31 - 1 # kernel computes bias element offsets in i32 # fp8 correctness gate (fixed; fp8 is lossy). FP8_MAX_ERR = 5e-2 FP8_MIN_COS = 0.98 @@ -187,7 +192,104 @@ def setup_seed(seed: int) -> None: torch.backends.cudnn.deterministic = True -def pytorch_ref_attention(q, k, v, causal=True): +def make_alibi_slopes(batch, num_heads, two_d=False, device="cuda"): + """Positive fp32 ALiBi slopes, [batch, num_heads] if two_d else [num_heads]. + + The canonical geometric ladder 2**(-8*(h+1)/H) with a per-batch jitter on the + 2D form, so a bug that broadcast or swapped the head/batch index cannot alias + into a pass. Deterministic: consumes no RNG, so Q/K/V stay bit-identical to a + run without ALiBi. + """ + base = torch.tensor( + [2.0 ** (-((h + 1) * 8.0 / num_heads)) for h in range(num_heads)], dtype=torch.float32, device=device + ) + if not two_d: + return base.contiguous() + jitter = 1.0 + 0.25 * torch.arange(batch, dtype=torch.float32, device=device)[:, None] + return (base[None, :] * jitter).contiguous() + + +def _alibi_term(alibi_slopes, q_lo, q_hi, delta, key_idx): + """-slope * |i + delta - j| for q rows [q_lo, q_hi) -> [B or 1, H, q_hi-q_lo, Skv]. + + Built per Q chunk rather than materialized whole: a full [B, H, Sq, Skv] fp32 + ALiBi matrix is the same size as the score matrix the caller is already + chunking to avoid. + """ + i = torch.arange(q_lo, q_hi, device=alibi_slopes.device)[:, None] + rel = (i + delta - key_idx.view(1, -1)).abs().to(torch.float32) + s = alibi_slopes.float() + if s.dim() == 1: + s = s.unsqueeze(0) + return -s[:, :, None, None] * rel + + +def _rows_logsumexp(q_t, k_t, causal, bias=None, alibi_slopes=None): + """Per-head sum and count of row logsumexp(scores), chunked over Q. + + q_t: [B, Sq, H, D], k_t: [B, Skv, Hkv, D]. Returns (sum_per_head, count), + both fp32 [H] / scalar, over finite rows only. + """ + q_h = q_t.transpose(1, 2).float() + k_h = k_t.transpose(1, 2).float() + B, H, Sq, D = q_h.shape + Skv = k_h.shape[2] + if H != k_h.shape[1]: + k_h = k_h.repeat_interleave(H // k_h.shape[1], dim=1) + delta = Skv - Sq + scale = 1.0 / math.sqrt(D) + k_trans = k_h.transpose(-1, -2).contiguous() + key_idx = torch.arange(Skv, device=q_t.device).view(1, 1, 1, Skv) + chunk = max(1, min(Sq, (64 * 1024 * 1024) // max(B * H * Skv, 1))) + total = torch.zeros(H, dtype=torch.float32, device=q_t.device) + count = 0 + for s0 in range(0, Sq, chunk): + s1 = min(s0 + chunk, Sq) + sc = torch.matmul(q_h[:, :, s0:s1, :], k_trans) * scale + if alibi_slopes is not None: + sc = sc + _alibi_term(alibi_slopes, s0, s1, delta, key_idx) + if bias is not None: + sc = sc + bias[s0:s1].float() + if causal: + q_idx = torch.arange(s0, s1, device=q_t.device).view(1, 1, -1, 1) + sc = sc.masked_fill(key_idx > q_idx + delta, float("-inf")) + lse = torch.logsumexp(sc, dim=-1) # [B, H, chunk] + finite = torch.isfinite(lse) + total += torch.where(finite, lse, torch.zeros_like(lse)).sum(dim=(0, 2)) + count += int(finite[:, 0, :].sum().item()) if finite.numel() else 0 + return total, max(count, 1) + + +def calibrate_sink(sum_lse, count, share): + """Per-head sink logit placing `share` of the softmax mass on the sink. + + share = sigmoid(sink - logsumexp(scores)), so invert it per head. Calibration + matters: logsumexp grows like ln(seqlen), so a sink drawn near 0 would own a + fraction of a percent of the mass -- below the bf16 noise floor, and a row + with a dropped sink would pass just as happily. + """ + return (sum_lse / count + math.log(share / (1.0 - share))).float().contiguous() + + +def _sink_softmax(scores, sink): + """softmax over [scores, sink], where the sink carries no value row. + + Returns probs summing to 1 - sink_share. A fully-masked row needs no + special-casing: the max collapses to the sink, every score term is + exp(-inf) = 0, and the row becomes all-sink -- probs 0, matching the kernel. + """ + s = sink.view(1, -1, 1, 1) + m = torch.maximum(scores.amax(dim=-1, keepdim=True), s) + e = torch.exp(scores - m) + return e / (e.sum(dim=-1, keepdim=True) + torch.exp(s - m)) + + +def pytorch_ref_attention(q, k, v, causal=True, bias=None, alibi_slopes=None, sink=None): + if bias is not None or alibi_slopes is not None or sink is not None: + # These must land after the sm_scale multiply and before the mask, and SDPA + # rejects an additive attn_mask together with is_causal, so defer to the + # explicit chunked path (identical result when Sq == Skv). + return pytorch_ref_attention_qkv_diff(q, k, v, causal=causal, bias=bias, alibi_slopes=alibi_slopes, sink=sink) q_t = q.transpose(1, 2).float() k_t = k.transpose(1, 2).float() v_t = v.transpose(1, 2).float() @@ -229,13 +331,7 @@ def pytorch_ref_attention_chunked(q_t, k_t, v_t, causal=True): @torch.no_grad() -def pytorch_ref_attention_qkv_diff(q, k, v, causal=True): - """Reference for seqlen_q != seqlen_kv with a BOTTOM-RIGHT aligned causal mask. - - q: [B,Sq,H,D]; k,v: [B,Skv,Hkv,D]. Row r keeps keys [0, r+delta] with - delta = Skv - Sq (so the mask hugs the bottom-right corner); an all-masked - row outputs 0. Chunked over Q to bound the score matrix memory. - """ +def pytorch_ref_attention_qkv_diff(q, k, v, causal=True, bias=None, alibi_slopes=None, sink=None): q_t = q.transpose(1, 2).float() k_t = k.transpose(1, 2).float() v_t = v.transpose(1, 2).float() @@ -256,11 +352,19 @@ def pytorch_ref_attention_qkv_diff(q, k, v, causal=True): for s0 in range(0, Sq, chunk): s1 = min(s0 + chunk, Sq) scores = torch.matmul(q_t[:, :, s0:s1, :], k_trans) * scale + if alibi_slopes is not None: + scores = scores + _alibi_term(alibi_slopes, s0, s1, delta, key_idx) + if bias is not None: + scores = scores + bias[s0:s1].float() if causal: q_idx = torch.arange(s0, s1, device=q_t.device).view(1, 1, -1, 1) scores = scores.masked_fill(key_idx > q_idx + delta, float("-inf")) - probs = torch.softmax(scores, dim=-1) - probs = torch.nan_to_num(probs, nan=0.0) # all-masked row -> 0 output + if sink is not None: + # The sink also fixes the all-masked row: it becomes all-sink, probs 0. + probs = _sink_softmax(scores, sink) + else: + probs = torch.softmax(scores, dim=-1) + probs = torch.nan_to_num(probs, nan=0.0) # all-masked row -> 0 output out[:, :, s0:s1, :] = torch.matmul(probs, v_t) return out.transpose(1, 2) @@ -269,6 +373,15 @@ def _ceil_div(a, b): return (a + b - 1) // b +def bias_fits(rows, cols, elem_size=2): + elems = rows * cols + if elems > BIAS_I32_MAX_ELEMS: + return False, f"bias {rows}x{cols} = {elems:.3g} elems exceeds i32 offset range" + if elems * elem_size > BIAS_BUFFER_MAX_BYTES: + return False, f"bias {elems * elem_size / 2**30:.1f} GiB exceeds 4 GiB buffer limit" + return True, "" + + def _block_table_from_indices(kv_indptr_cpu, kv_indices_cpu, batch_size, max_num_pages_per_seq): block_table_cpu = torch.zeros((batch_size, max_num_pages_per_seq), dtype=torch.int32) for b in range(batch_size): @@ -527,6 +640,12 @@ def _build_attn_inputs_for_config( page_size, kv_cache_layout, trigger_lazy_else, + use_bias=False, + use_alibi=False, + alibi_two_d=False, + use_sink=False, + sink_share=DEFAULT_SINK_SHARE, + causal=False, ): device = "cuda" H, D, H_KV = num_heads, head_dim, num_kv_heads @@ -562,8 +681,38 @@ def _build_attn_inputs_for_config( kv_cache = None k_t = torch.empty(total_kv, H_KV, D, dtype=dtype, device=device).uniform_(*UNIFORM_RANGE) v_t = torch.empty(total_kv, H_KV, D, dtype=dtype, device=device).uniform_(*UNIFORM_RANGE) + # Packed bias: row = global packed q token, column = per-batch-local key. + # Drawn after Q/K/V so those stay bit-identical to a no-bias run. + bias = ( + torch.empty(total_q, max(vl_kv), dtype=dtype, device=device).uniform_(*UNIFORM_RANGE) if use_bias else None + ) + alibi_slopes = make_alibi_slopes(B, H, alibi_two_d, device) if use_alibi else None + sink = None + if use_sink: + # One [H] table shared by every sequence, so calibrate over all their + # rows together -- matching how the kernel consumes it. + tot_lse = torch.zeros(H, dtype=torch.float32, device=device) + tot_n = 0 + for b in range(B): + sl, n = _rows_logsumexp( + q_t[cuq[b] : cuq[b + 1]].unsqueeze(0), + k_t[cukv[b] : cukv[b + 1]].unsqueeze(0), + causal, + bias=bias[cuq[b] : cuq[b + 1], : vl_kv[b]] if bias is not None else None, + alibi_slopes=( + (alibi_slopes[b] if alibi_slopes.dim() == 2 else alibi_slopes) + if alibi_slopes is not None + else None + ), + ) + tot_lse += sl + tot_n += n + sink = calibrate_sink(tot_lse, tot_n, sink_share) return { "varlen": True, + "sink": sink, + "bias": bias, + "alibi_slopes": alibi_slopes, "B": B, "Sq": Sq, "Skv": None, @@ -604,6 +753,16 @@ def _build_attn_inputs_for_config( k_t = torch.empty(B, Skv, H_KV, D, dtype=dtype, device=device).uniform_(*UNIFORM_RANGE) v_t = torch.empty(B, Skv, H_KV, D, dtype=dtype, device=device).uniform_(*UNIFORM_RANGE) + # Dense bias: (Sq, Skv), broadcast over batch and head. Drawn after Q/K/V so + # those stay bit-identical to a no-bias run. + bias = torch.empty(Sq, Skv, dtype=dtype, device=device).uniform_(*UNIFORM_RANGE) if use_bias else None + alibi_slopes = make_alibi_slopes(B, H, alibi_two_d, device) if use_alibi else None + sink = ( + calibrate_sink(*_rows_logsumexp(q_t, k_t, causal, bias=bias, alibi_slopes=alibi_slopes), sink_share) + if use_sink + else None + ) + if trigger_lazy_else: q_t.fill_(1.0) k_t.zero_() @@ -616,6 +775,9 @@ def _build_attn_inputs_for_config( return { "varlen": False, + "sink": sink, + "bias": bias, + "alibi_slopes": alibi_slopes, "B": B, "Sq": Sq, "Skv": Skv, @@ -639,6 +801,9 @@ def _build_attn_inputs_for_config( def _compute_reference_from_inputs(inputs, num_heads, head_dim, dtype, causal, seqlen_q, seqlen_kv): H, D = num_heads, head_dim q_t, k_t, v_t = inputs["q_t"], inputs["k_t"], inputs["v_t"] + bias = inputs["bias"] + alibi = inputs["alibi_slopes"] + sink = inputs["sink"] if inputs["varlen"]: ref_t = torch.empty(inputs["total_q"], H, D, dtype=dtype, device=q_t.device) @@ -648,20 +813,25 @@ def _compute_reference_from_inputs(inputs, num_heads, head_dim, dtype, causal, s qb = q_t[cuq[b] : cuq[b + 1]].unsqueeze(0).float() kb = k_t[cukv[b] : cukv[b + 1]].unsqueeze(0).float() vb = v_t[cukv[b] : cukv[b + 1]].unsqueeze(0).float() + bias_b = bias[cuq[b] : cuq[b + 1], : vl_kv[b]] if bias is not None else None + alibi_b = (alibi[b] if alibi.dim() == 2 else alibi) if alibi is not None else None ref_fn = pytorch_ref_attention if vl_q[b] == vl_kv[b] else pytorch_ref_attention_qkv_diff - ref_t[cuq[b] : cuq[b + 1]] = ref_fn(qb, kb, vb, causal=causal).to(dtype).squeeze(0) + ref_t[cuq[b] : cuq[b + 1]] = ( + ref_fn(qb, kb, vb, causal=causal, bias=bias_b, alibi_slopes=alibi_b, sink=sink).to(dtype).squeeze(0) + ) return ref_t self_attn = seqlen_kv is None or seqlen_kv == seqlen_q - if self_attn: - return pytorch_ref_attention(q_t.float(), k_t.float(), v_t.float(), causal=causal).to(dtype) - return pytorch_ref_attention_qkv_diff(q_t.float(), k_t.float(), v_t.float(), causal=causal).to(dtype) + ref_fn = pytorch_ref_attention if self_attn else pytorch_ref_attention_qkv_diff + return ref_fn(q_t.float(), k_t.float(), v_t.float(), causal=causal, bias=bias, alibi_slopes=alibi, sink=sink).to( + dtype + ) def _build_inputs_and_reference_for_config(**kwargs): setup_seed(kwargs.pop("seed")) causal = kwargs.pop("causal") - inputs = _build_attn_inputs_for_config(**kwargs) + inputs = _build_attn_inputs_for_config(causal=causal, **kwargs) ref_t = _compute_reference_from_inputs( inputs, kwargs["num_heads"], @@ -871,6 +1041,11 @@ def run_attn_config( use_block_table=False, page_size=64, kv_cache_layout="linear", + use_bias=False, + use_alibi=False, + alibi_two_d=False, + use_sink=False, + sink_share=DEFAULT_SINK_SHARE, ): """Unified flash-attention test/bench function. @@ -915,6 +1090,27 @@ def run_attn_config( if D not in (64, 128) or dtype_str not in ("bf16", "f16") or (seqlen_q is not None and seqlen_q < 384): return {"skip": True} + # ── bias addressing guard ──────────────────────────────────────────────── + if use_bias: + if varlen: + vl_q_cfg = list(varlen_seqlens_q) + vl_kv_cfg = list(varlen_seqlens_kv) if varlen_seqlens_kv is not None else vl_q_cfg + bias_rows, bias_cols = sum(vl_q_cfg), max(vl_kv_cfg) + if max(vl_q_cfg) > vl_q_cfg[-1]: + return { + "skip": True, + "skip_reason": ( + f"varlen bias reads OOB when the last seqlen ({vl_q_cfg[-1]}) " + f"is below max_seqlen_q ({max(vl_q_cfg)})" + ), + } + else: + bias_rows = seqlen_q + bias_cols = seqlen_kv if seqlen_kv is not None else seqlen_q + fits, why = bias_fits(bias_rows, bias_cols, torch.empty((), dtype=dtype).element_size()) + if not fits: + return {"skip": True, "skip_reason": why} + if use_block_table and (precomputed_inputs is None or precomputed_ref is None): return {"err": "block-table tests require precomputed_inputs and precomputed_ref"} @@ -934,6 +1130,12 @@ def run_attn_config( page_size=page_size, kv_cache_layout=kv_cache_layout, trigger_lazy_else=trigger_lazy_else, + use_bias=use_bias, + use_alibi=use_alibi, + alibi_two_d=alibi_two_d, + use_sink=use_sink, + sink_share=sink_share, + causal=causal, ) varlen = precomputed_inputs["varlen"] @@ -942,9 +1144,6 @@ def run_attn_config( Skv = precomputed_inputs["Skv"] vl_q = precomputed_inputs["vl_q"] vl_kv = precomputed_inputs["vl_kv"] - cuq = precomputed_inputs["cuq"] - cukv = precomputed_inputs["cukv"] - total_q = precomputed_inputs["total_q"] cu_q_t = precomputed_inputs["cu_q_t"] cu_kv_t = precomputed_inputs["cu_kv_t"] q_t = precomputed_inputs["q_t"] @@ -953,6 +1152,9 @@ def run_attn_config( cross = precomputed_inputs["cross"] max_seqlen_kv = precomputed_inputs["max_seqlen_kv"] kv_cache = precomputed_inputs["kv_cache"] + bias_t = precomputed_inputs["bias"] + alibi_t = precomputed_inputs["alibi_slopes"] + sink_t = precomputed_inputs["sink"] debug_counts = torch.zeros(2, dtype=torch.float32, device=device) if debug_lazy else None o_t = torch.zeros_like(q_t) @@ -993,6 +1195,9 @@ def run_attn_config( max_seqlen_kv=max_seqlen_kv if varlen else None, cross_seqlen=cross if varlen else None, num_kv_splits=int(num_kv_splits), + bias=bias_t, + alibi_slopes=alibi_t, + sink=sink_t, out=o_t, debug_counts=debug_counts, **_cfg_kw(), @@ -1017,21 +1222,12 @@ def run_attn_config( # ── reference ─────────────────────────────────────────────────────────── # precomputed_ref makes FlyDSL/aiter_ck/aiter_asm share one reference tensor. # Otherwise compute the cheapest reference path for the active mode. - _self_attn = not varlen and (seqlen_kv is None or seqlen_kv == seqlen_q) + # Delegated rather than inlined so the bias enters the reference in exactly one + # place; two copies of this dispatch would be free to drift apart. if precomputed_ref is not None: ref_t = precomputed_ref - elif varlen: - ref_t = torch.empty(total_q, H, D, dtype=dtype, device=device) - for b in range(B): - qb = q_t[cuq[b] : cuq[b + 1]].unsqueeze(0).float() - kb = k_t[cukv[b] : cukv[b + 1]].unsqueeze(0).float() - vb = v_t[cukv[b] : cukv[b + 1]].unsqueeze(0).float() - ref_fn = pytorch_ref_attention if vl_q[b] == vl_kv[b] else pytorch_ref_attention_qkv_diff - ref_t[cuq[b] : cuq[b + 1]] = ref_fn(qb, kb, vb, causal=causal).to(dtype).squeeze(0) - elif _self_attn: - ref_t = pytorch_ref_attention(q_t.float(), k_t.float(), v_t.float(), causal=causal).to(dtype) else: - ref_t = pytorch_ref_attention_qkv_diff(q_t.float(), k_t.float(), v_t.float(), causal=causal).to(dtype) + ref_t = _compute_reference_from_inputs(precomputed_inputs, H, D, dtype, causal, seqlen_q, seqlen_kv) o_f32 = o_t.contiguous().reshape(-1).float() ref_f32 = ref_t.contiguous().reshape(-1).float() @@ -1115,6 +1311,9 @@ def kernel_fn(): max_seqlen_kv=max_seqlen_kv if varlen else None, cross_seqlen=cross if varlen else None, num_kv_splits=int(num_kv_splits), + bias=bias_t, + alibi_slopes=alibi_t, + sink=sink_t, out=o_t, debug_counts=debug_counts, **_cfg_kw(), @@ -1156,6 +1355,9 @@ def run_aiter_bench( seqlen_kv=None, varlen_seqlens_q=None, varlen_seqlens_kv=None, + use_bias=False, + use_alibi=False, + use_sink=False, ): """Run true aiter_ck or true aiter_asm kernel via aiter and return {tflops, max_err, us}.""" try: @@ -1170,6 +1372,15 @@ def run_aiter_bench( return {"skip": True} if backend == "asm" and (varlen or (seqlen_kv is not None and seqlen_kv != seq_len)): return {"skip": True} + bias = precomputed_inputs["bias"] if precomputed_inputs is not None else None + if use_bias and (backend == "asm" or varlen or bias is None): + return {"skip": True} + alibi = precomputed_inputs["alibi_slopes"] if precomputed_inputs is not None else None + if use_alibi and (backend == "asm" or causal or use_bias or alibi is None): + return {"skip": True} + sink = precomputed_inputs["sink"] if precomputed_inputs is not None else None + if use_sink and (backend == "asm" or varlen or sink is None): + return {"skip": True} results = {} torch.cuda.empty_cache() @@ -1248,14 +1459,15 @@ def aiter_forward(): causal, # is_causal -1, # window_size_left -1, # window_size_right - 0, # sink_size + 1 if use_sink else 0, # sink_size True, # return_softmax_lse False, # return_dropout_randval + sink_ptr=sink if use_sink else None, cu_seqlens_q=cu_q_t, cu_seqlens_kv=cu_kv_t, out=None, - bias=None, - alibi_slopes=None, + bias=bias, + alibi_slopes=alibi, q_descale=None, k_descale=None, v_descale=None, @@ -2368,6 +2580,49 @@ def main(): action="store_true", help="Run additional varlen/cross-length configs from EXTRA_CONFIGS", ) + parser.add_argument( + "--bias", + action="store_true", + help="Add an additive attention bias to the scores: softmax(q@k^T * sm_scale + bias). " + "Dense bias is [Sq, Skv] broadcast over batch and head; varlen bias is packed " + "[total_q, max_seqlen_kv] with global q rows and batch-local key columns. " + "gfx950 bf16/f16 D=64/128 only; incompatible with --block-table and fp8. Rows whose " + "bias exceeds the i32 offset / 4 GiB buffer limits are SKIPped.", + ) + parser.add_argument( + "--alibi", + action="store_true", + help="Add a per-head ALiBi positional bias: score += -slope * |i + seqlen_kv - seqlen_q - j|, " + "applied after the 1/sqrt(D) scaling and bottom-right aligned like the causal mask. " + "Slopes are the canonical 2**(-8*(h+1)/H) ladder. gfx950 bf16/f16 D=64/128 only; " + "incompatible with --block-table and fp8. Combines with --bias.", + ) + parser.add_argument( + "--sink", + action="store_true", + help="Add a per-head attention sink: one extra softmax denominator logit with no matching V " + "row, O = sum_j exp(s_j-m) v_j / (exp(sink-m) + sum_j exp(s_j-m)). The sink is calibrated per " + "run to --sink-share of the softmax mass; an uncalibrated sink near 0 would sit below the bf16 " + "noise floor and pass even if dropped. gfx950 bf16/f16 D=64/128 only; incompatible with " + "--block-table and fp8. Combines with --bias and --alibi. Under --compare the aiter_ck " + "column runs a real sink baseline (mha_fwd sink_size/sink_ptr); aiter_asm has no sink " + "parameter and is SKIPped, as is the varlen path.", + ) + parser.add_argument( + "--sink-share", + type=float, + default=DEFAULT_SINK_SHARE, + dest="sink_share", + help=f"Fraction of the softmax mass the sink should take, in (0, 1). Default {DEFAULT_SINK_SHARE}. " + "Values in 0.25-0.95 keep the sink well above bf16 noise. Requires --sink.", + ) + parser.add_argument( + "--alibi-2d", + action="store_true", + dest="alibi_two_d", + help="Use a per-(batch, head) [B, H] slope table instead of the [H] form, exercising the " + "kernel's alibi_stride_b path. Requires --alibi.", + ) parser.add_argument( "--verbose", action="store_true", @@ -2444,6 +2699,26 @@ def main(): args = parser.parse_args() if not args.block_table and args.kv_cache_layout != "linear": parser.error("--kv-cache-layout requires --block-table") + # Paged KV and fp8 have no bias support in the kernel; reject rather than run + # a bias-free kernel against a biased reference. + if args.bias and args.block_table: + parser.error("--bias is not supported with --block-table (paged KV)") + if args.bias and args.dtype == "fp8": + parser.error("--bias is not supported with --dtype fp8") + if args.alibi and args.block_table: + parser.error("--alibi is not supported with --block-table (paged KV)") + if args.alibi and args.dtype == "fp8": + parser.error("--alibi is not supported with --dtype fp8") + if args.alibi_two_d and not args.alibi: + parser.error("--alibi-2d requires --alibi") + if args.sink and args.block_table: + parser.error("--sink is not supported with --block-table (paged KV)") + if args.sink and args.dtype == "fp8": + parser.error("--sink is not supported with --dtype fp8") + if args.sink_share != DEFAULT_SINK_SHARE and not args.sink: + parser.error("--sink-share requires --sink") + if not 0.0 < args.sink_share < 1.0: + parser.error(f"--sink-share must be in (0, 1), got {args.sink_share}") # Build kernel config from parsed args (no env-var reads). FLASH_ATTN_FUNC_KERNEL_CONFIG.update( @@ -2482,6 +2757,10 @@ def main(): causal_desc = {True: "causal", False: "non-causal", None: "causal+non-causal"}[args.causal] dtype_desc = args.dtype or "bf16+fp16" + _terms = [t for t, on in (("bias", args.bias), ("alibi", args.alibi), ("sink", args.sink)) if on] + bias_desc = ("; " + "+".join(_terms)) if _terms else "" + # Keep biased and unbiased baselines in separate CSVs so they stay diffable. + csv_tag = ("_" + "".join(_terms)) if _terms else "" extra_cases = ( [_extra_case_from_config(row) for row in EXTRA_CONFIGS] if args.extra and configs is DEFAULT_CONFIGS else [] ) @@ -2501,7 +2780,7 @@ def main(): if args.compare: # ---- Comparison mode: FlyDSL vs aiter_ck vs aiter_asm ---- print("=" * 130) - print(f"FlyDSL vs aiter_ck vs aiter_asm ({causal_desc}, {dtype_desc})") + print(f"FlyDSL vs aiter_ck vs aiter_asm ({causal_desc}, {dtype_desc}{bias_desc})") print(f"GPU: {torch.cuda.get_device_name(0)}") if args.num_kv_splits > 1: print( @@ -2566,10 +2845,9 @@ def main(): continue else: shared_ref = None - # Compute reference once for bf16/fp16 rows (fp8 helper owns quantization + reference). if dtype_str == "fp8": pass - elif args.trigger_lazy_else: + elif args.trigger_lazy_else or args.bias or args.alibi or args.sink: precomputed_inputs, shared_ref = _build_inputs_and_reference_for_config( batch=batch, seqlen_q=seq_len, @@ -2585,7 +2863,12 @@ def main(): use_block_table=False, page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", - trigger_lazy_else=True, + trigger_lazy_else=args.trigger_lazy_else, + use_bias=args.bias, + use_alibi=args.alibi, + use_sink=args.sink, + sink_share=args.sink_share, + alibi_two_d=args.alibi_two_d, ) else: # All three use the same seed -> same Q/K/V -> identical reference. @@ -2643,6 +2926,11 @@ def main(): use_block_table=args.block_table, page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", + use_bias=args.bias, + use_alibi=args.alibi, + use_sink=args.sink, + sink_share=args.sink_share, + alibi_two_d=args.alibi_two_d, ) except Exception as _fly_err: print(f" [FlyDSL unsupported] {_fmt_cfg(cfg)}: {_fly_err}", flush=True) @@ -2700,6 +2988,9 @@ def main(): num_kv_heads=nh_kv, precomputed_ref=shared_ref, precomputed_inputs=precomputed_inputs, + use_bias=args.bias, + use_alibi=args.alibi, + use_sink=args.sink, ) asm_r = run_aiter_bench( batch, @@ -2715,6 +3006,9 @@ def main(): num_kv_heads=nh_kv, precomputed_ref=shared_ref, precomputed_inputs=precomputed_inputs, + use_bias=args.bias, + use_alibi=args.alibi, + use_sink=args.sink, ) rows.append((cfg, fly_r, ck_r, asm_r)) @@ -2769,7 +3063,7 @@ def _cmp_avg(label, subset): _print_grouped_avgs(rows, lambda r: _tag_group(r[0]), _cmp_avg) print("=" * len(hdr2)) - csv_path = f"fmha_perf_compare_{_gpu_short_name()}.csv" + csv_path = f"fmha_perf_compare{csv_tag}_{_gpu_short_name()}.csv" _write_cmp_csv(csv_path, rows, cmp_avg_rows) print(f"Results saved to: {csv_path}") @@ -2856,6 +3150,11 @@ def _cmp_avg(label, subset): use_block_table=args.block_table, page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", + use_bias=args.bias, + use_alibi=args.alibi, + use_sink=args.sink, + sink_share=args.sink_share, + alibi_two_d=args.alibi_two_d, **kwargs, ) except Exception as _fly_err: @@ -2892,6 +3191,9 @@ def _cmp_avg(label, subset): seqlen_kv=kwargs.get("seqlen_kv"), varlen_seqlens_q=kwargs.get("varlen_seqlens_q"), varlen_seqlens_kv=kwargs.get("varlen_seqlens_kv"), + use_bias=args.bias, + use_alibi=args.alibi, + use_sink=args.sink, ) varlen_cmp_rows.append( ( @@ -2925,14 +3227,14 @@ def _extra_cmp_avg(label, subset): _print_grouped_avgs(varlen_cmp_rows, lambda r: (r[5], r[6]), _extra_cmp_avg) print("=" * len(xhdr2)) - varlen_csv_path = f"fmha_varlen_perf_compare_{_gpu_short_name()}.csv" + varlen_csv_path = f"fmha_varlen_perf_compare{csv_tag}_{_gpu_short_name()}.csv" _write_varlen_cmp_csv(varlen_csv_path, varlen_cmp_rows, varlen_cmp_avg_rows) print(f"Varlen results saved to: {varlen_csv_path}") else: # ---- Normal FlyDSL test mode ---- print("=" * 130) - print(f"FlyDSL flash_attn_func ({causal_desc}, {dtype_desc})") + print(f"FlyDSL flash_attn_func ({causal_desc}, {dtype_desc}{bias_desc})") print(f"GPU: {torch.cuda.get_device_name(0)}") print(f" Kernel opts: {FLASH_ATTN_FUNC_KERNEL_CONFIG}") if args.block_table: @@ -3021,6 +3323,11 @@ def _extra_cmp_avg(label, subset): precomputed_inputs=precomputed_inputs, page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", + use_bias=args.bias, + use_alibi=args.alibi, + use_sink=args.sink, + sink_share=args.sink_share, + alibi_two_d=args.alibi_two_d, ) if "err" in r: print(f" [FlyDSL unsupported] {_fmt_cfg(cfg)} {path}: {r['err']}", flush=True) @@ -3070,7 +3377,7 @@ def _normal_avg_fn(label, subset): _print_grouped_avgs(rows, lambda r: _tag_group(r[0]), _normal_avg_fn) print("=" * len(hdr)) - csv_path = f"fmha_perf_{_gpu_short_name()}.csv" + csv_path = f"fmha_perf{csv_tag}_{_gpu_short_name()}.csv" _write_normal_csv(csv_path, rows, normal_avg_rows) print(f"Results saved to: {csv_path}") @@ -3179,6 +3486,11 @@ def _normal_avg_fn(label, subset): precomputed_inputs=precomputed_inputs, page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", + use_bias=args.bias, + use_alibi=args.alibi, + use_sink=args.sink, + sink_share=args.sink_share, + alibi_two_d=args.alibi_two_d, **kwargs, ) except Exception as e: @@ -3276,7 +3588,7 @@ def _extra_normal_avg(label, subset): _print_grouped_avgs(varlen_rows, lambda r: (r[5], r[6]), _extra_normal_avg) print("=" * len(xhdr)) - varlen_csv_path = f"fmha_varlen_perf_{_gpu_short_name()}.csv" + varlen_csv_path = f"fmha_varlen_perf{csv_tag}_{_gpu_short_name()}.csv" _write_varlen_normal_csv(varlen_csv_path, varlen_rows, varlen_avg_rows) print(f"Varlen results saved to: {varlen_csv_path}") @@ -3428,6 +3740,180 @@ def test_lse_varlen(causal): _assert_lse_matches(lse[b, :, :n], ref, _ATOL_BF16) +# ── attention bias ─────────────────────────────────────────────────────────── + + +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("B,S,H,Hkv,D", [(1, 512, 8, 8, 128), (2, 384, 8, 4, 64)]) +def test_bias_dense(causal, B, S, H, Hkv, D): + """Dense bias is [Sq, Skv], broadcast over batch and head.""" + dtype = torch.bfloat16 + setup_seed(DEFAULT_SEED) + q = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(B, S, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(B, S, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + bias = torch.empty(S, S, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + + out = flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv, bias=bias) + torch.cuda.synchronize() + ref = pytorch_ref_attention(q.float(), k.float(), v.float(), causal=causal, bias=bias) + _, _, passed = _acc_metric(out.float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"biased output does not match the biased reference (B={B} S={S} causal={causal})" + + # The bias must actually change the result: an unbiased run must NOT match the + # biased reference, otherwise a silently-dropped bias would pass the check above. + out_nb = flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv) + torch.cuda.synchronize() + assert (out_nb.float() - ref.float()).abs().max().item() > 1e-2, "bias had no effect on the output" + + +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +def test_bias_varlen(causal): + """Varlen bias is packed [total_q, max_seqlen_kv]: global q rows, batch-local key columns.""" + dtype = torch.bfloat16 + D, H, Hkv = 128, 8, 4 + seqs = [512, 256, 384] + setup_seed(DEFAULT_SEED) + cu_list = [0] + for s in seqs: + cu_list.append(cu_list[-1] + s) + total, max_s = cu_list[-1], max(seqs) + cu = torch.tensor(cu_list, dtype=torch.int32, device="cuda") + q = torch.empty(total, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(total, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(total, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + bias = torch.empty(total, max_s, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + + out = flydsl_flash_attn_func( + q, + k, + v, + causal=causal, + num_kv_heads=Hkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_s, + max_seqlen_kv=max_s, + cross_seqlen=False, + bias=bias, + ) + torch.cuda.synchronize() + for b, n in enumerate(seqs): + s0, s1 = cu_list[b], cu_list[b + 1] + ref = pytorch_ref_attention( + q[s0:s1].unsqueeze(0).float(), + k[s0:s1].unsqueeze(0).float(), + v[s0:s1].unsqueeze(0).float(), + causal=causal, + bias=bias[s0:s1, :n], + ).squeeze(0) + _, _, passed = _acc_metric(out[s0:s1].float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"varlen batch {b} (seqlen {n}, causal={causal}) does not match the biased reference" + + +# ── ALiBi ──────────────────────────────────────────────────────────────────── + + +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("two_d", [False, True]) +@pytest.mark.parametrize("B,S,H,Hkv,D", [(2, 512, 8, 8, 128), (1, 384, 8, 4, 64)]) +def test_alibi_dense(causal, two_d, B, S, H, Hkv, D): + """score += -slope * |i + Skv - Sq - j|; slopes are [H] or [B, H] (alibi_stride_b).""" + dtype = torch.bfloat16 + setup_seed(DEFAULT_SEED) + q = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(B, S, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(B, S, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + slopes = make_alibi_slopes(B, H, two_d) + assert slopes.shape == ((B, H) if two_d else (H,)) + + out = flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv, alibi_slopes=slopes) + torch.cuda.synchronize() + ref = pytorch_ref_attention(q.float(), k.float(), v.float(), causal=causal, alibi_slopes=slopes) + _, _, passed = _acc_metric(out.float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"ALiBi output does not match the reference (B={B} S={S} two_d={two_d} causal={causal})" + + # Without slopes the result must differ, else a dropped ALiBi term would pass above. + out_nb = flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv) + torch.cuda.synchronize() + assert (out_nb.float() - ref.float()).abs().max().item() > 1e-2, "ALiBi had no effect on the output" + + +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +def test_alibi_varlen(causal): + """ALiBi positions are within-sequence: no packed-token base, per-batch lengths.""" + dtype = torch.bfloat16 + D, H, Hkv = 128, 8, 4 + seqs = [512, 256, 384] + setup_seed(DEFAULT_SEED) + cu_list = [0] + for s in seqs: + cu_list.append(cu_list[-1] + s) + total, max_s = cu_list[-1], max(seqs) + cu = torch.tensor(cu_list, dtype=torch.int32, device="cuda") + q = torch.empty(total, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(total, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(total, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + slopes = make_alibi_slopes(len(seqs), H, two_d=True) + + out = flydsl_flash_attn_func( + q, + k, + v, + causal=causal, + num_kv_heads=Hkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_s, + max_seqlen_kv=max_s, + cross_seqlen=False, + alibi_slopes=slopes, + ) + torch.cuda.synchronize() + for b, n in enumerate(seqs): + s0, s1 = cu_list[b], cu_list[b + 1] + ref = pytorch_ref_attention( + q[s0:s1].unsqueeze(0).float(), + k[s0:s1].unsqueeze(0).float(), + v[s0:s1].unsqueeze(0).float(), + causal=causal, + alibi_slopes=slopes[b], + ).squeeze(0) + _, _, passed = _acc_metric(out[s0:s1].float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"varlen batch {b} (seqlen {n}, causal={causal}) does not match the ALiBi reference" + + +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +def test_alibi_and_bias_combined(causal): + """ALiBi and bias are independent score terms and must both land.""" + dtype = torch.bfloat16 + B, S, H, D = 1, 512, 8, 128 + setup_seed(DEFAULT_SEED) + q = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + bias = torch.empty(S, S, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + slopes = make_alibi_slopes(B, H) + + out = flydsl_flash_attn_func(q, k, v, causal=causal, bias=bias, alibi_slopes=slopes) + torch.cuda.synchronize() + qf, kf, vf = q.float(), k.float(), v.float() + ref_both = pytorch_ref_attention(qf, kf, vf, causal=causal, bias=bias, alibi_slopes=slopes) + _, _, passed = _acc_metric(out.float().reshape(-1), ref_both.float().reshape(-1), D) + assert passed, "combined bias+ALiBi output does not match the combined reference" + + # Neither term alone explains the output. + ref_alibi = pytorch_ref_attention(qf, kf, vf, causal=causal, alibi_slopes=slopes) + ref_bias = pytorch_ref_attention(qf, kf, vf, causal=causal, bias=bias) + assert (out.float() - ref_alibi.float()).abs().max().item() > 1e-2, "bias term missing" + assert (out.float() - ref_bias.float()).abs().max().item() > 1e-2, "ALiBi term missing" + + def test_lse_fully_masked_rows(): """Cross-attention causal with Skv < Sq: leading query rows see no keys -> -inf.""" dtype = torch.bfloat16 @@ -3473,3 +3959,118 @@ def test_return_lse_rejects_fp8(): if __name__ == "__main__": main() + + +# ── attention sink ─────────────────────────────────────────────────────────── + + +def _sink_for(q, k, causal, share=DEFAULT_SINK_SHARE): + return calibrate_sink(*_rows_logsumexp(q, k, causal), share) + + +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("share", [0.25, 0.9]) +@pytest.mark.parametrize("B,S,H,Hkv,D", [(1, 512, 8, 8, 128), (2, 384, 8, 4, 64)]) +def test_sink_dense(causal, share, B, S, H, Hkv, D): + """One extra softmax denominator logit per head, with no matching V row.""" + dtype = torch.bfloat16 + setup_seed(DEFAULT_SEED) + q = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(B, S, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(B, S, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + sink = _sink_for(q, k, causal, share) + assert sink.shape == (H,) and sink.dtype == torch.float32 + + out = flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv, sink=sink) + torch.cuda.synchronize() + ref = pytorch_ref_attention(q.float(), k.float(), v.float(), causal=causal, sink=sink) + _, _, passed = _acc_metric(out.float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"sink output does not match the reference (B={B} S={S} share={share} causal={causal})" + + # Calibration is what makes this test meaningful: with the sink dropped the + # result must visibly differ. An uncalibrated sink near 0 would not. + out_ns = flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv) + torch.cuda.synchronize() + assert (out_ns.float() - ref.float()).abs().max().item() > 1e-2, "sink had no effect on the output" + + +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +def test_sink_varlen(causal): + dtype = torch.bfloat16 + D, H, Hkv = 128, 8, 4 + seqs = [512, 256, 384] + setup_seed(DEFAULT_SEED) + cu_list = [0] + for s in seqs: + cu_list.append(cu_list[-1] + s) + total, max_s = cu_list[-1], max(seqs) + cu = torch.tensor(cu_list, dtype=torch.int32, device="cuda") + q = torch.empty(total, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(total, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(total, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + # One [H] table shared by every sequence, calibrated over all their rows. + tot, cnt = torch.zeros(H, dtype=torch.float32, device="cuda"), 0 + for b in range(len(seqs)): + s0, s1 = cu_list[b], cu_list[b + 1] + sl, n = _rows_logsumexp(q[s0:s1].unsqueeze(0), k[s0:s1].unsqueeze(0), causal) + tot += sl + cnt += n + sink = calibrate_sink(tot, cnt, DEFAULT_SINK_SHARE) + + out = flydsl_flash_attn_func( + q, + k, + v, + causal=causal, + num_kv_heads=Hkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_s, + max_seqlen_kv=max_s, + cross_seqlen=False, + sink=sink, + ) + torch.cuda.synchronize() + for b, n in enumerate(seqs): + s0, s1 = cu_list[b], cu_list[b + 1] + ref = pytorch_ref_attention( + q[s0:s1].unsqueeze(0).float(), + k[s0:s1].unsqueeze(0).float(), + v[s0:s1].unsqueeze(0).float(), + causal=causal, + sink=sink, + ).squeeze(0) + _, _, passed = _acc_metric(out[s0:s1].float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"varlen batch {b} (seqlen {n}, causal={causal}) does not match the sink reference" + + +@_requires_gfx950 +@pytest.mark.parametrize("num_kv_splits", [2, 3, 4]) +def test_sink_splitk_counted_once(num_kv_splits): + """Split-K writes sink-free partials and folds the sink in once, in the combine. + + LSE is the sharp signal: it is the log denominator, so a sink counted + num_kv_splits times (or zero times) shows up directly instead of being + normalized away as it is in O. + """ + dtype = torch.bfloat16 + B, S, H, D = 1, 2048, 8, 128 + setup_seed(DEFAULT_SEED) + q = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(B, S, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + sink = _sink_for(q, k, True) + + out1, lse1 = flydsl_flash_attn_func(q, k, v, causal=True, sink=sink, return_lse=True) + outk, lsek = flydsl_flash_attn_func(q, k, v, causal=True, sink=sink, num_kv_splits=num_kv_splits, return_lse=True) + torch.cuda.synchronize() + # Split-K must agree with the single-split result it is meant to reproduce. + assert (lsek - lse1).abs().max().item() < 2e-2, f"split-K LSE diverges at {num_kv_splits} splits" + _, _, passed = _acc_metric(outk.float().reshape(-1), out1.float().reshape(-1), D) + assert passed, f"split-K output diverges at {num_kv_splits} splits" + + # A sink counted once per split would shift LSE by ~ln(num_kv_splits); assert + # we are nowhere near that, so the test cannot pass on a double-count. + assert (lsek - lse1).abs().max().item() < 0.5 * math.log(num_kv_splits) From e5e109de3e754dc5c78dfca4bdad3677a857fb3e Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Wed, 5 Aug 2026 11:17:22 +0000 Subject: [PATCH 02/12] add sink to split k case --- kernels/attention/flash_attn_utils.py | 30 +++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index ef4bbb6a4..78c4f84b1 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -3756,6 +3756,14 @@ def rescale_o(self, v_o, m_row, l_row, m_tile_max, v_p): l_row = _fmul(l_row, corr, self.fm_fast) return v_o, m_new, l_row, v_p + def fold_sink(self, v_o, m_row, l_row, sink_log2): + m_new = _fmax(m_row, sink_log2, self.fm_fast) + corr = rocdl.exp2(T.f32, as_mlir_value(_fsub(m_row, m_new, self.fm_fast))) + self.scale_o(v_o, corr) + sink_w = rocdl.exp2(T.f32, as_mlir_value(_fsub(sink_log2, m_new, self.fm_fast))) + l_row = _fadd(_fmul(l_row, corr, self.fm_fast), sink_w, self.fm_fast) + return m_new, l_row + def _lazy_rescale_o_rescale(self, _n, *_st, v_o, m_row, l_row, m_tile_max, v_p): corr = rocdl.exp2(T.f32, as_mlir_value(_fsub(m_row, m_tile_max, self.fm_fast))) scaled_accs = list(v_o) @@ -5289,6 +5297,7 @@ def __init__( seq_len=None, stride_q_n=None, LSE=None, + Sink=None, ): if isinstance(traits_or_ctx, DualwaveSplitKCombineContext): self.__dict__.update(traits_or_ctx.__dict__) @@ -5300,6 +5309,7 @@ def __init__( self.O = O self.WS = WS self.LSE = LSE + self.Sink = Sink self.batch_size = batch_size self.seq_len = seq_len self.stride_q_n = stride_q_n @@ -5404,6 +5414,26 @@ def reduce_m_max(self, m_s): m_max = _fmax(m_max, m_s[i + 1], self.fm_fast) return m_max + def fold_sink(self, m_max, bias_log2e): + sink_rsrc = buffer_ops.create_buffer_resource_from_addr( + as_mlir_value(fx.Int64(fx.ptrtoint(fx.get_iter(self.Sink)))), + num_records_bytes=as_mlir_value(fx.Int64(self.traits.NUM_HEADS_Q * 4)), + ) + sink_f32 = buffer_ops.buffer_load( + sink_rsrc, + as_mlir_value(fx.Int32(self.q_head_idx)), + vec_width=1, + dtype=T.f32, + ) + + sink_log2 = _fmul(sink_f32, fx.Float32(bias_log2e), self.fm_fast) + m_new = _fmax(m_max, sink_log2, self.fm_fast) + sink_w = rocdl.exp2(T.f32, as_mlir_value(_fsub(sink_log2, m_new, self.fm_fast))) + return m_new, sink_w + + def add_sink_den(self, den, sink_w): + return _fadd(den, sink_w, self.fm_fast) + def init_accumulators(self): return as_mlir_value(self.c_zero_v4f32), as_mlir_value(self.c_zero_f) From f49e07d943e7e0158de373c390283242587d885a Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Fri, 7 Aug 2026 12:00:14 +0000 Subject: [PATCH 03/12] update bias, add double buffering --- kernels/attention/flash_attn_gfx950.py | 142 +++++++++++++++------- kernels/attention/flash_attn_interface.py | 26 +++- kernels/attention/flash_attn_utils.py | 68 +++++++++++ tests/kernels/test_flash_attn_fwd.py | 62 +++++++++- 4 files changed, 249 insertions(+), 49 deletions(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 3590cc68a..680614a02 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -35,10 +35,16 @@ DualwaveStoreHelper, _anchor_v_o, _anchor_v_p, + _bias_buf_bytes, + _bias_dma_m0_base, + _bias_dma_src_elem, + _bias_lds_lane_base, + _buffer_load_lds_128, _dualwave_sync_barrier, - _get_q_pack, + _load_bias_frag_lds, + _make_bias_lds_ptr, _make_dualwave_swp_traits, - _mfma_acc, + _num_bias_dma, _s_barrier, _s_nop, _s_setprio, @@ -146,13 +152,31 @@ def build_flash_attn_dualwave_swp_module( # Shared-memory layout: one 16B-aligned K/V region (K0/V0/K1/V1). _lds_elem_dtype = dtype_to_elem_type(traits.DTYPE_STR) + # Bias adds a double-buffered [NUM_WAVES, ROWS_PER_WAVE, BLOCK_N] staging tile. + _BIAS_LDS_ELEMS = 2 * _bias_buf_bytes(traits) // traits.BF16_BYTES + _NUM_DMA_BIAS = _num_bias_dma(traits) - if const_expr(traits.PAGED): + if const_expr(traits.PAGED and HAS_BIAS): @fx.struct class SharedStorage: kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] bt: fx.Array[fx.Int32, traits.PAGED_BT_LDS_SIZE, 16] + bias: fx.Array[_lds_elem_dtype, _BIAS_LDS_ELEMS, 16] + + elif const_expr(traits.PAGED): + + @fx.struct + class SharedStorage: + kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] + bt: fx.Array[fx.Int32, traits.PAGED_BT_LDS_SIZE, 16] + + elif const_expr(HAS_BIAS): + + @fx.struct + class SharedStorage: + kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] + bias: fx.Array[_lds_elem_dtype, _BIAS_LDS_ELEMS, 16] else: @@ -232,19 +256,28 @@ def flash_attn_dualwave_swp_gfx950_kernel( # ---------------- attention bias ---------------- if const_expr(HAS_BIAS): - _bias_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(Bias), fx.make_layout(1, 1)) - if const_expr(traits.KV_VECTORIZED): - _BIAS_VEC = 8 - _bias_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), ctx.elem_dtype) - else: - _BIAS_VEC = 4 - _bias_atom = fx.make_copy_atom(fx.rocdl.BufferCopy64b(), ctx.elem_dtype) + _bias_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(Bias, max_size=False), fx.make_layout(1, 1)) + _BIAS_VEC = 8 if const_expr(traits.KV_VECTORIZED) else 4 _BIAS_GROUPS = 16 // _BIAS_VEC - _bias_frags = [ - fx.make_rmem_tensor(fx.make_layout(_BIAS_VEC, 1), ctx.elem_dtype) for _ in range_constexpr(_BIAS_GROUPS) - ] + _bias_frag_ty = Vec.make_type(_BIAS_VEC, ctx.elem_dtype) + _bias_frag_align = _BIAS_VEC * traits.BF16_BYTES _bias_log2e = Vec.filled(_BIAS_VEC, BIAS_LOG2E, fx.Float32) - _bias_row_base_i32 = fx.Int32(ctx.q_tok_base) if const_expr(VARLEN) else None + _bias_lds_base_idx = fx.Index(fx.ptrtoint(ctx.lds.bias.ptr)) + _bias_lds_base_ptr = _make_bias_lds_ptr(_bias_lds_base_idx) + _bias_gran_bytes = traits.DMA_BYTES + _bias_gran_elems = traits.DMA_BYTES // traits.BF16_BYTES + _bias_half_grans = traits.K_SUB_N // _bias_gran_elems + _BIAS_BUF_BYTES = _bias_buf_bytes(traits) + # bits [4, 7): `Swizzle(mask=3, base=4, shift=3)` computes + # i ^ ((i & (((1 << 3) - 1) << (4 + 3))) >> 3) + # which is exactly `(row % 8) * 16`. + _bias_swz_layout = fx.make_composed_layout( + fx.static(fx.SwizzleType.get(3, 4, 3)), + fx.make_layout(_BIAS_LDS_ELEMS * traits.BF16_BYTES, 1), + ) + _bias_lane_base = fx.Index( + _bias_lds_lane_base(traits, 0, ctx.wave_id, ctx.lane_mod_32, ctx.lane_div_32, _BIAS_VEC) + ) # ---------------- ALiBi ---------------- if const_expr(HAS_ALIBI): @@ -269,27 +302,51 @@ def _load_sink_log2(): return Vec(_sink_frag.load(), (1,), fx.Float32)[0] * fx.Float32(BIAS_LOG2E) - def _score_bias_half(s_h, tile_idx, h): - col_base_h = _seq_pad_col_base(traits, tile_idx, lane_div_32=ctx.lane_div_32) + fx.Int32(32 * h) + def _issue_bias_dma(tile_idx, buf): + """Coalesced global->LDS DMA of this wave's bias tile into buffer `buf`.""" + row_base = ctx.q_start + ctx.wave_q_offset + if const_expr(VARLEN): + row_base = row_base + ctx.q_tok_base + tile_col_base = tile_idx * fx.Index(traits.BLOCK_N) + for d in range_constexpr(_NUM_DMA_BIAS): + _buffer_load_lds_128( + _bias_div, + _bias_dma_m0_base(traits, buf, d, ctx.wave_id_uni, _bias_lds_base_idx), + _bias_dma_src_elem( + traits, + row_base, + tile_col_base, + d, + lane_in_warp=ctx.lane_in_warp, + bias_stride0_v=fx.Index(bias_stride0), + ), + 0, + _dma_atom=ctx.dma_atom, + _lds_ptr_ty=ctx.lds_ptr_ty, + ) + + def _read_bias_frag(h, g, buf): + """Read one 4-element bias fragment for score half `h`, group `g` from LDS.""" + gran = _seq_pad_score_threshold(traits, g * _BIAS_VEC) // _bias_gran_elems + _bias_half_grans * h + sel = gran * _bias_gran_bytes + buf * _BIAS_BUF_BYTES + crd = fx.make_int_tuple(fx.Int32(_bias_lane_base + fx.Index(sel))) + off = fx.get_scalar(fx.crd2idx(crd, _bias_swz_layout)) + return _load_bias_frag_lds(_bias_lds_base_ptr, fx.Int32(off), _bias_frag_ty, _bias_frag_align) + + def _score_bias_half(s_h, tile_idx, h, buf): src = Vec(s_h) out = [src[r] for r in range_constexpr(16)] if const_expr(HAS_ALIBI): + col_base_h = _seq_pad_col_base(traits, tile_idx, lane_div_32=ctx.lane_div_32) + fx.Int32(32 * h) rel0 = fx.Float32(ctx.q_row_i32 + ctx.delta_i32 - col_base_h) for r in range_constexpr(16): d = rel0 - fx.Float32(float(_seq_pad_score_threshold(traits, r))) out[r] = fmath.absf(d) * _alibi_neg_slope + out[r] if const_expr(HAS_BIAS): - bias_row_i32 = ctx.q_row_i32 - if const_expr(VARLEN): - bias_row_i32 = bias_row_i32 + _bias_row_base_i32 - base = bias_row_i32 * fx.Int32(bias_stride0) + col_base_h for g in range_constexpr(_BIAS_GROUPS): - col_off = _seq_pad_score_threshold(traits, g * _BIAS_VEC) - fx.copy(_bias_atom, fx.slice(_bias_div, (None, base + fx.Int32(col_off))), _bias_frags[g]) - for g in range_constexpr(_BIAS_GROUPS): - bv = Vec(Vec(_bias_frags[g].load(), (_BIAS_VEC,), ctx.elem_dtype).to(fx.Float32)) + bv = Vec(Vec(_read_bias_frag(h, g, buf), (_BIAS_VEC,), ctx.elem_dtype).to(fx.Float32)) r0 = g * _BIAS_VEC acc = Vec.from_elements([out[r0 + e] for e in range_constexpr(_BIAS_VEC)], fx.Float32) fused = bv * _bias_log2e + acc @@ -298,22 +355,16 @@ def _score_bias_half(s_h, tile_idx, h): return Vec.from_elements(out, fx.Float32) - def qk_scored(v_k, q_all_scaled_bf16, tile_idx): + def qk_scored(v_k, q_all_scaled_bf16, tile_idx, buf=0): if const_expr(not (HAS_BIAS or HAS_ALIBI)): return gemm_helper.qk(v_k, q_all_scaled_bf16) - out_halves = [] - for h, k_h in ((0, v_k[0]), (1, v_k[1])): - acc = gemm_helper.c_zero_v16f32 - for ks in range_constexpr(traits.K_STEPS_QK): - acc = _mfma_acc( - k_h[ks], - _get_q_pack(traits, q_all_scaled_bf16, ks), - acc, - gemm_helper.mma_atom, - gemm_helper.mfma_acc_vec_type, - ) - out_halves.append(_score_bias_half(acc, tile_idx, h)) - return (out_halves[0], out_halves[1]) + if const_expr(HAS_BIAS): + _issue_bias_dma(tile_idx + fx.Index(1), 1 - buf) + v_s_lo, v_s_hi = gemm_helper.qk(v_k, q_all_scaled_bf16) + return ( + _score_bias_half(v_s_lo, tile_idx, 0, buf), + _score_bias_half(v_s_hi, tile_idx, 1, buf), + ) def _main_body(): # Paged: stage the block-table row into LDS before any page-id ds_read. @@ -333,6 +384,9 @@ def _main_body(): _sched_barrier(0) _s_barrier() + if const_expr(HAS_BIAS): + _issue_bias_dma(ctx.split_tile(0), 0) + # Load this wave's Q rows and pre-scale by the 1/sqrt(D) softmax q_all_bf16 = q_loader.load_all() q_all_scaled_bf16 = q_loader.scale_all(q_all_bf16) @@ -362,7 +416,7 @@ def _main_body(): if const_expr(traits.PAGED): pro_pageid_2_lds = page_ids.load_page_id_lds(page_ids.split_tile(2)) - v_s_0 = qk_scored(v_k, q_all_scaled_bf16, ctx.split_tile(0)) + v_s_0 = qk_scored(v_k, q_all_scaled_bf16, ctx.split_tile(0), buf=0) _sched_barrier(0) if const_expr(traits.CAUSAL): @@ -439,7 +493,7 @@ def _main_body(): # Cluster 1 computes MMA0, finishes v_p_0 softmax, updates l_row, and casts P. if const_expr(traits.PAGED): c2_pageid_lds = page_ids.load_page_id_lds(j_idx) - v_s_1 = qk_scored(v_k, q_all_scaled_bf16, j_idx - 2) + v_s_1 = qk_scored(v_k, q_all_scaled_bf16, j_idx - 2, buf=1) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) v_p_0 = softmax_helper.cast_p(v_p_0) @@ -514,7 +568,7 @@ def _main_body(): # Cluster 5 mirrors C1: MMA0, finish v_p_1 softmax, update l_row, and cast P. if const_expr(traits.PAGED): _c6_kpid_lds = page_ids.load_page_id_lds(j_idx + 1) - v_s_0 = qk_scored(v_k, q_all_scaled_bf16, j_idx - 1) + v_s_0 = qk_scored(v_k, q_all_scaled_bf16, j_idx - 1, buf=0) v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_1) v_p_1 = softmax_helper.cast_p(v_p_1) @@ -605,7 +659,7 @@ def _main_body(): # Epilogue C1 (compute): MMA0 -> v_s_1; finish v_p_0 softmax (like C1). if const_expr(traits.PAGED): ec2_pageid_lds = page_ids.load_page_id_lds(max_m1) - v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m3) + v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m3, buf=1) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) v_p_0 = softmax_helper.cast_p(v_p_0) @@ -672,7 +726,7 @@ def _main_body(): _dualwave_sync_barrier() # Epilogue C5 computes MMA0, folds rescale_e3 into l_row, and finishes v_p_1 softmax. - v_s_0 = qk_scored(v_k, q_all_scaled_bf16, max_m2) + v_s_0 = qk_scored(v_k, q_all_scaled_bf16, max_m2, buf=0) l_row = softmax_helper.apply_l_rescale(l_row, rescale_e3) v_p_1 = softmax_helper.exp2(v_p_1, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_1) @@ -731,7 +785,7 @@ def _main_body(): _dualwave_sync_barrier() # Epilogue C9 computes the last-tile MMA0, folds rescale_e7 into l_row, and finishes v_p_0. - v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m1) + v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m1, buf=1) l_row = softmax_helper.apply_l_rescale(l_row, rescale_e7) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index 9d83cb092..ae149e39f 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -323,6 +323,7 @@ def _build_paged( varlen: bool = False, kv_cache_layout: str = "linear", return_lse: bool = False, + has_bias: bool = False, ): """Build (and cache) a paged-KV launcher (gfx950 DUALWAVE_SWP, paged=True). @@ -356,6 +357,7 @@ def _build_paged( dualwave_swp_setprio=setprio, dualwave_swp_enable_stagger=enable_stagger, return_lse=return_lse, + has_bias=has_bias, ) @@ -373,6 +375,7 @@ def _flydsl_flash_attn_paged( *, causal: bool, num_kv_heads: Optional[int], + bias: Optional[torch.Tensor], block_table: Optional[torch.Tensor], seqlen_k: Optional[torch.Tensor], max_seqlen_kv: Optional[int], @@ -487,6 +490,22 @@ def _flydsl_flash_attn_paged( cross = bool(cross_seqlen) if cross_seqlen is not None else True else: cross = skv != Sq + if bias is not None: + # Same convention as non-paged: rows are q tokens, columns are batch-local + # logical key positions (the block table only redirects the K/V fetch). + _bias_rows = int(q.shape[0]) if varlen else Sq + if bias.dim() != 2: + raise ValueError(f"flydsl_flash_attn_func: paged bias must be 2D, got {bias.dim()}D") + if bias.shape[0] != _bias_rows: + raise ValueError( + f"flydsl_flash_attn_func: paged bias must have {_bias_rows} rows " + f"({'total_q' if varlen else 'seq_len_q'}), got {tuple(bias.shape)}" + ) + if bias.shape[1] < skv: + raise ValueError( + f"flydsl_flash_attn_func: paged bias needs >= max_seqlen_kv={skv} columns, got {bias.shape[1]}" + ) + block_table_stride = int(block_table.shape[1]) # Flatten so the kernel's flat row-major index addresses block_table correctly. block_table_i32 = ( @@ -499,6 +518,7 @@ def _flydsl_flash_attn_paged( _arch = _gpu_arch(q.device) _paged_light_ok = ( (num_kv_splits <= 1) + and bias is None # the light paged kernel has no bias path and D in (64, 128) and dtype_str in ("bf16", "f16") and (not _arch.startswith("gfx950") or Sq <= _VARLEN_LIGHT_MAX_SEQ) @@ -536,6 +556,7 @@ def _flydsl_flash_attn_paged( num_kv_splits=int(num_kv_splits), varlen=varlen, kv_cache_layout=kv_cache_layout, + has_bias=bias is not None, ) if out is None: out = torch.empty_like(q) @@ -546,6 +567,8 @@ def _flydsl_flash_attn_paged( v_flat = v.contiguous() o_flat = out.contiguous() kwargs = dict(block_table=block_table_i32, block_table_stride=block_table_stride, stream=launch_stream) + if bias is not None: + kwargs["bias"] = bias if varlen: kwargs["cu_seqlens_q"] = cu_seqlens_q kwargs["cu_seqlens_kv"] = cu_seqlens_kv @@ -756,7 +779,7 @@ def flydsl_flash_attn_func( for _name, _t in (("bias", bias), ("alibi_slopes", alibi_slopes), ("sink", sink)): if _t is None: continue - if paged_kv: + if paged_kv and _name != "bias": raise NotImplementedError(f"flydsl_flash_attn_func: {_name} is not supported for paged KV") if dtype_str == "fp8": raise NotImplementedError(f"flydsl_flash_attn_func: {_name} is not supported for fp8") @@ -784,6 +807,7 @@ def flydsl_flash_attn_func( v, causal=causal, num_kv_heads=num_kv_heads, + bias=bias, block_table=block_table, seqlen_k=seqlen_k, max_seqlen_kv=max_seqlen_kv, diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 78c4f84b1..1a30ca29e 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -643,6 +643,74 @@ def _load_k_pack_aligned(traits, lds_kv_base_ptr, elem_idx, buf_id, kv_mfma_pack ).result +# ---------------- attention-bias LDS staging ---------------- +# The per-lane score layout wants 16 bias values from one q row, so a direct global +# read makes every lane in a wave touch a different bias row (one cache line per +# lane). Instead each wave DMAs its own [ROWS_PER_WAVE, BLOCK_N] tile with 8 lanes +# per row -- fully coalesced -- and reads the per-lane pattern back out of LDS. +# Column granules are XOR-swizzled by row so the read side spreads across banks; +# without it all 32 lanes of a half-wave would hit one bank (row stride is 128B). + + +def _bias_lanes_per_row(traits): + return traits.BLOCK_N * traits.BF16_BYTES // traits.DMA_BYTES + + +def _bias_gran_elems(traits): + return traits.DMA_BYTES // traits.BF16_BYTES + + +def _bias_wave_bytes(traits): + return traits.ROWS_PER_WAVE * traits.BLOCK_N * traits.BF16_BYTES + + +def _bias_buf_bytes(traits): + return traits.NUM_WAVES * _bias_wave_bytes(traits) + + +def _num_bias_dma(traits): + return _bias_wave_bytes(traits) // (traits.WARP_SIZE * traits.DMA_BYTES) + + +def _bias_dma_m0_base(traits, buf, d, wave_id_uni, lds_bias_base_idx): + """Wave-uniform M0 base for DMA batch `d` of this wave's bias tile.""" + lds_addr = ( + lds_bias_base_idx + + buf * _bias_buf_bytes(traits) + + wave_id_uni * _bias_wave_bytes(traits) + + d * (traits.WARP_SIZE * traits.DMA_BYTES) + ) + return rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(lds_addr))) + + +def _bias_dma_src_elem(traits, row_base, tile_col_base, d, lane_in_warp, bias_stride0_v): + """Per-lane global element index for DMA batch `d`; 8 lanes cover one 64-col row.""" + lanes_per_row = _bias_lanes_per_row(traits) + row_in_group = lane_in_warp // lanes_per_row + row_in_wave = d * (traits.WARP_SIZE // lanes_per_row) + row_in_group + gran = (lane_in_warp % lanes_per_row) ^ row_in_group + return (row_base + row_in_wave) * bias_stride0_v + tile_col_base + gran * _bias_gran_elems(traits) + + +def _bias_lds_lane_base(traits, buf, wave_id, lane_mod_32, lane_div_32, vec_elems): + """Per-lane byte base of this wave's bias tile in LDS buffer `buf` (granule 0).""" + return ( + buf * _bias_buf_bytes(traits) + + wave_id * _bias_wave_bytes(traits) + + lane_mod_32 * (traits.BLOCK_N * traits.BF16_BYTES) + + lane_div_32 * (vec_elems * traits.BF16_BYTES) + ) + + +def _make_bias_lds_ptr(lds_bias_base_idx): + return buffer_ops.create_llvm_ptr(lds_bias_base_idx, address_space=3) + + +def _load_bias_frag_lds(lds_bias_base_ptr, byte_offset, frag_type, align): + ptr = buffer_ops.get_element_ptr(lds_bias_base_ptr, byte_offset=byte_offset, elem_type=T.i8) + return llvm.LoadOp(frag_type, ptr, alignment=align).result + + def _ws_store_f32(f32_val, local_elem_index, rsrc): """32-bit f32 store into a per-split-z workspace region via raw buffer descriptor.""" f32_ir = as_mlir_value(fx.Float32(f32_val)) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 30458a511..400261f06 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -1179,6 +1179,7 @@ def run_attn_config( kv_cache_layout=kv_cache_layout, num_kv_splits=int(num_kv_splits), out=o_t, + bias=bias_t, **_paged_varlen_kw, **_cfg_kw(), ) @@ -1295,6 +1296,7 @@ def kernel_fn(): kv_cache_layout=kv_cache_layout, num_kv_splits=int(num_kv_splits), out=o_t, + bias=bias_t, **_paged_varlen_kw, **_cfg_kw(), ) @@ -2699,10 +2701,8 @@ def main(): args = parser.parse_args() if not args.block_table and args.kv_cache_layout != "linear": parser.error("--kv-cache-layout requires --block-table") - # Paged KV and fp8 have no bias support in the kernel; reject rather than run - # a bias-free kernel against a biased reference. - if args.bias and args.block_table: - parser.error("--bias is not supported with --block-table (paged KV)") + # fp8 has no bias support in the kernel; reject rather than run a bias-free + # kernel against a biased reference. Paged KV does support bias. if args.bias and args.dtype == "fp8": parser.error("--bias is not supported with --dtype fp8") if args.alibi and args.block_table: @@ -3813,6 +3813,60 @@ def test_bias_varlen(causal): assert passed, f"varlen batch {b} (seqlen {n}, causal={causal}) does not match the biased reference" +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("kv_cache_layout", ["linear", "vectorized"]) +def test_bias_paged(causal, kv_cache_layout): + """Paged bias is [Sq, max_seqlen_kv]: q rows, batch-local logical key columns. + + The block table only redirects the K/V fetch, so the bias column is still the + logical KV position -- the same index the causal mask already uses. + """ + dtype = torch.bfloat16 + B, Sq, H, Hkv, D = 2, 512, 8, 4, 128 + # Uniform KV lengths: ragged per-batch seqlen_k on the dense paged path already + # disagrees with this reference without any bias, so keep that out of scope here. + kv_lens = [Sq, Sq] + max_kv = max(kv_lens) + setup_seed(DEFAULT_SEED) + q = torch.empty(B, Sq, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + kv_cache = _build_paged_kv_for_test(B, max_kv, 64, Hkv, D, kv_lens, dtype, "cuda", kv_cache_layout) + bias = torch.empty(Sq, max_kv, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + + paged_kw = dict( + causal=causal, + num_kv_heads=Hkv, + max_seqlen_kv=max_kv, + block_table=kv_cache["block_table"], + seqlen_k=kv_cache["seqlen_k"], + kv_cache_layout=kv_cache_layout, + ) + out = flydsl_flash_attn_func(q, kv_cache["k_cache"], kv_cache["v_cache"], bias=bias, **paged_kw) + torch.cuda.synchronize() + + for b, n in enumerate(kv_lens): + kb, vb = _logical_kv_from_pages( + kv_cache["k_cache"][_page_ids_for_batch(kv_cache, b)], + kv_cache["v_cache"][_page_ids_for_batch(kv_cache, b)], + kv_cache_layout, + n, + ) + ref = pytorch_ref_attention( + q[b].unsqueeze(0).float(), + kb.unsqueeze(0).float(), + vb.unsqueeze(0).float(), + causal=causal, + bias=bias[:, :n], + ).squeeze(0) + _, _, passed = _acc_metric(out[b].float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"paged batch {b} ({kv_cache_layout}, causal={causal}) does not match the biased reference" + + # A bias-free paged run must NOT match, so a silently-dropped bias cannot pass. + out_nb = flydsl_flash_attn_func(q, kv_cache["k_cache"], kv_cache["v_cache"], **paged_kw) + torch.cuda.synchronize() + assert (out_nb.float() - out.float()).abs().max().item() > 1e-2, "bias had no effect on the paged output" + + # ── ALiBi ──────────────────────────────────────────────────────────────────── From 4252dbd076603c59418f7c715d913cf77d352331 Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Fri, 7 Aug 2026 15:06:58 +0000 Subject: [PATCH 04/12] fix tests --- kernels/attention/flash_attn_gfx950.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 680614a02..13d5eaaf2 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -256,7 +256,12 @@ def flash_attn_dualwave_swp_gfx950_kernel( # ---------------- attention bias ---------------- if const_expr(HAS_BIAS): - _bias_div = fx.logical_divide(fx.rocdl.make_buffer_tensor(Bias, max_size=False), fx.make_layout(1, 1)) + _bias_rec_bytes = ( + (fx.Int64(fx.get_scalar(fx.cosize(fx.get_layout(Bias)))) * traits.BF16_BYTES + 3) // 4 + ) * 4 + _bias_div = fx.logical_divide( + fx.rocdl.make_buffer_tensor(Bias, num_records_bytes=_bias_rec_bytes), fx.make_layout(1, 1) + ) _BIAS_VEC = 8 if const_expr(traits.KV_VECTORIZED) else 4 _BIAS_GROUPS = 16 // _BIAS_VEC _bias_frag_ty = Vec.make_type(_BIAS_VEC, ctx.elem_dtype) From 57a019baf822ea562bac8f26ec659426d23f1479 Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Mon, 10 Aug 2026 18:09:20 +0000 Subject: [PATCH 05/12] fix device in tests --- kernels/attention/flash_attn_gfx950.py | 2 +- tests/kernels/test_flash_attn_fwd.py | 23 +++++++++++++++-------- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 13d5eaaf2..02314cd29 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -335,7 +335,7 @@ def _read_bias_frag(h, g, buf): gran = _seq_pad_score_threshold(traits, g * _BIAS_VEC) // _bias_gran_elems + _bias_half_grans * h sel = gran * _bias_gran_bytes + buf * _BIAS_BUF_BYTES crd = fx.make_int_tuple(fx.Int32(_bias_lane_base + fx.Index(sel))) - off = fx.get_scalar(fx.crd2idx(crd, _bias_swz_layout)) + off = fx.get_scalar(fx.crd2idx(fx.make_int_tuple(fx.Int32(_bias_lane_base + fx.Index(sel))), _bias_swz_layout)) return _load_bias_frag_lds(_bias_lds_base_ptr, fx.Int32(off), _bias_frag_ty, _bias_frag_align) def _score_bias_half(s_h, tile_idx, h, buf): diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 400261f06..24686583b 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -383,7 +383,7 @@ def bias_fits(rows, cols, elem_size=2): def _block_table_from_indices(kv_indptr_cpu, kv_indices_cpu, batch_size, max_num_pages_per_seq): - block_table_cpu = torch.zeros((batch_size, max_num_pages_per_seq), dtype=torch.int32) + block_table_cpu = torch.zeros((batch_size, max_num_pages_per_seq), dtype=torch.int32, device="cpu") for b in range(batch_size): start = kv_indptr_cpu[b].item() end = kv_indptr_cpu[b + 1].item() @@ -469,8 +469,10 @@ def _build_paged_kv_for_test( kv_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32, device="cpu") kv_num_used_pages = torch.div(kv_lens_cpu + page_size - 1, page_size, rounding_mode="floor").int() - kv_indptr_cpu = torch.cumsum(torch.cat((torch.tensor([0], dtype=torch.int32), kv_num_used_pages)), dim=0).int() - kv_indices_cpu = torch.nn.functional.pad(torch.randperm(total_num_pages).int(), (0, 128), value=0) + kv_indptr_cpu = torch.cumsum( + torch.cat((torch.tensor([0], dtype=torch.int32, device="cpu"), kv_num_used_pages)), dim=0 + ).int() + kv_indices_cpu = torch.nn.functional.pad(torch.randperm(total_num_pages, device="cpu").int(), (0, 128), value=0) kv_last_page_len_cpu = ((kv_lens_cpu - 1) % page_size + 1).int() block_table_cpu = _block_table_from_indices( kv_indptr_cpu, @@ -564,7 +566,7 @@ def _build_paged_kv_from_logical_for_aiter(inputs, page_size=16): v_cache_4d = torch.zeros_like(k_cache_4d) kv_num_used_pages = [] kv_indices = [] - block_table_cpu = torch.zeros((batch_size, max_num_pages_per_seq), dtype=torch.int32) + block_table_cpu = torch.zeros((batch_size, max_num_pages_per_seq), dtype=torch.int32, device="cpu") for b, kv_len in enumerate(kv_lens): num_pages = _ceil_div(kv_len, page_size) @@ -573,6 +575,7 @@ def _build_paged_kv_from_logical_for_aiter(inputs, page_size=16): b * max_num_pages_per_seq, b * max_num_pages_per_seq + num_pages, dtype=torch.int32, + device="cpu", ) kv_indices.extend(page_ids.tolist()) block_table_cpu[b, :num_pages] = page_ids @@ -604,10 +607,14 @@ def _build_paged_kv_from_logical_for_aiter(inputs, page_size=16): else: k_cache, v_cache = k_cache_4d, v_cache_4d - kv_num_used_pages_cpu = torch.tensor(kv_num_used_pages, dtype=torch.int32) - kv_indptr_cpu = torch.cumsum(torch.cat((torch.tensor([0], dtype=torch.int32), kv_num_used_pages_cpu)), dim=0) - kv_indices_cpu = torch.nn.functional.pad(torch.tensor(kv_indices, dtype=torch.int32), (0, 128), value=0) - kv_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32) + kv_num_used_pages_cpu = torch.tensor(kv_num_used_pages, dtype=torch.int32, device="cpu") + kv_indptr_cpu = torch.cumsum( + torch.cat((torch.tensor([0], dtype=torch.int32, device="cpu"), kv_num_used_pages_cpu)), dim=0 + ) + kv_indices_cpu = torch.nn.functional.pad( + torch.tensor(kv_indices, dtype=torch.int32, device="cpu"), (0, 128), value=0 + ) + kv_lens_cpu = torch.tensor(kv_lens, dtype=torch.int32, device="cpu") kv_last_page_len_cpu = ((kv_lens_cpu - 1) % page_size + 1).int() return { "k_cache": k_cache, From 82d2cc06de3d03fa2bd32dbba1d7525d542cf171 Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Mon, 10 Aug 2026 18:15:48 +0000 Subject: [PATCH 06/12] fix codestyle --- kernels/attention/flash_attn_gfx950.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 02314cd29..13d5eaaf2 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -335,7 +335,7 @@ def _read_bias_frag(h, g, buf): gran = _seq_pad_score_threshold(traits, g * _BIAS_VEC) // _bias_gran_elems + _bias_half_grans * h sel = gran * _bias_gran_bytes + buf * _BIAS_BUF_BYTES crd = fx.make_int_tuple(fx.Int32(_bias_lane_base + fx.Index(sel))) - off = fx.get_scalar(fx.crd2idx(fx.make_int_tuple(fx.Int32(_bias_lane_base + fx.Index(sel))), _bias_swz_layout)) + off = fx.get_scalar(fx.crd2idx(crd, _bias_swz_layout)) return _load_bias_frag_lds(_bias_lds_base_ptr, fx.Int32(off), _bias_frag_ty, _bias_frag_align) def _score_bias_half(s_h, tile_idx, h, buf): From fe6db3fc2d643ef8956295b29d7fcd2bf7709219 Mon Sep 17 00:00:00 2001 From: root Date: Tue, 11 Aug 2026 15:07:59 +0000 Subject: [PATCH 07/12] improve performance --- kernels/attention/flash_attn_gfx950.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 13d5eaaf2..01274a111 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -162,7 +162,7 @@ def build_flash_attn_dualwave_swp_module( class SharedStorage: kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] bt: fx.Array[fx.Int32, traits.PAGED_BT_LDS_SIZE, 16] - bias: fx.Array[_lds_elem_dtype, _BIAS_LDS_ELEMS, 16] + bias: fx.Array[_lds_elem_dtype, _BIAS_LDS_ELEMS, 128] elif const_expr(traits.PAGED): @@ -176,7 +176,7 @@ class SharedStorage: @fx.struct class SharedStorage: kv: fx.Array[_lds_elem_dtype, traits.LDS_KV_TOTAL_SIZE, 16] - bias: fx.Array[_lds_elem_dtype, _BIAS_LDS_ELEMS, 16] + bias: fx.Array[_lds_elem_dtype, _BIAS_LDS_ELEMS, 128] else: @@ -268,7 +268,6 @@ def flash_attn_dualwave_swp_gfx950_kernel( _bias_frag_align = _BIAS_VEC * traits.BF16_BYTES _bias_log2e = Vec.filled(_BIAS_VEC, BIAS_LOG2E, fx.Float32) _bias_lds_base_idx = fx.Index(fx.ptrtoint(ctx.lds.bias.ptr)) - _bias_lds_base_ptr = _make_bias_lds_ptr(_bias_lds_base_idx) _bias_gran_bytes = traits.DMA_BYTES _bias_gran_elems = traits.DMA_BYTES // traits.BF16_BYTES _bias_half_grans = traits.K_SUB_N // _bias_gran_elems @@ -283,6 +282,9 @@ def flash_attn_dualwave_swp_gfx950_kernel( _bias_lane_base = fx.Index( _bias_lds_lane_base(traits, 0, ctx.wave_id, ctx.lane_mod_32, ctx.lane_div_32, _BIAS_VEC) ) + _bias_addr_base = _bias_lds_base_idx + fx.Index( + fx.get_scalar(fx.crd2idx(fx.make_int_tuple(fx.Int32(_bias_lane_base)), _bias_swz_layout)) + ) # ---------------- ALiBi ---------------- if const_expr(HAS_ALIBI): @@ -331,12 +333,9 @@ def _issue_bias_dma(tile_idx, buf): ) def _read_bias_frag(h, g, buf): - """Read one 4-element bias fragment for score half `h`, group `g` from LDS.""" gran = _seq_pad_score_threshold(traits, g * _BIAS_VEC) // _bias_gran_elems + _bias_half_grans * h - sel = gran * _bias_gran_bytes + buf * _BIAS_BUF_BYTES - crd = fx.make_int_tuple(fx.Int32(_bias_lane_base + fx.Index(sel))) - off = fx.get_scalar(fx.crd2idx(crd, _bias_swz_layout)) - return _load_bias_frag_lds(_bias_lds_base_ptr, fx.Int32(off), _bias_frag_ty, _bias_frag_align) + ptr = _make_bias_lds_ptr(_bias_addr_base ^ fx.Index(gran * _bias_gran_bytes)) + return _load_bias_frag_lds(ptr, fx.Int32(buf * _BIAS_BUF_BYTES), _bias_frag_ty, _bias_frag_align) def _score_bias_half(s_h, tile_idx, h, buf): src = Vec(s_h) From bd4e7bb325ce71a729b47d3071c5efcb78ae019e Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Wed, 12 Aug 2026 14:37:41 +0000 Subject: [PATCH 08/12] fix registers spills --- kernels/attention/flash_attn_gfx950.py | 6 +++-- kernels/attention/flash_attn_utils.py | 34 ++++++++++++++++++++------ 2 files changed, 31 insertions(+), 9 deletions(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 01274a111..6c2844db4 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -311,7 +311,9 @@ def _load_sink_log2(): def _issue_bias_dma(tile_idx, buf): """Coalesced global->LDS DMA of this wave's bias tile into buffer `buf`.""" - row_base = ctx.q_start + ctx.wave_q_offset + # Same value as ctx.q_start_pos_i32 (set later, by q_loader.load_all()), but + # built from the uniform wave_id_uni so the row base stays wave-uniform. + row_base = ctx.q_start + ctx.wave_id_uni * traits.ROWS_PER_WAVE if const_expr(VARLEN): row_base = row_base + ctx.q_tok_base tile_col_base = tile_idx * fx.Index(traits.BLOCK_N) @@ -325,7 +327,7 @@ def _issue_bias_dma(tile_idx, buf): tile_col_base, d, lane_in_warp=ctx.lane_in_warp, - bias_stride0_v=fx.Index(bias_stride0), + bias_stride0_v=bias_stride0, ), 0, _dma_atom=ctx.dma_atom, diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 86933f415..4ed5404ae 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -673,23 +673,43 @@ def _num_bias_dma(traits): def _bias_dma_m0_base(traits, buf, d, wave_id_uni, lds_bias_base_idx): - """Wave-uniform M0 base for DMA batch `d` of this wave's bias tile.""" + """Wave-uniform M0 base for DMA batch `d` of this wave's bias tile. + + Already scalar by construction: the LDS base is a compile-time global address and + `wave_id_uni` is `readfirstlane`-derived, so every term is uniform. A `readfirstlane` + here would be a round trip -- the intrinsic takes a VGPR operand, so the backend has + to copy SGPR->VGPR and back, and the copy stays live across the whole KV-tile loop. + """ lds_addr = ( lds_bias_base_idx + buf * _bias_buf_bytes(traits) + wave_id_uni * _bias_wave_bytes(traits) + d * (traits.WARP_SIZE * traits.DMA_BYTES) ) - return rocdl.readfirstlane(T.i32, as_mlir_value(fx.Int32(lds_addr))) + return fx.Int32(lds_addr) def _bias_dma_src_elem(traits, row_base, tile_col_base, d, lane_in_warp, bias_stride0_v): - """Per-lane global element index for DMA batch `d`; 8 lanes cover one 64-col row.""" + """Per-lane global element index for DMA batch `d`; 8 lanes cover one 64-col row. + + Kept in i32 end to end: the consumer (`_buffer_load_lds_128`) feeds this to a buffer + `voffset`, which is 32-bit, so a wider intermediate is truncated anyway. Bias tensors + above the i32 element range are rejected on the host side. + + The address also splits into a wave-uniform part (row base, batch offset, tile column) + pinned to an SGPR by `readfirstlane`, plus a per-lane part that is the same for every + `d`. The four DMA batches therefore share one VGPR instead of four full addresses. + `row_base` may be built from non-uniform-looking pieces (VARLEN's `q_tok_base`); the + `readfirstlane` is what makes the uniformity visible to the register allocator. + """ lanes_per_row = _bias_lanes_per_row(traits) - row_in_group = lane_in_warp // lanes_per_row - row_in_wave = d * (traits.WARP_SIZE // lanes_per_row) + row_in_group - gran = (lane_in_warp % lanes_per_row) ^ row_in_group - return (row_base + row_in_wave) * bias_stride0_v + tile_col_base + gran * _bias_gran_elems(traits) + stride = fx.Int32(bias_stride0_v) + lane_i32 = fx.Int32(lane_in_warp) + row_in_group = lane_i32 // fx.Int32(lanes_per_row) + gran = (lane_i32 % fx.Int32(lanes_per_row)) ^ row_in_group + uni = (fx.Int32(row_base) + fx.Int32(d * (traits.WARP_SIZE // lanes_per_row))) * stride + fx.Int32(tile_col_base) + uni_s = rocdl.readfirstlane(T.i32, as_mlir_value(uni)) + return fx.Int32(uni_s) + row_in_group * stride + gran * fx.Int32(_bias_gran_elems(traits)) def _bias_lds_lane_base(traits, buf, wave_id, lane_mod_32, lane_div_32, vec_elems): From 445137f5a1d0fdcd7873411f6f3fdd42545a15a1 Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Wed, 12 Aug 2026 16:10:08 +0000 Subject: [PATCH 09/12] fix wait --- kernels/attention/flash_attn_gfx950.py | 47 +++++++++++++++++++------- kernels/attention/flash_attn_utils.py | 12 ++++++- 2 files changed, 46 insertions(+), 13 deletions(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 6c2844db4..7e3e3f4f2 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -264,6 +264,7 @@ def flash_attn_dualwave_swp_gfx950_kernel( ) _BIAS_VEC = 8 if const_expr(traits.KV_VECTORIZED) else 4 _BIAS_GROUPS = 16 // _BIAS_VEC + _BIAS_READ_BATCH = min(4, _BIAS_GROUPS) _bias_frag_ty = Vec.make_type(_BIAS_VEC, ctx.elem_dtype) _bias_frag_align = _BIAS_VEC * traits.BF16_BYTES _bias_log2e = Vec.filled(_BIAS_VEC, BIAS_LOG2E, fx.Float32) @@ -309,6 +310,14 @@ def _load_sink_log2(): return Vec(_sink_frag.load(), (1,), fx.Float32)[0] * fx.Float32(BIAS_LOG2E) + # vmcnt budget for a bias fragment read: every qk_scored consumes the bias tile + # DMA'd by the previous one, and that call issues its own DMA only after its + # reads, so the only VMEM still in flight ahead of the read is the K and V + # prefetch of the two memory clusters in between. Epilogue C9 is the one site + # with no K prefetch between. + _BIAS_VMCNT = ctx.NUM_DMA_K + ctx.NUM_DMA_V + _BIAS_VMCNT_NO_K = ctx.NUM_DMA_V + def _issue_bias_dma(tile_idx, buf): """Coalesced global->LDS DMA of this wave's bias tile into buffer `buf`.""" # Same value as ctx.q_start_pos_i32 (set later, by q_loader.load_all()), but @@ -351,26 +360,40 @@ def _score_bias_half(s_h, tile_idx, h, buf): out[r] = fmath.absf(d) * _alibi_neg_slope + out[r] if const_expr(HAS_BIAS): - for g in range_constexpr(_BIAS_GROUPS): - bv = Vec(Vec(_read_bias_frag(h, g, buf), (_BIAS_VEC,), ctx.elem_dtype).to(fx.Float32)) - r0 = g * _BIAS_VEC - acc = Vec.from_elements([out[r0 + e] for e in range_constexpr(_BIAS_VEC)], fx.Float32) - fused = bv * _bias_log2e + acc - for e in range_constexpr(_BIAS_VEC): - out[r0 + e] = fused[e] + # Issue the fragment reads of a batch back to back so the backend can + # cover them with one counted lgkmcnt instead of draining after each + # ds_read. The batch is deliberately small: every extra fragment in + # flight is another live VGPR pair, and this kernel is at the wall. + for gb in range_constexpr(0, _BIAS_GROUPS, _BIAS_READ_BATCH): + raw = [_read_bias_frag(h, gb + k, buf) for k in range_constexpr(_BIAS_READ_BATCH)] + for k in range_constexpr(_BIAS_READ_BATCH): + bv = Vec(Vec(raw[k], (_BIAS_VEC,), ctx.elem_dtype).to(fx.Float32)) + r0 = (gb + k) * _BIAS_VEC + acc = Vec.from_elements([out[r0 + e] for e in range_constexpr(_BIAS_VEC)], fx.Float32) + fused = bv * _bias_log2e + acc + for e in range_constexpr(_BIAS_VEC): + out[r0 + e] = fused[e] return Vec.from_elements(out, fx.Float32) - def qk_scored(v_k, q_all_scaled_bf16, tile_idx, buf=0): + def qk_scored(v_k, q_all_scaled_bf16, tile_idx, buf=0, bias_vmcnt=_BIAS_VMCNT): if const_expr(not (HAS_BIAS or HAS_ALIBI)): return gemm_helper.qk(v_k, q_all_scaled_bf16) - if const_expr(HAS_BIAS): - _issue_bias_dma(tile_idx + fx.Index(1), 1 - buf) v_s_lo, v_s_hi = gemm_helper.qk(v_k, q_all_scaled_bf16) - return ( + if const_expr(HAS_BIAS): + # The bias tile read below was DMA'd one qk_scored call ago; everything + # issued since is the in-flight K/V prefetch. Retire exactly that batch + # instead of letting the backend fall back to vmcnt(0), which would also + # drain the prefetch pipeline. + _waitcnt_vm_n(bias_vmcnt) + scored = ( _score_bias_half(v_s_lo, tile_idx, 0, buf), _score_bias_half(v_s_hi, tile_idx, 1, buf), ) + if const_expr(HAS_BIAS): + # Issued after the reads so it is not part of the batch they wait on. + _issue_bias_dma(tile_idx + fx.Index(1), 1 - buf) + return scored def _main_body(): # Paged: stage the block-table row into LDS before any page-id ds_read. @@ -791,7 +814,7 @@ def _main_body(): _dualwave_sync_barrier() # Epilogue C9 computes the last-tile MMA0, folds rescale_e7 into l_row, and finishes v_p_0. - v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m1, buf=1) + v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m1, buf=1, bias_vmcnt=_BIAS_VMCNT_NO_K) l_row = softmax_helper.apply_l_rescale(l_row, rescale_e7) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 4ed5404ae..1bd54043a 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -727,8 +727,18 @@ def _make_bias_lds_ptr(lds_bias_base_idx): def _load_bias_frag_lds(lds_bias_base_ptr, byte_offset, frag_type, align): + # The alias scope is what stops SIInsertWaitcnts from treating this as a read of + # "some LDS-DMA destination" and forcing a full vmcnt(0) drain before it; the bias + # tile is instead guarded by the explicit counted wait the caller emits. Only the + # scope's presence matters here -- the LDS-DMA stores carry no AAInfo, so a finer + # per-double-buffer scope would not give the backend anything more to prove. ptr = buffer_ops.get_element_ptr(lds_bias_base_ptr, byte_offset=byte_offset, elem_type=T.i8) - return llvm.LoadOp(frag_type, ptr, alignment=align).result + return llvm.LoadOp( + frag_type, + ptr, + alignment=align, + alias_scopes=_dualwave_lds_alias_scopes(_dualwave_lds_scope("bias", 0)), + ).result def _ws_store_f32(f32_val, local_elem_index, rsrc): From af2cf88df61205d43bd923ff39892ca5cec7bf5b Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Wed, 12 Aug 2026 19:26:22 +0000 Subject: [PATCH 10/12] remove excess comments --- kernels/attention/flash_attn_gfx950.py | 15 --------------- kernels/attention/flash_attn_utils.py | 25 ------------------------- 2 files changed, 40 deletions(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 7e3e3f4f2..5679cf2c9 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -310,18 +310,11 @@ def _load_sink_log2(): return Vec(_sink_frag.load(), (1,), fx.Float32)[0] * fx.Float32(BIAS_LOG2E) - # vmcnt budget for a bias fragment read: every qk_scored consumes the bias tile - # DMA'd by the previous one, and that call issues its own DMA only after its - # reads, so the only VMEM still in flight ahead of the read is the K and V - # prefetch of the two memory clusters in between. Epilogue C9 is the one site - # with no K prefetch between. _BIAS_VMCNT = ctx.NUM_DMA_K + ctx.NUM_DMA_V _BIAS_VMCNT_NO_K = ctx.NUM_DMA_V def _issue_bias_dma(tile_idx, buf): """Coalesced global->LDS DMA of this wave's bias tile into buffer `buf`.""" - # Same value as ctx.q_start_pos_i32 (set later, by q_loader.load_all()), but - # built from the uniform wave_id_uni so the row base stays wave-uniform. row_base = ctx.q_start + ctx.wave_id_uni * traits.ROWS_PER_WAVE if const_expr(VARLEN): row_base = row_base + ctx.q_tok_base @@ -360,10 +353,6 @@ def _score_bias_half(s_h, tile_idx, h, buf): out[r] = fmath.absf(d) * _alibi_neg_slope + out[r] if const_expr(HAS_BIAS): - # Issue the fragment reads of a batch back to back so the backend can - # cover them with one counted lgkmcnt instead of draining after each - # ds_read. The batch is deliberately small: every extra fragment in - # flight is another live VGPR pair, and this kernel is at the wall. for gb in range_constexpr(0, _BIAS_GROUPS, _BIAS_READ_BATCH): raw = [_read_bias_frag(h, gb + k, buf) for k in range_constexpr(_BIAS_READ_BATCH)] for k in range_constexpr(_BIAS_READ_BATCH): @@ -381,10 +370,6 @@ def qk_scored(v_k, q_all_scaled_bf16, tile_idx, buf=0, bias_vmcnt=_BIAS_VMCNT): return gemm_helper.qk(v_k, q_all_scaled_bf16) v_s_lo, v_s_hi = gemm_helper.qk(v_k, q_all_scaled_bf16) if const_expr(HAS_BIAS): - # The bias tile read below was DMA'd one qk_scored call ago; everything - # issued since is the in-flight K/V prefetch. Retire exactly that batch - # instead of letting the backend fall back to vmcnt(0), which would also - # drain the prefetch pipeline. _waitcnt_vm_n(bias_vmcnt) scored = ( _score_bias_half(v_s_lo, tile_idx, 0, buf), diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 1bd54043a..fbf945d91 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -673,13 +673,6 @@ def _num_bias_dma(traits): def _bias_dma_m0_base(traits, buf, d, wave_id_uni, lds_bias_base_idx): - """Wave-uniform M0 base for DMA batch `d` of this wave's bias tile. - - Already scalar by construction: the LDS base is a compile-time global address and - `wave_id_uni` is `readfirstlane`-derived, so every term is uniform. A `readfirstlane` - here would be a round trip -- the intrinsic takes a VGPR operand, so the backend has - to copy SGPR->VGPR and back, and the copy stays live across the whole KV-tile loop. - """ lds_addr = ( lds_bias_base_idx + buf * _bias_buf_bytes(traits) @@ -690,18 +683,6 @@ def _bias_dma_m0_base(traits, buf, d, wave_id_uni, lds_bias_base_idx): def _bias_dma_src_elem(traits, row_base, tile_col_base, d, lane_in_warp, bias_stride0_v): - """Per-lane global element index for DMA batch `d`; 8 lanes cover one 64-col row. - - Kept in i32 end to end: the consumer (`_buffer_load_lds_128`) feeds this to a buffer - `voffset`, which is 32-bit, so a wider intermediate is truncated anyway. Bias tensors - above the i32 element range are rejected on the host side. - - The address also splits into a wave-uniform part (row base, batch offset, tile column) - pinned to an SGPR by `readfirstlane`, plus a per-lane part that is the same for every - `d`. The four DMA batches therefore share one VGPR instead of four full addresses. - `row_base` may be built from non-uniform-looking pieces (VARLEN's `q_tok_base`); the - `readfirstlane` is what makes the uniformity visible to the register allocator. - """ lanes_per_row = _bias_lanes_per_row(traits) stride = fx.Int32(bias_stride0_v) lane_i32 = fx.Int32(lane_in_warp) @@ -713,7 +694,6 @@ def _bias_dma_src_elem(traits, row_base, tile_col_base, d, lane_in_warp, bias_st def _bias_lds_lane_base(traits, buf, wave_id, lane_mod_32, lane_div_32, vec_elems): - """Per-lane byte base of this wave's bias tile in LDS buffer `buf` (granule 0).""" return ( buf * _bias_buf_bytes(traits) + wave_id * _bias_wave_bytes(traits) @@ -727,11 +707,6 @@ def _make_bias_lds_ptr(lds_bias_base_idx): def _load_bias_frag_lds(lds_bias_base_ptr, byte_offset, frag_type, align): - # The alias scope is what stops SIInsertWaitcnts from treating this as a read of - # "some LDS-DMA destination" and forcing a full vmcnt(0) drain before it; the bias - # tile is instead guarded by the explicit counted wait the caller emits. Only the - # scope's presence matters here -- the LDS-DMA stores carry no AAInfo, so a finer - # per-double-buffer scope would not give the backend anything more to prove. ptr = buffer_ops.get_element_ptr(lds_bias_base_ptr, byte_offset=byte_offset, elem_type=T.i8) return llvm.LoadOp( frag_type, From f7fb3d3750ac7f2390b9eb7702cf1dcc1341f1ae Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Thu, 13 Aug 2026 14:41:33 +0000 Subject: [PATCH 11/12] fix comments --- kernels/attention/flash_attn_gfx950.py | 63 ++-- kernels/attention/flash_attn_interface.py | 58 +++- kernels/attention/flash_attn_utils.py | 30 +- tests/kernels/test_flash_attn_fwd.py | 372 +++++++++++++++++++++- 4 files changed, 461 insertions(+), 62 deletions(-) diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index 5679cf2c9..4ba01862e 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -59,6 +59,7 @@ _v_pair_to_vec32, _v_vec32_to_pair, _waitcnt_vm_n, + bias_addressing_error, ) from kernels.common.kernels_common import dtype_to_elem_type @@ -365,7 +366,7 @@ def _score_bias_half(s_h, tile_idx, h, buf): return Vec.from_elements(out, fx.Float32) - def qk_scored(v_k, q_all_scaled_bf16, tile_idx, buf=0, bias_vmcnt=_BIAS_VMCNT): + def qk_scored(v_k, q_all_scaled_bf16, tile_idx, buf=0, bias_vmcnt=_BIAS_VMCNT, prefetch_next=True): if const_expr(not (HAS_BIAS or HAS_ALIBI)): return gemm_helper.qk(v_k, q_all_scaled_bf16) v_s_lo, v_s_hi = gemm_helper.qk(v_k, q_all_scaled_bf16) @@ -376,8 +377,10 @@ def qk_scored(v_k, q_all_scaled_bf16, tile_idx, buf=0, bias_vmcnt=_BIAS_VMCNT): _score_bias_half(v_s_hi, tile_idx, 1, buf), ) if const_expr(HAS_BIAS): - # Issued after the reads so it is not part of the batch they wait on. - _issue_bias_dma(tile_idx + fx.Index(1), 1 - buf) + if const_expr(prefetch_next): + _issue_bias_dma(tile_idx + fx.Index(1), 1 - buf) + else: + _sched_barrier(0) return scored def _main_body(): @@ -799,7 +802,7 @@ def _main_body(): _dualwave_sync_barrier() # Epilogue C9 computes the last-tile MMA0, folds rescale_e7 into l_row, and finishes v_p_0. - v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m1, buf=1, bias_vmcnt=_BIAS_VMCNT_NO_K) + v_s_1 = qk_scored(v_k, q_all_scaled_bf16, max_m1, buf=1, bias_vmcnt=_BIAS_VMCNT_NO_K, prefetch_next=False) l_row = softmax_helper.apply_l_rescale(l_row, rescale_e7) v_p_0 = softmax_helper.exp2(v_p_0, 16, 16) l_row = softmax_helper.reduce_sum(l_row, v_p_0) @@ -873,7 +876,9 @@ def _main_body(): output_store.store_splitk_partial_o(v_o, m_row, l_row, ctx.q_row) if const_expr(traits.CAUSAL and traits.CROSS_SEQLEN and not traits.SPLITK): - output_store.zero_o_block_if_needed() + output_store.zero_o_block_if_needed( + sink_log2=_load_sink_log2() if const_expr(HAS_SINK and traits.RETURN_LSE) else None + ) if active is None: _main_body() @@ -1017,6 +1022,24 @@ def launch_flash_attn_dualwave_swp( }, } + def _prep_bias(bias, bias_stride0, placeholder): + if bias is None: + if HAS_BIAS: + raise ValueError( + "flash_attn_dualwave_swp was built with has_bias=True but no `bias` tensor was " + "provided; pass a [total_q, max_seqlen_kv] bias (varlen) or a " + "[seq_len, seq_len_kv] bias (dense), with the same dtype as q." + ) + return placeholder, 0 + if HAS_BIAS: + bias_err = bias_addressing_error(bias.numel(), bias.element_size()) + if bias_err is not None: + raise ValueError(f"flash_attn_dualwave_swp: bias {tuple(bias.shape)} {bias_err}") + bias = bias.contiguous() + if bias_stride0 is None: + bias_stride0 = bias.stride(0) + return bias.view(-1), bias_stride0 + def _prep_alibi(alibi_slopes, placeholder): if alibi_slopes is None: if HAS_ALIBI: @@ -1126,20 +1149,7 @@ def _launch( block_table = O if block_table_stride is None: block_table_stride = 0 - if bias is None: - if HAS_BIAS: - raise ValueError( - "flash_attn_dualwave_swp was built with has_bias=True but no `bias` tensor was " - "provided; pass a [total_q, max_seqlen_kv] bias (varlen) or a " - "[seq_len, seq_len_kv] bias (dense), with the same dtype as q." - ) - bias = O - bias_stride0 = 0 - else: - bias = bias.contiguous() - if bias_stride0 is None: - bias_stride0 = bias.stride(0) - bias = bias.view(-1) + bias, bias_stride0 = _prep_bias(bias, bias_stride0, O) alibi_slopes, alibi_stride_b = _prep_alibi(alibi_slopes, O) sink = _prep_sink(sink, O) with CompilationContext.compile_hints(_dualwave_swp_compile_hints): @@ -1246,20 +1256,7 @@ def _compile( if block_table_stride is None: block_table_stride = 0 - if bias is None: - if HAS_BIAS: - raise ValueError( - "flash_attn_dualwave_swp was built with has_bias=True but no `bias` tensor was " - "provided; pass a [total_q, max_seqlen_kv] bias (varlen) or a " - "[seq_len, seq_len_kv] bias (dense), with the same dtype as q." - ) - bias = O - bias_stride0 = 0 - else: - bias = bias.contiguous() - if bias_stride0 is None: - bias_stride0 = bias.stride(0) - bias = bias.view(-1) + bias, bias_stride0 = _prep_bias(bias, bias_stride0, O) alibi_slopes, alibi_stride_b = _prep_alibi(alibi_slopes, O) sink = _prep_sink(sink, O) with CompilationContext.compile_hints(_dualwave_swp_compile_hints): diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index ae149e39f..442af8511 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -28,8 +28,7 @@ import torch import torch.nn.functional as F # noqa: F401 (imported for callers' convenience) -# Re-export so callers only need to import from this module. -from kernels.attention.flash_attn_utils import dualwave_splitk_workspace_elems +from kernels.attention.flash_attn_utils import bias_addressing_error, dualwave_splitk_workspace_elems __all__ = ["flydsl_flash_attn_func", "dualwave_splitk_workspace_elems"] @@ -476,7 +475,8 @@ def _flydsl_flash_attn_paged( # Per-batch KV lengths differ in general → bottom-right cross-length masking. Varlen # paged always uses cross masking (per-batch seqlen_q/seqlen_kv come from cu_seqlens). - skv = int(max_seqlen_kv) if max_seqlen_kv is not None else int(seqlen_k.max().item()) + _kv_lens = seqlen_k.reshape(-1).tolist() if max_seqlen_kv is None or (bias is not None and not varlen) else None + skv = int(max_seqlen_kv) if max_seqlen_kv is not None else int(max(_kv_lens)) max_kv_pages = (skv + page_size - 1) // page_size max_pages_per_split = (max_kv_pages + int(num_kv_splits) - 1) // int(num_kv_splits) if max_pages_per_split > _PAGED_BT_LDS_SIZE: @@ -491,6 +491,14 @@ def _flydsl_flash_attn_paged( else: cross = skv != Sq if bias is not None: + if not varlen: + if min(_kv_lens) != max(_kv_lens): + raise NotImplementedError( + f"flydsl_flash_attn_func: dense paged bias requires uniform seqlen_k, got lengths in " + f"[{min(_kv_lens)}, {max(_kv_lens)}]; the dense paged kernel receives only " + f"max_seqlen_kv. Use the varlen paged path (cu_seqlens_q/cu_seqlens_kv) for " + f"ragged KV lengths." + ) # Same convention as non-paged: rows are q tokens, columns are batch-local # logical key positions (the block table only redirects the K/V fetch). _bias_rows = int(q.shape[0]) if varlen else Sq @@ -652,13 +660,13 @@ def flydsl_flash_attn_func( # Split-K (gfx950 only, seq_len >= 384, D=64/128, bf16/f16). num_kv_splits: int = 1, # Additive attention bias, folded into the scores after sm_scale and before - # masking. gfx950 DUALWAVE_SWP only (dense / varlen / split-K). + # masking. gfx950 DUALWAVE_SWP only (dense / varlen / split-K / paged KV). bias: Optional[torch.Tensor] = None, # Per-head ALiBi slope table, computed analytically into the scores. Same - # path and restrictions as `bias`; the two may be combined. + # path as `bias` but no paged-KV support; may be combined with `bias`. alibi_slopes: Optional[torch.Tensor] = None, # Per-head attention-sink logit: one extra softmax denominator term with no - # matching V row. Same path and restrictions as `bias`; freely combinable. + # matching V row. Same path and restrictions as `alibi_slopes`; combinable. sink: Optional[torch.Tensor] = None, # fp8 dense ABI: per-tensor descales for pre-quantized e4m3fn Q/K/V. q_descale: Optional[torch.Tensor] = None, @@ -709,23 +717,37 @@ def flydsl_flash_attn_func( seqlen_q != seqlen_kv per batch. cross_seqlen: Whether seqlen_q and seqlen_kv differ. Required in varlen mode; dense mode infers it from ``q.shape[1] != k.shape[1]``. - block_table / seqlen_k: vLLM-style 2D block table metadata. + block_table / seqlen_k: vLLM-style 2D block table metadata. Enables the + native paged-KV path, which supports ``bias`` but not + ``alibi_slopes``, ``sink``, ``return_lse``, or fp8. num_kv_splits: Split-K factor (>1: gfx950 only, D=64/128, bf16/f16, seq>=384). bias: Additive attention bias with the same dtype as q, folded in as ``softmax(q @ k^T * sm_scale + bias)`` -- after the scale, before the causal/padding mask. Dense: ``[Sq, Skv]``, broadcast over batch and head. Varlen: ``[total_q, max_seqlen_kv]``, where the row is the *global* packed q token index and the column is the *per-batch-local* - key index, broadcast over head. Routes to the gfx950 DUALWAVE_SWP - kernel; paged KV and fp8 raise NotImplementedError rather than - silently dropping the bias. + key index, broadcast over head. Varlen self-attention leaves + ``max_seqlen_kv`` unset, so its column bound is ``max_seqlen_q``. + Routes to the gfx950 DUALWAVE_SWP kernel; fp8 raises + NotImplementedError rather than silently dropping the bias. + Paged KV is supported (dense, varlen, and paged split-K) with the + same row/column convention: rows are ``seq_len_q`` (dense) or + ``total_q`` (varlen) q tokens, columns are batch-local key indices + and must number at least ``max_seqlen_kv``. Dense paged + additionally requires a uniform ``seqlen_k`` across the batch -- + the dense paged launch only receives ``max_seqlen_kv``, so ragged + lengths would address the wrong bias columns and raise + ``NotImplementedError``; use the varlen paged path + (``cu_seqlens_q``/``cu_seqlens_kv``) for ragged KV. alibi_slopes: fp32 ALiBi slope table, ``[H]`` (broadcast over batch) or ``[B, H]``, values positive. Adds ``-slope * |i + seqlen_kv - seqlen_q - j|`` to the scores after the 1/sqrt(D) scaling (the slope is not divided by it), bottom-right aligned like the causal mask. Positions are measured *within* the sequence, so varlen does not offset by the packed-token base. Same - kernel path and restrictions as ``bias``; the two may be combined. + kernel path as ``bias`` and may be combined with it, but unlike + ``bias`` it is not supported with paged KV (raises + NotImplementedError), nor with fp8. sink: fp32 ``[H]`` per-head attention-sink logit -- one extra softmax denominator term that has no matching V row:: @@ -735,8 +757,9 @@ def flydsl_flash_attn_func( post-sm_scale logit space as the scores. Applied in the epilogue, so it touches no score element; under split-K the per-split partials stay sink-free and the combine pass folds it in exactly once. Same - kernel path and restrictions as ``bias``; freely combinable with it - and with ``alibi_slopes``. + kernel path and restrictions as ``alibi_slopes`` -- not supported + with paged KV or fp8 -- but freely combinable with ``bias`` and + ``alibi_slopes``. q_descale / k_descale / v_descale: fp32 shape-[1] descales required for dense fp8 e4m3fn inputs. out: Optional pre-allocated output tensor. For fp8, output is bf16; @@ -790,6 +813,9 @@ def flydsl_flash_attn_func( raise ValueError(f"flydsl_flash_attn_func: bias dtype must match q dtype {q.dtype}, got {bias.dtype}") if bias.dim() != 2: raise ValueError(f"flydsl_flash_attn_func: bias must be 2D, got {bias.dim()}D") + _bias_err = bias_addressing_error(bias.shape[0] * bias.shape[1], bias.element_size()) + if _bias_err is not None: + raise ValueError(f"flydsl_flash_attn_func: bias {tuple(bias.shape)} {_bias_err}") if has_alibi: if alibi_slopes.dtype != torch.float32: raise ValueError(f"flydsl_flash_attn_func: alibi_slopes must be float32, got {alibi_slopes.dtype}") @@ -890,9 +916,11 @@ def flydsl_flash_attn_func( f"flydsl_flash_attn_func: varlen bias must be [total_q, max_seqlen_kv] with " f"total_q={q.shape[0]}, got {tuple(bias.shape)}" ) - if max_seqlen_kv is not None and bias.shape[1] < int(max_seqlen_kv): + _bias_cols_min = int(max_seqlen_kv) if cross else Sq + if bias.shape[1] < _bias_cols_min: + _bound = "max_seqlen_kv" if cross else "max_seqlen_q, the self-attention KV maximum" raise ValueError( - f"flydsl_flash_attn_func: varlen bias needs >= max_seqlen_kv={int(max_seqlen_kv)} " + f"flydsl_flash_attn_func: varlen bias needs >= {_bound}={_bias_cols_min} " f"columns, got {bias.shape[1]}" ) elif tuple(bias.shape) != (Sq, Skv): diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index fbf945d91..2550bd7eb 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -692,6 +692,29 @@ def _bias_dma_src_elem(traits, row_base, tile_col_base, d, lane_in_warp, bias_st uni_s = rocdl.readfirstlane(T.i32, as_mlir_value(uni)) return fx.Int32(uni_s) + row_in_group * stride + gran * fx.Int32(_bias_gran_elems(traits)) +BIAS_MAX_OFFSET_ELEMS = 2**31 - 1 +BIAS_MAX_DESCRIPTOR_BYTES = 0xFFFFFFFF + + +def bias_addressing_error(elems, elem_size): + """Reason an `elems`-element bias cannot be addressed, or None if it fits. + + A bias past either limit wraps the i32 offset or overruns the descriptor, so + the kernel would silently read the wrong rows instead of failing. Callers + prefix their own entry-point name and the bias shape. + """ + if elems > BIAS_MAX_OFFSET_ELEMS: + return ( + f"has {elems} elements, exceeding the {BIAS_MAX_OFFSET_ELEMS} " + f"addressable by the kernel's i32 bias element offsets" + ) + if elems * elem_size > BIAS_MAX_DESCRIPTOR_BYTES: + return ( + f"occupies {elems * elem_size} bytes, exceeding the " + f"{BIAS_MAX_DESCRIPTOR_BYTES} byte range of the bias buffer descriptor" + ) + return None + def _bias_lds_lane_base(traits, buf, wave_id, lane_mod_32, lane_div_32, vec_elems): return ( @@ -4236,7 +4259,7 @@ def _store_splitk_partial_o_row(self, v_o, local_opart_row_base, opart_rsrc): for g in range_constexpr(2): self._store_splitk_partial_o_quad(v_o, dc, g, local_opart_row_base, opart_rsrc) - def zero_o_block_if_needed(self, causal_end_raw_i32=None): + def zero_o_block_if_needed(self, causal_end_raw_i32=None, sink_log2=None): if causal_end_raw_i32 is None: causal_end_raw_i32 = self.causal_end_raw_i32 traits = self.traits @@ -4245,6 +4268,9 @@ def zero_o_block_if_needed(self, causal_end_raw_i32=None): lane_mod_32 = self.lane_mod_32 seq_len_v = self.seq_len_v + lse_m_z = self.c_zero_f if sink_log2 is None else sink_log2 + lse_l_z = self.c_zero_f if sink_log2 is None else fx.Float32(1.0) + @flyc.jit def _zero_o_block_if_needed(): if causal_end_raw_i32 <= fx.Int32(0): @@ -4262,6 +4288,8 @@ def _zero_o_block_if_needed(): _store_atom_128=self.store_atom_128, o_div=self.o_div, ) + if const_expr(traits.RETURN_LSE): + self._store_lse_row(lse_m_z, lse_l_z, q_row_z) _zero_o_block_if_needed() diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 24686583b..ec89e419b 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -33,7 +33,13 @@ import pytest # noqa: E402 from flydsl.runtime.device import get_rocm_arch # noqa: E402 +from kernels.attention.flash_attn_gfx950 import build_flash_attn_dualwave_swp_module # noqa: E402 from kernels.attention.flash_attn_interface import flydsl_flash_attn_func # noqa: E402 +from kernels.attention.flash_attn_utils import ( # noqa: E402 + BIAS_MAX_DESCRIPTOR_BYTES, + BIAS_MAX_OFFSET_ELEMS, + bias_addressing_error, +) from tests.test_common import run_perftest # noqa: E402 UNIFORM_RANGE = (-1, 1) @@ -41,9 +47,6 @@ PAGED_KV_MIN_CONTEXT_LENGTH = 16384 # Target share of the softmax mass placed on the attention sink (see calibrate_sink). DEFAULT_SINK_SHARE = 0.5 -# Attention-bias addressing limits (see bias_fits). -BIAS_BUFFER_MAX_BYTES = 0xFFFFFFFF # buffer descriptor num_records clamp -BIAS_I32_MAX_ELEMS = 2**31 - 1 # kernel computes bias element offsets in i32 # fp8 correctness gate (fixed; fp8 is lossy). FP8_MAX_ERR = 5e-2 FP8_MIN_COS = 0.98 @@ -374,12 +377,9 @@ def _ceil_div(a, b): def bias_fits(rows, cols, elem_size=2): - elems = rows * cols - if elems > BIAS_I32_MAX_ELEMS: - return False, f"bias {rows}x{cols} = {elems:.3g} elems exceeds i32 offset range" - if elems * elem_size > BIAS_BUFFER_MAX_BYTES: - return False, f"bias {elems * elem_size / 2**30:.1f} GiB exceeds 4 GiB buffer limit" - return True, "" + """Skip predicate mirroring the kernel's own bias addressing limits.""" + why = bias_addressing_error(rows * cols, elem_size) + return why is None, why or "" def _block_table_from_indices(kv_indptr_cpu, kv_indices_cpu, batch_size, max_num_pages_per_seq): @@ -867,6 +867,7 @@ def _precompute_paged_kv_inputs_and_ref( page_size, kv_cache_layout, trigger_lazy_else=False, + use_bias=False, ): invalid_layout = _validate_kv_cache_layout(kv_cache_layout, page_size, head_dim, dtype) if invalid_layout is not None: @@ -888,7 +889,11 @@ def _precompute_paged_kv_inputs_and_ref( page_size=page_size, kv_cache_layout=kv_cache_layout, trigger_lazy_else=trigger_lazy_else, + use_bias=use_bias, ) + # ref_t folds in inputs["bias"]; a None here would compare the biased kernel + # run against an unbiased reference and report a meaningless PASS. + assert not use_bias or inputs["bias"] is not None, "use_bias=True produced no paged bias tensor" return inputs, ref_t, None @@ -1162,6 +1167,10 @@ def run_attn_config( bias_t = precomputed_inputs["bias"] alibi_t = precomputed_inputs["alibi_slopes"] sink_t = precomputed_inputs["sink"] + # ref_t is built from these same inputs, so a missing bias makes both sides + # unbiased and the comparison vacuous. Fail loudly instead. + if use_bias and bias_t is None: + return {"err": "use_bias=True but the precomputed inputs carry no bias"} debug_counts = torch.zeros(2, dtype=torch.float32, device=device) if debug_lazy else None o_t = torch.zeros_like(q_t) @@ -2594,8 +2603,9 @@ def main(): action="store_true", help="Add an additive attention bias to the scores: softmax(q@k^T * sm_scale + bias). " "Dense bias is [Sq, Skv] broadcast over batch and head; varlen bias is packed " - "[total_q, max_seqlen_kv] with global q rows and batch-local key columns. " - "gfx950 bf16/f16 D=64/128 only; incompatible with --block-table and fp8. Rows whose " + "[total_q, max_seqlen_kv] with global q rows and batch-local key columns. Combines " + "with --block-table, where columns stay the logical (batch-local) key positions. " + "gfx950 bf16/f16 D=64/128 only; incompatible with fp8. Rows whose " "bias exceeds the i32 offset / 4 GiB buffer limits are SKIPped.", ) parser.add_argument( @@ -2846,6 +2856,7 @@ def main(): page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", trigger_lazy_else=args.trigger_lazy_else, + use_bias=args.bias, ) if precompute_status is not None: rows.append((cfg, precompute_status, precompute_status, {"skip": True})) @@ -3121,6 +3132,7 @@ def _cmp_avg(label, subset): seed=args.seed, page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", + use_bias=args.bias, ) if precompute_status is not None: varlen_cmp_rows.append( @@ -3300,6 +3312,7 @@ def _extra_cmp_avg(label, subset): page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", trigger_lazy_else=args.trigger_lazy_else, + use_bias=args.bias, ) ) if precompute_status is not None: @@ -3437,6 +3450,7 @@ def _normal_avg_fn(label, subset): seed=args.seed, page_size=args.page_size, kv_cache_layout=kv_cache_layout or "linear", + use_bias=args.bias, ) ) if precompute_status is not None: @@ -3831,8 +3845,8 @@ def test_bias_paged(causal, kv_cache_layout): """ dtype = torch.bfloat16 B, Sq, H, Hkv, D = 2, 512, 8, 4, 128 - # Uniform KV lengths: ragged per-batch seqlen_k on the dense paged path already - # disagrees with this reference without any bias, so keep that out of scope here. + # Uniform KV lengths: the dense paged path forwards only max_seqlen_kv, so ragged + # lengths are rejected outright (see test_bias_paged_rejects_ragged_seqlen_k). kv_lens = [Sq, Sq] max_kv = max(kv_lens) setup_seed(DEFAULT_SEED) @@ -3874,6 +3888,303 @@ def test_bias_paged(causal, kv_cache_layout): assert (out_nb.float() - out.float()).abs().max().item() > 1e-2, "bias had no effect on the paged output" +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +def test_bias_paged_rejects_ragged_seqlen_k(causal): + """Dense paged bias rejects ragged per-batch seqlen_k instead of answering wrongly. + + The dense paged launch reduces seqlen_k to a single max_seqlen_kv and never + forwards the per-batch lengths, so a shorter batch would attend KV slots it + does not own and mask against the wrong bottom-right offset. + """ + dtype = torch.bfloat16 + B, Sq, H, Hkv, D = 2, 512, 8, 4, 128 + kv_lens = [256, 512] + max_kv = max(kv_lens) + setup_seed(DEFAULT_SEED) + q = torch.empty(B, Sq, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + ragged = _build_paged_kv_for_test(B, max_kv, 64, Hkv, D, kv_lens, dtype, "cuda", "linear") + bias = torch.empty(Sq, max_kv, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + paged_kw = dict(causal=causal, num_kv_heads=Hkv, max_seqlen_kv=max_kv, kv_cache_layout="linear") + + with pytest.raises(NotImplementedError, match="uniform seqlen_k"): + flydsl_flash_attn_func( + q, + ragged["k_cache"], + ragged["v_cache"], + bias=bias, + block_table=ragged["block_table"], + seqlen_k=ragged["seqlen_k"], + **paged_kw, + ) + + # The guard is about raggedness alone: identical shapes with uniform lengths run. + uniform = _build_paged_kv_for_test(B, max_kv, 64, Hkv, D, [max_kv] * B, dtype, "cuda", "linear") + flydsl_flash_attn_func( + q, + uniform["k_cache"], + uniform["v_cache"], + bias=bias, + block_table=uniform["block_table"], + seqlen_k=uniform["seqlen_k"], + **paged_kw, + ) + torch.cuda.synchronize() + + +@_requires_gfx950 +@pytest.mark.parametrize("causal", [False, True]) +def test_bias_paged_varlen_ragged_seqlen_k(causal): + """The varlen paged path -- what the dense rejection points callers at -- is correct. + + cu_seqlens_kv carries the per-batch KV lengths into the kernel, so ragged + lengths mask and bottom-right-align per batch instead of against one global max. + """ + dtype = torch.bfloat16 + H, Hkv, D = 8, 4, 128 + sq, skv = [512, 512], [256, 512] + setup_seed(DEFAULT_SEED) + cu_q = torch.tensor([0, sq[0], sum(sq)], dtype=torch.int32, device="cuda") + cu_kv = torch.tensor([0, skv[0], sum(skv)], dtype=torch.int32, device="cuda") + total_q, max_q, max_kv = sum(sq), max(sq), max(skv) + q = torch.empty(total_q, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + kv_cache = _build_paged_kv_for_test(len(sq), max_kv, 64, Hkv, D, skv, dtype, "cuda", "linear") + bias = torch.empty(total_q, max_kv, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + + out = flydsl_flash_attn_func( + q, + kv_cache["k_cache"], + kv_cache["v_cache"], + causal=causal, + num_kv_heads=Hkv, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_q=max_q, + max_seqlen_kv=max_kv, + cross_seqlen=True, + block_table=kv_cache["block_table"], + seqlen_k=kv_cache["seqlen_k"], + kv_cache_layout="linear", + bias=bias, + ) + torch.cuda.synchronize() + + for b, n in enumerate(skv): + s0, s1 = int(cu_q[b]), int(cu_q[b + 1]) + kb, vb = _logical_kv_from_pages( + kv_cache["k_cache"][_page_ids_for_batch(kv_cache, b)], + kv_cache["v_cache"][_page_ids_for_batch(kv_cache, b)], + "linear", + n, + ) + ref = pytorch_ref_attention_qkv_diff( + q[s0:s1].unsqueeze(0).float(), + kb.unsqueeze(0).float(), + vb.unsqueeze(0).float(), + causal=causal, + bias=bias[s0:s1, :n], + ).squeeze(0) + _, _, passed = _acc_metric(out[s0:s1].float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"varlen paged batch {b} (Sq={sq[b]}, Skv={n}, causal={causal}) does not match" + + +# ── attention bias: addressing limits ──────────────────────────────────────── +# +# The kernel computes bias element offsets as `row * stride + column` in signed +# i32 and describes the bias with a 32-bit-num_records buffer descriptor. A bias +# past either limit is unrepresentable, so it must be rejected up front instead +# of silently reading the wrong rows. + +# 2^31 elements at row 32768, and 4,295,098,368 bytes: over both limits at once. +_OVERSIZED_BIAS_SHAPE = (32769, 65536) +_OVERSIZED_BIAS_MATCH = "i32 bias element offsets" + + +def _unbacked_bias(rows, cols, dtype=torch.bfloat16): + """A [rows, cols] bias with zero-stride storage: shape without the allocation.""" + return torch.zeros(1, 1, dtype=dtype, device="cuda").expand(rows, cols) + + +@pytest.mark.parametrize( + "rows,cols,elem_size,expect", + [ + # The i32 element offset is the binding limit for the 2-byte bias dtypes. + (*_OVERSIZED_BIAS_SHAPE, 2, "i32 bias element offsets"), + (BIAS_MAX_OFFSET_ELEMS + 1, 1, 2, "i32 bias element offsets"), + (BIAS_MAX_OFFSET_ELEMS, 1, 2, None), # exactly at the limit still fits + (46340, 46340, 2, None), # ~4 GiB, the largest square bias that fits + (65536, 32768, 2, "i32 bias element offsets"), # exactly 2^31 elements: one over + # A 4-byte element trips the descriptor limit while the offset still fits. + (BIAS_MAX_OFFSET_ELEMS, 1, 4, "bias buffer descriptor"), + (BIAS_MAX_DESCRIPTOR_BYTES // 4, 1, 4, None), + ], +) +def test_bias_addressing_error_limits(rows, cols, elem_size, expect): + why = bias_addressing_error(rows * cols, elem_size) + if expect is None: + assert why is None, f"bias {rows}x{cols} ({elem_size}B) should fit, got: {why}" + else: + assert why is not None, f"bias {rows}x{cols} ({elem_size}B) should be rejected" + assert expect in why, f"unexpected reason for {rows}x{cols} ({elem_size}B): {why}" + + +def test_bias_dense_rejects_unaddressable(): + dtype = torch.bfloat16 + B, S, H, D = 1, 128, 4, 128 + q = torch.zeros(B, S, H, D, dtype=dtype, device="cuda") + bias = _unbacked_bias(*_OVERSIZED_BIAS_SHAPE, dtype=dtype) + with pytest.raises(ValueError, match=_OVERSIZED_BIAS_MATCH): + flydsl_flash_attn_func(q, q.clone(), q.clone(), causal=False, num_kv_heads=H, bias=bias) + + +def test_bias_varlen_rejects_unaddressable(): + dtype = torch.bfloat16 + seqs = [128, 128] + total, max_s = sum(seqs), max(seqs) + H, Hkv, D = 4, 4, 128 + cu = torch.tensor([0, seqs[0], total], dtype=torch.int32, device="cuda") + q = torch.zeros(total, H, D, dtype=dtype, device="cuda") + kv = torch.zeros(total, Hkv, D, dtype=dtype, device="cuda") + bias = _unbacked_bias(*_OVERSIZED_BIAS_SHAPE, dtype=dtype) + with pytest.raises(ValueError, match=_OVERSIZED_BIAS_MATCH): + flydsl_flash_attn_func( + q, + kv, + kv.clone(), + causal=False, + num_kv_heads=Hkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_s, + max_seqlen_kv=max_s, + cross_seqlen=False, + bias=bias, + ) + + +@_requires_gfx950 +def test_bias_varlen_self_attn_rejects_narrow_bias(): + """Varlen self-attention bounds bias columns by max_seqlen_q, not max_seqlen_kv. + + max_seqlen_kv is legitimately None when cross_seqlen=False, so a too-narrow + bias used to pass validation: the kernel then indexes key column j with + bias_stride0 = bias.shape[1], reading the following bias rows instead of failing. + """ + dtype = torch.bfloat16 + seqs = [512, 384] + total, max_s = sum(seqs), max(seqs) + H, Hkv, D = 8, 4, 128 + setup_seed(DEFAULT_SEED) + cu = torch.tensor([0, seqs[0], total], dtype=torch.int32, device="cuda") + q = torch.empty(total, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(total, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(total, Hkv, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + self_attn_kw = dict( + causal=False, + num_kv_heads=Hkv, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_s, + cross_seqlen=False, + ) + + for cols in (1, max_s - 1): + narrow = torch.zeros(total, cols, dtype=dtype, device="cuda") + with pytest.raises(ValueError, match="self-attention KV maximum"): + flydsl_flash_attn_func(q, k, v, bias=narrow, **self_attn_kw) + + # A bias exactly at the bound still runs, and matches the per-batch reference. + bias = torch.empty(total, max_s, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + out = flydsl_flash_attn_func(q, k, v, bias=bias, **self_attn_kw) + torch.cuda.synchronize() + for b, n in enumerate(seqs): + s0, s1 = int(cu[b]), int(cu[b + 1]) + ref = pytorch_ref_attention( + q[s0:s1].unsqueeze(0).float(), + k[s0:s1].unsqueeze(0).float(), + v[s0:s1].unsqueeze(0).float(), + causal=False, + bias=bias[s0:s1, :n], + ).squeeze(0) + _, _, passed = _acc_metric(out[s0:s1].float().reshape(-1), ref.float().reshape(-1), D) + assert passed, f"varlen self-attention batch {b} (seqlen {n}) does not match the biased reference" + + +def test_bias_paged_rejects_unaddressable(): + dtype = torch.bfloat16 + B, Sq, H, Hkv, D = 1, 128, 4, 4, 128 + q = torch.zeros(B, Sq, H, D, dtype=dtype, device="cuda") + kv_cache = _build_paged_kv_for_test(B, Sq, 64, Hkv, D, [Sq], dtype, "cuda", "linear") + bias = _unbacked_bias(*_OVERSIZED_BIAS_SHAPE, dtype=dtype) + with pytest.raises(ValueError, match=_OVERSIZED_BIAS_MATCH): + flydsl_flash_attn_func( + q, + kv_cache["k_cache"], + kv_cache["v_cache"], + causal=True, + num_kv_heads=Hkv, + max_seqlen_kv=Sq, + block_table=kv_cache["block_table"], + seqlen_k=kv_cache["seqlen_k"], + kv_cache_layout="linear", + bias=bias, + ) + + +def test_precompute_paged_bias_reaches_inputs_and_reference(): + """`--block-table --bias` must generate a bias AND fold it into the reference. + + The helper used to ignore use_bias, so the paged benchmark ran unbiased and + compared against an unbiased reference: a PASS that measured nothing. + """ + kw = dict( + batch=1, + seqlen_q=256, + seqlen_kv=None, + varlen_seqlens_q=None, + varlen_seqlens_kv=None, + num_heads=4, + head_dim=128, + num_kv_heads=4, + dtype=torch.bfloat16, + causal=False, + seed=DEFAULT_SEED, + page_size=64, + kv_cache_layout="linear", + ) + biased_inputs, biased_ref, biased_status = _precompute_paged_kv_inputs_and_ref(**kw, use_bias=True) + plain_inputs, plain_ref, plain_status = _precompute_paged_kv_inputs_and_ref(**kw) + assert biased_status is None and plain_status is None + assert biased_inputs["bias"] is not None, "use_bias=True must generate a paged bias" + assert plain_inputs["bias"] is None, "use_bias defaults to no bias" + + # The bias is drawn after Q/K/V, so the same seed leaves the inputs identical + # and any reference difference is the bias alone. + for key in ("q_t", "k_t", "v_t"): + assert torch.equal(biased_inputs[key], plain_inputs[key]), f"{key} must not depend on use_bias" + assert ( + biased_ref.float() - plain_ref.float() + ).abs().max().item() > 1e-2, "the paged reference must fold in the generated bias" + + +@_requires_gfx950 +def test_bias_launcher_rejects_unaddressable(): + """The kernel launcher guards too, for callers that bypass flydsl_flash_attn_func. + + The guard also has to fire before the launcher's ``bias.contiguous()``, which + would otherwise materialize gigabytes on the way to a guaranteed failure. + """ + dtype = torch.bfloat16 + B, S, H, D = 1, 384, 4, 128 + launch = build_flash_attn_dualwave_swp_module(num_heads=H, head_dim=D, causal=False, has_bias=True) + q, k, v, o = (torch.zeros(B, S, H, D, dtype=dtype, device="cuda") for _ in range(4)) + bias = _unbacked_bias(*_OVERSIZED_BIAS_SHAPE, dtype=dtype) + free_before = torch.cuda.mem_get_info()[0] + with pytest.raises(ValueError, match=_OVERSIZED_BIAS_MATCH): + launch(q, k, v, o, B, S, bias=bias) + assert torch.cuda.mem_get_info()[0] > free_before - 2**30, "rejected bias must not be materialized" + + # ── ALiBi ──────────────────────────────────────────────────────────────────── @@ -4135,3 +4446,38 @@ def test_sink_splitk_counted_once(num_kv_splits): # A sink counted once per split would shift LSE by ~ln(num_kv_splits); assert # we are nowhere near that, so the test cannot pass on a double-count. assert (lsek - lse1).abs().max().item() < 0.5 * math.log(num_kv_splits) + + +@_requires_gfx950 +@pytest.mark.parametrize("Sq,Skv", [(512, 128), (512, 160)]) +def test_sink_lse_cross_attn_skipped_blocks(Sq, Skv): + """Causal cross-attention with Skv < Sq skips whole q blocks that see no key. + + A skipped block never reaches the main body's fold_sink, so the skip path has + to write those rows' LSE itself: it is the per-head sink, not -inf and not + whatever the caller's output buffer happened to hold. Skv=160 also puts some + all-masked rows inside an active block, covering both paths at once. + """ + dtype = torch.bfloat16 + B, H, D = 2, 8, 128 + setup_seed(DEFAULT_SEED) + q = torch.empty(B, Sq, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + k = torch.empty(B, Skv, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + v = torch.empty(B, Skv, H, D, dtype=dtype, device="cuda").uniform_(*UNIFORM_RANGE) + sink = _sink_for(q, k, True) + + out, lse = flydsl_flash_attn_func(q, k, v, causal=True, sink=sink, return_lse=True) + torch.cuda.synchronize() + + # Sink-inclusive LSE = ln(exp(LSE_no_sink) + exp(sink)); -inf rows collapse to sink. + lse_ref_ns = _reference_lse(q, k, True, H) # B, H, Sq + assert bool((~torch.isfinite(lse_ref_ns)).any()), "test setup should produce fully-masked q rows" + lse_ref = torch.logaddexp(lse_ref_ns, sink.view(1, H, 1).expand_as(lse_ref_ns)) + diff = (lse.float() - lse_ref).abs().max().item() + assert diff <= _ATOL_BF16, f"sink-inclusive LSE max abs diff {diff:.3e} exceeds atol {_ATOL_BF16:.3e}" + + # The all-masked rows are the regression: their whole denominator is the sink. + n_masked = Sq - Skv + masked_lse = lse[:, :, :n_masked].float() + assert (masked_lse - sink.view(1, H, 1)).abs().max().item() <= 1e-4, "all-masked rows must carry the sink LSE" + assert out[:, :n_masked].abs().max().item() == 0.0, "all-masked rows must have zero output" From 44fc566667c3a07e4f55a7b06bcbe6fed7684fa8 Mon Sep 17 00:00:00 2001 From: Nikolai Protasov Date: Thu, 13 Aug 2026 14:50:05 +0000 Subject: [PATCH 12/12] fix codestyle --- kernels/attention/flash_attn_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index 2550bd7eb..3dc03e31b 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -692,6 +692,7 @@ def _bias_dma_src_elem(traits, row_base, tile_col_base, d, lane_in_warp, bias_st uni_s = rocdl.readfirstlane(T.i32, as_mlir_value(uni)) return fx.Int32(uni_s) + row_in_group * stride + gran * fx.Int32(_bias_gran_elems(traits)) + BIAS_MAX_OFFSET_ELEMS = 2**31 - 1 BIAS_MAX_DESCRIPTOR_BYTES = 0xFFFFFFFF