Google/diffusiongemma-26B-A4B-it

Overview

DiffusionGemma 26B-A4B is a block-diffusion language model built on Gemma 4's MoE backbone. Instead of generating tokens one at a time (autoregressive), it generates blocks of 256 tokens simultaneously via iterative denoising over a fixed canvas — trading higher time-to-first-token for significantly higher per-request generation throughput.

Key Architecture Features

• Block Diffusion: Generates tokens in 256-token canvas blocks via iterative denoising (up to 48 steps per block).
• MoE Backbone: Same Gemma 4 architecture — 128 fine-grained experts with top-8 routing, 26B total / 4B active parameters.
• Entropy-Bound Sampling: Uses an entropy-based sampler (diffusion_sampler: entropy_bound, entropy_bound: 0.1) for the denoising process.
• Multimodal: Supports text + images via the Gemma 4 vision encoder.
• Thinking Mode: Supports structured reasoning via <|channel>thought\n...<channel|> delimiters.
• Function Calling: Supports Gemma 4's tool-call protocol (works best in thinking mode).

Important Deployment Notes

DiffusionGemma requires several specific flags that differ from standard Gemma 4:

Flag	Why
--max-num-seqs 4	The diffusion state buffers (self_conditioning_probs) pre-allocate max_seqs × canvas_length × vocab_size tensors. With Gemma's 262K vocab and canvas_length=256, higher values cause OOM.
--generation-config vllm	The checkpoint's generation_config.json sets max_tokens: 256, which would override per-request limits. This flag ignores it.
--gpu-memory-utilization 0.85	Leaves headroom for activation memory during denoising.
--hf-overrides '{"diffusion_sampler":"entropy_bound","diffusion_entropy_bound":0.1}'	Configures the entropy-bound denoising sampler.
--diffusion-config '{"canvas_length": 256}'	Sets the canvas block size for generation.

Prerequisites

DiffusionGemma requires a vLLM build with diffusion model support, available in the Gemma docker image:

docker pull vllm/vllm-openai:gemma

Deployment

Single GPU (H100/H200, BF16)

docker run -itd --name diffusiongemma \
    --ipc=host --network host --gpus all \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    vllm/vllm-openai:gemma \
        --model google/diffusiongemma-26B-A4B-it \
        --max-model-len 262144 \
        --max-num-seqs 4 \
        --gpu-memory-utilization 0.85 \
        --host 0.0.0.0 --port 8000

Intel Xeon 6 (CPU, BF16)

DiffusionGemma runs on Intel Xeon 6 CPUs using vLLM's CPU backend with AMX BF16 acceleration. The generated example uses TP=1. Adjust --tensor-parallel-size based on the detected NUMA topology and deployment requirements.

docker run \
  --privileged --ipc=host -p 8000:8000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  vllm/vllm-openai-cpu:latest-x86_64 google/diffusiongemma-26B-A4B-it \
  --max-num-seqs 4 \
  --tensor-parallel-size 1 \
  --enable-auto-tool-choice \
  --tool-call-parser gemma4 \
  --chat-template examples/tool_chat_template_gemma4.jinja \
  --reasoning-parser gemma4

Full-Featured Server (text + image + thinking + tools)

docker run -itd --name diffusiongemma \
    --ipc=host --network host --gpus all \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    vllm/vllm-openai:gemma \
        --model google/diffusiongemma-26B-A4B-it \
        --max-model-len 262144 \
        --max-num-seqs 4 \
        --gpu-memory-utilization 0.85 \
        --mm-processor-kwargs '{"max_soft_tokens": 1120}' \
        --limit-mm-per-prompt '{"image": 7}' \
        --enable-auto-tool-choice \
        --tool-call-parser gemma4 \
        --reasoning-parser gemma4 \
        --host 0.0.0.0 --port 8000

Jev-Style Structured Decisions (Nightly)

vLLM's structured-read support for DiffusionGemma can answer yes/no, multiple-choice, and scored questions with a confidence value. The example server seeds a canvas with the response template, leaves only the answer slots noisy, and reads a probability distribution from every slot after one denoising step. It exposes the Jev-compatible /v1/systemone endpoint; this is an example server rather than a standard vLLM endpoint.

Use the pinned nightly image that contains the merged structured-read support:

docker run -d --name dgemma \
    --gpus '"device=0"' \
    --ipc=host \
    -p 8000:8000 \
    -p 8011:8011 \
    -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
    vllm/vllm-openai:nightly-e9757321527ca1ecd514c07c1418dd2c53da3d19 \
    google/diffusiongemma-26B-A4B-it \
    --served-model-name dgemma \
    --diffusion-config '{"canvas_length":64}' \
    --max-logprobs 32 \
    --enable-prefix-caching

Wait for the vLLM server to become healthy, then start the example server in the same container:

until curl -fsS localhost:8000/health >/dev/null; do sleep 2; done

docker exec -d dgemma \
    python /vllm-workspace/examples/features/structured_diffusion/structured_server.py \
    --upstream http://127.0.0.1:8000 \
    --tokenizer google/diffusiongemma-26B-A4B-it \
    --canvas 64

After the example server becomes healthy, send a Jev-style decision request:

until curl -fsS localhost:8011/health >/dev/null; do sleep 2; done

curl -sS localhost:8011/v1/systemone \
    -H 'content-type: application/json' \
    -d '{
      "model": "jev-latest",
      "state": {"ticket": "Everything is down, demo at noon."},
      "questions": {
        "urgent": {
          "type": "noul",
          "instructions": "Needs a reply within the hour?"
        }
      }
    }'

Client Usage

Text Generation

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
response = client.chat.completions.create(
    model="google/diffusiongemma-26B-A4B-it",
    messages=[{"role": "user", "content": "Write a poem about the ocean."}],
    max_tokens=512,
    temperature=0.7,
)
print(response.choices[0].message.content)

Offline Inference

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

model_path = "google/diffusiongemma-26B-A4B-it"
tokenizer = AutoTokenizer.from_pretrained(model_path)

llm = LLM(
    model=model_path,
    max_num_seqs=4,
    gpu_memory_utilization=0.85,
    hf_overrides={
        "diffusion_sampler": "entropy_bound",
        "diffusion_entropy_bound": 0.1,
    },
    diffusion_config={"canvas_length": 256},
)

messages = [{"role": "user", "content": "Explain quantum computing in simple terms."}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
outputs = llm.generate(prompt, SamplingParams(temperature=0.0, max_tokens=1024))

print(outputs[0].outputs[0].text)

Thinking Mode

Launch the server with --reasoning-parser gemma4, then enable per-request:

response = client.chat.completions.create(
    model="google/diffusiongemma-26B-A4B-it",
    messages=[{"role": "user", "content": "What is the derivative of x^3 * ln(x)?"}],
    max_tokens=32768,
    extra_body={"chat_template_kwargs": {"enable_thinking": True}},
)

Performance Characteristics

Compared to the autoregressive Gemma 4 26B-A4B baseline (single H100, concurrency=1, SPEED-Bench):

Metric	Gemma 4 26B-A4B (baseline)	DiffusionGemma 26B-A4B
Output TPS	199 tok/s	375 tok/s (1.9×)
E2E Request Time (mean)	2.87s	0.88s (3.3× faster)
TTFT (mean)	53ms	489ms (higher due to canvas denoising setup)
Per-request Gen TPS (mean)	205 tok/s	1,282 tok/s (6.2×)

The diffusion model trades higher time-to-first-token for significantly higher generation throughput. It generates fewer total tokens per prompt on average (shorter outputs) but completes requests much faster.

Known Limitations

• TTFT: ~10× higher than autoregressive baseline — the model must denoise an entire canvas block before emitting the first token.
• --max-num-seqs: Must be kept low (≤4) due to the large diffusion state tensors. Higher values cause CUDA OOM.
• Audio: Not supported — the diffusion checkpoints do not include an audio encoder.
• Function Calling: Works best in thinking mode. Non-thinking mode may answer directly instead of emitting tool calls.

References

• Model card
• vLLM structured-read support for DiffusionGemma
• djev-run Cloud Run deployment
• Gemma 4 26B-A4B (autoregressive baseline)
• Gemma docs