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) (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, ported) 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:

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

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]):

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.

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.