Repository navigation
Conversation
Sampled chains already score the vocabulary with a 3-bit copy of the lm_head and rescore the best 64 candidates with the real weights. Greedy chains still run the full-vocabulary head on every draft step. With OMLX_MTP_GREEDY_COARSE=1, greedy chains take the best exactly rescored coarse candidate instead. Verify is unchanged and still decides every emitted token. Logits processors, heads the coarse path does not support and sampled requests behave as before, and the default is off. The greedy path does not quantize the MTP fc; a sampled request in the same process still does.
1b0a083 to
aafb5ae
Compare
|
Numbers from a second machine and model, in case they help the review. One caveat first: this is not this branch. It is a smaller always-on variant of the same idea that I had measured before I found this PR. Greedy takes the argmax of Setup: M5 Max 128 GB, oMLX v0.7.4.dev1, Qwen3.8-Flash-Next oQ3e with MTP on (adaptive depth), temperature 0, one request at a time, 48 prompts from public benchmarks (HumanEval and others), up to 400 output tokens each. Two runs per arm, in the order A B B A. Decode tok/s, pooled over the 48 prompts (15,605 tokens per run):
That is +6.5%, with a 95% bootstrap interval over prompts of +5.7% to +7.2%. 47 of the 48 prompts were faster. Output: the two patched runs gave the same text as each other on 48 of 48 prompts, and the same text as one of the two main runs on 48 of 48. The two main runs differ from each other on 1 prompt, so the change is inside the run-to-run noise here. An earlier run with the oQ4e quant on a late-October base (16 prompts, 2,000 tokens each) gave +6.1% and the same text as MTP off on 16 of 16. Given that, I think this could be the default path rather than an opt-in. Verify still decides every emitted token, and a server that also gets sampled requests has already built the coarse copy. |
|
Thanks, that's useful, especially the 48/48 matching output. I'll keep it opt-in here and flip the default in a follow-up once this lands. @jundot, if you'd rather have it on by default in this PR, I can change it. |
|
Tested
|
|
Second data point for this PR: M5 Ultra, same probe and ABBA arm order as my M1 Max numbers above. The gain is smaller but it persists on hardware with a wide qmv path. Setup — Decode tok/s, mean of 3 hits after a warmup
Absolute throughput is 2.2–2.7× the M1 numbers, so the relative gain shrinking is expected if part of the win is amortising fixed per-step overhead. Losslessness: greedy SHA over all four arms is 11/12 identical. The one difference is Opt-in by default still looks right to me; the effect is real on two GPU generations, but it is a few percent, not a step change. |
Summary
Sampled MTP drafting already scores the vocabulary with a 3-bit copy of
lm_headand rescores the top 64 candidates with the real head. Greedy drafting still runs the full head on every draft step. This PR lets greedy use the coarse path too.OMLX_MTP_GREEDY_COARSE=1.Performance
M4 Max 64 GB, MLX 0.32.4, batched engine, Qwen3.8-27B oQ4e (8-bit
lm_head), MTP on, SpecPrefill off, temp 0, one request, cold prompt, 400 output tokens. Measured on this branch after rebasing onto165be1ff. Decode tok/s, pooled median, flag off vs on, rounds interleaved:Every round improved. Cycle time went from 57.8 to 55.7 ms at 4K and from 64.8 to 61.6 ms at 16K. Tokens per cycle also went up, from 2.81 to 3.04 and from 2.72 to 2.86. I haven't worked out why, so I'd only credit the cycle time to this change.
Raw decode tok/s (3 runs per round)
4K off: 47.85 48.67 49.40 / 47.44 50.55 46.65 / 51.17 49.88 52.13
4K on: 52.60 56.13 52.20 / 54.98 55.94 51.66 / 54.56 53.11 55.41
16K off: 44.81 41.49 43.99 / 46.63 41.22 42.09 / 46.16 43.50 45.57 / 45.91 42.33 43.29 / 46.35 44.33 41.27
16K on: 50.75 50.22 47.29 / 48.62 45.39 47.56 / 43.29 47.73 47.47 / 48.10 45.12 45.47 / 49.91 49.48 42.73
Trade-offs
mainhas the same property.fc. Sampled requests still do, and after one runs in the same processfcstays 4-bit.