Qwen3 EMLX / Bumblebee / Emily Benchmark
Mix.install(
[
{:emlx, path: Path.expand("../../emlx", __DIR__)},
{:emlx_axon, path: Path.expand("..", __DIR__)},
# {:emlx, github: "elixir-nx/emlx", branch: "main", sparse: "emlx", override: true},
# {:emlx_axon, github: "elixir-nx/emlx", branch: "main", sparse: "emlx_axon", override: true},
{:nx, github: "elixir-nx/nx", branch: "main", sparse: "nx", override: true},
{:axon, "~> 0.7"},
{:bumblebee, "~> 0.7"},
{:emily, "~> 0.7"},
{:kino, "~> 0.14"}
],
config: [nx: [default_backend: {EMLX.Backend, device: :gpu}]]
)
Overview
Dense (bf16) and MLX-4bit Qwen3-0.6B checkpoints; lanes per checkpoint:
- bb base — Bumblebee, stock Axon graph
-
bb+rewrite — Bumblebee +
EMLXAxon.rewrite/1 -
native —
EMLXAxon.TextGeneration - emily-eager / emily-native — dense only; Emily backend/compiler, unmodified Bumblebee graph
- emily-quantized — mlx4bit entry only; Emily int4-quantizes the dense checkpoint (not the lmstudio mlx4bit file on disk)
Each lane runs in a fresh :peer node (~150 ms overhead) so MLX allocator state does not bleed across lanes.
Peer helper
defmodule PeerBench.Registry do
@moduledoc false
use Agent
def start_link(_), do: Agent.start_link(fn -> %{} end, name: __MODULE__)
def put(module, binary) do
ensure_started()
Agent.update(__MODULE__, &Map.put(&1, module, binary))
end
def get(module) do
ensure_started()
Agent.get(__MODULE__, &Map.fetch!(&1, module))
end
defp ensure_started do
if is_nil(Process.whereis(__MODULE__)), do: start_link([])
end
end
defmodule PeerBench do
@moduledoc false
defmacro defshippable(name, do: block) do
quote do
{:module, unquote(name), bin, _} =
defmodule unquote(name) do
unquote(block)
end
PeerBench.Registry.put(unquote(name), bin)
end
end
@apps [:nx, :emlx, :emlx_axon, :axon, :bumblebee, :emily]
def run(modules, mod, fun, args) do
peer_args = Enum.flat_map(:code.get_path(), fn p -> [~c"-pa", p] end)
{:ok, pid, _node} =
:peer.start_link(%{name: :peer.random_name(), args: peer_args, connection: :standard_io})
try do
for m <- modules do
bin = PeerBench.Registry.get(m)
{:module, ^m} = :peer.call(pid, :code, :load_binary, [m, ~c"nofile", bin], :infinity)
end
for app <- @apps, do: :peer.call(pid, Application, :ensure_all_started, [app], :infinity)
:peer.call(pid, :erlang, :apply, [mod, fun, args], :infinity)
after
:peer.stop(pid)
end
end
end
:ok
Benchmark helpers
require PeerBench
PeerBench.defshippable Bench do
def run_serving(serving, prompt, extract_fn) do
t_start = System.monotonic_time(:millisecond)
result = Nx.Serving.run(serving, prompt)
t_end = System.monotonic_time(:millisecond)
{text, n_tokens} = extract_fn.(result)
elapsed_ms = t_end - t_start
tok_s = if elapsed_ms > 0, do: Float.round(n_tokens / elapsed_ms * 1000.0, 1), else: 0.0
{elapsed_ms, n_tokens, tok_s, text}
end
def warmup(label, serving, prompt, extract_fn, runs) do
IO.puts("\n==> [#{label}] Warmup (#{runs} run(s), not timed) ...")
for i <- 1..runs do
{ms, n, _, text} = run_serving(serving, prompt, extract_fn)
preview = text |> String.slice(0, 60) |> String.replace("\n", " ")
IO.puts(" warmup #{i}: #{n} tokens / #{ms} ms \"#{preview}...\"")
end
end
def bench(label, serving, prompt, extract_fn, runs) do
IO.puts("\n==> [#{label}] Benchmark (#{runs} run(s)) ...")
for i <- 1..runs do
{ms, n, tok_s, _} = run_serving(serving, prompt, extract_fn)
IO.puts(" run #{i}: #{n} tokens / #{ms} ms = #{tok_s} tok/s")
%{elapsed_ms: ms, n_tokens: n, tok_s: tok_s}
end
end
def stats(label, results) do
toks_list = Enum.map(results, & &1.tok_s)
sorted = Enum.sort(toks_list)
n = length(sorted)
median = Enum.at(sorted, div(n, 2))
mean = Float.round(Enum.sum(toks_list) / n, 1)
min_v = List.first(sorted)
max_v = List.last(sorted)
variance = Enum.reduce(toks_list, 0.0, fn x, acc -> acc + (x - mean) * (x - mean) end) / n
stddev = Float.round(:math.sqrt(variance), 1)
IO.puts(" #{label}: median=#{median} mean=#{mean}±#{stddev} min/max=#{min_v}/#{max_v} tok/s")
%{label: label, median: median, mean: mean, stddev: stddev, min: min_v, max: max_v}
end
end
PeerBench.defshippable Extract do
def bb(%{results: [%{text: text, token_summary: summary}]}) do
{text, summary.output}
end
def native(%{results: [%{generated_text: text, num_tokens: n}]}, _tokenizer) do
{text, n}
end
def native(%{results: [%{generated_text: text}]}, tokenizer) do
%{"input_ids" => ids} = Bumblebee.apply_tokenizer(tokenizer, text)
{text, elem(Nx.shape(ids), 1)}
end
end
PeerBench.defshippable ModelSource do
def resolve(path_raw) do
if File.exists?(Path.expand(path_raw)) do
{:local, Path.expand(path_raw)}
else
{:hf, path_raw}
end
end
end
:ok
Emily lanes
Quantize rewrite vendored from Emily’s livebooks/qwen3_quantized.livemd.
require PeerBench
PeerBench.defshippable EmilyQuantize do
alias Emily.Quantization.Layers
alias Emily.QuantizedWeight
def quantize(model, model_state, opts \\ []) do
{rewrite_graph(model), quantize_state(model, model_state, opts)}
end
defp rewrite_graph(model) do
Axon.rewrite_nodes(model, fn
%Axon.Node{op: :dense, meta: meta, name: name_fn} ->
fn [x], _output ->
quantized_dense_layer(x, meta[:units], use_bias: meta[:use_bias], name: name_fn)
end
_ ->
:skip
end)
end
defp quantized_dense_layer(x, units, opts) do
kernel_shape = fn input_shape ->
in_features = elem(input_shape, tuple_size(input_shape) - 1)
{in_features, units}
end
bias_shape = fn _input_shape -> {units} end
kernel = Axon.param("kernel", kernel_shape, initializer: :glorot_uniform)
{inputs, op} =
if opts[:use_bias] do
bias = Axon.param("bias", bias_shape, initializer: :zeros)
{[x, kernel, bias], &Layers.quantized_dense/4}
else
{[x, kernel], &Layers.quantized_dense/3}
end
Axon.layer(op, inputs,
name: opts[:name],
meta: %{units: units, use_bias: opts[:use_bias]},
op_name: :quantized_dense
)
end
defp quantize_state(model, state, opts) do
transpose = Keyword.fetch!(opts, :transpose)
dense_names =
model
|> Axon.properties()
|> Enum.filter(fn {_name, op} -> op == :dense end)
|> Enum.map(fn {name, _} -> name end)
Enum.reduce(dense_names, state, fn name, acc ->
update_in(acc, [Access.key!(:data), name, "kernel"], fn kernel ->
source = if transpose, do: Nx.transpose(kernel), else: kernel
QuantizedWeight.from_dense(source,
group_size: Keyword.fetch!(opts, :group_size),
bits: Keyword.fetch!(opts, :bits),
transpose: transpose
)
end)
end)
end
end
PeerBench.defshippable EmilyLane do
def run(kind, source, checkpoint_label, prompt, cfg) do
if cfg.skip_emily? do
nil
else
lane_label = "#{checkpoint_label} / emily-#{kind}"
Nx.default_backend(Emily.Backend)
try do
IO.puts("\n==> [#{lane_label}] Loading model on Emily.Backend ...")
{:ok, model_info} = Bumblebee.load_model(source, backend: Emily.Backend, type: :bf16)
{:ok, tokenizer} = Bumblebee.load_tokenizer(source)
{:ok, generation_config} = Bumblebee.load_generation_config(source)
generation_config =
Bumblebee.configure(generation_config,
max_new_tokens: cfg.max_new,
strategy: %{type: :greedy_search}
)
{defn_options, model_info} = build_defn_options(kind, model_info, cfg)
serving =
Bumblebee.Text.generation(model_info, tokenizer, generation_config, defn_options: defn_options)
Bench.warmup(lane_label, serving, prompt, &Extract.bb/1, cfg.warmup_runs)
Bench.bench(lane_label, serving, prompt, &Extract.bb/1, cfg.bench_runs)
rescue
exception ->
IO.puts("\n==> [#{lane_label}] FAILED: #{Exception.message(exception)}")
nil
end
end
end
defp build_defn_options(:eager, model_info, _cfg) do
{[compiler: Nx.Defn.Evaluator], model_info}
end
defp build_defn_options(:native, model_info, _cfg) do
{[compiler: Emily.Compiler, native: true, native_fallback: :raise], model_info}
end
defp build_defn_options(:quantized, model_info, cfg) do
IO.puts(
" quantizing every Axon.dense node " <>
"(bits=#{cfg.emily_quant_bits}, group_size=#{cfg.emily_quant_group_size}) ..."
)
{qmodel, qparams} =
EmilyQuantize.quantize(model_info.model, model_info.params,
bits: cfg.emily_quant_bits,
group_size: cfg.emily_quant_group_size,
transpose: true
)
{[compiler: Emily.Compiler, native: true, native_fallback: :raise],
%{model_info | model: qmodel, params: qparams}}
end
end
:ok
Checkpoint runner
require PeerBench
PeerBench.defshippable Runner do
@shipped_modules [Bench, Extract, ModelSource, EmilyQuantize, EmilyLane, Runner]
def run_checkpoint(%{kind: kind, label: label, path_raw: path_raw} = checkpoint, cfg) do
IO.puts("""
################################################################################
=== #{label} — #{inspect(ModelSource.resolve(path_raw))} ===
################################################################################
""")
stats =
for lane <- applicable_lanes(kind, cfg), into: %{} do
IO.puts("\n==> Spawning isolated peer for lane #{inspect(lane)} ...")
{lane, PeerBench.run(@shipped_modules, Runner, :run_lane, [lane, checkpoint, cfg])}
end
%{
label: label,
base: stats[:bb_base],
rewrite: stats[:bb_rewrite],
native: stats[:native],
emily_eager: stats[:emily_eager],
emily_native: stats[:emily_native],
emily_quantized: stats[:emily_quantized]
}
end
defp applicable_lanes(kind, cfg) do
[
bb_base: not cfg.skip_bb_base?,
bb_rewrite: not cfg.skip_bb_rewrite?,
native: true,
emily_eager: kind == :dense and not cfg.skip_emily?,
emily_native: kind == :dense and not cfg.skip_emily?,
emily_quantized: kind == :mlx4bit and not cfg.skip_emily?
]
|> Enum.filter(fn {_lane, enabled?} -> enabled? end)
|> Enum.map(fn {lane, _} -> lane end)
end
def run_lane(lane, %{kind: kind, label: label, path_raw: path_raw}, cfg) do
Nx.default_backend({EMLX.Backend, device: :gpu})
source = ModelSource.resolve(path_raw)
{model_info, tokenizer, generation_config} =
if lane in [:bb_base, :bb_rewrite] do
{:ok, model_info} = Bumblebee.load_model(source, backend: {EMLX.Backend, device: :gpu}, type: :bf16)
model_info = load_params(model_info, kind, path_raw)
{:ok, tokenizer} = Bumblebee.load_tokenizer(source)
{:ok, generation_config} = Bumblebee.load_generation_config(source)
generation_config =
Bumblebee.configure(generation_config,
max_new_tokens: cfg.max_new,
strategy: %{type: :greedy_search}
)
{model_info, tokenizer, generation_config}
else
tokenizer =
if lane == :native do
{:ok, tokenizer} = Bumblebee.load_tokenizer(source)
tokenizer
end
{nil, tokenizer, nil}
end
dispatch_lane(lane, kind, path_raw, label, source, model_info, tokenizer, generation_config, cfg)
end
defp dispatch_lane(:bb_base, _kind, _path_raw, label, _source, model_info, tokenizer, generation_config, cfg) do
IO.puts("==> Building Bumblebee.Text.generation serving (base — no rewrite) ...")
serving =
Bumblebee.Text.generation(model_info, tokenizer, generation_config,
compile: [batch_size: 1, sequence_length: cfg.seq_len],
defn_options: [compiler: EMLX]
)
lane_label = "#{label} / bb base"
Bench.warmup(lane_label, serving, cfg.prompt, &Extract.bb/1, cfg.warmup_runs)
results = Bench.bench(lane_label, serving, cfg.prompt, &Extract.bb/1, cfg.bench_runs)
Bench.stats("bb base (Bumblebee, no rewrite) ", results)
end
defp dispatch_lane(:bb_rewrite, _kind, _path_raw, label, _source, model_info, tokenizer, generation_config, cfg) do
IO.puts("==> Applying EMLXAxon.rewrite/1 (all rewrites — best performing path) ...")
model_rewritten = EMLXAxon.rewrite(model_info.model)
serving =
Bumblebee.Text.generation(%{model_info | model: model_rewritten}, tokenizer, generation_config,
compile: [batch_size: 1, sequence_length: cfg.seq_len],
defn_options: [compiler: EMLX]
)
lane_label = "#{label} / bb+rewrite"
Bench.warmup(lane_label, serving, cfg.prompt, &Extract.bb/1, cfg.warmup_runs)
results = Bench.bench(lane_label, serving, cfg.prompt, &Extract.bb/1, cfg.bench_runs)
Bench.stats("bb+rewrite (Bumblebee + EMLXAxon.rewrite) ", results)
end
defp dispatch_lane(:native, kind, path_raw, label, _source, _model_info, tokenizer, _generation_config, cfg) do
serving = build_native_serving(kind, path_raw, tokenizer, cfg)
extract = fn result -> Extract.native(result, tokenizer) end
lane_label = "#{label} / native"
Bench.warmup(lane_label, serving, cfg.prompt, extract, cfg.warmup_runs)
results = Bench.bench(lane_label, serving, cfg.prompt, extract, cfg.bench_runs)
Bench.stats("native (EMLXAxon.TextGeneration) ", results)
end
defp dispatch_lane(:emily_eager, _kind, _path_raw, label, source, _model_info, _tokenizer, _generation_config, cfg) do
EmilyLane.run(:eager, source, label, cfg.prompt, cfg)
|> wrap_stats("emily-eager (Evaluator + Emily.Backend) ")
end
defp dispatch_lane(:emily_native, _kind, _path_raw, label, source, _model_info, _tokenizer, _generation_config, cfg) do
EmilyLane.run(:native, source, label, cfg.prompt, cfg)
|> wrap_stats("emily-native (Emily.Compiler, native: true) ")
end
defp dispatch_lane(:emily_quantized, _kind, _path_raw, label, _source, _model_info, _tokenizer, _generation_config, cfg) do
dense_source = ModelSource.resolve(cfg.dense_path_raw)
EmilyLane.run(:quantized, dense_source, label, cfg.prompt, cfg)
|> wrap_stats("emily-quantized (Emily int4, native: true) ")
end
defp wrap_stats(nil, _label), do: nil
defp wrap_stats(results, label), do: Bench.stats(label, results)
defp load_params(model_info, :dense, _path_raw), do: model_info
defp load_params(model_info, :mlx4bit, path_raw) do
IO.puts("==> Loading MLX-4bit params (dequantize → Bumblebee layout → re-quantize) ...")
t0 = System.monotonic_time(:millisecond)
params = EMLXAxon.MLX4BitParams.load(Path.expand(path_raw))
params = EMLXAxon.QuantizeParams.quantize(params)
t1 = System.monotonic_time(:millisecond)
IO.puts(" params loaded in #{t1 - t0} ms")
%{model_info | params: params}
end
defp build_native_serving(kind, path_raw, tokenizer, cfg) do
IO.puts("\n==> Loading EMLXAxon.TextGeneration (native 28-layer bypass) ...")
t0 = System.monotonic_time(:millisecond)
serving =
case kind do
:mlx4bit ->
EMLXAxon.TextGeneration.from_mlx4bit(
Path.expand(path_raw),
tokenizer,
defn_options: [compiler: Nx.Defn.Evaluator],
max_new_tokens: cfg.max_new,
sampler: :greedy,
profile_timing: cfg.native_profile_timing?
)
:dense ->
{:ok, state} = EMLXAxon.Qwen3.DenseLoader.from_safetensors_dir(Path.expand(path_raw), type: :bf16)
EMLXAxon.TextGeneration.serving(
tokenizer,
state,
defn_options: [compiler: Nx.Defn.Evaluator],
max_new_tokens: cfg.max_new,
sampler: :greedy,
profile_timing: cfg.native_profile_timing?
)
end
IO.puts(" loaded in #{System.monotonic_time(:millisecond) - t0} ms")
serving
end
end
:ok
Summary rendering
defmodule Summary do
def ratio(nil, _), do: nil
def ratio(_, nil), do: nil
def ratio(row, col) when col.median > 0, do: Float.round(row.median / col.median, 2)
def ratio(_, _), do: nil
def to_markdown({:ok, result}) do
lanes = lanes(result)
matrix =
if length(lanes) < 2 do
"_fewer than two lanes ran — nothing to compare._"
else
matrix_markdown(lanes)
end
bullets =
Enum.map_join(lanes, "\n", fn {name, stats} ->
"- **#{name}**: median=#{stats.median} mean=#{stats.mean}±#{stats.stddev} min/max=#{stats.min}/#{stats.max} tok/s"
end)
"### #{result.label}\n\n#{matrix}\n\n#{bullets}"
end
def to_markdown({:error, label, formatted}) do
fence = String.duplicate("`", 3)
"### #{label} — FAILED\n\n#{fence}\n#{formatted}\n#{fence}"
end
def stats_rows(results) do
for {:ok, result} <- results,
{name, stats} <- lanes(result) do
%{
checkpoint: result.label,
lane: name,
median_tok_s: stats.median,
mean_tok_s: stats.mean,
stddev: stats.stddev,
min_tok_s: stats.min,
max_tok_s: stats.max
}
end
end
defp lanes(result) do
[
{"bb base", result.base},
{"bb+rewrite", result.rewrite},
{"native", result.native},
{"emily-eager", result.emily_eager},
{"emily-native", result.emily_native},
{"emily-quantized", result.emily_quantized}
]
|> Enum.reject(fn {_name, stats} -> is_nil(stats) end)
end
defp matrix_markdown(lanes) do
header = Enum.map_join(lanes, "", fn {name, _} -> " #{name} |" end)
sep = Enum.map_join(lanes, "", fn _ -> " --- |" end)
rows =
for {row_name, row_stats} <- lanes do
cells =
Enum.map_join(lanes, "", fn {col_name, col_stats} ->
" #{format_cell(row_name, col_name, row_stats, col_stats)} |"
end)
"| **#{row_name}** |#{cells}"
end
Enum.join(["| row ∖ col (row / col) |#{header}", "| --- |#{sep}" | rows], "\n")
end
defp format_cell(row_name, col_name, _row_stats, _col_stats) when row_name == col_name, do: "–"
defp format_cell(_row_name, _col_name, row_stats, col_stats) do
case ratio(row_stats, col_stats) do
nil -> "n/a"
value -> "#{value}×"
end
end
end
:ok
Configuration
dense_path_input =
Kino.Input.text("Dense model path",
default: System.get_env("EMLX_QWEN3_DENSE_MODEL_PATH") || "~/models/Qwen3-0.6B"
)
mlx4bit_path_input =
Kino.Input.text("MLX-4bit model path",
default: System.get_env("EMLX_QWEN3_4BIT_MODEL_PATH") || "~/models/Qwen3-0.6B-MLX-4bit"
)
skip_dense_input = Kino.Input.checkbox("Skip dense checkpoint", default: false)
skip_4bit_input = Kino.Input.checkbox("Skip MLX-4bit checkpoint", default: false)
skip_emily_input = Kino.Input.checkbox("Skip Emily lanes", default: false)
skip_bb_base_input = Kino.Input.checkbox("Skip bb base lane", default: false)
skip_bb_rewrite_input = Kino.Input.checkbox("Skip bb+rewrite lane", default: false)
native_profile_timing_input = Kino.Input.checkbox("Native per-token profile timing", default: true)
max_new_input = Kino.Input.number("Max new tokens", default: 60)
bench_runs_input = Kino.Input.number("Bench runs", default: 5)
warmup_runs_input = Kino.Input.number("Warmup runs", default: 2)
seq_len_input = Kino.Input.number("Sequence length", default: 1024)
emily_quant_bits_input = Kino.Input.number("Emily quant bits", default: 4)
emily_quant_group_size_input = Kino.Input.number("Emily quant group size", default: 64)
Kino.Layout.grid(
[
dense_path_input,
mlx4bit_path_input,
skip_dense_input,
skip_4bit_input,
skip_emily_input,
skip_bb_base_input,
skip_bb_rewrite_input,
native_profile_timing_input,
max_new_input,
bench_runs_input,
warmup_runs_input,
seq_len_input,
emily_quant_bits_input,
emily_quant_group_size_input
],
columns: 2
)
dense_path_raw = Kino.Input.read(dense_path_input)
cfg = %{
max_new: Kino.Input.read(max_new_input),
bench_runs: Kino.Input.read(bench_runs_input),
warmup_runs: Kino.Input.read(warmup_runs_input),
seq_len: Kino.Input.read(seq_len_input),
native_profile_timing?: Kino.Input.read(native_profile_timing_input),
skip_bb_rewrite?: Kino.Input.read(skip_bb_rewrite_input),
skip_bb_base?: Kino.Input.read(skip_bb_base_input),
skip_emily?: Kino.Input.read(skip_emily_input),
emily_quant_bits: Kino.Input.read(emily_quant_bits_input),
emily_quant_group_size: Kino.Input.read(emily_quant_group_size_input),
dense_path_raw: dense_path_raw,
prompt:
"<|im_start|>user\nList twenty programming languages with a one-line description each.<|im_end|>\n<|im_start|>assistant\n"
}
checkpoints = [
%{
kind: :dense,
label: "Qwen3-0.6B (dense bf16)",
skip?: Kino.Input.read(skip_dense_input),
path_raw: dense_path_raw
},
%{
kind: :mlx4bit,
label: "Qwen3-0.6B-MLX-4bit",
skip?: Kino.Input.read(skip_4bit_input),
path_raw: Kino.Input.read(mlx4bit_path_input)
}
]
IO.puts("""
max_new: #{cfg.max_new} bench_runs: #{cfg.bench_runs} warmup_runs: #{cfg.warmup_runs}
seq_len: #{cfg.seq_len}
""")
:ok
Run benchmarks
results =
checkpoints
|> Enum.reject(& &1.skip?)
|> Enum.shuffle()
|> Enum.map(fn checkpoint ->
try do
{:ok, Runner.run_checkpoint(checkpoint, cfg)}
rescue
exception -> {:error, checkpoint.label, Exception.format(:error, exception, __STACKTRACE__)}
end
end)
:ok
Results
Matrix cells: row median tok/s ÷ col median tok/s (> 1 ⇒ row faster).
results
|> Enum.map(&Summary.to_markdown/1)
|> Enum.join("\n\n---\n\n")
|> dbg()
|> Kino.Markdown.new()
results
|> Summary.stats_rows()
|> Kino.DataTable.new(name: "Per-lane tok/s stats")