|
1 | 1 | ### Added |
2 | 2 |
|
| 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 | + |
3 | 33 | - `Emily.async_eval/1` (and `Emily.Native.async_eval/2`) schedule evaluation of |
4 | 34 | one or more lazy graphs **without blocking on the GPU**, wrapping |
5 | 35 | `mlx::core::async_eval`. The work is handed to the device's command queue and |
|
0 commit comments