Metadata-Version: 2.4
Name: onecomp-runtime
Version: 0.1.0
Summary: Shared int4 inference runtime (fused Triton / GemLite / eager) for OneCompression GPTQ-packed checkpoints.
License: MIT
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.4
Requires-Dist: triton>=3.0
Requires-Dist: safetensors>=0.4
Provides-Extra: gemlite
Requires-Dist: gemlite>=0.5; extra == "gemlite"
Provides-Extra: diffusion
Requires-Dist: diffusers>=0.37; extra == "diffusion"
Requires-Dist: huggingface-hub>=0.24; extra == "diffusion"
Dynamic: license-file

# onecomp-runtime

Shared **int4 inference runtime** for [OneCompression](../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

```python
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 | kernel | I/O dtype | when |
|---|---|---|---|
| `gemlite` | GemLite Triton int4 | fp16 only | large M (FLUX), if installed |
| `fused` | bundled Triton dequant+GEMM | fp16/bf16/fp32 | default; bf16-safe |
| `eager` | dequant once to `nn.Linear` | any | groupsize≠32, actorder, odd shapes |

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