Run DiffusionGemma on vLLM

Current vLLM Support Status

vLLM has supported DiffusionGemma since day 0. Google collaborated with the vLLM team before launch, making vLLM the only production-ready serving path for DiffusionGemma today. No custom forks, no PR branches, no workarounds — just pip install vllm and --trust-remote-code.

vLLM loads models directly from HuggingFace format (safetensors), bypassing GGUF entirely. This is the recommended path for production deployments — multi-GPU, continuous batching, and OpenAI-compatible API out of the box.

Install vLLM with DiffusionGemma Support

Prerequisites

# Python 3.10+ required
python --version

# NVIDIA GPU with CUDA 12.x
nvidia-smi

# Install vLLM
pip install vllm

Start Serving — Single GPU

vllm serve google/diffusiongemma-26B-A4B-it \
  --trust-remote-code \
  --gpu-memory-utilization 0.92 \
  --max-model-len 65536 \
  --max-num-seqs 8 \
  --host 0.0.0.0 \
  --port 8000

Multi-GPU — Tensor Parallelism (shipped June 26, 2026)

TP=2 lets two 16GB GPUs (2×5060 Ti, 2×4080) run the full 26B model. TP=4 is for data center setups. vLLM #46177 fixed the vocab-parallel embedding mismatch — update to vLLM nightly (≥0.8.0.dev) for TP support.

TP=2 — 2 GPUs

vllm serve google/diffusiongemma-26B-A4B-it \
  --trust-remote-code \
  --tensor-parallel-size 2 \
  --gpu-memory-utilization 0.90 \
  --max-model-len 65536 \
  --max-num-seqs 8

TP=4 — 4 GPUs

vllm serve google/diffusiongemma-26B-A4B-it \
  --trust-remote-code \
  --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.88 \
  --max-model-len 65536 \
  --max-num-seqs 16

NVFP4 on Blackwell (RTX 5090)

vllm serve nvidia/diffusiongemma-26B-A4B-IT-NVFP4 \
  --trust-remote-code \
  --tensor-parallel-size 2 \
  --gpu-memory-utilization 0.85 \
  --max-model-len 65536

NVFP4 is Google's custom 4-bit floating point format. Best quality-to-memory ratio. Requires NVIDIA Blackwell (RTX 5090, B100, B200). Emulated on Ampere/Ada with ~10% overhead.

TP Troubleshooting

"matmul reduction-dim mismatch" — you're on a vLLM version before #46177 merged. Upgrade: pip install --upgrade vllm --pre (nightly).

OOM with TP=2 on 2×16GB — lower --gpu-memory-utilization to 0.80 and reduce --max-num-seqs to 4. The 26B model at Q4 needs ~16GB total across GPUs, but KV cache adds per-request overhead.

Pipeline parallel (PP) is not supported yet — only tensor parallel works for DiffusionGemma. The sampler output type doesn't match PP broadcast expectations.

Call the API

OpenAI-compatible endpoint

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "google/diffusiongemma-26B-A4B-it",
    "messages": [{"role": "user", "content": "Explain diffusion language models"}],
    "max_tokens": 256,
    "temperature": 0.7
  }'

Python client

from openai import OpenAI

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

response = client.chat.completions.create(
    model="google/diffusiongemma-26B-A4B-it",
    messages=[{"role": "user", "content": "Explain diffusion language models"}],
    max_tokens=256,
    temperature=0.7,
)
print(response.choices[0].message.content)

Related Guides