Merge pull request #21 from ausimian/m11-fast-kernels

M11: Fused MLX transformer kernels

Author: ausimian

RELEASE.md

@@ -2,6 +2,54 @@
 
 ## Added
 
+- M11 — MLX fused transformer kernels. Wires MLX's handwritten
+  `mx::fast::*` fused kernels (RMSNorm, LayerNorm, RoPE, scaled-dot-
+  product attention) into Emily as `defn`-callable helpers, and ships
+  a Bumblebee shim that swaps these in for the stock composed-defn
+  implementations when the Axon graph is rewritten with
+  `Emily.Bumblebee.FastKernels.apply/1`.
+  - **Native NIFs** (`c_src/ops/fast.cpp`) over `mx::fast::rms_norm`,
+    `mx::fast::layer_norm`, `mx::fast::rope`, and
+    `mx::fast::scaled_dot_product_attention`. Nullable weight / bias
+    / freqs arguments marshal via `std::optional`; the SDPA mask
+    argument list marshals via `std::vector<fine::ResourcePtr<Tensor>>`.
+  - **`Emily.Fast`** (`lib/emily/fast.ex`) — `rms_norm/3`,
+    `layer_norm/4`, `rope/3`, `rope_with_freqs/4`,
+    `scaled_dot_product_attention/4`,
+    `scaled_dot_product_attention_with_mask/5`. Each helper emits a
+    `Nx.Defn.Expr.optional/3` node whose op name matches a custom
+    callback on `Emily.Backend`; the Evaluator dispatches to the
+    fused kernel under Emily and falls back to a defn composition on
+    any other backend. This makes the helpers safe to drop into
+    Bumblebee inference paths without breaking BinaryBackend
+    conformance runs.
+  - **`Emily.Backend.fast_*`** — six custom callbacks (not part of the
+    `Nx.Backend` behaviour) that Evaluator picks up when the input
+    tensors carry Emily data. They unwrap refs and call the Native
+    NIF directly.
+  - **`Emily.Bumblebee.FastKernels`** (`test/support/`) — Axon graph
+    rewriter, mirroring the M10.5 `Emily.Quantization.Transform`
+    pattern. Rewrites `:rms_norm` and `:layer_norm` nodes via
+    `Axon.map_nodes`, `Bumblebee.Layers.apply_rotary_embedding/5` by
+    function-reference match, and coalesces
+    `attention_weights_impl + attention_output_impl` into one fused
+    SDPA layer via `Axon.rewrite_nodes`. RoPE handles all four
+    Bumblebee scaling strategies (`:linear`, `:dynamic`, `:longrope`,
+    `:llama3`) by precomputing the inverse-frequency table Elixir-
+    side and passing it to MLX via the `freqs`-override overload.
+  - **Tests**: per-kernel Native/defn/equivalence suites at
+    `test/emily/fast/`; shim unit tests at
+    `test/emily/bumblebee/fast_kernels_test.exs`; fused-kernel
+    variants of every `*_full` conformance suite tagged
+    `:fast_kernels_full` (excluded by default like the other
+    `*_full` tags). Run explicitly:
+    `mix test --only fast_kernels_full`.
+  - **Bench**: `bench/qwen3_tokens_per_sec.exs` gains an
+    `EMILY_BENCH_FAST_KERNELS=1` mode that benchmarks baseline vs
+    fused side-by-side, plus an `EMILY_BENCH_PIN=<multiplier>` flag
+    that fails with a non-zero exit when the fused mean throughput
+    doesn't clear the multiplier × baseline threshold.
+
 - M10 (partial) — Quantized inference primitives. Exposes MLX's affine
   int4/int8 group-wise quantization at the Native and Elixir levels, plus
   a direct-call helper for eager use. Enough to quantize a dense weight,

bench/qwen3_tokens_per_sec.exs

@@ -6,20 +6,29 @@
 #
 # Optional environment variables:
 #
-#     EMILY_BENCH_MODEL        HuggingFace repo id. Defaults to
-#                              "Qwen/Qwen3-0.6B".
-#     EMILY_BENCH_NEW_TOKENS   Number of tokens to greedy-decode per
-#                              run. Defaults to 64.
-#     EMILY_BENCH_PROMPT       Prompt text. Defaults to a fixed short
-#                              English sentence.
-#     EMILY_BENCH_WARMUP       Number of warm-up runs (not measured).
-#                              Defaults to 1.
-#     EMILY_BENCH_RUNS         Number of measured runs. Defaults to 3.
+#     EMILY_BENCH_MODEL          HuggingFace repo id. Defaults to
+#                                "Qwen/Qwen3-0.6B".
+#     EMILY_BENCH_NEW_TOKENS     Number of tokens to greedy-decode per
+#                                run. Defaults to 64.
+#     EMILY_BENCH_PROMPT         Prompt text. Defaults to a fixed short
+#                                English sentence.
+#     EMILY_BENCH_WARMUP         Number of warm-up runs (not measured).
+#                                Defaults to 1.
+#     EMILY_BENCH_RUNS           Number of measured runs. Defaults to 3.
+#     EMILY_BENCH_FAST_KERNELS   "1" → also benchmark the M11 fused
+#                                MLX kernels (RMSNorm, LayerNorm, RoPE,
+#                                SDPA) via `Emily.Bumblebee.FastKernels`.
+#                                Reports baseline vs fused side by side.
+#     EMILY_BENCH_PIN            "1.5" → fail with non-zero exit if the
+#                                fused mean tokens/sec doesn't beat
+#                                baseline mean by at least the given
+#                                multiplier. Implies
+#                                EMILY_BENCH_FAST_KERNELS=1.
 #
 # The first run downloads the model (~1.2 GB at f32, ~600 MB at f16).
 # We deliberately avoid `Benchee` — this benchmark has one workload and
-# one metric (tokens/sec). The whole script is standalone so a reader
-# can follow the generation flow without chasing macros.
+# one (or two) metrics. The whole script is standalone so a reader can
+# follow the generation flow without chasing macros.
 
 defmodule Emily.Bench.Qwen3 do
   @default_model "Qwen/Qwen3-0.6B"
@@ -34,19 +43,27 @@ defmodule Emily.Bench.Qwen3 do
     model_repo = System.get_env("EMILY_BENCH_MODEL", @default_model)
     prompt = System.get_env("EMILY_BENCH_PROMPT", @default_prompt)
 
-    new_tokens =
-      System.get_env("EMILY_BENCH_NEW_TOKENS")
-      |> env_int(@default_new_tokens)
-
+    new_tokens = System.get_env("EMILY_BENCH_NEW_TOKENS") |> env_int(@default_new_tokens)
     warmup = System.get_env("EMILY_BENCH_WARMUP") |> env_int(@default_warmup)
     runs = System.get_env("EMILY_BENCH_RUNS") |> env_int(@default_runs)
 
+    pin_threshold =
+      case System.get_env("EMILY_BENCH_PIN") do
+        nil -> nil
+        s -> elem(Float.parse(s), 0)
+      end
+
+    fast_kernels? =
+      System.get_env("EMILY_BENCH_FAST_KERNELS") == "1" or pin_threshold != nil
+
     IO.puts("Emily / Qwen3 throughput benchmark")
-    IO.puts("  model      : #{model_repo}")
-    IO.puts("  prompt     : #{inspect(prompt)}")
-    IO.puts("  new tokens : #{new_tokens}")
-    IO.puts("  warmup     : #{warmup}")
-    IO.puts("  runs       : #{runs}")
+    IO.puts("  model          : #{model_repo}")
+    IO.puts("  prompt         : #{inspect(prompt)}")
+    IO.puts("  new tokens     : #{new_tokens}")
+    IO.puts("  warmup         : #{warmup}")
+    IO.puts("  runs           : #{runs}")
+    IO.puts("  fused kernels  : #{fast_kernels?}")
+    if pin_threshold, do: IO.puts("  pin threshold  : #{pin_threshold}× baseline")
     IO.puts("")
 
     {:ok, model_info} = Bumblebee.load_model({:hf, model_repo})
@@ -60,6 +77,34 @@ defmodule Emily.Bench.Qwen3 do
         strategy: %{type: :greedy_search}
       )
 
+    IO.puts("=== baseline (composed defn kernels) ===")
+    baseline = bench(model_info, tokenizer, generation_config, prompt, new_tokens, warmup, runs)
+    {baseline_mean, _, _, _} = baseline
+
+    if fast_kernels? do
+      IO.puts("\n=== fused (Emily.Bumblebee.FastKernels) ===")
+
+      fused_model_info =
+        update_in(model_info.model, &Emily.Bumblebee.FastKernels.apply/1)
+
+      fused = bench(fused_model_info, tokenizer, generation_config, prompt, new_tokens, warmup, runs)
+      {fused_mean, _, _, _} = fused
+
+      speedup = fused_mean / baseline_mean
+      IO.puts("\nspeedup        : #{Float.round(speedup, 2)}× (fused mean / baseline mean)")
+
+      if pin_threshold do
+        if speedup >= pin_threshold do
+          IO.puts("PIN OK         : #{Float.round(speedup, 2)}× ≥ #{pin_threshold}×")
+        else
+          IO.puts("PIN FAIL       : #{Float.round(speedup, 2)}× < #{pin_threshold}×")
+          System.halt(1)
+        end
+      end
+    end
+  end
+
+  defp bench(model_info, tokenizer, generation_config, prompt, new_tokens, warmup, runs) do
     serving =
       Bumblebee.Text.generation(model_info, tokenizer, generation_config,
         defn_options: [compiler: Nx.Defn.Evaluator]
@@ -86,14 +131,16 @@ defmodule Emily.Bench.Qwen3 do
     mean = Enum.sum(tps_list) / length(tps_list)
     {min_tps, max_tps} = Enum.min_max(tps_list)
 
-    IO.puts("")
-    IO.puts("tokens/sec  mean=#{Float.round(mean, 2)}  min=#{Float.round(min_tps, 2)}  max=#{Float.round(max_tps, 2)}")
-    IO.puts("")
-    IO.puts("first completion:")
-    IO.puts(String.slice(sample, 0, 500))
+    IO.puts(
+      "tokens/sec     : mean=#{Float.round(mean, 2)}  min=#{Float.round(min_tps, 2)}  max=#{Float.round(max_tps, 2)}"
+    )
+
+    IO.puts("first completion:\n  #{String.slice(sample, 0, 500)}")
+    {mean, min_tps, max_tps, sample}
   end
 
   defp env_int(nil, default), do: default
+
   defp env_int(s, default) do
     case Integer.parse(s) do
       {n, ""} -> n

c_src/ops/fast.cpp

@@ -0,0 +1,161 @@
+// Fused transformer kernels from mlx::core::fast.
+//
+// These handwritten kernels beat the defn-composed equivalents on the
+// transformer hot paths: RMSNorm (one kernel vs rsqrt+mean+multiply
+// chain), LayerNorm (same for Welford+affine), RoPE (fused trig +
+// interleave), and Scaled-Dot-Product Attention (QK^T → scale → mask →
+// softmax → V as one dispatch instead of ~5). Elixir-side they're
+// surfaced as `Emily.Fast.*` helpers callable from inside `defn`.
+//
+// Nullable inputs (weight/bias, the RoPE `base` override, the precomp
+// `freqs`, per-tensor `offset`) marshal via `std::optional` — the same
+// pattern `ops/random.cpp` uses for PRNG keys.
+
+#include "../emily/tensor.hpp"
+
+#include <fine.hpp>
+#include <mlx/fast.h>
+#include <mlx/mlx.h>
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <variant>
+#include <vector>
+
+namespace mx = mlx::core;
+using emily::Tensor;
+using emily::unwrap_all;
+using emily::wrap;
+
+namespace {
+
+// -----------------------------------------------------------------
+// Nullable tensor helper
+// -----------------------------------------------------------------
+
+std::optional<mx::array> opt_array(
+    const std::optional<fine::ResourcePtr<Tensor>> &opt) {
+  if (opt) return (*opt)->array;
+  return std::nullopt;
+}
+
+// -----------------------------------------------------------------
+// fast_rms_norm/3 — mx::fast::rms_norm(x, weight?, eps)
+// -----------------------------------------------------------------
+//
+// Normalises the last axis of `x` by `rsqrt(mean(x^2) + eps)` and
+// optionally multiplies by `weight` (vector of size last-axis). The
+// weight is nil-able because some models (e.g. `pre_norm=False`
+// variants) use unit-scale RMSNorm.
+fine::ResourcePtr<Tensor> fast_rms_norm(
+    ErlNifEnv *,
+    fine::ResourcePtr<Tensor> x,
+    std::optional<fine::ResourcePtr<Tensor>> weight,
+    double eps) {
+  return wrap(mx::fast::rms_norm(
+      x->array, opt_array(weight), static_cast<float>(eps)));
+}
+FINE_NIF(fast_rms_norm, 0);
+
+// -----------------------------------------------------------------
+// fast_layer_norm/4 — mx::fast::layer_norm(x, weight?, bias?, eps)
+// -----------------------------------------------------------------
+//
+// Welford-style LayerNorm over the last axis with optional affine
+// (weight + bias). `weight` and `bias` are independently nullable to
+// match MLX — e.g. `elementwise_affine=False` PyTorch modules map to
+// both-nil.
+fine::ResourcePtr<Tensor> fast_layer_norm(
+    ErlNifEnv *,
+    fine::ResourcePtr<Tensor> x,
+    std::optional<fine::ResourcePtr<Tensor>> weight,
+    std::optional<fine::ResourcePtr<Tensor>> bias,
+    double eps) {
+  return wrap(mx::fast::layer_norm(
+      x->array,
+      opt_array(weight),
+      opt_array(bias),
+      static_cast<float>(eps)));
+}
+FINE_NIF(fast_layer_norm, 0);
+
+// -----------------------------------------------------------------
+// fast_rope/7 — mx::fast::rope(x, dims, traditional, base?, scale, offset, freqs?)
+// -----------------------------------------------------------------
+//
+// Fused rotary positional embedding. `dims` is the count of trailing
+// dimensions that carry rotated components (typically head_dim; the
+// remaining trailing dims, if any, are passed through). `traditional`
+// selects the paired-interleave layout (`true`, per Meta / MLX) vs
+// the split-half layout (`false`, per HuggingFace). `base` is the
+// theta override (nil → use the `freqs` argument instead); `freqs` is
+// a pre-computed 1-D tensor of inverse frequencies to support the
+// Llama-3 / LongRoPE / linear / dynamic scaling strategies that
+// Bumblebee implements outside of MLX.
+//
+// `offset` is a scalar integer tensor (Nx's canonical rep — Bumblebee
+// tracks the cumulative position offset as an %Nx.Tensor{} through
+// iterative decode), which matches the `array`-offset overload of
+// MLX's `rope`. The NIF always takes a tensor here and uses the
+// overload with `const array&` — users pass `Nx.tensor(0)` when
+// there's no KV-cache offset.
+fine::ResourcePtr<Tensor> fast_rope(
+    ErlNifEnv *,
+    fine::ResourcePtr<Tensor> x,
+    int64_t dims,
+    bool traditional,
+    std::optional<double> base,
+    double scale,
+    fine::ResourcePtr<Tensor> offset,
+    std::optional<fine::ResourcePtr<Tensor>> freqs) {
+  std::optional<float> base_f;
+  if (base) base_f = static_cast<float>(*base);
+
+  return wrap(mx::fast::rope(
+      x->array,
+      static_cast<int>(dims),
+      traditional,
+      base_f,
+      static_cast<float>(scale),
+      offset->array,
+      opt_array(freqs)));
+}
+FINE_NIF(fast_rope, 0);
+
+// -----------------------------------------------------------------
+// fast_scaled_dot_product_attention/6 —
+//   mx::fast::scaled_dot_product_attention(Q, K, V, scale, mask_mode, mask_arrs)
+// -----------------------------------------------------------------
+//
+// Computes `softmax((Q @ Kᵀ) * scale + mask) @ V` as a single fused
+// kernel over `[B, H, S, D]` inputs.
+//
+// `mask_mode` is the empty string, `"causal"`, or `"array"`:
+//   - `""`        — no mask.
+//   - `"causal"`  — upper-triangular -inf mask (no additional arrays).
+//   - `"array"`   — `mask_arrs` holds one broadcastable additive bias
+//                   tensor (Bumblebee's `bias = select(mask, 0, -inf)`
+//                   materialises this).
+//
+// MLX supports a handful of other modes (block-sparse etc.) — out of
+// scope for M11; add them when a model asks.
+fine::ResourcePtr<Tensor> fast_scaled_dot_product_attention(
+    ErlNifEnv *,
+    fine::ResourcePtr<Tensor> q,
+    fine::ResourcePtr<Tensor> k,
+    fine::ResourcePtr<Tensor> v,
+    double scale,
+    std::string mask_mode,
+    std::vector<fine::ResourcePtr<Tensor>> mask_arrs) {
+  return wrap(mx::fast::scaled_dot_product_attention(
+      q->array,
+      k->array,
+      v->array,
+      static_cast<float>(scale),
+      mask_mode,
+      unwrap_all(mask_arrs)));
+}
+FINE_NIF(fast_scaled_dot_product_attention, 0);
+
+} // namespace