OrbitQuant

OrbitQuant is a calibration-free post-training quantizer for transformer linear projections. It implements the method from OrbitQuant: Data-Agnostic Quantization for Image and Video Diffusion Transformers and exposes it through Hugging Face Transformers, Diffusers, and a direct PyTorch API.

The implementation is clean-room and Apache-2.0 licensed.

Features

• Automatic coverage of registered linear-compatible modules in transformer backbones, independent of model class or modality.
• Built-in support for torch.nn.Linear and Hugging Face Conv1D projections.
• Public adapters for custom F.linear-equivalent module types and transposed weight layouts.
• RPBH rotation, exact unit-sphere Lloyd-Max codebooks, packed 2/3/4/6/8-bit weights, and online activation quantization without calibration data.
• Model-specific policies for paper-sensitive AdaLN and output-layer handling.
• Compact safetensors artifacts and Hugging Face save_pretrained() / from_pretrained() integration.
• Packed-weight CUDA, Triton, and Metal inference paths that avoid a full dequantized weight matrix.
• A fused DiT runtime for Krea 2 (orbitquant.runtime.krea2): grouped INT8 GEMMs with SwiGLU/residual epilogues and an INT8 attention kernel, 2.4x faster than the eager path on an RTX 4060 Ti.
• torch.compile(fullgraph=True) support: quantized forwards run behind a registered custom op with a fake implementation, so compiled models match eager output exactly. The forward is also CUDA-Graph-capturable; torch.compile(mode="reduce-overhead") removes the per-launch overhead that dominates small decode batches.

Embeddings, timestep modules, task heads, and common final projections are kept in source precision by default. Every automatic decision is available as a machine-readable inventory before quantization.

INT8 output heads (opt-in)

Output heads stay in source precision by default. For memory-bound decode the head can be the largest per-token read (a 150k-row BF16 vocabulary head is 0.6 GB per step), so OrbitQuant also offers an explicit per-row INT8 head:

from orbitquant import quantize_output_heads

quantize_output_heads(model, names=("lm_head",))   # in-place Int8RowLinear replacement

Int8RowLinear stores per-row absmax INT8 weights plus FP32 scales (2x smaller than BF16), quantizes activations per token at forward time, and serves 1..8-row decode with the native matmul_int8_rows DP4A kernel (CUDA package 1.0.5+), falling back to torch._int_mm for larger batches or when the native package is unavailable. All paths share the epilogue float(sum) * (x_scale * w_scale), so results are identical across them. Head quantization changes logits at the INT8 rounding level (~0.4% relative); measure sampled outputs before adopting it for a model.

Install

pip install "orbitquant[hf]"

PyTorch's Linux wheels supply the matching Triton build used by the CUDA fallback. Do not install or upgrade Triton independently: doing so can force a different PyTorch/CUDA stack. The legacy kernels extra remains accepted for command compatibility, so this is also valid:

pip install "orbitquant[hf,kernels]"

The optimized native kernel package (orbitquant_packed_matmul) is provisioned automatically the first time a packed runtime path needs it. How it works: the kernel binary must match the runtime exactly — torch minor and CUDA version for CUDA builds, torch stable ABI (any torch>=2.11) for CPU builds, and the OS and architecture everywhere — so OrbitQuant derives that variant name from the running process and resolves it in order: an installed orbitquant_packed_matmul package, a locally built variant (LOCAL_KERNELS), the on-disk cache, and finally the matching prebuilt variant wheel from this repository's kernels-v1 GitHub release (checksum-verified against the release manifest, then cached). Everything happens once per process at model load, never inside the forward path, and any failure falls back to the Triton (CUDA) or reference paths with an actionable message. Provision explicitly (or offline) with:

orbitquant kernels-install          # download the matching prebuilt variant
orbitquant kernels-install --build  # or compile from bundled sources (needs a toolchain)
orbitquant kernels-status           # inspect what the resolver would do

Set ORBITQUANT_KERNELS_AUTOFETCH=0 to forbid downloads, ORBITQUANT_KERNELS_AUTOBUILD=1 to allow automatic source builds, and ORBITQUANT_KERNELS_CACHE to relocate the cache (default ~/.cache/orbitquant/kernels). Runtimes without a matching variant fall back to Triton (CUDA) or the reference paths; see docs/kernel-audit.md for the full contract.

Quantize A Transformers Model

Importing orbitquant registers the backend with supported Transformers and Diffusers versions. The default target_policy="auto" selects a known paper policy where applicable and otherwise uses the universal policy.

import orbitquant
from transformers import AutoModel
from orbitquant import OrbitQuantConfig

model_id = "google/vit-base-patch16-224"
model = AutoModel.from_pretrained(
    model_id,
    quantization_config=OrbitQuantConfig(target_policy="auto"),
    low_cpu_mem_usage=True,
)

This is checkpoint-level conversion, not post-load module replacement. For each target weight OrbitQuant reads aligned row tiles from safetensors, closes that source mapping, and feeds a temporary packed shard back to the normal Transformers loader. Peak model-side memory is bounded by packed resident state, one source row tile, quantization workspace, intentionally skipped state, and runtime overhead. The full source transformer is never resident at once.

Save an offline packed artifact without loading the source model again:

model.save_pretrained("./model-orbitquant-w4a4", max_shard_size="2GB")

Load the packed artifact after importing the backend:

import orbitquant
from transformers import AutoModel

model = AutoModel.from_pretrained(
    "./model-orbitquant-w4a4",
    device_map="auto",
)

Pre-quantized Diffusers artifacts default their remaining compute modules to BF16 when no dtype is supplied. This keeps activations on the packed W4A4 CUDA path instead of the much slower FP32 compatibility path. An explicit dtype is always honored, and on-the-fly source conversion keeps the source dtype.

Named recipes are w4a4, w3a3, w2a4, w2a3, and w4a6. They create a normal OrbitQuantConfig, so every field can be overridden.

For an unknown architecture resolved to target_policy="universal", W2 uses mixed-bit protection by default: out-of-block and the first/last four indexed blocks use W4, separate interior Q/K projections remain W2, and the other interior projections use W3. inspect_linear_module_policy() reports the effective weight_bits for every module. Set lowbit_interior_protection=False to retain the legacy uniform-W2 interior; known paper-model policies are unchanged unless protection is explicitly forced with lowbit_interior_protection=True.

Bounded on-the-fly conversion requires safetensors and the Transformers 5 weight-conversion APIs. Local paths and Hub model IDs use the normal revision, variant, token/auth, device_map, max_memory, CPU, and disk offload arguments. Pickle/.bin checkpoints fail with an actionable error; use the explicit post-load helper only when accepting that it has no checkpoint-level memory guarantee.

Inspect Coverage

Inspect a model before replacing modules:

from orbitquant import inspect_linear_module_policy, recipe
from transformers import AutoModel

model = AutoModel.from_pretrained("your-org/your-transformer")
report = inspect_linear_module_policy(model, recipe("w4a4"))

print(report["action_counts"])
print(report["quantized_modules"])
print(report["skipped_modules"])
print(report["unsupported_linear_modules"])

The universal policy quantizes every registered linear-compatible module except known embeddings, timestep modules, task/output heads, and explicit skips. It does not depend on names such as layers, blocks, or a particular model class.

Use modules_to_convert as an allowlist and define AdaLN/skips with exact names, substrings, or glob patterns:

from orbitquant import OrbitQuantConfig

config = OrbitQuantConfig(
    modules_to_convert=["backbone.*.projection"],
    modules_to_use_adaln=["backbone.*.modulation"],
    modules_to_not_convert=["*.sensitive_output"],
)

Explicit dtype overrides remain available through modules_dtype_dict.

Quantize An Instantiated Module

For ordinary PyTorch models or frameworks that do not use Hugging Face loading hooks, or as an explicit compatibility fallback:

from orbitquant import quantize_model, recipe

summary = quantize_model(
    model,
    recipe("w4a4"),
    quantization_device="cuda",
)
print(summary.quantized_modules)

Here staging_mode="streaming" limits device staging per module only. The source model is already loaded, so this helper does not make a bounded host-RAM claim.

The replacement supports arbitrary leading dimensions and treats the final dimension as in_features, including sequence, image-token, and video-token layouts.

Custom Linear Modules

Register a module whose forward operation is equivalent to F.linear. The adapter describes only its source weight layout and feature attributes:

from orbitquant import register_linear_adapter

register_linear_adapter(
    MyLinear,
    weight_layout="in_out",
    in_features_attr="input_size",
    out_features_attr="output_size",
)

OrbitQuant stores every replacement in canonical [out_features, in_features] layout. Modules with additional routing, tensor-parallel communication, sparse expert selection, or non-linear forward semantics need an architecture-aware adapter; the inspection report lists unregistered linear candidates instead of silently replacing them.

Diffusers

Quantize the denoiser while loading an arbitrary compatible pipeline:

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

model_id = "black-forest-labs/FLUX.1-schnell"
qconfig = build_diffusers_pipeline_quantization_config(
    OrbitQuantConfig(target_policy="auto"),
    components="transformer",
)
pipe = DiffusionPipeline.from_pretrained(
    model_id,
    quantization_config=qconfig,
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()

Use sequential offload without an OrbitQuant-specific hook or preparation step:

pipe = DiffusionPipeline.from_pretrained(
    model_id,
    quantization_config=qconfig,
    torch_dtype=torch.bfloat16,
)
pipe.enable_sequential_cpu_offload()

Save the packed pipeline for later prequantized loading:

pipe.save_pretrained("./pipeline-orbitquant", safe_serialization=True)

The default component is only transformer; text encoders require an explicit component opt-in. Diffusers uses the same safetensors row-sliced conversion and normal PipelineQuantizationConfig loading path. Unknown architectures receive structural target_policy="auto" coverage, not a quality guarantee; inspect the policy and validate output quality before publishing an artifact.

Published FLUX, Z-Image, and Wan repositories are compact Diffusers component artifacts. Their model cards contain the matching pipeline code and validation evidence.

Load a published component artifact together with its recorded source pipeline:

import torch
from huggingface_hub import snapshot_download
from orbitquant import load_quantized_pipeline_from_artifact

artifact_dir = snapshot_download(
    "WaveCut/FLUX.1-schnell-OrbitQuant-W4A4",
    repo_type="model",
)
pipe = load_quantized_pipeline_from_artifact(
    artifact_dir,
    torch_dtype=torch.bfloat16,
    runtime_mode="auto_fused",
)

CUDA Decode

Version 0.9.8 includes native 1.0.5 register-codebook GEMV for SM89 (Ada) and SM120 GPUs. For row-major W4A4 inputs with one to eight rows, input width 1024–16384, and at least 2048 output features, it decodes arbitrary 16-entry signed INT8 codebooks with register byte permutations. Two rows share each decoded weight vector. The kernel needs no shared lookup table or unpacked weight cache. Other devices and shapes keep the existing native dispatch.

On RTX PRO 4500 Blackwell with Torch 2.10.0+cu128, interleaved full YuE2 songs measured 22.61 → 22.16 s at CFG 1 and 26.04 → 24.57 s at CFG 1.5 versus native 1.0.3. Each setting used one warmup followed by two timed runs per kernel, with INT8 KV cache, fused RMS quantization, eight ODE steps, and FP16 VAE. All ten PCM outputs matched their corresponding control song. These are warm-generation measurements on one GPU, not a universal gain or a fresh-process memory comparison.

The Python package and native kernel are separate installations. Managed caches are partitioned by the minimum native release (currently 1.0.5), so an upgrade does not silently reuse an older managed binary. Older cache directories are preserved. Downloaded wheels must satisfy that version floor and match the release checksum; a current cached variant remains usable offline. Provision the new variant before taking an upgraded installation offline.

Explicitly installed packages and LOCAL_KERNELS retain priority. Upgrade an explicit installation with the matching wheel from the kernels-v1 release, or build the bundled source with kernel-builder. orbitquant kernels-status reports the native version floor and managed cache selection. ORBITQUANT_W4A4_DISABLE_GEMV=1 selects the Tensor Core fallback.

Reproduce Load Memory

Version 0.9.4 fixes a model-lifetime leak in the registry used by torch.compile. Earlier versions registered every OrbitQuantLinear and RTNInt4Linear in a strong-reference dictionary, so discarded layers remained alive even without compilation. The registry now holds weak values and uses non-reusable handles; live compiled models keep working, while stale handles fail explicitly. This changes ownership, not quantization or kernel arithmetic.

The memory harness runs ordinary, on-the-fly, and prequantized loads in separate processes and reports RSS separately from mmap virtual size:

python scripts/measure_streaming_load_memory.py \
  --framework transformers \
  --model-id google/vit-base-patch16-224 \
  --prequantized-model-id ./vit-base-orbitquant \
  --torch-dtype bfloat16 \
  --output ./memory-results.json

One CPU/MPS host run with Torch 2.12.1 measured:

The streaming run processed 324.0 MiB of selected source tensors with a 9.0 MiB largest source tensor, retained no full dequantized cache, and reported no source-release failure. CUDA allocated/reserved and NVML fields are emitted when CUDA and NVML are available; they were unavailable for this Mac run.

Packed Runtime

runtime_mode="auto_fused" is the default:

CUDA and MPS do not silently materialize a full BF16/FP16 weight matrix in auto_fused. If no packed backend is available, the error includes the missing backend and installation guidance. CPU retains a compatibility fallback when the optional native package is absent.

End-to-end FLUX.2 Klein 9B numbers for the released package are recorded in docs/flux2-klein-9b-sdnq-vs-orbitquant.md: on a single RTX PRO 6000, OrbitQuant 0.5.0's hot median (1.164 s at 1024x1024/4 steps) is 3.2% faster than unquantized BF16 and 8.6% faster than SDNQ UINT4, with 12.51 GB peak CUDA allocated memory versus BF16's 34.72 GB.

Use the explicit reference path for compatibility or numerical debugging:

config = orbitquant.recipe("w4a4", runtime_mode="dequant_bf16")

On CUDA compute capability 8.0 or newer, the W4A4 fast path fuses token norm, RPBH/FWHT, and codebook assignment in the native package, decodes only a bounded output-channel chunk of packed weights to INT8, and uses the Torch CUTLASS tensor-core matmul. It never materializes the full BF16/FP16 weight matrix. The existing direct packed CUDA MMA kernel remains the fallback for compatible W4A4 shapes when CUTLASS INT8 matmul is unavailable.

The optimized CUDA path maps the fixed Lloyd-Max centroids to a symmetric INT8 surrogate plus one scalar per codebook. Packed checkpoint indices and artifact size are unchanged. Use dequant_bf16 when exact Lloyd-Max centroid evaluation is required.

Build the ABI3 native package locally without Kernel Hub:

cd native-kernels/orbitquant-packed-matmul
nix --option sandbox relaxed run .#build-and-copy -L
export PYTHONPATH="$PWD/build/<matching-torch-backend-platform-variant>:$PYTHONPATH"

PyTorch 2.9 CUDA users can reduce allocator reservation during native diffusion inference by setting the allocator before Python starts:

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py

CUDA and Metal variants must match the Torch/backend tuple. The CPU variant targets the LibTorch Stable ABI 2.11, but remains specific to its operating system, C++ runtime, and architecture. See docs/kernel-audit.md for tested shapes, benchmark methodology, and local package verification.

Fused DiT Runtime (Krea 2)

Version 0.10 speeds up every W4A4 model on CUDA without API changes:

• Triton kernels take row counts as runtime arguments, so a new prompt length or image size no longer recompiles them.
• Large-row W4A4 layers use an INT8 GEMM that applies the scale epilogue in registers; the output is bit-identical to the previous torch._int_mm path.
• Projections that read the same tensor (attention Q/K/V, SwiGLU gate/up) share one activation quantization. The reuse needs version counters, so run inference under torch.no_grad(); torch.inference_mode() disables it.

On an RTX 4060 Ti the Qwen-Image 2.1 W4A4 pipeline went from 15.0 to 10.3 s per 1024x1024 image with pixel-identical output.

orbitquant.runtime.krea2 goes further for Krea 2 Turbo: each block runs one grouped Q|K|V|gate GEMM with the sigmoid gate in its epilogue, a fused Q/K RMSNorm + RoPE kernel, SwiGLU and gated residual updates in GEMM epilogues, W8A8 down projections, and optionally an INT8 Q.K^T / FP16 P.V attention kernel (the SageAttention v1 scheme). Blocks whose Q/K RMSNorm scales one channel far above the rest keep BF16 attention.

import torch
from diffusers import Krea2Pipeline
from orbitquant.runtime.krea2 import Krea2FastRunner, install, save_fused

pipe = Krea2Pipeline.from_pretrained(model_dir, dtype=torch.bfloat16)
fused = install(pipe.transformer, attention="int8")  # on the host, before .to("cuda")
save_fused(fused, "fused_blocks.safetensors")  # optional: map it next time
pipe.to("cuda")
runner = Krea2FastRunner(pipe, fused)
with torch.no_grad():
    states = runner.encode(prompt, torch.device("cuda"))
    latents = runner.denoise(states, width=1024, height=1024)
    image = runner.decode(latents, width=1024, height=1024)

Building the fused buffers copies the block weights into new host memory; install(..., fused_path="fused_blocks.safetensors") maps a saved file instead. On an RTX 4060 Ti a 1024x1024 eight-step image takes 10.8 s (25.9 s with the stock pipeline), and the DiT needs 7.3 GB instead of 10 GB.

Validated Architecture Coverage

The integration suite instantiates and inventories encoder-only, decoder-only, encoder-decoder, causal LM, and vision transformer families:

The paper-aligned release artifacts remain FLUX.1-schnell, Z-Image-Turbo, and Wan 2.1 T2V. FLUX.2 Klein is an additional validated diffusion target.

Architecture coverage means the model can be discovered, quantized, executed, saved, and restored through the same public quantization API. It does not guarantee a quality-preserving bit setting. OrbitQuant was evaluated in the paper on image and video diffusion transformers; language and classification models can be more sensitive, and their quality must be measured before publishing a checkpoint. The library exposes module overrides for such recipes but does not silently substitute a different quantization algorithm.

Method Conformance

The implementation follows the paper's shared data-agnostic basis:

• RPBH permutation, Rademacher signs, block FWHT, and orthonormal scaling.
• Offline folded weight rotation with BF16 row norms and quantized unit directions.
• Online per-token norm, normalized activation rotation, nearest-centroid quantization, and rescaling.
• One fixed codebook per (input dimension, bit width, algorithm version) and no prompt, timestep, or calibration statistics.
• INT4 group-64 RTN for model policies that identify dynamic AdaLN projections.

The detailed requirement matrix is in docs/paper-methodology-audit.md.

Citation

If you use OrbitQuant in research, cite the original paper describing the method:

Donghyun Lee, Jitesh Chavan, Duy Nguyen, Sam Huang, Liming Jiang, Priyadarshini Panda, Timo Mertens, and Saurabh Shukla. "OrbitQuant: Data-Agnostic Quantization for Image and Video Diffusion Transformers." arXiv:2607.02461, 2026. https://doi.org/10.48550/arXiv.2607.02461

@misc{lee2026orbitquant,
  title         = {OrbitQuant: Data-Agnostic Quantization for Image and Video Diffusion Transformers},
  author        = {Lee, Donghyun and Chavan, Jitesh and Nguyen, Duy and Huang, Sam and Jiang, Liming and Panda, Priyadarshini and Mertens, Timo and Shukla, Saurabh},
  year          = {2026},
  eprint        = {2607.02461},
  archiveprefix = {arXiv},
  primaryclass  = {cs.CV},
  doi           = {10.48550/arXiv.2607.02461},
  url           = {https://arxiv.org/abs/2607.02461},
}

Development

uv sync --extra hf --extra dev
uv run pytest -q
uv run ruff check .

License

OrbitQuant is licensed under Apache-2.0. Model artifacts retain the license and provenance of their source checkpoints.