11# Qwen3-0.6B greedy-decode throughput on `Emily.Backend`.
22#
3+ # Reports two lanes by default and prints the speedup between them:
4+ #
5+ # * baseline — the op-by-op `Nx.Defn.Evaluator` walk.
6+ # * native — the single-NIF `Emily.Compiler` (`native: true`): the whole
7+ # generation graph, including the `defn while` decode loop,
8+ # replays as one NIF call per step. Run with
9+ # `native_fallback: :raise`, so the number is a true
10+ # full-native measurement, not a silent fallback.
11+ #
12+ # (The opt-in `EMILY_BENCH_FAST_KERNELS` lane is orthogonal — it benches the
13+ # fused MLX kernels under the evaluator.)
14+ #
315# Usage:
416#
517# elixir bench/qwen3_tokens_per_sec.exs
618#
7- # Standalone Mix.install so the bench doesn't need the github-ref
8- # Bumblebee pinned in Emily's own mix.exs — same rationale as
9- # `scripts/qwen3_conformance.exs`. See that script's header.
19+ # Standalone via Mix.install so a reader can run it without the project's
20+ # test setup; the dep versions below track Emily's own (keep in sync).
1021#
1122# Optional environment variables:
1223#
4051
4152Mix . install ( [
4253 { :emily , path: Path . expand ( ".." , __DIR__ ) } ,
43- # `override: true` because Emily's mix.exs declares
44- # `{:bumblebee, "~> 0.6", optional: true}` — Mix would otherwise
45- # refuse the github ref here as a child-dep conflict.
46- { :bumblebee ,
47- github: "elixir-nx/bumblebee" ,
48- ref: "273805e95507dc7866b958d90e0012a3abad1761" ,
49- override: true } ,
50- { :axon , "~> 0.7" } ,
54+ # Versions track Emily's own deps (Bumblebee 0.7 / Axon 0.8 / Nx 0.12);
55+ # keep them in sync with mix.exs so the standalone install resolves.
56+ { :bumblebee , "~> 0.7" } ,
57+ { :axon , "~> 0.8" } ,
5158 { :tokenizers , "~> 0.5" } ,
52- { :nx , "~> 0.10 " }
59+ { :nx , "~> 0.12 " }
5360] )
5461
5562defmodule Emily.Bench.Qwen3 do
@@ -99,9 +106,32 @@ defmodule Emily.Bench.Qwen3 do
99106 strategy: % { type: :greedy_search }
100107 )
101108
102- IO . puts ( "=== baseline (composed defn kernels) ===" )
103- baseline = bench ( model_info , tokenizer , generation_config , prompt , new_tokens , warmup , runs )
104- { baseline_mean , _ , _ , _ } = baseline
109+ cfg = % {
110+ tokenizer: tokenizer ,
111+ generation_config: generation_config ,
112+ prompt: prompt ,
113+ new_tokens: new_tokens ,
114+ warmup: warmup ,
115+ runs: runs
116+ }
117+
118+ # Baseline: the op-by-op Evaluator walk (~one BEAM↔worker round-trip per
119+ # op). Native: the single-NIF replay — the whole generation graph,
120+ # including the `defn while` decode loop, runs as one NIF call per step
121+ # (`native_fallback: :raise` so this is a true full-native measurement,
122+ # never a silent fallback to the evaluator).
123+ { baseline_mean , _ , _ , _ } =
124+ bench ( "baseline (evaluator, op-by-op)" , model_info , [ compiler: Nx.Defn.Evaluator ] , cfg )
125+
126+ { native_mean , _ , _ , _ } =
127+ bench (
128+ "native (single-NIF Emily.Compiler)" ,
129+ model_info ,
130+ [ compiler: Emily.Compiler , native: true , native_fallback: :raise ] ,
131+ cfg
132+ )
133+
134+ IO . puts ( "\n native speedup : #{ Float . round ( native_mean / baseline_mean , 2 ) } × over the evaluator" )
105135
106136 if fast_kernels? do
107137 unless Code . ensure_loaded? ( Emily.Bumblebee.FastKernels ) do
@@ -113,16 +143,18 @@ defmodule Emily.Bench.Qwen3 do
113143 System . halt ( 1 )
114144 end
115145
116- IO . puts ( " \n === fused ( Emily.Bumblebee.FastKernels) ===" )
146+ fused_model_info = update_in ( model_info . model , & Emily.Bumblebee.FastKernels . apply / 1 )
117147
118- fused_model_info =
119- update_in ( model_info . model , & Emily.Bumblebee.FastKernels . apply / 1 )
120-
121- fused = bench ( fused_model_info , tokenizer , generation_config , prompt , new_tokens , warmup , runs )
122- { fused_mean , _ , _ , _ } = fused
148+ { fused_mean , _ , _ , _ } =
149+ bench (
150+ "fused (Emily.Bumblebee.FastKernels)" ,
151+ fused_model_info ,
152+ [ compiler: Nx.Defn.Evaluator ] ,
153+ cfg
154+ )
123155
124156 speedup = fused_mean / baseline_mean
125- IO . puts ( "\n speedup : #{ Float . round ( speedup , 2 ) } × (fused mean / baseline mean)" )
157+ IO . puts ( "\n fused speedup : #{ Float . round ( speedup , 2 ) } × (fused mean / baseline mean)" )
126158
127159 if pin_threshold do
128160 if speedup >= pin_threshold do
@@ -135,10 +167,21 @@ defmodule Emily.Bench.Qwen3 do
135167 end
136168 end
137169
138- defp bench ( model_info , tokenizer , generation_config , prompt , new_tokens , warmup , runs ) do
170+ defp bench ( label , model_info , defn_options , cfg ) do
171+ % {
172+ tokenizer: tokenizer ,
173+ generation_config: generation_config ,
174+ prompt: prompt ,
175+ new_tokens: new_tokens ,
176+ warmup: warmup ,
177+ runs: runs
178+ } = cfg
179+
180+ IO . puts ( "=== #{ label } ===" )
181+
139182 serving =
140183 Bumblebee.Text . generation ( model_info , tokenizer , generation_config ,
141- defn_options: [ compiler: Nx.Defn.Evaluator ]
184+ defn_options: defn_options
142185 )
143186
144187 for _ <- Stream . duplicate ( :ok , warmup ) do
0 commit comments