Stable wrapper for causal self attention with an owned KV cache.

This is the public semantic API for decode paths like Qwen and Llama that need compact GQA KV heads, fused cache update, valid prefix slicing, and causal SDPA through EMLX's native MLX kernels.

Inputs use Bumblebee convention:

• query — {B, T_q, N_q, D}
• new_key / new_value — {B, T_new, N_kv, D}
• key_cache / value_cache — {B, T_max, N_kv, D}
• offset — integer count of positions already present in the cache

Options:

• :scale — required float, usually 1 / sqrt(head_dim)
• :key_mask — optional {B, offset + T_new} mask with 1 = attend and 0 = skip padding

Returns {attention, updated_key_cache, updated_value_cache} as raw EMLX {device, ref} pairs. Keeping the caches in this representation lets callers store and reuse device resident cache buffers without converting them back to Nx.Tensor between decode steps.

Clears the MLX memory cache, releasing unused GPU memory back to the system.

Useful after inference batches to prevent memory growth. Does not affect tensors that are still referenced.

Examples

EMLX.clear_cache()
#=> :ok

Returns the default MLX device for this process.

Reads :default_device from the :emlx application environment, falling back to :gpu. Override in tests or config via:

Application.put_env(:emlx, :default_device, :cpu)

Dequantize a quantized Nx.Tensor (created by EMLX.quantize/2) to a dense float tensor by calling mx::dequantize. Supports every mode EMLX.quantize/2 accepts ("affine" and the microscaled variants).

Dequantizes packed weights to floating point.

Converts quantized weights back to their original floating point representation. Useful for debugging and verification.

Parameters

• w - Quantized weights as uint32 (packed int4 values)
• scales - Per-group scale factors (or u8 for microscaled modes)
• biases - Per-group zero points, or nil for microscaled modes
• group_size - Number of weights per group (default: 64)
• bits - Quantization bits (default: 4)
• mode - "affine" (default), "mxfp4", "mxfp8", or "nvfp4"

Evaluates a (possibly lazy) MLX tensor by routing the work through an EMLX.CommandQueue. Blocks the caller until the worker thread has finished mlx::core::eval/1 for this tensor.

Resolves the queue via resolve_worker/1:

1. If the calling process has bound a queue with EMLX.CommandQueue.with_queue/2, that queue is used.
2. Otherwise the application-default worker for the tensor's device (CPU or GPU) is used — see EMLX.Application.

Fused layer normalisation (mlx::fast::layer_norm).

• x — input tensor; normalised over the last axis.
• weight — {hidden} scale vector (gamma).
• bias — {hidden} bias vector (beta).
• eps — numerical stability constant (e.g. 1.0e-5).

Prefer EMLX.Fast.layer_norm/4 inside defn.

Fused layer normalisation without bias (mlx::fast::layer_norm, weight-only variant).

Prefer EMLX.Fast.layer_norm/3 inside defn.

Fused RMS normalisation (mlx::fast::rms_norm).

Single Metal shader. Normalises over the last axis of x and scales by weight. Output shape and type match x.

Prefer EMLX.Fast.rms_norm/3 inside defn; call this directly only from eager (non-defn) code.

Fused rotary position embedding (mlx::fast::rope).

Single Metal shader. Applies RoPE with a scalar position offset.

• a — input {B, ..., T, D}; T must be the second-to-last axis — e.g. {B, T, D} or, with a heads axis, {B, H, T, D} (heads-transposed). If H (or any other middle axis) is placed before T instead — e.g. {B, T, H, D} — mlx::core::fast::rope silently rotates row h at angle position + h instead of position for every h > 0 (only row 0 is correct); see elixir-nx/emlx#121. For Bumblebee's heads-not-yet-transposed {B, T, H, D} convention, use EMLX.fast_rope_ids/6 or EMLX.fast_rope_positions/6 instead, which guard against this.
• dims — number of feature dims to rotate (≤ last-axis size, must be even)
• traditional — false for split-half (Qwen3); true for interleaved
• base — angular frequency base (e.g. 10_000 or 1_000_000)
• scale — position scale (1.0 unless using NTK-aware scaling)
• offset — integer token position (length of KV cache already filled)

Prefer EMLX.Fast.rope/6 inside defn.

Fused RoPE with a per-batch offset array (mlx::fast::rope, array-offset overload).

Calls the const array& offset overload of mlx::fast::rope, where offset has shape {B} — one starting position per batch example. Positions within each example are assumed to be sequential: [offset[b], offset[b]+1, ..., offset[b]+T-1].

Typically you build offset by slicing position_ids[:, 0] (first token's position for each batch example) before calling this function.

Prefer EMLX.Fast.rope_with_positions/6 inside defn.

Fused RoPE for arbitrary per-token position_ids.

Uses full {B, T} position IDs (not just an offset) and mirrors Bumblebee's per-token cos/sin lookup, avoiding the sequential-offset assumption of fast_rope_ids.

• a — input tensor {B, T, H, D}.
• dims — number of feature dims to rotate.
• traditional — currently only false is supported.
• base — angular frequency base (e.g. 1_000_000).
• scale — position scale.
• position_ids — {B, T} integer tensor.

Fused RoPE with precomputed inv-frequency vector (mlx::fast::rope, freqs overload).

• a — input tensor; shape {B, T, ..., D}.
• dims — number of feature dims to rotate.
• traditional— false for split-half (Qwen3/Bumblebee); true for interleaved.
• scale — position scale (typically 1.0 when using precomputed freqs).
• offset — {B} per-batch starting position tensor.
• freqs — {dims/2} precomputed inverse-frequency tensor.

Prefer EMLX.Fast.rope_with_freqs/6 inside defn.

Flash-attention style SDPA, no mask (mlx::fast::scaled_dot_product_attention).

GQA-native: k/v may have fewer heads than q — no pre-tiling needed.

• q — {B, N_q, T_q, D}
• k — {B, N_kv, T_kv, D}
• v — {B, N_kv, T_kv, D}
• scale — scalar (typically 1 / sqrt(D))
• sinks — optional learned per-head attention-sink logits tensor, or nil

Prefer EMLX.Fast.scaled_dot_product_attention/4 inside defn.

Flash-attention SDPA with built-in causal mask (mlx::fast::scaled_dot_product_attention, mask_mode="causal").

MLX constructs the upper-triangular causal mask internally — no explicit mask tensor required. GQA-native: k/v may have fewer heads than q.

• q — {B, N_q, T_q, D}
• k — {B, N_kv, T_kv, D}
• v — {B, N_kv, T_kv, D}
• scale — scalar (typically 1 / sqrt(D))

sinks — optional learned per-head attention-sink logits tensor, or nil.

Prefer EMLX.Fast.scaled_dot_product_attention_causal/4 inside defn.

Causal SDPA with a runtime key_mask check performed at the C++ level.

At the NIF level: evaluates all(key_mask == 1) (cheap for small/constant tensors). If true → uses the pure causal Metal kernel (no mask allocation). If false → builds a combined causal + key_mask additive float mask and calls the masked kernel.

• q — {B, N_q, T_q, D}
• k — {B, N_kv, T_kv, D}
• v — {B, N_kv, T_kv, D}
• scale — scalar (typically 1 / sqrt(D))
• key_mask — {B, T_kv} boolean/int tensor (1 = attend, 0 = padding)
• sinks — optional learned per-head attention-sink logits tensor, or nil

Prefer EMLX.Fast.scaled_dot_product_attention_causal_key_masked/5 inside defn.

Flash-attention SDPA with an additive or boolean mask.

mask must be broadcast-compatible with {B, N_q, T_q, T_kv}. Boolean false entries are treated as -∞.

sinks — optional learned per-head attention-sink logits tensor, or nil.

Prefer EMLX.Fast.scaled_dot_product_attention/5 inside defn.

Fused SwiGLU activation: silu(gate) * up where silu(x) = x * sigmoid(x).

• gate — gate tensor; silu is applied element-wise.
• up — up-projection tensor; same shape as gate.

Output has the same shape and dtype as gate.

Prefer EMLX.Fast.swiglu/2 inside defn.

Returns the scalar value of a 0-d tensor as a number.

Worker-routed: the NIF body calls mlx::core::eval(*t) and t->item<T>(), both of which require running on the OS thread that owns the tensor's stream encoder.

Fused KV cache update + variable-length SDPA in a single Metal command buffer.

Receives tensors in Bumblebee {B, T, N, D} convention. Internally transposes to MLX {B, N, T, D} for mlx::fast::scaled_dot_product_attention, then transposes the result back. Returns a 3-tuple of EMLX {device, ref} pairs.

• q — {B, T_q, N_q, D} post-RoPE query
• new_k — {B, T_new, N_kv, D} current key projection (post-RoPE)
• new_v — {B, T_new, N_kv, D} current value projection
• k_cache — {B, T_max, N_kv, D} preallocated key buffer
• v_cache — {B, T_max, N_kv, D} preallocated value buffer
• offset — integer, number of positions already in cache
• scale — float, 1 / sqrt(head_dim)

Returns {{dev, attn_ref}, {dev, k_upd_ref}, {dev, v_upd_ref}}.

Like kv_cache_attention/7 but also applies key_mask to exclude padding positions from attention in both prefill and decode steps.

• key_mask — {B, T_kv} integer or boolean tensor with 1 = attend, 0 = skip (padding). Must cover exactly valid_len = offset + T_new positions.

The combined additive mask applies causal AND key_mask constraints without calling mlx::core::all(), avoiding Metal sort-kernel compilation issues for small tensor shapes.

Returns {{dev, attn_ref}, {dev, k_upd_ref}, {dev, v_upd_ref}}.

Fused KV cache update + SDPA for the native NIF loop (BNHD layout).

Accepts q, new_k, new_v already transposed to {B, N, T, D} and a pre-allocated cache in {B, N_kv, T_max, D} layout.

Internally, the cache arrays are move-extracted from their ENIF resources before slice_update so that MLX's donation optimisation fires at eval time: the existing Metal buffer is reused in-place — no new allocation.

Returns {{dev, attn_ref}, {dev, k_upd_ref}, {dev, v_upd_ref}}.

Returns a map with current memory usage information.

Keys:

• :active_memory - bytes currently allocated and in use
• :peak_memory - highest active memory since last reset
• :cache_memory - bytes in the allocator cache (freed but not returned to OS)

Examples

iex> info = EMLX.memory_info()
iex> is_integer(info.active_memory) and is_integer(info.peak_memory) and is_integer(info.cache_memory)
true

Starts a Metal GPU frame capture, writing the trace to path (a .gputrace bundle openable in Xcode's GPU Debugger via File → Open).

Requires MTL_CAPTURE_ENABLED=1 to be set in the process environment before the BEAM starts — Apple's Metal framework reads this gate once at process startup (macOS 14+), and setting it later (e.g. from within Elixir) has no effect. On MLX 0.31.2, forgetting to set it raises a clear EMLX.NIFError ("[metal::start_capture] Failed to start: Capture layer is not inserted.") rather than silently no-op'ing — confirmed by direct test, contradicting this NIF's original design assumption that start_capture/stop_capture never throw on ordinary misconfiguration. Still worth confirming path exists and is non-empty after metal_stop_capture/0 as a belt-and-suspenders check, since a future MLX version could change this behavior back to a silent no-op.

Examples

# export MTL_CAPTURE_ENABLED=1   (before starting the BEAM)
EMLX.metal_start_capture("/tmp/trace.gputrace")
# ... run some MLX work ...
EMLX.metal_stop_capture()

Stops a Metal GPU frame capture started with metal_start_capture/1.

See metal_start_capture/1 for the MTL_CAPTURE_ENABLED precondition.

Quantize a dense 2-D Nx.Tensor and return an annotated quantized tensor.

The returned tensor carries the original logical shape and type (e.g. {:s, 4}). Its backend stores the packed uint32 data and a EMLX.Quantization.Config with scales, biases, group_size, and bits.

Options

• :type — storage type: {:s, 2}, {:s, 4} (default), or {:s, 8}.
• :group_size — 32, 64, or 128 (default 64). Must evenly divide the last dimension of tensor. Microscaled modes pin this to a specific value (see :mode below).
• :mode — "affine" (default, real biases), or a microscaled mode — "mxfp4" (group_size 32, bits 4), "mxfp8" (group_size 32, bits 8), or "nvfp4" (group_size 16, bits 4). Microscaled modes have no biases (mx::fp_quantize returns only (wq, scales)); the returned tensor's EMLX.Quantization.Config.biases is nil.

Quantizes a floating point tensor to packed format.

Returns a tuple of {quantized_weights, scales, biases} where:

• quantized_weights - Packed uint32 tensor (8 int4 values per uint32)
• scales - Per-group scale factors (or u8 for microscaled modes)
• biases - Per-group zero points, or nil for microscaled modes ("mxfp4"/"mxfp8"/"nvfp4" — mx::fp_quantize doesn't emit biases)

Parameters

• w - Float tensor to quantize
• group_size - Number of weights per group (default: 64)
• bits - Quantization bits (default: 4)
• mode - "affine" (default), "mxfp4", "mxfp8", or "nvfp4"

Run activation @ dequantize(qw) using mx::quantized_matmul.

qw must be a quantized tensor produced by EMLX.quantize/2. Raises ArgumentError if both arguments are quantized.

Performs quantized matrix multiplication.

This is the key operation for efficient 4-bit inference. It multiplies x with quantized weights w (packed as uint32), using scales and (for "affine") biases for dequantization during the computation.

Parameters

• x - Input tensor (e.g., {batch, seq, hidden})
• w - Quantized weights as uint32 (8 int4 values packed per uint32)
• scales - Per-group scale factors (bfloat16, or u8 for microscaled modes)
• biases - Per-group zero points (bfloat16), or nil for microscaled modes
• transpose - Whether to transpose weights (default: true)
• group_size - Number of weights per scale/bias group (default: 64)
• bits - Quantization bits (default: 4)
• mode - "affine" (default), "mxfp4", "mxfp8", or "nvfp4"

Resets the peak memory counter to zero.

Examples

EMLX.reset_peak_memory()
#=> :ok

Sets the cache limit in bytes. Returns the previous limit.

When cached memory exceeds this limit, it will be reclaimed on the next allocation. Set to 0 to disable caching entirely.

Examples

prev = EMLX.set_cache_limit(500_000_000)
EMLX.set_cache_limit(prev)

Sets the memory limit in bytes. Returns the previous limit.

The memory limit is a guideline for maximum memory usage during graph evaluation. Defaults to 1.5× the device's recommended working set size.

Examples

prev = EMLX.set_memory_limit(8_000_000_000)
EMLX.set_memory_limit(prev)

Unlinks a POSIX shared-memory segment by its handle name.

Call this if the receiver never opens the %Nx.Pointer{kind: :ipc} returned by Nx.to_pointer/2 — otherwise the shm name persists until the next reboot. Safe to call even if the segment has already been unlinked (ENOENT is ignored).

Returns {address, byte_size} for the tensor's raw GPU buffer.

Evals the tensor first (same pattern as to_blob/1). The pointer is valid as long as no further MLX evaluation is triggered on the array and the Elixir tensor term is kept alive. On Apple Silicon the address is accessible from both CPU and GPU due to unified memory.

Copies tensor data into a new POSIX shared-memory segment and returns {shm_name, byte_size}.

Note: this involves a memcpy — MLX arrays are immutable so zero-copy cross-process sharing is not possible. permissions is a Unix mode integer (e.g. 0o400 for owner-read-only).

The shm name persists until the receiver opens and unlinks it (which EMLX.NIF.array_from_shm/4 does automatically).

Moves a tensor to target_device (:cpu or :gpu) without deallocating the source. This is a no-op if the tensor is already on the target device.

Unlike Nx.backend_transfer/2, this does not round-trip through a binary and does not call backend_deallocate on the original tensor. Internally it schedules mlx::core::contiguous(arr, target_device) on the target device's worker thread, which on Apple Silicon (unified memory) avoids any physical data copy.

Converts an EMLX device ref back to an Nx.Tensor.

Example

result_ref = EMLX.some_operation(input)
result_tensor = EMLX.to_nx(result_ref)