fix(attention): guard fused backend dispatch on device and dtype

[wayrise]

Follow-up to #7 (thanks @rakhimovv). That PR fixed flash_attention() in the Wan video DiT; this carries the same fix across the remaining dispatch sites and adds the regression coverage.

Rebased onto main after #7 merged, so the diff here no longer touches that guard.

The problem

FA2/FA3/sage are CUDA half-precision kernels. They raise rather than degrade, so any dispatch site that selects one from import availability alone breaks the moment flash-attn is installed — for every fp32 or CPU input, including the CPU tensors this suite runs on.

flash-attn is not a declared dependency and CI installs CPU-only torch, so CI never sees this. It does, however, break the make all gate CONTRIBUTING.md asks contributors to run, on essentially every training box.

input	before
cuda / bf16	OK
cuda / fp32	RuntimeError: FlashAttention only support fp16 and bf16 data type
cpu / fp32	NotImplementedError: Could not run 'flash_attn::_flash_attn_forward' ... 'CPU' backend

What changed

• action_backbone/components.py — the existing guards checked device but not dtype, so cuda/fp32 still raised here. This is the site fix(video_backbone/wan): fall back to SDPA when flash-attn kernels cannot run #7's description cited as already correct; it was only half-guarded. xformers is deliberately left device-only: memory_efficient_attention does support fp32, and guarding it on half precision would downgrade it needlessly.
• video_backbone/wan/shared/core/attention/attention.py — ATTENTION_IMPLEMENTATION was resolved once at import; resolve_implementation() now narrows it per call. fix(video_backbone/wan): fall back to SDPA when flash-attn kernels cannot run #7 left this one to maintainer judgement.
• deploy/server.py — reported the DiT backend from the availability flags alone, so it could log flash_attention_2 for a run that correctly used SDPA. It now asks dit.py, through a new fused_backend_name() helper, and qualifies the answer:

  Video DiT          : flash_attention_2 (CUDA fp16/bf16 only; torch_sdpa otherwise)

No configuration that works today changes behaviour: any call that currently reaches a fused kernel without raising is necessarily CUDA + half, which is exactly what the guards admit.

Tests

tests/test_attention_backend_dispatch.py fakes the backends, so it covers the "flash-attn is installed" configuration without flash-attn installed — the CI configuration, where this class of bug otherwise hides. It also covers the dit.py path from #7, which merged without a regression test, so that guard is now protected too.

Twelve cases. Two of them pin behaviour that was already correct (the fused fast path on cuda/bf16, and the sage CPU fallback) so a future guard can't over-correct. Three need a GPU and skip on CI.

Checked out against origin/main sources, 10 of the 12 fail.

Verification

Python 3.10, torch 2.7.1+cu128, 8xH200. Full suite run twice - once with the official flash_attn 2.8.3 cu12torch2.7cxx11abiTRUE-cp310 wheel present, once without. Identical both ways:

with flash-attn   : 1983 passed, 4 skipped, 1 failed
without flash-attn: 1983 passed, 4 skipped, 1 failed

That one failure is pre-existing and unrelated to this PR: test_cosmos_predict25_joint_cross_attn.py::test_joint_cross_attn_forward_on_real_cosmos fails identically with openwam/ checked out at origin/main. It overrides model/video_backbone=cosmos_predict25, but the repo ships configs/model/video_backbone/cosmos_predict25_2b.yaml, so Hydra raises MissingConfigException before any model code runs. It surfaces only where the Cosmos-Predict2.5-2B checkpoint is mounted, which is why CI and most boxes skip it. Left alone here - happy to open a separate issue.

Spot checks on the fast path: cuda/bf16 still reaches the fused kernel (verified with a spy on flash_attn_func), and its output is bit-identical to SDPA against an fp64 reference. cuda/fp32 backward, which died in the kernel before, now produces gradients.

ruff check, ruff format --check and compileall are clean.

🤖 Generated with Claude Code

[KraHsu]

Looks good to me — this closes the gaps #7 left open, and the evidence below matches the PR description.

What I checked

• All dispatch sites are covered. Grepping openwam/ for flash_attn_func( / flash3_fn( / sageattn( / memory_efficient_attention( gives exactly the sites this PR touches: wan/models/dit.py, the four wrappers in action_backbone/components.py, and the four backend functions in wan/shared/core/attention/attention.py. Cosmos backbones go through torch_sdpa directly. Nothing is left unguarded.
• The components.py gap is real. With a real flash_attn (2.7.4.post1 cp312/cu12.8/torch2.7 wheel, torch 2.7.1+cu128, 8×H20): on main (which already has #7) the ActionDiT fn resolves to _flash2 and cuda/fp32 still raises RuntimeError: FlashAttention only support fp16 and bf16 data type; on this branch the same call returns (1, 4, 8, 64). dit.flash_attention on cuda/fp32 runs forward + backward on both (that part is #7).
• Tests. tests/test_attention_backend_dispatch.py against main (post-#7): 8/12 fail, all in the components.py / shared-core / diagnostics groups — i.e. 10/12 against pre-#7 sources as you state, with the two dit.py cases now covered upstream. On this branch: 12/12 pass with the three GPU cases actually running (not skipped). make test without flash-attn: 1928 passed; with the real wheel present: 1928 passed as well.
• Diagnostics. With flash-attn installed the deploy report now reads Video DiT : flash_attention_2 (CUDA fp16/bf16 only; torch_sdpa otherwise) and the same for the Wan shared core. ruff check / ruff format --check clean.

Two small things (non-blocking)

1. Rebase now that #7 is merged. main is at fe0b01e (#7 in), so the dit.py guard hunk here is a no-op and the "Note on #7" paragraph is moot; a rebase leaves only fused_backend_name() in that file and keeps git blame honest about where the guard came from.
2. The ActionDiT report line is the one still unqualified. _log_attention_backends now prints ActionDiT : _flash2 while the other two lines carry the (CUDA fp16/bf16 only; torch_sdpa otherwise) suffix — and after this PR _flash2 falls back exactly the same way. Worth the same suffix (xformers → CUDA only), or the line reads as if the ActionDiT were the odd one out.

Nit, take it or leave it: the predicate now lives in three places (fused_ok inline in dit.py, _fused_kernel_usable in components.py, inline again in resolve_implementation). Fine to leave as-is to keep dit.py byte-identical to #7 for now; a single shared helper could be a follow-up once the rebase lands.

[KraHsu]

Re-approving after the force-push (e3a6e2b → b9f0b1f).

Checked what changed: the branch is rebased onto current main (fe0b01e, #7 in), so the dit.py guard hunk is gone and only fused_backend_name() remains there; the components.py, attention.py, server.py and test hunks are byte-identical to the version I reviewed, and the PR description drops the now-moot "Note on #7". Re-ran on the new head (8×H20, torch 2.7.1+cu128): ruff check / ruff format --check clean, tests/test_attention_backend_dispatch.py 12/12 with the GPU cases executing, make test 1928 passed; CI green.

The ActionDiT report line is still unqualified (ActionDiT : _flash2 vs the suffixed Video DiT / shared-core lines) — fine as a follow-up.