# `EMLX.Fast`
[🔗](https://github.com/elixir-nx/emlx/blob/v0.4.2/emlx/lib/emlx/fast.ex#L1)

Single-kernel Metal shaders from `mlx::fast`, exposed as `deftransform`
functions.

Every function is defn-safe: call inside `defn`, `Nx.Defn.jit`, or from
`Axon.rewrite_nodes/2` rewrite callbacks without restriction — **except
`einsum/2`**, which is eager-only (see its docs).

## Two execution paths

Each `deftransform` here dispatches on whether its tensor arguments are
concrete (eager — a real `EMLX.Backend`-backed tensor) or expression-backed
(traced — building an `Nx.Defn.Expr` graph):

- **Eager**: the fused NIF (e.g. `EMLX.fast_rms_norm/3`) runs immediately.
- **Traced**: the function returns an `Nx.Defn.Expr.metadata/2` node
  carrying a `:__EMLX__` key — `%{op: opcode, operands: [...], attrs: [...]}`
  — naming the native EMLX opcode, its operand tensors, and its int-encoded
  attributes (see `EMLX.Native.Expr.f64_bits/1`). `EMLX.Native.Expr`'s
  `:metadata` `expand_node` clause recognizes this key and lowers straight
  to the native op — no graph split. The metadata node wraps an
  `Nx.runtime_call/4` of the *same* eager callback used above (`inner`) —
  this compiler never evaluates it (it's discarded in favor of the
  `:__EMLX__` payload); it exists only so (a) the operand tensors are
  ordinary reachable dependencies for `EMLX.Defn.Tree.post_order/2` to
  visit, and (b) any other `Nx.Defn.Compiler` — notably the default
  `Nx.Defn.Evaluator` and `Nx.Defn.Grad` — still gets a correct fallback:
  `runtime_call` just runs the real NIF against concrete tensors, so it's
  both exact (not a slower plain-`Nx` approximation) and free to build
  (unlike a full composite reference formula, `Nx.runtime_call/4` is a
  single lightweight node — no per-op sub-expression tracing cost). The
  `:__EMLX__` node is itself wrapped once more in a
  `Nx.Defn.Kernel.custom_grad/3` annotation (`with_reference_grad/3`) so
  `Nx.Defn.grad` differentiates through each op's plain-`Nx` `*_reference/N`
  formula (VJP via `Nx.Defn.Grad.transform/3` — see `with_reference_grad/3`)
  instead of hitting the non-differentiable `runtime_call` forward pass.

A bare `Nx.runtime_call` (anything *not* wrapped in `:__EMLX__` metadata)
always forces a graph split — see `EMLX.split_point?/1`.

## Functions

- `rms_norm/3` — fused RMS normalisation
- `layer_norm/4` — fused layer normalisation (with bias)
- `layer_norm/3` — fused layer normalisation (weight-only, no bias)
- `rope/6` — fused RoPE with scalar integer offset
- `rope_with_positions/6` — fused RoPE accepting a `position_ids` tensor
- `rope_with_freqs/6` — fused RoPE with precomputed inv-frequency tensor (for `:llama3` scaling)
- `scaled_dot_product_attention/4` — flash-attention SDPA (no mask)
- `scaled_dot_product_attention/5` — flash-attention SDPA, either an additive/bool
  `mask` tensor, or an `opts` keyword list (`:sinks`) when there's no mask
- `scaled_dot_product_attention/6` — flash-attention SDPA with `mask` and `opts` (`:sinks`)
- `scaled_dot_product_attention_causal/4` — flash-attention SDPA with built-in causal mask
- `scaled_dot_product_attention_causal/5` — same, plus `opts` (`:sinks`)
- `scaled_dot_product_attention_causal_key_masked/5` — causal SDPA; checks key_mask at C++ level, fast-paths to pure causal when all-ones
- `scaled_dot_product_attention_causal_key_masked/6` — same, plus `opts` (`:sinks`)
- `swiglu/2` — fused SwiGLU: `silu(gate) * up`
- `kv_cache_sdpa_update/8` — fused decode-time KV-cache write + causal SDPA
  read, 3 outputs (`{attn_out, k_cache_updated, v_cache_updated}`) —
  inference/decode-loop only, no grad support
- `einsum/2` — variadic-operand Einstein summation (`mlx::core::einsum`);
  **eager-only, not defn-safe** (see its docs)

## Axon graph rewrite example

    Axon.rewrite_nodes(model, fn
      %Axon.Node{op: :rms_norm, opts: [eps: eps]} ->
        fn [x, weight], _output -> EMLX.Fast.rms_norm(x, weight, eps) end
      _ -> :skip
    end)

# `einsum`

Variadic-operand einsum computed by MLX's path-optimised
`mlx::core::einsum` kernel.

`subscripts` is a standard Einstein-summation equation (e.g.
`"ij,jk->ik"`, `"bij,bjk->bik"`, `"bhid,bhjd->bhij"`,
`"ij,jk,kl->il"`). `operands` is the corresponding list of 2+ tensors.

## Eager-only, not defn-callable

Unlike the other helpers in this module, `einsum/2` does **not** emit an
`Nx.Defn.Expr` node — it takes refs directly off `EMLX.Backend`-backed
tensors and calls the NIF eagerly, in the same "direct-call helper" style
as `EMLX.quantized_matmul/2`. Every operand must live on `EMLX.Backend`;
anything else raises `ArgumentError`.

## Examples

    iex> a = Nx.iota({2, 3}, backend: EMLX.Backend, type: :f32)
    iex> b = Nx.iota({3, 4}, backend: EMLX.Backend, type: :f32)
    iex> y = EMLX.Fast.einsum("ij,jk->ik", [a, b])
    iex> Nx.shape(y)
    {2, 4}

# `kv_cache_sdpa_update`

Fused decode-time KV-cache write (2× `put_slice`) + causal-masked SDPA read,
as one native op — the traced-graph analog of the hand-written
`EMLX.kv_cache_sdpa_update/7` eager NIF, but operating on the *full*
fixed-size preallocated cache with a *dynamic* (array-valued) `offset`
instead of a host int, so it can sit inside a compiled `:while` decode
loop without forcing a host readback per iteration.

All tensors use **Bumblebee-native, heads-NOT-transposed** layout (matching
`Bumblebee.Layers.Decoder.attention_cache/4`'s `{batch, seq, heads, head_dim}`
cache shape and `EMLXAxon`'s `build_sdpa_layer/5` pre-transpose inputs) — the
head-transpose SDPA itself needs is done internally, once, on `q` and on the
already-written `k_upd`/`v_upd`; the returned caches stay untransposed so
they can feed a `:while` carry directly with no extra transpose.

- `q`        — `{B, T_q,  N_q,  D}` (T_q = 1 for decode; not enforced)
- `new_k`    — `{B, T_q,  N_kv, D}`
- `new_v`    — `{B, T_q,  N_kv, D}`
- `k_cache`  — `{B, T_max, N_kv, D}`
- `v_cache`  — `{B, T_max, N_kv, D}`
- `offset`   — scalar/`{1}` int tensor — write position along the `T_max` axis
- `key_mask` — `{B, T_max}` — 1 = valid (readable), 0 = masked
- `scale`    — pre-computed scalar

Returns `{attn_out, k_cache_updated, v_cache_updated}` — `attn_out` is
**head-transposed**, `{B, N_q, T_q, D}` (matching what the unfused
`scaled_dot_product_attention_causal_key_masked/6` metadata node itself
already returns; callers transpose it back same as today), while the
updated caches keep `k_cache`/`v_cache`'s own (untransposed) shape.

Eager: composes plain `Nx.put_slice` ×2 (clamped host offset) +
`scaled_dot_product_attention_causal_key_masked/6` — deliberately
independent of the traced path's C++ formula below, so it doubles as an
correctness oracle for the native `multi_op_registry["kv_cache_sdpa_update"]`
kernel (same combined causal+key_mask formula, reimplemented once in
Elixir/Nx and once in C++).

Traced: lowers to one `multi_op_registry["kv_cache_sdpa_update"]`
(`emlx/compiler.cpp`) native IR instruction.

# `layer_norm`

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

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

Output shape and type match `x`.

# `layer_norm`

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`).

Output shape and type match `x`.

# `rms_norm`

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

- `x`      — input tensor; normalised over the last axis.
- `weight` — `{hidden}` scale vector (same size as last axis of `x`).
- `eps`    — numerical stability constant (e.g. `1.0e-6`).

Output shape and type match `x`.

# `rope`

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

- `a`           — input `{B, ..., T, D}`; `...` dims are passed through, but
  **`T` must be the second-to-last axis** (e.g. `{B, H, T, D}`,
  heads-transposed) — see `EMLX.fast_rope/6`'s `@doc` for why a heads-axis
  placed *before* `T` (Bumblebee's `{B, T, H, D}`) silently miscomputes for
  more than one head ([elixir-nx/emlx#121](https://github.com/elixir-nx/emlx/issues/121)).
  Use `rope_with_positions/6` or `rope_with_freqs/6` for that layout instead.
- `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 position offset (tokens already in the KV cache).

**`traditional` must match the model checkpoint's convention.**
For Qwen3 (split-half): `traditional: false`.

Output shape and type match `a`.

# `rope_with_freqs`

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

Use this variant when the model's RoPE scaling strategy produces a fixed
`{dims/2}` inv-frequency tensor that can be baked at graph-rewrite time
(e.g. `:llama3` smooth-interpolation). Strategies that are seq-len conditional
or require cos/sin post-multiply (`:linear`, `:dynamic`, `:longrope`) should
use `rope_with_positions/6` instead.

- `a`            — input `{B, T, ..., D}` (Bumblebee convention: heads NOT yet transposed)
- `position_ids` — `{B, T}` integer tensor. For **decode** (`T = 1`) the fast path uses
  `position_ids[b,0]` as the per-batch offset into `freqs` (same contract as
  `mlx::fast::rope` with a scalar offset per batch). For **prefill** (`T > 1`) a
  per-token path runs so arbitrary positions (e.g. left-padded
  `[0,…,0,1,2,…]`) are correct; the offset-only entry point cannot represent
  that.
- `dims`         — number of feature dims to rotate.
- `traditional`  — `false` for split-half (Bumblebee / Qwen3); `true` for interleaved.
- `scale`        — position scale (`1.0` for most strategies with precomputed freqs).
- `freqs`        — `{dims/2}` tensor of precomputed inverse frequencies.

Output shape and type match `a`.

# `rope_with_positions`

Fused RoPE accepting a `position_ids` tensor (`mlx::fast::rope`, array-offset overload).

Use this variant when the calling convention provides `position_ids` as a tensor
(e.g. from Bumblebee's rotary embedding layer) rather than a scalar integer offset.

- `a`            — input `{B, T, ..., D}` (Bumblebee convention: heads NOT yet transposed)
- `position_ids` — `{B, T}` integer tensor; each row holds the token positions for
                   one batch example. **Positions must be sequential within each row**
                   (standard causal LM). The starting offset for batch item `b` is
                   taken as `position_ids[b, 0]`; subsequent positions are inferred
                   by MLX as `offset + 0, offset + 1, ...`.
- `dims`         — number of feature dims to rotate.
- `traditional`  — `false` for split-half (Bumblebee / Qwen3); `true` for interleaved.
- `base`         — angular frequency base (e.g. `10_000`).
- `scale`        — position scale (`1.0` unless using NTK-aware scaling).

Output shape and type match `a`.

> ### Sequential positions only (fast T=1 path) {: .warning}
> For **decode** with `T = 1` and `base` below about `1.0e5`, the `fast_rope_ids`
> opcode is used; it assumes sequential positions from `position_ids[b, 0]`. For
> **larger** `base` (e.g. Qwen3 `rope_theta` 1M) or **prefill** (`T > 1`), the
> per-token `fast_rope_positions` opcode is used, matching Bumblebee for
> arbitrary per-token `position_ids`.

# `scaled_dot_product_attention`

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

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

- `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)`)

Output: `{B, N_q, T_q, D}` — same dtype as `q`.
Softmax accumulates in float32 internally regardless of input dtype.

# `scaled_dot_product_attention`

Flash-attention SDPA — either an additive/boolean `mask` tensor (5th arg),
or, with no mask, an `opts` keyword list (disambiguated by `is_list/1`):

* `:sinks` — optional learned per-head attention-sink logits tensor,
  appended as an extra key/value pair the softmax normalises against (see
  `mlx::fast::scaled_dot_product_attention`'s `sinks` parameter). Shape
  must broadcast against `{B, N_q}` (typically `{N_q}`).

`mask` must be broadcast-compatible with `{B, N_q, T_q, T_kv}`.
Boolean `false` entries are masked out (`-∞`); float entries are added to
the pre-softmax scores.

For causal masking in decode (single query token), prefer the no-mask arity
since `T_q=1` is always trivially causal.

# `scaled_dot_product_attention`

Flash-attention SDPA with an additive/boolean `mask` and `opts` (`:sinks`
— see `scaled_dot_product_attention/5`).

# `scaled_dot_product_attention_causal`

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

MLX constructs the upper-triangular causal mask internally without materialising it,
making this equivalent to `scaled_dot_product_attention/5` with a causal boolean mask
but cheaper: no mask tensor allocation, and the mask is fused into the Metal kernel.

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

Input/output layout matches `scaled_dot_product_attention/4`:
- `q`     — `{B, N_q,  T_q,  D}`
- `k`     — `{B, N_kv, T_kv, D}`
- `v`     — `{B, N_kv, T_kv, D}`
- `scale` — pre-computed scalar (typically `1 / sqrt(D)`)
- Output  — `{B, N_q, T_q, D}`, same dtype as `q`

# `scaled_dot_product_attention_causal`

Causal flash-attention SDPA with `opts` (`:sinks` — see
`scaled_dot_product_attention/5`).

# `scaled_dot_product_attention_causal_key_masked`

Causal SDPA with the key_mask check delegated to the C++ NIF (eager) or
folded directly into the compiled graph (traced).

When eager, the NIF evaluates `all(key_mask == 1)`:
- **true** (no padding, e.g. single-sequence decode) → pure causal SDPA,
  no mask tensor allocated.
- **false** (padded batch or multi-sequence) → builds a combined
  causal + key_mask additive mask and calls the masked SDPA kernel.

This avoids the `Nx.cond` double-evaluation problem: the NIF forces eval
of only the small `{B, T_kv}` key_mask subgraph, then branches in C++.

When traced, the compiled `:fast_sdpa_causal_key_masked*` opcode always
builds the combined causal+key_mask additive mask in-graph (a compiled
program can't branch on a runtime `all(key_mask)` check).

Input/output layout matches `scaled_dot_product_attention_causal/4`:
- `q`        — `{B, N_q,  T_q,  D}`
- `k`        — `{B, N_kv, T_kv, D}`
- `v`        — `{B, N_kv, T_kv, D}`
- `scale`    — pre-computed scalar
- `key_mask` — `{B, T_kv}` boolean/int tensor (1 = attend, 0 = masked)
- Output     — `{B, N_q, T_q, D}`, same dtype as `q`

An optional 6th `opts` keyword list accepts `:sinks` (see
`scaled_dot_product_attention/5`).

# `swiglu`

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

Eliminates the two-op `silu(gate_proj) * up_proj` pattern that appears in
Qwen3's FFN layers (28× per decode step).

- `gate` — gate-projection output; silu is applied element-wise.
- `up`   — up-projection output; same shape as `gate`.

Output has the same shape and dtype as `gate`.

---

*Consult [api-reference.md](api-reference.md) for complete listing*