One Wrong Parameter Was Costing 71% of My Prefill Throughput

The Setup

I've been running mlx-local-server — an OpenAI-compatible Rust/PyO3 inference server for Apple Silicon — and decided to systematically benchmark four optimization techniques across three sub-2B quantized models: Qwen2.5-0.5B, Llama-3.2-1B, and Qwen2.5-Coder-1.5B. Hardware: M4 Pro, 24 GB unified memory.

I measured prefill throughput (prompt tokens / TTFT), decode throughput (tok/s), and time-to-first-token at p50. Four iterations, each with a clear hypothesis. Here's what happened.

Iteration 1: The Embarrassing Win

The server had prefill_step_size=512. The mlx_lm upstream default is 2048. That's it — that's the whole story for iteration 1.

Changing 512 to 2048 gave:

• +33–71% prefill throughput across all models on a 480-token context
• −7–22% TTFT on medium and long prompts
• Zero decode regression, zero accuracy impact

Why does this matter? Apple Silicon's GPU reaches much higher utilization when processing 2048 tokens in one dispatch vs. 512. The 4× larger chunk means fewer round-trips between the CPU scheduler and the Metal command queue. The prefill path is compute-bound, not memory-bandwidth-bound, so bigger batches win.

The lesson is obvious in hindsight: verify framework defaults before adding any complexity. I added Metal memory management at the same time (set_cache_limit, set_memory_limit), which eliminated some OOM panics in long sessions but didn't affect throughput numbers.

Iteration 2: KV Quantization Backfired

4-bit KV cache quantization is widely recommended for inference servers. Compress the stored key/value tensors by 4×, save memory bandwidth, go faster. On CUDA hardware with HBM, this math works. On Apple Silicon, it doesn't — at least not for short contexts.

Results with kv_bits=4 as the server default:

• −23–38% decode throughput across all models
• +9–41% TTFT (worse)

The problem: on unified memory, the same bus serves both model weights and KV cache. Dequantizing the KV cache at each decode step costs real ALU cycles on the Metal GPU. For contexts under ~500 tokens — which is most interactive use — that cost exceeds the bandwidth saved. The crossover point is probably around 800–1,200 tokens based on the throughput curves.

I kept KV quantization as an opt-in per-request parameter (kv_bits field). It's genuinely useful for long-document tasks. Just not as a default.

Iteration 3: Speculative Decoding Was Basically Noise

Speculative decoding — using a small draft model to propose multiple tokens that the target model verifies in one pass — can give 2–3× speedup on CUDA systems. I benchmarked Qwen2.5-0.5B as draft for Qwen2.5-Coder-1.5B on three code-generation prompts (highest expected acceptance rate).

Result: +1% average decode throughput. Within measurement noise.

The math explains it. For a 1.5B/0.5B pair, the target/draft cost ratio is roughly 3:1. With num_draft_tokens=4, the draft runs 4 times per verification step at (1/3)× target cost = 1.33× overhead. Even at 90% token acceptance, the bandwidth savings barely exceed the draft overhead on a shared memory bus.

Speculative decoding is a bet on the target model being bottlenecked by sequential forward passes. For sub-2B models on unified memory, that bottleneck doesn't dominate. For 7B+ target models with a well-matched 1B draft, the speedup should be real.

Iteration 4: Auto Warm-Up

After every model load, I now fire a 1-token background inference ("Hi") before any user traffic arrives. This pre-compiles Metal compute shaders and sizes the KV allocator pool.

The first real user request normally pays 200–1,000 ms of Metal JIT compilation overhead. The warm-up eliminates that spike. In a benchmark averaged across N=3 requests, the effect is diluted — but in production where users see the very first request, it's a meaningful latency reduction.

Also increased the Metal allocator cache from 512 MB to 1 GB, which lets KV cache buffers from one request be recycled by the next without re-allocation.

Summary

Four iterations, two shipped, two reverted:

• ✅ prefill_step_size=2048 — the single biggest gain, just fixing a default
• ✅ Metal memory management + auto warm-up — eliminates OOM and cold-start spikes
• ❌ KV-4bit default — -38% decode at short ctx; keep as opt-in
• ❌ Speculative decoding (sub-2B pair) — +1%; try with 7B+ targets

The full benchmark data, scripts, and the academic paper are on GitHub.