For years, diffusion was how AI generated images, not language. When open-weights text diffusion models like DiffusionGemma arrived, watching a paragraph materialize out of blank mask tokens felt like watching a digital Polaroid develop.
Yet almost immediately, developers hit a practical puzzle: what are you supposed to use a text diffusion model for in everyday production?
Built on Google’s Gemma 4 architecture, DiffusionGemma is an open-weights Mixture-of-Experts model designed for text diffusion. In open-ended generation, it fills a canvas of placeholder tokens by refining them in parallel across dozens of denoising passes, whereas traditional models decode text strictly one token at a time. But treating it like a standard chatbot misses what makes text diffusion powerful: the ability to evaluate an entire output structure in a single forward pass.
Return of the classifier
Before large language models took over, production machine learning was built around classification: fast, dedicated models that returned clean probability distributions over fixed classes. Replacing them with generative LLMs made simple routing slow and chatty. In production pipelines, most steps do not need an essay. They just need a fast, calibrated decision.
That tension led TypeSafe AI to launch Jev, founded by Diogo Almeida (one of the co-inventors of RLHF and InstructGPT). Almeida called this category System 1 models, referencing Daniel Kahneman’s Thinking, Fast and Slow. While System 2 models spend seconds generating step-by-step reasoning tokens, a System 1 model focuses on fast, calibrated judgment across probabilities, categorical choices, and rubric scores.
https://medium.com/media/da95bed8e93eadb3bec06bbdb62eb48b/href
The open-source community immediately scrambled to recreate that speed on open weights. Projects like Kev on Qwen 3.5 showed how much effort that took on causal architectures: developers had to build custom block-causal attention masks to keep questions from leaking into one another, bolt on a dedicated <decide> pointer readout head, and run LoRA fine-tuning jobs on cloud GPUs.
Then Matt Mastracci noticed something hiding in plain sight, which @googlegemma shared shortly after: you don’t need to train custom pointer heads or rewrite attention masks if your model is already a discrete diffusion model. Single-pass structured judgment is the exact job text diffusion is wired to do.
How 1-step canvas diffusion works
As Google’s DiffusionGemma developer guide explains, the model splits inference into two halves: a causal Encoder that ingests the prompt context and writes to the KV-cache, and a bidirectional Denoiser that iteratively updates the output canvas.
Traditional autoregressive language models are constrained by GPU memory bandwidth. To generate a single token, the GPU must stream gigabytes of model weights from memory into registers, leaving tensor cores under-utilized. DiffusionGemma shifts that workload from memory bandwidth to parallel compute by evaluating an entire canvas simultaneously.
In an autoregressive model, the causal attention mask forces tokens to look only backward. Generating a JSON payload like {"urgent": "yes", "team": "compute", "severity": "critical"} requires predicting every brace, quotation mark, and field name sequentially. Each field's probability is conditioned on whatever token happened to win the greedy sample before it.
DiffusionGemma’s denoiser works differently. Because it uses non-causal, bidirectional attention across the output canvas, every token on the canvas can attend to the entire prompt context in the KV-cache, the surrounding JSON keys, and every other answer slot at the exact same moment.
Mastracci’s vLLM pull request turns that mechanism into a structured inference engine by configuring three parameters:
- Set diffusion_canvas_length to 16 or 32 tokens instead of 256, just wide enough to fit your target JSON object.
- Seed the canvas with pre-tokenized JSON syntax, placing mask tokens only at the answer slots and end-of-sequence tokens as padding.
- Set diffusion_max_steps: 1 and diffusion_read_only: true so vLLM stops right after the first denoising pass.
Instead of spending dozens of steps discovering syntax, the model runs a single forward pass over the prompt and the 16 seeded tokens. It returns temperature-1 log-probabilities directly from the unmasked vocabulary distributions at each slot index.
Because the pull request includes a lightweight proxy (structured_server.py), it normalizes probabilities across your allowed choices, merges whitespace variants (like "yes" and " yes"), and calculates the Shannon entropy for each field to measure how divided the predictions are.
When you configure samples="auto", this entropy calculation acts as an automated uncertainty detector. If the model is confident and entropy stays below the threshold, it returns the decision in a single pass. If the scores on any field are divided, the server automatically fires three parallel follow-up reads, averages the distributions, and returns confidence intervals with standard error.
Deploying on Cloud Run GPUs
To test this on production infrastructure, I deployed an INT4 AWQ checkpoint to Google Cloud Run with an NVIDIA L4 GPU. Quantizing DiffusionGemma to 4-bit reduces the weights from 52 GB down to 16.1 GB, fitting the entire model and KV-cache inside the L4’s 24 GB of VRAM.
Running a 16 GB model on serverless container infrastructure introduces an operational hurdle: Cloud Run’s default startup TCP probe expects the container to bind its listening port within 240 seconds. Streaming 16 GB of model shards from Hugging Face and warming up Triton attention kernels takes roughly three minutes on cold starts. If unauthenticated requests from shared Cloud Run egress IPs hit Hugging Face rate limits, startup times can exceed the platform limit.
To solve this, entrypoint.sh starts the lightweight Python proxy on port 8080 in the background before launching vLLM. Because structured_server.py only loads the tokenizer on boot, it binds port 8080 and passes Cloud Run's startup probe in under three seconds. A /ready endpoint on the proxy polls vLLM's internal /health route and returns HTTP 503 until the GPU finishes loading the model shards:
# 1. Bind port 8080 immediately so Cloud Run's startup probe passes
MODEL_ID="cyankiwi/diffusiongemma-26B-A4B-it-AWQ-INT4"
python3 /app/structured_server.py \
--host 0.0.0.0 --port "${PORT:-8080}" \
--upstream "http://127.0.0.1:8000" \
--model "${MODEL_ID}" --tokenizer "${MODEL_ID}" \
--canvas 64 &
With port 8080 open, entrypoint.sh starts vllm serve on localhost using Triton attention, prefix caching, and a 64-token diffusion canvas:
# 2. Load the 4-bit AWQ weights into the L4 GPU on localhost:8000
exec vllm serve "${MODEL_ID}" \
--host 127.0.0.1 --port 8000 \
--attention-backend TRITON_ATTN \
--gpu-memory-utilization 0.92 \
--max-model-len 2048 --max-num-seqs 16 --max-logprobs 32 \
--enforce-eager --enable-prefix-caching \
--diffusion-config '{"canvas_length": 64}'
Finally, deploy the container image to Cloud Run in us-central1 with an NVIDIA L4 GPU and HF_TOKEN configured to avoid Hugging Face download throttling:
gcloud run deploy diffusiongemma-triage \
--image="${IMAGE_URI}" --region="us-central1" \
--gpu=1 --gpu-type=nvidia-l4 --no-gpu-zonal-redundancy \
--cpu=8 --memory=32Gi --no-cpu-throttling \
--max-instances=1 --concurrency=16 --timeout=600 --port=8080 \
--set-env-vars="HF_TOKEN=${HF_TOKEN}" \
--no-allow-unauthenticated
Benchmarking production incident triage
With the service running on Cloud Run, I tested it against eight realistic Google Cloud operations alerts: four straightforward single-system outages and four ambiguous cross-team incidents, such as an HTTP 502 spike following both a Cloud Armor WAF update and a container release.
Each request evaluates three typed questions simultaneously inside a 16-token JSON canvas ({"urgent": "[MASK]", "team": "[MASK]", "severity": "[MASK]"}):
- Urgent: Should we wake an on-call engineer (binary probability between 0 and 1)?
- Team: Which team owns the incident (compute, networking, security, storage, or billing)?
- Severity: How severe is customer impact on a 1-to-4 rubric (from none to critical)?
On the four clear-cut outages, the model was confident on its first pass and returned in about 180 ms, hitting Jev’s sub-200ms target on a single serverless L4 GPU.
On the four cross-team incidents, the probabilities softened to reflect the ambiguity instead of forcing a false 100% match. In the HTTP 502 spike, where both a WAF change and a container rollout had just landed, the model leaned toward networking (0.87) over compute and hedged on whether to page on-call (0.67). That uncertainty automatically triggered the extra verification passes, which still finished in under 500 ms because vLLM's prefix cache kept the prompt in GPU memory.
Production takeaways
One schema tip if you try this yourself: stick to single-word labels (yes, compute, critical) so every option fits into a single SentencePiece token, or map longer internal IDs to simple aliases in your prompt.
Not every step in an AI pipeline needs a conversation. Sometimes you just need a fast decision that understands natural language and tells you when an edge case is uncertain. Best of all, you don’t need to fine-tune a custom model to get it. The capability is already built into open-weights diffusion models today.
Give it a try on Cloud Run, and let me know where you end up using it on X, LinkedIn, or Bluesky.
How to build a Jev-style classifier with DiffusionGemma and vLLM was originally published in Google Cloud – Community on Medium, where people are continuing the conversation by highlighting and responding to this story.
Source Credit: https://medium.com/google-cloud/how-to-build-a-jev-style-classifier-with-diffusiongemma-and-vllm-ef2e0bfa9ad7?source=rss—-e52cf94d98af—4
About the Author
