diff --git a/README.md b/README.md index 5a23861..dba35e4 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,15 @@ The mlx-c submodule sits inside `mlxcore-sys` rather than at the repository root because `cargo package` only ships files under the crate directory — from the root it would be missing from the published crate. +Within `mlxcore`, the modules follow the mlx-c headers they wrap: + +| Module | Wraps | Holds | +| --- | --- | --- | +| `Array` (crate root) | `array.h`, `ops.h` | array construction, arithmetic, reductions, shapes, indexing | +| `mlxcore::fast` | `fast.h` | fused kernels: `layer_norm`, `rms_norm`, `rope`, `scaled_dot_product_attention` | +| `mlxcore::io` | `io.h` | `load_safetensors` / `save_safetensors` | +| `mlxcore::random` | `random.h` | seeded keys, `uniform`, `normal` | + ## Building from source ```sh @@ -67,8 +76,9 @@ The first build compiles MLX from source and takes several minutes. Runnable examples live in `crates/mlxcore/examples`: ```sh -cargo run --example hello # arrays, shapes, streams -cargo run --example relu # y = relu(x @ W + b) with random weights +cargo run --example hello # arrays, shapes, streams +cargo run --example relu # y = relu(x @ W + b) with random weights +cargo run --example attention # a transformer attention block: QKV, RoPE, masked SDPA ``` ## Note: suffix float literals with `f32` diff --git a/crates/mlxcore/examples/attention.rs b/crates/mlxcore/examples/attention.rs new file mode 100644 index 0000000..f6604ab --- /dev/null +++ b/crates/mlxcore/examples/attention.rs @@ -0,0 +1,117 @@ +//! One transformer attention block, the way an encoder model builds it. +//! +//! A fused-QKV projection, rotary position embeddings, sliding-window masked +//! attention, and an output projection — the shape of a ModernBERT layer. This +//! is the composition [`mlxcore::fast`] exists for; everything here is a real op, +//! not a sketch. +//! +//! Run with: +//! ```sh +//! cargo run --example attention +//! ``` + +use mlxcore::fast::{self, Mask}; +use mlxcore::{Array, Result, Stream, random}; + +const BATCH: i32 = 1; +const LENGTH: i32 = 6; +const HIDDEN: i32 = 8; +const HEADS: i32 = 2; +const HEAD_DIM: i32 = HIDDEN / HEADS; + +/// Local attention window, as a total width. Each query sees keys within half of +/// it on either side, which is how ModernBERT keeps most of its layers cheap. +const WINDOW: i32 = 4; + +fn main() -> Result<()> { + // The CPU stream keeps this reproducible and lets float64 through; a real + // model would use `Stream::default()` for the GPU. + let stream = Stream::cpu(); + + let key = random::Key::new(0)?; + let (x_key, w_key) = key.split(&stream)?; + let x = random::normal::(&[BATCH, LENGTH, HIDDEN], 0.0, 1.0, Some(&x_key), &stream)?; + + // Checkpoints store a linear layer's weight as (out_features, in_features) + // and compute `x @ w.T`, so every projection swaps the last two axes. + let (qkv_key, out_key) = w_key.split(&stream)?; + let w_qkv = random::normal::(&[3 * HIDDEN, HIDDEN], 0.0, 0.1, Some(&qkv_key), &stream)?; + let w_out = random::normal::(&[HIDDEN, HIDDEN], 0.0, 0.1, Some(&out_key), &stream)?; + + // --- 1. Fused QKV projection ------------------------------------------ + // One matmul for all three tensors, then split by reshaping the output into + // an explicit axis of 3 rather than slicing it apart. + let qkv = x + .matmul(&w_qkv.swapaxes(-2, -1, &stream)?, &stream)? + .reshape(&[BATCH, LENGTH, 3, HEADS, HEAD_DIM], &stream)?; + + // `take_axis` with a 0-d index drops the axis it selects on, so each of these + // is (batch, length, heads, head_dim). Transposing brings `heads` forward, + // which is the layout the attention kernel requires. + let project = |i: i32| -> Result { + qkv.take_axis(&Array::from_scalar(i), 2, &stream)? + .transpose_axes(&[0, 2, 1, 3], &stream) + }; + let (q, k, v) = (project(0)?, project(1)?, project(2)?); + println!("q/k/v: {:?} (batch, heads, length, head_dim)", q.shape()); + + // --- 2. Rotary position embeddings ------------------------------------ + // Applied to queries and keys only — values carry no position. `false` is the + // split-halves layout every Hugging Face checkpoint uses. + let rope = |a: &Array| fast::rope(a, HEAD_DIM, false, Some(10_000.0), 1.0, 0, None, &stream); + let (q, k) = (rope(&q)?, rope(&k)?); + + // --- 3. The sliding-window mask --------------------------------------- + // `|i - j| <= window / 2`, built by broadcasting a position vector against + // its own transpose. A bool mask drops the false positions. + let positions = Array::arange::(0.0, LENGTH as f64, 1.0, &stream)?; + let distance = positions + .expand_dims(1, &stream)? + .subtract(&positions.expand_dims(0, &stream)?, &stream)? + .abs(&stream)?; + let mask = distance + .less_equal(&Array::from_scalar(WINDOW / 2), &stream)? + // (length, length) -> (batch, heads, length, length) by broadcasting. + .expand_dims(0, &stream)? + .expand_dims(0, &stream)?; + + // --- 4. Attention ------------------------------------------------------ + let scale = 1.0 / (HEAD_DIM as f32).sqrt(); + let attended = + fast::scaled_dot_product_attention(&q, &k, &v, scale, Mask::Array(&mask), &stream)?; + + // --- 5. Merge the heads and project out -------------------------------- + // Undo the transpose from step 1, then collapse (heads, head_dim) back into + // one hidden axis. `contiguous` because `reshape` needs dense input and the + // transpose left this strided. + let merged = attended + .transpose_axes(&[0, 2, 1, 3], &stream)? + .contiguous(&stream)? + .reshape(&[BATCH, LENGTH, HIDDEN], &stream)?; + let y = merged.matmul(&w_out.swapaxes(-2, -1, &stream)?, &stream)?; + + // A residual connection and a layer norm, the rest of the block. The norm has + // no affine parameters here, so both optional operands are `None`. + let normed = fast::layer_norm(&x.add(&y, &stream)?, None, None, 1e-5, &stream)?; + normed.eval(); + + println!("attention out: {:?}", y.shape()); + println!("block out: {:?}", normed.shape()); + + // Every row is normalized, so each one has mean 0 and unit variance. + let means = normed.mean_axes(&[-1], false, &stream)?; + println!("row means: {:?}", means.to_vec::()); + + // Proof the mask did something: position 0 and position 5 are more than + // `WINDOW / 2` apart, so neither attends to the other. Widening the window to + // cover the whole sequence changes the answer. + let full = fast::scaled_dot_product_attention(&q, &k, &v, scale, Mask::None, &stream)?; + let drift = full + .subtract(&attended, &stream)? + .abs(&stream)? + .max(false, &stream)? + .item::(); + println!("windowed vs. full attention, max difference: {drift:.4}"); + + Ok(()) +} diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index a5189d7..f1a9186 100644 --- a/crates/mlxcore/src/array.rs +++ b/crates/mlxcore/src/array.rs @@ -9,6 +9,7 @@ use crate::dtype::{ArrayElement, Dtype}; use crate::error::{self, Result}; use crate::ffi::as_ffi_ptr; use crate::stream::Stream; +use crate::vector::VectorArray; /// An N-dimensional MLX array. /// @@ -297,14 +298,29 @@ impl Array { /// `bool` becomes 0 or 1, and numeric-to-`bool` tests `!= 0`. Converting a /// float that is out of the target integer's range is *unspecified*. /// - /// Only dtypes with an [`ArrayElement`] impl can be targeted, so MLX's - /// `float16`, `bfloat16`, and `complex64` are out of reach for now. + /// Only dtypes with an [`ArrayElement`] impl can be targeted this way; + /// [`astype_dtype`](Self::astype_dtype) covers the rest. pub fn astype(&self, stream: &Stream) -> Result { + self.astype_dtype(T::DTYPE, stream) + } + + /// Converts the elements to `dtype`. + /// + /// The runtime counterpart to [`astype`](Self::astype), and the only way to + /// reach `float16` and `bfloat16` — the dtypes with no Rust equivalent, and + /// the ones most published checkpoints store their weights in. + /// + /// An `Array` only holds a handle, so a `float16` array is fully usable + /// without a Rust `f16` type: arithmetic, `matmul`, `shape`, and `dtype` all + /// work. Only [`to_vec`](Self::to_vec) and [`item`](Self::item) need the + /// element type, and those still reject it — so cast to `float32` before + /// reading results back out. + pub fn astype_dtype(&self, dtype: Dtype, stream: &Stream) -> Result { error::install(); let mut out = unsafe { sys::mlx_array_new() }; // SAFETY: handle/stream are valid; the result is written into `out`. let status = - unsafe { sys::mlx_astype(&mut out, self.handle, T::DTYPE.as_raw(), stream.as_raw()) }; + unsafe { sys::mlx_astype(&mut out, self.handle, dtype.as_raw(), stream.as_raw()) }; Self::from_op(out, status) } @@ -383,6 +399,11 @@ impl Array { self.unary_op(stream, sys::mlx_tanh) } + /// Elementwise Gauss error function. + pub fn erf(&self, stream: &Stream) -> Result { + self.unary_op(stream, sys::mlx_erf) + } + /// Elementwise power: `self ** other`. pub fn power(&self, other: &Array, stream: &Stream) -> Result { self.binary_op(other, stream, sys::mlx_power) @@ -542,6 +563,40 @@ impl Array { self.binary_op(other, stream, sys::mlx_matmul) } + /// Fused `alpha * (a @ b) + beta * c`. + /// + /// The same value as `a.matmul(b)?.add(c)` with `alpha == beta == 1.0`, but + /// in a single kernel — which is what makes it the right primitive for a + /// linear layer with a bias. `c` broadcasts against the product, so a + /// per-output bias vector works directly. + /// + /// An associated function because there is no reason for any one of the + /// three operands to be the receiver. + pub fn addmm( + c: &Array, + a: &Array, + b: &Array, + alpha: f32, + beta: f32, + stream: &Stream, + ) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; the result is written into `out`. + let status = unsafe { + sys::mlx_addmm( + &mut out, + c.as_raw(), + a.as_raw(), + b.as_raw(), + alpha, + beta, + stream.as_raw(), + ) + }; + Self::from_op(out, status) + } + /// Sum of all elements, returning a scalar array. /// /// With `keepdims == false` the result is 0-dimensional. @@ -649,6 +704,59 @@ impl Array { self.reduce_axes_op(axes, keepdims, stream, sys::mlx_any_axes) } + /// Softmax over every element, as one flat distribution. + /// + /// Unlike the reductions above the shape is unchanged — this normalizes + /// rather than reduces. In a model you almost always want + /// [`softmax_axes`](Self::softmax_axes) with the last axis instead. + /// + /// `precise` accumulates in `float32` for half-precision inputs. It costs a + /// little speed and avoids overflow in `exp`; on a `float32` array it makes + /// no difference. + pub fn softmax(&self, precise: bool, stream: &Stream) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: handle/stream are valid; the result is written into `out`. + let status = unsafe { sys::mlx_softmax(&mut out, self.handle, precise, stream.as_raw()) }; + Self::from_op(out, status) + } + + /// Softmax over the given axes, leaving the shape unchanged. + /// + /// `a.softmax_axes(&[-1], true, &stream)` is the usual per-row distribution + /// over logits. See [`softmax`](Self::softmax) for `precise`. + pub fn softmax_axes(&self, axes: &[i32], precise: bool, stream: &Stream) -> Result { + error::install(); + let axes_ptr = as_ffi_ptr(axes); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: `axes_ptr`/`axes.len()` describe a valid slice (or null/0) for + // the duration of the call; handle/stream are valid. + let status = unsafe { + sys::mlx_softmax_axes( + &mut out, + self.handle, + axes_ptr, + axes.len(), + precise, + stream.as_raw(), + ) + }; + Self::from_op(out, status) + } + + /// Every element sorted ascending, as a flat 1-dimensional array. + pub fn sort(&self, stream: &Stream) -> Result { + self.unary_op(stream, sys::mlx_sort) + } + + /// Sorts along `axis`, ascending, keeping the shape. + /// + /// There is no descending option and no `k`-largest shortcut: take the tail + /// with [`slice_axis`](Self::slice_axis), which is what a top-k reads as. + pub fn sort_axis(&self, axis: i32, stream: &Stream) -> Result { + self.axis_op(axis, stream, sys::mlx_sort_axis) + } + /// Returns a new array with the same data reinterpreted as `shape`. /// /// The product of `shape` must equal [`size`](Self::size). @@ -666,11 +774,46 @@ impl Array { self.unary_op(stream, sys::mlx_transpose) } + /// Permutes the axes into the order given by `axes`. + /// + /// `axes` must be a permutation of `0..ndim()`, and names the *source* axis + /// for each output position: `x.transpose_axes(&[0, 2, 1, 3], &s)` on a + /// `(batch, len, heads, head_dim)` array gives `(batch, heads, len, + /// head_dim)`, the layout attention kernels expect. Negative entries count + /// from the end. + pub fn transpose_axes(&self, axes: &[i32], stream: &Stream) -> Result { + self.shape_op(axes, stream, sys::mlx_transpose_axes) + } + + /// Exchanges two axes, leaving the rest in place. + /// + /// `w.swapaxes(-2, -1, &s)` is the transpose a linear layer needs, where + /// weights are stored `(out_features, in_features)` but the product wants + /// `x @ w.T`. Cheaper to read than the full permutation + /// [`transpose_axes`](Self::transpose_axes) would spell out. + pub fn swapaxes(&self, axis1: i32, axis2: i32, stream: &Stream) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: handle/stream are valid; the result is written into `out`. + let status = + unsafe { sys::mlx_swapaxes(&mut out, self.handle, axis1, axis2, stream.as_raw()) }; + Self::from_op(out, status) + } + /// Removes all axes of length 1. pub fn squeeze(&self, stream: &Stream) -> Result { self.unary_op(stream, sys::mlx_squeeze) } + /// Removes the given axes, each of which must have length 1. + /// + /// The targeted counterpart to [`squeeze`](Self::squeeze): dropping a known + /// trailing axis with `a.squeeze_axes(&[-1], &s)` cannot accidentally also + /// collapse a batch dimension that happens to be 1. + pub fn squeeze_axes(&self, axes: &[i32], stream: &Stream) -> Result { + self.shape_op(axes, stream, sys::mlx_squeeze_axes) + } + /// Inserts a new axis of length 1 at position `axis`. pub fn expand_dims(&self, axis: i32, stream: &Stream) -> Result { error::install(); @@ -680,6 +823,249 @@ impl Array { Self::from_op(out, status) } + // Indexing. Rust cannot overload `[]` to return a new `Array` — `Index` must + // hand back a reference to something that already exists — so NumPy-style + // subscripting is spelled out as these methods instead. + + /// Gathers elements at `indices` from the flattened array. + /// + /// `self` is treated as 1-dimensional, so the result takes `indices`'s + /// shape. To index a single axis and keep the others, use + /// [`take_axis`](Self::take_axis). + /// + /// `indices` must be an integer array; negative indices count from the end. + pub fn take(&self, indices: &Array, stream: &Stream) -> Result { + self.binary_op(indices, stream, sys::mlx_take) + } + + /// Gathers slices along `axis` at `indices`. + /// + /// `axis` is replaced by `indices`'s shape, so taking `(2, 3)` indices from + /// a `(10, 4)` array along axis 0 gives `(2, 3, 4)`. A 0-dimensional + /// `indices` therefore *removes* the axis, which is how a single row is + /// selected. + /// + /// This is the embedding lookup: `weight.take_axis(&ids, 0, &s)` turns + /// `(vocab, dims)` weights and `(batch, len)` token ids into + /// `(batch, len, dims)`. + /// + /// Negative `axis` counts from the end. + pub fn take_axis(&self, indices: &Array, axis: i32, stream: &Stream) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; the result is written into `out`. + let status = unsafe { + sys::mlx_take_axis( + &mut out, + self.handle, + indices.as_raw(), + axis, + stream.as_raw(), + ) + }; + Self::from_op(out, status) + } + + /// Gathers one element per index along `axis`, pairing `self` and `indices` + /// elementwise on every other axis. + /// + /// `indices` must have the same rank as `self` and match it on every axis + /// but `axis`; the result has `indices`'s shape. Unlike + /// [`take_axis`](Self::take_axis) the index *varies with position*, which + /// makes this the tool for "row `i` of batch element `i`" gathers that would + /// otherwise need NumPy fancy indexing: broadcast a `(batch, count)` index + /// array to `(batch, count, dims)` and take along axis 1 to pull `count` + /// chosen positions out of each sequence. + pub fn take_along_axis(&self, indices: &Array, axis: i32, stream: &Stream) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; the result is written into `out`. + let status = unsafe { + sys::mlx_take_along_axis( + &mut out, + self.handle, + indices.as_raw(), + axis, + stream.as_raw(), + ) + }; + Self::from_op(out, status) + } + + /// Extracts the strided region between `start` and `stop`. + /// + /// All three slices need one entry per dimension — there is no `:` + /// shorthand, so pass `0` and the axis length for axes you want whole, or + /// use [`slice_axis`](Self::slice_axis) when only one axis is interesting. + /// Bounds follow Python's `a[start:stop:step]`: negative values count from + /// the end and out-of-range values clamp rather than erroring. + pub fn slice( + &self, + start: &[i32], + stop: &[i32], + strides: &[i32], + stream: &Stream, + ) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: each ptr/len pair describes a valid slice (or null/0) for the + // duration of the call; handle/stream are valid. + let status = unsafe { + sys::mlx_slice( + &mut out, + self.handle, + as_ffi_ptr(start), + start.len(), + as_ffi_ptr(stop), + stop.len(), + as_ffi_ptr(strides), + strides.len(), + stream.as_raw(), + ) + }; + Self::from_op(out, status) + } + + /// Slices a single axis, keeping every other axis whole. + /// + /// [`slice`](Self::slice) without spelling out the axes you are not + /// touching. `stop` of `None` runs to the end, so the last two entries of + /// the final axis — a two-way top-k after [`sort_axis`](Self::sort_axis) — + /// are `p.slice_axis(-1, -2, None, &s)`. + /// + /// The axis keeps its dimension even when one element is selected; combine + /// with [`squeeze_axes`](Self::squeeze_axes) to drop it, or use + /// [`take_axis`](Self::take_axis) with a 0-dimensional index. + /// + /// # Errors + /// Returns an error if `axis` is out of range for this array. + pub fn slice_axis( + &self, + axis: i32, + start: i32, + stop: Option, + stream: &Stream, + ) -> Result { + let shape = self.shape(); + let ndim = shape.len() as i32; + let axis = if axis < 0 { axis + ndim } else { axis }; + if axis < 0 || axis >= ndim { + return Err(crate::Error::new(format!( + "axis {axis} out of range for an array with {ndim} dimension(s)" + ))); + } + + // Whole-array bounds, then narrow the one axis. mlx clamps a `stop` + // beyond the axis length, so the dimension itself is a fine "to the end". + let axis = axis as usize; + let starts = { + let mut s = vec![0; shape.len()]; + s[axis] = start; + s + }; + let stops = { + let mut s = shape.clone(); + s[axis] = stop.unwrap_or(shape[axis]); + s + }; + self.slice(&starts, &stops, &vec![1; shape.len()], stream) + } + + /// Selects from `on_true` where `condition` is true and `on_false` elsewhere. + /// + /// All three operands broadcast against each other, so masking a row of + /// logits down to `-inf` is + /// `Array::where_cond(&mask, &logits, &Array::from_scalar(f32::MIN), &s)`. + /// + /// An associated function, not a method: MLX names this after the condition + /// (`mx.where`), and `where` is a Rust keyword. + pub fn where_cond( + condition: &Array, + on_true: &Array, + on_false: &Array, + stream: &Stream, + ) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; the result is written into `out`. + let status = unsafe { + sys::mlx_where( + &mut out, + condition.as_raw(), + on_true.as_raw(), + on_false.as_raw(), + stream.as_raw(), + ) + }; + Self::from_op(out, status) + } + + /// Splits into `parts` equally-sized pieces along `axis`. + /// + /// `self.shape()[axis]` must be divisible by `parts`. This is the gated-MLP + /// idiom: one `(.., 2 * d)` projection split into value and gate halves. + pub fn split(&self, parts: i32, axis: i32, stream: &Stream) -> Result> { + error::install(); + let mut out = VectorArray::new(); + // SAFETY: handle/stream are valid; the pieces are written into `out`, + // which frees them on drop. + let status = + unsafe { sys::mlx_split(out.as_mut_ptr(), self.handle, parts, axis, stream.as_raw()) }; + error::check(status)?; + out.to_arrays() + } + + /// Stacks `arrays` along a *new* axis inserted at `axis`. + /// + /// Every array must have the same shape, and the result gains one dimension + /// of length `arrays.len()`. Compare [`concatenate`](Self::concatenate), + /// which joins along an axis that already exists. + pub fn stack(arrays: &[&Array], axis: i32, stream: &Stream) -> Result { + Self::vector_op(arrays, axis, stream, sys::mlx_stack_axis) + } + + /// Joins `arrays` along the existing `axis`. + /// + /// Shapes must agree on every axis but `axis`, whose lengths add up. + pub fn concatenate(arrays: &[&Array], axis: i32, stream: &Stream) -> Result { + Self::vector_op(arrays, axis, stream, sys::mlx_concatenate_axis) + } + + /// Shared plumbing for `res = op(a, axis, stream)` single-axis ops. + fn axis_op( + &self, + axis: i32, + stream: &Stream, + op: unsafe extern "C" fn(*mut sys::mlx_array, sys::mlx_array, i32, sys::mlx_stream) -> i32, + ) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: handle/stream are valid; `op` writes the result into `out`. + let status = unsafe { op(&mut out, self.handle, axis, stream.as_raw()) }; + Self::from_op(out, status) + } + + /// Shared plumbing for `res = op(arrays, axis, stream)` combining ops. + fn vector_op( + arrays: &[&Array], + axis: i32, + stream: &Stream, + op: unsafe extern "C" fn( + *mut sys::mlx_array, + sys::mlx_vector_array, + i32, + sys::mlx_stream, + ) -> i32, + ) -> Result { + error::install(); + let inputs = VectorArray::from_arrays(arrays); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: `inputs` holds handles to the borrowed arrays and outlives the + // call; the result is written into `out`. + let status = unsafe { op(&mut out, inputs.as_raw(), axis, stream.as_raw()) }; + Self::from_op(out, status) + } + /// Shared plumbing for `res = op(shape, shape_num, dtype, stream)` /// constructors. fn fill_op( @@ -2059,4 +2445,407 @@ mod tests { let b = Array::from_slice(&[1.0f32, 2.0], &[2]); let _ = &a + &b; } + + #[test] + fn astype_dtype_reaches_half_precision() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.5f32, 2.5], &[2]); + + // The dtypes `astype::` cannot name, because Rust has no `f16`. + for dtype in [Dtype::Float16, Dtype::Bfloat16] { + let half = a.astype_dtype(dtype, &s).unwrap(); + assert_eq!(half.dtype(), dtype); + // Still a usable array: shape, arithmetic, and a cast back all work. + assert_eq!(half.shape(), vec![2]); + let doubled = half.add(&half, &s).unwrap(); + assert_eq!(doubled.dtype(), dtype); + assert_eq!( + doubled.astype::(&s).unwrap().to_vec::(), + vec![3.0, 5.0] + ); + } + } + + #[test] + fn erf_is_the_gelu_building_block() { + let s = Stream::cpu(); + let x = Array::from_slice(&[0.0f32, 1.0, -1.0], &[3]); + + let out = x.erf(&s).unwrap().to_vec::(); + + // erf(0) = 0 and erf is odd, with erf(1) = 0.8427. + assert!(out[0].abs() < 1e-6, "{out:?}"); + assert!((out[1] - 0.8427).abs() < 1e-3, "{out:?}"); + assert!((out[2] + 0.8427).abs() < 1e-3, "{out:?}"); + } + + #[test] + fn addmm_matches_matmul_plus_bias() { + let s = Stream::cpu(); + let x = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + let w = Array::from_slice(&[1.0f32, 0.0, 0.0, 1.0], &[2, 2]); + let bias = Array::from_slice(&[10.0f32, 20.0], &[2]); + + let fused = Array::addmm(&bias, &x, &w, 1.0, 1.0, &s).unwrap(); + + // `w` is the identity, so this is `x` plus a broadcast bias row. + assert_eq!(fused.shape(), vec![2, 2]); + assert_eq!(fused.to_vec::(), vec![11.0, 22.0, 13.0, 24.0]); + // And it agrees with spelling the two steps out. + let separate = x.matmul(&w, &s).unwrap().add(&bias, &s).unwrap(); + assert_eq!(fused.to_vec::(), separate.to_vec::()); + } + + #[test] + fn addmm_scales_both_terms() { + let s = Stream::cpu(); + let x = Array::from_slice(&[1.0f32, 1.0, 1.0, 1.0], &[2, 2]); + let c = Array::from_slice(&[1.0f32], &[1]); + + // alpha * (x @ x) + beta * c, with each row of `x @ x` summing to 2. + let out = Array::addmm(&c, &x, &x, 3.0, 10.0, &s).unwrap(); + + assert_eq!(out.to_vec::(), vec![16.0; 4]); + } + + #[test] + fn softmax_axes_normalizes_each_row() { + let s = Stream::cpu(); + // Equal logits within each row, so each row is a uniform distribution — + // and the rows differ, which a whole-array softmax would blur together. + let a = Array::from_slice(&[1.0f32, 1.0, 5.0, 5.0], &[2, 2]); + + let p = a.softmax_axes(&[-1], true, &s).unwrap(); + + assert_eq!(p.shape(), vec![2, 2]); + assert_eq!(p.to_vec::(), vec![0.5, 0.5, 0.5, 0.5]); + } + + #[test] + fn softmax_over_all_elements_sums_to_one() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + + let p = a.softmax(true, &s).unwrap(); + + // Normalizing, not reducing: the shape survives. + assert_eq!(p.shape(), vec![2, 2]); + let total = p.sum(false, &s).unwrap().item::(); + assert!((total - 1.0).abs() < 1e-5, "{total}"); + } + + #[test] + fn sort_axis_orders_rows_independently() { + let s = Stream::cpu(); + let a = Array::from_slice(&[3.0f32, 1.0, 2.0, 9.0, 7.0, 8.0], &[2, 3]); + + let sorted = a.sort_axis(-1, &s).unwrap(); + + assert_eq!(sorted.shape(), vec![2, 3]); + assert_eq!(sorted.to_vec::(), vec![1.0, 2.0, 3.0, 7.0, 8.0, 9.0]); + } + + #[test] + fn sort_flattens_the_whole_array() { + let s = Stream::cpu(); + let a = Array::from_slice(&[3.0f32, 1.0, 2.0, 0.0], &[2, 2]); + + let sorted = a.sort(&s).unwrap(); + + assert_eq!(sorted.shape(), vec![4]); + assert_eq!(sorted.to_vec::(), vec![0.0, 1.0, 2.0, 3.0]); + } + + #[test] + fn sorting_then_slicing_gives_a_top_k() { + let s = Stream::cpu(); + let p = Array::from_slice(&[0.1f32, 0.7, 0.2], &[1, 3]); + + // The idiom the docs point at: ascending sort, then take the tail. + let top2 = p + .sort_axis(-1, &s) + .unwrap() + .slice_axis(-1, -2, None, &s) + .unwrap(); + + assert_eq!(top2.shape(), vec![1, 2]); + assert_eq!(top2.to_vec::(), vec![0.2, 0.7]); + } + + #[test] + fn transpose_axes_permutes_to_attention_layout() { + let s = Stream::cpu(); + // (batch, length, heads, head_dim) -> (batch, heads, length, head_dim). + let a = Array::zeros::(&[2, 8, 4, 16], &s).unwrap(); + + let out = a.transpose_axes(&[0, 2, 1, 3], &s).unwrap(); + + assert_eq!(out.shape(), vec![2, 4, 8, 16]); + } + + #[test] + fn transpose_axes_moves_the_data_too() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]); + + let out = a.transpose_axes(&[1, 0], &s).unwrap(); + + assert_eq!(out.shape(), vec![3, 2]); + assert_eq!(out.to_vec::(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]); + } + + #[test] + fn swapaxes_transposes_the_last_two() { + let s = Stream::cpu(); + let w = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]); + + // The transpose a linear layer needs: (out, in) -> (in, out). + let out = w.swapaxes(-2, -1, &s).unwrap(); + + assert_eq!(out.shape(), vec![3, 2]); + assert_eq!(out.to_vec::(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]); + } + + #[test] + fn squeeze_axes_drops_only_the_named_axis() { + let s = Stream::cpu(); + // A batch of 1 that must survive, and a trailing axis that must not. + let a = Array::zeros::(&[1, 3, 1], &s).unwrap(); + + assert_eq!(a.squeeze_axes(&[-1], &s).unwrap().shape(), vec![1, 3]); + // Where the untargeted `squeeze` would have taken both. + assert_eq!(a.squeeze(&s).unwrap().shape(), vec![3]); + } + + #[test] + fn squeeze_axes_rejects_a_longer_axis() { + let s = Stream::cpu(); + let a = Array::zeros::(&[2, 3], &s).unwrap(); + + assert!(a.squeeze_axes(&[0], &s).is_err()); + } + + #[test] + fn take_gathers_from_the_flattened_array() { + let s = Stream::cpu(); + let a = Array::from_slice(&[10.0f32, 20.0, 30.0, 40.0], &[2, 2]); + let indices = Array::from_slice(&[3i32, 0], &[2]); + + let out = a.take(&indices, &s).unwrap(); + + // Row-major order, so index 3 is the last element. + assert_eq!(out.shape(), vec![2]); + assert_eq!(out.to_vec::(), vec![40.0, 10.0]); + } + + #[test] + fn take_axis_is_an_embedding_lookup() { + let s = Stream::cpu(); + // A (vocab=3, dims=2) embedding table and a (batch=1, len=3) id array. + let weight = Array::from_slice(&[0.0f32, 0.1, 1.0, 1.1, 2.0, 2.1], &[3, 2]); + let ids = Array::from_slice(&[2i32, 0, 2], &[1, 3]); + + let out = weight.take_axis(&ids, 0, &s).unwrap(); + + assert_eq!(out.shape(), vec![1, 3, 2]); + assert_eq!(out.to_vec::(), vec![2.0, 2.1, 0.0, 0.1, 2.0, 2.1]); + } + + #[test] + fn take_axis_with_a_scalar_index_removes_the_axis() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + + // The `h[:, 0]` of a Python model: select one position, drop the axis. + let out = a.take_axis(&Array::from_scalar(0i32), 1, &s).unwrap(); + + assert_eq!(out.shape(), vec![2]); + assert_eq!(out.to_vec::(), vec![1.0, 3.0]); + } + + #[test] + fn take_along_axis_varies_the_index_per_row() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]); + // Row 0 wants column 2, row 1 wants column 0 — something `take_axis` + // cannot express, since its indices apply to every row alike. + let indices = Array::from_slice(&[2i32, 0], &[2, 1]); + + let out = a.take_along_axis(&indices, 1, &s).unwrap(); + + assert_eq!(out.shape(), vec![2, 1]); + assert_eq!(out.to_vec::(), vec![3.0, 4.0]); + } + + #[test] + fn take_along_axis_gathers_marker_positions() { + let s = Stream::cpu(); + // The real use: pull chosen positions out of each sequence of a batch. + // (batch=2, len=3, dims=2), with hidden state `position * 10 + feature`. + let h = Array::from_slice( + &[ + 0.0f32, 1.0, 10.0, 11.0, 20.0, 21.0, // batch 0 + 0.0, 1.0, 10.0, 11.0, 20.0, 21.0, // batch 1 + ], + &[2, 3, 2], + ); + // Batch 0 wants positions (2, 0); batch 1 wants (1, 1). + let positions = Array::from_slice(&[2i32, 0, 1, 1], &[2, 2]); + + let indices = positions + .expand_dims(-1, &s) + .unwrap() + .broadcast_to(&[2, 2, 2], &s) + .unwrap(); + let out = h.take_along_axis(&indices, 1, &s).unwrap(); + + assert_eq!(out.shape(), vec![2, 2, 2]); + assert_eq!( + out.to_vec::(), + vec![ + 20.0, 21.0, 0.0, 1.0, // batch 0: positions 2 then 0 + 10.0, 11.0, 10.0, 11.0, // batch 1: position 1 twice + ] + ); + } + + #[test] + fn slice_extracts_a_strided_region() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]); + + // Every row, every other column. + let out = a.slice(&[0, 0], &[2, 3], &[1, 2], &s).unwrap(); + + assert_eq!(out.shape(), vec![2, 2]); + assert_eq!(out.to_vec::(), vec![1.0, 3.0, 4.0, 6.0]); + } + + #[test] + fn slice_axis_keeps_other_axes_whole() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]); + + // The last two columns, via a negative start and an open end. + let tail = a.slice_axis(-1, -2, None, &s).unwrap(); + assert_eq!(tail.shape(), vec![2, 2]); + assert_eq!(tail.to_vec::(), vec![2.0, 3.0, 5.0, 6.0]); + + // A single row, with the axis retained. + let first = a.slice_axis(0, 0, Some(1), &s).unwrap(); + assert_eq!(first.shape(), vec![1, 3]); + assert_eq!(first.to_vec::(), vec![1.0, 2.0, 3.0]); + } + + #[test] + fn slice_axis_rejects_an_out_of_range_axis() { + let s = Stream::cpu(); + let a = Array::zeros::(&[2, 3], &s).unwrap(); + + let err = a.slice_axis(2, 0, None, &s).unwrap_err(); + + assert!(err.message().contains("out of range"), "{err}"); + } + + #[test] + fn where_cond_selects_elementwise() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3]); + let mask = Array::from_slice(&[true, false, true], &[3]); + + // The masked-logits idiom: keep where true, sink to a floor elsewhere. + let out = Array::where_cond(&mask, &a, &Array::from_scalar(-1e4f32), &s).unwrap(); + + assert_eq!(out.to_vec::(), vec![1.0, -1e4, 3.0]); + } + + #[test] + fn where_cond_broadcasts_all_three_operands() { + let s = Stream::cpu(); + // A (2, 1) condition against (1, 3) branches. + let mask = Array::from_slice(&[true, false], &[2, 1]); + let yes = Array::from_slice(&[1.0f32, 2.0, 3.0], &[1, 3]); + let no = Array::from_scalar(0.0f32); + + let out = Array::where_cond(&mask, &yes, &no, &s).unwrap(); + + assert_eq!(out.shape(), vec![2, 3]); + assert_eq!(out.to_vec::(), vec![1.0, 2.0, 3.0, 0.0, 0.0, 0.0]); + } + + #[test] + fn split_halves_a_gated_projection() { + let s = Stream::cpu(); + // The GLU shape: one (2, 4) projection holding a value and a gate half. + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &[2, 4]); + + let parts = a.split(2, -1, &s).unwrap(); + + assert_eq!(parts.len(), 2); + assert_eq!(parts[0].shape(), vec![2, 2]); + assert_eq!(parts[0].to_vec::(), vec![1.0, 2.0, 5.0, 6.0]); + assert_eq!(parts[1].to_vec::(), vec![3.0, 4.0, 7.0, 8.0]); + } + + #[test] + fn split_rejects_an_uneven_division() { + let s = Stream::cpu(); + let a = Array::zeros::(&[2, 3], &s).unwrap(); + + assert!(a.split(2, -1, &s).is_err()); + } + + #[test] + fn stack_adds_an_axis_and_concatenate_extends_one() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0], &[2]); + let b = Array::from_slice(&[3.0f32, 4.0], &[2]); + + // The feature-vector idiom: separate scalars gathered into one axis. + let stacked = Array::stack(&[&a, &b], -1, &s).unwrap(); + assert_eq!(stacked.shape(), vec![2, 2]); + assert_eq!(stacked.to_vec::(), vec![1.0, 3.0, 2.0, 4.0]); + + // Whereas concatenating reuses the axis that is already there. + let joined = Array::concatenate(&[&a, &b], 0, &s).unwrap(); + assert_eq!(joined.shape(), vec![4]); + assert_eq!(joined.to_vec::(), vec![1.0, 2.0, 3.0, 4.0]); + } + + #[test] + fn concatenate_joins_unequal_lengths() { + let s = Stream::cpu(); + // (1, 2) and (1, 1) along the last axis — the pooled-plus-features + // concatenation a decision head does. + let pooled = Array::from_slice(&[1.0f32, 2.0], &[1, 2]); + let extra = Array::from_slice(&[9.0f32], &[1, 1]); + + let out = Array::concatenate(&[&pooled, &extra], -1, &s).unwrap(); + + assert_eq!(out.shape(), vec![1, 3]); + assert_eq!(out.to_vec::(), vec![1.0, 2.0, 9.0]); + } + + #[test] + fn stack_rejects_mismatched_shapes() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0], &[2]); + let b = Array::from_slice(&[3.0f32], &[1]); + + assert!(Array::stack(&[&a, &b], 0, &s).is_err()); + } + + #[test] + fn combining_ops_accept_a_single_array() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0], &[2]); + + // A one-element vector is a legitimate degenerate case, and exercises + // the borrowed-handle bookkeeping without a second operand. + assert_eq!(Array::stack(&[&a], 0, &s).unwrap().shape(), vec![1, 2]); + assert_eq!( + Array::concatenate(&[&a], 0, &s).unwrap().to_vec::(), + vec![1.0, 2.0] + ); + } } diff --git a/crates/mlxcore/src/fast.rs b/crates/mlxcore/src/fast.rs new file mode 100644 index 0000000..623115e --- /dev/null +++ b/crates/mlxcore/src/fast.rs @@ -0,0 +1,338 @@ +//! Fused kernels from MLX's `fast` namespace. +//! +//! Each of these is expressible with the ops on [`Array`], but MLX ships a +//! single fused kernel that is faster and, for the normalizations, better +//! behaved numerically. They are also exactly the pieces a transformer needs, so +//! a model built on this crate should reach for these rather than composing them +//! by hand. +//! +//! These are stateless kernels: the parameters are arguments, not fields. The +//! layers that own those parameters belong a level up. + +use mlxcore_sys as sys; + +use crate::array::Array; +use crate::error::{self, Result}; +use crate::ffi::absent_array; +use crate::stream::Stream; + +/// How [`scaled_dot_product_attention`] should mask its attention scores. +/// +/// An enum rather than the string-plus-optional-array pair mlx-c takes, so an +/// array mask cannot be requested without supplying one. +pub enum Mask<'a> { + /// Every query attends to every key. + None, + /// Each query attends only to keys at or before its own position. + /// + /// The decoder mask. MLX builds it inside the kernel, so nothing is + /// materialized. + Causal, + /// An explicit mask, broadcast against `(batch, heads, queries, keys)`. + /// + /// A `bool` array masks out the `false` positions; any other dtype is *added* + /// to the scores, which is how an additive bias is applied. Rank must be 4 + /// or less. + /// + /// This is what a bidirectional encoder needs: padding masks, and + /// sliding-window masks that are not simply causal. + Array(&'a Array), +} + +impl Mask<'_> { + /// The `(mask_mode, mask_arr)` pair mlx-c expects. + fn as_ffi(&self) -> (&'static [u8], sys::mlx_array) { + match self { + // MLX validates the mode string and accepts only "", "causal", and + // "array". Written as NUL-terminated byte literals so no `CString` + // allocation is needed to cross the boundary. + Mask::None => (b"\0", absent_array()), + Mask::Causal => (b"causal\0", absent_array()), + Mask::Array(mask) => (b"array\0", mask.as_raw()), + } + } +} + +/// Layer normalization over the last axis. +/// +/// Rescales each row to zero mean and unit variance, then applies the optional +/// affine `weight` and `bias`. Both must be 1-dimensional and match the last +/// axis of `x`; omitting them gives the bare normalization. `eps` is added to +/// the variance before the reciprocal square root, so it must be positive. +pub fn layer_norm( + x: &Array, + weight: Option<&Array>, + bias: Option<&Array>, + eps: f32, + stream: &Stream, +) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; the optional operands are passed as null + // handles, which the C shim turns into `std::nullopt`. + let status = unsafe { + sys::mlx_fast_layer_norm( + &mut out, + x.as_raw(), + weight.map_or_else(absent_array, Array::as_raw), + bias.map_or_else(absent_array, Array::as_raw), + eps, + stream.as_raw(), + ) + }; + Array::from_op(out, status) +} + +/// Root-mean-square normalization over the last axis. +/// +/// Divides each row by its RMS without centering it first — the normalization +/// most recent decoder models use in place of [`layer_norm`]. `weight` is the +/// optional 1-dimensional scale; there is no bias. +pub fn rms_norm(x: &Array, weight: Option<&Array>, eps: f32, stream: &Stream) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; an omitted `weight` is a null handle. + let status = unsafe { + sys::mlx_fast_rms_norm( + &mut out, + x.as_raw(), + weight.map_or_else(absent_array, Array::as_raw), + eps, + stream.as_raw(), + ) + }; + Array::from_op(out, status) +} + +/// Rotary position embedding, applied to the last axis of `x`. +/// +/// Rotates the first `dims` features of each position by an angle that grows +/// with the position, which is how attention learns relative distances without a +/// position embedding table. `dims` must be even and no larger than the last +/// axis; features past it pass through untouched. +/// +/// - `traditional` interleaves the rotated pairs as in the original RoPE paper. +/// Hugging Face checkpoints — and so nearly every published model — use the +/// split-halves layout, which is `false`. +/// - `base` is the geometric base for the frequencies: often `10000.0`, and much +/// larger for long-context models. Pass `None` only together with `freqs`. +/// - `scale` divides the positions, the usual knob for stretching a model past +/// the context it was trained on. `1.0` leaves them alone. +/// - `offset` is the position of `x`'s first element. Nonzero when decoding one +/// token at a time against a cache; `0` for a whole sequence at once. +/// - `freqs` supplies per-dimension frequencies directly, for the scaling +/// schemes `base` alone cannot express. +// Eight parameters, because that is what the kernel takes — MLX's own Python and +// Swift bindings have the same arity. Bundling them into a config struct would +// put a layer between this crate and mlx-c for no gain. +#[allow(clippy::too_many_arguments)] +pub fn rope( + x: &Array, + dims: i32, + traditional: bool, + base: Option, + scale: f32, + offset: i32, + freqs: Option<&Array>, + stream: &Stream, +) -> Result { + error::install(); + let base = sys::mlx_optional_float { + value: base.unwrap_or_default(), + has_value: base.is_some(), + }; + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; an omitted `freqs` is a null handle. + let status = unsafe { + sys::mlx_fast_rope( + &mut out, + x.as_raw(), + dims, + traditional, + base, + scale, + offset, + freqs.map_or_else(absent_array, Array::as_raw), + stream.as_raw(), + ) + }; + Array::from_op(out, status) +} + +/// Scaled dot-product attention: `softmax(scale * q @ k.T + mask) @ v`. +/// +/// All three operands must be rank 4, shaped `(batch, heads, length, head_dim)` +/// — the layout [`Array::transpose_axes`] produces from a +/// `(batch, length, heads, head_dim)` projection. `queries` and `keys` must agree +/// on `head_dim`, and `keys` and `values` on `length`. Grouped-query attention +/// works: `keys` and `values` may carry fewer heads than `queries`, as long as +/// the count divides evenly. +/// +/// `scale` multiplies the scores before the softmax, conventionally +/// `1.0 / (head_dim as f32).sqrt()`. +/// +/// The result is `(batch, heads, length, head_dim)`, so it needs transposing +/// back before the output projection. +/// +/// MLX's attention-sink parameter is not exposed. +pub fn scaled_dot_product_attention( + queries: &Array, + keys: &Array, + values: &Array, + scale: f32, + mask: Mask<'_>, + stream: &Stream, +) -> Result { + error::install(); + let (mode, mask_array) = mask.as_ffi(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; `mode` is a NUL-terminated literal with + // 'static lifetime, and `mask_array` is null unless `Mask::Array` was given. + let status = unsafe { + sys::mlx_fast_scaled_dot_product_attention( + &mut out, + queries.as_raw(), + keys.as_raw(), + values.as_raw(), + scale, + mode.as_ptr().cast::(), + mask_array, + absent_array(), + stream.as_raw(), + ) + }; + Array::from_op(out, status) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn layer_norm_centers_and_scales() { + let s = Stream::cpu(); + let x = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + + let y = layer_norm(&x, None, None, 1e-5, &s).unwrap(); + + // Each row becomes (-1, 1): mean 0, unit variance, up to `eps`. + let out = y.to_vec::(); + assert_eq!(out.len(), 4); + for pair in out.chunks(2) { + assert!((pair[0] + 1.0).abs() < 1e-3, "{pair:?}"); + assert!((pair[1] - 1.0).abs() < 1e-3, "{pair:?}"); + } + } + + #[test] + fn layer_norm_applies_affine() { + let s = Stream::cpu(); + let x = Array::from_slice(&[1.0f32, 2.0], &[1, 2]); + let weight = Array::from_slice(&[2.0f32, 2.0], &[2]); + let bias = Array::from_slice(&[1.0f32, 1.0], &[2]); + + let y = layer_norm(&x, Some(&weight), Some(&bias), 1e-5, &s).unwrap(); + + // (-1, 1) * 2 + 1. + let out = y.to_vec::(); + assert!((out[0] + 1.0).abs() < 1e-3, "{out:?}"); + assert!((out[1] - 3.0).abs() < 1e-3, "{out:?}"); + } + + #[test] + fn rms_norm_divides_by_root_mean_square() { + let s = Stream::cpu(); + // The RMS of (3, 4) is sqrt((9 + 16) / 2) = 3.5355, so the row scales to + // (0.8485, 1.1314). A `layer_norm` would have centered it to (-1, 1). + let x = Array::from_slice(&[3.0f32, 4.0], &[1, 2]); + + let y = rms_norm(&x, None, 1e-5, &s).unwrap(); + + let out = y.to_vec::(); + assert!((out[0] - 0.8485).abs() < 1e-3, "{out:?}"); + assert!((out[1] - 1.1314).abs() < 1e-3, "{out:?}"); + } + + #[test] + fn rope_leaves_position_zero_alone() { + let s = Stream::cpu(); + // (batch, heads, length, head_dim). The rotation angle at position 0 is + // zero, so the first position passes through while the second does not. + let x = Array::ones::(&[1, 1, 2, 4], &s).unwrap(); + + let y = rope(&x, 4, false, Some(10000.0), 1.0, 0, None, &s).unwrap(); + + assert_eq!(y.shape(), vec![1, 1, 2, 4]); + let out = y.to_vec::(); + assert_eq!(&out[..4], &[1.0; 4]); + assert!(out[4..].iter().any(|v| (v - 1.0).abs() > 1e-3), "{out:?}"); + } + + #[test] + fn rope_offset_shifts_positions() { + let s = Stream::cpu(); + let x = Array::ones::(&[1, 1, 2, 4], &s).unwrap(); + + // With `offset = 1` the first element is position 1, not 0, so nothing is + // left unrotated. + let y = rope(&x, 4, false, Some(10000.0), 1.0, 1, None, &s).unwrap(); + + let out = y.to_vec::(); + assert!(out[..4].iter().any(|v| (v - 1.0).abs() > 1e-3), "{out:?}"); + } + + #[test] + fn attention_with_one_key_returns_that_value() { + let s = Stream::cpu(); + // A single key means a softmax over one score, which is 1 whatever the + // scale, so the output is exactly `values`. + let q = Array::ones::(&[1, 1, 1, 4], &s).unwrap(); + let k = Array::ones::(&[1, 1, 1, 4], &s).unwrap(); + let v = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[1, 1, 1, 4]); + + let out = scaled_dot_product_attention(&q, &k, &v, 0.5, Mask::None, &s).unwrap(); + + assert_eq!(out.shape(), vec![1, 1, 1, 4]); + assert_eq!(out.to_vec::(), vec![1.0, 2.0, 3.0, 4.0]); + } + + #[test] + fn attention_mask_excludes_keys() { + let s = Stream::cpu(); + let q = Array::ones::(&[1, 1, 1, 2], &s).unwrap(); + let k = Array::ones::(&[1, 1, 2, 2], &s).unwrap(); + let v = Array::from_slice(&[1.0f32, 1.0, 5.0, 5.0], &[1, 1, 2, 2]); + + // Unmasked, the two identical keys score the same and the output is the + // average of the values. Masking the second key leaves only the first. + let even = scaled_dot_product_attention(&q, &k, &v, 0.5, Mask::None, &s).unwrap(); + assert_eq!(even.to_vec::(), vec![3.0, 3.0]); + + let mask = Array::from_slice(&[true, false], &[1, 1, 1, 2]); + let masked = scaled_dot_product_attention(&q, &k, &v, 0.5, Mask::Array(&mask), &s).unwrap(); + assert_eq!(masked.to_vec::(), vec![1.0, 1.0]); + } + + #[test] + fn causal_attention_hides_later_keys() { + let s = Stream::cpu(); + let q = Array::ones::(&[1, 1, 2, 2], &s).unwrap(); + let k = Array::ones::(&[1, 1, 2, 2], &s).unwrap(); + let v = Array::from_slice(&[1.0f32, 1.0, 5.0, 5.0], &[1, 1, 2, 2]); + + let out = scaled_dot_product_attention(&q, &k, &v, 0.5, Mask::Causal, &s).unwrap(); + + // Query 0 sees only value 0; query 1 sees both and averages them. + assert_eq!(out.to_vec::(), vec![1.0, 1.0, 3.0, 3.0]); + } + + #[test] + fn attention_rejects_low_rank_inputs() { + let s = Stream::cpu(); + let a = Array::ones::(&[2, 2], &s).unwrap(); + + let err = scaled_dot_product_attention(&a, &a, &a, 1.0, Mask::None, &s).unwrap_err(); + + assert!(err.message().contains("rank 4"), "{err}"); + } +} diff --git a/crates/mlxcore/src/io.rs b/crates/mlxcore/src/io.rs new file mode 100644 index 0000000..8ac4800 --- /dev/null +++ b/crates/mlxcore/src/io.rs @@ -0,0 +1,322 @@ +//! Reading and writing arrays on disk, in the `safetensors` format. +//! +//! This is how a model's weights come in: one file mapping parameter names to +//! arrays. The arrays load lazily like any other MLX result — nothing touches +//! the disk until they are evaluated or used in an operation. + +use std::collections::HashMap; +use std::ffi::{CStr, CString, c_char}; +use std::os::unix::ffi::OsStrExt; +use std::path::Path; +use std::ptr; + +use mlxcore_sys as sys; + +use crate::array::Array; +use crate::error::{self, Error, Result}; +use crate::stream::Stream; + +/// Loads every tensor in a `safetensors` file, keyed by name. +/// +/// Arrays keep whatever dtype the file stores, which for a published checkpoint +/// is usually `float16` or `bfloat16`. Neither has a Rust element type, so +/// [`Array::to_vec`] cannot read them directly — use +/// [`Array::astype_dtype`](Array::astype_dtype) to pick the dtype you want to +/// compute in. +/// +/// The file's `__metadata__` header is not returned. +/// +/// # Errors +/// Returns an error if the file is missing, is not a valid `safetensors` +/// container, or holds a dtype MLX does not support. +pub fn load_safetensors(path: impl AsRef, stream: &Stream) -> Result> { + error::install(); + let path = path.as_ref(); + let file = c_path(path)?; + + let mut arrays = ArrayMap::new(); + // mlx-c returns the metadata whether or not we want it, so it needs a + // destination to be freed from. + let mut metadata = StringMap::new(); + // SAFETY: both maps are freshly allocated and outlive the call; `file` is a + // NUL-terminated path; the results are written into the two maps. + let status = unsafe { + sys::mlx_load_safetensors( + arrays.as_mut_ptr(), + metadata.as_mut_ptr(), + file.as_ptr(), + stream.as_raw(), + ) + }; + error::check(status)?; + arrays.to_hash_map() +} + +/// Writes `arrays` to a `safetensors` file, keyed by name. +/// +/// The arrays are evaluated as part of writing, so a lazy graph does not need +/// [`Array::eval`] first. An existing file at `path` is overwritten. No +/// `__metadata__` header is written. +/// +/// # Errors +/// Returns an error if the path cannot be opened for writing. +pub fn save_safetensors(path: impl AsRef, arrays: &HashMap) -> Result<()> { + error::install(); + let path = path.as_ref(); + let file = c_path(path)?; + + let mut map = ArrayMap::new(); + for (name, array) in arrays { + map.insert(name, array)?; + } + let metadata = StringMap::new(); + // SAFETY: `file` is NUL-terminated; both containers outlive the call. + let status = + unsafe { sys::mlx_save_safetensors(file.as_ptr(), map.as_raw(), metadata.as_raw()) }; + error::check(status) +} + +/// Converts a path into the NUL-terminated string mlx-c takes. +/// +/// Paths are bytes on macOS, so this does not go through `str` and non-UTF-8 +/// names work. Only an interior NUL is rejected, which no real path has. +fn c_path(path: &Path) -> Result { + CString::new(path.as_os_str().as_bytes()).map_err(|_| { + Error::new(format!( + "path contains an interior NUL byte: {}", + path.display() + )) + }) +} + +/// An owned `mlx_map_string_to_array`. +struct ArrayMap { + handle: sys::mlx_map_string_to_array, +} + +impl ArrayMap { + fn new() -> Self { + // SAFETY: allocates a fresh empty map, whose ownership moves into `self`. + Self { + handle: unsafe { sys::mlx_map_string_to_array_new() }, + } + } + + fn as_raw(&self) -> sys::mlx_map_string_to_array { + self.handle + } + + fn as_mut_ptr(&mut self) -> *mut sys::mlx_map_string_to_array { + &mut self.handle + } + + /// Adds `array` under `name`, copying the handle in. + fn insert(&mut self, name: &str, array: &Array) -> Result<()> { + let key = CString::new(name) + .map_err(|_| Error::new(format!("tensor name contains a NUL byte: {name:?}")))?; + // SAFETY: `key` is NUL-terminated and outlives the call; the map stores a + // copy of the handle, sharing the buffer rather than taking it over, so + // `array` stays the owner. + let status = unsafe { + sys::mlx_map_string_to_array_insert(self.handle, key.as_ptr(), array.as_raw()) + }; + error::check(status) + } + + /// Copies every entry out into an owned Rust map. + fn to_hash_map(&self) -> Result> { + let iter = ArrayMapIter::new(self.handle); + let mut out = HashMap::new(); + loop { + let mut key: *const c_char = ptr::null(); + let mut value = unsafe { sys::mlx_array_new() }; + // SAFETY: the iterator and map are alive; on success mlx writes a + // borrowed key pointer and a handle sharing the value's buffer. + let status = unsafe { + sys::mlx_map_string_to_array_iterator_next(&mut key, &mut value, iter.handle) + }; + + if status != 0 { + // SAFETY: `value` was allocated above and mlx did not write to + // it, so it is ours to free. + unsafe { sys::mlx_array_free(value) }; + // mlx-c signals "past the last entry" with 2, reserving 1 for a + // real failure — so this is not something `error::check` can be + // handed directly. + if status == 2 { + return Ok(out); + } + error::check(status)?; + } + + if key.is_null() { + unsafe { sys::mlx_array_free(value) }; + return Err(Error::new("safetensors entry has no name")); + } + // SAFETY: mlx points `key` at the map's own NUL-terminated string, + // valid while the map lives; the name is copied out here. + let name = unsafe { CStr::from_ptr(key) } + .to_string_lossy() + .into_owned(); + // SAFETY: `value` holds a handle mlx wrote and nothing else owns. + out.insert(name, unsafe { Array::from_raw(value) }); + } + } +} + +impl Drop for ArrayMap { + fn drop(&mut self) { + // SAFETY: `handle` was created by mlx and is owned solely by `self`. + unsafe { sys::mlx_map_string_to_array_free(self.handle) }; + } +} + +/// An owned iterator over an [`ArrayMap`]. +/// +/// Borrows the map it walks; mlx keeps a raw pointer to the underlying container, +/// so the map must outlive this. +struct ArrayMapIter { + handle: sys::mlx_map_string_to_array_iterator, +} + +impl ArrayMapIter { + fn new(map: sys::mlx_map_string_to_array) -> Self { + // SAFETY: `map` is a valid handle; the iterator it returns is owned here. + Self { + handle: unsafe { sys::mlx_map_string_to_array_iterator_new(map) }, + } + } +} + +impl Drop for ArrayMapIter { + fn drop(&mut self) { + // SAFETY: `handle` was created by mlx and is owned solely by `self`. + unsafe { sys::mlx_map_string_to_array_iterator_free(self.handle) }; + } +} + +/// An owned `mlx_map_string_to_string`, used only to hold the metadata header. +struct StringMap { + handle: sys::mlx_map_string_to_string, +} + +impl StringMap { + fn new() -> Self { + // SAFETY: allocates a fresh empty map, whose ownership moves into `self`. + Self { + handle: unsafe { sys::mlx_map_string_to_string_new() }, + } + } + + fn as_raw(&self) -> sys::mlx_map_string_to_string { + self.handle + } + + fn as_mut_ptr(&mut self) -> *mut sys::mlx_map_string_to_string { + &mut self.handle + } +} + +impl Drop for StringMap { + fn drop(&mut self) { + // SAFETY: `handle` was created by mlx and is owned solely by `self`. + unsafe { sys::mlx_map_string_to_string_free(self.handle) }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Dtype; + + /// A unique path under the temp directory, removed when the guard drops. + struct TempFile { + path: std::path::PathBuf, + } + + impl TempFile { + fn new(name: &str) -> Self { + // Thread id keeps concurrent test binaries from colliding; the tests + // themselves run serialized (see the Makefile). + let unique = format!("mlxcore-{}-{name}.safetensors", std::process::id()); + Self { + path: std::env::temp_dir().join(unique), + } + } + } + + impl Drop for TempFile { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.path); + } + } + + #[test] + fn round_trips_named_arrays() { + let s = Stream::cpu(); + let file = TempFile::new("round-trip"); + + let mut weights = HashMap::new(); + weights.insert( + "encoder.weight".to_string(), + Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]), + ); + weights.insert( + "encoder.bias".to_string(), + Array::from_slice(&[5.0f32], &[1]), + ); + save_safetensors(&file.path, &weights).unwrap(); + + let loaded = load_safetensors(&file.path, &s).unwrap(); + + assert_eq!(loaded.len(), 2); + let weight = &loaded["encoder.weight"]; + assert_eq!(weight.shape(), vec![2, 2]); + assert_eq!(weight.to_vec::(), vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(loaded["encoder.bias"].to_vec::(), vec![5.0]); + } + + #[test] + fn preserves_half_precision_dtype() { + let s = Stream::cpu(); + let file = TempFile::new("half"); + + // The dtype checkpoints actually ship in. There is no Rust `f16`, so the + // array is built in f32 and cast — and must come back as float16. + let half = Array::from_slice(&[1.0f32, 2.0], &[2]) + .astype_dtype(Dtype::Float16, &s) + .unwrap(); + let mut weights = HashMap::new(); + weights.insert("half".to_string(), half); + save_safetensors(&file.path, &weights).unwrap(); + + let loaded = load_safetensors(&file.path, &s).unwrap(); + + let array = &loaded["half"]; + assert_eq!(array.dtype(), Dtype::Float16); + // Unreadable as f16, but castable back to something Rust can see. + let recovered = array.astype::(&s).unwrap(); + assert_eq!(recovered.to_vec::(), vec![1.0, 2.0]); + } + + #[test] + fn empty_map_round_trips() { + let s = Stream::cpu(); + let file = TempFile::new("empty"); + + save_safetensors(&file.path, &HashMap::new()).unwrap(); + + // Exercises the iterator's very first `next` returning end-of-iteration. + assert!(load_safetensors(&file.path, &s).unwrap().is_empty()); + } + + #[test] + fn missing_file_is_an_error() { + let s = Stream::cpu(); + let missing = std::env::temp_dir().join("mlxcore-does-not-exist.safetensors"); + + let err = load_safetensors(&missing, &s).unwrap_err(); + + assert!(!err.message().is_empty()); + } +} diff --git a/crates/mlxcore/src/lib.rs b/crates/mlxcore/src/lib.rs index a0d0713..b26dec7 100644 --- a/crates/mlxcore/src/lib.rs +++ b/crates/mlxcore/src/lib.rs @@ -18,7 +18,10 @@ mod dtype; mod error; mod ffi; mod stream; +mod vector; +pub mod fast; +pub mod io; pub mod random; pub use array::Array; diff --git a/crates/mlxcore/src/vector.rs b/crates/mlxcore/src/vector.rs new file mode 100644 index 0000000..43ff3d0 --- /dev/null +++ b/crates/mlxcore/src/vector.rs @@ -0,0 +1,81 @@ +//! An owned `mlx_vector_array`, the container mlx-c uses for multi-array +//! arguments and results. + +use mlxcore_sys as sys; + +use crate::array::Array; +use crate::error::{self, Result}; + +/// A vector of arrays, owning its `mlx_vector_array` handle. +/// +/// A handful of mlx-c entry points speak in vectors rather than single arrays: +/// [`Array::split`] returns one, [`Array::stack`] and [`Array::concatenate`] +/// take one. This is a thin RAII holder so those call sites cannot leak the +/// handle on an early return. +/// +/// Internal on purpose — the public API takes `&[&Array]` and returns +/// `Vec`, so this type never appears in a signature. +pub(crate) struct VectorArray { + handle: sys::mlx_vector_array, +} + +impl VectorArray { + /// An empty vector. + pub(crate) fn new() -> Self { + error::install(); + // SAFETY: allocates a fresh empty vector, whose ownership moves into `self`. + Self { + handle: unsafe { sys::mlx_vector_array_new() }, + } + } + + /// A vector referring to the same arrays as `arrays`. + /// + /// mlx-c pushes a *copy* of each `mlx_array` handle, and the copy shares the + /// underlying buffer rather than taking it over. The borrowed `Array`s stay + /// the owners, so this vector can be dropped on its own. + pub(crate) fn from_arrays(arrays: &[&Array]) -> Self { + let vec = Self::new(); + for array in arrays { + // SAFETY: both handles are valid; `append_value` copies the handle in. + unsafe { sys::mlx_vector_array_append_value(vec.handle, array.as_raw()) }; + } + vec + } + + /// Returns the raw handle. The `VectorArray` retains ownership. + pub(crate) fn as_raw(&self) -> sys::mlx_vector_array { + self.handle + } + + /// A pointer for the `res` out-parameter of an mlx-c call. + pub(crate) fn as_mut_ptr(&mut self) -> *mut sys::mlx_vector_array { + &mut self.handle + } + + /// Number of arrays held. + pub(crate) fn len(&self) -> usize { + // SAFETY: handle is valid for the lifetime of `self`. + unsafe { sys::mlx_vector_array_size(self.handle) } + } + + /// Copies every element out as an owned [`Array`]. + pub(crate) fn to_arrays(&self) -> Result> { + (0..self.len()) + .map(|i| { + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: `i < len`; `get` writes into `out` a handle sharing the + // element's buffer, which `from_op` then owns. + let status = unsafe { sys::mlx_vector_array_get(&mut out, self.handle, i) }; + Array::from_op(out, status) + }) + .collect() + } +} + +impl Drop for VectorArray { + fn drop(&mut self) { + // SAFETY: `handle` was created by mlx and is owned solely by `self`. + unsafe { sys::mlx_vector_array_free(self.handle) }; + } +}