Commit aad28dd
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
- utils
- tests/attention
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
373 | 373 | | |
374 | 374 | | |
375 | 375 | | |
| 376 | + | |
| 377 | + | |
| 378 | + | |
| 379 | + | |
| 380 | + | |
| 381 | + | |
| 382 | + | |
| 383 | + | |
376 | 384 | | |
377 | | - | |
378 | | - | |
| 385 | + | |
| 386 | + | |
379 | 387 | | |
380 | 388 | | |
381 | 389 | | |
| |||
0 commit comments