Materialized differentiable attention (matmul → mask → softmax → matmul),
GQA-aware. A correct-vjp fallback for shapes/masks the flash kernel doesn't
take. (ops.sdpa's dK vjp is CORRECT — FD-verified, the training default;
the old "broken dK" note here predated the flash-kernel root-cause: the
bugs were flash's own dK transpose + dQ barrier, both fixed.) O(L²) —
pair with gradient checkpointing.
Materialized differentiable attention (matmul → mask → softmax → matmul), GQA-aware. A correct-vjp fallback for shapes/masks the flash kernel doesn't take. (ops.sdpa's dK vjp is CORRECT — FD-verified, the training default; the old "broken dK" note here predated the flash-kernel root-cause: the bugs were flash's own dK transpose + dQ barrier, both fixed.) O(L²) — pair with gradient checkpointing.