feat(attention): complete Flash Attention VJP implementation

[Copilot]

Proposed changes

Implements Flash Attention VJP (Vector-Jacobian Product) backward pass for scaled dot-product attention, enabling memory-efficient training with automatic differentiation. Includes vector VJP for short sequences (≤8) and STEEL VJP for longer sequences, with GQA support, causal masking, and numeric stability.

Core implementation:

• Two-kernel STEEL VJP approach (dQ and dKV kernels) with K-row ownership model to eliminate atomic operations
• Vector VJP kernels with float32 accumulator support for half/bfloat16 dtypes
• Cached logsumexp mechanism to enable VJP access without materializing attention matrices
• Forward kernels updated to output logsumexp during training; training mode fallback removed

Correctness fixes:

• Fixed CMakeLists.txt variable ordering for STEEL_ATTN_HEADERS
• Fixed vector VJP shared memory overflow (BD vs D confusion in indexing)
• Fixed GQA race condition in dKV dispatch (use num_kv_heads not H)
• Fixed K/V buffer sharing (reload K after V load)
• Fixed K transpose loading strides for dQ computation
• Zero-initialize dK/dV outputs before vector VJP

Numeric stability:

• Added exp2 clamping in dQ kernel
• Added isfinite() check on P_T_val in dKV kernel
• Added float32 accumulator kernel for half/bfloat16
• Added mask/sinks fallback guard for STEEL VJP

Memory safety:

• Added static_assert for O_smem buffer reuse safety
• Added simdgroup_barrier for BD==128 after V loads
• Added alignas for O_smem alignment
• Removed D=256 support (exceeds 32KB threadgroup limit)

CUDA integration:

• Integrated cuDNN SDPA backward graph with caching
• Wired into ScaledDotProductAttentionVJP::eval_gpu
• Forward/backward use stats (LSE) when requested

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

---

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.