Module ops 

GPU op wrappers: bind tensors, compute grid sizes, dispatch WGSL kernels. Semantics mirror backend::cpu exactly; see the shader sources for the kernel-side contracts.

Functions

adamw

AdamW step, updating param/m/v in place.

add

alloc_storage

ce_bwd

dropout

embedding

Gather from row-chunked embedding tables. chunks[i] holds rows [i * chunk_rows, …) of the full table; each token is written by exactly the dispatch whose chunk owns its id.

gather_nll

gelu

gelu_bwd

kv_append

Append src [h, t, hd] into cache [h, cap, hd] at row offset len.

layernorm

layernorm_bwd

matmul

matmul_into

Compute n columns of a wider [m, n_out] output, writing at column offset n_off — used to keep each weight-chunk binding under max_storage_buffer_binding_size.

merge_heads

scale

scatter_add_rows

dst[ids[r]] += src[r], in place (CAS-loop f32 atomics).

softmax

softmax_bwd

split_heads

sum_rows

sumsq_partials

Per-workgroup partial sums of squares; the caller reads back and sums. The partial count covers the whole (possibly 2-D) dispatch grid; groups past the data contribute zeros.

unmerge_heads

unsplit_head