dflash-mlx: block diffusion speculative decoding for Apple Silicon

Speculative decoding here is a draft, a verify, and a rollback

The premise of DFlash is described in two sentences at the top of the documentation. Block diffusion draft generates 16 tokens in one pass, and the target model verifies in one pass. The output is lossless, because every emitted token is verified against the target model before it is committed.

The mechanics are the standard speculative decoding contract with an unusual draft strategy. A small draft model of roughly one billion parameters generates sixteen tokens in parallel using block diffusion. The target model then verifies those sixteen tokens in a single forward pass. Greedy acceptance keeps the correct prefix and rejects the rest.

The losslessness claim is stated with a careful caveat, and that caveat matters. Every emitted token is the target model's greedy argmax at verification time, so no unverified token is ever emitted. But the output can still differ from pure autoregressive decoding because of MLX dispatch divergence. In other words, the guarantee is about the tokens, not about bitwise identity with a non speculative run.

The implementation is built on stock MLX, with a small number of targeted Metal kernels added where rollback and long context verification need tighter numerical control. The paper it implements is DFlash: Block Diffusion for Flash Speculative Decoding, by Chen et al. from 2026.

Why block diffusion beats one token at a time

The choice of draft mechanism is what separates this from ordinary speculative decoding.

Ordinary speculative decoding drafts tokens sequentially, so the draft model itself becomes the bottleneck: if the draft runs at a fraction of the target's speed, the whole scheme slows down. Block diffusion removes that dependency. The draft model produces sixteen tokens in parallel in a single pass, so the draft cost is amortised across the whole block rather than paid sixteen times.

That is the whole reason the acceptance rates in the benchmark table are as high as they are. In the table, acceptance sits between 83 and 91 percent across every model and context length measured. Those are high figures for speculative decoding, and they are what converts the arithmetic speedup into the throughput numbers.

The comparison target is `mlx_lm.stream_generate`, the stock MLX language model streaming path, run on the same machine with the same model. The protocol is stated in the documentation: sequential, three repeats, median, with a 60 second cooldown between runs, on an Apple M5 Max with 64 GB of unified memory and MLX 0.31.1.

What the table does not show is a cold start, a batched serving scenario, or a machine other than that one. The gains are single stream, and the baseline is the stock path rather than a heavily tuned serving stack.

The numbers fall off as contexts grow, and the table is honest about it

The benchmark table covers four models across context lengths from 1024 up to 16384 tokens, and the pattern is consistent: the speedup is highest at short contexts and decays as the sequence grows.

For Qwen3.5-9B the speedup runs 4.37x at 1024 tokens, then 3.65x at 2048, 3.06x at 4096 and 2.22x at 8192. For Qwen3.6-27B-4bit it goes from 2.95x at 1024 down to 2.78x at 16384. The smallest figure in the table is 1.22x, for Qwen3.5-35B-A3B-4bit at 8192 tokens.

There is a clear pattern in which models benefit least. The 35B-A3B mixture of experts model is the weakest case throughout, sitting between 1.22x and 2.20x, while the 9B dense model is the strongest at up to 4.37x. That makes sense once you look at the absolute throughput: the MoE model already runs at 133 to 143 tokens per second in the baseline, so there is comparatively little headroom for a draft model to recover, while the 9B model at 29 to 31 tokens per second has a lot of room.

Acceptance is the number to watch alongside the speedup, and it is remarkably stable, staying between 83 and 91 percent for every configuration. The one place it dips most is Qwen3.6-35B-A3B-4bit at 16384 tokens, where it is 84.40 percent.

Per run JSON output is kept in `benchmark/results/`, and the documentation says the results can be reproduced on your own hardware with `dflash benchmark`.

The engineering detail is rollback, and it is done with a tape

The technical details section is where the interesting engineering is, and the first item is tape replay rollback.

The problem is that verifying sixteen tokens in one forward pass means the target model's state has already advanced past what might be accepted. Reverting that state is the expensive part. Instead of snapshotting and restoring the full GatedDeltaNet state, dflash-mlx records an innovation tape during verify and replays only the accepted steps through a custom Metal kernel. The stated benefit is that it keeps rollback cost low and preserves acceptance over long generations, which connects directly to the acceptance rate staying high in the table at 8192 and 16384 tokens.

Two more pieces of the same kind. Target owned long context attention routing means the Qwen and Gemma adapters route verify blocks through the appropriate MLX or GQA-reshape SDPA path internally, which keeps the public command line free of attention kernel switches. And a verify specialised int4 qmm provides custom Metal kernels for the small-M quantized matmuls that dominate the target verify step.

The int4 path is hardware aware in a specific way. M5-class Apple GPUs, identified as `applegpu_g17*`, use the Metal 4 NAX M=16 path automatically when the shape is supported, while older Apple GPUs keep the steel simdgroup-MMA fallbacks. The feature is auto-enabled on MoE targets and on dense models with 40 or more layers.

Numerical coherence is called out separately, with bf16-sensitive paths including recurrent state replay and small projections stabilised across speculative cycles so that accepted tokens stay consistent. That is the property that would silently break losslessness if it were not handled, since a small numerical drift in the replayed state would change which tokens verify.

Two caching paths that cannot be used together

Two features sit on top of the core mechanism, and the documentation is explicit that they are mutually exclusive.

The prefix cache is described as an L1 plus L2 hierarchy. It stores RAM snapshots of the target KV cache, the GDN recurrent state, captured hidden states and the last logits, with optional SSD spill, byte and entry budgets, and automatic eviction. A hit skips prefill on revisited prompts. The design is inspired by the tiered KV cache work in oMLX, though dflash-mlx stores DFlash prefix snapshots rather than active paged KV blocks.

Positional sparse prefill, exposed as `prompt_token_positions`, is the other. `stream_dflash_generate` accepts an optional list of original positions alongside a compacted `prompt_tokens_override`, so a caller can prefill the target on a selected token subset placed at its true positions. The mechanics differ by layer type: full attention layers apply RoPE to each token at its original position, with contiguous runs delegated to the native rope and staying bitwise-exact; GDN layers run over the compacted sequence; and Gemma 4 sliding window layers rebuild their window mask from the true positions.

The two are mutually exclusive, so a caller has to choose between reusing a cached prefix and prefill on a sparse subset. The documentation also warns about how to use the sparse path: use chunk based selection, because incoherent selection such as every-other-token shreds local structure and collapses draft acceptance.

That warning is worth heeding. The feature exists so a caller such as oMLX's specprefill can prefill only part of a long prompt, and the difference between contiguous chunks and a scattered selection is the difference between a working acceleration and one that makes things slower.

Installing and running it, and what the CLI exposes

The install is a pip install, and the requirements are narrow. The package is `dflash-mlx` on PyPI, version 0.1.10, Apache-2.0 licensed, requiring Python 3.10 or newer.

pip install dflash-mlx

There is an optional extra for benchmark dataset support:

pip install "dflash-mlx[bench]"

The runtime dependencies are just `mlx>=0.32.1` and `mlx-lm>=0.31.3`, and the extra adds `datasets`. The console script is `dflash`, mapped to `dflash_mlx.cli:main`.

The quick start shows the two modes. One-shot generation, where the draft is resolved automatically:

dflash generate --model Qwen/Qwen3.5-9B --prompt "$PROMPT"

And a server with an OpenAI compatible interface, where you name the target and the draft explicitly:

dflash serve \
  --model mlx-community/Qwen3.6-27B-4bit \
  --draft z-lab/Qwen3.6-27B-DFlash \
  --port 8000

Note that the draft is a separate checkpoint in the serve path, with `z-lab` publishing DFlash versions of the same models. The badge on the README calls the implementation stock MLX with targeted Metal kernels, and the `examples/` directory contains a `demo.py`.

The repository is not archived and the last push was on 2026-08-20, with releases v0.1.8, v0.1.9 and v0.1.10 all landing between May and June 2026.

Editorial conclusion

Adopt dflash-mlx if you serve open models on Apple Silicon and want higher generation throughput without changing the output distribution, since the verification step guarantees every emitted token is the target model's own argmax. Do not adopt it on hardware other than Apple Silicon, since the implementation is built on MLX and custom Metal kernels, and do not expect the speedup to hold at long contexts, where the published table shows it falling from 4.37x to 2.22x. Verify first that your target model has a matching DFlash draft checkpoint, and read the acceptance figures rather than the speedup column, because acceptance is what determines whether the method pays off on your prompts.

Frequently asked questions

Official sources

1. bstnxbt/dflash-mlx on GitHub
2. Issues
3. License: Apache-2.0
4. README
5. Releases