DiffusionGemma

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

Setup

Sign in to Baseten

uvx truss login --browser

Install the OpenAI SDK

uv pip install openai

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

Concurrency

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:

uvx truss push

You should see output similar to:

✨ Model diffusiongemma-26B-A4B-it-latency was successfully pushed ✨

   Model ID:      abc1d2ef
   Deployment ID: xyz123
   Endpoint:      model-abc1d2ef.api.baseten.co
   Logs:          https://app.baseten.co/models/abc1d2ef/logs/xyz123

truss 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
cURL

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 -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?"}
    ]
  }'

Examples

Models

Engines

Custom Docker servers

Custom Python models

Chains

Setup

Hardware

Engine

Context

Concurrency

Write the config

Flags

Deploy

Call the model

Next steps

Call your model

Autoscaling

​Setup

Hardware

Engine

Context

Concurrency

​Write the config

​Flags

​Deploy

​Call the model

​Next steps

Call your model

Autoscaling

Setup

Write the config

Flags

Deploy

Call the model

Next steps