[samsja]

Summary

Rework of the FA4 custom op proposed in #3676 (stacked on its head commit; not measured in isolation from it). Same goal: FA4's flash_attn_varlen_func is a Python autograd.Function that calls torch._guards.active_fake_mode(), which Dynamo refuses to trace, so with activation checkpointing on, the graph break drops the whole block to eager.

What differs from #3676:

• The op schema mirrors FA2/FA3's varlen convention instead of a minimal non-CP-only signature: flash_attn_4_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal, window_size_left, window_size_right) -> (out, lse), and likewise for the backward op. causal and both window edges are schema args, not hardcoded.
• A thin adapter exposes the op with FA2/FA3's exact call convention (positional cu/max, causal=/window_size=(-1, -1) kwargs, -1 -> None, returns out only) and _funcs[4] dispatches to it.
• _compute_attention collapses to a single version-free path; the FA4 branch (and the qv-misalignment comment) are gone.
• The raw flash_attn_4_varlen_func re-export stays because gpt_oss and afmoe import it from attn.py.

Only the non-CP path changes, as in #3676: ring and Ulysses CP replace _compute_attention and keep their own FA4 calls. The wide schema is the shape Ulysses' local FA would need if it ever routes through the op. prime_rl_attn is already on the selective AC save list via #3676.

Validation

• ruff check and ruff format pass.
• Wiring validated locally with stubbed FA4 kernels (the cu13 FA4 wheel cannot load on this machine's CUDA 12.9 driver): op schema registration, fake impls (out (t, h, d), lse (h, t) fp32), autograd plumbing with exact forward and matching dq/dk/dv, sliding-window and non-causal cases, _funcs dispatch through _compute_attention, and torch.compile(fullgraph=True) traces forward+backward while a raw autograd.Function calling active_fake_mode() graph-breaks with Unsupported.
• Still pending on a cu13-driver GPU box: numeric parity vs the real FA4 kernels (the fix(trainer): run FA4 through a custom op so compile can trace attention #3676 bar: out/dq exact, dk/dv within 1 bf16 ulp) and perf/VRAM numbers. No perf claims are made here.

🤖 Generated with Claude Code