[refactor]: move FA default custom-op registration to attention/utils/

Addresses the S2 finding from the post-merge review on hao-ai-lab#1373: per FastVideo
precedent (fastvideo/attention/utils/flash_attn_cute.py — the FP4 cute
custom-op wrapper hao-ai-lab#1388 says it mirrors) and vLLM's similar pattern
(vllm/utils/flashinfer.py), per-kernel custom-op wrappers live in
attention/utils/, not inline in the backend dispatcher.

Move the entire `if fa_version == "2": ... elif "3": ... elif "4": ... else:`
block — including the FA version probe, custom_op + register_fake +
setup_context + register_autograd, and `flash_attn_func_compilable` —
from `attention/backends/flash_attn.py` to a new sibling at
`attention/utils/flash_attn_default.py`. The backend now imports
`fa_version` and `flash_attn_func_compilable` from there; the rest of the
backend (FlashAttentionBackend, FlashAttentionImpl, FP4 quantize helpers,
masked-branch routing) is unchanged.

Benefits:
  - Internal consistency with flash_attn_cute.py (the FP4 template the
    file-level comment claims to mirror) and flash_attn_no_pad.py (the
    masked/varlen wrappers, already in utils/).
  - A future backend that wants the same traceable FA2/FA3/FA4 op can
    import from utils/ without depending on the FA backend's dispatcher.
  - Cleaner separation between backend registry/dispatch logic and
    per-kernel custom-op plumbing.

No semantic change. Verified there is exactly one registration of
`fastvideo::_flash_attn_default_forward` after the move (grep clean).
Update `test_flash_attn_default_custom_op.py` to import from the utils
module (the test referenced `fa_backend._fa_default`, a name that lived
inside the moved block).

Author: Mister-Raggs

fastvideo/attention/backends/flash_attn.py

@@ -1,13 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import importlib.util
 import os
-
 import torch
 import torch.nn.functional as F
 from dataclasses import dataclass
 
-from fastvideo import envs
+from fastvideo.attention.utils.flash_attn_default import (
+    fa_version,
+    flash_attn_func_compilable,
+)
+
 from fastvideo.attention.backends.abstract import (
     AttentionBackend,
     AttentionImpl,
@@ -17,233 +19,6 @@
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
-
-# FA4 (flash_attn.cute) is explicit opt-in via FASTVIDEO_FA4=1, mirroring the
-# kernel package's FASTVIDEO_VSA_CUTEDSL: its CuTeDSL kernels JIT-compile per
-# shape family and can fail at runtime on some arch/shape combinations, so it
-# is never auto-selected just because it is installed. Below sm90 a capability
-# gate in flash_attn_cute routes to FA2 the calls FA4 cannot serve there:
-# grad-enabled (its backward asserts sm90+) and GQA (pack_gqa fails CuTeDSL
-# JIT, observed on sm_89).
-if envs.FASTVIDEO_FA4:
-    try:
-        from fastvideo.attention.utils.flash_attn_cute import flash_attn_func
-    except ImportError as e:
-        raise RuntimeError(f"FASTVIDEO_FA4=1 but flash_attn.cute (FA4) is not usable ({e}); "
-                           "fix the FA4 install (see the flash-attn-4 pin in pyproject.toml) "
-                           "or unset FASTVIDEO_FA4.") from e
-    fa_version = "4"
-else:
-    try:
-        from flash_attn_interface import flash_attn_func as flash_attn_3_func
-
-        # flash_attn 3 no longer have a different API, see following commit:
-        # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7
-        flash_attn_func = flash_attn_3_func
-        fa_version = "3"
-    except ImportError:
-        from flash_attn import flash_attn_func as flash_attn_2_func
-        flash_attn_func = flash_attn_2_func
-        fa_version = "2"
-    try:
-        if importlib.util.find_spec("flash_attn.cute") is not None:
-            logger.info("flash_attn.cute (FA4) is installed but not enabled; "
-                        "set FASTVIDEO_FA4=1 to use it for inference.")
-    except ImportError:
-        pass
-
-# torch.compile traceability: the FA4/cute path (fa_version=="4") is
-# already a registered torch.library custom op, so dynamo treats it as a
-# graph node. The external FA2/FA3 `flash_attn_func` is NOT — dynamo
-# breaks the graph at the call site (observed: wanvideo.py self-attn,
-# once per layer every step), which fragments the compiled region and
-# blocks CUDA-graph capture. Wrap the FA2/FA3 default call in a custom
-# op (mirrors the FP4 `flash_attn_cute` template) so it becomes an
-# 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).
-#
-# 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
-    # 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)
-        # `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
-        # 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=(),
-        device_types="cuda",
-    )
-    def _flash_attn_default_forward(
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        softmax_scale: float | None,
-        causal: bool,
-    ) -> torch.Tensor:
-        return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
-
-    @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,
-    ) -> torch.Tensor:
-        del softmax_scale, causal
-        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):
-        # Autograd carve-out. The custom op above registers a forward + fake
-        # kernel but NO backward (register_autograd), so it is opaque to
-        # autograd. Inference runs under no_grad / inference_mode and routes
-        # through the traceable custom op — that is the torch.compile win, and
-        # the only path this PR claims. Training backprops through attention,
-        # so route grad-enabled calls to the original FA2/FA3 `flash_attn_func`
-        # (itself an autograd.Function, so backward is correct) at the cost of a
-        # dynamo graph break on the training path — i.e. pre-PR behavior, no
-        # regression. Full autograd parity for the custom op (mirroring the FP4
-        # cute template) is a tracked follow-up.
-        if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad):
-            return _fa_default(q, k, v, softmax_scale=softmax_scale, causal=causal)
-        return torch.ops.fastvideo._flash_attn_default_forward(q, k, v, softmax_scale, causal)
-elif fa_version == "4":
-    # FA4 path: `flash_attn_func` (from `flash_attn_cute`) goes through a
-    # registered torch.library custom op (with an FA4 backward on sm90+;
-    # grad-enabled and GQA calls below sm90 route to FA2), so a passthrough
-    # is enough — no extra registration needed.
-    def flash_attn_func_compilable(q, k, v, softmax_scale=None, causal=False):
-        return flash_attn_func(q, k, v, softmax_scale=softmax_scale, causal=causal)
-else:
-    # Defensive: the probe above only ever sets fa_version to "2", "3",
-    # or "4"; an unexpected value means an import/probe regression and
-    # we want a loud error at import, not a silent NameError later.
-    raise RuntimeError(f"Unsupported FlashAttention version: {fa_version!r} — expected "
-                       f"'2', '3', or '4' from the import probe above.")
-
 logger.info("Using FlashAttention-%s backend", fa_version)
 
 # FP4 FA4 support: quantize Q/K to NVFP4 E2M1 for block-scaled MMA on Blackwell.