
I asked Codex to optimize DeepSeek V4 Flash 8-bit MLX on oMLX. Got ~1.6x prefill and ~3x decode speedup.
Follow-up to my earlier posts:
- Should I sell my Mac Studio? https://www.reddit.com/r/MacStudio/s/GK7QP8Lg87
- Kimi benchmark: https://www.reddit.com/r/LocalLLaMA/s/ujBsYLYmpd
Short version: my Mac Studio was sitting mostly idle, and from those Reddit threads I learned about DS4 and then oMLX. DS4 got me running DeepSeek V4 locally, but I wanted the 8-bit MLX version because I worry about accuracy loss in 4-bit variants.
So I tried mlx-community/deepseek-ai-DeepSeek-V4-Flash-8bit, the 302GB 8-bit affine MLX model, and asked Codex to optimize oMLX for the model.
I am not an oMLX/Metal kernel expert, so I am sharing this partly to sanity-check the work. Codex claims the changes should not reduce accuracy, and my Hermes/tool-calling runs look fine so far, but I would appreciate review from people who understand this stack better.
Base oMLX work
DeepSeek V4 support/tool calling came from oMLX DeepSeek V4 DSML/template/parser work, especially:
https://github.com/jundot/omlx/pull/2048
There were also follow-up fixes for DSML/tool-call stopping, parser-side stop behavior, prompt/prefix-cache determinism, shared expert SwiGLU clamp behavior, and native DeepSeek V4 2-bit/3-bit Metal paths.
The work below was separate: it focused on making the 8-bit affine model faster while keeping the same 302GB model format.
What changed for 8-bit affine
The issue was that DeepSeek V4 Flash 8-bit affine MoE was falling back to slower generic affine paths instead of using native DeepSeek MoE Metal kernels.
Codex changed:
- Enabled native DeepSeek affine MoE kernels for
bits=8,group_size=64 - Added 8-bit affine Metal kernel instantiations
- Replaced some generic route sorting with bucket/counting route paths
- Set
route_sort_min_routes=1so the native route path is used earlier - Added route-indexed decode kernels to avoid route sort/materialization overhead during decode
- Tuned affine8 dequant/load with a
uint32load specialization - Verified DeepSeek V4 parser/template and OpenAI-style tool calling still worked
Current config:
affine8_variant = 7
route_sort = bucket
route_sort_min_routes = 1
affine8_route_decode = 1
Results
| Metric | Before patches | After patches | Notes |
|---|---|---|---|
| Prefill | ~300-321 tok/s | ~533 tok/s | 12K prompt, salted uncached |
| Prefill | ~300-321 tok/s | ~528-530 tok/s | 30K prompt, salted uncached |
| Decode | ~7.31 tok/s | ~20-22 tok/s | Controlled benchmark |
| Decode | ~7.31 tok/s | ~19.5-20.7 tok/s | Real Hermes runs, ~80K-120K context |
Recent real Hermes/oMLX runs:
| Prompt size | Output | Result | Notes |
|---|---|---|---|
| 79K | 1,175 tokens | 19.8 tok/s | Long-context run |
| 80K | Tool call | 20.7 tok/s | Tool-calling run |
| 95K | 443 tokens | 19.5 tok/s | Long-context run |
| 117K | 1,680 tokens | 19.2 tok/s | Tool-calling run |
| 119K | 431 tokens | 19.3 tok/s | Tool-calling run |
Accuracy / correctness question
Codex says there should be no meaningful accuracy loss because:
- weights and quantization format were unchanged
- router/top-k/expert selection was not intentionally changed
- no experts were dropped
- optimized affine8 MoE outputs were compared against gather/qmm/native reference paths
- focused affine8 tests, DeepSeek V4 parser/template tests, and live tool-call smoke tests passed
I understand tiny numerical drift may still happen because kernel/load/order changed, but Codex claims this is not the same as a model-level accuracy drop.
Is that reasoning sound? What evals/tests would you run to verify no meaningful accuracy or tool-calling regression?
Next optimization direction?
Codex suggested these possible next directions:
- More affine8 dequant/load tuning: better vectorized loads, memory coalescing, fewer scale/bias reloads, less threadgroup-memory pressure
- Fewer kernel launches / less intermediate movement in routing and MoE buffers
- More fused MoE work, although this seems harder and riskier
- More single-token decode profiling, since real runs are still around ~20 tok/s
- Better instrumentation around routing, bucket/sort, block-plan build, native kernel time, down-projection, affine8 dequant/load, and decode costs
Questions:
- Is affine8 dequant/load tuning the right next direction for prefill?
- Has anyone done similar DeepSeek MoE route-indexed/fused/affine8 work in MLX, oMLX, llama.cpp, vLLM, or another runtime?
- Is ~530 tok/s prefill and ~20 tok/s decode on a Mac Studio M3 Ultra 512GB close to the ceiling for this 302GB 8-bit model, or is there obvious headroom?
Again, I am mostly asking the community to verify whether the result and next direction make sense.