Differentiable flash attention matching optiq's training kernel. q,k,v are
[B,Hq,Tq,D]/[B,Hkv,Tkv,D]; bf16 is cast to fp16 (kernel dtype) and the
output cast back. Only causal/full attention (sliding-window masks fall
back to ops.sdpa at the call site, mirroring optiq's router).
Differentiable flash attention matching optiq's training kernel. q,k,v are [B,Hq,Tq,D]/[B,Hkv,Tkv,D]; bf16 is cast to fp16 (kernel dtype) and the output cast back. Only causal/full attention (sliding-window masks fall back to ops.sdpa at the call site, mirroring optiq's router).