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>

crates/xtrain-autodiff/tests/autograd.rs

@@ -625,6 +625,104 @@ fn attention_batched_bwd() {
     );
 }
 
+// ---- fused FLASH causal attention (the T14 op) ----
+// Same structure as attention_batched_bwd, but exercises ops::flash_attention.
+// q,k,v: [bh, seq, hd]. Grad-check dq/dk/dv against finite-diff of L=sum(W∘out).
+// seq=40 > FA_TILE=32 so the online-softmax tile-rescale path is exercised (not
+// just a single KV tile).
+#[test]
+fn flash_attention_batched_bwd() {
+    require_gpu();
+    let (bh, seq, hd) = (2, 40, 16);
+    let n = bh * seq * hd;
+    let scale = 1.0 / (hd as f32).sqrt();
+    let q_h = fill(n, 241);
+    let k_h = fill(n, 242);
+    let v_h = fill(n, 243);
+    let w = fill(n, 244);
+
+    let q = Var::leaf(cuda(&q_h, &[bh, seq, hd]));
+    let k = Var::leaf(cuda(&k_h, &[bh, seq, hd]));
+    let v = Var::leaf(cuda(&v_h, &[bh, seq, hd]));
+    let out = ops::flash_attention(&q, &k, &v, scale);
+    scalar_loss(&out, &w).backward();
+
+    let dq = q.grad().unwrap().to_device(Device::Cpu);
+    let dk = k.grad().unwrap().to_device(Device::Cpu);
+    let dv = v.grad().unwrap().to_device(Device::Cpu);
+
+    let fwd = move |qh: &[f32], kh: &[f32], vh: &[f32]| -> f32 {
+        let qv = cuda(qh, &[bh, seq, hd]);
+        let kv = cuda(kh, &[bh, seq, hd]);
+        let vv = cuda(vh, &[bh, seq, hd]);
+        let (o, _) = qv.flash_attention(&kv, &vv, scale);
+        weighted_sum(&o, &w)
+    };
+    let (kf, vf, ff) = (k_h.clone(), v_h.clone(), fwd.clone());
+    let lq = move |x: &[f32], _s: &[usize]| ff(x, &kf, &vf);
+    report(
+        "flash dQ",
+        &grad_check(
+            &q_h,
+            &[bh, seq, hd],
+            &lq,
+            dq.as_slice::<f32>(),
+            cfg_nonlinear(),
+        ),
+    );
+    let (qf, vf, ff) = (q_h.clone(), v_h.clone(), fwd.clone());
+    let lk = move |x: &[f32], _s: &[usize]| ff(&qf, x, &vf);
+    report(
+        "flash dK",
+        &grad_check(
+            &k_h,
+            &[bh, seq, hd],
+            &lk,
+            dk.as_slice::<f32>(),
+            cfg_nonlinear(),
+        ),
+    );
+    let (qf, kf, ff) = (q_h.clone(), k_h.clone(), fwd.clone());
+    let lv = move |x: &[f32], _s: &[usize]| ff(&qf, &kf, x);
+    report(
+        "flash dV",
+        &grad_check(
+            &v_h,
+            &[bh, seq, hd],
+            &lv,
+            dv.as_slice::<f32>(),
+            cfg_linear(),
+        ),
+    );
+}
+
+// flash forward must equal the composed attention forward (same SDPA math).
+#[test]
+fn flash_matches_composed_fwd() {
+    require_gpu();
+    let (bh, seq, hd) = (2, 40, 16);
+    let n = bh * seq * hd;
+    let scale = 1.0 / (hd as f32).sqrt();
+    let q = cuda(&fill(n, 341), &[bh, seq, hd]);
+    let k = cuda(&fill(n, 342), &[bh, seq, hd]);
+    let v = cuda(&fill(n, 343), &[bh, seq, hd]);
+    let (oc, _) = q.attention(&k, &v, scale);
+    let (of, _) = q.flash_attention(&k, &v, scale);
+    let oc = oc.to_device(Device::Cpu);
+    let of = of.to_device(Device::Cpu);
+    let max_rel = oc
+        .as_slice::<f32>()
+        .iter()
+        .zip(of.as_slice::<f32>())
+        .map(|(c, f)| (c - f).abs() / (c.abs() + 1e-6))
+        .fold(0.0f32, f32::max);
+    println!("flash-vs-composed fwd max rel: {max_rel:.3e}");
+    assert!(
+        max_rel < 1e-4,
+        "flash fwd diverges from composed: {max_rel:.3e}"
+    );
+}
+
 // --- test helpers ---
 
 // Scalar loss node L = sum(W ∘ out): wraps a fixed-weight Var and reduces. We