Module cpu 

CPU reference backend. Mathematically identical semantics to the WGSL kernels; used for testing, verification, and gradient checking.

Functions

adamw

AdamW step (decoupled weight decay), in place. step is 1-based.

add

ce_bwd

Cross-entropy backward: dlogits = (probs - onehot(ids)) * scale.

dropout

Inverted dropout with a counter-based RNG: element i is kept when hash(seed, i) maps to u >= p, and scaled by scale (the caller-computed 1/(1-p)). Applying the same (seed, p, scale) to dy gives the backward pass. scale is passed in (rather than recomputed here) so the CPU and WGSL kernel multiply by the exact same bits — see ops::dropout.

embedding

Fused token + positional embedding gather.

gather_nll

-log(probs[r, ids[r]]) per row.

gelu

GELU, tanh approximation (“gelu_new”) — the variant GPT-2 was trained with.

gelu_bwd

d/dx of GELU (tanh approximation), applied to upstream dy.

kv_append

Append src [h, t, hd] into dst [h, cap, hd] at row offset len per head.

layernorm

LayerNorm over the last dim (biased variance, like PyTorch).

layernorm_bwd_dparams

LayerNorm gamma/beta gradients: dgamma = sum_r dy * xhat, dbeta = sum_r dy.

layernorm_bwd_dx

LayerNorm backward for x: returns dx; gamma/beta gradients are computed by layernorm_bwd_dparams.

matmul

Batched matmul matching shaders/matmul.wgsl: C[b] = alpha * A[b] @ B[b] (+ bias broadcast over rows). a_stride/b_stride are per-batch element strides (0 broadcasts). With trans_a, A[b] is stored [k, m]; with trans_b, B[b] is stored [n, k] instead of [k, n].

merge_heads

x: [h, t, hd] -> [t, c], c = h * hd.

pcg_hash

PCG output hash — identical to shaders/dropout.wgsl so masks match bit-for-bit across backends.

scatter_add_rows

dwte scatter-add: dst[ids[r]] += src[r] for rows of width c.

softmax

Stable softmax over the last dim with optional causal masking, matching shaders/softmax.wgsl. Row r’s query position is r % q_len; key j is visible when j <= query + off. Masked entries produce 0.

softmax_bwd

Softmax backward: dx = y * (dy - sum(y * dy)) per row. The causal-masked forward writes exact zeros at masked entries, so no mask is needed here.

split_heads

qkv: [t, 3c] -> (q, k, v) each [h, t, hd], hd = c / h.

sum_rows

Column sums: [rows, cols] -> [cols]. Bias gradients.

unmerge_heads

merge_heads backward: dy [t, c] -> [h, t, hd].

unsplit_head

split_heads backward for one of q/k/v: place d [h, t, hd] into the which third of a zeroed [t, 3c] buffer.