Machine Learning (std/ml/) #

A CPU tensor/autograd core, a wide attention/diffusion/architecture library, GPU backends, and an LLM orchestration layer (LangChain/LangGraph/LangSmith-inspired).

#include <std/ml/tensor.h>

Every function below takes/returns FloatTensor (and the other structs here) by a region-less reference (&FloatTensor/const &FloatTensor — see Memory & Regions's "Outliving References" section): nothing in this library cares whether the caller's tensor is heap/stack/static/ arena-backed, so no function pins one specific region. A raw pointer (or a reference of any region) converts to these implicitly, with no unsafe.

Stateful types (models, layers, optimizers, data loaders, agents, device nets and buffers) are used through methods: net.forward(x), v.loss(x, y), opt.step(), buf.free(). Constructors are static methods: Mlp::new(...), FloatTensor::new_2d(...), Dataset::moons(...), SparseTensor::from_dense(&d, 0.0f, 1), TrainConfig::default(). An operation on one tensor is a method of that tensor — x.relu(), logits.log_softmax(), loss.backward(), t.fill(0.0f) — const unless it changes the tensor in place; operations on two tensors stay functions (tensor_matmul(a, b)). What every model can do — add its parameters to an optimizer, count them, save and load them, copy them, train on a built-in loss — is written once as generic functions over traits the models satisfy: model_add_params(&net, &opt), model_save_file(&net, path), module_fit(&net, TRAIN_LOSS_CE, &opt, &data, &cfg), model_to(&net, DEVICE_ROCM).

Tensor #

#include <std/ml/tensor.h>

CPU tensors (1D/2D) with reverse-mode autograd — PyTorch-style ergonomics: build a computation graph implicitly by calling ops on requiresGrad tensors, then call FloatTensor::backward().

A tensor is a DenseTensor<E>: values stored as E, computed in float32, gradients always float32. Three element types are provided:

The core ops below (tensor_add/sub/mul/matmul, scale, sum), their autograd and the backward pass work on all three. A FloatTensor op runs the float32 kernel on the tensor itself; a 16-bit one loads its operands to float32, runs the same kernel and rounds its result once. The backward pass is the float32 one either way, and a graph may mix element types (to_float() / from_float() are graph nodes). The other ops in this library (activations, losses, layers) take a FloatTensor. See Half-precision tensors for conversions, batch helpers and mixed-precision training.

&FloatTensor x = FloatTensor::from_1d(((const float*)values)[0..3UL], /*requiresGrad=*/1);
&FloatTensor y = tensor_mul(x, x).sum();  // y = sum(x^2)
y.backward();                               // x.grad now holds 2x

Ops: tensor_add/tensor_sub/tensor_mul/FloatTensor::scale/tensor_matmul/ FloatTensor::relu/FloatTensor::sum. Verified against hand-computed gradients and a real finite-difference check on a matmul→relu→sum chain.

tensor_matmul and its backward pass run on sgemm: packed SIMD kernels at the CPU's best SIMD level (AVX-512 down to SSE2, SVE or Neon), split across the thread pool for larger products. With a CPU vendor BLAS linked — AOCL-BLAS on AMD, oneMKL on Intel, Accelerate on Apple, which safeguard's [build] blas does by default when it finds one — the vendor's kernels do the multiplying. The backward pass accumulates straight into the existing gradients. Set SAFEC_NUM_THREADS to limit the thread count.

Elementwise ops (tensor_add/sub/mul/scale/fill, ReLU and every backward accumulation) are vectorized loops that split across the pool from 65,536 elements. The transcendental activations — sigmoid, tanh, GELU, SiLU, ELU, GLU, SwiGLU — and softmax use vmath.h's vectorized exp/tanh/sigmoid, so they no longer call libm per element.

generic<E>
struct DenseTensor {
    &heap E             data;          // flat row-major, length == size
    &heap unsigned long shape;
    unsigned long       ndim;          // 1 or 2
    unsigned long       size;
    int                 requiresGrad;
    &heap float         grad;          // float32; allocated only when requiresGrad
    // ... autograd bookkeeping (gradFn, parents, visited)

    float at1(unsigned long i) const;                      // as float32
    float at2(unsigned long r, unsigned long c) const;
    void  set1(unsigned long i, float v);                  // rounded to E
    void  set2(unsigned long r, unsigned long c, float v);
    void  free() consume;
    // ... and the methods below
};

typedef struct DenseTensor<float>       FloatTensor;
typedef struct DenseTensor<struct F16>  HalfTensor;
typedef struct DenseTensor<struct BF16> BF16Tensor;

// What generic ML code asks of a tensor; all three satisfy it.
trait Tensor {
    float at1(unsigned long i) const;
    float at2(unsigned long r, unsigned long c) const;
    void  set1(unsigned long i, float v);
    void  set2(unsigned long r, unsigned long c, float v);
    void  fill(float v);
    void  backward();
    void  zero_grad();
    void  free();
}

// Shown for FloatTensor; HalfTensor:: and BF16Tensor:: are the same.
static &FloatTensor FloatTensor::new_1d(unsigned long n, int requiresGrad);
static &FloatTensor FloatTensor::new_2d(unsigned long rows, unsigned long cols, int requiresGrad);
static &FloatTensor FloatTensor::from_1d([]const float values, int requiresGrad);
static &FloatTensor FloatTensor::from_2d([]const float values, unsigned long rows, unsigned long cols, int requiresGrad);   // values.len >= rows * cols
&FloatTensor FloatTensor::zeros_like() const;
void    FloatTensor::fill(float v);   // in place

// Generic over the element type: both operands are the same DenseTensor<E>.
generic<E> &DenseTensor<E> tensor_add(const &DenseTensor<E> a, const &DenseTensor<E> b);
generic<E> &DenseTensor<E> tensor_sub(const &DenseTensor<E> a, const &DenseTensor<E> b);
generic<E> &DenseTensor<E> tensor_mul(const &DenseTensor<E> a, const &DenseTensor<E> b);   // elementwise
generic<E> &DenseTensor<E> tensor_matmul(const &DenseTensor<E> a, const &DenseTensor<E> b);
&FloatTensor FloatTensor::scale(float k) const;
&FloatTensor FloatTensor::sum() const;

// Between element types, as graph nodes (gradients pass through unchanged).
&FloatTensor FloatTensor::to_float() const;                       // a FloatTensor: itself
static &HalfTensor HalfTensor::from_float(const &FloatTensor x);  // rounded

void FloatTensor::backward();    // t: a scalar (a loss); accumulates every ancestor's grad
void FloatTensor::zero_grad();

// Gradient recording: the program's autograd state. Disabled (inference),
// composed ops record no graph (leaf flags unchanged).
struct GradGuard { int was; void restore(); };
struct Autograd {
    int  enabled() const;
    void set_enabled(int enabled);
    struct GradGuard no_grad();   // disables recording; restore() puts the mode back
};
static &Autograd Autograd::get();

struct GradGuard g = Autograd::get().no_grad();
// ... forward passes record nothing ...
g.restore();

Tensor arenas #

Every op allocates its output, so a training loop does a malloc (and, for fresh memory, page faults) per op per step. On a small MLP that cost more than the arithmetic. A TensorArena is a chunked bump allocator, 64-byte aligned, for one step's temporaries. Between TensorArena::begin and tensor_arena_end, every tensor created on that thread comes from the arena: its data, shape and gradient. TensorArena::reset drops them all at once and keeps the memory, already faulted in, for the next step.

static struct TensorArena TensorArena::new(unsigned long chunkBytes);   // 0 = 16 MiB chunks
void TensorArena::begin();                                       // this thread's new tensors come from this arena
void tensor_arena_end();                                         // back to the heap
void TensorArena::reset();                                       // drop every tensor in it, keep the memory
void TensorArena::free();
unsigned long TensorArena::used() const;   // bytes since the last reset

struct TensorArena step = TensorArena::new(0UL);
while (s < steps) {
    step.begin();
    // ... forward, FloatTensor::backward, optimizer update ...
    tensor_arena_end();
    step.reset();   // this step's tensors are gone
    s = s + 1;
}
step.free();

Tensors created outside the arena, such as parameters, stay on the heap, and so do their gradients, so they survive resets. FeatureMaps created while an arena is active come from it too, every intermediate inside UNet::forward included; FeatureMap::free leaves those to the reset. Each tensor records its owner. FloatTensor::free() on an arena tensor does nothing, and a tensor must not be used after its arena is reset.

Autograd ops and optimizer — tensor_nn.h #

#include <std/ml/tensor_nn.h>

Activations with backward passes (the ones every GPU backend uses), 1D pooling, and Adam.

&FloatTensor FloatTensor::relu() const;
&FloatTensor FloatTensor::sigmoid() const;
&FloatTensor FloatTensor::tanh() const;
&FloatTensor FloatTensor::softmax() const;
&FloatTensor FloatTensor::elu(float alpha) const;
&FloatTensor FloatTensor::gelu() const;
&FloatTensor FloatTensor::silu() const;
&FloatTensor FloatTensor::glu() const;      // split last axis into [a1, a2]: a1 * sigmoid(a2)
&FloatTensor FloatTensor::swiglu() const;   // silu(a1) * a2 — output has half the columns
&FloatTensor FloatTensor::maxpool1d(unsigned long kernel, unsigned long stride) const;
&FloatTensor FloatTensor::avgpool1d(unsigned long kernel, unsigned long stride) const;

struct AdamState {
    &heap float m; &heap float v;
    unsigned long size;
    float lr; float beta1; float beta2; float eps;
    unsigned long t;
};
static struct AdamState AdamState::new(unsigned long size, float lr, float beta1, float beta2, float eps);
static struct AdamState AdamState::new_default(unsigned long size, float lr);   // beta1 0.9, beta2 0.999, eps 1e-8 (paper defaults)
void AdamState::step(&FloatTensor param);                               // uses param.grad
void AdamState::free();

Activations #

#include <std/ml/activations.h>

Forward-only elementwise activations (no autograd — for inference layers; the autograd versions are in tensor_nn.h above); FloatTensor::layernorm_rows (row-wise mean-0/var-1 normalization, no affine — apply your own scale/shift for AdaLN-style conditioning, see DiT below); tensor_residual_add (a named tensor_add alias for residual/skip- connection call sites).

&FloatTensor FloatTensor::relu_fwd() const;
&FloatTensor FloatTensor::sigmoid_fwd() const;
&FloatTensor FloatTensor::tanh_fwd() const;
&FloatTensor FloatTensor::silu_fwd() const;
&FloatTensor FloatTensor::gelu_fwd() const;
&FloatTensor FloatTensor::layernorm_rows(float eps) const;
&FloatTensor tensor_residual_add(const &FloatTensor x, const &FloatTensor sublayerOut);

Attention #

#include <std/ml/attention.h>

softmax_rows() and attention_forward() (single-head scaled dot-product attention) build mha_forward() (standard multi-head attention — splits Q/K/V's columns into heads, runs attention per head, concatenates results). Also in this file:

• gated_attention_forward/gated_mha_forward — sigmoid output gating, elementwise-multiplied into the attention result.
• linear_attention_forward — the phi(x) = elu(x)+1 kernel-trick reformulation (Katharopoulos et al.), O((seqLen+kvLen)·headDim²) instead of the full O(seqLen·kvLen·headDim) score matrix.
• flash_attention_forward — tiled, online-softmax attention (Dao et al.): processes K/V in blocks, never materializing the full [seqLen,kvLen] score matrix. Verified to match attention_forward's output within tolerance for a block size that doesn't evenly divide kvLen.

Every one of them records its backward edges like tensor_matmul does, so attention layers train, and runs where its operands live (devices). Q/K/V are [seqLen, dim] tensors: no batch dimension (call once per sequence) and no positional encoding baked in.

&FloatTensor FloatTensor::softmax_rows() const;
&FloatTensor FloatTensor::softmax_rows_mask(int causal, unsigned long window) const;
&FloatTensor FloatTensor::transpose() const;
&FloatTensor FloatTensor::slice_cols(unsigned long startCol, unsigned long numCols) const;
void    tensor_set_cols(&FloatTensor dst, unsigned long startCol, const &FloatTensor src);   // in place, no gradient

&FloatTensor attention_forward(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V);
&FloatTensor mha_forward(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V, unsigned long numHeads);
&FloatTensor attention_forward_mask(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V,
                                    int causal, unsigned long window);
&FloatTensor mha_forward_mask(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V,
                              unsigned long numHeads, int causal, unsigned long window);
&FloatTensor gated_attention_forward(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V, const &FloatTensor gateLogits);
&FloatTensor gated_mha_forward(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V,
                          unsigned long numHeads, const &FloatTensor gateLogits);
&FloatTensor linear_attention_forward(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V);
&FloatTensor flash_attention_forward(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V, unsigned long kvBlockSize);
int attention_set_fused(int on);   // GPU multi-head ops: 1 fused kernels (default), 0 composed ops; returns the previous

Masks. Query row i sits at position pos = i + (kvLen - seqLen): the last seqLen keys are the queries' own tokens, any before them a KV cache. causal hides keys j > pos; a sliding window w > 0 hides keys j <= pos - w, so each query sees its last w positions. A query with no visible key gets zeros. These are the masks of the GPU inference kernels below (ROCm, CUDA, Metal).

How it runs.

The fused op is FlashAttention-2's tiling (Dao 2023) in device_kernels.cu, shared by ROCm and CUDA: a block of 256 threads takes one head's tile of 16 query rows and walks the keys in tiles of 16 through shared memory with an online softmax, so the [seqLen, kvLen] score matrix is never formed; tiles a causal or windowed mask hides entirely are skipped. It keeps each row's log-sum-exp, and its backward is two more launches (dQ by query tiles, then dK and dV by key tiles) that recompute the scores from it. On Metal, and on the CPU for flash_attention_forward, the same arithmetic runs one row at a time (device_ref.sc). With 16-bit activations (precision) the kernels read and write fp16 / bf16 and compute in fp32.

On an MI100, a causal 4-head forward + backward pass:

The inference kernels in gpu_rocm.h, gpu_cuda_attention.h and gpu_mps.h (grouped-query K/V, bf16, matrix cores) stay the faster choice for a forward pass alone on their own [heads][seqLen][headDim] buffers; they keep no log-sum-exp and have no backward.

MultiHeadAttention is a self-attention layer with its own projections, mha_forward_mask(X Wq, X Wk, X Wv, numHeads, causal, window) Wo. Its weights require grad, and it is a Module. Set causal for a decoder.

struct MultiHeadAttention { &FloatTensor Wq; &FloatTensor Wk; &FloatTensor Wv; &FloatTensor Wo;
                            unsigned long numHeads; int causal; unsigned long window; };
static struct MultiHeadAttention MultiHeadAttention::new(unsigned long dModel, unsigned long numHeads, &Rng rng);
&FloatTensor MultiHeadAttention::forward(const &FloatTensor X) const;   // [seqLen, dModel] -> [seqLen, dModel]
void    MultiHeadAttention::free();

A small decoder language model, trained on a GPU through trainer_fit (a batch is one sequence: the dataset's rows are consecutive tokens, x the token and y the next one, iterated in order):

struct TinyLm {
    &FloatTensor E;      // [vocab, dim] token embedding
    &FloatTensor P;      // [seq, dim] learned positions
    &FloatTensor Wout;   // [dim, vocab]
    struct LayerNorm ln1; struct LayerNorm ln2;
    struct MultiHeadAttention mha;   // mha.causal = 1
    struct Mlp ffn;
};

static &FloatTensor lm_loss(&TinyLm m, const &FloatTensor x, const &FloatTensor y) {
    &FloatTensor h = tensor_add(tensor_gather_rows(m.E, x), m.P);
    &FloatTensor h1 = tensor_add(h, m.mha.forward(m.ln1.forward(h)));
    &FloatTensor h2 = tensor_add(h1, m.ffn.forward(m.ln2.forward(h1)));
    return tensor_cross_entropy_loss(tensor_matmul(h2, m.Wout), y);
}

// ... add every parameter to opt, then:
struct TrainConfig cfg = TrainConfig::default();
cfg.batchSize = seq; cfg.shuffle = 0; cfg.device = DEVICE_AUTO;
struct TrainReport rep = trainer_fit(&lm, lm_loss, &opt, &tokens, &cfg);

Advanced attention #

#include <std/ml/attention_advanced.h>

• Mla/EgMla — Multi-head Latent Attention (DeepSeek-V2): compresses K/V into a shared low-rank latent before up-projecting back to full width, saving KV-cache memory. EgMla shares one latent across independent per-group up-projections (grouped-query attention's KV-sharing layered onto MLA's compression, the groups joined with tensor_concat_cols). Both are layers holding their weights, which require grad, and Modules: they train like MultiHeadAttention, on a GPU too.
• windowed_attention_forward — Shifted-Window (Swin) attention adapted to 1D sequences: partitions into fixed-size local windows, with a cyclic shift (shift param) so the next layer's windows straddle this layer's boundaries.
• GatedDeltaNet::forward — Gated DeltaNet (Yang et al. 2024): a linear-recurrent attention replacement using a delta-rule state update (write the residual between the new value and what the state already predicts) plus a scalar decay gate, computed via a sequential scan.

Verified via self-consistency checks against manual per-group/per-window reconstruction, and a hand-computed 2-step state trace for Gated DeltaNet. windowed_attention_forward and GatedDeltaNet are forward-only host loops (for a trainable local window, use mha_forward_mask's window).

struct Mla { &FloatTensor Wq; &FloatTensor Wdkv; &FloatTensor Wuk; &FloatTensor Wuv; unsigned long numHeads; };
static struct Mla Mla::new(unsigned long dModel, unsigned long dLatent, unsigned long numHeads, &Rng rng);
&FloatTensor Mla::forward(const &FloatTensor X) const;           // [seqLen, dModel] -> [seqLen, dModel]

struct EgMla { &FloatTensor Wq; &FloatTensor Wdkv; struct Vec<&FloatTensor> Wuk; struct Vec<&FloatTensor> Wuv;  // one per group
               unsigned long headsPerGroup; };
static struct EgMla EgMla::new(unsigned long dModel, unsigned long dLatent, unsigned long numGroups, unsigned long headsPerGroup, &Rng rng);
&FloatTensor EgMla::forward(const &FloatTensor X) const;
&FloatTensor windowed_attention_forward(const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V,
                                   unsigned long windowSize, unsigned long shift);

struct GatedDeltaNet { &FloatTensor Walpha; &FloatTensor Wbeta; unsigned long keyDim; unsigned long valueDim; };
static struct GatedDeltaNet GatedDeltaNet::new(unsigned long dModel, unsigned long keyDim, unsigned long valueDim);
&FloatTensor GatedDeltaNet::forward(const &FloatTensor X, const &FloatTensor Q, const &FloatTensor K, const &FloatTensor V) const;
void    GatedDeltaNet::free();

The weights are public fields: to load pretrained ones, assign them (an EgMla's per-group up-projections are Vecs whose length is the group count). free() releases them.

RNN family #

#include <std/ml/rnn.h>

RNNCell/GRUCell/LSTMCell/XLSTMCell — plain RNN, GRU, LSTM, and the xLSTM sLSTM variant (Beck et al. 2024: exponential input/forget gating with a log-domain stabilizer and a normalizer state, replacing LSTM's sigmoid gates). Row-vector convention (x_t is [1, inputSize]). xLSTM's mLSTM (matrix-memory) variant is not implemented — sLSTM only. Each cell verified against hand-computed gate values for a real forward pass.

static struct RNNCell  RNNCell::new(unsigned long inputSize, unsigned long hiddenSize);
&FloatTensor         RNNCell::forward(const &FloatTensor x_t, const &FloatTensor h_prev) const;
void            RNNCell::free();

static struct GRUCell  GRUCell::new(unsigned long inputSize, unsigned long hiddenSize);
&FloatTensor         GRUCell::forward(const &FloatTensor x_t, const &FloatTensor h_prev) const;
void            GRUCell::free();

static struct LSTMCell LSTMCell::new(unsigned long inputSize, unsigned long hiddenSize);
// Returns h_t; the new cell state is written through c_out.
&FloatTensor         LSTMCell::forward(const &FloatTensor x_t, const &FloatTensor h_prev, const &FloatTensor c_prev, &&FloatTensor c_out) const;
void            LSTMCell::free();

struct XLSTMState { &FloatTensor h; &FloatTensor c; &FloatTensor n; &FloatTensor m; };   // hidden, cell, normalizer, stabilizer
static struct XLSTMCell  XLSTMCell::new(unsigned long inputSize, unsigned long hiddenSize);
struct XLSTMState XLSTMCell::forward(const &FloatTensor x_t, const &XLSTMState prev) const;
void              XLSTMCell::free();
void              XLSTMState::free();

Each cell's weight tensors are public fields (W_ih, W_hh, … — named per gate) so you can load trained weights into them.

State space models (Mamba) #

#include <std/ml/ssm.h>

A state space model mixes a sequence through a linear recurrence per channel instead of attention. Each of D channels carries an N-value hidden state that every token updates and reads:

h_t = exp(Δ_t·A)·h_{t−1} + Δ_t·B_t·x_t, y_t = C_t·h_t + D·x_t

A sequence costs O(T·D·N). Decoding one token costs O(D·N) with a state of fixed size, however long the context, so there is no KV cache. Mamba (Gu & Dao 2023, "S6") makes the step Δ and the B and C projections functions of the input. This lets the model choose, token by token, what to keep and what to forget.

ssm_selective_scan is the scan alone. Mamba is one Mamba-1 mixer block, without the surrounding RMSNorm and residual add:

1. (x, z) = X·inProj.
2. A causal depthwise conv1d (kernel dConv) and SiLU on x.
3. (dt, B, C) = x·xProj.
4. Δ = softplus(dt·dtProj + dtBias).
5. The scan with A = −exp(Alog).
6. Gate with silu(z).
7. outProj.

Weights are [in, out] like the rest of std::ml. To load a PyTorch nn.Linear, transpose its weight. The Conv1d weight [dInner, 1, dConv] loads as [dInner, dConv]. Mamba::new follows the paper's initialization: Alog[d][n] = log(n+1), D = 1, and dtBias chosen so that softplus(dtBias) spans [0.001, 0.1].

The convolution and the scan run the channels in parallel on the thread pool. The projections go through tensor_matmul. Mamba::step decodes one token and advances a MambaState (the last dConv − 1 convolution inputs plus the SSM state). Feeding tokens one at a time through Mamba::step gives the same rows as Mamba::forward on the whole sequence, and the tests check this. With weights copied into HF transformers' MambaMixer, the outputs match to 2e-8 absolute, on outputs of about 0.05. Forward only, no autograd.

// Selective scan alone: x, delta [T, D]; A [D, N]; B, C [T, N]; Dskip [D] -> y [T, D]
&FloatTensor ssm_selective_scan(const &FloatTensor x, const &FloatTensor delta, const &FloatTensor A,
                           const &FloatTensor B, const &FloatTensor C, const &FloatTensor Dskip);

struct Mamba {
    &FloatTensor inProj;   // [dModel, 2 * dInner]
    &FloatTensor convW;    // [dInner, dConv]
    &FloatTensor convB;    // [dInner]
    &FloatTensor xProj;    // [dInner, dtRank + 2 * dState]
    &FloatTensor dtProj;   // [dtRank, dInner]
    &FloatTensor dtBias;   // [dInner]
    &FloatTensor Alog;     // [dInner, dState]; A = -exp(Alog)
    &FloatTensor Dskip;    // [dInner]
    &FloatTensor outProj;  // [dInner, dModel]
    unsigned long dModel; unsigned long dInner; unsigned long dState; unsigned long dConv; unsigned long dtRank;
};
static struct Mamba Mamba::new(unsigned long dModel, unsigned long dState, unsigned long dConv, unsigned long expand);
&FloatTensor      Mamba::forward(const &FloatTensor X) const;   // [T, dModel] -> [T, dModel]
void         Mamba::free();

struct MambaState { &FloatTensor conv; &FloatTensor h; };   // [dInner, dConv-1], [dInner, dState]
static struct MambaState MambaState::new(const &Mamba layer);
&FloatTensor           Mamba::step(&MambaState state, const &FloatTensor x) const;   // [1, dModel]
void              MambaState::reset();
void              MambaState::free();

struct Mamba m = Mamba::new(768UL, 16UL, 4UL, 2UL);
&FloatTensor y = m.forward(prompt);   // prefill
struct MambaState st = MambaState::new(&m);
// ... or decode token by token with a fixed-size state
&FloatTensor yt = m.step(&st, token);

CNN #

#include <std/ml/cnn.h>

These layers run forward only. To train convolutional networks, use convnet.h, which has the same operations on autograd tensors.

Conv2D (stride + zero-padding, no dilation/grouped conv), FeatureMap::maxpool2d/FeatureMap::avgpool2d, FeatureMap::upsample2x_nearest (nearest- neighbor 2x spatial upsample), concat_channels (channel-axis concatenation) — all over a channels-first FeatureMap (not FloatTensor, which tops out at 2D). Verified via hand-computed 3×3-kernel and pooling-window checks.

3×3, stride-1 layers (with at least 8 input and output channels) use Winograd F(2×2, 3×3) (Lavin & Gray, arXiv:1509.09308). Each 2×2 block of outputs comes from a 4×4 input tile with 16 multiplies per channel pair instead of 36. The transformed tiles and filters turn the work into 16 GEMMs over channels. The input transform, the GEMMs and the output transform each run on the thread pool, with the input and output transforms working on eight tiles per vector. Results match the direct sum to about 1e-6 relative. Other layers use im2col:

Conv2D::forward is im2col + sgemm. Each block of output pixels gathers its receptive fields into a per-thread column buffer and multiplies it by the weight matrix. Blocks run on the thread pool, so the full im2col matrix is never built. 1×1 convolutions skip the gather. Conv2D::relu_forward applies the ReLU in the same pass. Pooling, upsampling and concatenation are split across the pool by channel once a map is large enough.

struct FeatureMap {
    &heap float   data;      // channels × height × width, channels-first
    unsigned long channels;
    unsigned long height;
    unsigned long width;
};
static struct FeatureMap FeatureMap::new(unsigned long channels, unsigned long height, unsigned long width);
void              FeatureMap::free();

static struct Conv2D Conv2D::new(unsigned long inChannels, unsigned long outChannels, unsigned long kH, unsigned long kW, unsigned long stride, unsigned long padding);
struct FeatureMap Conv2D::forward(const &FeatureMap input) const;
struct FeatureMap Conv2D::relu_forward(const &FeatureMap input) const;   // relu(conv), one pass
void              Conv2D::free();

struct FeatureMap FeatureMap::maxpool2d(unsigned long kernel, unsigned long stride) const;
struct FeatureMap FeatureMap::avgpool2d(unsigned long kernel, unsigned long stride) const;
struct FeatureMap FeatureMap::upsample2x_nearest() const;
struct FeatureMap concat_channels(const &FeatureMap a, const &FeatureMap b);

Padded and global pooling, residual blocks #

The _ex pooling functions add zero padding on every side. The output side is (dim + 2·padding − kernel)/stride + 1. Max pooling skips the padded positions, as PyTorch's MaxPool2d does. Average pooling counts them as zeros in the divisor when countIncludePad is set (PyTorch's AvgPool2d default). Otherwise it divides by the number of positions inside the plane. FeatureMap::global_avgpool2d / FeatureMap::global_maxpool2d reduce each channel's whole plane to one value, giving channels × 1 × 1: the head of a ResNet-style classifier.

feature_map_add is the residual connection, a + b element by element, with an optional ReLU in the same pass. ResidualBlock is the ResNet basic block (He et al. 2016):

relu(conv2(relu(conv1(x))) + shortcut(x))

conv1 is 3×3 with the block's stride. conv2 is 3×3 with stride 1. Both use padding 1. The shortcut is x itself, or a 1×1 convolution with the stride when the stride isn't 1 or the channel count changes. There is no batch norm: fold a trained model's batch norms into the convolutions' weights and biases before loading them. The tests check padded pooling against a direct reference, and the block against the composition of its convolutions.

struct FeatureMap FeatureMap::maxpool2d_ex(unsigned long kernel, unsigned long stride, unsigned long padding) const;
struct FeatureMap FeatureMap::avgpool2d_ex(unsigned long kernel, unsigned long stride, unsigned long padding, int countIncludePad) const;
struct FeatureMap FeatureMap::global_avgpool2d() const;   // channels x 1 x 1
struct FeatureMap FeatureMap::global_maxpool2d() const;
struct FeatureMap feature_map_add(const &FeatureMap a, const &FeatureMap b, int relu);

struct ResidualBlock { struct Conv2D conv1; struct Conv2D conv2; struct Conv2D projection; int hasProjection; };
static struct ResidualBlock ResidualBlock::new(unsigned long inChannels, unsigned long outChannels, unsigned long stride);
struct FeatureMap    ResidualBlock::forward(const &FeatureMap input) const;
void                 ResidualBlock::free();

// ResNet-18 stem and first stage
struct FeatureMap s = stem7x7s2.relu_forward(&image);
struct FeatureMap p = s.maxpool2d_ex(3UL, 2UL, 1UL);
struct FeatureMap r = layer1a.forward(&p);

U-Net #

#include <std/ml/unet.h>

A fixed 2-level encoder/decoder CNN with skip connections (Ronneberger et al. 2015) — the classic DDPM-style diffusion denoiser backbone. Verified via output-shape checks and full self-consistency against a manual replication of UNet::forward's call sequence.

static struct UNet UNet::new(unsigned long inChannels, unsigned long c1, unsigned long c2, unsigned long outChannels);
struct FeatureMap UNet::forward(const &FeatureMap input);
void UNet::free();

Diffusion #

#include <std/ml/diffusion.h>

• edm_karras_sigmas — the Karras et al. ("Elucidating the Design Space...") noise schedule.
• ddpm_linear_schedule/ddpm_sampler_step — DDPM ancestral sampling.
• ddim_sampler_step — deterministic DDIM (eta=0).
• dpm_solver_1_step/dpm_solver_2_step — DPM-Solver (Lu et al. 2022), eps-prediction exponential integrator, order 1 and a midpoint order-2 variant (dpm_solver_2_step takes your model, any DenoiseModel, and evaluates it a second time at the log-SNR midpoint).
• dpm_solver_pp_1_step/dpm_solver_pp_2_step — DPM-Solver++ (Lu et al. 2022b), the data(x0)-prediction reformulation. Order 1 is provably identical to ddim_sampler_step — verified numerically, not just by derivation.

DPM-Solver-v3 is not implemented (its distinguishing idea needs "empirical model statistics" calibrated from a real trained model over real data).

void edm_karras_sigmas(unsigned long numSteps, float sigmaMin, float sigmaMax, float rho,
                       &float outSigmas);                         // numSteps values
void ddpm_linear_schedule(unsigned long numSteps, float betaStart, float betaEnd,
                          &float outBeta, &float outAlphaBar);    // numSteps values each

&FloatTensor ddpm_sampler_step(const &FloatTensor x_t, const &FloatTensor epsPred,
                          float beta_t, float alphaBar_t, const &FloatTensor noise);
&FloatTensor ddim_sampler_step(const &FloatTensor x_t, const &FloatTensor epsPred,
                          float alphaBar_t, float alphaBar_prev);
&FloatTensor dpm_solver_1_step(const &FloatTensor x_cur, const &FloatTensor epsPred,
                          float alpha_cur, float sigma_cur, float alpha_next, float sigma_next);
&FloatTensor dpm_solver_pp_1_step(const &FloatTensor x_cur, const &FloatTensor x0Pred,
                             float alpha_cur, float sigma_cur, float alpha_next, float sigma_next);

// Order-2 solvers call your model again at the midpoint: any struct with
// a predict method returning the predicted noise (eps) or clean sample (x0).
trait DenoiseModel {
    &FloatTensor predict(const &FloatTensor x, float sigma);
}
generic<M: DenoiseModel>
&FloatTensor dpm_solver_2_step(&M model, const &FloatTensor x_cur, const &FloatTensor epsPred,
                          float alpha_cur, float sigma_cur, float alpha_next, float sigma_next);
generic<M: DenoiseModel>
&FloatTensor dpm_solver_pp_2_step(&M model, const &FloatTensor x_cur, const &FloatTensor x0Pred,
                             float alpha_cur, float sigma_cur, float alpha_next, float sigma_next);

The model is your own struct, holding its weights; the solver calls its predict method directly, with no callback or type-erased context.

Transformer (DiT / JiT) #

#include <std/ml/transformer.h>

• DiTBlock — Diffusion Transformer (Peebles & Xie 2023), adaLN- style: a conditioning vector produces 6 modulation vectors per block (shift/scale/gate for each of the attention and FFN sublayers). Verified via a gate=0 exact-identity check (with gates zeroed, the block's output must equal its input exactly) and full self-consistency.
• JiTBlock/JiT — a plain pre-LN transformer with conditioning added once, up front, to the token embeddings (as opposed to DiT's per-layer adaptive modulation) — the architectural distinction this implementation draws between the two.

Both are built from mha_forward, layer_norm, tensor_mul_row / tensor_add_bias (the adaLN modulation) and a ReLU feed-forward block. Their weights require grad, so they train through trainer_fit — on a GPU with TrainConfig.device, where the attention runs as the fused kernels:

struct DitModel { struct DiTBlock blk; &FloatTensor c; };   // c: [1, condDim]
static &FloatTensor dit_loss(&DitModel m, const &FloatTensor x, const &FloatTensor y) {
    return tensor_mse_loss(m.blk.forward(x, m.c), y);
}
struct DitModel m;
m.blk = DiTBlock::new(64UL, 256UL, 16UL, 4UL);
// ... m.c, then:
struct Optimizer opt = Optimizer::new(OPT_ADAM, (float)0.001);
model_add_params(&m.blk, &opt);
struct TrainConfig cfg = TrainConfig::default();
cfg.batchSize = 16UL; cfg.shuffle = 0; cfg.device = DEVICE_AUTO;   // a batch: one 16-token sequence
struct TrainReport rep = trainer_fit(&m, dit_loss, &opt, &tokens, &cfg);

static struct DiTBlock DiTBlock::new(unsigned long dModel, unsigned long dHidden, unsigned long dModelCond, unsigned long numHeads);
&FloatTensor DiTBlock::forward(const &FloatTensor x, const &FloatTensor c) const;   // c = conditioning
void    DiTBlock::free();

static struct JiTBlock JiTBlock::new(unsigned long dModel, unsigned long dHidden, unsigned long numHeads);
&FloatTensor JiTBlock::forward(const &FloatTensor x) const;
void    JiTBlock::free();
// The model: patch embedding + conditioning, then the blocks.
struct JiT { &FloatTensor WPatchEmbed; struct Vec<struct JiTBlock> blocks; };
static struct JiT JiT::new(unsigned long patchDim, unsigned long dModel, unsigned long dHidden, unsigned long numHeads, unsigned long depth);
&FloatTensor JiT::forward(const &FloatTensor patches, const &FloatTensor cond) const;   // cond: [1, dModel]
void    JiT::free();

Training: layers, losses, optimizers #

#include <std/ml/nn.h>        // Rng, more autograd ops, losses, Mlp, Optimizer
#include <std/ml/datasets.h>  // Dataset, generators, MNIST / CSV loaders, DataLoader
#include <std/ml/trainer.h>   // trainer_fit, trainer_evaluate, trainer_accuracy

These three headers train a model end to end on tensor.h's autograd. A batch is a [batch, features] tensor. Class labels are tensors of class indices stored as floats.

• Rng is a seeded splitmix64 generator with uniform, normal and bounded-integer draws. Every initialization, shuffle and sample in these modules goes through one, so runs are reproducible.

• Differentiable ops, in addition to tensor.h / tensor_nn.h:

  • Arithmetic: tensor_add_bias (row broadcast) and FloatTensor::add_scalar.
  • Elementwise: FloatTensor::exp, FloatTensor::log, FloatTensor::abs, FloatTensor::leaky_relu, FloatTensor::clamp and tensor_minimum.
  • Reductions and reshaping: FloatTensor::mean, FloatTensor::sum_rows ([m,n] → [m,1]), tensor_concat_cols and tensor_gather_rows (embedding or codebook lookup, with gradients scatter-added back).
  • FloatTensor::log_softmax.
  • FloatTensor::detach, which stops gradients.

• Losses, each returning one value averaged over elements:

  • tensor_mse_loss.
  • tensor_huber_loss.
  • tensor_bce_logits_loss: binary cross-entropy taken from logits, so it never computes log(0).
  • tensor_cross_entropy_loss: softmax cross-entropy against class indices.

  tensor_accuracy scores a batch's predictions.

• Mlp is a stack of fully connected layers: inDim → hidden × hiddenLayers → outDim, with ACT_RELU, ACT_LEAKY, ACT_TANH, ACT_SIGMOID, ACT_GELU or ACT_SILU between layers. It is initialized like PyTorch's nn.Linear, with weights stored [in, out]. model_copy / model_soft_update keep target networks in sync.

• Optimizer holds a list of parameters and offers these kinds:

  Kind	Update
  OPT_SGD	Plain gradient descent.
  OPT_MOMENTUM / OPT_NESTEROV	Heavy-ball or Nesterov momentum (momentum 0.9).
  OPT_RMSPROP	Running average of squared gradients (beta2 0.99).
  OPT_ADAGRAD	Learning rate divided by the root of the summed squared gradients (Duchi et al. 2011).
  OPT_ADAM	Adam with L2 weight decay (beta1 0.9, beta2 0.999).
  OPT_ADAMW	Adam with decoupled weight decay (weightDecay 0.01).
  OPT_LION	Sign of the interpolated momentum (Chen et al. 2023), with decoupled decay. Use a learning rate 3–10× smaller than Adam's.

  Optimizer::clip_grad_norm clips the global gradient norm. The fields (lr, betas, weightDecay) are public, so a schedule can change them between steps.

• Saving and loading. tensors_save / tensors_load write a list of tensors (shapes and float32 values) to a file and read it back into tensors of the same shapes. model_save_file / model_load_file do this for an Mlp. Loading returns 0, and changes nothing, when the file is missing or doesn't match.

trainer_fit is generic over the model type, so the forward pass and loss are resolved at compile time. It runs epochs of shuffled mini-batches: loss, backward, optional clipping, step, and a per-epoch lr decay. Each step's tensors come from a TensorArena that is reset after the step. Parameters and optimizer state always stay on the heap.

static &FloatTensor clf_loss(&Mlp m, const &FloatTensor x, const &FloatTensor y) {
    return tensor_cross_entropy_loss(m.forward(x), y);
}
static &FloatTensor clf_logits(&Mlp m, const &FloatTensor x) { return m.forward(x); }

struct Rng rng = Rng::new(1UL);
struct Dataset data = Dataset::moons(400UL, (float)0.1, &rng);
struct DatasetSplit sp = data.split((float)0.25, &rng);
struct Mlp net = Mlp::new(2UL, 32UL, 2UL, 2UL, ACT_RELU, ACT_NONE, &rng);
struct Optimizer opt = Optimizer::new(OPT_ADAM, (float)0.01);
model_add_params(&net, &opt);
struct TrainConfig cfg = TrainConfig::default();     // 10 epochs, batch 32
cfg.epochs = 60UL;
struct TrainReport rep = trainer_fit(&net, clf_loss, &opt, &sp.train, &cfg);
float acc = trainer_accuracy(&net, clf_logits, &sp.test, 64UL);   // > 0.95 in the tests

Datasets. A Dataset is x [n, features] plus y [n, targets], with classes set for classification. The synthetic generators are:

• Dataset::moons, Dataset::circles, Dataset::spirals, Dataset::blobs and Dataset::xor: 2D toy classification.
• Dataset::sine: regression.
• Dataset::gaussian_ring: the usual GAN mode-coverage test.
• Dataset::shapes: small binary images of four shape classes.

Dataset::load_mnist reads MNIST / Fashion-MNIST IDX files, with pixels scaled to [0, 1]. Dataset::load_csv reads numeric CSV, with an optional header and label column. Both return n == 0 on error. Dataset::split, Dataset::standardize and Dataset::standardize_pair (which applies the train set's statistics to both sets) prepare the data. A DataLoader hands out batches in a fresh shuffled order every epoch.

Every op's gradient is checked against central differences in the tests. The optimizers are checked to converge, and an MLP trained on moons reaches more than 95% accuracy.

The same calls train on a GPU: set TrainConfig.device, or place the model with model_to, and every op runs where its operands live (Training on a GPU).

Models as traits: generic functions for every model #

#include <std/ml/module.h>

Every model is a struct with methods. module.h names what a model can do as small traits, one capability each, and composes them into the roles generic code asks for. The functions after the table are written once and apply to every model that has the capability. Each call is resolved at compile time.

struct Mlp net = Mlp::new(2UL, 32UL, 2UL, 2UL, ACT_RELU, ACT_NONE, &rng);
struct Optimizer opt = Optimizer::new(OPT_ADAM, (float)0.01);
model_add_params(&net, &opt);
struct TrainReport rep = module_fit(&net, TRAIN_LOSS_CE, &opt, &train, &cfg);
float acc = module_accuracy(&net, &test, 64UL);
model_save_file(&net, "net.bin");

The same five lines train a ConvNet, a LoraMlp (its collect_params returns only the adapters, so only they train and only they are saved) or a model of your own once it has the methods. A model can also be assembled from others with a bounded generic struct that is itself a Module:

generic<M: Module>
struct Avg2 {
    M a;
    M b;
    &FloatTensor forward(const &FloatTensor x);
    void collect_params(&Vec<&FloatTensor> out) const;
    void collect_state(&Vec<&FloatTensor> out) const;
};
generic<M: Module>
&FloatTensor Avg2::forward(const &FloatTensor x) {
    return tensor_add(self.a.forward(x), self.b.forward(x)).scale((float)0.5);
}
generic<M: Module>
void Avg2::collect_params(&Vec<&FloatTensor> out) const { self.a.collect_params(out); self.b.collect_params(out); }
generic<M: Module>
void Avg2::collect_state(&Vec<&FloatTensor> out) const { self.a.collect_state(out); self.b.collect_state(out); }

struct Avg2<struct Mlp> pair;   // ... pair.a, pair.b = mlp_new(...)
module_fit(&pair, TRAIN_LOSS_CE, &opt, &train, &cfg);

trainer_fit remains for losses of your own: it takes the loss as a function of the model and a batch. tests/stdlib/tests/test_ml_modules.sc runs one generic check (forward, parameter count, save and load, copy) over ten Modules, checks module_fit against trainer_fit to the bit, and trains the Avg2 pair above.

Normalization: LayerNorm and AdaLN #

#include <std/ml/norm.h>

Trainable normalization layers, with autograd, over [batch, dim] rows. activations.h's FloatTensor::layernorm_rows is the forward-only version.

• LayerNorm (Ba et al. 2016) computes y = LN(x) · gamma + beta, with learned gamma (initialized to 1) and beta (initialized to 0).
• AdaLN is the adaptive LayerNorm of DiT (Peebles & Xie 2023). A linear map of a per-sample conditioning vector, such as a timestep or class embedding, gives (shift, scale[, gate]), and y = LN(x) · (1 + scale) + shift. The map starts at zero (adaLN-Zero), so a fresh AdaLN is a plain LayerNorm. With gated, AdaLNOut.gate also starts at 0, so a residual branch x + gate · f(y) starts as the identity.
• Building blocks: FloatTensor::layer_norm, tensor_mul_row (multiply by a vector on every row), and attention.h's FloatTensor::slice_cols, which is now differentiable.

struct AdaLN ad = AdaLN::new(dim, condDim, 1);   // gated
struct AdaLNOut o = ad.forward(x, cond);        // x [m, dim], cond [m, condDim]
&FloatTensor h = tensor_add(x, tensor_mul(o.gate, ffn.forward(o.out)));
model_add_params(&ad, &opt);

The tests check both layers' gradients against finite differences, and check that a fresh AdaLN equals LayerNorm. They also check that AdaLN learns a class-dependent affine map that a plain LayerNorm cannot represent, bringing its loss from above 0.5 to below 0.001.

Fine-tuning: LoRA and QLoRA #

#include <std/ml/lora.h>

LoRA (Hu et al. 2021) freezes a pretrained Mlp and learns a low-rank update for each layer:

y = x·W + b + (alpha / r) · (dropout(x)·A)·B

A [in, r] starts random and B [r, out] starts at zero, so the adapted network begins identical to the base. Only A and B train, which is r · (in + out) values per layer instead of in · out.

• LoraMlp::attach sets the base's requiresGrad to 0, and LoraMlp::detach restores it.
• LoraMlp::merge folds scale · A·B into W, so the base alone then computes the adapted output at no extra inference cost.

QLoRA (Dettmers et al. 2023) does the same on a base stored as 4-bit NormalFloat (NF4). Each block of 64 weights keeps one absmax scale and a 4-bit index into the 16 NF4 levels, which are quantiles of a normal distribution, so Gaussian-like weights use the levels evenly. Double quantization also stores the scales in 8 bits, with one float per 256 scales, bringing the base to about 4.1 bits per weight.

QLoraMlp::attach copies the base, so the original Mlp can be freed. Each forward pass dequantizes the weights to float32, and gradients flow through the frozen 4-bit base to the adapters. Nf4FloatTensor::quantize / Nf4FloatTensor::dequantize are also usable on their own.

struct LoraConfig cfg = LoraConfig::default(8UL);    // rank 8, alpha 8
struct QLoraMlp ft = QLoraMlp::attach(&pretrained, &cfg, NF4_BLOCK, 1);   // NF4 + double quant
pretrained.free();
struct Optimizer opt = Optimizer::new(OPT_ADAMW, (float)0.001);
model_add_params(&ft, &opt);                   // adapters only
trainer_fit(&ft, my_loss, &opt, &data, &tc);   // my_loss calls QLoraMlp::forward

In the tests, an MLP pretrained on moons is fine-tuned to a rotated, label-swapped task that the frozen base gets wrong (accuracy below 70%):

• LoRA reaches more than 95% accuracy. The base weights stay bit-identical, and the merged base reproduces the adapter's output to 1e-3.
• NF4 reproduces values on its grid exactly. On Gaussian weights the relative RMS error is below 12%, at 4.5 bits per weight, or under 4.2 bits with double quantization at almost the same error.
• QLoRA works on a base stored in under 1/6 of its float32 size. The 4-bit base still scores above 95% on the original task, and the adapters reach more than 95% on the new one.

Convolutional networks #

#include <std/ml/convnet.h>

cnn.h runs convolutions forward on FeatureMaps. convnet.h makes them differentiable on FloatTensors, so they train with FloatTensor::backward, the optimizers and trainer_fit like Mlp does. Images are rows, as in every other model here: a batch is a [batch, C*H*W] tensor, and each row holds one image in CHW order. A 2D tensor doesn't carry the image geometry, so every op takes C, H, W as arguments. Each op records gradients for its input and its trained parameters.

Convolution and pooling #

tensor_conv2d_ex and tensor_conv_transpose2d_ex take a Conv2dSpec with a separate kernel size, stride, padding and dilation for each axis, plus groups and, for the transposed convolution, outPadH / outPadW. Conv2dSpec::new(kernel, stride, padding) fills in a square kernel with dilation 1, one group and no output padding; change the fields you need. tensor_conv2d / tensor_conv_transpose2d are the common case with one stride and padding for both axes.

groups == inC gives a depthwise convolution. Output padding must be below the stride or the dilation, as in PyTorch; it picks among the output sizes that the same stride maps back to one input size.

Pooling takes a Pool2dSpec, built with pool2d_spec(kernel, stride, padding): a separate window, stride and padding per axis, ceilMode, and, for avg pooling, countIncludePad.

• FloatTensor::maxpool2d_ex / FloatTensor::avgpool2d_ex clip their windows to the plane. Max pooling therefore ignores the padding, as PyTorch's MaxPool2d does, and sends each window's gradient to its first largest input.
• With countIncludePad, avg pooling divides by the window's extent inside the padded plane, as PyTorch's AvgPool2d does by default. Without it, avg pooling divides by the positions inside the plane.
• ceilMode rounds the output side up. As in PyTorch, the last window must still start inside the input or its left padding.
• FloatTensor::maxpool2d / FloatTensor::avgpool2d take a square window. FloatTensor::global_avgpool2d / FloatTensor::global_maxpool2d reduce each channel to one value, [batch, C*H*W] → [batch, C].

conv_out_side_ex, conv_transpose_out_side_ex and pool_out_side compute the sizes.

Both convolutions use im2col + sgemm, per group. A 3×3, stride-1 layer with at least 8 input and output channels runs its forward pass through cnn.h's Winograd F(2×2, 3×3) kernel, and a 1×1 layer through its direct path. A large enough batch is split across the thread pool, one chunk of images per thread. Weight and bias gradients are summed per chunk and then added up in chunk order, so the result doesn't depend on scheduling. Pooling, resizing and the norms split large inputs by plane, row or channel. Those units never share an output, so the threaded result is bit-for-bit the single-threaded one.

Resizing #

x.resize2d(C, H, W, outH, outW, mode, alignCorners) resizes to any size, up or down, as PyTorch's interpolate(size=...):

With alignCorners, the corner pixels' centres map onto each other. Without it, pixel areas line up. Edges are clamped, and gradients scatter back with the same weights. FloatTensor::upsample_nearest is a fast path for an integer factor. FloatTensor::upsample_bilinear resizes by an integer factor without aligning corners.

Normalization and dropout #

• Batch norm (tensor_batch_norm; Ioffe & Szegedy 2015) normalizes each channel over the batch and the plane. Pass plane = H*W for feature maps, or 1 for [batch, features] rows. In training it uses the batch's statistics and moves the running ones toward them by momentum. The running variance uses the unbiased estimate, as in PyTorch. In inference it uses the running statistics.
• Group norm (tensor_group_norm; Wu & He 2018) normalizes each image's channel groups over their channels and planes. It doesn't depend on the batch, so it behaves the same in training and inference. One group is layer norm over (C, H, W); one group per channel is instance norm.
• Spectral norm (tensor_spectral_norm; Miyato et al. 2018) divides a weight by its largest singular value. It estimates that value by power iteration from a vector u kept between calls. The gradient treats the power-iteration vectors as constants, as PyTorch does.
• Dropout: FloatTensor::dropout zeroes each value with probability p and scales the rest by 1/(1 − p) in training, and passes x unchanged in inference. FloatTensor::dropout2d drops whole channel planes.

Conditioning. FloatTensor::label_planes turns class indices into one constant plane per class: 1 in the plane of the row's label and 0 elsewhere. tensor_append_label_channels adds those planes after an image's channels. This is how the conditional ConvVae and Dcgan pass the label to their image-side network.

ConvNet #

ConvNet is the convolutional counterpart of Mlp. Start it with the input geometry and append layers; each layer's input is the previous layer's output.

Place a norm between a convolution with ACT_NONE and the activation. Weights and biases are initialized like PyTorch's, from U(−1/√fanIn, 1/√fanIn), and the norms start at gamma 1 and beta 0. params holds the trained tensors. buffers holds what the optimizer doesn't train: running statistics and spectral norms' u. Both stay on the heap.

The training field is 1 from ConvNet::new. In training mode, batch norm uses batch statistics and updates the running ones, spectral norms iterate, and dropout drops. Call n.set_training(0) for inference; it returns the previous mode. The models in autoencoder.h and gan.h switch modes themselves: their loss and training-step functions run in training mode, and encode, decode, sample and score run in inference mode.

Other operations on a ConvNet:

• ConvNet::forward maps [batch, in_dim()] to [batch, out_dim()]. The outC, outH and outW fields give the current output geometry.
• model_add_params registers the parameters (not the buffers).
• ConvNet::clone, model_copy and model_soft_update handle target and averaged networks, buffers included.
• ConvNet::fold_batchnorm folds each batch norm that directly follows a convolution or dense layer into that layer's weight and bias, for inference. The output is unchanged, but don't train the net afterwards.
• model_save_file / model_load_file write the parameters and buffers to a file and read them back into a net built with the same layers. Loading returns 0 on a mismatch, and nothing is overwritten. tensors_save / tensors_load and an Mlp's files in nn.h use the same format.

static &FloatTensor cnn_loss(&ConvNet m, const &FloatTensor x, const &FloatTensor y) {
    return tensor_cross_entropy_loss(m.forward(x), y);
}

struct Rng rng = Rng::new(1UL);
struct ConvNet net = ConvNet::new(1UL, 28UL, 28UL);   // MNIST
net.conv(16UL, 3UL, 1UL, 1UL, ACT_NONE, &rng);       // 16 x 28 x 28
net.batchnorm(ACT_RELU);
net.residual_block(16UL, 1UL, 1, ACT_RELU, &rng);   // 16 x 28 x 28
net.residual_block(32UL, 2UL, 1, ACT_RELU, &rng);   // 32 x 14 x 14
net.residual_block(64UL, 2UL, 1, ACT_RELU, &rng);   // 64 x 7 x 7
net.dropout((float)0.1, 1, &rng);
net.global_avgpool();              // 64 x 1 x 1
net.dense(10UL, ACT_NONE, &rng);   // logits
struct Optimizer opt = Optimizer::new(OPT_ADAM, (float)0.001);
model_add_params(&net, &opt);
struct TrainConfig cfg = TrainConfig::default();
trainer_fit(&net, cnn_loss, &opt, &mnist, &cfg);
net.set_training(0);   // inference from here on
net.fold_batchnorm();
model_save_file(&net, "mnist_resnet.bin");

The tests check every op's input and parameter gradients against central differences. The cases cover:

• conv2d with stride 1 and 2, a dilated grouped convolution, a depthwise convolution, and a 2×3 kernel with different strides and padding per axis;
• transposed convolution with stride 1 and 2, grouped with output padding, and dilated;
• max pooling with disjoint, overlapping, padded, and 2×3 ceil-mode windows;
• avg pooling, with and without padding counted, and in ceil mode;
• global average and max pooling;
• nearest upsampling, bilinear upsampling by 2 and 3, aligned-corner bilinear to 5×4, bicubic to 7×5, and nearest downsampling;
• batch norm in training and inference mode, layer norm, and group norm;
• spectral norm;
• dropout and channel dropout.

Further checks against reference values:

• Sizes: the output sides of padded and ceil-mode pooling and of dilated convolutions match PyTorch's.
• cnn.h parity: tensor_conv2d, the padded pooling and global pooling match Conv2D::forward, *pool2d_forward_ex and global_*pool2d. The Winograd path matches a direct sum.
• Resizing: bilinear upsampling reproduces PyTorch's output for a 2×2 input, and bicubic resizing keeps a constant image constant.

ConvNet checks:

• An input-gradient check runs through a projection skip, a concat skip, group norm and dense layers.
• The residual blocks' parameter counts are checked.
• Batch-norm folding leaves inference output unchanged within 1e-4.
• A saved net loads into a differently initialized one and gives the same output exactly; a mismatched net refuses to load.

Other checks:

• Batch norm's running statistics follow the batch's.
• Dropout drops the right fraction, and channel dropout drops whole planes.
• Spectral norm leaves a largest singular value of 1.
• Group norm doesn't depend on the rest of the batch.
• Threaded convolutions match one thread within 1e-4, and pooling, upsampling and batch norm match bit for bit.
• End to end, a small residual CNN with batch norm and channel dropout classifies held-out 12×12 shapes with more than 90% accuracy.

Autoencoders #

#include <std/ml/autoencoder.h>

MLP encoders and decoders over flattened rows. Each *_loss function has trainer_fit's loss signature. Reconstruction is RECON_MSE for real values, or RECON_BCE for data in [0, 1]. With BCE the decoder outputs logits, and *_decode / *_reconstruct apply the sigmoid.

struct VaeConfig cfg = VaeConfig::default(4UL);   // 256 hidden, BCE, beta 1
cfg.classes = 4UL;                               // conditional
struct Vae v = Vae::new(100UL, &cfg);
struct Optimizer opt = Optimizer::new(OPT_ADAM, (float)0.002);
model_add_params(&v, &opt);
model_fit(&v, &opt, &shapes, &tc);   // the VAE's own loss (Vae::loss)
&FloatTensor samples = v.sample(64UL, (?&FloatTensor)labels);   // [64, 100] in (0, 1)

The tests check each variant's reconstruction against its untrained error. They also check that a classifier trained on the real shapes recognizes the conditional VAE's samples as their requested class (accuracy above 70%), and that the VQ-VAE keeps more than 4 codes in use.

Convolutional autoencoders #

ConvAutoencoder and ConvVae are the plain AE and the VAE with ConvNet encoders and decoders that you build. The encoder maps the image to features. The decoder takes the code as a latent × 1 × 1 input and must output the image's size; start it with ConvNet::dense + ConvNet::reshape, or with a transposed convolution from 1×1. In ConvVae, linear heads on the flattened features produce μ and log σ². Its loss is Vae's, and it reads latent, recon, beta, classes and seed from a VaeConfig.

With classes > 0, ConvVae is a conditional VAE. The encoder gets the label as classes extra input channels (one constant plane per class), and the decoder gets the one-hot label appended to the code. So for C-channel images the encoder takes C + classes channels, and the decoder takes latent + classes inputs and outputs C channels. As with Vae, y carries the class indices, and ConvVae::sample takes labels or null.

Both models take ownership of the nets, so free the model and not the nets. Loss functions run the nets in training mode. Encode, decode, reconstruct and sample run them in inference mode, so batch norm uses its running statistics there.

unsigned long classes = 10UL;                                // conditional, MNIST digits
struct ConvNet enc = ConvNet::new(1UL + classes, 28UL, 28UL);
enc.conv(32UL, 4UL, 2UL, 1UL, ACT_LEAKY, &rng);   // 32 x 14 x 14
enc.conv(64UL, 4UL, 2UL, 1UL, ACT_NONE, &rng);    // 64 x 7 x 7
enc.batchnorm(ACT_LEAKY);
struct ConvNet dec = ConvNet::new(16UL + classes, 1UL, 1UL);  // latent 16 + label
dec.dense(64UL * 49UL, ACT_RELU, &rng);
dec.reshape(64UL, 7UL, 7UL);
dec.deconv(32UL, 4UL, 2UL, 1UL, ACT_NONE, &rng);   // 32 x 14 x 14
dec.batchnorm(ACT_RELU);
dec.deconv(1UL, 4UL, 2UL, 1UL, ACT_NONE, &rng);   // 1 x 28 x 28 logits
struct VaeConfig cfg = VaeConfig::default(16UL);
cfg.classes = classes;
struct ConvVae v = ConvVae::new(enc, dec, &cfg);
struct Optimizer opt = Optimizer::new(OPT_ADAM, (float)0.001);
model_add_params(&v, &opt);
model_fit(&v, &opt, &mnist, &tc);
&FloatTensor samples = v.sample(64UL, (?&FloatTensor)labels);   // [64, 784] in (0, 1)

The tests run on 8×8 shapes, with batch norm in both nets:

• A conv AE must cut its reconstruction error to below 30% of the untrained error.
• A conv VAE must cut its loss to below 60% of the first epoch's, keep KL above 0.1, and reconstruct with MSE below 0.1.
• The conditional conv VAE's samples must be recognized as their requested class more than 70% of the time, by a classifier trained on the real images.

GANs #

#include <std/ml/gan.h>

An MLP generator and discriminator (leaky ReLU), trained in alternation. The objective is set in the config:

The penalty's gradient with respect to the critic's parameters is a mixed second derivative. Rather than differentiating the backward pass, it is computed from a central difference of the critic along each row's normalized input gradient. That costs two extra critic passes and is exact to O(ε²), with ε = 0.01; the tests check it against finite differences of the penalty itself. Because the penalty treats rows independently, the critic must not use batch norm (group norm or layer norm are fine).

classes > 0 makes it a conditional GAN: the label goes to both networks. emaDecay keeps an exponential moving average of the generator's weights (Karras et al. 2018), and Gan::sample samples from it. The game oscillates, and the average sits nearer its centre than any single step does. The GAN owns its two optimizers (Adam with β1 = 0.5, or RMSProp for WGAN) and the arena for its steps. Gan::train runs epochs over a Dataset, and Gan::train_step runs one step on your own batch.

struct GanConfig cfg = GanConfig::default(4UL);   // latent 4
cfg.loss = GAN_HINGE;
struct Gan g = Gan::new(2UL, &cfg);
g.train(&data, 80UL, 64UL);
&FloatTensor fake = g.sample(1000UL, null);   // heap tensor, caller frees

In the tests, every objective learns a 2D Gaussian, with the mean within 0.25. The WGAN-GP critic's input gradient ends up with a norm between 0.6 and 1.4. The spread is bounded loosely, because GANs tend to under-cover it. On a 4-mode ring, a conditional GAN puts more than 75% of its samples on the mode of their label.

DCGAN #

Dcgan plays the same game with a ConvNet generator and discriminator (Radford et al. 2016). G takes the latent batch as a latent × 1 × 1 input and must output the data's size. D takes a data row and must output one score. It reads the objective, learning rates, nCritic, clip, classes, emaDecay and seed from a GanConfig.

With classes > 0, it is a conditional DCGAN. G gets the one-hot label appended to its noise, so it takes latent + classes inputs. D gets the label as classes extra input channels, one constant plane per class. Dcgan::train_step, Dcgan::sample and Dcgan::score take labels or null, and Dcgan::train takes them from the dataset.

The DCGAN owns the nets, the optimizers, the averaged generator and the step arena. Training steps run the nets in training mode, so batch norm uses batch statistics, as DCGAN prescribes. Dcgan::sample and Dcgan::score run in inference mode.

struct ConvNet gen = ConvNet::new(16UL, 1UL, 1UL);   // latent 16
gen.deconv(64UL, 4UL, 1UL, 0UL, ACT_NONE, &rng);    // 64 x 4 x 4
gen.batchnorm(ACT_RELU);
gen.deconv(32UL, 4UL, 2UL, 1UL, ACT_NONE, &rng);   // 32 x 8 x 8
gen.batchnorm(ACT_RELU);
gen.deconv(1UL, 4UL, 2UL, 1UL, ACT_SIGMOID, &rng);   // 1 x 16 x 16
struct ConvNet disc = ConvNet::new(1UL, 16UL, 16UL);
disc.conv(32UL, 4UL, 2UL, 1UL, ACT_LEAKY, &rng);   // 32 x 8 x 8
disc.conv(64UL, 4UL, 2UL, 1UL, ACT_NONE, &rng);    // 64 x 4 x 4
disc.batchnorm(ACT_LEAKY);
disc.conv(1UL, 4UL, 1UL, 0UL, ACT_NONE, &rng);   // one score
struct GanConfig cfg = GanConfig::default(16UL);
struct Dcgan g = Dcgan::new(gen, disc, &cfg);
g.train(&images, 20UL, 64UL);
&FloatTensor fake = g.sample(64UL, null);   // [64, 256], heap tensor, caller frees

The tests train on 8×8 shapes. A DCGAN must move its samples' mean intensity from the untrained sigmoid's 0.5 to within 0.1 of the data's. A conditional DCGAN's samples must be recognized as their requested class more than 55% of the time, by a classifier trained on the real images (chance is 25%).

Reinforcement learning #

#include <std/ml/rl.h>

Environments. Bandit is a k-armed Gaussian bandit. GridWorld has walls, pits, goals, slip probability and truncation; GridWorld::cliff() is Sutton & Barto's cliff walk. CartPole uses Gym's CartPole-v1 dynamics, its 12° / 2.4 m limits and 500-step truncation. Steps return an EnvStep that separates terminated (no future value) from truncated (time limit, so learners bootstrap). Observations are tensors.

struct Dqn d = Dqn::new(CARTPOLE_OBS, CARTPOLE_ACTIONS, &cfg);
struct CartPole env = CartPole::new(1UL);
env.reset(obs);
while (1) {
    unsigned long a = d.act(obs, 1);   // epsilon-greedy
    struct EnvStep s = env.step(a, next);
    d.observe(obs, a, s.reward, next, s.terminated);   // store + train
    // obs = next ...
    if (s.terminated || s.truncated) break;
}

The tests run each technique on its problem:

• The four bandit agents pick the best of 10 arms more than 85% of the time.
• On the cliff walk, Q-learning learns the optimal 13-step path along the edge, and SARSA a safer, longer one (Sutton & Barto's Example 6.6).
• DQN, REINFORCE, A2C and PPO each reach CartPole's "solved" mean greedy return of 195 within their step budget.

Training on a GPU #

#include <std/ml/device.h>   // devices, precision; tensor.h, trainer.h and module.h include it

Every tensor has a device: DEVICE_CPU, or a GPU backend (DEVICE_ROCM for AMD, DEVICE_CUDA for NVIDIA, DEVICE_MPS for Apple Silicon, and the cross-vendor DEVICE_VULKAN — any Vulkan 1.2 GPU, Apple's through MoltenVK — and DEVICE_WEBGPU). An op runs on the device of its operands. If any operand lives on a GPU, the op runs there, uploading the others, and its result lives there too. A model whose parameters are on a GPU therefore trains there through the same code that trains it on the CPU: every layer, loss, optimizer, trainer_fit, module_fit / model_fit, the autoencoders, GANs, RL agents and LoRA adapters above. The backward pass runs each node where its forward pass ran.

struct Mlp net = Mlp::new(784UL, 256UL, 2UL, 10UL, ACT_RELU, ACT_NONE, &rng);
struct Optimizer opt = Optimizer::new(OPT_ADAM, (float)0.001);
model_add_params(&net, &opt);
struct TrainConfig cfg = TrainConfig::default();
cfg.device = DEVICE_AUTO;                  // the first usable GPU, else the CPU
cfg.precision = DEVICE_FP16;               // optional: 16-bit activations
struct TrainReport rep = trainer_fit(&net, clf_loss, &opt, &train, &cfg);
model_sync(&net);                          // host copies current: save, inspect, evaluate on the CPU

Choosing the device.

The build's features link the backends in, and safeguard adds each one's libraries to the link (the Link column; by hand, pass them yourself). A program built without any runs everything on the CPU and links no GPU library:

device_compiled(d), device_available(d), device_name(d), device_sync(d) and device_stats(d) (launches, GEMMs, transfers and bytes since device_stats_reset) answer what a program runs on. Device memory comes from a cache of size classes that steps reuse; device_bytes_reserved(d) and device_cache_clear(d) manage it.

Host and device copies. A tensor keeps its host values and a device copy coherent lazily: an op on the device uploads a host-written tensor once, and reading values on the host (at1, at2, a CPU-only op) downloads them once. t.sync() / model_sync(&m) make the host copies current explicitly, for code that reads t->data directly. Gradients work the same way: zero_grad() on a device marks the gradient zero without touching it, and the first backward op that reaches it writes instead of adding.

What runs on the device. Every elementwise op and activation, matmul, bias and row ops, sums and means, concat / slice / gather, transpose, softmax (with attention's causal / sliding-window masks) and log-softmax, multi-head attention as one fused launch for every head (and its backward), the fused losses (cross-entropy, BCE on logits, Huber), layer norm, batch norm and group norm (each fused with a following ReLU or leaky ReLU), max / average / global pooling, nearest upsampling, bilinear and bicubic resizing, convolutions and transposed convolutions, sparse products, every optimizer (one launch for all parameters) and gradient-norm clipping. Small GEMMs with a long reduction run split-K. An op without a device kernel computes on the CPU, copying its operands home first, so every op works on every device; the GPU only changes the speed.

Traits. DenseTensor and SparseTensor conform to trait DeviceTensor : Tensor (device(), place(d), sync()), so generic code can take either. RocmTensor, CudaTensor, MpsTensor, VulkanTensor and WebGpuTensor each inherit DeviceTensor and add device_buffer(), the values' buffer on that backend (a hipDeviceptr, a CUdeviceptr, an MTLBuffer's GPU address, a Vulkan buffer device address, an offset into the WebGPU backend's heap), for code that calls a vendor library itself.

generic<T: DeviceTensor>
int to_host(&T t) { return t.sync(); }   // a DenseTensor or a SparseTensor

Backends. A backend is a struct of static functions (alloc, release, upload, download, copy, sync, gemm, gemm_batched, conv, launch, offset, storage16), registered with device_register_backend_(device, &tag); every call resolves at compile time, and the runtime only picks the device from the tensors. Every backend runs one kernel set, device_kernels.cu: ROCm, CUDA and Metal compile it at run time (hiprtc, NVRTC, Metal), Vulkan runs it as SPIR-V compiled ahead of time, and WebGPU runs a WGSL port of it that a script checks against the source.

Precision: fp16 and bf16 on the device #

device_set_precision(d, DEVICE_FP16) (or DEVICE_BF16), or TrainConfig.precision for one fit, trains in 16 bits on d:

• Storage. On every GPU backend, activations and their gradients are stored in 16 bits. Parameters, their gradients and the optimizer state stay fp32. A buffer's element type travels with its device handle, so every kernel, copy, GEMM and convolution reads it from there; kernels compute in fp32 and round once per stored value. Each kernel source compiles twice, a plain fp32 build and one for 16-bit operands, and a launch picks the plain build when no operand is 16-bit, so fp32 training pays nothing.
• Compute. Matmuls run as 16-bit GEMMs with fp32 accumulation (rocBLAS / cuBLAS gemm_ex), convolutions in MIOpen's / cuDNN's 16-bit kernels; the Metal, Vulkan and WebGPU backends use their own GEMM and convolution kernels. The first time each convolution shape and direction runs, the 16-bit and fp32 versions are timed and the faster kept for that shape (device_set_precision_pick(DEVICE_PICK_ALWAYS) runs every one in 16 bits). A parameter's 16-bit copy is made once per optimizer step and reused.
• fp16 loss scale. fp16 adds a dynamic loss scale on the device, like PyTorch's GradScaler: it starts at 65536 and multiplies the loss's gradient, and every leaf gradient is unscaled when backward() ends, so user code never sees a scaled gradient. A step whose gradients overflowed raises a flag the optimizer reads on the device, skipping the update; the scale halves, and 2000 clean steps in a row double it. device_loss_scale(d) reads it. bf16 shares fp32's exponent range and needs no scaling.

On an MI100, fp16 is the faster 16-bit type: MIOpen's fp16 kernels use the matrix cores, while its bf16 convolutions run at about fp32 speed or slower (GPU training benchmarks).

Sparse tensors #

#include <std/ml/sparse.h>

A SparseTensor is a rows × cols matrix in CSR form that stores only its nonzeros. Its values are an ordinary FloatTensor of length nnz, so a sparse matrix can be a trained parameter (a sparse layer, a graph's learned edge weights). Gradients reach the values; the pattern is fixed. The index arrays are float tensors too (indices exact up to 2^24), so place(d) moves the whole matrix to a device and its products run there.

struct SparseTensor A = SparseTensor::from_dense(adj, (float)0.0, 1);   // learned edge weights
&FloatTensor H = tensor_spmm(&A, X);           // [rows, n] = A [rows, cols] . X [cols, n]
// ... loss ... loss.backward();               // A.values->grad, X->grad

SparseTensor conforms to Tensor and DeviceTensor, so generic code written against those traits takes it next to a dense tensor.

GPU backends #

#include <std/ml/gpu_mps.h>   // Apple Silicon/macOS — [features] "mps"
#include <std/ml/gpu_cuda.h>  // NVIDIA — [features] "cuda"
#include <std/ml/gpu_rocm.h>  // AMD — [features] "rocm"
#include <std/ml/gpu_spirv.h> // Vulkan (SPIR-V), cross-vendor — [features] "vulkan"
#include <std/ml/gpu_webgpu.h>// WebGPU

The "unified memory" piece of the ML brief: MPS's default buffer storage mode is genuine CPU+GPU shared memory on Apple Silicon (MLX's own selling point) — no explicit host↔device copy the way CUDA/ROCm's discrete memory model needs.

ROCm / HIP (AMD) #

#include <std/ml/gpu_rocm.h>   // links -lamdhip64 -lhiprtc -lrocblas (safeguard adds them)

Data lives in RocmBufs — GPU buffers described by a device address and an element count that every op checks — and ops are queued in order on one HIP stream, so a whole training step runs on the GPU and only the values you ask for (rocm_read) come back. Matmuls go to rocBLAS. Everything else runs this library's own HIP kernels — elementwise math, ReLU and its gradient, in-place SGD, sum reduction, and a fused MSE gradient (dy = 2(y − t) plus the summed loss in one pass) — compiled with hiprtc at first use for the GPU that is present, so one build runs on any ROCm GPU. On a machine with several GPUs the backend picks the one with the most compute units.

One SGD step of a two-layer MLP, entirely on the GPU:

// forward: Z1 = X W1, H = relu(Z1), Y = H W2
std::rocm_gemm(X, 0, W1, 0, Z1, B, H_DIM, I);
std::rocm_relu(Z1, H);
std::rocm_gemm(H, 0, W2, 0, Y, B, O, H_DIM);
// backward for loss = sum((Y - T)^2)
std::rocm_mse_grad(Y, T, dY, loss);
std::rocm_gemm(H, 1, dY, 0, dW2, H_DIM, O, B);     // dW2 = H^T dY
std::rocm_gemm(dY, 0, W2, 1, dH, B, H_DIM, O);     // dH  = dY W2^T
std::rocm_relu_grad(Z1, dH, dZ1);
std::rocm_gemm(X, 1, dZ1, 0, dW1, I, H_DIM, B);    // dW1 = X^T dZ1
std::rocm_sgd(W2, dW2, lr);
std::rocm_sgd(W1, dW1, lr);

Record such a sequence once between rocm_capture_begin() and rocm_capture_end() and replay it with rocm_graph_launch() to launch the whole step at once; this pays off when the host doesn't wait on the GPU every step.

GEMMs (fp32 and 16-bit) go through rocBLAS gemm_ex. The default solution rocBLAS picks for a shape can be several times slower than its best: on an MI100, 3.5× for a 128×256×1024 fp16 product and 2–3× for small bf16 ones. So the first call for each (types, transposes, size) times every solution rocBLAS offers on that call's buffers, like MIOpen's Find, and caches the fastest. That takes about 0.1 s per shape. Run each shape once before rocm_capture_begin(), since an untuned shape inside a capture uses the default solution. Because the choice comes from timing, two runs can pick different solutions and differ in the last bits. Autograd tensors use the GPU through their device, not through these host-array calls.

On rocBLAS versus a hand-written kernel: this library's own tiled SGEMM (rocm_matmul_f32) reaches about 9.5 TFLOPS at 4096³ on an MI100, while rocBLAS reaches 33.7 TFLOPS using the matrix cores, and rocBLAS is also faster at the small sizes a training loop uses (≈5 µs, close to the launch floor). So matmuls use rocBLAS, and the custom kernels cover the work rocBLAS doesn't do.

Attention #

Both attention ops take q as [heads][seqLen][headDim] and k, v as [heads][kvLen][headDim], all fp32, and write out as [heads][seqLen][headDim]. headDim can be at most 128. A scale of 0 or less means 1/√headDim. Each op returns 1 once the work is queued, or 0 when the shape is rejected or no ROCm GPU is present.

• rocm_flash_attention computes softmax(scale · QKᵀ) V exactly. It streams K/V tiles through shared memory with an online softmax (FlashAttention), so memory stays O(seqLen · headDim) and no seqLen × kvLen score matrix is ever stored. With causal, query i sees keys j ≤ i + kvLen − seqLen (the last query rows line up with the last keys, as in a KV cache), so kvLen ≥ seqLen is required. Fully masked tiles are skipped, which makes causal attention about twice as fast.
• rocm_linear_attention computes φ(Q)(φ(K)ᵀV) / φ(Q)(φ(K)ᵀ1) with φ(x) = elu(x) + 1 (Katharopoulos et al.). Time and memory are linear in sequence length. With causal, each query sees only the keys up to its own position (kvLen == seqLen); this uses per-chunk running states and a masked term inside each chunk. It approximates softmax attention, so its results differ from flash attention. It matches std::linear_attention_forward on the CPU.

On GPUs with matrix cores (CDNA: MI100, MI200, MI300) both ops run on them in fp32 (MFMA). Other AMD GPUs use CUDA-core kernels that give the same results up to rounding. rocm_attention_use_mfma(0) selects the CUDA-core kernels on a CDNA GPU too, for comparison; rocm_attention_use_mfma(1) restores the default, and both return whether the matrix-core kernels are in use.

MI100, batch 4 × 16 heads, fp32, ms per call, against PyTorch 2.9.1 on ROCm on the same GPU (full results):

CUDA (NVIDIA) #

#include <std/ml/gpu_cuda.h>   // links -lcuda -lcublas -lcublasLt -lnvrtc (safeguard adds them)

The same model as the ROCm backend, with the same function names under cuda_: CudaBufs are GPU buffers carrying their element count, ops queue in order on one CUDA stream, and only what you cuda_read comes back. Matmuls go to cuBLAS in full fp32 (no TF32). The rest runs this library's own CUDA kernels, compiled by NVRTC at first use into a cubin for the GPU that is present (-arch=sm_90 on an H100/H200), so one build runs on any GPU the installed CUDA knows. The backend uses the device's primary context, the one the CUDA runtime and PyTorch use, so buffers can coexist with other CUDA code in the process.

On top of the ROCm op set, two fusions cut the launches per training step from 10 to 7:

• cuda_gemm_relu — relu(op(a) op(b)) as one cuBLASLt matmul with the ReLU epilogue. The ReLU gradient can then test the activation instead of the pre-activation (h > 0 exactly where z > 0), so Z1 is never stored.
• cuda_sgd2 — both weight updates of a two-tensor model in one kernel.

Other CUDA-specific choices: reads of up to 64 KB (a loss, one output) go through a pinned staging buffer, and so take one DMA. Copies of 1 MiB and up in either direction run as four parallel lanes (the caller plus three helper threads started on first use). Each lane stages its slice through pinned memory on its own stream, so the DMA of one chunk overlaps the copy of the next. On the H200 host that moves 16 MB in 0.45–0.7 ms, against 1.05 ms for the driver's own pageable-memory path, which copies on one thread. A one-block reduction stores its result directly, so cuda_mse_grad/cuda_sum need no zeroing memset. cuBLAS gets a fixed 32 MB workspace, so its calls can be captured into graphs.

The MLP step above, fused:

std::cuda_gemm_relu(X, 0, W1, 0, H, B, H_DIM, I);    // H = relu(X W1)
std::cuda_gemm(H, 0, W2, 0, Y, B, O, H_DIM);
std::cuda_mse_grad(Y, T, dY, loss);
std::cuda_gemm(H, 1, dY, 0, dW2, H_DIM, O, B);
std::cuda_gemm(dY, 0, W2, 1, dH, B, H_DIM, O);
std::cuda_relu_grad(H, dH, dZ1);                      // H > 0 where Z1 > 0
std::cuda_gemm(X, 1, dZ1, 0, dW1, I, H_DIM, B);
std::cuda_sgd2(W2, dW2, W1, dW1, lr);

Before this backend was verified, each host-array call created and destroyed a CUDA context and a cuBLAS handle. Through tensor_matmul_cuda_blas, the small MLP benchmark took 83 s to train and 816 s for inference on the H200. The host-array functions now share the backend's context, stream and kernels.

Convolution (cuDNN) and FFT (cuFFT) on CUDA #

#include <std/ml/gpu_cudnn.h>   // link: -lcudnn + the gpu_cuda libraries
#include <std/dsp/fft_cuda.h>   // link: -lcufft + the gpu_cuda libraries

Two libraries on top of gpu_cuda.h. They queue their work on the same stream, so they order with the CudaBuf ops and can be recorded into a CudaGraph. Link them only when you use them: each is its own object file in the std archive.

gpu_cudnn.h runs CNN and U-Net on the GPU. Tensors are CudaBufs in NCHW order, which is FeatureMap's CHW layout with a batch axis in front. The first convolution of each shape (plain or with ReLU) builds the fused graph conv → + bias (→ ReLU) with cuDNN's graph API. It times the engine configs cuDNN's heuristics suggest, and the legacy API's best algorithm (cudnnFindConvolutionForwardAlgorithm), on that call's buffers, then caches the fastest one with its own workspace. Make that first call outside graph capture. Later calls cost a few microseconds of host time. The legacy cudnnConvolutionForward spends about 40 µs of CPU per call on cuDNN 9. Convolutions run in plain fp32 unless cuda_set_tf32(1) is on, the same switch the cuBLAS GEMMs use.

fft_cuda.h is the GPU counterpart of fft.h. Unlike fft.h, it accepts any length, not only powers of two. Forward transforms are unnormalized and inverse ones include 1/N, as in fft.h.

tests/stdlib/cuda checks both libraries against the CPU code: Conv2D::forward at stride 1 and 2, batched; pooling, upsample and concat; a whole U-Net; fft/ifft, dft for a non-power-of-two length, batched and 2D transforms; and the fp32 device transforms. It also replays a convolution from a graph, runs one with TF32 on and one on other buffers, and checks large and odd-sized host transfers, including a write queued after kernels that still read the old contents.

On the H200, the queued ops beat PyTorch-CUDA on batched convolutions (32 images: 0.164 ms vs 0.257 ms in fp32) and on the U-Net (batch 1: 0.175 ms vs 0.206 ms, or 0.164 ms as one graph). PyTorch is faster for a single one-image convolution (0.022 ms vs 0.034 ms). The device FFTs match cuFFT through PyTorch. See the benchmarks.

Attention on CUDA #

#include <std/ml/gpu_cuda_attention.h>   // link: the gpu_cuda libraries (cuDNN is loaded at run time)

Batched multi-head versions of flash_attention_forward and linear_attention_forward on CudaBufs, queued on gpu_cuda's stream (so they order with the other ops and can be recorded into a CudaGraph). Forward only.

• Flash attention, softmax(Q Kᵀ · scale) V without the score matrix. fp16/bf16 go to cuDNN's fused attention engine when libcudnn.so.9 can be loaded with dlopen — SafeC builds the same graph cudnn-frontend's sdpa() does, one cached plan per shape — and otherwise to this library's FlashAttention-2 kernel on tensor cores (mma.sync, sm_80+, head dim ≤ 128). fp32 runs that kernel in 3xTF32 (fp32-level error, ~2e-7). Head dim 256 and older GPUs use a CUDA-core kernel with the same algorithm. cuda_attention_use_cudnn(0) forces the own kernels.
• Linear attention, φ(x) = elu(x) + 1, in O(n d²): up to head dim 128, a per-head pass walks the 64-key chunks with the state φ(K)ᵀV in registers (storing it before each chunk when causal), then every chunk's outputs run at once. Exact fp32 arithmetic for every data type.
• Shapes: the cuda_flash_attention / cuda_linear_attention / _f32 forms take the same arguments as ROCm's rocm_flash_attention / rocm_linear_attention (fp32, [heads][seq][dim]). The _ex and cuda_half_ forms take a CudaAttnShape: batch, grouped-query K/V (kvHeads dividing heads), and a sliding window. Head dims 32, 64, 96, 128, 256.
• Masking: query i sits at position kvLen - seqLen + i (the last seqLen keys are the queries' own tokens, any before them a cache, as in mps_flash_attn_bf16). causal hides later keys; window (flash only) also hides keys window or more positions back.
• Decoding: when (heads / kvHeads) · seqLen ≤ 16 (one or a few new tokens against a KV cache), flash attention in any dtype runs a split-KV (flash-decoding) kernel pair instead: the keys are cut into 32-key-multiple splits, about eight blocks per SM, each block reading every key and value once for all the query heads that share it, and a second kernel merges the partial softmaxes. On the H200, against the cuDNN path it replaces, a bf16 32k-token decode with 8 query heads over 1 kv head (head dim 256) drops from 1.40 ms to 36 µs and 32/8 heads at head dim 128 from 533 to 54 µs. PyTorch's scaled_dot_product_attention (FlashAttention-2's split-KV kernel on tensor cores) is still 1.1–1.8× faster. cuda_attention_use_decode(0) turns the path off.
• Each (kernel, dtype, head dim) variant is compiled by NVRTC on its first call, and cuDNN plans and the linear-attention and decode workspaces are made then too, so make one call per shape before capturing a graph.

On the H200, fp16/bf16 flash attention matches PyTorch's scaled_dot_product_attention (both use cuDNN), fp32 is 1.34× faster at head dim 64 and level at 128 (0.97–1.05×), and linear attention is up to 2.6× faster than PyTorch and up to 2.1× faster than flash-linear-attention (level with PyTorch's TF32 einsums at head dim 128 fp32 non-causal). See the benchmarks. gpu_cuda.h gained cuda_compute_capability(), cuda_sm_count() and cuda_load_module(source, name) for libraries like this one that bring their own NVRTC kernels.

Convolution (MIOpen) and FFT (hipFFT) on ROCm #

#include <std/ml/gpu_miopen.h>   // link: -lMIOpen + the gpu_rocm libraries
#include <std/dsp/fft_rocm.h>    // link: -lhipfft + the gpu_rocm libraries

The AMD counterparts of gpu_cudnn.h and fft_cuda.h: the same functions with rocm_ in place of cuda_, on RocmBufs, queued on gpu_rocm's stream so they order with its ops and can be recorded into a RocmGraph.

gpu_miopen.h runs CNN and U-Net through MIOpen. The first convolution of each shape asks MIOpen to time its algorithms on that call's buffers (miopenFindConvolutionForwardAlgorithm) and caches the fastest one together with its own workspace. Make that first call outside graph capture. MIOpen has no fused convolution + bias + ReLU for arbitrary shapes, so the bias add and ReLU run as one HIP kernel after the convolution (rocm_bias_channels).

fft_rocm.h runs the transforms through hipFFT (rocFFT underneath), for any length. Forward transforms are unnormalized and inverse ones include 1/N.

tests/stdlib/rocm/rocm_dnn_fft_check.sc runs the same checks as the CUDA suite against the CPU code (73 checks on an MI100): convolutions at stride 1 and 2, batched, with ReLU and replayed from a graph; pooling, upsample and concat; a whole U-Net; every FFT form; and fp16 and bf16 convolutions and pooling against cnn.h on the same rounded inputs. 16-bit pooling is a HIP kernel in gpu_rocm.h, since MIOpen has no bfloat16 pooling solver on every GPU (none on the MI100).

On the MI100, one 64-channel convolution takes 0.027 ms against PyTorch-ROCm's 0.063 ms and the batch-1 U-Net 0.227 ms against 0.315 ms; at batch 16–32 the two are within 2%. Device-resident FFTs are 1.4–1.7× faster than through PyTorch. See the benchmarks.

Host-array API (MPS, CUDA, SPIR-V, WebGPU) #

Each backend exposes the same family of functions over plain float arrays, prefixed mps_, cuda_, spirv_ or webgpu_. Each uploads, runs, and copies back in one call, and returns 1 on success or 0 on failure (no device, compile error), leaving out untouched. ROCm's equivalents (rocm_*_f32) are listed in its table above.

// <p> = mps | cuda | spirv | webgpu
int <p>_available();

int <p>_add_f32([]const float a, []const float b, []float out);   // also sub, mul, div, pow; n = out.len
int <p>_relu_f32([]const float a, []float out);                    // also log, exp, sqrt
int <p>_scale_f32([]const float a, float k, []float out);
int <p>_sum_f32([]const float a, []float out);                     // out[0] = sum of a
int <p>_matmul_f32([]const float a, []const float b, []float out,
                   unsigned long M, unsigned long K, unsigned long N);          // [M,K] x [K,N]
int cuda_matmul_f32_blas([]const float a, []const float b, []float out,
                         unsigned long M, unsigned long K, unsigned long N);    // cuBLAS

Host arrays are passed as slices ([]const float in, []float out); n is out.len, and an op returns 0 when an input is shorter. Device-resident data lives in typed buffer objects instead: CudaBuf (CUDA), RocmBuf (ROCm), SpirvBuf (SPIR-V), WebGpuBuf (WebGPU) and MpsBuf (Metal, below), each made with alloc(n, dtype) or upload(host) and read back with read(bytes) / read_f32(floats) (cuda_read / rocm_read for CUDA and ROCm).

Persistent buffers keep weights on the device across calls, so a training loop doesn't re-upload them: upload once, then run the buffer ops (cuda_gemm, webgpu_matmul, webgpu_sgd_update, mps_matmul_f32_persistent, …) on the buffers.

On Metal a device buffer is a struct MpsBuf: the MTLBuffer handle, its size in bytes and its element type (0 = fp32, MPS_FP16, MPS_BF16). Its memory is shared with the CPU, so read and write are copies without a transfer. A zero-initialized MpsBuf is invalid, every op returns 0 for an invalid one, and free() on it does nothing.

struct MpsBuf {
    ?&MtlObject handle; unsigned long bytes; int dtype;
    int  valid() const;
    void free();                                          // back to the buffer pool
    int  read([]unsigned char dst) const;                 // valid once the batches writing it have ended
    int  write([]const unsigned char src);
    static struct MpsBuf alloc(unsigned long n, int dtype);                 // n elements
    static struct MpsBuf upload([]const float host);          // fp32 values
    static struct MpsBuf upload_f16([]const float host);    // rounded on the CPU; also upload_bf16
    static struct MpsBuf upload_bits([]const unsigned short bits, int dtype);  // raw fp16 / bf16
    static struct MpsBuf linear_attn_state(unsigned long numHeads, unsigned long headDim);
    static struct MpsBuf batch_last_output();             // owned by the open batch: don't free()
};
int mps_matmul_f32_persistent(struct MpsBuf bufA, struct MpsBuf bufB, &float out,
                              unsigned long M, unsigned long K, unsigned long N);
int mps_sgd_update_f32_chained(struct MpsBuf bufW, struct MpsBuf bufGrad, float lr, unsigned long n);  // in place

Half precision (MPS, CUDA, WebGPU): upload as fp16/bf16, compute natively, read back raw 16-bit values (convert with float16.h below). ROCm keeps 16-bit data in RocmHalfBufs instead; see its table above.

static struct CudaHalfBuf CudaHalfBuf::upload([]const float host, int dtype);   // also RocmHalfBuf; Metal: MpsBuf::upload_f16 / _bf16; WebGPU: WebGpuBuf::upload_f16
int    mps_matmul_f16_persistent(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out,
                                unsigned long M, unsigned long K, unsigned long N); // also _bf16
int    cuda_half_sgd(struct CudaHalfBuf w, struct CudaHalfBuf g, float lr);   // also rocm_half_sgd
int   mps_matmul_f16(const &float a, const &float b, &float out,
                     unsigned long M, unsigned long K, unsigned long N);   // also mps_matmul_bf16
int mps_relu_f16([]const float a, []float out);           // also mps_relu_bf16

MPS batching. Between mps_batch_begin() and mps_batch_end(), MPS ops encode onto one shared command buffer instead of each committing and waiting, so a fused forward + loss + backward step runs as one GPU submission. *_chained variants take the previous op's output buffer (MpsBuf::batch_last_output()) as input without a round trip through host memory. mps_batch_submit() sends what the open batch holds so far to the GPU without waiting, so long batches overlap encoding with execution.

void  mps_batch_begin();
void  mps_batch_end();                      // commit, wait, run finalizers
void  mps_batch_submit();                   // send what's encoded so far, don't wait
int   mps_batch_is_active();
generic<T> void mps_batch_register_finalize(fn void(?&T) finalizeFn, ?&T ctx);  // runs after mps_batch_end's wait

// Transposed matmuls for backward passes: abt = A·Bᵀ, atb = Aᵀ·B.
int mps_matmul_abt_f32(const &float a, const &float b, &float out, unsigned long M, unsigned long N, unsigned long K);
int mps_matmul_atb_f32(const &float a, const &float b, &float out, unsigned long M, unsigned long K, unsigned long N);
int mps_relu_backward_f32([]const float a, []const float selfGrad, []float out);

// Chained forms (first operand already on the GPU), plus fp16/bf16 versions:
//   mps_matmul_f32_chained, mps_relu_f32_chained, mps_sub_f32_chained, mps_mul_f32_chained,
//   mps_square_f32_chained, mps_sum_f32_chained, mps_scale_f32_chained, mps_gated_gelu_f32_chained,
//   mps_matmul_abt_f32_chained, mps_matmul_atb_f32_chained, mps_matmul_atb_f32_chained_ab,
//   mps_relu_backward_f32_chained, mps_relu_backward_f32_chained_ab, mps_matmul_abt_f32_persistent,
//   mps_sgd_update_f32_chained;
//   mps_relu_{f16,bf16}_chained, mps_sub_{f16,bf16}_chained, mps_scale_{f16,bf16}_chained,
//   mps_matmul_atb_{f16,bf16}_chained_ab, mps_matmul_abt_{f16,bf16}_persistent,
//   mps_relu_backward_{f16,bf16}_chained_ab, mps_sgd_update_{f16,bf16}_chained

The _ab suffix means both inputs are already GPU buffers (MpsBufs, passed by value). See std/ml/gpu_mps.h for each one's argument order.

Attention on MPS #

#include <std/ml/gpu_mps.h>   // Apple GPUs; links Metal and Foundation (safeguard adds them)

The same batched multi-head flash_attention_forward and linear_attention_forward as CUDA's, on Metal buffers (MpsBuf::upload, MpsBuf::upload_f16 / upload_bf16, MpsBuf::alloc). Forward only; all math in fp32.

• Flash attention, softmax(Q Kᵀ · scale) V without the score matrix: FlashAttention-2 on the GPU's 8×8 simdgroup matrix units for head dims up to 256 (Q block in registers, K/V tiles in threadgroup memory, online softmax on the fragments), a simdgroup-reduction kernel at 512. fp16 multiplies in half precision; bf16 is converted to fp32 as tiles are staged, since the M1 family has no bf16 matrix units.
• Linear attention, φ(x) = elu(x) + 1, in O(n d²). Head dims 64 and 128 run chunk-parallel on the matrix units — per-segment states φ(K)ᵀV, a prefix over 64-row chunks (causal) or their total (non-causal), then every chunk's outputs at once, with the states in a scratch buffer the library keeps. Other head dims, a KV-cache offset with causal, and the recurrent form below walk the sequence per head.
• Shapes: mps_flash_attention / mps_linear_attention / _f32 take the same arguments as rocm_ and cuda_'s (fp32, [heads][seq][dim]); the _ex forms take an MpsAttnShape (batch, grouped-query K/V, sliding window) and a dtype (0 = fp32, MPS_FP16, MPS_BF16). Head dims 32, 64, 96, 128, 256, 512.
• Masking: as on CUDA — query i sits at position kvLen - seqLen + i; causal hides later keys; window (flash only) also hides keys window or more positions back. scale ≤ 0 means 1/√headDim.
• Batching: inside mps_batch_begin()/mps_batch_end() the calls are encoded into the batch's one command buffer; outside one, each call runs and waits.

On an M1 Pro, fp32 flash attention is 1.01–1.42× faster than MLX's fused attention and 1.13–1.70× faster than PyTorch's mps device, bf16 at head dim 128 is level with MLX, fp16 and bf16 at 64 reach 0.80–1.02× of them (2.6–3.4 TFLOPS), and linear attention is 2.0–4.4× faster than vectorized PyTorch and MLX. See the benchmarks.

tests/stdlib/mps/run.sh builds mps_attention_check, which compares every path (both flash kernels, the chunked and sequential linear kernels, all three data types, GQA, KV-cache offsets, sliding windows, the host-array forms) with a double-precision reference and attention.h's CPU functions: 95 checks.

MPS attention and transformer ops #

bf16 building blocks for running a transformer layer on an Apple GPU. Tensors are raw bf16 bits (unsigned short, convert with float16.h), row-major, with heads side by side in each row — the layout a layer's projections produce: Q and outputs are [rows, numHeads*headDim], K/V [rows, numKvHeads*headDim]. numHeads must be a multiple of numKvHeads (grouped-query attention). The flash and linear functions run on the same kernels as mps_flash_attention_ex (head dims 32–512); the decode kernel takes 64, 128, 256 or 512. Like every MPS entry point, each returns 0 when it can't run (no device, unsupported shape) so the caller can fall back to the CPU.

int mps_attn_batch_bf16(const &unsigned short q, const &unsigned short kAll, const &unsigned short vAll,
                        const &unsigned int cacheLens, &unsigned short out,
                        unsigned long batchSize, unsigned long numHeads, unsigned long numKvHeads,
                        unsigned long headDim, unsigned long maxKvLen,
                        int isSliding, unsigned long slidingWindow);

int mps_flash_attn_bf16(const &unsigned short q, const &unsigned short k, const &unsigned short v,
                        &unsigned short out, unsigned long seqLen, unsigned long kvLen,
                        unsigned long numHeads, unsigned long numKvHeads, unsigned long headDim,
                        float scale, int causal, int isSliding, unsigned long slidingWindow);

int mps_linear_attn_bf16(const &unsigned short q, const &unsigned short k, const &unsigned short v,
                         &unsigned short out, unsigned long seqLen, unsigned long kvLen,
                         unsigned long numHeads, unsigned long numKvHeads, unsigned long headDim,
                         int causal);

// Recurrent linear-attention state (S and z for every head), GPU-resident.
struct MpsBuf MpsBuf::linear_attn_state(unsigned long numHeads, unsigned long headDim);  // zeroed
unsigned long mps_linear_attn_state_bytes(unsigned long numHeads, unsigned long headDim);
int mps_linear_attn_bf16_chained(struct MpsBuf bufQ, struct MpsBuf bufK, struct MpsBuf bufV, struct MpsBuf bufState,
                                 &unsigned short out,
                                 unsigned long seqLen, unsigned long numHeads,
                                 unsigned long numKvHeads, unsigned long headDim);

// Chained bf16 elementwise/normalization ops for the rest of the layer:
int mps_gated_gelu_bf16_chained(struct MpsBuf bufRaw, &unsigned short out, unsigned long rows, unsigned long intermediateSize);
int mps_add_bf16_chained(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out, unsigned long n);       // residuals
int mps_gelu_mul_bf16_chained(struct MpsBuf bufA, struct MpsBuf bufB, &unsigned short out, unsigned long n);  // gelu(a) * b
int mps_rmsnorm_bf16_chained(struct MpsBuf bufX, struct MpsBuf weightBuf, &unsigned short out,
                             unsigned long rows, unsigned long cols, float eps);

• Masking. With causal, query i sits at position kvLen - seqLen + i — the last seqLen keys are the queries' own tokens and any before them are a cache — and sees keys at or before it. isSliding also hides keys slidingWindow or more positions behind the query. mps_attn_batch_bf16's cacheLens[b] is each sequence's KV length including the new token.
• Scale. mps_flash_attn_bf16 multiplies q·k by scale: pass 1/sqrt(headDim) to match flash_attention_forward, or 1.0 for models that pre-scale q. The decode kernel applies no scale. Linear attention needs none, since dividing by phi(q)·z cancels any scaling of q.
• Batching. The one-shot functions work inside an open batch (results land in out at mps_batch_end()) or on their own (filled on return). mps_attn_batch_bf16_chained and mps_flash_attn_bf16_chained take device buffers and leave their output as MpsBuf::batch_last_output() — so attention can feed the output projection (mps_matmul_abt_bf16_persistent) in the same submission. mps_matmul_abt_bf16_persistent with M == 1 uses a dedicated matrix-vector kernel.

Linear attention decode loop:

struct MpsBuf state = MpsBuf::linear_attn_state(numHeads, headDim);
if (state.valid()) {
    mps_batch_begin();
    mps_linear_attn_bf16_chained(promptQ, promptK, promptV, state, promptOut, promptLen,
                                 numHeads, numKvHeads, headDim);
    mps_batch_end();
    // each generated token:
    mps_batch_begin();
    mps_linear_attn_bf16_chained(tokQ, tokK, tokV, state, tokOut, 1, numHeads, numKvHeads, headDim);
    mps_batch_end();
    // ...
    state.free();
} else {
    // no Metal device: use linear_attention_forward on the CPU
}

On an M1 Pro, a decode step over 8 sequences × 2048 cached tokens (8 heads, head dim 256) takes about 0.3–0.5 ms and a 1024-token causal flash-attention prefill (8 heads, head dim 256) about 7 ms.

Tensor wrappers #

Autograd FloatTensor ops that run one forward pass on a GPU or a CPU library and reuse the CPU backward. Each falls back to the CPU when its device is unavailable. On ROCm, CUDA, Metal, Vulkan and WebGPU, tensors placed on the device run every op there, backward included.

tensor_auto.h is CPU-only on purpose: across every size benchmarked (matmul from 2×2×2 to 128×128×128, and a 128×512×1024×256 training step) BLAS never lost to the core loop (measured before tensor_matmul moved to sgemm) and MPS never beat BLAS, so "auto" doesn't pick the GPU. To use a GPU, place the tensors on it.

Backend status #

The SPIR-V is generated from device_kernels.cu by gen_device_kernels.sh (Slang) and checked in, so a program needs only the Vulkan loader. WGSL has no pointers, 64-bit integers or float atomics, so WebGPU runs device_kernels.wgsl, a port gen_device_kernels_wgsl.sh checks against the CUDA source.

fp16 / bf16 #

No native 16-bit float in SafeC's type system (Type.h's TypeKind has only Float32/Float64 — a real compiler feature, not a stdlib change). Instead, the standard workaround: fp16/bf16 carried as raw bits in unsigned short, with explicit correctly-rounded (round-to-nearest-even) conversion to/from float. On top of that sit half-precision tensors and mixed-precision training (tensor_half.h, below), native half/bfloat GPU compute on MPS, CUDA, WebGPU and ROCm (see Host-array API), and 16-bit training on a device (precision).

#include <std/ml/float16.h>

unsigned short f32_to_fp16(float v);          // round to nearest even
float          fp16_to_f32(unsigned short h);
unsigned short f32_to_bf16(float v);
float          bf16_to_f32(unsigned short b);

void f32_to_fp16_bulk([]const float src, []unsigned short dst);
void fp16_to_f32_bulk([]const unsigned short src, []float dst);
void f32_to_bf16_bulk([]const float src, []unsigned short dst);
void bf16_to_f32_bulk([]const unsigned short src, []float dst);
void fp16_to_bf16_bulk([]const unsigned short src, []unsigned short dst);
void bf16_to_fp16_bulk([]const unsigned short src, []unsigned short dst);

int         half_cpu_features();        // HALF_CPU_* bits, detected once
const &char half_cpu_features_name();   // "F16C, AVX512-BF16", ..., "none"

What the CPU does natively. Every x86-64 CPU since 2012 converts fp16 in hardware (F16C), and so does every AArch64 core (FCVTL / FCVTN). bf16 is newer and optional: AVX512-BF16 on Cooper Lake, Sapphire Rapids and Zen 4 or later, and FEAT_BF16 on Armv8.6 cores (Neoverse V1 / N2, Apple M2 and later). fp16 arithmetic (AVX512-FP16, FEAT_FP16) is optional too. half_cpu_features() reports what this CPU has:

Detection runs once: cpuid and xgetbv on x86-64, getauxval on Linux AArch64, sysctl on Apple. SAFEC_HALF=soft reports none and makes the bulk conversions use the portable code, for comparison.

Bulk conversions.

• fp16: F16C on x86-64 and FCVTL / FCVTN on AArch64.
• fp32 → bf16: AVX512-BF16 where present, 16 values per instruction. VCVTNEPS2BF16 flushes fp32 subnormals to zero, so a block containing one is spotted with VFPCLASSPS and converted in software instead.
• Everywhere else: bf16 → fp32 is a shift, and fp32 → bf16 an integer rounding, both in loops the compiler vectorizes.
• Arrays of 256K values or more are split across the thread pool.

The results are bit-identical to the scalar functions, except that NaNs may keep more of their payload. The tests compare every one of the 65,536 fp16 patterns, and random fp32 patterns of every class, in both modes.

On a Ryzen 7 7800X3D (one thread, 16M values, read + write bandwidth):

Across all 16 threads, every conversion reaches 58–99 GB/s with the hardware paths.

Half-precision tensors — tensor_half.h #

#include <std/ml/tensor_half.h>

Two layers:

• HalfTensor (fp16) and BF16Tensor (bf16) — DenseTensor stored in 16 bits, half the memory of a FloatTensor. Their ops and autograd are the core tensor ops: operands load to float32, the float32 kernel runs, the result rounds once (tensor_matmul accumulates each dot product in float32, like tensor cores), and gradients are float32. This header adds conversions between element types, the batch helpers mixed-precision training uses, and gemv for decoding with 16-bit weights.
• Mixed precision on float32 tensors. FloatTensor::cast rounds values to fp16/bf16 precision in the forward pass and passes gradients through unchanged (the straight-through estimator), so a graph can compute in reduced precision while weights and gradients stay float32. LossScaler is fp16 training's dynamic loss scaling; bf16 shares float32's exponent range and doesn't need it.

#define TENSOR_FP16 1
#define TENSOR_BF16 2

float          half_round(float v, int dtype);          // nearest dtype value, as float32
unsigned short half_encode(float v, int dtype);
float          half_decode(unsigned short bits, int dtype);

// dst = src rounded to dst's element type (any pair, incl. fp16 <-> bf16);
// sizes must match (0 otherwise). Not a graph node -- see to_float / from_float.
generic<S, D> int tensor_convert_into(const &DenseTensor<S> src, &DenseTensor<D> dst);

// Mixed-precision batches: fp32 masters <-> 16-bit copies, 16-bit gradients back
generic<S, D> void tensors_convert(const &Vec<&DenseTensor<S>> src, &Vec<&DenseTensor<D>> out);  // appends copies
generic<S, D> int  tensors_convert_into(const &Vec<&DenseTensor<S>> src, &Vec<&DenseTensor<D>> dst);
generic<E>    void tensors_free(&Vec<&DenseTensor<E>> ts);
generic<E>    int  half_grads_accumulate(const &Vec<&DenseTensor<E>> grads, &Vec<&FloatTensor> params,
                                         float invScale);  // 1 = all finite

int HalfTensor::gemv([]const float x, []float y) const;     // y = w x, float32 x (k) / y (n)
int HalfTensor::gemv_batch([]const float x, unsigned long m, []float y) const; // m vectors: x [m][k], y [m][n]
int BF16Tensor::gemv([]const float x, []float y) const;
int BF16Tensor::gemv_batch([]const float x, unsigned long m, []float y) const;

&FloatTensor FloatTensor::cast(int dtype) const;                            // straight-through gradient
&FloatTensor tensor_matmul_lowp(const &FloatTensor a, const &FloatTensor b, int dtype);    // cast(cast(a) @ cast(b))

struct LossScaler {
    float scale;            // multiply the loss by this before FloatTensor::backward()
    float growthFactor;     // 2.0
    float backoffFactor;    // 0.5
    int   growthInterval;   // 2000 clean steps
    int   goodSteps;
};
static struct LossScaler LossScaler::new();                  // scale 65536 (PyTorch GradScaler's defaults)
int  LossScaler::unscale(&FloatTensor p);                  // divide p.grad by scale; 0 if any grad is inf/NaN
int  LossScaler::unscale_all(&Vec<&FloatTensor> params);   // every tensor in a list
void LossScaler::update(int allFinite);               // grow after clean steps, back off after overflow

A HalfTensor's or BF16Tensor's data has exactly the layout the GPU half-precision functions read and write (mps_upload_f16_persistent and friends).

gemv is the matrix-vector product of decoding one token: w is a linear layer's 2D fp16/bf16 weight ([n][k]), x and y are float32. Each weight is widened in registers, so there is no float32 copy of the matrix; it is vectorized for the CPU's SIMD tier and split over rows on the thread pool. On a Ryzen 7 7800X3D (16 threads) a bf16 24576 × 1536 product takes 1.0 ms (76 GB/s of weights, near the memory's bandwidth), 6× faster than sgemm with one row on widened float32 weights; fp16 reaches 35 GB/s (its widening is done in integer arithmetic). It took the Gemma 4 E2B CPU decode from 4.43 s to 0.225 s per token. On AArch64 a NEON kernel loads eight weights as four 32-bit lanes and widens both halves in place. On an M1 Pro (8 performance cores) it reaches 95–113 GB/s for bf16 and 40–55 GB/s for fp16. gemv_batch does the same for m vectors at once, for a short prompt's prefill or speculative verification, and reads the weights from memory once for all of them. Six vectors over a bf16 24576 × 1536 weight take 1.65 ms on the M1 Pro, against 4.35 ms for six gemv calls. With it, the Gemma 4 E2B CPU path on the M1 Pro takes 136 ms per token.

Storage and arithmetic in bf16, with autograd:

#include <std/ml/tensor_half.h>
#include <std/collections/slice.h>
#include <std/io.h>

int main() {
    // Store weights in bf16 (half the memory), compute a matmul in bf16.
    &BF16Tensor w = BF16Tensor::new_2d(2UL, 2UL, 1);
    &BF16Tensor x = BF16Tensor::new_2d(2UL, 1UL, 0);
    w.set2(0UL, 0UL, 0.5f);  w.set2(0UL, 1UL, -1.0f);
    w.set2(1UL, 0UL, 2.0f);  w.set2(1UL, 1UL, 0.25f);
    x.set1(0UL, 3.0f);       x.set1(1UL, 4.0f);
    &BF16Tensor y = tensor_matmul(w, x);      // f32 accumulate, bf16 result
    std::println_float((double)y.at1(0UL));   // -2.5
    std::println_float((double)y.at1(1UL));   // 7

    y.sum().backward();                       // w.grad: float32
    struct Slice<float> g = std::Slice::of(w.grad, w.size);
    std::println_float((double)*g.at(1UL));   // 4 (= x[1])
    x.free(); w.free();
    return 0;
}

Mixed-precision training — fp16 forward, float32 master weights, skipped steps on overflow:

#include <std/ml/tensor_half.h>
#include <std/collections/slice.h>

// One mixed-precision SGD step: fp16 forward, float32 master weights,
// dynamic loss scaling so small fp16 gradients don't flush to zero.
int train_step(&FloatTensor w, const &FloatTensor x, const &FloatTensor target, &LossScaler scaler, float lr) {
    &FloatTensor y = std::tensor_matmul_lowp(x, w, TENSOR_FP16);
    &FloatTensor diff = std::tensor_sub(y, target);
    &FloatTensor loss = std::tensor_mul(diff, diff).sum();
    loss.scale(scaler.scale).backward();

    int finite = scaler.unscale(w);
    if (finite) {
        struct Slice<float> g = std::Slice::of(w.grad, w.size);
        for (unsigned long i = 0UL; i < w.size; i = i + 1UL) {
            w.set1(i, w.at1(i) - lr * *g.at(i));
        }
    }
    scaler.update(finite);
    w.zero_grad();
    return finite;
}

int main() {
    float wv[2] = { 0.1f, 0.2f };
    float xv[2] = { 1.0f, 2.0f };
    float tv[1] = { 1.0f };
    &FloatTensor w = FloatTensor::from_2d(wv[0..2UL * 1UL], 2UL, 1UL, 1);
    &FloatTensor x = FloatTensor::from_2d(xv[0..1UL * 2UL], 1UL, 2UL, 0);
    &FloatTensor t = FloatTensor::from_2d(tv[0..1UL * 1UL], 1UL, 1UL, 0);
    struct LossScaler scaler = LossScaler::new();
    int ok = 1;
    for (int i = 0; i < 20; i = i + 1) ok = ok & train_step(w, x, t, &scaler, 0.05f);
    float pred = w.at1(0UL) + w.at1(1UL) * (float)2.0;
    return (ok && pred > (float)0.95 && pred < (float)1.05) ? 0 : 1;
}

Mixed precision on the GPU #

#include <std/ml/gpu_rocm.h>     // ROCm; gpu_cuda.h has the same functions with cuda_ names
#include <std/ml/gpu_miopen.h>   // 16-bit convolution (gpu_cudnn.h on CUDA)

Mixed-precision training normally keeps:

• fp32 master weights and optimizer state;
• a 16-bit copy of the weights for the forward and backward passes;
• gradients that carry the loss scale.

ROCm and CUDA have the same queued ops for every part of the step:

Because the skip happens on the GPU, an overflowed step needs no host round trip before the update. Read the 4-byte flag once per step, whenever it's convenient, to grow or back off the LossScaler. bf16 shares fp32's exponent range and usually trains without loss scaling; pass an invalid flag (dev == 0) to rocm_mp_* to never skip.

Every kernel computes in fp32 and rounds once per stored value. Updates run without FMA contraction (fp contract(off) in HIP, __fmul_rn-style intrinsics in CUDA), so the results match the CPU arithmetic. tests/stdlib/rocm/rocm_check.sc checks this on an MI100:

• fp16 ↔ bf16 over every 16-bit pattern;
• 40 buffers through bf16 and back;
• unscaling, and the overflow flag;
• three Adam, AdamW and momentum-SGD steps against the CPU (within 1e-5) with their bf16 copies, and a flagged step changing nothing;
• every 16-bit elementwise op and the NCHW bias bit-for-bit.

The CUDA versions have the same tests in tests/stdlib/cuda/cuda_check.sc. Their kernels compile (clang, sm_80) and the tests build, but they haven't yet run on an NVIDIA GPU.

Device selection #

A tensor's device picks where its ops run: t.place(d), t.to(d), model_to(&m, d), TrainConfig.device, device_set_default(d) or device_auto() (Training on a GPU). Every layer built from tensor ops, transformer and diffusion blocks included, runs on the device of its parameters and input.

LLM orchestration #

#include <std/ml/llm.h>

• LlmClient — speaks the OpenAI-chat-completions wire shape (what vLLM's own server, and most local/hosted LLM servers, implement). client.chat() POSTs to <host>:<port>/v1/chat/completions.
• PromptTemplate — prompt_template_render(tmpl, vars) does {name} substitution against a Value object.
• Chain — a fixed Value -> Value step sequence.
• Graph executor — LangGraph-style: named nodes transform a shared Value state, connected by fixed or conditional (router-function) edges, run until "__end__".
• Tracer — an in-memory, JSON-dumpable log of named input/output/timing events (LangSmith-style).

struct LlmMessage { struct String role; struct String content; };
struct LlmClient  { struct String host; unsigned short port; struct String model; struct String apiKey; };

static struct LlmClient LlmClient::new(const &char host, unsigned short port, const &char model, const &char apiKey);
// The assistant's reply text; *ok = 0 on transport or response-shape failure.
struct String LlmClient::chat(&Vec<struct LlmMessage> messages, ?&int ok);

struct String prompt_template_render(const &char tmpl, const &Value vars);   // {name} -> vars.name

typedef fn struct Value(const &Value input) ChainStepFn;
static struct Chain Chain::new();
void         Chain::add_step(ChainStepFn step);
struct Value Chain::run(const &Value input);
void         Chain::free();

typedef fn struct Value(const &Value state) GraphNodeFn;
typedef fn struct String(const &Value state) GraphRouterFn;   // name of the next node, or "__end__"
static struct Graph Graph::new(const &char entryNode);
void         Graph::add_node(const &char name, GraphNodeFn nodeFn, const &char nextNode);
void         Graph::add_conditional_node(const &char name, GraphNodeFn nodeFn, GraphRouterFn router);
struct Value Graph::run(struct Value initialState, int maxSteps);
void         Graph::free();

struct TraceEvent { struct String name; struct String inputJson; struct String outputJson; double durationMs; };
static struct Tracer Tracer::new();
void          Tracer::record(const &char name, const &Value input, const &Value output, double durationMs);
struct String Tracer::to_json() const;
void          Tracer::free();

All five verified end-to-end, including LlmClient against a real local mock HTTP server standing in for an LLM endpoint (the same technique std/rpc/server_fn.h's own tests use — nothing here calls out to a real hosted LLM API).