[perf]: address hao-ai-lab#1388 bot review — mark lse non-diff + reuse indices in backward

Two improvements, both from the gemini-code-assist + Copilot reviews on hao-ai-lab#1388:

(a) Mark `softmax_lse` non-differentiable in all 3 custom-op setup_contexts
    (default + masked + varlen_qk). We return lse alongside out so it can be
    saved for backward; nobody should differentiate through it. `del grad_lse`
    in backward silently drops grads if a caller wires lse into a loss —
    `ctx.mark_non_differentiable(lse)` makes autograd error loudly instead.
    The public `*_compilable` dispatchers already drop lse (`out, _ = op(...)`),
    so this only matters for callers using `torch.ops.fastvideo._*_forward`
    directly, but it's free defensive hygiene.

(b) Stop re-running `unpad_input` for tensors that share a mask. The first
    `unpad_input(mask)` call returns `indices` + `cu_seqlens` + `max_s`;
    subsequent unpads for out/dout/lse on the same mask just re-derived
    the same indices (and the `.max().item()` call inside each unpad does
    a GPU→CPU sync). Replace those with `tensor.flatten(0,1)[indices].view(...)`,
    which is what `unpad_input` does internally anyway:

      - flash_attn_no_pad backward: 4 unpad_input calls -> 1
      - flash_attn_varlen_qk_no_pad backward: 6 unpad_input calls -> 2
        (one per distinct mask; v reuses k_indices, out/dout/lse reuse
        q_indices, the final repad of dk/dv reuses k_indices instead of
        recomputing).

Semantics are preserved — `unpad_input(x, mask)[0]` is equivalent to
`x.flatten(0,1)[indices]` for the same mask. The existing 28 unit tests
(inference parity atol=0 rtol=0, training-backward parity, opcheck with
and without grad inputs, all on FA2 + A100) cover the regressions.

Author: Mister-Raggs

fastvideo/attention/backends/flash_attn.py

@@ -93,6 +93,12 @@ def _flash_attn_default_setup_context(ctx, inputs, output):
         q, k, v, softmax_scale, causal = inputs
         out, lse = output
         ctx.save_for_backward(q, k, v, out, lse)
+        # `lse` is an auxiliary output we save to feed FA2's backward; nobody
+        # should differentiate through it. Mark it non-differentiable so
+        # autograd errors loudly if a caller wires it into a loss, rather
+        # than silently producing zero/None grads through the `del grad_lse`
+        # in our backward.
+        ctx.mark_non_differentiable(lse)
         # FA2's *forward* substitutes `1 / sqrt(head_dim)` for `softmax_scale=None`
         # internally; FA2's *backward* (`_flash_attn_backward`) demands a concrete
         # float in its C++ schema and rejects None at the binding boundary. Resolve

fastvideo/attention/utils/flash_attn_no_pad.py

@@ -270,6 +270,8 @@ def _flash_attn_no_pad_setup_context(ctx, inputs, output):
         qkv, key_padding_mask, causal, dropout_p, softmax_scale, deterministic = inputs
         out, lse = output
         ctx.save_for_backward(qkv, out, lse, key_padding_mask)
+        # Auxiliary output, not differentiable — see default-path note.
+        ctx.mark_non_differentiable(lse)
         # FA2's varlen backward requires a concrete float for softmax_scale.
         if softmax_scale is None:
             softmax_scale = qkv.shape[-1] ** -0.5  # head_dim from qkv's last dim
@@ -284,23 +286,20 @@ def _flash_attn_no_pad_backward(ctx, grad_out, grad_lse):
         qkv, out_padded, lse_padded, key_padding_mask = ctx.saved_tensors
         b, s, _three, h, d = qkv.shape
 
-        # Re-unpad qkv -> q, k, v unpadded ([nnz, h, d] each).
+        # One `unpad_input` call (on qkv) gives us indices + cu_seqlens + max_s;
+        # reuse those for out / dout / lse below via direct indexing instead
+        # of redundant `unpad_input` calls (each of which would re-run
+        # `nonzero` + `cumsum` + a `.max().item()` GPU→CPU sync).
         x = rearrange(qkv, "b s three h d -> b s (three h d)")
         x_unpad, indices, cu_seqlens, max_s, _ = unpad_input(x, key_padding_mask)
         x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=h)
         q_unpad, k_unpad, v_unpad = (t.contiguous() for t in x_unpad.unbind(dim=1))
 
-        # Re-unpad out and dout using the same mask.
-        out_unpad = rearrange(
-            unpad_input(rearrange(out_padded, "b s h d -> b s (h d)"), key_padding_mask)[0],
-            "nnz (h d) -> nnz h d", h=h).contiguous()
-        dout_unpad = rearrange(
-            unpad_input(rearrange(grad_out, "b s h d -> b s (h d)"), key_padding_mask)[0],
-            "nnz (h d) -> nnz h d", h=h).contiguous()
-
-        # Re-unpad lse: [b, h, s] -> [b, s, h] -> [nnz, h] -> [h, nnz].
-        lse_unpad = unpad_input(lse_padded.permute(0, 2, 1).contiguous(),
-                                key_padding_mask)[0].t().contiguous()
+        # Direct-index variants reuse `indices` (computed above).
+        out_unpad = out_padded.flatten(0, 1)[indices].view(-1, h, d).contiguous()
+        dout_unpad = grad_out.flatten(0, 1)[indices].view(-1, h, d).contiguous()
+        # lse_padded [b, h, s] -> [b, s, h] -> [nnz, h] -> [h, nnz].
+        lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[indices].t().contiguous()
 
         dq_unpad = torch.empty_like(q_unpad)
         dk_unpad = torch.empty_like(k_unpad)
@@ -402,6 +401,8 @@ def _flash_attn_varlen_qk_no_pad_setup_context(ctx, inputs, output):
         out, lse = output
         ctx.save_for_backward(query, key, value, out, lse,
                               query_padding_mask, key_padding_mask)
+        # Auxiliary output, not differentiable — see default-path note.
+        ctx.mark_non_differentiable(lse)
         if softmax_scale is None:
             softmax_scale = query.shape[-1] ** -0.5
         ctx.softmax_scale = softmax_scale
@@ -416,26 +417,24 @@ def _flash_attn_varlen_qk_no_pad_backward(ctx, grad_out, grad_lse):
         b, sq, h, d = query.shape
         sk = key.shape[1]
 
-        # Re-unpad q with q_mask; k, v with k_mask.
+        # One `unpad_input` call per distinct mask; reuse the returned
+        # indices via direct indexing for everything else that shares
+        # the same mask (v with k_mask; out/dout/lse with q_mask; the
+        # final repad of dk/dv also reuses k_indices). Avoids ~4
+        # redundant `unpad_input` calls + their GPU→CPU `.max().item()`
+        # syncs.
         q_unpad, q_indices, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
             rearrange(query, "b s h d -> b s (h d)"), query_padding_mask)
         k_unpad, k_indices, cu_seqlens_k, max_seqlen_k, _ = unpad_input(
             rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
-        v_unpad, _, _, _, _ = unpad_input(
-            rearrange(value, "b s h d -> b s (h d)"), key_padding_mask)
         q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
         k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
-        v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
-
-        # out and lse follow q's shape, so re-unpad with q_mask.
-        out_unpad = rearrange(
-            unpad_input(rearrange(out_padded, "b s h d -> b s (h d)"), query_padding_mask)[0],
-            "nnz (h d) -> nnz h d", h=h).contiguous()
-        dout_unpad = rearrange(
-            unpad_input(rearrange(grad_out, "b s h d -> b s (h d)"), query_padding_mask)[0],
-            "nnz (h d) -> nnz h d", h=h).contiguous()
-        lse_unpad = unpad_input(lse_padded.permute(0, 2, 1).contiguous(),
-                                query_padding_mask)[0].t().contiguous()
+        v_unpad = value.flatten(0, 1)[k_indices].view(-1, h, d).contiguous()
+
+        # out / dout / lse follow q's shape, so index with q_indices.
+        out_unpad = out_padded.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
+        dout_unpad = grad_out.flatten(0, 1)[q_indices].view(-1, h, d).contiguous()
+        lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[q_indices].t().contiguous()
 
         dq_unpad = torch.empty_like(q_unpad)
         dk_unpad = torch.empty_like(k_unpad)
@@ -454,6 +453,8 @@ def _flash_attn_varlen_qk_no_pad_backward(ctx, grad_out, grad_lse):
             rng_state=None,
         )
 
+        # k_indices is already available from the unpad_input above —
+        # no need to recompute it for the dk/dv repad.
         def _repad(dt_unpad, indices, batch, seqlen):
             padded = pad_input(rearrange(dt_unpad, "nnz h d -> nnz (h d)"), indices, batch, seqlen)
             return rearrange(padded, "b s (h d) -> b s h d", h=h)