Module ops 

Backend-agnostic tensor operations. Each op validates shapes, dispatches to the CPU reference or the WGPU kernels, and returns a tensor on the same device as its inputs.

Structs

MatmulSpec

Functions

adamw

AdamW step with decoupled weight decay, updating param/m/v in place. step is 1-based.

add

Elementwise a + b; shapes must match exactly.

ce_bwd

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

dropout

Inverted dropout with a deterministic counter RNG (identical masks on both backends). Apply the same (p, seed) to dy for the backward pass.

embedding

Fused token + positional embedding: out[t] = wte[ids[t]] + wpe[t + pos].

embedding_chunked

Embedding gather from a row-chunked table. chunks[i] holds rows [i * chunk_rows, …) of the full [vocab, c] table — chunking keeps every GPU binding under max_storage_buffer_binding_size (GPT-2’s wte alone is ~147 MiB, above the WebGPU default and llvmpipe’s hard 128 MiB limit).

gather_nll

Per-row NLL: out[r] = -log(probs[r, ids[r]]).

gelu

GELU (tanh approximation).

gelu_bwd

GELU backward: dx = gelu’(x) * dy.

kv_append

Append src [h, t, hd] into the preallocated cache [h, cap, hd] at row offset len within each head, in place. Used by KV-cache decode.

layernorm

LayerNorm over the last dim.

layernorm_bwd

LayerNorm backward: (dx, dgamma, dbeta).

matmul

Batched matmul with optional A/B-transpose, scaling, and bias.

matmul_chunked_transb

a [m, k] @ concat(chunks)^T -> [m, sum n_i], chunks each [n_i, k]. The column-chunked equivalent of matmul with trans_b — used for the weight-tied LM head over a chunked wte.

merge_heads

x [h, t, hd] -> [t, h*hd].

scale

y = x * alpha.

scatter_add_rows

dst[ids[r]] += src[r], in place. Embedding-table gradient scatter.

softmax

Softmax over the last dim. With causal, row r (query position r % q_len, where q_len is the second-to-last dim) sees key j only when j <= query + off; off is key_len - query_len when using a KV cache.

softmax_bwd

Softmax backward from the forward output y: dx = y * (dy - sum(y*dy)) per row. Causal masking needs no parameters here — masked forward outputs are exact zeros.

split_heads

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

sum_rows

Column sums over all leading dims: [.., cols] -> [cols]. Bias gradients.

sumsq

Sum of squares of all elements (for the global gradient norm). Training-path op with a sync readback — WebGPU tensors are native-only here (browser training is out of scope for 1.0).

unmerge_heads

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

unsplit_head

split_heads backward for one of q/k/v (which in 0..3): d [h, t, hd] -> [t, 3c] with the other thirds zero.