Fuse the transformer blocks of six diffusion families into stored checkpoints (0.11.0)

README.md

@@ -22,9 +22,10 @@ The implementation is clean-room and Apache-2.0 licensed.
   `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.
+- Fused transformer checkpoints for Krea 2, FLUX.2, Ideogram 4, Qwen-Image 2.1,
+  MiniMax-H3 and Boogu-Image (`orbitquant.fused`): each block runs grouped INT8
+  GEMMs with norm/modulation prologues and SwiGLU/gated-residual epilogues plus
+  an INT8 attention kernel, and the checkpoint stores only the fused groups.
 - `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;
@@ -446,48 +447,75 @@ system, C++ runtime, and architecture. See
 [`docs/kernel-audit.md`](docs/kernel-audit.md) for tested shapes, benchmark
 methodology, and local package verification.
 
-## Fused DiT Runtime (Krea 2)
+## Fused Transformers
 
-Version 0.10 speeds up every W4A4 model on CUDA without API changes:
+Version 0.11 runs the transformer blocks of supported diffusion models fused:
+the projections that read one input (attention Q|K|V, SwiGLU gate|up, the two
+streams of a joint attention) form one weight matrix and one INT8-surrogate
+GEMM, the RMSNorm/LayerNorm and AdaLN modulation in front of them run in the
+activation-quantization prologue, and SwiGLU, sigmoid gates and gated residual
+updates run in the GEMM epilogue. Q/K normalization and RoPE are one kernel,
+and attention uses an INT8 Q.K^T / FP16-accumulated P.V kernel (the
+SageAttention v1 scheme); blocks whose Q/K RMSNorm scales one channel far above
+the rest keep BF16 attention.
 
-- 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.
+A fused checkpoint stores only the fused groups and records the layout in its
+`quantization_config`, so the usual loader returns the fused transformer:
 
-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.
+```python
+import torch
+import orbitquant  # registers the loader
+from diffusers import Krea2Pipeline
+
+pipe = Krea2Pipeline.from_pretrained(
+    "WaveCut/Krea-2-Turbo-OrbitQuant-W4A4", torch_dtype=torch.bfloat16
+).to("cuda")
+```
 
-`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.
+Fuse a per-projection checkpoint on the host and save it:
 
 ```python
 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)
+from diffusers import Flux2Transformer2DModel
+from orbitquant import fused
+
+transformer = Flux2Transformer2DModel.from_pretrained(
+    model_dir, subfolder="transformer", torch_dtype=torch.bfloat16
+)
+fused.fuse(transformer)  # attention="flash" keeps BF16 attention everywhere
+fused.save_pretrained(transformer, "fused/transformer")
 ```
 
-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.
+`orbitquant.fused.fuse_component_artifact(model, source_dir, output_dir)` does
+the same for an `orbitquant-v1` component artifact, and
+`load_orbitquant_artifact` loads the result fused.
+
+| Family | Transformer class | Fused blocks |
+| --- | --- | --- |
+| Krea 2 | `Krea2Transformer2DModel` | Q\|K\|V\|gate, gated output, SwiGLU; W8A8 down projections |
+| FLUX.2 | `Flux2Transformer2DModel` | double-stream (per-stream Q\|K\|V, gated residuals) and single-stream blocks; reference-image KV cache |
+| Ideogram 4 | `Ideogram4Transformer2DModel` | sandwich-norm blocks with row-indexed modulation |
+| Qwen-Image 2.1 | `QwenImage21Transformer2DModel` | block-causal prefill and KV-cached decode |
+| MiniMax-H3 | `MiniMaxH3Transformer3DModel` | per-row (timestep, modality) AdaLN tables, partial rotate-half RoPE |
+| Boogu-Image | `BooguImageTransformer2DModel` | refiner, single- and double-stream blocks; 8-bit activations |
+
+Groups take the activation width of their projections unless the layout
+overrides it: `fuse(..., activation_bits={"qkv": 8})` gives a group 8-bit
+activations (per-token absmax INT8 of the rotated input). Boogu-Image runs all
+its groups this way because its 3360 channels only allow a 32-wide rotation
+block; the GEMMs stay INT8, so the wider codes cost almost no speed. Low-bit
+artifacts with mixed weight widths (boundary and interior protection) fuse into
+row segments of one GEMM.
+
+The fused runtime needs CUDA and Triton. Run inference under `torch.no_grad()`
+or `torch.inference_mode()`. Fused weights are frozen parameters, so model CPU
+offload and Diffusers block-level group offloading (including streamed
+`use_stream=True` prefetching) move them with their blocks; leaf-level group
+offloading does not, because a fused block reads its groups directly instead of
+calling them. Stream without `record_stream`: recorded streams keep every
+streamed block reserved until the compute stream has used it, which grew
+MiniMax-H3's GPU memory to its 12 GiB cap (5.1 GiB without).
+`orbitquant.runtime.krea2` from 0.10 remains for existing deployments.
 
 ## Validated Architecture Coverage
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "orbitquant"
-version = "0.10.0"
+version = "0.11.0"
 description = "Calibration-free OrbitQuant for transformer linear projections"
 readme = "README.md"
 requires-python = ">=3.11"

src/orbitquant/__init__.py

@@ -24,7 +24,7 @@
 from orbitquant.quantizer import register_hf_quantizers
 from orbitquant.recipes import recipe
 
-__version__ = "0.10.0"
+__version__ = "0.11.0"
 
 register_hf_quantizers()
 

src/orbitquant/artifacts/checksums.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import hashlib
+from collections.abc import Iterable
 from pathlib import Path
 
 _IGNORED_ARTIFACT_PATH_PARTS = {
@@ -25,6 +26,16 @@ def sha256_file(path: str | Path) -> str:
     return digest.hexdigest()
 
 
+def existing_file_checksums(root: str | Path, relative_paths: Iterable[str]) -> dict[str, str]:
+    """SHA-256 of each listed path that is a file under ``root``; missing paths are left out."""
+    root_path = Path(root)
+    return {
+        relative_path: sha256_file(root_path / relative_path)
+        for relative_path in relative_paths
+        if (root_path / relative_path).is_file()
+    }
+
+
 def write_sha256sums(root: str | Path, output: str | Path | None = None) -> Path:
     root_path = Path(root)
     output_path = root_path / "SHA256SUMS" if output is None else Path(output)

src/orbitquant/artifacts/loader.py

@@ -106,11 +106,20 @@ def load_orbitquant_artifact(
         parent, child_name = _parent_and_child(model, name)
         _set_child(parent, child_name, replacement)
 
+    if config.fused_layout:
+        from orbitquant.fused import prepare_skeleton
+
+        prepare_skeleton(model, config.fused_layout, config=config)
+
     state_dict = load_file(
         artifact_path / "model.safetensors",
         device=_safetensors_load_device(device, model),
     )
     missing, unexpected = model.load_state_dict(state_dict, strict=strict)
     if strict and (missing or unexpected):
         raise RuntimeError(f"artifact state mismatch: missing={missing}, unexpected={unexpected}")
+    if config.fused_layout:
+        from orbitquant.fused import finalize
+
+        finalize(model)
     return manifest

src/orbitquant/config.py

@@ -112,6 +112,9 @@ class OrbitQuantConfig(QuantizationConfigMixin):
     # Known model policies keep their existing paper-aligned recipes. True
     # forces this policy for any target; False disables it.
     lowbit_interior_protection: bool | str = "auto"
+    # Fused-block checkpoints (orbitquant.fused): family, layout version and group kinds. The
+    # loader rebuilds the fused groups from it before the weights load.
+    fused_layout: dict[str, Any] | None = None
 
     def __post_init__(self) -> None:
         if self.weight_bits not in _SUPPORTED_BITS:

src/orbitquant/fused/__init__.py

@@ -0,0 +1,32 @@
+"""Fused transformer blocks for OrbitQuant checkpoints.
+
+``fuse(transformer)`` regroups each block's OrbitQuant projections (Q|K|V, gate|up, ...) into
+matrices that run as one INT8-surrogate GEMM with the following elementwise work in its
+epilogue, and replaces the block forward accordingly; the transformer stays a drop-in module
+for its Diffusers pipeline. Fused checkpoints store only the fused groups and load directly
+into the fused form. The runtime needs the CUDA Triton backend; importing this package does
+not, so the quantizer can recognize fused checkpoints anywhere.
+"""
+
+from importlib import import_module
+
+_EXPORTS = {
+    "FusedRuntime": "orbitquant.fused.runtime",
+    "finalize": "orbitquant.fused.install",
+    "fuse": "orbitquant.fused.install",
+    "fuse_component_artifact": "orbitquant.fused.convert",
+    "prepare_skeleton": "orbitquant.fused.install",
+    "release": "orbitquant.fused.install",
+    "save_pretrained": "orbitquant.fused.install",
+}
+
+__all__ = sorted(_EXPORTS)
+
+
+def __getattr__(name):
+    module = _EXPORTS.get(name)
+    if module is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(import_module(module), name)
+    globals()[name] = value
+    return value

src/orbitquant/fused/convert.py

@@ -0,0 +1,86 @@
+"""Fused copies of published ``orbitquant-v1`` component artifacts.
+
+``fuse_component_artifact`` rewrites such an artifact so that ``model.safetensors`` holds the
+fused groups instead of the per-projection tensors and ``quantization_config.json`` carries the
+``fused_layout``; every other file is copied. The manifest keeps the projection list and the
+per-module widths the loader rebuilds the groups from. Diffusers-format components need no
+helper: load the transformer with ``from_pretrained``, ``fuse`` it and save it with
+``orbitquant.fused.save_pretrained``.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+import shutil
+from pathlib import Path
+
+from safetensors.torch import save_file
+from torch import nn
+
+from orbitquant.artifacts.checksums import (
+    existing_file_checksums,
+    is_ignored_artifact_relative_path,
+    write_sha256sums,
+)
+from orbitquant.artifacts.loader import load_orbitquant_artifact
+from orbitquant.artifacts.manifest import OrbitQuantManifest
+from orbitquant.config import OrbitQuantConfig
+from orbitquant.fused.install import fuse
+
+_REWRITTEN = ("model.safetensors", "quantization_config.json", "orbitquant_manifest.json")
+
+
+def fuse_component_artifact(
+    model: nn.Module,
+    source_dir: str | Path,
+    output_dir: str | Path,
+    *,
+    dense: str | None = None,
+    validate_checksums: bool = True,
+) -> OrbitQuantManifest:
+    """Load the artifact at ``source_dir`` into ``model`` (the unquantized component skeleton),
+    fuse it and write the fused artifact to ``output_dir``. Rewrite the model card and refresh
+    ``SHA256SUMS`` (``write_sha256sums``) after editing it."""
+    source, output = Path(source_dir), Path(output_dir)
+    if output.resolve() == source.resolve():
+        raise ValueError("write the fused artifact to a new directory")
+    config = OrbitQuantConfig.from_dict(
+        json.loads((source / "quantization_config.json").read_text(encoding="utf-8"))
+    )
+    if config.fused_layout:
+        raise ValueError(f"{source} is already fused")
+    manifest = load_orbitquant_artifact(model, source, validate_checksums=validate_checksums)
+    config.fused_layout = fuse(model, dense=dense, config=config)
+
+    output.mkdir(parents=True, exist_ok=True)
+    for path in sorted(source.rglob("*")):
+        relative = path.relative_to(source).as_posix()
+        if (
+            path.is_dir()
+            or relative in _REWRITTEN
+            or relative == "SHA256SUMS"
+            or is_ignored_artifact_relative_path(relative)
+        ):
+            continue
+        target = output / relative
+        target.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(path, target)
+
+    state = {name: value.detach().cpu().contiguous() for name, value in model.state_dict().items()}
+    save_file(state, output / "model.safetensors")
+    (output / "quantization_config.json").write_text(
+        json.dumps(config.to_dict(), indent=2) + "\n", encoding="utf-8"
+    )
+    fused_manifest = dataclasses.replace(
+        manifest,
+        module_shapes={name: list(tensor.shape) for name, tensor in state.items()},
+        # A published manifest can list files that a later repository edit removed (accepted
+        # with validate_checksums=False); the fused manifest lists what the fused copy ships.
+        checksums=existing_file_checksums(output, manifest.checksums),
+    )
+    (output / "orbitquant_manifest.json").write_text(
+        json.dumps(fused_manifest.to_dict(), indent=2) + "\n", encoding="utf-8"
+    )
+    write_sha256sums(output)
+    return fused_manifest

src/orbitquant/fused/families/__init__.py

@@ -0,0 +1,74 @@
+"""Block adapters: which projections each block fuses and the forward that runs them.
+
+An adapter module defines:
+
+* ``FAMILY``: the name recorded in fused checkpoints;
+* ``MODEL_CLASSES``: transformer class names it serves;
+* ``blocks(model)``: ``(path, kind)`` of every fused block;
+* ``groups(kind, block)``: ``{name: GroupSpec}`` for a block of that kind;
+* ``forward(kind)``: the replacement ``forward`` (same signature as the original block);
+* ``prepare(block, kind)``: per-block settings derived from the loaded weights (optional);
+* ``install(model)``: model-level changes outside the blocks (optional).
+"""
+
+from __future__ import annotations
+
+import importlib
+from dataclasses import dataclass, field
+
+from orbitquant.fused.groups import RowSource
+
+_MODULES = ("krea2", "flux2", "ideogram4", "minimax_h3", "qwenimage21", "boogu")
+
+
+@dataclass(frozen=True)
+class GroupSpec:
+    sources: tuple[RowSource, ...]
+    swiglu: bool = False
+    # Dense (unquantized) sources become per-row INT8 unless the checkpoint asks for BF16.
+    dense: str | None = None
+    extra: dict = field(default_factory=dict)
+
+
+# INT8 Q/K quantization loses the small channels of heads whose Q/K RMSNorm scales one channel
+# far above the rest (Krea 2 block 0: 36x the median); such blocks keep BF16 attention.
+INT8_ATTENTION_MAX_NORM_SPREAD = 3.0
+
+
+def int8_attention_allowed(*norm_weights, offset: float = 0.0) -> bool:
+    """Whether a block's Q/K RMSNorm scales (``weight + offset``) are flat enough for INT8 Q/K."""
+    import torch
+
+    scales = torch.cat([w.reshape(-1) for w in norm_weights]).float().add(offset).abs()
+    return bool(scales.max() <= INT8_ATTENTION_MAX_NORM_SPREAD * scales.median())
+
+
+def specs(**groups) -> dict[str, GroupSpec]:
+    """``specs(qkv=["attn.to_q", ...], gateup=(["ff.gate", "ff.up"], {"swiglu": True}))``."""
+    out = {}
+    for name, value in groups.items():
+        options = {}
+        if isinstance(value, tuple):
+            value, options = value
+        sources = tuple(s if isinstance(s, RowSource) else RowSource(s) for s in value)
+        out[name] = GroupSpec(sources, **options)
+    return out
+
+
+def adapter_for_class(class_name: str):
+    for name in _MODULES:
+        try:
+            module = importlib.import_module(f"orbitquant.fused.families.{name}")
+        except ModuleNotFoundError:
+            continue
+        if class_name in module.MODEL_CLASSES:
+            return module
+    raise ValueError(f"no fused block adapter for {class_name}")
+
+
+def adapter_for_family(family: str):
+    try:
+        module = importlib.import_module(f"orbitquant.fused.families.{family}")
+    except ModuleNotFoundError as exc:
+        raise ValueError(f"unknown fused family {family!r}") from exc
+    return module