onecomp-runtime

Shared int4 inference runtime for OneCompression packed checkpoints. OneCompression produces GPTQ/RTN-packed safetensors; this is the consumer side — the int4 GEMM layers, GPTQ unpack/dequant helpers, backend selection, and a generic diffusion loader that every per-model runtime builds a thin adapter on top of.

pip install onecomp-runtime              # import as onecomp_runtime
pip install onecomp-runtime[gemlite]     # + GemLite int4 kernels
pip install onecomp-runtime[diffusion]   # + diffusers (generic loader builds diffusers classes)

Why

The int4 leaf machinery was copy-pasted across the FLUX.2 / LTX-2.3 / FireRed / Irodori runtimes — fused_int4_linear.py was byte-identical in three of them. A fix to the kernel (K-padding, warmup buckets, dtype safety) had to be hand- propagated to every repo. This package is the single source of truth.

Layout

onecomp_runtime/
  layers/
    fused_int4_linear.py   # Triton dequant+GEMM (AutoGPTQ-v1, gs=32)
    gemlite_int4_linear.py # GemLite kernel wrapper (fp16 I/O)
    packed_linear.py       # PackedRTNLinear, PackedEmbedding (RTN uint8-nibble)
    packed_conv.py         # int4 Conv1d / ConvTranspose1d (DAC-VAE codecs)
  quant_utils.py           # GPTQ + RTN unpack/dequant helpers
  backend.py               # resolve_backend / can_use_fused / build_{gemlite,fused,eager}
  diffusion.py             # load_int4_model(build_meta_model, ...) — generic GPTQ loader

Usage — a per-model runtime adapter

from onecomp_runtime.diffusion import load_int4_model
from diffusers import Flux2Transformer2DModel

def load_int4_transformer(path, **kw):
    return load_int4_model(
        path,
        lambda cfg: Flux2Transformer2DModel.from_config(cfg),
        label="flux2-klein-lite",
        **kw,
    )

The only per-model code is build_meta_model (construct the bare module from the checkpoint's config_json) and an optional post_load(model) hook for buffer fixups (e.g. FireRed/Qwen-Image RoPE tables that meta-init leaves uninitialised).

Backends

backend="auto" → gemlite if importable, else fused. bf16 is the safe default for Qwen-Image / LTX (fp16 overflows to NaN); fp16 is only required on the GemLite path.

Checkpoint contract

A single safetensors with metadata keys config_json, quant_layers_json (per-layer manifest: name, wbits, groupsize, actorder, in_features, out_features), and checkpoint_format (gptq v1 / gptq_v2). The RTN tier (packed_linear, packed_conv) consumes the encoder/embedding/conv extras that Irodori-style checkpoints add — those runtimes drive the layers directly rather than through load_int4_model.