Merge pull request #20 from ausimian/m10.5-bumblebee-quantized-inference

M10.5: Bumblebee quantized inference integration

Author: ausimian

PLAN.md

@@ -411,46 +411,74 @@ comparison; quantized memory soak clean.
 
 ### M10.5 — Bumblebee quantized inference integration
 
-Closes the gap M10 left open: getting `Native.quantized_matmul`
-reachable from `Nx.Defn.jit`-traced Axon forward passes so Bumblebee's
-AWQ-loading (when it lands) routes through the fused kernel.
-
-Approach choices (pick before starting):
-
-1. **Defn-native dequantize** — implement MLX's int4/int8 affine
-   dequantize using Nx bit primitives (right-shift + mask + multiply +
-   add). `Emily.Quantization.Layers.quantized_dense/3` becomes
-   `Nx.dot(x, dequantize_defn(qw))`. Correct and unblocks the full
-   Axon/Bumblebee path, but uses two kernels (dequantize + matmul)
-   instead of MLX's fused one — M11's fast-kernel work subsumes the
-   perf gap.
-2. **Emily.Compiler custom-op intercept** — fork `Nx.Defn.Evaluator`
-   under Emily to recognise a sentinel `Expr` node and route to
-   `Native.quantized_matmul`. Full fused-kernel story but large
-   surface; fragile against upstream Nx evolution.
-3. **Upstream Nx extension** — add a custom-backend-op hook to
-   `Nx.Defn.Compiler` / `Nx.Defn.Evaluator`. Cleanest long-term
-   solution; slowest to land because it needs upstream review/merge.
-
-Also in scope for M10.5:
-
-- **Test-only AWQ loader** (`test/support/awq_loader.ex`) — reads
-  `Qwen/Qwen3-0.6B-AWQ` safetensors, extracts `qweight`, `scales`,
-  `qzeros`, maps to MLX's `(w_q, scales, biases)` layout. The
-  trickiest bit is the AWQ zero-point → MLX bias conversion
-  (`biases = -scales * zero_points`) and the AWQ `[in, out/pack]` vs.
-  MLX `[out, in]` layout difference.
-- **`:qwen3_quant_full` conformance test** — greedy-decode
-  Qwen3-0.6B-AWQ on Emily, assert the completion matches a checked-in
-  reference produced by MLX's Python bindings on the same quantized
-  weights.
-- **Bumblebee upstream contribution (optional, follow-up to M10.5)** —
-  upstream the AWQ loader into `deps/bumblebee` so the test-only path
-  becomes unnecessary.
-
-**Exit**: Qwen3-0.6B-AWQ greedy-decodes end-to-end on Emily under
-`Nx.Defn.jit`; conformance test green; Axon-integrated quantization
-documented.
+Closes the gap M10 left open: `Native.quantized_matmul` is now
+reachable from `Nx.Defn.jit`-traced Axon forward passes, and a
+quantized Qwen3-0.6B greedy-decodes end-to-end under Bumblebee's
+standard `Bumblebee.Text.generation/4` serving.
+
+**Shipped**:
+
+- **`Emily.Quantization.dequantize_defn/1`** (`lib/emily/quantization.ex`)
+  — defn-native analogue of `QuantizedWeight.to_dense/1`, built from
+  `Nx.right_shift` / `Nx.bitwise_and` / multiply / add. Supports
+  `bits ∈ {2, 4, 8}` (lanes-per-u32 is integral). `bits ∈ {3, 6}`
+  (cross-u32 packing) is out of scope — the Native path remains for
+  those.
+- **`Emily.Quantization.Layers.quantized_dense/4`**
+  (`lib/emily/quantization/layers.ex`) — Axon-compatible layer op
+  (`deftransform` → `defnp`). Pattern-matches on `%QuantizedWeight{}`,
+  dispatches `Nx.dot(x, dequantize_defn(qw))` (`transpose=false`) or
+  `Nx.dot(x, Nx.transpose(dequantize_defn(qw)))` (`transpose=true`).
+  Two kernel dispatches per matmul instead of MLX's fused one —
+  accepted trade-off for integration without forking
+  `Nx.Defn.Evaluator`. M11's fast-kernel work closes the gap.
+- **`Emily.Quantization.Transform`** (`test/support/quantization_transform.ex`)
+  — graph rewriter + model-state quantizer, modeled on
+  `Axon.Quantization`. `quantize/3` takes a dense Axon model + dense
+  `Axon.ModelState` and returns the pair with every `:dense` node
+  rewritten to `:quantized_dense` and every dense kernel replaced
+  with `%QuantizedWeight{}`. Lives in `test/support/` because Axon is
+  an `only: :test` dep of Emily; graduates to `lib/` when an upstream
+  Bumblebee AWQ-loading path lands.
+- **`:qwen3_quant_full` conformance test** (`test/emily/conformance/qwen3_quant_full_test.exs`)
+  — loads dense Qwen3-0.6B via Bumblebee, quantizes via `Transform.quantize/3`
+  (`bits=4, group_size=128, transpose=true`), runs
+  `Bumblebee.Text.generation/4` greedy decode for 32 tokens. Pins the
+  continuation string as a regression gate. Opt-in via `mix test --only
+  qwen3_quant_full` (mirrors `:qwen3_full`'s model-size discipline).
+
+**Approach chosen**: Option 1 (defn-native dequantize). Option 2
+(`Emily.Compiler` custom-op intercept) and Option 3 (upstream Nx
+extension) remain available if the two-kernel-vs-fused gap materially
+hurts a real workload after M11.
+
+**Scope reductions from original PLAN.md M10.5**:
+
+- **AWQ safetensors loader deferred.** The original plan called for a
+  test-only loader that reads `Qwen/Qwen3-0.6B-AWQ` and maps to
+  `%QuantizedWeight{}`. On closer inspection the AWQ→MLX conversion is
+  meaningfully more involved than first thought: AWQ groups along the
+  `in` axis while MLX's `transpose=false` path expects groups along the
+  stored last axis, so correct conversion requires transposing the
+  packed tensor, unpacking `qzeros` into per-group zero-points,
+  computing `biases = -scales * zero_points`, and mapping HF param
+  names to Bumblebee's internal naming. All tractable but substantial.
+  The from-dense path above exercises the same defn-integration
+  pipeline (graph rewrite + QW params + defn-native dequantize + JIT'd
+  forward) and produces a useful regression gate; AWQ-specific loading
+  is now a proper follow-up milestone rather than M10.5-scope.
+- **Conformance oracle adjusted.** PLAN.md originally envisioned
+  asserting against MLX Python's output on the same quantized
+  weights. Since we're now quantizing Qwen3-0.6B ourselves (not
+  loading AWQ), the reference is what this pipeline produces on a
+  clean checkout — same discipline as `qwen3_full_test.exs`.
+
+**Follow-ups** (out of M10.5 scope):
+
+- AWQ safetensors loader + Bumblebee param-name mapping. When it
+  lands, adds a second conformance test that loads real
+  `Qwen/Qwen3-0.6B-AWQ` weights end-to-end.
+- Optional upstream contribution to `deps/bumblebee` for AWQ loading.
 
 ### M11 — `mlx::fast::*` fused kernels
 

lib/emily/quantization.ex

@@ -2,37 +2,47 @@ defmodule Emily.Quantization do
   @moduledoc """
   Quantized inference primitives.
 
-  M10 ships the Native bindings (`Emily.Native.quantize/3`,
-  `Emily.Native.dequantize/5`, `Emily.Native.quantized_matmul/7`), the
-  `Emily.QuantizedWeight` container, and this direct-call helper. This is
-  enough to:
-
-    * quantize a dense weight, store it packed at rest, and later dispatch a
-      fused quantized matmul against it;
-    * benchmark the quantized path against an `Nx.dot`-on-dequantized oracle;
-    * run quantized inference in *eager* code (plain Elixir, outside `defn`).
-
-  ## What M10 does NOT ship
-
-  Integration with `Nx.Defn`-traced Axon forward passes (and therefore
-  `Bumblebee` serving with AWQ checkpoints) is deferred to a follow-up.
-  `Nx.Defn.Evaluator` walks `Nx.Defn.Expr` and dispatches via the
-  `Nx.Backend` behaviour; there is no public hook to inject a custom op like
-  `Native.quantized_matmul`, and none of `deftransform`, `hook`, or
-  `metadata` can call NIFs on tensors that are still `Expr` nodes at trace
-  time. Closing that gap requires either (a) a defn-native dequantize built
-  from Nx bit primitives (skips the fused kernel), or (b) an Emily-specific
-  compiler variant that recognizes a sentinel `Expr` node. Both are
-  meaningful scope; M10.5 will tackle one of them.
-
-  Use `quantized_matmul/2` below for any eager quantized inference today.
+  Two entry points:
+
+    * `quantized_matmul/2` — eager-mode fused kernel (materialized
+      tensors only). Extracts refs from a `%QuantizedWeight{}` and
+      calls `Native.quantized_matmul/7`.
+    * `dequantize_defn/1` — defn-native analogue of
+      `QuantizedWeight.to_dense/1`, composed from `Nx.right_shift` /
+      `Nx.bitwise_and` / multiply / add. Use inside `Nx.Defn.jit`-traced
+      Axon forward passes; `Nx.dot(x, dequantize_defn(qw))` replaces
+      the fused `quantized_matmul` kernel with two dispatches — M11's
+      fast-kernel work closes that gap.
+
+  `dequantize_defn/1` supports `bits ∈ #{inspect([2, 4, 8])}`;
+  `bits ∈ {3, 6}` use cross-u32 lane packing (out of scope here — use
+  `QuantizedWeight.to_dense/1`).
+
+  See `Emily.Quantization.Layers.quantized_dense/4` for the
+  Axon-compatible layer op built on `dequantize_defn/1`.
   """
 
+  import Nx.Defn
+
   alias Emily.Backend, as: B
   alias Emily.Native
   alias Emily.QuantizedWeight
   alias Nx.Tensor, as: T
 
+  @defn_supported_bits [2, 4, 8]
+
+  @doc """
+  Bit widths supported by `dequantize_defn/1` (and therefore by
+  `Emily.Quantization.Layers.quantized_dense/4` /
+  `Emily.Quantization.Transform`).
+
+  `bits ∈ {3, 6}` use cross-u32 lane packing and aren't supported by
+  the defn-native path; `QuantizedWeight.to_dense/1` (the Native path)
+  still handles them.
+  """
+  @spec defn_supported_bits() :: [pos_integer()]
+  def defn_supported_bits, do: @defn_supported_bits
+
   @doc """
   Compute `x @ W^T` where `W` is represented as a `QuantizedWeight`.
 
@@ -96,4 +106,105 @@ defmodule Emily.Quantization do
             "match scales dtype #{inspect(s_type)}. Cast the input with " <>
             "`Nx.as_type/2` before calling."
   end
+
+  # ================================================================
+  # Defn-native dequantize
+  # ================================================================
+
+  @doc """
+  Reconstruct a dense tensor from a `QuantizedWeight`, built entirely
+  from Nx primitives so it composes inside `defn` traces.
+
+  This is the defn-compatible analogue of `QuantizedWeight.to_dense/1`.
+  The math is identical to MLX's `dequantize`:
+
+      w[i] = (w_q_packed >> ((i mod lpu) * bits)) & mask * scales[g] + biases[g]
+
+  where `lpu = div(32, bits)` (lanes per u32), `mask = (1 <<< bits) - 1`,
+  and `g = div(i, group_size)` is the group index along the last axis.
+
+  Supported: `bits ∈ #{inspect(@defn_supported_bits)}`. `bits ∈ {3, 6}`
+  pack across u32 boundaries and are out of scope here.
+
+  ## Example
+
+      qw = Emily.QuantizedWeight.from_dense(w, group_size: 64, bits: 4)
+      dense_defn = Emily.Quantization.dequantize_defn(qw)
+      dense_native = Emily.QuantizedWeight.to_dense(qw)
+      # element-wise identical
+  """
+  @spec dequantize_defn(QuantizedWeight.t()) :: Nx.Tensor.t()
+  deftransform dequantize_defn(qw) do
+    %QuantizedWeight{
+      value: q,
+      scales: s,
+      biases: b,
+      group_size: group_size,
+      bits: bits
+    } = qw
+
+    validate_defn_bits!(bits)
+
+    dequantize_impl(q, s, b, group_size: group_size, bits: bits)
+  end
+
+  defp validate_defn_bits!(bits) when bits in @defn_supported_bits, do: :ok
+
+  defp validate_defn_bits!(bits) do
+    raise ArgumentError,
+          "Emily.Quantization.dequantize_defn/1: bits=#{bits} uses cross-u32 " <>
+            "lane packing, which is out of scope for the defn-native path. " <>
+            "Supported: #{inspect(@defn_supported_bits)}. Use " <>
+            "`Emily.QuantizedWeight.to_dense/1` (the Native path) for unsupported " <>
+            "bit widths."
+  end
+
+  # Expects `opts` to carry compile-time `:group_size` and `:bits`. Both
+  # are used for shape arithmetic (lanes-per-u32, per-group reshape) so
+  # they must be trace-time constants.
+  defnp dequantize_impl(w_q, scales, biases, opts \\ []) do
+    opts = keyword!(opts, [:group_size, :bits])
+    group_size = opts[:group_size]
+    bits = opts[:bits]
+
+    # Unpack: (..., packed) → (..., packed, lpu) via broadcast-shift
+    # with a length-lpu shift vector, then mask to `bits`-width nibbles.
+    shifts = build_shifts(bits)
+    mask = build_mask(bits)
+
+    # new_axis appends a length-1 axis; right_shift broadcasts against shifts.
+    shifted = Nx.right_shift(Nx.new_axis(w_q, -1), shifts)
+    masked = Nx.bitwise_and(shifted, mask)
+
+    # Flatten (packed, lpu) → orig_last, then regroup to (..., groups,
+    # group_size) so per-group scale/bias broadcast trivially.
+    grouped = masked |> Nx.flatten(axes: [-2, -1]) |> group_last_axis(group_size)
+
+    # Cast u32 → scales dtype, then per-group affine recombine; flatten
+    # back to (..., orig_last).
+    grouped_f = Nx.as_type(grouped, Nx.type(scales))
+    dequantized = grouped_f * Nx.new_axis(scales, -1) + Nx.new_axis(biases, -1)
+
+    Nx.flatten(dequantized, axes: [-2, -1])
+  end
+
+  deftransformp build_shifts(bits) do
+    lpu = div(32, bits)
+    shifts = for i <- 0..(lpu - 1), do: i * bits
+    Nx.tensor(shifts, type: :u32)
+  end
+
+  deftransformp build_mask(bits) do
+    import Bitwise
+    Nx.tensor((1 <<< bits) - 1, type: :u32)
+  end
+
+  # Reshape `(..., n)` → `(..., n / group_size, group_size)`. Uses
+  # `:auto` so we don't recompute the quotient.
+  deftransformp group_last_axis(t, group_size) do
+    shape = Nx.shape(t)
+    rank = tuple_size(shape)
+    new_shape = shape |> put_elem(rank - 1, :auto) |> Tuple.insert_at(rank, group_size)
+    Nx.reshape(t, new_shape)
+  end
 end

lib/emily/quantization/layers.ex

@@ -0,0 +1,72 @@
+defmodule Emily.Quantization.Layers do
+  @moduledoc """
+  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 and
+  `Emily.Quantization.Transform.quantize_dense_layers/2` for the
+  graph-rewrite that wires it in.
+  """
+
+  import Nx.Defn
+
+  alias Emily.QuantizedWeight
+
+  @doc """
+  Axon layer op: `x @ W (+ bias)` where `W` is a `%QuantizedWeight{}`.
+
+  Mirrors the signature of `Axon.Quantization.Layers.weight_only_quantized_dense/4`:
+
+    * `input` — activation tensor, shape `(..., in)`.
+    * `kernel` — `%QuantizedWeight{}`. The stored layout is determined
+      by `kernel.transpose`:
+        * `transpose: false` (the AWQ / Axon-native layout) — packed
+          representation of a `[in, out]` weight; the layer computes
+          `Nx.dot(x, dense)`.
+        * `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))`.
+    * `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
+      signature for drop-in use under `Axon.layer/3`.
+    * `opts` — reserved for Axon-layer metadata; not used by this
+      implementation directly (all state lives on the
+      `%QuantizedWeight{}`).
+  """
+  deftransform quantized_dense(input, kernel, bias \\ 0, opts \\ []) 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} =
+      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)
+  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)
+  end
+end

lib/emily/quantized_weight.ex

@@ -30,12 +30,14 @@ defmodule Emily.QuantizedWeight do
   tensor — so the direct-call helper is the supported path.
 
   > #### Defn-traced Axon forward passes {: .info}
-  > Routing through `Nx.Defn.jit`-traced Axon models (and therefore
-  > Bumblebee AWQ checkpoints) is not yet shipped; the layer-op dispatch
-  > path requires either a defn-native dequantize or an
-  > `Nx.Defn.Evaluator`-level custom-op hook. Tracked as M10.5 in
-  > `PLAN.md`. Use `quantized_matmul/2` directly on materialized tensors
-  > today.
+  > `Emily.Quantization.dequantize_defn/1` (M10.5) is the defn-native
+  > analogue of `to_dense/1`; pair it with
+  > `Emily.Quantization.Layers.quantized_dense/4` to splice a quantized
+  > linear into any `Nx.Defn.jit`-traced Axon forward pass. The layer
+  > performs `Nx.dot(x, dequantize_defn(qw))` instead of MLX's single
+  > fused `quantized_matmul` — two kernels vs. one, but fully
+  > integrated with the rest of Bumblebee's defn graph. M11's fast-kernel
+  > work will close the perf gap.
   """
 
   @derive {Nx.Container,