dflash-mlx
An implementation of DFlash block-diffusion speculative decoding for the MLX framework on Apple Silicon, generating multiple tokens per pass while guaranteeing output identical to standard autoregressive decoding.
dflash-mlx implements DFlash, a speculative-decoding technique, for MLX on Apple Silicon. A small draft model generates sixteen tokens at once using block diffusion, and the target model verifies all sixteen in a single forward pass; only tokens that match the target model's own greedy output are kept, so the result is lossless even though generation is faster.
It is built on stock MLX with a small number of custom Metal kernels for cases that need tighter numerical control, including a tape-replay approach to rolling back rejected tokens without snapshotting the full model state, routing for long-context attention on Qwen and Gemma model families, and quantized matrix multiplication kernels tuned for the verification step. A prefix cache keeps target and recurrent state in memory, with optional overflow to disk, so revisited prompts skip recomputation.
It is written in Python under the Apache-2.0 license.