[Metal] Add Flash Attention VJP for training#2995

Closed

Brooooooklyn wants to merge 1 commit intoml-explore:mainfrom

mlx-node:flash-attn

Commits on Mar 5, 2026

feat(attention): fused Flash Attention VJP kernels for Metal with smart dispatch

Brooooooklyn
and
claude
committed