gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

wip: T10 batched forward (validation)

Browse Source
This commit is contained in:
Gahow Wang
2026-06-16 00:19:26 +08:00
parent d2a585c5cb
commit 67687ec8fe
17 changed files with 959 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
crates/xtrain-autodiff/src/ops.rs
View File

@@ -120,15 +120,17 @@ pub fn swiglu(gate: &Var, up: &Var) -> Var {
mul(&silu(gate), up)
}
/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]`. Orthogonal map, so the
/// backward is the inverse rotation of `dy` — no cached forward values needed.
pub fn rope(x: &Var, theta: f32) -> Var {
let out = x.value().rope(theta);
/// RoPE (rotate_half) over `x:[tokens,heads,head_dim]` with per-sequence position
/// `row % period` (`period` = sequence length; `period == tokens` for a single
/// sequence). Orthogonal map, so the backward is the inverse rotation of `dy` — no
/// cached forward values needed.
pub fn rope(x: &Var, theta: f32, period: usize) -> Var {
let out = x.value().rope(theta, period);
Var::from_op(
out,
vec![x.clone()],
Box::new(move |dy, parents| {
Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta));
Var::push_grad(&parents[0], Tensor::rope_backward(dy, theta, period));
}),
)
}
@@ -190,6 +192,20 @@ pub fn transpose_3d01(x: &Var) -> Var {
)
}
/// 4D axis-(1,2) transpose `[a,b,c,d] -> [a,c,b,d]`. Self-inverse structure: the
/// backward is the same transpose applied to the grad. Lays out the batched
/// multi-head attention `[B,S,nh,hd] <-> [B,nh,S,hd]`.
pub fn transpose_4d12(x: &Var) -> Var {
let out = x.value().transpose_4d12();
Var::from_op(
out,
vec![x.clone()],
Box::new(|d, parents| {
Var::push_grad(&parents[0], d.transpose_4d12());
}),
)
}
/// 2D transpose `[r,c] -> [c,r]` as an autograd node (backward transposes the
/// grad back). Used for `Kᵀ` in attention scores.
pub fn transpose_2d(x: &Var) -> Var {
@@ -266,6 +282,29 @@ pub fn merge_heads(heads_v: &[Var]) -> Var {
)
}
/// Batched causal scaled-dot-product attention. `q`,`k`,`v` are each
/// `[bh, seq, head_dim]` (bh = batch·n_heads). Returns `[bh, seq, head_dim]`.
/// One fused op (2 batched GEMMs + 1 causal-softmax kernel forward; 4 batched
/// GEMMs + 1 softmax-backward kernel in backward) — replaces the per-(batch,head)
/// matmul/softmax loop, so attention is a handful of launches regardless of bh.
/// Caches the softmax `probs` for backward.
pub fn attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
let (out, probs) = q.value().attention(&k.value(), &v.value(), scale);
Var::from_op(
out,
vec![q.clone(), k.clone(), v.clone()],
Box::new(move |dout, parents| {
let q = parents[0].value();
let k = parents[1].value();
let v = parents[2].value();
let (dq, dk, dv) = Tensor::attention_backward(&q, &k, &v, &probs, dout, scale);
Var::push_grad(&parents[0], dq);
Var::push_grad(&parents[1], dk);
Var::push_grad(&parents[2], dv);
}),
)
}
/// Cross-entropy mean loss over logits `x:[rows,cols]` with one I32 target per
/// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
/// scaled by the upstream scalar grad.

128
crates/xtrain-autodiff/tests/autograd.rs
View File

@@ -327,12 +327,12 @@ fn rope_bwd() {
let w = fill(n, 82);
let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim]));
let out = ops::rope(&x, theta);
let out = ops::rope(&x, theta, tokens);
scalar_loss(&out, &w).backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let wf = w.clone();
let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).rope(theta), &wf);
let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).rope(theta, tokens), &wf);
report(
"rope dX",
&grad_check(
@@ -345,6 +345,38 @@ fn rope_bwd() {
);
}
// ---- rope batched (per-sequence position = row % period) ----
// tokens = B*S laid end to end; period = S. Sequences 2 and 3 re-use positions
// 0..S, so the kernel's `tok % period` must reset RoPE per sequence.
#[test]
fn rope_batched_bwd() {
require_gpu();
let (b, s, heads, head_dim) = (3, 4, 2, 8);
let tokens = b * s;
let n = tokens * heads * head_dim;
let theta = 10000.0;
let x_h = fill(n, 83);
let w = fill(n, 84);
let x = Var::leaf(cuda(&x_h, &[tokens, heads, head_dim]));
let out = ops::rope(&x, theta, s);
scalar_loss(&out, &w).backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let wf = w.clone();
let lx = move |v: &[f32], sh: &[usize]| weighted_sum(&cuda(v, sh).rope(theta, s), &wf);
report(
"rope batched dX",
&grad_check(
&x_h,
&[tokens, heads, head_dim],
&lx,
dx.as_slice::<f32>(),
cfg_linear(),
),
);
}
// ---- softmax ----
#[test]
fn softmax_bwd() {
@@ -501,6 +533,98 @@ fn attention_composed_bwd() {
);
}
// ---- transpose_4d12 ([a,b,c,d] -> [a,c,b,d]) ----
#[test]
fn transpose_4d12_bwd() {
require_gpu();
let (a, b, c, d) = (2, 3, 4, 5);
let n = a * b * c * d;
let x_h = fill(n, 131);
let w = fill(n, 132);
let x = Var::leaf(cuda(&x_h, &[a, b, c, d]));
let out = ops::transpose_4d12(&x);
scalar_loss(&out, &w).backward();
let dx = x.grad().unwrap().to_device(Device::Cpu);
let wf = w.clone();
let lx = move |v: &[f32], s: &[usize]| weighted_sum(&cuda(v, s).transpose_4d12(), &wf);
report(
"transpose_4d12 dX",
&grad_check(&x_h, &[a, b, c, d], &lx, dx.as_slice::<f32>(), cfg_linear()),
);
}
// ---- fused batched causal attention (the T10 op) ----
// q,k,v: [bh, seq, hd]. Grad-check dq/dk/dv against finite-diff of L = sum(W∘out).
// bh = 2 (e.g. batch 1 × 2 heads, or 2 sequences × 1 head) exercises the batched
// GEMM stride; the causal mask is applied inside the op.
#[test]
fn attention_batched_bwd() {
require_gpu();
let (bh, seq, hd) = (2, 5, 6);
let n = bh * seq * hd;
let scale = 1.0 / (hd as f32).sqrt();
let q_h = fill(n, 141);
let k_h = fill(n, 142);
let v_h = fill(n, 143);
let w = fill(n, 144);
let q = Var::leaf(cuda(&q_h, &[bh, seq, hd]));
let k = Var::leaf(cuda(&k_h, &[bh, seq, hd]));
let v = Var::leaf(cuda(&v_h, &[bh, seq, hd]));
let out = ops::attention(&q, &k, &v, scale);
scalar_loss(&out, &w).backward();
let dq = q.grad().unwrap().to_device(Device::Cpu);
let dk = k.grad().unwrap().to_device(Device::Cpu);
let dv = v.grad().unwrap().to_device(Device::Cpu);
let fwd = move |qh: &[f32], kh: &[f32], vh: &[f32]| -> f32 {
let qv = cuda(qh, &[bh, seq, hd]);
let kv = cuda(kh, &[bh, seq, hd]);
let vv = cuda(vh, &[bh, seq, hd]);
let (o, _) = qv.attention(&kv, &vv, scale);
weighted_sum(&o, &w)
};
let (kf, vf, ff) = (k_h.clone(), v_h.clone(), fwd.clone());
let lq = move |x: &[f32], _s: &[usize]| ff(x, &kf, &vf);
report(
"attn(batched) dQ",
&grad_check(
&q_h,
&[bh, seq, hd],
&lq,
dq.as_slice::<f32>(),
cfg_nonlinear(),
),
);
let (qf, vf, ff) = (q_h.clone(), v_h.clone(), fwd.clone());
let lk = move |x: &[f32], _s: &[usize]| ff(&qf, x, &vf);
report(
"attn(batched) dK",
&grad_check(
&k_h,
&[bh, seq, hd],
&lk,
dk.as_slice::<f32>(),
cfg_nonlinear(),
),
);
let (qf, kf, ff) = (q_h.clone(), k_h.clone(), fwd.clone());
let lv = move |x: &[f32], _s: &[usize]| ff(&qf, &kf, x);
report(
"attn(batched) dV",
&grad_check(
&v_h,
&[bh, seq, hd],
&lv,
dv.as_slice::<f32>(),
cfg_linear(),
),
);
}
// --- test helpers ---
// Scalar loss node L = sum(W ∘ out): wraps a fixed-weight Var and reduces. We

1
crates/xtrain-cuda/build.rs
View File

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

66
crates/xtrain-cuda/src/cublas.rs
View File

@@ -93,3 +93,69 @@ pub fn sgemm(
assert_eq!(status, 0, "cublasSgemm failed: {status}");
});
}
/// Strided-batched row-major SGEMM: for each `i` in `0..batch`,
/// `C_i[m,n] = alpha·opA(A_i)·opB(B_i) + beta·C_i`, where `A_i`/`B_i`/`C_i` are
/// consecutive matrices laid `stride_*` elements apart in one contiguous buffer.
/// Same row-major⟺col-major trick as [`sgemm`] (compute col-major `Cᵀ`), applied
/// per batch element. Used for the batched attention `QKᵀ` / `PV` GEMMs (and their
/// backwards), so the whole attention runs as 2 batched-GEMM launches, not a
/// per-(batch,head) Python loop. `A`/`B`/`C` are device pointers to the first
/// matrix; strides are in ELEMENTS.
#[allow(clippy::too_many_arguments)]
pub fn sgemm_strided_batched(
trans_a: bool,
trans_b: bool,
m: usize,
n: usize,
k: usize,
alpha: f32,
a: *const f32,
stride_a: usize,
b: *const f32,
stride_b: usize,
beta: f32,
c: *mut f32,
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
};
with_handle(|handle| {
let status = unsafe {
ffi::cublasSgemmStridedBatched(
handle,
op_b,
op_a,
n as i32,
m as i32,
k as i32,
&alpha,
b,
ldb as i32,
stride_b as i64,
a,
lda as i32,
stride_a as i64,
&beta,
c,
ldc as i32,
stride_c as i64,
batch as i32,
)
};
assert_eq!(status, 0, "cublasSgemmStridedBatched failed: {status}");
});
}

52
crates/xtrain-cuda/src/ffi.rs
View File

@@ -125,7 +125,9 @@ unsafe extern "C" {
pub fn launch_silu_f32(x: *const f32, y: *mut f32, n: i32, s: CudaStream);
pub fn launch_silu_dx_f32(x: *const f32, dy: *const f32, dx: *mut f32, n: i32, s: CudaStream);
// RoPE (rotate_half), x:[tokens,heads,head_dim], position = token index.
// RoPE (rotate_half), x:[tokens,heads,head_dim], position = (token index %
// period). `period` = sequence length, so a flattened batch of sequences gets
// per-sequence positions; period == tokens reproduces the single-sequence case.
pub fn launch_rope_f32(
x: *const f32,
y: *mut f32,
@@ -133,6 +135,7 @@ unsafe extern "C" {
heads: i32,
head_dim: i32,
theta: f32,
period: i32,
s: CudaStream,
);
pub fn launch_rope_dx_f32(
@@ -142,6 +145,7 @@ unsafe extern "C" {
heads: i32,
head_dim: i32,
theta: f32,
period: i32,
s: CudaStream,
);
@@ -211,6 +215,31 @@ unsafe extern "C" {
c: i32,
s: CudaStream,
);
// 4D axis-(1,2) transpose: in:[a,b,c,d] -> out:[a,c,b,d]. out[i,k,j,l]=in[i,j,k,l].
pub fn launch_transpose_4d12_f32(
input: *const f32,
out: *mut f32,
a: i32,
b: i32,
c: i32,
d: i32,
s: CudaStream,
);
}
// Batched attention helper (csrc/ops/attention.cu): causal row-wise softmax over
// score rows [rows, seq] with query position = (row % seq); scales logits by
// `scale` (= 1/sqrt(head_dim)) and masks future columns to probability 0.
#[cfg(not(no_cuda))]
unsafe extern "C" {
pub fn launch_softmax_causal_f32(
x: *const f32,
y: *mut f32,
rows: i32,
seq: i32,
scale: f32,
s: CudaStream,
);
}
// GPU-side optimizer kernels (csrc/ops/optim.cu): AdamW step (m/v on device) and
@@ -267,6 +296,27 @@ unsafe extern "C" {
c: *mut f32,
ldc: i32,
) -> i32;
#[allow(clippy::too_many_arguments)]
pub fn cublasSgemmStridedBatched(
handle: CublasHandle,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const f32,
a: *const f32,
lda: i32,
stride_a: i64,
b: *const f32,
ldb: i32,
stride_b: i64,
beta: *const f32,
c: *mut f32,
ldc: i32,
stride_c: i64,
batch_count: i32,
) -> i32;
}
#[cfg(not(no_cuda))]

48
crates/xtrain-distributed/src/ddp.rs
View File

@@ -17,7 +17,7 @@ use std::thread;
use std::time::Instant;
use xtrain_autodiff::tape::Var;
use xtrain_model::{Config, TinyTransformer, ids_tensor};
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;
use xtrain_train::checkpoint;
@@ -91,10 +91,11 @@ pub fn train_rank(
let mut losses = Vec::with_capacity(cfg.steps);
let mut evals = Vec::new();
let mut best_val: Option<f32> = None;
// Each rank reaches the global batch mean as (Σ_global / world) · (1/b_local),
// where b_local = batch_size / world (see DdpContext::all_reduce_average_grads).
// Each rank runs ONE batched forward over its b_local = batch_size/world
// sequences → backward grad = local mean (Σ_local / b_local). all_reduce_average
// (sum across ranks, /world) then gives Σ_global/(world·b_local) = Σ_global/
// B_global — already the global-batch mean — so the clip pre-scale is 1.0.
let batch_local = cfg.batch_size / ctx.world;
let inv_batch_local = 1.0 / batch_local as f32;
let start = Instant::now();
let mut tokens_seen: u64 = 0;
// Rank 0 owns the held-out eval + best-val checkpoint (params are identical
@@ -105,31 +106,36 @@ pub fn train_rank(
let lr = cfg.schedule.lr(step);
// Draw the whole global batch from the shared RNG (same on every rank);
// run forward+backward only on this rank's shard. The tape SUMs the
// shard's grads; the union of shards == the single-GPU batch.
let mut local_loss_sum = 0.0f32;
// collect only this rank's shard (global index % world == rank) and run it
// as ONE batched forward/backward. The union of shards == the single-GPU
// batch; each rank's backward yields its local mean (Σ_local / b_local).
let mut inputs = Vec::with_capacity(batch_local);
let mut targets_v = Vec::with_capacity(batch_local);
for i in 0..cfg.batch_size {
let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
if i % ctx.world != ctx.rank {
continue; // not this rank's sequence
if i % ctx.world == ctx.rank {
inputs.push(input);
targets_v.push(target);
}
let ids = ids_tensor(&input, device);
let targets = ids_tensor(&target, device);
let loss = model.loss(&ids, &targets);
local_loss_sum += read_scalar(&loss);
loss.backward();
tokens_seen += cfg.seq_len as u64;
}
let ids = batched_ids_tensor(&inputs, device);
let targets = batched_ids_tensor(&targets_v, device);
let loss = model.loss_batched(&ids, &targets, batch_local);
let local_mean = read_scalar(&loss); // Σ_local / b_local
loss.backward();
tokens_seen += (batch_local * cfg.seq_len) as u64;
// AllReduce(sum) + /world the grads → every rank holds Σ_global/world.
// AllReduce(sum) + /world the grads → every rank holds Σ_global/B_global
// (local means summed over ranks, /world = global mean). See note above.
ctx.all_reduce_average_grads(&params);
// The reported loss is the global mean: average local sums across ranks.
let step_loss = all_reduce_loss(ctx, local_loss_sum) / cfg.batch_size as f32;
// Reported loss = global mean: sum the per-rank local sums (= mean·b_local)
// across ranks, /B_global. With equal b_local this is mean over ranks.
let step_loss =
all_reduce_loss(ctx, local_mean * batch_local as f32) / cfg.batch_size as f32;
losses.push(step_loss);
// clip pre_scale = 1/b_local finishes the average to Σ_global/B_global,
// identical to the single-GPU clip(pre_scale = 1/B_global).
let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, inv_batch_local);
// Grads are already the global-batch mean — just clip (pre-scale 1.0).
let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, 1.0);
opt.step(lr, &params);
for p in &params {
p.zero_grad();

23
crates/xtrain-distributed/tests/ddp_correctness.rs
View File

@@ -13,7 +13,7 @@ use std::time::Instant;
use xtrain_cuda::device;
use xtrain_distributed::{DdpConfig, DdpContext, build_model, get_unique_id, launch, train_rank};
use xtrain_model::{Config, ids_tensor};
use xtrain_model::{Config, batched_ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;
use xtrain_train::clip::clip_grad_norm_gpu;
@@ -47,22 +47,25 @@ fn run_single_gpu(cfg: Config, corpus: &Corpus, dcfg: &DdpConfig) -> (Vec<f32>,
let params = model.params();
let mut opt = GpuAdamW::new(dcfg.weight_decay);
let mut rng = dcfg.seed;
let inv_batch = 1.0 / dcfg.batch_size as f32;
let mut losses = Vec::new();
for step in 0..dcfg.steps {
let lr = dcfg.schedule.lr(step);
let mut loss_sum = 0.0f32;
// Sample the whole global batch and run it as ONE batched forward/backward
// (matches the T10 DDP path: backward yields the global-batch mean grad).
let mut inputs = Vec::with_capacity(dcfg.batch_size);
let mut targets_v = Vec::with_capacity(dcfg.batch_size);
for _ in 0..dcfg.batch_size {
let (input, target) = corpus.sample(dcfg.seq_len, &mut rng);
let ids = ids_tensor(&input, device);
let targets = ids_tensor(&target, device);
let loss = model.loss(&ids, &targets);
loss_sum += loss.value().to_device(Device::Cpu).as_slice::<f32>()[0];
loss.backward();
inputs.push(input);
targets_v.push(target);
}
losses.push(loss_sum * inv_batch);
clip_grad_norm_gpu(&params, dcfg.max_grad_norm, inv_batch);
let ids = batched_ids_tensor(&inputs, device);
let targets = batched_ids_tensor(&targets_v, device);
let loss = model.loss_batched(&ids, &targets, dcfg.batch_size);
losses.push(loss.value().to_device(Device::Cpu).as_slice::<f32>()[0]);
loss.backward();
clip_grad_norm_gpu(&params, dcfg.max_grad_norm, 1.0);
opt.step(lr, &params);
for p in &params {
p.zero_grad();

2
crates/xtrain-model/src/lib.rs
View File

@@ -24,4 +24,4 @@ pub use config::Config;
#[cfg(not(no_cuda))]
mod model;
#[cfg(not(no_cuda))]
pub use model::{TinyTransformer, ids_tensor, param_to_host};
pub use model::{TinyTransformer, batched_ids_tensor, ids_tensor, param_to_host};

138
crates/xtrain-model/src/model.rs
View File

@@ -30,7 +30,6 @@ pub struct TinyTransformer {
blocks: Vec<Block>,
final_norm: Var, // [dim]
lm_head: Var, // [dim, vocab]
device: Device,
}
impl TinyTransformer {
@@ -72,7 +71,6 @@ impl TinyTransformer {
blocks,
final_norm,
lm_head,
device,
}
}
@@ -106,16 +104,34 @@ impl TinyTransformer {
}
/// Forward over a single sequence of token `ids` (`[seq]` I32 on this
/// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`.
/// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`. This
/// is the batch-1 special case of [`forward_batched`](Self::forward_batched)
/// (used by the autoregressive sampler / inference path).
pub fn forward(&self, ids: &Tensor) -> Var {
let seq = ids.shape()[0];
let mask = self.causal_mask(seq);
self.forward_batched(ids, 1)
}
let mut h = ops::embedding(&self.embed, ids); // [seq, dim]
/// Batched forward over `batch` sequences of equal length `seq`, flattened to
/// `[batch*seq]` I32 ids in sequence-major order (sequence 0's `seq` tokens,
/// then sequence 1's, …). Returns logits `[batch*seq, vocab]` in the SAME flat
/// layout. The whole graph runs on the flattened tokens so every linear
/// projection is ONE big `[batch*seq, dim] × [dim, out]` GEMM (the
/// GPU-filling win); only attention is sequence-aware (per-sequence causal
/// mask + RoPE position, NO cross-sequence attention).
pub fn forward_batched(&self, ids: &Tensor, batch: usize) -> Var {
let total = ids.shape()[0];
assert_eq!(
total % batch,
0,
"ids len {total} not divisible by batch {batch}"
);
let seq = total / batch;
let mut h = ops::embedding(&self.embed, ids); // [batch*seq, dim]
for b in &self.blocks {
// --- Attention sub-block (pre-norm + residual) ---
let normed = ops::rms_norm(&h, &b.attn_norm, self.cfg.eps);
let attn = self.attention(b, &normed, &mask, seq);
let attn = self.attention(b, &normed, batch, seq);
h = ops::add(&h, &attn);
// --- MLP sub-block (pre-norm + residual) ---
@@ -125,7 +141,7 @@ impl TinyTransformer {
}
let h = ops::rms_norm(&h, &self.final_norm, self.cfg.eps);
ops::matmul(&h, &self.lm_head) // [seq, vocab]
ops::matmul(&h, &self.lm_head) // [batch*seq, vocab]
}
/// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
@@ -134,76 +150,76 @@ impl TinyTransformer {
ops::cross_entropy(&logits, targets)
}
/// Multi-head causal self-attention. `x`:[seq,dim] (already normed).
fn attention(&self, b: &Block, x: &Var, mask: &Var, seq: usize) -> Var {
/// Batched cross-entropy mean loss: `forward_batched(ids, batch)` against
/// flat `targets` (`[batch*seq]` I32, same sequence-major layout). The CE mean
/// is over all `batch*seq` rows — identical to averaging the per-sequence
/// losses, so the loss value matches the looped single-sequence path.
pub fn loss_batched(&self, ids: &Tensor, targets: &Tensor, batch: usize) -> Var {
let logits = self.forward_batched(ids, batch);
ops::cross_entropy(&logits, targets)
}
/// Multi-head causal self-attention over a flattened batch. `x`:[batch*seq,dim]
/// (already normed), laid out sequence-major. The Q/K/V/O projections are big
/// `[batch*seq, dim]` GEMMs; the scaled-dot-product attention itself runs as a
/// fused BATCHED op over the `batch·n_heads` (sequence,head) blocks — each
/// attends within its own `[seq,seq]` causal window (NO cross-sequence
/// attention), with RoPE positions reset per sequence (`period = seq`). Causal
/// masking is applied inside the fused op's softmax kernel (no additive
/// `[seq,seq]` mask tensor).
fn attention(&self, b: &Block, x: &Var, batch: usize, seq: usize) -> Var {
let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
let total = batch * seq;
let bh = batch * nh;
let scale = 1.0 / (hd as f32).sqrt();
// Project, then lay out as per-head [seq, head_dim] tensors.
// [seq,dim] @ [dim,dim] = [seq,dim]
// reshape [seq, nh, hd]
// Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
// [B*S,dim] @ [dim,dim] = [B*S,dim]
// reshape [B*S, nh, hd]
// qk-norm per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
// rope (kernel expects exactly [tokens, heads, head_dim])
// transpose [nh, seq, hd] → split into nh × [seq, hd]
let to_heads = |proj: Var, norm: Option<&Var>| -> Vec<Var> {
let r = ops::reshape(&proj, &[seq, nh, hd]);
// rope [B*S, nh, hd] with per-sequence position (period = seq)
// reshape [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
let r = ops::reshape(&proj, &[total, nh, hd]);
let r = match norm {
// Per-head RMSNorm: flatten the (seq,nh) head rows, norm over hd,
// Per-head RMSNorm: flatten the (B*S,nh) head rows, norm over hd,
// restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
Some(gamma) => {
let flat = ops::reshape(&r, &[seq * nh, hd]);
let flat = ops::reshape(&r, &[total * nh, hd]);
let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
let r = ops::reshape(&normed, &[seq, nh, hd]);
ops::rope(&r, self.cfg.rope_theta)
let r = ops::reshape(&normed, &[total, nh, hd]);
ops::rope(&r, self.cfg.rope_theta, seq)
}
None => r,
};
let t = ops::transpose_3d01(&r); // [nh, seq, hd]
ops::split_heads(&t)
let r = ops::reshape(&r, &[batch, seq, nh, hd]);
let t = ops::transpose_4d12(&r); // [B, nh, S, hd]
ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
};
let q = to_heads(ops::matmul(x, &b.wq), Some(&b.q_norm));
let k = to_heads(ops::matmul(x, &b.wk), Some(&b.k_norm));
let v = to_heads(ops::matmul(x, &b.wv), None);
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);
// Per-head scaled-dot-product attention with causal mask.
let heads_out: Vec<Var> = (0..nh)
.map(|i| {
let kt = ops::transpose_2d(&k[i]); // [hd, seq]
let scores = ops::scale(&ops::matmul(&q[i], &kt), scale); // [seq,seq]
let scores = ops::add(&scores, mask); // causal
let probs = ops::softmax(&scores);
ops::matmul(&probs, &v[i]) // [seq, hd]
})
.collect();
// 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).
let out = ops::attention(&q, &k, &v, scale); // [B*nh, S, hd]
// Stack heads back: nh × [seq,hd] → [nh,seq,hd] → [seq,nh,hd] → [seq,dim].
let merged = ops::merge_heads(&heads_out); // [nh, seq, hd]
let t = ops::transpose_3d01(&merged); // [seq, nh, hd]
let concat = ops::reshape(&t, &[seq, nh * hd]); // [seq, dim]
// Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
// [B,S,nh,hd] → [B*S, dim].
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
}
/// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[seq,dim].
/// 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 act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
ops::matmul(&act, &b.w_down) // [seq, dim]
}
/// Additive causal mask `[seq,seq]`: 0 on/below the diagonal, −1e9 above it
/// (so softmax zeros out future positions). A constant leaf (no grad needed,
/// but harmless if it accumulates one — it has no consumers downstream of x).
fn causal_mask(&self, seq: usize) -> Var {
let mut m = vec![0.0f32; seq * seq];
for i in 0..seq {
for j in (i + 1)..seq {
m[i * seq + j] = -1.0e9;
}
}
Var::leaf(Tensor::from_slice(&m, &[seq, seq]).to_device(self.device))
}
}
/// Materialise a parameter's value back to a host `Vec<f32>` (for the GD step
@@ -216,3 +232,17 @@ pub fn param_to_host(v: &Var) -> Vec<f32> {
pub fn ids_tensor(ids: &[i32], device: Device) -> Tensor {
Tensor::from_slice(ids, &[ids.len()]).to_device(device)
}
/// Flatten `batch` equal-length sequences into one `[batch*seq]` I32 tensor in
/// sequence-major order (the layout `forward_batched` expects). Each row of
/// `seqs` is one sequence; all must have the same length.
pub fn batched_ids_tensor(seqs: &[Vec<i32>], device: Device) -> Tensor {
assert!(!seqs.is_empty(), "empty batch");
let seq = seqs[0].len();
let mut flat = Vec::with_capacity(seqs.len() * seq);
for s in seqs {
assert_eq!(s.len(), seq, "ragged batch: sequences must be equal length");
flat.extend_from_slice(s);
}
Tensor::from_slice(&flat, &[flat.len()]).to_device(device)
}

142
crates/xtrain-model/tests/batched.rs Normal file
View File

@@ -0,0 +1,142 @@
// T10 batched-forward equivalence: a batched forward over B sequences must equal
// the old single-sequence path (run each sequence on its own, concatenate the
// logits) — both for the forward logits AND every parameter's gradient.
//
// This is THE on-GPU correctness gate for batching (no PyTorch needed): if the
// per-sequence RoPE position, per-sequence causal masking, or any flattened op
// were wrong, the batched logits/grads would drift from the looped reference.
//
// Forward equivalence: batched logits[b*S+i] == single-seq-b logits[i].
// Gradient equivalence: the batched loss is the mean over all B*S rows, i.e.
// (1/B)·Σ_b mean_i(loss_b); summing the B single-sequence losses and scaling by
// 1/B gives the SAME scalar, so their summed grads (tape fan-out) ×1/B match the
// batched grads. We check that.
#![cfg(not(no_cuda))]
use xtrain_cuda::device;
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor, ids_tensor};
use xtrain_tensor::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 batched_matches_looped_single_sequence() {
assert!(device::device_count().unwrap() > 0, "no CUDA device");
device::set_device(0).unwrap();
let device = Device::Cuda(0);
let mut cfg = Config::tiny();
cfg.vocab = 16;
let batch = 3usize;
let seq = 5usize;
// B distinct sequences (sequence-major), within vocab.
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();
// --- Batched forward: ONE pass over [B*S]. ---
let bmodel = build(cfg, device);
let bids = batched_ids_tensor(&seqs, device);
let blogits = host(&bmodel.forward_batched(&bids, batch).value());
// --- Looped reference: each sequence on its own, concatenate logits. ---
let smodel = build(cfg, device);
let mut slogits = Vec::with_capacity(batch * seq * cfg.vocab);
for s in &seqs {
let ids = ids_tensor(s, device);
slogits.extend(host(&smodel.forward(&ids).value()));
}
// Forward equivalence (fp GEMM rounding only differs in summation order).
let max_rel = blogits
.iter()
.zip(&slogits)
.map(|(b, s)| (b - s).abs() / s.abs().max(1e-4))
.fold(0.0f32, f32::max);
println!("batched vs looped: logits max rel err = {max_rel:.3e}");
assert!(max_rel < 1e-3, "batched logits diverged: {max_rel:.3e}");
// --- Gradient equivalence. ---
// Batched: loss = mean over B*S rows; one backward.
let bparams = bmodel.params();
let btgt = batched_ids_tensor(&tgts, device);
let bloss = bmodel.loss_batched(&bids, &btgt, batch);
let bloss_val = host(&bloss.value())[0];
bloss.backward();
// Looped: Σ_b loss_b (each a per-sequence mean), then grad ×(1/B) == batched.
let sparams = smodel.params();
let mut sloss_sum = 0.0f32;
for (s, t) in seqs.iter().zip(&tgts) {
let ids = ids_tensor(s, device);
let tg = ids_tensor(t, device);
let l = smodel.loss(&ids, &tg);
sloss_sum += host(&l.value())[0];
l.backward();
}
println!(
"batched loss = {bloss_val:.6} looped mean = {:.6}",
sloss_sum / batch as f32
);
assert!(
(bloss_val - sloss_sum / batch as f32).abs() < 1e-4,
"batched loss != looped mean"
);
let mut max_grad_rel = 0.0f32;
for (bp, sp) in bparams.iter().zip(&sparams) {
let bg = host(&bp.grad().expect("batched grad"));
let sg = host(&sp.grad().expect("looped grad"));
for (g_b, g_s) in bg.iter().zip(&sg) {
// looped grad is the SUM over B sequences; ×(1/B) recovers the mean.
let g_s = g_s / batch as f32;
let rel = (g_b - g_s).abs() / g_s.abs().max(1e-4);
max_grad_rel = max_grad_rel.max(rel);
}
}
println!("batched vs looped: grad max rel err = {max_grad_rel:.3e}");
assert!(
max_grad_rel < 5e-3,
"batched grads diverged: {max_grad_rel:.3e}"
);
}

42
crates/xtrain-model/tests/parity.py
View File

@@ -55,10 +55,13 @@ NH = int(cfg["n_heads"])
HD = int(cfg["head_dim"])
EPS = float(cfg["eps"])
THETA = float(cfg["rope_theta"])
# Batched: B sequences of length SEQ, flattened sequence-major to [B*SEQ] ids.
B = int(cfg.get("batch", "1"))
SEQ = int(cfg["seq"])
ids = read_ids("ids.txt")
targets = read_ids("targets.txt")
SEQ = len(ids)
assert len(ids) == B * SEQ, f"ids {len(ids)} != B*SEQ {B*SEQ}"
# Load params as leaf tensors requiring grad (float64 for a clean reference).
P = {}
@@ -76,15 +79,16 @@ def rms_norm(x, gamma):
return x * torch.rsqrt(ms + EPS) * gamma
def rope(x): # x: [seq, nh, hd], position = token index, matching the kernel
def rope(x): # x: [B*SEQ, nh, hd], position = (row % SEQ) — resets per sequence
half = HD // 2
out = torch.empty_like(x)
i = torch.arange(half, dtype=torch.float64)
freq = THETA ** (-(2.0 * i) / HD) # [half]
pos = torch.arange(SEQ, dtype=torch.float64).reshape(SEQ, 1) # [seq,1]
ang = pos * freq # [seq, half]
c = torch.cos(ang).reshape(SEQ, 1, half)
s = torch.sin(ang).reshape(SEQ, 1, half)
# Position within each sequence: rows 0..SEQ for seq 0, 0..SEQ for seq 1, ...
pos = (torch.arange(B * SEQ, dtype=torch.float64) % SEQ).reshape(B * SEQ, 1)
ang = pos * freq # [B*SEQ, half]
c = torch.cos(ang).reshape(B * SEQ, 1, half)
s = torch.sin(ang).reshape(B * SEQ, 1, half)
x0 = x[..., :half]
x1 = x[..., half:]
out[..., :half] = x0 * c - x1 * s
@@ -102,26 +106,30 @@ for l in range(NL):
"ffn_norm", "w_gate", "w_up", "w_down"]})
idx = torch.tensor(ids, dtype=torch.long)
# Per-sequence causal mask (broadcast over the batch); NO cross-sequence attention.
mask = torch.triu(torch.full((SEQ, SEQ), -1.0e9, dtype=torch.float64), diagonal=1)
h = emb[idx] # [seq, dim]
h = emb[idx] # [B*SEQ, dim] (everything stays flattened, matching the Rust path)
for L in layers:
# Attention
x = rms_norm(h, L["attn_norm"])
q = (x @ L["wq"]).reshape(SEQ, NH, HD)
k = (x @ L["wk"]).reshape(SEQ, NH, HD)
v = (x @ L["wv"]).reshape(SEQ, NH, HD)
q = (x @ L["wq"]).reshape(B * SEQ, NH, HD)
k = (x @ L["wk"]).reshape(B * SEQ, NH, HD)
v = (x @ L["wv"]).reshape(B * SEQ, NH, HD)
# Per-head QK-norm (Qwen3-style), before RoPE.
q = rms_norm(q, L["q_norm"])
k = rms_norm(k, L["k_norm"])
q = rope(q).transpose(0, 1) # [nh, seq, hd]
k = rope(k).transpose(0, 1)
v = v.transpose(0, 1)
q = rope(q) # [B*SEQ, nh, hd]
k = rope(k)
# Reshape to [B, NH, SEQ, HD] so attention runs within each sequence.
q = q.reshape(B, SEQ, NH, HD).transpose(1, 2) # [B, nh, seq, hd]
k = k.reshape(B, SEQ, NH, HD).transpose(1, 2)
v = v.reshape(B, SEQ, NH, HD).transpose(1, 2)
scale = 1.0 / math.sqrt(HD)
scores = (q @ k.transpose(-1, -2)) * scale + mask # [nh, seq, seq]
scores = (q @ k.transpose(-1, -2)) * scale + mask # [B, nh, seq, seq]
probs = torch.softmax(scores, dim=-1)
out = probs @ v # [nh, seq, hd]
out = out.transpose(0, 1).reshape(SEQ, DIM) # [seq, dim]
out = probs @ v # [B, nh, seq, hd]
out = out.transpose(1, 2).reshape(B * SEQ, DIM) # [B*SEQ, dim]
attn = out @ L["wo"]
h = h + attn
# MLP
@@ -133,7 +141,7 @@ for L in layers:
h = h + mlp
h = rms_norm(h, final_norm)
logits = h @ lm_head # [seq, vocab]
logits = h @ lm_head # [B*SEQ, vocab]
loss = torch.nn.functional.cross_entropy(
logits, torch.tensor(targets, dtype=torch.long), reduction="mean")

19
crates/xtrain-model/tests/parity_dump.rs
View File

@@ -53,12 +53,17 @@ fn dump_for_parity() {
);
fs::create_dir_all(&dir).unwrap();
// Fixed config + ids (independent of any text, for reproducibility).
// Fixed config + ids (independent of any text, for reproducibility). B>1 so
// the batched forward is exercised: 2 sequences of length 4, flattened
// sequence-major to [B*S]=8 ids. Per-sequence RoPE position (resets at the
// sequence boundary) + per-sequence causal masking (no cross-sequence
// attention) are both checked against PyTorch.
let mut cfg = Config::tiny();
cfg.vocab = 12;
let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6];
let batch = 2usize;
let seq = 4usize;
let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6]; // [B*S], sequence-major
let targets: Vec<i32> = vec![1, 4, 1, 5, 9, 2, 6, 0];
let seq = ids.len();
// Same deterministic init as the overfit test.
let mut seed = 1u64;
@@ -83,6 +88,7 @@ fn dump_for_parity() {
writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
writeln!(f, "eps {:e}", cfg.eps).unwrap();
writeln!(f, "rope_theta {:e}", cfg.rope_theta).unwrap();
writeln!(f, "batch {batch}").unwrap();
writeln!(f, "seq {seq}").unwrap();
}
{
@@ -105,10 +111,11 @@ fn dump_for_parity() {
write_vec(&dir, &format!("w_{name}.txt"), &param_to_host(p), &shape);
}
// Forward logits + loss, then backward → per-param grads.
// Batched forward logits + loss (B sequences as one forward), then backward
// → per-param grads.
let ids_t = ids_tensor(&ids, device);
let targets_t = ids_tensor(&targets, device);
let logits = model.forward(&ids_t);
let logits = model.forward_batched(&ids_t, batch);
write_vec(
&dir,
"logits.txt",
@@ -116,7 +123,7 @@ fn dump_for_parity() {
logits.value().shape(),
);
let loss = model.loss(&ids_t, &targets_t);
let loss = model.loss_batched(&ids_t, &targets_t, batch);
let loss_val = param_to_host(&loss)[0];
{
let mut f = fs::File::create(dir.join("loss.txt")).unwrap();

215
crates/xtrain-tensor/src/tensor.rs
View File

@@ -454,13 +454,20 @@ impl Tensor {
dx
}
/// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; the position
/// of each token is its row index. Returns the rotated tensor.
/// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; each token's
/// position is `row % period`. `period` = sequence length, so a flattened
/// batch `[B*S,heads,head_dim]` gets per-sequence positions (pass `period=S`);
/// pass `period=tokens` for a single sequence (position = row). Returns the
/// rotated tensor.
#[cfg(not(no_cuda))]
pub fn rope(&self, theta: f32) -> Self {
pub fn rope(&self, theta: f32, period: usize) -> Self {
assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]");
let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
assert_eq!(head_dim % 2, 0, "head_dim must be even");
assert!(
period > 0 && tokens % period == 0,
"tokens must be a multiple of period"
);
let out = Tensor::zeros(&self.shape, DType::F32, self.device());
unsafe {
xtrain_cuda::ffi::launch_rope_f32(
@@ -470,6 +477,7 @@ impl Tensor {
heads as i32,
head_dim as i32,
theta,
period as i32,
std::ptr::null_mut(),
);
}
@@ -477,9 +485,9 @@ impl Tensor {
}
/// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
/// orthogonal map, so it needs no cached forward values, only `theta`.
/// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
#[cfg(not(no_cuda))]
pub fn rope_backward(dy: &Tensor, theta: f32) -> Self {
pub fn rope_backward(dy: &Tensor, theta: f32, period: usize) -> Self {
let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]);
let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device());
unsafe {
@@ -490,6 +498,7 @@ impl Tensor {
heads as i32,
head_dim as i32,
theta,
period as i32,
std::ptr::null_mut(),
);
}
@@ -667,6 +676,202 @@ impl Tensor {
out
}
// --- Batched attention (the T10 fused op) ---
/// Batched causal scaled-dot-product attention. `self`=Q, `k`, `v` are each
/// `[bh, seq, head_dim]` (bh = batch·n_heads), contiguous F32 on one GPU.
/// Computes, per batch element, `out = softmax(causal(Q·Kᵀ / √hd)) · V`. The
/// two GEMMs run as `cublasSgemmStridedBatched` and the softmax+scale+causal
/// mask is one kernel, so the whole attention is 3 launches regardless of bh.
/// Returns `(out, probs)` where `probs`:[bh,seq,seq] is cached for backward.
#[cfg(not(no_cuda))]
pub fn attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> (Tensor, 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");
let (bh, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
let dev = self.device();
// 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(
false,
true,
seq,
seq,
hd,
1.0,
self.data_ptr() as *const f32,
seq * hd,
k.data_ptr() as *const f32,
seq * hd,
0.0,
scores.data_ptr() as *mut f32,
seq * seq,
bh,
);
// probs = softmax(causal(scores · scale)), one block per [bh·seq] row.
let probs = 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,
(bh * seq) as i32,
seq as i32,
scale,
std::ptr::null_mut(),
);
}
// 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(
false,
false,
seq,
hd,
seq,
1.0,
probs.data_ptr() as *const f32,
seq * seq,
v.data_ptr() as *const f32,
seq * hd,
0.0,
out.data_ptr() as *mut f32,
seq * hd,
bh,
);
(out, probs)
}
/// Backward of [`attention`](Self::attention). Inputs: forward `q`,`k`,`v`,
/// the cached `probs`, the upstream `dout` (all batched `[bh,seq,*]`), and the
/// same `scale`. Returns `(dq, dk, dv)`.
///
/// dP = dOut · Vᵀ ; dV = Pᵀ · dOut
/// dScores = softmax_jacobian(P, dP) · scale (scale folded back in)
/// dQ = dScores · K ; dK = dScoresᵀ · Q
///
/// Masked (future) entries of P are 0, so the softmax Jacobian zeros their
/// gradient — the causal mask needs no special handling here.
#[cfg(not(no_cuda))]
pub fn attention_backward(
q: &Tensor,
k: &Tensor,
v: &Tensor,
probs: &Tensor,
dout: &Tensor,
scale: f32,
) -> (Tensor, Tensor, Tensor) {
let (bh, seq, hd) = (q.shape[0], q.shape[1], q.shape[2]);
let dev = q.device();
// 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(
false,
true,
seq,
seq,
hd,
1.0,
dout.data_ptr() as *const f32,
seq * hd,
v.data_ptr() as *const f32,
seq * hd,
0.0,
dp.data_ptr() as *mut f32,
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(
true,
false,
seq,
hd,
seq,
1.0,
probs.data_ptr() as *const f32,
seq * seq,
dout.data_ptr() as *const f32,
seq * hd,
0.0,
dv.data_ptr() as *mut f32,
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].
let dscores = Tensor::softmax_backward(
&probs.reshape(&[bh * seq, seq]),
&dp.reshape(&[bh * seq, seq]),
)
.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(
false,
false,
seq,
hd,
seq,
1.0,
dscores.data_ptr() as *const f32,
seq * seq,
k.data_ptr() as *const f32,
seq * hd,
0.0,
dq.data_ptr() as *mut f32,
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(
true,
false,
seq,
hd,
seq,
1.0,
dscores.data_ptr() as *const f32,
seq * seq,
q.data_ptr() as *const f32,
seq * hd,
0.0,
dk.data_ptr() as *mut f32,
seq * hd,
bh,
);
(dq, dk, dv)
}
/// 4D axis-(1,2) transpose: `self`:[a,b,c,d] → [a,c,b,d],
/// `out[i,k,j,l]=self[i,j,k,l]`. Lays out batched multi-head attention
/// (`[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");
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 {
xtrain_cuda::ffi::launch_transpose_4d12_f32(
self.data_ptr() as *const f32,
out.data_ptr() as *mut f32,
a as i32,
b as i32,
c as i32,
d as i32,
std::ptr::null_mut(),
);
}
out
}
// Shared validation for same-shape binary elementwise ops.
#[cfg(not(no_cuda))]
fn check_binary(&self, other: &Tensor, op: &str) {

45
crates/xtrain-train/src/train_loop.rs
View File

@@ -1,18 +1,20 @@
//! The training loop: sample sequences → forward `loss` → backward → grad clip
//! (with batch averaging) → AdamW step → zero grads; with an LR schedule,
//! periodic loss logging, and periodic checkpointing.
//! The training loop: sample a batch of sequences → ONE batched forward `loss` →
//! backward → grad clip → AdamW step → zero grads; with an LR schedule, periodic
//! loss logging, and periodic checkpointing.
//!
//! The T5 model is single-sequence, so a "batch" of `batch_size` sequences is
//! handled by running forward+backward on each and letting the tape SUM their
//! grads (its fan-out rule); the clip pass then multiplies by `1/batch_size` to
//! recover the batch-mean gradient before clipping + the optimizer step.
//! Since T10 the model is batched (`loss_batched`): `batch_size` sequences are
//! flattened to `[batch*seq]` and run as a SINGLE forward/backward, so the linear
//! projections become big `[batch*seq, dim]` GEMMs that fill the GPU. The
//! cross-entropy mean is over all `batch*seq` rows — already the batch-mean loss,
//! so backward yields the batch-mean gradient directly (clip pre-scale = 1.0; no
//! more "loop B times + SUM + ×1/batch" hack).
#![cfg(not(no_cuda))]
use std::path::PathBuf;
use std::time::Instant;
use xtrain_model::{TinyTransformer, ids_tensor};
use xtrain_model::{TinyTransformer, batched_ids_tensor, ids_tensor};
use xtrain_optim::GpuAdamW;
use xtrain_tensor::Device;
@@ -67,7 +69,6 @@ pub fn train(
let mut losses = Vec::with_capacity(cfg.steps);
let mut evals = Vec::new();
let mut best_val: Option<f32> = None;
let inv_batch = 1.0 / cfg.batch_size as f32;
let start = Instant::now();
let mut tokens_seen: u64 = 0;
// Best-val checkpointing only kicks in when we actually evaluate.
@@ -76,22 +77,26 @@ pub fn train(
for step in 0..cfg.steps {
let lr = cfg.schedule.lr(step);
// Accumulate grads over `batch_size` sequences (tape SUMs them).
let mut step_loss = 0.0f32;
// Sample `batch_size` sequences and run them as ONE batched forward/
// backward. The CE mean over all batch*seq rows is the batch-mean loss, so
// backward already yields the batch-mean gradient (clip pre-scale = 1.0).
let mut inputs = Vec::with_capacity(cfg.batch_size);
let mut targets_v = Vec::with_capacity(cfg.batch_size);
for _ in 0..cfg.batch_size {
let (input, target) = corpus.sample(cfg.seq_len, &mut rng);
let ids = ids_tensor(&input, device);
let targets = ids_tensor(&target, device);
let loss = model.loss(&ids, &targets);
step_loss += read_scalar(&loss);
loss.backward();
tokens_seen += cfg.seq_len as u64;
inputs.push(input);
targets_v.push(target);
}
step_loss *= inv_batch;
let ids = batched_ids_tensor(&inputs, device);
let targets = batched_ids_tensor(&targets_v, device);
let loss = model.loss_batched(&ids, &targets, cfg.batch_size);
let step_loss = read_scalar(&loss);
loss.backward();
tokens_seen += (cfg.batch_size * cfg.seq_len) as u64;
losses.push(step_loss);
// Average the summed grads (×1/batch) and clip to the global norm.
let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, inv_batch);
// Backward already produced the batch-mean gradient — just clip it.
let gnorm = clip_grad_norm_gpu(&params, cfg.max_grad_norm, 1.0);
opt.step(lr, &params);
for p in &params {
p.zero_grad();

93
csrc/ops/attention.cu Normal file
View File

@@ -0,0 +1,93 @@
// Batched scaled-dot-product attention helpers (Phase T10).
//
// The QKᵀ and PV matmuls run as cublasSgemmStridedBatched in Rust; the only
// kernel attention needs of its own is a CAUSAL row-wise softmax over the score
// rows. Scores are [B*nh, S, S] flattened to rows of length S; for a flat row r
// the query position within its sequence is `r % S`, so columns j > r%S are
// future positions and get probability 0 (no additive -1e9 mask tensor needed).
//
// The forward also folds in the 1/sqrt(head_dim) scale (applied to logits before
// the max/exp) so we don't need a separate scale pass. Backward is the ordinary
// softmax Jacobian (csrc/ops/nn.cu launch_softmax_dx_f32): masked entries have
// y=0, so their contribution vanishes — no causal-specific backward needed.
//
// All F32, row-major, contiguous. Reduction helpers mirror nn.cu (inlined so the
// file is self-contained, matching the csrc/ layout).
#include <math.h>
extern "C" {
__device__ __forceinline__ float att_warp_sum(float v) {
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
v += __shfl_down_sync(0xffffffff, v, off);
return v;
}
__device__ __forceinline__ float att_warp_max(float v) {
#pragma unroll
for (int off = 16; off > 0; off >>= 1)
v = fmaxf(v, __shfl_down_sync(0xffffffff, v, off));
return v;
}
__device__ __forceinline__ float att_block_sum(float v) {
__shared__ float sh[32];
int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
v = att_warp_sum(v);
if (lane == 0) sh[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : 0.0f;
if (warp == 0) v = att_warp_sum(v);
__shared__ float bc;
if (threadIdx.x == 0) bc = v;
__syncthreads();
return bc;
}
__device__ __forceinline__ float att_block_max(float v) {
__shared__ float sh[32];
int lane = threadIdx.x & 31, warp = threadIdx.x >> 5;
int nwarps = (blockDim.x + 31) >> 5;
v = att_warp_max(v);
if (lane == 0) sh[warp] = v;
__syncthreads();
v = (threadIdx.x < nwarps) ? sh[threadIdx.x] : -INFINITY;
if (warp == 0) v = att_warp_max(v);
__shared__ float bc;
if (threadIdx.x == 0) bc = v;
__syncthreads();
return bc;
}
// One block per score row. rows = B*nh*S total; the query position within its
// sequence is (blockIdx.x % seq). Logits are scaled by `scale` (= 1/sqrt(hd))
// before softmax; columns j > qpos are masked to probability 0.
__global__ void softmax_causal_k(const float* x, float* y, int seq, float scale) {
int r = blockIdx.x;
int qpos = r % seq;
const float* xr = x + (size_t)r * seq;
float* yr = y + (size_t)r * seq;
int valid = qpos + 1; // attend to columns [0, qpos]
float m = -INFINITY;
for (int c = threadIdx.x; c < valid; c += blockDim.x)
m = fmaxf(m, xr[c] * scale);
m = att_block_max(m);
float sum = 0.0f;
for (int c = threadIdx.x; c < valid; c += blockDim.x) {
float e = expf(xr[c] * scale - m);
yr[c] = e;
sum += e;
}
sum = att_block_sum(sum);
float inv = 1.0f / sum;
for (int c = threadIdx.x; c < seq; c += blockDim.x)
yr[c] = (c < valid) ? yr[c] * inv : 0.0f;
}
void launch_softmax_causal_f32(const float* x, float* y, int rows, int seq,
float scale, void* s) {
int blk = seq < 1024 ? seq : 1024;
if (blk < 32) blk = 32;
softmax_causal_k<<<rows, blk, 0, (cudaStream_t)s>>>(x, y, seq, scale);
}
} // extern "C"

22
csrc/ops/model.cu
View File

@@ -63,4 +63,26 @@ void launch_transpose_3d01_f32(const float* in, float* out, int a, int b, int c,
transpose_3d01_k<<<grid, blk, 0, (cudaStream_t)s>>>(in, out, a, b, c);
}
// =====================================================================
// 4D axis-(1,2) transpose: in:[a,b,c,d] -> out:[a,c,b,d]. out[i,k,j,l]=in[i,j,k,l].
// Lays out batched multi-head attention: [B,S,nh,hd] <-> [B,nh,S,hd], so a
// flattened [B*nh, S, hd] view feeds the strided-batched-GEMM attention. Its own
// backward is the same op (swap b,c), so one kernel suffices.
// =====================================================================
__global__ void transpose_4d12_k(const float* in, float* out, int a, int b, int c, int d) {
int idx = blockIdx.x * blockDim.x + threadIdx.x; // over a*b*c*d
if (idx >= a * b * c * d) return;
int l = idx % d;
int k = (idx / d) % c;
int j = (idx / (d * c)) % b;
int i = idx / (d * c * b);
// out[i,k,j,l] at ((i*c + k)*b + j)*d + l
out[(((i * c + k) * b) + j) * d + l] = in[idx];
}
void launch_transpose_4d12_f32(const float* in, float* out, int a, int b, int c, int d, void* s) {
int n = a * b * c * d, blk = 256, grid = (n + blk - 1) / blk;
transpose_4d12_k<<<grid, blk, 0, (cudaStream_t)s>>>(in, out, a, b, c, d);
}
} // extern "C"

24
csrc/ops/nn.cu
View File

@@ -215,14 +215,20 @@ void launch_silu_dx_f32(const float* x, const float* dy, float* dx, int n, void*
// dx[i+h] = dy[i+h]*cos - dy[i]*sin
// =====================================================================
__global__ void rope_k(const float* x, float* y, int heads, int head_dim, float theta) {
// `period` is the sequence length: a flattened batch lays B sequences end to end
// along the `tokens` axis, so each token's RoPE position is its index WITHIN its
// own sequence, `tok % period`. With period == tokens (single sequence) this is
// the original position = row.
__global__ void rope_k(const float* x, float* y, int heads, int head_dim,
float theta, int period) {
int tok = blockIdx.x;
int head = blockIdx.y;
int half = head_dim / 2;
int i = threadIdx.x;
if (i >= half) return;
int pos = tok % period;
float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)tok * freq;
float angle = (float)pos * freq;
float c = cosf(angle), sn = sinf(angle);
int base = (tok * heads + head) * head_dim;
float x0 = x[base + i], x1 = x[base + i + half];
@@ -230,20 +236,22 @@ __global__ void rope_k(const float* x, float* y, int heads, int head_dim, float
y[base + i + half] = x1 * c + x0 * sn;
}
void launch_rope_f32(const float* x, float* y, int tokens, int heads,
int head_dim, float theta, void* s) {
int head_dim, float theta, int period, void* s) {
dim3 grid(tokens, heads);
int blk = head_dim / 2;
rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta);
rope_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, period);
}
__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, float theta) {
__global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim,
float theta, int period) {
int tok = blockIdx.x;
int head = blockIdx.y;
int half = head_dim / 2;
int i = threadIdx.x;
if (i >= half) return;
int pos = tok % period;
float freq = powf(theta, -(float)(2 * i) / (float)head_dim);
float angle = (float)tok * freq;
float angle = (float)pos * freq;
float c = cosf(angle), sn = sinf(angle);
int base = (tok * heads + head) * head_dim;
float d0 = dy[base + i], d1 = dy[base + i + half];
@@ -251,10 +259,10 @@ __global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, f
dx[base + i + half] = d1 * c - d0 * sn;
}
void launch_rope_dx_f32(const float* dy, float* dx, int tokens, int heads,
int head_dim, float theta, void* s) {
int head_dim, float theta, int period, void* s) {
dim3 grid(tokens, heads);
int blk = head_dim / 2;
rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta);
rope_dx_k<<<grid, blk, 0, (cudaStream_t)s>>>(dy, dx, heads, head_dim, theta, period);
}
// =====================================================================