Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 25 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,34 @@ endif

MAKE_JOBS ?= $(JOBS)

.PHONY: all clean
.PHONY: all clean bench-native

all: $(NIF_SO)

# ------------------------------------------------------------------
# bench-native: standalone C++ microbenchmarks under bench/native/.
#
# Invoked by `mix bench.native` (which sets the same env elixir_make
# uses for the NIF build). Links against the vendored libmlx via the
# same rpath the NIF does, so the binary finds its own shared library
# without relying on global DYLD/LD paths.
# ------------------------------------------------------------------

BENCH_NATIVE_SRC := bench/native/compile_microbench.cpp
BENCH_NATIVE_BIN := $(BUILD_DIR)/compile_microbench

$(BENCH_NATIVE_BIN): $(BENCH_NATIVE_SRC) | $(BUILD_DIR)
$(CXX) -std=c++17 -O3 -Wall -Wextra \
-isystem $(MLX_INCLUDE_DIR) \
$(BENCH_NATIVE_SRC) \
-L$(MLX_LIB_DIR) -lmlx \
-Wl,-rpath,$(MLX_LIB_DIR) \
-o $(BENCH_NATIVE_BIN)

bench-native: $(BENCH_NATIVE_BIN)
@echo "Running $(BENCH_NATIVE_BIN)"
@$(BENCH_NATIVE_BIN) $(BENCH_NATIVE_ARGS)

$(BUILD_DIR):
@mkdir -p $(BUILD_DIR)

Expand Down
54 changes: 40 additions & 14 deletions PLAN.md
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,9 @@ into BEAM.
## Core design decisions

1. **Backend-first; compiler layered on top.** The Backend is enough to
run Bumblebee. `mlx::core::compile` is an opt-in optimisation added
last.
run Bumblebee. Wrapping `mlx::core::compile` was planned as an
opt-in optimisation on top, but was dropped after de-risking — see
M6 for the measurement and reasoning.
2. **Trace in Elixir, not in C++.** `Nx.Defn.Expr` is already a fully
traced tree; we walk it from Elixir and emit one `Emily.Native` call
per node. No C++→BEAM callbacks.
Expand Down Expand Up @@ -152,7 +153,9 @@ throughput number.
rejected once we accounted for the per-call ETS deep-copy cost on a
Qwen3-sized expression tree. The closure-capture path avoids the copy
and matches the upstream Evaluator pattern.*
- **Do not use `mlx::core::compile` yet.** Lazy eval at the Backend layer suffices.
- **Do not use `mlx::core::compile`.** M6 de-risked this and dropped
it — the fusion win on transformer-shaped workloads is below the
1.20× gate. Lazy eval at the Backend layer is the shipping story.

**Testing — Layer 3 (Compiler):**

Expand All @@ -173,20 +176,43 @@ throughput number.
match `Nx.Defn.Evaluator` running on the same backend within float
tolerance. (Training is out of scope for v1.)

### M6 — `mlx::core::compile` wrapping
### M6 — `mlx::core::compile` wrapping — **dropped**

- After Compiler has built the lazy op sequence, optionally wrap it in
`mlx::core::compile` for shape-pinned specialisation
- Purely an optimisation; Backend-only path remains the default
- Thunk constructed in a single NIF call from a recorded op list; MLX
traces a closure that replays ops against placeholders — no BEAM
callbacks
De-risked in pure C++ before paying the Backend/Compiler integration
cost, per the PLAN gate ("If <20% win, drop"). Full results:
[`bench/compile_microbench.md`](bench/compile_microbench.md).

**Testing:** equivalence tests rerun with `mlx_compile: true`; benchmark
speedup on a Qwen3 forward pass. If <20% win, drop.
Summary of findings on MLX 0.25.1, Apple Silicon:

**Exit:** compile mode is off by default, opt-in, demonstrably faster,
zero regressions.
- **Pure elementwise workload (harness validation):** 2.78× on GPU,
1.47× on CPU — confirms `mx::compile` does what it advertises when
fusion is available.
- **Transformer block (Qwen3-0.6B-shaped, seq ∈ {128, 512}):**
1.04–1.07× on GPU, **regression** (0.82–0.88×) on CPU. Fails the
1.20× gate across every workload shape tested.

Why: transformer inference is matmul-dominated, and MLX's compile does
not fuse matmul kernels with adjacent elementwise ops. The fusion
surface (RMSNorm chains, softmax neighbourhood, SwiGLU's silu×up) is a
small fraction of block runtime, bounding whole-block speedup to
single-digit percent. On CPU the tape-replay overhead exceeds the
fusion gain.

The BEAM-integrated compile path could not outperform this C++ ceiling,
so shipping M6 would deliver a <20% speedup at best — and a regression
at worst if a user selects the CPU device.

**Artefacts retained** so the decision can be re-measured against
future MLX releases without rebuilding the harness:

- `bench/native/compile_microbench.cpp` — the microbench source
- `lib/mix/tasks/bench.native.ex` — `mix bench.native` task
- `bench-native` target in the root `Makefile`
- `bench/compile_microbench.md` — results + reproduction instructions

If MLX gains matmul-adjacent fusion (bias-fused matmul, attention
fusion outside `fast::scaled_dot_product_attention`), re-run the bench
and revisit.

### M7 — 1.0 release

Expand Down
32 changes: 32 additions & 0 deletions RELEASE.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,37 @@
# Release notes for next release

## Changed

- M6 — `mlx::core::compile` wrapping: **dropped** after Phase-1
de-risk. A pure-C++ microbenchmark on MLX 0.25.1 against an Apple
Silicon GPU showed the fusion win on a Qwen3-0.6B-shaped transformer
block is 1.04–1.07× on GPU (below the PLAN's 1.20× gate) and a
regression on CPU (0.82–0.88×). A sanity workload (pure elementwise
chain) in the same harness shows the expected 2.78× GPU / 1.47× CPU
wins, confirming the measurement is trustworthy — the limiting factor
is that MLX compile doesn't fuse matmul with surrounding elementwise
ops, and transformer inference is matmul-dominated. The BEAM-
integrated compile path could not exceed this C++ ceiling, so Phase 2
and 3 were not built.
- **`bench/native/compile_microbench.cpp`** — standalone C++
microbench (hand-written RMSNorm + GQA-lite attention + SwiGLU
block, plus an 8-op elementwise sanity test). Links against the
vendored libmlx via the same rpath the NIF uses.
- **`mix bench.native`** (`lib/mix/tasks/bench.native.ex`) — Mix task
that invokes the new `bench-native` Makefile target with the same
env `elixir_make` sets, ensuring the bench uses the project's
pinned MLX without a second fetch. Supports `--seq`, `--warmup`,
`--iters` args via `mix bench.native -- <args>`.
- **`bench-native`** target added to the root `Makefile`, producing
`$(BUILD_DIR)/compile_microbench`.
- **`bench/compile_microbench.md`** — full results table +
reproduction instructions. Retained so the decision can be
re-measured against future MLX releases without rebuilding the
harness.
- **`PLAN.md`** updated: M6 section rewritten to record the drop,
core design decision #1 and the M5 section footnote updated to
match.

## Added

- M5 — `Emily.Compiler`, an `Nx.Defn.Compiler` implementation that runs
Expand Down
112 changes: 112 additions & 0 deletions bench/compile_microbench.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
# M6 de-risk — `mlx::core::compile` microbenchmark

Phase-1 measurement for PLAN milestone M6. The PLAN gates the whole
milestone on a ≥20% win from wrapping `mlx::core::compile`:

> Equivalence tests rerun with `mlx_compile: true`; benchmark speedup on
> a Qwen3 forward pass. **If <20% win, drop.**
>
> — `PLAN.md:186-189`

Rather than pay the full Backend/Compiler integration cost to find out,
we answer the question in pure C++ against the vendored MLX 0.25.1 on the
same Apple Silicon target Emily runs on. If compile doesn't help a
transformer block in raw C++, it can't help under BEAM.

## Setup

- Binary: `bench/native/compile_microbench.cpp`
- Harness: `mix bench.native` (Mix task that invokes the `bench-native`
target in the root `Makefile` with the same env elixir_make sets)
- MLX: 0.25.1 (cocoa-xu prebuilt)
- Host: M-series Mac (Metal GPU)
- Each benchmark runs 50-iteration warmup + 1000 measured iterations
(500 for seq=512), reporting min/median/p95 wall-time per iteration.
- Both variants call `mx::eval(out); mx::synchronize()` at the end of
every iteration so compile vs. uncompiled are compared apples-to-apples.

## Results

### Sanity: 8-op elementwise chain (1M elements)

Validates the harness: a pure elementwise workload is exactly what
`mx::compile` is designed to fuse.

| Device | Variant | min (ms) | median (ms) | p95 (ms) |
|--------|------------|---------:|------------:|---------:|
| GPU | uncompiled | 1.656 | 1.746 | 2.016 |
| GPU | compiled | 0.552 | 0.628 | 0.727 |
| CPU | uncompiled | 1.423 | 1.463 | 1.625 |
| CPU | compiled | 0.986 | 0.997 | 1.036 |

**GPU speedup: 2.78× median** (fusion collapses 8 kernel launches into 1).
**CPU speedup: 1.47× median**. Harness verified.

### Transformer block — Qwen3-0.6B-shaped (seq=128)

RMSNorm → Q/K/V proj → SDPA (matmul, scale, softmax, matmul) → output
proj → residual → RMSNorm → SwiGLU FFN → residual. hidden=1024,
heads=16, head_dim=64, intermediate=2816, batch=1.

| Device | Variant | min (ms) | median (ms) | p95 (ms) |
|--------|------------|---------:|------------:|---------:|
| GPU | uncompiled | 2.588 | 3.072 | 3.964 |
| GPU | compiled | 2.624 | 2.943 | 3.536 |
| CPU | uncompiled | 6.835 | 7.151 | 8.109 |
| CPU | compiled | 7.792 | 8.082 | 8.810 |

**GPU speedup: 1.04× median — FAILS 1.20× gate.**
**CPU speedup: 0.88× median — compile is slower on CPU.**

### Transformer block — longer seq (seq=512)

Tests whether scaling attention (which grows O(seq²)) shifts the fusion
ratio. It does not.

| Device | Variant | min (ms) | median (ms) | p95 (ms) |
|--------|------------|---------:|------------:|---------:|
| GPU | uncompiled | 10.804 | 11.463 | 12.068 |
| GPU | compiled | 10.206 | 10.758 | 11.240 |
| CPU | uncompiled | 21.067 | 21.744 | 22.524 |
| CPU | compiled | 25.750 | 26.544 | 27.340 |

**GPU speedup: 1.07× median — FAILS 1.20× gate.**
**CPU speedup: 0.82× median.**

## Interpretation

1. The harness is correct: a pure-elementwise sanity workload yields
the expected 2-3× compile win.
2. A transformer block is matmul-dominated. MLX's `mx::compile` fuses
elementwise chains but does **not** fuse matmul kernels with their
surrounding elementwise ops. The fusion surface (RMSNorm chains,
softmax neighbourhood, SwiGLU's silu×up) is a small fraction of
block runtime, bounding the whole-block speedup to single-digit
percent on GPU.
3. On CPU, compile is a **regression**. Tape-replay overhead exceeds
fusion gains for workloads that aren't Metal-kernel-launch-bound.
4. Scaling sequence length (128 → 512) does not materially change the
ratio. This isn't a "small workload" problem; it's a workload-shape
problem.

## Decision

**Drop M6.** The Phase-1 gate is not met and the measurement explains
why in a way that Phase-2/3 BEAM integration cannot change: the
BEAM-integrated compile path cannot outperform its C++ ceiling, and
that ceiling is 1.04–1.10× on the target workload (transformer
inference).

The microbench source and harness remain in `bench/native/` so this
result can be re-measured against future MLX releases — if MLX adds
matmul-adjacent fusion (e.g. bias-fused matmul or attention fusion
outside `fast::scaled_dot_product_attention`), M6 becomes worth
revisiting.

## Reproduce

```bash
mix bench.native # default: warmup 50, iters 1000, seq 128
mix bench.native -- --seq 512 --iters 500
mix bench.native -- --warmup 20 --iters 200 # quick smoke run
```
Loading
Loading