Skip to content

Commit 795cc86

Browse files
committed
docs: add a Native compilation section to the README
The Features list and Usage section only mentioned the op-by-op compiler mode; the single-NIF native replay and `:fuse` opt-in are now shipping but the README didn't reflect either. Expand the Defn-compiler Features bullet to name the three modes briefly, and add a Native compilation section between Usage and Livebooks (matching the depth of Concurrency model — a peer "advanced use" topic) that shows all three modes in one code block, calls out the f32-reassociation trade-off for `:fuse` and the sampling-divergence caveat, and points readers at `Emily.Generation` for the autoregressive-decode case.
1 parent 0f25ef8 commit 795cc86

1 file changed

Lines changed: 51 additions & 1 deletion

File tree

README.md

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,11 @@ Emily backend smoke test.
5454
with a `[:emily, :fallback, *]` telemetry event. See the
5555
`Fallbacks` section of `Emily.Backend` for the per-op catalogue.
5656
- **Defn compiler.** `Emily.Compiler` runs `defn` / `Nx.Serving` /
57-
Bumblebee inference on MLX. Backs the results with lazy MLX graphs.
57+
Bumblebee inference on MLX. The default mode dispatches op-by-op
58+
through `Emily.Backend`; opt in to `native: true` for a single-NIF
59+
replay of the whole graph (~5× decode throughput on Qwen3-0.6B), and
60+
`fuse: true` for `mx::compile` kernel fusion on top. See
61+
[Native compilation](#native-compilation).
5862
- **Fused transformer kernels.** `Emily.Fast` exposes
5963
`mx::fast::rms_norm`, `layer_norm`, `rope`, and scaled-dot-product
6064
attention as defn-callable helpers with composed-defn fallbacks for
@@ -198,6 +202,52 @@ The low-level tensor API (`Emily.from_binary/3`, `to_binary/1`,
198202
`shape/1`, `dtype/1`, `eval/1`) remains available for diagnostics
199203
and direct MLX round-trips, but most users should go through Nx.
200204

205+
## Native compilation
206+
207+
`Emily.Compiler` has three modes; opt in per-call or globally via
208+
`Nx.Defn.global_default_options/1`:
209+
210+
```elixir
211+
# Default: op-by-op dispatch through Emily.Backend. Bit-identical to
212+
# the Nx evaluator.
213+
Nx.Defn.jit(&forward/1, compiler: Emily.Compiler).(input)
214+
215+
# Native single-NIF replay. Lowers the traced Nx.Defn.Expr to a flat
216+
# IR and replays the whole forward graph in one NIF call per
217+
# invocation — ~5× decode throughput on Qwen3-0.6B, bit-identical to
218+
# the evaluator. Safe to install globally: anything the IR can't lower
219+
# routes through Nx.Defn.Evaluator under the default
220+
# `native_fallback: :eval` (with a `[:emily, :compiler, :fallback]`
221+
# telemetry event).
222+
Nx.Defn.jit(&forward/1, compiler: Emily.Compiler, native: true).(input)
223+
224+
# Native + mx::compile kernel fusion. Fuses elementwise runs the plain
225+
# replay leaves as separate kernels (RMSNorm / softmax / SiLU gating /
226+
# residual adds) — ~1.1× over the plain native lane on Qwen3-0.6B
227+
# greedy decode (~5.4× the evaluator overall). For a `defn while`,
228+
# the loop body is fused under mx::compile and cached per stream so
229+
# it cache-hits across iterations.
230+
Nx.Defn.jit(&forward/1, compiler: Emily.Compiler, native: true, fuse: true).(input)
231+
```
232+
233+
Trade-off for `:fuse`: `mx::compile` reassociates f32, so logits drift
234+
by a few ULP and the output is **not** bit-identical to the evaluator.
235+
Greedy argmax is empirically robust to that drift (Qwen3-0.6B token
236+
ids matched the evaluator's exactly in our benchmarks), but
237+
**sampling strategies will diverge** from the evaluator even with a
238+
fixed seed.
239+
240+
Pass `native_fallback: :raise` to fail rather than silently degrade to
241+
the Evaluator — the conformance suites use this to prove a model
242+
lowers fully native. See the `Emily.Compiler` moduledoc for the full
243+
option list (`:device`, `:hooks`, `:max_concurrency`, etc.) and the
244+
trade-offs around `:fuse` in depth.
245+
246+
For autoregressive decode specifically, `Emily.Generation` is a
247+
model-agnostic loop driver that JIT-compiles a caller-supplied
248+
shape-stable per-token forward and runs the loop from Elixir with
249+
KV-cache threading, stop conditions, and per-token streaming.
250+
201251
## Livebooks
202252

203253
End-to-end Livebooks under `livebooks/`. Each one declares its own

0 commit comments

Comments
 (0)