[perf]: register a real backward for FA2 default + masked/varlen custom ops (training-under-compile)

[Mister-Raggs]

Summary

Adds register_autograd parity to the FA2 default-path and masked/varlen-path custom ops, so dynamo sees a traceable node on both inference and training paths through attention. Mirrors the FP4 cute template's 4-piece pattern (forward + register_fake + setup_context + register_autograd) and follows up on the direction confirmed in Slack ("training-under-compile is definitely wanted").

Prerequisites

Stacks on #1373 (the FA2/FA3 default custom-op wrapper + the CI-unblocking autograd carve-out). The diff against main currently includes #1373's commits; will rebase onto main once #1373 merges.

Scope (FA2 only — FA3 is a separate follow-up)

Path	Status before	Status after
FA2 default flash_attn_func	custom op, no backward → carve-out routes grads to original autograd.Function (graph break on training)	full register_autograd; training also traceable
FA2 masked flash_attn_no_pad	unwrapped (graph break on every call)	custom op with full register_autograd
FA2 varlen flash_attn_varlen_qk_no_pad (cross-attn)	unwrapped (graph break)	custom op with full register_autograd
FA3 default / masked / varlen	same as before #1373 / carve-out	unchanged — separate PR gated on Hopper validation
FA4 (cute)	already custom op	unchanged

The carve-out (flash_attn_func_compilable checking is_grad_enabled and requires_grad) is removed for FA2 — autograd flows through the op like inference does.

Mechanism (the non-trivial bit)

FA2's varlen kernels return softmax_lse in unpadded form [nheads, total_q]. To keep the custom op's outputs statically-shaped (so register_fake matches dynamo's expectation), we pad lse to [batch, nheads, seqlen] on the way out and re-unpad it in backward using the saved mask. Backward then re-unpads qkv / out / dout, allocates dq / dk / dv on the unpadded form, calls flash_attn.flash_attn_interface._flash_attn_varlen_backward, and re-pads gradients back. softmax_scale=None is resolved to head_dim ** -0.5 in setup_context (FA2's backward demands a concrete float).

Validation matrix

Gate	Result
FA2 default-path unit tests (16 cases)	✅ A100, fp16+bf16, causal T/F, inference-parity atol=0 rtol=0, grad through op matches original autograd.Function, opcheck × {with grad, without grad}
FA2 masked/varlen unit tests (12 cases)	✅ A100, both ops × inference-parity + grad-through-op + opcheck × grad+no-grad
Wan2.1-T2V-1.3B end-to-end SSIM (default path)	✅ MIN=1.000000 49/49 frames
HunyuanVideo-1.5 end-to-end SSIM (masked path)	⚠️ Deferred — blocked by upstream bug #1387 (text_encoding.py:221 crashes on HV15's dual-encoder preprocess for both t2v and i2v variants); HV15 is the only model in FastVideo that exercises the masked branch. Unit-test bit-exactness (atol=0 rtol=0) is the proxy until that issue is fixed.

Test command:

pytest fastvideo/tests/attention/test_flash_attn_default_custom_op.py \
       fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py -v
# 28/28 expected on FA2 + CUDA boxes; skips on FA3/FA4

Inference perf A/B (per maintainer request)

Wan2.1-T2V-1.3B, RTX A5000, bf16, sp=1, 480×832 / 49 frames / 20 steps. Per condition: 3 eager + 4 compile runs, 1 warmup discarded, median of measured.

Compile (enable_torch_compile=True):

	Warmup	Median wall
Pre-#1373 (3a67319c)	60.4 s	50.8 s
Post-#1373 main (2f3ca8aa)	53.8 s	49.8 s
This PR (3626db1b)	52.3 s	49.3 s

Cumulative (pre-#1373 → this PR): median −3.0%, warmup −13.5%.

The per-iteration delta is small — Wan's e2e is already compile-squeezed (Kuan's profile bounds compile-payoff at ~few-%). The clearer signal is the −13.5% warmup: removing the graph break at the FA call site shrinks the compiled region, so dynamo has less to trace.

Eager (no compile): 59.2 / 59.0 / 59.1 s — flat. The wrapping adds zero measurable overhead on the no-compile path.

SSIM (this PR vs pre-#1373): MEAN = MIN = MAX = 1.000000 in both eager and compile modes. Bit-identical output across conditions.

This PR's marginal benefit over #1373 on Wan is sub-noise (~1%) — expected. Wan doesn't take the masked/varlen branch this PR wraps; that benefit fires only on HunyuanVideo-1.5, which is blocked by upstream #1387. The full register_autograd for the default path is invisible to inference (autograd doesn't run under inference_mode) — its payoff is on the training-under-compile path, validated structurally by the 28-unit-test suite (grad-through-op parity) but not measured here.

What this PR does NOT measure end-to-end

• Training-under-compile inference wall — the FA2 register_autograd is the centerpiece of this PR for the training path, but a training-loop wall measurement is out of scope here. The 28-unit-test suite validates grad-through-op parity + opcheck with grad inputs, which is the correctness floor.
• HunyuanVideo-1.5 e2e perf — the masked/varlen wrapping only fires on HV15, currently blocked by upstream issue [bug]: HunyuanVideo-1.5 generation crashes at text_encoding.py:221 — 'list' object has no attribute 'strip' #1387. The unit-test bit-exactness (atol=0 rtol=0) is the proxy.

Files

• fastvideo/attention/backends/flash_attn.py — FA2 default custom op gets backward + setup_context + register_autograd; FA3 path keeps carve-out; FA4 path unchanged; masked-branch call site routes through *_compilable dispatchers.
• fastvideo/attention/utils/flash_attn_no_pad.py — wraps flash_attn_no_pad and flash_attn_varlen_qk_no_pad as custom ops with full FA2 register_autograd (lse padding + re-unpad on backward); FA3/FA4 keep carve-out.
• fastvideo/tests/attention/test_flash_attn_default_custom_op.py — extended.
• fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py — new.

[mergify]

Merge Protections

🟢 Merge protection satisfied — ready to merge.

Show 1 satisfied protection

🟢 PR merge requirements

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces torch.library.custom_op wrappers for FlashAttention (FA2/FA3) and masked/varlen attention paths to improve torch.compile traceability and prevent graph breaks. For FA2, full autograd support is registered, while FA3/FA4 utilize a carve-out pattern for training. New regression tests are added to ensure numerical consistency and autograd correctness. Feedback focuses on optimizing the backward passes by reusing computed indices to avoid redundant unpad_input calls.

[gemini-code-assist]

The code calls unpad_input multiple times for out_padded, grad_out, and lse_padded. Since indices is already computed in line 282 using the same key_padding_mask, you can avoid the overhead of unpad_input (which involves finding non-zeros and computing prefix sums) by directly indexing the flattened tensors. Additionally, the .contiguous() calls on the results of unpad_input are redundant as indexing already returns a contiguous copy.

        out_unpad = out_padded.flatten(0, 1)[indices].view(-1, h, d)
        dout_unpad = grad_out.flatten(0, 1)[indices].view(-1, h, d)

        # Re-unpad lse: [b, h, s] -> [b, s, h] -> [nnz, h] -> [h, nnz].
        lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[indices].t().contiguous()

[gemini-code-assist]

key_unpad and v_unpad use the same key_padding_mask. You can compute k_indices once and reuse it to unpad value, avoiding a redundant unpad_input call.

        k_unpad, k_indices, cu_seqlens_k, max_seqlen_k, _ = unpad_input(
            rearrange(key, "b s h d -> b s (h d)"), key_padding_mask)
        v_unpad = rearrange(value, "b s h d -> (b s) (h d)")[k_indices]

[gemini-code-assist]

In the cross-attention backward, v_unpad can reuse k_indices, and out_unpad, dout_unpad, and lse_unpad can all reuse q_indices. This avoids multiple redundant calls to unpad_input and its internal non-zero finding logic.

        v_unpad = rearrange(value, "b s h d -> (b s) (h d)")[k_indices]
        q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
        k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()
        v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h).contiguous()

        # out and lse follow q's shape, so re-unpad with q_mask.
        out_unpad = out_padded.flatten(0, 1)[q_indices].view(-1, h, d)
        dout_unpad = grad_out.flatten(0, 1)[q_indices].view(-1, h, d)
        lse_unpad = lse_padded.permute(0, 2, 1).contiguous().flatten(0, 1)[q_indices].t().contiguous()

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds torch.compile-friendly custom-op wrappers around FlashAttention default and masked/varlen paths, with FA2-only register_autograd to keep training graphable, and introduces regression tests to lock in inference/training equivalence + op schema/fake consistency.

Changes:

• Wrap FA2/FA3 default-path flash_attn_func behind a torch.library.custom_op (FA2: real backward; FA3: carve-out for grads).
• Add masked/varlen custom ops (flash_attn_no_pad*) with FA2 register_autograd and FA3/FA4 carve-out behavior.
• Add CUDA-only regression tests validating inference equality, backward correctness (FA2), and torch.library.opcheck.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 12 comments.

File	Description
fastvideo/attention/backends/flash_attn.py	Introduces default-path custom op + FA2 autograd registration; routes backend calls through compilable wrappers.
fastvideo/attention/utils/flash_attn_no_pad.py	Adds masked/varlen custom ops, FA2 backward wiring via private FA2 varlen backward, and compilable dispatchers.
fastvideo/tests/attention/test_flash_attn_default_custom_op.py	New CUDA regression tests for default-path custom op behavior (inference, training, opcheck; FA2-only backward tests).
fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py	New CUDA regression tests for masked/varlen custom ops (inference, FA2 backward, opcheck incl. autograd registration).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Mister-Raggs]

Two threads from the automated review worth addressing in narrative, since per-line replies don't capture the full reasoning. (The other threads — mark_non_differentiable(lse) and reusing indices across unpad_input calls in backward — are addressed in commit c58b736.)

Version-dependent op schema

The same custom-op name registers different schemas across FA versions — FA2 returns (out, lse), FA3 returns out. Only one branch ever imports per process, so within any given process the op has a single stable schema and there's no within-process ambiguity. The brittleness is real only for exported graphs (torch.export / AOTInductor artifacts) that cross FA-version boundaries — an artifact cooked on FA2 with (out, lse) wouldn't load on an FA3 box where the op declares out.

Three options weighed:

1. Versioned op names — fastvideo::_flash_attn_default_fa2_forward, …_fa3_forward. Cleanest separation, more code, more churn.
2. Unified schema with FA3 returning a placeholder lse — keeps one schema, but forces FA3 to expose an lse it doesn't use (FA3's path is still the autograd carve-out from [perf]: register FA2/FA3 default flash_attn_func as a torch.library custom op #1373; nothing consumes lse there).
3. Accept the version-dependent schema — what's checked in. Documented in the file-level comment; the only brittleness is for graph-export use cases that aren't current FastVideo workflows.

Picked (3). If the FA3 leg lands a real backward later (gated on a Hopper box for grad-check), it'll converge on (out, lse) and the schema becomes uniform across FA2/FA3 naturally. Happy to refactor to (1) up front if cleaner separation is preferred — let me know.

Inference-parity test asserts atol=0, rtol=0

This is bit-identical by construction, not aspirationally. The custom-op forward calls the same flash_attn_func as the reference, only toggling return_attn_probs=True to extract softmax_lse for backward. FA's forward kernel computes softmax_lse as a side-product of the online softmax regardless — return_attn_probs controls whether it's returned, not whether it's computed. Same kernel, same inputs, same dtype → identical out.

Validated on A100 + FA2 v2.8.1 across fp16/bf16 × causal {True, False} — 16 cases on the default path, 4 on each masked/varlen op, all pass with atol=0, rtol=0.

The training-backward parity tests do allow numerical tolerance (atol = rtol = 6e-3 for fp16, 2e-2 for bf16) — that's where independent dq/dk/dv allocations + ordering differences between the registered backward and the original autograd.Function legitimately introduce small drift. Forward bit-equality is the stronger regression alarm to keep tight; if it ever loses bits, the wrapping changed something it wasn't meant to and we want to know.

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[Mister-Raggs]

Force-pushed 3626db1b — rebased onto main now that #1373 has merged, and addressed the remaining S2 finding from Gob's post-merge review on #1373.

S2 (file placement): Moved the entire if fa_version == "2": ... elif "3": ... elif "4": block (probe + custom_op + register_fake + setup_context + register_autograd + dispatcher) 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; everything from FlashAttentionBackend onwards is unchanged. Verified single registration of fastvideo::_flash_attn_default_forward after the move (grep clean — no duplicate-registration risk). Updated test_flash_attn_default_custom_op.py to import from where the symbols now live.

Rebase conflict: tiny — took main's version (long carve-out comment + yapf-clean single-line if) for the FA3 carve-out block; no semantic change.

Commits on top of clean main:

• ce044a77 FA2 default real-backward
• aad28dd0 FA2 masked/varlen real-backward
• 6044df30 mark_non_differentiable(lse) + reuse indices in backward (the prior gemini/Copilot suggestions)
• 3626db1b move FA default custom-op registration to attention/utils/

Re-validation on GPU is pending. The 28-unit-test suite was last green on A100 + FA2 v2.8.1 at the pre-rebase tip; 3626db1b is a pure file-move + import-path change with byte-identical kernel-level logic, so I'd expect a re-run to stay green.

[Mister-Raggs]

Updated the PR body with the inference perf A/B numbers (per the perf request). tl;dr:

• Compile median, pre-[perf]: register FA2/FA3 default flash_attn_func as a torch.library custom op #1373 → this PR: −3.0% (50.83 → 49.30 s on Wan-1.3B / RTX A5000); warmup −13.5%.
• [perf]: register a real backward for FA2 default + masked/varlen custom ops (training-under-compile) #1388 marginal over [perf]: register FA2/FA3 default flash_attn_func as a torch.library custom op #1373 on Wan: ~1% (sub-noise) — expected since Wan doesn't take the masked/varlen branch this PR wraps; that benefit fires only on HunyuanVideo-1.5, which is blocked by [bug]: HunyuanVideo-1.5 generation crashes at text_encoding.py:221 — 'list' object has no attribute 'strip' #1387.
• Eager flat (~59 s ±0.2 s across all three conditions).
• SSIM = 1.0 across conditions in both modes.

Full numbers + interpretation in the body. The wrapping does what it claims structurally — graph break removed, compile region shrinks (the warmup drop is the cleanest signal of that), output bit-identical. The per-iter delta is small because Wan inference is already compile-squeezed (per Kuan's profile).

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[alexzms]

Did an independent pass on this — ran the two new test files on an H200. Forcing FA2 (blocking the FA3/cute probes so fa_version == "2"), I get 28/28 passing, so the core register_autograd work checks out end-to-end. Two things I'd like addressed before approving 👇

1. register_fake for the varlen-qk op uses the query's head_dim instead of the value's

In fastvideo/attention/utils/flash_attn_no_pad.py, _flash_attn_varlen_qk_no_pad_forward_fake does:

out = query.new_empty(query.shape)

but the real forward's out takes its head_dim from value (the kernel returns out_unpad with d_v, and out_padded is [b, sq, h, d_v]). The default-path fake gets this right — q.new_empty(b, sq, hq, v.shape[-1]) — so the two are inconsistent. With d_q == d_v (all current models) it's harmless, but if d_v != d_q the fake reports the wrong output shape, which under torch.compile means dynamo traces a wrong shape for this op (downstream shape mismatch / guard failure) — exactly the case register_fake exists to cover. Suggested fix:

b, sq, h, _ = query.shape
out = query.new_empty(b, sq, h, value.shape[-1])   # don't `del value` before reading its shape

2. Inconsistent FA2 gating in the new tests — they fail on Hopper + recent flash_attn_3

The _FA_VARLEN_VERSION != "2" skip is only applied to the autograd tests; test_varlen_qk_inference_matches_original and test_varlen_qk_forward_opcheck aren't gated. On an H200 with flash_attn_3==3.0.0 installed (probe resolves to FA3), those two run instead of skipping and fail:

TypeError: flash_attn_varlen_func() got an unexpected keyword argument 'dropout_p'
  fastvideo/attention/utils/flash_attn_no_pad.py:166

To be clear: the root cause is pre-existing on main — the FA3 varlen path passes dropout_p, which recent flash_attn_3 dropped — not introduced here. But the new ungated tests surface it, so any Hopper box on a current FA3 will go red (CI stayed green presumably because it has no / an older FA3). Could you gate the inference + opcheck tests on FA2 the same way the autograd ones are, so they skip cleanly off-FA2? The underlying FA3 dropout_p incompatibility is worth a separate fix, but it's out of scope for this PR.

Everything else looks good — nice work on the lse pad/unpad round-trip and the non-diff handling.

[Mister-Raggs]

Thanks for the careful pass! Both addressed in b1846a1:

1. Fixed — the varlen-qk register_fake now builds out with value's head_dim (query.new_empty(b, sq, h, value.shape[-1])) in both the FA2 and FA3 fakes, and no longer dels value before reading its shape.
2. Gated test_varlen_qk_inference_matches_original and test_varlen_qk_forward_opcheck on _fa2_only, matching the autograd tests, so they skip cleanly off FA2. Agreed the FA3 varlen dropout_p incompatibility is a separate, pre-existing fix — out of scope here.

Also typed the two new _repad backward closures (-> Tensor) while in here so pre-commit mypy is clean on the new code.

[alexzms]

LGTM. Thanks for iterative contribution!

[alexzms]

/merge

[alexzms]

Both fixes verified on H200 — the varlen-qk fake now uses value's head_dim (and the FA3 fake too), and the inference/opcheck tests gate on FA2 cleanly. Re-ran the suite: default probe (FA3) now 0 failed / 13 passed / 15 skipped, and forced-FA2 still 28/28. Thanks for the quick turnaround! LGTM.

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease

[Mister-Raggs]

Rebased onto current main to clear the merge conflicts. Most of the 6 commits replayed as-is; the one thing that's genuinely new since your approval is the FA4 backend selection.

Main's #1540 made FA4 (flash_attn.cute) an explicit opt-in via FASTVIDEO_FA4=1 and deleted the auto-detect fallback. This branch had moved that resolver into attention/utils/flash_attn_default.py and flash_attn_no_pad.py, so I folded #1540's opt-in gating into both — they now match main's behavior (FA4 only when FASTVIDEO_FA4=1, else FA3→FA2) instead of auto-selecting FA4 when installed. Main's _forward_impl addition from #1447 is preserved.

@alexzms since that resolver is new vs the version you approved, worth a quick re-look at flash_attn_default.py before re-merging. yapf/ruff/codespell are clean locally on the changed files.