[Metal] Add Flash Attention VJP for training#2995

Open

Brooooooklyn wants to merge 1 commit intoml-explore:mainfrom

mlx-node:flash-attn

Commits on Jan 23, 2026

feat(attention): add Flash Attention VJP for vector path (L≤8)
Brooooooklyn
committed