diffusiongemma-26B-A4B-w4a16

DiffusionGemma-26B-A4B W4A16

2.3 to 2.7x faster than the bf16 base at comparable accuracy, on a single GB10.

This repo is two things: a 4-bit weight-quantized google/diffusiongemma-26B-A4B-it, and the Triton kernels that make it fast. Quantization alone gets you about half the win. The rest comes from the kernels, so they ship together and from_pretrained wires up both.

What the model is

DiffusionGemma is a block-diffusion language model, not autoregressive. It denoises a 256-token canvas over up to 48 steps, committing tokens once their entropy drops below a bound, instead of emitting one token at a time. 26B total parameters with 4B active: a 128-expert MoE routing top-8 per token, across 30 encoder and 30 decoder layers.

That architecture is why off-the-shelf quantization recipes don't transfer directly. They assume an autoregressive decode loop, where weight bit-width and sampling are independent knobs. Here they are not: quantization noise feeds the entropy bound that decides how many tokens commit per step, so granularity and throughput are coupled. The kernels in this repo target the denoising loop.

What's in here

The expert weights are asymmetric uint4 with an fp16 scale and zero-point per group. The GEMM reads packed nibbles and dequantizes inline, so it streams real 4-bit traffic rather than fake-quant. Experts drop from 85 GiB to 11.5 GiB, and since the expert path is bandwidth-bound that accounts for most of the speedup.

On top of that: one fused Triton launch handles all 128 experts instead of a 128-launch Python loop, with the activation folded into the first GEMM's epilogue. RMSNorm collapses from roughly six elementwise launches into one, 180 times per forward. The weighted combine uses a fixed-order reduction rather than index_add_, because in block diffusion a 1-ULP difference from atomics flips an acceptance threshold and forks the entire trajectory.

The largest single kernel win came from the sampler. The stock generation loop computes Categorical(logits).entropy() twice per denoising step on the same 268 MiB tensor. Computing it once, in a single streaming pass, and replacing softmax plus multinomial with one Gumbel-max kernel is worth about 1.4x by itself. That finding is upstreamed as transformers#47582.

Quickstart

pip install "transformers>=5.12.1" torch triton accelerate safetensors huggingface_hub

from transformers import AutoModelForCausalLM, AutoTokenizer

REPO  = "GoedelMachines/diffusiongemma-26B-A4B-w4a16"
model = AutoModelForCausalLM.from_pretrained(REPO, trust_remote_code=True)
tok   = AutoTokenizer.from_pretrained(REPO)

ids = tok.apply_chat_template([{"role": "user", "content": "What is 17*23? Show your work."}],
                              add_generation_prompt=True, return_tensors="pt",
                              return_dict=True)["input_ids"].cuda()
out = model.generate(input_ids=ids, max_new_tokens=512)
print(tok.decode(out.sequences[0][ids.shape[1]:], skip_special_tokens=True))

trust_remote_code=True is required. It routes through this repo's modeling_dg_w4.py, which unpacks the 4-bit experts and installs the Triton kernels. Without it transformers cannot read the packed expert buffers at all. What comes back is an ordinary DiffusionGemmaForBlockDiffusion, so .generate() and everything else behave normally. You just get the fast path by default.

First forward JIT-compiles the kernels, about 45 seconds.

Three tiers are available via tier= or the DG_W4_TIER env var. turbo is the default and uses every kernel, at accuracy parity with a different RNG stream. bitexact applies only the entropy dedup, which cannot change output, and runs roughly 30% slower. off gives plain W4 with no sampler or norm kernels.

If you'd rather skip the Auto classes, load_fast() in kernels/load_w4_checkpoint.py is what they call into. Clone the repo and python run.py "your prompt" also works.

Benchmarks

Same GB10 box, same prompts, same generation config, single stream, batch size 1. Throughput is total tokens divided by total wall clock.

Accuracy deltas of +1, -1, -1 at n=100 sit well inside a binomial standard error of about 3 points, so read them as indistinguishable from parity rather than evidence of it. One MATH failure in the W4 arm is a 2048-token truncation the bf16 arm didn't hit, so the real quality gap is probably smaller than the table suggests. Quantization is round-to-nearest min/max with no calibration; GPTQ or AWQ would likely recover some of it.

Outputs are not token-identical to bf16. Per-question agreement is 99%, 95%, and 91%. The model reaches comparable scores along different trajectories, which is what 4-bit weights do.

Reproduced independently on the shipped checkpoint (rather than by packing experts from the bf16 base): GSM8K 49/50 at 190.0 tok/s, HumanEval 47/50 at 237.6, and MATH-500 431/500 at 175.5 on a clean GPU. The first 100 MATH-500 questions scored 87/100, matching the packed-path run exactly.

If you prefer mean per-request decode rate over aggregate throughput, the W4 numbers are 186.2, 224.9, and 168.4. The two conventions differ by up to 13% and papers disagree on which to use, so both are here.

RTX 5090

Runs in 18 GB. The bf16 model is 48.1 GiB and does not fit in 32 GB, so there is no baseline to compare against on this card.

Latency on everyday questions

The benchmarks above are long chain-of-thought. Short interactive questions behave differently, because the model denoises a fixed 256-token canvas whether the answer is 12 tokens or 250. Seventeen ordinary questions on a 5090:

Tokens per second badly understates short replies: a 12-token answer reads 79 tok/s but comes back in 0.15 s. For conversational use, time to answer is the metric that means anything.

Requirements

Tile configs ship for GB10 (sm_121) and RTX 5090 (sm_120). Anything else falls back to the GB10 config, which is safe rather than tuned. Needs about 18 GB of VRAM for weights and 20 GB peak; the fp32 logits tensor alone is 268 MiB per denoising step. Download is roughly 17 GB. Software: transformers>=5.12.1, torch>=2.11, triton, accelerate, safetensors.

A CUDA moe_align_block_size extension is used if present at ~/.cache/torch_extensions/*/marlin_aux_gb10/marlin_aux_gb10.so, worth 3 to 7% of forward time. It is not shipped, and without it the pure-PyTorch alignment path runs automatically. Every benchmark above was produced with it available, so expect to land slightly below these figures without it.

Checkpoint format

dg-w4a16-v1. Asymmetric uint4, fp16 scale and zero-point per group, two nibbles per byte (byte j = q[2j] | q[2j+1] << 4). gate_up_proj uses group size 128 along the contraction dim. down_proj uses 64, because its contraction dim is 704 and 128 doesn't divide it.

Group size is pinned to the GEMM's K-tile so one scale and zero load per K-chunk. We tried padding 704 to 768 to allow group 128. The isolated kernel got 1.27x faster, but the coarser quantization raised per-token entropy enough that fewer tokens cleared the acceptance bound each step, and end-to-end throughput dropped. Quantization granularity is coupled to throughput here in a way it simply isn't for autoregressive models.

Layer N of the encoder and decoder hold byte-identical expert weights, verified across all 30 pairs. The release ships 30 shards and binds each to both modules, halving the download and the expert VRAM.

Limitations

Tuned for one GPU; configs for other architectures are untested end to end. Accuracy was validated on 100 questions per benchmark, enough to rule out a large regression but not to resolve a 1-point difference. No calibration-based quantization, RTN only. The vision tower is carried at bf16 and unquantized, and multimodal paths are untested here. MOE_NS, the preshuffled-transposed down pack, is reachable only through the offline packer and not the checkpoint loader.

License

Derivative of google/diffusiongemma-26B-A4B-it, governed by the Gemma Terms of Use and the Gemma Prohibited Use Policy. The Triton kernels in kernels/ are Apache-2.0. The weights are not. See LICENSE-NOTICE.md.

Citation

@misc{diffusiongemma-w4a16,
  title  = {DiffusionGemma-26B-A4B W4A16: 4-bit block-diffusion inference on GB10},
  year   = {2026},
  note   = {https://huggingface.co/GoedelMachines/diffusiongemma-26B-A4B-w4a16}
}