cuda: fused flash-attention kernel (fwd + flash-style bwd)

csrc/ops/flash_attention.cu: a single fused fwd kernel (one block per
query row, streams KV in tiles of 32, online softmax — running max/sum
+ rescaled V accumulator, causal mask inlined, never materializes the
[bh,S,S] scores) writing out[bh,S,hd] + the per-row logsumexp L (O(N),
saved for backward). flash-style bwd: recompute scores from Q/K/V + L,
collapse the softmax Jacobian with D[i]=ΣdO·O, dQ owned per row, dK/dV
atomicAdd across rows. Tensor::flash_attention / flash_attention_backward
wrap them (bf16 upcasts Q/K/V→f32 for the kernel, same fp32-softmax
policy as composed).

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

crates/xtrain-cuda/src/ffi.rs

@@ -243,6 +243,59 @@ unsafe extern "C" {
     );
 }
 
+// Fused flash-attention (csrc/ops/flash_attention.cu, Phase T14). A SINGLE kernel
+// each for forward/backward that streams over KV tiles with an online softmax and
+// NEVER materializes the [bh,S,S] score matrix. Q/K/V/out are [bh,S,hd] row-major
+// F32; the forward saves only the per-row logsumexp `l` ([bh*S], O(N)) for backward.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    // Forward: o[bh,S,hd] = softmax(causal(Q·Kᵀ·scale))·V, online over KV tiles.
+    // Also writes l[bh*S] = per-row logsumexp (saved for backward, not the scores).
+    #[allow(clippy::too_many_arguments)]
+    pub fn launch_flash_attention_fwd_f32(
+        q: *const f32,
+        k: *const f32,
+        v: *const f32,
+        o: *mut f32,
+        l: *mut f32,
+        bh: i32,
+        seq: i32,
+        hd: i32,
+        scale: f32,
+        s: CudaStream,
+    );
+    // Per-row D[i]=Σ_d dO[i,d]·O[i,d] over `rows`=bh*S rows of width `hd`. Must run
+    // before the backward kernel (which takes the precomputed D, not O).
+    pub fn launch_flash_attention_rowdot_f32(
+        d_o: *const f32,
+        o: *const f32,
+        d_d: *mut f32,
+        rows: i32,
+        hd: i32,
+        s: CudaStream,
+    );
+    // Backward: recomputes scores from Q/K/V + saved logsumexp `l` (NO cached probs)
+    // and the precomputed `d_d` (= D), produces dq/dk/dv. dq/dk/dv must be PRE-ZEROED
+    // (dk/dv are accumulated across query rows via atomicAdd).
+    #[allow(clippy::too_many_arguments)]
+    pub fn launch_flash_attention_bwd_f32(
+        q: *const f32,
+        k: *const f32,
+        v: *const f32,
+        d_o: *const f32,
+        l: *const f32,
+        d_d: *mut f32,
+        dq: *mut f32,
+        dk: *mut f32,
+        dv: *mut f32,
+        bh: i32,
+        seq: i32,
+        hd: i32,
+        scale: f32,
+        s: CudaStream,
+    );
+}
+
 // GPU-side optimizer kernels (csrc/ops/optim.cu): AdamW step (m/v on device) and
 // the global grad-norm reduction + in-place rescale (Phase T7).
 #[cfg(not(no_cuda))]