autodiff+model: flash-attention op + --flash opt-in wiring

ops::flash_attention autograd node (fwd caches O(N) logsumexp instead of
O(N²) probs; bwd via Tensor::flash_attention_backward). Model gets a
use_flash bool + with_flash(bool) builder; the SDPA core in attention()
picks ops::flash_attention vs ops::attention. flash threads through
block_forward so the recompute (T13) segment also runs flash. Default
off = composed path, graph unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-autodiff/src/ops.rs

@@ -325,6 +325,32 @@ pub fn attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
     )
 }
 
+/// Fused FLASH causal scaled-dot-product attention (Phase T14). Same interface as
+/// [`attention`] (`q`,`k`,`v` each `[bh, seq, head_dim]`), but the forward is a
+/// SINGLE fused kernel with an online softmax over KV tiles — the `[bh,seq,seq]`
+/// score matrix is NEVER materialized, and backward caches only the per-row
+/// logsumexp (O(N)) instead of the whole probs (O(N²)). Mathematically the same
+/// SDPA, so it matches the composed [`attention`] within fp/bf16 tolerance.
+/// Opt-in via the model's `--flash` flag; the composed path stays the default.
+pub fn flash_attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
+    let (out, lse) = q.value().flash_attention(&k.value(), &v.value(), scale);
+    let out_cache = out.clone();
+    Var::from_op(
+        out,
+        vec![q.clone(), k.clone(), v.clone()],
+        Box::new(move |dout, parents| {
+            let q = parents[0].value();
+            let k = parents[1].value();
+            let v = parents[2].value();
+            let (dq, dk, dv) =
+                Tensor::flash_attention_backward(&q, &k, &v, &out_cache, &lse, dout, scale);
+            Var::push_grad(&parents[0], dq);
+            Var::push_grad(&parents[1], dk);
+            Var::push_grad(&parents[2], dv);
+        }),
+    )
+}
+
 /// Cross-entropy mean loss over logits `x:[rows,cols]` with one I32 target per
 /// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
 /// scaled by the upstream scalar grad.