diff --git a/README.md b/README.md index bbbd8fb..1ae38f7 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,75 @@ system, C++ runtime, and architecture. See [`docs/kernel-audit.md`](docs/kernel-audit.md) for tested shapes, benchmark methodology, and local package verification. -## Fused DiT Runtime (Krea 2) +## Fused Transformers -Version 0.10 speeds up every W4A4 model on CUDA without API changes: +Version 0.11 runs the transformer blocks of supported diffusion models fused: +the projections that read one input (attention Q|K|V, SwiGLU gate|up, the two +streams of a joint attention) form one weight matrix and one INT8-surrogate +GEMM, the RMSNorm/LayerNorm and AdaLN modulation in front of them run in the +activation-quantization prologue, and SwiGLU, sigmoid gates and gated residual +updates run in the GEMM epilogue. Q/K normalization and RoPE are one kernel, +and attention uses an INT8 Q.K^T / FP16-accumulated P.V kernel (the +SageAttention v1 scheme); blocks whose Q/K RMSNorm scales one channel far above +the rest keep BF16 attention. -- Triton kernels take row counts as runtime arguments, so a new prompt length - or image size no longer recompiles them. -- Large-row W4A4 layers use an INT8 GEMM that applies the scale epilogue in - registers; the output is bit-identical to the previous `torch._int_mm` path. -- Projections that read the same tensor (attention Q/K/V, SwiGLU gate/up) - share one activation quantization. The reuse needs version counters, so run - inference under `torch.no_grad()`; `torch.inference_mode()` disables it. +A fused checkpoint stores only the fused groups and records the layout in its +`quantization_config`, so the usual loader returns the fused transformer: -On an RTX 4060 Ti the Qwen-Image 2.1 W4A4 pipeline went from 15.0 to 10.3 s -per 1024x1024 image with pixel-identical output. +```python +import torch +import orbitquant # registers the loader +from diffusers import Krea2Pipeline + +pipe = Krea2Pipeline.from_pretrained( + "WaveCut/Krea-2-Turbo-OrbitQuant-W4A4", torch_dtype=torch.bfloat16 +).to("cuda") +``` -`orbitquant.runtime.krea2` goes further for Krea 2 Turbo: each block runs one -grouped Q|K|V|gate GEMM with the sigmoid gate in its epilogue, a fused Q/K -RMSNorm + RoPE kernel, SwiGLU and gated residual updates in GEMM epilogues, -W8A8 down projections, and optionally an INT8 Q.K^T / FP16 P.V attention -kernel (the SageAttention v1 scheme). Blocks whose Q/K RMSNorm scales one -channel far above the rest keep BF16 attention. +Fuse a per-projection checkpoint on the host and save it: ```python import torch -from diffusers import Krea2Pipeline -from orbitquant.runtime.krea2 import Krea2FastRunner, install, save_fused - -pipe = Krea2Pipeline.from_pretrained(model_dir, dtype=torch.bfloat16) -fused = install(pipe.transformer, attention="int8") # on the host, before .to("cuda") -save_fused(fused, "fused_blocks.safetensors") # optional: map it next time -pipe.to("cuda") -runner = Krea2FastRunner(pipe, fused) -with torch.no_grad(): - states = runner.encode(prompt, torch.device("cuda")) - latents = runner.denoise(states, width=1024, height=1024) - image = runner.decode(latents, width=1024, height=1024) +from diffusers import Flux2Transformer2DModel +from orbitquant import fused + +transformer = Flux2Transformer2DModel.from_pretrained( + model_dir, subfolder="transformer", torch_dtype=torch.bfloat16 +) +fused.fuse(transformer) # attention="flash" keeps BF16 attention everywhere +fused.save_pretrained(transformer, "fused/transformer") ``` -Building the fused buffers copies the block weights into new host memory; -`install(..., fused_path="fused_blocks.safetensors")` maps a saved file instead. -On an RTX 4060 Ti a 1024x1024 eight-step image takes 10.8 s (25.9 s with the -stock pipeline), and the DiT needs 7.3 GB instead of 10 GB. +`orbitquant.fused.fuse_component_artifact(model, source_dir, output_dir)` does +the same for an `orbitquant-v1` component artifact, and +`load_orbitquant_artifact` loads the result fused. + +| Family | Transformer class | Fused blocks | +| --- | --- | --- | +| Krea 2 | `Krea2Transformer2DModel` | Q\|K\|V\|gate, gated output, SwiGLU; W8A8 down projections | +| FLUX.2 | `Flux2Transformer2DModel` | double-stream (per-stream Q\|K\|V, gated residuals) and single-stream blocks; reference-image KV cache | +| Ideogram 4 | `Ideogram4Transformer2DModel` | sandwich-norm blocks with row-indexed modulation | +| Qwen-Image 2.1 | `QwenImage21Transformer2DModel` | block-causal prefill and KV-cached decode | +| MiniMax-H3 | `MiniMaxH3Transformer3DModel` | per-row (timestep, modality) AdaLN tables, partial rotate-half RoPE | +| Boogu-Image | `BooguImageTransformer2DModel` | refiner, single- and double-stream blocks; 8-bit activations | + +Groups take the activation width of their projections unless the layout +overrides it: `fuse(..., activation_bits={"qkv": 8})` gives a group 8-bit +activations (per-token absmax INT8 of the rotated input). Boogu-Image runs all +its groups this way because its 3360 channels only allow a 32-wide rotation +block; the GEMMs stay INT8, so the wider codes cost almost no speed. Low-bit +artifacts with mixed weight widths (boundary and interior protection) fuse into +row segments of one GEMM. + +The fused runtime needs CUDA and Triton. Run inference under `torch.no_grad()` +or `torch.inference_mode()`. Fused weights are frozen parameters, so model CPU +offload and Diffusers block-level group offloading (including streamed +`use_stream=True` prefetching) move them with their blocks; leaf-level group +offloading does not, because a fused block reads its groups directly instead of +calling them. Stream without `record_stream`: recorded streams keep every +streamed block reserved until the compute stream has used it, which grew +MiniMax-H3's GPU memory to its 12 GiB cap (5.1 GiB without). +`orbitquant.runtime.krea2` from 0.10 remains for existing deployments. ## Validated Architecture Coverage 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/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/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..d4051a4 --- /dev/null +++ b/src/orbitquant/fused/convert.py @@ -0,0 +1,86 @@ +"""Fused copies of published ``orbitquant-v1`` component artifacts. + +``fuse_component_artifact`` rewrites such an artifact so that ``model.safetensors`` holds the +fused groups instead of the per-projection tensors and ``quantization_config.json`` carries the +``fused_layout``; every other file is copied. The manifest keeps the projection list and the +per-module widths the loader rebuilds the groups from. Diffusers-format components need no +helper: load the transformer with ``from_pretrained``, ``fuse`` it and save it with +``orbitquant.fused.save_pretrained``. +""" + +from __future__ import annotations + +import dataclasses +import json +import shutil +from pathlib import Path + +from safetensors.torch import save_file +from torch import nn + +from orbitquant.artifacts.checksums import ( + existing_file_checksums, + is_ignored_artifact_relative_path, + write_sha256sums, +) +from orbitquant.artifacts.loader import load_orbitquant_artifact +from orbitquant.artifacts.manifest import OrbitQuantManifest +from orbitquant.config import OrbitQuantConfig +from orbitquant.fused.install import fuse + +_REWRITTEN = ("model.safetensors", "quantization_config.json", "orbitquant_manifest.json") + + +def fuse_component_artifact( + model: nn.Module, + source_dir: str | Path, + output_dir: str | Path, + *, + dense: str | None = None, + validate_checksums: bool = True, +) -> OrbitQuantManifest: + """Load the artifact at ``source_dir`` into ``model`` (the unquantized component skeleton), + fuse it and write the fused artifact to ``output_dir``. Rewrite the model card and refresh + ``SHA256SUMS`` (``write_sha256sums``) after editing it.""" + source, output = Path(source_dir), Path(output_dir) + if output.resolve() == source.resolve(): + raise ValueError("write the fused artifact to a new directory") + config = OrbitQuantConfig.from_dict( + json.loads((source / "quantization_config.json").read_text(encoding="utf-8")) + ) + if config.fused_layout: + raise ValueError(f"{source} is already fused") + manifest = load_orbitquant_artifact(model, source, validate_checksums=validate_checksums) + config.fused_layout = fuse(model, dense=dense, config=config) + + output.mkdir(parents=True, exist_ok=True) + for path in sorted(source.rglob("*")): + relative = path.relative_to(source).as_posix() + if ( + path.is_dir() + or relative in _REWRITTEN + or relative == "SHA256SUMS" + or is_ignored_artifact_relative_path(relative) + ): + continue + target = output / relative + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, target) + + state = {name: value.detach().cpu().contiguous() for name, value in model.state_dict().items()} + save_file(state, output / "model.safetensors") + (output / "quantization_config.json").write_text( + json.dumps(config.to_dict(), indent=2) + "\n", encoding="utf-8" + ) + fused_manifest = dataclasses.replace( + manifest, + module_shapes={name: list(tensor.shape) for name, tensor in state.items()}, + # A published manifest can list files that a later repository edit removed (accepted + # with validate_checksums=False); the fused manifest lists what the fused copy ships. + checksums=existing_file_checksums(output, manifest.checksums), + ) + (output / "orbitquant_manifest.json").write_text( + json.dumps(fused_manifest.to_dict(), indent=2) + "\n", encoding="utf-8" + ) + write_sha256sums(output) + return fused_manifest diff --git a/src/orbitquant/fused/families/__init__.py b/src/orbitquant/fused/families/__init__.py new file mode 100644 index 0000000..dfbefc9 --- /dev/null +++ b/src/orbitquant/fused/families/__init__.py @@ -0,0 +1,74 @@ +"""Block adapters: which projections each block fuses and the forward that runs them. + +An adapter module defines: + +* ``FAMILY``: the name recorded in fused checkpoints; +* ``MODEL_CLASSES``: transformer class names it serves; +* ``blocks(model)``: ``(path, kind)`` of every fused block; +* ``groups(kind, block)``: ``{name: GroupSpec}`` for a block of that kind; +* ``forward(kind)``: the replacement ``forward`` (same signature as the original block); +* ``prepare(block, kind)``: per-block settings derived from the loaded weights (optional); +* ``install(model)``: model-level changes outside the blocks (optional). +""" + +from __future__ import annotations + +import importlib +from dataclasses import dataclass, field + +from orbitquant.fused.groups import RowSource + +_MODULES = ("krea2", "flux2", "ideogram4", "minimax_h3", "qwenimage21", "boogu") + + +@dataclass(frozen=True) +class GroupSpec: + sources: tuple[RowSource, ...] + swiglu: bool = False + # Dense (unquantized) sources become per-row INT8 unless the checkpoint asks for BF16. + dense: str | None = None + extra: dict = field(default_factory=dict) + + +# INT8 Q/K quantization loses the small channels of heads whose Q/K RMSNorm scales one channel +# far above the rest (Krea 2 block 0: 36x the median); such blocks keep BF16 attention. +INT8_ATTENTION_MAX_NORM_SPREAD = 3.0 + + +def int8_attention_allowed(*norm_weights, offset: float = 0.0) -> bool: + """Whether a block's Q/K RMSNorm scales (``weight + offset``) are flat enough for INT8 Q/K.""" + import torch + + scales = torch.cat([w.reshape(-1) for w in norm_weights]).float().add(offset).abs() + return bool(scales.max() <= INT8_ATTENTION_MAX_NORM_SPREAD * scales.median()) + + +def specs(**groups) -> dict[str, GroupSpec]: + """``specs(qkv=["attn.to_q", ...], gateup=(["ff.gate", "ff.up"], {"swiglu": True}))``.""" + out = {} + for name, value in groups.items(): + options = {} + if isinstance(value, tuple): + value, options = value + sources = tuple(s if isinstance(s, RowSource) else RowSource(s) for s in value) + out[name] = GroupSpec(sources, **options) + return out + + +def adapter_for_class(class_name: str): + for name in _MODULES: + try: + module = importlib.import_module(f"orbitquant.fused.families.{name}") + except ModuleNotFoundError: + continue + if class_name in module.MODEL_CLASSES: + return module + raise ValueError(f"no fused block adapter for {class_name}") + + +def adapter_for_family(family: str): + try: + module = importlib.import_module(f"orbitquant.fused.families.{family}") + except ModuleNotFoundError as exc: + raise ValueError(f"unknown fused family {family!r}") from exc + return module 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..2766999 --- /dev/null +++ b/src/orbitquant/fused/families/krea2.py @@ -0,0 +1,186 @@ +"""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 weakref + +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 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 + + +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]) + ] + ) + _, 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) + ) + 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] + 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 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..d06ffbe --- /dev/null +++ b/src/orbitquant/fused/groups.py @@ -0,0 +1,381 @@ +"""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) + + +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" + + 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}") + _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 + ) + 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) + _register_weights(self, q=q, scales=scales, 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) + _register_weights(self, weight=weight, 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..05a34b8 --- /dev/null +++ b/src/orbitquant/fused/install.py @@ -0,0 +1,231 @@ +"""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 + _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] + 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 + _install_model(model, adapter) + + +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..c482a94 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._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 @@ -308,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] @@ -544,6 +555,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 +570,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_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() 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" },