Skip to content

Commit 5e3e4d2

Browse files
authored
Merge pull request #154 from ausimian/feat/expr-compiler-release-notes
docs: release note for the single-NIF native compiler
2 parents c0b9ff0 + 6a96c64 commit 5e3e4d2

1 file changed

Lines changed: 30 additions & 0 deletions

File tree

RELEASE.md

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,35 @@
11
### Added
22

3+
- **Single-NIF native compiler — `Nx.Defn.jit`/`compile` with
4+
`compiler: Emily.Compiler, native: true`.** A real `Nx.Defn.Compiler` path
5+
that lowers a traced `Nx.Defn.Expr` to a flat IR **once** and replays the
6+
whole forward graph in a **single NIF call per invocation**, instead of one
7+
BEAM↔worker round-trip per op. For dispatch-bound workloads — autoregressive
8+
decode, where the structurally-identical graph is otherwise rebuilt op-by-op
9+
every token — this collapses the per-token dispatch cost (a 100-op microbench
10+
shows a >15× build/dispatch collapse). Weights cross the NIF boundary once
11+
(captured by the compiled program) and are never re-serialized per call.
12+
Opt in per call:
13+
14+
Nx.Defn.jit(&forward/1, compiler: Emily.Compiler, native: true).(input)
15+
16+
Coverage is **no-fallback**: the full Nx primitive set (with `Emily.Backend`'s
17+
dtype-coercion and op-composition semantics ported into the lowering), the
18+
fused `Emily.Fast.*` / `Nx.Block.*` kernels (RMSNorm, LayerNorm, RoPE, scaled
19+
dot-product attention and its mask/sink variants, the LinAlg blocks),
20+
quantized matmul (now an `Nx.block` node so it fuses under the compiler too),
21+
dynamic KV-cache writes (`put_slice` at a runtime offset), container/tuple
22+
outputs, and `cond` (lowered to a select-chain). DistilBERT and ViT forwards
23+
run end-to-end under the compiler with `config :emily, :fallback, :raise`.
24+
Unsupported constructs — `while` loops and arbitrary BEAM `reduce` functions —
25+
raise a clear compile-time error rather than silently falling back, so
26+
generation loops stay driven from Elixir.
27+
28+
An opt-in compiled eval mode additionally wraps the replay in
29+
`mlx::core::compile`, fusing the elementwise runs (rms-norm, softmax, SiLU
30+
gating, residual adds) the replay leaves as separate kernels — measured at
31+
~1.5–1.6× over the plain replay on a decode-shaped transformer block.
32+
333
- `Emily.async_eval/1` (and `Emily.Native.async_eval/2`) schedule evaluation of
434
one or more lazy graphs **without blocking on the GPU**, wrapping
535
`mlx::core::async_eval`. The work is handed to the device's command queue and

0 commit comments

Comments
 (0)