docs: Phase T12 — bf16 mixed precision design

docs/11-bf16-mixed-precision.md: the AMP split (bf16 linears +
activations, fp32 master / norms / softmax / RoPE / CE, no loss
scaling), the cast-op bridge, module layout, and the dual
verification gate (fp32 unchanged + bf16 looser-tol + convergence +
mem/throughput). Memory/throughput before->after to be filled from
the dash5 bench.

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

crates/xtrain-autodiff/src/ops.rs

@@ -13,7 +13,27 @@
 #![cfg(not(no_cuda))]
 
 use crate::tape::Var;
-use xtrain_tensor::Tensor;
+use xtrain_tensor::{DType, Tensor};
+
+/// dtype cast as an autograd node (Phase T12 — the AMP bridge between fp32 master
+/// weights / fp32 reductions and the bf16 compute stream). Forward casts `x` to
+/// `target`; **backward casts the upstream grad back to `x`'s dtype**. So a fp32
+/// master-weight leaf fed through `cast(w, BF16)` into a bf16 matmul accumulates
+/// an **fp32** grad — AdamW / clip / DDP all-reduce stay fp32, untouched.
+pub fn cast(x: &Var, target: DType) -> Var {
+    let src = x.value().dtype();
+    if src == target {
+        return x.clone();
+    }
+    let out = x.value().to_dtype(target);
+    Var::from_op(
+        out,
+        vec![x.clone()],
+        Box::new(move |d, parents| {
+            Var::push_grad(&parents[0], d.to_dtype(src));
+        }),
+    )
+}
 
 /// `C = A @ B` (2D). Backward: `dA = dC @ Bᵀ`, `dB = Aᵀ @ dC`.
 pub fn matmul(a: &Var, b: &Var) -> Var {

crates/xtrain-cuda/build.rs

@@ -36,6 +36,7 @@ fn main() {
         .file("../../csrc/ops/model.cu")
         .file("../../csrc/ops/optim.cu")
         .file("../../csrc/ops/attention.cu")
+        .file("../../csrc/ops/cast.cu")
         .compile("xtrain_cuda_kernels");
 }
 

crates/xtrain-cuda/src/cublas.rs

@@ -19,6 +19,7 @@
 
 use crate::ffi::{self, CublasHandle};
 use std::cell::RefCell;
+use std::ffi::c_void;
 
 thread_local! {
     static HANDLE: RefCell<Option<CublasHandle>> = const { RefCell::new(None) };
@@ -159,3 +160,131 @@ pub fn sgemm_strided_batched(
         assert_eq!(status, 0, "cublasSgemmStridedBatched failed: {status}");
     });
 }
+
+/// bf16 row-major GEMM `C[m,n] = opA(A)·opB(B)` via `cublasGemmEx`: bf16 in/out,
+/// **fp32 accumulation** (`CUBLAS_COMPUTE_32F`) — the standard AMP matmul (Phase
+/// T12). `a`/`b`/`c` are device pointers to row-major **bf16** matrices; the
+/// row-major⟺col-major transpose algebra is identical to [`sgemm`] (we compute
+/// the col-major `Cᵀ`). `alpha`/`beta` are fp32 host scalars (compute is fp32).
+#[allow(clippy::too_many_arguments)]
+pub fn gemm_ex(
+    trans_a: bool,
+    trans_b: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    alpha: f32,
+    a: *const c_void,
+    b: *const c_void,
+    beta: f32,
+    c: *mut c_void,
+) {
+    let lda = if trans_a { m } else { k };
+    let ldb = if trans_b { k } else { n };
+    let ldc = n;
+    let op_a = if trans_a {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let op_b = if trans_b {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let bf16 = ffi::CUDA_R_16BF;
+
+    with_handle(|handle| {
+        let status = unsafe {
+            ffi::cublasGemmEx(
+                handle,
+                op_b,
+                op_a,
+                n as i32,
+                m as i32,
+                k as i32,
+                &alpha as *const f32 as *const c_void,
+                b,
+                bf16,
+                ldb as i32,
+                a,
+                bf16,
+                lda as i32,
+                &beta as *const f32 as *const c_void,
+                c,
+                bf16,
+                ldc as i32,
+                ffi::CUBLAS_COMPUTE_32F,
+                ffi::CUBLAS_GEMM_DEFAULT,
+            )
+        };
+        assert_eq!(status, 0, "cublasGemmEx failed: {status}");
+    });
+}
+
+/// Strided-batched bf16 GEMM (Phase T12) — the [`gemm_ex`] analogue of
+/// [`sgemm_strided_batched`] for the batched attention GEMMs. bf16 in/out, fp32
+/// accumulation; strides are in ELEMENTS.
+#[allow(clippy::too_many_arguments)]
+pub fn gemm_ex_strided_batched(
+    trans_a: bool,
+    trans_b: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    alpha: f32,
+    a: *const c_void,
+    stride_a: usize,
+    b: *const c_void,
+    stride_b: usize,
+    beta: f32,
+    c: *mut c_void,
+    stride_c: usize,
+    batch: usize,
+) {
+    let lda = if trans_a { m } else { k };
+    let ldb = if trans_b { k } else { n };
+    let ldc = n;
+    let op_a = if trans_a {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let op_b = if trans_b {
+        ffi::CUBLAS_OP_T
+    } else {
+        ffi::CUBLAS_OP_N
+    };
+    let bf16 = ffi::CUDA_R_16BF;
+
+    with_handle(|handle| {
+        let status = unsafe {
+            ffi::cublasGemmStridedBatchedEx(
+                handle,
+                op_b,
+                op_a,
+                n as i32,
+                m as i32,
+                k as i32,
+                &alpha as *const f32 as *const c_void,
+                b,
+                bf16,
+                ldb as i32,
+                stride_b as i64,
+                a,
+                bf16,
+                lda as i32,
+                stride_a as i64,
+                &beta as *const f32 as *const c_void,
+                c,
+                bf16,
+                ldc as i32,
+                stride_c as i64,
+                batch as i32,
+                ffi::CUBLAS_COMPUTE_32F,
+                ffi::CUBLAS_GEMM_DEFAULT,
+            )
+        };
+        assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
+    });
+}

crates/xtrain-cuda/src/ffi.rs

@@ -324,3 +324,126 @@ unsafe extern "C" {
 pub const CUBLAS_OP_N: i32 = 0;
 #[cfg(not(no_cuda))]
 pub const CUBLAS_OP_T: i32 = 1;
+
+// --- bf16 mixed precision (Phase T12) ---
+//
+// cudaDataType / cublasComputeType enum values (same as xserv's gemm.rs). The
+// bf16 GEMM uses bf16 in/out with fp32 accumulation (CUBLAS_COMPUTE_32F).
+#[cfg(not(no_cuda))]
+pub const CUDA_R_32F: i32 = 0;
+#[cfg(not(no_cuda))]
+pub const CUDA_R_16BF: i32 = 14;
+#[cfg(not(no_cuda))]
+pub const CUBLAS_COMPUTE_32F: i32 = 68;
+/// CUBLAS_GEMM_DEFAULT — let cuBLAS pick the algorithm.
+#[cfg(not(no_cuda))]
+pub const CUBLAS_GEMM_DEFAULT: i32 = -1;
+
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    // General GEMM with explicit in/out + compute types (bf16 path). `alpha`/
+    // `beta` are fp32 host scalars (compute type is fp32). Pointers are void* so
+    // the same FFI serves bf16 / fp32.
+    #[allow(clippy::too_many_arguments)]
+    pub fn cublasGemmEx(
+        handle: CublasHandle,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: *const std::ffi::c_void,
+        a: *const std::ffi::c_void,
+        a_type: i32,
+        lda: i32,
+        b: *const std::ffi::c_void,
+        b_type: i32,
+        ldb: i32,
+        beta: *const std::ffi::c_void,
+        c: *mut std::ffi::c_void,
+        c_type: i32,
+        ldc: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+    #[allow(clippy::too_many_arguments)]
+    pub fn cublasGemmStridedBatchedEx(
+        handle: CublasHandle,
+        transa: i32,
+        transb: i32,
+        m: i32,
+        n: i32,
+        k: i32,
+        alpha: *const std::ffi::c_void,
+        a: *const std::ffi::c_void,
+        a_type: i32,
+        lda: i32,
+        stride_a: i64,
+        b: *const std::ffi::c_void,
+        b_type: i32,
+        ldb: i32,
+        stride_b: i64,
+        beta: *const std::ffi::c_void,
+        c: *mut std::ffi::c_void,
+        c_type: i32,
+        ldc: i32,
+        stride_c: i64,
+        batch_count: i32,
+        compute_type: i32,
+        algo: i32,
+    ) -> i32;
+}
+
+// bf16 cast + elementwise kernels (csrc/ops/cast.cu). Pointers are void* (bf16
+// buffers); f32 sides are typed. The activation stream flows bf16; the math
+// accumulates in fp32 inside each kernel.
+#[cfg(not(no_cuda))]
+unsafe extern "C" {
+    pub fn launch_cast_f32_to_bf16(input: *const f32, out: *mut c_void, n: i32, s: CudaStream);
+    pub fn launch_cast_bf16_to_f32(input: *const c_void, out: *mut f32, n: i32, s: CudaStream);
+
+    pub fn launch_add_bf16(
+        a: *const c_void,
+        b: *const c_void,
+        out: *mut c_void,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_mul_bf16(
+        a: *const c_void,
+        b: *const c_void,
+        out: *mut c_void,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_scale_bf16(
+        input: *const c_void,
+        out: *mut c_void,
+        alpha: f32,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_silu_bf16(x: *const c_void, y: *mut c_void, n: i32, s: CudaStream);
+    pub fn launch_silu_dx_bf16(
+        x: *const c_void,
+        dy: *const c_void,
+        dx: *mut c_void,
+        n: i32,
+        s: CudaStream,
+    );
+    pub fn launch_add_bias_bf16(
+        x: *const c_void,
+        bias: *const c_void,
+        out: *mut c_void,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+    pub fn launch_sum_rows_bf16(
+        dout: *const c_void,
+        dbias: *mut c_void,
+        rows: i32,
+        cols: i32,
+        s: CudaStream,
+    );
+}

crates/xtrain-distributed/src/bin/train_ddp.rs

@@ -82,6 +82,9 @@ fn main() {
     let val_tokens: usize = flag(&args, "--val-tokens", 0);
     let eval_every: usize = flag(&args, "--eval-every", 0);
     let eval_batches: usize = flag(&args, "--eval-batches", 64);
+    // bf16 mixed precision (Phase T12): fp32 master weights, bf16 linears +
+    // activations. Opt-in; default fp32 reproduces v0–v4 numerics.
+    let bf16 = args.iter().any(|a| a == "--bf16");
     let ckpt: Option<PathBuf> = args
         .iter()
         .position(|a| a == "--ckpt")
@@ -161,12 +164,22 @@ fn main() {
          eval every {eval_every}"
     );
 
+    if bf16 {
+        println!("bf16 mixed precision: ON (fp32 master weights)");
+    }
     let results = launch(
         &devices,
         &train_corpus,
         valid.as_ref(),
         &dcfg,
-        move |device| build_model(cfg, device),
+        move |device| {
+            let m = build_model(cfg, device);
+            if bf16 {
+                m.with_compute_dtype(xtrain_tensor::DType::BF16)
+            } else {
+                m
+            }
+        },
     );
     let r0 = &results[0];
     let start = r0.losses.first().copied().unwrap_or(0.0);

crates/xtrain-model/src/model.rs

@@ -5,7 +5,7 @@
 use crate::config::Config;
 use xtrain_autodiff::ops;
 use xtrain_autodiff::tape::Var;
-use xtrain_tensor::{Device, Tensor};
+use xtrain_tensor::{DType, Device, Tensor};
 
 /// One decoder block's learnable tensors.
 struct Block {
@@ -30,6 +30,13 @@ pub struct TinyTransformer {
     blocks: Vec<Block>,
     final_norm: Var, // [dim]
     lm_head: Var,    // [dim, vocab]
+    /// Compute dtype for the forward graph (Phase T12). `F32` (default) = the
+    /// original path, bit-identical to T10/T11. `BF16` = mixed precision: the
+    /// parameter leaves stay fp32 (master), but each linear's weight is cast to
+    /// bf16 on the fly and the activation stream flows bf16 (see
+    /// `docs/11-bf16-mixed-precision.md`). The cast op's backward upcasts the bf16
+    /// weight grad back to fp32, so AdamW/clip/DDP stay fp32 and unchanged.
+    compute_dtype: DType,
 }
 
 impl TinyTransformer {
@@ -71,6 +78,7 @@ impl TinyTransformer {
             blocks,
             final_norm,
             lm_head,
+            compute_dtype: DType::F32,
         }
     }
 
@@ -78,6 +86,35 @@ impl TinyTransformer {
         &self.cfg
     }
 
+    /// Set the forward compute dtype (Phase T12). `BF16` enables mixed precision
+    /// (fp32 master weights, bf16 linears + activations); `F32` (the default) is
+    /// the unchanged full-precision path. Builder-style so existing call sites
+    /// that don't opt in keep the fp32 numerics bit-for-bit.
+    pub fn with_compute_dtype(mut self, dtype: DType) -> Self {
+        assert!(
+            matches!(dtype, DType::F32 | DType::BF16),
+            "compute_dtype must be F32 or BF16"
+        );
+        self.compute_dtype = dtype;
+        self
+    }
+
+    pub fn compute_dtype(&self) -> DType {
+        self.compute_dtype
+    }
+
+    /// Project `x` (activation, in the compute dtype) by weight `w` (an fp32
+    /// master leaf). In bf16 mode the weight is cast to bf16 via the autograd
+    /// `cast` op (whose backward upcasts the grad to fp32); in fp32 mode this is
+    /// just `matmul(x, w)`. The activation `x` already carries `compute_dtype`.
+    fn linear(&self, x: &Var, w: &Var) -> Var {
+        match self.compute_dtype {
+            DType::F32 => ops::matmul(x, w),
+            DType::BF16 => ops::matmul(x, &ops::cast(w, DType::BF16)),
+            _ => unreachable!(),
+        }
+    }
+
     /// 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()`.
@@ -127,21 +164,42 @@ impl TinyTransformer {
         );
         let seq = total / batch;
 
-        let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim]
+        // Embedding gathers from the fp32 master table; in bf16 mode cast the
+        // activation stream to bf16 here (norms are cast to bf16 gammas too).
+        let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim], fp32
+        if self.compute_dtype == DType::BF16 {
+            h = ops::cast(&h, DType::BF16);
+        }
         for b in &self.blocks {
             // --- Attention sub-block (pre-norm + residual) ---
-            let normed = ops::rms_norm(&h, &b.attn_norm, self.cfg.eps);
+            let normed = ops::rms_norm(&h, &self.norm_gamma(&b.attn_norm), self.cfg.eps);
             let attn = self.attention(b, &normed, batch, 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 normed = ops::rms_norm(&h, &self.norm_gamma(&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) // [batch*seq, vocab]
+        let h = ops::rms_norm(&h, &self.norm_gamma(&self.final_norm), self.cfg.eps);
+        // lm_head matmul in compute dtype; cast logits back to fp32 for CE.
+        let logits = self.linear(&h, &self.lm_head); // [batch*seq, vocab]
+        if self.compute_dtype == DType::BF16 {
+            ops::cast(&logits, DType::F32)
+        } else {
+            logits
+        }
+    }
+
+    /// A norm/QK-norm gamma in the compute dtype. fp32 master leaf → bf16 (cast
+    /// op, grad upcast) in bf16 mode; identity in fp32 mode.
+    fn norm_gamma(&self, gamma: &Var) -> Var {
+        match self.compute_dtype {
+            DType::F32 => gamma.clone(),
+            DType::BF16 => ops::cast(gamma, DType::BF16),
+            _ => unreachable!(),
+        }
     }
 
     /// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
@@ -186,7 +244,7 @@ impl TinyTransformer {
                 // restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
                 Some(gamma) => {
                     let flat = ops::reshape(&r, &[total * nh, hd]);
-                    let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
+                    let normed = ops::rms_norm(&flat, &self.norm_gamma(gamma), self.cfg.eps);
                     let r = ops::reshape(&normed, &[total, nh, hd]);
                     ops::rope(&r, self.cfg.rope_theta, seq)
                 }
@@ -197,9 +255,9 @@ impl TinyTransformer {
             ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
         };
 
-        let q = to_bh(ops::matmul(x, &b.wq), Some(&b.q_norm));
-        let k = to_bh(ops::matmul(x, &b.wk), Some(&b.k_norm));
-        let v = to_bh(ops::matmul(x, &b.wv), None);
+        let q = to_bh(self.linear(x, &b.wq), Some(&b.q_norm));
+        let k = to_bh(self.linear(x, &b.wk), Some(&b.k_norm));
+        let v = to_bh(self.linear(x, &b.wv), None);
 
         // Fused batched causal SDPA over all B*nh (sequence,head) blocks at once
         // (2 batched GEMMs + 1 causal-softmax kernel; no per-head/per-seq loop).
@@ -210,15 +268,15 @@ impl TinyTransformer {
         let out = ops::reshape(&out, &[batch, nh, seq, hd]);
         let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
         let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
-        ops::matmul(&concat, &b.wo) // out projection
+        self.linear(&concat, &b.wo) // out projection
     }
 
     /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*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 gate = self.linear(x, &b.w_gate); // [seq, ffn_hidden]
+        let up = self.linear(x, &b.w_up); // [seq, ffn_hidden]
         let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
-        ops::matmul(&act, &b.w_down) // [seq, dim]
+        self.linear(&act, &b.w_down) // [seq, dim]
     }
 }
 

crates/xtrain-model/tests/bf16.rs

@@ -0,0 +1,145 @@
+// T12 bf16 mixed-precision correctness gate (on-GPU, no PyTorch).
+//
+// The SAME model (identical fp32 master weights) run in fp32 vs bf16 compute
+// mode must agree within a LOOSE bf16 tolerance (bf16 = 7-bit mantissa ≈ 2-3
+// decimal digits → ~1e-2 relative error is expected and acceptable), both for
+// the forward loss/logits AND every parameter's gradient. We also assert no
+// NaN/Inf leaks and that the fp32 grads are fp32 (the cast op upcast the bf16
+// weight grad back to the fp32 master, so AdamW/clip/DDP stay fp32).
+//
+// This is the "bf16 within looser tol vs fp32 reference" gate; the short-run
+// convergence comparison is the train_loop-level bench on dash5.
+#![cfg(not(no_cuda))]
+
+use xtrain_cuda::device;
+use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
+use xtrain_tensor::{DType, Device};
+
+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 build(cfg: Config, device: Device) -> TinyTransformer {
+    let mut seed = 1u64;
+    TinyTransformer::new(cfg, device, |shape| {
+        seed = seed.wrapping_add(1);
+        let n: usize = shape.iter().product();
+        if shape.len() == 1 {
+            fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
+        } else {
+            fill(n, seed, 0.08)
+        }
+    })
+}
+
+fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
+    t.to_device(Device::Cpu).as_slice::<f32>().to_vec()
+}
+
+#[test]
+fn bf16_matches_fp32_within_loose_tol() {
+    assert!(device::device_count().unwrap() > 0, "no CUDA device");
+    device::set_device(0).unwrap();
+    let device = Device::Cuda(0);
+
+    // A few layers / heads so the bf16 rounding accumulates through the depth
+    // the real model has (not just a single matmul).
+    let mut cfg = Config::tiny();
+    cfg.vocab = 32;
+    cfg.n_layers = 3;
+    let batch = 2usize;
+    let seq = 8usize;
+
+    let seqs: Vec<Vec<i32>> = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let tgts: Vec<Vec<i32>> = (0..batch)
+        .map(|b| {
+            (0..seq)
+                .map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
+                .collect()
+        })
+        .collect();
+    let ids = batched_ids_tensor(&seqs, device);
+    let tgt = batched_ids_tensor(&tgts, device);
+
+    // fp32 reference.
+    let fp32 = build(cfg, device);
+    let f_logits = host(&fp32.forward_batched(&ids, batch).value());
+    let f_loss = fp32.loss_batched(&ids, &tgt, batch);
+    let f_loss_val = host(&f_loss.value())[0];
+    f_loss.backward();
+    let f_params = fp32.params();
+
+    // bf16 — SAME init (build re-runs the same deterministic fill).
+    let bf16 = build(cfg, device).with_compute_dtype(DType::BF16);
+    let b_logits = host(&bf16.forward_batched(&ids, batch).value());
+    let b_loss = bf16.loss_batched(&ids, &tgt, batch);
+    let b_loss_val = host(&b_loss.value())[0];
+    b_loss.backward();
+    let b_params = bf16.params();
+
+    // No NaN/Inf in the bf16 forward.
+    assert!(
+        b_logits.iter().all(|v| v.is_finite()) && b_loss_val.is_finite(),
+        "bf16 forward produced non-finite values"
+    );
+
+    // Forward loss within loose bf16 tol.
+    let loss_rel = (b_loss_val - f_loss_val).abs() / f_loss_val.abs().max(1e-4);
+    println!("bf16 vs fp32: loss {b_loss_val:.5} vs {f_loss_val:.5} (rel {loss_rel:.3e})");
+    assert!(
+        loss_rel < 2e-2,
+        "bf16 loss too far from fp32: {loss_rel:.3e}"
+    );
+
+    // Logits: bf16 has ~2-3 decimal digits → compare on a robust (median-style)
+    // basis, requiring the bulk to be within ~3e-2 and the mean error small.
+    let n = f_logits.len();
+    let mut rels: Vec<f32> = f_logits
+        .iter()
+        .zip(&b_logits)
+        .map(|(f, b)| (b - f).abs() / f.abs().max(1.0))
+        .collect();
+    rels.sort_by(|a, b| a.partial_cmp(b).unwrap());
+    let p99 = rels[(n as f32 * 0.99) as usize];
+    let mean: f32 = rels.iter().sum::<f32>() / n as f32;
+    println!("bf16 vs fp32 logits: mean rel {mean:.3e}, p99 rel {p99:.3e}");
+    assert!(mean < 1e-2, "bf16 logits mean rel err too high: {mean:.3e}");
+    assert!(p99 < 5e-2, "bf16 logits p99 rel err too high: {p99:.3e}");
+
+    // Gradients: fp32 master grads must be fp32 (cast op upcast), finite, and
+    // within loose bf16 tol of the fp32 reference (mean over each param tensor).
+    let mut worst_param_mean = 0.0f32;
+    for (fp, bp) in f_params.iter().zip(&b_params) {
+        let bg = bp.grad().expect("bf16 grad");
+        assert_eq!(bg.dtype(), DType::F32, "bf16-mode grad must be fp32 master");
+        let fg = host(&fp.grad().expect("fp32 grad"));
+        let bg = host(&bg);
+        assert!(bg.iter().all(|v| v.is_finite()), "bf16 grad has non-finite");
+        // Scale-relative mean error over the tensor (robust to a few small entries).
+        let scale = fg.iter().map(|v| v.abs()).fold(0.0f32, f32::max).max(1e-6);
+        let mean_err: f32 =
+            fg.iter().zip(&bg).map(|(f, b)| (f - b).abs()).sum::<f32>() / fg.len() as f32 / scale;
+        worst_param_mean = worst_param_mean.max(mean_err);
+    }
+    println!("bf16 vs fp32 grads: worst per-tensor scaled-mean err = {worst_param_mean:.3e}");
+    assert!(
+        worst_param_mean < 3e-2,
+        "bf16 grads too far from fp32: {worst_param_mean:.3e}"
+    );
+}

crates/xtrain-tensor/src/dtype.rs

@@ -1,12 +1,16 @@
 //! Tensor data types.
 //!
-//! T2 only needs `F32`, but the enum + `TensorDType` trait are structured so
-//! half-precision types (F16/BF16) can be added later (T7 mixed precision)
-//! without touching call sites.
+//! T2 only needs `F32`; `BF16` was added in T12 for mixed-precision training
+//! (bf16 linears / activations, fp32 master weights — see
+//! `docs/11-bf16-mixed-precision.md`). The enum + `TensorDType` trait keep call
+//! sites dtype-polymorphic.
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum DType {
     F32,
+    /// bfloat16: 1 sign / 8 exponent / 7 mantissa. Same exponent range as f32
+    /// (so no loss scaling needed), ~2-3 decimal digits. The T12 AMP compute type.
+    BF16,
     /// 32-bit signed integers. Used for cross-entropy targets (token ids).
     I32,
 }
@@ -15,6 +19,7 @@ impl DType {
     pub fn size_bytes(self) -> usize {
         match self {
             DType::F32 => 4,
+            DType::BF16 => 2,
             DType::I32 => 4,
         }
     }
@@ -22,6 +27,7 @@ impl DType {
     pub fn name(self) -> &'static str {
         match self {
             DType::F32 => "f32",
+            DType::BF16 => "bf16",
             DType::I32 => "i32",
         }
     }
@@ -50,6 +56,16 @@ impl TensorDType for f32 {
     }
 }
 
+impl TensorDType for half::bf16 {
+    const DTYPE: DType = DType::BF16;
+    fn to_f64(self) -> f64 {
+        self.to_f64()
+    }
+    fn from_f64(v: f64) -> Self {
+        half::bf16::from_f64(v)
+    }
+}
+
 impl TensorDType for i32 {
     const DTYPE: DType = DType::I32;
     fn to_f64(self) -> f64 {

crates/xtrain-tensor/src/tensor.rs

@@ -107,6 +107,45 @@ impl Tensor {
         }
     }
 
+    // --- dtype cast (Phase T12, bf16 mixed precision) ---
+
+    /// Cast between F32 and BF16 (the AMP bridge: fp32 master ↔ bf16 compute).
+    /// Same dtype returns a cheap clone. Requires a contiguous CUDA tensor.
+    /// I32 is not castable here (only used for token-id targets).
+    #[cfg(not(no_cuda))]
+    pub fn to_dtype(&self, target: DType) -> Self {
+        if self.dtype == target {
+            return self.clone();
+        }
+        assert!(
+            matches!(self.device(), Device::Cuda(_)),
+            "to_dtype requires a CUDA tensor"
+        );
+        assert!(self.is_contiguous(), "to_dtype requires contiguous tensor");
+        let n = self.numel() as i32;
+        let out = Tensor::zeros(&self.shape, target, self.device());
+        match (self.dtype, target) {
+            (DType::F32, DType::BF16) => unsafe {
+                xtrain_cuda::ffi::launch_cast_f32_to_bf16(
+                    self.data_ptr() as *const f32,
+                    out.data_ptr() as *mut std::ffi::c_void,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            (DType::BF16, DType::F32) => unsafe {
+                xtrain_cuda::ffi::launch_cast_bf16_to_f32(
+                    self.data_ptr() as *const std::ffi::c_void,
+                    out.data_ptr() as *mut f32,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            (a, b) => panic!("unsupported cast {a} -> {b}"),
+        }
+        out
+    }
+
     // --- Host data access (CPU only) ---
 
     /// Typed read-only view of the data. Requires a contiguous CPU tensor.
@@ -136,7 +175,10 @@ impl Tensor {
     /// GPU. Available only when CUDA was compiled in (`not(no_cuda)`).
     #[cfg(not(no_cuda))]
     pub fn scale(&self, alpha: f32) -> Self {
-        assert_eq!(self.dtype, DType::F32, "scale only supports F32 in T2");
+        assert!(
+            matches!(self.dtype, DType::F32 | DType::BF16),
+            "scale supports F32/BF16"
+        );
         assert!(self.is_contiguous(), "scale requires contiguous tensor");
         assert!(
             matches!(self.device(), Device::Cuda(_)),
@@ -144,14 +186,27 @@ impl Tensor {
         );
 
         let out = Tensor::zeros(&self.shape, self.dtype, self.device());
-        unsafe {
-            xtrain_cuda::ffi::launch_scale_f32(
-                self.data_ptr() as *const f32,
-                out.data_ptr() as *mut f32,
-                alpha,
-                self.numel() as i32,
-                std::ptr::null_mut(), // default stream
-            );
+        let n = self.numel() as i32;
+        match self.dtype {
+            DType::F32 => unsafe {
+                xtrain_cuda::ffi::launch_scale_f32(
+                    self.data_ptr() as *const f32,
+                    out.data_ptr() as *mut f32,
+                    alpha,
+                    n,
+                    std::ptr::null_mut(), // default stream
+                );
+            },
+            DType::BF16 => unsafe {
+                xtrain_cuda::ffi::launch_scale_bf16(
+                    self.data_ptr() as *const std::ffi::c_void,
+                    out.data_ptr() as *mut std::ffi::c_void,
+                    alpha,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            _ => unreachable!(),
         }
         out
     }
@@ -165,8 +220,11 @@ impl Tensor {
     /// 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.dtype, other.dtype, "matmul dtype mismatch");
+        assert!(
+            matches!(self.dtype, DType::F32 | DType::BF16),
+            "matmul supports F32/BF16"
+        );
         assert_eq!(self.ndim(), 2, "matmul requires 2D lhs");
         assert_eq!(other.ndim(), 2, "matmul requires 2D rhs");
         assert_eq!(
@@ -187,19 +245,36 @@ impl Tensor {
         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());
-        xtrain_cuda::cublas::sgemm(
-            false,
-            false,
-            m,
-            n,
-            k,
-            1.0,
-            self.data_ptr() as *const f32,
-            other.data_ptr() as *const f32,
-            0.0,
-            out.data_ptr() as *mut f32,
-        );
+        let out = Tensor::zeros(&[m, n], self.dtype, self.device());
+        match self.dtype {
+            // fp32 path — unchanged (bit-identical to T7/T10/T11).
+            DType::F32 => xtrain_cuda::cublas::sgemm(
+                false,
+                false,
+                m,
+                n,
+                k,
+                1.0,
+                self.data_ptr() as *const f32,
+                other.data_ptr() as *const f32,
+                0.0,
+                out.data_ptr() as *mut f32,
+            ),
+            // bf16 path — GemmEx, bf16 in/out, fp32 accumulation.
+            DType::BF16 => xtrain_cuda::cublas::gemm_ex(
+                false,
+                false,
+                m,
+                n,
+                k,
+                1.0,
+                self.data_ptr() as *const std::ffi::c_void,
+                other.data_ptr() as *const std::ffi::c_void,
+                0.0,
+                out.data_ptr() as *mut std::ffi::c_void,
+            ),
+            _ => unreachable!(),
+        }
         out
     }
 
@@ -207,9 +282,15 @@ impl Tensor {
     /// 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");
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .transpose_2d()
+                .to_dtype(DType::BF16);
+        }
+        assert_eq!(self.dtype, DType::F32, "transpose supports F32/BF16");
         assert!(
             matches!(self.device(), Device::Cuda(_)),
             "transpose requires a CUDA tensor"
@@ -245,35 +326,69 @@ impl Tensor {
         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)");
 
+        assert_eq!(a.dtype, b.dtype, "matmul_backward dtype mismatch");
+        assert_eq!(a.dtype, dc.dtype, "matmul_backward dtype mismatch");
         let (m, k, n) = (a.shape[0], a.shape[1], b.shape[1]);
+        let dt = a.dtype;
         // dA[M,K] = dC[M,N] · Bᵀ  (B stored [K,N], transposed by cuBLAS)
-        let da = Tensor::zeros(&[m, k], DType::F32, a.device());
-        xtrain_cuda::cublas::sgemm(
-            false,
-            true,
-            m,
-            k,
-            n,
-            1.0,
-            dc.data_ptr() as *const f32,
-            b.data_ptr() as *const f32,
-            0.0,
-            da.data_ptr() as *mut f32,
-        );
+        let da = Tensor::zeros(&[m, k], dt, a.device());
         // dB[K,N] = Aᵀ · dC[M,N]  (A stored [M,K], transposed by cuBLAS)
-        let db = Tensor::zeros(&[k, n], DType::F32, a.device());
-        xtrain_cuda::cublas::sgemm(
-            true,
-            false,
-            k,
-            n,
-            m,
-            1.0,
-            a.data_ptr() as *const f32,
-            dc.data_ptr() as *const f32,
-            0.0,
-            db.data_ptr() as *mut f32,
-        );
+        let db = Tensor::zeros(&[k, n], dt, a.device());
+        match dt {
+            DType::F32 => {
+                xtrain_cuda::cublas::sgemm(
+                    false,
+                    true,
+                    m,
+                    k,
+                    n,
+                    1.0,
+                    dc.data_ptr() as *const f32,
+                    b.data_ptr() as *const f32,
+                    0.0,
+                    da.data_ptr() as *mut f32,
+                );
+                xtrain_cuda::cublas::sgemm(
+                    true,
+                    false,
+                    k,
+                    n,
+                    m,
+                    1.0,
+                    a.data_ptr() as *const f32,
+                    dc.data_ptr() as *const f32,
+                    0.0,
+                    db.data_ptr() as *mut f32,
+                );
+            }
+            DType::BF16 => {
+                xtrain_cuda::cublas::gemm_ex(
+                    false,
+                    true,
+                    m,
+                    k,
+                    n,
+                    1.0,
+                    dc.data_ptr() as *const std::ffi::c_void,
+                    b.data_ptr() as *const std::ffi::c_void,
+                    0.0,
+                    da.data_ptr() as *mut std::ffi::c_void,
+                );
+                xtrain_cuda::cublas::gemm_ex(
+                    true,
+                    false,
+                    k,
+                    n,
+                    m,
+                    1.0,
+                    a.data_ptr() as *const std::ffi::c_void,
+                    dc.data_ptr() as *const std::ffi::c_void,
+                    0.0,
+                    db.data_ptr() as *mut std::ffi::c_void,
+                );
+            }
+            _ => panic!("matmul_backward supports F32/BF16"),
+        }
         (da, db)
     }
 
@@ -287,15 +402,28 @@ impl Tensor {
     #[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(),
-            );
+        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
+        let n = self.numel() as i32;
+        match self.dtype {
+            DType::F32 => 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,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            DType::BF16 => unsafe {
+                xtrain_cuda::ffi::launch_add_bf16(
+                    self.data_ptr() as *const std::ffi::c_void,
+                    other.data_ptr() as *const std::ffi::c_void,
+                    out.data_ptr() as *mut std::ffi::c_void,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            _ => unreachable!(),
         }
         out
     }
@@ -304,15 +432,28 @@ impl Tensor {
     #[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(),
-            );
+        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
+        let n = self.numel() as i32;
+        match self.dtype {
+            DType::F32 => 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,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            DType::BF16 => unsafe {
+                xtrain_cuda::ffi::launch_mul_bf16(
+                    self.data_ptr() as *const std::ffi::c_void,
+                    other.data_ptr() as *const std::ffi::c_void,
+                    out.data_ptr() as *mut std::ffi::c_void,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            _ => unreachable!(),
         }
         out
     }
@@ -324,17 +465,31 @@ impl Tensor {
         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");
+        assert_eq!(self.dtype, bias.dtype, "add_bias dtype mismatch");
         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(),
-            );
+        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
+        match self.dtype {
+            DType::F32 => 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(),
+                );
+            },
+            DType::BF16 => unsafe {
+                xtrain_cuda::ffi::launch_add_bias_bf16(
+                    self.data_ptr() as *const std::ffi::c_void,
+                    bias.data_ptr() as *const std::ffi::c_void,
+                    out.data_ptr() as *mut std::ffi::c_void,
+                    rows as i32,
+                    cols as i32,
+                    std::ptr::null_mut(),
+                );
+            },
+            _ => panic!("add_bias supports F32/BF16"),
         }
         out
     }
@@ -346,15 +501,27 @@ impl Tensor {
     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(),
-            );
+        let out = Tensor::zeros(&[cols], self.dtype, self.device());
+        match self.dtype {
+            DType::F32 => 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(),
+                );
+            },
+            DType::BF16 => unsafe {
+                xtrain_cuda::ffi::launch_sum_rows_bf16(
+                    self.data_ptr() as *const std::ffi::c_void,
+                    out.data_ptr() as *mut std::ffi::c_void,
+                    rows as i32,
+                    cols as i32,
+                    std::ptr::null_mut(),
+                );
+            },
+            _ => panic!("sum_rows supports F32/BF16"),
         }
         out
     }
@@ -367,6 +534,14 @@ impl 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");
+        // bf16: compute the reduction in fp32 (standard AMP), downcast y back to
+        // bf16. inv_rms stays fp32 (the cache the fp32 backward kernel consumes).
+        if self.dtype == DType::BF16 {
+            let (y, inv_rms) = self
+                .to_dtype(DType::F32)
+                .rms_norm(&gamma.to_dtype(DType::F32), eps);
+            return (y.to_dtype(DType::BF16), inv_rms);
+        }
         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());
@@ -394,6 +569,17 @@ impl Tensor {
         dy: &Tensor,
         inv_rms: &Tensor,
     ) -> (Tensor, Tensor) {
+        // bf16: upcast (x, gamma, dy) to fp32, run the fp32 backward, downcast the
+        // grads back to bf16 (inv_rms is already the fp32 cache).
+        if x.dtype == DType::BF16 {
+            let (dx, dgamma) = Tensor::rms_norm_backward(
+                &x.to_dtype(DType::F32),
+                &gamma.to_dtype(DType::F32),
+                &dy.to_dtype(DType::F32),
+                inv_rms,
+            );
+            return (dx.to_dtype(DType::BF16), dgamma.to_dtype(DType::BF16));
+        }
         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());
@@ -424,15 +610,30 @@ impl Tensor {
     /// 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(),
-            );
+        assert!(
+            matches!(self.dtype, DType::F32 | DType::BF16),
+            "silu supports F32/BF16"
+        );
+        let out = Tensor::zeros(&self.shape, self.dtype, self.device());
+        let n = self.numel() as i32;
+        match self.dtype {
+            DType::F32 => unsafe {
+                xtrain_cuda::ffi::launch_silu_f32(
+                    self.data_ptr() as *const f32,
+                    out.data_ptr() as *mut f32,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            DType::BF16 => unsafe {
+                xtrain_cuda::ffi::launch_silu_bf16(
+                    self.data_ptr() as *const std::ffi::c_void,
+                    out.data_ptr() as *mut std::ffi::c_void,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            _ => unreachable!(),
         }
         out
     }
@@ -441,15 +642,28 @@ impl Tensor {
     /// 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(),
-            );
+        let dx = Tensor::zeros(&x.shape, x.dtype, x.device());
+        let n = x.numel() as i32;
+        match x.dtype {
+            DType::F32 => 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,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            DType::BF16 => unsafe {
+                xtrain_cuda::ffi::launch_silu_dx_bf16(
+                    x.data_ptr() as *const std::ffi::c_void,
+                    dy.data_ptr() as *const std::ffi::c_void,
+                    dx.data_ptr() as *mut std::ffi::c_void,
+                    n,
+                    std::ptr::null_mut(),
+                );
+            },
+            _ => panic!("silu_backward supports F32/BF16"),
         }
         dx
     }
@@ -468,6 +682,12 @@ impl Tensor {
             period > 0 && tokens % period == 0,
             "tokens must be a multiple of period"
         );
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .rope(theta, period)
+                .to_dtype(DType::BF16);
+        }
         let out = Tensor::zeros(&self.shape, DType::F32, self.device());
         unsafe {
             xtrain_cuda::ffi::launch_rope_f32(
@@ -488,6 +708,10 @@ impl Tensor {
     /// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
     #[cfg(not(no_cuda))]
     pub fn rope_backward(dy: &Tensor, theta: f32, period: usize) -> Self {
+        if dy.dtype == DType::BF16 {
+            return Tensor::rope_backward(&dy.to_dtype(DType::F32), theta, period)
+                .to_dtype(DType::BF16);
+        }
         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 {
@@ -509,6 +733,9 @@ impl Tensor {
     #[cfg(not(no_cuda))]
     pub fn softmax(&self) -> Self {
         assert_eq!(self.ndim(), 2, "softmax requires 2D input");
+        if self.dtype == DType::BF16 {
+            return self.to_dtype(DType::F32).softmax().to_dtype(DType::BF16);
+        }
         let (rows, cols) = (self.shape[0], self.shape[1]);
         let out = Tensor::zeros(&self.shape, DType::F32, self.device());
         unsafe {
@@ -527,6 +754,10 @@ impl Tensor {
     /// Inputs are the forward output `y` and upstream `dy`.
     #[cfg(not(no_cuda))]
     pub fn softmax_backward(y: &Tensor, dy: &Tensor) -> Self {
+        if y.dtype == DType::BF16 {
+            return Tensor::softmax_backward(&y.to_dtype(DType::F32), &dy.to_dtype(DType::F32))
+                .to_dtype(DType::BF16);
+        }
         let (rows, cols) = (y.shape[0], y.shape[1]);
         let dx = Tensor::zeros(&y.shape, DType::F32, y.device());
         unsafe {
@@ -550,6 +781,11 @@ impl 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");
+        // CE math (log-sum-exp) is fp32 (probs/loss cached fp32). The model casts
+        // logits→fp32 before CE; this guard keeps the op robust to bf16 logits.
+        if self.dtype == DType::BF16 {
+            return self.to_dtype(DType::F32).cross_entropy(target);
+        }
         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());
@@ -658,9 +894,15 @@ impl Tensor {
     /// 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");
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .transpose_3d01()
+                .to_dtype(DType::BF16);
+        }
+        assert_eq!(self.dtype, DType::F32, "transpose_3d01 supports F32/BF16");
         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 {
@@ -689,54 +931,59 @@ impl Tensor {
         assert_eq!(self.ndim(), 3, "attention Q must be [bh,seq,head_dim]");
         assert_eq!(self.shape(), k.shape(), "Q/K shape mismatch");
         assert_eq!(self.shape(), v.shape(), "Q/V shape mismatch");
+        assert_eq!(self.dtype, k.dtype, "Q/K dtype mismatch");
+        assert_eq!(self.dtype, v.dtype, "Q/V dtype mismatch");
         let (bh, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
         let dev = self.device();
+        let dt = self.dtype;
 
-        // scores[bh,seq,seq] = Q[bh,seq,hd] · Kᵀ[bh,hd,seq]
-        let scores = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
-        xtrain_cuda::cublas::sgemm_strided_batched(
+        // scores[bh,seq,seq] = Q[bh,seq,hd] · Kᵀ[bh,hd,seq]  (GEMM in self dtype)
+        let scores = Tensor::zeros(&[bh, seq, seq], dt, dev);
+        strided_batched_gemm(
+            dt,
             false,
             true,
             seq,
             seq,
             hd,
-            1.0,
-            self.data_ptr() as *const f32,
+            self.data_ptr(),
             seq * hd,
-            k.data_ptr() as *const f32,
+            k.data_ptr(),
             seq * hd,
-            0.0,
-            scores.data_ptr() as *mut f32,
+            scores.data_ptr(),
             seq * seq,
             bh,
         );
-        // probs = softmax(causal(scores · scale)), one block per [bh·seq] row.
-        let probs = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
+        // probs = softmax(causal(scores · scale)). Softmax math is fp32 (stable);
+        // for bf16 we upcast scores → f32 → kernel → downcast probs back to bf16
+        // (so the cached probs activation is half-size). One block per [bh·seq] row.
+        let scores_f32 = scores.to_dtype(DType::F32);
+        let probs_f32 = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
         unsafe {
             xtrain_cuda::ffi::launch_softmax_causal_f32(
-                scores.data_ptr() as *const f32,
-                probs.data_ptr() as *mut f32,
+                scores_f32.data_ptr() as *const f32,
+                probs_f32.data_ptr() as *mut f32,
                 (bh * seq) as i32,
                 seq as i32,
                 scale,
                 std::ptr::null_mut(),
             );
         }
+        let probs = probs_f32.to_dtype(dt);
         // out[bh,seq,hd] = probs[bh,seq,seq] · V[bh,seq,hd]
-        let out = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
-        xtrain_cuda::cublas::sgemm_strided_batched(
+        let out = Tensor::zeros(&[bh, seq, hd], dt, dev);
+        strided_batched_gemm(
+            dt,
             false,
             false,
             seq,
             hd,
             seq,
-            1.0,
-            probs.data_ptr() as *const f32,
+            probs.data_ptr(),
             seq * seq,
-            v.data_ptr() as *const f32,
+            v.data_ptr(),
             seq * hd,
-            0.0,
-            out.data_ptr() as *mut f32,
+            out.data_ptr(),
             seq * hd,
             bh,
         );
@@ -764,45 +1011,44 @@ impl Tensor {
     ) -> (Tensor, Tensor, Tensor) {
         let (bh, seq, hd) = (q.shape[0], q.shape[1], q.shape[2]);
         let dev = q.device();
+        let dt = q.dtype;
 
         // dP[bh,seq,seq] = dOut[bh,seq,hd] · Vᵀ[bh,hd,seq]
-        let dp = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
-        xtrain_cuda::cublas::sgemm_strided_batched(
+        let dp = Tensor::zeros(&[bh, seq, seq], dt, dev);
+        strided_batched_gemm(
+            dt,
             false,
             true,
             seq,
             seq,
             hd,
-            1.0,
-            dout.data_ptr() as *const f32,
+            dout.data_ptr(),
             seq * hd,
-            v.data_ptr() as *const f32,
+            v.data_ptr(),
             seq * hd,
-            0.0,
-            dp.data_ptr() as *mut f32,
+            dp.data_ptr(),
             seq * seq,
             bh,
         );
         // dV[bh,seq,hd] = Pᵀ[bh,seq,seq] · dOut[bh,seq,hd]
-        let dv = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
-        xtrain_cuda::cublas::sgemm_strided_batched(
+        let dv = Tensor::zeros(&[bh, seq, hd], dt, dev);
+        strided_batched_gemm(
+            dt,
             true,
             false,
             seq,
             hd,
             seq,
-            1.0,
-            probs.data_ptr() as *const f32,
+            probs.data_ptr(),
             seq * seq,
-            dout.data_ptr() as *const f32,
+            dout.data_ptr(),
             seq * hd,
-            0.0,
-            dv.data_ptr() as *mut f32,
+            dv.data_ptr(),
             seq * hd,
             bh,
         );
         // dScores = softmax Jacobian (per row) applied to dP, then ×scale.
-        // Reuse the row-wise softmax backward over the flattened [bh·seq, seq].
+        // softmax_backward + scale are dtype-aware (fp32 math inside for bf16).
         let dscores = Tensor::softmax_backward(
             &probs.reshape(&[bh * seq, seq]),
             &dp.reshape(&[bh * seq, seq]),
@@ -810,38 +1056,36 @@ impl Tensor {
         .reshape(&[bh, seq, seq]);
         let dscores = dscores.scale(scale);
         // dQ[bh,seq,hd] = dScores[bh,seq,seq] · K[bh,seq,hd]
-        let dq = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
-        xtrain_cuda::cublas::sgemm_strided_batched(
+        let dq = Tensor::zeros(&[bh, seq, hd], dt, dev);
+        strided_batched_gemm(
+            dt,
             false,
             false,
             seq,
             hd,
             seq,
-            1.0,
-            dscores.data_ptr() as *const f32,
+            dscores.data_ptr(),
             seq * seq,
-            k.data_ptr() as *const f32,
+            k.data_ptr(),
             seq * hd,
-            0.0,
-            dq.data_ptr() as *mut f32,
+            dq.data_ptr(),
             seq * hd,
             bh,
         );
         // dK[bh,seq,hd] = dScoresᵀ[bh,seq,seq] · Q[bh,seq,hd]
-        let dk = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
-        xtrain_cuda::cublas::sgemm_strided_batched(
+        let dk = Tensor::zeros(&[bh, seq, hd], dt, dev);
+        strided_batched_gemm(
+            dt,
             true,
             false,
             seq,
             hd,
             seq,
-            1.0,
-            dscores.data_ptr() as *const f32,
+            dscores.data_ptr(),
             seq * seq,
-            q.data_ptr() as *const f32,
+            q.data_ptr(),
             seq * hd,
-            0.0,
-            dk.data_ptr() as *mut f32,
+            dk.data_ptr(),
             seq * hd,
             bh,
         );
@@ -853,9 +1097,15 @@ impl Tensor {
     /// (`[B,S,nh,hd] <-> [B,nh,S,hd]`). Its own backward is the same op (swap b,c).
     #[cfg(not(no_cuda))]
     pub fn transpose_4d12(&self) -> Self {
-        assert_eq!(self.dtype, DType::F32, "transpose_4d12 only supports F32");
         assert_eq!(self.ndim(), 4, "transpose_4d12 requires a 4D tensor");
         assert!(self.is_contiguous(), "transpose_4d12 requires contiguous");
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .transpose_4d12()
+                .to_dtype(DType::BF16);
+        }
+        assert_eq!(self.dtype, DType::F32, "transpose_4d12 supports F32/BF16");
         let (a, b, c, d) = (self.shape[0], self.shape[1], self.shape[2], self.shape[3]);
         let out = Tensor::zeros(&[a, c, b, d], DType::F32, self.device());
         unsafe {
@@ -875,8 +1125,11 @@ impl Tensor {
     // 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!(
+            matches!(self.dtype, DType::F32 | DType::BF16),
+            "{op} supports F32/BF16"
+        );
+        assert_eq!(self.dtype, other.dtype, "{op} dtype mismatch");
         assert_eq!(self.shape(), other.shape(), "{op} shape mismatch");
         assert_eq!(self.device(), other.device(), "{op} device mismatch");
         assert!(
@@ -886,6 +1139,64 @@ impl Tensor {
     }
 }
 
+/// Dispatch a strided-batched GEMM on `dt`: fp32 → `sgemm_strided_batched`,
+/// bf16 → `gemm_ex_strided_batched` (bf16 in/out, fp32 accum). Pointers are the
+/// raw `data_ptr()` bytes of contiguous same-dtype tensors. `alpha=1, beta=0`.
+/// The fp32 path is bit-identical to the inlined T10 call it replaces.
+#[cfg(not(no_cuda))]
+#[allow(clippy::too_many_arguments)]
+fn strided_batched_gemm(
+    dt: DType,
+    trans_a: bool,
+    trans_b: bool,
+    m: usize,
+    n: usize,
+    k: usize,
+    a: *const u8,
+    stride_a: usize,
+    b: *const u8,
+    stride_b: usize,
+    c: *const u8,
+    stride_c: usize,
+    batch: usize,
+) {
+    match dt {
+        DType::F32 => xtrain_cuda::cublas::sgemm_strided_batched(
+            trans_a,
+            trans_b,
+            m,
+            n,
+            k,
+            1.0,
+            a as *const f32,
+            stride_a,
+            b as *const f32,
+            stride_b,
+            0.0,
+            c as *mut f32,
+            stride_c,
+            batch,
+        ),
+        DType::BF16 => xtrain_cuda::cublas::gemm_ex_strided_batched(
+            trans_a,
+            trans_b,
+            m,
+            n,
+            k,
+            1.0,
+            a as *const std::ffi::c_void,
+            stride_a,
+            b as *const std::ffi::c_void,
+            stride_b,
+            0.0,
+            c as *mut std::ffi::c_void,
+            stride_c,
+            batch,
+        ),
+        _ => panic!("strided_batched_gemm supports F32/BF16"),
+    }
+}
+
 impl std::fmt::Debug for Tensor {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(

crates/xtrain-train/src/bin/train.rs

@@ -31,6 +31,8 @@ use xtrain_cuda::device;
 #[cfg(not(no_cuda))]
 use xtrain_model::{Config, TinyTransformer};
 #[cfg(not(no_cuda))]
+use xtrain_tensor::DType;
+#[cfg(not(no_cuda))]
 use xtrain_tensor::Device;
 #[cfg(not(no_cuda))]
 use xtrain_train::data::Corpus;
@@ -107,6 +109,9 @@ fn main() {
     let val_tokens: usize = flag(&args, "--val-tokens", 0);
     let eval_every: usize = flag(&args, "--eval-every", 0);
     let eval_batches: usize = flag(&args, "--eval-batches", 64);
+    // bf16 mixed precision (Phase T12): fp32 master weights, bf16 linears +
+    // activations. Opt-in; default fp32 reproduces v0–v4 numerics.
+    let bf16 = args.iter().any(|a| a == "--bf16");
     let ckpt: PathBuf = PathBuf::from(
         args.iter()
             .position(|a| a == "--ckpt")
@@ -155,7 +160,7 @@ fn main() {
     );
 
     let mut seed = 1u64;
-    let model = TinyTransformer::new(cfg, device, |shape| {
+    let mut model = TinyTransformer::new(cfg, device, |shape| {
         seed = seed.wrapping_add(1);
         let n: usize = shape.iter().product();
         if shape.len() == 1 {
@@ -166,6 +171,10 @@ fn main() {
             fill(n, seed, 0.04)
         }
     });
+    if bf16 {
+        model = model.with_compute_dtype(DType::BF16);
+        println!("bf16 mixed precision: ON (fp32 master weights)");
+    }
 
     // Eval-only mode: load a checkpoint and score it on the held-out val set, then
     // exit. Used to put an EXISTING model (e.g. v0) and a new one on the same

csrc/ops/cast.cu

@@ -0,0 +1,141 @@
+// bf16 mixed-precision kernels (Phase T12, KI-2).
+//
+// Two groups:
+//   1. f32 <-> bf16 cast — the bridge between fp32 master weights / fp32
+//      reductions and the bf16 compute/activation stream.
+//   2. bf16 elementwise ops (add / mul / silu / scale + their backwards) — the
+//      residual-stream ops that flow bf16 activations. Each loads bf16 -> float,
+//      computes in fp32, stores bf16 (so the math accumulates in fp32 while the
+//      stored activation is half-size). Matmuls go through cuBLAS GemmEx
+//      (cublas.rs); norm / softmax / rope / cross-entropy stay fp32 (the Rust
+//      wrappers upcast around the existing fp32 kernels).
+//
+// bf16 is __nv_bfloat16; __float2bfloat16 / __bfloat162float round-trip via fp32.
+
+#include <cuda_bf16.h>
+
+extern "C" {
+
+// --- f32 <-> bf16 cast ---
+
+__global__ void cast_f32_to_bf16_k(const float* in, __nv_bfloat16* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = __float2bfloat16(in[i]);
+}
+void launch_cast_f32_to_bf16(const float* in, void* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    cast_f32_to_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        in, (__nv_bfloat16*)out, n);
+}
+
+__global__ void cast_bf16_to_f32_k(const __nv_bfloat16* in, float* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = __bfloat162float(in[i]);
+}
+void launch_cast_bf16_to_f32(const void* in, float* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    cast_bf16_to_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)in, out, n);
+}
+
+// --- bf16 elementwise (load->fp32->compute->store bf16) ---
+
+__global__ void add_bf16_k(const __nv_bfloat16* a, const __nv_bfloat16* b,
+                           __nv_bfloat16* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n)
+        out[i] = __float2bfloat16(__bfloat162float(a[i]) + __bfloat162float(b[i]));
+}
+void launch_add_bf16(const void* a, const void* b, void* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    add_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
+}
+
+__global__ void mul_bf16_k(const __nv_bfloat16* a, const __nv_bfloat16* b,
+                           __nv_bfloat16* out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n)
+        out[i] = __float2bfloat16(__bfloat162float(a[i]) * __bfloat162float(b[i]));
+}
+void launch_mul_bf16(const void* a, const void* b, void* out, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    mul_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)a, (const __nv_bfloat16*)b, (__nv_bfloat16*)out, n);
+}
+
+__global__ void scale_bf16_k(const __nv_bfloat16* in, __nv_bfloat16* out,
+                             float alpha, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = __float2bfloat16(__bfloat162float(in[i]) * alpha);
+}
+void launch_scale_bf16(const void* in, void* out, float alpha, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    scale_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)in, (__nv_bfloat16*)out, alpha, n);
+}
+
+// SiLU: y = x*sigmoid(x). Backward: dx = dy * (sig + x*sig*(1-sig)).
+__global__ void silu_bf16_k(const __nv_bfloat16* x, __nv_bfloat16* y, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float v = __bfloat162float(x[i]);
+        float sig = 1.0f / (1.0f + expf(-v));
+        y[i] = __float2bfloat16(v * sig);
+    }
+}
+void launch_silu_bf16(const void* x, void* y, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    silu_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)x, (__nv_bfloat16*)y, n);
+}
+
+__global__ void silu_dx_bf16_k(const __nv_bfloat16* x, const __nv_bfloat16* dy,
+                               __nv_bfloat16* dx, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        float v = __bfloat162float(x[i]);
+        float sig = 1.0f / (1.0f + expf(-v));
+        float g = sig + v * sig * (1.0f - sig);
+        dx[i] = __float2bfloat16(__bfloat162float(dy[i]) * g);
+    }
+}
+void launch_silu_dx_bf16(const void* x, const void* dy, void* dx, int n, void* s) {
+    int blk = 256, grid = (n + blk - 1) / blk;
+    silu_dx_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)dy, (__nv_bfloat16*)dx, n);
+}
+
+// Broadcast bias add: out[r,c] = x[r,c] + bias[c]. x:[rows,cols], bias:[cols].
+__global__ void add_bias_bf16_k(const __nv_bfloat16* x, const __nv_bfloat16* bias,
+                                __nv_bfloat16* out, int rows, int cols) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < rows * cols)
+        out[i] = __float2bfloat16(__bfloat162float(x[i]) +
+                                  __bfloat162float(bias[i % cols]));
+}
+void launch_add_bias_bf16(const void* x, const void* bias, void* out, int rows,
+                          int cols, void* s) {
+    int blk = 256, grid = (rows * cols + blk - 1) / blk;
+    add_bias_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)bias, (__nv_bfloat16*)out,
+        rows, cols);
+}
+
+// Column-sum over rows: dbias[c] = sum_r dout[r,c] (bias backward), fp32 accum.
+__global__ void sum_rows_bf16_k(const __nv_bfloat16* dout, __nv_bfloat16* dbias,
+                                int rows, int cols) {
+    int c = blockIdx.x * blockDim.x + threadIdx.x;
+    if (c < cols) {
+        float acc = 0.0f;
+        for (int r = 0; r < rows; ++r) acc += __bfloat162float(dout[r * cols + c]);
+        dbias[c] = __float2bfloat16(acc);
+    }
+}
+void launch_sum_rows_bf16(const void* dout, void* dbias, int rows, int cols, void* s) {
+    int blk = 256, grid = (cols + blk - 1) / blk;
+    sum_rows_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
+        (const __nv_bfloat16*)dout, (__nv_bfloat16*)dbias, rows, cols);
+}
+
+}  // extern "C"

docs/11-bf16-mixed-precision.md

@@ -0,0 +1,83 @@
+# Phase T12: bf16 混合精度（fp32 master）— Design Document
+
+> KI-2 的具体落地。v4（dim768, fp32）在单卡 32GB 下 per-rank batch 32（global 256）**OOM**，被迫降到 batch 16 训练。bf16 把激活显存减半（找回 batch-256 甜点区），并在 dim768 这个已 compute-bound 的规模上加速 tensor-core GEMM。附带收益：xserv 推理是 **BF16-only**，bf16 训练让闭环更贴。
+
+## Goal
+
+在**不动 fp32 路径任何数值**的前提下，新增一个 **opt-in 的 bf16 混合精度模式**（标准 AMP，fp32 master weights）：
+
+1. **正确性硬闸门**：fp32 全套回归（T3 GEMM / T4 12 算子 grad-check / T5 结构+overfit+PyTorch 对拍 / T6 AdamW+checkpoint / T8 DDP / T10 batched / xserv 闭环）在**同样紧容差**下保持绿。bf16 是**加法、可选**的，绝不扰动 fp32。
+2. **bf16 正确性**：bf16 前向/梯度在**更松的 bf16 容差**（≈2–3 位十进制有效数字 → ~1e-2 相对误差）内对住 fp32 参考；一段**短 bf16 训练收敛对住 fp32**（loss 曲线接近、无 NaN/发散）。
+3. **显存+吞吐（payoff）**：dim768 bf16 能跑 per-rank batch 32（解 OOM）；测 dim768 bf16 vs fp32 的显存+tok/s。
+
+## 什么是 bf16、什么是 fp32（标准 AMP split）
+
+| 组件 | 精度 | 理由 |
+|---|---|---|
+| **master weights** + AdamW state(m/v) + 优化器更新 | **fp32** | 小步长更新需要 fp32 精度，否则被 bf16 的 8-bit 尾数吃掉 |
+| **linear GEMM**（q/k/v/o、gate/up/down、lm_head） | **bf16 in/out + fp32 accum** | compute+memory 主体；tensor-core 走 `cublasGemmEx`（`CUDA_R_16BF` in/out，`CUBLAS_COMPUTE_32F` 累加） |
+| **激活流**（残差流、attention Q/K/V/probs/out、MLP 中间） | **bf16** | 激活显存减半——这是解 OOM 的关键，不只是 GEMM 提速 |
+| **RMSNorm / QK-norm** | **fp32**（bf16→fp32 算 reduction→bf16） | 求和/rsqrt 数值敏感 |
+| **softmax（attention）/ RoPE / cross-entropy** | **fp32** | softmax 的 exp/求和、CE 的 log、RoPE 的 sin/cos 都数值敏感 |
+| **梯度 → AdamW** | **fp32** | 见下「cast 算子」——grad 在 fp32 master leaf 上累加，AdamW/clip/DDP all-reduce 全程 fp32、**完全不改** |
+
+**无 loss scaling**：bf16 是 8-bit 指数（与 fp32 同动态范围），不像 fp16（5-bit 指数易下溢）。所以梯度不会下溢到 0，**不需要** loss scaling。
+
+## Module Layout（surgical：fp32 路径逐字节不动）
+
+核心思路：**所有 op 按 `self.dtype()` 分派**。fp32 分支跑原 kernel（一字不改）；bf16 分支是新增代码。
+
+### 1. `xtrain-tensor::dtype` — 加 `BF16`
+- `DType::BF16`，`size_bytes()=2`，`half::bf16` 实现 `TensorDType`（`half` crate 已是依赖）。
+
+### 2. `xtrain-cuda` — bf16 GEMM + cast kernel
+- `ffi.rs`：声明 `cublasGemmEx` / `cublasGemmStridedBatchedEx`（void* 指针、`a_type/b_type/c_type/compute_type`），常量 `CUDA_R_16BF=14`、`CUDA_R_32F=0`、`CUBLAS_COMPUTE_32F=68`（数值同 xserv `gemm.rs`）。
+- `cublas.rs`：`gemm_ex(...)` / `gemm_ex_strided_batched(...)`——和 `sgemm` 同样的 row-major⟺col-major 转置代数，只是走 `GemmEx`、in/out=bf16、accum=fp32。
+- `csrc/ops/cast.cu` + ffi：`launch_cast_f32_to_bf16` / `launch_cast_bf16_to_f32`（逐元素 `__float2bfloat16` / `__bfloat162float`）。
+
+### 3. `xtrain-tensor::tensor` — dtype-polymorphic ops
+- `to_dtype(target)`：f32↔bf16 cast（CUDA），同 dtype 直接 clone。
+- `matmul` / `matmul_backward` / `attention` / `attention_backward`：按 dtype 分派——fp32 走原 `sgemm`（**不动**），bf16 走 `gemm_ex`，输出同 dtype。
+- 逐元素 op（`add`/`mul`/`silu`/`scale`…）+ `embedding`：允许 bf16 输入。逐元素 kernel 对 bf16 走「load→fp32→算→store bf16」（新增 bf16 kernel）或对 norm/softmax/CE 在 wrapper 里 upcast→fp32 kernel→downcast。**fp32 调用走原 f32 kernel 不变。**
+
+### 4. `xtrain-autodiff::ops` — `cast` 算子 + bf16 透传
+- **`cast(x, target_dtype)`**：前向 `x.to_dtype(target)`；**反向把 grad cast 回 `x` 的 dtype**。这是 AMP 的关键钩子：
+  - fp32 master weight leaf → `cast(w, BF16)` 喂给 matmul；matmul 的 bf16 grad 经 cast 反向**升回 fp32**，累加在 fp32 leaf 上。
+  - ⟹ `.grad()` 是 **fp32**，AdamW / `clip_grad_norm_gpu` / DDP `all_reduce_average_grads` **一行不改**，全程 fp32 master。
+- 其它 op 的 backward 自然按张量 dtype 流转（softmax/rms_norm wrapper 内部 upcast→fp32→downcast，对外是 bf16）。
+
+### 5. `xtrain-model::TinyTransformer` — `compute_dtype` 开关
+- `new_amp(cfg, device, dtype, init)` 或 `forward_batched` 接 `compute_dtype: DType`（默认 `F32` = 原路径，逐字节同）。
+- bf16 模式：embedding 输出 `cast→bf16` 进入残差流；每个 weight matmul 前 `cast(w, BF16)`；norm/softmax/rope 对 bf16 激活自动 fp32 内算；最后 logits `cast→fp32` 给 cross_entropy。**fp32 模式 `compute_dtype==F32` 时跳过所有 cast，graph 与 T10/T11 完全一致。**
+
+### 6. `xtrain-train` / `xtrain-distributed` — `--bf16` flag
+- `TrainConfig`/`DdpConfig` 加 `compute_dtype: DType`；`train.rs`/`train_ddp.rs` 加 `--bf16` flag。
+- AdamW / clip / checkpoint / DDP all-reduce **不改**（master 永远 fp32，grad 永远 fp32）。
+
+## Key Design Decisions
+
+- **cast 算子承载 fp32 master ↔ bf16 compute 的桥**：不需要在优化器里维护一份独立的 bf16 weight 副本——fp32 leaf 即 master，前向临时 cast 出 bf16，反向 grad 自动升回 fp32。最小改动、零优化器侵入。
+- **按 dtype 分派而非新类型**：fp32 路径走的还是同一个函数的 `F32` 分支 → 原 kernel、原 cuBLAS 调用、原 launch 顺序，数值逐字节不变（满足硬闸门）。
+- **norm/softmax/CE 不写 bf16 reduction kernel**：wrapper 里 `to_dtype(F32)` → 复用现有 fp32 kernel → `to_dtype(BF16)`。多两个 cast launch，但**复用已验证的 fp32 数值**，且这些不是显存/算力大头。
+- **无 loss scaling**：bf16 8-bit 指数，省掉 fp16 那套 scale/unscale/inf-check。
+
+## 验证方法（双闸门）
+
+### 闸门 ① fp32 不回归（hard gate）
+全套现有测试在原紧容差下保持绿（bf16 是 opt-in，默认 dtype=F32）：
+- `cargo test` 全 crate：grad-check（rel ≤2e-2）、structural、GEMM 对 cuBLAS（~1e-7）、batched==looped、overfit 27/27、AdamW GPU bit-exact + host 对 torch、checkpoint 逐位、DDP loss 对单卡 <1e-6、**PyTorch 对拍**（loss/logits/grad）。
+- **xserv 闭环**：v4 ckpt（fp32 训）重导 safetensors md5 一致 + xserv 贪心逐 token 对住。
+
+### 闸门 ② bf16 正确性 + 收敛
+- **bf16 looser-tol 数值**：同一组随机权重/输入，bf16 forward logits 与 bf16 grad 对 fp32 参考在 **rel ~1e-2**（bf16 2–3 位有效数字）内。
+- **短训练收敛**：dim768（或缩小代理）bf16 跑数百步，loss 曲线对住 fp32，无 NaN/发散，end loss 接近。
+
+### 闸门 ③ 显存 + 吞吐（payoff）
+- **dim768 bf16 能跑 per-rank batch 32**（v4 OOM 的触发点）。
+- 测 dim768 **bf16 vs fp32** 的峰值显存 + steady-state tok/s（预期：显存↓、tok/s↑）。
+
+## 实测结果（dash5, 8× RTX 5090, sm_120）
+
+> 见 `docs/known-issues.md` KI-2 的 before→after 表（fp32 batch32 OOM → bf16 batch32 fit；显存 A→B；tok/s A→B）。
+</content>
+</invoke>