Add diffusion-gemma block-diffusion support

[lnigam]

Overview

This PR adds initial diffusion-gemma support for Gemma 4 based block-diffusion checkpoints. This is just a draft PR to get feedback on multiple design aspects of diffusion model like block diffusion, approximation of soft embeddings, separate vs single encoder-decoder, prefix KV-cache reuse, diffusion server utility etc.

git clone https://github.com/ggml-org/llama.cpp
cd llama.cpp
gh pr checkout 24427
cmake -B build -DGGML_CUDA=ON
cmake --build build -j --config Release

then use a GGUF (any can work but for eg)

pip install -U "huggingface_hub[cli]"
hf download unsloth/diffusiongemma-26B-A4B-it-GGUF \
    --local-dir unsloth/diffusiongemma-26B-A4B-it-GGUF \
    --include "*Q8_0*" # Use "*Q4_K_M*" for a smaller 16 GB download

Diffusion Gemma benchmark commands

Benchmark scripts:

• https://github.com/lnigam/llama.cpp-scripts/tree/diffusion-gemma

The optimized CUDA path uses the current binary defaults for fused full-softmax,
fused top-k sampling, fused self-conditioning embedding, final softcap fusion,
direct self-conditioning, and final token copy on stop.

The only extra tuning flag passed explicitly is:

--diffusion-cuda-mmq-max-x 64

--run-max-denoising-step is not enabled. Top-k is not passed, so the default
full-softmax path is used (top_k=0).

CLI

.\build\bin\Release\llama-diffusion-gemma-cli.exe `
  -m "<path-to-model.gguf>" `
  -p "Answer in about 1000 words: explain block diffusion generation and CUDA sampling optimizations." `
  -n 1024 `
  -c 8096 `
  -ngl 999 `
  --diffusion-steps 48 `
  --diffusion-cuda-mmq-max-x 64

Server

.\build\bin\Release\llama-diffusion-gemma-server.exe `
  -m "<path-to-model.gguf>" `
  --host 127.0.0.1 `
  --port 18081 `
  -c 8096 `
  -ngl 999 `
  --diffusion-steps 48 `
  --diffusion-cuda-mmq-max-x 64 `
  --metrics `
  --slots

Example request:

curl.exe http://127.0.0.1:18081/v1/chat/completions `
  -H "Content-Type: application/json" `
  -d "{\"model\":\"diffusion-gemma\",\"messages\":[{\"role\":\"user\",\"content\":\"Answer in about 1000 words: explain block diffusion generation.\"}],\"max_tokens\":1024}"

Benchmark scripts

CLI benchmark:

python .\bench-diffusion-gemma-cli.py `
  --binary <path-to-llama-diffusion-gemma-cli> `
  --model <path-to-model.gguf> `
  --prompt-file .\diffusion-gemma-prompts.txt `
  --output-dir .\benchmark-results `
  --repeat 1 `
  --warmup 0 `
  --n-predict 1024

Server benchmark:

python .\bench-diffusion-gemma-server.py `
  --binary <path-to-llama-diffusion-gemma-server> `
  --model <path-to-model.gguf> `
  --prompt-file .\diffusion-gemma-prompts.txt `
  --output-dir .\benchmark-results `
  --repeat 1 `
  --warmup 0 `
  --max-tokens 1024

The benchmark scripts default to:

• --ctx-size 8096
• --diffusion-steps 48
• --diffusion-cuda-mmq-max-x 64
• no request-level top-k override, so the binary uses top_k=0
• no --ignore-eos, so EOS is respected by default

Additional information

• Add GGUF conversion support for diffusion Gemma checkpoints, including self-conditioning tensors and multimodal Gemma 4 vision/mmproj export.
• Register the diffusion-gemma architecture and model implementation.
• Implement the diffusion Gemma graph using Gemma 4 decoder blocks with bidirectional canvas attention, prompt-prefix conditioning, KV-cache reuse, and self-conditioning.
• Add sparse top-k self-conditioning through on-device embedding gather.
• Add llama-diffusion-gemma-cli for block-diffusion generation.
• Add llama-diffusion-gemma-server, an HTTP server with /v1/completions, /v1/chat/completions, /health, /props, /metrics, and /slots.
• Add CUDA backend support for diffusion top-k sampling, entropy/stability decisions, self-conditioning buffers, device-side canvas updates, and device-loop early stopping.
• Enable CUDA graph friendly execution by keeping persistent diffusion inputs/output state on device and avoiding inter-step host sampling copies.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: yes, for generating the initial architecture support related changes and some bug fixes, merge conflicts and code review

[Ajay9o9]

Tried the latest nvidia-diffusion-gemma branch (e1fc5359f) with diffusiongemma-26B-A4B-it-Q4_K_M.gguf and hit a load error:

i'm on
Branch: nvidia-diffusion-gemma
Commit: e1fc535

Model from Unsloth
diffusiongemma-26B-A4B-it-Q4_K_M.gguf

loading fails with

/repos/llama-diffusion-gemma$ ./build/bin/llama-diffusion-gemma-cli \
  -m "/media/gemma4/diffusiongemma-26B-A4B-it-Q4_K_M.gguf" \
  -p "Hello"
  
ggml_cuda_init: found 1 CUDA devices (Total VRAM: 11889 MiB):
  Device 0: NVIDIA GeForce RTX 3060, compute capability 8.6, VMM: yes, VRAM: 11889 MiB
0.00.624.210 W load: control-looking token:    212 '</s>' was not control-type; this is probably a bug in the model. its type will be overridden
0.00.624.556 W load: control-looking token:     50 '<|tool_response>' was not control-type; this is probably a bug in the model. its type will be overridden
0.00.642.077 W load: special_eog_ids contains '<|tool_response>', removing '</s>' token from EOG list
0.00.861.710 E llama_model_load: error loading model: missing tensor 'self_cond_norm.weight'
0.00.861.719 E llama_model_load_from_file_impl: failed to load model
0.00.861.720 E error: failed to load model '/media/gemma4/diffusiongemma-26B-A4B-it-Q4_K_M.gguf'

missing tensor 'self_cond_norm.weight'

GGUF, inspection shows :

self_cond_down.weight
self_cond_gate.weight
self_cond_pre_norm.weight
self_cond_up.weight

but not:

self_cond_norm.weight

The runtime appears to expect:

{ LLM_TENSOR_SELF_COND_NORM, "self_cond_norm" }

from llama-arch.cpp

I also checked the current branch and couldn't find any references to self_cond_pre_norm in the runtime:

grep -R "self_cond_pre_norm" src/ -n

returns nothing.

Is this a naming mismatch?

[Jakeshadow]

Nice, been waiting for this. For anyone building from this PR — heads up that Apple Silicon won't see the speedup. It's a compute-bound model, needs a high-CUDA-core GPU. Detailed GPU requirements here: https://diffrun.dev

[lnigam]

@Ajay9o9 Now this PR works for unsloth model.

[mohamed-em2m]

llama-diffusion-gemma-server don't appear in builds

[lnigam]

@mohamed-em2m they are building in my local and I also can see it in the CI logs. Are there any errors while building?

[theIvanR]

Awesome job! I will have a look later at the code. Proposal: make the code modular with a clear fetcher, and builder(s).

For the fetcher do something like this:

git clone --depth 1 https://github.com/ggml-org/llama.cpp
cd llama.cpp

git fetch --depth 1 origin pull/24427/head:pr24427
git checkout pr24427

While for the rest, something similar to how I did it here:
https://github.com/theIvanR/lmstudio-unlocked-backend/tree/main/Generate%20Backends/Windows

EDIT:
I tried with this cpu only builder and something is broken:

C:\Users\Admin\source>talk_to_gemma.cmd
warning: no usable GPU found, --gpu-layers option will be ignored
warning: one possible reason is that llama.cpp was compiled without GPU support
warning: consult docs/build.md for compilation instructions
0.00.000.889 I diffusion-gemma: GGML_CUDA_MMQ_MAX_X=64
0.01.195.906 W load: control-looking token:     50 '<|tool_response>' was not control-type; this is probably a bug in the model. its type will be overridden
0.01.196.466 W load: control-looking token:    212 '</s>' was not control-type; this is probably a bug in the model. its type will be overridden
0.01.252.519 W load: special_eog_ids contains '<|tool_response>', removing '</s>' token from EOG list
[SUCCESS] Done

@echo off
setlocal

set "EXE=C:\Users\Admin\source\llama.cpp\build_cpu\bin\llama-diffusion-gemma-cli.exe"
set "MODEL=C:\Users\Admin\Documents\LLM Models\unsloth\diffusiongemma-26B-A4B-it-GGUF\diffusiongemma-26B-A4B-it-Q4_K_M.gguf"

if not exist "%EXE%" (
    echo [ERROR] CLI not found: "%EXE%"
    exit /b 1
)

if not exist "%MODEL%" (
    echo [ERROR] Model not found: "%MODEL%"
    exit /b 1
)

"%EXE%" ^
  -m "%MODEL%" ^
  -p "Answer in about 1000 words: explain block diffusion generation and CUDA sampling optimizations." ^
  -n 1024 ^
  -c 8096 ^
  -ngl 999 ^
  --diffusion-steps 48 ^
  --diffusion-cuda-mmq-max-x 64

if errorlevel 1 (
    echo [ERROR] Command failed
    exit /b 1
)

echo [SUCCESS] Done
exit /b 0

EDIT 2: trying on gpu with my builder and this directory:

set "EXE=C:\Users\Admin\source\llama.cpp\build_gpu_cuda\bin\llama-diffusion-gemma-cli.exe"

It did something:

PS C:\Users\Admin\source> ./talk_to_gemma.cmd
ggml_cuda_init: found 1 CUDA devices (Total VRAM: 8191 MiB):
  Device 0: NVIDIA GeForce GTX 1070 Ti, compute capability 6.1, VMM: yes, VRAM: 8191 MiB
0.00.000.655 I diffusion-gemma: GGML_CUDA_MMQ_MAX_X=64
0.01.348.419 W load: control-looking token:     50 '<|tool_response>' was not control-type; this is probably a bug in the model. its type will be overridden
0.01.349.167 W load: control-looking token:    212 '</s>' was not control-type; this is probably a bug in the model. its type will be overridden
0.01.434.171 W load: special_eog_ids contains '<|tool_response>', removing '</s>' token from EOG list
0.17.762.433 I formatted prompt: <|turn>system
<|think|>
<turn|>
<|turn>user
Answer in about 1000 words: explain block diffusion generation and CUDA sampling optimizations.<turn|>
<|turn>model

0.17.762.850 W llama_context: n_ctx_seq (8192) < n_ctx_train (262144) -- the full capacity of the model will not be utilized
0.17.763.676 W llama_kv_cache_iswa: using full-size SWA cache (ref: https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
0.18.941.717 I diffusion-gemma: prefix=34 canvas=256 max_canvases=4 steps=48 entropy_bound=0.100 temp=[0.40,0.80] n_ctx=8096 mm=0
0.18.941.728 I diffusion-gemma: gpu sampling: on | device self-cond: on | device loop: on
0.18.941.728 I diffusion-gemma: device early-stop interval=1
0.23.787.601 I prefill (encoder, no self-cond): 34 tokens in 4.846 s (7.0 tok/s)

[lnigam]

Latest benchmark numbers on RTX 5090 using my benchmark

Server diffusion benchmark

ready ms: 9080.63
runs kept: 10
request wall ms avg/min/max: 5461.17 / 5142.71 / 5563.61
predicted ms avg: 5402.404
canvas tok/s avg: 189.650
steps/s avg: 34.984
ms/step avg: 28.585
mean denoising steps/block: 47.250

Using aiperf (speedbench)
command: aiperf profile -m "diffusiongemma-26B-A4B-it-Q4_K_M" --concurrency 1 --tokenizer "google/gemma-4-26B-A4B-it" --endpoint-type chat -u http://127.0.0.1:18081 --artifact-dir ./gemma-diffsion-5090-topk0 --public-dataset speed_bench_throughput_2k --osl 1024 --extra-inputs "max_tokens:1024,min_tokens:1024,ignore_eos:true,top_k:0" --osl-stddev 0 --ui-type none --streaming --request-count 100 --warmup-request-count 10 --use-server-token-count --request-timeout-seconds 1200

[coder543]

Google reported 700 tokens per second on RTX 5090, so I guess there is a lot of room for optimization here

[lnigam]

@theIvanR Can you try the CPU only build with -DGGML_CPU_REPACK=OFF
I dont have 1070 Ti with me. For GPU build can you try following:

%EXE% ^
-m "%MODEL%" ^
-p "Briefly explain block diffusion." ^
-n 64 ^
-c 1024 ^
-ngl 8 ^
--no-kv-offload ^
--diffusion-steps 2 ^
--diffusion-cuda-mmq-max-x 0 ^
--top-k 256 ^
--no-diffusion-device-denoise-loop

[lnigam]

@coder543 It depends mostly on the number of denoising steps needed for the canvas to converge. For coding prompts, model converges generally in around 20 steps while the above benchmark took on an average 34 steps to converge.
diffusiongemma-Q8_0 model is performing relatively better than Q4_KM. and with coding speedbench it is reaching around 300 t/s.

For some coding prompts it is reaching upto 534 t/s

[mohamed-em2m]

@mohamed-em2m they are building in my local and I also can see it in the CI logs. Are there any errors while building?

it's appear now

[theIvanR]

@lnigam Thanks for getting back, I will try it later when I get some time. In the mean time I would also suggest adding a verbose mode and or a debug one to filter out exactly these issues when launching, say with sensible fallbacks and try catch type blocks.

[Hcl192088]

How fast can it run if partial-offload? e.g. on 12gvram like 4070
I know MoE model can still be decent fast with partial offload does it work on this one as well?

[Fantou-Arthur]

How fast can it run if partial-offload? e.g. on 12gvram like 4070 I know MoE model can still be decent fast with partial offload does it work on this one as well?

I don't have a proper answer but probably really slow :

This is a diffusion model, to be efficient the prompt/output tokens needs to go throught the different layers several times (I don't know how many for a proper answer, but it's ~20 according to gemini). So basically the tokens are going into all of the layers at the same time, the ones going in your gpu will be quick but the other one will goes throught your ram Which would slow down by a lot the answers. I'm not sure that even with a very good gpu doing some offloading on diffusion llm would be efficient.
You could try to use diffusiongemma q2 or q3 but I never had satisfying experiences with models under q4.

(On a 4070 I'ld recommend you to use qwen3.6 a3b in q4. I have a similar setup and it's what I use to code with Cline.)

TLDR : I don't know how fast it would run but probably slow because the entire model needs to be loaded simultaneously.

[lnigam]

Latest performance numbers on RTX6000 pro-

[KorigamiK]

Thanks for all the development on this pull request, I compiled this branch on my hardware

• GPU: NVIDIA GeForce RTX 3090
• VRAM: 24576 MiB
• Driver: 550.142
• CUDA: 12.4

These were my results on running a few tests while also testing out the server on Zed locally for refactoring some code with diffusiongemma-26B-A4B IQ4_XS quant

llama-diffusion-gemma-cli -m model.gguf \
  -n 512 \
  --diffusion-steps 64 \
  --temp 0.7 \
  -p "Write a Python function to merge two sorted lists, include type hints, docstring, and a small test block."

• perf: 47 effective denoising steps, 256 answer tokens
• generation: 3.57 s
• throughput: 71.8 canvas tok/s
• step time: 0.076 s/step
• wall clock including load: 22.39 s

What would be nicer would be to have the model integrated into the llama-server with similar api schema!

[lockewerks]

diffusion-gemma-cli.cpp doesn't set inputs.enable_thinking in format_chat, so it takes the true default from common_chat_templates_inputs and the template emits the <|think|> preamble. diffusion-gemma-server.cpp:227 sets it to false, so the two binaries format the same prompt differently.

In the CLI that matters most at small -n: with a single canvas the thought channel fills all 256 tokens, no <channel|> close is ever emitted, and the extraction near the end of main falls through to returning the raw thought text.

     inputs.add_generation_prompt = true;
+    inputs.enable_thinking = false;
     return common_chat_templates_apply(tmpls.get(), inputs).prompt;

With that, -n 256 -c 4096 on unsloth Q4_K_M returns a clean answer in one canvas. RTX 4090, full offload, for a datapoint between your 5090 and 1070 Ti numbers: 148 canvas tok/s at -c 4096, dropping to 70 at -c 8192 where it runs out of VRAM at 22.4G/23G.

Happy to send this as a PR against your branch if that's easier.