Merge pull request #11 from ausimian/feat/m4-qwen3

M4: Bumblebee Qwen3 end-to-end on Emily.Backend

Author: ausimian

RELEASE.md

@@ -110,6 +110,69 @@
   last-axis fast path stays on MLX; interior-axis usage is rare on
   our M3/M4 critical path (transformer inference doesn't need it).
 
+- M4 — Qwen3 inference. `Qwen/Qwen3-0.6B` greedy-decodes end-to-end on
+  `Emily.Backend` through Bumblebee's causal-LM serving. Everything on
+  Qwen3's critical path (QK-norm, rotary embeddings, GQA, SwiGLU FFN,
+  RMSNorm, tied embeddings, KV-cache `put_slice` in a `defn` while
+  loop) runs correctly.
+  - **Native `put_slice/4`** in `Emily.Backend`, backed by a new
+    `Native.slice_update/3` NIF over `mx::slice_update`. Replaces the
+    BinaryBackend round-trip — autoregressive decoding calls
+    `put_slice` per layer per token to append into the KV cache, and
+    the old fallback transferred ~1 MB of cache state through the
+    allocator on every call. Also fixes a latent bug in the old
+    implementation: dynamic scalar-tensor `start_indices` on
+    `Emily.Backend` used to slip through unconverted and crash inside
+    BinaryBackend. `slice_start` is now applied to every start
+    index, matching the `slice/5` callback.
+  - **Operand-type promotion in `put_slice`.** `Nx.put_slice`
+    promotes the output type across tensor/update (an s32 pad buffer
+    clashing with an s64 decoder input becomes s64), but the backend
+    callback still receives the original-type operands. We cast both
+    `t` and `slice` to `out.type` via `Native.astype` before
+    dispatching to `slice_update`. Without this the MLX buffer
+    silently disagrees with the Nx shape metadata — the first symptom
+    is `Nx.to_binary` returning a half-sized binary and
+    `BinaryBackend.bitstring_part` raising a match error deep inside
+    the tokenizer decode. Mirrors the arithmetic-op promotion fix
+    landed in M3.
+  - **`test/emily/conformance/qwen3_test.exs`**
+    (`@moduletag :conformance`) — ports `Bumblebee.Text.Qwen3Test`
+    verbatim (three architectures: `:base`,
+    `:for_causal_language_modeling`, `:for_sequence_classification`),
+    with HF reference slices checked in, plus a `greedy generation`
+    describe block that drives
+    `Bumblebee.Text.Generation.build_generate` on the tiny-random
+    causal LM. That smoke test feeds synthetic `input_ids` in
+    `[0, 1024)` (tokenizer vocab is 151 k but the tiny checkpoint's
+    embedding is 1024 rows), greedy-decodes 16 tokens through the
+    full generation pipeline (`Axon.predict` + logit processing +
+    `Nx.argmax` + `put_slice` KV-cache update + `defn while`), and
+    asserts bit-exact equality against both `Nx.BinaryBackend` run
+    on the same inputs *and* a checked-in 16-token reference.
+  - **`test/emily/conformance/qwen3_full_test.exs`**
+    (`@moduletag :qwen3_full`, excluded from `--only conformance`
+    because the checkpoint is ~1.5 GB) — loads `Qwen/Qwen3-0.6B`
+    proper, greedy-decodes 32 tokens from a fixed prompt through
+    `Nx.Serving`, and asserts the completion string matches a
+    checked-in reference. Run with `mix test --only qwen3_full`.
+  - **`bench/qwen3_tokens_per_sec.exs`** — standalone wall-clock
+    throughput harness. Loads `Qwen/Qwen3-0.6B`, runs N warmup
+    iterations + M measured iterations of greedy decode, reports
+    tokens/sec. Prompt, token count, and iteration counts are
+    overridable via `EMILY_BENCH_*` env vars. Baseline observed on a
+    dev M3 host: ~13.8 tok/s at 16 new tokens under the
+    `Nx.Defn.Evaluator` compiler (no `mlx::core::compile` wrap yet —
+    that lands in M6). Intended as a regression gate, not a headline
+    number.
+  - **Bumblebee dependency** bumped from Hex 0.6.3 to a pinned `main`
+    commit (`273805e9…`) so `Bumblebee.Text.Qwen3` is available — the
+    text port is on main but not yet in a Hex release. Revert to a
+    Hex version as soon as one ships Qwen3 support.
+  - **`test_helper.exs`** extended the exclude list with
+    `:qwen3_full` so the weights-heavy test stays out of
+    `mix test --only conformance`.
+
 - M3 — DistilBERT end-to-end on Bumblebee. Every Nx op on the
   transformer critical path now runs natively on MLX; the full
   forward pass matches HuggingFace Transformers (PyTorch) reference

bench/qwen3_tokens_per_sec.exs

@@ -0,0 +1,105 @@
+# Qwen3-0.6B greedy-decode throughput on `Emily.Backend`.
+#
+# Usage:
+#
+#     MIX_ENV=test mix run bench/qwen3_tokens_per_sec.exs
+#
+# 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.
+#
+# 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.
+
+defmodule Emily.Bench.Qwen3 do
+  @default_model "Qwen/Qwen3-0.6B"
+  @default_prompt "The quick brown fox jumps over the lazy dog."
+  @default_new_tokens 64
+  @default_warmup 1
+  @default_runs 3
+
+  def run do
+    Nx.global_default_backend(Emily.Backend)
+
+    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)
+
+    warmup = System.get_env("EMILY_BENCH_WARMUP") |> env_int(@default_warmup)
+    runs = System.get_env("EMILY_BENCH_RUNS") |> env_int(@default_runs)
+
+    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("")
+
+    {:ok, model_info} = Bumblebee.load_model({:hf, model_repo})
+    {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, model_repo})
+
+    {:ok, generation_config} = Bumblebee.load_generation_config({:hf, model_repo})
+
+    generation_config =
+      Bumblebee.configure(generation_config,
+        max_new_tokens: new_tokens,
+        strategy: %{type: :greedy_search}
+      )
+
+    serving =
+      Bumblebee.Text.generation(model_info, tokenizer, generation_config,
+        defn_options: [compiler: Nx.Defn.Evaluator]
+      )
+
+    for _ <- Stream.duplicate(:ok, warmup) do
+      IO.puts("[warmup] generating…")
+      %{results: [_]} = Nx.Serving.run(serving, prompt)
+    end
+
+    measurements =
+      for n <- 1..runs//1 do
+        {elapsed_us, %{results: [%{text: text}]}} =
+          :timer.tc(fn -> Nx.Serving.run(serving, prompt) end)
+
+        secs = elapsed_us / 1_000_000
+        tps = new_tokens / secs
+        IO.puts("[run #{n}] #{Float.round(secs, 3)} s, #{Float.round(tps, 2)} tok/s")
+        {secs, tps, text}
+      end
+
+    tps_list = Enum.map(measurements, fn {_, tps, _} -> tps end)
+    [{_, _, sample} | _] = measurements
+    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))
+  end
+
+  defp env_int(nil, default), do: default
+  defp env_int(s, default) do
+    case Integer.parse(s) do
+      {n, ""} -> n
+      _ -> default
+    end
+  end
+end
+
+Emily.Bench.Qwen3.run()

c_src/ops/index.cpp

@@ -31,6 +31,27 @@ fine::ResourcePtr<Tensor> slice(
 }
 FINE_NIF(slice, 0);
 
+// slice_update/3: write `update` into `src` starting at `start`. `stop`
+// is derived as `start + shape(update)` and strides default to 1 on
+// every axis (Nx.put_slice has no stride parameter). Output shape
+// equals `src.shape`.
+fine::ResourcePtr<Tensor> slice_update(
+    ErlNifEnv *,
+    fine::ResourcePtr<Tensor> src,
+    fine::ResourcePtr<Tensor> update,
+    std::vector<int64_t> start) {
+  const auto &update_shape = update->array.shape();
+  mx::Shape start_shape = to_mlx_shape(start);
+  mx::Shape stop_shape;
+  stop_shape.reserve(start_shape.size());
+  for (size_t i = 0; i < start_shape.size(); ++i) {
+    stop_shape.push_back(start_shape[i] + update_shape[i]);
+  }
+  return wrap(mx::slice_update(
+      src->array, update->array, std::move(start_shape), std::move(stop_shape)));
+}
+FINE_NIF(slice_update, 0);
+
 // take/3: gather along `axis` using integer indices.
 fine::ResourcePtr<Tensor> take(
     ErlNifEnv *,

lib/emily/backend.ex

@@ -592,13 +592,21 @@ defmodule Emily.Backend do
   defp slice_start(i) when is_integer(i), do: i
   defp slice_start(%T{} = t), do: t |> Nx.backend_copy(Nx.BinaryBackend) |> Nx.to_number()
 
-  # put_slice: MLX has no direct primitive; route via BinaryBackend.
-  # Nx's Backend contract order is (out, tensor, start_indices, slice).
-  # start_indices arrive as scalar tensors on Emily.Backend; Nx auto-
-  # transfers them when the BinaryBackend call goes through `to_indices`.
+  # put_slice: implemented natively via MLX `slice_update`. Nx promotes
+  # operand types at the API layer — `Nx.put_slice(s32_buf, _, s64_upd)`
+  # declares an s64 output — but our callback arguments still carry the
+  # original backend types. We cast both `t` and `slice` to `out.type`
+  # before dispatching so the MLX buffer matches Nx's shape/type view.
+  # Scalar-tensor starts are materialised to integers here (dynamic
+  # indices show up when autoregressive loops dispatch put_slice from
+  # within `defn`).
   @impl true
-  def put_slice(out, t, start_indices, slice),
-    do: via_binary(out, [t, slice], &Nx.put_slice(&1, start_indices, &2))
+  def put_slice(%T{type: type} = out, %T{} = t, start_indices, %T{} = slice) do
+    starts = Enum.map(start_indices, &slice_start/1)
+    src_ref = Native.astype(ref(t), type)
+    update_ref = Native.astype(ref(slice), type)
+    Native.slice_update(src_ref, update_ref, starts) |> wrap(out)
+  end
 
   @impl true
   def select(%T{} = out, pred, on_true, on_false) do

lib/emily/native.ex

@@ -221,6 +221,9 @@ defmodule Emily.Native do
   @spec slice(tensor(), [integer()], [integer()], [integer()]) :: tensor()
   def slice(_a, _start, _stop, _strides), do: nif()
 
+  @spec slice_update(tensor(), tensor(), [integer()]) :: tensor()
+  def slice_update(_src, _update, _start), do: nif()
+
   @spec take(tensor(), tensor(), integer()) :: tensor()
   def take(_a, _indices, _axis), do: nif()
 

mix.exs

@@ -52,7 +52,12 @@ defmodule Emily.MixProject do
       {:elixir_make, "~> 0.9"},
       {:fine, "~> 0.1"},
       {:nx, "~> 0.10"},
-      {:bumblebee, "~> 0.6", only: :test},
+      # Bumblebee >= 0.6.3 (the latest Hex release) lacks Qwen3 support.
+      # Pinned to a `main` commit that contains `Bumblebee.Text.Qwen3` so
+      # M4 can exercise Qwen3-0.6B end-to-end. Bump deliberately when a
+      # newer release lands on Hex.
+      {:bumblebee,
+       github: "elixir-nx/bumblebee", ref: "273805e95507dc7866b958d90e0012a3abad1761", only: :test},
       {:tokenizers, "~> 0.5", only: :test},
       {:stream_data, "~> 1.1", only: [:dev, :test]},
       {:credo, "~> 1.7", only: [:dev, :test], runtime: false},

mix.lock

@@ -1,6 +1,6 @@
 %{
   "axon": {:hex, :axon, "0.7.0", "2e2c6d93b4afcfa812566b8922204fa022b60081e86ebd411df4db7ea30f5457", [:mix], [{:kino, "~> 0.7", [hex: :kino, repo: "hexpm", optional: true]}, {:kino_vega_lite, "~> 0.1.7", [hex: :kino_vega_lite, repo: "hexpm", optional: true]}, {:nx, "~> 0.9", [hex: :nx, repo: "hexpm", optional: false]}, {:polaris, "~> 0.1", [hex: :polaris, repo: "hexpm", optional: false]}, {:table_rex, "~> 3.1.1", [hex: :table_rex, repo: "hexpm", optional: true]}], "hexpm", "ee9857a143c9486597ceff434e6ca833dc1241be6158b01025b8217757ed1036"},
-  "bumblebee": {:hex, :bumblebee, "0.6.3", "c0028643c92de93258a9804da1d4d48797eaf7911b702464b3b3dd2cc7f938f1", [:mix], [{:axon, "~> 0.7.0", [hex: :axon, repo: "hexpm", optional: false]}, {:jason, "~> 1.4.0", [hex: :jason, repo: "hexpm", optional: false]}, {:nx, "~> 0.9.0 or ~> 0.10.0", [hex: :nx, repo: "hexpm", optional: false]}, {:nx_image, "~> 0.1.0", [hex: :nx_image, repo: "hexpm", optional: false]}, {:nx_signal, "~> 0.2.0", [hex: :nx_signal, repo: "hexpm", optional: false]}, {:progress_bar, "~> 3.0", [hex: :progress_bar, repo: "hexpm", optional: false]}, {:safetensors, "~> 0.1.3", [hex: :safetensors, repo: "hexpm", optional: false]}, {:tokenizers, "~> 0.4", [hex: :tokenizers, repo: "hexpm", optional: false]}, {:unpickler, "~> 0.1.0", [hex: :unpickler, repo: "hexpm", optional: false]}, {:unzip, "~> 0.12.0", [hex: :unzip, repo: "hexpm", optional: false]}], "hexpm", "c619197787561f8e5fb2ffba269c341654accaec9d591999b7fddd55761dd079"},
+  "bumblebee": {:git, "https://github.com/elixir-nx/bumblebee.git", "273805e95507dc7866b958d90e0012a3abad1761", [ref: "273805e95507dc7866b958d90e0012a3abad1761"]},
   "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"},
   "castore": {:hex, :castore, "1.0.18", "5e43ef0ec7d31195dfa5a65a86e6131db999d074179d2ba5a8de11fe14570f55", [:mix], [], "hexpm", "f393e4fe6317829b158fb74d86eb681f737d2fe326aa61ccf6293c4104957e34"},
   "complex": {:hex, :complex, "0.6.0", "b0130086a7a8c33574d293b2e0e250f4685580418eac52a5658a4bd148f3ccf1", [:mix], [], "hexpm", "0a5fa95580dcaf30fcd60fe1aaf24327c0fe401e98c24d892e172e79498269f9"},
@@ -27,5 +27,5 @@
   "telemetry": {:hex, :telemetry, "1.4.1", "ab6de178e2b29b58e8256b92b382ea3f590a47152ca3651ea857a6cae05ac423", [:rebar3], [], "hexpm", "2172e05a27531d3d31dd9782841065c50dd5c3c7699d95266b2edd54c2dafa1c"},
   "tokenizers": {:hex, :tokenizers, "0.5.1", "b0975d92b4ee5b18e8f47b5d65b9d5f1e583d9130189b1a2620401af4e7d4b35", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:rustler, ">= 0.0.0", [hex: :rustler, repo: "hexpm", optional: true]}, {:rustler_precompiled, "~> 0.6", [hex: :rustler_precompiled, repo: "hexpm", optional: false]}], "hexpm", "5f08d97cc7f2ed3d71d370d68120da6d3de010948ccf676c9c0eb591ba4bacc9"},
   "unpickler": {:hex, :unpickler, "0.1.0", "c2262c0819e6985b761e7107546cef96a485f401816be5304a65fdd200d5bd6a", [:mix], [], "hexpm", "e2b3f61e62406187ac52afead8a63bfb4e49394028993f3c4c42712743cab79e"},
-  "unzip": {:hex, :unzip, "0.12.0", "beed92238724732418b41eba77dcb7f51e235b707406c05b1732a3052d1c0f36", [:mix], [], "hexpm", "95655b72db368e5a84951f0bed586ac053b55ee3815fd96062fce10ce4fc998d"},
+  "unzip": {:hex, :unzip, "0.13.0", "bf5ec6ac6063c69e6ec54c8b4a3b8dcd7a2719d28d10d7025776ab107957cde9", [:mix], [], "hexpm", "4bcb9892ecbf2042606b43ab685a1bffe03c14003e6246f5453db2c829237fd9"},
 }

test/emily/conformance/qwen3_full_test.exs

@@ -0,0 +1,55 @@
+defmodule Emily.Conformance.Qwen3FullTest do
+  @moduledoc """
+  Full `Qwen/Qwen3-0.6B` end-to-end conformance test.
+
+  This test is excluded even from `mix test --only conformance`: the
+  model is ~1.5 GB on first fetch, so blowing it out into CI on every
+  push is the wrong default. Run explicitly:
+
+      mix test --only qwen3_full
+
+  The reference text pinned below is the greedy decode produced by
+  `Emily.Backend` on an Apple-Silicon host. A failure means the
+  backend has drifted, Bumblebee's Qwen3 port has changed, or the HF
+  checkpoint has been republished — all of which are real signals.
+  """
+
+  use ExUnit.Case, async: false
+
+  @moduletag :qwen3_full
+  @moduletag capture_log: true
+  @moduletag timeout: 600_000
+
+  @prompt "The quick brown fox jumps over the lazy dog."
+  @reference_text " The quick brown fox is a character in the story. The quick brown fox is a character in the story. The quick brown fox is a character in the story"
+
+  setup_all do
+    prev = Nx.default_backend()
+    Nx.global_default_backend(Emily.Backend)
+    on_exit(fn -> Nx.global_default_backend(prev) end)
+    :ok
+  end
+
+  test "Qwen/Qwen3-0.6B greedy decodes the pinned continuation" do
+    {:ok, model_info} = Bumblebee.load_model({:hf, "Qwen/Qwen3-0.6B"})
+    {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "Qwen/Qwen3-0.6B"})
+    {:ok, generation_config} = Bumblebee.load_generation_config({:hf, "Qwen/Qwen3-0.6B"})
+
+    config =
+      Bumblebee.configure(generation_config,
+        max_new_tokens: 32,
+        strategy: %{type: :greedy_search}
+      )
+
+    serving =
+      Bumblebee.Text.generation(model_info, tokenizer, config,
+        defn_options: [compiler: Nx.Defn.Evaluator]
+      )
+
+    %{results: [%{text: text, token_summary: summary}]} =
+      Nx.Serving.run(serving, @prompt)
+
+    assert summary.output == 32
+    assert text == @reference_text
+  end
+end