Reproducing the 2.24× speedup: documented MLX patch alone doesn't get there — please publish the actual fork commit

[benjamin-levin]

Reproducing the 2.24× speedup: documented MLX patch alone doesn't get there — please publish the actual fork commit

Hi @youssofal — first, thank you for the clear, comprehensive README and the verified benchmark numbers. The single-model native-MTP architecture is genuinely impressive work, and the public default model Youssofal/Qwen3.6-27B-MTPLX-Optimized-Speed runs cleanly out of the box.

I spent ~6 hours this week trying to reproduce the headline 2.24× / ~63 tok/s number on M4 Max 36 GB by following the public source path:

1. ✅ Cloned and built MTPLX/native_extensions/verify_mlp (the gdn-tail + gate-up Metal kernels)
2. ✅ Cloned ml-explore/mlx v0.31.2 and applied the documented qmv_fast_impl patch:
   • constexpr int num_simdgroups = 2; → 4
   • int bn = 8; / MTL::Size group_dims(bk, 2, 1); → bn = 16; / MTL::Size group_dims(bk, 4, 1); in qmv() launch config
   • #pragma clang loop unroll_count(4) on the inner row loop in qmv_fast_impl
3. ✅ Wheel built (mlx-0.31.2.dev20260513+68cf2fd-cp313-cp313-macosx_15_0_arm64.whl), installed in the brew mtplx venv
4. ✅ mtplx doctor --deep confirms patched MLX loads cleanly, variant: native_full_rowwise + enabled: True when MTPLX_NATIVE_MLP_ROWWISE=1

But the end-to-end result is identical to stock:

ctx	tok/s, sustained	tok/s, performance-cold
512	—	20.8
1k	15.4	15.5
4k	6.5	6.4
16k	1.9	—

(Native kernels are confirmed active via mlp_variant_stats.native_calls > 0 in the health endpoint. Patched MLX is confirmed loaded via mlx_fork.version = 0.31.2.dev20260513+68cf2fd. mlx_fork.ok = False because our build commit ≠ 2377a99f.)

Direct kernel microbench (Mac Studio M4 Max, mx.quantized_matmul, K=5120, N=5120, group_size=64, bits=4, 200 iters median):

M	Stock MLX 0.31.2	Patched (with unroll)	Patched (no unroll)
1	0.134 ms	0.151 (+12%)	0.142 (+6%)
3	0.173 ms	0.210 (+21%)	0.159 (-8%)
4	0.172 ms	0.218 (+27%)	0.168 (-2%)
6	0.200 ms	0.238 (+19%)	0.200 (0%)
8	0.233 ms	0.264 (+13%)	0.236 (+1%)
16	0.272 ms	0.291 (+7%)	0.266 (-2%)

Two observations:

1. The unroll_count(4) pragma is a regression on M4 Max at every M. Probably an M5 Max-specific tuning that doesn't transfer.
2. Even without the pragma, the patch is at best 2-8% faster at M=3,4 — far from the kernel-level speedup the headline 2.24× would require if qmv is the dominant verify-MLP cost.

Request: would you be willing to publish the actual contents of mlx-mtplx-0.31.2-qmm at commit 2377a99f? Either as a git remote URL, a public branch on a fork, or even a patch series attached to the repo. Specifically, if there are changes beyond what README.md lines 228-230 describe — e.g. additional kernel restructuring, register-pressure tuning, different reduction strategies, simdgroup-shared-memory work, anything in qmv_impl or related kernels — those are the bits the community can't reverse-engineer from the README.

Why this matters: MTPLX is currently the only documented native-MTP MLX runtime, and Youssofal/Qwen3.6-27B-MTPLX-Optimized-Speed is genuinely the best Apple-Silicon-local thinking-agent model I've tested. But the gap between "I followed the README" and "I get the headline number" is a UX cliff for anyone on M4 / M3 Max hardware. Publishing the actual fork — even as a sealed read-only branch — would let users on non-M5 hardware extract a tractable subset of the speedup.

If publication isn't possible (commercial, sponsored, in-flight upstream PR, etc.), that's a fair answer too — just knowing for sure would save reproducers a lot of time.

Hardware: Mac Studio M4 Max, 36 GB unified memory, iogpu.wired_limit_mb=32000, mx.set_wired_limit(31 GB). macOS 15.7.4. MTPLX 0.3.4 from your homebrew tap, Youssofal/Qwen3.6-27B-MTPLX-Optimized-Speed model. Stack baseline (plain mlx_lm.server, no MTP, no patches) gets 50.5 tok/s @ 128k context on the same hardware with Qwen3.6-35B-A3B-4bit, so the hardware itself isn't the constraint.

Thank you again for the project — happy to share any additional diagnostic data (mtplx doctor --deep --json, per-cycle traces, etc.) that would help triage.

[youssofal]

Thank you for the comprehensive analysis, to be completely honest the readme is outdated from the pre-launch era. Recently I tested both my custom QMV and stock MLX. the QMV saved around 10% aggregate only on short generations (168 tokens) and actually caused regressions on long generations. It achieved the headline figure but I did not use it for the current sustained mode (replacing Burst mode) because it was poor on long generations. The current MTPLX runs primarily on stock MLX. The key point is, I still achieve a 2.1x decode TPS increase without the QMV patch on short generations. The big discrepancy is on verify times.

I have noticed older chips (M2,M3,M4) suffer worse verify times compared to the M5. It is to the point I have had to reccomend users of M2 and M1 chips to lower MTP heads to 2 down from 3 to prevent verify time domination.

The key lies in solving and understanding the cause of poor verify times on older chips, My understanding is it is compute bound but I am hoping there are other optimization channels.

I would like to see your acceptance figures and verify time figures to diagnose this further, in the latest update I have added a new command.

Could you please run: MTPLX bench tune and paste your output here? It will show me your verify times and hardware utilisation.

(@AirRunner) tagging you since you have asked me the same thing previously)

[benjamin-levin]

Done — uninstalled the patched mlx-0.31.2.dev+68cf2fd wheel, reinstalled stock mlx==0.31.2, bumped the tap to mtplx 0.3.6, and installed thermalforge. Posting the official mtplx tune numbers below.

mtplx tune (performance-cold, thermalforge max-fan verified, stock MLX 0.31.2) — M4 Max 36 GB, Youssofal/Qwen3.6-27B-MTPLX-Optimized-Speed, 192 tokens/candidate, isolated subprocesses, fans ramped to 3473-3475 RPM (96% of 3625 capacity) and confirmed before each candidate:

mode	tok/s	× AR	verify_ms/call	verify_hidden_eval_ms	verify_forward_ms	accept (per depth)
AR	18.93	1.00	—	—	—	—
D1	30.83	1.63	44.47	42.42	2.05	91.0%
D2	32.98	1.74 (BEST)	58.38	56.25	2.13	92.75% / 86.76%
D3	29.94	1.58	74.81	72.66	2.14	86.67% / 71.67% / 62.71%

Tune picked D2 and saved it (Saved: Web UI starts will use depth 2 for this model.). Matches your "drop M1-M4 to depth 2" guidance exactly.

The clean split between verify_hidden_eval_ms_per_call (linear in depth: 42 → 56 → 73 ms) and verify_forward_ms_per_call (basically constant at ~2.1 ms) is the clearest picture I've gotten of where the verify cost lives. Verify-forward looks compute-bound and depth-invariant; the per-depth hidden-eval is what scales. So on stock MLX 0.31.2 the verify path is dominated by the hidden-state evaluation, not the LM-head forward — i.e. the speedup ceiling on this Mac is set by how fast you can run the GDN-tail + gate-up hidden-eval lanes, which is exactly the surface QMV targeted.

verify-profile (per-stage attribution; sync inflates absolute times, composition is real). At 6 verify tokens, mean GDN-layer = 167 ms vs attention-layer = 48 ms (per-layer GDN ~3.5× attention). Per-stage stock_elapsed for v=6: mlp=82ms, gdn_projections=24ms, gdn_out_proj=13ms, gdn_recurrent=9ms, attn_projections=7ms. Heaviest GDN terms are gdn_projections + gdn_out_proj (the qmatmul lanes) — same spot QMV targeted, consistent with it only paying off short-gen.

mlx_fork.ok = false, optional_fast_mlx_fork_active = false, stock_mlx_likely = true — so the 1.74× above is the pure stock-MLX MTP envelope on M4 Max. Nice to have a clean number to anchor to.

Full artifacts (tune.json + per-candidate AR/D1/D2/D3 JSONs): https://gist.github.com/benjamin-levin/986842d01bf8e8dfb12bf5822699eaea

Happy to share mtplx doctor --deep --json too if useful.

[tienlbhoc]

@benjamin-levin , how much ram when you use around 100 150k context ?

[wwadge]

M4 Max 48GB here, getting 45-50 tps

[benjamin-levin]

@tienlbhoc — measured numbers below. M4 Max 36 GB, Youssofal/Qwen3.6-27B-MTPLX-Optimized-Speed, mlx_lm.stream_generate, MTP enabled (default depth), prefill via mlx-lm's chunked path, 4-token decode after prefill (peak is dominated by prefill anyway). Baseline active memory after model load is 14.09 GB.

context	wall (s)	prefill tok/s	decode tok/s	peak unified memory
16 k	82	195.3	27.3	18.84 GB
32 k	174	184.3	25.7	21.34 GB
65 k	398	163.4	22.6	26.36 GB
100 k	689	145.3	20.0	31.87 GB
150 k	—	—	—	extrapolated ~39 GB → OOM on 36 GB Mac

So short answer for the 100–150k window:

• 100 k fits, but only barely. Peak (31.87 GB) lands right against the wired-memory cap (max_recommended_working_set_size ≈ 32.71 GB on my 36 GB Mac). I would not recommend running another big process alongside it.
• 150 k does not fit on a 36 GB M4 Max without modification. The peak-vs-context slope is ~155 KB/token in the regime above 16 k (KV per token alone is ~64 KB for this model's 16 full-attention layers; the rest is prefill intermediates / activations). Extrapolating to 150 k gives ~39 GB peak, which is well past the wired limit.

To get to 150 k on 36 GB you'd need one of:

• 8-bit (or 4-bit) KV-cache quantization on the full-attention layers (would knock ~3-5 GB off at 150 k),
• a smaller prefill chunk size so fewer in-flight activations pile up at the peak,
• or just 48 GB+ of unified memory.

@wwadge — your 48 GB Mac would clear 150 k comfortably (extrapolated ~39 GB peak, leaving ~9 GB headroom).

A few other things from the same run worth noting since they relate to the verify-time discussion:

• Prefill tps drops with context: 195 → 145 tok/s from 16 k → 100 k. That's the per-chunk attention cost scaling with T_kv. The full-attention layers (16 of 64) are where this lives.
• Decode tps drops correspondingly: 27.3 → 20.0 tok/s. Same root cause — KV-bandwidth at long context. Consistent with your verify-profile observation that the heaviest stage is the GDN-tail / gate-up qmatmul lanes.
• GDN state is fixed-size (1 × 48 heads × 128 × 128 × bf16 = ~1.5 MB per layer × 48 layers = ~72 MB total). All the context-scaling memory comes from the 16 full-attention layers.

I can re-run with mtplx ask --max-tokens N directly if you want sustained-decode numbers at long context rather than the prefill-then-4-decode profile. Just let me know what budget you'd like (it took ~11.5 min to do 100 k prefill).

[tienlbhoc]

I think i need kv q8 q6 q4 options for this app

[youssofal]

I think i need kv q8 q6 q4 options for this app

Adding soon.

[benjamin-levin]

(Replacing my earlier deleted comment with corrected numbers — verified against Apple's official Mac Studio spec page.)

A possible explanation for the gap between my 20 tok/s @ 100 k on 36 GB and @wwadge's 45–50 tok/s on 48 GB:

The 36 GB Mac Studio is structurally tied to the binned M4 Max — 14-core CPU (10P+4E) / 32-core GPU / 410 GB/s unified-memory bandwidth. The 48 GB and 64 GB Mac Studio configs ship the full M4 Max — 16-core CPU (12P+4E) / 40-core GPU / 546 GB/s. Apple does not offer the full M4 Max with only 36 GB. So a 36 GB result and a 48 GB result aren't just different memory capacities — they're different bandwidth tiers.

Mac Studio unified-memory bandwidth (verified, Apple spec page)

Mac Studio 2025 ships M4 Max + M3 Ultra — there is no M4 Ultra; Apple kept the older Ultra die.

Chip variant	CPU	GPU	Memory bandwidth	Mac Studio memory
M4 Max (binned)	14c (10P+4E)	32c	410 GB/s	36 GB only
M4 Max (full)	16c (12P+4E)	40c	546 GB/s	up to 64 GB
M3 Ultra (binned)	28c (20P+8E)	60c	819 GB/s	from 96 GB
M3 Ultra (full)	32c	80c	819 GB/s	higher BTO tiers

The M3 Ultra row's bandwidth is the same across binned and full because chip binning disables cores, not memory channels.

Hypothesis for the tps gap

MTPLX decode at long context is dominated by KV-cache traffic on the 16 full-attention layers (the linear-attention layers have fixed-size state). On this model and prompt, peak unified-memory at 100 k is 31.87 GB on my 36 GB Mac Studio — right against the wired-memory cap (max_recommended_working_set_size ≈ 32.71 GB). Two things plausibly compound to produce the observed 20 → 45–50 tok/s gap between 36 GB and 48 GB:

1. Raw bandwidth: 410 → 546 GB/s = +33%. If decode is bandwidth-bound at long context, that alone moves 20 tok/s toward ~27 tok/s.
2. Wired-cap headroom: at 31.87 GB peak on a 32.71 GB cap, prefill is operating right at the boundary. Anything that nudges over the cap forces eviction/recompute on subsequent KV reads. A 48 GB Mac Studio with the full M4 Max has a ~40–41 GB wired cap, well above the 31.87 GB peak — no boundary pressure during prefill.

I can't decompose the contribution of (1) vs (2) without a 48 GB box to A/B against, but both plausibly land in the same direction and explain the drastic difference. Would be interesting to see @wwadge's mtplx bench tune output on the same model — verify times in particular — to confirm whether the gap is bandwidth-dominated or also reflects something architectural about the binned-vs-full M4 Max chip beyond memory channels.

[wwadge]

Just for completing my data point: I'm with the 16c so 546 GB/s bandwidth, I just typed "hello" as as context there. This is the first time I'm getting decent speeds with the qwen-27b so many thanks for that!

[youssofal]

@benjamin-levin Thank you for the detailed analysis.

Yes, and it gets deeper than that. With MTP unlike traditional models speedup isnt only bandwidth bound but compute bound. Having more GPU cores (and higher frequencies) leads to a significant speedup, paired with lower bandwidth is why there is such a discrepancy.

Untill I figure out ways to further bring down verify times in an approach similar to Flash infer all reduce on CUDA that is better than mx compile there is nothing to be done. Experience with MTP is hardware based.

That being said, one contribution I can make is this:

It is incredible how much just slight background GPU load impacts decode TPS. When I ran my headline benchmarks I had every single app closed on my Mac. Even a chrome tab can lead to a 5% decrease in TPS.

I will add into the MTPLX bench command a check for background app usage.

[benjamin-levin]

For anyone wanting to A/B these MTPLX numbers against the same workload on a stacked-optimization mlx-lm — I made my forks public. Single pip install pulls both:

pip install git+https://github.com/benjamin-levin/mlx-lm.git@main

Requires Xcode CLT + CMake; the mlx fork builds from source (~5–10 min on M-series). setup.py auto-pulls the matching mlx fork. These will all be upstreamed to ml-explore/* once each PR clears review — the forks are interim.

What's stacked (M4 Max 36 GB, Qwen3.6-35B-A3B-4bit unless noted)

mlx — C++/Metal:

• mx.fast.fused_qsdpa — 2-pass FlashAttention on quantized KV with cooperative tile-shared dequant across GQA heads. +18% @ 32k N=1, +27% @ 96k N=1. left_padding extension recovers +22.8% on hetero N=3 batches.
• mx.fast.fused_swiglu_gather_qmv — fuses silu(gate)*up inline with the quantized matmul (one Metal dispatch instead of three). ~1.05× geomean across 27 shapes.

mlx-lm — Python:

• prefer-prefill scheduler (opt-in) — pauses decode while prefill is queued. 2.58× N=3 @ 8k, 1.71× N=3 @ 32k, no-op at N=1.
• bf16 GDN state (opt-in env var) — kernel −8.9%; e2e modest on this baseline, larger stacked.
• MoE compile-block — bit-exact, +0.4–1.1% e2e in isolation (within noise; structural cleanup).
• Depth-1 MTP spec decode — 1.13–1.21× across echo / code / qa workloads.
• Persistent disk-backed prompt cache — 1477–4585× TTFT on cached prefixes (4k–96k), bit-exact.
• SnapKV long-context KV compression — 1.21× @ 32k → 1.94× @ 128k, no-op below 49k.
• Generic cache snap/restore + bit-exact PLD — unlocks PLD on GDN/Mamba ArraysCache (1.40–1.68× GDN echo/code-edit).
• (Auto-speculative-router is on the auto-speculative-router branch only — conflicts with the bit-exact PLD chosen for main.)

Why fusions dominate this stack

The M4 series and below have no native int4/int8 GPU matmul instructions — M5's Neural Accelerators added these, but anything M4 or earlier dequantizes every quantized weight to bf16 in registers before the matmul fires. Storage saves; on-chip bandwidth during compute is bf16-equivalent. KV-quant in particular only pays off when the dequant is fused into the attention kernel (the stock 3-launch quantized-SDPA path is actually slower than bf16 SDPA at long context). Most of the fusions here exist to amortize that per-op dequant cost.

Per-PR isolated numbers + any honest corrections vs originally claimed wins live in each draft PR's body on the fork.

[AirRunner]

@benjamin-levin have you seen ml-explore/mlx#3026 ? It's adding a fused quantized KV attention kernel to MLX stock, targeting a similar problem as your fused_qsdpa.

There is an interesting architectural difference though: #3026 dequantizes in registers per GQA head relying on L2 sharing, while you dequantize into threadgroup memory once per KV row in your approach. So your approach may amortize the cost across all GQA heads, that may explain why their benchmarks on D=256 GQA 6:1 don't show speedup at long context.

[wwadge]

Some more data points:

After running mtplx tune on my 16c M4 Max (48GB), "samuelfaj/Qwen3.6-35B-A3B-4bit-MTPLX-Optimized-Speed":

AR 60.4 tok/s 1.00x
D1 68.2 tok/s 1.13x BEST
D2 56.1 tok/s 0.93x
D3 46.4 tok/s 0.77x

[wwadge]

After #77 ; 35b is actually 98 tok/s now.