Merge pull request #13 from ausimian/feat/m6-compile-derisk

M6: drop mlx::core::compile wrapping after Phase-1 de-risk

Author: ausimian

Makefile

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

PLAN.md

@@ -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.
@@ -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):**
 
@@ -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
 

RELEASE.md

@@ -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

bench/compile_microbench.md

@@ -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
+```