Merge pull request #200 from ausimian/perf/fused-quantized-matmul

perf: fuse quantized_dense onto mx::quantized_matmul kernel

Author: ausimian

RELEASE.md

@@ -0,0 +1,10 @@
+### Changed
+
+- Quantized dense layers now use the fused `mx::quantized_matmul` kernel
+  instead of dequantizing the full weight to bf16 and running a dense
+  matmul. The packed low-bit weights are streamed directly, so a decode
+  step no longer re-dequantizes the entire model on every token. On a
+  4-bit Qwen3-0.6B this makes native quantized generation roughly 13×
+  faster end-to-end — and quantized inference is now *faster* than dense,
+  as it should be, rather than slower. Non-Emily backends keep the
+  composed dequantize + `Nx.dot` fallback.

bench/qmm_microbench.exs

@@ -0,0 +1,91 @@
+# Micro-benchmark: fused quantized_matmul (mx::quantized_matmul) vs the
+# current quantized_dense path (dequantize_defn + Nx.dot), on GPU.
+# No Bumblebee / model download needed. Single-token (batch=1) decode-shaped.
+#
+#   mix run bench/qmm_microbench.exs
+alias Emily.Quantization
+alias Emily.Quantization.Layers
+alias Emily.QuantizedWeight
+
+Nx.default_backend(Emily.Backend)
+
+native = [compiler: Emily.Compiler, native: true]
+
+dtype = :bf16
+group_size = 64
+bits = 4
+warmup = 100
+iters = 2000
+
+# Qwen3-0.6B-shaped projections. Weight is [out, in] (transpose: true, the
+# from_dense default); activation is [1, in] (one decode token).
+shapes = [
+  {"q_proj  [2048,1024]", 2048, 1024},
+  {"kv_proj [1024,1024]", 1024, 1024},
+  {"o_proj  [1024,2048]", 1024, 2048},
+  {"mlp_up  [3072,1024]", 3072, 1024},
+  {"mlp_dn  [1024,3072]", 1024, 3072}
+]
+
+# Force a full worker sync on the result (native :sync already blocks on
+# mx::eval, but realizing the bytes is belt-and-suspenders).
+sync = fn t -> Nx.to_binary(t) end
+
+time_fn = fn compiled, x, qw ->
+  Enum.each(1..warmup, fn _ -> compiled.(x, qw) end)
+  sync.(compiled.(x, qw))
+  t0 = System.monotonic_time(:microsecond)
+  Enum.each(1..iters, fn _ -> compiled.(x, qw) end)
+  sync.(compiled.(x, qw))
+  t1 = System.monotonic_time(:microsecond)
+  iters / ((t1 - t0) / 1_000_000)
+end
+
+IO.puts("dtype=#{dtype} group_size=#{group_size} bits=#{bits} warmup=#{warmup} iters=#{iters}\n")
+IO.puts("  before = old quantized_dense (dequantize_defn + Nx.dot)")
+IO.puts("  after  = quantized_dense now (#197: fused mx::quantized_matmul)\n")
+
+IO.puts(String.pad_trailing("shape", 22) <> "  before(it/s)  after(it/s)  speedup  maxΔ")
+IO.puts(String.duplicate("-", 68))
+
+for {label, out_f, in_f} <- shapes do
+  {w, _} = Nx.Random.normal(Nx.Random.key(0), shape: {out_f, in_f}, type: dtype)
+  qw = QuantizedWeight.from_dense(w, group_size: group_size, bits: bits)
+  {x, _} = Nx.Random.normal(Nx.Random.key(1), shape: {1, in_f}, type: dtype)
+
+  # Pass the QuantizedWeight (an Nx.Container) as a jit ARGUMENT, not a
+  # closure — its tensors become Expr params and its keep-metadata
+  # (group_size/bits/transpose/mode) stays available at trace time. This
+  # mirrors how Bumblebee threads quantized model params into the forward.
+  #
+  # `before` = the old layer body (dequantize the full bf16 weight, then
+  # dense Nx.dot). `after` = the shipped layer, which now lowers to the
+  # fused mx::quantized_matmul kernel.
+  before = fn x, qw -> Nx.dot(x, Nx.transpose(Quantization.dequantize_defn(qw))) end
+  after_fn = fn x, qw -> Layers.quantized_dense(x, qw) end
+
+  before_compiled = Nx.Defn.jit(before, native)
+  after_compiled = Nx.Defn.jit(after_fn, native)
+
+  # correctness: after (fused) vs before (dequant), same math up to fp reorder
+  b = before_compiled.(x, qw)
+  a = after_compiled.(x, qw)
+  max_delta = Nx.subtract(b, a) |> Nx.abs() |> Nx.reduce_max() |> Nx.to_number()
+
+  before_rate = time_fn.(before_compiled, x, qw)
+  after_rate = time_fn.(after_compiled, x, qw)
+
+  IO.puts(
+    String.pad_trailing(label, 22) <>
+      "  " <>
+      String.pad_trailing(:erlang.float_to_binary(before_rate, decimals: 0), 12) <>
+      "  " <>
+      String.pad_trailing(:erlang.float_to_binary(after_rate, decimals: 0), 11) <>
+      "  " <>
+      String.pad_trailing(
+        :erlang.float_to_binary(after_rate / before_rate, decimals: 2) <> "x",
+        7
+      ) <>
+      "  " <> :erlang.float_to_binary(max_delta, decimals: 4)
+  )
+end

bench/qwen3_quantized_tps.exs

@@ -0,0 +1,102 @@
+# End-to-end quantized Qwen3-0.6B greedy-decode throughput on Emily's
+# native lane. Loads the dense model, rewrites every dense layer to
+# `Emily.Quantization.Layers.quantized_dense/4` + quantizes the params
+# (affine 4-bit) via `Emily.Quantization.Transform`, then measures
+# tokens/sec through a compiled `Bumblebee.Text.generation` serving.
+#
+# Run in TEST env so `Emily.Quantization.Transform` (test/support) is on
+# the compile path, and in-project (NOT Mix.install) so it exercises the
+# LOCAL build:
+#
+#     MIX_ENV=test mix run bench/qwen3_quantized_tps.exs
+#
+# Optional env: EMILY_BENCH_MODEL, EMILY_BENCH_NEW_TOKENS (64),
+# EMILY_BENCH_WARMUP (1), EMILY_BENCH_RUNS (3), EMILY_BENCH_GROUP_SIZE (64).
+#
+# Greedy decode is deterministic, so before/after (dequant+dot vs fused)
+# generate the identical token sequence — the wall-clock RATIO is exact
+# even if generation stops before max_new_tokens.
+
+Nx.global_default_backend(Emily.Backend)
+
+env_int = fn name, default ->
+  case System.get_env(name) do
+    nil -> default
+    s -> case Integer.parse(s), do: ({n, _} -> n; _ -> default)
+  end
+end
+
+repo = System.get_env("EMILY_BENCH_MODEL", "Qwen/Qwen3-0.6B")
+prompt = System.get_env("EMILY_BENCH_PROMPT", "The quick brown fox jumps over the lazy dog.")
+new_tokens = env_int.("EMILY_BENCH_NEW_TOKENS", 64)
+warmup = env_int.("EMILY_BENCH_WARMUP", 1)
+runs = env_int.("EMILY_BENCH_RUNS", 3)
+group_size = env_int.("EMILY_BENCH_GROUP_SIZE", 64)
+
+IO.puts("Emily / Qwen3 QUANTIZED (affine 4-bit, group_size=#{group_size}) throughput")
+IO.puts("  model      : #{repo}")
+IO.puts("  new tokens : #{new_tokens}   warmup: #{warmup}   runs: #{runs}")
+IO.puts("  lane       : native (Emily.Compiler, native: true, native_fallback: :raise)\n")
+
+{:ok, model_info} = Bumblebee.load_model({:hf, repo})
+{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, repo})
+{:ok, gen_config} = Bumblebee.load_generation_config({:hf, repo})
+
+gen_config =
+  Bumblebee.configure(gen_config,
+    max_new_tokens: new_tokens,
+    strategy: %{type: :greedy_search}
+  )
+
+# EMILY_BENCH_QUANTIZE=0 skips quantization → a dense-native calibration
+# lane (anchors what "as fast as it should be" looks like on this host).
+quantize? = System.get_env("EMILY_BENCH_QUANTIZE", "1") != "0"
+
+serving_model_info =
+  if quantize? do
+    IO.puts("Quantizing dense layers -> quantized_dense (%QuantizedWeight{})…")
+
+    {qmodel, qparams} =
+      Emily.Quantization.Transform.quantize(model_info.model, model_info.params,
+        bits: 4,
+        group_size: group_size,
+        transpose: true
+      )
+
+    %{model_info | model: qmodel, params: qparams}
+  else
+    IO.puts("Dense (no quantization) calibration lane")
+    model_info
+  end
+
+serving =
+  Bumblebee.Text.generation(serving_model_info, tokenizer, gen_config,
+    defn_options: [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+  )
+
+for _ <- 1..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)
+secs_list = Enum.map(measurements, fn {secs, _, _} -> secs end)
+[{_, _, sample} | _] = measurements
+mean = Enum.sum(tps_list) / length(tps_list)
+{min_tps, max_tps} = Enum.min_max(tps_list)
+median_secs = secs_list |> Enum.sort() |> Enum.at(div(length(secs_list), 2))
+
+IO.puts("\ntokens/sec  : mean=#{Float.round(mean, 2)}  min=#{Float.round(min_tps, 2)}  max=#{Float.round(max_tps, 2)}")
+IO.puts("median secs : #{Float.round(median_secs, 4)} (use the ratio of this across before/after runs)")
+IO.puts("sample      :\n  #{String.slice(sample, 0, 300)}")

lib/emily/quantization/layers.ex

@@ -3,10 +3,13 @@ defmodule Emily.Quantization.Layers do
   Defn-traceable quantized layer op for use inside Axon graphs.
 
   `quantized_dense/4` is the drop-in replacement for `Axon.Layers.dense/4`
-  on a `%Emily.QuantizedWeight{}` kernel. See `Emily.Quantization` for
-  the defn-integration trade-offs; the `qwen3_quantized` notebook walks
-  through a concrete `Axon.rewrite_nodes/2`-based graph rewrite that
-  swaps every `:dense` for a layer calling this op.
+  on a `%Emily.QuantizedWeight{}` kernel. It lowers to the fused
+  `mx::quantized_matmul` kernel (via `Emily.Quantization.quantized_matmul_defn/2`),
+  which streams the packed low-bit weights directly rather than
+  materializing a dense weight per call — the single-kernel path decode
+  loops want. The `qwen3_quantized` notebook walks through a concrete
+  `Axon.rewrite_nodes/2`-based graph rewrite that swaps every `:dense`
+  for a layer calling this op.
   """
 
   import Nx.Defn
@@ -20,14 +23,15 @@ defmodule Emily.Quantization.Layers do
 
     * `input` — activation tensor, shape `(..., in)`.
     * `kernel` — `%QuantizedWeight{}`. The stored layout is determined
-      by `kernel.transpose`:
+      by `kernel.transpose` (passed straight through to the fused
+      kernel):
         * `transpose: false` (the AWQ / Axon-native layout) — packed
           representation of a `[in, out]` weight; the layer computes
-          `Nx.dot(x, dense)`.
+          `x @ W`.
         * `transpose: true` (the MLX / PyTorch-native layout, i.e. fresh
           output of `QuantizedWeight.from_dense/2` on a `[out, in]`
           weight) — packed representation of a `[out, in]` weight; the
-          layer computes `Nx.dot(x, Nx.transpose(dense))`.
+          layer computes `x @ Wᵀ`.
     * `bias` — either an `Nx.Tensor`, a number, or a keyword list (in
       which case it's treated as `opts` and bias defaults to 0). Matches
       `Axon.Quantization.Layers.weight_only_quantized_dense/4`'s
@@ -50,33 +54,33 @@ defmodule Emily.Quantization.Layers do
     # When Axon.dense registers `use_bias: false`, the generated op call
     # is arity-3 with layer opts as the third arg (matches
     # `Axon.Quantization.Layers.weight_only_quantized_dense/4`'s contract).
-    {bias, opts} =
+    # Axon also injects `:mode` (`:inference` / `:train`); weight-only
+    # quantization has no mode-dependent behaviour, so `opts` is absorbed
+    # and ignored here.
+    {bias, _opts} =
       case bias do
         b when is_list(b) -> {Nx.tensor(0), Keyword.merge(opts, b)}
         b -> {b, opts}
       end
 
-    %QuantizedWeight{transpose: transpose} = kernel
-    opts = Keyword.put(opts, :transpose, transpose)
-    quantized_dense_impl(input, kernel, bias, opts)
+    # Assert the kernel is a %QuantizedWeight{} at the layer boundary so a
+    # bad kernel fails here rather than deep inside the fused kernel. Its
+    # layout/mode/bits/group_size are read off the struct by
+    # `quantized_matmul_defn/2`, so nothing extra needs threading through.
+    %QuantizedWeight{} = kernel
+    quantized_dense_impl(input, kernel, bias)
   end
 
-  # `transpose` is threaded through `opts` as a compile-time constant so
-  # the branch selects at trace time (no runtime `if` over booleans).
-  defnp quantized_dense_impl(x, kernel, bias, opts \\ []) do
-    # `:mode` is injected by Axon's compiler (`:inference` / `:train`)
-    # for every layer op; accept-and-ignore here since weight-only
-    # quantization has no mode-dependent behavior.
-    opts = keyword!(opts, [:transpose, mode: :inference])
-    dense = Emily.Quantization.dequantize_defn(kernel)
-
-    y =
-      if opts[:transpose] do
-        Nx.dot(x, Nx.transpose(dense))
-      else
-        Nx.dot(x, dense)
-      end
-
-    Nx.add(y, bias)
+  defnp quantized_dense_impl(x, kernel, bias) do
+    # Fused single-kernel `mx::quantized_matmul` — streams the 4-bit
+    # packed weights directly — instead of dequantizing the full weight
+    # to bf16 and then running a dense `Nx.dot`. Decode is
+    # memory-bandwidth bound on the weight, so this is ~2-4x faster per
+    # matmul (larger gains on the fatter MLP projections). `transpose`,
+    # `mode`, `bits`, and `group_size` are read off the `%QuantizedWeight{}`
+    # by `quantized_matmul_defn/2`; non-Emily backends still get the
+    # composed `dequantize_defn/1` + `Nx.dot/2` via the block's fallback.
+    Emily.Quantization.quantized_matmul_defn(x, kernel)
+    |> Nx.add(bias)
   end
 end

test/emily/quantization/layers_test.exs

@@ -137,4 +137,56 @@ defmodule Emily.Quantization.LayersTest do
       assert_close(actual, expected, tol: 1.0e-3)
     end
   end
+
+  describe "quantized_dense/4 — native lane lowers to the fused kernel" do
+    # #197: under the native single-NIF compiler (the decode-loop path),
+    # the layer must lower to the fused `mx::quantized_matmul` opcode, not
+    # fall back to op-by-op eval. `native_fallback: :raise` fails the test
+    # if any node can't be lowered natively; the eager `quantized_matmul/2`
+    # (same C++ kernel) is the numeric oracle.
+    @native_raise [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+
+    test "transpose=true: fused native lowering matches eager quantized_matmul" do
+      w =
+        Nx.iota({4, 128}, backend: Emily.Backend, type: :f32)
+        |> Nx.divide(4 * 128 / 2)
+        |> Nx.subtract(1.0)
+
+      x =
+        Nx.iota({3, 128}, backend: Emily.Backend, type: :f32)
+        |> Nx.divide(128.0)
+        |> Nx.subtract(0.5)
+
+      qw = QuantizedWeight.from_dense(w, group_size: 64, bits: 4, transpose: true)
+
+      fun = fn x, qw -> Layers.quantized_dense(x, qw) end
+      actual = Nx.Defn.jit(fun, @native_raise).(x, qw)
+      expected = Quantization.quantized_matmul(x, qw)
+
+      assert Nx.shape(actual) == {3, 4}
+      assert_close(actual, expected, tol: 1.0e-3)
+    end
+
+    test "transpose=false + bias: fused native lowering matches eager" do
+      w =
+        Nx.iota({2, 128}, backend: Emily.Backend, type: :f32)
+        |> Nx.divide(2 * 128 / 2)
+        |> Nx.subtract(1.0)
+
+      x =
+        Nx.iota({3, 2}, backend: Emily.Backend, type: :f32)
+        |> Nx.divide(2.0)
+        |> Nx.subtract(0.5)
+
+      qw = QuantizedWeight.from_dense(w, group_size: 64, bits: 4, transpose: false)
+      b = Nx.iota({128}, backend: Emily.Backend, type: :f32) |> Nx.divide(128.0)
+
+      fun = fn x, qw, b -> Layers.quantized_dense(x, qw, b) end
+      actual = Nx.Defn.jit(fun, @native_raise).(x, qw, b)
+      expected = Quantization.quantized_matmul(x, qw) |> Nx.add(b)
+
+      assert Nx.shape(actual) == {3, 128}
+      assert_close(actual, expected, tol: 1.0e-3)
+    end
+  end
 end