Skip to main content

​
Setup

Sign in to Baseten with Truss, then install the OpenAI SDK.
Sign in to Baseten
uvx truss login --browser
Install the OpenAI SDK
uv pip install openai
RedHatAI/diffusiongemma-26B-A4B-it-FP8-dynamic is a 26B-parameter diffusion transformer model (4B active per token) with up to 8K context. 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:
FlagValueWhat it does
--tensor-parallel-size1Number of GPUs to shard the model across.
--attention-backendTRITON_ATTNAttention kernel backend vLLM uses. TRITON_ATTN: Triton-based attention kernels, required by some model architectures not yet supported by the default backend.
--generation-configvllmSource of the default generation (sampling) settings. vllm: Use vLLM’s own defaults instead of the checkpoint’s generation_config.json.
--hf-overrides.diffusion_samplerentropy_boundSampler the diffusion language model uses to unmask tokens during denoising.
--hf-overrides.diffusion_entropy_bound0.1Entropy threshold for the entropy-bound diffusion sampler.
--diffusion-config.canvas_length256Number 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-len8192Maximum context length (tokens) the server accepts per request.
--max-num-seqs8Maximum number of concurrent sequences in the batch.
--gpu-memory-utilization0.85Fraction 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
Your model ID is printed in the truss push output (abcd1234 in the example). Use it wherever you see {model_id} in the next section.

​
Call the model

Your deployment serves an OpenAI-compatible API. Replace {model_id} with your model ID and make sure BASETEN_API_KEY is set. Now call your deployment to run inference:
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)