From 1a96a059157c2e6fcf10e953b33b7d054958de43 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:55:03 +0200 Subject: [PATCH] Add the Fibo 2.0 fused family; pad SwiGLU halves to whole interleave chunks (0.15.0) - orbitquant.fused.families.fibo2 fuses BriaFibo2Transformer2DModel: the Z-Image blocks with per-token modulation tables (noisy image tokens vs clean context tokens), sandwich post-norms and tanh gates as row kernels, the gated cross-attention of the injection blocks, and the Perceiver layers; 8-bit activations; one sample at a time. - SwiGLU groups whose halves are not a whole number of 128-row interleave chunks (Fibo 2.0: 9536) get zero rows up to the next chunk; the down projection reads only the real width. - Tests: a tiny Fibo 2.0 transformer, fused BF16 groups against the eager model (t2i, edit) and a quantized W4A8 model. - README: Fibo 2.0 in the fused-family table. Version 0.15.0. --- README.md | 3 +- pyproject.toml | 2 +- src/orbitquant/__init__.py | 2 +- src/orbitquant/fused/families/__init__.py | 1 + src/orbitquant/fused/families/fibo2.py | 283 ++++++++++++++++++++++ src/orbitquant/fused/groups.py | 16 +- src/orbitquant/fused/install.py | 11 +- tests/test_fused_fibo2.py | 150 ++++++++++++ uv.lock | 2 +- 9 files changed, 462 insertions(+), 8 deletions(-) create mode 100644 src/orbitquant/fused/families/fibo2.py create mode 100644 tests/test_fused_fibo2.py diff --git a/README.md b/README.md index e556570..672f2d3 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ The implementation is clean-room and Apache-2.0 licensed. - Packed-weight CUDA, Triton, and Metal inference paths that avoid a full dequantized weight matrix. - Fused transformer checkpoints for Krea 2, FLUX.2, Ideogram 4, Qwen-Image 2.1, - MiniMax-H3, Boogu-Image, FLUX.3 Klein and MOSS-Image (`orbitquant.fused`): each block runs + MiniMax-H3, Boogu-Image, FLUX.3 Klein, MOSS-Image and Fibo 2.0 (`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. @@ -554,6 +554,7 @@ the same for an `orbitquant-v1` component artifact, and | Boogu-Image | `BooguImageTransformer2DModel` | refiner, single- and double-stream blocks; 8-bit activations | | FLUX.3 Klein | `flux3.dit.DiT` (Black Forest Labs' `flux3` package) | per-mode and single-stream blocks: Q\|K\|V, SwiGLU, gated residuals; one sample at a time | | MOSS-Image | `MOSSImageTransformer` (the model's `pipeline_moss_image.py`) | joint image/text blocks: RMSNorm and shared AdaLN prologues, Q\|K\|V, SiTU gate\|up, gated residuals; one sample at a time | +| Fibo 2.0 | `BriaFibo2Transformer2DModel` | Z-Image blocks with per-token (noisy/clean) modulation tables, sandwich post-norms and tanh gates as row kernels; gated cross-attention of the five injection blocks; Perceiver layers; SwiGLU halves padded to whole interleave chunks; 8-bit activations; one sample at a time | Groups take the activation width of their projections unless the layout overrides it: `fuse(..., activation_bits={"qkv": 8})` gives a group 8-bit diff --git a/pyproject.toml b/pyproject.toml index d2c384f..8307f51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.14.0" +version = "0.15.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 b2d2d44..5259c75 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.14.0" +__version__ = "0.15.0" register_hf_quantizers() diff --git a/src/orbitquant/fused/families/__init__.py b/src/orbitquant/fused/families/__init__.py index 758b1ca..6b877c4 100644 --- a/src/orbitquant/fused/families/__init__.py +++ b/src/orbitquant/fused/families/__init__.py @@ -27,6 +27,7 @@ "boogu", "flux3", "moss_image", + "fibo2", ) diff --git a/src/orbitquant/fused/families/fibo2.py b/src/orbitquant/fused/families/fibo2.py new file mode 100644 index 0000000..e658f08 --- /dev/null +++ b/src/orbitquant/fused/families/fibo2.py @@ -0,0 +1,283 @@ +"""Fibo 2.0 (Bria ``BriaFibo2Transformer2DModel``) blocks. + +The noise refiner and the 44 main blocks are Z-Image blocks: RMSNorm x (1 + scale) prologue, one +Q|K|V GEMM, Q/K RMSNorm with interleaved complex RoPE in one kernel, sandwich post-norms with +tanh gates as row kernels and a SwiGLU feed-forward. The main blocks modulate per token: gist +and context tokens take the clean (t = 0) rows, image tokens the noisy rows, so the prologue and +the post-norm kernels read their modulation from a two-row table by token index. Five injection +blocks add a gated cross-attention to their own Qwen3-VL bundle (7680 channels) between the +attention and the feed-forward. The two Perceiver layers turn the first bundle into the gist +tokens with plain residuals. + +Every group takes 8-bit activations (per-token absmax INT8 of the rotated input): the +feed-forward output projection reads 9536 channels, whose RPBH block is 64 wide, too narrow to +spread outliers for 4-bit codes; the GEMMs are INT8 either way. +""" + +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 = "fibo2" +MODEL_CLASSES = ("BriaFibo2Transformer2DModel",) + +_BLOCK_GROUPS = ("qkv", "attn_out", "ff_in", "ff_out") +_CROSS_GROUPS = ("xq", "xkv", "xout") +ACTIVATION_BITS = dict.fromkeys(_BLOCK_GROUPS + _CROSS_GROUPS, 8) + + +def blocks(model): + paths = [(f"noise_refiner.{i}", "block") for i in range(len(model.noise_refiner))] + paths += [(f"perceiver.layers.{i}", "perceiver") for i in range(len(model.perceiver.layers))] + for i, block in enumerate(model.layers): + paths.append((f"layers.{i}", "inject" if block.cross_attention is not None else "block")) + return paths + + +def _cross_groups(prefix="cross_attention"): + return dict( + xq=[f"{prefix}.to_q"], + xkv=[f"{prefix}.to_k", f"{prefix}.to_v"], + xout=[f"{prefix}.to_out.0"], + ) + + +def groups(kind, block): + if kind == "perceiver": + return specs( + **_cross_groups(), + qkv=["self_attention.to_q", "self_attention.to_k", "self_attention.to_v"], + attn_out=["self_attention.to_out.0"], + ff_in=(["ffn.w1", "ffn.w3"], {"swiglu": True}), + ff_out=["ffn.w2"], + ) + out = dict( + qkv=["attention.to_q", "attention.to_k", "attention.to_v"], + attn_out=["attention.to_out.0"], + ff_in=(["feed_forward.w1", "feed_forward.w3"], {"swiglu": True}), + ff_out=["feed_forward.w2"], + ) + if kind == "inject": + out.update(_cross_groups()) + return specs(**out) + + +def prepare(block, kind): + attn = block.self_attention if kind == "perceiver" else block.attention + block.oq_int8_attention = int8_attention_allowed(attn.norm_q.weight, attn.norm_k.weight) + + +def forward(kind): + return _perceiver if kind == "perceiver" else _block + + +def _cached(runtime, name, key, build, slots=2): + """A small per-forward cache keyed by tensor identity (rope tables, token indices): the + tensors are shared by every block of one transformer call.""" + cache = runtime.__dict__.setdefault(name, []) + for entry in cache: + if entry[0] is key: + return entry[1] + value = build() + cache.insert(0, (key, value)) + del cache[slots:] + return value + + +def _rope(runtime, freqs_cis, rows): + """Complex ``[1, seq, head_dim / 2]`` frequencies -> interleaved fp32 cos/sin rows.""" + + def build(): + freqs = freqs_cis.reshape(-1, freqs_cis.shape[-1]) + return ( + freqs.real.float().repeat_interleave(2, -1).contiguous(), + freqs.imag.float().repeat_interleave(2, -1).contiguous(), + ) + + cos, sin = _cached(runtime, "fibo2_rope", freqs_cis, build) + return cos[:rows], sin[:rows] + + +def _token_index(runtime, noise_mask, sample): + """Row ``i`` of the modulation tables for every token: 0 noisy, 1 clean.""" + if noise_mask is None: + return None + index = _cached(runtime, "fibo2_index", noise_mask, lambda: (noise_mask == 0).to(torch.int32)) + return index[sample].contiguous() + + +def _modulation(block, adaln): + """``(scale_msa, gate_msa, scale_mlp, gate_mlp)`` rows of one timestep embedding, with the + eager ``1 +`` and ``tanh`` applied.""" + scale_msa, gate_msa, scale_mlp, gate_mlp = block.adaLN_modulation(adaln).chunk(4, dim=1) + return 1.0 + scale_msa, gate_msa.tanh(), 1.0 + scale_mlp, gate_mlp.tanh() + + +def _tables(block, adaln_input, noise_mask, adaln_noisy, adaln_clean, sample): + """Modulation tables ``[modes, dim]`` of one sample: one row (global) or noisy | clean.""" + if noise_mask is None: + rows = _modulation(block, adaln_input[sample : sample + 1]) + else: + noisy = _modulation(block, adaln_noisy[sample : sample + 1]) + clean = _modulation(block, adaln_clean[sample : sample + 1]) + rows = [torch.cat(pair) for pair in zip(noisy, clean, strict=True)] + return [r.to(torch.bfloat16).contiguous() for r in rows] + + +def _valid(mask, rows): + """Indices of the real keys of one sample's ``[seq]`` padding mask (``None``: all real).""" + if mask is None: + return None + valid = mask.reshape(-1)[:rows].bool() + return None if bool(valid.all()) else valid + + +def _heads(src, norm, heads, cos=None, sin=None): + return act.qk_norm_rope_with_triton( + src, + heads, + src.shape[1] // heads, + norm.weight, + cos, + sin, + norm.eps, + weight_mode=act.QK_WEIGHT_DIFFUSERS, + style=act.ROPE_INTERLEAVED, + ) + + +def _attend(block, attn, query_src, kv_src, inner, *, cos=None, sin=None, valid=None, int8=True): + """Q/K RMSNorm (+ RoPE), attention; ``query_src`` holds Q, ``kv_src`` K | V.""" + runtime = block.oq_runtime + heads = attn.heads + head_dim = inner // heads + query = _heads(query_src, attn.norm_q, heads, cos, sin) + key = _heads(kv_src[:, :inner], attn.norm_k, heads, cos, sin) + value = kv_src[:, inner : 2 * inner].unflatten(-1, (heads, head_dim))[None] + if valid is not None: + key, value = key[:, valid], value[:, valid] + out = runtime.attention(query, key, value, int8_allowed=int8 and block.oq_int8_attention) + return out.reshape(query.shape[1], inner) + + +def _self_attention(block, attn, fused, source, cos, sin, valid): + runtime = block.oq_runtime + qkv = runtime.linear(fused.qkv, source) + inner = qkv.shape[1] // 3 + return _attend( + block, attn, qkv[:, :inner], qkv[:, inner:], inner, cos=cos, sin=sin, valid=valid + ) + + +def _cross_attention(block, attn, fused, x, norm, text, text_valid): + """``attn(norm(x), text)`` of one sample, before its output projection.""" + runtime = block.oq_runtime + dim = x.shape[1] + query = runtime.linear( + fused.xq, + runtime.input(x, dim, norm=act.NORM_RMS, norm_weight=norm.weight, rms_eps=norm.eps), + ) + kv = runtime.linear(fused.xkv, runtime.input(text, text.shape[1])) + return _attend(block, attn, query, kv, query.shape[1], valid=text_valid, int8=False) + + +def _swiglu(runtime, fused, source): + """SwiGLU hidden state as the input of ``ff_out``. The hidden width (9536) is not a whole + number of interleave chunks: the gated GEMM pads it with zero columns, which the strided + view leaves out.""" + mid = runtime.linear(fused.ff_in, source, epilogue=gemm.EPILOGUE_SWIGLU) + width = fused.ff_out.in_features + return runtime.input(mid[:, :width], width) + + +def _block( + self, + x, + attn_mask, + freqs_cis, + adaln_input=None, + noise_mask=None, + adaln_noisy=None, + adaln_clean=None, + encoder_hidden_states=None, + encoder_attention_mask=None, +): + runtime, fused = self.oq_runtime, self.oq_fused + batch, seq, dim = x.shape + cos, sin = _rope(runtime, freqs_cis, seq) + out = torch.empty_like(x, dtype=torch.bfloat16) + for sample in range(batch): + h = out[sample] + h.copy_(x[sample]) + index = _token_index(runtime, noise_mask, sample) + scale_msa, gate_msa, scale_mlp, gate_mlp = _tables( + self, adaln_input, noise_mask, adaln_noisy, adaln_clean, sample + ) + prologue = dict(norm=act.NORM_RMS, mod=act.MOD_SCALE, mod_index=index) + norm = self.attention_norm1 + source = runtime.input( + h, dim, norm_weight=norm.weight, rms_eps=norm.eps, mod_scale=scale_msa, **prologue + ) + valid = _valid(None if attn_mask is None else attn_mask[sample], seq) + att = _self_attention(self, self.attention, fused, source, cos, sin, valid) + att_out = runtime.linear(fused.attn_out, runtime.input(att, att.shape[1])) + norm = self.attention_norm2 + act.postnorm_residual_with_triton( + h, att_out, norm.weight, gate_msa, norm.eps, gate_index=index + ) + if self.cross_attention is not None: + text = encoder_hidden_states[sample] + text_valid = _valid( + None if encoder_attention_mask is None else encoder_attention_mask[sample], + text.shape[0], + ) + cross = _cross_attention( + self, self.cross_attention, fused, h, self.cross_attn_norm1, text, text_valid + ) + cross_out = runtime.linear(fused.xout, runtime.input(cross, cross.shape[1])) + gate = self.cross_attn_gate.tanh().to(torch.bfloat16).expand(dim) + norm = self.cross_attn_norm2 + act.postnorm_residual_with_triton(h, cross_out, norm.weight, gate, norm.eps) + norm = self.ffn_norm1 + source = runtime.input( + h, dim, norm_weight=norm.weight, rms_eps=norm.eps, mod_scale=scale_mlp, **prologue + ) + ff_out = runtime.linear(fused.ff_out, _swiglu(runtime, fused, source)) + norm = self.ffn_norm2 + act.postnorm_residual_with_triton( + h, ff_out, norm.weight, gate_mlp, norm.eps, gate_index=index + ) + return out + + +def _perceiver( + self, latents, encoder_hidden_states, encoder_attention_mask=None, attention_mask=None +): + runtime, fused = self.oq_runtime, self.oq_fused + batch, rows, dim = latents.shape + out = torch.empty_like(latents, dtype=torch.bfloat16) + for sample in range(batch): + h = out[sample] + h.copy_(latents[sample]) + text = encoder_hidden_states[sample] + text_valid = _valid( + None if encoder_attention_mask is None else encoder_attention_mask[sample], + text.shape[0], + ) + cross = _cross_attention( + self, self.cross_attention, fused, h, self.cross_attn_norm, text, text_valid + ) + h.add_(runtime.linear(fused.xout, runtime.input(cross, cross.shape[1]))) + norm = self.self_attn_norm + source = runtime.input(h, dim, norm=act.NORM_RMS, norm_weight=norm.weight, rms_eps=norm.eps) + valid = _valid(None if attention_mask is None else attention_mask[sample], rows) + att = _self_attention(self, self.self_attention, fused, source, None, None, valid) + h.add_(runtime.linear(fused.attn_out, runtime.input(att, att.shape[1]))) + norm = self.ffn_norm + source = runtime.input(h, dim, norm=act.NORM_RMS, norm_weight=norm.weight, rms_eps=norm.eps) + h.add_(runtime.linear(fused.ff_out, _swiglu(runtime, fused, source))) + return out diff --git a/src/orbitquant/fused/groups.py b/src/orbitquant/fused/groups.py index b91e1d2..d0e7001 100644 --- a/src/orbitquant/fused/groups.py +++ b/src/orbitquant/fused/groups.py @@ -50,9 +50,20 @@ def _rows(source: RowSource, layer: nn.Module) -> tuple[int, int]: return source.start, stop +def swiglu_half_rows(rows: int, chunk: int = SWIGLU_INTERLEAVE) -> int: + """Rows of one SwiGLU half rounded up to whole interleave chunks.""" + return -(-rows // chunk) * chunk + + 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}") + if gate.shape != up.shape: + raise ValueError("SwiGLU halves need equal row counts") + pad = swiglu_half_rows(gate.shape[0], chunk) - gate.shape[0] + if pad: + # Zero rows (zero scale, weight and bias) complete the last chunk: their gated outputs + # are 0, past the width the next projection reads. + zeros = gate.new_zeros((pad, *gate.shape[1:])) + gate, up = torch.cat([gate, zeros]), torch.cat([up, zeros]) pieces = [] for start in range(0, gate.shape[0], chunk): pieces += [gate[start : start + chunk], up[start : start + chunk]] @@ -188,6 +199,7 @@ def source_segments(layers, sources, *, swiglu: bool = False) -> tuple[tuple[int 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") + sizes = [(swiglu_half_rows(rows), bits) for rows, bits in sizes] return _segments(sizes) @classmethod diff --git a/src/orbitquant/fused/install.py b/src/orbitquant/fused/install.py index 05a34b8..9296f18 100644 --- a/src/orbitquant/fused/install.py +++ b/src/orbitquant/fused/install.py @@ -24,7 +24,14 @@ 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.groups import ( + GROUP_TYPES, + Bf16Group, + Int8Group, + PackedGroup, + _resolve, + swiglu_half_rows, +) from orbitquant.fused.runtime import FusedRuntime from orbitquant.layers import OrbitQuantLinear @@ -154,7 +161,7 @@ def _empty_group(block, spec, group_kind: str, config, device, activation_bits=N 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 + rows += swiglu_half_rows(stop - source.start) if spec.swiglu else stop - source.start return GROUP_TYPES[group_kind].empty(rows, first.in_features, bias=bias, device=device) diff --git a/tests/test_fused_fibo2.py b/tests/test_fused_fibo2.py new file mode 100644 index 0000000..11358f3 --- /dev/null +++ b/tests/test_fused_fibo2.py @@ -0,0 +1,150 @@ +import copy + +import pytest +import torch +import torch.nn.functional as F + +pytest.importorskip("triton") +diffusers = pytest.importorskip("diffusers") +if not hasattr(diffusers, "BriaFibo2Transformer2DModel"): + pytest.skip("diffusers without Fibo 2.0", allow_module_level=True) + +from orbitquant.fused.families import adapter_for_class, adapter_for_family, fibo2 # noqa: E402 +from orbitquant.fused.groups import swiglu_half_rows # noqa: E402 +from orbitquant.kernels import available_backends # noqa: E402 + +# Head width 128 matches the RoPE axes (32 + 48 + 48); 7 heads make the SwiGLU width 2368, which is +# not a whole number of 128-row interleave chunks, as in the real model (9536). +CONFIG = dict( + dim=896, + n_heads=7, + n_layers=3, + n_refiner_layers=1, + cap_feat_dim=768, + injection_layer_ids=(1,), + perceiver_num_layers=1, + min_num_gist_tokens=64, + max_num_gist_tokens=128, + gist_step=64, + gist_min_text_len=16, + gist_max_text_len=64, +) +TARGETS = [ + "noise_refiner.*.attention.to_*", + "noise_refiner.*.feed_forward.w*", + "layers.*.attention.to_*", + "layers.*.feed_forward.w*", + "layers.*.cross_attention.to_*", + "perceiver.layers.*.cross_attention.to_*", + "perceiver.layers.*.self_attention.to_*", + "perceiver.layers.*.ffn.w*", +] + + +def _require_cuda(): + if not torch.cuda.is_available() or not available_backends()["triton_cuda"]: + pytest.skip("CUDA/Triton backend is not available") + + +def _model(): + torch.manual_seed(0) + model = diffusers.BriaFibo2Transformer2DModel(**CONFIG).to(torch.bfloat16) + with torch.no_grad(): + for name, p in model.named_parameters(): + if name.endswith("cross_attn_gate"): + p.fill_(0.5) + elif "norm" in name: + p.copy_(torch.rand_like(p) * 0.5 + 0.75) + elif "adaLN_modulation" in name: + p.mul_(0.05) + return model.eval() + + +def _inputs(edit): + g = torch.Generator("cuda").manual_seed(1) + latents = torch.randn(1, 32, 16, 16, device="cuda", generator=g).to(torch.bfloat16) + text = torch.randn(1, 6, 40, 768, device="cuda", generator=g).to(torch.bfloat16) + context = ( + [torch.randn(1, 32, 16, 12, device="cuda", generator=g).to(torch.bfloat16)] + if edit + else None + ) + return dict( + hidden_states=latents, + timestep=torch.tensor([0.4], device="cuda"), + encoder_hidden_states=text, + context_latents=context, + return_dict=False, + ) + + +def _cos(a, b): + return F.cosine_similarity(a.float().flatten(), b.float().flatten(), dim=0).item() + + +def test_family_is_registered_for_the_transformer_class(): + assert adapter_for_family("fibo2") is fibo2 + assert adapter_for_class("BriaFibo2Transformer2DModel") is fibo2 + + +def test_blocks_and_groups(): + model = _model() + blocks = dict(fibo2.blocks(model)) + assert blocks == { + "noise_refiner.0": "block", + "perceiver.layers.0": "perceiver", + "layers.0": "block", + "layers.1": "inject", + "layers.2": "block", + } + assert list(fibo2.groups("block", model.layers[0])) == ["qkv", "attn_out", "ff_in", "ff_out"] + inject = fibo2.groups("inject", model.layers[1]) + assert [s.path for s in inject["xkv"].sources] == [ + "cross_attention.to_k", + "cross_attention.to_v", + ] + perceiver = fibo2.groups("perceiver", model.perceiver.layers[0]) + assert perceiver["ff_in"].swiglu and [s.path for s in perceiver["ff_in"].sources] == [ + "ffn.w1", + "ffn.w3", + ] + + +def test_swiglu_halves_pad_to_whole_chunks(): + assert swiglu_half_rows(2368) == 2432 + assert swiglu_half_rows(9536) == 9600 + assert swiglu_half_rows(1024) == 1024 + + +@pytest.mark.parametrize("edit", [False, True], ids=["t2i", "edit"]) +def test_fused_bf16_groups_match_the_eager_model(edit): + _require_cuda() + from orbitquant import OrbitQuantConfig + from orbitquant.fused import fuse + + model = _model().cuda() + reference = copy.deepcopy(model) + fuse(model, dense="bf16", config=OrbitQuantConfig(), attention="sdpa") + with torch.inference_mode(): + expected = reference(**_inputs(edit))[0] + actual = model(**_inputs(edit))[0] + assert actual.shape == expected.shape + assert _cos(actual, expected) > 0.999 + + +def test_fused_w4a8_model_tracks_the_eager_model(): + _require_cuda() + from orbitquant import OrbitQuantConfig, quantize_model + from orbitquant.fused import fuse + + model = _model().cuda() + reference = copy.deepcopy(model) + config = OrbitQuantConfig(target_policy="universal", modules_to_convert=TARGETS) + summary = quantize_model(model, config, quantization_device="cuda") + assert len(summary.quantized_modules) == 7 * 4 + 4 + 11 + fuse(model, config=config, attention="sdpa") + assert model.layers[0].oq_fused.ff_in.out_features == 2 * 2432 + with torch.inference_mode(): + expected = reference(**_inputs(True))[0] + actual = model(**_inputs(True))[0] + assert _cos(actual, expected) > 0.95 diff --git a/uv.lock b/uv.lock index be721b1..57eb81b 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.14.0" +version = "0.15.0" source = { editable = "." } dependencies = [ { name = "huggingface-hub" },