You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit 8f05d8d
Browse filesBrowse the repository at this point in the historyBrowse files
#109 bundled a sliding-window KV cap, a Gemma4-specific MTP-assistant
auto-resolution path, and this benchmark tooling into one PR. The first two
are not part of this extraction:
- the KV cap was independently ported and evaluated earlier this session —
no measurable benefit at 150 tokens or ~9k context, parked
- the auto-resolution wiring (Gemma4MTPRegistry, mtpAsstRef reusing the
draft-model path) duplicates what #137 already shipped more generally as
the explicit --mtp-assistant-model flag; rebasing it forward would
reintroduce a second, narrower implementation of a feature that already
exists on main
This PR is only the tooling, and it rebases clean because it is genuinely
orthogonal: mtp_bench.py drives the server with --mtp --num-mtp-tokens N
--turbo-kv, the single-checkpoint MTP path, not the Gemma4 dual-model
wiring in conflict. Verified none of these three files reference
Gemma4MTPRegistry, mtp-assistant-model, or anything else from the withheld
part of #109.
run_benchmark.sh's Test 13 previously shelled out to `swift run Gemma4MTPBench`,
a product that no longer exists in Package.swift — that path was already
broken on main before this PR. It now drives mtp_bench.py against the regular
SwiftLM binary instead.
README's benchmark numbers (Gemma4 26B, 4-bit and 8-bit) are carried over from
the original PR's measurements, not reproduced in this extraction — a 40K/100K
context benchmark run is multi-hour. The 8-bit table matches data already used
as reference in mlx-swift-lm#46 and the #137 comment thread earlier this
session, so it is not new to this repo's history, just newly landing in the
README.
Verified: run_benchmark.sh syntax checked, mtp_bench.py compiles and its
--help output parses correctly; confirmed --mtp, --num-mtp-tokens and
--turbo-kv all already exist on main independent of --mtp-assistant-model.
Refs #109
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@@ -93,26 +92,51 @@ Benchmarked with `gemma-4-26b-a4b-it-4bit` running three configurations across 5
93
92
94
93
## 📊 Performance: Gemma 4-26B on Apple Silicon
95
94
95
+
Benchmark results for `gemma-4-26b-a4b-it-4bit` (**26B MoE, ~4B active params/token**, 4-bit) on M5 Pro 64 GB.
96
96
97
-
Benchmark results for `gemma-4-26b-a4b-it-4bit` (26B MoE, 4-bit) on M5 Pro 64 GB.
97
+
> ⚠️ This is a **Mixture-of-Experts (MoE)** model, not a dense model. Each token activates ~4B of the 26B parameters. "Vanilla" = all experts loaded into unified RAM (no SSD streaming).
> Values shown as `generation speed · GPU memory allocated`
113
+
> GPU Peak physical RAM (from `ioreg`): Vanilla 100K peaks at **21.8 GB** in-use · TurboQuant 100K stays at **17.1 GB** in-use
109
114
110
-
**Key takeaways:**
111
-
- 🚀 **Speed Doubled**: The newer MLX backend modifications have more than doubled raw `SSD Stream` inference speed (from 4.5 -> **10.8 tok/s**) while maintaining streaming stability.
112
-
- 📄 **40K context on 24 GB MacBook Pro**: SSD + TurboQuant effortlessly fits a 26B model in **22.5 GB** of memory footprint.
113
-
- 📚 **100K context on 24 GB MacBook Pro**: Due to hyper-efficient 3-bit KV compression paired with SSD weight streaming, you can process 100,000 tokens of context on a 24 GB machine — only utilizing **22.3 GB** total. (Previously required a 64 GB Mac Studio).
115
+
**Key takeaways (4-bit):**
116
+
- 🚀 **TurboQuant is the headline win**: At 100K context, `Vanilla + TurboQuant` delivers **66.9 tok/s** vs **27.5 tok/s** Vanilla — a **2.43× speedup**.
117
+
- 💾 **Massive memory savings**: OS RAM at 40K context drops from **48.7 GB → 18.2 GB** with TurboQuant (63% reduction).
118
+
- ⚡ **MTP neutral on 4-bit MoE**: The 4-bit model is compute-bound (MoE expert dispatch). Batch verification scales linearly with token count, so MTP provides no net throughput gain over vanilla at 4-bit.
119
+
- ⚠️ **TQ + MTP undercuts TQ alone**: Adding MTP to a TurboQuant server removes the bandwidth bottleneck (KV is now tiny) but not the FFN compute — MTP then adds overhead without proportional gains.
120
+
- 🖥️ **SSD Stream for 24 GB Macs**: Enables long-context inference with only ~22–27 GB RAM across all context depths.
- 🎯 **MTP works at 8-bit**: The 8-bit model is **bandwidth-bound** (2× heavier weights than 4-bit). The KV reads in the 3-token verification batch amortize across all 3 queries — so batch verification costs ~1.3× vanilla while producing 2+ tokens. Net: **+20% at 40K, +51% at 100K**.
135
+
- 🚀 **TurboQuant still wins outright**: At 100K, TurboQuant alone gives **48.3 tok/s** vs MTP's **22.5 tok/s** — TQ eliminates memory pressure and achieves 3.24× speedup over vanilla.
136
+
- ❌ **Don't combine TQ + MTP on 8-bit**: TurboQuant compresses the KV cache 4×, removing the bandwidth bottleneck and making MTP compute-bound again — TQ+MTP (31.0) is slower than TQ alone (50.1) at 40K.
137
+
- 💡 **Precision-dependent guidance**: Use `--mtp` for 8-bit models at 40K+ contexts without TurboQuant. Use `--turbo-kv` (without MTP) when maximum throughput or memory efficiency is the priority.
114
138
115
-
> Run `./run_benchmark.sh` to generate these metrics on your own device. (See **Benchmarks & Testing** below).
139
+
> Run `./run_benchmark.sh`→ Option 13 to reproduce these metrics on your own device.
read -p "Enter assistant model Hub ID: " ASST_MODEL
1352
-
fi
1353
1344
1354
1345
echo""
1355
1346
read -p "Enter context lengths to test [default: 512,40000,100000]: " CONTEXTS
1356
1347
CONTEXTS=${CONTEXTS:-"512,40000,100000"}
1357
1348
1358
1349
echo""
1359
-
echo"Building benchmark binary..."
1360
-
swift build -c release --product Gemma4MTPBench
1350
+
echo"Note: Test 13 uses the pre-built SwiftLM binary (same as Test 1)."
1351
+
echo" Make sure you have run ./build.sh first."
1352
+
echo""
1353
+
1354
+
python3 -u scripts/profiling/mtp_bench.py \
1355
+
--model "$FULL_MODEL" \
1356
+
--contexts "$CONTEXTS" \
1357
+
--max-tokens 60
1361
1358
1362
-
IFS=','read -ra ADDR <<<"$CONTEXTS"
1363
-
forctxin"${ADDR[@]}";do
1364
-
ctx=$(echo "$ctx"| tr -d '')
1365
-
echo""
1366
-
echo"--- Test 13: Context (max-kv-size=$ctx) on $FULL_MODEL ---"
1367
-
swift run -c release Gemma4MTPBench \
1368
-
--main-model "$FULL_MODEL" \
1369
-
--asst-model "$ASST_MODEL" \
1370
-
--prompt "Write a detailed 3-paragraph essay on the impact of the Industrial Revolution on modern supply chain logistics. Ensure you include dates and specific technological advancements." \
0 commit comments