optimizing-omlx-mtp-under-contention.mdx
Post File

When speculative decoding lost to ordinary batching

by discostew

How I traced an oMLX concurrency slowdown to mixed decode paths, replaced it with a small scheduler policy, and got the fix merged upstream.

  • oMLX
  • inference
  • performance engineering
  • Apple Silicon
  • open source

Why choosing speculative decoding by default can be the wrong policy under contention.

TL;DR

I contributed a small scheduler change to oMLX after finding that its VLM MTP path could reduce total throughput under concurrent load.

On my measured Qwen3.8-27B 4-bit workload on an Apple M3 Ultra, the old mixed route produced 69.34 tokens per second for a greedy batch of four. The new policy produced 108.14 tokens per second. The sampled batch improved from 64.68 to 103.79 tokens per second.

The fix did not change MTP generation. It changed when oMLX chooses MTP: keep it for an uncontended request, but use the existing batch generator when peers are already ready or admitted.

The two-file change was merged into oMLX as PR #2752.

The performance problem

MTP speculative decoding was useful for a single request. The problem appeared when several requests became ready together.

The scheduler could send the first request down the MTP path while its peers fell back to ordinary BatchGenerator batching. The outputs remained correct, but the two decode paths competed for the same model and GPU. The result was a mixed group that ran much more slowly than keeping all four requests on the ordinary batched path.

That distinction mattered. The problem was not “MTP is slow.” The evidence showed a routing problem:

  • one request could benefit from MTP;
  • a ready group benefited from ordinary batching; and
  • mixing the paths under contention lost aggregate throughput.

The smaller fix

A broad response would have been to build true batched MTP or migrate an active MTP request into a different decode path. Both would require larger scheduler and cache-lifecycle changes.

I chose a narrower policy that matched the measured failure mode.

Before oMLX performs the MTP-only final target-model forward, the scheduler now checks for peers in its waiting, running, and prefilling collections. If a peer is present, the route declines MTP and lets the existing BatchGenerator fallback handle the group. A chunked-prefill request is excluded from counting itself as contention.

The merged scheduler change added no configuration and no scheduler state. It left MTP generation, verification, rollback, cache ownership, and sampling unchanged.

That was the engineering decision: change the routing seam, not the decoding machinery.

What I measured

I tested a Qwen3.8-27B 4-bit target with its 4-bit MTP drafter on an Apple M3 Ultra with 512 GB of unified memory. Each request used a 1,024-token prompt and generated 100 tokens. The matrix covered greedy decoding and sampled decoding with temperature=0.7, top_p=0.95, and top_k=20.

Workload Plain batching Previous mixed MTP New contention policy
Batch 1, greedy 35.98 tok/s 47.53 tok/s 45.69 tok/s
Batch 1, sampled 35.51 tok/s 44.06 tok/s 42.00 tok/s
Batch 4, greedy 107.40 tok/s 69.34 tok/s 108.14 tok/s
Batch 4, sampled 102.93 tok/s 64.68 tok/s 103.79 tok/s

In this setup, the new policy improved batch-4 aggregate decode throughput by about 56% for greedy decoding and 60% for sampled decoding compared with the previous mixed route. It also landed within 1% of the plain-batching baseline in both batch-4 profiles.

The single-request policy cases still used MTP and remained faster than plain decoding in this run. I treat that as evidence that the routing guard preserved the intended single-request behavior, not as a universal MTP speedup claim.

The aggregate decode metric was calculated as:

total completion tokens / (batch wall time - longest request TTFT)

This excludes the serialized prefill interval and measures the shared decode window consistently across the compared runs. Plain and mixed-MTP baselines used three repeats; the new-policy confirmation used two.

The complete methodology, sanitized JSONL records, provenance hashes, and checksums are in the public evidence package.

Correctness and regression coverage

Throughput was only half of the acceptance bar. Concurrent generation also needed evidence that requests were not leaking output into one another.

The policy run used distinct row markers and required answers across eight chat-correctness cases and 20 response rows. Every row contained its own marker and answer, no row contained another request’s marker, and every case completed normally. The published artifacts remove generated text while retaining timings, token counts, output hashes, finish reasons, MTP statistics, correctness checks, memory measurements, and run-completion records.

The merged regression tests verify four important boundaries:

  • a waiting peer declines MTP before the final model forward;
  • a running peer does the same;
  • a prefilling peer does the same; and
  • the request currently finishing chunked prefill does not count itself as a peer.

The PR reported 291 focused scheduler and VLM-MTP tests passing, plus focused Ruff fatal-error checks and git diff --check. Upstream CI then passed on Python 3.11, 3.12, and 3.13 before the merge. Those results are recorded on the merged PR.

What the result does and does not show

The measured conclusion is narrow: on this model, host, prompt length, generation length, and concurrency pattern, avoiding mixed MTP and ordinary batching recovered batch-4 throughput without disabling MTP for an uncontended request.

The broader engineering lesson is a hypothesis worth carrying to other systems: an optimization should be selected for the workload the scheduler can see, not enabled merely because it is faster in isolation.

This change does not solve every arrival pattern. If another request arrives after an MTP request has already begun decoding, the original request continues with MTP and the late arrival still falls back to BatchGenerator. Moving active decode state, serializing late arrivals, or implementing true batched MTP remains outside the scope of this fix.

Why this contribution matters to me

The code was small: one routing guard and focused regression coverage, with 99 additions and 15 deletions across two files.

The work around it was the larger contribution:

  • isolate the contention-specific failure mode;
  • compare the mixed route against plain batching instead of assuming speculative decoding should win;
  • choose a change at the smallest useful seam;
  • test greedy and sampled behavior;
  • preserve correctness evidence and provenance in a reviewable package; and
  • state what the patch intentionally does not solve.

That is the kind of performance engineering I want to keep doing: measure the real system, make the smallest change that explains the result, and leave enough evidence for someone else to audit it.

The final implementation is in oMLX merge commit fc640a8a.