diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index e72e7e163..4ba01862e 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, @@ -31,8 +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, + _load_bias_frag_lds, + _make_bias_lds_ptr, _make_dualwave_swp_traits, + _num_bias_dma, _s_barrier, _s_nop, _s_setprio, @@ -40,11 +52,14 @@ _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, _v_vec32_to_pair, _waitcnt_vm_n, + bias_addressing_error, ) from kernels.common.kernels_common import dtype_to_elem_type @@ -67,13 +82,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 +122,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,18 +148,37 @@ 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) + # 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 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, 128] - if const_expr(traits.PAGED): + 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, 128] + else: @fx.struct @@ -149,12 +196,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 +255,134 @@ def flash_attn_dualwave_swp_gfx950_kernel( gemm_helper = DualwaveGemmHelper(ctx) softmax_helper = DualwaveSoftmaxHelper(ctx) + # ---------------- attention bias ---------------- + if const_expr(HAS_BIAS): + _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_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) + _bias_lds_base_idx = fx.Index(fx.ptrtoint(ctx.lds.bias.ptr)) + _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) + ) + _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): + _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) + + _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`.""" + 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) + 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=bias_stride0, + ), + 0, + _dma_atom=ctx.dma_atom, + _lds_ptr_ty=ctx.lds_ptr_ty, + ) + + def _read_bias_frag(h, g, buf): + gran = _seq_pad_score_threshold(traits, g * _BIAS_VEC) // _bias_gran_elems + _bias_half_grans * h + 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) + 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): + 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, 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) + if const_expr(HAS_BIAS): + _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): + if const_expr(prefetch_next): + _issue_bias_dma(tile_idx + fx.Index(1), 1 - buf) + else: + _sched_barrier(0) + return scored + def _main_body(): # Paged: stage the block-table row into LDS before any page-id ds_read. if const_expr(traits.PAGED): @@ -221,6 +401,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) @@ -249,7 +432,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), buf=0) _sched_barrier(0) if const_expr(traits.CAUSAL): @@ -326,7 +510,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, 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) @@ -401,7 +585,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, 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) @@ -492,7 +676,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, 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) @@ -559,7 +743,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, 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) @@ -618,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 = gemm_helper.qk(v_k, q_all_scaled_bf16) + 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) @@ -671,6 +855,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) @@ -689,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() @@ -715,11 +904,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 +919,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 +941,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 +951,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 +984,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 +1007,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 +1022,74 @@ 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: + 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 +1108,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 +1149,9 @@ def _launch( block_table = O if block_table_stride is None: block_table_stride = 0 + 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): if stream is None: return launch_flash_attn_dualwave_swp( @@ -884,6 +1164,9 @@ def _launch( cu_seqlens_q, cu_seqlens_kv, block_table, + bias, + alibi_slopes, + sink, batch_size, seq_len, seq_len_kv, @@ -891,6 +1174,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 +1187,9 @@ def _launch( cu_seqlens_q, cu_seqlens_kv, block_table, + bias, + alibi_slopes, + sink, batch_size, seq_len, seq_len_kv, @@ -909,6 +1197,8 @@ def _launch( stride_kv_n, head_dim_runtime, block_table_stride, + bias_stride0, + alibi_stride_b, stream=stream, ) @@ -930,6 +1220,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 +1255,10 @@ def _compile( block_table = O if block_table_stride is None: block_table_stride = 0 + + 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): return flyc.compile( launch_flash_attn_dualwave_swp, @@ -973,6 +1271,9 @@ def _compile( cu_seqlens_q, cu_seqlens_kv, block_table, + bias, + alibi_slopes, + sink, batch_size, seq_len, seq_len_kv, @@ -980,6 +1281,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..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"] @@ -133,6 +132,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 +153,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 +202,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 +224,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 +279,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 +299,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, ) @@ -305,6 +322,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). @@ -338,6 +356,7 @@ def _build_paged( dualwave_swp_setprio=setprio, dualwave_swp_enable_stagger=enable_stagger, return_lse=return_lse, + has_bias=has_bias, ) @@ -355,6 +374,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], @@ -455,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: @@ -469,6 +490,30 @@ 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: + 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 + 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 = ( @@ -481,6 +526,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) @@ -518,6 +564,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) @@ -528,6 +575,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 @@ -610,6 +659,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 / paged KV). + bias: Optional[torch.Tensor] = None, + # Per-head ALiBi slope table, computed analytically into the scores. Same + # 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 `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, k_descale: Optional[torch.Tensor] = None, @@ -659,8 +717,49 @@ 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. 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 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:: + + 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 ``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; @@ -697,6 +796,36 @@ 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 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") + 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") + _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}") + 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, @@ -704,6 +833,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, @@ -778,6 +908,39 @@ 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)}" + ) + _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 >= {_bound}={_bias_cols_min} " + 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 +972,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 +1020,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 +1045,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 +1073,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 +1121,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 +1142,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/kernels/attention/flash_attn_utils.py b/kernels/attention/flash_attn_utils.py index c8b5a7b47..3dc03e31b 100644 --- a/kernels/attention/flash_attn_utils.py +++ b/kernels/attention/flash_attn_utils.py @@ -643,6 +643,103 @@ 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): + 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 fx.Int32(lds_addr) + + +def _bias_dma_src_elem(traits, row_base, tile_col_base, d, lane_in_warp, bias_stride0_v): + lanes_per_row = _bias_lanes_per_row(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)) + + +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 ( + 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, + alias_scopes=_dualwave_lds_alias_scopes(_dualwave_lds_scope("bias", 0)), + ).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)) @@ -3756,6 +3853,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) @@ -4155,7 +4260,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 @@ -4164,6 +4269,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): @@ -4181,6 +4289,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() @@ -5291,6 +5401,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__) @@ -5302,6 +5413,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 @@ -5406,6 +5518,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) diff --git a/tests/kernels/test_flash_attn_fwd.py b/tests/kernels/test_flash_attn_fwd.py index 982e8876d..ec89e419b 100644 --- a/tests/kernels/test_flash_attn_fwd.py +++ b/tests/kernels/test_flash_attn_fwd.py @@ -33,12 +33,20 @@ 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) 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 # fp8 correctness gate (fixed; fp8 is lossy). FP8_MAX_ERR = 5e-2 FP8_MIN_COS = 0.98 @@ -187,7 +195,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 +334,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 +355,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,8 +376,14 @@ def _ceil_div(a, b): return (a + b - 1) // b +def bias_fits(rows, cols, elem_size=2): + """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): - 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() @@ -356,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, @@ -451,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) @@ -460,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 @@ -491,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, @@ -527,6 +647,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 +688,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 +760,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 +782,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 +808,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 +820,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"], @@ -690,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: @@ -711,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 @@ -871,6 +1053,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 +1102,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 +1142,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 +1156,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 +1164,13 @@ 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"] + # 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) @@ -977,6 +1195,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(), ) @@ -993,6 +1212,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 +1239,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() @@ -1099,6 +1312,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(), ) @@ -1115,6 +1329,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 +1373,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 +1390,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 +1477,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 +2598,50 @@ 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. 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( + "--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 +2718,24 @@ 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") + # 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: + 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 +2774,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 +2797,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( @@ -2560,16 +2856,16 @@ 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})) 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 +2881,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 +2944,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 +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, ) asm_r = run_aiter_bench( batch, @@ -2715,6 +3024,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 +3081,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}") @@ -2820,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( @@ -2856,6 +3169,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 +3210,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 +3246,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: @@ -2991,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: @@ -3021,6 +3343,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 +3397,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}") @@ -3123,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: @@ -3179,6 +3507,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 +3609,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 +3761,531 @@ 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" + + +@_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: 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) + 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" + + +@_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 ──────────────────────────────────────────────────────────────────── + + +@_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 +4331,153 @@ 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) + + +@_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"