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.

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.

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.

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.

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.

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]) written into their slots: slots[t] (integers, [T]) is block slots[t] / block_size, position slots[t] % block_size. 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.

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.

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.