From a0529a1cec6996c062ed449e7cf245cc78f2c795 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Fri, 2 Oct 2026 02:24:34 +0200 Subject: [PATCH 1/4] Fuse the transformer blocks of six diffusion families into stored checkpoints (0.11.0) orbitquant.fused groups the projections that read one input into one INT8-surrogate GEMM per block (Q|K|V, SwiGLU gate|up, both streams of a joint attention), runs norm and AdaLN modulation in the activation-quantization prologue and SwiGLU, sigmoid gates and gated residual updates in the GEMM epilogue, fuses Q/K normalization with RoPE and uses the INT8 Q.K^T / FP16 P.V attention kernel where a block's Q/K norm weights allow it. Families: Krea 2, FLUX.2, Ideogram 4, Qwen-Image 2.1, MiniMax-H3 and Boogu-Image. A fused checkpoint stores only the fused groups and records the layout in its quantization_config; from_pretrained and load_orbitquant_artifact rebuild the groups before the weights load. Groups may mix weight widths (low-bit boundary/interior protection) as row segments of one GEMM and may take 8-bit activations (per-token absmax INT8 of the rotated input). The activation quantizer covers RPBH blocks up to 16384. --- README.md | 94 +-- pyproject.toml | 2 +- src/orbitquant/__init__.py | 2 +- src/orbitquant/artifacts/loader.py | 9 + src/orbitquant/config.py | 3 + src/orbitquant/fused/__init__.py | 32 + src/orbitquant/fused/convert.py | 84 +++ src/orbitquant/fused/families/__init__.py | 73 +++ src/orbitquant/fused/families/boogu.py | 357 ++++++++++ src/orbitquant/fused/families/flux2.py | 346 ++++++++++ src/orbitquant/fused/families/ideogram4.py | 139 ++++ src/orbitquant/fused/families/krea2.py | 124 ++++ src/orbitquant/fused/families/minimax_h3.py | 145 +++++ src/orbitquant/fused/families/qwenimage21.py | 230 +++++++ src/orbitquant/fused/groups.py | 377 +++++++++++ src/orbitquant/fused/install.py | 223 +++++++ src/orbitquant/fused/runtime.py | 220 +++++++ src/orbitquant/kernels/triton_activation.py | 648 ++++++++++++++++--- src/orbitquant/kernels/triton_attention.py | 41 +- src/orbitquant/kernels/triton_int8_gemm.py | 37 +- src/orbitquant/quantizer.py | 16 + tests/test_fused_family_kernels.py | 302 +++++++++ tests/test_fused_layout.py | 25 + uv.lock | 2 +- 24 files changed, 3383 insertions(+), 148 deletions(-) create mode 100644 src/orbitquant/fused/__init__.py create mode 100644 src/orbitquant/fused/convert.py create mode 100644 src/orbitquant/fused/families/__init__.py create mode 100644 src/orbitquant/fused/families/boogu.py create mode 100644 src/orbitquant/fused/families/flux2.py create mode 100644 src/orbitquant/fused/families/ideogram4.py create mode 100644 src/orbitquant/fused/families/krea2.py create mode 100644 src/orbitquant/fused/families/minimax_h3.py create mode 100644 src/orbitquant/fused/families/qwenimage21.py create mode 100644 src/orbitquant/fused/groups.py create mode 100644 src/orbitquant/fused/install.py create mode 100644 src/orbitquant/fused/runtime.py create mode 100644 tests/test_fused_family_kernels.py create mode 100644 tests/test_fused_layout.py diff --git a/README.md b/README.md index bbbd8fb..03f024f 100644 --- a/README.md +++ b/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,69 @@ 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 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()`. `orbitquant.runtime.krea2` from 0.10 remains for +existing deployments. ## Validated Architecture Coverage diff --git a/pyproject.toml b/pyproject.toml index 5d60b16..561d981 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/orbitquant/__init__.py b/src/orbitquant/__init__.py index 013187c..f33e0db 100644 --- a/src/orbitquant/__init__.py +++ b/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() diff --git a/src/orbitquant/artifacts/loader.py b/src/orbitquant/artifacts/loader.py index e89bdc8..60df220 100644 --- a/src/orbitquant/artifacts/loader.py +++ b/src/orbitquant/artifacts/loader.py @@ -106,6 +106,11 @@ 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), @@ -113,4 +118,8 @@ def load_orbitquant_artifact( 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 diff --git a/src/orbitquant/config.py b/src/orbitquant/config.py index 9072024..bab0969 100644 --- a/src/orbitquant/config.py +++ b/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: diff --git a/src/orbitquant/fused/__init__.py b/src/orbitquant/fused/__init__.py new file mode 100644 index 0000000..58d4de6 --- /dev/null +++ b/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 diff --git a/src/orbitquant/fused/convert.py b/src/orbitquant/fused/convert.py new file mode 100644 index 0000000..d2ea130 --- /dev/null +++ b/src/orbitquant/fused/convert.py @@ -0,0 +1,84 @@ +"""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 ( + is_ignored_artifact_relative_path, + sha256_file, + 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()}, + checksums={name: sha256_file(output / name) for name in 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 diff --git a/src/orbitquant/fused/families/__init__.py b/src/orbitquant/fused/families/__init__.py new file mode 100644 index 0000000..819b6de --- /dev/null +++ b/src/orbitquant/fused/families/__init__.py @@ -0,0 +1,73 @@ +"""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). +""" + +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 diff --git a/src/orbitquant/fused/families/boogu.py b/src/orbitquant/fused/families/boogu.py new file mode 100644 index 0000000..b947972 --- /dev/null +++ b/src/orbitquant/fused/families/boogu.py @@ -0,0 +1,357 @@ +"""Boogu-Image (``BooguImageTransformer2DModel``) blocks. + +Refiner and single-stream blocks share the Lumina recipe: RMSNorm x (1 + scale) prologue, one +Q|K|V GEMM for grouped-query attention (28/7 heads of 120 channels, zero-padded to 128 for the +INT8 kernel), interleaved complex RoPE, sandwich post-norms with tanh gates as row kernels and a +SwiGLU feed-forward. Double-stream blocks run a joint instruction|image attention with +per-stream Q|K|V and output projections, then an image self-attention; their MLP inputs are +modulated norms of the block input, not of the attention-updated streams. + +Attention masks only mark the padding of batched packed sequences. Each sample runs alone, so +padded keys are dropped instead and attention stays mask-free. + +Every group takes 8-bit activations (per-token absmax INT8 of the rotated input): at 3360 +channels the RPBH block is 32 wide, too narrow to spread outliers for 4-bit codes. On the Q|K|V +inputs the 4-bit error turns into knit-like texture across the image; with 8 bits everywhere the +images stay closer to BF16 than the SDNQ W4A8 checkpoint, at the speed of 4-bit codes (the GEMMs +are INT8 either way). +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from orbitquant.fused.families import int8_attention_allowed, specs +from orbitquant.kernels import triton_activation as act +from orbitquant.kernels import triton_int8_gemm as gemm + +FAMILY = "boogu" +MODEL_CLASSES = ("BooguImageTransformer2DModel",) + +_SINGLE_LISTS = ("noise_refiner", "ref_image_refiner", "context_refiner", "single_stream_layers") +_JOINT = "img_instruct_attn.processor" +ACTIVATION_BITS = dict.fromkeys( + ( + "qkv", + "out", + "ff_in", + "ff_out", + "joint_img_qkv", + "joint_txt_qkv", + "joint_img_out", + "joint_txt_out", + "joint_out", + "self_qkv", + "self_out", + "img_ff_in", + "img_ff_out", + "txt_ff_in", + "txt_ff_out", + ), + 8, +) + + +def blocks(model): + paths = [ + (f"{name}.{i}", "single") + for name in _SINGLE_LISTS + for i in range(len(getattr(model, name, ()))) + ] + return paths + [ + (f"double_stream_layers.{i}", "double") for i in range(len(model.double_stream_layers)) + ] + + +def groups(kind, block): + if kind == "single": + return specs( + qkv=["attn.to_q", "attn.to_k", "attn.to_v"], + out=["attn.to_out.0"], + ff_in=(["feed_forward.linear_1", "feed_forward.linear_3"], {"swiglu": True}), + ff_out=["feed_forward.linear_2"], + ) + return specs( + joint_img_qkv=[f"{_JOINT}.img_to_q", f"{_JOINT}.img_to_k", f"{_JOINT}.img_to_v"], + joint_txt_qkv=[ + f"{_JOINT}.instruct_to_q", + f"{_JOINT}.instruct_to_k", + f"{_JOINT}.instruct_to_v", + ], + joint_img_out=[f"{_JOINT}.img_out"], + joint_txt_out=[f"{_JOINT}.instruct_out"], + joint_out=["img_instruct_attn.to_out.0"], + self_qkv=["img_self_attn.to_q", "img_self_attn.to_k", "img_self_attn.to_v"], + self_out=["img_self_attn.to_out.0"], + img_ff_in=(["img_feed_forward.linear_1", "img_feed_forward.linear_3"], {"swiglu": True}), + img_ff_out=["img_feed_forward.linear_2"], + txt_ff_in=( + ["instruct_feed_forward.linear_1", "instruct_feed_forward.linear_3"], + {"swiglu": True}, + ), + txt_ff_out=["instruct_feed_forward.linear_2"], + ) + + +def prepare(block, kind): + attns = [block.attn] if kind == "single" else [block.img_instruct_attn, block.img_self_attn] + weights = [w for attn in attns for w in (attn.norm_q.weight, attn.norm_k.weight)] + block.oq_int8_attention = int8_attention_allowed(*weights) + + +def forward(kind): + return _single if kind == "single" else _double + + +def _modulation(norm_zero, temb): + """The four ``[dim]`` rows a ``LuminaRMSNormZero`` projects from the timestep embedding.""" + emb = norm_zero.linear(F.silu(temb)) + return [t.reshape(-1).contiguous() for t in emb.chunk(4, dim=1)] + + +def _rope_tables(freqs_cis, rows): + """Complex ``[1, seq, head_dim / 2]`` frequencies -> interleaved fp32 cos/sin rows.""" + freqs = freqs_cis.reshape(-1, freqs_cis.shape[-1])[:rows] + return ( + freqs.real.float().repeat_interleave(2, -1).contiguous(), + freqs.imag.float().repeat_interleave(2, -1).contiguous(), + ) + + +def _valid_keys(mask, rows): + if mask is None: + return None + if mask.dim() != 2: + raise ValueError("fused Boogu blocks take [batch, seq] padding masks") + valid = mask.reshape(-1)[:rows].bool() + return None if bool(valid.all()) else valid + + +class _Heads: + """Query/key/value buffers of one GQA attention, heads zero-padded to a power of two.""" + + def __init__(self, attn, qkv_width, rows, device): + self.attn = attn + self.heads = attn.heads + self.head_dim = attn.norm_q.weight.shape[0] + self.kv_heads = (qkv_width - self.heads * self.head_dim) // (2 * self.head_dim) + padded = 1 << (self.head_dim - 1).bit_length() + kv_shape = (1, rows, self.kv_heads, padded) + self.query = torch.empty((1, rows, self.heads, padded), device=device, dtype=torch.bfloat16) + self.key = torch.empty(kv_shape, device=device, dtype=torch.bfloat16) + self.value = torch.zeros(kv_shape, device=device, dtype=torch.bfloat16) + + def fill(self, qkv, cos, sin, rows): + """RMS-normed, rotated Q/K and V of ``qkv`` rows into buffer rows ``rows``.""" + attn, head_dim = self.attn, self.head_dim + q_width, kv_width = self.heads * head_dim, self.kv_heads * head_dim + for src, heads, norm, out in ( + (qkv[:, :q_width], self.heads, attn.norm_q, self.query), + (qkv[:, q_width : q_width + kv_width], self.kv_heads, attn.norm_k, self.key), + ): + act.qk_norm_rope_with_triton( + src, + heads, + head_dim, + norm.weight, + cos, + sin, + norm.eps, + weight_mode=act.QK_WEIGHT_DIFFUSERS, + style=act.ROPE_INTERLEAVED, + out=out[:, rows], + out_dim=out.shape[-1], + ) + self.value[0, rows, :, :head_dim] = qkv[:, q_width + kv_width :].unflatten( + -1, (self.kv_heads, head_dim) + ) + + def attend(self, block, valid=None): + key, value = self.key, self.value + if valid is not None: + key, value = key[:, valid], value[:, valid] + out = block.oq_runtime.attention( + self.query, + key, + value, + int8_allowed=block.oq_int8_attention, + sm_scale=float(self.attn.scale), + ) + return out[0, :, :, : self.head_dim].reshape(out.shape[1], self.heads * self.head_dim) + + +def _self_attention(block, attn, qkv, freqs_cis, mask): + rows = qkv.shape[0] + heads = _Heads(attn, qkv.shape[1], rows, qkv.device) + cos, sin = _rope_tables(freqs_cis, rows) + heads.fill(qkv, cos, sin, slice(0, rows)) + return heads.attend(block, _valid_keys(mask, rows)) + + +def _gated(h, x, norm, gate): + act.postnorm_residual_with_triton(h, x, norm.weight, gate, norm.eps, norm=act.NORM_RMS_TORCH) + + +def _ffn(runtime, h, x, ff_in, ff_out, norm_in, norm_out, gate, mod=None): + source = runtime.input( + x, + x.shape[1], + norm=act.NORM_RMS_TORCH, + norm_weight=norm_in.weight, + rms_eps=norm_in.eps, + **(mod or {}), + ) + mid = runtime.linear(ff_in, source, epilogue=gemm.EPILOGUE_SWIGLU) + _gated(h, runtime.linear(ff_out, runtime.input(mid, mid.shape[1])), norm_out, gate) + + +def _single(self, hidden_states, attention_mask, image_rotary_emb, temb=None): + if hidden_states.shape[0] != 1: + return torch.cat( + [ + _single( + self, + hidden_states[i : i + 1], + None if attention_mask is None else attention_mask[i : i + 1], + image_rotary_emb[i : i + 1], + None if temb is None else temb[i : i + 1], + ) + for i in range(hidden_states.shape[0]) + ] + ) + runtime, fused = self.oq_runtime, self.oq_fused + dim = hidden_states.shape[2] + h = hidden_states[0].clone() + if self.modulation: + scale_msa, gate_msa, scale_mlp, gate_mlp = _modulation(self.norm1, temb) + gate_msa, gate_mlp = torch.tanh(gate_msa), torch.tanh(gate_mlp) + norm1 = self.norm1.norm + attn_mod = dict(mod=act.MOD_ONE_PLUS_SCALE, mod_scale=scale_msa) + mlp_mod = dict(mod=act.MOD_ONE_PLUS_SCALE, mod_scale=scale_mlp) + else: + gate_msa = gate_mlp = torch.ones(dim, device=h.device, dtype=h.dtype) + norm1, attn_mod, mlp_mod = self.norm1, {}, {} + source = runtime.input( + h, dim, norm=act.NORM_RMS_TORCH, norm_weight=norm1.weight, rms_eps=norm1.eps, **attn_mod + ) + qkv = runtime.linear(fused.qkv, source) + att = _self_attention(self, self.attn, qkv, image_rotary_emb, attention_mask) + _gated(h, runtime.linear(fused.out, runtime.input(att, att.shape[1])), self.norm2, gate_msa) + _ffn( + runtime, h, h, fused.ff_in, fused.ff_out, self.ffn_norm1, self.ffn_norm2, gate_mlp, mlp_mod + ) + return h[None] + + +def _double( + self, + img_hidden_states, + instruct_hidden_states, + img_attention_mask, + joint_attention_mask, + image_rotary_emb, + rotary_emb, + temb=None, + encoder_seq_lengths=None, + seq_lengths=None, +): + if not self.modulation: + raise NotImplementedError("fused Boogu double-stream blocks are modulated") + if img_hidden_states.shape[0] != 1: + outs = [ + _double( + self, + img_hidden_states[i : i + 1], + instruct_hidden_states[i : i + 1], + None if img_attention_mask is None else img_attention_mask[i : i + 1], + None if joint_attention_mask is None else joint_attention_mask[i : i + 1], + image_rotary_emb[i : i + 1], + rotary_emb[i : i + 1], + temb[i : i + 1], + [encoder_seq_lengths[i]], + [seq_lengths[i]], + ) + for i in range(img_hidden_states.shape[0]) + ] + return torch.cat([o[0] for o in outs]), torch.cat([o[1] for o in outs]) + runtime, fused = self.oq_runtime, self.oq_fused + l_txt, l_all = int(encoder_seq_lengths[0]), int(seq_lengths[0]) + l_img = l_all - l_txt + img_in, txt_in = img_hidden_states[0, :l_img], instruct_hidden_states[0, :l_txt] + dim = img_in.shape[1] + + i_scale1, i_gate_msa, i_scale_mlp, i_gate_mlp = _modulation(self.img_norm1, temb) + i_scale2, i_shift_mlp, _, _ = _modulation(self.img_norm2, temb) + i_scale3, i_gate_self, _, _ = _modulation(self.img_norm3, temb) + t_scale1, t_gate_msa, t_scale_mlp, t_gate_mlp = _modulation(self.instruct_norm1, temb) + t_scale2, t_shift_mlp, _, _ = _modulation(self.instruct_norm2, temb) + + def modulated(x, norm_zero, scale): + return runtime.input( + x, + dim, + norm=act.NORM_RMS_TORCH, + mod=act.MOD_ONE_PLUS_SCALE, + norm_weight=norm_zero.norm.weight, + mod_scale=scale, + rms_eps=norm_zero.norm.eps, + ) + + def mlp_input(x, norm_zero, scale, scale_mlp, shift_mlp): + # Two norms in a row (RMSNormZero, then ffn_norm1): the first runs as a prologue-only + # pass, the shift-scale between them with the eager rounding of each step. + return (1 + scale_mlp) * modulated(x, norm_zero, scale).dense() + shift_mlp + + joint_attn = self.img_instruct_attn + txt_qkv = runtime.linear(fused.joint_txt_qkv, modulated(txt_in, self.instruct_norm1, t_scale1)) + img_qkv = runtime.linear(fused.joint_img_qkv, modulated(img_in, self.img_norm1, i_scale1)) + heads = _Heads(joint_attn, img_qkv.shape[1], l_all, img_in.device) + cos, sin = _rope_tables(rotary_emb, l_all) + heads.fill(txt_qkv, cos[:l_txt], sin[:l_txt], slice(0, l_txt)) + heads.fill(img_qkv, cos[l_txt:], sin[l_txt:], slice(l_txt, l_all)) + att = heads.attend(self, _valid_keys(joint_attention_mask, l_all)) + merged = torch.empty((l_all, dim), device=img_in.device, dtype=torch.bfloat16) + runtime.linear(fused.joint_txt_out, runtime.input(att[:l_txt], dim), out=merged[:l_txt]) + runtime.linear(fused.joint_img_out, runtime.input(att[l_txt:], dim), out=merged[l_txt:]) + joint = runtime.linear(fused.joint_out, runtime.input(merged, dim)) + + self_qkv = runtime.linear(fused.self_qkv, modulated(img_in, self.img_norm3, i_scale3)) + self_att = _self_attention( + self, self.img_self_attn, self_qkv, image_rotary_emb, img_attention_mask + ) + self_out = runtime.linear(fused.self_out, runtime.input(self_att, dim)) + + img_mlp = mlp_input(img_in, self.img_norm2, i_scale2, i_scale_mlp, i_shift_mlp) + txt_mlp = mlp_input(txt_in, self.instruct_norm2, t_scale2, t_scale_mlp, t_shift_mlp) + + img = img_in.clone() + _gated(img, joint[l_txt:], self.img_attn_norm, torch.tanh(i_gate_msa)) + _gated(img, self_out, self.img_self_attn_norm, torch.tanh(i_gate_self)) + _ffn( + runtime, + img, + img_mlp, + fused.img_ff_in, + fused.img_ff_out, + self.img_ffn_norm1, + self.img_ffn_norm2, + torch.tanh(i_gate_mlp), + ) + txt = txt_in.clone() + _gated(txt, joint[:l_txt], self.instruct_attn_norm, torch.tanh(t_gate_msa)) + _ffn( + runtime, + txt, + txt_mlp, + fused.txt_ff_in, + fused.txt_ff_out, + self.instruct_ffn_norm1, + self.instruct_ffn_norm2, + torch.tanh(t_gate_mlp), + ) + + img_out, txt_out = img_hidden_states.clone(), instruct_hidden_states.clone() + img_out[0, :l_img] = img + txt_out[0, :l_txt] = txt + return img_out, txt_out diff --git a/src/orbitquant/fused/families/flux2.py b/src/orbitquant/fused/families/flux2.py new file mode 100644 index 0000000..df3baff --- /dev/null +++ b/src/orbitquant/fused/families/flux2.py @@ -0,0 +1,346 @@ +"""FLUX.2 (``Flux2Transformer2DModel``) double- and single-stream blocks. + +Double-stream block, per stream: LayerNorm + AdaLN prologue, one Q|K|V GEMM, Q/K RMSNorm + +interleaved RoPE into a joint text|image sequence, the gated output projection and the SwiGLU +feed-forward with their residual updates in GEMM epilogues. + +Single-stream block: the fused ``to_qkv_mlp_proj`` splits into a Q|K|V group and a SwiGLU +group whose activation lands next to the attention output, so ``to_out`` reads both from one +buffer. + +The reference-image KV cache modes keep their attention pattern (diffusers' +``_flux2_kv_causal_attention``) on SDPA. +""" + +from __future__ import annotations + +import torch + +from orbitquant.fused.families import int8_attention_allowed, specs +from orbitquant.fused.groups import RowSource +from orbitquant.kernels import triton_activation as act +from orbitquant.kernels import triton_int8_gemm as gemm + +FAMILY = "flux2" +MODEL_CLASSES = ("Flux2Transformer2DModel",) + + +def blocks(model): + return [(f"transformer_blocks.{i}", "double") for i in range(len(model.transformer_blocks))] + [ + (f"single_transformer_blocks.{i}", "single") + for i in range(len(model.single_transformer_blocks)) + ] + + +def groups(kind, block): + if kind == "double": + hidden = block.ff.linear_out.in_features + return specs( + img_qkv=["attn.to_q", "attn.to_k", "attn.to_v"], + txt_qkv=["attn.add_q_proj", "attn.add_k_proj", "attn.add_v_proj"], + img_out=["attn.to_out.0"], + txt_out=["attn.to_add_out"], + img_ff_in=( + [ + RowSource("ff.linear_in", 0, hidden), + RowSource("ff.linear_in", hidden, 2 * hidden), + ], + {"swiglu": True}, + ), + img_ff_out=["ff.linear_out"], + txt_ff_in=( + [ + RowSource("ff_context.linear_in", 0, hidden), + RowSource("ff_context.linear_in", hidden, 2 * hidden), + ], + {"swiglu": True}, + ), + txt_ff_out=["ff_context.linear_out"], + ) + attn = block.attn + inner, hidden = attn.inner_dim, attn.mlp_hidden_dim + return specs( + qkv=[RowSource("attn.to_qkv_mlp_proj", 0, 3 * inner)], + mlp=( + [ + RowSource("attn.to_qkv_mlp_proj", 3 * inner, 3 * inner + hidden), + RowSource("attn.to_qkv_mlp_proj", 3 * inner + hidden, 3 * inner + 2 * hidden), + ], + {"swiglu": True}, + ), + out=["attn.to_out"], + ) + + +def prepare(block, kind): + attn = block.attn + weights = [attn.norm_q.weight, attn.norm_k.weight] + if kind == "double": + weights += [attn.norm_added_q.weight, attn.norm_added_k.weight] + block.oq_int8_attention = int8_attention_allowed(*weights) + + +def forward(kind): + return _double if kind == "double" else _single + + +def _modulation(mod, sets): + """``Flux2Modulation.split`` as flat ``[rows, dim]`` tensors plus a row index (or None).""" + if mod.ndim == 2: + mod = mod.unsqueeze(1) + chunks = torch.chunk(mod, 3 * sets, dim=-1) + rows = mod.shape[1] + out = [] + for i in range(sets): + shift, scale, gate = (c.reshape(rows, -1).contiguous() for c in chunks[3 * i : 3 * i + 3]) + out.append((shift, scale, gate)) + return out, rows + + +def _row_index(rows, seq, device): + if rows == 1: + return None + if rows != seq: + raise ValueError(f"modulation has {rows} rows for {seq} tokens") + return torch.arange(seq, device=device, dtype=torch.int32) + + +def _ln_input(runtime, x, dim, shift, scale, index, eps): + return runtime.input( + x, + dim, + norm=act.NORM_LAYER, + mod=act.MOD_SHIFT_SCALE, + mod_scale=scale if index is not None else scale.reshape(-1), + mod_shift=shift if index is not None else shift.reshape(-1), + mod_index=index, + rms_eps=eps, + ) + + +def _gate(gate, index): + return (gate, index) if index is not None else (gate.reshape(-1), None) + + +def _qk(attn, qkv, rows, query, key, cos, sin, added): + """Q/K RMSNorm + RoPE of ``qkv`` into rows ``rows`` of ``[1, seq, heads, head_dim]``.""" + inner, heads, head_dim = attn.inner_dim, attn.heads, attn.head_dim + norm_q = attn.norm_added_q if added else attn.norm_q + norm_k = attn.norm_added_k if added else attn.norm_k + for src, norm, out in ( + (qkv[:, :inner], norm_q, query), + (qkv[:, inner : 2 * inner], norm_k, key), + ): + act.qk_norm_rope_with_triton( + src, + heads, + head_dim, + norm.weight, + cos, + sin, + norm.eps, + weight_mode=act.QK_WEIGHT_TORCH, + style=act.ROPE_INTERLEAVED, + out=out[:, rows], + ) + + +def _kv_cache_attention(query, key, value, kwargs, num_txt): + from diffusers.models.transformers.transformer_flux2 import _flux2_kv_causal_attention + + mode, cache = kwargs.get("kv_cache_mode"), kwargs.get("kv_cache") + num_ref = int(kwargs.get("num_ref_tokens", 0) or 0) + if mode == "extract" and cache is not None and num_ref > 0: + cache.store( + key[:, num_txt : num_txt + num_ref].clone(), + value[:, num_txt : num_txt + num_ref].clone(), + ) + if mode == "extract" and num_ref > 0: + return _flux2_kv_causal_attention(query, key, value, num_txt, num_ref) + return _flux2_kv_causal_attention(query, key, value, num_txt, 0, kv_cache=cache) + + +def _attention(self, query, key, value, kwargs, num_txt): + runtime = self.oq_runtime + if kwargs and kwargs.get("kv_cache_mode") is not None: + return _kv_cache_attention(query, key, value, kwargs, num_txt) + return runtime.attention(query, key, value, int8_allowed=self.oq_int8_attention) + + +def _double( + self, + hidden_states, + encoder_hidden_states, + temb_mod_img, + temb_mod_txt, + image_rotary_emb=None, + joint_attention_kwargs=None, +): + if hidden_states.shape[0] != 1: + outs = [ + _double( + self, + hidden_states[i : i + 1], + encoder_hidden_states[i : i + 1], + temb_mod_img[i : i + 1], + temb_mod_txt[i : i + 1], + image_rotary_emb, + joint_attention_kwargs, + ) + for i in range(hidden_states.shape[0]) + ] + return torch.cat([o[0] for o in outs]), torch.cat([o[1] for o in outs]) + runtime, fused, attn = self.oq_runtime, self.oq_fused, self.attn + kwargs = joint_attention_kwargs or {} + img, txt = hidden_states[0], encoder_hidden_states[0] + s_img, s_txt, dim = img.shape[0], txt.shape[0], img.shape[1] + seq = s_txt + s_img + eps = float(self.norm1.eps) + (i_msa, i_mlp), i_rows = _modulation(temb_mod_img, 2) + (t_msa, t_mlp), t_rows = _modulation(temb_mod_txt, 2) + i_index = _row_index(i_rows, s_img, img.device) + t_index = _row_index(t_rows, s_txt, img.device) + + inner, heads, head_dim = attn.inner_dim, attn.heads, attn.head_dim + qkv = torch.empty((seq, 3 * inner), device=img.device, dtype=torch.bfloat16) + runtime.linear( + fused.txt_qkv, + _ln_input(runtime, txt, dim, t_msa[0], t_msa[1], t_index, eps), + out=qkv[:s_txt], + ) + runtime.linear( + fused.img_qkv, + _ln_input(runtime, img, dim, i_msa[0], i_msa[1], i_index, eps), + out=qkv[s_txt:], + ) + cos, sin = image_rotary_emb + query = torch.empty((1, seq, heads, head_dim), device=img.device, dtype=torch.bfloat16) + key = torch.empty_like(query) + _qk(attn, qkv[:s_txt], slice(0, s_txt), query, key, cos[:s_txt], sin[:s_txt], True) + _qk(attn, qkv[s_txt:], slice(s_txt, seq), query, key, cos[s_txt:], sin[s_txt:], False) + value = qkv[:, 2 * inner :].unflatten(-1, (heads, head_dim))[None] + att = _attention(self, query, key, value, kwargs, s_txt).reshape(seq, inner) + + gate, index = _gate(i_msa[2], i_index) + img = img.clone() + runtime.linear( + fused.img_out, + runtime.input(att[s_txt:], inner), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=img, + gate=gate, + gate_index=index, + out=img, + ) + shift, scale, gate = i_mlp + mid = runtime.linear( + fused.img_ff_in, + _ln_input(runtime, img, dim, shift, scale, i_index, eps), + epilogue=gemm.EPILOGUE_SWIGLU, + ) + gate, index = _gate(gate, i_index) + runtime.linear( + fused.img_ff_out, + runtime.input(mid, mid.shape[1]), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=img, + gate=gate, + gate_index=index, + out=img, + ) + + gate, index = _gate(t_msa[2], t_index) + txt = txt.clone() + runtime.linear( + fused.txt_out, + runtime.input(att[:s_txt], inner), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=txt, + gate=gate, + gate_index=index, + out=txt, + ) + shift, scale, gate = t_mlp + mid = runtime.linear( + fused.txt_ff_in, + _ln_input(runtime, txt, dim, shift, scale, t_index, eps), + epilogue=gemm.EPILOGUE_SWIGLU, + ) + gate, index = _gate(gate, t_index) + runtime.linear( + fused.txt_ff_out, + runtime.input(mid, mid.shape[1]), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=txt, + gate=gate, + gate_index=index, + out=txt, + ) + return txt[None], img[None] + + +def _single( + self, + hidden_states, + encoder_hidden_states, + temb_mod, + image_rotary_emb=None, + joint_attention_kwargs=None, + split_hidden_states=False, + text_seq_len=None, +): + if encoder_hidden_states is not None: + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + if hidden_states.shape[0] != 1: + out = torch.cat( + [ + _single( + self, + hidden_states[i : i + 1], + None, + temb_mod[i : i + 1], + image_rotary_emb, + joint_attention_kwargs, + ) + for i in range(hidden_states.shape[0]) + ] + ) + else: + out = _single_one(self, hidden_states, temb_mod, image_rotary_emb, joint_attention_kwargs) + if split_hidden_states: + return out[:, :text_seq_len], out[:, text_seq_len:] + return out + + +def _single_one(self, hidden_states, temb_mod, image_rotary_emb, joint_attention_kwargs): + runtime, fused, attn = self.oq_runtime, self.oq_fused, self.attn + kwargs = joint_attention_kwargs or {} + h = hidden_states[0].clone() + seq, dim = h.shape + eps = float(self.norm.eps) + ((shift, scale, gate),), rows = _modulation(temb_mod, 1) + index = _row_index(rows, seq, h.device) + source = _ln_input(runtime, h, dim, shift, scale, index, eps) + inner, heads, head_dim = attn.inner_dim, attn.heads, attn.head_dim + qkv = runtime.linear(fused.qkv, source) + joined = torch.empty((seq, inner + attn.mlp_hidden_dim), device=h.device, dtype=torch.bfloat16) + runtime.linear(fused.mlp, source, epilogue=gemm.EPILOGUE_SWIGLU, out=joined, col_offset=inner) + cos, sin = image_rotary_emb + query = torch.empty((1, seq, heads, head_dim), device=h.device, dtype=torch.bfloat16) + key = torch.empty_like(query) + _qk(attn, qkv, slice(0, seq), query, key, cos, sin, False) + value = qkv[:, 2 * inner :].unflatten(-1, (heads, head_dim))[None] + num_txt = int(kwargs.get("num_txt_tokens", 0) or 0) + joined[:, :inner] = _attention(self, query, key, value, kwargs, num_txt).reshape(seq, inner) + gate, gate_index = _gate(gate, index) + runtime.linear( + fused.out, + runtime.input(joined, joined.shape[1]), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate, + gate_index=gate_index, + out=h, + ) + return h[None] diff --git a/src/orbitquant/fused/families/ideogram4.py b/src/orbitquant/fused/families/ideogram4.py new file mode 100644 index 0000000..7236b07 --- /dev/null +++ b/src/orbitquant/fused/families/ideogram4.py @@ -0,0 +1,139 @@ +"""Ideogram 4 (``Ideogram4Transformer2DModel``) blocks. + +Sandwich-norm block: RMSNorm x (1 + scale) prologue, one Q|K|V GEMM, Q/K RMSNorm + rotate-half +MRoPE on 256-wide heads, attention, then ``h + tanh(gate) * RMSNorm(out)`` as a row kernel (the +post-norm needs the whole row, so it cannot ride in a GEMM epilogue); the SwiGLU feed-forward +mirrors it. The AdaLN modulation projection keeps its own module. +""" + +from __future__ import annotations + +import torch + +from orbitquant.fused.families import int8_attention_allowed, specs +from orbitquant.kernels import triton_activation as act +from orbitquant.kernels import triton_int8_gemm as gemm + +FAMILY = "ideogram4" +MODEL_CLASSES = ("Ideogram4Transformer2DModel",) + + +def blocks(model): + return [(f"layers.{i}", "main") for i in range(len(model.layers))] + + +def groups(kind, block): + return specs( + qkv=["attention.to_q", "attention.to_k", "attention.to_v"], + out=["attention.to_out.0"], + ff_in=(["feed_forward.w1", "feed_forward.w3"], {"swiglu": True}), + ff_out=["feed_forward.w2"], + ) + + +def prepare(block, kind): + attn = block.attention + block.oq_int8_attention = int8_attention_allowed(attn.norm_q.weight, attn.norm_k.weight) + + +def forward(kind): + return _forward + + +def _per_row(t, seq): + """``[1, rows, dim]`` modulation -> (table, index) addressing it per token.""" + t = t.reshape(-1, t.shape[-1]) + if t.shape[0] == 1: + return t.reshape(-1), None + if t.shape[0] != seq: + raise ValueError(f"modulation has {t.shape[0]} rows for {seq} tokens") + return t.contiguous(), torch.arange(seq, device=t.device, dtype=torch.int32) + + +def _forward(self, hidden_states, attention_mask, image_rotary_emb, adaln_input): + if hidden_states.shape[0] != 1: + return torch.cat( + [ + _forward( + self, + hidden_states[i : i + 1], + None if attention_mask is None else attention_mask[i : i + 1], + tuple(t[i : i + 1] for t in image_rotary_emb), + adaln_input[i : i + 1], + ) + for i in range(hidden_states.shape[0]) + ] + ) + runtime, fused, attn = self.oq_runtime, self.oq_fused, self.attention + _, seq, dim = hidden_states.shape + mod = self.adaln_modulation(adaln_input) + scale_msa, gate_msa, scale_mlp, gate_mlp = mod.chunk(4, dim=-1) + gate_msa, gate_mlp = torch.tanh(gate_msa), torch.tanh(gate_mlp) + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + scale_msa, index = _per_row(scale_msa, seq) + gate_msa, _ = _per_row(gate_msa, seq) + scale_mlp, _ = _per_row(scale_mlp, seq) + gate_mlp, _ = _per_row(gate_mlp, seq) + eps = float(self.attention_norm1.eps) + h = hidden_states[0].clone() + + source = runtime.input( + h, + dim, + norm=act.NORM_RMS, + mod=act.MOD_SCALE, + norm_weight=self.attention_norm1.weight, + mod_scale=scale_msa, + mod_index=index, + rms_eps=eps, + ) + qkv = runtime.linear(fused.qkv, source) + heads, head_dim = attn.num_heads, attn.head_dim + cos, sin = (t.reshape(seq, -1) for t in image_rotary_emb) + query, key = ( + act.qk_norm_rope_with_triton( + qkv[:, offset : offset + dim], + heads, + head_dim, + norm.weight, + cos, + sin, + norm.eps, + weight_mode=act.QK_WEIGHT_DIFFUSERS, + style=act.ROPE_HALF, + rope_bf16=True, + ) + for offset, norm in ((0, attn.norm_q), (dim, attn.norm_k)) + ) + value = qkv[:, 2 * dim :].unflatten(-1, (heads, head_dim))[None] + single_segment = attention_mask is None or bool(attention_mask.all()) + if single_segment: + att = runtime.attention(query, key, value, int8_allowed=self.oq_int8_attention) + else: + att = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=attention_mask, + ).transpose(1, 2) + out = runtime.linear(fused.out, runtime.input(att.reshape(seq, dim), dim)) + act.postnorm_residual_with_triton( + h, out, self.attention_norm2.weight, gate_msa, eps, norm=act.NORM_RMS, gate_index=index + ) + + source = runtime.input( + h, + dim, + norm=act.NORM_RMS, + mod=act.MOD_SCALE, + norm_weight=self.ffn_norm1.weight, + mod_scale=scale_mlp, + mod_index=index, + rms_eps=eps, + ) + mid = runtime.linear(fused.ff_in, source, epilogue=gemm.EPILOGUE_SWIGLU) + out = runtime.linear(fused.ff_out, runtime.input(mid, mid.shape[1])) + act.postnorm_residual_with_triton( + h, out, self.ffn_norm2.weight, gate_mlp, eps, norm=act.NORM_RMS, gate_index=index + ) + return h[None] diff --git a/src/orbitquant/fused/families/krea2.py b/src/orbitquant/fused/families/krea2.py new file mode 100644 index 0000000..c7f4017 --- /dev/null +++ b/src/orbitquant/fused/families/krea2.py @@ -0,0 +1,124 @@ +"""Krea 2 Turbo (``Krea2Transformer2DModel``) main blocks. + +Per block: one Q|K|V|gate GEMM with the sigmoid gate in its epilogue, Q/K RMSNorm + interleaved +RoPE, attention, the gated output projection and the SwiGLU feed-forward with their residual +updates in GEMM epilogues. The protected BF16 ``ff.down`` runs as per-row INT8 (W8A8). +""" + +from __future__ import annotations + +import torch + +from orbitquant.fused.families import int8_attention_allowed, specs +from orbitquant.kernels import triton_activation as act +from orbitquant.kernels import triton_int8_gemm as gemm + +FAMILY = "krea2" +MODEL_CLASSES = ("Krea2Transformer2DModel",) + + +def blocks(model): + return [(f"transformer_blocks.{i}", "main") for i in range(len(model.transformer_blocks))] + + +def groups(kind, block): + return specs( + qkvg=["attn.to_q", "attn.to_k", "attn.to_v", "attn.to_gate"], + out=["attn.to_out.0"], + gateup=(["ff.gate", "ff.up"], {"swiglu": True}), + down=["ff.down"], + ) + + +def prepare(block, kind): + # Krea 2 stores its Q/K RMSNorm weights zero-centered. + attn = block.attn + block.oq_int8_attention = int8_attention_allowed( + attn.norm_q.weight, attn.norm_k.weight, offset=1.0 + ) + + +def forward(kind): + return _forward + + +def _forward(self, hidden_states, temb, image_rotary_emb, attention_mask=None): + if hidden_states.shape[0] != 1: + return torch.cat( + [ + _forward( + self, + hidden_states[i : i + 1], + temb[i : i + 1], + image_rotary_emb, + (None if attention_mask is None else attention_mask[i : i + 1]), + ) + for i in range(hidden_states.shape[0]) + ] + ) + runtime, fused, attn = self.oq_runtime, self.oq_fused, self.attn + _, seq, dim = hidden_states.shape + h = hidden_states.reshape(seq, dim) + modulation = temb.reshape(-1).unflatten(-1, (6, -1)) + self.scale_shift_table + prescale, preshift, pregate, postscale, postshift, postgate = ( + t.reshape(dim).contiguous() for t in modulation.unbind(-2) + ) + eps = float(self.norm1.eps) + + source = runtime.input( + h, + dim, + norm=act.NORM_RMS_ONE_PLUS, + mod=act.MOD_SHIFT_SCALE, + norm_weight=self.norm1.weight, + mod_scale=prescale, + mod_shift=preshift, + rms_eps=eps, + ) + qkvg = runtime.linear( + fused.qkvg, + source, + epilogue=gemm.EPILOGUE_SIGMOID_TAIL, + sig_from=fused.qkvg.out_features - dim, + ) + heads, kv_heads, head_dim = attn.num_heads, attn.num_kv_heads, attn.head_dim + qw, kw = heads * head_dim, kv_heads * head_dim + cos, sin = image_rotary_emb + query = act.qk_norm_rope_with_triton( + qkvg[:, :qw], heads, head_dim, attn.norm_q.weight, cos, sin, attn.norm_q.eps + ) + key = act.qk_norm_rope_with_triton( + qkvg[:, qw : qw + kw], kv_heads, head_dim, attn.norm_k.weight, cos, sin, attn.norm_k.eps + ) + value = qkvg[:, qw + kw : qw + 2 * kw].unflatten(-1, (kv_heads, head_dim))[None] + if attention_mask is not None and not bool(attention_mask.all()): + valid = attention_mask.reshape(-1).bool() + key, value = key[:, valid], value[:, valid] + att = runtime.attention(query, key, value, int8_allowed=self.oq_int8_attention).reshape( + seq, dim + ) + + source = runtime.input(att, dim, prologue=act.PROLOGUE_MULTIPLY, x2=qkvg[:, qw + 2 * kw :]) + h = runtime.linear( + fused.out, source, epilogue=gemm.EPILOGUE_RESIDUAL_GATE, residual=h, gate=pregate, out=h + ) + source = runtime.input( + h, + dim, + norm=act.NORM_RMS_ONE_PLUS, + mod=act.MOD_SHIFT_SCALE, + norm_weight=self.norm2.weight, + mod_scale=postscale, + mod_shift=postshift, + rms_eps=eps, + ) + mid = runtime.linear(fused.gateup, source, epilogue=gemm.EPILOGUE_SWIGLU) + h = runtime.linear( + fused.down, + runtime.input(mid, mid.shape[1]), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=postgate, + out=h, + ) + return h.reshape(1, seq, dim) diff --git a/src/orbitquant/fused/families/minimax_h3.py b/src/orbitquant/fused/families/minimax_h3.py new file mode 100644 index 0000000..acb250a --- /dev/null +++ b/src/orbitquant/fused/families/minimax_h3.py @@ -0,0 +1,145 @@ +"""MiniMax-H3 (``MiniMaxH3Transformer3DModel``) main blocks. + +One packed video/audio/text sequence; the AdaLN modulation of every row is a row of a +(timestep, modality) table. The prologues and residual epilogues read that table through the +per-row index directly instead of materializing ``[seq, hidden]`` modulation tensors. Q/K get +RMSNorm and a rotate-half RoPE on the leading 96 channels of each 128-wide head. The token +refiner blocks (text only, once per forward) keep their modules. + +Padding rows form their own attention document (the model's mask pairs live rows with live rows +and padding rows with padding rows); the two documents run as separate mask-free attentions. +""" + +from __future__ import annotations + +import torch + +from orbitquant.fused.families import int8_attention_allowed, specs +from orbitquant.fused.groups import RowSource +from orbitquant.kernels import triton_activation as act +from orbitquant.kernels import triton_int8_gemm as gemm + +FAMILY = "minimax_h3" +MODEL_CLASSES = ("MiniMaxH3Transformer3DModel",) + + +def blocks(model): + return [(f"transformer_blocks.{i}", "main") for i in range(len(model.transformer_blocks))] + + +def groups(kind, block): + ffn = block.ff.net[2].in_features + return specs( + qkv=["attn.to_q", "attn.to_k", "attn.to_v"], + out=["attn.to_out.0"], + # diffusers' SwiGLU computes first_half * silu(second_half): the gate is the second half. + ff_in=( + [RowSource("ff.net.0.proj", ffn, 2 * ffn), RowSource("ff.net.0.proj", 0, ffn)], + {"swiglu": True}, + ), + ff_out=["ff.net.2"], + ) + + +def prepare(block, kind): + attn = block.attn + block.oq_int8_attention = int8_attention_allowed(attn.norm_q.weight, attn.norm_k.weight) + + +def forward(kind): + return _forward + + +def _forward(self, hidden_states, temb, adaln_indices, rotary_emb, attention_mask=None): + if hidden_states.shape[0] != 1: + return torch.cat( + [ + _forward( + self, hidden_states[i : i + 1], temb, adaln_indices, rotary_emb, attention_mask + ) + for i in range(hidden_states.shape[0]) + ] + ) + runtime, fused, attn = self.oq_runtime, self.oq_fused, self.attn + _, seq, dim = hidden_states.shape + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + t.contiguous() for t in self.adaln_proj(temb) + ) + index = adaln_indices.to(torch.int32).contiguous() + h = hidden_states[0].clone() + source = runtime.input( + h, + dim, + norm=act.NORM_RMS_TORCH, + mod=act.MOD_SHIFT_SCALE, + norm_weight=self.norm1.weight, + mod_scale=scale_msa, + mod_shift=shift_msa, + mod_index=index, + rms_eps=float(self.norm1.eps), + ) + qkv = runtime.linear(fused.qkv, source) + heads, head_dim = attn.heads, attn.head_dim + inner = heads * head_dim + cos, sin = rotary_emb + query, key = ( + act.qk_norm_rope_with_triton( + qkv[:, offset : offset + inner], + heads, + head_dim, + norm.weight, + cos, + sin, + norm.eps, + weight_mode=act.QK_WEIGHT_TORCH, + style=act.ROPE_HALF, + rope_bf16=True, + ) + for offset, norm in ((0, attn.norm_q), (inner, attn.norm_k)) + ) + value = qkv[:, 2 * inner :].unflatten(-1, (heads, head_dim))[None] + if attention_mask is None: + att = runtime.attention(query, key, value, int8_allowed=self.oq_int8_attention) + else: + att = torch.empty_like(query) + first_document = attention_mask.reshape(seq, seq)[0] + for rows in (first_document, ~first_document): + rows = rows.nonzero().flatten() + if rows.numel(): + att[:, rows] = runtime.attention( + query[:, rows], + key[:, rows], + value[:, rows], + int8_allowed=self.oq_int8_attention, + ) + runtime.linear( + fused.out, + runtime.input(att.reshape(seq, inner), inner), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate_msa, + gate_index=index, + out=h, + ) + source = runtime.input( + h, + dim, + norm=act.NORM_RMS_TORCH, + mod=act.MOD_SHIFT_SCALE, + norm_weight=self.norm2.weight, + mod_scale=scale_mlp, + mod_shift=shift_mlp, + mod_index=index, + rms_eps=float(self.norm2.eps), + ) + mid = runtime.linear(fused.ff_in, source, epilogue=gemm.EPILOGUE_SWIGLU) + runtime.linear( + fused.ff_out, + runtime.input(mid, mid.shape[1]), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate_mlp, + gate_index=index, + out=h, + ) + return h[None] diff --git a/src/orbitquant/fused/families/qwenimage21.py b/src/orbitquant/fused/families/qwenimage21.py new file mode 100644 index 0000000..bc6400b --- /dev/null +++ b/src/orbitquant/fused/families/qwenimage21.py @@ -0,0 +1,230 @@ +"""Qwen-Image 2.1 (``QwenImage21Transformer2DModel``) single-stream blocks. + +Per block: LayerNorm x (1 + scale) prologue, one Q|K|V GEMM, Q/K RMSNorm + interleaved complex +RoPE, attention over the block-causal joint sequence, the tanh-gated output projection and the +SwiGLU feed-forward with their residual updates in GEMM epilogues. The shared modulation holds +one row per sample plus, with ``causal_condition``, a ``t = 0`` row for the text and condition +tokens; the prologues and gates address it per token. + +Attention follows ``QwenImage21AttnProcessor``: in the prefill every text segment is causal over +its own keys after full access to everything before it (SDPA with an explicit mask; text is +short), every image segment attends to all keys up to its end, and the target attends to all +keys; the decode reuses the cached post-RoPE prefix keys and values. Padded text keys +(``key_valid``) are dropped rather than masked. The checkpoint keeps some ``img_mlp.out`` +projections in BF16 for quality; they stay BF16 here too. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from orbitquant.fused.families import int8_attention_allowed, specs +from orbitquant.kernels import triton_activation as act +from orbitquant.kernels import triton_int8_gemm as gemm + +FAMILY = "qwenimage21" +MODEL_CLASSES = ("QwenImage21Transformer2DModel",) +DENSE = "bf16" + + +def blocks(model): + return [(f"transformer_blocks.{i}", "main") for i in range(len(model.transformer_blocks))] + + +def groups(kind, block): + return specs( + qkv=["attn.to_q", "attn.to_k", "attn.to_v"], + out=["attn.to_out.0"], + ff_in=(["img_mlp.gate_layer", "img_mlp.proj"], {"swiglu": True}), + ff_out=["img_mlp.out"], + ) + + +def prepare(block, kind): + attn = block.attn + block.oq_int8_attention = int8_attention_allowed(attn.norm_q.weight, attn.norm_k.weight) + + +def forward(kind): + return _forward + + +def _rows(params, target_token_mask, sample): + """Modulation table of one sample and the per-token row index into it (or None).""" + if target_token_mask is None: + return params[sample : sample + 1].contiguous(), None + table = torch.cat([params[sample : sample + 1], params[-1:]]).contiguous() + index = torch.where(target_token_mask, 0, 1).to(torch.int32) + return table, index + + +def _rope_tables(freqs_cis): + """Complex ``[seq, head_dim / 2]`` frequencies -> interleaved fp32 cos/sin rows.""" + return ( + freqs_cis.real.float().repeat_interleave(2, -1).contiguous(), + freqs_cis.imag.float().repeat_interleave(2, -1).contiguous(), + ) + + +def _attend(runtime, block, query, key, value, valid): + if valid is not None and not bool(valid.all()): + key, value = key[:, valid], value[:, valid] + return runtime.attention(query, key, value, int8_allowed=block.oq_int8_attention) + + +def _text_segment(query, key, value, start, valid): + """Causal over the segment's own keys after full access to the keys before it.""" + rows = query.shape[1] + mask = torch.ones(rows, start + rows, dtype=torch.bool, device=query.device) + mask[:, start:] = torch.tril(mask[:, start:]) + if valid is not None: + mask = mask & valid[None, : start + rows] + out = F.scaled_dot_product_attention( + query.transpose(1, 2), + key[:, : start + rows].transpose(1, 2), + value[:, : start + rows].transpose(1, 2), + attn_mask=mask[None, None], + ) + return out.transpose(1, 2) + + +def _forward( + self, + hidden_states, + modulation, + rotary_emb=None, + attention_mask=None, + target_token_mask=None, + layer_cache=None, + kv_cache_mode=None, + cache_write_slice=None, + segments=None, + key_valid=None, +): + batch = hidden_states.shape[0] + if batch != 1 and layer_cache is not None: + raise NotImplementedError("fused Qwen-Image 2.1 blocks cache one sample at a time") + if key_valid is None and isinstance(attention_mask, torch.Tensor): + # The decode passes key validity as a [batch, 1, 1, keys] mask instead. + key_valid = attention_mask.reshape(batch, -1).bool() + outs = [] + for sample in range(batch): + valid = None if key_valid is None else key_valid[sample] + outs.append( + _sample( + self, + hidden_states[sample], + modulation, + sample, + rotary_emb, + target_token_mask, + layer_cache, + kv_cache_mode, + cache_write_slice, + segments, + valid, + ) + ) + return torch.stack(outs) + + +def _sample( + self, + hidden, + modulation, + sample, + rotary_emb, + target_token_mask, + layer_cache, + kv_cache_mode, + cache_write_slice, + segments, + valid, +): + runtime, fused, attn = self.oq_runtime, self.oq_fused, self.attn + seq, dim = hidden.shape + mod1, mod2 = modulation.chunk(2, dim=-1) + scale1, gate1 = mod1.chunk(2, dim=-1) + scale2, gate2 = mod2.chunk(2, dim=-1) + scale1, index = _rows(scale1, target_token_mask, sample) + scale2, _ = _rows(scale2, target_token_mask, sample) + gate1, _ = _rows(torch.tanh(gate1), target_token_mask, sample) + gate2, _ = _rows(torch.tanh(gate2), target_token_mask, sample) + h = hidden.clone() + + def modulated(x, norm, scale): + return runtime.input( + x, + dim, + norm=act.NORM_LAYER, + mod=act.MOD_ONE_PLUS_SCALE, + mod_scale=scale, + mod_index=index, + rms_eps=norm.eps, + ) + + qkv = runtime.linear(fused.qkv, modulated(h, self.img_norm1, scale1)) + inner, heads = attn.inner_dim, attn.heads + head_dim = inner // heads + query = torch.empty((1, seq, heads, head_dim), device=h.device, dtype=torch.bfloat16) + key = torch.empty_like(query) + cos, sin = _rope_tables(rotary_emb) if rotary_emb is not None else (None, None) + for offset, norm, out in ((0, attn.norm_q, query), (inner, attn.norm_k, key)): + act.qk_norm_rope_with_triton( + qkv[:, offset : offset + inner], + heads, + head_dim, + norm.weight, + cos, + sin, + norm.eps, + weight_mode=act.QK_WEIGHT_DIFFUSERS, + style=act.ROPE_INTERLEAVED, + out=out, + ) + value = qkv[:, 2 * inner :].unflatten(-1, (heads, head_dim))[None] + if layer_cache is not None and kv_cache_mode == "extract" and cache_write_slice is not None: + layer_cache.store(key[:, cache_write_slice].clone(), value[:, cache_write_slice].clone()) + elif layer_cache is not None and kv_cache_mode == "cached": + cached_key, cached_value = layer_cache.get() + key = torch.cat([cached_key, key], dim=1) + value = torch.cat([cached_value, value], dim=1) + + if segments is None: + att = _attend(runtime, self, query, key, value, valid) + else: + prefix = segments[-1][1] if segments else 0 + att = torch.empty_like(query) + for start, end, is_text in segments: + if is_text: + att[:, start:end] = _text_segment(query[:, start:end], key, value, start, valid) + else: + part_valid = None if valid is None else valid[:end] + att[:, start:end] = _attend( + runtime, self, query[:, start:end], key[:, :end], value[:, :end], part_valid + ) + att[:, prefix:] = _attend(runtime, self, query[:, prefix:], key, value, valid) + + runtime.linear( + fused.out, + runtime.input(att.reshape(seq, inner), inner), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate1, + gate_index=index, + out=h, + ) + mid = runtime.linear( + fused.ff_in, modulated(h, self.img_norm2, scale2), epilogue=gemm.EPILOGUE_SWIGLU + ) + runtime.linear( + fused.ff_out, + runtime.input(mid, mid.shape[1]), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate2, + gate_index=index, + out=h, + ) + return h diff --git a/src/orbitquant/fused/groups.py b/src/orbitquant/fused/groups.py new file mode 100644 index 0000000..f81780a --- /dev/null +++ b/src/orbitquant/fused/groups.py @@ -0,0 +1,377 @@ +"""Fused projection weights: projections that read the same activation, stored as one matrix. + +* ``PackedGroup``: OrbitQuant low-bit rows (W2/W3/W4) of several projections concatenated, run + as one INT8-surrogate GEMM. Rows destined for a SwiGLU epilogue interleave gate and up chunks. + Projections of different weight widths (low-bit boundary and interior protection) become row + segments that decode with their own codebooks; the per-row scale absorbs the codebook scale + difference, so a mixed group is still one GEMM. +* ``Int8Group``: per-row INT8 weights (W8A8, ``Int8RowLinear`` numerics). +* ``Bf16Group``: source-precision weights kept as shipped. + +Groups are built from the layers of a loaded model (``from_sources``) or created empty with the +same shapes (``empty``) so a checkpoint saved from them loads straight into the buffers. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from torch import nn + +from orbitquant.codebooks import get_codebook +from orbitquant.int8_head import quantize_int8_rows +from orbitquant.layers import OrbitQuantLinear + +SWIGLU_INTERLEAVE = 128 + + +@dataclass(frozen=True) +class RowSource: + """Rows ``[start, stop)`` of the projection at ``path`` (``stop=None``: to the end).""" + + path: str + start: int = 0 + stop: int | None = None + + +def _resolve(root: nn.Module, path: str) -> nn.Module: + module = root + for part in path.split("."): + module = module[int(part)] if part.isdigit() else getattr(module, part) + return module + + +def _rows(source: RowSource, layer: nn.Module) -> tuple[int, int]: + out_features = layer.out_features + stop = out_features if source.stop is None else source.stop + if not 0 <= source.start < stop <= out_features: + raise ValueError(f"rows [{source.start}, {stop}) outside {source.path} ({out_features})") + return source.start, stop + + +def _interleave(gate: torch.Tensor, up: torch.Tensor, chunk: int) -> torch.Tensor: + if gate.shape != up.shape or gate.shape[0] % chunk: + raise ValueError(f"SwiGLU halves need equal row counts divisible by {chunk}") + pieces = [] + for start in range(0, gate.shape[0], chunk): + pieces += [gate[start : start + chunk], up[start : start + chunk]] + return torch.cat(pieces, dim=0) + + +def surrogate_constants( + in_features: int, weight_bits: int, activation_bits: int, codebook_version: int +) -> tuple[torch.Tensor, float]: + """INT8 surrogate codes of the weight codebook and ``alpha = a_scale * w_scale`` (8-bit + activations carry their absolute scale per token: ``a_scale = 1``).""" + from orbitquant.kernels.triton_cuda import fit_int8_centroid_surrogate + + act_scale = 1.0 + if activation_bits != 8: + _, act_scale = fit_int8_centroid_surrogate( + get_codebook(in_features, activation_bits, codebook_version).centroids + ) + codes, w_scale = fit_int8_centroid_surrogate( + get_codebook(in_features, weight_bits, codebook_version).centroids + ) + return codes, float(act_scale * w_scale) + + +def _segments(sizes_and_bits) -> tuple[tuple[int, int], ...]: + """Consecutive ``(rows, bits)`` pieces merged into runs of one weight width.""" + runs: list[list[int]] = [] + for rows, bits in sizes_and_bits: + if runs and runs[-1][1] == bits: + runs[-1][0] += rows + else: + runs.append([rows, bits]) + return tuple((rows, bits) for rows, bits in runs) + + +class PackedGroup(nn.Module): + kind = "packed" + + def __init__( + self, + packed: torch.Tensor, + row_norms: torch.Tensor, + bias: torch.Tensor | None, + *, + in_features: int, + segments, + activation_bits: int, + codebook_version: int, + ): + super().__init__() + self.in_features = int(in_features) + self.segments = tuple((int(rows), int(bits)) for rows, bits in segments) + self.out_features = sum(rows for rows, _ in self.segments) + self.activation_bits = int(activation_bits) + self.codebook_version = int(codebook_version) + if tuple(row_norms.shape) != (self.out_features,): + raise ValueError(f"{tuple(row_norms.shape)} row norms for {self.out_features} rows") + if tuple(packed.shape) != self.packed_shape(self.segments, self.in_features): + raise ValueError(f"packed weights {tuple(packed.shape)} do not match {self.segments}") + self.register_buffer("packed", packed) + self.register_buffer("row_norms", row_norms) + self.register_buffer("bias", bias) + constants = { + bits: surrogate_constants( + self.in_features, bits, self.activation_bits, self.codebook_version + ) + for bits in dict.fromkeys(bits for _, bits in self.segments) + } + for bits, (codes, _) in constants.items(): + self.register_buffer(f"codes_w{bits}", codes, persistent=False) + self.alpha = constants[self.segments[0][1]][1] + factor = None + if len(constants) > 1: + factor = torch.cat( + [ + torch.full((rows,), constants[bits][1] / self.alpha, dtype=torch.float32) + for rows, bits in self.segments + ] + ) + self.register_buffer("scale_factor", factor, persistent=False) + + @staticmethod + def packed_shape(segments, in_features: int) -> tuple[int, ...]: + """``[rows, row_bytes]`` for one weight width, a flat byte run for mixed widths.""" + if len(segments) == 1: + rows, bits = segments[0] + return (rows, in_features * bits // 8) + return (sum(rows * in_features * bits // 8 for rows, bits in segments),) + + @property + def weight_bits(self) -> int | None: + return self.segments[0][1] if len(self.segments) == 1 else None + + @classmethod + def empty( + cls, + segments, + in_features: int, + *, + activation_bits: int, + codebook_version: int, + bias: bool = False, + device: torch.device | str = "meta", + ) -> PackedGroup: + segments = tuple(segments) + rows = sum(count for count, _ in segments) + return cls( + torch.empty(cls.packed_shape(segments, in_features), dtype=torch.uint8, device=device), + torch.empty(rows, dtype=torch.float32, device=device), + torch.empty(rows, dtype=torch.bfloat16, device=device) if bias else None, + in_features=in_features, + segments=segments, + activation_bits=activation_bits, + codebook_version=codebook_version, + ) + + @staticmethod + def source_segments(layers, sources, *, swiglu: bool = False) -> tuple[tuple[int, int], ...]: + """Row segments of a group built from ``layers`` (loaded or skeleton projections).""" + sizes = [] + for source, layer in zip(sources, layers, strict=True): + start, stop = _rows(source, layer) + sizes.append((stop - start, int(layer.weight_bits))) + if swiglu: + if len(sizes) != 2: + raise ValueError("a SwiGLU group takes exactly a gate and an up source") + if sizes[0][1] != sizes[1][1]: + raise ValueError("the gate and up halves of a SwiGLU group need one weight width") + return _segments(sizes) + + @classmethod + def from_sources( + cls, + root: nn.Module, + sources: list[RowSource], + *, + swiglu: bool = False, + activation_bits: int | None = None, + ) -> PackedGroup: + """Concatenate the rows of OrbitQuant projections; with ``swiglu`` the two sources are + the gate and up halves, interleaved in ``SWIGLU_INTERLEAVE`` row chunks. + ``activation_bits`` overrides the projections' activation width (the weights do not + depend on it).""" + layers = [_resolve(root, source.path) for source in sources] + first = layers[0] + + def signature(layer): + return ( + layer.in_features, + layer.activation_bits, + layer.weight_codebook.algorithm_version, + layer.rotation, + ) + + for layer in layers: + if not isinstance(layer, OrbitQuantLinear) or layer.packed_weight_indices is None: + raise ValueError("packed groups need loaded OrbitQuant projections") + if signature(layer) != signature(first): + raise ValueError("projections of one group must share width and rotation") + segments = cls.source_segments(layers, sources, swiglu=swiglu) + packed, norms, biases = [], [], [] + has_bias = [layer.bias is not None for layer in layers] + if any(has_bias) and not all(has_bias): + raise ValueError("either every projection of a group has a bias or none has") + for source, layer in zip(sources, layers, strict=True): + start, stop = _rows(source, layer) + row_bytes = layer.in_features * layer.weight_bits // 8 + rows = layer.packed_weight_indices.reshape(layer.out_features, row_bytes) + packed.append(rows[start:stop]) + norms.append(layer.row_norms.float().reshape(-1)[start:stop]) + if layer.bias is not None: + biases.append(layer.bias.to(torch.bfloat16).reshape(-1)[start:stop]) + if swiglu: + packed = [_interleave(packed[0], packed[1], SWIGLU_INTERLEAVE)] + norms = [_interleave(norms[0][:, None], norms[1][:, None], SWIGLU_INTERLEAVE)[:, 0]] + if biases: + biases = [ + _interleave(biases[0][:, None], biases[1][:, None], SWIGLU_INTERLEAVE)[:, 0] + ] + if len(segments) > 1: + packed = [rows.reshape(-1) for rows in packed] + return cls( + torch.cat(packed, dim=0).contiguous(), + torch.cat(norms, dim=0).contiguous(), + torch.cat(biases, dim=0).contiguous() if biases else None, + in_features=first.in_features, + segments=segments, + activation_bits=activation_bits or first.activation_bits, + codebook_version=first.weight_codebook.algorithm_version, + ) + + def row_scales(self) -> torch.Tensor: + """Per-row GEMM scale relative to ``alpha`` (the row norms, rescaled for segments whose + codebook surrogate has a different scale than the first one).""" + if self.scale_factor is None: + return self.row_norms + return self.row_norms * self.scale_factor.to(self.row_norms.device) + + def decode(self, workspace: torch.Tensor) -> torch.Tensor: + from orbitquant.kernels.triton_activation import decode_lowbit_to_int8_with_triton + + view = workspace[: self.out_features * self.in_features].view( + self.out_features, self.in_features + ) + flat = self.packed.reshape(-1) + row = offset = 0 + for rows, bits in self.segments: + nbytes = rows * self.in_features * bits // 8 + codes = getattr(self, f"codes_w{bits}") + decode_lowbit_to_int8_with_triton( + flat[offset : offset + nbytes], + codes.to(flat.device), + bits, + rows, + self.in_features, + out=view[row : row + rows], + ) + row += rows + offset += nbytes + return view + + +class Int8Group(nn.Module): + kind = "int8" + + def __init__(self, q: torch.Tensor, scales: torch.Tensor, bias: torch.Tensor | None): + super().__init__() + self.out_features, self.in_features = (int(v) for v in q.shape) + self.register_buffer("q", q) + self.register_buffer("scales", scales) + self.register_buffer("bias", bias) + + @classmethod + def empty(cls, rows: int, in_features: int, *, bias: bool = False, device="meta") -> Int8Group: + return cls( + torch.empty(rows, in_features, dtype=torch.int8, device=device), + torch.empty(rows, dtype=torch.float32, device=device), + torch.empty(rows, dtype=torch.bfloat16, device=device) if bias else None, + ) + + @classmethod + def from_sources( + cls, + root: nn.Module, + sources: list[RowSource], + *, + swiglu: bool = False, + activation_bits: int | None = None, + ) -> Int8Group: + weights, biases = [], [] + for source in sources: + layer = _resolve(root, source.path) + weight = _dense_weight(layer) + start, stop = _rows(source, layer) + weights.append(weight[start:stop]) + if getattr(layer, "bias", None) is not None: + biases.append(layer.bias.to(torch.bfloat16).reshape(-1)[start:stop]) + if swiglu: + weights = [_interleave(weights[0], weights[1], SWIGLU_INTERLEAVE)] + if biases: + biases = [ + _interleave(biases[0][:, None], biases[1][:, None], SWIGLU_INTERLEAVE)[:, 0] + ] + q, scales = quantize_int8_rows(torch.cat(weights, dim=0)) + return cls(q, scales, torch.cat(biases).contiguous() if biases else None) + + +class Bf16Group(nn.Module): + kind = "bf16" + + def __init__(self, weight: torch.Tensor, bias: torch.Tensor | None): + super().__init__() + self.out_features, self.in_features = (int(v) for v in weight.shape) + self.register_buffer("weight", weight) + self.register_buffer("bias", bias) + + @classmethod + def empty(cls, rows: int, in_features: int, *, bias: bool = False, device="meta") -> Bf16Group: + return cls( + torch.empty(rows, in_features, dtype=torch.bfloat16, device=device), + torch.empty(rows, dtype=torch.bfloat16, device=device) if bias else None, + ) + + @classmethod + def from_sources( + cls, + root: nn.Module, + sources: list[RowSource], + *, + swiglu: bool = False, + activation_bits: int | None = None, + ) -> Bf16Group: + weights, biases = [], [] + for source in sources: + layer = _resolve(root, source.path) + start, stop = _rows(source, layer) + weights.append(_dense_weight(layer)[start:stop].to(torch.bfloat16)) + if getattr(layer, "bias", None) is not None: + biases.append(layer.bias.to(torch.bfloat16).reshape(-1)[start:stop]) + if swiglu: + weights = [_interleave(weights[0], weights[1], SWIGLU_INTERLEAVE)] + if biases: + biases = [ + _interleave(biases[0][:, None], biases[1][:, None], SWIGLU_INTERLEAVE)[:, 0] + ] + return cls( + torch.cat(weights, dim=0).contiguous(), + torch.cat(biases).contiguous() if biases else None, + ) + + +def _dense_weight(layer: nn.Module) -> torch.Tensor: + """Source-precision weight of a projection the artifact kept unquantized.""" + if isinstance(layer, OrbitQuantLinear): + raise ValueError("an OrbitQuant projection belongs in a packed group") + weight = getattr(layer, "weight", None) + if weight is None: + raise ValueError(f"{type(layer).__name__} has no dense weight") + return weight.detach() + + +GROUP_TYPES = {cls.kind: cls for cls in (PackedGroup, Int8Group, Bf16Group)} diff --git a/src/orbitquant/fused/install.py b/src/orbitquant/fused/install.py new file mode 100644 index 0000000..8050471 --- /dev/null +++ b/src/orbitquant/fused/install.py @@ -0,0 +1,223 @@ +"""Turn a loaded OrbitQuant transformer into its fused form, or prepare one to load as fused. + +``fuse(model)`` regroups the projections of every block into ``block.oq_fused`` (a +``ModuleDict`` of groups), replaces the source projections with ``nn.Identity`` and swaps in the +family's block forward. ``model.oq_fused_layout`` then describes the result; saved next to the +weights it lets ``prepare_skeleton`` rebuild the same empty groups before a checkpoint loads, +so a fused checkpoint never holds the unfused projections. + +A group's type and row segments follow from its source projections: OrbitQuant projections form +a packed group (segmented by weight width), projections the artifact kept dense become INT8 rows +or stay BF16 as the layout's ``dense`` mode says. The quantization config already determines the +projection types and widths, so the skeleton rebuilds the same groups from it. The layout's +``activation_bits`` maps group names to an activation width other than the config's (8-bit +activations for groups whose inputs 4-bit codes cannot carry). +""" + +from __future__ import annotations + +import json +import types +from pathlib import Path + +import torch +from torch import nn + +from orbitquant.fused.families import adapter_for_class, adapter_for_family +from orbitquant.fused.groups import GROUP_TYPES, Bf16Group, Int8Group, PackedGroup, _resolve +from orbitquant.fused.runtime import FusedRuntime +from orbitquant.layers import OrbitQuantLinear + +LAYOUT_VERSION = 1 + + +def _set(root: nn.Module, path: str, module: nn.Module) -> None: + parent_path, _, name = path.rpartition(".") + parent = _resolve(root, parent_path) if parent_path else root + if name.isdigit(): + parent[int(name)] = module + else: + setattr(parent, name, module) + + +DENSE_MODES = ("int8", "bf16") + + +def _group_kind(block: nn.Module, spec, dense: str) -> str: + packed = [isinstance(_resolve(block, s.path), OrbitQuantLinear) for s in spec.sources] + if all(packed): + return PackedGroup.kind + if any(packed): + raise ValueError("a group mixes OrbitQuant and dense projections") + mode = spec.dense or dense + if mode not in DENSE_MODES: + raise ValueError(f"unknown dense group mode {mode!r}") + return Int8Group.kind if mode == "int8" else Bf16Group.kind + + +def _quant_config(model, config=None): + """The OrbitQuant config: given, on the model, on its HF quantizer or in its config.""" + if config is None: + config = getattr(model, "quantization_config", None) + if config is None: + config = getattr(getattr(model, "hf_quantizer", None), "quantization_config", None) + if config is None: + model_config = getattr(model, "config", None) + payload = getattr(model_config, "quantization_config", None) + if payload is None and isinstance(model_config, dict): + payload = model_config.get("quantization_config") + if isinstance(payload, dict): + from orbitquant.config import OrbitQuantConfig + + config = OrbitQuantConfig.from_dict(payload) + else: + config = payload + if config is None: + raise ValueError("the model carries no OrbitQuant quantization_config") + return config + + +def _attach(model, block, adapter, kind, group_modules, runtime): + specs = adapter.groups(kind, block) + for spec in specs.values(): + for source in spec.sources: + _set(block, source.path, nn.Identity()) + block.oq_fused = group_modules + block.__dict__["oq_runtime"] = runtime + block.oq_kind = kind + block.forward = types.MethodType(adapter.forward(kind), block) + + +def fuse( + model: nn.Module, + *, + dense: str | None = None, + attention: str = "int8", + activation_bits: dict[str, int] | None = None, + config=None, +) -> dict: + """Fuse a loaded OrbitQuant transformer in place (weights on the host or the GPU). + + ``dense`` picks INT8 rows or BF16 for projections the artifact kept unquantized (default: + the family's ``DENSE``, else INT8); ``activation_bits`` maps group names to their activation + width (default: the family's ``ACTIVATION_BITS``, else the config's); ``config`` is the + OrbitQuant config when the model does not carry one (``orbitquant-v1`` component + artifacts).""" + adapter = adapter_for_class(type(model).__name__) + dense = dense or getattr(adapter, "DENSE", "int8") + if dense not in DENSE_MODES: + raise ValueError(f"unknown dense group mode {dense!r}") + widths = dict( + getattr(adapter, "ACTIVATION_BITS", {}) if activation_bits is None else activation_bits + ) + runtime = FusedRuntime.from_config(_quant_config(model, config), attention=attention) + for path, kind in adapter.blocks(model): + block = _resolve(model, path) + modules = nn.ModuleDict() + for name, spec in adapter.groups(kind, block).items(): + group_cls = GROUP_TYPES[_group_kind(block, spec, dense)] + modules[name] = group_cls.from_sources( + block, list(spec.sources), swiglu=spec.swiglu, activation_bits=widths.get(name) + ) + prepare = getattr(adapter, "prepare", None) + if prepare is not None: + prepare(block, kind) + _attach(model, block, adapter, kind, modules, runtime) + layout = {"family": adapter.FAMILY, "version": LAYOUT_VERSION, "dense": dense} + if widths: + layout["activation_bits"] = widths + model.oq_fused_layout = layout + model.__dict__["oq_runtime"] = runtime + return layout + + +def _empty_group(block, spec, group_kind: str, config, device, activation_bits=None): + layers = [_resolve(block, source.path) for source in spec.sources] + first = layers[0] + bias = getattr(first, "bias", None) is not None + if group_kind == PackedGroup.kind: + return PackedGroup.empty( + PackedGroup.source_segments(layers, spec.sources, swiglu=spec.swiglu), + first.in_features, + activation_bits=activation_bits or first.activation_bits, + codebook_version=config.codebook_version, + bias=bias, + device=device, + ) + rows = 0 + for source, layer in zip(spec.sources, layers, strict=True): + stop = layer.out_features if source.stop is None else source.stop + rows += stop - source.start + return GROUP_TYPES[group_kind].empty(rows, first.in_features, bias=bias, device=device) + + +def prepare_skeleton( + model: nn.Module, layout: dict, *, attention: str = "int8", config=None +) -> None: + """Install empty fused groups (and placeholders for their sources) before a fused + checkpoint loads. Call after the quantized projections were turned into OrbitQuant + skeletons; ``finalize`` completes the blocks once the weights are in.""" + if int(layout.get("version", 0)) != LAYOUT_VERSION: + raise ValueError(f"unsupported fused layout version {layout.get('version')!r}") + adapter = adapter_for_family(layout["family"]) + config = _quant_config(model, config) + dense = layout.get("dense", "int8") + widths = layout.get("activation_bits", {}) + runtime = FusedRuntime.from_config(config, attention=attention) + for path, kind in adapter.blocks(model): + block = _resolve(model, path) + specs = adapter.groups(kind, block) + reference = _resolve(block, next(iter(specs.values())).sources[0].path) + tensors = [*reference.buffers(), *reference.parameters()] + device = tensors[0].device if tensors else torch.device("meta") + modules = nn.ModuleDict( + { + name: _empty_group( + block, spec, _group_kind(block, spec, dense), config, device, widths.get(name) + ) + for name, spec in specs.items() + } + ) + _attach(model, block, adapter, kind, modules, runtime) + model.oq_fused_layout = layout + model.__dict__["oq_runtime"] = runtime + + +def finalize(model: nn.Module) -> None: + """Derive the per-block settings that depend on loaded weights (e.g. INT8 attention).""" + layout = model.oq_fused_layout + adapter = adapter_for_family(layout["family"]) + prepare = getattr(adapter, "prepare", None) + if prepare is None: + return + for path, kind in adapter.blocks(model): + prepare(_resolve(model, path), kind) + + +def release(model: nn.Module) -> None: + """Drop the runtime's device buffers (call before moving the weights off the GPU).""" + runtime = model.__dict__.get("oq_runtime") + if runtime is not None: + runtime.release() + + +def is_fused_tensor_name(name: str) -> bool: + return ".oq_fused." in f".{name}" + + +def save_pretrained(model: nn.Module, directory, **kwargs) -> None: + """``model.save_pretrained`` for a fused transformer, recording the layout in the + ``quantization_config`` of ``config.json`` so ``from_pretrained`` loads it fused.""" + layout = model.oq_fused_layout + config = _quant_config(model) + config.fused_layout = layout + model.save_pretrained(directory, **kwargs) + config_path = Path(directory) / "config.json" + data = json.loads(config_path.read_text()) + quantization = data.get("quantization_config") + if quantization is None: + quantization = config.to_dict() + data["quantization_config"] = quantization + if quantization.get("fused_layout") != layout: + quantization["fused_layout"] = layout + config_path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n") diff --git a/src/orbitquant/fused/runtime.py b/src/orbitquant/fused/runtime.py new file mode 100644 index 0000000..3bb603b --- /dev/null +++ b/src/orbitquant/fused/runtime.py @@ -0,0 +1,220 @@ +"""Per-device state and the primitive operations fused blocks are written with. + +A ``FusedRuntime`` belongs to one transformer. It owns the INT8 workspace the packed groups +decode into right before their GEMM, one activation quantizer per input width, and the +attention dispatch. Block forwards call: + +* ``input(x, width, ...)``: the prologue (norm, modulation, product) of a block input; its + RPBH/Lloyd-Max codes are computed once per activation width and shared by every group that + reads the same input; +* ``linear(group, source, ...)``: one fused GEMM with its epilogue; +* ``attention(q, k, v, ...)``: INT8 Q.K^T / FP16 P.V where the block allows it, else SDPA. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from orbitquant.codebooks import get_codebook +from orbitquant.fused.groups import Bf16Group, Int8Group, PackedGroup +from orbitquant.kernels import triton_activation as act +from orbitquant.kernels import triton_int8_gemm as gemm +from orbitquant.kernels.triton_attention import int8_attention_with_triton +from orbitquant.rotations import get_rpbh_rotation + +ATTENTION_BACKENDS = ("int8", "flash", "cudnn", "sdpa") + + +class FusedInput: + """One block input prepared for its consumers: INT8 codes per activation width for packed + groups, the bf16 prologue output for INT8-row and BF16 groups (each computed on first use).""" + + __slots__ = ("_runtime", "_x", "_width", "_kwargs", "_codes", "_dense") + + def __init__(self, runtime, x, width, kwargs): + self._runtime, self._x, self._width, self._kwargs = runtime, x, width, kwargs + self._codes: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + self._dense = None + + def codes(self, bits: int): + codes = self._codes.get(bits) + if codes is None: + quantizer = self._runtime.quantizer(self._width, bits, self._x.device) + codes = self._codes[bits] = quantizer(self._x, **self._kwargs) + return codes + + def dense(self): + if self._dense is None: + if not self._kwargs: + self._dense = self._x.reshape(-1, self._width) + else: + # The prologue does not depend on the activation width. + bits = next(iter(self._codes), 4) + quantizer = self._runtime.quantizer(self._width, bits, self._x.device) + self._dense, _ = quantizer(self._x, prologue_only=True, **self._kwargs) + return self._dense + + +class FusedRuntime: + def __init__( + self, + *, + rotation_seed: int, + block_size, + codebook_version: int, + activation_eps: float, + attention: str = "int8", + ): + if attention not in ATTENTION_BACKENDS: + raise ValueError(f"unknown attention backend {attention!r}") + self.rotation_seed = rotation_seed + self.block_size = block_size + self.codebook_version = codebook_version + self.activation_eps = activation_eps + self.attention_backend = attention + self._quantizers: dict[tuple, act.ActivationQuantizer] = {} + self._workspace: torch.Tensor | None = None + + @classmethod + def from_config(cls, config, *, attention: str = "int8") -> FusedRuntime: + return cls( + rotation_seed=config.rotation_seed, + block_size=config.block_size, + codebook_version=config.codebook_version, + activation_eps=config.activation_eps, + attention=attention, + ) + + def quantizer(self, width: int, bits: int, device: torch.device) -> act.ActivationQuantizer: + key = (width, bits, str(device)) + quantizer = self._quantizers.get(key) + if quantizer is None: + rotation = get_rpbh_rotation( + dim=width, seed=self.rotation_seed, block_size=self.block_size + ) + # 8-bit activations are per-token absmax INT8 of the rotated input, no codebook. + codebook = None if bits == 8 else get_codebook(width, bits, self.codebook_version) + quantizer = act.ActivationQuantizer(rotation, codebook, self.activation_eps, device) + self._quantizers[key] = quantizer + return quantizer + + def release(self) -> None: + """Drop device buffers before the weights leave the GPU.""" + self._quantizers.clear() + self._workspace = None + + def _decode(self, group: PackedGroup) -> torch.Tensor: + need = group.out_features * group.in_features + device = group.packed.device + if ( + self._workspace is None + or self._workspace.numel() < need + or (self._workspace.device != device) + ): + self._workspace = torch.empty(need, device=device, dtype=torch.int8) + return group.decode(self._workspace) + + def input(self, x: torch.Tensor, width: int, **prologue) -> FusedInput: + return FusedInput(self, x, width, prologue) + + def linear( + self, + group, + source: FusedInput, + *, + epilogue: int = gemm.EPILOGUE_PLAIN, + out: torch.Tensor | None = None, + col_offset: int = 0, + residual: torch.Tensor | None = None, + gate: torch.Tensor | None = None, + gate_index: torch.Tensor | None = None, + sig_from: int = 0, + ) -> torch.Tensor: + common = dict( + bias=group.bias, + out=out, + col_offset=col_offset, + epilogue=epilogue, + residual=residual, + gate=gate, + gate_index=gate_index, + sig_from=sig_from, + ) + if isinstance(group, PackedGroup): + codes, norms = source.codes(group.activation_bits) + return gemm.matmul_int8_scaled_with_triton( + codes, self._decode(group), norms, group.row_scales(), alpha=group.alpha, **common + ) + if isinstance(group, Int8Group): + codes, scales = act.quantize_rows_int8_with_triton(source.dense()) + return gemm.matmul_int8_scaled_with_triton( + codes, group.q, scales, group.scales, scale_mode=1, **common + ) + if isinstance(group, Bf16Group): + return _dense_epilogue(F.linear(source.dense(), group.weight, group.bias), **common) + raise TypeError(f"unknown fused group {type(group).__name__}") + + def attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + int8_allowed: bool = True, + sm_scale: float | None = None, + ) -> torch.Tensor: + """``[1, S, H, D]`` query, ``[1, S, Hkv, D]`` key/value -> ``[1, S, H, D]``.""" + if self.attention_backend == "int8" and int8_allowed: + return int8_attention_with_triton(query, key, value, sm_scale=sm_scale) + backend = { + "cudnn": SDPBackend.CUDNN_ATTENTION, + "sdpa": None, + }.get(self.attention_backend, SDPBackend.FLASH_ATTENTION) + kwargs = dict(scale=sm_scale, enable_gqa=query.shape[2] != key.shape[2]) + q, k, v = (t.transpose(1, 2) for t in (query, key, value)) + if backend is None: + out = F.scaled_dot_product_attention(q, k, v, **kwargs) + else: + with sdpa_kernel(backend): + out = F.scaled_dot_product_attention(q, k, v, **kwargs) + return out.transpose(1, 2) + + +def _dense_epilogue( + y, + *, + bias, + out, + col_offset, + epilogue, + residual, + gate, + gate_index, + sig_from, +): + """Eager equivalents of the GEMM epilogues for BF16 groups (bias is already applied).""" + del bias + if epilogue == gemm.EPILOGUE_SWIGLU: + rows = y.shape[0] + chunks = y.reshape(rows, -1, 2, gemm.SWIGLU_INTERLEAVE) + g, u = chunks[:, :, 0].reshape(rows, -1), chunks[:, :, 1].reshape(rows, -1) + y = F.silu(g) * u + elif epilogue == gemm.EPILOGUE_GELU_TANH: + y = F.gelu(y, approximate="tanh") + elif epilogue == gemm.EPILOGUE_SIGMOID_TAIL: + y = torch.cat([y[:, :sig_from], torch.sigmoid(y[:, sig_from:])], dim=1) + elif epilogue == gemm.EPILOGUE_RESIDUAL_GATE: + g = gate[gate_index] if gate_index is not None else gate + y = residual[:, col_offset : col_offset + y.shape[1]] + g * y + if out is None: + return y + out[:, col_offset : col_offset + y.shape[1]] = y + return out + + +def default_sm_scale(head_dim: int) -> float: + return 1.0 / math.sqrt(head_dim) diff --git a/src/orbitquant/kernels/triton_activation.py b/src/orbitquant/kernels/triton_activation.py index 76542d1..6b87f4a 100644 --- a/src/orbitquant/kernels/triton_activation.py +++ b/src/orbitquant/kernels/triton_activation.py @@ -1,12 +1,15 @@ """Activation-side Triton kernels for fused DiT blocks. * RPBH + Lloyd-Max activation quantization straight to INT8 surrogate codes, optionally fed by - a prologue that replaces the elementwise ops in front of a projection (RMSNorm with a - ``1 + weight`` scale followed by AdaLN modulation, or an elementwise product such as an - attention output gate). Every value eager code rounds to bf16 is rounded at the same point. + a prologue that replaces the elementwise ops in front of a projection: an RMSNorm or + LayerNorm followed by AdaLN modulation (per tensor or selected per row from a table), or an + elementwise product such as an attention output gate. Every value eager code rounds to bf16 + is rounded at the same point. * per-token absmax INT8 for W8A8 projections (``Int8RowLinear`` numerics), -* per-head RMSNorm + interleaved rotary embedding for Q/K, -* packed W4 -> INT8 surrogate weight decode. +* per-head RMSNorm + rotary embedding for Q/K (interleaved pairs or rotate-half, optionally on + a leading slice of the head, optionally zero-padded to a wider head), +* post-norm gated residual updates (``res + gate * norm(x)``), +* packed W2/W3/W4 -> INT8 surrogate weight decode. """ from __future__ import annotations @@ -20,15 +23,56 @@ from triton.language.extra import libdevice # noqa: E402 PROLOGUE_NONE = 0 -PROLOGUE_NORM_MODULATE = 1 +PROLOGUE_NORM_MODULATE = 1 # NORM_RMS_ONE_PLUS + MOD_SHIFT_SCALE (Krea 2) PROLOGUE_MULTIPLY = 2 +# Norm in front of a projection. The RMSNorm variants differ in where eager code rounds: +# torch.nn.RMSNorm rounds once (``x * inv * w``), diffusers' RMSNorm rounds ``x * inv`` before +# the weight multiply. +NORM_NONE = 0 +NORM_RMS_ONE_PLUS = 1 # x * inv * (1 + w), one rounding +NORM_RMS = 2 # bf16(x * inv) * w (diffusers RMSNorm) +NORM_RMS_PLAIN = 3 # x * inv, no weight +NORM_LAYER = 4 # (x - mean) * rsqrt(var + eps), no affine +NORM_RMS_TORCH = 5 # x * inv * w, one rounding (torch.nn.RMSNorm) + +# AdaLN modulation after the norm. +MOD_NONE = 0 +MOD_SHIFT_SCALE = 1 # (1 + scale) * n + shift +MOD_SCALE = 2 # n * scale (the caller already folded the 1 +) +MOD_ONE_PLUS_SCALE = 3 # (1 + scale) * n + +_PROLOGUES = { + PROLOGUE_NONE: (NORM_NONE, MOD_NONE), + PROLOGUE_NORM_MODULATE: (NORM_RMS_ONE_PLUS, MOD_SHIFT_SCALE), +} + @triton.jit def _bf16(x): return x.to(tl.bfloat16).to(tl.float32) +@triton.jit +def _normalize(x, w_ptr, offs, mask, D: tl.constexpr, rms_eps, NORM: tl.constexpr): + if NORM == 4: + mean = libdevice.div_rn(tl.sum(x, axis=0), D * 1.0) + centered = tl.where(mask, x - mean, 0.0) + var = libdevice.div_rn(tl.sum(centered * centered, axis=0), D * 1.0) + return _bf16(centered * libdevice.rsqrt(var + rms_eps)) + if NORM == 0: + return x + mean_sq = libdevice.div_rn(tl.sum(x * x, axis=0), D * 1.0) + inv = libdevice.rsqrt(mean_sq + rms_eps) + if NORM == 1: + return _bf16(x * inv * (tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) + 1.0)) + if NORM == 2: + return _bf16(_bf16(x * inv) * tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32)) + if NORM == 5: + return _bf16(x * inv * tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32)) + return _bf16(x * inv) + + @triton.jit def _prologue_kernel( x_ptr, @@ -36,32 +80,39 @@ def _prologue_kernel( w_ptr, scale_ptr, shift_ptr, + index_ptr, y_ptr, norms_ptr, x_stride, x2_stride, + mod_stride, D: tl.constexpr, DP2: tl.constexpr, rms_eps, - PROLOGUE: tl.constexpr, + MULTIPLY: tl.constexpr, + NORM: tl.constexpr, + MOD: tl.constexpr, + ROW_INDEXED: tl.constexpr, ): row = tl.program_id(0).to(tl.int64) offs = tl.arange(0, DP2) mask = offs < D x = tl.load(x_ptr + row * x_stride + offs, mask=mask, other=0.0).to(tl.float32) - if PROLOGUE == 1: - mean_sq = libdevice.div_rn(tl.sum(x * x, axis=0), D * 1.0) - inv = libdevice.rsqrt(mean_sq + rms_eps) - w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) + 1.0 - normed = _bf16(x * inv * w) - scale = tl.load(scale_ptr + offs, mask=mask, other=0.0).to(tl.float32) - shift = tl.load(shift_ptr + offs, mask=mask, other=0.0).to(tl.float32) - y = _bf16(_bf16(_bf16(1.0 + scale) * normed) + shift) - elif PROLOGUE == 2: + if MULTIPLY: x2 = tl.load(x2_ptr + row * x2_stride + offs, mask=mask, other=0.0).to(tl.float32) y = _bf16(x * x2) else: - y = x + y = _normalize(x, w_ptr, offs, mask, D, rms_eps, NORM) + if MOD != 0: + mod_row = tl.load(index_ptr + row).to(tl.int64) * mod_stride if ROW_INDEXED else 0 + scale = tl.load(scale_ptr + mod_row + offs, mask=mask, other=0.0).to(tl.float32) + if MOD == 1: + shift = tl.load(shift_ptr + mod_row + offs, mask=mask, other=0.0).to(tl.float32) + y = _bf16(_bf16(_bf16(1.0 + scale) * y) + shift) + elif MOD == 2: + y = _bf16(y * scale) + else: + y = _bf16(_bf16(1.0 + scale) * y) tl.store(y_ptr + row * D + offs, y.to(tl.bfloat16), mask=mask) tl.store(norms_ptr + row, tl.sqrt(tl.sum(y * y, axis=0))) @@ -77,31 +128,7 @@ def _fwht_local_stage(values, block_size: tl.constexpr, stage_width: tl.constexp @triton.jit -def _rpbh_int8_kernel( - y_ptr, - norms_ptr, - perm_ptr, - signs_ptr, - bounds_ptr, - codes_ptr, - q_ptr, - D: tl.constexpr, - NUM_BLOCKS: tl.constexpr, - OB: tl.constexpr, - STAGES: tl.constexpr, - LEVELS: tl.constexpr, - eps, - inv_sqrt_block, -): - row_block = tl.program_id(0) - row = (row_block // NUM_BLOCKS).to(tl.int64) - block = row_block % NUM_BLOCKS - cols = block * OB + tl.arange(0, OB) - src = tl.load(perm_ptr + cols).to(tl.int32) - signs = tl.load(signs_ptr + cols).to(tl.float32) - norm = tl.load(norms_ptr + row).to(tl.float32) - values = tl.load(y_ptr + row * D + src).to(tl.float32) - values = values * signs / (norm + eps) +def _fwht(values, OB: tl.constexpr, STAGES: tl.constexpr): if STAGES > 0: values = _fwht_local_stage(values, OB, 1) if STAGES > 1: @@ -126,30 +153,261 @@ def _rpbh_int8_kernel( values = _fwht_local_stage(values, OB, 1024) if STAGES > 11: values = _fwht_local_stage(values, OB, 2048) + if STAGES > 12: + values = _fwht_local_stage(values, OB, 4096) + if STAGES > 13: + values = _fwht_local_stage(values, OB, 8192) + return values + + +# Widest transform held in one tensor. Wider RPBH blocks (8192, 16384) are rotated as 4096-wide +# chunks whose last one or two butterfly stages pair whole chunks, which are plain elementwise +# ops in registers; one 16384-wide tensor needs more shared memory than consumer GPUs have. +FWHT_CHUNK = 4096 +_CHUNK = tl.constexpr(FWHT_CHUNK) +_CHUNK_STAGES = tl.constexpr(FWHT_CHUNK.bit_length() - 1) + + +@triton.jit +def _rotated_chunk( + y_ptr, perm_ptr, signs_ptr, row, start, denom, D: tl.constexpr, NORMALIZE: tl.constexpr +): + cols = start + tl.arange(0, _CHUNK) + src = tl.load(perm_ptr + cols).to(tl.int32) + signs = tl.load(signs_ptr + cols).to(tl.float32) + values = tl.load(y_ptr + row * D + src).to(tl.float32) * signs + if NORMALIZE: + values = values / denom + return _fwht(values, _CHUNK, _CHUNK_STAGES) + + +@triton.jit +def _rotated_chunks( + y_ptr, + perm_ptr, + signs_ptr, + row, + start, + denom, + D: tl.constexpr, + CHUNKS: tl.constexpr, + CROSS: tl.constexpr, + NORMALIZE: tl.constexpr, +): + # CHUNKS consecutive chunks from ``start``; CROSS stages (widths 4096, 8192) pair chunks. + v0 = _rotated_chunk(y_ptr, perm_ptr, signs_ptr, row, start, denom, D, NORMALIZE) + v1 = _rotated_chunk(y_ptr, perm_ptr, signs_ptr, row, start + _CHUNK, denom, D, NORMALIZE) + v2 = v0 + v3 = v1 + if CHUNKS == 4: + v2 = _rotated_chunk( + y_ptr, perm_ptr, signs_ptr, row, start + 2 * _CHUNK, denom, D, NORMALIZE + ) + v3 = _rotated_chunk( + y_ptr, perm_ptr, signs_ptr, row, start + 3 * _CHUNK, denom, D, NORMALIZE + ) + if CROSS > 0: + v0, v1 = v0 + v1, v0 - v1 + if CHUNKS == 4: + v2, v3 = v2 + v3, v2 - v3 + if CROSS > 1: + v0, v2 = v0 + v2, v0 - v2 + v1, v3 = v1 + v3, v1 - v3 + return v0, v1, v2, v3 + + +@triton.jit +def _absmax_int8_store(values, q_ptr, row, start, scale, D: tl.constexpr): + q = libdevice.rint(libdevice.div_rn(values, scale)) + q = tl.minimum(tl.maximum(q, -127.0), 127.0) + tl.store(q_ptr + row * D + start + tl.arange(0, _CHUNK), q.to(tl.int8)) + + +@triton.jit +def _rpbh_absmax_int8_kernel( + y_ptr, + perm_ptr, + signs_ptr, + q_ptr, + scale_ptr, + D: tl.constexpr, + DP2: tl.constexpr, + STAGES: tl.constexpr, + CHUNKS: tl.constexpr, + inv_sqrt_block, +): + # One row per program: the INT8 scale needs the absolute maximum of the whole rotated row. + # D is a multiple of the RPBH block, so the butterflies of the padded row never pair a + # channel with another block. Blocks wider than a chunk imply D == DP2 (no padding). + row = tl.program_id(0).to(tl.int64) + if CHUNKS == 1: + cols = tl.arange(0, DP2) + mask = cols < D + src = tl.load(perm_ptr + cols, mask=mask, other=0).to(tl.int32) + signs = tl.load(signs_ptr + cols, mask=mask, other=0).to(tl.float32) + values = tl.load(y_ptr + row * D + src, mask=mask, other=0.0).to(tl.float32) * signs + values = _fwht(values, DP2, STAGES) * inv_sqrt_block + scale = libdevice.div_rn(tl.maximum(tl.max(tl.abs(values), axis=0), 1e-12), 127.0) + q = libdevice.rint(libdevice.div_rn(values, scale)) + q = tl.minimum(tl.maximum(q, -127.0), 127.0) + tl.store(q_ptr + row * D + cols, q.to(tl.int8), mask=mask) + else: + v0, v1, v2, v3 = _rotated_chunks( + y_ptr, perm_ptr, signs_ptr, row, 0, 1.0, D, CHUNKS, STAGES - _CHUNK_STAGES, False + ) + v0 *= inv_sqrt_block + v1 *= inv_sqrt_block + v2 *= inv_sqrt_block + v3 *= inv_sqrt_block + peak = tl.maximum(tl.max(tl.abs(v0), axis=0), tl.max(tl.abs(v1), axis=0)) + if CHUNKS == 4: + peak = tl.maximum( + peak, tl.maximum(tl.max(tl.abs(v2), axis=0), tl.max(tl.abs(v3), axis=0)) + ) + scale = libdevice.div_rn(tl.maximum(peak, 1e-12), 127.0) + _absmax_int8_store(v0, q_ptr, row, 0, scale, D) + _absmax_int8_store(v1, q_ptr, row, _CHUNK, scale, D) + if CHUNKS == 4: + _absmax_int8_store(v2, q_ptr, row, 2 * _CHUNK, scale, D) + _absmax_int8_store(v3, q_ptr, row, 3 * _CHUNK, scale, D) + tl.store(scale_ptr + row, scale) + + +@triton.jit +def _codes_store( + values, + bounds_ptr, + codes_ptr, + q_ptr, + row, + start, + inv_sqrt_block, + D: tl.constexpr, + WIDTH: tl.constexpr, + LEVELS: tl.constexpr, +): values *= inv_sqrt_block - indices = tl.zeros((OB,), dtype=tl.int32) + indices = tl.zeros((WIDTH,), dtype=tl.int32) for level in tl.static_range(0, LEVELS - 1): indices += (values > tl.load(bounds_ptr + level)).to(tl.int32) - tl.store(q_ptr + row * D + cols, tl.load(codes_ptr + indices).to(tl.int8)) + tl.store( + q_ptr + row * D + start + tl.arange(0, WIDTH), tl.load(codes_ptr + indices).to(tl.int8) + ) + + +@triton.jit +def _rpbh_int8_kernel( + y_ptr, + norms_ptr, + perm_ptr, + signs_ptr, + bounds_ptr, + codes_ptr, + q_ptr, + D: tl.constexpr, + NUM_BLOCKS: tl.constexpr, + OB: tl.constexpr, + STAGES: tl.constexpr, + LEVELS: tl.constexpr, + eps, + inv_sqrt_block, +): + row_block = tl.program_id(0) + row = (row_block // NUM_BLOCKS).to(tl.int64) + block = row_block % NUM_BLOCKS + norm = tl.load(norms_ptr + row).to(tl.float32) + if OB <= _CHUNK: + cols = block * OB + tl.arange(0, OB) + src = tl.load(perm_ptr + cols).to(tl.int32) + signs = tl.load(signs_ptr + cols).to(tl.float32) + values = tl.load(y_ptr + row * D + src).to(tl.float32) + values = _fwht(values * signs / (norm + eps), OB, STAGES) + _codes_store( + values, bounds_ptr, codes_ptr, q_ptr, row, block * OB, inv_sqrt_block, D, OB, LEVELS + ) + else: + start = block * OB + v0, v1, v2, v3 = _rotated_chunks( + y_ptr, + perm_ptr, + signs_ptr, + row, + start, + norm + eps, + D, + OB // _CHUNK, + STAGES - _CHUNK_STAGES, + True, + ) + _codes_store( + v0, bounds_ptr, codes_ptr, q_ptr, row, start, inv_sqrt_block, D, _CHUNK, LEVELS + ) + _codes_store( + v1, + bounds_ptr, + codes_ptr, + q_ptr, + row, + start + _CHUNK, + inv_sqrt_block, + D, + _CHUNK, + LEVELS, + ) + if OB // _CHUNK == 4: + _codes_store( + v2, + bounds_ptr, + codes_ptr, + q_ptr, + row, + start + 2 * _CHUNK, + inv_sqrt_block, + D, + _CHUNK, + LEVELS, + ) + _codes_store( + v3, + bounds_ptr, + codes_ptr, + q_ptr, + row, + start + 3 * _CHUNK, + inv_sqrt_block, + D, + _CHUNK, + LEVELS, + ) class ActivationQuantizer: """RPBH/Lloyd-Max INT8 surrogate activation quantizer of one OrbitQuantLinear input. Projections that share an input (Q/K/V/gate, gate/up) share one quantizer: the codes and - token norms are identical, so they are produced once. + token norms are identical, so they are produced once. Without a codebook the rotated input + is quantized to INT8 with one absmax scale per token (8-bit activations) and the call + returns that scale instead of the token norm. """ def __init__(self, rotation, codebook, eps: float, device: torch.device): self.dim = int(rotation.dim) self.block = int(rotation.block_size) self.num_blocks = int(rotation.num_blocks) - if self.block > 4096: - raise ValueError("fused activation quantization supports RPBH blocks up to 4096") + if self.block > 16384: + raise ValueError("fused activation quantization supports RPBH blocks up to 16384") self.inv_sqrt_block = float(rotation.normalization) self.eps = float(eps) self.perm = rotation.permutation.to(device=device, dtype=torch.int32).contiguous() self.signs = rotation.signs.to(device=device, dtype=torch.int8).contiguous() + self.absmax = codebook is None + if self.absmax: + if self.block > FWHT_CHUNK and self.dim not in (2 * FWHT_CHUNK, 4 * FWHT_CHUNK): + raise ValueError( + "8-bit activations with RPBH blocks above 4096 need an 8192 or 16384 wide input" + ) + self.levels, self.scale = 256, 1.0 + return self.bounds = codebook.boundaries.to(device=device, dtype=torch.float32).contiguous() self.levels = int(codebook.centroids.numel()) codes, self.scale = fit_int8_centroid_surrogate(codebook.centroids) @@ -160,39 +418,89 @@ def __call__( x, *, prologue=PROLOGUE_NONE, + norm=None, + mod=None, x2=None, norm_weight=None, mod_scale=None, mod_shift=None, + mod_index=None, rms_eps=1e-5, + prologue_only=False, ): - """Return ``(codes int8 [rows, dim], token_norms fp32 [rows])``.""" + """Return ``(codes int8 [rows, dim], token_norms fp32 [rows])``, or with + ``prologue_only`` the bf16 prologue output instead of the codes. + + ``prologue`` selects a preset (``PROLOGUE_MULTIPLY`` multiplies by ``x2``); ``norm`` and + ``mod`` override its norm and modulation. With ``mod_index`` the modulation vectors are + rows ``mod_index[row]`` of the ``mod_scale``/``mod_shift`` tables. + """ if x.dim() != 2 or x.stride(1) != 1: x = x.reshape(-1, self.dim).contiguous() rows = x.shape[0] + multiply = prologue == PROLOGUE_MULTIPLY + if multiply: + norm_mode, mod_mode = NORM_NONE, MOD_NONE + else: + norm_mode, mod_mode = _PROLOGUES[prologue] + if norm is not None: + norm_mode = norm + if mod is not None: + mod_mode = mod if x2 is not None and (x2.dim() != 2 or x2.stride(1) != 1): x2 = x2.reshape(rows, self.dim).contiguous() + if mod_mode != MOD_NONE: + # Per-tensor vectors are [dim]; indexed tables are [rows, dim] sharing one stride. + mod_scale = mod_scale.reshape(-1, self.dim).contiguous() + if mod_shift is not None: + mod_shift = mod_shift.reshape(-1, self.dim).contiguous() y = torch.empty((rows, self.dim), device=x.device, dtype=torch.bfloat16) norms = torch.empty(rows, device=x.device, dtype=torch.float32) q = torch.empty((rows, self.dim), device=x.device, dtype=torch.int8) if rows == 0: - return q, norms + return (y if prologue_only else q), norms _prologue_kernel[(rows,)]( x, x2 if x2 is not None else x, norm_weight if norm_weight is not None else norms, mod_scale if mod_scale is not None else x, mod_shift if mod_shift is not None else x, + mod_index if mod_index is not None else norms, y, norms, x.stride(0), x2.stride(0) if x2 is not None else 0, + mod_scale.stride(0) if mod_index is not None else 0, D=self.dim, DP2=triton.next_power_of_2(self.dim), rms_eps=float(rms_eps), - PROLOGUE=int(prologue), + MULTIPLY=multiply, + NORM=int(norm_mode), + MOD=int(mod_mode), + ROW_INDEXED=mod_index is not None, num_warps=8, + # Under FP contraction the compiler folds the bf16 round trips that mirror the + # eager ops' rounding (extf(truncf(x)) -> x) and fuses the multiply-adds. + enable_fp_fusion=False, ) + if prologue_only: + return y, norms + if self.absmax: + row_width = triton.next_power_of_2(self.dim) + _rpbh_absmax_int8_kernel[(rows,)]( + y, + self.perm, + self.signs, + q, + norms, + D=self.dim, + DP2=row_width, + STAGES=self.block.bit_length() - 1, + CHUNKS=row_width // FWHT_CHUNK if self.block > FWHT_CHUNK else 1, + inv_sqrt_block=self.inv_sqrt_block, + num_warps=16 if row_width >= 16384 else 8, + ) + return q, norms _rpbh_int8_kernel[(rows * self.num_blocks,)]( y, norms, @@ -208,7 +516,7 @@ def __call__( LEVELS=self.levels, eps=self.eps, inv_sqrt_block=self.inv_sqrt_block, - num_warps=8 if self.block >= 512 else 4, + num_warps=16 if self.block >= 8192 else 8 if self.block >= 512 else 4, ) return q, norms @@ -240,67 +548,198 @@ def quantize_rows_int8_with_triton(x: torch.Tensor) -> tuple[torch.Tensor, torch return q, scales +QK_WEIGHT_NONE = 0 +QK_WEIGHT_ONE_PLUS = 1 # x * inv * (1 + w), one rounding (Krea 2) +QK_WEIGHT_DIFFUSERS = 2 # bf16(x * inv) * w (diffusers RMSNorm) +QK_WEIGHT_TORCH = 3 # x * inv * w, one rounding (torch.nn.RMSNorm) + +ROPE_INTERLEAVED = 0 # rotates channel pairs (2i, 2i + 1) +ROPE_HALF = 1 # rotate_half: channel i pairs with i + rotary_dim / 2 + + +@triton.jit +def _qk_normalize(x, inv, w_ptr, offs, mask, WEIGHT: tl.constexpr): + if WEIGHT == 1: + return _bf16(x * inv * (tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) + 1.0)) + if WEIGHT == 2: + return _bf16(_bf16(x * inv) * tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32)) + if WEIGHT == 3: + return _bf16(x * inv * tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32)) + return _bf16(x * inv) + + @triton.jit def _qk_norm_rope_kernel( src_ptr, src_stride, out_ptr, + out_row_stride, + out_head_stride, w_ptr, cos_ptr, sin_ptr, + rope_stride, rms_eps, HEADS: tl.constexpr, HEAD_DIM: tl.constexpr, + OUT_DIM: tl.constexpr, + DP2: tl.constexpr, + ROT: tl.constexpr, + WEIGHT: tl.constexpr, + STYLE: tl.constexpr, + ROPE_BF16: tl.constexpr, ): row = tl.program_id(0).to(tl.int64) head = tl.program_id(1) - half = tl.arange(0, HEAD_DIM // 2) + offs = tl.arange(0, DP2) + mask = offs < HEAD_DIM base = src_ptr + row * src_stride + head * HEAD_DIM - even = tl.load(base + 2 * half).to(tl.float32) - odd = tl.load(base + 2 * half + 1).to(tl.float32) - mean_sq = libdevice.div_rn( - tl.sum(even * even, axis=0) + tl.sum(odd * odd, axis=0), HEAD_DIM * 1.0 - ) + x = tl.load(base + offs, mask=mask, other=0.0).to(tl.float32) + mean_sq = libdevice.div_rn(tl.sum(x * x, axis=0), HEAD_DIM * 1.0) inv = libdevice.rsqrt(mean_sq + rms_eps) - even = _bf16(even * inv * (tl.load(w_ptr + 2 * half).to(tl.float32) + 1.0)) - odd = _bf16(odd * inv * (tl.load(w_ptr + 2 * half + 1).to(tl.float32) + 1.0)) - cos_even = tl.load(cos_ptr + row * HEAD_DIM + 2 * half) - cos_odd = tl.load(cos_ptr + row * HEAD_DIM + 2 * half + 1) - sin_even = tl.load(sin_ptr + row * HEAD_DIM + 2 * half) - sin_odd = tl.load(sin_ptr + row * HEAD_DIM + 2 * half + 1) - out_even = even * cos_even + (-odd) * sin_even - out_odd = odd * cos_odd + even * sin_odd - dst = out_ptr + (row * HEADS + head) * HEAD_DIM - tl.store(dst + 2 * half, out_even.to(tl.bfloat16)) - tl.store(dst + 2 * half + 1, out_odd.to(tl.bfloat16)) - - -def qk_norm_rope_with_triton(src, heads, head_dim, weight, cos, sin, rms_eps): - """Per-head RMSNorm (``1 + weight`` scale) then interleaved rotary embedding. - - ``src`` is a ``[rows, >= heads * head_dim]`` strided view; returns - ``[1, rows, heads, head_dim]`` bf16. Multiplies and adds are not contracted into FMAs, as - in the eager elementwise ops. + y = _qk_normalize(x, inv, w_ptr, offs, mask, WEIGHT) + if ROT > 0: + rot_mask = offs < ROT + if STYLE == 0: + partner = offs ^ 1 + sign = tl.where(offs % 2 == 0, -1.0, 1.0) + else: + partner = tl.where(offs < ROT // 2, offs + ROT // 2, offs - ROT // 2) + sign = tl.where(offs < ROT // 2, -1.0, 1.0) + xp = tl.load(base + partner, mask=rot_mask, other=0.0).to(tl.float32) + yp = _qk_normalize(xp, inv, w_ptr, partner, rot_mask, WEIGHT) + cos = tl.load(cos_ptr + row * rope_stride + offs, mask=rot_mask, other=0.0).to(tl.float32) + sin = tl.load(sin_ptr + row * rope_stride + offs, mask=rot_mask, other=0.0).to(tl.float32) + if ROPE_BF16: + rotated = _bf16(_bf16(y * _bf16(cos)) + _bf16(sign * yp * _bf16(sin))) + else: + rotated = y * cos + sign * yp * sin + y = tl.where(rot_mask, rotated, y) + y = tl.where(mask, y, 0.0) + dst = out_ptr + row * out_row_stride + head * out_head_stride + tl.store(dst + offs, y.to(tl.bfloat16), mask=offs < OUT_DIM) + + +def qk_norm_rope_with_triton( + src, + heads, + head_dim, + weight, + cos, + sin, + rms_eps, + *, + weight_mode=QK_WEIGHT_ONE_PLUS, + style=ROPE_INTERLEAVED, + rotary_dim=None, + rope_bf16=False, + out=None, + out_dim=None, +): + """Per-head RMSNorm then rotary embedding. + + ``src`` is a ``[rows, >= heads * head_dim]`` strided view; ``cos``/``sin`` are + ``[rows, rotary_dim]`` (``None`` skips the rotation). Channels past ``rotary_dim`` pass + through. Returns ``[1, rows, heads, out_dim]`` bf16 (``out_dim >= head_dim`` zero-pads each + head); ``out`` may be a row slice of a larger buffer. Multiplies and adds are not contracted + into FMAs, as in the eager elementwise ops. """ rows = src.shape[0] - out = torch.empty((1, rows, heads, head_dim), device=src.device, dtype=torch.bfloat16) + out_dim = head_dim if out_dim is None else int(out_dim) + if out is None: + out = torch.empty((1, rows, heads, out_dim), device=src.device, dtype=torch.bfloat16) + rotary = 0 if cos is None else int(rotary_dim or cos.shape[-1]) if rows: _qk_norm_rope_kernel[(rows, heads)]( src, src.stride(0), out, - weight, - cos, - sin, + out.stride(1), + out.stride(2), + weight if weight is not None else src, + cos if cos is not None else src, + sin if sin is not None else src, + cos.stride(0) if cos is not None else 0, float(rms_eps), HEADS=heads, HEAD_DIM=head_dim, - num_warps=1, + OUT_DIM=out_dim, + DP2=triton.next_power_of_2(max(head_dim, out_dim)), + ROT=rotary, + WEIGHT=int(weight_mode) if weight is not None else QK_WEIGHT_NONE, + STYLE=int(style), + ROPE_BF16=bool(rope_bf16), + num_warps=1 if head_dim <= 128 else 2, enable_fp_fusion=False, ) return out +@triton.jit +def _postnorm_residual_kernel( + res_ptr, + x_ptr, + w_ptr, + gate_ptr, + index_ptr, + out_ptr, + res_stride, + x_stride, + gate_stride, + D: tl.constexpr, + DP2: tl.constexpr, + rms_eps, + NORM: tl.constexpr, + ROW_INDEXED: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, DP2) + mask = offs < D + x = tl.load(x_ptr + row * x_stride + offs, mask=mask, other=0.0).to(tl.float32) + n = _normalize(x, w_ptr, offs, mask, D, rms_eps, NORM) + gate_row = tl.load(index_ptr + row).to(tl.int64) * gate_stride if ROW_INDEXED else 0 + gate = tl.load(gate_ptr + gate_row + offs, mask=mask, other=0.0).to(tl.float32) + res = tl.load(res_ptr + row * res_stride + offs, mask=mask, other=0.0).to(tl.float32) + tl.store(out_ptr + row * res_stride + offs, (res + _bf16(gate * n)).to(tl.bfloat16), mask=mask) + + +def postnorm_residual_with_triton( + residual, x, norm_weight, gate, rms_eps, *, norm=NORM_RMS, gate_index=None, out=None +): + """``residual + gate * norm(x)`` row by row (sandwich-norm blocks), in place by default. + + ``gate`` is a ``[dim]`` vector, or a ``[rows, dim]`` table addressed by ``gate_index``. + """ + dim = x.shape[-1] + rows = x.numel() // dim + x = x.reshape(rows, dim) + if x.stride(1) != 1: + x = x.contiguous() + res = residual.reshape(rows, dim) + out = res if out is None else out.reshape(rows, dim) + gate = gate.reshape(-1, dim).contiguous() + if rows: + _postnorm_residual_kernel[(rows,)]( + res, + x, + norm_weight if norm_weight is not None else x, + gate, + gate_index if gate_index is not None else gate, + out, + res.stride(0), + x.stride(0), + gate.stride(0), + D=dim, + DP2=triton.next_power_of_2(dim), + rms_eps=float(rms_eps), + NORM=int(norm), + ROW_INDEXED=gate_index is not None, + num_warps=8, + enable_fp_fusion=False, + ) + return residual + + @triton.jit def _decode_w4_kernel(packed_ptr, codes_ptr, out_ptr, total_bytes, BLOCK: tl.constexpr): offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) @@ -313,13 +752,48 @@ def _decode_w4_kernel(packed_ptr, codes_ptr, out_ptr, total_bytes, BLOCK: tl.con tl.store(out_ptr + out_offs, pair.to(tl.int8), mask=out_offs < 2 * total_bytes) -def decode_w4_to_int8_with_triton(packed, codes, out_features, in_features, out=None): - """Packed W4 indices (low nibble first) -> INT8 surrogate weight ``[out, in]``.""" +@triton.jit +def _decode_lowbit_kernel( + packed_ptr, + codes_ptr, + out_ptr, + total_values, + total_bytes, + BITS: tl.constexpr, + BLOCK: tl.constexpr, +): + offs = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offs < total_values + bit = offs * BITS + first = bit >> 3 + shift = (bit & 7).to(tl.int32) + low = tl.load(packed_ptr + first, mask=mask, other=0).to(tl.int32) + high = tl.load(packed_ptr + first + 1, mask=mask & (first + 1 < total_bytes), other=0) + word = low | (high.to(tl.int32) << 8) + index = (word >> shift) & ((1 << BITS) - 1) + tl.store(out_ptr + offs, tl.load(codes_ptr + index).to(tl.int8), mask=mask) + + +def decode_lowbit_to_int8_with_triton(packed, codes, bits, out_features, in_features, out=None): + """Packed low-bit indices (LSB-first bit stream, ``orbitquant.packing``) -> INT8 surrogate + weight ``[out, in]``.""" if out is None: out = torch.empty((out_features, in_features), device=packed.device, dtype=torch.int8) - total = out_features * in_features // 2 - block = 1024 - _decode_w4_kernel[(triton.cdiv(total, block),)]( - packed.reshape(-1), codes, out, total, BLOCK=block + if bits == 4: + total = out_features * in_features // 2 + block = 1024 + _decode_w4_kernel[(triton.cdiv(total, block),)]( + packed.reshape(-1), codes, out, total, BLOCK=block + ) + return out + values = out_features * in_features + block = 2048 + _decode_lowbit_kernel[(triton.cdiv(values, block),)]( + packed.reshape(-1), codes, out, values, packed.numel(), BITS=int(bits), BLOCK=block ) return out + + +def decode_w4_to_int8_with_triton(packed, codes, out_features, in_features, out=None): + """Packed W4 indices (low nibble first) -> INT8 surrogate weight ``[out, in]``.""" + return decode_lowbit_to_int8_with_triton(packed, codes, 4, out_features, in_features, out=out) diff --git a/src/orbitquant/kernels/triton_attention.py b/src/orbitquant/kernels/triton_attention.py index 69df6a5..7f854fb 100644 --- a/src/orbitquant/kernels/triton_attention.py +++ b/src/orbitquant/kernels/triton_attention.py @@ -10,10 +10,12 @@ kernel exponentiates with ``exp2``. * The online softmax runs in FP32. * P and V are FP16; each ``BLOCK_N``-token tile of ``P @ V`` accumulates in FP16 and is added - to an FP32 accumulator. + to an FP32 accumulator. V is converted once up front: converting each tile inside the loop + costs about 10% of the kernel time. Tensors use the ``[batch=1, tokens, heads, head_dim]`` layout of DiT projections; rows may be -strided (e.g. a column slice of a fused Q|K|V GEMM output). +strided (e.g. a column slice of a fused Q|K|V GEMM output). Keys and values may be a different +(e.g. compacted, padding-free) sequence than the queries. """ from __future__ import annotations @@ -103,6 +105,7 @@ def _attention_kernel( k_scale_ptr, o_ptr, seq, + seq_kv, stride_qr, stride_qh, stride_kr, @@ -125,7 +128,7 @@ def _attention_kernel( q_rows = offs_m[:, None].to(tl.int64) * stride_qr + head * stride_qh q = tl.load(q_ptr + q_rows + offs_d[None, :], mask=offs_m[:, None] < seq, other=0) q_scale = tl.load(q_scale_ptr + head * tl.cdiv(seq, BLOCK_M) + start_m) - k_scales = k_scale_ptr + kv_head * tl.cdiv(seq, BLOCK_N) + k_scales = k_scale_ptr + kv_head * tl.cdiv(seq_kv, BLOCK_N) k_ptrs = k_ptr + offs_n[None, :].to(tl.int64) * stride_kr + kv_head * stride_kh k_ptrs += offs_d[:, None] v_ptrs = v_ptr + offs_n[:, None].to(tl.int64) * stride_vr + kv_head * stride_vh @@ -133,7 +136,7 @@ def _attention_kernel( m_i = tl.full([BLOCK_M], float("-inf"), tl.float32) l_i = tl.zeros([BLOCK_M], tl.float32) acc = tl.zeros([BLOCK_M, D], tl.float32) - full = seq // BLOCK_N + full = seq_kv // BLOCK_N for block in range(0, full): qk_scale = q_scale * tl.load(k_scales + block) acc, l_i, m_i = _attention_tile( @@ -141,10 +144,10 @@ def _attention_kernel( ) k_ptrs += BLOCK_N * stride_kr v_ptrs += BLOCK_N * stride_vr - if full * BLOCK_N < seq: + if full * BLOCK_N < seq_kv: qk_scale = q_scale * tl.load(k_scales + full) acc, l_i, m_i = _attention_tile( - acc, l_i, m_i, q, qk_scale, k_ptrs, v_ptrs, seq - full * BLOCK_N, True, BLOCK_N + acc, l_i, m_i, q, qk_scale, k_ptrs, v_ptrs, seq_kv - full * BLOCK_N, True, BLOCK_N ) out = acc / l_i[:, None] o_rows = offs_m[:, None].to(tl.int64) * stride_or + head * stride_oh @@ -157,10 +160,15 @@ def _attention_kernel( _attention = autotune( [triton.Config({}, num_warps=w, num_stages=s) for w in (4, 8) for s in (2, 3, 4)], - ["GROUPS", "D"], + ["GROUPS", "D", "BLOCK_M"], )(_attention_kernel) +def _query_block(dim: int) -> int: + # The FP32 accumulator is BLOCK_M x D: wide heads take shorter query tiles. + return BLOCK_Q if dim <= 128 else 64 + + def _quantize_blocks(x, block, multiplier, mean=None): rows, heads, dim = x.shape out = torch.empty((rows, heads, dim), device=x.device, dtype=torch.int8) @@ -191,26 +199,28 @@ def int8_attention_with_triton( *, sm_scale: float | None = None, ) -> torch.Tensor: - """``softmax(q k^T * sm_scale) v`` for ``[1, S, Hq, D]`` queries and ``[1, S, Hkv, D]`` + """``softmax(q k^T * sm_scale) v`` for ``[1, S, Hq, D]`` queries and ``[1, Skv, Hkv, D]`` keys/values (``Hq`` a multiple of ``Hkv``, ``D`` a power of two, last dim contiguous). - Returns ``[1, S, Hq, D]`` in the query dtype.""" + Returns ``[1, S, Hq, D]`` in the query dtype. A head padded with zero channels to a power of + two gives the unpadded result in its leading channels when ``sm_scale`` is the unpadded one.""" if query.dim() != 4 or query.shape[0] != 1: raise ValueError("expected [1, tokens, heads, head_dim] tensors") q, k, v = query[0], key[0], value[0] seq, heads, dim = q.shape - kv_heads = k.shape[1] + seq_kv, kv_heads = k.shape[0], k.shape[1] if heads % kv_heads: raise ValueError(f"{heads} query heads are not a multiple of {kv_heads} kv heads") if dim & (dim - 1) or q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1: raise ValueError("head_dim must be a power of two and contiguous") if sm_scale is None: sm_scale = 1.0 / math.sqrt(dim) - q_codes, q_scales = _quantize_blocks(q, BLOCK_Q, sm_scale * LOG2E) - k_codes, k_scales = _quantize_blocks(k, BLOCK_KV, 1.0, k.float().mean(dim=0).contiguous()) - v16 = v.to(torch.float16) + block_q = _query_block(dim) + q_codes, q_scales = _quantize_blocks(q, block_q, sm_scale * LOG2E) + k_codes, k_scales = _quantize_blocks(k, BLOCK_KV, 1.0, k.mean(dim=0, dtype=torch.float32)) + v16 = v if v.dtype == torch.float16 else v.to(torch.float16) out = torch.empty((1, seq, heads, dim), device=q.device, dtype=query.dtype) o = out[0] - _attention[(triton.cdiv(seq, BLOCK_Q), heads)]( + _attention[(triton.cdiv(seq, block_q), heads)]( q_codes, k_codes, v16, @@ -218,6 +228,7 @@ def int8_attention_with_triton( k_scales, o, seq, + seq_kv, q_codes.stride(0), q_codes.stride(1), k_codes.stride(0), @@ -228,7 +239,7 @@ def int8_attention_with_triton( o.stride(1), GROUPS=heads // kv_heads, D=dim, - BLOCK_M=BLOCK_Q, + BLOCK_M=block_q, BLOCK_N=BLOCK_KV, ) return out diff --git a/src/orbitquant/kernels/triton_int8_gemm.py b/src/orbitquant/kernels/triton_int8_gemm.py index 559546d..2a46fa1 100644 --- a/src/orbitquant/kernels/triton_int8_gemm.py +++ b/src/orbitquant/kernels/triton_int8_gemm.py @@ -14,8 +14,10 @@ * ``EPILOGUE_SWIGLU``: rows of ``b`` interleave gate and up in chunks of ``BLOCK_N // 2``; the kernel writes ``bf16(silu(bf16(gate))) * bf16(up)`` (half as many columns). -* ``EPILOGUE_RESIDUAL_GATE``: ``out = res + bf16(gate[n] * bf16(y))`` (in place allowed). +* ``EPILOGUE_RESIDUAL_GATE``: ``out = res + bf16(gate[n] * bf16(y))`` (in place allowed); the + gate may be selected per row from a table (``gate_index``). * ``EPILOGUE_SIGMOID_TAIL``: columns ``>= sig_from`` store ``sigmoid(bf16(y))``. +* ``EPILOGUE_GELU_TANH``: ``gelu_tanh(bf16(y))``. """ from __future__ import annotations @@ -34,6 +36,7 @@ EPILOGUE_SWIGLU = 1 EPILOGUE_RESIDUAL_GATE = 2 EPILOGUE_SIGMOID_TAIL = 3 +EPILOGUE_GELU_TANH = 4 # The SwiGLU epilogue pairs gate/up halves inside one N tile; weights packed for it interleave # rows in chunks of this size. @@ -89,17 +92,20 @@ def _int8_scaled_gemm_body( bias_ptr, res_ptr, gate_ptr, + gate_index_ptr, M, N, K, stride_om, col_offset, + gate_stride, M_BUCKET, alpha, sig_from, HAS_BIAS: tl.constexpr, SCALE_MODE: tl.constexpr, EPILOGUE: tl.constexpr, + GATE_INDEXED: tl.constexpr, EVEN_K: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, @@ -158,13 +164,23 @@ def _int8_scaled_gemm_body( elif EPILOGUE == 2: cols = (col_offset + offs_n)[None, :] mask = mask_m[:, None] & mask_n[None, :] - gate = tl.load(gate_ptr + col_offset + offs_n, mask=mask_n, other=0.0).to(tl.float32) + if GATE_INDEXED: + gate_rows = tl.load(gate_index_ptr + offs_m, mask=mask_m, other=0).to(tl.int64) + gate = tl.load( + gate_ptr + gate_rows[:, None] * gate_stride + cols, mask=mask, other=0.0 + ).to(tl.float32) + else: + gate = tl.load(gate_ptr + col_offset + offs_n, mask=mask_n, other=0.0).to(tl.float32) + gate = gate[None, :] res = tl.load(res_ptr + rows + cols, mask=mask, other=0.0).to(tl.float32) - out = res + _bf16(gate[None, :] * _bf16(y)) + out = res + _bf16(gate * _bf16(y)) tl.store(out_ptr + rows + cols, out.to(out_ptr.dtype.element_ty), mask=mask) - elif EPILOGUE == 3: + elif EPILOGUE == 3 or EPILOGUE == 4: y = _bf16(y) - if pid_n * BLOCK_N >= sig_from: + if EPILOGUE == 4: + inner = 0.7978845608028654 * (y + 0.044715 * y * y * y) + y = 0.5 * y * (1.0 + libdevice.tanh(inner)) + elif pid_n * BLOCK_N >= sig_from: y = libdevice.div_rn(1.0, 1.0 + libdevice.exp(-y)) tl.store( out_ptr + rows + (col_offset + offs_n)[None, :], @@ -179,7 +195,7 @@ def _int8_scaled_gemm_body( ) -_KEY = ["M_BUCKET", "N", "K", "SCALE_MODE", "HAS_BIAS", "EPILOGUE"] +_KEY = ["M_BUCKET", "N", "K", "SCALE_MODE", "HAS_BIAS", "EPILOGUE", "GATE_INDEXED"] # The residual-gate epilogue usually updates its output in place: every benchmark run of the # autotuner would add the projection to the residual again. _int8_scaled_gemm = autotune(_configs(), _KEY, restore_value=["out_ptr"])(_int8_scaled_gemm_body) @@ -207,6 +223,7 @@ def matmul_int8_scaled_with_triton( epilogue: int = EPILOGUE_PLAIN, residual: torch.Tensor | None = None, gate: torch.Tensor | None = None, + gate_index: torch.Tensor | None = None, sig_from: int = 0, ) -> torch.Tensor: """Scaled ``a @ b.T`` for INT8 ``a [M, K]`` and ``b [N, K]``. @@ -214,7 +231,8 @@ def matmul_int8_scaled_with_triton( ``out`` may be a wider row-major buffer; the result lands in columns ``[col_offset, col_offset + N)`` (``N // 2`` for SwiGLU) so chunked weight decodes write in place. The residual-gate epilogue reads ``residual`` with the same layout as ``out`` and may - alias it. + alias it; its ``gate`` is a vector over the output columns, or with ``gate_index`` a row-major + table whose row ``gate_index[m]`` gates output row ``m``. """ if a.dtype != torch.int8 or b.dtype != torch.int8: raise ValueError("matmul_int8_scaled_with_triton expects INT8 operands") @@ -243,6 +261,8 @@ def matmul_int8_scaled_with_triton( def grid(meta): return (triton.cdiv(rows, meta["BLOCK_M"]) * triton.cdiv(n, meta["BLOCK_N"]),) + if gate_index is not None and gate.stride(-1) != 1: + gate = gate.contiguous() kernel[grid]( a, b, @@ -252,17 +272,20 @@ def grid(meta): bias if bias is not None else a_scale, residual if residual is not None else out, gate if gate is not None else a_scale, + gate_index if gate_index is not None else a_scale, rows, n, k, out.stride(0), int(col_offset), + gate.stride(0) if gate_index is not None else 0, _row_bucket(rows), float(alpha), int(sig_from), HAS_BIAS=bias is not None, SCALE_MODE=int(scale_mode), EPILOGUE=int(epilogue), + GATE_INDEXED=gate_index is not None, EVEN_K=k % 128 == 0, ) return out diff --git a/src/orbitquant/quantizer.py b/src/orbitquant/quantizer.py index cc247ce..c6d7ac5 100644 --- a/src/orbitquant/quantizer.py +++ b/src/orbitquant/quantizer.py @@ -91,6 +91,10 @@ def _module_and_tensor_name(model: Any, param_name: str) -> tuple[Any, str]: def _is_prequantized_state_tensor(module: Any, tensor_name: str) -> bool: + from orbitquant.fused.groups import GROUP_TYPES + + if isinstance(module, tuple(GROUP_TYPES.values())): + return tensor_name in module._buffers if isinstance(module, OrbitQuantLinear): return tensor_name in _ORBITQUANT_STATE_TENSORS and ( tensor_name in module._parameters or tensor_name in module._buffers @@ -544,6 +548,14 @@ def _process_model_before_weight_loading(self, model: Any, *args: Any, **kwargs: model, self.quantization_config, ) + if self.pre_quantized and self.quantization_config.fused_layout: + from orbitquant.fused import prepare_skeleton + + prepare_skeleton( + model, + self.quantization_config.fused_layout, + config=self.quantization_config, + ) _mark_quantized_gpt_projection_parents_initialized(model) if self._diffusers_streaming_quantization: self._diffusers_model = model @@ -551,6 +563,10 @@ def _process_model_before_weight_loading(self, model: Any, *args: Any, **kwargs: return model def _process_model_after_weight_loading(self, model: Any, *args: Any, **kwargs: Any) -> Any: + if self.pre_quantized and self.quantization_config.fused_layout: + from orbitquant.fused import finalize + + finalize(model) if ( not self.pre_quantized and not self._transformers_streaming_quantization diff --git a/tests/test_fused_family_kernels.py b/tests/test_fused_family_kernels.py new file mode 100644 index 0000000..f4b15f1 --- /dev/null +++ b/tests/test_fused_family_kernels.py @@ -0,0 +1,302 @@ +import pytest +import torch +import torch.nn.functional as F + +from orbitquant.kernels import available_backends + + +def _require_cuda(): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + +def _quantizer(dim, bits=4): + from orbitquant.codebooks import get_codebook + from orbitquant.kernels import triton_activation as act + from orbitquant.rotations import get_rpbh_rotation + + rotation = get_rpbh_rotation(dim=dim, seed=0, block_size="paper") + return act.ActivationQuantizer( + rotation, get_codebook(dim, bits, 2), 1e-10, torch.device("cuda") + ) + + +def _bf(t): + return t.to(torch.bfloat16) + + +@pytest.mark.parametrize("norm", ["layer", "rms", "rms_torch"]) +@pytest.mark.parametrize("indexed", [False, True]) +def test_norm_modulation_prologues_match_eager(norm, indexed): + _require_cuda() + from orbitquant.kernels import triton_activation as act + + torch.manual_seed(0) + rows, dim = 77, 1024 + x = torch.randn(rows, dim, device="cuda", dtype=torch.bfloat16) + weight = _bf(torch.randn(dim, device="cuda") * 0.1 + 1.0) + table_rows = 3 if indexed else 1 + scale = _bf(torch.randn(table_rows, dim, device="cuda") * 0.1) + shift = _bf(torch.randn(table_rows, dim, device="cuda") * 0.1) + index = (torch.arange(rows, device="cuda") % 3).to(torch.int32) if indexed else None + if norm == "layer": + normed = F.layer_norm(x, (dim,), eps=1e-6) + mode, eps = act.NORM_LAYER, 1e-6 + elif norm == "rms": + inv = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-5) + normed = _bf(x * inv) * weight + mode, eps = act.NORM_RMS, 1e-5 + else: + normed = F.rms_norm(x, (dim,), weight, eps=1e-5) + mode, eps = act.NORM_RMS_TORCH, 1e-5 + s = scale[index.long()] if indexed else scale + t = shift[index.long()] if indexed else shift + expected = (1 + s) * normed + t + quantizer = _quantizer(dim) + actual, _ = quantizer( + x, + norm=mode, + mod=act.MOD_SHIFT_SCALE, + norm_weight=weight, + mod_scale=scale if indexed else scale.reshape(-1), + mod_shift=shift if indexed else shift.reshape(-1), + mod_index=index, + rms_eps=eps, + prologue_only=True, + ) + diff = (actual.float() - expected.float()).abs() + # One bf16 ulp where the reduction order differs from the eager kernels. + assert diff.max().item() <= 2**-6 * expected.float().abs().max().item() + assert (diff > 0).float().mean().item() < 0.01 + + +@pytest.mark.parametrize( + ("style", "head_dim", "rotary", "out_dim"), + [ + ("interleaved", 128, 128, 128), + ("half", 256, 256, 256), + ("half", 128, 96, 128), + ("half", 120, 120, 128), + ], +) +def test_qk_norm_rope_styles_match_eager(style, head_dim, rotary, out_dim): + _require_cuda() + from orbitquant.kernels import triton_activation as act + + torch.manual_seed(1) + rows, heads = 65, 3 + src = torch.randn(rows, heads * head_dim + 16, device="cuda", dtype=torch.bfloat16) + weight = _bf(torch.randn(head_dim, device="cuda") * 0.1 + 1.0) + angles = torch.rand(rows, rotary // 2, device="cuda") * 6.0 + if style == "interleaved": + cos = angles.cos().repeat_interleave(2, -1) + sin = angles.sin().repeat_interleave(2, -1) + else: + cos = torch.cat([angles.cos(), angles.cos()], -1) + sin = torch.cat([angles.sin(), angles.sin()], -1) + x = src[:, : heads * head_dim].unflatten(-1, (heads, head_dim)) + normed = F.rms_norm(x, (head_dim,), weight, eps=1e-6) + rot, rest = normed[..., :rotary].float(), normed[..., rotary:] + c, s = cos[:, None], sin[:, None] + if style == "interleaved": + pairs = rot.reshape(*rot.shape[:-1], -1, 2) + rotated = torch.stack([-pairs[..., 1], pairs[..., 0]], -1).flatten(-2) + out = rot * c + rotated * s + else: + half = rotary // 2 + rotated = torch.cat([-rot[..., half:], rot[..., :half]], -1) + out = rot * c + rotated * s + expected = torch.cat([out.to(torch.bfloat16), rest], -1) + actual = act.qk_norm_rope_with_triton( + src, + heads, + head_dim, + weight, + cos, + sin, + 1e-6, + weight_mode=act.QK_WEIGHT_TORCH, + style=act.ROPE_INTERLEAVED if style == "interleaved" else act.ROPE_HALF, + out_dim=out_dim, + )[0] + assert actual.shape == (rows, heads, out_dim) + torch.testing.assert_close( + actual[..., :head_dim].float(), expected.float(), atol=2e-2, rtol=1e-2 + ) + if out_dim > head_dim: + assert torch.count_nonzero(actual[..., head_dim:]) == 0 + + +def test_postnorm_residual_matches_eager(): + _require_cuda() + from orbitquant.kernels import triton_activation as act + + torch.manual_seed(2) + rows, dim = 50, 768 + res = torch.randn(rows, dim, device="cuda", dtype=torch.bfloat16) + x = torch.randn(rows, dim, device="cuda", dtype=torch.bfloat16) + weight = _bf(torch.randn(dim, device="cuda") * 0.1 + 1.0) + gate = torch.tanh(torch.randn(dim, device="cuda")).to(torch.bfloat16) + inv = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-5) + expected = res + gate * (_bf(x * inv) * weight) + actual = act.postnorm_residual_with_triton(res.clone(), x, weight, gate, 1e-5) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=1e-2) + + +@pytest.mark.parametrize("bits", [2, 3, 4]) +def test_lowbit_decode_matches_unpack(bits): + _require_cuda() + from orbitquant.kernels import triton_activation as act + from orbitquant.packing import pack_lowbit, unpack_lowbit + + torch.manual_seed(3) + out_features, in_features = 96, 512 + indices = torch.randint(0, 1 << bits, (out_features * in_features,), dtype=torch.uint8) + packed = pack_lowbit(indices, bits=bits).cuda() + codes = torch.randint(-127, 128, (1 << bits,), dtype=torch.int8, device="cuda") + decoded = act.decode_lowbit_to_int8_with_triton(packed, codes, bits, out_features, in_features) + expected = codes[unpack_lowbit(packed.cpu(), bits=bits, length=indices.numel()).long().cuda()] + assert torch.equal(decoded.reshape(-1), expected) + + +def test_gelu_and_indexed_gate_epilogues_match_eager(): + _require_cuda() + from orbitquant.kernels import triton_int8_gemm as gemm + + torch.manual_seed(4) + rows, k, n = 70, 256, 384 + a = torch.randint(-127, 128, (rows, k), device="cuda", dtype=torch.int8) + b = torch.randint(-127, 128, (n, k), device="cuda", dtype=torch.int8) + a_s = torch.rand(rows, device="cuda") * 1e-3 + 1e-4 + b_s = torch.rand(n, device="cuda") + 0.5 + y = ((a.float() @ b.float().T) * a_s[:, None] * b_s[None, :] * 0.5).to(torch.bfloat16) + gelu = gemm.matmul_int8_scaled_with_triton( + a, b, a_s, b_s, alpha=0.5, epilogue=gemm.EPILOGUE_GELU_TANH + ) + torch.testing.assert_close( + gelu.float(), F.gelu(y.float(), approximate="tanh"), atol=2e-2, rtol=2e-2 + ) + table = torch.randn(4, n, device="cuda", dtype=torch.bfloat16) + index = (torch.arange(rows, device="cuda") % 4).to(torch.int32) + residual = torch.randn(rows, n, device="cuda", dtype=torch.bfloat16) + out = gemm.matmul_int8_scaled_with_triton( + a, + b, + a_s, + b_s, + alpha=0.5, + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=residual.clone(), + gate=table, + gate_index=index, + ) + expected = residual + table[index.long()] * y + torch.testing.assert_close(out.float(), expected.float(), atol=5e-2, rtol=2e-2) + + +@pytest.mark.parametrize(("head_dim", "kv_len"), [(128, 300), (256, 700)]) +def test_int8_attention_handles_wide_heads_and_separate_kv(head_dim, kv_len): + _require_cuda() + from orbitquant.kernels.triton_attention import int8_attention_with_triton + + torch.manual_seed(5) + q = torch.randn(1, 513, 4, head_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, kv_len, 2, head_dim, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, kv_len, 2, head_dim, device="cuda", dtype=torch.bfloat16) + out = int8_attention_with_triton(q, k, v) + ref = F.scaled_dot_product_attention( + q.transpose(1, 2).float(), + k.transpose(1, 2).float(), + v.transpose(1, 2).float(), + enable_gqa=True, + ).transpose(1, 2) + error = (out.float() - ref).norm() / ref.norm() + assert error.item() < 5e-2 + + +def _projection(name, rows, bits, dim=1024): + from orbitquant import OrbitQuantConfig + from orbitquant.layers import OrbitQuantLinear + + linear = torch.nn.Linear(dim, rows, bias=False, dtype=torch.bfloat16) + config = OrbitQuantConfig(weight_bits=bits, activation_bits=4, codebook_version=2) + return OrbitQuantLinear.from_linear(linear, config=config, module_name=name).cuda() + + +def test_mixed_width_group_is_one_gemm_matching_its_projections(): + _require_cuda() + from orbitquant.fused.groups import PackedGroup, RowSource + from orbitquant.fused.runtime import FusedRuntime + + torch.manual_seed(6) + root = torch.nn.Module() + root.q, root.k, root.v = ( + _projection("q", 256, 2), + _projection("k", 128, 2), + _projection("v", 128, 3), + ) + sources = [RowSource("q"), RowSource("k"), RowSource("v")] + group = PackedGroup.from_sources(root, sources).cuda() + assert group.segments == ((384, 2), (128, 3)) + runtime = FusedRuntime( + rotation_seed=0, block_size="paper", codebook_version=2, activation_eps=1e-10 + ) + x = torch.randn(64, 1024, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + fused = runtime.linear(group, runtime.input(x, 1024)) + reference = torch.cat([root.q(x), root.k(x), root.v(x)], dim=1) + error = (fused.float() - reference.float()).norm() / reference.float().norm() + assert error.item() < 2e-2 + + skeleton = PackedGroup.empty( + PackedGroup.source_segments([root.q, root.k, root.v], sources), + 1024, + activation_bits=4, + codebook_version=2, + device="meta", + ) + skeleton.load_state_dict(group.state_dict(), assign=True) + skeleton.cuda() + with torch.no_grad(): + again = runtime.linear(skeleton, runtime.input(x, 1024)) + assert torch.equal(again, fused) + + +@pytest.mark.parametrize("dim", [3360, 13568, 6144, 8192, 16384]) +def test_absmax_int8_activations_match_rotated_reference(dim): + _require_cuda() + from orbitquant.kernels import triton_activation as act + from orbitquant.rotations import get_rpbh_rotation + + torch.manual_seed(7) + rotation = get_rpbh_rotation(dim=dim, seed=0, block_size="paper") + quantizer = act.ActivationQuantizer(rotation, None, 1e-10, torch.device("cuda")) + x = torch.randn(37, dim, device="cuda", dtype=torch.bfloat16) + x[:, 5] *= 40 # an outlier channel + codes, scales = quantizer(x) + rotated = rotation.apply_to_activations(x.float()) + expected_scales = rotated.abs().amax(-1) / 127 + torch.testing.assert_close(scales, expected_scales, rtol=1e-5, atol=0) + expected = torch.round(rotated / expected_scales[:, None]).clamp(-127, 127) + assert (codes.float() - expected).abs().max().item() <= 1 + assert (codes.float() != expected).float().mean().item() < 1e-3 + + +@pytest.mark.parametrize("dim", [3360, 8192, 16384]) +def test_lloyd_max_activation_codes_match_the_rotated_reference(dim): + _require_cuda() + from orbitquant.codebooks import get_codebook + from orbitquant.kernels import triton_activation as act + from orbitquant.rotations import get_rpbh_rotation + + torch.manual_seed(8) + rotation = get_rpbh_rotation(dim=dim, seed=0, block_size="paper") + codebook = get_codebook(dim, 4, 2) + quantizer = act.ActivationQuantizer(rotation, codebook, 1e-10, torch.device("cuda")) + x = torch.randn(9, dim, device="cuda", dtype=torch.bfloat16) + codes, _ = quantizer(x) + rotated = rotation.apply_to_activations(x.float()) + unit = rotated / (rotated.norm(dim=-1, keepdim=True) + 1e-10) + index = torch.bucketize(unit, codebook.boundaries.to("cuda", torch.float32)) + expected = quantizer.codes[index] + assert (codes != expected).float().mean().item() < 1e-3 diff --git a/tests/test_fused_layout.py b/tests/test_fused_layout.py new file mode 100644 index 0000000..73595da --- /dev/null +++ b/tests/test_fused_layout.py @@ -0,0 +1,25 @@ +import subprocess +import sys + +from orbitquant.fused.groups import PackedGroup, _segments + + +def test_segments_merge_runs_of_one_weight_width(): + assert _segments([(256, 2), (128, 2), (128, 3)]) == ((384, 2), (128, 3)) + assert _segments([(128, 4), (128, 4)]) == ((256, 4),) + + +def test_packed_rows_are_byte_rows_for_one_width_and_a_flat_run_for_mixed_widths(): + assert PackedGroup.packed_shape(((384, 4),), 1024) == (384, 512) + assert PackedGroup.packed_shape(((384, 2), (128, 3)), 1024) == (384 * 256 + 128 * 384,) + + +def test_recognizing_fused_checkpoints_does_not_need_the_triton_runtime(): + # The HF quantizer imports the group types for every checkpoint tensor it checks. + code = ( + "import sys, orbitquant.fused, orbitquant.fused.groups;" + "print(any(name.startswith('orbitquant.kernels.triton') for name in sys.modules))" + ) + command = [sys.executable, "-c", code] + result = subprocess.run(command, capture_output=True, text=True, check=True) + assert result.stdout.strip() == "False" diff --git a/uv.lock b/uv.lock index 841db14..c2ed311 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.10.0" +version = "0.11.0" source = { editable = "." } dependencies = [ { name = "huggingface-hub" }, From 2c473760a109f908576fb361ba560371c4862dd7 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Fri, 2 Oct 2026 02:53:40 +0200 Subject: [PATCH 2/4] Stream fused weights with diffusers group offloading and run Krea 2 blocks on valid rows Fused groups keep their weights as frozen parameters instead of buffers: diffusers' streamed group offloading moves only parameters back to the host, so block-level streaming kept every fused block on the GPU and ran out of memory under a VRAM cap. Krea 2 blocks run on the rows of valid tokens only (padded text rows are never attended to and are dropped at the output), find those rows once per forward, and the text fusion stack runs once per prompt instead of in every denoising step. --- src/orbitquant/fused/families/__init__.py | 3 +- src/orbitquant/fused/families/krea2.py | 72 +++++++++++++++++++++-- src/orbitquant/fused/groups.py | 20 ++++--- src/orbitquant/fused/install.py | 8 +++ src/orbitquant/quantizer.py | 2 +- 5 files changed, 90 insertions(+), 15 deletions(-) diff --git a/src/orbitquant/fused/families/__init__.py b/src/orbitquant/fused/families/__init__.py index 819b6de..dfbefc9 100644 --- a/src/orbitquant/fused/families/__init__.py +++ b/src/orbitquant/fused/families/__init__.py @@ -7,7 +7,8 @@ * ``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). +* ``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 diff --git a/src/orbitquant/fused/families/krea2.py b/src/orbitquant/fused/families/krea2.py index c7f4017..2766999 100644 --- a/src/orbitquant/fused/families/krea2.py +++ b/src/orbitquant/fused/families/krea2.py @@ -7,6 +7,8 @@ from __future__ import annotations +import weakref + import torch from orbitquant.fused.families import int8_attention_allowed, specs @@ -38,6 +40,38 @@ def prepare(block, kind): ) +def install(model): + """The text fusion stack does not depend on the timestep: run it once per prompt instead of + in every denoising step. The last two results are kept (a CFG pipeline alternates two + prompts), each for the same embeddings and mask tensor objects at the same versions, so a new + prompt always recomputes.""" + text_fusion = model.text_fusion + compute = text_fusion.forward + entries = [] + + def stamp(tensor): + return None if tensor is None else (weakref.ref(tensor), tensor._version) + + def matches(held, tensor): + if held is None: + return tensor is None + ref, version = held + return tensor is not None and ref() is tensor and tensor._version == version + + def forward(encoder_hidden_states, attention_mask=None): + if torch.is_grad_enabled(): + return compute(encoder_hidden_states, attention_mask=attention_mask) + for hidden, mask, out in entries: + if matches(hidden, encoder_hidden_states) and matches(mask, attention_mask): + return out + out = compute(encoder_hidden_states, attention_mask=attention_mask) + entries.insert(0, (stamp(encoder_hidden_states), stamp(attention_mask), out)) + del entries[2:] + return out + + text_fusion.forward = forward + + def forward(kind): return _forward @@ -56,9 +90,40 @@ def _forward(self, hidden_states, temb, image_rotary_emb, attention_mask=None): for i in range(hidden_states.shape[0]) ] ) - runtime, fused, attn = self.oq_runtime, self.oq_fused, self.attn _, seq, dim = hidden_states.shape h = hidden_states.reshape(seq, dim) + rows = _valid_rows(attention_mask) + if rows is not None: + # Padded text rows are never attended to and the model drops them at the output, so + # only the valid rows run; the padded ones keep their input values. + cos, sin = image_rotary_emb + h[rows] = _block(self, h[rows], temb, (cos[rows], sin[rows])) + return h.reshape(1, seq, dim) + return _block(self, h, temb, image_rotary_emb).reshape(1, seq, dim) + + +# The last mask seen: (weak reference, version, valid rows). +_LAST_MASK: list = [None] + + +def _valid_rows(attention_mask): + """Indices of the valid rows, or None when every row is valid. Every block of a forward + gets the same mask tensor, so the host synchronization of finding them happens once.""" + if attention_mask is None: + return None + held = _LAST_MASK[0] + if held is not None and held[0]() is attention_mask and held[1] == attention_mask._version: + return held[2] + mask = attention_mask.reshape(-1) + rows = None if bool(mask.all()) else mask.nonzero().squeeze(1) + _LAST_MASK[0] = (weakref.ref(attention_mask), attention_mask._version, rows) + return rows + + +def _block(self, h, temb, image_rotary_emb): + """One block over all-valid rows ``h`` ([rows, dim], updated in place).""" + runtime, fused, attn = self.oq_runtime, self.oq_fused, self.attn + seq, dim = h.shape modulation = temb.reshape(-1).unflatten(-1, (6, -1)) + self.scale_shift_table prescale, preshift, pregate, postscale, postshift, postgate = ( t.reshape(dim).contiguous() for t in modulation.unbind(-2) @@ -91,9 +156,6 @@ def _forward(self, hidden_states, temb, image_rotary_emb, attention_mask=None): qkvg[:, qw : qw + kw], kv_heads, head_dim, attn.norm_k.weight, cos, sin, attn.norm_k.eps ) value = qkvg[:, qw + kw : qw + 2 * kw].unflatten(-1, (kv_heads, head_dim))[None] - if attention_mask is not None and not bool(attention_mask.all()): - valid = attention_mask.reshape(-1).bool() - key, value = key[:, valid], value[:, valid] att = runtime.attention(query, key, value, int8_allowed=self.oq_int8_attention).reshape( seq, dim ) @@ -121,4 +183,4 @@ def _forward(self, hidden_states, temb, image_rotary_emb, attention_mask=None): gate=postgate, out=h, ) - return h.reshape(1, seq, dim) + return h diff --git a/src/orbitquant/fused/groups.py b/src/orbitquant/fused/groups.py index f81780a..d06ffbe 100644 --- a/src/orbitquant/fused/groups.py +++ b/src/orbitquant/fused/groups.py @@ -88,6 +88,15 @@ def _segments(sizes_and_bits) -> tuple[tuple[int, int], ...]: return tuple((rows, bits) for rows, bits in runs) +def _register_weights(module: nn.Module, **tensors: torch.Tensor | None) -> None: + # Frozen parameters, not buffers: diffusers' streamed group offloading moves only the + # parameters of an offloaded module back to the host. + for name, tensor in tensors.items(): + module.register_parameter( + name, None if tensor is None else nn.Parameter(tensor, requires_grad=False) + ) + + class PackedGroup(nn.Module): kind = "packed" @@ -112,9 +121,7 @@ def __init__( raise ValueError(f"{tuple(row_norms.shape)} row norms for {self.out_features} rows") if tuple(packed.shape) != self.packed_shape(self.segments, self.in_features): raise ValueError(f"packed weights {tuple(packed.shape)} do not match {self.segments}") - self.register_buffer("packed", packed) - self.register_buffer("row_norms", row_norms) - self.register_buffer("bias", bias) + _register_weights(self, packed=packed, row_norms=row_norms, bias=bias) constants = { bits: surrogate_constants( self.in_features, bits, self.activation_bits, self.codebook_version @@ -281,9 +288,7 @@ class Int8Group(nn.Module): def __init__(self, q: torch.Tensor, scales: torch.Tensor, bias: torch.Tensor | None): super().__init__() self.out_features, self.in_features = (int(v) for v in q.shape) - self.register_buffer("q", q) - self.register_buffer("scales", scales) - self.register_buffer("bias", bias) + _register_weights(self, q=q, scales=scales, bias=bias) @classmethod def empty(cls, rows: int, in_features: int, *, bias: bool = False, device="meta") -> Int8Group: @@ -326,8 +331,7 @@ class Bf16Group(nn.Module): def __init__(self, weight: torch.Tensor, bias: torch.Tensor | None): super().__init__() self.out_features, self.in_features = (int(v) for v in weight.shape) - self.register_buffer("weight", weight) - self.register_buffer("bias", bias) + _register_weights(self, weight=weight, bias=bias) @classmethod def empty(cls, rows: int, in_features: int, *, bias: bool = False, device="meta") -> Bf16Group: diff --git a/src/orbitquant/fused/install.py b/src/orbitquant/fused/install.py index 8050471..05a34b8 100644 --- a/src/orbitquant/fused/install.py +++ b/src/orbitquant/fused/install.py @@ -128,9 +128,16 @@ def fuse( layout["activation_bits"] = widths model.oq_fused_layout = layout model.__dict__["oq_runtime"] = runtime + _install_model(model, adapter) return layout +def _install_model(model: nn.Module, adapter) -> None: + install = getattr(adapter, "install", None) + if install is not None: + install(model) + + def _empty_group(block, spec, group_kind: str, config, device, activation_bits=None): layers = [_resolve(block, source.path) for source in spec.sources] first = layers[0] @@ -181,6 +188,7 @@ def prepare_skeleton( _attach(model, block, adapter, kind, modules, runtime) model.oq_fused_layout = layout model.__dict__["oq_runtime"] = runtime + _install_model(model, adapter) def finalize(model: nn.Module) -> None: diff --git a/src/orbitquant/quantizer.py b/src/orbitquant/quantizer.py index c6d7ac5..605a9e8 100644 --- a/src/orbitquant/quantizer.py +++ b/src/orbitquant/quantizer.py @@ -94,7 +94,7 @@ def _is_prequantized_state_tensor(module: Any, tensor_name: str) -> bool: from orbitquant.fused.groups import GROUP_TYPES if isinstance(module, tuple(GROUP_TYPES.values())): - return tensor_name in module._buffers + return tensor_name in module._parameters or tensor_name in module._buffers if isinstance(module, OrbitQuantLinear): return tensor_name in _ORBITQUANT_STATE_TENSORS and ( tensor_name in module._parameters or tensor_name in module._buffers From b98a589f69540861b6981cf7aa4fde6723a70a92 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Fri, 2 Oct 2026 03:08:11 +0200 Subject: [PATCH 3/4] Load fused checkpoints with their stored tensor dtypes Diffusers casts every floating tensor of a checkpoint to the model dtype before the quantizer places it, which rounded the FP32 per-row scales of INT8 groups (and would round BF16 dense rows of an FP16 model). The quantizer now places fused group tensors from the stored values, so a fused checkpoint loads bit-identical to the model it was saved from. --- src/orbitquant/quantizer.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/orbitquant/quantizer.py b/src/orbitquant/quantizer.py index 605a9e8..c482a94 100644 --- a/src/orbitquant/quantizer.py +++ b/src/orbitquant/quantizer.py @@ -312,6 +312,13 @@ def create_quantized_param( module, tensor_name = _module_and_tensor_name(model, param_name) if not _is_prequantized_state_tensor(module, tensor_name): raise ValueError(f"{param_name} is not an OrbitQuant pre-quantized tensor") + from orbitquant.fused.groups import GROUP_TYPES + + stored = state_dict.get(param_name) if isinstance(state_dict, dict) else None + if isinstance(module, tuple(GROUP_TYPES.values())) and isinstance(stored, torch.Tensor): + # Diffusers casts floating tensors to the model dtype before this call; fused + # groups keep their stored dtypes (FP32 row scales, BF16 dense rows). + param_value = stored if tensor_name in module._parameters: old_value = module._parameters[tensor_name] From ac83ece6dc724139839faa39392cd79978ea19d9 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Fri, 2 Oct 2026 03:33:04 +0200 Subject: [PATCH 4/4] Fuse artifacts whose manifest lists removed files; document streamed offloading of fused weights --- README.md | 16 +++++++++++----- src/orbitquant/artifacts/checksums.py | 11 +++++++++++ src/orbitquant/fused/convert.py | 6 ++++-- tests/test_artifact_writer.py | 14 ++++++++++++++ 4 files changed, 40 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 03f024f..1ae38f7 100644 --- a/README.md +++ b/README.md @@ -503,13 +503,19 @@ 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 no speed. Low-bit artifacts -with mixed weight widths (boundary and interior protection) fuse into row -segments of one GEMM. +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()`. `orbitquant.runtime.krea2` from 0.10 remains for -existing deployments. +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 diff --git a/src/orbitquant/artifacts/checksums.py b/src/orbitquant/artifacts/checksums.py index 0135928..be0b190 100644 --- a/src/orbitquant/artifacts/checksums.py +++ b/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) diff --git a/src/orbitquant/fused/convert.py b/src/orbitquant/fused/convert.py index d2ea130..d4051a4 100644 --- a/src/orbitquant/fused/convert.py +++ b/src/orbitquant/fused/convert.py @@ -19,8 +19,8 @@ from torch import nn from orbitquant.artifacts.checksums import ( + existing_file_checksums, is_ignored_artifact_relative_path, - sha256_file, write_sha256sums, ) from orbitquant.artifacts.loader import load_orbitquant_artifact @@ -75,7 +75,9 @@ def fuse_component_artifact( fused_manifest = dataclasses.replace( manifest, module_shapes={name: list(tensor.shape) for name, tensor in state.items()}, - checksums={name: sha256_file(output / name) for name in manifest.checksums}, + # 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" diff --git a/tests/test_artifact_writer.py b/tests/test_artifact_writer.py index 86e300d..afb12da 100644 --- a/tests/test_artifact_writer.py +++ b/tests/test_artifact_writer.py @@ -18,6 +18,7 @@ validate_orbitquant_artifact, ) from orbitquant.artifacts.checksums import ( + existing_file_checksums, read_sha256sums, validate_sha256sums, write_sha256sums, @@ -319,6 +320,19 @@ def test_sha256sums_ignore_huggingface_local_dir_cache_metadata(tmp_path): assert validated[".cache/huggingface/download/.gitattributes.metadata"] == "0" * 64 +def test_existing_file_checksums_leave_out_missing_paths(tmp_path): + (tmp_path / "benchmark").mkdir() + (tmp_path / "benchmark" / "summary.json").write_text("{}", encoding="utf-8") + + checksums = existing_file_checksums( + tmp_path, ["benchmark/summary.json", "benchmark/comparison.metrics.csv"] + ) + + assert checksums == { + "benchmark/summary.json": sha256_file(tmp_path / "benchmark" / "summary.json") + } + + def test_load_orbitquant_artifact_restores_quantized_modules_into_matching_model(tmp_path): torch.manual_seed(0) source = TinyArtifactModel()