diff --git a/kernels/attention/flash_attn_generic.py b/kernels/attention/flash_attn_generic.py index 20bb63356..872538890 100644 --- a/kernels/attention/flash_attn_generic.py +++ b/kernels/attention/flash_attn_generic.py @@ -94,6 +94,7 @@ def build_flash_attn_func_module_primary( varlen=False, paged=False, kv_cache_layout="linear", + return_lse=False, ): """Build the flash_attn_func launcher using the post-refactor FlyDSL API. @@ -211,6 +212,11 @@ def _guarded(*args, **kwargs): GQA_GROUP_SIZE = NUM_HEADS_Q // NUM_HEADS_KV HEAD_DIM = head_dim CAUSAL = causal + # Emit per-row log-sum-exp (LSE). Convention (A1<->A3 contract): natural log, + # softmax scale folded in, i.e. LSE_i = ln(sum_j exp(sm_scale * (q_i . k_j))). + # Output tensor is fp32 [batch, num_heads, seq_len] (varlen: [batch, num_heads, + # max_seqlen_q], padded). See docs/flash_attn_lse_contract.md. + RETURN_LSE = bool(return_lse) # Varlen uses packed QKV via launch-time cu_seqlens. VARLEN = bool(varlen) CROSS_SEQLEN = bool(cross_seqlen) @@ -300,6 +306,7 @@ def flash_attn_generic_kernel( K: fx.Tensor, V: fx.Tensor, O: fx.Tensor, # noqa: E741 + LSE: fx.Tensor, seq_len: fx.Int32, seq_len_kv: fx.Int32, CuSeqQ: fx.Tensor, @@ -736,6 +743,31 @@ def coop_dma_v_nomajor(tile_start, buf_id=0): O, max_size=False, num_records_bytes=_q_nrec_bytes, base_byte_offset=_q_batch_byte_off ) + # LSE output: fp32 [batch, num_heads, seq_len] (varlen: padded to max_seqlen_q). + # Per-batch descriptor folds the batch offset into the base so the flat index + # stays 0-based within the batch. Only built when RETURN_LSE. + if const_expr(RETURN_LSE): + _lse_per_batch_elems = fx.Index(NUM_HEADS_Q) * seq_len_v + _lse_nrec_bytes = _raw(_lse_per_batch_elems * fx.Index(4)) + _lse_batch_byte_off = _raw(batch_idx * _lse_per_batch_elems * fx.Index(4)) + lse_rsrc = buffer_ops.create_buffer_resource( + LSE, max_size=False, num_records_bytes=_lse_nrec_bytes, base_byte_offset=_lse_batch_byte_off + ) + # Row index within the batch region and an out-of-bounds sentinel (== num + # records) that the hardware buffer bound drops, used to predicate the store. + _lse_oob_off = _lse_per_batch_elems + + def _store_lse(m_final, l_final): + # LSE = sm_scale * m_raw + ln(l); natural log, scale folded (A1<->A3). + lse_val = _fadd(_fmul(m_final, fx.Float32(sm_scale)), fmath.log(_raw(l_final), fastmath=fm_fast)) + lse_local = q_head_idx * seq_len_v + q_row + # One writer per row: low half-wave (lane_div_32 == 0) stores; the high + # half-wave and any partial-tile OOB row (q_row >= seqlen_q_b) are + # redirected to the OOB sentinel offset and dropped by the buffer bound. + off_row = ArithValue(q_row < seqlen_q_b).select(lse_local, _lse_oob_off) + off = fx.Index(ArithValue(lane_div_32 == fx.Index(0)).select(off_row, _lse_oob_off)) + buffer_ops.buffer_store(_raw(fx.Float32(lse_val)), lse_rsrc, _raw(fx.Int32(off))) + # ---- DMA loading for K (buffer_load_dwordx4 ... lds) ---- if const_expr(ENABLE_DMA): k_rsrc = buffer_ops.create_buffer_resource( @@ -1414,6 +1446,9 @@ def _read_v_pack(step_idx): loop_results = yield _yield_args def _zero_o_block(): + if const_expr(RETURN_LSE): + # No valid keys for this q-tile: LSE = ln(0) = -inf for every row. + _store_lse(loop_results[0], loop_results[1]) if const_expr(USE_PERMLANE_OSTORE): zero_pack = Vec.filled(4, 0, fx.Int32) for dc in range_constexpr(D_CHUNKS): @@ -1434,6 +1469,9 @@ def _normalize_and_store_o(): l_final = loop_results[1] o_finals = [loop_results[2 + dc] for dc in range_constexpr(D_CHUNKS)] + if const_expr(RETURN_LSE): + _store_lse(loop_results[0], l_final) + inv_l_rcp = rocdl.rcp(T.f32, l_final) if const_expr(CAUSAL and CROSS_SEQLEN): inv_l = ArithValue(fx.Float32(l_final) > c_zero_f).select(inv_l_rcp, c_zero_f) @@ -1515,6 +1553,7 @@ def launch_flash_attn_generic( K: fx.Tensor, V: fx.Tensor, O: fx.Tensor, # noqa: E741 + LSE: fx.Tensor, CuSeqQ: fx.Tensor, CuSeqKv: fx.Tensor, BlockTable: fx.Tensor, @@ -1548,6 +1587,7 @@ def launch_flash_attn_generic( K, V, O, + LSE, seq_len, seq_len_kv, CuSeqQ, @@ -1593,25 +1633,42 @@ def _launch( block_table=None, block_table_stride=0, seq_len_kv=None, + lse=None, stream=None, ): # Dense/non-paged pass the output tensor as a placeholder for the unused # cu_seqlens / block_table slots; the kernel only reads them under - # const_expr(VARLEN) / const_expr(PAGED). + # const_expr(VARLEN) / const_expr(PAGED). LSE is likewise a placeholder + # unless the kernel was built with return_lse=True (RETURN_LSE store gate). cq = cu_seqlens_q if cu_seqlens_q is not None else Out ck = cu_seqlens_kv if cu_seqlens_kv is not None else Out bt = block_table if block_table is not None else Out + ls = lse if lse is not None else Out skv = seq_len if seq_len_kv is None else seq_len_kv with CompilationContext.compile_hints(_fmha_compile_hints): return launch_flash_attn_generic( - Q, K, V, Out, cq, ck, bt, block_table_stride, batch_size, seq_len, skv, fx.Stream(stream) + Q, K, V, Out, ls, cq, ck, bt, block_table_stride, batch_size, seq_len, skv, fx.Stream(stream) ) - def _compile(Q, K, V, Out, batch_size, seq_len, seq_len_kv=None, stream=None): + def _compile(Q, K, V, Out, batch_size, seq_len, seq_len_kv=None, lse=None, stream=None): + ls = lse if lse is not None else Out skv = seq_len if seq_len_kv is None else seq_len_kv with CompilationContext.compile_hints(_fmha_compile_hints): return flyc.compile( - launch_flash_attn_generic, Q, K, V, Out, Out, Out, Out, 0, batch_size, seq_len, skv, fx.Stream(stream) + launch_flash_attn_generic, + Q, + K, + V, + Out, + ls, + Out, + Out, + Out, + 0, + batch_size, + seq_len, + skv, + fx.Stream(stream), ) _launch.compile = _compile diff --git a/kernels/attention/flash_attn_gfx950.py b/kernels/attention/flash_attn_gfx950.py index ce20a0443..f91e769ba 100644 --- a/kernels/attention/flash_attn_gfx950.py +++ b/kernels/attention/flash_attn_gfx950.py @@ -66,6 +66,7 @@ def build_flash_attn_dualwave_swp_module( cross_seqlen=False, paged=False, kv_cache_layout="linear", + return_lse=False, ): """Build an DUALWAVE_SWP flash_attn launcher for D=64/128 bf16/f16 on gfx950. @@ -194,6 +195,11 @@ class SharedStorage: DUALWAVE_SWP_SETPRIO = bool(dualwave_swp_setprio) DUALWAVE_SWP_DEBUG_LAZY_COUNTS = bool(dualwave_swp_debug_lazy_counts) DUALWAVE_SWP_ENABLE_STAGGER = bool(dualwave_swp_enable_stagger) + # Emit per-row log-sum-exp (LSE). Convention (A1<->A3 contract): natural log, + # softmax scale folded in, i.e. LSE_i = ln(sum_j exp(sm_scale * (q_i . k_j))). + # Output is fp32 [batch, num_heads, seq_len]. For split-K the LSE is finalized + # in the combine kernel from the per-split (m, l). See docs/flash_attn_lse_contract.md. + RETURN_LSE = bool(return_lse) VARLEN = bool(varlen) # Cross-length (seqlen_q != seqlen_kv): emit the extra in-loop v_s_1 causal mask # so a diagonal kv-tile landing on the v_s_1 slot is masked. Off by default so @@ -218,6 +224,7 @@ def flash_attn_dualwave_swp_gfx950_kernel( K: fx.Tensor, V: fx.Tensor, O: fx.Tensor, # noqa: E741 + LSE: fx.Tensor, DebugCounts: fx.Tensor, CuSeqQ: fx.Tensor, CuSeqKv: fx.Tensor, @@ -543,6 +550,7 @@ def _buffer_store_128(pack_i32_vec, elem_index): c_zero_i = fx.Int32(0) head_dim_f32 = fx.Float32(fx.Int32(head_dim_runtime)) c_log2e_f = fx.Float32(_LOG2E) + c_ln2_f = fx.Float32(1.0 / _LOG2E) c_sm_scale_log2e = fx.Float32( arith.mulf( _raw(fmath.rsqrt(head_dim_f32, fastmath=fm_fast)), @@ -2064,6 +2072,26 @@ def _swap_halves(dw): o_pack = Vec.from_elements([fx.Int32(w0), fx.Int32(w1), fx.Int32(w2), fx.Int32(w3)], fx.Int32) o_global = o_base + (dc * D_CHUNK + 2 * g * 8) _buffer_store_128(o_pack, o_global) + + if const_expr(RETURN_LSE): + # LSE = m_row * ln2 + ln(l_row); natural log, scale folded (m_row is + # already sm_scale*log2e-scaled). fp32 [batch, num_heads, seq_len]; + # per-batch descriptor folds the batch offset into the 48-bit base. + _lse_base_i64 = fx.Int64(fx.ptrtoint(fx.get_iter(LSE))) + _lse_per_batch_elems = fx.Index(NUM_HEADS_Q) * seq_len_v + _lse_per_batch_bytes = _lse_per_batch_elems * fx.Index(4) + _lse_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(_lse_base_i64 + fx.Int64(batch_idx * _lse_per_batch_bytes)), + num_records_bytes=_raw(fx.Int64(_lse_per_batch_bytes)), + ) + _lse_val = _fadd(_fmul(m_row, c_ln2_f), fmath.log(_raw(l_row), fastmath=fm_fast)) + _lse_local = q_head_idx * seq_len_v + q_row + # One writer per row (low half-wave); partial-tile OOB rows and the + # high half-wave redirect to the sentinel offset (== num records), + # which the buffer bound drops. + _lse_off_row = ArithValue(q_row < seqlen_q_v).select(_lse_local, _lse_per_batch_elems) + _lse_off = fx.Index(ArithValue(lane < fx.Index(32)).select(_lse_off_row, _lse_per_batch_elems)) + _ws_store_f32(_lse_val, _lse_off, _lse_rsrc) else: # Split-K stores normalized O_partial plus this row's fp32 m/l. # Per-split descriptors fold split offsets into the 48-bit base. @@ -2142,12 +2170,14 @@ def _swap_halves(dw): def flash_attn_splitk_combine_kernel( O: fx.Tensor, # noqa: E741 WS: fx.Tensor, + LSE: fx.Tensor, batch_size: fx.Int32, seq_len: fx.Int32, stride_q_n: fx.Int32, ): elem_dtype = dtype_to_elem_type(dtype_str) fm_fast = fx.arith.FastMathFlags.fast + c_ln2_f = fx.Float32(1.0 / _LOG2E) seq_v = fx.Index(seq_len) stride_v = fx.Index(stride_q_n) bs_v = fx.Index(batch_size) @@ -2245,6 +2275,24 @@ def _accum_split(acc, den): acc, den = _accum_split(acc, den) + if const_expr(RETURN_LSE): + # Combined LSE = m_max * ln2 + ln(den); den = sum_s 2^(m_s - m_max) * l_s is + # the global base-2 denominator relative to m_max, so ln(den) completes the + # natural-log, scale-folded LSE. One lane per (b,h,s) row (col == 0) writes; + # others redirect to the sentinel offset dropped by the buffer bound. + _lse_base_i64_c = fx.Int64(fx.ptrtoint(fx.get_iter(LSE))) + _lse_per_batch_elems_c = fx.Index(NUM_HEADS_Q) * seq_v + _lse_per_batch_bytes_c = _lse_per_batch_elems_c * fx.Index(4) + _lse_rsrc_c = buffer_ops.create_buffer_resource_from_addr( + _raw(_lse_base_i64_c + fx.Int64(b * _lse_per_batch_bytes_c)), + num_records_bytes=_raw(fx.Int64(_lse_per_batch_bytes_c)), + ) + _lse_val_c = _fadd(_fmul(m_max, c_ln2_f), fmath.log(_raw(den), fastmath=fm_fast)) + _lse_off_c = fx.Index( + ArithValue(col == fx.Index(0)).select(local_ml_idx_c, _lse_per_batch_elems_c) + ) + buffer_ops.buffer_store(_raw(fx.Float32(_lse_val_c)), _lse_rsrc_c, _raw(fx.Int32(_lse_off_c))) + inv_rcp = rocdl.rcp(T.f32, den) inv = ArithValue(fx.Float32(den) > fx.Float32(0.0)).select(inv_rcp, fx.Float32(0.0)) inv4 = Vec.from_elements([fx.Float32(inv)], fx.Float32).broadcast_to(4) @@ -2273,6 +2321,7 @@ def launch_flash_attn_dualwave_swp( K: fx.Tensor, V: fx.Tensor, O: fx.Tensor, # noqa: E741 + LSE: fx.Tensor, DebugCounts: fx.Tensor, CuSeqQ: fx.Tensor, CuSeqKv: fx.Tensor, @@ -2308,6 +2357,7 @@ def launch_flash_attn_dualwave_swp( K, V, O, + LSE, DebugCounts, CuSeqQ, CuSeqKv, @@ -2330,7 +2380,7 @@ def launch_flash_attn_dualwave_swp( ) if const_expr(SPLITK): combine_rows = bs_idx * NUM_HEADS_Q * sl_idx - flash_attn_splitk_combine_kernel(O, DebugCounts, batch_size, seq_len, stride_q_n).launch( + flash_attn_splitk_combine_kernel(O, DebugCounts, LSE, batch_size, seq_len, stride_q_n).launch( grid=(combine_rows // COMBINE_ROWS_PER_BLOCK, 1, 1), block=(COMBINE_BLOCK, 1, 1), stream=stream, @@ -2363,6 +2413,7 @@ def _launch( cu_seqlens_kv=None, block_table=None, block_table_stride=None, + lse=None, stream=None, ): if stride_kv_n is None: @@ -2392,6 +2443,9 @@ def _launch( block_table = O if block_table_stride is None: block_table_stride = 0 + # LSE is only written under const_expr(RETURN_LSE); O placeholder otherwise. + if lse is None: + lse = O with CompilationContext.compile_hints(_dualwave_swp_compile_hints): if stream is None: return launch_flash_attn_dualwave_swp( @@ -2399,6 +2453,7 @@ def _launch( K, V, O, + lse, debug_counts, cu_seqlens_q, cu_seqlens_kv, @@ -2416,6 +2471,7 @@ def _launch( K, V, O, + lse, debug_counts, cu_seqlens_q, cu_seqlens_kv, @@ -2448,6 +2504,7 @@ def _compile( cu_seqlens_kv=None, block_table=None, block_table_stride=None, + lse=None, stream=None, ): if stride_kv_n is None: @@ -2472,6 +2529,8 @@ def _compile( block_table = O if block_table_stride is None: block_table_stride = 0 + if lse is None: + lse = O with CompilationContext.compile_hints(_dualwave_swp_compile_hints): return flyc.compile( launch_flash_attn_dualwave_swp, @@ -2479,6 +2538,7 @@ def _compile( K, V, O, + lse, debug_counts, cu_seqlens_q, cu_seqlens_kv, diff --git a/kernels/attention/flash_attn_interface.py b/kernels/attention/flash_attn_interface.py index af8a81082..c135af90c 100644 --- a/kernels/attention/flash_attn_interface.py +++ b/kernels/attention/flash_attn_interface.py @@ -97,6 +97,7 @@ def _build_dense( path_tag: str, waves_per_eu: int, daz: bool, + return_lse: bool = False, ): """Build (and cache) one dense generic launcher variant.""" from kernels.attention.flash_attn_generic import build_flash_attn_func_module @@ -113,6 +114,7 @@ def _build_dense( path_tag=path_tag, waves_per_eu=waves_per_eu, daz=daz, + return_lse=return_lse, ) @@ -130,6 +132,7 @@ def _build_dense_dualwave( setprio: bool, debug_lazy_counts: bool, enable_stagger: bool, + return_lse: bool = False, ): """Build (and cache) the dense gfx950 DUALWAVE_SWP launcher.""" from kernels.attention.flash_attn_gfx950 import build_flash_attn_dualwave_swp_module @@ -147,6 +150,7 @@ def _build_dense_dualwave( dualwave_swp_setprio=setprio, dualwave_swp_debug_lazy_counts=debug_lazy_counts, dualwave_swp_enable_stagger=enable_stagger, + return_lse=return_lse, ) @@ -192,6 +196,7 @@ def _build_varlen( setprio: bool, debug_lazy_counts: bool, enable_stagger: bool, + return_lse: 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 @@ -210,6 +215,7 @@ def _build_varlen( dualwave_swp_setprio=setprio, dualwave_swp_debug_lazy_counts=debug_lazy_counts, dualwave_swp_enable_stagger=enable_stagger, + return_lse=return_lse, ) @@ -227,6 +233,7 @@ def _build_varlen_light( setprio: bool, debug_lazy_counts: bool, enable_stagger: bool, + return_lse: bool = False, ): """Build a lightweight packed-varlen launcher for short attention.""" from kernels.attention.flash_attn_generic import build_flash_attn_func_module @@ -243,6 +250,7 @@ def _build_varlen_light( flat_work_group_size=128, waves_per_eu=waves_per_eu, daz=daz, + return_lse=return_lse, ) @@ -259,6 +267,7 @@ def _build_splitk( lazy_rescale: bool, setprio: bool, enable_stagger: bool, + return_lse: 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 @@ -275,6 +284,7 @@ def _build_splitk( dualwave_swp_lazy_rescale=lazy_rescale, dualwave_swp_setprio=setprio, dualwave_swp_enable_stagger=enable_stagger, + return_lse=return_lse, ) @@ -294,6 +304,7 @@ def _build_paged( num_kv_splits: int = 1, varlen: bool = False, kv_cache_layout: str = "linear", + return_lse: bool = False, ): """Build (and cache) a paged-KV launcher (gfx950 DUALWAVE_SWP, paged=True). @@ -326,6 +337,7 @@ def _build_paged( dualwave_swp_lazy_rescale=lazy_rescale, dualwave_swp_setprio=setprio, dualwave_swp_enable_stagger=enable_stagger, + return_lse=return_lse, ) @@ -546,6 +558,7 @@ def _build_paged_light( setprio: bool, debug_lazy_counts: bool, enable_stagger: bool, + return_lse: bool = False, ): """Build a lightweight paged-varlen launcher for short attention.""" from kernels.attention.flash_attn_generic import build_flash_attn_func_module @@ -565,6 +578,7 @@ def _build_paged_light( path_tag="N32", waves_per_eu=waves_per_eu, daz=daz, + return_lse=return_lse, ) @@ -602,6 +616,12 @@ def flydsl_flash_attn_func( v_descale: Optional[torch.Tensor] = None, # Output tensor; allocated if None. out: Optional[torch.Tensor] = None, + # Also return the per-row log-sum-exp (LSE) needed by the backward pass. + # LSE convention (A1<->A3 contract): natural log, softmax scale folded in, + # i.e. LSE_i = ln(sum_j exp(sm_scale * (q_i . k_j))). fp32 tensor of shape + # [B, num_heads, Sq] (varlen: [B, num_heads, max_seqlen_q], padded). Not + # supported for fp8. See docs/flash_attn_lse_contract.md. + return_lse: bool = False, # Kernel build options. waves_per_eu: int = 2, daz: bool = True, @@ -659,7 +679,10 @@ def flydsl_flash_attn_func( Returns: Output tensor with the same shape as q. The dtype is bf16 for fp8 - inputs, otherwise the same dtype as q. + inputs, otherwise the same dtype as q. When ``return_lse=True`` returns + ``(out, lse)`` where ``lse`` is fp32 ``[B, num_heads, Sq]`` (varlen: + ``[B, num_heads, max_seqlen_q]``, padded) holding the per-row + natural-log, scale-folded log-sum-exp. """ # ── validation ────────────────────────────────────────────────────────── if not (q.is_cuda and k.is_cuda and v.is_cuda): @@ -670,9 +693,13 @@ def flydsl_flash_attn_func( raise ValueError(f"flydsl_flash_attn_func: q/k/v must share dtype; got {q.dtype}/{k.dtype}/{v.dtype}") dtype_str = _dtype_str(q) + if return_lse and dtype_str == "fp8": + raise NotImplementedError("flydsl_flash_attn_func: return_lse is not supported for fp8") paged_kv = any(x is not None for x in (block_table, seqlen_k)) if dtype_str == "fp8" and paged_kv: 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") if paged_kv: return _flydsl_flash_attn_paged( q, @@ -786,6 +813,7 @@ def flydsl_flash_attn_func( lazy_rescale=dualwave_swp_lazy_rescale, setprio=dualwave_swp_setprio, enable_stagger=dualwave_swp_enable_stagger, + return_lse=return_lse, ) elif varlen: # Short varlen attention uses generic light; long/debug stays on dualwave. @@ -810,6 +838,7 @@ def flydsl_flash_attn_func( setprio=dualwave_swp_setprio, debug_lazy_counts=debug_lazy, enable_stagger=dualwave_swp_enable_stagger, + return_lse=return_lse, ) else: exe = _build_varlen( @@ -825,6 +854,7 @@ def flydsl_flash_attn_func( setprio=dualwave_swp_setprio, debug_lazy_counts=debug_lazy, enable_stagger=dualwave_swp_enable_stagger, + return_lse=return_lse, ) else: _arch = _gpu_arch(q.device) @@ -861,6 +891,7 @@ def flydsl_flash_attn_func( setprio=dualwave_swp_setprio, debug_lazy_counts=debug_lazy, enable_stagger=dualwave_swp_enable_stagger, + return_lse=return_lse, ) else: block_m, flat_work_group_size, path_tag = _dense_generic_tile(B, Sq, H, D, dtype_str, q.device) @@ -876,6 +907,7 @@ def flydsl_flash_attn_func( path_tag=path_tag, waves_per_eu=waves_per_eu, daz=daz, + return_lse=return_lse, ) # ── allocate output ───────────────────────────────────────────────── @@ -901,12 +933,15 @@ def flydsl_flash_attn_func( v_flat = v.contiguous() o_flat = out.contiguous() + # ── allocate LSE (fp32 [B, num_heads, Sq]) ─────────────────────────── + lse = torch.empty((B, H, Sq), dtype=torch.float32, device=q.device) if return_lse else None + # ── launch ────────────────────────────────────────────────────────── 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, stream=launch_stream) + exe(q_flat, k_flat, v_flat, o_flat, B, Sq, workspace=_ws, lse=lse, stream=launch_stream) elif varlen: - kwargs = dict(cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, stream=launch_stream) + kwargs = dict(cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, lse=lse, stream=launch_stream) if cross: kwargs["seq_len_kv"] = int(max_seqlen_kv) if debug_lazy: @@ -914,7 +949,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(lse=lse, stream=launch_stream) if cross: kwargs["seq_len_kv"] = Skv if debug_lazy: @@ -923,4 +958,6 @@ def flydsl_flash_attn_func( kwargs.update(q_descale=q_descale, k_descale=k_descale, v_descale=v_descale) exe(q_flat, k_flat, v_flat, o_flat, B, Sq, **kwargs) + if return_lse: + return out, lse return out diff --git a/tests/kernels/test_flash_attn_lse.py b/tests/kernels/test_flash_attn_lse.py new file mode 100644 index 000000000..90939cee0 --- /dev/null +++ b/tests/kernels/test_flash_attn_lse.py @@ -0,0 +1,209 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Forward LSE (log-sum-exp) correctness tests for FlyDSL flash attention. + +Validates the per-row LSE emitted by ``flydsl_flash_attn_func(..., return_lse=True)`` +against a float32 PyTorch reference. The A1<->A3 interface contract (see +``docs/flash_attn_lse_contract.md``) is: + + LSE_i = ln( sum_j exp( sm_scale * (q_i . k_j) ) ) (masked keys excluded) + +i.e. natural log, with the softmax scale (sm_scale = 1/sqrt(head_dim)) folded in, +and a fully-masked row (no visible keys) storing -inf. Output layout is fp32 +``[B, num_heads, Sq]`` (varlen: ``[B, num_heads, max_seqlen_q]``, padded rows are +left undefined and not checked). + +Covers the dense (gfx950 dualwave / generic), split-K combine, and varlen store +paths, plus GQA, cross-attention (Skv != Sq) and fully-masked rows. +""" + +import math + +import pytest + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +try: + import torch +except ImportError: + torch = None +if torch is None or not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) + +# Imported after the torch guard so a torch-less collection skips rather than errors. +from kernels.attention.flash_attn_interface import flydsl_flash_attn_func # noqa: E402 + + +def _sm_scale(head_dim: int) -> float: + return 1.0 / math.sqrt(head_dim) + + +def _reference_lse(q, k, causal, num_kv_heads): + """float32 reference LSE for dense inputs. + + q: [B, Sq, H, D], k: [B, Skv, Hkv, D] -> lse: [B, H, Sq] (natural log, scale + folded). Causal is bottom-right aligned: key j is visible to query i iff + ``j <= i + (Skv - Sq)`` (matches the kernel's bottom-right convention). + """ + B, Sq, H, D = q.shape + Skv, Hkv = k.shape[1], k.shape[2] + sm = _sm_scale(D) + qf = q.float().permute(0, 2, 1, 3) # B, H, Sq, D + kf = k.float().permute(0, 2, 1, 3) # B, Hkv, Skv, D + kf = kf.repeat_interleave(H // Hkv, dim=1) # B, H, Skv, D + scores = torch.einsum("bhqd,bhkd->bhqk", qf, kf) * sm + if causal: + qi = torch.arange(Sq, device=q.device).view(Sq, 1) + ki = torch.arange(Skv, device=q.device).view(1, Skv) + mask = (ki > (qi + (Skv - Sq))).view(1, 1, Sq, Skv) + scores = scores.masked_fill(mask, float("-inf")) + return torch.logsumexp(scores, dim=-1) # B, H, Sq + + +def _assert_lse_matches(lse, lse_ref, atol): + """Finite rows match within atol; -inf (fully-masked) rows match exactly.""" + lse = lse.float() + lse_ref = lse_ref.float() + finite = torch.isfinite(lse_ref) + # Fully-masked rows: reference is -inf, kernel must be non-finite (i.e. -inf). + if (~finite).any(): + assert bool(((~torch.isfinite(lse)) == (~finite)).all()), "masked (-inf) rows mismatch" + if finite.any(): + diff = (lse[finite] - lse_ref[finite]).abs().max().item() + assert diff <= atol, f"LSE max abs diff {diff:.3e} exceeds atol {atol:.3e}" + + +def _rand(*shape, dtype, device="cuda"): + return torch.randn(*shape, device=device, dtype=dtype) + + +# bf16 inputs: the dot-product accumulates in the kernel with slightly different +# ordering than the fp32 reference, so allow a small absolute tolerance. +_ATOL_BF16 = 8e-3 +_ATOL_F16 = 4e-3 + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize( + "B,S,H,Hkv,D", + [ + (2, 256, 8, 8, 128), # MHA + (2, 256, 8, 2, 128), # GQA + (1, 256, 4, 1, 128), # MQA + (2, 192, 4, 4, 64), # D=64 + (3, 128, 6, 3, 128), # GQA, odd batch + ], +) +def test_lse_dense(dtype, causal, B, S, H, Hkv, D): + torch.manual_seed(S + H + D + int(causal)) + q = _rand(B, S, H, D, dtype=dtype) + k = _rand(B, S, Hkv, D, dtype=dtype) + v = _rand(B, S, Hkv, D, dtype=dtype) + _, lse = flydsl_flash_attn_func(q, k, v, causal=causal, num_kv_heads=Hkv, return_lse=True) + torch.cuda.synchronize() + assert lse.shape == (B, H, S) + assert lse.dtype == torch.float32 + atol = _ATOL_BF16 if dtype == torch.bfloat16 else _ATOL_F16 + _assert_lse_matches(lse, _reference_lse(q, k, causal, Hkv), atol) + + +@pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("num_kv_splits", [2, 3]) +@pytest.mark.parametrize("D", [64, 128]) +def test_lse_splitk(causal, num_kv_splits, D): + # Split-K requires seq_len >= 384, D in {64,128}, bf16/f16. + B, S, H = 1, 512, 8 + torch.manual_seed(S + D + num_kv_splits + int(causal)) + dtype = torch.bfloat16 + q = _rand(B, S, H, D, dtype=dtype) + k = _rand(B, S, H, D, dtype=dtype) + v = _rand(B, S, H, D, dtype=dtype) + _, lse = flydsl_flash_attn_func( + q, k, v, causal=causal, num_kv_splits=num_kv_splits, return_lse=True + ) + torch.cuda.synchronize() + assert lse.shape == (B, H, S) + _assert_lse_matches(lse, _reference_lse(q, k, causal, H), _ATOL_BF16) + + +@pytest.mark.parametrize("causal", [False, True]) +def test_lse_varlen(causal): + dtype = torch.bfloat16 + D, H, Hkv = 128, 4, 4 + seqs = [100, 200, 60] + torch.manual_seed(sum(seqs) + int(causal)) + cu = torch.tensor([0, 100, 300, 360], dtype=torch.int32, device="cuda") + total = int(cu[-1].item()) + max_seqlen_q = max(seqs) + q = _rand(total, H, D, dtype=dtype) + k = _rand(total, Hkv, D, dtype=dtype) + v = _rand(total, Hkv, D, dtype=dtype) + _, lse = flydsl_flash_attn_func( + q, k, v, + causal=causal, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + max_seqlen_q=max_seqlen_q, + cross_seqlen=False, + return_lse=True, + ) + torch.cuda.synchronize() + assert lse.shape == (len(seqs), H, max_seqlen_q) + for b, (s0, s1) in enumerate(zip(cu[:-1].tolist(), cu[1:].tolist())): + n = s1 - s0 + qb = q[s0:s1].unsqueeze(0) # 1, n, H, D + kb = k[s0:s1].unsqueeze(0) + ref = _reference_lse(qb, kb, causal, Hkv)[0] # H, n + # Only the valid [:, :n] region is defined; padded rows are ignored. + _assert_lse_matches(lse[b, :, :n], ref, _ATOL_BF16) + + +def test_lse_fully_masked_rows(): + """Cross-attention causal with Skv < Sq: leading query rows see no keys -> -inf.""" + dtype = torch.bfloat16 + B, Sq, Skv, H, D = 2, 128, 32, 4, 128 + torch.manual_seed(Sq + Skv) + q = _rand(B, Sq, H, D, dtype=dtype) + k = _rand(B, Skv, H, D, dtype=dtype) + v = _rand(B, Skv, H, D, dtype=dtype) + _, lse = flydsl_flash_attn_func(q, k, v, causal=True, return_lse=True) + torch.cuda.synchronize() + ref = _reference_lse(q, k, True, H) + assert (~torch.isfinite(ref)).any(), "test setup should produce fully-masked rows" + _assert_lse_matches(lse, ref, _ATOL_BF16) + + +def test_return_lse_false_returns_only_out(): + """Backwards-compat: default return_lse=False returns a bare tensor.""" + dtype = torch.bfloat16 + q = _rand(2, 128, 4, 128, dtype=dtype) + k = _rand(2, 128, 4, 128, dtype=dtype) + v = _rand(2, 128, 4, 128, dtype=dtype) + out = flydsl_flash_attn_func(q, k, v, causal=False) + torch.cuda.synchronize() + assert isinstance(out, torch.Tensor) + assert out.shape == q.shape + + +def test_return_lse_rejects_fp8(): + dtype = torch.bfloat16 + q = _rand(2, 128, 4, 128, dtype=dtype) + with pytest.raises(NotImplementedError): + flydsl_flash_attn_func( + q.to(torch.float8_e4m3fn), q.to(torch.float8_e4m3fn), q.to(torch.float8_e4m3fn), + causal=False, + q_descale=torch.ones(1, device="cuda"), + k_descale=torch.ones(1, device="cuda"), + v_descale=torch.ones(1, device="cuda"), + return_lse=True, + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-s"]))