[perf]: register a real backward for the FA2 default custom op

Follow-up to hao-ai-lab#1373 + its CI carve-out (commit fae1cd1). The custom op
fastvideo::_flash_attn_default_forward shipped with a forward + fake
kernel but no register_autograd, so it was opaque to autograd; the
carve-out routed grad-enabled calls back to the original FA2
flash_attn_func (an autograd.Function) at the cost of a graph break on
the training path. This commit closes that gap for FA2 by mirroring the
FP4 cute template's 4-piece pattern:

  - forward returns (out, softmax_lse), obtained via flash_attn_func's
    return_attn_probs=True path so we don't touch private FA2 fwd APIs;
  - register_fake returns the matching tuple, with softmax_lse fixed at
    [batch, nheads, seqlen_q] fp32;
  - setup_context saves q,k,v,out,lse + softmax_scale, causal;
  - the backward calls flash_attn.flash_attn_interface._flash_attn_backward
    (FA2's private bwd) with all version-fragile kwargs pinned to the
    flash-attn==2.8.1 defaults (the version FastVideo pins).

With autograd registered, flash_attn_func_compilable can drop its carve-
out — grad-enabled calls go through the op like inference calls, so the
training path is also dynamo-traceable. Numerics unchanged (lse is saved-
for-backward only, never differentiated; backward writes the same dq/dk/
dv FA2's autograd.Function would produce).

FA3 keeps the carve-out from fae1cd1 untouched. FA3 exposes a different
private _flash_attn_backward signature that wants validation on a real
Hopper box; once Kuan-Hao's Modal FA3 setup PR lands we mirror the FA2
pattern there.

Tests: extend test_flash_attn_default_custom_op.py with
  - test_default_op_backward_through_registered_autograd: calls
    torch.ops.fastvideo._flash_attn_default_forward directly with
    requires_grad inputs, autograd.grad through the op, asserts grads
    match the original flash_attn_func to dtype-appropriate tolerance.
  - test_default_op_opcheck_with_grad_inputs: torch.library.opcheck
    with requires_grad inputs (exercises test_autograd_registration —
    the gap that let hao-ai-lab#1373's first revision ship without a backward).
Both gate on fa_version=="2"; FA3 retains the carve-out tests as-is.

Directional confirmation from Will Lin (FastVideo maintainer) 2026-05-22:
"it definitely is [wanted], but testing the correctness and performance
is more involved".

Author: Mister-Raggs

fastvideo/attention/backends/flash_attn.py

@@ -32,19 +32,129 @@
 # opaque-but-traceable node. The kernel still runs eager inside the op
 # (correct — flash-attn must run eager); only dynamo's treatment of the
 # boundary changes, so numerics are unchanged (SSIM-gate to confirm).
-if fa_version in ("2", "3"):
-    _fa_default = flash_attn_func
-
+#
+# Autograd: FA2 has full register_autograd parity — the custom op's
+# backward calls flash_attn's `_flash_attn_backward` directly, so
+# training backprops *through* the op (no graph break on the training
+# path either). FA3 currently keeps the no-backward + carve-out pattern
+# from PR #1373 because FA3's private backward signature wants
+# validation on a real Hopper box (gated on Kuan-Hao's Modal FA3 setup
+# PR). Once that lands the FA3 path can mirror FA2.
+if fa_version == "2":
     # Scope: this op covers exactly the q/k/v + softmax_scale + causal
     # call shape used by FlashAttentionImpl.forward's default branch
     # (see `flash_attn_func_compilable(...)` call site below). The
-    # masked/no-pad and varlen / cross-attn paths use different
-    # entry points (`flash_attn_no_pad`, `flash_attn_varlen_*`) which
-    # are intentionally out of scope for this PR — wrapping them is a
+    # masked/no-pad and varlen / cross-attn paths use different entry
+    # points (`flash_attn_no_pad`, `flash_attn_varlen_*`) which are
+    # intentionally out of scope for this PR — wrapping them is a
     # natural follow-up. The wrapper's signature is the contract: any
     # extra kwarg (dropout_p, window_size, alibi_slopes, deterministic,
     # return_attn_probs, ...) raises TypeError at the call site, so
     # silent loss of kwargs is not a failure mode.
+    from flash_attn.flash_attn_interface import _flash_attn_backward as _fa2_backward
+    _fa_default = flash_attn_func
+
+    @torch.library.custom_op(
+        "fastvideo::_flash_attn_default_forward",
+        mutates_args=(),
+        device_types="cuda",
+    )
+    def _flash_attn_default_forward(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: float | None,
+        causal: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # `return_attn_probs=True` asks FA2 to also return softmax_lse +
+        # S_dmask. We need softmax_lse to feed the backward; S_dmask is the
+        # dropout mask (always None here since dropout_p is fixed at 0).
+        out, softmax_lse, _ = _fa_default(q, k, v, softmax_scale=softmax_scale,
+                                          causal=causal, return_attn_probs=True)
+        return out, softmax_lse
+
+    @torch.library.register_fake("fastvideo::_flash_attn_default_forward")
+    def _flash_attn_default_forward_fake(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        softmax_scale: float | None,
+        causal: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del softmax_scale, causal
+        # FA2 default path: out = [batch, seqlen_q, nheads, head_dim_v],
+        # softmax_lse = [batch, nheads, seqlen_q], fp32 regardless of q dtype.
+        b, sq, hq = q.shape[0], q.shape[1], q.shape[2]
+        out = q.new_empty(b, sq, hq, v.shape[-1])
+        lse = q.new_empty(b, hq, sq, dtype=torch.float32)
+        return out, lse
+
+    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)
+        # 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
+        # the default here so the value saved on ctx (and passed to backward) is
+        # always a real float — matches what FA2's own autograd.Function does.
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** -0.5
+        ctx.softmax_scale = softmax_scale
+        ctx.causal = causal
+
+    def _flash_attn_default_backward(ctx, grad_out, grad_lse):
+        # We only differentiate `out`; softmax_lse is saved-for-backward, not
+        # a real differentiable output. (Mirrors the FP4 cute template.)
+        del grad_lse
+        q, k, v, out, lse = ctx.saved_tensors
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        # FA2's `_flash_attn_backward` writes into dq/dk/dv in place. The
+        # extra kwargs (window_size_*, softcap, alibi_slopes, deterministic,
+        # rng_state) are pinned to the same defaults the forward wrapper
+        # uses — flash-attn==2.8.1 (the version FastVideo pins) requires
+        # all of them explicitly. `rng_state=None` is correct for our
+        # `dropout_p=0` configuration.
+        _fa2_backward(
+            grad_out, q, k, v, out, lse,
+            dq, dk, dv,
+            dropout_p=0.0,
+            softmax_scale=ctx.softmax_scale,
+            causal=ctx.causal,
+            window_size_left=-1,
+            window_size_right=-1,
+            softcap=0.0,
+            alibi_slopes=None,
+            deterministic=False,
+            rng_state=None,
+        )
+        return dq, dk, dv, None, None
+
+    torch.library.register_autograd(
+        "fastvideo::_flash_attn_default_forward",
+        _flash_attn_default_backward,
+        setup_context=_flash_attn_default_setup_context,
+    )
+
+    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
+        # Backward is registered: autograd flows through the op (training
+        # path is also traceable; no carve-out needed). Public API matches
+        # `flash_attn_func` — returns just `out`; we drop the saved-for-
+        # backward `lse` here so callers see the original single-tensor
+        # contract.
+        out, _ = torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
+        return out
+elif fa_version == "3":
+    # FA3 path: same forward+fake custom op as the original PR #1373, with
+    # the autograd carve-out kept. The full backward (mirroring the FA2
+    # leg above) wants a Hopper box for grad-check validation, which we
+    # don't have until Kuan-Hao's Modal FA3 setup PR lands. Until then
+    # this keeps inference traceable + training correct (via the original
+    # autograd.Function path + a pre-PR-style graph break on training).
+    _fa_default = flash_attn_func
+
     @torch.library.custom_op(
         "fastvideo::_flash_attn_default_forward",
         mutates_args=(),
@@ -68,8 +178,6 @@ def _flash_attn_default_forward_fake(
         causal: bool,
     ) -> torch.Tensor:
         del softmax_scale, causal
-        # FA2/FA3 default path: [batch, seqlen_q, nheads, head_dim_v],
-        # same dtype/device as q (head dim taken from v).
         return q.new_empty(q.shape[0], q.shape[1], q.shape[2], v.shape[-1])
 
     def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):

fastvideo/tests/attention/test_flash_attn_default_custom_op.py

@@ -42,12 +42,14 @@ def fa_default_impls():
             f"got {fa_backend.fa_version!r}"
         )
 
-    # compilable dispatcher, the original FA wrapper it falls back to, and
-    # the raw custom op for opcheck.
+    # compilable dispatcher, the original FA wrapper it falls back to, the
+    # raw custom op for opcheck, and the fa_version (FA2 has full register_
+    # autograd; FA3 keeps the carve-out so some tests gate on this).
     return (
         fa_backend.flash_attn_func_compilable,
         fa_backend._fa_default,
         torch.ops.fastvideo._flash_attn_default_forward,
+        fa_backend.fa_version,
     )
 
 
@@ -71,7 +73,7 @@ def test_default_compilable_inference_matches_original(fa_default_impls, dtype,
     """No-grad path routes through the custom op and is numerically identical."""
     if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
         pytest.skip("bfloat16 is not supported on this GPU")
-    compilable, original, _ = fa_default_impls
+    compilable, original, _, _ = fa_default_impls
 
     torch.manual_seed(0)
     q, k, v = _qkv(dtype, requires_grad=False)
@@ -91,7 +93,7 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau
     """
     if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
         pytest.skip("bfloat16 is not supported on this GPU")
-    compilable, original, _ = fa_default_impls
+    compilable, original, _, _ = fa_default_impls
 
     torch.manual_seed(0)
     q_ref, k_ref, v_ref = _qkv(dtype, requires_grad=True)
@@ -115,7 +117,65 @@ def test_default_compilable_training_backward_flows(fa_default_impls, dtype, cau
 @pytest.mark.parametrize("causal", [False, True])
 def test_default_forward_opcheck(fa_default_impls, causal):
     """Schema / fake-kernel consistency for the custom op (forward only)."""
-    _, _, op = fa_default_impls
+    _, _, op, _ = fa_default_impls
     torch.manual_seed(0)
     q, k, v = _qkv(torch.float16, requires_grad=False)
     torch.library.opcheck(op, (q, k, v, None, causal))
+
+
+# --------------------------------------------------------------------------- #
+# FA2-only: backward is registered on the custom op itself. Exercises the     #
+# `register_autograd` wiring directly (not via the dispatcher's carve-out).   #
+# Skipped on FA3 until Kuan-Hao's Modal FA3 setup PR lands and we mirror the  #
+# pattern there.                                                              #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("causal", [False, True])
+def test_default_op_backward_through_registered_autograd(fa_default_impls, dtype, causal):
+    """FA2: gradients flow through ``torch.ops.fastvideo._flash_attn_default_forward``
+    itself (no dispatcher carve-out involved) and match the original
+    ``flash_attn_func``'s gradients."""
+    if dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
+        pytest.skip("bfloat16 is not supported on this GPU")
+    _, original, op, fa_version = fa_default_impls
+    if fa_version != "2":
+        pytest.skip(f"register_autograd is only wired for FA2 right now; got {fa_version!r}")
+
+    torch.manual_seed(0)
+    q_ref, k_ref, v_ref = _qkv(dtype, requires_grad=True)
+    q_test, k_test, v_test = _clone(q_ref, k_ref, v_ref)
+
+    # Reference grads via the original autograd.Function.
+    out_ref = original(q_ref, k_ref, v_ref, softmax_scale=None, causal=causal)
+    # Custom-op grads via the registered backward — unpack (out, lse), discard lse.
+    out_test, _ = op(q_test, k_test, v_test, None, causal)
+
+    torch.testing.assert_close(out_test, out_ref,
+                               atol=0 if dtype == torch.float16 else 1e-3,
+                               rtol=0 if dtype == torch.float16 else 1e-3)
+
+    dout = torch.randn_like(out_ref)
+    dq_ref, dk_ref, dv_ref = torch.autograd.grad(
+        (out_ref * dout).sum(), (q_ref, k_ref, v_ref))
+    dq_test, dk_test, dv_test = torch.autograd.grad(
+        (out_test * dout).sum(), (q_test, k_test, v_test))
+
+    atol = rtol = 6e-3 if dtype == torch.float16 else 2e-2
+    torch.testing.assert_close(dq_test, dq_ref, atol=atol, rtol=rtol)
+    torch.testing.assert_close(dk_test, dk_ref, atol=atol, rtol=rtol)
+    torch.testing.assert_close(dv_test, dv_ref, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("causal", [False, True])
+def test_default_op_opcheck_with_grad_inputs(fa_default_impls, causal):
+    """FA2: full ``opcheck`` including ``test_autograd_registration`` —
+    catches a missing/inconsistent backward at unit-test time, which was
+    exactly the gap that #1373's first revision shipped."""
+    _, _, op, fa_version = fa_default_impls
+    if fa_version != "2":
+        pytest.skip(f"autograd registration only wired for FA2; got {fa_version!r}")
+    torch.manual_seed(0)
+    q, k, v = _qkv(torch.float16, requires_grad=True)
+    torch.library.opcheck(op, (q, k, v, None, causal))