DiffusionGemma

Google’s DiffusionGemma diffusion language model (26B total, 4B active), served from an FP8 quantized checkpoint.

Google’s DiffusionGemma diffusion language model (26B total, 4B active), served from an FP8 quantized checkpoint.

Setup

Install the Baseten CLI and sign in, then install the OpenAI SDK.

Install and sign in to Baseten

macOS or Linux

Terminal

brew tap basetenlabs/baseten
brew install baseten

Windows

Download and extract the binary, then move baseten.exe to a directory on your PATH:

Terminal

Invoke-WebRequest `
  https://github.com/basetenlabs/baseten-cli/releases/download/v0.4.0/baseten_0.4.0_windows_amd64.zip `
  -OutFile baseten.zip; Expand-Archive -Force baseten.zip .

For other platforms or a specific version, see the Baseten CLI install reference.

baseten auth login

Install the OpenAI SDK

uv pip install openai

Prefer not to install? Sign in with uvx truss login --browser and deploy with uvx truss push.

This preset serves DiffusionGemma 26B A4B on a single H100 with FP8 weights and dynamic activation quantization, optimized for low-latency diffusion decoding.

Hardware

H100

Engine

vLLM (nightly-2c9c07c8… build)

Context

8K

Concurrency

8

Write the config

Create and move into the project directory:

mkdir diffusiongemma-26B-A4B-it-latency && cd diffusiongemma-26B-A4B-it-latency

Then create a file named config.yaml and paste the following:

config.yaml

model_name: "model:diffusiongemma-26B-A4B-it preset:latency"
base_image:
  image: vllm/vllm-openai:nightly-2c9c07c85e56c799afffd5a671a8a0bace377a39
model_metadata:
  repo_id: RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic
  example_model_input:
    model: google/diffusiongemma-26B-A4B-it
    messages:
      - role: user
        content: Explain how diffusion language models differ from autoregressive ones.
    stream: true
    max_tokens: 512
  tags:
    - openai-compatible
weights:
  - source: "hf://RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic@main"
    mount_location: "/app/checkpoint/diffusiongemma"
    auth_secret_name: "hf_access_token"
build_commands:
  - apt-get update && apt-get install -y --no-install-recommends git ca-certificates && git clone --filter=blob:none https://github.com/vllm-project/vllm.git /opt/vllm-dgemma && cd /opt/vllm-dgemma && git checkout d25326b1fcfbcdfdc4133e7263b0d95ec31c9b87
  - cd /opt/vllm-dgemma && VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_LOCATION='https://wheels.vllm.ai/3d300aecb1e6639872b698bd74ed38fb81d9603e/vllm-0.22.1rc1.dev373%2Bg3d300aecb-cp38-abi3-manylinux_2_28_x86_64.whl' pip install --no-deps --force-reinstall .
docker_server:
  start_command: >-
    sh -c "vllm serve /app/checkpoint/diffusiongemma
    --served-model-name google/diffusiongemma-26B-A4B-it
    --tensor-parallel-size 1
    --attention-backend TRITON_ATTN
    --generation-config vllm
    --hf-overrides.diffusion_sampler entropy_bound
    --hf-overrides.diffusion_entropy_bound 0.1
    --diffusion-config.canvas_length 256
    --enable-chunked-prefill
    --enable-prefix-caching
    --max-model-len 8192
    --max-num-seqs 8
    --gpu-memory-utilization 0.85
    --trust-remote-code"
  readiness_endpoint: /health
  liveness_endpoint: /health
  predict_endpoint: /v1/chat/completions
  server_port: 8000
environment_variables:
  VLLM_USE_V2_MODEL_RUNNER: "1"
  VLLM_LOGGING_LEVEL: INFO
  PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
resources:
  accelerator: H100
  use_gpu: true
secrets:
  hf_access_token: null
runtime:
  health_checks:
    restart_check_delay_seconds: 300
    restart_threshold_seconds: 300
    stop_traffic_threshold_seconds: 120
  predict_concurrency: 8

This config serves the RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic checkpoint on a single H100 with vLLM built from the DiffusionGemma pull-request branch, because diffusion support has not yet landed in a vLLM release. Setting max-num-seqs to 8 and gpu-memory-utilization to 0.85 leaves the headroom that diffusion warmup needs for its logits buffers, and the deployment exposes an OpenAI-compatible chat completions endpoint under the served name google/diffusiongemma-26B-A4B-it.

Flags

The start_command passes these flags to the engine. Each one controls a runtime or serving behavior:

Flag	Value	What it does
--tensor-parallel-size	1	Number of GPUs to shard the model across.
--attention-backend	TRITON_ATTN	Attention kernel backend vLLM uses. TRITON_ATTN: Triton-based attention kernels, required by some model architectures not yet supported by the default backend.
--generation-config	vllm	Source of the default generation (sampling) settings. vllm: Use vLLM’s own defaults instead of the checkpoint’s generation_config.json.
--hf-overrides.diffusion_sampler	entropy_bound	Sampler the diffusion language model uses to unmask tokens during denoising.
--hf-overrides.diffusion_entropy_bound	0.1	Entropy threshold for the entropy-bound diffusion sampler.
--diffusion-config.canvas_length	256	Number of tokens in the diffusion canvas, the block the model denoises per step.
--enable-chunked-prefill	(no value)	Process long prompts in chunks so decode requests keep running.
--enable-prefix-caching	(no value)	Reuse KV cache across requests that share a prefix.
--max-model-len	8192	Maximum context length (tokens) the server accepts per request.
--max-num-seqs	8	Maximum number of concurrent sequences in the batch.
--gpu-memory-utilization	0.85	Fraction of GPU memory vLLM may use for weights and KV cache.
--trust-remote-code	(no value)	Execute model-specific Python from the checkpoint (required for many Qwen, Phi, and custom architectures).

Deploy

Push the config to Baseten with the Baseten CLI, or with the Truss CLI if you prefer it:

Baseten CLI

baseten model push

You should see output similar to:

Pushing model "diffusiongemma-26B-A4B-it-latency"...
Uploading model...
Uploaded model in 0s
✨ Model diffusiongemma-26B-A4B-it-latency was successfully pushed ✨

  Model:       diffusiongemma-26B-A4B-it-latency (abc1d2ef)
  Deployment:  xyz123
  Environment: production

🪵 View logs:
   deployment:   baseten model deployment logs --model-id abc1d2ef --deployment-id xyz123
   environment:  baseten model environment logs --model-id abc1d2ef --environment production  (once deployed)
   app:          https://app.baseten.co/models/abc1d2ef/logs/xyz123

🚀 Invoke your model:
   URL:  https://model-abc1d2ef.api.baseten.co/deployment/xyz123/predict
   CLI:  baseten model predict --model-id abc1d2ef

baseten model push prints your model ID (abc1d2ef in the example). The examples below use it wherever you see {model_id}, and read your API key from the BASETEN_API_KEY environment variable.

Call the model

Your deployment serves an OpenAI-compatible API.

Now call your deployment to run inference:

Python

main.py

import os
from openai import OpenAI

client = OpenAI(
    api_key=os.environ["BASETEN_API_KEY"],
    base_url="https://model-{model_id}.api.baseten.co/environments/production/sync/v1",
)

response = client.chat.completions.create(
    model="google/diffusiongemma-26B-A4B-it",
    messages=[
        {"role": "user", "content": "What is machine learning?"}
    ],
)

print(response.choices[0].message.content)

cURL

curl -s https://model-{model_id}.api.baseten.co/environments/production/sync/v1/chat/completions \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $BASETEN_API_KEY" \
  -d '{
    "model": "google/diffusiongemma-26B-A4B-it",
    "messages": [
      {"role": "user", "content": "What is machine learning?"}
    ]
  }'

Next steps

Call your model

Endpoint anatomy, authentication, and sync versus async inference

Autoscaling

Scale replicas with traffic, including scale to zero