[HIP] [GFX950]OPUS MHA Support LSE

aiter/ops/mha.py

@@ -334,6 +334,7 @@ def gen_fmha_fwd_bf16_opus_fwd_fake(
     out: Tensor,
     causal: bool,
     softmax_scale: float,
+    lse: Tensor | None = None,
     seqstart_q: Tensor | None = None,
     seqstart_k: Tensor | None = None,
     seqstart_q_pad: Tensor | None = None,
@@ -347,7 +348,7 @@ def gen_fmha_fwd_bf16_opus_fwd_fake(
 # OPUS gfx950 bf16 forward (shared entry point): low-level @compile_ops stub bound to
 # the pybind symbol via fc_name. Dispatches by head dim in C++ to the symmetric D=128
 # kernel (batch only) or the asymmetric D_QK=192/D_V=128 kernel (batch + group/varlen).
-# Writes `out` in place, returns None.
+# Writes `out` (and `lse`, when given) in place, returns None.
 @compile_ops(
     "module_fmha_fwd_bf16_opus",
     fc_name="fmha_fwd_bf16_opus_fwd",
@@ -360,6 +361,7 @@ def _fmha_fwd_bf16_opus_fwd(
     out: Tensor,
     causal: bool,
     softmax_scale: float,
+    lse: Tensor | None = None,
     seqstart_q: Tensor | None = None,
     seqstart_k: Tensor | None = None,
     seqstart_q_pad: Tensor | None = None,
@@ -376,26 +378,38 @@ def fmha_fwd_bf16_opus_fwd(
     softmax_scale: float,
     causal: bool,
     out: Tensor | None = None,
-) -> Tensor:
+    return_lse: bool = False,
+    lse: Tensor | None = None,
+) -> Tensor | tuple[Tensor, Tensor]:
     """Public wrapper for the OPUS gfx950 bf16 dense (batch) forward (D=128 and
     D_QK=192/D_V=128). q/k/v are dense bshd [B, S, H, D]; allocates `out`
     ([B, S, H_q, D_v]) if needed and forwards. The kernel applies `softmax_scale`
-    to Q·K^T internally, handles GQA fan-out, and produces no LSE.
+    to Q·K^T internally and handles GQA fan-out.
+
+    `lse` is an output buffer for the log-sum-exp of the scaled scores ([B, H_q, S]
+    float32, natural log; rows that see no keys get -inf), filled when supplied and
+    allocated here when `return_lse` is set. Like `out` it does not change the return
+    type on its own: only `return_lse` does, and then the return is `(out, lse)`.
 
     Varlen / packed inputs go through `fmha_fwd_bf16_opus_varlen_fwd` instead.
     """
     v_head_dim = v.size(-1)
+    batch, q_seq_len, q_head_num = q.size(0), q.size(1), q.size(2)
 
     if out is None:
-        batch, q_seq_len, q_head_num = q.size(0), q.size(1), q.size(2)
         out = torch.empty(
             (batch, q_seq_len, q_head_num, v_head_dim),
             dtype=q.dtype,
             device=q.device,
         )
 
-    _fmha_fwd_bf16_opus_fwd(q, k, v, out, bool(causal), float(softmax_scale))
-    return out
+    if return_lse and lse is None:
+        lse = torch.empty(
+            (batch, q_head_num, q_seq_len), dtype=torch.float32, device=q.device
+        )
+
+    _fmha_fwd_bf16_opus_fwd(q, k, v, out, bool(causal), float(softmax_scale), lse=lse)
+    return (out, lse) if return_lse else out
 
 
 def fmha_fwd_bf16_opus_varlen_fwd(
@@ -411,11 +425,19 @@ def fmha_fwd_bf16_opus_varlen_fwd(
     out: Tensor | None = None,
     seqstart_q_pad: Tensor | None = None,
     seqstart_k_pad: Tensor | None = None,
-) -> Tensor:
+    return_lse: bool = False,
+    lse: Tensor | None = None,
+) -> Tensor | tuple[Tensor, Tensor]:
     """Public wrapper for the OPUS gfx950 bf16 group/varlen forward (D_QK=192/D_V=128
     only). q/k/v are packed [total, H, D]; allocates `out` ([total_q, H_q, D_v]) if
-    needed and forwards. The kernel applies `softmax_scale` to Q·K^T internally,
-    handles GQA fan-out, and produces no LSE.
+    needed and forwards. The kernel applies `softmax_scale` to Q·K^T internally and
+    handles GQA fan-out.
+
+    `lse` is an output buffer for the log-sum-exp of the scaled scores ([H_q, total_q]
+    float32, natural log), filled when supplied and allocated here when `return_lse` is
+    set. Like `out` it does not change the return type on its own: only `return_lse`
+    does, and then the return is `(out, lse)`. Rows that see no keys get -inf; rows in
+    the padding gaps of a KV-padded layout are left untouched.
 
     seqstart_q / seqstart_k          : cumulative REAL sequence lengths (int32, len
                                        num_groups+1; drive masks / tile counts).
@@ -424,13 +446,16 @@ def fmha_fwd_bf16_opus_varlen_fwd(
     max_seqlen_q / max_seqlen_k      : upper bounds driving the grid.
     """
     v_head_dim = v.size(-1)
+    total_q, q_head_num = q.size(0), q.size(1)
 
     if out is None:
-        total_q, q_head_num = q.size(0), q.size(1)
         out = torch.empty(
             (total_q, q_head_num, v_head_dim), dtype=q.dtype, device=q.device
         )
 
+    if return_lse and lse is None:
+        lse = torch.empty((q_head_num, total_q), dtype=torch.float32, device=q.device)
+
     seqstart_q = seqstart_q.to(torch.int32).contiguous()
     seqstart_k = seqstart_k.to(torch.int32).contiguous()
     if seqstart_q_pad is not None:
@@ -445,14 +470,15 @@ def fmha_fwd_bf16_opus_varlen_fwd(
         out,
         bool(causal),
         float(softmax_scale),
+        lse,
         seqstart_q,
         seqstart_k,
         seqstart_q_pad if seqstart_q_pad is not None else seqstart_q,
         seqstart_k_pad if seqstart_k_pad is not None else seqstart_k,
         int(max_seqlen_q),
         int(max_seqlen_k),
     )
-    return out
+    return (out, lse) if return_lse else out
 
 
 # ---------------------------------------------------------------------------
@@ -1958,10 +1984,11 @@ def _can_impl_fmha_fwd_hd192_v128_bf16_opus():
         return hdim_q == 192 and hdim_v == 128
 
     def can_impl_fmha_fwd_bf16_opus():
-        # Shared eligibility for the OPUS gfx950 bf16 forward kernels (inference-only:
-        # no LSE/dropout mask, so it must never capture return_lse / the autograd path).
-        # Cheapest / most-selective gates first so the per-head-dim helpers (which read
-        # env vars) are only evaluated once the common conditions already hold.
+        # Shared eligibility for the OPUS gfx950 bf16 forward kernels. LSE is supported
+        # ([B, H_q, S] fp32, natural log), so return_lse no longer disqualifies the path;
+        # the dropout mask (return_softmax) still does. Cheapest / most-selective gates
+        # first so the per-head-dim helpers (which read env vars) are only evaluated once
+        # the common conditions already hold.
         ret = get_gfx() == "gfx950"
         ret = ret and (q.dtype == dtypes.bf16)
         ret = ret and (nhead_q % nhead_k == 0)
@@ -1971,7 +1998,7 @@ def can_impl_fmha_fwd_bf16_opus():
         ret = ret and (window_size_left == -1 and window_size_right == -1)
         ret = ret and (sink_size == 0 and sink_ptr is None)
         ret = ret and (q_descale is None and k_descale is None and v_descale is None)
-        ret = ret and (not return_lse) and (not return_softmax)
+        ret = ret and (not return_softmax)
         ret = ret and (
             _can_impl_fmha_fwd_hd128_bf16_opus()
             or _can_impl_fmha_fwd_hd192_v128_bf16_opus()
@@ -2059,17 +2086,22 @@ def _validate_cu(name: str, x: torch.Tensor | None):
         rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
     elif can_impl_fmha_fwd_bf16_opus():
         # OPUS gfx950 dense forward (shared entry point; dispatches D=128 vs
-        # D_QK=192/D_V=128 in C++ by head dim). Inference-only: the lse/S_dmask/rng
-        # slots are unused placeholders (gate guarantees not return_lse/return_softmax).
+        # D_QK=192/D_V=128 in C++ by head dim). The S_dmask/rng slots stay unused
+        # placeholders (the gate guarantees no dropout mask).
+        softmax_lse = torch.empty(
+            (batch_size, nhead_q, seqlen_q) if return_lse else (0,),
+            dtype=torch.float32,
+            device=q.device,
+        )
         out_ = fmha_fwd_bf16_opus_fwd(
             q,
             k,
             v,
             softmax_scale=float(softmax_scale),
             causal=bool(causal),
             out=out,
+            lse=softmax_lse if return_lse else None,
         )
-        softmax_lse = torch.empty((0,), dtype=torch.float32, device=q.device)
         S_dmask = torch.empty((0,), dtype=torch.float32, device=q.device)
         rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
     elif can_impl_fmha_v3_fwd() and seqlen_q > 128:  # Prefer CK for decode cases
@@ -2893,9 +2925,9 @@ def can_impl_fmha_fwd_with_sink_varlen_asm():
 
     def can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen():
         # OPUS gfx950 group/varlen D_QK=192 / D_V=128 bf16 forward. Enabled by DEFAULT
-        # (no env). Packed THD q/k/v; supports KV padding (cu_seqlens_*_padded) and
-        # cross-attention (causal bottom-right aligned). Inference-only: no LSE / dropout
-        # / bias / alibi / swa / sink / quant / paged.
+        # (no env). Packed THD q/k/v; supports KV padding (cu_seqlens_*_padded),
+        # cross-attention (causal bottom-right aligned) and LSE ([H_q, total_q] fp32,
+        # natural log). No dropout / bias / alibi / swa / sink / quant / paged.
         # AITER_DISABLE_FMHA_OPUS=1 force-disables it (fall back to v3/CK; for A/B).
         if int(os.environ.get("AITER_DISABLE_FMHA_OPUS", "0")) != 0:
             return False
@@ -2910,7 +2942,7 @@ def can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen():
         ret = ret and (sink_size == 0 and sink_ptr is None)
         ret = ret and (q_descale is None and k_descale is None and v_descale is None)
         ret = ret and (block_table is None)
-        ret = ret and (not return_lse) and (not return_softmax)
+        ret = ret and (not return_softmax)
         return ret
 
     q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
@@ -2919,6 +2951,11 @@ def can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen():
         # OPUS gfx950 group/varlen D=192 path. cu_seqlens_* are the REAL cumulative
         # lengths (masks / tile counts); cu_seqlens_*_padded are the PHYSICAL row
         # offsets (KV padding). When no padded arrays are given, physical == real.
+        softmax_lse = torch.empty(
+            (nhead_q, q.size(0)) if return_lse else (0,),
+            dtype=torch.float32,
+            device=q.device,
+        )
         out = fmha_fwd_bf16_opus_varlen_fwd(
             q,
             k,
@@ -2932,8 +2969,8 @@ def can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen():
             seqstart_k_pad=cu_seqlens_k_padded,
             max_seqlen_q=int(max_seqlen_q),
             max_seqlen_k=int(max_seqlen_k),
+            lse=softmax_lse if return_lse else None,
         )
-        softmax_lse = torch.empty((0,), dtype=torch.float32, device=q.device)
         S_dmask = torch.empty((0,), dtype=torch.float32, device=q.device)
         rng_state = torch.empty((2,), dtype=torch.int64, device=q.device)
     elif can_impl_fmha_fwd_with_sink_varlen_asm():

aiter/test_mha_common.py

@@ -678,3 +678,49 @@ def _tol(ref_val, pt_val, is_forward=False):
     bwd_tols = [_tol(dq, dq_pt), _tol(dk, dk_pt), _tol(dv, dv_pt)]
 
     return out, (dq, dk, dv), fwd_tol, bwd_tols
+
+
+def opus_ref_lse(q, k, causal, budget=1 << 23):
+    """fp32 logsumexp of the scaled scores, bottom-right causal, GQA-aware.
+
+    attention_ref downcasts its lse to the input dtype, too coarse to check against.
+    Chunked over query rows (`budget` score elements) to bound peak memory.
+    """
+    batch, seqlen_q, nheads, d = q.shape
+    seqlen_k, nheads_k = k.shape[1], k.shape[2]
+    group = nheads // nheads_k
+    scale = d**-0.5
+
+    k_f = k.float()
+    lse = torch.empty((batch, nheads, seqlen_q), dtype=torch.float32, device=q.device)
+    col = torch.arange(seqlen_k, device=q.device)
+    off = seqlen_k - seqlen_q
+    rows = max(1, budget // max(1, batch * nheads * seqlen_k))
+
+    for lo in range(0, seqlen_q, rows):
+        hi = min(lo + rows, seqlen_q)
+        q_c = q[:, lo:hi].float().reshape(batch, hi - lo, nheads_k, group, d)
+        scores = torch.einsum("bthgd,bshd->bhgts", q_c, k_f) * scale
+        if causal:
+            row = torch.arange(lo, hi, device=q.device)[:, None]
+            scores = scores.masked_fill(col > row + off, float("-inf"))
+        lse[:, :, lo:hi] = torch.logsumexp(scores, dim=-1).reshape(
+            batch, nheads, hi - lo
+        )
+    return lse
+
+
+def opus_check_lse(tag, lse, lse_ref):
+    """Compare fp32 LSE against opus_ref_lse.
+
+    0.01 accommodates D=128 folding softmax_scale into bf16 Q (~4e-3 off an fp32
+    post-scale reference); D=192 lands at ~1e-6.
+    """
+    assert lse.dtype == torch.float32, f"{tag}: lse dtype {lse.dtype}, expected float32"
+    assert torch.equal(
+        torch.isneginf(lse), torch.isneginf(lse_ref)
+    ), f"{tag}: -inf (fully-masked row) pattern differs from the reference"
+    finite = ~torch.isneginf(lse_ref)
+    diff = (lse[finite] - lse_ref[finite]).abs().max().item() if finite.any() else 0.0
+    print(f"[{tag}] lse max diff: {diff}")
+    assert diff <= 0.01, f"{tag}: lse diff {diff} > 0.01"

csrc/include/fmha_fwd_hd128_bf16_opus_defs.h

@@ -22,6 +22,11 @@ struct opus_gqa_kargs {
     int stride_kv_n;
     int stride_kv_h;
     float softmax_scale;  // QK^T scale (host passes 1/sqrt(D) by default)
+    // Optional fp32 log-sum-exp (natural log) output, [B, H, N] with unit stride along
+    // the query dim. nullptr => not produced (the kernel skips the store).
+    void* __restrict__ ptr_lse;
+    int stride_lse_b;
+    int stride_lse_h;
 };
 
 // Configuration traits for the GQA kernel (tile sizes, data types, vector lengths,

csrc/include/fmha_fwd_hd128_bf16_opus_kernel.hpp

@@ -559,6 +559,9 @@ __device__ __attribute__((always_inline)) void gqa_d128_impl(opus_gqa_kargs karg
     };
     const unsigned int kv_num_records = rec_bytes((int64_t)kargs.N_KV * kargs.stride_kv_n);
     const unsigned int qo_num_records = rec_bytes((int64_t)(kargs.N - q_block_start) * kargs.stride_q_n);
+    // Same bound for the fp32 LSE buffer (different element size than D_ATTN).
+    const unsigned int lse_num_records =
+        (unsigned int)((int64_t)(kargs.N - q_block_start) * (int64_t)sizeof(D_ACC));
 
     // Global memory tensors
     auto g_q = make_gmem(reinterpret_cast<const D_ATTN*>(kargs.ptr_q) + qo_gmem_offset, qo_num_records);
@@ -888,6 +891,18 @@ __device__ __attribute__((always_inline)) void gqa_d128_impl(opus_gqa_kargs karg
     else                          do_epilogue(v_s0, 0);
     __builtin_amdgcn_sched_barrier(0);
 
+    // ──── Optional LSE (fp32, natural log) ────
+    if (kargs.ptr_lse != nullptr && lane_id < T::W_M) {
+        constexpr D_ACC LN2 = D_ACC(0.69314718055994531f);   // 1 / log2(e)
+        const D_ACC lse = (l_row > D_ACC(0.0f))
+                              ? ((m_row + __builtin_amdgcn_logf(l_row)) * LN2)
+                              : -opus::numeric_limits<D_ACC>::infinity();
+        auto g_lse = make_gmem(reinterpret_cast<D_ACC*>(kargs.ptr_lse) +
+                                   (int64_t)b * kargs.stride_lse_b +
+                                   (int64_t)h * kargs.stride_lse_h + q_block_start,
+                               lse_num_records);
+        g_lse.store(lse, warp_id * T::Q_TILE_SIZE + lane_id);
+    }
 
     // ──── Normalize O and store to gmem ────
     D_ACC l_inv = (l_row > D_ACC(0.0f)) ? (D_ACC(1.0f) / l_row) : D_ACC(0.0f);

csrc/include/fmha_fwd_hd192_v128_bf16_opus_defs.h

@@ -54,6 +54,15 @@ struct opus_gqa_d192_kargs {
     // Runtime option bits (see OPT_* below). Decided once by the host and read by the
     // kernel, so the head/tail-merge decision is NOT recomputed on both sides.
     int opt;
+    // ── optional log-sum-exp output (fp32, natural log) ──
+    // nullptr => not produced (the kernel skips the store entirely; scalar branch).
+    // One value per (head, query row), unit stride along the query dim:
+    //   batch mode: [B, H, N]        (stride_lse_b, stride_lse_h)
+    //   group mode: [H, total_q]     (stride_lse_h; the group's row offset comes from
+    //                                seqstart_q_pad, same as Q/O)
+    void* __restrict__ ptr_lse;
+    int stride_lse_b;
+    int stride_lse_h;
 };
 
 // opus_gqa_d192_kargs::opt bit flags.