[perf] FA3 custom_op + functional RoPE: compile-traceable DistributedAttention - #1384
[perf] FA3 custom_op + functional RoPE: compile-traceable DistributedAttention#1384rich7420 wants to merge 3 commits into
Conversation
mergify
Bot
added
type: perf
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enables torch.compile support for attention layers by refactoring in-place RoPE assignments into functional operations and removing @torch.compiler.disable decorators. It also introduces a new utility module that wraps Flash Attention 3 using torch.library.custom_op to ensure traceability. Feedback was provided regarding the Flash Attention 3 wrapper, noting that it currently lacks autograd registration and discards the softmax LSE tensor, which limits its functionality to inference-only when used with torch.compile.
| out = _raw_flash_attn_3_func(q, k, v, softmax_scale=softmax_scale, causal=causal) | ||
| if isinstance(out, tuple): | ||
| out = out[0] | ||
| return out |
There was a problem hiding this comment.
The current implementation of _flash_attn_3_forward discards the softmax_lse tensor returned by flash_attn_func and does not register an autograd implementation for the custom op. This effectively makes this wrapper inference-only when used with torch.compile. If training or gradient-based optimization is attempted with this op under torch.compile, it will fail or produce incorrect results because the backward pass is not defined for the custom op.
Consider adding a comment documenting this limitation or implementing register_autograd if training support is intended.
892081e to
2cf6935
Compare
|
Hi @rich7420 — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings aren't all human-verified; ping @SolitaryThinker if anything looks off. TL;DRThe FA3 custom op has a fake implementation and a non-mutating forward schema, but it is wired into the general FlashAttention backend without an autograd formula, so FA3 training/gradient use will fail. The PR also removes the distributed-attention compile guard while the SP all-to-all path still reaches unwrapped distributed collectives, and no compile/custom-op regression test was added. Verdict: Changes requested
Findings (formatted for upload)[S1] FA3 custom op is selected for training but has no autograd formulaWhat: Why it matters: FastVideo uses this backend from both inference and training stacks on CUDA. With requires-grad Q/K/V, a Suggested fix: Either implement FA3 Evidence: [S1] DistributedAttention compile guard was removed before SP collectives became compile-safeWhat: Why it matters: The previous disable made the distributed wrapper an explicit eager boundary. Removing it can turn that into hidden graph breaks or compile failures around NCCL collectives, so users may think distributed attention is compiled when it is not, or see runtime failures only under SP/VSA configurations. Suggested fix: Keep the compile disable around the distributed wrapper, or at least around the SP communication helper, until the collectives are wrapped in a compile-safe custom op with fake/autograd support. If the path is actually supported today, add a rank-2 Evidence: [S2] No compile/custom-op regression test was addedWhat: The PR changes only Why it matters: This PR is specifically changing Suggested fix: Add CUDA tests for FA3 wrapper forward parity against raw FA3 for fp16/bf16 and causal/non-causal, run Evidence: — Gob (@SolitaryThinker's AI reviewer). Full review (including S3 items) is archived locally. |
…9.3% e2e on H100 sp=1) PR hao-ai-lab#684 ("[Feature] Optionally enable torch compile") added @torch.compiler.disable on DistributedAttention.forward / _VSA.forward so torch.compile would not crash on the bare flash_attn_3 call + in-place RoPE slice assignment. Side effect: every attention call became a graph break, fragmenting the per-block compiled region and leaving 84% of denoising direct_copy_kernel_cuda launches in eager mode. Three changes that together unlock the compile region for sp=1: 1. fastvideo/attention/utils/flash_attn_3_compile.py (new): Register flash_attn_3 as torch.library.custom_op + register_fake — mirrors the pattern flash_attn_cute.py already uses for FA4. Inductor now treats FA3 as an opaque-but-known boundary and can fuse adjacent ops into a single triton kernel. 2. fastvideo/attention/backends/flash_attn.py: Import flash_attn_func from the new wrapper instead of bare flash_attn_interface.flash_attn_func. 3. fastvideo/attention/layer.py: - Drop @torch.compiler.disable from DistributedAttention.forward and DistributedAttention_VSA.forward. With FA3 now a custom_op, and sequence_model_parallel_all_to_all_4D already short-circuiting on world_size==1, the disable is no longer required. - Replace qkv[:batch_size * 2] = _apply_rotary_emb(qkv[:batch_size * 2], ...) with the functional form qk_rope = _apply_rotary_emb(qkv[:batch_size * 2], ...) qkv = torch.cat([qk_rope, qkv[batch_size * 2:]], dim=0) so torch.compile can trace it (same pattern LocalAttention.forward already uses). Wan T2V 1.3B inference, H100 sp=1+FA3, 30 steps, 720x1280x77: avg_generation_time_s: 127.314s -> 115.499s (-11.815s, -9.28%) denoising_stage: 117.72s -> 104.80s (-12.92s, -11.0%) individual_times: [127.30, 127.33] -> [115.40, 115.60] (198ms variance, well below the ~2% Modal H100 jitter floor) GPU compute_only_ms: 372,726ms -> 335,984ms (-36,743ms) Hardware-normalized: -8.76% (calibrated against unchanged FA3 device_kernel timing on different Modal H100 cards) nsys-ai diff confirms the expected fusion landed: - New: triton_red_fused__flash_attn_3_forward__to_copy_add_addmm_mean_mul_ pow_rsqr * 5,400 instances (self-attn FA3 fused with QK RMSNorm). - New: triton_poi_fused__to_copy_add_cat_mul_neg_slice_stack_unbind_unsqueeze_ view * 5,482 (RoPE as a single kernel). - CatArrayBatchedCopy_vectorized: 5,400 -> 0 (CFG concat folded into fusion). - elementwise_kernel: 19,455 -> 8,655; unrolled_elementwise_kernel: 11,486 -> 686. Many small eager kernels collapsed. - Eliminated: triton_red_fused__to_copy_add_addmm_mul_native_layer_norm_view_0 / _3 (10k+ instances combined), refused into wider kernels. Quality verified: VBench 5-dimension A/B on 5 fixed-seed prompts vs the pre-PR baseline, all metrics within +/-0.0013: subject_consistency: 0.9526 -> 0.9524 (-0.0002) background_consistency: 0.9558 -> 0.9572 (+0.0013) aesthetic_quality: 0.6050 -> 0.6042 (-0.0009) imaging_quality: 0.6035 -> 0.6024 (-0.0010) temporal_flickering: 0.9787 -> 0.9786 (-0.0000) Outputs are algorithm-identical; small per-pixel deltas (avg-frame SSIM 0.946) come from inductor fusion changing fp32 accumulator ordering inside the new triton kernels, amplified by 30-step iterative denoising. sp>1 path unchanged in behavior — sequence_model_parallel_all_to_all_4D short-circuits on world_size==1 only; for >1 the dist collective runs as before, just inside a (possibly broken) compile graph. Validated only at sp=1 on this branch; sp>1 should be regression-tested before relying on fusion in distributed runs.
d482434 to
f4e5b8d
Compare
The custom_op wrapper registers no autograd, so a backward through it fails. Since flash_attn.py wires this wrapper in for both inference and training, training with the FA3 backend would break (raw FA3 supports backward). Route grad-requiring inputs to the raw FA3 function and keep the no-grad inference path on the compile-traceable custom op.
f4e5b8d to
2360b99
Compare
SolitaryThinker
left a comment
SolitaryThinker
left a comment
There was a problem hiding this comment.
Hi @rich7420 — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings aren't all human-verified; ping @SolitaryThinker if anything looks off.
Summary
Reviewing at head 2360b997 (open, mergeable, fastcheck green; kernel/attention scope). The real perf win (−9.28% e2e, H100 sp=1+FA3) comes from layer.py: dropping @torch.compiler.disable and rewriting the in-place RoPE slice-assign to a functional cat. That part is sound and output-equivalent. No blockers — outputs are correct and the speedup is real. Two things to resolve before merge.
Verdict: COMMENT
No in-scope blockers. The perf win is legitimate; the concerns are the redundant new module and the unverified sp>1 path.
Major concerns
MAJOR — fastvideo/attention/utils/flash_attn_3_compile.py duplicates the on-main FA2/FA3 custom-op wrapper from #1373; the FA3 inference path now double-wraps.
flash_attn_func_compilable + torch.ops.fastvideo._flash_attn_default_forward already live on main (backends/flash_attn.py:49-89, landed in #1373 2f3ca8aa, which is an ancestor of this head). That wrapper already (a) wraps the FA2/FA3 default flash_attn_func in a torch.library.custom_op, (b) registers a fake, and (c) has the same autograd carve-out routing grad inputs to the raw func. The only live hot-path call is FlashAttentionImpl.forward:307 → flash_attn_func_compilable — not the module-level flash_attn_func this PR redirects, which flash_attn.py:37 captures as _fa_default and re-wraps. So at runtime FA3 inference becomes flash_attn_func_compilable → _flash_attn_default_forward (custom op) → _fa_default → _flash_attn_3_forward (new custom op) → raw FA3: a custom op nested inside a custom op. Dynamo treats the outer op as opaque and never traces into the inner one, so the new op's traceability/fake are dead under every current call path (grep confirms no other live importer of the backend flash_attn_func). The body's "Inductor can now graph through FA3" already landed in #1373.
Suggested fix: drop flash_attn_3_compile.py + the flash_attn.py:14-19 import change, keep only the layer.py disable removal (the actual perf source). If #1373's wrapper is somehow insufficient for FA3, please document why in the PR body — the diff doesn't show a reason.
MAJOR — sp>1 / VSA compile path is unverified after removing @torch.compiler.disable (layer.py:59,151).
With the guard gone, at sp>1 the compiled region now spans sequence_model_parallel_all_to_all_4D → DistributedAutograd.AllToAll4D.apply → dist.all_to_all_single (NCCL inside an autograd.Function). At sp=1 the collective short-circuits (base_device_communicator.py:142 returns input unchanged), so the claimed path is clean — but sp>1 isn't tested. Mitigating: compile is applied without fullgraph=True by default (composed_pipeline_base.py:115,153), so dynamo graph-breaks gracefully at the collective — a perf non-improvement, not a correctness regression. Real risk only if a user sets fullgraph=True (then sp>1 hard-errors at compile) or assumes sp>1 is fused when it silently isn't. You already flag this ("sp > 1 paths were not re-verified").
Suggested fix: run /test full on multi-GPU before merge; if sp>1 + fullgraph is meant to work, add a rank-2 torch.compile(..., fullgraph=True) smoke over DistributedAttention/_VSA. Otherwise note in the body that sp>1 falls back to a graph break.
Minor
MINOR — if the new module is kept despite the above, the fastvideo::_flash_attn_3_forward op + its register_fake have no test (#1373's test_flash_attn_default_custom_op.py covers _flash_attn_default_forward only). Removing the module (preferred) moots this; otherwise mirror that test (forward parity vs raw FA3 fp16/bf16 × causal + torch.library.opcheck). The fake itself is correct.
Test plan
Solid evidence for the claimed path: H100 sp=1+FA3 before/after (127.31→115.50 s), nsys-ai diff showing the expected fusion, and a 5-prompt VBench A/B (all dims within ±0.0013). The reported avg-frame SSIM 0.946 would pass the repo's Wan T2V gate (tests/ssim/test_wan_t2v_similarity.py:94, threshold 0.93) — worth a /test ssim to confirm in CI. Gaps: no sp>1/VSA regression and no committed test/benchmark for the new op. Recommend /test full + /test ssim before merge.
(Note: the autograd-fallback you added in 2360b997 resolves the training-breakage flagged in the earlier review on this PR — that item is fixed at this head.)
— Gob (@SolitaryThinker's AI reviewer).
|
Thanks @SolitaryThinker 's AI Gob— verified on a real H100, quick notes:
|
|
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 |
…compile-trace # Conflicts: # fastvideo/attention/layer.py
rich7420
changed the title
|
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 |
|
Rebase triage: a pure rebase onto main is no longer possible — main has since absorbed this PR's core contribution. #1373 registered the FA2/FA3 What's still novel here: the two functional-RoPE hunks in Two ways forward:
Leaving the branch untouched until you weigh in. |
2 participants
Purpose
Make
DistributedAttention.forwardcompile-traceable so inductor can fuseacross the Wan T2V block forward. −9.28 % e2e on H100 sp=1 + FA3.
Changes
attention/utils/flash_attn_3_compile.py— wrapflash_attn_3astorch.library.custom_op+register_fake, mirroring the FA4 wrapper inflash_attn_cute.py. Inductor can now graph through FA3. The custom op hasno autograd, so grad-requiring inputs (training) route to the raw FA3 func
(its own backward); inference (no-grad) takes the compile-traceable op.
attention/backends/flash_attn.py— importflash_attn_funcfrom the newwrapper instead of the bare
flash_attn_interfacesymbol.attention/layer.py— rewrite the in-placeqkv[:bs*2] = _apply_rotary_emb(qkv[:bs*2], …)to a functional cat-basedform in both
DistributedAttention.forwardandDistributedAttention_VSA.forward(matches
LocalAttention.forward), so torch.compile can trace RoPE. The@_maybe_compiler_disabledecorator from upstream is left in place.sequence_model_parallel_all_to_all_4Dalready short-circuits onworld_size == 1, so the sp=1 path is fully compileable. sp > 1 paths werenot re-verified — recommend
/test fullon multi-GPU before merge.Test Results
Measured with
FASTVIDEO_DISABLE_ATTENTION_COMPILE=0(attention traced intothe compile graph).
Perf — Wan T2V 1.3B, H100 sp=1 + FA3, 30 steps, 720×1280×77:
avg_generation_time_sdenoising_stagensys-ai diffconfirms expected fusion:triton_red_fused__flash_attn_3_forward__to_copy_add_addmm_mean_mul_pow_rsqrt× 5 400triton_poi_fused__to_copy_add_cat_mul_neg_slice_stack_unbind_…× 5 482CatArrayBatchedCopy_vectorized) 5 754 → 354 launcheselementwise_kernel19 455 → 8 655;unrolled_elementwise_kernel11 486 → 686Quality — VBench A/B on 5 fixed-seed prompts vs pre-PR baseline:
vbench.subject_consistencyvbench.background_consistencyvbench.aesthetic_qualityvbench.imaging_qualityvbench.temporal_flickeringAll 5 dimensions within ±0.0013. Visually identical on side-by-side
inspection.
Outputs are algorithm-identical; small per-pixel deltas (avg-frame SSIM
0.946) come from inductor fusion changing fp32 accumulator ordering inside
the new triton kernels, amplified by 30-step iterative denoising. Perceptual
quality (VBench, eye test) preserved.
Checklist
pre-commit run --all-filesFASTVIDEO_DISABLE_ATTENTION_COMPILE=0)_maybe_compiler_disableenv-gate/test fullbefore merge