[BUG][CK][MHA] Fix for MHA with softmax-sink

csrc/include/mha_bwd.h

@@ -48,7 +48,7 @@ struct mha_bwd_args
     void* dk_ptr;
     void* dv_ptr;
     void* dbias_ptr;
-    const void* sink_ptr   = nullptr; // sink scores [batch, nhead] log-space (LSEDataType=float); nullptr disables sink
+    const void* sink_ptr   = nullptr; // sink scores [nhead] log-space (LSEDataType=float); nullptr disables sink
     void*       d_sink_ptr = nullptr; // sink gradient accumulator [nhead] (LSEDataType=float); nullptr disables sink grad
     // Usage notes for sequence length pointer parameters:
     //

csrc/include/torch/mha_bwd.h

@@ -1,6 +1,6 @@
 #pragma once
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 #include <torch/extension.h>
 
 namespace aiter {
@@ -25,7 +25,7 @@ std::vector<at::Tensor> mha_bwd(const at::Tensor& dout, // [b, sq, hq, d]
                                 std::optional<const at::Tensor> alibi_slopes, // [hq] or [b, hq]
                                 std::optional<const at::Tensor> rng_state,
                                 std::optional<at::Generator> gen,
-                                std::optional<const at::Tensor> sink,  // [b, hq] log-space sink scores (float)
+                                std::optional<const at::Tensor> sink,  // [hq] log-space sink scores (float)
                                 std::optional<at::Tensor> d_sink);     // [hq] sink gradient output (float)
 } // namespace torch_itfs
 } // namespace aiter

csrc/include/torch/mha_varlen_bwd.h

@@ -1,6 +1,6 @@
 #pragma once
 // SPDX-License-Identifier: MIT
-// Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 #include <torch/extension.h>
 
 namespace aiter {
@@ -31,7 +31,7 @@ mha_varlen_bwd(const at::Tensor& dout,         // [total_q, hq, d]
                std::optional<at::Generator> gen,
                std::optional<const at::Tensor> cu_seqlens_q_padded, // [b+1]
                std::optional<const at::Tensor> cu_seqlens_k_padded, // [b+1]
-               std::optional<const at::Tensor> sink,                // [b, hq] log-space sink scores (float)
+               std::optional<const at::Tensor> sink,                // [hq] log-space sink scores (float)
                std::optional<at::Tensor> d_sink                     // [hq] sink gradient output (float)
 );
 } // namespace torch_itfs

csrc/py_itfs_ck/mha_bwd_kernels.cu

@@ -32,7 +32,7 @@ mha_bwd(const at::Tensor &dout,         // [b, sq, hq, d_v]
         std::optional<const at::Tensor> alibi_slopes_, // [hq] or [b, hq]
         std::optional<const at::Tensor> rng_state_,
         std::optional<at::Generator> gen_,
-        std::optional<const at::Tensor> sink_,         // [b, hq] log-space sink scores (float)
+        std::optional<const at::Tensor> sink_,         // [hq] log-space sink scores (float)
         std::optional<at::Tensor> d_sink_)             // [hq] sink gradient output (float)
 {
     if (is_causal) { window_size_right = 0; }
@@ -298,8 +298,8 @@ mha_bwd(const at::Tensor &dout,         // [b, sq, hq, d_v]
                 CHECK_DEVICE(sink);
                 TORCH_CHECK(sink.dtype() == torch::kFloat32, "sink must be float32");
                 TORCH_CHECK(sink.is_contiguous(), "sink must be contiguous");
-                TORCH_CHECK(sink.dim() == 2 && sink.size(0) == batch_size && sink.size(1) == num_heads,
-                            "sink must have shape [batch_size, num_heads]");
+                TORCH_CHECK(sink.dim() == 1 && sink.size(0) == num_heads,
+                            "sink must have shape [num_heads]");
                 sink_data_ptr = sink.data_ptr();
             }
             if (d_sink_.has_value() && d_sink_.value().defined()) {
@@ -343,7 +343,7 @@ mha_bwd(const at::Tensor &dout,         // [b, sq, hq, d_v]
                                 dk_expanded.data_ptr(),
                                 dv_expanded.data_ptr(),
                                 dbias_ptr,
-                                sink_data_ptr,   // sink_ptr [b, hq]
+                                sink_data_ptr,   // sink_ptr [hq]
                                 d_sink_data_ptr, // d_sink_ptr [hq]
                                 nullptr, // seqstart_q_ptr
                                 nullptr, // seqstart_k_ptr

csrc/py_itfs_ck/mha_varlen_bwd_kernels.cu

@@ -37,7 +37,7 @@ mha_varlen_bwd(const at::Tensor &dout,         // [total_q, hq, d_v]
                std::optional<at::Generator> gen_,
                std::optional<const at::Tensor> cu_seqlens_q_padded, // [b+1]
                std::optional<const at::Tensor> cu_seqlens_k_padded, // [b+1]
-               std::optional<const at::Tensor> sink_,               // [b, hq] log-space sink scores (float)
+               std::optional<const at::Tensor> sink_,               // [hq] log-space sink scores (float)
                std::optional<at::Tensor> d_sink_)                   // [hq] sink gradient output (float)
 {
     if (is_causal) { window_size_right = 0; }
@@ -306,8 +306,8 @@ mha_varlen_bwd(const at::Tensor &dout,         // [total_q, hq, d_v]
                 CHECK_DEVICE(sink);
                 TORCH_CHECK(sink.dtype() == torch::kFloat32, "sink must be float32");
                 TORCH_CHECK(sink.is_contiguous(), "sink must be contiguous");
-                TORCH_CHECK(sink.dim() == 2 && sink.size(0) == batch_size && sink.size(1) == num_heads,
-                            "sink must have shape [batch_size, num_heads]");
+                TORCH_CHECK(sink.dim() == 1 && sink.size(0) == num_heads,
+                            "sink must have shape [num_heads]");
                 sink_data_ptr = sink.data_ptr();
             }
             if (d_sink_.has_value() && d_sink_.value().defined()) {
@@ -351,7 +351,7 @@ mha_varlen_bwd(const at::Tensor &dout,         // [total_q, hq, d_v]
                                 dk_expanded.data_ptr(),
                                 dv_expanded.data_ptr(),
                                 nullptr, // dbias
-                                sink_data_ptr,   // sink_ptr [b, hq]
+                                sink_data_ptr,   // sink_ptr [hq]
                                 d_sink_data_ptr, // d_sink_ptr [hq]
                                 seqstart_q_ptr, // seqstart_q_ptr (physical cumulative)
                                 seqstart_k_ptr, // seqstart_k_ptr (physical cumulative)

op_tests/test_mha.py

@@ -882,9 +882,14 @@ def test_flash_attn_seq_padding(
 # ---------------------------------------------------------------------------
 # Sink backward tests (mha_bwd with sink / d_sink)
 #
-# Reference formula (derived from kernel block_fmha_bwd_dot_do_o.hpp):
+# `sink` is a learnable per-head softmax offset: a virtual extra logit column
+# that enters the softmax denominator but contributes nothing to the output
+# (its V row is zero).  It therefore has shape [H] -- one scalar per Q head,
+# shared by every batch entry -- and its gradient d_sink has the same shape.
+#
+# Reference formula:
 #   D[b, h, q]      = sum_j(dout[b, q, h, j] * out[b, q, h, j]) * p_undrop
-#   P_sink[b, h, q] = exp(sink[b, h] - lse_fwd[b, h, q])
+#   P_sink[b, h, q] = exp(sink[h] - lse[b, h, q])
 #   d_sink[h]       = sum_{b, q} (-P_sink[b, h, q] * D[b, h, q])
 # ---------------------------------------------------------------------------
 
@@ -906,8 +911,15 @@ def _sink_make_qkvo(
     return q, k, v, dout
 
 
-def _sink_run_fwd(q, k, v, softmax_scale, causal):
-    """Run mha_fwd and return (out, lse)."""
+def _sink_run_fwd(q, k, v, softmax_scale, causal, sink=None):
+    """
+    Run mha_fwd and return (out, lse).
+    When `sink` is given it is forwarded as `sink_ptr`, so the returned LSE is
+    log(exp(lse_without_sink) + exp(sink)) and `out` is normalised by that same
+    denominator.  Feeding a sink-free LSE to the backward instead would make
+    P_sink = exp(sink - lse) unbounded, i.e. a softmax state no forward pass can
+    produce; the pair (out, lse) produced here is the physically reachable one.
+    """
     out, lse, _, _ = aiter.mha_fwd(
         q,
         k,
@@ -920,6 +932,7 @@ def _sink_run_fwd(q, k, v, softmax_scale, causal):
         sink_size=0,
         return_softmax_lse=True,
         return_dropout_randval=False,
+        sink_ptr=sink,
     )
     return out, lse
 
@@ -930,13 +943,13 @@ def _sink_reference_d_sink(dout, out, lse, sink, p_undrop=1.0):
 
     dout : [B, Sq, H, Dv]
     out  : [B, Sq, H, Dv]
-    lse  : [B, H, Sq]       (forward LSE without sink)
-    sink : [B, H]
+    lse  : [B, H, Sq]       (forward LSE, sink included in the denominator)
+    sink : [H]
     returns d_sink : [H]
     """
     D_bsh = (dout.float() * out.float()).sum(dim=-1) * p_undrop  # [B, Sq, H]
     D_bhs = D_bsh.permute(0, 2, 1)  # [B, H, Sq]
-    sink_bhs = sink.unsqueeze(-1)  # [B, H, 1]
+    sink_bhs = sink.view(1, -1, 1)  # [1, H, 1], shared across the batch
     p_sink = torch.exp(sink_bhs.float() - lse.float())  # [B, H, Sq]
     d_sink = (-p_sink * D_bhs).sum(dim=(0, 2))  # [H]
     return d_sink.float()
@@ -948,6 +961,7 @@ def _sink_reference_d_sink(dout, out, lse, sink, p_undrop=1.0):
     # (batch, seqlen_q, seqlen_k, nhead, nhead_k, hdim)
     (2, 128, 128, 4, 4, 64),
     (1, 64, 64, 6, 2, 128),
+    (5, 512, 512, 8, 8, 64),
 ]
 
 
@@ -965,11 +979,12 @@ def test_mha_bwd_sink_dsink(
     q, k, v, dout = _sink_make_qkvo(
         batch, seqlen_q, seqlen_k, nhead, nhead_k, hdim, hdim_v, dtype, device
     )
-    out, lse = _sink_run_fwd(q.detach(), k.detach(), v.detach(), softmax_scale, causal)
-
-    sink = torch.empty(batch, nhead, device=device, dtype=torch.float32).uniform_(
-        30.0, 60.0
-    )
+    sink = torch.empty(nhead, device=device, dtype=torch.float32).uniform_(
+            -1.0, 1.0
+            #3.0, 6.0
+            #30.0,60.0
+        )
+    out, lse = _sink_run_fwd(q.detach(), k.detach(), v.detach(), softmax_scale, causal, sink)
     d_sink = torch.zeros(nhead, device=device, dtype=torch.float32)
 
     _dq, _dk, _dv, _softmax_d = aiter.mha_bwd(
@@ -990,13 +1005,14 @@ def test_mha_bwd_sink_dsink(
     )
 
     assert d_sink.abs().max() > 0, "d_sink was not updated by mha_bwd"
-
+    #print(f"sink: {sink}")
+    #print(f"d_sink: {d_sink}")
     d_sink_ref = _sink_reference_d_sink(dout, out, lse, sink)
     torch.testing.assert_close(
         d_sink,
         d_sink_ref,
-        rtol=0.02,
-        atol=0.5,
+        rtol=1e-3,
+        atol=1e-3,
         msg=f"d_sink mismatch for dtype={dtype}, causal={causal}, B={batch}, Sq={seqlen_q}, H={nhead}",
     )
 
@@ -1015,7 +1031,9 @@ def test_mha_bwd_with_sink_dq_dk_dv(
     q, k, v, dout = _sink_make_qkvo(
         batch, seqlen_q, seqlen_k, nhead, nhead_k, hdim, hdim_v, dtype, device
     )
-    out, lse = _sink_run_fwd(q.detach(), k.detach(), v.detach(), softmax_scale, causal)
+    sink_small = torch.full((nhead,), -1000.0, device=device, dtype=torch.float32)
+
+    out, lse = _sink_run_fwd(q.detach(), k.detach(), v.detach(), softmax_scale, causal, sink_small)
 
     common_bwd_args = {
         "dropout_p": 0.0,
@@ -1030,7 +1048,6 @@ def test_mha_bwd_with_sink_dq_dk_dv(
         dout, q.detach(), k.detach(), v.detach(), out, lse, **common_bwd_args
     )
 
-    sink_small = torch.full((batch, nhead), -1000.0, device=device, dtype=torch.float32)
     d_sink = torch.zeros(nhead, device=device, dtype=torch.float32)
 
     dq_sink, dk_sink, dv_sink, _ = aiter.mha_bwd(

op_tests/test_mha_varlen.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MIT
-# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
 import argparse
 import itertools
@@ -1107,11 +1107,25 @@ def varlen_flash_attn_seq_padding_benchmark(
 
 # ---------------------------------------------------------------------------
 # Sink backward tests (mha_varlen_bwd with sink / d_sink)
+#
+# `sink` is a learnable per-head softmax offset -- one scalar per Q head, shared
+# by every sequence in the batch -- so it has shape [H], matching its gradient
+# d_sink, which the kernel accumulates into a single cell per head across all
+# batches. See the comment block above the sink tests in test_mha.py for the
+# derivation of the reference formula.
 # ---------------------------------------------------------------------------
 
 
-def _vsink_run_fwd(q, k, v, softmax_scale, causal):
-    """Run mha_fwd and return (out, lse)."""
+def _vsink_run_fwd(q, k, v, softmax_scale, causal, sink=None):
+    """
+    Run mha_fwd and return (out, lse).
+
+    When `sink` is given it is forwarded as `sink_ptr`, so the returned LSE is
+    log(exp(lse_without_sink) + exp(sink)) and `out` is normalised by that same
+    denominator -- the state the backward reference assumes. Passing a sink-free
+    LSE instead would make P_sink = exp(sink - lse) exceed 1, which no softmax
+    can produce.
+    """
     out, lse, _, _ = aiter.mha_fwd(
         q,
         k,
@@ -1124,6 +1138,7 @@ def _vsink_run_fwd(q, k, v, softmax_scale, causal):
         sink_size=0,
         return_softmax_lse=True,
         return_dropout_randval=False,
+        sink_ptr=sink,
     )
     return out, lse
 
@@ -1134,30 +1149,31 @@ def _vsink_reference_d_sink_varlen(dout, out, lse_group, sink, seqlens_q):
 
     dout       : [total_q, H, Dv]
     out        : [total_q, H, Dv]
-    lse_group  : [H, total_q]   – group-mode LSE (flattened across batches)
-    sink       : [B, H]
+    lse_group  : [H, total_q]   - group-mode LSE (flattened across batches)
+    sink       : [H]            - shared by every sequence in the batch
     seqlens_q  : list of per-batch sequence lengths
     returns d_sink : [H]
     """
-    nhead = sink.shape[1]
+    nhead = sink.shape[0]
     d_sink = torch.zeros(nhead, device=sink.device, dtype=torch.float32)
 
     offset = 0
-    for b, sq in enumerate(seqlens_q):
+    for sq in seqlens_q:
         dout_b = dout[offset : offset + sq].float()
         out_b = out[offset : offset + sq].float()
         lse_b = lse_group[:, offset : offset + sq]
 
         D_qh = (dout_b * out_b).sum(dim=-1)
         D_hq = D_qh.permute(1, 0)
-        p_sink = torch.exp(sink[b].float().unsqueeze(-1) - lse_b)
+        p_sink = torch.exp(sink.float().unsqueeze(-1) - lse_b)
         d_sink += (-p_sink * D_hq).sum(dim=-1)
         offset += sq
 
     return d_sink
 
 
 _VSINK_DTYPES = [dtypes.fp16, dtypes.bf16]
+_VSINK_RANGE = (-1.0, 1.0)
 
 
 @pytest.mark.parametrize("dtype", _VSINK_DTYPES)
@@ -1181,17 +1197,18 @@ def test_mha_varlen_bwd_sink_dsink(dtype):
     v = torch.randn(total_k, nhead, hdim_v, device=device, dtype=dtype)
     dout = torch.randn(total_q, nhead, hdim_v, device=device, dtype=dtype)
 
+    sink = torch.empty(nhead, device=device, dtype=torch.float32).uniform_(
+        *_VSINK_RANGE
+    )
+
     q_b = q.view(batch, seqlen, nhead, hdim)
     k_b = k.view(batch, seqlen, nhead, hdim)
     v_b = v.view(batch, seqlen, nhead, hdim_v)
-    out_b, lse_b = _vsink_run_fwd(q_b, k_b, v_b, softmax_scale, causal=False)
+    out_b, lse_b = _vsink_run_fwd(q_b, k_b, v_b, softmax_scale, causal=False, sink=sink)
 
     out = out_b.view(total_q, nhead, hdim_v)
     lse = lse_b.permute(1, 0, 2).reshape(nhead, total_q).contiguous()
 
-    sink = torch.empty(batch, nhead, device=device, dtype=torch.float32).uniform_(
-        30.0, 60.0
-    )
     d_sink = torch.zeros(nhead, device=device, dtype=torch.float32)
 
     dq, dk, dv, _ = aiter.mha_varlen_bwd(
@@ -1223,11 +1240,14 @@ def test_mha_varlen_bwd_sink_dsink(dtype):
     assert dv.shape == v.shape
 
     d_sink_ref = _vsink_reference_d_sink_varlen(dout, out, lse, sink, seqlens_q)
+    # See the matching assert in test_mha.py: with a physically reachable LSE
+    # the two sides differ only by fp32 accumulation order (~1e-6 relative),
+    # so the old rtol=0.02/atol=0.5 is no longer meaningful.
     torch.testing.assert_close(
         d_sink,
         d_sink_ref,
-        rtol=0.02,
-        atol=0.5,
+        rtol=1e-3,
+        atol=1e-3,
         msg="varlen d_sink mismatch vs reference",
     )
 
@@ -1242,7 +1262,6 @@ def test_mha_varlen_bwd_sink_variable_lengths(dtype):
 
     seqlens_q = [48, 80]
     seqlens_k = [48, 80]
-    batch = len(seqlens_q)
     max_seqlen_q = max(seqlens_q)
     max_seqlen_k = max(seqlens_k)
     total_q = sum(seqlens_q)
@@ -1264,13 +1283,19 @@ def test_mha_varlen_bwd_sink_variable_lengths(dtype):
     v = torch.randn(total_k, nhead, hdim_v, device=device, dtype=dtype)
     dout = torch.randn(total_q, nhead, hdim_v, device=device, dtype=dtype)
 
+    sink = torch.empty(nhead, device=device, dtype=torch.float32).uniform_(
+        *_VSINK_RANGE
+    )
+
     out_parts, lse_parts = [], []
     offset_q, offset_k = 0, 0
     for sq, sk in zip(seqlens_q, seqlens_k):
         q_b = q[offset_q : offset_q + sq].unsqueeze(0)
         k_b = k[offset_k : offset_k + sk].unsqueeze(0)
         v_b = v[offset_k : offset_k + sk].unsqueeze(0)
-        out_b, lse_b = _vsink_run_fwd(q_b, k_b, v_b, softmax_scale, causal=False)
+        out_b, lse_b = _vsink_run_fwd(
+            q_b, k_b, v_b, softmax_scale, causal=False, sink=sink
+        )
         out_parts.append(out_b.squeeze(0))
         lse_parts.append(lse_b.squeeze(0).permute(1, 0))
         offset_q += sq
@@ -1279,9 +1304,6 @@ def test_mha_varlen_bwd_sink_variable_lengths(dtype):
     out = torch.cat(out_parts, dim=0)
     lse = torch.cat(lse_parts, dim=0).permute(1, 0).contiguous()
 
-    sink = torch.empty(batch, nhead, device=device, dtype=torch.float32).uniform_(
-        30.0, 60.0
-    )
     d_sink = torch.zeros(nhead, device=device, dtype=torch.float32)
 
     _dq, _dk, _dv, _ = aiter.mha_varlen_bwd(
@@ -1313,7 +1335,7 @@ def test_mha_varlen_bwd_sink_variable_lengths(dtype):
     torch.testing.assert_close(
         d_sink,
         d_sink_ref,
-        rtol=0.02,
-        atol=0.5,
+        rtol=1e-3,
+        atol=1e-3,
         msg="varlen variable-length d_sink mismatch",
     )