model: tiny RoPE+RMSNorm+SwiGLU transformer + overfit test

New crate xtrain-model: a from-scratch decoder built entirely from the
autodiff op set.
- Config (tiny: dim=32, 2 layers, 2 heads, head_dim=16, ffn=64).
- TinyTransformer: embedding -> N x {pre-RMSNorm -> multi-head causal
  attention (RoPE, additive causal mask, per-head SDPA) -> residual;
  pre-RMSNorm -> SwiGLU MLP -> residual} -> final RMSNorm -> LM head.
  x@W weight convention (engine GEMM is plain A@B); dim=n_heads*head_dim.
- params()/zero_grad-able leaves for the optimizer; param_to_host export.
- overfit test: char-level bring-up (embedded text -> vocab -> shifted
  targets), minimal hand-written GD (p -= lr*grad) memorises one fixed
  batch -> loss ~0 + greedy argmax matches targets. End-to-end fwd+bwd
  correctness signal. Gated #![cfg(not(no_cuda))].

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

Cargo.lock

@@ -103,6 +103,15 @@ dependencies = [
  "cc",
 ]
 
+[[package]]
+name = "xtrain-model"
+version = "0.1.0"
+dependencies = [
+ "xtrain-autodiff",
+ "xtrain-cuda",
+ "xtrain-tensor",
+]
+
 [[package]]
 name = "xtrain-tensor"
 version = "0.1.0"

Cargo.toml

@@ -4,6 +4,7 @@ members = [
     "crates/xtrain-cuda",
     "crates/xtrain-tensor",
     "crates/xtrain-autodiff",
+    "crates/xtrain-model",
 ]
 
 [workspace.package]

crates/xtrain-autodiff/src/ops.rs

@@ -146,6 +146,126 @@ pub fn softmax(x: &Var) -> Var {
     )
 }
 
+/// Token embedding gather: `out[s,:] = table[ids[s], :]`. `table`:[vocab,dim]
+/// (a learnable [`Var`]), `ids`:[seq] I32 (a constant index, not a `Var`).
+/// Backward scatter-adds the upstream grad back into the table rows.
+pub fn embedding(table: &Var, ids: &Tensor) -> Var {
+    let out = table.value().embedding(ids);
+    let vocab = table.value().shape()[0];
+    let ids = ids.clone();
+    Var::from_op(
+        out,
+        vec![table.clone()],
+        Box::new(move |dout, parents| {
+            let dtable = Tensor::embedding_backward(dout, &ids, vocab);
+            Var::push_grad(&parents[0], dtable);
+        }),
+    )
+}
+
+/// Reshape (contiguous, metadata-only). Backward reshapes the grad back to the
+/// input shape. Used for the multi-head layout swap `[seq, h*hd] <-> [seq, h, hd]`.
+pub fn reshape(x: &Var, new_shape: &[usize]) -> Var {
+    let in_shape: Vec<usize> = x.value().shape().to_vec();
+    let out = x.value().reshape(new_shape);
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            Var::push_grad(&parents[0], d.reshape(&in_shape));
+        }),
+    )
+}
+
+/// 3D axis-(0,1) transpose `[a,b,c] -> [b,a,c]`. Self-inverse structure: the
+/// backward is the same transpose applied to the grad.
+pub fn transpose_3d01(x: &Var) -> Var {
+    let out = x.value().transpose_3d01();
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(|d, parents| {
+            Var::push_grad(&parents[0], d.transpose_3d01());
+        }),
+    )
+}
+
+/// 2D transpose `[r,c] -> [c,r]` as an autograd node (backward transposes the
+/// grad back). Used for `Kᵀ` in attention scores.
+pub fn transpose_2d(x: &Var) -> Var {
+    let out = x.value().transpose_2d();
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(|d, parents| {
+            Var::push_grad(&parents[0], d.transpose_2d());
+        }),
+    )
+}
+
+/// Split a `[heads, seq, head_dim]` tensor into one `[seq, head_dim]` [`Var`] per
+/// head. Each head block is contiguous in this layout, so the forward copies the
+/// head block into its own contiguous tensor; the backward scatters each head's
+/// grad back into a zero `[heads, seq, head_dim]` grad (the engine then SUMs the
+/// `heads` contributions on the shared parent — fan-out).
+pub fn split_heads(x: &Var) -> Vec<Var> {
+    let v = x.value();
+    assert_eq!(v.ndim(), 3, "split_heads requires [heads,seq,head_dim]");
+    let (heads, seq, hd) = (v.shape()[0], v.shape()[1], v.shape()[2]);
+    let dev = v.device();
+    let flat_host = v.to_device(xtrain_tensor::Device::Cpu);
+    let flat = flat_host.as_slice::<f32>();
+    (0..heads)
+        .map(|h| {
+            let base = h * seq * hd;
+            let block = Tensor::from_slice(&flat[base..base + seq * hd], &[seq, hd]).to_device(dev);
+            Var::from_op(
+                block,
+                vec![x.clone()],
+                Box::new(move |d, parents| {
+                    let mut host = vec![0.0f32; heads * seq * hd];
+                    let dvals = d.to_device(xtrain_tensor::Device::Cpu);
+                    let base = h * seq * hd;
+                    host[base..base + seq * hd].copy_from_slice(dvals.as_slice::<f32>());
+                    let g = Tensor::from_slice(&host, &[heads, seq, hd]).to_device(dev);
+                    Var::push_grad(&parents[0], g);
+                }),
+            )
+        })
+        .collect()
+}
+
+/// Inverse of [`split_heads`]: stack per-head `[seq, head_dim]` outputs into a
+/// `[heads, seq, head_dim]` tensor. Backward hands each head its own slice of the
+/// grad.
+pub fn merge_heads(heads_v: &[Var]) -> Var {
+    let heads = heads_v.len();
+    let v0 = heads_v[0].value();
+    let (seq, hd) = (v0.shape()[0], v0.shape()[1]);
+    let dev = v0.device();
+    let mut host = vec![0.0f32; heads * seq * hd];
+    for (h, hv) in heads_v.iter().enumerate() {
+        let block = hv.value().to_device(xtrain_tensor::Device::Cpu);
+        let base = h * seq * hd;
+        host[base..base + seq * hd].copy_from_slice(block.as_slice::<f32>());
+    }
+    let out = Tensor::from_slice(&host, &[heads, seq, hd]).to_device(dev);
+    Var::from_op(
+        out,
+        heads_v.to_vec(),
+        Box::new(move |d, parents| {
+            let dhost = d.to_device(xtrain_tensor::Device::Cpu);
+            let dflat = dhost.as_slice::<f32>();
+            for (h, parent) in parents.iter().enumerate() {
+                let base = h * seq * hd;
+                let g =
+                    Tensor::from_slice(&dflat[base..base + seq * hd], &[seq, hd]).to_device(dev);
+                Var::push_grad(parent, g);
+            }
+        }),
+    )
+}
+
 /// 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.

crates/xtrain-autodiff/src/tape.rs

@@ -68,6 +68,19 @@ impl Var {
         self.0.borrow().grad.clone()
     }
 
+    /// Clear the accumulated gradient. Call on every parameter between training
+    /// steps so the next `backward` accumulates from zero (grads SUM otherwise).
+    pub fn zero_grad(&self) {
+        self.0.borrow_mut().grad = None;
+    }
+
+    /// Overwrite this node's value tensor in place. Used by the optimizer to
+    /// apply a parameter update (`p ← p − lr·grad`) while keeping the leaf's
+    /// identity stable across steps.
+    pub fn set_value(&self, value: Tensor) {
+        self.0.borrow_mut().value = value;
+    }
+
     /// Pointer identity, used to dedup nodes during the topological sort.
     fn id(&self) -> *const RefCell<VarNode> {
         Rc::as_ptr(&self.0)

crates/xtrain-autodiff/tests/structural.rs

@@ -0,0 +1,220 @@
+// GPU grad-checks for the Phase T5 structural ops added on top of the T4 set:
+// embedding (gather fwd / scatter-add bwd), reshape, transpose_3d01,
+// transpose_2d, and split/merge_heads. Same harness as autograd.rs:
+// L = sum(W ∘ out), W fixed random ⇒ upstream dOut = W; run backward(), then
+// grad-check each leaf's .grad() against central finite differences.
+//
+// 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};
+
+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))
+}
+
+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()
+}
+
+// Structural ops are exactly linear in their input → a large eps just sharpens
+// f32 resolution (same as add/mul/transpose in autograd.rs).
+fn cfg_linear() -> GradCheckConfig {
+    GradCheckConfig {
+        eps: 1e-2,
+        rel_tol: 2e-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:?}");
+}
+
+// L = sum(W ∘ out): a constant-W leaf mul + sum-to-scalar reduction.
+fn scalar_loss(out: &Var, w: &[f32]) -> Var {
+    let wt = Var::leaf(cuda(w, out.value().shape()));
+    sum_all(&ops::mul(out, &wt))
+}
+
+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| {
+            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);
+        }),
+    )
+}
+
+// ---- embedding (gather fwd / scatter-add bwd) ----
+// Includes a repeated id so the atomic scatter-add accumulation is exercised.
+#[test]
+fn embedding_bwd() {
+    require_gpu();
+    let (vocab, dim) = (5, 7);
+    let ids_host: Vec<i32> = vec![0, 3, 1, 3, 2, 0]; // 0 and 3 repeat
+    let seq = ids_host.len();
+    let table_h = fill(vocab * dim, 201);
+    let w = fill(seq * dim, 202);
+
+    let ids = Tensor::from_slice(&ids_host, &[seq]).to_device(Device::Cuda(0));
+    let table = Var::leaf(cuda(&table_h, &[vocab, dim]));
+    let out = ops::embedding(&table, &ids);
+    scalar_loss(&out, &w).backward();
+
+    let dtable = table.grad().unwrap().to_device(Device::Cpu);
+    let idf = ids_host.clone();
+    let wf = w.clone();
+    let lt = move |v: &[f32], s: &[usize]| {
+        let ids = Tensor::from_slice(&idf, &[seq]).to_device(Device::Cuda(0));
+        weighted_sum(&cuda(v, s).embedding(&ids), &wf)
+    };
+    report(
+        "embedding dTable",
+        &grad_check(
+            &table_h,
+            &[vocab, dim],
+            &lt,
+            dtable.as_slice::<f32>(),
+            cfg_linear(),
+        ),
+    );
+}
+
+// ---- reshape ----
+#[test]
+fn reshape_bwd() {
+    require_gpu();
+    let (rows, cols) = (6, 8);
+    let x_h = fill(rows * cols, 211);
+    let w = fill(rows * cols, 212);
+
+    let x = Var::leaf(cuda(&x_h, &[rows, cols]));
+    let out = ops::reshape(&x, &[rows * 2, cols / 2]);
+    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).reshape(&[rows * 2, cols / 2]), &wf);
+    report(
+        "reshape dX",
+        &grad_check(&x_h, &[rows, cols], &lx, dx.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- transpose_3d01 ([a,b,c] -> [b,a,c]) ----
+#[test]
+fn transpose_3d01_bwd() {
+    require_gpu();
+    let (a, b, c) = (3, 4, 5);
+    let x_h = fill(a * b * c, 221);
+    let w = fill(a * b * c, 222);
+
+    let x = Var::leaf(cuda(&x_h, &[a, b, c]));
+    let out = ops::transpose_3d01(&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).transpose_3d01(), &wf);
+    report(
+        "transpose_3d01 dX",
+        &grad_check(&x_h, &[a, b, c], &lx, dx.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- transpose_2d ----
+#[test]
+fn transpose_2d_bwd() {
+    require_gpu();
+    let (r, c) = (5, 7);
+    let x_h = fill(r * c, 231);
+    let w = fill(r * c, 232);
+
+    let x = Var::leaf(cuda(&x_h, &[r, c]));
+    let out = ops::transpose_2d(&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).transpose_2d(), &wf);
+    report(
+        "transpose_2d dX",
+        &grad_check(&x_h, &[r, c], &lx, dx.as_slice::<f32>(), cfg_linear()),
+    );
+}
+
+// ---- split_heads + merge_heads round-trip (identity reshuffle of [nh,seq,hd]) ----
+// out = merge_heads(split_heads(x)) must equal x, and its grad must be dOut=W
+// reshuffled identically — i.e. dx grad-checks against the identity composition.
+#[test]
+fn split_merge_heads_bwd() {
+    require_gpu();
+    let (nh, seq, hd) = (3, 4, 5);
+    let x_h = fill(nh * seq * hd, 241);
+    let w = fill(nh * seq * hd, 242);
+
+    let x = Var::leaf(cuda(&x_h, &[nh, seq, hd]));
+    let heads = ops::split_heads(&x);
+    let out = ops::merge_heads(&heads); // back to [nh,seq,hd]
+    scalar_loss(&out, &w).backward();
+
+    let dx = x.grad().unwrap().to_device(Device::Cpu);
+    // forward is identity, so grad-check the identity map.
+    let wf = w.clone();
+    let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s), &wf);
+    report(
+        "split/merge_heads dX",
+        &grad_check(
+            &x_h,
+            &[nh, seq, hd],
+            &lx,
+            dx.as_slice::<f32>(),
+            cfg_linear(),
+        ),
+    );
+}

crates/xtrain-cuda/build.rs

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

crates/xtrain-cuda/src/ffi.rs

@@ -177,6 +177,41 @@ unsafe extern "C" {
     );
 }
 
+// Structural ops for the tiny transformer (csrc/ops/model.cu): token embedding
+// (gather fwd / scatter-add bwd) and a 3D axis-(0,1) transpose for the multi-head
+// attention layout. F32 values, I32 ids, row-major contiguous.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    // Embedding: out[s,:] = table[ids[s], :]. table:[vocab,dim], ids:[seq] (I32).
+    pub fn launch_embedding_fwd_f32(
+        table: *const f32,
+        ids: *const i32,
+        out: *mut f32,
+        seq: i32,
+        dim: i32,
+        s: CudaStream,
+    );
+    // Scatter-add: dtable[ids[s],:] += dout[s,:] (dtable pre-zeroed; atomic).
+    pub fn launch_embedding_bwd_f32(
+        dout: *const f32,
+        ids: *const i32,
+        dtable: *mut f32,
+        seq: i32,
+        dim: i32,
+        s: CudaStream,
+    );
+
+    // 3D axis-(0,1) transpose: in:[a,b,c] -> out:[b,a,c]. out[j,i,k]=in[i,j,k].
+    pub fn launch_transpose_3d01_f32(
+        input: *const f32,
+        out: *mut f32,
+        a: i32,
+        b: i32,
+        c: i32,
+        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-model/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "xtrain-model"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xtrain-tensor = { path = "../xtrain-tensor" }
+xtrain-autodiff = { path = "../xtrain-autodiff" }
+
+[dev-dependencies]
+# Acceptance tests drive the GPU (device selection) directly.
+xtrain-cuda = { path = "../xtrain-cuda" }

crates/xtrain-model/build.rs

@@ -0,0 +1,26 @@
+use std::env;
+use std::path::Path;
+use std::process::Command;
+
+// Same per-crate convention as the other crates: this crate's tiny-transformer
+// forward/backward calls GPU ops (via xtrain-autodiff / xtrain-tensor), so it
+// gates GPU code + tests behind `not(no_cuda)`. cfg does not propagate across
+// crates, so each crate re-detects nvcc. No CUDA is compiled here.
+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-model/src/config.rs

@@ -0,0 +1,54 @@
+//! Tiny-transformer hyperparameters. Host-only (no GPU), always compiled.
+
+/// Architecture config for [`crate::TinyTransformer`]. Keep it tiny — T5 is a
+/// correctness bring-up, not a real training run.
+#[derive(Debug, Clone, Copy)]
+pub struct Config {
+    /// Vocabulary size (char-level in the bring-up).
+    pub vocab: usize,
+    /// Model / residual width. Must equal `n_heads * head_dim`.
+    pub dim: usize,
+    /// Number of decoder blocks.
+    pub n_layers: usize,
+    /// Number of attention heads.
+    pub n_heads: usize,
+    /// Per-head dimension (`dim / n_heads`).
+    pub head_dim: usize,
+    /// SwiGLU hidden width (gate/up project to this, down projects back).
+    pub ffn_hidden: usize,
+    /// RMSNorm epsilon.
+    pub eps: f32,
+    /// RoPE base frequency (theta).
+    pub rope_theta: f32,
+}
+
+impl Config {
+    /// A minimal config used by the bring-up / overfit test.
+    pub fn tiny() -> Self {
+        let n_heads = 2;
+        let head_dim = 16;
+        Config {
+            vocab: 0, // set by the caller from the char vocab
+            dim: n_heads * head_dim,
+            n_layers: 2,
+            n_heads,
+            head_dim,
+            ffn_hidden: 64,
+            eps: 1e-5,
+            rope_theta: 10000.0,
+        }
+    }
+
+    /// Total learnable parameter count (for logging / sanity).
+    pub fn num_params(&self) -> usize {
+        let per_layer = 2 * self.dim                       // 2 rmsnorm gammas
+            + 3 * self.dim * self.dim                       // q/k/v proj
+            + self.dim * self.dim                           // out proj
+            + 2 * self.dim * self.ffn_hidden                // gate/up proj
+            + self.ffn_hidden * self.dim; // down proj
+        self.vocab * self.dim          // embedding
+            + self.n_layers * per_layer
+            + self.dim                 // final norm
+            + self.dim * self.vocab // lm head
+    }
+}

crates/xtrain-model/src/lib.rs

@@ -0,0 +1,26 @@
+//! Tiny modern-architecture transformer (Phase T5).
+//!
+//! A from-scratch decoder built entirely from the [`xtrain_autodiff`] op set:
+//! token embedding → `n_layers` × {pre-RMSNorm → multi-head causal attention
+//! (RoPE) → residual; pre-RMSNorm → SwiGLU MLP → residual} → final RMSNorm →
+//! LM-head matmul. The forward builds an autograd graph; calling `.backward()`
+//! on the cross-entropy loss fills every parameter's `.grad()`.
+//!
+//! Conventions (matching the engine, not HuggingFace):
+//! - Linear weights are `[in, out]` and applied as `x @ W` (no transpose), since
+//!   the engine's GEMM is plain `A @ B`.
+//! - `dim == n_heads * head_dim` (no separate attention projection size).
+//! - RoPE position = token row index (the kernel's built-in convention).
+//! - Causal masking is an additive `[seq,seq]` constant (−1e9 above the diagonal)
+//!   added to the attention scores before softmax.
+//!
+//! Everything GPU-facing is gated behind `not(no_cuda)`; on a GPU-less host the
+//! crate still `cargo check`s (only [`Config`] is visible there).
+
+mod config;
+pub use config::Config;
+
+#[cfg(not(no_cuda))]
+mod model;
+#[cfg(not(no_cuda))]
+pub use model::{TinyTransformer, ids_tensor, param_to_host};

crates/xtrain-model/src/model.rs

@@ -0,0 +1,205 @@
+//! The tiny transformer forward graph + parameter container (Phase T5).
+
+#![cfg(not(no_cuda))]
+
+use crate::config::Config;
+use xtrain_autodiff::ops;
+use xtrain_autodiff::tape::Var;
+use xtrain_tensor::{Device, Tensor};
+
+/// One decoder block's learnable tensors.
+struct Block {
+    attn_norm: Var, // [dim]
+    wq: Var,        // [dim, dim]
+    wk: Var,        // [dim, dim]
+    wv: Var,        // [dim, dim]
+    wo: Var,        // [dim, dim]
+    ffn_norm: Var,  // [dim]
+    w_gate: Var,    // [dim, ffn_hidden]
+    w_up: Var,      // [dim, ffn_hidden]
+    w_down: Var,    // [ffn_hidden, dim]
+}
+
+/// A tiny RoPE+RMSNorm+SwiGLU decoder. Holds every parameter as a leaf [`Var`];
+/// `forward` builds an autograd graph over them.
+pub struct TinyTransformer {
+    cfg: Config,
+    embed: Var, // [vocab, dim]
+    blocks: Vec<Block>,
+    final_norm: Var, // [dim]
+    lm_head: Var,    // [dim, vocab]
+    device: Device,
+}
+
+impl TinyTransformer {
+    /// Build a model with parameters initialised from `init(shape) -> host data`.
+    /// The caller controls initialisation (deterministic for tests / PyTorch
+    /// parity). `init` receives the logical shape and returns row-major data.
+    pub fn new(cfg: Config, device: Device, mut init: impl FnMut(&[usize]) -> Vec<f32>) -> Self {
+        let leaf = |data: Vec<f32>, shape: &[usize]| -> Var {
+            Var::leaf(Tensor::from_slice(&data, shape).to_device(device))
+        };
+        let mut mk = |shape: &[usize]| -> Var {
+            let data = init(shape);
+            assert_eq!(data.len(), shape.iter().product::<usize>(), "init size");
+            leaf(data, shape)
+        };
+
+        let embed = mk(&[cfg.vocab, cfg.dim]);
+        let blocks = (0..cfg.n_layers)
+            .map(|_| Block {
+                attn_norm: mk(&[cfg.dim]),
+                wq: mk(&[cfg.dim, cfg.dim]),
+                wk: mk(&[cfg.dim, cfg.dim]),
+                wv: mk(&[cfg.dim, cfg.dim]),
+                wo: mk(&[cfg.dim, cfg.dim]),
+                ffn_norm: mk(&[cfg.dim]),
+                w_gate: mk(&[cfg.dim, cfg.ffn_hidden]),
+                w_up: mk(&[cfg.dim, cfg.ffn_hidden]),
+                w_down: mk(&[cfg.ffn_hidden, cfg.dim]),
+            })
+            .collect();
+        let final_norm = mk(&[cfg.dim]);
+        let lm_head = mk(&[cfg.dim, cfg.vocab]);
+
+        Self {
+            cfg,
+            embed,
+            blocks,
+            final_norm,
+            lm_head,
+            device,
+        }
+    }
+
+    pub fn config(&self) -> &Config {
+        &self.cfg
+    }
+
+    /// All learnable parameters, in a stable order. The optimizer (a hand-written
+    /// GD step in T5, AdamW in T6) iterates this; each holds its `.grad()` after
+    /// `backward()`.
+    pub fn params(&self) -> Vec<Var> {
+        let mut ps = vec![self.embed.clone()];
+        for b in &self.blocks {
+            ps.extend([
+                b.attn_norm.clone(),
+                b.wq.clone(),
+                b.wk.clone(),
+                b.wv.clone(),
+                b.wo.clone(),
+                b.ffn_norm.clone(),
+                b.w_gate.clone(),
+                b.w_up.clone(),
+                b.w_down.clone(),
+            ]);
+        }
+        ps.push(self.final_norm.clone());
+        ps.push(self.lm_head.clone());
+        ps
+    }
+
+    /// Forward over a single sequence of token `ids` (`[seq]` I32 on this
+    /// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`.
+    pub fn forward(&self, ids: &Tensor) -> Var {
+        let seq = ids.shape()[0];
+        let mask = self.causal_mask(seq);
+
+        let mut h = ops::embedding(&self.embed, ids); // [seq, dim]
+        for b in &self.blocks {
+            // --- Attention sub-block (pre-norm + residual) ---
+            let normed = ops::rms_norm(&h, &b.attn_norm, self.cfg.eps);
+            let attn = self.attention(b, &normed, &mask, seq);
+            h = ops::add(&h, &attn);
+
+            // --- MLP sub-block (pre-norm + residual) ---
+            let normed = ops::rms_norm(&h, &b.ffn_norm, self.cfg.eps);
+            let mlp = self.swiglu_mlp(b, &normed);
+            h = ops::add(&h, &mlp);
+        }
+
+        let h = ops::rms_norm(&h, &self.final_norm, self.cfg.eps);
+        ops::matmul(&h, &self.lm_head) // [seq, vocab]
+    }
+
+    /// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
+    pub fn loss(&self, ids: &Tensor, targets: &Tensor) -> Var {
+        let logits = self.forward(ids);
+        ops::cross_entropy(&logits, targets)
+    }
+
+    /// Multi-head causal self-attention. `x`:[seq,dim] (already normed).
+    fn attention(&self, b: &Block, x: &Var, mask: &Var, seq: usize) -> Var {
+        let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
+        let scale = 1.0 / (hd as f32).sqrt();
+
+        // Project, then lay out as per-head [seq, head_dim] tensors.
+        //   [seq,dim] @ [dim,dim] = [seq,dim]
+        //   reshape   [seq, nh, hd]
+        //   rope      (kernel expects exactly [tokens, heads, head_dim])
+        //   transpose [nh, seq, hd]  → split into nh × [seq, hd]
+        let to_heads = |proj: Var, rope: bool| -> Vec<Var> {
+            let r = ops::reshape(&proj, &[seq, nh, hd]);
+            let r = if rope {
+                ops::rope(&r, self.cfg.rope_theta)
+            } else {
+                r
+            };
+            let t = ops::transpose_3d01(&r); // [nh, seq, hd]
+            ops::split_heads(&t)
+        };
+
+        let q = to_heads(ops::matmul(x, &b.wq), true);
+        let k = to_heads(ops::matmul(x, &b.wk), true);
+        let v = to_heads(ops::matmul(x, &b.wv), false);
+
+        // Per-head scaled-dot-product attention with causal mask.
+        let heads_out: Vec<Var> = (0..nh)
+            .map(|i| {
+                let kt = ops::transpose_2d(&k[i]); // [hd, seq]
+                let scores = ops::scale(&ops::matmul(&q[i], &kt), scale); // [seq,seq]
+                let scores = ops::add(&scores, mask); // causal
+                let probs = ops::softmax(&scores);
+                ops::matmul(&probs, &v[i]) // [seq, hd]
+            })
+            .collect();
+
+        // Stack heads back: nh × [seq,hd] → [nh,seq,hd] → [seq,nh,hd] → [seq,dim].
+        let merged = ops::merge_heads(&heads_out); // [nh, seq, hd]
+        let t = ops::transpose_3d01(&merged); // [seq, nh, hd]
+        let concat = ops::reshape(&t, &[seq, nh * hd]); // [seq, dim]
+        ops::matmul(&concat, &b.wo) // out projection
+    }
+
+    /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[seq,dim].
+    fn swiglu_mlp(&self, b: &Block, x: &Var) -> Var {
+        let gate = ops::matmul(x, &b.w_gate); // [seq, ffn_hidden]
+        let up = ops::matmul(x, &b.w_up); // [seq, ffn_hidden]
+        let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
+        ops::matmul(&act, &b.w_down) // [seq, dim]
+    }
+
+    /// Additive causal mask `[seq,seq]`: 0 on/below the diagonal, −1e9 above it
+    /// (so softmax zeros out future positions). A constant leaf (no grad needed,
+    /// but harmless if it accumulates one — it has no consumers downstream of x).
+    fn causal_mask(&self, seq: usize) -> Var {
+        let mut m = vec![0.0f32; seq * seq];
+        for i in 0..seq {
+            for j in (i + 1)..seq {
+                m[i * seq + j] = -1.0e9;
+            }
+        }
+        Var::leaf(Tensor::from_slice(&m, &[seq, seq]).to_device(self.device))
+    }
+}
+
+/// Materialise a parameter's value back to a host `Vec<f32>` (for the GD step
+/// and PyTorch parity export).
+pub fn param_to_host(v: &Var) -> Vec<f32> {
+    v.value().to_device(Device::Cpu).as_slice::<f32>().to_vec()
+}
+
+/// Build an I32 id tensor on `device` from token ids.
+pub fn ids_tensor(ids: &[i32], device: Device) -> Tensor {
+    Tensor::from_slice(ids, &[ids.len()]).to_device(device)
+}

crates/xtrain-model/tests/overfit.rs

@@ -0,0 +1,133 @@
+// End-to-end acceptance for the Phase T5 tiny transformer: overfit one fixed
+// char-level batch with a hand-written gradient-descent step and assert the loss
+// collapses toward 0. This is THE signal that the whole fwd+bwd graph (embedding,
+// RMSNorm, RoPE, multi-head attention, SwiGLU, LM head, cross-entropy) is wired
+// correctly — a single buggy backward would stall the loss.
+//
+// The optimizer here is deliberately minimal (`p ← p − lr·grad`); AdamW / LR
+// schedule / real data are T6. Gated behind `not(no_cuda)` (runs on dash5).
+#![cfg(not(no_cuda))]
+
+use xtrain_autodiff::tape::Var;
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, ids_tensor};
+use xtrain_tensor::Device;
+
+// Deterministic LCG fill in [-scale, scale).
+fn fill(n: usize, seed: u64, scale: f32) -> 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) * 2.0 * scale
+        })
+        .collect()
+}
+
+fn require_gpu() {
+    assert!(
+        device::device_count().expect("device count") > 0,
+        "no CUDA device"
+    );
+    device::set_device(0).unwrap();
+}
+
+// One GD step over every parameter: p ← p − lr·grad, then zero the grad.
+fn gd_step(params: &[Var], lr: f32) {
+    for p in params {
+        if let Some(g) = p.grad() {
+            let updated = p.value().add(&g.scale(-lr));
+            p.set_value(updated);
+        }
+        p.zero_grad();
+    }
+}
+
+#[test]
+fn overfit_tiny_batch() {
+    require_gpu();
+    let device = Device::Cuda(0);
+
+    // --- Char-level bring-up: tiny embedded text → vocab → (input, target). ---
+    let text = "hello tiny transformer world";
+    let mut vocab_chars: Vec<char> = text.chars().collect();
+    vocab_chars.sort_unstable();
+    vocab_chars.dedup();
+    let vocab = vocab_chars.len();
+    let stoi = |c: char| vocab_chars.iter().position(|&x| x == c).unwrap() as i32;
+
+    let tokens: Vec<i32> = text.chars().map(stoi).collect();
+    // Next-token prediction: input = tokens[..n-1], target = tokens[1..].
+    let input: Vec<i32> = tokens[..tokens.len() - 1].to_vec();
+    let target: Vec<i32> = tokens[1..].to_vec();
+    let ids = ids_tensor(&input, device);
+    let targets = ids_tensor(&target, device);
+
+    // --- Tiny model with small-scale deterministic init. ---
+    let mut cfg = Config::tiny();
+    cfg.vocab = vocab;
+    let mut seed = 1u64;
+    let model = TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        // RMSNorm gammas ([dim]) init to ~1; everything else small random.
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.08)
+        }
+    });
+    let params = model.params();
+    println!(
+        "overfit: vocab={vocab} seq={} params={}",
+        input.len(),
+        cfg.num_params()
+    );
+
+    let read_loss = |l: &Var| -> f32 { l.value().to_device(Device::Cpu).as_slice::<f32>()[0] };
+
+    let lr = 0.3f32;
+    let steps = 200;
+    let start = read_loss(&model.loss(&ids, &targets));
+    let mut last = start;
+    for step in 0..steps {
+        let loss = model.loss(&ids, &targets);
+        last = read_loss(&loss);
+        if step % 20 == 0 || step == steps - 1 {
+            println!("step {step:3}: loss = {last:.6}");
+        }
+        loss.backward();
+        gd_step(&params, lr);
+    }
+
+    println!("overfit: start loss = {start:.6} → final loss = {last:.6} ({steps} steps)");
+    // A correct fwd+bwd memorises this tiny fixed batch: loss → ~0.
+    assert!(
+        last < 0.05,
+        "overfit failed to drive loss to ~0: start {start:.4} final {last:.4}"
+    );
+    assert!(last < start, "loss did not decrease");
+
+    // Sanity: greedy argmax should reproduce the target sequence after overfit.
+    let logits = model.forward(&ids).value().to_device(Device::Cpu);
+    let lg = logits.as_slice::<f32>();
+    let mut correct = 0;
+    for (r, &t) in target.iter().enumerate() {
+        let row = &lg[r * vocab..(r + 1) * vocab];
+        let argmax = row
+            .iter()
+            .enumerate()
+            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+            .unwrap()
+            .0 as i32;
+        if argmax == t {
+            correct += 1;
+        }
+    }
+    println!("overfit: greedy match {correct}/{}", target.len());
+    assert_eq!(correct, target.len() as i32, "did not memorise the batch");
+}

crates/xtrain-tensor/src/tensor.rs

@@ -563,6 +563,98 @@ impl Tensor {
         dx
     }
 
+    // --- Structural / model ops (the T5 kernels) ---
+
+    /// Reshape to `new_shape` (must keep `numel`). Pure metadata change on a
+    /// contiguous tensor — no data movement, shares the same storage. The
+    /// multi-head layout `[seq, n_heads*head_dim] <-> [seq, n_heads, head_dim]`
+    /// is exactly this.
+    pub fn reshape(&self, new_shape: &[usize]) -> Self {
+        assert!(self.is_contiguous(), "reshape requires a contiguous tensor");
+        assert_eq!(
+            shape::num_elements(new_shape),
+            self.numel(),
+            "reshape numel mismatch: {:?} -> {:?}",
+            self.shape.as_slice(),
+            new_shape
+        );
+        Self {
+            storage: self.storage.clone(),
+            shape: Dims::from_slice(new_shape),
+            strides: shape::contiguous_strides(new_shape),
+            offset: self.offset,
+            dtype: self.dtype,
+        }
+    }
+
+    /// Embedding gather: `out[s,:] = self[ids[s], :]`. `self`:[vocab,dim] table,
+    /// `ids`:[seq] I32 → out:[seq,dim].
+    #[cfg(not(no_cuda))]
+    pub fn embedding(&self, ids: &Tensor) -> Self {
+        assert_eq!(self.dtype, DType::F32, "embedding table must be F32");
+        assert_eq!(self.ndim(), 2, "embedding table must be [vocab,dim]");
+        assert_eq!(ids.dtype, DType::I32, "embedding ids must be I32");
+        assert_eq!(ids.ndim(), 1, "embedding ids must be 1D");
+        let (seq, dim) = (ids.shape[0], self.shape[1]);
+        let out = Tensor::zeros(&[seq, dim], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_embedding_fwd_f32(
+                self.data_ptr() as *const f32,
+                ids.data_ptr() as *const i32,
+                out.data_ptr() as *mut f32,
+                seq as i32,
+                dim as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("embedding sync failed");
+        out
+    }
+
+    /// Embedding backward (scatter-add): `dtable[ids[s],:] += dout[s,:]`, where
+    /// `dout`:[seq,dim], `ids`:[seq] I32. `vocab` sizes the output table.
+    #[cfg(not(no_cuda))]
+    pub fn embedding_backward(dout: &Tensor, ids: &Tensor, vocab: usize) -> Self {
+        let (seq, dim) = (dout.shape[0], dout.shape[1]);
+        let dtable = Tensor::zeros(&[vocab, dim], DType::F32, dout.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_embedding_bwd_f32(
+                dout.data_ptr() as *const f32,
+                ids.data_ptr() as *const i32,
+                dtable.data_ptr() as *mut f32,
+                seq as i32,
+                dim as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("embedding_backward sync failed");
+        dtable
+    }
+
+    /// 3D axis-(0,1) transpose: `self`:[a,b,c] → [b,a,c], `out[j,i,k]=self[i,j,k]`.
+    /// Lays out multi-head attention (`[seq,heads,hd] <-> [heads,seq,hd]`). Its
+    /// own backward is the same op (swap a,b).
+    #[cfg(not(no_cuda))]
+    pub fn transpose_3d01(&self) -> Self {
+        assert_eq!(self.dtype, DType::F32, "transpose_3d01 only supports F32");
+        assert_eq!(self.ndim(), 3, "transpose_3d01 requires a 3D tensor");
+        assert!(self.is_contiguous(), "transpose_3d01 requires contiguous");
+        let (a, b, c) = (self.shape[0], self.shape[1], self.shape[2]);
+        let out = Tensor::zeros(&[b, a, c], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_transpose_3d01_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                a as i32,
+                b as i32,
+                c as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("transpose_3d01 sync failed");
+        out
+    }
+
     // Shared validation for same-shape binary elementwise ops.
     #[cfg(not(no_cuda))]
     fn check_binary(&self, other: &Tensor, op: &str) {

csrc/ops/model.cu

@@ -0,0 +1,66 @@
+// Structural ops the tiny transformer (Phase T5) needs on top of the T4 op set:
+// token embedding (gather forward / scatter-add backward) and a 3D axis-(0,1)
+// transpose used to lay out multi-head attention ([seq,heads,hd] <-> [heads,seq,hd]).
+//
+// reshape is a pure metadata change (no data movement) and so has no kernel — it
+// lives entirely in the Rust Tensor layer. All kernels here are F32 row-major
+// contiguous; ids are I32. Each launcher matches the existing csrc/ style.
+
+extern "C" {
+
+// =====================================================================
+// Embedding: gather rows of a table by integer ids.
+//   table:[vocab, dim], ids:[seq] (I32)  ->  out[s,:] = table[ids[s], :]
+// Backward (scatter-add): dtable[ids[s], :] += dout[s, :]. Multiple positions
+// may map to the same id, so the accumulation must be atomic.
+// =====================================================================
+
+__global__ void embedding_fwd_k(const float* table, const int* ids, float* out,
+                                int seq, int dim) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x; // over seq*dim
+    if (i >= seq * dim) return;
+    int s = i / dim, c = i % dim;
+    out[i] = table[ids[s] * dim + c];
+}
+void launch_embedding_fwd_f32(const float* table, const int* ids, float* out,
+                              int seq, int dim, void* s) {
+    int n = seq * dim, blk = 256, grid = (n + blk - 1) / blk;
+    embedding_fwd_k<<<grid, blk, 0, (cudaStream_t)s>>>(table, ids, out, seq, dim);
+}
+
+// dtable is assumed pre-zeroed (Tensor::zeros). Scatter-add with atomics so
+// repeated ids accumulate correctly.
+__global__ void embedding_bwd_k(const float* dout, const int* ids, float* dtable,
+                                int seq, int dim) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x; // over seq*dim
+    if (i >= seq * dim) return;
+    int s = i / dim, c = i % dim;
+    atomicAdd(&dtable[ids[s] * dim + c], dout[i]);
+}
+void launch_embedding_bwd_f32(const float* dout, const int* ids, float* dtable,
+                              int seq, int dim, void* s) {
+    int n = seq * dim, blk = 256, grid = (n + blk - 1) / blk;
+    embedding_bwd_k<<<grid, blk, 0, (cudaStream_t)s>>>(dout, ids, dtable, seq, dim);
+}
+
+// =====================================================================
+// 3D axis-(0,1) transpose: in:[a,b,c] -> out:[b,a,c]  (last dim contiguous).
+//   out[j, i, k] = in[i, j, k]
+// Its own backward is the same op with (a,b) swapped, so one kernel suffices.
+// =====================================================================
+
+__global__ void transpose_3d01_k(const float* in, float* out, int a, int b, int c) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x; // over a*b*c
+    if (idx >= a * b * c) return;
+    int k = idx % c;
+    int j = (idx / c) % b;
+    int i = idx / (b * c);
+    // out index: ((j*a) + i)*c + k
+    out[(j * a + i) * c + k] = in[idx];
+}
+void launch_transpose_3d01_f32(const float* in, float* out, int a, int b, int c, void* s) {
+    int n = a * b * c, blk = 256, grid = (n + blk - 1) / blk;
+    transpose_3d01_k<<<grid, blk, 0, (cudaStream_t)s>>>(in, out, a, b, c);
+}
+
+} // extern "C"