perf: fuse quantized_dense onto mx::quantized_matmul kernel

[ausimian]

Summary

Emily.Quantization.Layers.quantized_dense/4 — the Axon layer backing every quantized projection — now lowers to the fused mx::quantized_matmul kernel (via Emily.Quantization.quantized_matmul_defn/2) instead of dequantizing the full weight to bf16 and running a dense Nx.dot. Inside the compiled decode :while loop the old path re-dequantized the entire model on every token, making quantized inference slower than dense.

Closes #197. This graduates the already-built-but-unwired fused quantized matmul (see PLAN.md M10.5 / M11 — "Option 2 remains available if the two-kernel-vs-fused gap materially hurts a real workload"). The motivating workload is the quantized gap highlighted in elixir-nx/emlx#124.

Results

Same machine, native lane (compiler: Emily.Compiler, native: true, native_fallback: :raise), Qwen/Qwen3-0.6B, 64-token greedy decode:

Lane	tok/s	vs dense
Dense f32	16.7	1.0×
Quantized 4-bit — before	2.74	0.16× (6× slower)
Quantized 4-bit — after	36.1	2.2× faster

• ~13× faster end-to-end on quantized decode; quantization flips from slower than dense to faster than dense — the memory-bandwidth win it's supposed to deliver.
• Per-matmul micro-bench: fused is 2.3–4.6× faster than dequant+dot (largest on the bandwidth-bound MLP projections).
• Greedy output is byte-identical before/after.

(Absolute tok/s are from a box slower than the emlx PR's M4 Max, so they aren't directly comparable to that PR's numbers; the ratios and the "quantized now beats dense" result are the machine-independent findings.)

Changes

• lib/emily/quantization/layers.ex — wire the layer to quantized_matmul_defn/2; layout/mode/bits/group_size are read off the %QuantizedWeight{}. Non-Emily backends keep the composed dequantize_defn + Nx.dot fallback via the block's default fun.
• test/emily/quantization/layers_test.exs — two native-lane tests (native_fallback: :raise) proving quantized_dense genuinely lowers to the fused opcode rather than falling back to op-by-op eval.
• bench/qmm_microbench.exs — per-matmul before/after (no model download needed).
• bench/qwen3_quantized_tps.exs — end-to-end quantized Qwen3-0.6B tok/s (run MIX_ENV=test mix run …).

Testing

• mix precommit clean: compile --warnings-as-errors, format, credo --strict, docs build, and 40 doctests / 79 properties / 737 tests, 0 failures.
• Existing layers_test already used the eager fused kernel as its oracle, so no tolerance changes were needed.

[ausimian]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. 🎉

Reviewed commit: 61b36f4582

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".