Block Diffusion¶

Autoregressive models make one irreversible left-to-right choice at a time. Diffusion language models instead maintain a region of uncertain tokens and repeatedly revise or reveal positions. Block diffusion bounds that region: generate one block, commit it to the prefix, then generate the next block. This combines parallel work inside a block with a causal order between blocks.

mlxcel implements two related but distinct families:

• DiffusionGemma (diffusion_gemma, diffusion_gemma_text): initializes a token canvas and denoises it with entropy- or confidence-based acceptance. The multimodal form supports images.
• LLaDA-2 MoE (llada2_moe): initializes generation positions with <|mask|> and iteratively unmasks confident positions in fixed-size text blocks.

These are native target decoders, not DFlash drafters. DFlash uses diffusion only to propose tokens that an autoregressive target verifies; see MTP and DFlash.

From Autoregression to Denoising¶

An autoregressive target estimates p(x_t | x_<t) and cannot revise x_t after moving to t+1. A discrete diffusion or masked model instead considers a corrupted state z_s and predicts cleaner token values as the noise/mask level decreases:

z_S -> z_(S-1) -> ... -> z_0 = x

Several positions can change in one model evaluation because they exist simultaneously on the canvas. The possible advantage is not “one forward pass generates the whole answer”; it is that a useful amount of parallel token work can be done per denoising step. The actual cost depends on block length, number of model evaluations, how quickly confidence rises, and whether the attention shape uses the hardware efficiently.

Block generation also solves two practical problems of a full-answer canvas:

• memory and attention cost are bounded by the configured block/canvas size;
• completed blocks become a stable prefix, so long answers need not be regenerated globally.

The trade-off is that a committed block cannot normally be revised by a later block. Small blocks are more causal and stream sooner; large blocks offer more within-block parallelism and context but delay commitment.

Two mlxcel Generation Loops¶

DiffusionGemma Canvas Denoising¶

1. Encode the prompt into a read-only prefix held in dense FP16 KV caches.
2. Allocate a canvas with length min(max_canvas, max(remaining, min_canvas)) and initialize it from the seeded MLX random generator.
3. Embed the current canvas with learned self-conditioning. Canvas positions attend bidirectionally within the canvas and causally to the committed prefix.
4. Use the selected sampler to accept new predictions.
5. Stop when the canvas is stable and confident for the configured stability window or when the denoising-step budget is exhausted.
6. Commit the canvas to the prefix and start another canvas if the output budget remains.

LLaDA-2 Block Unmasking¶

1. Split the prompt and generation region into blocks. The future generation region begins as <|mask|> tokens.
2. Commit prompt blocks to each layer's KV cache.
3. Repeatedly forward the active generation block against the frozen prefix.
4. Reveal every still-masked position whose chosen-token confidence exceeds the current threshold.
5. If no position qualifies, reveal the single highest-confidence active position so every iteration makes progress.
6. Commit the fully revealed block, stream its decoded tokens in sequence order, and continue with the next block.

LLaDA-2's default generation block length is 32, with 32 steps and an initial/floor confidence threshold of 0.95. The engine supports a linearly decaying threshold internally, but the public CLI currently exposes only the step override; block length and threshold schedule use model/runtime defaults.

Samplers and Reveal Policies¶

DiffusionGemma exposes two per-step samplers through --diffusion-sampler:

• entropy-bound (default): calculate per-position entropy, order positions from low to high entropy, and accept a prefix until the cumulative entropy rule exceeds the checkpoint's entropy_bound. At least one position is accepted. This follows the mlx-vlm reference rule.
• confidence-threshold: accept unrevealed positions whose top-1 probability exceeds --diffusion-threshold (default 0.9). If none qualify, accept the highest-confidence unrevealed position to guarantee progress.

The confidence policy can reveal many easy positions early, reducing steps at a possible quality cost. The entropy policy adapts work to uncertainty using the checkpoint's calibrated bound. Do not compare them only by “tokens accepted per step”; evaluate final task quality and end-to-end latency.

LLaDA-2 uses its own confidence-based transfer mask. --diffusion-sampler and --diffusion-threshold do not change it. Ordinary --temp, --top-k, and --top-p settings affect its chosen-token sampling when temperature is above zero.

CLI Reference¶

DiffusionGemma¶

mlxcel generate \
  --model diffusiongemma-26B-A4B-it-4bit \
  --prompt "Explain TCP and UDP in one paragraph." \
  --max-tokens 256 \
  --max-denoising-steps 48 \
  --diffusion-sampler entropy-bound \
  --seed 42 \
  --profile

LLaDA-2 MoE¶

--max-denoising-steps doubles as the per-block LLaDA-2 step override. The other --diffusion-* sampler, threshold, and canvas controls are DiffusionGemma controls and do not configure LLaDA-2.

mlxcel generate \
  --model ./models/LLaDA2.0-mini \
  --prompt "Explain why retrieval and reranking are separate stages." \
  --max-tokens 256 \
  --max-denoising-steps 32 \
  --temp 0 \
  --seed 42 \
  --profile

LLaDA-2 rejects --draft-model: its native decoder already unmasks a block and the runtime does not nest an additional speculative protocol around it.

Multimodal Input¶

DiffusionGemma's multimodal checkpoint accepts repeated --image <path> arguments:

mlxcel generate \
  --model google/diffusiongemma-26B-A4B-it \
  --image photo.jpg \
  --prompt "Describe what you see." \
  --max-tokens 256 \
  --seed 42

It reuses the Gemma 4 vision pipeline: each image is resized and padded to 768×768, encoded as 256 soft tokens, and inserted between image sentinel tokens. Image blocks receive bidirectional attention during prefill. Audio and video are rejected.

LLaDA-2 is text-only and rejects image, audio, and video input.

Server Mode¶

Both families use a dedicated single-stream worker. They do not join the continuous-batching scheduler, so requests are queued and processed one at a time. Streaming is still supported, but internal positions can converge out of order; mlxcel emits text when a block is committed in normal sequence order rather than exposing unstable intermediate tokens.

mlxcel serve \
  --model diffusiongemma-26B-A4B-it-4bit \
  --diffusion-sampler entropy-bound \
  --max-denoising-steps 48 \
  --port 8080

The server exposes chat and completion routes. For DiffusionGemma image requests, use the normal OpenAI image_url content part. Unsupported media returns an error for that request without terminating the worker.

DiffusionGemma Response Semantics¶

DiffusionGemma is a thinking-channel model. Give it enough max_tokens: a 64-token budget can be consumed by reasoning and leave visible content empty, while 256 is a more practical starting point for short prompts. Streaming chat deltas can carry reasoning_content; non-streaming chat responses follow mlxcel's convention of omitting reasoning. Raw /v1/completions works mechanically, but chat templating is strongly preferred.

Cancellation is checked once per denoising step, so cancellation latency is bounded by the current step rather than being instantaneous.

Reproducibility¶

Pass --seed N to initialize the MLX random generator. With the same checkpoint, prompt, generation options, and temperature 0, repeated DiffusionGemma runs use the same canvas initialization and commit stream; LLaDA sampling uses the same seeded source. Kernel/version changes can still alter floating-point ties.

MLXCEL_DIFFUSION_DEBUG_CANVAS=1 replaces DiffusionGemma random initialization with a fixed debug pattern for Python-reference parity work. It is not a normal-quality generation mode.

Performance Reasoning¶

For a block of B useful tokens and S denoising forwards, a first approximation is

useful rate ~= B / (S * t_forward(B) + t_commit)

This is deliberately different from autoregressive tokens/s. A larger block increases useful work but also attention and each forward's cost. A lower step budget reduces time but can commit less-converged tokens. Early stopping or revealing multiple positions per iteration changes effective S.

Measured on an M1 Ultra with a 4-bit DiffusionGemma 26B-A4B checkpoint:

• about 20 tok/s for short answers using a 64-token canvas that converged in roughly 12 of 48 steps;
• about 8 tok/s for longer outputs that used most of 48 steps on a 256-token canvas;
• approximately 13.5 GiB resident model memory.

These are hardware- and prompt-specific observations, not universal model claims. Re-run the benchmark with the same model revision and the workload you will serve.

Quality and Tuning Checklist¶

• Start from the checkpoint sampler and step defaults.
• Evaluate answer quality, formatting, EOS behavior, and repetition as well as speed.
• Compare short and long outputs; the final partial block can behave differently.
• Record actual steps and early stops, not only the maximum setting.
• With stochastic sampling, compare multiple seeds.
• For a server, include queueing in latency because requests are serial.
• Do not compare a diffusion model and an autoregressive model only by raw tokens/s; include model size, task quality, TTFT, and output-length semantics.

Current Limitations¶

Further Reading¶

• Google AI for Developers, DiffusionGemma model card
• Google AI for Developers, How DiffusionGemma works
• Nie et al., Large Language Diffusion Models