OUI-1-NVFP4

OUI-1-NVFP4

An unofficial NVFP4 quantization of thesysdev/OUI-1, the generative-UI finetune of Google's DiffusionGemma 26B-A4B-it. Not affiliated with or endorsed by Thesys, Google, or NVIDIA.

It fits on two 16 GB GPUs (tensor parallel 2) and keeps most of OUI-1's accuracy on the Generative UI Benchmark.

Why attention and the dense MLP stay in BF16

The exclusion set matches NVIDIA's own DiffusionGemma NVFP4 recipe, and it matters. Earlier bakes that also quantized decoder attention and the dense MLP to W4A4 produced corrupted output: repeated and fused tokens (Card Cardheaderheader, = = = Cardamountamo), with an average output of 334 bytes where OUI-1 writes about 1,800. Keeping those layers in BF16 fixed it. Only the experts, which hold most of the parameters, are 4-bit.

Results on the Generative UI Benchmark

generative-ui-bench, OpenUI format, 46 briefs × 4 generations, thinking off, scored with the benchmark's own score.ts.

Read this caveat before quoting the 121: the calibration set for this quantization is the benchmark's own 46 briefs (with OUI-1's published outputs for them), so the 121/184 is measured on the prompts the activation scales were calibrated on. Calibration only sets activation ranges — it is not training — and the base/finetune gap looks the same without any calibration (24 vs 132 at FP8), but this is not a held-out number.

Hardware: 2× NVIDIA RTX PRO 2000 Blackwell (16 GB, sm_120), vLLM 0.26.0. The 184 generations took 1,061 s at concurrency 4 (the NVIDIA base checkpoint took 1,264 s under identical settings; OUI-1 writes about 40% less, and a diffusion model's time scales with output length).

Serving (vLLM 0.26, two 16 GB cards)

vllm serve <this-repo> \
  --tensor-parallel-size 2 --disable-custom-all-reduce \
  --kernel-config '{"moe_backend": "MARLIN"}' \
  --limit-mm-per-prompt '{"image":0,"video":0}' \
  --max-model-len 16384 --max-num-seqs 4 \
  --kv-cache-dtype fp8 --enforce-eager \
  --kv-cache-memory 698351616 \
  --diffusion-config '{"canvas_length":256,"max_denoising_steps":48}' \
  --default-chat-template-kwargs '{"enable_thinking": false}'

Each flag was needed on this hardware:

• moe_backend: MARLIN — moe_intermediate_size is 704, which the FP4 MoE kernels (FLASHINFER_CUTLASS and the other CUTLASS-family backends) do not accept; vLLM's automatic choice fails to start. MARLIN and HUMMING accept 704.
• --limit-mm-per-prompt — vLLM registers DiffusionGemma as multimodal; this is a text-only benchmark setting.
• --kv-cache-memory (a fixed KV size instead of --gpu-memory-utilization) — with the profiler sizing KV greedily, the canvas-256 warmup buffer did not fit on a 16 GB card.
• Do not send temperature or seed. vLLM 0.26 rejects them for diffusion models (0.24 silently ignored them). The checkpoint applies its own sampling schedule. On a streamed request the refusal arrives as an error event inside a 200 response, not as a status code.

Chat template. The chat_template.jinja here is OUI-1's own. Compared with Google's DiffusionGemma template it lacks three lines that close an empty thought channel when thinking is off:

{%- if not enable_thinking | default(false) -%}
    {{- '<|channel>thought\n<channel|>' -}}
{%- endif -%}

On vLLM 0.26 without them, responses can begin with a literal thought line. The benchmark numbers above were measured with those three lines restored (Google's template), for both models.

One change to config.json after quantization

DiffusionGemma's global-attention layers (5, 11, 17, 23, 29) share K and V, so they have no v_proj weight. compressed-tensors writes the ignore list expanded to the modules that exist, so those five v_proj entries were missing. vLLM fuses q/k/v and requires one scheme for all three, and refused to load (Found a different quantization schemes for ['q_proj', 'k_proj', 'v_proj']). The five model.decoder.layers.{5,11,17,23,29}.self_attn.v_proj entries were added to quantization_config.ignore. No weight was changed; V is K there, and K is BF16.

How it was made

• llm-compressor 0.12, QuantizationModifier(targets="Linear", scheme="NVFP4"), CPU-only (the sequential pipeline would not trace this model with a GPU attached), about 4.5 hours.
• The fused 3-D expert parameters were first rewritten as per-expert nn.Linear modules (named experts.N.{gate,up,down}_proj, the layout vLLM's FusedMoE loads) so targets="Linear" reaches them.
• Calibration: 184 samples, max sequence length 8,192 (system prompt alone is about 5,200 tokens) — see the caveat above.

License

Apache 2.0, as for OUI-1 and the base model; use is also subject to the Gemma license. The base weights are copyright Google DeepMind and the finetune is by Thesys; this repository contains a quantized derivative of those weights, modified as described above.