[perf]: extend FA2 register_autograd parity to the masked/varlen paths

Builds on the FA2 default-path real-backward (commit 7dbc1e0). Wraps
the two remaining flash-attn entry points FlashAttentionImpl.forward
calls — flash_attn_no_pad (masked self-attn) and
flash_attn_varlen_qk_no_pad (cross-attn / unequal q-k seqlen) — as
torch.library.custom_ops with full register_autograd on FA2. Same
trade as the default path: dynamo sees one traceable node, the
internal unpad/pad bookkeeping runs eager inside, and training
backprops through the op (no graph break on either the inference or
training path).

The non-trivial part vs the default leg is `softmax_lse`. FA2 varlen
returns lse in the unpadded form ([nheads, total_q]); to keep the
custom op's outputs statically-shaped (so register_fake matches),
we pad lse out to [batch, nheads, seqlen] before returning and
re-unpad in backward using the saved mask. The backward then re-unpads
qkv/out/dout via unpad_input and calls FA2's
flash_attn.flash_attn_interface._flash_attn_varlen_backward on the
unpadded form, then re-pads d{q,k,v} back to the input shape.

softmax_scale=None is resolved to `head_dim**-0.5` in setup_context
(FA2's varlen backward demands a concrete float, same as the default
leg).

FA3 / FA4 keep the autograd carve-out pattern from hao-ai-lab#1373: forward+fake
only, dispatcher falls back to the original autograd.Function for
grad-enabled calls. Those legs ship as separate follow-ups gated on
Hopper / Blackwell box validation.

Wire FlashAttentionImpl.forward's masked branch to call the
*_compilable wrappers instead of the originals, so the graph-break
elimination + autograd parity actually takes effect at the call site.

Add fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py:
  - inference parity (atol=0, rtol=0) for both ops
  - backward-through-registered-autograd for both ops (FA2-only;
    dq/dk/dv match the original autograd.Function within dtype tol)
  - torch.library.opcheck with and without grad inputs (the
    with-grad case exercises test_autograd_registration —
    structurally identical to the gap that let hao-ai-lab#1373's first
    revision ship without a backward).

GPU-gated on CUDA + FA2 v2.8.1. Heavy SSIM gate (HunyuanVideo-1.5)
on a 40-80 GB box is the integration check for the masked path
end-to-end.

Author: Mister-Raggs

fastvideo/attention/backends/flash_attn.py

@@ -417,9 +417,17 @@ def _forward_impl(
         attn_metadata: FlashAttnMetadata,
     ):
         if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None):
+            # Route through the *_compilable wrappers so dynamo sees one
+            # traceable node for each masked entry point (the unpad/pad
+            # bookkeeping runs eager inside the custom op). On FA2 these
+            # wrappers go through ops with full register_autograd, so
+            # training also backprops through the op (no graph break on
+            # the training path); on FA3/FA4 they carve out to the
+            # autograd.Function for grad-enabled calls — see
+            # fastvideo/attention/utils/flash_attn_no_pad.py.
             from fastvideo.attention.utils.flash_attn_no_pad import (
-                flash_attn_no_pad,
-                flash_attn_varlen_qk_no_pad,
+                flash_attn_no_pad_compilable as flash_attn_no_pad,
+                flash_attn_varlen_qk_no_pad_compilable as flash_attn_varlen_qk_no_pad,
             )
 
             attn_mask = attn_metadata.attn_mask