docs: Phase T3 — GEMM fwd/bwd + finite-diff

Design doc covering the tiled forward, the dA/dB math + how transpose is
handled (materialize + reuse forward), the cuBLAS row-major reference, and
the finite-diff harness design + how T4 reuses it per-op.

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

Cargo.lock

@@ -88,6 +88,13 @@ version = "1.0.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
 
+[[package]]
+name = "xtrain-autodiff"
+version = "0.1.0"
+dependencies = [
+ "xtrain-tensor",
+]
+
 [[package]]
 name = "xtrain-cuda"
 version = "0.1.0"
@@ -101,6 +108,7 @@ version = "0.1.0"
 dependencies = [
  "half",
  "smallvec",
+ "xtrain-autodiff",
  "xtrain-cuda",
 ]
 

Cargo.toml

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

crates/xtrain-autodiff/Cargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "xtrain-autodiff"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+xtrain-tensor = { path = "../xtrain-tensor" }

crates/xtrain-autodiff/src/finite_diff.rs

@@ -0,0 +1,126 @@
+//! Central finite-difference gradient check.
+
+/// A scalar loss as a function of a flat parameter vector with a given shape.
+/// `(data, shape) -> loss`. The closure owns how `data` becomes a `Tensor`
+/// (e.g. `Tensor::from_slice(data, shape).to_device(...)`) and runs forward.
+pub type ParamFn<'a> = dyn Fn(&[f32], &[usize]) -> f32 + 'a;
+
+#[derive(Debug, Clone, Copy)]
+pub struct GradCheckConfig {
+    /// Perturbation magnitude per element.
+    pub eps: f32,
+    /// Max allowed relative error: `|num - ana| / (|num| + |ana| + atol)`.
+    pub rel_tol: f32,
+    /// Absolute floor in the denominator, so near-zero grads don't blow up the
+    /// relative error.
+    pub atol: f32,
+}
+
+impl Default for GradCheckConfig {
+    fn default() -> Self {
+        // eps=1e-3 balances truncation error (∝ eps²) against the ~1e-7 f32
+        // rounding noise on f(x±eps); rel_tol=2e-2 is the usual slack for an
+        // f32 GPU GEMM checked against an f32 central difference.
+        Self {
+            eps: 1e-3,
+            rel_tol: 2e-2,
+            atol: 1e-4,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct GradCheckResult {
+    pub passed: bool,
+    pub max_rel_err: f32,
+    /// Index of the worst element (largest relative error).
+    pub worst_index: usize,
+    pub worst_numeric: f32,
+    pub worst_analytic: f32,
+}
+
+/// Check `analytic_grad` against the central finite difference of `f` at `x`.
+///
+/// - `x`: flat parameter values (the point at which the gradient is taken).
+/// - `shape`: logical shape passed through to `f`.
+/// - `f`: scalar loss; called `2 * x.len()` times.
+/// - `analytic_grad`: candidate gradient, same length as `x`.
+pub fn grad_check(
+    x: &[f32],
+    shape: &[usize],
+    f: &ParamFn,
+    analytic_grad: &[f32],
+    cfg: GradCheckConfig,
+) -> GradCheckResult {
+    assert_eq!(
+        x.len(),
+        analytic_grad.len(),
+        "param/grad length mismatch: {} vs {}",
+        x.len(),
+        analytic_grad.len()
+    );
+
+    let mut perturbed = x.to_vec();
+    let mut max_rel_err = 0.0f32;
+    let mut worst_index = 0;
+    let mut worst_numeric = 0.0f32;
+    let mut worst_analytic = 0.0f32;
+
+    for i in 0..x.len() {
+        let orig = x[i];
+
+        perturbed[i] = orig + cfg.eps;
+        let f_plus = f(&perturbed, shape);
+
+        perturbed[i] = orig - cfg.eps;
+        let f_minus = f(&perturbed, shape);
+
+        perturbed[i] = orig; // restore for the next element
+
+        let numeric = (f_plus - f_minus) / (2.0 * cfg.eps);
+        let analytic = analytic_grad[i];
+
+        let rel_err = (numeric - analytic).abs() / (numeric.abs() + analytic.abs() + cfg.atol);
+        if rel_err > max_rel_err {
+            max_rel_err = rel_err;
+            worst_index = i;
+            worst_numeric = numeric;
+            worst_analytic = analytic;
+        }
+    }
+
+    GradCheckResult {
+        passed: max_rel_err <= cfg.rel_tol,
+        max_rel_err,
+        worst_index,
+        worst_numeric,
+        worst_analytic,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Host-only sanity check (no GPU): loss = sum(x²), grad = 2x.
+    #[test]
+    fn quadratic_grad_check() {
+        let x = vec![1.0f32, -2.0, 3.0, 0.5];
+        let f = |v: &[f32], _shape: &[usize]| v.iter().map(|t| t * t).sum::<f32>();
+        let grad: Vec<f32> = x.iter().map(|t| 2.0 * t).collect();
+
+        let res = grad_check(&x, &[4], &f, &grad, GradCheckConfig::default());
+        assert!(res.passed, "max_rel_err = {}", res.max_rel_err);
+    }
+
+    // A deliberately wrong gradient must be rejected.
+    #[test]
+    fn wrong_grad_is_rejected() {
+        let x = vec![1.0f32, 2.0, 3.0];
+        let f = |v: &[f32], _shape: &[usize]| v.iter().map(|t| t * t).sum::<f32>();
+        let bad_grad = vec![0.0f32, 0.0, 0.0];
+
+        let res = grad_check(&x, &[3], &f, &bad_grad, GradCheckConfig::default());
+        assert!(!res.passed);
+    }
+}

crates/xtrain-autodiff/src/lib.rs

@@ -0,0 +1,15 @@
+//! Reusable numerical-gradient checking for xtrain (Phase T3+).
+//!
+//! Given a scalar loss `f(x)` and an analytic gradient `g`, verify that `g`
+//! matches the central finite-difference estimate
+//! `(f(x+ε·eᵢ) - f(x-ε·eᵢ)) / 2ε` for every element `i`, within a relative
+//! tolerance. Later phases (T4 autograd) reuse this per-op: wrap each op's
+//! forward as the loss, run its backward to get `g`, and `grad_check`.
+//!
+//! The harness is host-only and dtype-agnostic at this layer: it works on a
+//! flat `&[f32]` parameter vector + shape and a closure. The closure is free to
+//! push the data to the GPU and run kernels — that detail stays out of here.
+
+pub mod finite_diff;
+
+pub use finite_diff::{GradCheckConfig, GradCheckResult, ParamFn, grad_check};

crates/xtrain-cuda/build.rs

@@ -22,6 +22,8 @@ fn main() {
     println!("cargo:rustc-link-search=native={cuda_path}/lib64");
     println!("cargo:rustc-link-lib=dylib=cudart");
     println!("cargo:rustc-link-lib=dylib=cuda");
+    // cuBLAS is used only as a correctness reference for the hand-written GEMM.
+    println!("cargo:rustc-link-lib=dylib=cublas");
 
     cc::Build::new()
         .cuda(true)
@@ -29,6 +31,7 @@ fn main() {
         .flag("-gencode=arch=compute_120,code=sm_120")
         .file("../../csrc/test/vecadd.cu")
         .file("../../csrc/ops/elementwise.cu")
+        .file("../../csrc/ops/gemm.cu")
         .compile("xtrain_cuda_kernels");
 }
 

crates/xtrain-cuda/src/ffi.rs

@@ -39,4 +39,56 @@ unsafe extern "C" {
         n: i32,
         stream: CudaStream,
     );
+
+    // Tiled GEMM: C = A @ B, row-major F32. A:[M,K] B:[K,N] C:[M,N]
+    // (csrc/ops/gemm.cu).
+    pub fn launch_gemm_tiled_f32(
+        a: *const f32,
+        b: *const f32,
+        c: *mut f32,
+        m: i32,
+        n: i32,
+        k: i32,
+        stream: CudaStream,
+    );
+
+    // Out-of-place 2D transpose: out[j,i] = in[i,j]. in:[rows,cols] row-major,
+    // out:[cols,rows] row-major (csrc/ops/gemm.cu).
+    pub fn launch_transpose_f32(
+        input: *const f32,
+        out: *mut f32,
+        rows: i32,
+        cols: i32,
+        stream: 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))]
+pub type CublasHandle = *mut c_void;
+
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    pub fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
+    pub fn cublasDestroy_v2(handle: CublasHandle) -> i32;
+    pub fn cublasSgemm_v2(
+        handle: CublasHandle,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: *const f32,
+        a: *const f32,
+        lda: i32,
+        b: *const f32,
+        ldb: i32,
+        beta: *const f32,
+        c: *mut f32,
+        ldc: i32,
+    ) -> i32;
+}
+
+#[cfg(not(no_cuda))]
+pub const CUBLAS_OP_N: i32 = 0;

crates/xtrain-tensor/Cargo.toml

@@ -7,3 +7,6 @@ edition.workspace = true
 xtrain-cuda = { path = "../xtrain-cuda" }
 half.workspace = true
 smallvec.workspace = true
+
+[dev-dependencies]
+xtrain-autodiff = { path = "../xtrain-autodiff" }

crates/xtrain-tensor/src/tensor.rs

@@ -156,6 +156,97 @@ impl Tensor {
         xtrain_cuda::device::synchronize().expect("scale kernel sync failed");
         out
     }
+
+    // --- GEMM (the T3 kernels) ---
+
+    /// Matrix multiply: `C = self @ other`. `self`:[M,K], `other`:[K,N] → [M,N].
+    ///
+    /// Runs the tiled `gemm_tiled_f32` CUDA kernel. Requires contiguous F32
+    /// tensors on the same GPU. Available only when CUDA is compiled in.
+    #[cfg(not(no_cuda))]
+    pub fn matmul(&self, other: &Tensor) -> Self {
+        assert_eq!(self.dtype, DType::F32, "matmul only supports F32");
+        assert_eq!(other.dtype, DType::F32, "matmul only supports F32");
+        assert_eq!(self.ndim(), 2, "matmul requires 2D lhs");
+        assert_eq!(other.ndim(), 2, "matmul requires 2D rhs");
+        assert_eq!(
+            self.shape[1], other.shape[0],
+            "inner dimension mismatch: {:?} @ {:?}",
+            self.shape, other.shape
+        );
+        assert!(
+            self.is_contiguous() && other.is_contiguous(),
+            "matmul requires contiguous tensors"
+        );
+        assert_eq!(self.device(), other.device(), "matmul device mismatch");
+        assert!(
+            matches!(self.device(), Device::Cuda(_)),
+            "matmul requires CUDA tensors"
+        );
+
+        let m = self.shape[0];
+        let k = self.shape[1];
+        let n = other.shape[1];
+        let out = Tensor::zeros(&[m, n], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_gemm_tiled_f32(
+                self.data_ptr() as *const f32,
+                other.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                m as i32,
+                n as i32,
+                k as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        xtrain_cuda::device::synchronize().expect("matmul kernel sync failed");
+        out
+    }
+
+    /// Out-of-place 2D transpose: returns a new contiguous tensor `out[j,i] =
+    /// self[i,j]`. Requires a contiguous F32 CUDA tensor.
+    #[cfg(not(no_cuda))]
+    pub fn transpose_2d(&self) -> Self {
+        assert_eq!(self.dtype, DType::F32, "transpose only supports F32");
+        assert_eq!(self.ndim(), 2, "transpose_2d requires 2D tensor");
+        assert!(self.is_contiguous(), "transpose requires contiguous tensor");
+        assert!(
+            matches!(self.device(), Device::Cuda(_)),
+            "transpose requires a CUDA tensor"
+        );
+
+        let rows = self.shape[0];
+        let cols = self.shape[1];
+        let out = Tensor::zeros(&[cols, rows], DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_transpose_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("transpose kernel sync failed");
+        out
+    }
+
+    /// Backward of `C = A @ B` given the upstream gradient `dC` (shape [M,N]).
+    /// Returns `(dA, dB)` where `dA = dC @ Bᵀ` ([M,K]) and `dB = Aᵀ @ dC`
+    /// ([K,N]). All tensors contiguous F32 on the same GPU.
+    #[cfg(not(no_cuda))]
+    pub fn matmul_backward(a: &Tensor, b: &Tensor, dc: &Tensor) -> (Tensor, Tensor) {
+        assert_eq!(a.ndim(), 2, "matmul_backward requires 2D A");
+        assert_eq!(b.ndim(), 2, "matmul_backward requires 2D B");
+        assert_eq!(dc.ndim(), 2, "matmul_backward requires 2D dC");
+        assert_eq!(a.shape[1], b.shape[0], "A/B inner dim mismatch");
+        assert_eq!(dc.shape[0], a.shape[0], "dC rows != A rows (M)");
+        assert_eq!(dc.shape[1], b.shape[1], "dC cols != B cols (N)");
+
+        let da = dc.matmul(&b.transpose_2d()); // [M,N] @ [N,K] = [M,K]
+        let db = a.transpose_2d().matmul(dc); // [K,M] @ [M,N] = [K,N]
+        (da, db)
+    }
 }
 
 impl std::fmt::Debug for Tensor {

crates/xtrain-tensor/tests/gemm.rs

@@ -0,0 +1,183 @@
+// GPU acceptance tests for the hand-written GEMM forward + backward (Phase T3).
+// Gated behind `not(no_cuda)`: on a GPU-less machine these compile out so host
+// `cargo check` stays green; they run on dash5.
+#![cfg(not(no_cuda))]
+
+use xtrain_autodiff::{GradCheckConfig, grad_check};
+use xtrain_cuda::device;
+use xtrain_tensor::{Device, Tensor};
+
+// Deterministic pseudo-random fill in [-0.5, 0.5), seeded by a linear
+// congruential generator so tests are reproducible without an RNG dep.
+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();
+}
+
+// --- cuBLAS reference (correctness oracle for the hand-written kernel) ---
+
+/// Row-major `C = A @ B` via cuBLAS `Sgemm` (which is column-major).
+/// Identity: row-major C = A@B  ⟺  column-major Cᵀ = Bᵀ @ Aᵀ. We hand cuBLAS
+/// our row-major B and A as-is (it reads them as the col-major transposes) with
+/// OP_N, swapped order, and m=N, n=M, k=K. Output lands row-major in `c`.
+fn cublas_matmul(a: &Tensor, b: &Tensor) -> Tensor {
+    let m = a.shape()[0];
+    let k = a.shape()[1];
+    let n = b.shape()[1];
+    let c = Tensor::zeros(&[m, n], xtrain_tensor::DType::F32, a.device());
+
+    let alpha = 1.0f32;
+    let beta = 0.0f32;
+    unsafe {
+        let mut handle = std::ptr::null_mut();
+        assert_eq!(xtrain_cuda::ffi::cublasCreate_v2(&mut handle), 0);
+        let status = xtrain_cuda::ffi::cublasSgemm_v2(
+            handle,
+            xtrain_cuda::ffi::CUBLAS_OP_N,
+            xtrain_cuda::ffi::CUBLAS_OP_N,
+            n as i32,
+            m as i32,
+            k as i32,
+            &alpha,
+            b.data_ptr() as *const f32,
+            n as i32,
+            a.data_ptr() as *const f32,
+            k as i32,
+            &beta,
+            c.data_ptr() as *mut f32,
+            n as i32,
+        );
+        assert_eq!(status, 0, "cublasSgemm failed: {status}");
+        device::synchronize().unwrap();
+        xtrain_cuda::ffi::cublasDestroy_v2(handle);
+    }
+    c
+}
+
+fn max_rel_err(got: &[f32], reference: &[f32]) -> f32 {
+    got.iter()
+        .zip(reference)
+        .map(|(g, r)| (g - r).abs() / (g.abs() + r.abs() + 1e-6))
+        .fold(0.0f32, f32::max)
+}
+
+// --- Forward: hand-written tiled GEMM vs cuBLAS sgemm ---
+
+fn run_fwd(m: usize, k: usize, n: usize) {
+    require_gpu();
+    let a = Tensor::from_slice(&fill(m * k, 1), &[m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&fill(k * n, 2), &[k, n]).to_device(Device::Cuda(0));
+
+    let mine = a.matmul(&b).to_device(Device::Cpu);
+    let reference = cublas_matmul(&a, &b).to_device(Device::Cpu);
+
+    let rel = max_rel_err(mine.as_slice::<f32>(), reference.as_slice::<f32>());
+    println!("fwd GEMM [{m}x{k}]@[{k}x{n}] vs cuBLAS: max_rel_err = {rel:.3e}");
+    assert!(rel < 1e-3, "fwd rel-err {rel} too high for {m}x{k}x{n}");
+}
+
+#[test]
+fn fwd_square() {
+    run_fwd(64, 64, 64);
+}
+
+#[test]
+fn fwd_rect() {
+    run_fwd(65, 97, 33); // non-tile-aligned dims exercise the boundary masking
+}
+
+#[test]
+fn fwd_large() {
+    run_fwd(256, 256, 256);
+}
+
+// --- Backward: dA, dB vs the finite-difference harness ---
+//
+// Scalar loss L = sum(W ∘ C) with C = A @ B and W fixed random weights.
+// Then dC = W, dA = dC @ Bᵀ, dB = Aᵀ @ dC (matmul_backward). We check each of
+// dA and dB against central differences of L w.r.t. that input.
+
+fn run_bwd(m: usize, k: usize, n: usize) {
+    require_gpu();
+
+    let a_host = fill(m * k, 11);
+    let b_host = fill(k * n, 22);
+    let w_host = fill(m * n, 33); // loss weights, == dC
+
+    let a = Tensor::from_slice(&a_host, &[m, k]).to_device(Device::Cuda(0));
+    let b = Tensor::from_slice(&b_host, &[k, n]).to_device(Device::Cuda(0));
+    let dc = Tensor::from_slice(&w_host, &[m, n]).to_device(Device::Cuda(0));
+
+    let (da, db) = Tensor::matmul_backward(&a, &b, &dc);
+    let da_host = da.to_device(Device::Cpu);
+    let db_host = db.to_device(Device::Cpu);
+
+    let cfg = GradCheckConfig::default();
+
+    // Check dA: vary A, hold B fixed.
+    let b_fixed = b_host.clone();
+    let w_fixed = w_host.clone();
+    let loss_a = move |a_vals: &[f32], a_shape: &[usize]| -> f32 {
+        let av = Tensor::from_slice(a_vals, a_shape).to_device(Device::Cuda(0));
+        let bv = Tensor::from_slice(&b_fixed, &[k, n]).to_device(Device::Cuda(0));
+        let c = av.matmul(&bv).to_device(Device::Cpu);
+        c.as_slice::<f32>()
+            .iter()
+            .zip(&w_fixed)
+            .map(|(c, w)| c * w)
+            .sum()
+    };
+    let res_a = grad_check(&a_host, &[m, k], &loss_a, da_host.as_slice::<f32>(), cfg);
+    println!(
+        "bwd dA [{m}x{k}]: max_rel_err = {:.3e} (worst num={:.5} ana={:.5} @ {})",
+        res_a.max_rel_err, res_a.worst_numeric, res_a.worst_analytic, res_a.worst_index
+    );
+    assert!(res_a.passed, "dA grad-check failed: {:?}", res_a);
+
+    // Check dB: vary B, hold A fixed.
+    let a_fixed = a_host.clone();
+    let w_fixed2 = w_host.clone();
+    let loss_b = move |b_vals: &[f32], b_shape: &[usize]| -> f32 {
+        let av = Tensor::from_slice(&a_fixed, &[m, k]).to_device(Device::Cuda(0));
+        let bv = Tensor::from_slice(b_vals, b_shape).to_device(Device::Cuda(0));
+        let c = av.matmul(&bv).to_device(Device::Cpu);
+        c.as_slice::<f32>()
+            .iter()
+            .zip(&w_fixed2)
+            .map(|(c, w)| c * w)
+            .sum()
+    };
+    let res_b = grad_check(&b_host, &[k, n], &loss_b, db_host.as_slice::<f32>(), cfg);
+    println!(
+        "bwd dB [{k}x{n}]: max_rel_err = {:.3e} (worst num={:.5} ana={:.5} @ {})",
+        res_b.max_rel_err, res_b.worst_numeric, res_b.worst_analytic, res_b.worst_index
+    );
+    assert!(res_b.passed, "dB grad-check failed: {:?}", res_b);
+}
+
+#[test]
+fn bwd_square() {
+    run_bwd(16, 16, 16);
+}
+
+#[test]
+fn bwd_rect() {
+    run_bwd(12, 20, 8);
+}

csrc/ops/gemm.cu

@@ -0,0 +1,75 @@
+extern "C" {
+
+// Tiled GEMM (shared memory). C = A @ B, all row-major F32.
+// A: [M, K], B: [K, N], C: [M, N].
+// Each block computes a TILE_SIZE x TILE_SIZE tile of C, cooperatively loading
+// tiles of A and B into shared memory. FP32 accumulation.
+#define TILE_SIZE 32
+
+__global__ void gemm_tiled_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K
+) {
+    __shared__ float As[TILE_SIZE][TILE_SIZE];
+    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
+
+    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
+    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
+
+    float sum = 0.0f;
+
+    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        int a_col = t * TILE_SIZE + threadIdx.x;
+        if (row < M && a_col < K) {
+            As[threadIdx.y][threadIdx.x] = A[row * K + a_col];
+        } else {
+            As[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        int b_row = t * TILE_SIZE + threadIdx.y;
+        if (b_row < K && col < N) {
+            Bs[threadIdx.y][threadIdx.x] = B[b_row * N + col];
+        } else {
+            Bs[threadIdx.y][threadIdx.x] = 0.0f;
+        }
+
+        __syncthreads();
+
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum += As[threadIdx.y][k] * Bs[k][threadIdx.x];
+        }
+
+        __syncthreads();
+    }
+
+    if (row < M && col < N) {
+        C[row * N + col] = sum;
+    }
+}
+
+void launch_gemm_tiled_f32(
+    const float* A, const float* B, float* C,
+    int M, int N, int K, void* stream
+) {
+    dim3 block(TILE_SIZE, TILE_SIZE);
+    dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);
+    gemm_tiled_f32<<<grid, block, 0, (cudaStream_t)stream>>>(A, B, C, M, N, K);
+}
+
+// Out-of-place transpose: out[j, i] = in[i, j].
+// in: [rows, cols] row-major, out: [cols, rows] row-major.
+__global__ void transpose_f32(const float* in, float* out, int rows, int cols) {
+    int col = blockIdx.x * blockDim.x + threadIdx.x; // over cols of `in`
+    int row = blockIdx.y * blockDim.y + threadIdx.y; // over rows of `in`
+    if (row < rows && col < cols) {
+        out[col * rows + row] = in[row * cols + col];
+    }
+}
+
+void launch_transpose_f32(const float* in, float* out, int rows, int cols, void* stream) {
+    dim3 block(16, 16);
+    dim3 grid((cols + block.x - 1) / block.x, (rows + block.y - 1) / block.y);
+    transpose_f32<<<grid, block, 0, (cudaStream_t)stream>>>(in, out, rows, cols);
+}
+
+}

docs/02-gemm-autodiff.md

@@ -0,0 +1,130 @@
+# Phase: GEMM Forward/Backward + Finite-Diff — Design Document
+
+## Goal
+
+在 T2 的 `Tensor` 之上交付**手写 GEMM 的前向 + 反向**，以及一个**可复用的有限差分梯度检查 harness**。
+这是 autograd（T4）的地基：T4 的每个算子 backward 都会用同一个 harness 对拍。本 Phase 只做三件事：
+
+1. **GEMM 前向**：手写 tiled CUDA kernel，`C = A @ B`，row-major F32，正确性优先；
+2. **GEMM 反向**：给上游梯度 `dC`，算 `dA = dC · Bᵀ`、`dB = Aᵀ · dC`（复用前向 + 一个 transpose kernel）；
+3. **finite-diff harness**：给标量 loss `f(x)` 和解析梯度 `g`，逐元素用中心差分核对。
+
+**明确不做**（留给 T4+）：autograd tape / 计算图、其它算子、broadcast、非 contiguous 输入、半精度、性能调优（T7）。
+
+## Module Layout
+
+```
+csrc/ops/gemm.cu                       # gemm_tiled_f32 + transpose_f32 + launch_*（由 xtrain-cuda/build.rs 编）
+
+crates/xtrain-cuda/
+├── build.rs                           # 新增 gemm.cu；链接 cublas（仅作前向对拍参考）
+└── src/ffi.rs                         # 新增 launch_gemm_tiled_f32 / launch_transpose_f32（no_cuda 门控）
+                                       #   + cuBLAS sgemm FFI（no_cuda 门控，仅测试用）
+
+crates/xtrain-tensor/
+├── Cargo.toml                         # dev-dep: xtrain-autodiff
+├── src/tensor.rs                      # matmul / transpose_2d / matmul_backward（no_cuda 门控）
+└── tests/gemm.rs                      # 前向对 cuBLAS、反向对 finite-diff（#![cfg(not(no_cuda))]）
+
+crates/xtrain-autodiff/               # 新 crate（reusable harness）
+├── Cargo.toml
+└── src/
+    ├── lib.rs
+    └── finite_diff.rs                 # grad_check + GradCheckConfig/Result
+```
+
+## Key Design Decisions
+
+### 前向：tiled GEMM（沿用 xserv 风格）
+
+`csrc/ops/gemm.cu` 的 `gemm_tiled_f32`：`TILE_SIZE=32`，每个 block 协作把 A/B 的 tile 载入 shared
+memory，FP32 累加。这就是 xserv `csrc/gemm/tiled.cu` 的 F32 版（去掉 BF16）——正确性已被 xserv 验证过，
+边界 mask（`row<M && col<N`、tile 越界填 0）让非 32 对齐的 M/N/K 也对。
+
+张量层 `Tensor::matmul(&self, other)` 镜像 T2 `scale` 的接线方式：校验 2D / contiguous / F32 / 同卡 →
+`zeros` 输出 → `launch_gemm_tiled_f32` → `synchronize`。
+
+### 反向：dA / dB 的数学与 transpose 处理
+
+`C = A·B`，`A:[M,K]`、`B:[K,N]`、`C:[M,N]`。给上游 `dC:[M,N]`：
+
+```
+dA = dC · Bᵀ      # [M,N] · [N,K] = [M,K]
+dB = Aᵀ · dC      # [K,M] · [M,N] = [K,N]
+```
+
+推导（标量 loss `L`，`C_ij = Σ_k A_ik B_kj`）：
+`∂L/∂A_ik = Σ_j (∂L/∂C_ij)(∂C_ij/∂A_ik) = Σ_j dC_ij · B_kj = (dC·Bᵀ)_ik`；
+`∂L/∂B_kj = Σ_i dC_ij · A_ik = (Aᵀ·dC)_kj`。
+
+**transpose 怎么处理**：T2 没有 transpose/reshape。这里加一个**最小的 out-of-place transpose
+kernel**（`transpose_f32`：`out[j,i]=in[i,j]`，写出连续的 `[cols,rows]`），`matmul_backward` 先把
+`B`/`A` 物化转置成连续张量，再复用 `matmul`：
+
+```rust
+let da = dc.matmul(&b.transpose_2d());   // [M,N] @ [N,K]
+let db = a.transpose_2d().matmul(dc);    // [K,M] @ [M,N]
+```
+
+选「物化转置 + 复用前向」而非「带 transpose flag 的 GEMM 变体」：代码量最小、复用已对拍过的前向 kernel、
+逻辑一眼能看懂。代价是多一次拷贝 + 临时显存——属于 T7 性能范畴（届时可改成 transpose-flag GEMM 或
+non-contiguous GEMM 来省掉）。
+
+### cuBLAS 仅作前向参考（row-major vs col-major）
+
+cuBLAS `Sgemm` 是 column-major，而我们 row-major。用恒等式
+`row-major C=A·B ⟺ col-major Cᵀ=Bᵀ·Aᵀ`：把 row-major 的 B、A 原样喂给 cuBLAS（它当成各自的
+col-major 转置），传 `OP_N, OP_N`，参数 `m=N, n=M, k=K, lda=N, ldb=K, ldc=N`，输出原地就是 row-major
+的 `C`，**不做任何显式转置**。cuBLAS FFI（`cublasSgemm_v2` 等）声明在 `xtrain-cuda/src/ffi.rs`，和
+kernel FFI 一样 `#[cfg(not(no_cuda))]` 门控，`build.rs` 加 `-lcublas`。它只在测试里当 oracle，不进
+production 路径。
+
+### finite-diff harness（T4 的复用核心）
+
+`xtrain-autodiff::grad_check`：
+
+```rust
+pub type ParamFn<'a> = dyn Fn(&[f32], &[usize]) -> f32 + 'a;   // (flat 参数, shape) -> 标量 loss
+
+pub fn grad_check(x: &[f32], shape: &[usize], f: &ParamFn,
+                  analytic_grad: &[f32], cfg: GradCheckConfig) -> GradCheckResult;
+```
+
+- 逐元素中心差分：`(f(x+ε·eᵢ) − f(x−ε·eᵢ)) / 2ε`，共调 `f` 两次 `x.len()` 遍。
+- 相对误差判据：`|num − ana| / (|num| + |ana| + atol)`，`atol` 防近零梯度炸比值。
+- 默认 `eps=1e-3 / rel_tol=2e-2 / atol=1e-4`：`eps` 平衡截断误差(∝eps²)与 f32 在 `f(x±eps)` 上
+  ~1e-7 的舍入噪声；`rel_tol=2e-2` 是 f32 GPU GEMM 对 f32 中心差分的常规裕度。
+- **host-only、不依赖 CUDA**：harness 只碰 `&[f32]`+shape+闭包；GPU 前向藏在闭包里（`from_slice →
+  to_device → matmul → sum`）。所以这个 crate 本地能 build/test，自带两个纯 host 单测（`sum(x²)` 梯度
+  `2x` 通过；错误梯度被拒）。
+
+**T4 怎么用**：每个算子 wrap forward 为标量 loss（`L = sum(W∘out)`），跑该算子 backward 拿解析梯度，
+对每个输入张量调一次 `grad_check`。本 Phase 的 GEMM 反向测试就是这个模式的样板。
+
+## 验证方法
+
+GPU 测试 `#![cfg(not(no_cuda))]` 门控，dash5 实跑：
+
+```sh
+ssh dash5
+export PATH=/usr/local/cuda/bin:/opt/wjh/.cargo/bin:$PATH
+cd ~/projects/xtrain
+cargo test -p xtrain-autodiff                       # host-only harness 自检
+cargo test -p xtrain-tensor --test gemm -- --nocapture
+```
+
+- **前向**：手写 tiled GEMM vs cuBLAS sgemm，随机矩阵（square / rect 非对齐 / 256³），`max_rel_err < 1e-3`。
+- **反向**：标量 loss `L = sum(W∘C)`（`dC = W`），`matmul_backward` 出 `dA`/`dB`，各自对 finite-diff
+  harness 核对，`rel_tol = 2e-2` 内通过。
+
+本地（无 GPU）：`cargo check --workspace --all-targets` + `cargo fmt --all -- --check` 绿；GPU 测试编译出局。
+
+## Takeaways
+
+1. **反向 = 前向 + 转置的复用**：`dA=dC·Bᵀ`、`dB=Aᵀ·dC` 全靠已对拍过的前向 kernel + 一个 transpose
+   kernel 拼出来，不必为 backward 单写 kernel。代价（多余拷贝）留给 T7。
+2. **harness 与设备解耦**：把 finite-diff 设计成「只吃 `&[f32]`+闭包」，GPU 细节进闭包，于是 harness
+   本地可测、且 T4 任何算子（不限 GEMM）都能直接复用。
+3. **cuBLAS 收口在 FFI + 门控**：参考实现也守 `no_cuda` 约定，本地 check 不被 cublas 链接拖累。
+4. **row-major/col-major 是 GEMM 最大坑**：靠 `C=AB ⟺ Cᵀ=BᵀAᵀ` 交换 A/B 传入顺序 + 调 ld，零显式转置。
+```