DEV Community

oooocean66
oooocean66

Posted on

Hands-On with Diffusion Gemma: Another Gemma with a Diffusion Mechanism

What Is Diffusion Gemma?

Diffusion Gemma is an LLM built on top of Gemma-4-26B-A4B. Gemma-4-26B-A4B is what's known as a "causal model," but Diffusion Gemma is called a "non-causal model," and it carries a distinctive mechanism in how it outputs tokens.

This is easier to grasp by looking at the model's architecture diagram first, so let's start there.

The DiffusionGemma-26B-A4B-it Model Configuration

The configuration of the DiffusionGemma-26B-A4B-it model is shown below.

Architecture diagram of Diffusion Gemma-26B-A4B-it
Figure 1: Architecture diagram of Diffusion Gemma-26B-A4B-it

As you may already suspect, this model does use a diffusion model, but its causal-model portion is also used as an encoder — so it's more accurate to call it a causal/non-causal hybrid model.

That said, the causal-model portion is only ever used to ingest the input query, so it doesn't loop repeatedly the way a normal causal model does. It exists solely to build the state-space representation of the input query, so it's traversed at most once.

How Is This Different from Previous Models?

For causal models, the common structure today is usually what's called Decoder-Only Transformers — a neural network that looks like a single straight line, essentially the Decoder split off from the full Transformer architecture.

In this setup, the network takes the input string (the input query) and predicts "the token that will likely come next," one token at a time. It appends the predicted token to the end of the sequence and feeds that back into the same network — and repeats this loop until a control token called End of String appears.

Processing flow in a causal model
Figure 2: Processing flow in a causal model

Diffusion Gemma, by contrast, uses the causal-model portion as an Encoder, and attaches an entirely different mechanism as the Decoder — making the overall setup considerably more complex.

The Decoder portion is built as a diffusion model. It inherits Gemma-4's attention structure, and even its parameter settings are identical to the original model — but the direction information flows between blocks, and the structure of the data fed in, are completely different. In fact, it uses a generation method similar to Stable Diffusion — the model from 2022 often called the "ancestor of image-generation AI" — to output a fixed batch of tokens all at once.

The field in which this batch of 256 tokens is drawn is called the Canvas. In its initial state, the Canvas is filled entirely with masked tokens across all 256 token slots (shown as "noisy tokens" in the diagram below; these are typically represented as [MASK]). The Canvas and the notion of timestep are features that a causal model simply doesn't have.

After a decoding pass finishes, the model checks whether any noise still remains. If it does, the model flags the locations that most need further decoding and runs another decoding pass — and this is the part that's distinctive about running a diffusion model as an LLM.

Processing flow in a Diffusion LLM
Figure 3: Processing flow in a Diffusion LLM

After 256 tokens have been output, the model returns to causal-model-like behavior to build context for the next Canvas, in order to generate the following block — and it repeats this cycle of causal-style context building followed by non-causal refinement until the full text is complete. Because it also carries causal-model characteristics in this way, it helps the model maintain output quality.

The Output Process Is Quite Unusual

Fittingly for a model that generates tokens the way image-generation models do, its output process is genuinely unusual.

(Note: the prompts I actually tested with were in Japanese, so several of the screenshots in this article — including this one, the one below it, and the generation-speed screenshot further down — show Japanese input and output.)

Diffusion Gemma generating tokens, example 1, showing a denoising progress window over a Japanese prompt and response
Figure 4: Diffusion Gemma generating tokens (example 1)

Looking near the bottom of the output, you'll notice the same words and various symbols repeating. The following example is interesting too.

Diffusion Gemma generating tokens, example 2, showing denoising progress with partially resolved tokens
Figure 5: Diffusion Gemma generating tokens (example 2)

The denoising process is indicated by the green window at the top, and the spots that haven't finished denoising appear to be marked with "*". Since this model seems to process diffusion in 256-token units, it's genuinely fascinating to watch the process unfold until the sequence of tokens becomes complete.

Detailed Specs

  • Maximum context size: 256k tokens
  • 30 blocks total (5-LA + 1-GA) × 5 groups
  • Hidden size: 2,816
  • FFN type: MoE (omitted from the diagram above for simplicity)
    • Activation function: Gelu_Pytorch_tanh
    • Activation dimensions: 4,304
  • Attention structure: A hybrid of Local Attention (LA) and Global Attention (GA). The parameter values for these structures are shared between the Encoder and Decoder.
    • Local Attention (LA): Uses Sliding Window Attention
    • Handles local context analysis via the sliding window
    • Low compute cost, fast, and uses less memory
    • Global Attention (GA): Uses Group Query Attention for accurate whole-sequence understanding
    • Grasps the full sequence
    • Acts as the high-performance coordinator that handles long context and complex reasoning
  • Encoder and Decoder:
    • The Encoder is the mechanism that processes the input query, and it has exactly the same structure as the original Gemma-4 — a unidirectional causal attention structure. After passing through the attention structure once, its vector information is stored in the KV cache and relayed to the Decoder.
    • The Decoder has a non-causal, bidirectional network structure and operates on 256-token units, doing the same kind of denoising work you'd see with images. It moves back and forth through this layer, denoising, until the answer for each token is known.
    • Once a full pass of solutions is obtained, the probability distribution is used to judge whether every token has been finalized. If not, the timestep is shifted by one and the denoising process runs again.
    • Self-Conditioning This mechanism, unique to Diffusion Gemma, decides where to focus denoising effort during the denoising process. It corresponds to what's labeled "Memory FFNN" in the diagram, and takes a GatedMLP structure. It's involved in how the Memory variable is handled after passing through the Decoder.
    • Weight sharing This model shares the actual weight values between the Encoder and Decoder. The shared components are as follows:
    • Token embedding layer (embed_tokens)
      • encoder.language_model.embed_tokens.weight ↔ decoder.embed_tokens.weight
      • Also shared with lm_head.weight
    • Compute weights across all transformer layers
      • Self-attention layers: all weights for q_proj, k_proj, v_proj, o_proj
      • MLP / MoE layers: all weights for gate_proj, up_proj, down_proj, plus the MoE router and expert weights
      • LayerNorm layers: the scale weights for each normalization layer, such as input_layernorm, post_attention_layernorm, and pre_feedforward_layernorm
      • Various scale coefficients: scale, per_expert_scale, etc.
    • Final LayerNorm layer (norm)
      • encoder.language_model.norm.weight ↔ decoder.norm.weight
  • Other notable techniques Since this model is based on Gemma-4-26B-A4B, it naturally inherits the following features as well. KV Cache Sharing in particular acts as the bridge connecting the Encoder and Decoder, and its importance is significant.
    • Per-Layer Embedding (PLE): Keeps embedding information independently at the layer level
    • Beyond the embedding information at the model's input, this retains information specific to each layer, contributing to greater representational power at inference time.
    • KV Cache Sharing: A mechanism for sharing the KV cache
    • Normally kept independent, this mechanism shares and synchronizes the KV cache instead.
    • It's used to achieve more consistent inference.
    • When running multi-GPU dispatch, this shared-state synchronization causes frequent inter-GPU communication, so caution is needed in setups without NVLink.
    • Dual RoPE configuration:
    • Local Attention layers use RoPE This embeds token position information using a rotation matrix, ensuring positional information is reliably preserved for long text.
    • Global Attention layers use p-RoPE1 This is a modified version of RoPE for embedding token position information, in which RoPE is applied only to the top p% of K/V matrix pairs.
      • Because it's effective at suppressing the impact of noise, it's applied to the Global Attention layer that acts as the coordinator, contributing to clearer, less complicated positional information over long text.

What's Good About Diffusion Gemma?

Simply Put: It's Fast

Simply put, its token output speed is on an entirely different level. Take a look at the figure below. As discussed later, this is one example of output produced using a tool called Unsloth Studio.

Screenshot of Diffusion Gemma's output showing a generation speed of 2069.5 tokens per second
Figure 6: Diffusion Model's remarkable generation speed

Look at the number at the very bottom: it reads 2069.5 tok/s.
This is the speed that was actually observed — it has not been altered in any way. The GPU used was an NVIDIA A100, and even the smaller Qwen3.5-9B model only manages around 112 tok/s on the same hardware2, so you can see that this output speed is beyond what you'd expect.

How Do You Actually Run It?

As of August 2026, llama.cpp Still Doesn't Support It

My first thought was "surely llama.cpp can run this," so I gave it a try — but it turns out the llama.cpp pull request for this is still marked as a Draft.

GitHub: Diffusion Gemma #24423

Unsloth opened this pull request, but for whatever reason it ended up being treated as a Draft, so being able to run this on official llama.cpp still looks like it's some way off. Because of that, I used Unsloth Studio for this article instead.3

Running Diffusion Gemma with Unsloth Studio

I had already downloaded the GGUF file available from the page below and placed it under <homedir>/llama.cpp/models ahead of time.

unsloth/diffusiongemma-26B-A4B-it-GGUF
https://huggingface.co/unsloth/diffusiongemma-26B-A4B-it-GGUF

Run the following commands to download and install the tool.
Once that's done, reload your .bashrc.

curl -fsSL https://unsloth.ai/install.sh | sh
source ~/.bashrc
Enter fullscreen mode Exit fullscreen mode

At the end, it will ask whether you want to launch now — either Yes or No is fine. You can also launch it later with the command below.

unsloth studio -H 0.0.0.0 -p <port number>
Enter fullscreen mode Exit fullscreen mode

For more detail on this process, the following page is a good reference.

Unsloth Studio Installation

Doing so brings up a screen like this one:

Unsloth Studio screen for setting up an account password

If, instead of the screen above, you're shown a form asking for your "Current Password," stop Unsloth Studio once, then run the following command to reset your password.

unsloth studio reset-password
Enter fullscreen mode Exit fullscreen mode

That will print out a message like the one below, giving you an initial password. Start Unsloth Studio again, and log in with the username "unsloth" and the "".

New password for 'unsloth': <reset password string>

Sessions and API keys revoked. A running Unsloth takes it on the
next request, though repeated failed logins can hold the rate limit
shut for up to a minute.
Enter fullscreen mode Exit fullscreen mode

You'll then see a screen like this. Click "Model Hub."

Unsloth Studio main screen with the Model Hub option

Click "On Device" at the top of the screen, then click "Add folder" and add the path where you placed the model's GGUF file. It will scan that directory for GGUF files. Find the model name starting with "diffusiongemma" that you downloaded earlier.

Unsloth Studio On Device tab for selecting local models
Unsloth Studio Add Folder dialog for pointing to the GGUF model directory
Unsloth Studio showing the diffusiongemma model found in the selected folder

Clicking it brings up a screen like the one below — click the "▶ Run" button.

Unsloth Studio screen with the Run button to launch the model

After about a minute of loading, it'll be ready to use in chat.

Go ahead and experience the blazing-fast world of Diffusion Gemma for yourself.

Unsloth Studio chat interface with Diffusion Gemma loaded and ready to use

A Note on Memory Usage

Running nvidia-smi at this point showed the following memory usage:

Mon Aug 3 22:47:44 2026
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.86.10              Driver Version: 570.86.10      CUDA Version: 12.8      |
|--------------------------------+------------------------+----------------------+
| GPU  Name        Persistence-M | Bus-Id        Disp.A   | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                |                         |              MIG M. |
|================================+=========================+======================|
|   0  NVIDIA A100-PCIE-40GB Off |   00000000:C1:00.0 Off  |                    0 |
| N/A   30C    P0    40W / 250W |   19335MiB / 40960MiB  |      0%      Default |
|                                |                         |             Disabled |
+--------------------------------+-------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                               |
|  GPU   GI   CI        PID   Type   Process name                        GPU Memory        |
|        ID   ID                                                         Usage             |
|===========================================================================================|
|    0   N/A  N/A     87348      C   ...dio/unsloth_studio/bin/python        414MiB         |
|    0   N/A  N/A     87384      C   ...diffusion-gemma-visual-server      18404MiB         |
+-----------------------------------------------------------------------------------------+
Enter fullscreen mode Exit fullscreen mode

As shown above, the 4-bit quantized model consumed roughly 18,404 MiB of memory — about the same as a 4-bit-quantized GGUF of the causal version, Gemma-4-26B-A4B. The fact that memory usage stays this low despite such a complex structure comes down mainly to the fact that the parameter settings are shared between the Encoder and Decoder sides.

That said, this won't run on my own RTX 3060/3050 at home, nor on the RTX A4000 available for individual use on GPUSOROBAN, so you'll need to prepare a reasonably capable GPU for it. It still feels like a model that has a ways to go.

What About Its Reasoning Ability?

Looking at Artificial Analysis, which publishes scores for various LLMs, the score unfortunately isn't very high. Disappointingly, it scores lower than the Gemma-4-26B-A4B model it's based on.

Overall score comparison from Artificial Analysis
Figure 7: Overall score comparison from Artificial Analysis

Even though this is technically a hybrid with a causal model, my own read is that this gap comes from the difference between models that scrutinize content character-by-character and Diffusion Gemma, which scrutinizes a whole 256-token Canvas as a single block — leading to less careful content checking.

That said, compared on the same site against OpenAI's GPT-4o (2024-May), which scores 11, Diffusion Gemma appears to score somewhat higher in quality. In that sense, its handling of text does seem to be at a reasonably solid level, and there may be some room for it in lighter-duty tasks. The detailed breakdown is as follows:

Additional detailed score comparisons from Artificial Analysis
Figure 8: Additional detailed score comparisons

Summary

In this article, I introduced Diffusion Gemma, likely the most polished implementation yet of a diffusion mechanism built into an LLM.

Techniques like MTP (speculative decoding), which I covered previously, and the Diffusion LLM approach covered here both seem like answers to the same demand: making LLMs faster.

On the upside, its remarkable speed hints at a real shift ahead for LLMs. On the downside, its weaker reasoning quality shows this is still very much a work in progress. Considering that even llama.cpp — an OSS inference engine with a huge user base — still doesn't support it since its release, there are clearly still plenty of hurdles before this approach can stand on the front line.

Diffusion LLMs themselves have existed since 2021, and a model called LLaDa drew significant attention in 2025. In Japan, companies like ELYZA have taken up this model as well, and the momentum behind this approach is clearly picking up fast.

With that momentum, model reasoning ability is also improving rapidly, and it seems plausible that models capable of reasoning at a level not far off from causal models will start to appear.

Causal models are gradually hitting a point of saturation in their own approach, with reasoning ability showing signs of plateauing. Against that backdrop, I find myself genuinely curious whether this kind of bold new approach is what breaks through the next wall and delivers the next big leap forward.

References


  1. https://www.linkedin.com/posts/xuan-thai-nguyen-b31b0b3a_little-notes-on-p-rope-in-gemma-4-first-activity-7446051476784951296-vamg ↩

  2. Based on inference measurements I ran myself on an NVIDIA A100 setup. ↩

  3. This — the fact that the implementation approach for diffusion models themselves keeps changing — may well be the biggest challenge facing Diffusion LLMs today. ↩

Top comments (0)