[Cute, bwd, sm90/100/110] Support learnable sink in backward by henrylhtsang · Pull Request #2706 · Dao-AILab/flash-attention

henrylhtsang · 2026-07-13T19:32:40Z

Summary

add FP16, BF16, and FP32 learnable sink gradients to supported CuTe backward kernels on SM90, SM100, and SM110; the sink dtype is independent of Q/K/V and dSink preserves it
reuse one dQ postprocess CTA per head to reduce dSink from FP32 dPsum and LSE, avoiding atomics and separate zero-initialization
support fixed-length and varlen autograd paths, including gradients from a returned LSE
return zero output and LSE equal to the sink when K is empty; learnable sink backward requires non-empty Q and K
fix SM90 fully masked sink-only rows producing NaN LSE

Test plan

CUDA_VISIBLE_DEVICES=0 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -x -q \
  'tests/cute/test_flash_attn.py::test_flash_attn_output[128-128-64-False-0-0.0-False-False-True-mha-dtype0]' \
  'tests/cute/test_flash_attn.py::test_flash_attn_output[128-128-128-True-0-0.0-False-False-True-gqa-dtype0]' \
  'tests/cute/test_flash_attn.py::test_flash_attn_output[128-128-192-False-1-15.0-True-False-True-mqa-dtype0]'

CUDA_VISIBLE_DEVICES=0 FLASH_ATTENTION_FAKE_TENSOR=1 \
  FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -x -q \
  'tests/cute/test_flash_attn.py::test_flash_attn_learnable_sink_backward_dtype' \
  'tests/cute/test_flash_attn.py::test_flash_attn_varlen_learnable_sink_backward_with_lse'

The first three cases passed on SM100. The focused fake-tensor run passed all four cases: FP16, BF16, and FP32 sink dtypes plus the varlen/LSE regression.

henrylhtsang · 2026-07-16T01:39:29Z

cc @Johnsonms @drisspg @jayhshah would appreciate a review

drisspg · 2026-07-29T00:43:23Z

        mCuSeqlensQ: Optional[cute.Tensor],
        mSeqUsedQ: Optional[cute.Tensor],
+        mdPsum: Optional[cute.Tensor],
+        mLSE: Optional[cute.Tensor],


ditto down ehre

drisspg · 2026-07-29T00:43:23Z

    batchp1 = cute.sym_int()
    mCuSeqlensQ = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cuseqlens_q else None
    mSeqUsedQ = fake_tensor(Int32, (batch,), divisibility=1) if has_seqused_q else None
+    mLearnableSink = fake_tensor(cutlass.BFloat16, (mQ.shape[-2],), divisibility=1) if has_learnable_sink else None


side note; do we not support fp16 for our sinks on purpose here? feels like a good follow up

drisspg · 2026-07-29T00:43:23Z

    arch, dtype, hdim, block_size, num_threads,
    atom_layout, swap_ab,
    use_2cta_instrs=False, cluster_size=1,
+    dpsum=None, lse=None, learnable_sink=None, dsink=None,


same deal could be a tuple of types since they need to come together

drisspg · 2026-07-29T00:43:23Z

                aux_scalars=ctx.aux_scalars,
                mask_mod=ctx.mask_mod,
                dlse=dlse,
+                learnable_sink=learnable_sink,


codex: It doesn't look like learnable-sink backward is tested for varlen—the varlen test hardcodes has_learnable_sink=False and excludes it from backward. Could you add a numerical varlen regression, ideally with unequal and zero sequence lengths?

drisspg · 2026-07-29T00:43:23Z

            # do_o = ((g.float() * out.float()).sum(-1)).transpose(1, 2)
-            dq, dk, dv = torch.autograd.grad(out, (q, k, v), g)
+            grad_tensors = (q, k, v, learnable_sink) if has_learnable_sink else (q, k, v)
+            grads = torch.autograd.grad(out, grad_tensors, g)


CODEX:

Can we also add a return_lse=True case with learnable_sink, using a loss that depends on both out and lse, and compare dSink against the reference? This currently differentiates only out, so it tests the dO contribution to dSink but not the dLSE contribution. I want both features exercised together.

drisspg · 2026-07-29T00:43:23Z

+                        lse_val = Float32(lse_val)
+                        sink_prob = (
+                            Float32(1.0)
+                            if lse_val == -Float32.inf


Codex now: SM90 sink-only rows can reach this code with LSE = NaN, so this -inf special case does not protect them and sink_prob/dSink become NaN. flash_fwd_sm90.py leaves row_max=-inf before adding the finite sink, while SM100 replaces the empty-row max with the sink. Could SM90 use the same handling and add a causal seqlen_q > seqlen_k sink-gradient regression? The current backward test excludes that case.

drisspg · 2026-07-29T00:43:23Z

        scale: cutlass.Float32,
        mCuSeqlensQ: Optional[cute.Tensor],
        mSeqUsedQ: Optional[cute.Tensor],
+        mdPsum: Optional[cute.Tensor],


Could we bundle these as sink_tensors: tuple[cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor] | None, or preferably a small NamedTuple? These four tensors are all-or-none, so grouping them would encode that invariant and avoid threading four optional arguments through both entry points.

henrylhtsang · 2026-07-29T23:35:25Z

-            lse.fill_(float("-inf"))
+            if learnable_sink is None:
+                lse.fill_(float("-inf"))
+            else:


NOTE: this is a behavior change for forward when sink is present.

this makes sense to me ; feels consistent

henrylhtsang · 2026-07-29T23:45:26Z

-                row_sum[r] += cute.math.exp2(
-                    sink_val_cur * LOG2_E - row_max[r] * scale_log2, fastmath=True
-                )
+                if row_max[r] == -Float32.inf:


afaik this is for sm90 only.

henrylhtsang · 2026-07-29T23:49:31Z

addressed feedback

henrylhtsang · 2026-07-30T00:17:52Z

Tested this on an H100 / SM90. We have to relax the dSink test tolerance for the SM90 result to pass, but the relaxation is narrow and dtype-aware rather than a blanket threshold.

The kernel accumulates dSink in FP32 and casts it to the sink dtype. On SM90, the varlen + LSE case differed from the PyTorch reference by exactly one BF16 ULP (0.00390625), while the previous cancellation-based tolerance evaluated to zero. FP16 also needed an output-dtype rounding allowance; FP32 passed the original comparison.

drisspg

Looks pretty good; two comments

drisspg · 2026-07-31T00:19:43Z

+    if learnable_sink is not None:
+        assert arch // 10 in [9, 10, 11], "Learnable sink backward is supported on SM90 and SM100/SM110"
+        assert lse is not None, "learnable_sink backward requires LSE"
+        if q.numel() == 0 or k.numel() == 0:
+            dsink = (
+                dlse.sum(dim=(0, 2) if dlse.ndim == 3 else 1).to(learnable_sink.dtype)
+                if dlse is not None
+                else torch.zeros_like(learnable_sink)
+            )
+            return torch.zeros_like(q), torch.zeros_like(k), torch.zeros_like(v), dsink
    sparse_q = None


Could we preserve the preallocated dq/dk/dv outputs in this early-return path? Other paths in _flash_attn_bwd respect these buffers (and test_flash_attn_bwd_preallocated_outputs checks their identity). For example:

Suggested change

if learnable_sink is not None:

assert arch // 10 in [9, 10, 11], "Learnable sink backward is supported on SM90 and SM100/SM110"

assert lse is not None, "learnable_sink backward requires LSE"

if q.numel() == 0 or k.numel() == 0:

dsink = (

dlse.sum(dim=(0, 2) if dlse.ndim == 3 else 1).to(learnable_sink.dtype)

if dlse is not None

else torch.zeros_like(learnable_sink)

)

return torch.zeros_like(q), torch.zeros_like(k), torch.zeros_like(v), dsink

sparse_q = None

if q.numel() == 0 or k.numel() == 0:

dq = torch.zeros_like(q) if dq is None else dq.zero_()

dk = torch.zeros_like(k) if dk is None else dk.zero_()

dv = torch.zeros_like(v) if dv is None else dv.zero_()

dsink = (

dlse.sum(dim=(0, 2) if dlse.ndim == 3 else 1).to(learnable_sink.dtype)

if dlse is not None

else torch.zeros_like(learnable_sink)

)

return dq, dk, dv, dsink

drisspg · 2026-07-31T00:19:44Z

+        # Reuse one existing dQ postprocess CTA per head to reduce dSink from
+        # the per-row dPsum and LSE written by backward preprocess. This avoids
+        # both a global atomic accumulator and a separate zero-initialization.
+        if const_expr(sink_tensors is not None):
+            mdPsum, mLSE, mLearnableSink, mdSink = sink_tensors
+            block_x, block_y, block_z = cute.arch.block_idx()
+            sink_head_idx = head_idx if const_expr(mCuSeqlensQ is None) else block_x
+            # Varlen uses block_x to select one CTA per head. block_y and block_z
+            # are currently always zero, but check them defensively.
+            reduce_sink = (
+                m_block == 0 and batch_idx == 0
+                if const_expr(mCuSeqlensQ is None)
+                else block_x < mdSink.shape[0] and block_y == 0 and block_z == 0
+            )
+            if reduce_sink:
+                sink_sum = Float32(0.0)
+                num_batch = (
+                    mdQ.shape[0] if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ.shape[0] - 1
+                )
+                sink_val = Float32(mLearnableSink[sink_head_idx])
+                sink_batch = 0
+                while sink_batch < num_batch:
+                    sink_seqlen = SeqlenInfoQK.create(
+                        sink_batch,
+                        mdQ.shape[1],
+                        0,
+                        mCuSeqlensQ=mCuSeqlensQ,
+                        mSeqUsedQ=mSeqUsedQ,
+                        tile_m=self.tile_m * self.cluster_size,
+                    )
+                    sink_row = tidx
+                    while sink_row < sink_seqlen.seqlen_q:
+                        if const_expr(mCuSeqlensQ is None):
+                            dpsum_val = mdPsum[sink_batch, sink_head_idx, sink_row]
+                            lse_val = mLSE[sink_batch, sink_head_idx, sink_row]
+                        else:
+                            dpsum_val = mdPsum[
+                                sink_head_idx, sink_seqlen.padded_offset_q + sink_row
+                            ]
+                            lse_val = mLSE[sink_head_idx, sink_seqlen.offset_q + sink_row]
+                        lse_val = Float32(lse_val)
+                        sink_prob = (
+                            Float32(1.0)
+                            if lse_val == -Float32.inf
+                            else cute.math.exp2(
+                                (sink_val - lse_val) * utils.LOG2_E,
+                                fastmath=True,
+                            )
+                        )
+                        sink_sum += -sink_prob * Float32(dpsum_val)
+                        sink_row += self.num_threads
+                    sink_batch += 1
+
+                sink_sum = utils.warp_reduce(sink_sum, operator.add)
+                lane_idx = cute.arch.lane_idx()
+                warp_idx = tidx // cute.arch.WARP_SIZE
+                num_warps = self.num_threads // cute.arch.WARP_SIZE
+                if lane_idx == 0:
+                    sdQaccum_flat[warp_idx] = sink_sum
+                cute.arch.barrier(barrier_id=5, number_of_threads=self.num_threads)
+                if warp_idx == 0:
+                    sink_sum = sdQaccum_flat[lane_idx] if lane_idx < num_warps else Float32(0.0)
+                    sink_sum = utils.warp_reduce(sink_sum, operator.add)
+                    if lane_idx == 0:
+                        mdSink[sink_head_idx] = sink_sum.to(mdSink.element_type)
+                cute.arch.barrier(barrier_id=5, number_of_threads=self.num_threads)


How is performance here? Do we get enough parallelism?

drisspg · 2026-07-31T00:26:48Z

+                num_warps = self.num_threads // cute.arch.WARP_SIZE
+                if lane_idx == 0:
+                    sdQaccum_flat[warp_idx] = sink_sum
+                cute.arch.barrier(barrier_id=5, number_of_threads=self.num_threads)


Good follow up; we should switch to named barriers

wait actually can this just be: cute.arch.sync_threads()?

henrylhtsang changed the title ~~Support learnable sink backward in CuTe FA~~ [Cute, bwd, sm90/100/110] Support learnable sink backward in CuTe FA Jul 13, 2026

henrylhtsang marked this pull request as ready for review July 16, 2026 01:39

henrylhtsang changed the title ~~[Cute, bwd, sm90/100/110] Support learnable sink backward in CuTe FA~~ [Cute, bwd, sm90/100/110] Support learnable sink in backward Jul 16, 2026

henrylhtsang marked this pull request as draft July 16, 2026 02:04

henrylhtsang marked this pull request as ready for review July 16, 2026 04:23

henrylhtsang closed this Jul 28, 2026

henrylhtsang reopened this Jul 29, 2026

henrylhtsang and others added 16 commits July 28, 2026 17:29

Support learnable sink backward

9d1ca23

Simplify learnable sink backward plumbing

b406cda

Format learnable sink postprocess

44872b3

Support learnable sink with hd256 and frozen QKV

813f41d

Format standalone sink reduction

010975d

Tighten learnable sink backward checks

057cb43

Keep learnable sink scope lean

d7139b7

Remove standalone learnable sink varlen test

c059c07

Always test learnable sink backward

e910c11

Simplify learnable sink test setup

fc53958

Detach learnable sink test tensors

9ec06ef

Use Tuple return type for backward

840fc75

Remove sink backward tensor wrapper

619ecac

Clarify sink reduction CTA selection

1c69b4d

Refine learnable sink backward handling

4a0d5ba

Handle sink-only rows in backward

c193150

henrylhtsang marked this pull request as draft July 29, 2026 00:35

drisspg reviewed Jul 29, 2026

View reviewed changes

Address learnable sink review feedback

0d49703

henrylhtsang force-pushed the learnable-sink-backward branch from da6a8b0 to 0d49703 Compare July 29, 2026 01:10

henrylhtsang commented Jul 29, 2026

View reviewed changes

Separate learnable sink dtype coverage

37a3b54

henrylhtsang marked this pull request as ready for review July 29, 2026 23:49

henrylhtsang marked this pull request as draft July 29, 2026 23:54

Use dtype-aware learnable sink gradient tolerance

77024cd

henrylhtsang marked this pull request as ready for review July 30, 2026 00:34

Handle empty learnable sink backward

7662df9

henrylhtsang requested a review from drisspg July 30, 2026 22:22

drisspg reviewed Jul 31, 2026

View reviewed changes

Uh oh!

Conversation

henrylhtsang commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

henrylhtsang commented Jul 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

henrylhtsang commented Jul 29, 2026

Uh oh!

henrylhtsang commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

drisspg left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

henrylhtsang commented Jul 13, 2026 •

edited

Loading

henrylhtsang commented Jul 30, 2026 •

edited

Loading