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) and norm_rope (norm, RoPE and a cache write); and, through crate::softmax and crate::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.

use nuvec_core::{Array, Result, fast};

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

Structs

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

AttentionMask

The mask of scaled_dot_product_attention.

FastOp

A fused operation (see the module documentation).

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 crate::add 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_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.

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.