Add forward LSE output to FlyDSL flash attention kernels

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

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,13 +2443,17 @@ 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(
                     Q,
                     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,13 +2529,16 @@ 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,
                 Q,
                 K,
                 V,
                 O,
+                lse,
                 debug_counts,
                 cu_seqlens_q,
                 cu_seqlens_kv,