[perf] Add FA4 CuTe backward support for VSA-256 - #1639
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Welcome to FastVideo! Thanks for your first pull request.
How our CI works:
PRs run a two-tier CI system:
- Pre-commit — formatting (yapf), linting (ruff), type checking (mypy). Runs immediately on every PR.
- Fastcheck — core GPU tests (encoders, VAEs, transformers, kernels, unit tests). Runs automatically via Buildkite on relevant file changes (~10-15 min).
- Full Suite — integration tests, training pipelines, SSIM regression. Runs only when a reviewer adds the
readylabel.
Before your PR is reviewed:
-
pre-commit run --all-filespasses locally - You've added or updated tests for your changes
- The PR description explains what and why
If pre-commit fails, a bot comment will explain how to fix it. Fastcheck and Full Suite results appear in the Checks section below.
Useful links:
|
Merge Protections🟠 1 of 1 protections blocking · waiting on 🤖 CI
🟠 PR merge requirementsWaiting for
Waiting checks:
|
ita9naiwa
changed the title
|
@alexzms @SolitaryThinker, since this is a follow-up to #1354, would you mind taking a look when you have a chance? |
|
The CI failure appears to be unrelated to this PR. It comes from the existing |
|
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 |
…SL requirement The VSA-256 FA4 CuTe install instructions pointed at the wrong revision and steered users into a broken CuTe DSL: - The README pinned `940cd9680f...` and claimed it was the revision the repo-root `pyproject.toml` uses. It has been `82d6441eec...` since the pin was bumped, so anyone following the README installed a different FA4 than CI and the Docker image run. - The README told users to `pip install "nvidia-cutlass-dsl>=4.5.0"` first. The pinned revision requires `nvidia-cutlass-dsl==4.6.0.dev0`; installing `flash-attn-4` under its distribution name resolves that correctly on its own, so the manual floor only creates a chance to end up on 4.5.x. - A CuTe DSL mismatch surfaces only when the kernel JIT-compiles, as errors from inside CuTe (`fmax() missing 1 required positional argument`, `cute.core has no attribute 'ThrMma'`, `cannot import name 'alloc_reserved_mbarrier'`), none of which point back at the install. The README now maps each symptom to its cause. Verified on Blackwell (sm_100, GB200). Running the documented install end to end resolves nvidia-cutlass-dsl 4.6.0.dev0 and quack-kernels 0.5.3, and against that environment all 67 VSA kernel tests pass, including `tests/test_vsa256_forward*.py`, `tests/test_vsa256_triton.py` and `tests/test_vsa256_backward.py`. Substituting nvidia-cutlass-dsl 4.5.2 fails the same forward tests at JIT time, which is how each row of the symptom table was observed. Docs only; no runtime code changes.
Resolves the yapf-formatting conflicts introduced by hao-ai-lab#1700 / #6eb95693 by keeping this branch's logic and adopting main's formatting. No behaviour change.
Follow-ups found while reproducing this PR on GB200 (sm_100). video_sparse_attn_h3.py composed the compression branch onto the attention output with an in-place addcmul_. On the CuTe backend that output *is* the tensor FA4's autograd node saved for its backward, so the moment this path has a backward at all, training dies with "one of the variables needed for gradient computation has been modified by an inplace operation ... output 0 of FlashAttnFuncBackward". Same defect the BSHD branch of video_sparse_attn_bshd already fixed here; H3 landed after this branch's base and so never got it. Now out-of-place. The KV-owned backward metadata was built unconditionally. It is a pair of dense [B, H, kv_blocks, q_blocks] int32 index tensors that FA4 keeps alive on its ctx until backward runs, and inference never reads it. Gated on requires_grad: at Wan-14B 720p shape (B=1, H=40, D=128, S=92160, topk 45/360) a no_grad forward drops from 17.23 ms / 1998 MiB transient to 15.92 ms / 1918 MiB. Training keeps the metadata and is unchanged. The aux (lse) return was transposed to [B, S, H]. FA4 hands back [B, H, S] already (interface.py builds lse_shape as (batch, num_head, seqlen_q) when qv is None), which is what the Triton path's aux contract is, so the transpose made the two backends disagree and cost a 14 MiB fp32 copy per call at the shape above. Dropped, and detach happens first now. Tests - gated compression branch (compress_attn_weight is not None) backward: the branch Wan and H3 actually run, and the one the in-place bug lives in. - partially filled Q tiles, and q_len != kv_len; forward had cross coverage, backward had none. - inference forward is bitwise identical to the training forward, pinning the requires_grad gate. - aux is [B, H, S] on both entrypoints. - fastvideo/tests/attention/test_vsa_h3_backward.py covers the H3 backend on both backends and cross-checks CuTe gradients against Triton. - Gradient tolerances tightened from avg_abs 2e-2 / max_rel 0.5 to 1e-3 / 0.25; measured error across every case above is <= 1.2e-4 and <= 0.11. test_vsa_varlen.py seeds the RNG now. It draws every tensor and every variable block size from the global RNG and then asserts max_rel < 0.05, so adding any test that runs before it shifts its inputs and it fails on unlucky data. That is the "unrelated" CI failure reported on this PR: it reproduces in a full-suite run and passes in isolation. Verified on GB200 (sm_100), FA4 82d6441 + nvidia-cutlass-dsl 4.6.0.dev0 + quack-kernels 0.5.3: fastvideo-kernel/tests/ and fastvideo/tests/attention/ give 318 passed, stable across repeated runs. The one remaining failure, test_fa4_quantize_op_fake_matches_real, fails identically on an unmodified main. Reverting either fix above turns the matching new test red. pre-commit (yapf, ruff, codespell) passes on every changed file; the mypy hook only errors on the worktree's directory name, and does so identically on an unmodified main.
|
Thank you for this — and welcome. This is a genuinely well-executed piece of work, and it reproduces cleanly: on GB200 (sm_100) I measure 3.06x / 2.76x backward over the Triton fallback at your 480p / 720p shapes, right in line with your B300 numbers. A few things I want to call out specifically, because they are the parts that are easy to get wrong and you got right. Delegating the autograd bridge to FA4's public I've pushed three commits on top rather than round-tripping them to you:
One more thing worth knowing, since you flagged it: the 318 passed locally on GB200; the one remaining failure fails identically on an unmodified main. Merging once CI is green — thanks again for taking this on. |
mergify
Bot
added
scope: attention
a1c75bd to
b565da4
Compare
alexzms
left a comment
alexzms
left a comment
There was a problem hiding this comment.
Verified on GB200 (sm_100), the only arch that actually executes this path: 324 passed, CuTe-vs-Triton parity across 128 shape/head/batch/dtype combinations, and backward at 2.6-3.0x over the Triton fallback (61-66% of measured dense bf16 peak). H3 dispatch confirmed end to end with FASTVIDEO_VSA_CUTEDSL=1.
One deployment note before this goes out: the FA4_CUTE_REF bump in docker/Dockerfile means the published fastvideo-dev image keeps the old FA4 until it is rebuilt from main. VSA-128 backward needs the new pin — the old revision only reads block_size[0] on sm_90, so on sm_100 q_subtile_factor comes out as 2 and the 128 metadata is rejected. VSA-256 works on either pin, so H3 is fine on the current image; just don't expect the 128 path to work in a container until the rebuild lands.
…ling (bf16 K pre-scaling) The backward pre-scaled K by sm_scale*log2(e) in bf16 before recomputing logits, so the recompute drifted from the forward's saved LSE proportionally to |logit|; exp2 amplifies that into exponentially wrong probabilities and garbage dK/dV (plus a mis-scaled dQ) at real activation magnitudes. The forward is exact and unit-scale tests passed, which is how it survived — and the Triton backward is the default training path on every GPU except sm_90 with the compiled extension (the FA4 CuTe backward from #1639 is opt-in and sm_100-only). Fix: pass raw K, apply sm_scale*RCP_LN2 in fp32 after tl.dot in both _attn_bwd_dkdv and _attn_bwd_dq, and scale dq by sm_scale instead of LN2, matching the forward's rounding exactly. Adds a seeded regression test sweeping input scale against an fp32 masked-dense SDPA reference: on GB200 the unfixed kernel passes at unit scale but fails at scale=16 (dq rel L2 err 5.9e-1); the fixed kernel passes all scales (rel err <= 6.9e-3, grad-norm ratios 0.9995-1.0001). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
3 participants
Purpose
Enable FA4 backward support for VSA-256 on Blackwell. Related issue/PR: #1354
Changes
_flash_attn_bwdinto the existing_flash_attn_fwdpath.Test Plan
Test Results
Benchmark Summary
Compared with the Triton fallback, FA4 CuTe provides:
Checklist
pre-commit run --all-filesand fixed all issues