FLUX.2 Klein 9B OrbitQuant W4A4

This is a complete Diffusers pipeline derived from black-forest-labs/FLUX.2-klein-9B at revision 92196c8e11f7b6cf2b7493e037d8c5345c559216.

Both compute-heavy components are compressed:

transformer: 144 OrbitQuant W4A4 projections plus 3 AdaLN INT4 weight-only projections, stored as fused blocks (see Fused transformer);
Qwen3 text encoder: 252 OrbitQuant W4A4 projections; lm_head remains BF16.

The text encoder is quantized to permit a controlled comparison with the published FLUX.2-klein-9B-SDNQ-uint4-static pipeline. This is an extension of OrbitQuant's architecture-independent adapter; the OrbitQuant paper itself leaves text encoders in BF16.

Install

pip install "orbitquant[hf,kernels]>=0.11.0"

The transformer is a fused checkpoint and needs OrbitQuant 0.11 or newer. The per-projection layout read by earlier versions stays available at revision 1bdb332c.

OrbitQuant uses packed low-bit inference by default and does not silently materialize all weights in BF16. The optimized native kernel package is provisioned automatically at first model load: OrbitQuant derives the exact runtime variant (torch minor and CUDA version for CUDA, the torch stable ABI for CPU, plus OS and architecture), downloads the matching prebuilt wheel from the OrbitQuant kernels release with checksum verification, and caches it under ~/.cache/orbitquant/kernels. When no variant matches, CUDA inference falls back to the Triton packed path. Provision explicitly (or inspect the resolution) with:

orbitquant kernels-install
orbitquant kernels-status

For ComfyUI, install the ComfyUI-OrbitQuant node pack; its install hook provisions the same kernels automatically. The full contract lives in the provisioning documentation.

With PyTorch 2.9, launch CUDA inference with expandable allocator segments to minimize reserved memory:

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py

Diffusers

import torch
import orbitquant  # registers the quantizer
from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained(
    "WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4",
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(
    prompt="An intricate orbital conservatory above Earth, documentary realism",
    height=1024,
    width=1024,
    num_inference_steps=4,
    guidance_scale=1.0,
    generator=torch.Generator(device="cuda").manual_seed(0),
).images[0]
image.save("orbitquant.png")

Fused transformer

OrbitQuant 0.11 runs every transformer block fused, and this checkpoint stores the transformer in that layout, so from_pretrained loads it fused with no conversion step:

Double-stream blocks: per stream, the LayerNorm and AdaLN modulation run in the activation-quantization prologue of one Q|K|V GEMM; Q/K RMSNorm and RoPE are one kernel; the output projection and the SwiGLU feed-forward apply their gated residual updates in the GEMM epilogue.
Single-stream blocks: the joint QKV+MLP projection splits into a Q|K|V GEMM and a SwiGLU GEMM that writes next to the attention output, so the output projection reads both from one buffer.
Attention computes Q·Kᵀ in INT8 and P·V in FP16 (the SageAttention v1 scheme) in blocks whose Q/K RMSNorm weights stay within a narrow range, and BF16 Flash attention in the others.
Reference-image editing keeps diffusers' KV-cache attention pattern.

RTX 4060 Ti 16 GB, torch 2.10.0+cu128, Diffusers 80c7ed26 (0.41.0.dev0), whole pipeline resident, 1024×1024, four steps, guidance 1.0, seed 0, the ten prompts of prompts.json, hot:

Transformer	Denoise (4 steps)	Image	Reference edit	CUDA allocated peak
Per-projection W4A4 (revision `1bdb332c`)	5.44 s	6.24 s	13.4 s	12.93 GiB
Fused W4A4 (this revision)	3.15 s	3.98 s	8.3 s	12.60 GiB

The fused kernels round differently from the per-projection path, so a seed's four-step trajectory moves: images stay as detailed and as legible as before but are not identical to the per-projection ones (mean PSNR 17.2 dB between the two over the ten prompts). Reference edits, which the reference image anchors, agree closely (PSNR 32.7 dB, SSIM 0.985). The benchmark and visual comparison below were made with the per-projection transformer.

Convert the source checkpoint on load

To create a fresh transformer-only OrbitQuant pipeline directly from the source safetensors checkpoint, use the normal Diffusers loader. This row-streams source weights into packed tensors instead of keeping the complete BF16 transformer and the quantized transformer resident together:

import torch
import orbitquant
from diffusers import DiffusionPipeline
from orbitquant import (
    OrbitQuantConfig,
    build_diffusers_pipeline_quantization_config,
)

qconfig = build_diffusers_pipeline_quantization_config(
    OrbitQuantConfig(target_policy="auto"),
    components="transformer",
)
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B",
    quantization_config=qconfig,
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()

Use pipe.enable_sequential_cpu_offload() instead for sequential offload, and orbitquant.fused.fuse(pipe.transformer) (before moving it to the GPU) to run the converted transformer fused. The published artifact also quantizes its text encoder for the SDNQ comparison; on-the-fly conversion keeps text encoders in source precision unless they are explicitly included in components. Guaranteed bounded-memory conversion requires a safetensors source checkpoint.

Transformers Component

The quantized Qwen3 component can also be loaded through Transformers:

import torch
import orbitquant
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4"
tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder="tokenizer")
text_encoder = AutoModelForCausalLM.from_pretrained(
    repo_id,
    subfolder="text_encoder",
    torch_dtype=torch.bfloat16,
).to("cuda")

Quantization

Setting	Value
Weight / activation bits	W4A4
Rotation	RPBH, seed 0
Block policy	Largest power-of-two divisor of the input dimension
Codebook	Lloyd-Max version 2
Row norms	BF16
AdaLN	INT4 RTN, group size 64, BF16 activations
Runtime	transformer: fused blocks (`orbitquant.fused`, family `flux2`); text encoder: `auto_fused` native RPBH/INT8 surrogate plus CUTLASS W4A4 on supported CUDA GPUs
Calibration data	None

The packed transformer and text-encoder weight payload is 10.67 GB; including the VAE, the weight payload is 10.84 GB. The complete pipeline before this card asset is 10.85 GB.

Benchmark (RTX PRO 6000, OrbitQuant 0.5.0, per-projection transformer)

All three variants ran back-to-back on one NVIDIA RTX PRO 6000 Blackwell Workstation Edition (96 GB) in one session, one separate process per variant, same protocol and pinned revisions (BF16 92196c8, SDNQ ed71b3f, OrbitQuant ee3a38f), Torch 2.9.1+cu128, Diffusers 0.39.0, Transformers 5.13.0, BF16 arithmetic, no CPU offload, 1024x1024, four steps, guidance 1.0, seed 0, ten prompts, PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, orbitquant==0.5.0 with a locally built sm_120 native kernel package:

Variant	Load	Cold image	Hot mean	Hot median	Hot NVML peak	CUDA allocated peak	CUDA reserved peak
BF16 (no quantization)	3.06 s	1.54 s	1.2014 s	1.2019 s	36.40 GB	34.72 GB	35.12 GB
SDNQ UINT4	3.02 s	5.81 s	1.2735 s	1.2730 s	15.42 GB	13.75 GB	14.14 GB
OrbitQuant W4A4 (0.5.0)	1.95 s	4.96 s	1.1633 s	1.1640 s	14.17 GB	12.51 GB	12.89 GB

OrbitQuant's hot median is 3.2% faster than unquantized BF16 and 8.6% faster than SDNQ, it loads fastest of the three, and it uses 22.2 GB less peak NVML memory than BF16 and 1.25 GB less than SDNQ. Machine-readable results: benchmark/summary.json.

Visual Comparison

The matrix uses the complete ten-prompt stress pack with full 1024x1024 tiles and WebP quality 95. BF16 is the full-precision reference from the controlled visual run; the SDNQ and OrbitQuant columns use the recorded benchmark outputs. Every row uses the same prompt, seed, resolution, step count and guidance.

Visual Assessment

No collapse: all thirty outputs are finite, coherent and detailed. OrbitQuant did not produce blank, noisy or structurally broken images.
Micro-detail and materials: all three variants preserve gears, filigree, architectural interiors, paper grain, metal, resin and reflected surfaces. OrbitQuant remains competitive with BF16 and SDNQ in these cases.
Dense composition: all variants retain foreground/background separation and the main hierarchy in the architectural cutaway and orbital-banquet prompts. Individual requested objects move or disappear because quantization changes the denoising trajectory.
Counting: none of the variants reliably renders exactly nine performers or every exact repeated motif. This is a base-model limitation in the tested setting rather than an OrbitQuant-only collapse.
English typography: OrbitQuant is strongest on this row: it preserves the headline, subtitle and all four specification lines. SDNQ preserves the headline and three table lines but omits or corrupts some requested text.
Russian typography: all variants render the large headline, subtitle and archive stamp well; small contents text contains errors in every column.
Japanese and Chinese typography: visual glyph quality is plausible, but exact requested strings are not reliably reproduced by any variant.
Trajectory fidelity: both quantizers change the denoising trajectory at the same seed; neither remains consistently closer to BF16 across all ten prompts.

This subjective paired inspection demonstrates non-collapse and exposes concrete failure modes; it is not a substitute for GenEval or another task-specific objective metric. In this controlled FLUX.2 Klein 9B comparison, OrbitQuant produces the smaller complete 4-bit pipeline, reaches SDNQ hot-generation parity with lower runtime memory, and produces the strongest English fine-print result in the prompt pack.

Limitations

The fused transformer runs on CUDA through Triton kernels only; there is no CPU or MPS path.
The text encoder's fastest native path needs the matching prebuilt (or locally built) kernel package; Triton packed matmul is a compatible fallback.
The optimized CUDA W4A4 tensor-core path maps each fixed Lloyd-Max codebook to symmetric INT8 surrogate values and one scalar. Use runtime_mode="dequant_bf16" for exact-centroid reference evaluation.
W4A4 activation quantization can move the denoising trajectory farther from BF16 than weight-only UINT4.
Quantizing the text encoder is outside the paper's default layer policy.
This model inherits the FLUX Non-Commercial License from the source checkpoint.