2x Decode From Block-Diffusion Drafting

After an engine swap failed to beat our incumbent, the decode win came from changing how tokens are drafted rather than how they're computed.

Summary

We'd concluded that engine rewrites recover at most 1.35× on prefill and nothing on decode, because decode is bandwidth-bound. The remaining lever was speculative decoding — specifically block-diffusion drafting, where a small draft model proposes several tokens at once and the target verifies them in a batch. Measured on the same prompt as every other benchmark in this series: decode went from 28–34 tok/s to 62–67 tok/s at 26k context, with prefill also improving. Roughly 2×, and the first genuinely large decode win we've found.

Why this was the only lever left

Some accounting from earlier in this series.

Our incumbent server uses the checkpoint's multi-token-prediction head at depth 3, yielding 2.10 accepted tokens per round and 28–34 tok/s decode at 26k context.

Engine replacement is a dead end: the alternative we evaluated (the engine-evaluation post) won prefill by 4–8% and lost decode by 30–40%. A kernel-level port (the kernel-unfusing post) got to 98.3% of the mature framework, which is to say there was nothing left there either.

Decode is memory-bandwidth-bound. You cannot make a bandwidth-bound loop faster by computing more efficiently. You can only make it run fewer times — which is what speculative decoding does. Every accepted draft token is a decode step you didn't pay for.

So: better drafting, or nothing.

What we tested

A block-diffusion draft model — 5 layers, 3.85 GB, tapping intermediate layers of the target — paired with an MLX inference implementation that does lossless verification, state rollback for the recurrent layers, and verify-specialized quantized matmuls.

Published numbers from the implementation's author on a newer chip: 2.37× at 1k context, 2.12× at 2k, 1.77× at 4k, with 88–90% draft acceptance.

Note the decay. The multiplier falls with context length while acceptance does not fall — meaning the cost of verification grows with context, not the quality of the drafts. And no published number existed anywhere near our 25–48k working range, which is exactly the regime we care about.

That gap is why we measured rather than adopting on the strength of the README.

The result

Same box, exclusive, same 100,000-byte prompt used in every other benchmark in this series:

About 2× on decode at our actual context length, and prefill is up too.

The extrapolation from the published numbers would have predicted decay to something well under 1.77× by 26k. We got roughly 2×. Measuring in your own regime is worth it in both directions — this time the surprise was favorable.

Acceptance ran 70–75% with an adaptive block size that settled between 4 and 5, reaching 82.5% at block size 5. Lower than the 88–90% published on shorter contexts, and still plenty: at 70% acceptance with a block of 5 you're getting several tokens per verification pass.

The honest caveat about that table

The two rows are not a clean comparison, and it would be easy to overstate the result.

The new configuration is 4-bit; the incumbent is 8-bit. Four-bit reads half the bytes per token. On a bandwidth-bound decode loop, that alone is a large speedup independent of any drafting.

We isolated it. A clean baseline of the same 4-bit model with no speculative decoding at all, on the same machine at correct process priority: 40.6 tok/s short-prompt decode.

So the honest decomposition is roughly: 28–34 → 40.6 from the quantization change, and 40.6 → 80–86 from the drafting. The drafting is worth about 2×; the quantization is worth the rest. Both are real, they're separate, and the quantization half comes with a quality cost that the drafting half does not — block-diffusion verification is lossless, so the output distribution is unchanged.

Publishing the combined number without that split would have been the easy thing to do and would have misattributed roughly half the win.

Related: an earlier baseline for this model measured 1.87 tok/s and is worthless — it was taken on a machine that was both thermally degraded (part 2) and running at the wrong process priority (part 3). Two independent measurement defects in one number. It stayed in our notes for weeks as a plausible-looking data point.

Packaging traps

Two, both of which cost real time.

The published package version does not know this model. The release on the package index carries a model registry that stops one generation earlier, so the server exits with "specify a draft model" and no useful hint. The project's main branch carries the same version string but adds the newer architecture to its registry.

Same version number, different capabilities. Install from source, and pass the draft model explicitly even when you think the registry should resolve it.

Two different projects install a command with the same name. The research group's reference package and the community MLX implementation both provide a CLI entry point called dflash. Installed in the same environment, one shadows the other and you get confusing errors from a program you didn't think you were running.

Separate virtual environments, one per implementation. Which is good practice generally and mandatory here.

What's next, and what's uncertain

No draft model exists for our production checkpoint. Our served model is a community-modified variant; the public draft targets the stock weights. Drafts are trained against a specific target, so acceptance will drop by some unknown amount when the target is a modified checkpoint. Measuring that drop is the next experiment, and the size of it determines whether this ships or whether we need a draft trained against our actual target.

The mature server we run in production has an open pull request adding this as a first-class backend. Their preliminary numbers on a smaller machine show 19.3 tok/s against 11.05 for plain autoregressive and 9.80 for the native multi-token head — a similar ratio to ours. If that merges, we get this without changing servers, and keep the session-bank prefix caching that made the incumbent win the engine comparison in the first place.

That's the actual strategic position: the drafting technique won, and the right way to adopt it is probably to wait for it to land in the server we already run, rather than switching servers to get it sooner and losing the cache.

Drop-in serving configs for tricks like this — with the knobs explained — are the kind of thing we put in the shop rather than re-derive per rig.

Key takeaways

1. A bandwidth-bound loop can only be sped up by running it fewer times. Speculative decoding is the lever; better kernels aren't.
2. Measure in your own context regime. The published decay curve predicted well under 1.77× at our length; we measured about 2×.
3. Decompose combined wins. Ours was quantization plus drafting, and reporting the total would have misattributed half of it.
4. Lossless and lossy speedups are not interchangeable. Verification-based drafting doesn't change the output distribution; quantization does.
5. A package index release and the project's main branch can share a version string and differ in capability. Ours did, and the failure message pointed nowhere.
6. Two packages can install the same CLI name. Separate environments per implementation.
7. A benchmark taken on a degraded machine survives in your notes for weeks. Ours had two independent measurement defects and looked entirely plausible.