docs: Phase T4 — autograd engine

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

crates/xtrain-autodiff/build.rs

@@ -0,0 +1,27 @@
+use std::env;
+use std::path::Path;
+use std::process::Command;
+
+// xtrain-autodiff's `Var` tape calls GPU ops (via xtrain-tensor), so it gates
+// those call sites behind `not(no_cuda)` — the same per-crate convention as
+// xtrain-cuda / xtrain-tensor (cfg does not propagate across crates). The
+// grad_check harness itself is host-only and always compiles. This script only
+// detects nvcc and emits the cfg; it compiles no CUDA.
+fn main() {
+    println!("cargo:rustc-check-cfg=cfg(no_cuda)");
+
+    let cuda_path = env::var("CUDA_HOME")
+        .or_else(|_| env::var("CUDA_PATH"))
+        .unwrap_or_else(|_| "/usr/local/cuda".to_string());
+
+    if !nvcc_available(&cuda_path) {
+        println!("cargo:rustc-cfg=no_cuda");
+    }
+}
+
+fn nvcc_available(cuda_path: &str) -> bool {
+    if Command::new("nvcc").arg("--version").output().is_ok() {
+        return true;
+    }
+    Path::new(&format!("{cuda_path}/bin/nvcc")).exists()
+}

crates/xtrain-autodiff/src/lib.rs

@@ -13,3 +13,14 @@
 pub mod finite_diff;
 
 pub use finite_diff::{GradCheckConfig, GradCheckResult, ParamFn, grad_check};
+
+// Tape-based autograd engine + differentiable ops (Phase T4). These call GPU
+// kernels via xtrain-tensor, so they are gated behind `not(no_cuda)` (the
+// per-crate convention); the grad_check harness above stays host-only.
+#[cfg(not(no_cuda))]
+pub mod ops;
+#[cfg(not(no_cuda))]
+pub mod tape;
+
+#[cfg(not(no_cuda))]
+pub use tape::Var;

crates/xtrain-autodiff/src/ops.rs

@@ -0,0 +1,172 @@
+//! Differentiable ops as autograd nodes (Phase T4).
+//!
+//! Each function runs the forward [`Tensor`] kernel, then builds a [`Var`] whose
+//! backward closure computes the analytic gradient (see
+//! `docs/03-autograd-engine.md` for the math) and pushes it to each parent via
+//! [`Var::push_grad`] (which SUMs — correct under fan-out). Forward outputs that
+//! the backward needs (softmax `y`, rms `inv_rms`, cross-entropy `probs`) are
+//! cached by moving them into the closure.
+//!
+//! Attention is NOT a node here: it is composed from `matmul` + `scale` +
+//! `softmax` in user code, and its backward falls out of theirs.
+
+#![cfg(not(no_cuda))]
+
+use crate::tape::Var;
+use xtrain_tensor::Tensor;
+
+/// `C = A @ B` (2D). Backward: `dA = dC @ Bᵀ`, `dB = Aᵀ @ dC`.
+pub fn matmul(a: &Var, b: &Var) -> Var {
+    let out = a.value().matmul(&b.value());
+    Var::from_op(
+        out,
+        vec![a.clone(), b.clone()],
+        Box::new(|dc, parents| {
+            let a = parents[0].value();
+            let b = parents[1].value();
+            let (da, db) = Tensor::matmul_backward(&a, &b, dc);
+            Var::push_grad(&parents[0], da);
+            Var::push_grad(&parents[1], db);
+        }),
+    )
+}
+
+/// Elementwise `out = a + b` (same shape). Backward: grad flows unchanged to both.
+pub fn add(a: &Var, b: &Var) -> Var {
+    let out = a.value().add(&b.value());
+    Var::from_op(
+        out,
+        vec![a.clone(), b.clone()],
+        Box::new(|d, parents| {
+            Var::push_grad(&parents[0], d.clone());
+            Var::push_grad(&parents[1], d.clone());
+        }),
+    )
+}
+
+/// Elementwise `out = a * b` (Hadamard). Backward: `da = d∘b`, `db = d∘a`.
+pub fn mul(a: &Var, b: &Var) -> Var {
+    let out = a.value().mul(&b.value());
+    Var::from_op(
+        out,
+        vec![a.clone(), b.clone()],
+        Box::new(|d, parents| {
+            let a = parents[0].value();
+            let b = parents[1].value();
+            Var::push_grad(&parents[0], d.mul(&b));
+            Var::push_grad(&parents[1], d.mul(&a));
+        }),
+    )
+}
+
+/// Broadcast bias add: `out[r,c] = x[r,c] + bias[c]`. Backward: `dx = d`,
+/// `dbias[c] = sum_r d[r,c]` (sum over the broadcast dim).
+pub fn add_bias(x: &Var, bias: &Var) -> Var {
+    let out = x.value().add_bias(&bias.value());
+    Var::from_op(
+        out,
+        vec![x.clone(), bias.clone()],
+        Box::new(|d, parents| {
+            Var::push_grad(&parents[0], d.clone());
+            Var::push_grad(&parents[1], d.sum_rows());
+        }),
+    )
+}
+
+/// Scale by a constant: `out = x * alpha`. Backward: `dx = d * alpha`.
+pub fn scale(x: &Var, alpha: f32) -> Var {
+    let out = x.value().scale(alpha);
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            Var::push_grad(&parents[0], d.scale(alpha));
+        }),
+    )
+}
+
+/// RMSNorm: `y = x * rsqrt(mean(x²)+eps) * gamma`. Caches `inv_rms` for backward.
+pub fn rms_norm(x: &Var, gamma: &Var, eps: f32) -> Var {
+    let (y, inv_rms) = x.value().rms_norm(&gamma.value(), eps);
+    Var::from_op(
+        y,
+        vec![x.clone(), gamma.clone()],
+        Box::new(move |dy, parents| {
+            let x = parents[0].value();
+            let gamma = parents[1].value();
+            let (dx, dgamma) = Tensor::rms_norm_backward(&x, &gamma, dy, &inv_rms);
+            Var::push_grad(&parents[0], dx);
+            Var::push_grad(&parents[1], dgamma);
+        }),
+    )
+}
+
+/// SiLU: `y = x * sigmoid(x)`. Backward uses the forward `x`.
+pub fn silu(x: &Var) -> Var {
+    let out = x.value().silu();
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(|dy, parents| {
+            let x = parents[0].value();
+            Var::push_grad(&parents[0], Tensor::silu_backward(&x, dy));
+        }),
+    )
+}
+
+/// SwiGLU (SiLU-gated GLU): `out = silu(gate) ∘ up`. Composed from `silu` + `mul`
+/// so its backward comes from theirs — no dedicated kernel needed.
+pub fn swiglu(gate: &Var, up: &Var) -> Var {
+    mul(&silu(gate), up)
+}
+
+/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]`. Orthogonal map, so the
+/// backward is the inverse rotation of `dy` — no cached forward values needed.
+pub fn rope(x: &Var, theta: f32) -> Var {
+    let out = x.value().rope(theta);
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |dy, parents| {
+            Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta));
+        }),
+    )
+}
+
+/// Row-wise softmax. Caches the output `y` for the Jacobian backward.
+pub fn softmax(x: &Var) -> Var {
+    let y = x.value().softmax();
+    let y_cache = y.clone();
+    Var::from_op(
+        y,
+        vec![x.clone()],
+        Box::new(move |dy, parents| {
+            Var::push_grad(&parents[0], Tensor::softmax_backward(&y_cache, dy));
+        }),
+    )
+}
+
+/// 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.
+pub fn cross_entropy(x: &Var, target: &Tensor) -> Var {
+    let (probs, per_row) = x.value().cross_entropy(target);
+    let rows = x.value().shape()[0];
+    // Mean loss as a host scalar wrapped back into a [1] tensor.
+    let mean = per_row.to_device(xtrain_tensor::Device::Cpu);
+    let mean_val: f32 = mean.as_slice::<f32>().iter().sum::<f32>() / rows as f32;
+    let loss = Tensor::from_slice(&[mean_val], &[1]).to_device(x.value().device());
+
+    let target = target.clone();
+    Var::from_op(
+        loss,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            // `d` is the scalar upstream grad (1.0 when this is the loss root).
+            let upstream = d.to_device(xtrain_tensor::Device::Cpu).as_slice::<f32>()[0];
+            let scale = upstream / rows as f32;
+            let dx = Tensor::cross_entropy_backward(&probs, &target, scale);
+            Var::push_grad(&parents[0], dx);
+        }),
+    )
+}

crates/xtrain-autodiff/src/tape.rs

@@ -0,0 +1,164 @@
+//! Tape-based reverse-mode autograd (Phase T4).
+//!
+//! A [`Var`] is a reference-counted node wrapping a value [`Tensor`], an optional
+//! accumulated gradient, and — for non-leaf nodes — the parents it was computed
+//! from plus a backward closure. Forward ops (see [`crate::ops`]) build these
+//! nodes; [`Var::backward`] walks the graph in reverse topological order and
+//! pushes gradients to parents, **summing** on fan-out (a tensor consumed by
+//! several ops).
+//!
+//! The graph is dynamic (define-by-run) and the design favours clarity over
+//! speed: each op synchronizes its own kernels (T3 has no streams yet), and
+//! gradient accumulation is an explicit elementwise add.
+
+#![cfg(not(no_cuda))]
+
+use std::cell::RefCell;
+use std::rc::Rc;
+use xtrain_tensor::Tensor;
+
+/// Backward closure: given this node's accumulated grad and its parents, compute
+/// and accumulate each parent's gradient contribution.
+pub type BackwardFn = Box<dyn Fn(&Tensor, &[Var])>;
+
+pub struct VarNode {
+    pub value: Tensor,
+    pub grad: Option<Tensor>,
+    parents: Vec<Var>,
+    backward: Option<BackwardFn>,
+}
+
+/// A node in the autograd tape. Cheap to clone (bumps the `Rc`); clones share
+/// the same underlying node, which is how fan-out is detected.
+#[derive(Clone)]
+pub struct Var(Rc<RefCell<VarNode>>);
+
+impl Var {
+    /// A leaf node (parameter / input). After `backward`, its `.grad()` holds the
+    /// accumulated gradient of the loss w.r.t. this tensor.
+    pub fn leaf(value: Tensor) -> Self {
+        Var(Rc::new(RefCell::new(VarNode {
+            value,
+            grad: None,
+            parents: Vec::new(),
+            backward: None,
+        })))
+    }
+
+    /// Build a non-leaf node from a forward `value`, its `parents`, and a
+    /// `backward` closure. Used by the op constructors in [`crate::ops`] and by
+    /// callers that compose custom nodes (e.g. a loss reduction or a transpose
+    /// the built-in op set doesn't cover).
+    pub fn from_op(value: Tensor, parents: Vec<Var>, backward: BackwardFn) -> Self {
+        Var(Rc::new(RefCell::new(VarNode {
+            value,
+            grad: None,
+            parents,
+            backward: Some(backward),
+        })))
+    }
+
+    /// Clone of the value tensor (cheap: storage is `Arc`-shared).
+    pub fn value(&self) -> Tensor {
+        self.0.borrow().value.clone()
+    }
+
+    /// The accumulated gradient, if any (populated after `backward`).
+    pub fn grad(&self) -> Option<Tensor> {
+        self.0.borrow().grad.clone()
+    }
+
+    /// Pointer identity, used to dedup nodes during the topological sort.
+    fn id(&self) -> *const RefCell<VarNode> {
+        Rc::as_ptr(&self.0)
+    }
+
+    /// Accumulate `g` into this node's grad slot (SUM — the fan-out rule).
+    fn accumulate(&self, g: Tensor) {
+        let mut node = self.0.borrow_mut();
+        match node.grad.take() {
+            None => node.grad = Some(g),
+            Some(existing) => node.grad = Some(existing.add(&g)),
+        }
+    }
+
+    /// Reverse-mode backward from this node (treated as a scalar loss: its
+    /// upstream grad is seeded to ones). Populates `.grad` on every node that
+    /// transitively feeds it.
+    ///
+    /// `loss` must be a scalar (single element) so the seed `dL/dL = 1` is
+    /// unambiguous, matching the `L = sum(W∘out)` grad-check convention.
+    pub fn backward(&self) {
+        assert_eq!(
+            self.value().numel(),
+            1,
+            "backward() expects a scalar loss; got shape {:?}",
+            self.value().shape()
+        );
+
+        // 1. Topological order (post-order DFS), parents before children.
+        let mut topo: Vec<Var> = Vec::new();
+        let mut visited: Vec<*const RefCell<VarNode>> = Vec::new();
+        build_topo(self, &mut topo, &mut visited);
+
+        // 2. Seed the loss gradient with ones.
+        let seed = ones_like(&self.value());
+        self.accumulate(seed);
+
+        // 3. Walk in reverse: each node hands its grad to its parents' closures.
+        for var in topo.iter().rev() {
+            let (grad, parents, backward) = {
+                let node = var.0.borrow();
+                (
+                    node.grad.clone(),
+                    node.parents.clone(),
+                    node.backward.is_some(),
+                )
+            };
+            let grad = match grad {
+                Some(g) => g,
+                None => continue, // node didn't contribute to the loss
+            };
+            if backward {
+                // Borrow the closure out, run it, then drop the borrow.
+                let node = var.0.borrow();
+                let bw = node.backward.as_ref().unwrap();
+                bw(&grad, &parents);
+            }
+        }
+    }
+
+    /// Drop the parents/closure so the graph can be freed, keeping value+grad.
+    /// (Not needed for tests; provided for completeness of the engine.)
+    pub fn detach_graph(&self) {
+        let mut node = self.0.borrow_mut();
+        node.parents.clear();
+        node.backward = None;
+    }
+
+    /// Distribute `g` to a parent (called from op backward closures). Every node
+    /// accumulates its grad: intermediates need it for the chain rule, leaves
+    /// expose it as the result. This SUM is what makes fan-out correct.
+    pub fn push_grad(parent: &Var, g: Tensor) {
+        parent.accumulate(g);
+    }
+}
+
+/// Post-order DFS: append a node only after all its parents, dedup by identity.
+fn build_topo(var: &Var, topo: &mut Vec<Var>, visited: &mut Vec<*const RefCell<VarNode>>) {
+    if visited.contains(&var.id()) {
+        return;
+    }
+    visited.push(var.id());
+    let parents = var.0.borrow().parents.clone();
+    for p in &parents {
+        build_topo(p, topo, visited);
+    }
+    topo.push(var.clone());
+}
+
+/// A ones tensor matching `t`'s shape/device (the backward seed for a scalar).
+fn ones_like(t: &Tensor) -> Tensor {
+    let host = vec![1.0f32; t.numel()];
+    Tensor::from_slice(&host, t.shape()).to_device(t.device())
+}

crates/xtrain-autodiff/tests/autograd.rs

@@ -0,0 +1,546 @@
+// GPU acceptance tests for the Phase T4 autograd engine + per-op backward.
+// Pattern (from xtrain-tensor/tests/gemm.rs `run_bwd`): build a scalar loss
+// L = sum(W ∘ out) with W fixed random ⇒ the upstream grad dOut = W. Run the op
+// through the tape, call backward(), and grad-check each input's .grad() against
+// central finite differences of L.
+//
+// Gated behind `not(no_cuda)`: compiles out on a GPU-less host, runs on dash5.
+#![cfg(not(no_cuda))]
+
+use xtrain_autodiff::ops;
+use xtrain_autodiff::tape::Var;
+use xtrain_autodiff::{GradCheckConfig, grad_check};
+use xtrain_cuda::device;
+use xtrain_tensor::{Device, Tensor};
+
+// Deterministic LCG fill in [-0.5, 0.5), same as the gemm tests.
+fn fill(n: usize, seed: u64) -> Vec<f32> {
+    let mut state = seed
+        .wrapping_mul(2862933555777941757)
+        .wrapping_add(3037000493);
+    (0..n)
+        .map(|_| {
+            state = state
+                .wrapping_mul(6364136223846793005)
+                .wrapping_add(1442695040888963407);
+            ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
+        })
+        .collect()
+}
+
+fn require_gpu() {
+    assert!(
+        device::device_count().expect("device count") > 0,
+        "no CUDA device"
+    );
+    device::set_device(0).unwrap();
+}
+
+fn cuda(data: &[f32], shape: &[usize]) -> Tensor {
+    Tensor::from_slice(data, shape).to_device(Device::Cuda(0))
+}
+
+// L = sum(W ∘ out) for fixed weights W over the op output.
+fn weighted_sum(out: &Tensor, w: &[f32]) -> f32 {
+    out.to_device(Device::Cpu)
+        .as_slice::<f32>()
+        .iter()
+        .zip(w)
+        .map(|(o, w)| o * w)
+        .sum()
+}
+
+// Tolerances: ops with elementwise/linear forwards (add, mul, scale, bias, rope)
+// are exactly linear in each input, so a large eps just sharpens f32 resolution.
+// Nonlinear ops (rms_norm, silu, softmax, cross_entropy) carry O(eps²) truncation
+// → smaller eps. atol floors near-zero grads.
+fn cfg_linear() -> GradCheckConfig {
+    GradCheckConfig {
+        eps: 1e-2,
+        rel_tol: 2e-2,
+        atol: 1e-3,
+    }
+}
+fn cfg_nonlinear() -> GradCheckConfig {
+    GradCheckConfig {
+        eps: 1e-3,
+        rel_tol: 3e-2,
+        atol: 1e-3,
+    }
+}
+
+fn report(name: &str, res: &xtrain_autodiff::GradCheckResult) {
+    println!(
+        "{name}: max_rel_err = {:.3e} (worst num={:.5} ana={:.5} @ {})",
+        res.max_rel_err, res.worst_numeric, res.worst_analytic, res.worst_index
+    );
+    assert!(res.passed, "{name} grad-check failed: {res:?}");
+}
+
+// ---- add ----
+#[test]
+fn add_bwd() {
+    require_gpu();
+    let (m, n) = (8, 6);
+    let a_h = fill(m * n, 1);
+    let b_h = fill(m * n, 2);
+    let w = fill(m * n, 3);
+
+    let a = Var::leaf(cuda(&a_h, &[m, n]));
+    let b = Var::leaf(cuda(&b_h, &[m, n]));
+    let out = ops::add(&a, &b);
+    let loss = scalar_loss(&out, &w);
+    loss.backward();
+
+    let da = a.grad().unwrap().to_device(Device::Cpu);
+    let db = b.grad().unwrap().to_device(Device::Cpu);
+    let bf = b_h.clone();
+    let wf = w.clone();
+    let la = move |v: &[f32], s: &[usize]| {
+        let o = cuda(v, s).add(&cuda(&bf, &[m, n]));
+        weighted_sum(&o, &wf)
+    };
+    report(
+        "add dA",
+        &grad_check(&a_h, &[m, n], &la, da.as_slice::<f32>(), cfg_linear()),
+    );
+    let af = a_h.clone();
+    let wf = w.clone();
+    let lb = move |v: &[f32], s: &[usize]| {
+        let o = cuda(&af, &[m, n]).add(&cuda(v, s));
+        weighted_sum(&o, &wf)
+    };
+    report(
+        "add dB",
+        &grad_check(&b_h, &[m, n], &lb, db.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- mul ----
+#[test]
+fn mul_bwd() {
+    require_gpu();
+    let (m, n) = (8, 6);
+    let a_h = fill(m * n, 11);
+    let b_h = fill(m * n, 22);
+    let w = fill(m * n, 33);
+
+    let a = Var::leaf(cuda(&a_h, &[m, n]));
+    let b = Var::leaf(cuda(&b_h, &[m, n]));
+    let out = ops::mul(&a, &b);
+    scalar_loss(&out, &w).backward();
+
+    let da = a.grad().unwrap().to_device(Device::Cpu);
+    let db = b.grad().unwrap().to_device(Device::Cpu);
+    let bf = b_h.clone();
+    let wf = w.clone();
+    let la = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).mul(&cuda(&bf, &[m, n])), &wf);
+    report(
+        "mul dA",
+        &grad_check(&a_h, &[m, n], &la, da.as_slice::<f32>(), cfg_linear()),
+    );
+    let af = a_h.clone();
+    let wf = w.clone();
+    let lb = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(&af, &[m, n]).mul(&cuda(v, s)), &wf);
+    report(
+        "mul dB",
+        &grad_check(&b_h, &[m, n], &lb, db.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- add_bias (broadcast) ----
+#[test]
+fn add_bias_bwd() {
+    require_gpu();
+    let (m, n) = (10, 7);
+    let x_h = fill(m * n, 5);
+    let b_h = fill(n, 6);
+    let w = fill(m * n, 7);
+
+    let x = Var::leaf(cuda(&x_h, &[m, n]));
+    let bias = Var::leaf(cuda(&b_h, &[n]));
+    let out = ops::add_bias(&x, &bias);
+    scalar_loss(&out, &w).backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    let dbias = bias.grad().unwrap().to_device(Device::Cpu);
+    let bf = b_h.clone();
+    let wf = w.clone();
+    let lx =
+        move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).add_bias(&cuda(&bf, &[n])), &wf);
+    report(
+        "add_bias dX",
+        &grad_check(&x_h, &[m, n], &lx, dx.as_slice::<f32>(), cfg_linear()),
+    );
+    let xf = x_h.clone();
+    let wf = w.clone();
+    let lb =
+        move |v: &[f32], s: &[usize]| weighted_sum(&cuda(&xf, &[m, n]).add_bias(&cuda(v, s)), &wf);
+    report(
+        "add_bias dBias",
+        &grad_check(&b_h, &[n], &lb, dbias.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- matmul (sanity through the Var layer; T3 already checks the kernel) ----
+#[test]
+fn matmul_bwd() {
+    require_gpu();
+    let (m, k, n) = (6, 5, 4);
+    let a_h = fill(m * k, 41);
+    let b_h = fill(k * n, 42);
+    let w = fill(m * n, 43);
+
+    let a = Var::leaf(cuda(&a_h, &[m, k]));
+    let b = Var::leaf(cuda(&b_h, &[k, n]));
+    let out = ops::matmul(&a, &b);
+    scalar_loss(&out, &w).backward();
+
+    let da = a.grad().unwrap().to_device(Device::Cpu);
+    let db = b.grad().unwrap().to_device(Device::Cpu);
+    let bf = b_h.clone();
+    let wf = w.clone();
+    let la =
+        move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).matmul(&cuda(&bf, &[k, n])), &wf);
+    report(
+        "matmul dA",
+        &grad_check(&a_h, &[m, k], &la, da.as_slice::<f32>(), cfg_linear()),
+    );
+    let af = a_h.clone();
+    let wf = w.clone();
+    let lb =
+        move |v: &[f32], s: &[usize]| weighted_sum(&cuda(&af, &[m, k]).matmul(&cuda(v, s)), &wf);
+    report(
+        "matmul dB",
+        &grad_check(&b_h, &[k, n], &lb, db.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- rms_norm ----
+#[test]
+fn rms_norm_bwd() {
+    require_gpu();
+    let (rows, cols) = (5, 16);
+    let eps = 1e-5;
+    let x_h = fill(rows * cols, 51);
+    let g_h: Vec<f32> = fill(cols, 52).iter().map(|v| v + 1.0).collect(); // gamma ~1
+    let w = fill(rows * cols, 53);
+
+    let x = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let gamma = Var::leaf(cuda(&g_h, &[cols]));
+    let out = ops::rms_norm(&x, &gamma, eps);
+    scalar_loss(&out, &w).backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    let dg = gamma.grad().unwrap().to_device(Device::Cpu);
+    let gf = g_h.clone();
+    let wf = w.clone();
+    let lx = move |v: &[f32], s: &[usize]| {
+        let (o, _) = cuda(v, s).rms_norm(&cuda(&gf, &[cols]), eps);
+        weighted_sum(&o, &wf)
+    };
+    report(
+        "rms_norm dX",
+        &grad_check(
+            &x_h,
+            &[rows, cols],
+            &lx,
+            dx.as_slice::<f32>(),
+            cfg_nonlinear(),
+        ),
+    );
+    let xf = x_h.clone();
+    let wf = w.clone();
+    let lg = move |v: &[f32], s: &[usize]| {
+        let (o, _) = cuda(&xf, &[rows, cols]).rms_norm(&cuda(v, s), eps);
+        weighted_sum(&o, &wf)
+    };
+    report(
+        "rms_norm dGamma",
+        &grad_check(&g_h, &[cols], &lg, dg.as_slice::<f32>(), cfg_nonlinear()),
+    );
+}
+
+// ---- silu ----
+#[test]
+fn silu_bwd() {
+    require_gpu();
+    let n = 64;
+    let x_h = fill(n, 61);
+    let w = fill(n, 62);
+
+    let x = Var::leaf(cuda(&x_h, &[n]));
+    let out = ops::silu(&x);
+    scalar_loss(&out, &w).backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    let wf = w.clone();
+    let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).silu(), &wf);
+    report(
+        "silu dX",
+        &grad_check(&x_h, &[n], &lx, dx.as_slice::<f32>(), cfg_nonlinear()),
+    );
+}
+
+// ---- swiglu (composed: silu(gate) ∘ up) ----
+#[test]
+fn swiglu_bwd() {
+    require_gpu();
+    let n = 48;
+    let g_h = fill(n, 71);
+    let u_h = fill(n, 72);
+    let w = fill(n, 73);
+
+    let gate = Var::leaf(cuda(&g_h, &[n]));
+    let up = Var::leaf(cuda(&u_h, &[n]));
+    let out = ops::swiglu(&gate, &up);
+    scalar_loss(&out, &w).backward();
+
+    let dg = gate.grad().unwrap().to_device(Device::Cpu);
+    let du = up.grad().unwrap().to_device(Device::Cpu);
+    let uf = u_h.clone();
+    let wf = w.clone();
+    let lg =
+        move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).silu().mul(&cuda(&uf, &[n])), &wf);
+    report(
+        "swiglu dGate",
+        &grad_check(&g_h, &[n], &lg, dg.as_slice::<f32>(), cfg_nonlinear()),
+    );
+    let gf = g_h.clone();
+    let wf = w.clone();
+    let lu =
+        move |v: &[f32], s: &[usize]| weighted_sum(&cuda(&gf, &[n]).silu().mul(&cuda(v, s)), &wf);
+    report(
+        "swiglu dUp",
+        &grad_check(&u_h, &[n], &lu, du.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- rope ----
+#[test]
+fn rope_bwd() {
+    require_gpu();
+    let (tokens, heads, head_dim) = (4, 2, 8);
+    let n = tokens * heads * head_dim;
+    let theta = 10000.0;
+    let x_h = fill(n, 81);
+    let w = fill(n, 82);
+
+    let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim]));
+    let out = ops::rope(&x, theta);
+    scalar_loss(&out, &w).backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    let wf = w.clone();
+    let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).rope(theta), &wf);
+    report(
+        "rope dX",
+        &grad_check(
+            &x_h,
+            &[tokens, heads, head_dim],
+            &lx,
+            dx.as_slice::<f32>(),
+            cfg_linear(),
+        ),
+    );
+}
+
+// ---- softmax ----
+#[test]
+fn softmax_bwd() {
+    require_gpu();
+    let (rows, cols) = (4, 10);
+    let x_h = fill(rows * cols, 91);
+    let w = fill(rows * cols, 92);
+
+    let x = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let out = ops::softmax(&x);
+    scalar_loss(&out, &w).backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    let wf = w.clone();
+    let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).softmax(), &wf);
+    report(
+        "softmax dX",
+        &grad_check(
+            &x_h,
+            &[rows, cols],
+            &lx,
+            dx.as_slice::<f32>(),
+            cfg_nonlinear(),
+        ),
+    );
+}
+
+// ---- cross_entropy (scalar loss; backward = (softmax - onehot)/rows) ----
+#[test]
+fn cross_entropy_bwd() {
+    require_gpu();
+    let (rows, cols) = (5, 8);
+    let x_h = fill(rows * cols, 101);
+    let targets: Vec<i32> = (0..rows).map(|r| (r * 3 % cols) as i32).collect();
+    let target = Tensor::from_slice(&targets, &[rows]).to_device(Device::Cuda(0));
+
+    let x = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let loss = ops::cross_entropy(&x, &target);
+    loss.backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    // Loss is already scalar (mean NLL) — grad-check it directly, no W weighting.
+    let tgt = targets.clone();
+    let lx = move |v: &[f32], s: &[usize]| {
+        let t = Tensor::from_slice(&tgt, &[rows]).to_device(Device::Cuda(0));
+        let (_, per_row) = cuda(v, s).cross_entropy(&t);
+        per_row
+            .to_device(Device::Cpu)
+            .as_slice::<f32>()
+            .iter()
+            .sum::<f32>()
+            / rows as f32
+    };
+    report(
+        "cross_entropy dX",
+        &grad_check(
+            &x_h,
+            &[rows, cols],
+            &lx,
+            dx.as_slice::<f32>(),
+            cfg_nonlinear(),
+        ),
+    );
+}
+
+// ---- FAN-OUT: a tensor feeding two consumers must SUM grads ----
+// y = x*x + x*x via two separate mul nodes on the same Var x → dL/dx must be the
+// sum of both branches. With W=1, out=2x², so dOut=W=1 and dx (numeric) = 4x.
+#[test]
+fn fanout_grad_accumulation() {
+    require_gpu();
+    let n = 12;
+    let x_h = fill(n, 111);
+    let w = vec![1.0f32; n];
+
+    let x = Var::leaf(cuda(&x_h, &[n]));
+    let sq1 = ops::mul(&x, &x); // x∘x  (x consumed twice within one node)
+    let sq2 = ops::mul(&x, &x); // x∘x  (x consumed again across nodes)
+    let out = ops::add(&sq1, &sq2); // 2x²
+    scalar_loss(&out, &w).backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    let wf = w.clone();
+    let lx = move |v: &[f32], s: &[usize]| {
+        let t = cuda(v, s);
+        let o = t.mul(&t).add(&t.mul(&t));
+        weighted_sum(&o, &wf)
+    };
+    // Analytic dx should be 4x; fan-out summed all four uses of x.
+    report(
+        "fanout dX",
+        &grad_check(&x_h, &[n], &lx, dx.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- COMPOSED ATTENTION: attn = matmul(softmax(matmul(Q,Kᵀ)·scale), V) ----
+// Single head, single batch. Backward falls out of matmul+scale+softmax nodes.
+#[test]
+fn attention_composed_bwd() {
+    require_gpu();
+    let (s, d) = (5, 6); // seq_len, head_dim
+    let scale = 1.0 / (d as f32).sqrt();
+    let q_h = fill(s * d, 121);
+    let k_h = fill(s * d, 122);
+    let v_h = fill(s * d, 123);
+    let w = fill(s * d, 124); // weights over the [s,d] attention output
+
+    let attn = |q: &Var, k: &Var, v: &Var| -> Var {
+        let kt = transpose_var(k); // [d,s]  (manual transpose node)
+        let scores = ops::scale(&ops::matmul(q, &kt), scale); // [s,s]
+        let probs = ops::softmax(&scores);
+        ops::matmul(&probs, v) // [s,d]
+    };
+
+    let q = Var::leaf(cuda(&q_h, &[s, d]));
+    let k = Var::leaf(cuda(&k_h, &[s, d]));
+    let v = Var::leaf(cuda(&v_h, &[s, d]));
+    let out = attn(&q, &k, &v);
+    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);
+
+    // Re-run the same forward inside the loss closures (host-side) per input.
+    let fwd = move |qh: &[f32], kh: &[f32], vh: &[f32]| -> f32 {
+        let qv = cuda(qh, &[s, d]);
+        let kv = cuda(kh, &[s, d]);
+        let vv = cuda(vh, &[s, d]);
+        let scores = qv.matmul(&kv.transpose_2d()).scale(scale);
+        let probs = scores.softmax();
+        weighted_sum(&probs.matmul(&vv), &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(
+        "attn dQ",
+        &grad_check(&q_h, &[s, d], &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(
+        "attn dK",
+        &grad_check(&k_h, &[s, d], &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(
+        "attn dV",
+        &grad_check(&v_h, &[s, d], &lv, dv.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// --- test helpers ---
+
+// Scalar loss node L = sum(W ∘ out): wraps a fixed-weight Var and reduces. We
+// implement it as: elementwise mul by a constant-W leaf, then sum-to-scalar.
+fn scalar_loss(out: &Var, w: &[f32]) -> Var {
+    let wt = Var::leaf(cuda(w, out.value().shape()));
+    let prod = ops::mul(out, &wt);
+    sum_all(&prod)
+}
+
+// Sum-to-scalar node: out = sum(x). Backward broadcasts the scalar grad to a
+// ones-shaped tensor over x. Implemented here (test-local) since the engine's
+// op set doesn't include a generic reduction; cross_entropy is the only loss op.
+fn sum_all(x: &Var) -> Var {
+    let xv = x.value();
+    let total: f32 = xv.to_device(Device::Cpu).as_slice::<f32>().iter().sum();
+    let scalar = Tensor::from_slice(&[total], &[1]).to_device(xv.device());
+    let shape: Vec<usize> = xv.shape().to_vec();
+    Var::from_op(
+        scalar,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            // d is [1]; broadcast d to a same-shape tensor over the input.
+            let dval = d.to_device(Device::Cpu).as_slice::<f32>()[0];
+            let ones = vec![dval; shape.iter().product()];
+            let g = Tensor::from_slice(&ones, &shape).to_device(Device::Cuda(0));
+            Var::push_grad(&parents[0], g);
+        }),
+    )
+}
+
+// Manual transpose node for the composed-attention test (the engine has no
+// transpose op; xserv does the equivalent host-side reshape around RoPE).
+fn transpose_var(x: &Var) -> Var {
+    let xt = x.value().transpose_2d();
+    Var::from_op(
+        xt,
+        vec![x.clone()],
+        Box::new(|d, parents| {
+            Var::push_grad(&parents[0], d.transpose_2d());
+        }),
+    )
+}

crates/xtrain-cuda/build.rs

@@ -32,6 +32,7 @@ fn main() {
         .file("../../csrc/test/vecadd.cu")
         .file("../../csrc/ops/elementwise.cu")
         .file("../../csrc/ops/gemm.cu")
+        .file("../../csrc/ops/nn.cu")
         .compile("xtrain_cuda_kernels");
 }
 

crates/xtrain-cuda/src/ffi.rs

@@ -63,6 +63,120 @@ unsafe extern "C" {
     );
 }
 
+// Transformer / autograd op kernels (csrc/ops/nn.cu). Forward + backward for the
+// ops the Phase T4 tape engine needs. All F32, row-major, contiguous.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    // Elementwise: out = a + b ; out = a * b.
+    pub fn launch_add_f32(a: *const f32, b: *const f32, out: *mut f32, n: i32, s: CudaStream);
+    pub fn launch_mul_f32(a: *const f32, b: *const f32, out: *mut f32, n: i32, s: CudaStream);
+
+    // Broadcast bias add: out[r,c] = x[r,c] + bias[c]. x:[rows,cols], bias:[cols].
+    pub fn launch_add_bias_f32(
+        x: *const f32,
+        bias: *const f32,
+        out: *mut f32,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+    // Column-sum (over rows): dbias[c] = sum_r dout[r,c]. Bias backward.
+    pub fn launch_sum_rows_f32(
+        dout: *const f32,
+        dbias: *mut f32,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+
+    // RMSNorm forward: writes y[rows,cols] and inv_rms[rows] (cached for bwd).
+    pub fn launch_rms_norm_f32(
+        x: *const f32,
+        gamma: *const f32,
+        y: *mut f32,
+        inv_rms: *mut f32,
+        rows: i32,
+        cols: i32,
+        eps: f32,
+        s: CudaStream,
+    );
+    pub fn launch_rms_norm_dx_f32(
+        x: *const f32,
+        gamma: *const f32,
+        dy: *const f32,
+        inv_rms: *const f32,
+        dx: *mut f32,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+    pub fn launch_rms_norm_dgamma_f32(
+        x: *const f32,
+        dy: *const f32,
+        inv_rms: *const f32,
+        dgamma: *mut f32,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+
+    // SiLU: y = x*sigmoid(x); backward dx.
+    pub fn launch_silu_f32(x: *const f32, y: *mut f32, n: i32, s: CudaStream);
+    pub fn launch_silu_dx_f32(x: *const f32, dy: *const f32, dx: *mut f32, n: i32, s: CudaStream);
+
+    // RoPE (rotate_half), x:[tokens,heads,head_dim], position = token index.
+    pub fn launch_rope_f32(
+        x: *const f32,
+        y: *mut f32,
+        tokens: i32,
+        heads: i32,
+        head_dim: i32,
+        theta: f32,
+        s: CudaStream,
+    );
+    pub fn launch_rope_dx_f32(
+        dy: *const f32,
+        dx: *mut f32,
+        tokens: i32,
+        heads: i32,
+        head_dim: i32,
+        theta: f32,
+        s: CudaStream,
+    );
+
+    // Row-wise softmax + Jacobian backward.
+    pub fn launch_softmax_f32(x: *const f32, y: *mut f32, rows: i32, cols: i32, s: CudaStream);
+    pub fn launch_softmax_dx_f32(
+        y: *const f32,
+        dy: *const f32,
+        dx: *mut f32,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+
+    // Cross-entropy: fwd writes probs[rows,cols] + per-row loss[rows];
+    // bwd dx = scale*(probs - onehot).
+    pub fn launch_cross_entropy_fwd_f32(
+        x: *const f32,
+        target: *const i32,
+        probs: *mut f32,
+        loss: *mut f32,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+    pub fn launch_cross_entropy_dx_f32(
+        probs: *const f32,
+        target: *const i32,
+        dx: *mut f32,
+        rows: i32,
+        cols: i32,
+        scale: f32,
+        s: CudaStream,
+    );
+}
+
 // cuBLAS — used ONLY as a correctness reference for the hand-written GEMM in
 // tests. Declared (and linked, see build.rs) only when CUDA is compiled in.
 #[cfg(not(no_cuda))]

crates/xtrain-tensor/src/dtype.rs

@@ -7,18 +7,22 @@
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum DType {
     F32,
+    /// 32-bit signed integers. Used for cross-entropy targets (token ids).
+    I32,
 }
 
 impl DType {
     pub fn size_bytes(self) -> usize {
         match self {
             DType::F32 => 4,
+            DType::I32 => 4,
         }
     }
 
     pub fn name(self) -> &'static str {
         match self {
             DType::F32 => "f32",
+            DType::I32 => "i32",
         }
     }
 }
@@ -45,3 +49,13 @@ impl TensorDType for f32 {
         v as f32
     }
 }
+
+impl TensorDType for i32 {
+    const DTYPE: DType = DType::I32;
+    fn to_f64(self) -> f64 {
+        self as f64
+    }
+    fn from_f64(v: f64) -> Self {
+        v as i32
+    }
+}

crates/xtrain-tensor/src/tensor.rs

@@ -247,6 +247,334 @@ impl Tensor {
         let db = a.transpose_2d().matmul(dc); // [K,M] @ [M,N] = [K,N]
         (da, db)
     }
+
+    // --- Transformer / autograd op primitives (the T4 kernels) ---
+    //
+    // Each is a thin, contiguous-F32-on-GPU wrapper over a kernel in
+    // csrc/ops/nn.cu. The autograd `Var` layer (xtrain-autodiff) builds nodes on
+    // top of these; the analytic backwards are derived in docs/03-autograd-engine.md.
+
+    /// Elementwise `out = self + other` (same shape).
+    #[cfg(not(no_cuda))]
+    pub fn add(&self, other: &Tensor) -> Self {
+        self.check_binary(other, "add");
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_add_f32(
+                self.data_ptr() as *const f32,
+                other.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                self.numel() as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("add sync failed");
+        out
+    }
+
+    /// Elementwise `out = self * other` (same shape, Hadamard product).
+    #[cfg(not(no_cuda))]
+    pub fn mul(&self, other: &Tensor) -> Self {
+        self.check_binary(other, "mul");
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_mul_f32(
+                self.data_ptr() as *const f32,
+                other.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                self.numel() as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("mul sync failed");
+        out
+    }
+
+    /// Broadcast bias add: `out[r,c] = self[r,c] + bias[c]`.
+    /// `self`:[rows,cols], `bias`:[cols].
+    #[cfg(not(no_cuda))]
+    pub fn add_bias(&self, bias: &Tensor) -> Self {
+        assert_eq!(self.ndim(), 2, "add_bias requires 2D input");
+        assert_eq!(bias.ndim(), 1, "bias must be 1D");
+        assert_eq!(self.shape[1], bias.shape[0], "bias len != cols");
+        let (rows, cols) = (self.shape[0], self.shape[1]);
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_add_bias_f32(
+                self.data_ptr() as *const f32,
+                bias.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("add_bias sync failed");
+        out
+    }
+
+    /// Column-sum over rows: `out[c] = sum_r self[r,c]`. This is the bias
+    /// backward (sum the upstream grad over the broadcast dim). `self`:[rows,cols]
+    /// → [cols].
+    #[cfg(not(no_cuda))]
+    pub fn sum_rows(&self) -> Self {
+        assert_eq!(self.ndim(), 2, "sum_rows requires 2D input");
+        let (rows, cols) = (self.shape[0], self.shape[1]);
+        let out = Tensor::zeros(&[cols], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_sum_rows_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("sum_rows sync failed");
+        out
+    }
+
+    /// RMSNorm forward: `y[r,c] = x[r,c] * inv_rms[r] * gamma[c]` with
+    /// `inv_rms = rsqrt(mean(x²) + eps)`. `self`:[rows,cols], `gamma`:[cols].
+    /// Returns `(y, inv_rms)`; `inv_rms`:[rows] is cached for backward.
+    #[cfg(not(no_cuda))]
+    pub fn rms_norm(&self, gamma: &Tensor, eps: f32) -> (Tensor, Tensor) {
+        assert_eq!(self.ndim(), 2, "rms_norm requires 2D input");
+        assert_eq!(gamma.ndim(), 1, "gamma must be 1D");
+        assert_eq!(self.shape[1], gamma.shape[0], "gamma len != cols");
+        let (rows, cols) = (self.shape[0], self.shape[1]);
+        let y = Tensor::zeros(&self.shape, DType::F32, self.device());
+        let inv_rms = Tensor::zeros(&[rows], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_rms_norm_f32(
+                self.data_ptr() as *const f32,
+                gamma.data_ptr() as *const f32,
+                y.data_ptr() as *mut f32,
+                inv_rms.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                eps,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("rms_norm sync failed");
+        (y, inv_rms)
+    }
+
+    /// RMSNorm backward. Inputs are the forward `x`, `gamma`, upstream `dy`, and
+    /// the cached `inv_rms`. Returns `(dx, dgamma)`.
+    #[cfg(not(no_cuda))]
+    pub fn rms_norm_backward(
+        x: &Tensor,
+        gamma: &Tensor,
+        dy: &Tensor,
+        inv_rms: &Tensor,
+    ) -> (Tensor, Tensor) {
+        let (rows, cols) = (x.shape[0], x.shape[1]);
+        let dx = Tensor::zeros(&[rows, cols], DType::F32, x.device());
+        let dgamma = Tensor::zeros(&[cols], DType::F32, x.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_rms_norm_dx_f32(
+                x.data_ptr() as *const f32,
+                gamma.data_ptr() as *const f32,
+                dy.data_ptr() as *const f32,
+                inv_rms.data_ptr() as *const f32,
+                dx.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+            xtrain_cuda::ffi::launch_rms_norm_dgamma_f32(
+                x.data_ptr() as *const f32,
+                dy.data_ptr() as *const f32,
+                inv_rms.data_ptr() as *const f32,
+                dgamma.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("rms_norm_backward sync failed");
+        (dx, dgamma)
+    }
+
+    /// SiLU forward: `y = x * sigmoid(x)`, elementwise.
+    #[cfg(not(no_cuda))]
+    pub fn silu(&self) -> Self {
+        assert_eq!(self.dtype, DType::F32, "silu only supports F32");
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_silu_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                self.numel() as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("silu sync failed");
+        out
+    }
+
+    /// SiLU backward: `dx = dy * (sig + x*sig*(1-sig))`, `sig = sigmoid(x)`.
+    /// Inputs are the forward `x` and upstream `dy`.
+    #[cfg(not(no_cuda))]
+    pub fn silu_backward(x: &Tensor, dy: &Tensor) -> Self {
+        let dx = Tensor::zeros(&x.shape, DType::F32, x.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_silu_dx_f32(
+                x.data_ptr() as *const f32,
+                dy.data_ptr() as *const f32,
+                dx.data_ptr() as *mut f32,
+                x.numel() as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("silu_backward sync failed");
+        dx
+    }
+
+    /// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; the position
+    /// of each token is its row index. Returns the rotated tensor.
+    #[cfg(not(no_cuda))]
+    pub fn rope(&self, theta: f32) -> Self {
+        assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]");
+        let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
+        assert_eq!(head_dim % 2, 0, "head_dim must be even");
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_rope_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                tokens as i32,
+                heads as i32,
+                head_dim as i32,
+                theta,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("rope sync failed");
+        out
+    }
+
+    /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
+    /// orthogonal map, so it needs no cached forward values, only `theta`.
+    #[cfg(not(no_cuda))]
+    pub fn rope_backward(dy: &Tensor, theta: f32) -> Self {
+        let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]);
+        let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_rope_dx_f32(
+                dy.data_ptr() as *const f32,
+                dx.data_ptr() as *mut f32,
+                tokens as i32,
+                heads as i32,
+                head_dim as i32,
+                theta,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("rope_backward sync failed");
+        dx
+    }
+
+    /// Row-wise safe softmax over the last dim. `self`:[rows,cols].
+    #[cfg(not(no_cuda))]
+    pub fn softmax(&self) -> Self {
+        assert_eq!(self.ndim(), 2, "softmax requires 2D input");
+        let (rows, cols) = (self.shape[0], self.shape[1]);
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_softmax_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("softmax sync failed");
+        out
+    }
+
+    /// Softmax backward (Jacobian): `dx[r,c] = y[r,c]*(dy[r,c] - sum_c'(dy*y))`.
+    /// Inputs are the forward output `y` and upstream `dy`.
+    #[cfg(not(no_cuda))]
+    pub fn softmax_backward(y: &Tensor, dy: &Tensor) -> Self {
+        let (rows, cols) = (y.shape[0], y.shape[1]);
+        let dx = Tensor::zeros(&y.shape, DType::F32, y.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_softmax_dx_f32(
+                y.data_ptr() as *const f32,
+                dy.data_ptr() as *const f32,
+                dx.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("softmax_backward sync failed");
+        dx
+    }
+
+    /// Cross-entropy forward over logits `self`:[rows,cols] with one I32 target
+    /// per row. Returns `(probs, loss)` where `probs`:[rows,cols] is the softmax
+    /// (cached for backward) and `loss`:[rows] is the per-row negative log-likelihood.
+    #[cfg(not(no_cuda))]
+    pub fn cross_entropy(&self, target: &Tensor) -> (Tensor, Tensor) {
+        assert_eq!(self.ndim(), 2, "cross_entropy requires 2D logits");
+        assert_eq!(target.dtype, DType::I32, "target must be I32");
+        assert_eq!(target.numel(), self.shape[0], "one target per row");
+        let (rows, cols) = (self.shape[0], self.shape[1]);
+        let probs = Tensor::zeros(&self.shape, DType::F32, self.device());
+        let loss = Tensor::zeros(&[rows], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_cross_entropy_fwd_f32(
+                self.data_ptr() as *const f32,
+                target.data_ptr() as *const i32,
+                probs.data_ptr() as *mut f32,
+                loss.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("cross_entropy sync failed");
+        (probs, loss)
+    }
+
+    /// Cross-entropy backward: `dx = scale * (probs - onehot(target))`. With
+    /// `scale = upstream / rows`, this is the gradient of the mean per-row loss.
+    #[cfg(not(no_cuda))]
+    pub fn cross_entropy_backward(probs: &Tensor, target: &Tensor, scale: f32) -> Self {
+        let (rows, cols) = (probs.shape[0], probs.shape[1]);
+        let dx = Tensor::zeros(&probs.shape, DType::F32, probs.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_cross_entropy_dx_f32(
+                probs.data_ptr() as *const f32,
+                target.data_ptr() as *const i32,
+                dx.data_ptr() as *mut f32,
+                rows as i32,
+                cols as i32,
+                scale,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("cross_entropy_backward sync failed");
+        dx
+    }
+
+    // Shared validation for same-shape binary elementwise ops.
+    #[cfg(not(no_cuda))]
+    fn check_binary(&self, other: &Tensor, op: &str) {
+        assert_eq!(self.dtype, DType::F32, "{op} only supports F32");
+        assert_eq!(other.dtype, DType::F32, "{op} only supports F32");
+        assert_eq!(self.shape(), other.shape(), "{op} shape mismatch");
+        assert_eq!(self.device(), other.device(), "{op} device mismatch");
+        assert!(
+            self.is_contiguous() && other.is_contiguous(),
+            "{op} requires contiguous tensors"
+        );
+    }
 }
 
 impl std::fmt::Debug for Tensor {

csrc/ops/nn.cu

@@ -0,0 +1,358 @@
+// Forward + backward CUDA kernels for the transformer ops the autograd engine
+// (Phase T4) needs: elementwise add/mul, broadcast bias add + its row-sum
+// backward, RMSNorm, SiLU, RoPE, row-wise softmax, and cross-entropy.
+//
+// All F32, row-major, contiguous. Forward kernels mirror xserv
+// (docs/04-transformer-kernels.md, docs/05-attention.md); the backward kernels
+// are new (xserv is inference-only). Reduction helpers are inlined here so this
+// file is self-contained (no shared header), matching the existing csrc/ layout.
+
+#include <math.h>
+
+extern "C" {
+
+// --- Warp / block reductions (sum + max), block handles one row ---
+
+__device__ __forceinline__ float warp_reduce_sum(float v) {
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+        v += __shfl_down_sync(0xffffffff, v, off);
+    return v;
+}
+
+__device__ __forceinline__ float warp_reduce_max(float v) {
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1)
+        v = fmaxf(v, __shfl_down_sync(0xffffffff, v, off));
+    return v;
+}
+
+__device__ __forceinline__ float block_reduce_sum(float v) {
+    __shared__ float shared[32];
+    int lane = threadIdx.x & 31;
+    int warp = threadIdx.x >> 5;
+    int nwarps = (blockDim.x + 31) >> 5;
+    v = warp_reduce_sum(v);
+    if (lane == 0) shared[warp] = v;
+    __syncthreads();
+    v = (threadIdx.x < nwarps) ? shared[threadIdx.x] : 0.0f;
+    if (warp == 0) v = warp_reduce_sum(v);
+    // broadcast warp-0 lane-0 result to whole block
+    __shared__ float bcast;
+    if (threadIdx.x == 0) bcast = v;
+    __syncthreads();
+    return bcast;
+}
+
+__device__ __forceinline__ float block_reduce_max(float v) {
+    __shared__ float shared[32];
+    int lane = threadIdx.x & 31;
+    int warp = threadIdx.x >> 5;
+    int nwarps = (blockDim.x + 31) >> 5;
+    v = warp_reduce_max(v);
+    if (lane == 0) shared[warp] = v;
+    __syncthreads();
+    v = (threadIdx.x < nwarps) ? shared[threadIdx.x] : -INFINITY;
+    if (warp == 0) v = warp_reduce_max(v);
+    __shared__ float bcast;
+    if (threadIdx.x == 0) bcast = v;
+    __syncthreads();
+    return bcast;
+}
+
+// =====================================================================
+// Elementwise add / mul (same-shape)
+// =====================================================================
+
+__global__ void add_k(const float* a, const float* b, float* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = a[i] + b[i];
+}
+void launch_add_f32(const float* a, const float* b, float* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    add_k<<<grid, blk, 0, (cudaStream_t)s>>>(a, b, out, n);
+}
+
+__global__ void mul_k(const float* a, const float* b, float* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = a[i] * b[i];
+}
+void launch_mul_f32(const float* a, const float* b, float* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    mul_k<<<grid, blk, 0, (cudaStream_t)s>>>(a, b, out, n);
+}
+
+// =====================================================================
+// Broadcast bias add: out[r,c] = x[r,c] + bias[c]   (x:[rows,cols])
+// Backward for bias is a column-sum (sum over rows): dbias[c] = sum_r dout[r,c].
+// =====================================================================
+
+__global__ void add_bias_k(const float* x, const float* bias, float* out,
+                           int rows, int cols) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < rows * cols) out[i] = x[i] + bias[i % cols];
+}
+void launch_add_bias_f32(const float* x, const float* bias, float* out,
+                         int rows, int cols, void* s) {
+    int n = rows * cols, blk = 256, grid = (n + blk - 1) / blk;
+    add_bias_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, bias, out, rows, cols);
+}
+
+// dbias[c] = sum_r dout[r,c]. One block per column, threads stride over rows.
+__global__ void sum_rows_k(const float* dout, float* dbias, int rows, int cols) {
+    int col = blockIdx.x;
+    float acc = 0.0f;
+    for (int r = threadIdx.x; r < rows; r += blockDim.x)
+        acc += dout[r * cols + col];
+    acc = block_reduce_sum(acc);
+    if (threadIdx.x == 0) dbias[col] = acc;
+}
+void launch_sum_rows_f32(const float* dout, float* dbias, int rows, int cols, void* s) {
+    int blk = 256;
+    sum_rows_k<<<cols, blk, 0, (cudaStream_t)s>>>(dout, dbias, rows, cols);
+}
+
+// =====================================================================
+// RMSNorm: y[r,c] = x[r,c] * inv_rms[r] * gamma[c],  inv_rms = rsqrt(mean(x²)+eps)
+// x:[rows,cols], gamma:[cols]. Forward also writes inv_rms[rows] for backward.
+// =====================================================================
+
+__global__ void rms_norm_k(const float* x, const float* gamma, float* y,
+                           float* inv_rms, int rows, int cols, float eps) {
+    int r = blockIdx.x;
+    const float* xr = x + r * cols;
+    float* yr = y + r * cols;
+    float ss = 0.0f;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x) ss += xr[c] * xr[c];
+    ss = block_reduce_sum(ss);
+    float ir = rsqrtf(ss / cols + eps);
+    if (threadIdx.x == 0) inv_rms[r] = ir;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x)
+        yr[c] = xr[c] * ir * gamma[c];
+}
+void launch_rms_norm_f32(const float* x, const float* gamma, float* y,
+                         float* inv_rms, int rows, int cols, float eps, void* s) {
+    int blk = cols < 1024 ? cols : 1024;
+    if (blk < 32) blk = 32;
+    rms_norm_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, gamma, y, inv_rms, rows, cols, eps);
+}
+
+// RMSNorm backward.
+// Let g[c] = dy[r,c]*gamma[c], ir = inv_rms[r], n = cols.
+// dx[r,c] = ir*g[c] - x[r,c]*ir³/n * sum_c(g[c]*x[r,c])
+// dgamma[c] = sum_r dy[r,c] * x[r,c] * ir   (accumulated across rows)
+__global__ void rms_norm_dx_k(const float* x, const float* gamma, const float* dy,
+                              const float* inv_rms, float* dx, int rows, int cols) {
+    int r = blockIdx.x;
+    const float* xr = x + r * cols;
+    const float* dyr = dy + r * cols;
+    float* dxr = dx + r * cols;
+    float ir = inv_rms[r];
+    float dot = 0.0f; // sum_c g[c]*x[c]
+    for (int c = threadIdx.x; c < cols; c += blockDim.x)
+        dot += dyr[c] * gamma[c] * xr[c];
+    dot = block_reduce_sum(dot);
+    float coeff = ir * ir * ir / (float)cols * dot;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x)
+        dxr[c] = ir * dyr[c] * gamma[c] - xr[c] * coeff;
+}
+void launch_rms_norm_dx_f32(const float* x, const float* gamma, const float* dy,
+                            const float* inv_rms, float* dx, int rows, int cols, void* s) {
+    int blk = cols < 1024 ? cols : 1024;
+    if (blk < 32) blk = 32;
+    rms_norm_dx_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, gamma, dy, inv_rms, dx, rows, cols);
+}
+
+// dgamma[c] = sum_r dy[r,c] * x[r,c] * inv_rms[r]. One block per column.
+__global__ void rms_norm_dgamma_k(const float* x, const float* dy, const float* inv_rms,
+                                  float* dgamma, int rows, int cols) {
+    int col = blockIdx.x;
+    float acc = 0.0f;
+    for (int r = threadIdx.x; r < rows; r += blockDim.x)
+        acc += dy[r * cols + col] * x[r * cols + col] * inv_rms[r];
+    acc = block_reduce_sum(acc);
+    if (threadIdx.x == 0) dgamma[col] = acc;
+}
+void launch_rms_norm_dgamma_f32(const float* x, const float* dy, const float* inv_rms,
+                                float* dgamma, int rows, int cols, void* s) {
+    int blk = 256;
+    rms_norm_dgamma_k<<<cols, blk, 0, (cudaStream_t)s>>>(x, dy, inv_rms, dgamma, rows, cols);
+}
+
+// =====================================================================
+// SiLU: y = x * sigmoid(x). Backward: dx = dy * (sig + x*sig*(1-sig)).
+// =====================================================================
+
+__global__ void silu_k(const float* x, float* y, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) { float xv = x[i]; y[i] = xv / (1.0f + expf(-xv)); }
+}
+void launch_silu_f32(const float* x, float* y, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    silu_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, n);
+}
+
+__global__ void silu_dx_k(const float* x, const float* dy, float* dx, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float xv = x[i];
+        float sig = 1.0f / (1.0f + expf(-xv));
+        dx[i] = dy[i] * (sig + xv * sig * (1.0f - sig));
+    }
+}
+void launch_silu_dx_f32(const float* x, const float* dy, float* dx, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    silu_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, dy, dx, n);
+}
+
+// =====================================================================
+// RoPE (rotate_half layout). x:[tokens, heads, head_dim]; position = token index.
+// y[i]      = x[i]*cos - x[i+h]*sin
+// y[i+h]    = x[i+h]*cos + x[i]*sin     (i in [0,half), h=half_dim)
+// freq[i] = theta^(-2i/head_dim); angle = pos*freq[i].
+// Backward is the inverse (transpose) rotation: apply +angle's transpose ≡ -angle.
+// dx[i]   = dy[i]*cos + dy[i+h]*sin
+// dx[i+h] = dy[i+h]*cos - dy[i]*sin
+// =====================================================================
+
+__global__ void rope_k(const float* x, float* y, int heads, int head_dim, float theta) {
+    int tok = blockIdx.x;
+    int head = blockIdx.y;
+    int half = head_dim / 2;
+    int i = threadIdx.x;
+    if (i >= half) return;
+    float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
+    float angle = (float)tok * freq;
+    float c = cosf(angle), sn = sinf(angle);
+    int base = (tok * heads + head) * head_dim;
+    float x0 = x[base + i], x1 = x[base + i + half];
+    y[base + i]        = x0 * c - x1 * sn;
+    y[base + i + half] = x1 * c + x0 * sn;
+}
+void launch_rope_f32(const float* x, float* y, int tokens, int heads,
+                     int head_dim, float theta, void* s) {
+    dim3 grid(tokens, heads);
+    int blk = head_dim / 2;
+    rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta);
+}
+
+__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, float theta) {
+    int tok = blockIdx.x;
+    int head = blockIdx.y;
+    int half = head_dim / 2;
+    int i = threadIdx.x;
+    if (i >= half) return;
+    float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
+    float angle = (float)tok * freq;
+    float c = cosf(angle), sn = sinf(angle);
+    int base = (tok * heads + head) * head_dim;
+    float d0 = dy[base + i], d1 = dy[base + i + half];
+    dx[base + i]        = d0 * c + d1 * sn;
+    dx[base + i + half] = d1 * c - d0 * sn;
+}
+void launch_rope_dx_f32(const float* dy, float* dx, int tokens, int heads,
+                        int head_dim, float theta, void* s) {
+    dim3 grid(tokens, heads);
+    int blk = head_dim / 2;
+    rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta);
+}
+
+// =====================================================================
+// Row-wise safe softmax. x:[rows,cols] → y. Backward (Jacobian):
+// dx[r,c] = y[r,c] * (dy[r,c] - sum_c'(dy[r,c']*y[r,c']))
+// =====================================================================
+
+__global__ void softmax_k(const float* x, float* y, int rows, int cols) {
+    int r = blockIdx.x;
+    const float* xr = x + r * cols;
+    float* yr = y + r * cols;
+    float m = -INFINITY;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x) m = fmaxf(m, xr[c]);
+    m = block_reduce_max(m);
+    float sum = 0.0f;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x) {
+        float e = expf(xr[c] - m);
+        yr[c] = e;
+        sum += e;
+    }
+    sum = block_reduce_sum(sum);
+    float inv = 1.0f / sum;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x) yr[c] *= inv;
+}
+void launch_softmax_f32(const float* x, float* y, int rows, int cols, void* s) {
+    int blk = cols < 1024 ? cols : 1024;
+    if (blk < 32) blk = 32;
+    softmax_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, y, rows, cols);
+}
+
+__global__ void softmax_dx_k(const float* y, const float* dy, float* dx,
+                             int rows, int cols) {
+    int r = blockIdx.x;
+    const float* yr = y + r * cols;
+    const float* dyr = dy + r * cols;
+    float* dxr = dx + r * cols;
+    float dot = 0.0f; // sum_c dy*y
+    for (int c = threadIdx.x; c < cols; c += blockDim.x) dot += dyr[c] * yr[c];
+    dot = block_reduce_sum(dot);
+    for (int c = threadIdx.x; c < cols; c += blockDim.x)
+        dxr[c] = yr[c] * (dyr[c] - dot);
+}
+void launch_softmax_dx_f32(const float* y, const float* dy, float* dx,
+                           int rows, int cols, void* s) {
+    int blk = cols < 1024 ? cols : 1024;
+    if (blk < 32) blk = 32;
+    softmax_dx_k<<<rows, blk, 0, (cudaStream_t)s>>>(y, dy, dx, rows, cols);
+}
+
+// =====================================================================
+// Cross-entropy over logits x:[rows,cols] with int target per row.
+// Forward writes per-row loss[r] = -log(softmax(x)[target]) and the softmax
+// probs[r,:] (cached for backward). Backward: dx[r,c] = (probs[r,c]-onehot)/rows
+// (mean reduction; the *rows scale folds the 1/rows of mean loss into dx).
+// =====================================================================
+
+__global__ void cross_entropy_fwd_k(const float* x, const int* target,
+                                    float* probs, float* loss, int rows, int cols) {
+    int r = blockIdx.x;
+    const float* xr = x + r * cols;
+    float* pr = probs + r * cols;
+    float m = -INFINITY;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x) m = fmaxf(m, xr[c]);
+    m = block_reduce_max(m);
+    float sum = 0.0f;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x) {
+        float e = expf(xr[c] - m);
+        pr[c] = e;
+        sum += e;
+    }
+    sum = block_reduce_sum(sum);
+    float inv = 1.0f / sum;
+    for (int c = threadIdx.x; c < cols; c += blockDim.x) pr[c] *= inv;
+    if (threadIdx.x == 0) {
+        int t = target[r];
+        loss[r] = -logf(pr[t]);
+    }
+}
+void launch_cross_entropy_fwd_f32(const float* x, const int* target,
+                                  float* probs, float* loss, int rows, int cols, void* s) {
+    int blk = cols < 1024 ? cols : 1024;
+    if (blk < 32) blk = 32;
+    cross_entropy_fwd_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, target, probs, loss, rows, cols);
+}
+
+// dx[r,c] = scale * (probs[r,c] - [c==target]).  scale = upstream/rows.
+__global__ void cross_entropy_dx_k(const float* probs, const int* target,
+                                   float* dx, int rows, int cols, float scale) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= rows * cols) return;
+    int r = i / cols, c = i % cols;
+    float g = probs[i] - (c == target[r] ? 1.0f : 0.0f);
+    dx[i] = g * scale;
+}
+void launch_cross_entropy_dx_f32(const float* probs, const int* target,
+                                 float* dx, int rows, int cols, float scale, void* s) {
+    int n = rows * cols, blk = 256, grid = (n + blk - 1) / blk;
+    cross_entropy_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(probs, target, dx, rows, cols, scale);
+}
+
+} // extern "C"

docs/03-autograd-engine.md

@@ -0,0 +1,136 @@
+# Phase: Autograd Engine + Op Backward — Design Document
+
+## Goal
+
+在 T3 的 `Tensor`（matmul/transpose/finite-diff harness）之上，交付 **tape-based 动态 autograd 引擎** + 一个 tiny 现代 transformer 所需算子的**前向 kernel + 解析 backward**，每个 backward 都用 T3 的有限差分 harness 对拍通过。
+
+具体三件事：
+
+1. **autograd 引擎**：define-by-run 的反向自动微分。`Var` 包一个 `Tensor` + 可选 grad；每个 op 在 tape 上记一个节点（父节点 + backward 闭包）；`backward()` 按逆拓扑序遍历，把梯度推给父节点。**关键正确性点：梯度累加**——一个张量被多个 op 消费（扇出）时，各路梯度必须**求和**（T3 没有累加路径，在这里实现）。
+2. **算子节点**：`matmul` / `add` / `mul` / `add_bias`(broadcast) / `scale` / `rms_norm` / `silu` / `swiglu` / `rope` / `softmax` / `cross_entropy`，各带前向 CUDA kernel（需要时）+ 解析 backward。
+3. **Attention 用组合**：`attn = matmul(softmax(matmul(Q,Kᵀ)·scale), V)`。一旦 matmul/softmax/scale 是 autograd 节点，attention 的 backward 自动成立——**不写 fused attention backward kernel**，只加一个端到端 grad-check 测试。
+
+**明确不做**（留给 T5/T6）：组装 transformer / 训练 loop / 优化器 / embedding / KV-cache / GQA 重复。本 Phase 只到「算子 backward 逐个对拍通过」。
+
+## Module Layout
+
+```
+csrc/ops/nn.cu                         # 所有 T4 算子的 fwd+bwd kernel + launch_*（含 inlined warp/block reduce）
+
+crates/xtrain-cuda/
+├── build.rs                           # 新增 nn.cu
+└── src/ffi.rs                         # 新增 launch_* 声明（no_cuda 门控）
+
+crates/xtrain-tensor/
+├── src/dtype.rs                       # 新增 I32（cross-entropy target 用）
+└── src/tensor.rs                      # add/mul/add_bias/sum_rows/rms_norm(+bwd)/silu(+bwd)/
+                                       #   rope(+bwd)/softmax(+bwd)/cross_entropy(+bwd)（no_cuda 门控）
+
+crates/xtrain-autodiff/               # 引擎落在这里（已含 grad_check harness，自然归宿）
+├── build.rs                           # 新增：检测 nvcc → no_cuda cfg（cfg 不跨 crate 传播）
+├── src/
+│   ├── lib.rs                         # 导出 tape::Var + ops（no_cuda 门控）
+│   ├── finite_diff.rs                 # T3 既有 harness（不动）
+│   ├── tape.rs                        # Var / VarNode / backward / 梯度累加
+│   └── ops.rs                         # 各算子的 Var 节点构造器
+└── tests/autograd.rs                  # 每算子 grad-check + 扇出累加 + 组合 attention（#![cfg(not(no_cuda))]）
+```
+
+为什么引擎放 `xtrain-autodiff` 而不是新 crate：该 crate 本就是「自动微分」语义的归宿，且已持有 `grad_check`。前向 kernel/`Tensor` 方法仍按 T2/T3 约定落在 `xtrain-tensor`（与 `scale`/`matmul` 一致），引擎只是在其上叠 tape。
+
+## Key Design Decisions
+
+### Tape 设计：`Rc<RefCell<VarNode>>` + 逆拓扑遍历
+
+```rust
+pub struct VarNode {
+    value: Tensor,            // 前向输出
+    grad: Option<Tensor>,     // 反向累加的梯度
+    parents: Vec<Var>,        // 计算来源
+    backward: Option<BackwardFn>,  // None=叶子
+}
+pub struct Var(Rc<RefCell<VarNode>>);
+type BackwardFn = Box<dyn Fn(&Tensor, &[Var])>;
+```
+
+- `Var` clone 只是 bump `Rc`，**clone 共享同一节点**——这正是「扇出」的识别方式（同一 `Rc::as_ptr` 在多处出现）。
+- `backward()`：① post-order DFS 建拓扑序（按指针去重）；② 把 loss（必须是标量）的 grad 种子设为 1；③ 逆序遍历，每个节点把自己的 grad 传给父节点的 backward 闭包。
+- 闭包签名 `Fn(&grad, &parents)`：给本节点已累加的 grad 和父节点列表，闭包算出各父的梯度贡献并 `push_grad` 回去。前向需要 cache 的中间量（softmax 的 `y`、rms 的 `inv_rms`、ce 的 `probs`）用 `move` 闭包捕获。
+
+### 梯度累加（扇出求和）——本 Phase 的正确性核心
+
+`push_grad(parent, g)` 一律走 `accumulate`：
+
+```rust
+fn accumulate(&self, g: Tensor) {
+    match self.grad.take() {
+        None => self.grad = Some(g),                 // 首次
+        Some(prev) => self.grad = Some(prev.add(&g)),// 扇出：SUM
+    }
+}
+```
+
+任何节点（叶子或中间）都累加：中间节点需要完整 grad 才能继续链式；叶子的累加结果就是输出。一个张量喂多个消费者时，多路 `push_grad` 自动求和。`mul(&x, &x)` 这类「同一 `Var` 进同一节点两次」也正确：`parents=[x,x]`（同指针），两次 `push_grad` 累加，拓扑去重保证 x 只遍历一次但收齐两路。测试 `fanout_grad_accumulation` 专门验证：`y=x∘x + x∘x`，`dL/dx` 须 = 4x（四处 x 全部求和）。
+
+### 各算子 backward 数学
+
+记上游梯度为 `d`（=本节点输出的梯度）。
+
+| op | forward | backward |
+|----|---------|----------|
+| `matmul` | `C=A@B` | `dA=d@Bᵀ`, `dB=Aᵀ@d`（复用 T3 `matmul_backward`）|
+| `add` | `a+b` | `da=d`, `db=d` |
+| `mul` | `a∘b` | `da=d∘b`, `db=d∘a` |
+| `add_bias` | `x[r,c]+bias[c]` | `dx=d`, `dbias[c]=Σ_r d[r,c]`（沿广播维求和）|
+| `scale` | `x·α` | `dx=d·α` |
+| `silu` | `x·σ(x)` | `dx=d·(σ + x·σ·(1−σ))`, `σ=σ(x)` |
+| `swiglu` | `silu(g)∘u` | 由 `silu`+`mul` 组合自动得 |
+| `rope` | rotate_half 旋转 | RoPE 是正交变换，`dx` = 用**逆（转置）旋转**作用于 `d`（角度 +θ 的转置 ≡ −θ）|
+| `softmax` | row-wise safe softmax → `y` | Jacobian：`dx[r,c]=y[r,c]·(d[r,c] − Σ_c' d·y)` |
+| `cross_entropy` | mean NLL(softmax(x), tgt) | `dx = (probs − onehot)/rows`，再乘上游标量 grad |
+
+**RMSNorm**（`y[r,c]=x[r,c]·ir·γ[c]`, `ir=rsqrt(mean(x²)+eps)`）：
+设 `g[c]=d[r,c]·γ[c]`，`n=cols`，
+```
+dx[r,c]   = ir·g[c] − x[r,c]·ir³/n·Σ_c'(g[c']·x[r,c'])
+dγ[c]     = Σ_r d[r,c]·x[r,c]·ir
+```
+前向 cache 每行 `inv_rms[r]`，backward 直接复用，避免重算 reduce。
+
+**RoPE 反向推导**：前向是 2×2 旋转矩阵 `R(θ)`，正交 ⇒ `Rᵀ = R(−θ)`。故
+```
+dx[i]      = d[i]·cos + d[i+h]·sin
+dx[i+h]    = d[i+h]·cos − d[i]·sin
+```
+position=0 时旋转是恒等，backward 也恒等（sanity check）。
+
+**Softmax 反向推导**：`∂y_i/∂x_j = y_i(δ_ij − y_j)`，链式后
+`dx_i = Σ_j d_j·y_i(δ_ij − y_j) = y_i(d_i − Σ_j d_j y_j)`，即每行减去 `Σ(d∘y)` 后乘 `y`。
+
+**Cross-entropy 反向推导**：`L=−log softmax(x)[t]`，softmax+NLL 的经典结果 `∂L/∂x_c = softmax_c − [c=t]`；取 batch 平均 ⇒ 除以 rows。kernel 把 `scale=upstream/rows` 折进去。
+
+### Attention 用组合，不写 fused kernel
+
+```
+Kᵀ      = transpose(K)
+scores  = scale(matmul(Q, Kᵀ), 1/√d)   # [s,s]
+probs   = softmax(scores)
+out     = matmul(probs, V)              # [s,d]
+```
+
+每一步都是已有 autograd 节点，`backward()` 自动沿 matmul→softmax→scale→matmul 链回传，得到 `dQ/dK/dV`，无需手写 attention backward。测试 `attention_composed_bwd` 单头单 batch 端到端 grad-check Q/K/V 三者。（transpose 在测试里用一个临时 `Var::from_op` 节点包，因为引擎暂未把 transpose 列为 op——T5 若需要再补。）
+
+### kernel 实现要点
+
+- `nn.cu` 自带 inlined `warp/block_reduce_sum/max`（不引外部头文件，与现有 csrc/ 单文件风格一致）；block-reduce 末尾广播到全 block，便于 softmax/rms 的「标量广播」模式。
+- 每个 op 各自 `cudaDeviceSynchronize()`（T3 约定，无 stream）。
+- 全 F32、row-major、contiguous；cross-entropy target 用新增的 `DType::I32`。
+
+## 验证方法
+
+模板沿用 T3 `gemm.rs::run_bwd`：标量 loss `L = sum(W∘out)`，`W` 固定随机 ⇒ 上游 `dOut = W`；跑 op 的 `backward()` 拿 `.grad()`，对每个输入用 `grad_check` 与中心差分对拍。
+
+- **每算子**一个 grad-check（线性/双线性 op 用大 eps=1e-2、rel_tol=2e-2；非线性 op 用 eps=1e-3、rel_tol=3e-2、atol=1e-3 压住近零梯度）。
+- **扇出累加**：`fanout_grad_accumulation`，验证 `dL/dx=4x`。
+- **组合 attention**：`attention_composed_bwd`，端到端 grad-check `dQ/dK/dV`。
+- 全部 `#![cfg(not(no_cuda))]` 门控；本地只 `cargo check`/`fmt`，构建+实跑在 dash5（8× RTX 5090, sm_120），capture 每 op 的 pass + max rel-err。