Module fast 

Fused operations with dedicated GPU kernels, as MLX’s mx.fast: rms_norm, layer_norm, rope, cross_entropy, scaled_dot_product_attention and gated_delta (mlx-lm’s Gated DeltaNet kernels); nuvec’s add_rms_norm (a residual connection and its norm), norm_rope (norm, RoPE and a cache write) and paged_attention with paged_write (a block-table KV cache); and, through crate::Array::softmax and crate::Array::logsumexp over the last axis, MLX’s softmax and log-sum-exp kernels.

Each operation is one GPU kernel (MLX’s, or nuvec’s own) with an equivalent graph of ordinary operations, its fallback. As in MLX, the fallback runs on the CPU stream, and it defines the forward-mode derivative and batching; reverse mode uses dedicated kernels where MLX has them.

Like every operation, these return arrays: invalid arguments give a failed array that carries the error and the caller’s location (see crate::Array::check).

Memory

The norms, RoPE, norm_rope, cross entropy, softmax, log-sum-exp, the gated delta rule and the paged kernels hold no scratch beyond their inputs and outputs on the GPU (the fallback graphs on the CPU hold their intermediates). Attention’s scratch depends on its plan: see scaled_dot_product_attention. Rotation frequencies are computed per call, never cached.

use nuvec_core::{array, fast};

let x = array![[1.0, 2.0], [3.0, 4.0]];
let y = fast::rms_norm(&x, None, 1e-5); // x / sqrt(mean(x²) + eps), per row
let v = y.to_vec::<f32>().unwrap();
assert!((v[0] - 1.0 / 2.5f32.sqrt()).abs() < 1e-5);

Structs

AttentionOptions

The options of scaled_dot_product_attention and plan_attention: the mask, attention sinks, a sliding window and a cap on score memory. Build them with AttentionOptions::causal or Default and the chained setters:

AttentionPlan

How an attention call runs and the memory it holds (see plan_attention).

KernelCall

A recorded call of a custom kernel.

KernelConfig

How a MetalKernel is launched.

MetalKernel

A Metal kernel built from source, as MLX’s mx.fast.metal_kernel: a name, the names of its inputs and outputs, and the body of a kernel function. Calling it generates the function’s signature as MLX does:

Rotation

The rotation of norm_rope: rope’s parameters.

StepOptions

Options of gated_delta_steps_with.

Enums

AttentionKernel

The kernel family an attention call runs (see plan_attention).

AttentionMask

The mask of scaled_dot_product_attention.

FastOp

A fused operation (see crate::fast).

TemplateArg

A template argument of a custom kernel.

Functions

add_rms_norm

A residual connection followed by RMS normalisation: h = x + residual and rms_norm(h, weight, eps), returned as (h, normed), from one GPU kernel. The results equal x + residual then rms_norm exactly; a weight that would widen the sum’s dtype runs the two operations.

causal_conv1d_slots

A causal depthwise convolution for a batch of sequences whose last inputs live in a pool (Gated DeltaNet’s convolution state): sequence b’s window is pool[slots[b]] (pool is [slots, K − 1, C], slots [B] distinct integers in range) followed by its new inputs x[b] (x is [B, L, C]), convolved with w ([C, K, 1], one filter per channel) and, with silu, passed through silu. Returns the outputs ([B, L, C]) and the pool with each slot holding its window’s last K − 1 inputs, written in place when the evaluation owns it.

cross_entropy

Cross entropy over the last axis of logits against integer class indices targets (the logits’ shape without the last axis; negative indices count from the end), as MLX’s fast.cross_entropy: logsumexp(logits) - logits[target] per row, in float32.

gated_delta

The gated delta rule, Gated DeltaNet’s recurrence (as mlx-lm’s gated_delta_update once its gates are computed). For each step t and value head, with S its state ([Dv, Dk]):

gated_delta_slots

gated_delta for a batch of sequences whose states live in a pool: sequence b starts from pool[slots[b]] (pool is [slots, Hv, Dv, Dk], slots [B] integers, distinct) and its final state replaces it. Returns the outputs and the updated pool, which is written in place when the evaluation owns it, so decoding many sequences moves no states. Each sequence’s results equal gated_delta’s on its state alone.

gated_delta_steps

gated_delta returning the state after every step, [B, T, Hv, Dv, Dk] (step t’s state at index t; the last is gated_delta’s final state), so that a caller can roll the recurrence back to any step, as speculative decoding does after a rejected draft. The outputs equal gated_delta’s bit for bit; writing every step’s state costs T times the state’s bytes.

gated_delta_steps_with

gated_delta_steps for a runtime that stores its states in another dtype than float32 (bfloat16, say): each step’s output is computed from the state the previous step stored, and its state is then stored in options.storage, so states[t] and the outputs equal t + 1 serial gated_delta calls on a state kept in that dtype, bit for bit (each step runs the kernel serial decode runs). With float32 storage it is gated_delta_steps.

layer_norm

Layer normalisation over the last axis, (x - mean) / sqrt(var + eps) · weight + bias, as MLX’s fast.layer_norm: statistics in float32, the result in the promoted type.

moe_route

A mixture of experts’ routing from its router’s probabilities probs ([..., E]): the indices (uint32) of each row’s k largest probabilities and those probabilities ([..., k] each), divided by their sum with normalize (mlx-lm’s norm_topk_prob).

norm_rope

An attention layer’s query or key preparation in one kernel: each head’s features of x ([B, H, L, D], any strides, as the transposed projection) optionally normalised as rms_norm with norm = (weight, eps) (Qwen 3’s q_norm and k_norm), then rotated as rope at positions offset..offset + L.

norm_rope_at

norm_rope for sequences at their own positions (a batched decoding step): each [b, h, l] row of x ([B, H, L, D], any strides) optionally normalised as rms_norm with norm = (weight, eps), then rotated as rope at position positions[b] + l (positions: [B] integers).

norm_rope_kv

Gemma 4’s keys and values from one projection (attention_k_eq_v): K is x ([B, H, L, D], any strides) normalised with k_norm = (weight, eps) and rotated at positions offset.. (as norm_rope); V is x RMS-normalised without a weight (v_eps; mlx-lm’s RMSNormNoScale). With caches = (k_cache, v_cache) both are written at offset (in place when the evaluation owns them), as norm_rope writes one.

paged_attention

Attention over a paged KV cache, for a batch of sequences of different lengths.

paged_write

cache ([blocks, kv_heads, block_size, D]) with the positions new holds ([T, kv_heads, D], or [B, T, kv_heads, D] for several sequences) written into their slots: slots (integers, new’s leading shape) holds block slot / block_size, position slot % block_size of each. Slots must be distinct; out-of-range slots (negative, or at least blocks · block_size) are skipped. In place when the evaluation owns the cache; on the GPU, new is read through its strides (a transposed view of a contiguous cache needs no copy).

plan_attention

How scaled_dot_product_attention will run for these shapes ([B, H, L, D] each), dtype and options, on the calling thread’s default device, without evaluating or allocating anything. Masks and sinks in options are read for their shape and dtype only. Errors are those the attention call would report.

residual_rms_norm

A normalised residual branch: s = (h + rms_norm(x, weight, eps)) · scale (no multiply without a scale) and, with next = (weight, eps), rms_norm of s, returned as (s, normed). Gemma 4’s sandwich norm, residual add and layer scalar, and the next layer’s input norm. On the GPU it is one kernel equal to that composition bit for bit when every input has h’s floating-point dtype, the weights are 1-D and the scale has one element (otherwise the composition, with MLX’s type promotion).

rms_norm

Root-mean-square normalisation over the last axis, x / sqrt(mean(x²) + eps) · weight, as MLX’s fast.rms_norm: statistics in float32, the result in the promoted type of x and weight (1-D, the last axis’s size).

rope

Rotary position embeddings, as MLX’s fast.rope, of x (at least 3-D: [B, ..., T, D], positions along T): the first dims features (even) are rotated in pairs, the two halves (or interleaved pairs, traditional) by angle (offset + t) · scale · θᵢ, with θᵢ = base^(-2i/dims) or 1 / freqs[i]. Exactly one of base and freqs is given. offset is an integer scalar or one per batch entry.

scaled_dot_product_attention

Scaled dot-product attention, as MLX’s fast.scaled_dot_product_attention: softmax(scale · q kᵀ + mask) v over [B, H, L, D] arrays, with the mask, sinks, window and score cap in AttentionOptions.

weighted_sum

The experts’ outputs x ([..., k, D]) weighted by w ([..., k]) and summed over the k experts: (x * w[..., None]).sum(-2), products and sum in x’s dtype as mlx-lm’s MoE blocks compute them. On the GPU it is one kernel equal to that composition bit for bit when both have one floating-point dtype and k < 32 (otherwise the composition).