test+bins: flash grad-check, flash==composed, PyTorch parity, --flash flag

autograd: flash_attention_batched_bwd (dQ/dK/dV finite-diff, seq>tile)
+ flash_matches_composed_fwd. model/tests/flash.rs: flash==composed
on-vs-off (logits/loss/every param grad), fp32 + bf16. parity_dump:
XTRAIN_PARITY_FLASH dumps the flash path for the same parity.py oracle
(PyTorch SDPA parity at B>1). train + train_ddp get the --flash flag.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-17 23:10:39 +08:00
parent 0e20821633
commit 5f3b81ac96
5 changed files with 280 additions and 1 deletions

View File

@@ -116,6 +116,9 @@ fn main() {
// exact grads, lower peak activation memory (lets dim1024 batch32 fit). Opt-in;
// default off stores every activation (unchanged numerics).
let recompute = args.iter().any(|a| a == "--recompute");
// Fused flash-attention (Phase T14): single fused SDPA kernel, online softmax,
// no materialized [bh,S,S] scores. Opt-in; default off keeps the composed path.
let flash = args.iter().any(|a| a == "--flash");
let ckpt: PathBuf = PathBuf::from(
args.iter()
.position(|a| a == "--ckpt")
@@ -183,6 +186,10 @@ fn main() {
model = model.with_recompute(true);
println!("activation recompute: ON (per-block gradient checkpointing)");
}
if flash {
model = model.with_flash(true);
println!("flash-attention: ON (fused SDPA kernel, no materialized scores)");
}
// Eval-only mode: load a checkpoint and score it on the held-out val set, then
// exit. Used to put an EXISTING model (e.g. v0) and a new one on the same