Enum FastOp 

#[non_exhaustive]
pub enum FastOp {
    RmsNorm {
        eps: u32,
    },
    AddRmsNorm {
        eps: u32,
    },
    ResidualRmsNorm {
        eps: u32,
        scaled: bool,
        next: Option<u32>,
    },
    RmsNormVjp {
        eps: u32,
    },
    LayerNorm {
        eps: u32,
    },
    LayerNormVjp {
        eps: u32,
    },
    Rope {
        dims: usize,
        traditional: bool,
        base: u32,
        scale: u32,
        forward: bool,
    },
    NormRope {
        eps: Option<u32>,
        dims: usize,
        traditional: bool,
        base: u32,
        scale: u32,
        offset: usize,
        cache: bool,
        positions: bool,
        paged: bool,
    },
    DepthwiseConv1d {
        stride: usize,
        padding: usize,
        dilation: usize,
    },
    PagedAttention {
        scale: u32,
        window: u32,
        sinks: bool,
    },
    PagedWrite,
    MoeRoute {
        k: usize,
        normalize: bool,
    },
    WeightedSum,
    ConvSlots {
        silu: bool,
    },
    Softmax,
    CausalSoftmax {
        rows: u32,
    },
    LogSumExp,
    CrossEntropy,
    CrossEntropyVjp,
    Attention {
        scale: u32,
        causal: bool,
        mask: bool,
        sinks: bool,
        window: u32,
        passes: u8,
        blocks: u32,
        gqa_kernel: bool,
        tk1: bool,
    },
    GatedDelta {
        steps: bool,
    },
    GatedDeltaSlots,
    GatherMm,
    GatherQmm {
        quant: Quant,
        transpose: bool,
        global_scale: bool,
    },
}

A fused operation (see crate::fast).

Variants (Non-exhaustive)

Non-exhaustive enums could have additional variants added in future. Therefore, when matching against variants of non-exhaustive enums, an extra wildcard arm must be added to account for any future variants.

RmsNorm

(x, w): x / sqrt(mean(x²) + eps) · w over the last axis; w is 1-D or a scalar 1.

Fields

eps: u32

f32::to_bits of epsilon.

AddRmsNorm

(x, r, w): h = x + r and FastOp::RmsNorm of h (see add_rms_norm).

Fields

eps: u32

f32::to_bits of epsilon.

ResidualRmsNorm

(h, x, w[, scale][, w2]): s = (h + RmsNorm(x, w)) · scale and, with next, FastOp::RmsNorm of s with w2 (see residual_rms_norm).

Fields

eps: u32

f32::to_bits of the first norm’s epsilon.

scaled: bool

Whether a 0-D scale follows the weight.

next: Option<u32>

f32::to_bits of the second norm’s epsilon, when there is one.

RmsNormVjp

(x, w, g): the gradients of FastOp::RmsNorm for x and w.

Fields

eps: u32

f32::to_bits of epsilon.

LayerNorm

(x, w, b): (x - mean) / sqrt(var + eps) · w + b over the last axis; w and b are 1-D or the scalars 1 and 0.

Fields

eps: u32

f32::to_bits of epsilon.

LayerNormVjp

(x, w, b, g): the gradients of FastOp::LayerNorm for x, w and b.

Fields

eps: u32

f32::to_bits of epsilon.

Rope

(x, offset[, freqs]): rotary position embeddings of the first dims features.

Fields

dims: usize

Rotated features (even).

traditional: bool

Rotate interleaved pairs instead of the two halves.

base: u32

f32::to_bits of the frequency base (unused with freqs).

scale: u32

f32::to_bits of the position scale.

forward: bool

Rotate forwards (false: the inverse rotation, the gradient).

NormRope

([cache | pool, slots,] x, w[, positions][, freqs]): norm_rope and norm_rope_at: each [b, h, l] row of x RMS-normalised (with eps) and rotated as FastOp::Rope at position offset + l (positions[b] + l with positions); written into cache at the same positions, into pool at slots[b, l], or returned as a new [B, H, L, D] array.

Fields

eps: Option<u32>

f32::to_bits of epsilon, or no normalisation.

dims: usize

Rotated features (even).

traditional: bool

Rotate interleaved pairs instead of the two halves.

base: u32

f32::to_bits of the frequency base (unused with freqs).

scale: u32

f32::to_bits of the position scale.

offset: usize

The position of the first row.

cache: bool

Whether the first input is a cache written at the rows’ positions (or, with paged, a pool written at their slots).

positions: bool

Whether per-row positions ([B], int32) follow the weight.

paged: bool

Whether the first inputs are a pool [blocks, H, block_size, D] and slots [B, L] (int32; out-of-range slots skipped). Implies cache.

DepthwiseConv1d

(x, w): conv1d with groups equal to the channels (depthwise).

Fields

stride: usize

Positions between outputs.

padding: usize

Zeros before and after the input.

dilation: usize

Positions between taps.

PagedAttention

(q, k_cache, v_cache, block_tables, lengths[, sinks]): paged_attention.

Fields

scale: u32

f32::to_bits of the score scale.

window: u32

A sliding window of this many keys (0: none).

sinks: bool

Whether per-head attention sinks ([Hq]) are the last input.

PagedWrite

(cache, new, slots): paged_write.

MoeRoute

(probs): moe_route: the indices (uint32) of the k largest probabilities along the last axis in the stable ascending order (as argpartition), and those probabilities, divided by their sum if normalize.

Fields

k: usize

Experts per row.

normalize: bool

Whether the scores are divided by their sum.

WeightedSum

(x, w): weighted_sum: x ([..., k, D]) times w ([..., k]), summed over k.

ConvSlots

(x, w, pool, slots): causal_conv1d_slots: each sequence’s state pool[slots[b]] then its new inputs, convolved depthwise (and silu); the outputs and the pool with each slot’s last K − 1 inputs.

Fields

silu: bool

Whether silu follows the convolution.

Softmax

(x): softmax over the last axis, accumulated in float32.

CausalSoftmax

(x): FastOp::Softmax of causally masked scores: row r of the second-to-last axis (a multiple of rows long) keeps columns up to r % rows + L − rows (L the last axis), the others taken as the dtype’s lowest finite value (MLX’s blocked attention mask) and so 0 in the result.

Fields

rows: u32

Rows per block.

LogSumExp

(x): log(sum(exp(x))) over the last axis, which becomes size 1.

CrossEntropy

(logits, targets): logsumexp(logits) - logits[target] per row, in float32.

CrossEntropyVjp

(logits, targets, loss, g): the gradient of FastOp::CrossEntropy for the logits.

Attention

(q, k, v[, mask][, sinks]): scaled dot-product attention (see scaled_dot_product_attention).

Fields

scale: u32

f32::to_bits of the score scale.

causal: bool

Causal masking (keys after the query’s position, aligned at the end, are hidden).

mask: bool

Whether an array mask (bool or additive, [B, Hq, Lq, Lk]) is the fourth input.

sinks: bool

Whether per-head attention sinks ([Hq]) are the last input.

window: u32

A sliding window of this many keys (0: none; needs causal).

passes: u8

The vector kernels’ passes over the keys (1 or 2; at most 8 queries), or 0 for the tiled kernel. Set by the attention plan.

blocks: u32

Key partitions of the first of two vector passes.

gqa_kernel: bool

The decode-only grouped-query two-pass kernel (one query, no mask or sinks).

tk1: bool

The one-key-per-step variant of the two-pass kernel.

GatedDelta

(q, k, v, g, beta, state): the gated delta rule (see gated_delta), with the outputs and the final state (or every step’s, gated_delta_steps).

Fields

steps: bool

The state after every step ([B, T, Hv, Dv, Dk]) instead of the final one.

GatedDeltaSlots

(q, k, v, g, beta, pool, slots): gated_delta_slots, with the outputs and the updated pool.

GatherMm

(a, b, lhs_indices, rhs_indices[, offsets]): matrix products of gathered batches (see gather_mm), as FastOp::GatherQmm in full precision.

GatherQmm

(x, w, scales[, biases][, global_scale], lhs_indices, rhs_indices[, offsets]): quantized matrix products of gathered batches (see gather_qmm); biases for affine only, a global scale (float32, one per weight batch) for nvfp4 only. The indices are uint32 and broadcast to the output’s batch shape. offsets (int32, one per batch of w) marks where each weight batch’s rows start when the rows are sorted by it.

Fields

quant: Quant

The mode, group size and bits.

transpose: bool

x · wᵀ (else x · w).

global_scale: bool

Whether a global scale is given.

Trait Implementations

impl Clone for FastOp

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for FastOp

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Eq for FastOp

impl Hash for FastOp

fn hash<__H: Hasher>(&self, state: &mut __H)

Feeds this value into the given Hasher.

fn hash_slice<H>(data: &[Self], state: &mut H)

where H: Hasher, Self: Sized,

Feeds a slice of this type into the given Hasher.

impl PartialEq for FastOp

fn eq(&self, other: &Self) -> bool

Equality operator ==.

fn ne(&self, other: &Rhs) -> bool

Inequality operator !=.

impl StructuralPartialEq for FastOp