WaveCut/FLUX.2-klein-9B-OrbitQuant-W4A4

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;
• 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.4.0"

OrbitQuant uses packed low-bit inference by default. CUDA first uses an importable native package and otherwise uses the Triton packed fallback. It does not silently materialize all weights in BF16. The native CUDA package can be built locally without publishing to Kernel Hub; see the OrbitQuant kernel instructions.

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")

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. 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	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.

L40S Benchmark

Both variants ran as separate processes on the same NVIDIA L40S with Torch 2.9.1+cu128, CUDA 12.8, Diffusers 0.39.0, Transformers 5.13.0, BF16 output arithmetic, no CPU offload, 1024x1024 output, four steps, guidance 1.0, seed 0, and the same ten prompts. OrbitQuant used PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True.

Variant	Load	Cold image	Hot mean	Hot median	Hot NVML peak	CUDA allocated peak	CUDA reserved peak
SDNQ UINT4	5.918 s	16.956 s	2.0885 s	2.0875 s	17.564 GB	14.844 GB	16.377 GB
OrbitQuant W4A4	2.543 s	4.193 s	2.0907 s	2.0920 s	15.731 GB	13.942 GB	14.544 GB

All 396 OrbitQuant projections used native token norm, RPBH/FWHT and codebook assignment followed by chunked packed-weight decode, CUTLASS INT8 matmul and a fused Triton epilogue. No full BF16/FP16 weight matrix was materialized. OrbitQuant's hot mean was within 0.11% of SDNQ while using 1.833 GB less peak NVML memory and 0.902 GB less CUDA allocated memory. Full machine-readable results are in benchmark/summary.json.

Metal Kernel Reference

The packed Metal operator was measured on an Apple M2 Max with Torch 2.12.1, FP16 activations, W4 weights, in_features=768, and out_features=2304.

Activation rows	Packed Metal	Resident FP16 F.linear	Materialize + F.linear	Packed vs materialize
1	0.0470 ms	0.0411 ms	0.2310 ms	4.92x
2	0.0439 ms	0.0422 ms	0.2371 ms	5.40x
3	0.0451 ms	0.0404 ms	0.2348 ms	5.20x
8	0.0420 ms	0.0503 ms	0.2512 ms	5.97x
9	0.0446 ms	0.0501 ms	0.2422 ms	5.43x
16	0.0459 ms	0.0581 ms	0.2512 ms	5.47x
31	0.0428 ms	0.0659 ms	0.2623 ms	6.12x

One-row projections use the SIMD-group packed matvec. Aligned FP16/BF16 projections with two or more rows use the padded matrix path. The packed weight payload, row norms and centroids occupy 25.26% of the materialized FP16 weight size for this shape.

For a complete 4096-coordinate RPBH block with constants resident on MPS, the 512-thread fused activation path measured 1.17-1.45x faster than the 256-thread path across 1-4096 activation rows. Smaller or multi-block RPBH dimensions retain 256 threads.

These are operator-level references, not a full FLUX.2 pipeline benchmark. Packed execution consumes low-bit indices directly and does not allocate a complete floating-point weight matrix.

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 corresponding L40S 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

• Optimized native CUDA measurements require a locally built matching ABI3 kernel package.
• Triton packed matmul is a compatible fallback, not the fastest measured backend.
• 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.

References

• OrbitQuant paper
• OrbitQuant library
• Source FLUX.2 Klein 9B model
• Controlled SDNQ comparison checkpoint