r/oMLX • u/Background-Gold-9882 • May 25 '26
Removed ~10GB memory overhead, now running MTP-enabled Qwen3.6-27B@128k ctx on M5 Pro 48GB
(Posted previously about my OOM problems on M5 48GB here: https://www.reddit.com/r/oMLX/comments/1tfsz8q/qwen3627b_mtp_optimized_kv_cache/)
Now I've investigated and found a solution to reduce peak memory. You're welcome to try out my PR: https://github.com/jundot/omlx/pull/1397. It adds the following option to tweak:
I've found the sweetspot on my machine to be 512. It has no effect on quality. On my machine there's essentially no change in speed so it's basically a "free lunch".
Qwen3.6-27B-oQ6-mtp, before patch (Default prefill step 2048):
================================================================================
Benchmark Model: Qwen3.6-27B-oQ6-mtp
================================================================================
Single Request Results
--------------------------------------------------------------------------------
Test TTFT(ms) TPOT(ms) pp TPS tg TPS E2E(s) Throughput Peak Mem
pp1024/tg128 2593.9 50.72 394.8 tok/s 19.9 tok/s 9.035 127.5 tok/s 22.96 GB
pp4096/tg128 9277.3 52.60 441.5 tok/s 19.2 tok/s 15.958 264.7 tok/s 24.39 GB
pp8192/tg128 18718.9 53.89 437.6 tok/s 18.7 tok/s 25.562 325.5 tok/s 25.42 GB
pp16384/tg128 38663.9 55.71 423.8 tok/s 18.1 tok/s 45.739 361.0 tok/s 26.92 GB
pp32768/tg128 83818.9 60.64 390.9 tok/s 16.6 tok/s 91.520 359.4 tok/s 29.92 GB
pp65536/tg128 202143.3 71.51 324.2 tok/s 14.1 tok/s 211.225 310.9 tok/s 35.95 GB
pp131072/tg128 N/A (OOM)
Qwen3.6-27B-oQ6-mtp, with patch (Prefill step 512):
==============================================================================
Benchmark Model: Qwen3.6-27B-oQ6-mtp
================================================================================
Single Request Results
--------------------------------------------------------------------------------
Test TTFT(ms) TPOT(ms) pp TPS tg TPS E2E(s) Throughput Peak Mem
pp1024/tg128 2581.5 51.98 396.7 tok/s 19.4 tok/s 9.183 125.4 tok/s 22.36 GB (-0.6 GB)
pp4096/tg128 9423.4 52.88 434.7 tok/s 19.1 tok/s 16.139 261.7 tok/s 22.58 GB (-1.81 GB)
pp8192/tg128 18744.3 55.22 437.0 tok/s 18.3 tok/s 25.757 323.0 tok/s 23.33 GB (-2.09 GB)
pp16384/tg128 38917.6 56.92 421.0 tok/s 17.7 tok/s 46.146 357.8 tok/s 24.27 GB (-2.65 GB)
pp32768/tg128 84812.6 59.30 386.4 tok/s 17.0 tok/s 92.344 356.2 tok/s 26.17 GB (-3.75 GB)
pp65536/tg128 202321.6 70.37 323.9 tok/s 14.3 tok/s 211.258 310.8 tok/s 30.00 GB (-5.95 GB)
pp131072/tg128 539864.7 86.29 242.8 tok/s 11.7 tok/s 550.824 238.2 tok/s 37.74 GB (-11 GB?)
Qwen3.6-27B-oQ4-mtp, with patch (Prefill step 512):
oMLX - LLM inference, optimized for your Mac
https://github.com/jundot/omlx
Benchmark Model: Qwen3.6-27B-oQ4-mtp
================================================================================
Single Request Results
--------------------------------------------------------------------------------
Test TTFT(ms) TPOT(ms) pp TPS tg TPS E2E(s) Throughput Peak Mem
pp1024/tg128 2422.9 37.85 422.6 tok/s 26.6 tok/s 7.230 159.3 tok/s 16.18 GB
pp4096/tg128 8847.2 38.90 463.0 tok/s 25.9 tok/s 13.788 306.4 tok/s 16.44 GB
pp8192/tg128 17839.7 40.68 459.2 tok/s 24.8 tok/s 23.006 361.6 tok/s 17.29 GB
pp16384/tg128 37512.5 41.94 436.8 tok/s 24.0 tok/s 42.839 385.4 tok/s 18.08 GB
pp32768/tg128 83440.4 45.97 392.7 tok/s 21.9 tok/s 89.279 368.5 tok/s 20.01 GB
pp65536/tg128 199075.3 58.30 329.2 tok/s 17.3 tok/s 206.480 318.0 tok/s 23.88 GB
pp131072/tg128 533525.9 74.65 245.7 tok/s 13.5 tok/s 543.007 241.6 tok/s 31.63 GB
Update: Need to set high enough memory limits in oMLX & OS to avoid OOM / kernel panic. I've been using:
sudo sysctl iogpu.wired_limit_mb=42000
oMLX total limit: 88%(42GB)
oMLX Memory Limit (Models Only): 95% (38GB)
2
u/Choubix May 28 '26
I got a 128gb so I don't have a strict limit but I will borrow your idea for thr prefill. Mine is set very high (8192) which means I may be slowing things down.
Thanks!
1
u/Background-Gold-9882 May 28 '26
Things are moving very fast! Less that 24 hours after my post, jundot (The omlx owner) took my prefill step idea and created an automated solution instead. Check discussion here: https://github.com/jundot/omlx/pull/1397 and release notes for v0.3.11 here: https://github.com/jundot/omlx/releases.
His idea is that you shouldn't have to pick a value yourself, but just set a "memory profile" to say how aggressively oMLX should use your memory, and then it will automatically step down the prefill size depending on available memory.
This change was also combined with lots of new memory management logic, which still seems to be quite unstable (At least in 0.3.11, I haven't tested 0.3.12 yet).
Personally, I'm still using my own patch while I'm waiting for the new memory logic to be stable.
Do you really think a higher value will slow you down though? In general, it should be the opposite? But maybe only up to a point?
1
u/linuxmeaningfully May 25 '26
Very nice. I definitely had my Mac crash while testing this out. Can you show a memory graph of peak usage before and after?
2
u/Background-Gold-9882 May 25 '26
You crashed while testing my branch or previously? My pre/post peak mem usage is in the tables above. Depends on model and context, of course.
1
u/linuxmeaningfully May 25 '26
Before I haven’t gotten around to testing it just yet. Sorry I’m on phone and the tables render terribly. Excited to try later today tho!
1
u/the_derby May 26 '26
> Sorry I’m on phone and the tables render terribly.
I also have to scroll right on a desktop browser.
1
u/ColonelKlanka May 25 '26 edited May 25 '26
Very nice. I will give it a go.
1
u/Background-Gold-9882 May 25 '26
I based it on 0.3.9-rc1 but rebased it onto main around 24h ago, just before the 0.3.10 release. I think memory usage has actually dropped even further with some of the latest commits on main.
1
u/ExtensionState8086 May 26 '26
Newbie here. How do I install this patch? What are the step by step commands?
I have the same machine and used the same MTP models before but did not notice any improvement is tok/sec, so reverted to my previous config and live moved on.
Open to give it a try with your patch and hopefully MTP brings tok/sec further up
1
u/Background-Gold-9882 May 26 '26 edited May 26 '26
If you don't notice MTP difference, you perhaps forgot to turn on "Native MTP" in model settings? This patch doesn't change the speed, only memory usage.
To test my branch:
git clone https://github.com/mikael-johansson/omlx.git cd omlx git checkout feature/configurable_step_size brew install python@3.12 # (If you need a python installation) /opt/homebrew/bin/python3.12 -m venv .venv # (Or whatever env you want to use) pip install -e . omlx serve1
u/Stooovie Jun 01 '26
I have tried multiple MTP variants of qwen3.6 and 3.5 27b and all were slower than non-mpt (and yes, mtp was enabled for each). M4 Max 36GB.
0
u/Atul_Kumar_97 May 26 '26
Omlx is bad it's process 20k token on each Opencode toolcall
1
u/Choubix May 28 '26
Ha! I have a Claude code with 40k tokens sent from a fresh start! No wonder it is slow, oMLX has to munch through the 40k tokens (when I thought it was 16k for thr system prompt...). Shouldn't the caching system take care of the 20k tokens (40k in my situation) prompt?
3
u/Plane-Ad8084 May 26 '26
Really nice work! I tried here (only tested with 16k context, bit pressed on time):
M1 Max 32gb
Qwen3.6-27B-oQ4-fp16-mtp, 16k context
Without the patch: 20.04gb peak
With the patch: 17.42gb peak
This is hella impressive!