Hero model · day-1 MLX support

DiffusionGemma 26B-A4B-it

Google's discrete-diffusion text generation model — parallel block decoding instead of one token at a time — runs day-1 on Rapid-MLX, while upstream mlx-lm still reports it as an unsupported model type.

TL;DR. rapid-mlx serve diffusion-gemma-26b-4bit and you get an OpenAI-compatible endpoint backed by a non-autoregressive diffusion decoder. The model emits whole blocks of tokens in parallel per refinement step instead of looping one token at a time. Quants live at mlx-community/diffusiongemma-26B-A4B-it-{4bit,8bit}.

What "discrete diffusion text gen" means

A standard autoregressive LLM emits text one token at a time — sample a token, append it to the context, run the model again. A discrete-diffusion model instead starts from a fully masked block, runs a fixed number of "denoising" passes, and at each pass replaces the most-confident masked positions with real tokens. Because every position in the block can be updated in the same forward pass, the decoder is naturally parallel.

The practical effect on a Mac: each forward pass produces many tokens of output rather than one, which on Apple Silicon's bandwidth-bound regime amortises memory traffic across the whole block.

Why Rapid-MLX ships it

Upstream mlx-lm currently rejects this checkpoint with "Model type diffusion_gemma not supported" (tracked as mlx-lm#1391). Rapid-MLX vendors the diffusion sampler path and the block-aware KV / state handling needed to make the OpenAI /v1/chat/completions contract behave the same as a standard AR model from the client's point of view — the streaming delta chunks just arrive in larger groups instead of one token each.

Honest caveat. We have not benchmarked this against every other MLX-flavoured serving stack head-to-head; calling Rapid-MLX "the fastest local serve" would be a claim we cannot substantiate right now. What we will commit to: day-1 support, OpenAI-shape responses, and an alias that just works. If you have a reproducible comparison against another local runtime, please send it our way for the community performance page.

Quants we publish

aliashf repocache footprintrecommended for
diffusion-gemma-26b-4bit mlx-community/diffusiongemma-26B-A4B-it-4bit ~15 GB 32 GB Macs — the recommended quant
diffusion-gemma-26b-8bit mlx-community/diffusiongemma-26B-A4B-it-8bit ~28 GB 64 GB Macs — highest quality

Run it

$ rapid-mlx serve diffusion-gemma-26b-4bit

Streaming chat completion

The deltas come in larger batches than an AR model — every iteration of the diffusion loop reveals a block of tokens — but the wire format is identical SSE chunks. 

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="not-used")

stream = client.chat.completions.create(
    model="diffusion-gemma-26b-4bit",
    messages=[
        {"role": "user", "content": "Write a haiku about Apple Silicon."}
    ],
    stream=True,
)
for chunk in stream:
    delta = chunk.choices[0].delta.content
    if delta:
        print(delta, end="", flush=True)

Engineering notes

References