Skip to content

Commit aad28dd

Browse files
committed
[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.
1 parent ce044a7 commit aad28dd

3 files changed

Lines changed: 664 additions & 2 deletions

File tree

fastvideo/attention/backends/flash_attn.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -373,9 +373,17 @@ def forward(
373373
attn_metadata: FlashAttnMetadata,
374374
):
375375
if (attn_metadata is not None and hasattr(attn_metadata, "attn_mask") and attn_metadata.attn_mask is not None):
376+
# Route through the *_compilable wrappers so dynamo sees one
377+
# traceable node for each masked entry point (the unpad/pad
378+
# bookkeeping runs eager inside the custom op). On FA2 these
379+
# wrappers go through ops with full register_autograd, so
380+
# training also backprops through the op (no graph break on
381+
# the training path); on FA3/FA4 they carve out to the
382+
# autograd.Function for grad-enabled calls — see
383+
# fastvideo/attention/utils/flash_attn_no_pad.py.
376384
from fastvideo.attention.utils.flash_attn_no_pad import (
377-
flash_attn_no_pad,
378-
flash_attn_varlen_qk_no_pad,
385+
flash_attn_no_pad_compilable as flash_attn_no_pad,
386+
flash_attn_varlen_qk_no_pad_compilable as flash_attn_varlen_qk_no_pad,
379387
)
380388

381389
attn_mask = attn_metadata.attn_mask

0 commit comments

Comments
 (0)