From 517ea3e3758cc246afda168849b6f96a385a9b47 Mon Sep 17 00:00:00 2001 From: Valeriy Selitskiy <239034+iamwavecut@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:55:46 +0200 Subject: [PATCH] Add the FLUX.3 Klein fused family; stop per-GEMM host copies of group constants (0.13.0) - orbitquant.fused.families.flux3 fuses Black Forest Labs' native flux3.dit.DiT: per-mode and single-stream blocks with Q|K|V and SwiGLU GEMMs over one set of codes, Q/K RMSNorm + RoPE in one kernel and gated residuals in the GEMM epilogues. A mode whose tokens share one time embedding computes its AdaLN modulation on one row. - Fused groups built on the host kept their codebook surrogate codes on the CPU and copied them to the GPU in every decode: a blocking transfer per GEMM. The constants now move to the weight device once. On an RTX 5090 this cuts FLUX.3 Klein 1024x1024 generation from 2.44 s to 2.06 s with bit-identical images. - The activation quantizer no longer copies a plain bf16 input before rotating it; it writes only the token norms. --- README.md | 8 +- pyproject.toml | 2 +- src/orbitquant/__init__.py | 2 +- src/orbitquant/fused/families/__init__.py | 2 +- src/orbitquant/fused/families/flux3.py | 170 ++++++++++++++++++++ src/orbitquant/fused/groups.py | 15 +- src/orbitquant/kernels/triton_activation.py | 18 ++- tests/test_fused_flux3.py | 135 ++++++++++++++++ tests/test_fused_layout.py | 11 ++ uv.lock | 2 +- 10 files changed, 353 insertions(+), 12 deletions(-) create mode 100644 src/orbitquant/fused/families/flux3.py create mode 100644 tests/test_fused_flux3.py diff --git a/README.md b/README.md index eced5fc..2bbbf25 100644 --- a/README.md +++ b/README.md @@ -23,9 +23,10 @@ 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 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. + MiniMax-H3, Boogu-Image and FLUX.3 Klein (`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; @@ -512,6 +513,7 @@ the same for an `orbitquant-v1` component artifact, and | 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 | +| 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 | 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 b417b6d..27647bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "orbitquant" -version = "0.12.0" +version = "0.13.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 60dcf61..f5f489f 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.12.0" +__version__ = "0.13.0" register_hf_quantizers() diff --git a/src/orbitquant/fused/families/__init__.py b/src/orbitquant/fused/families/__init__.py index dfbefc9..c89057b 100644 --- a/src/orbitquant/fused/families/__init__.py +++ b/src/orbitquant/fused/families/__init__.py @@ -18,7 +18,7 @@ from orbitquant.fused.groups import RowSource -_MODULES = ("krea2", "flux2", "ideogram4", "minimax_h3", "qwenimage21", "boogu") +_MODULES = ("krea2", "flux2", "ideogram4", "minimax_h3", "qwenimage21", "boogu", "flux3") @dataclass(frozen=True) diff --git a/src/orbitquant/fused/families/flux3.py b/src/orbitquant/fused/families/flux3.py new file mode 100644 index 0000000..2e88917 --- /dev/null +++ b/src/orbitquant/fused/families/flux3.py @@ -0,0 +1,170 @@ +"""FLUX.3 Klein (Black Forest Labs' native ``flux3.dit.DiT``) per-mode and single-stream blocks. + +Every block is one parallel attention + SwiGLU block: LayerNorm and AdaLN modulation run in the +activation-quantization prologue, Q|K|V and the SwiGLU gate|up rows are two GEMMs over one set of +codes, Q/K RMSNorm and rotary embedding are one kernel, and both output projections add their +gated residual in the GEMM epilogue. The blocks run one sample at a time, as the distilled model +does. The modulation of a mode is computed once per forward, on one row when every token of the +mode carries the same time embedding (the sampler fills a mode's timesteps with one value). +""" + +from __future__ import annotations + +import types + +import torch +import torch.nn.functional as F + +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 = "flux3" +MODEL_CLASSES = ("DiT",) +EPS = 1e-6 + + +def blocks(model): + paths = [ + f"content_mode_blocks.{mode}.{i}" + for mode in model.content_mode_blocks + for i in range(len(model.content_mode_blocks[mode])) + ] + paths += [f"txt_mode_blocks.{i}" for i in range(len(model.txt_mode_blocks))] + paths += [f"single_blocks.{i}" for i in range(len(model.single_blocks))] + return [(path, "block") for path in paths] + + +def groups(kind, block): + hidden = block.mlp_out.in_features + return specs( + qkv=["q_proj", "k_proj", "v_proj"], + mlp=( + [RowSource("mlp_in", 0, hidden), RowSource("mlp_in", hidden, 2 * hidden)], + {"swiglu": True}, + ), + attn_out=["attn_out"], + mlp_out=["mlp_out"], + ) + + +def prepare(block, kind): + norm = block.norm + block.oq_int8_attention = int8_attention_allowed(norm.query_norm.scale, norm.key_norm.scale) + + +def forward(kind): + return _block + + +def install(model): + """Per-mode modulation returns contiguous bf16 shift, scale and gate once per forward.""" + for module in model.modules(): + if type(module).__name__ == "ModeModulation": + module.forward = types.MethodType(_modulation, module) + + +def _modulation(self, vec): + # One host sync per mode, before the blocks run. A uniform mode costs one row, returned + # expanded over the tokens: the single-stream concatenation still sees one row per token. + rows = vec.shape[0] + uniform = rows > 1 and torch.equal(vec, vec[:1].expand_as(vec)) + chunks = self.lin(F.silu(vec[:1] if uniform else vec)).chunk(3, dim=-1) + chunks = [c.to(torch.bfloat16).contiguous() for c in chunks] + return tuple(c.expand(rows, -1) if uniform else c for c in chunks) + + +def _index(runtime, seq, rows, device): + """Modulation row of every token: all zeros for a one-row modulation, else the token index.""" + cache = runtime.__dict__.setdefault("flux3_index", {}) + key = (seq, rows == 1, device) + index = cache.get(key) + if index is None: + make = torch.zeros if rows == 1 else torch.arange + index = cache[key] = make(seq, device=device, dtype=torch.int32) + return index + + +def _rope(runtime, freqs): + """Rotation matrices ``[1, 1, L, 64, 2, 2]`` as per-channel cos and sin ``[L, 128]``.""" + cached = runtime.__dict__.get("flux3_rope") + if cached is not None and cached[0] is freqs: + return cached[1], cached[2] + matrices = freqs[0, 0] + cos = matrices[..., 0, 0].repeat_interleave(2, dim=-1).contiguous() + sin = matrices[..., 1, 0].repeat_interleave(2, dim=-1).contiguous() + runtime.__dict__["flux3_rope"] = (freqs, cos, sin) + return cos, sin + + +def _block(self, x, freqs, bounds, mod): + if len(bounds) != 1: + raise NotImplementedError("fused FLUX.3 blocks run one sample at a time") + start, end = bounds[0] + if start == end: + return x # the phantom token of an empty mode; its result is never read + runtime, fused = self.oq_runtime, self.oq_fused + shift, scale, gate = mod + if scale.stride(0) == 0: # one modulation row expanded over the tokens + shift, scale, gate = shift[:1], scale[:1], gate[:1] + h = x.to(torch.bfloat16).clone() + seq, dim = h.shape + index = _index(runtime, seq, shift.shape[0], h.device) + source = runtime.input( + h, + dim, + norm=act.NORM_LAYER, + mod=act.MOD_SHIFT_SCALE, + mod_scale=scale, + mod_shift=shift, + mod_index=index, + rms_eps=EPS, + ) + inner = fused.qkv.out_features // 3 + heads = self.heads + head_dim = inner // heads + qkv = runtime.linear(fused.qkv, source) + mid = runtime.linear(fused.mlp, source, epilogue=gemm.EPILOGUE_SWIGLU) + cos, sin = _rope(runtime, freqs) + query = torch.empty((1, seq, heads, head_dim), device=h.device, dtype=torch.bfloat16) + key = torch.empty_like(query) + for src, weight, out in ( + (qkv[:, :inner], self.norm.query_norm.scale, query), + (qkv[:, inner : 2 * inner], self.norm.key_norm.scale, key), + ): + act.qk_norm_rope_with_triton( + src, + heads, + head_dim, + weight, + cos, + sin, + EPS, + weight_mode=act.QK_WEIGHT_TORCH, + style=act.ROPE_INTERLEAVED, + out=out, + ) + value = qkv[:, 2 * inner :].unflatten(-1, (heads, head_dim))[None] + attention = runtime.attention(query, key, value, int8_allowed=self.oq_int8_attention) + attention = attention.reshape(seq, inner) + runtime.linear( + fused.attn_out, + runtime.input(attention, inner), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate, + gate_index=index, + out=h, + ) + runtime.linear( + fused.mlp_out, + runtime.input(mid, mid.shape[1]), + epilogue=gemm.EPILOGUE_RESIDUAL_GATE, + residual=h, + gate=gate, + gate_index=index, + out=h, + ) + return h + diff --git a/src/orbitquant/fused/groups.py b/src/orbitquant/fused/groups.py index d06ffbe..b91e1d2 100644 --- a/src/orbitquant/fused/groups.py +++ b/src/orbitquant/fused/groups.py @@ -256,7 +256,16 @@ def row_scales(self) -> torch.Tensor: 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) + return self.row_norms * self._on_device("scale_factor", self.row_norms.device) + + def _on_device(self, name: str, device: torch.device) -> torch.Tensor: + """A host-built constant buffer, moved to ``device`` once. A copy per call would be a + blocking host-to-device transfer in every GEMM (and illegal under CUDA graph capture).""" + tensor = getattr(self, name) + if tensor.device != device: + tensor = tensor.to(device) + setattr(self, name, tensor) + return tensor def decode(self, workspace: torch.Tensor) -> torch.Tensor: from orbitquant.kernels.triton_activation import decode_lowbit_to_int8_with_triton @@ -268,10 +277,10 @@ def decode(self, workspace: torch.Tensor) -> torch.Tensor: row = offset = 0 for rows, bits in self.segments: nbytes = rows * self.in_features * bits // 8 - codes = getattr(self, f"codes_w{bits}") + codes = self._on_device(f"codes_w{bits}", flat.device) decode_lowbit_to_int8_with_triton( flat[offset : offset + nbytes], - codes.to(flat.device), + codes, bits, rows, self.in_features, diff --git a/src/orbitquant/kernels/triton_activation.py b/src/orbitquant/kernels/triton_activation.py index 6b87f4a..5e8aadb 100644 --- a/src/orbitquant/kernels/triton_activation.py +++ b/src/orbitquant/kernels/triton_activation.py @@ -93,6 +93,7 @@ def _prologue_kernel( NORM: tl.constexpr, MOD: tl.constexpr, ROW_INDEXED: tl.constexpr, + STORE_Y: tl.constexpr, ): row = tl.program_id(0).to(tl.int64) offs = tl.arange(0, DP2) @@ -113,7 +114,8 @@ def _prologue_kernel( 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) + if STORE_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))) @@ -454,7 +456,18 @@ def __call__( 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) + # Without a prologue a bf16 input row is its own prologue output; only norms are written. + passthrough = ( + not multiply + and norm_mode == NORM_NONE + and mod_mode == MOD_NONE + and not prologue_only + and x.dtype == torch.bfloat16 + and x.stride(0) == self.dim + ) + y = x + if not passthrough: + 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: @@ -478,6 +491,7 @@ def __call__( NORM=int(norm_mode), MOD=int(mod_mode), ROW_INDEXED=mod_index is not None, + STORE_Y=not passthrough, 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. diff --git a/tests/test_fused_flux3.py b/tests/test_fused_flux3.py new file mode 100644 index 0000000..c264436 --- /dev/null +++ b/tests/test_fused_flux3.py @@ -0,0 +1,135 @@ +import types + +import pytest +import torch +from torch import nn + +pytest.importorskip("triton") +from orbitquant.fused.families import adapter_for_class, adapter_for_family, flux3 # noqa: E402 + + +def test_family_is_registered_for_the_native_transformer_class(): + assert adapter_for_family("flux3") is flux3 + assert adapter_for_class("DiT") is flux3 + + +def _stub_model(modes=("image", "image_cond"), depth=2, single=3): + def block(): + return nn.Module() + + model = nn.Module() + model.content_mode_blocks = nn.ModuleDict( + {mode: nn.ModuleList([block() for _ in range(depth)]) for mode in modes} + ) + model.txt_mode_blocks = nn.ModuleList([block() for _ in range(depth)]) + model.single_blocks = nn.ModuleList([block() for _ in range(single)]) + return model + + +def test_every_block_of_every_mode_is_listed_once(): + paths = [path for path, _ in flux3.blocks(_stub_model())] + assert paths == [ + "content_mode_blocks.image.0", + "content_mode_blocks.image.1", + "content_mode_blocks.image_cond.0", + "content_mode_blocks.image_cond.1", + "txt_mode_blocks.0", + "txt_mode_blocks.1", + "single_blocks.0", + "single_blocks.1", + "single_blocks.2", + ] + + +def test_groups_split_mlp_in_into_gate_and_up_halves(): + block = nn.Module() + block.mlp_out = nn.Linear(9216, 3072, bias=False) + specs = flux3.groups("block", block) + assert list(specs) == ["qkv", "mlp", "attn_out", "mlp_out"] + gate, up = specs["mlp"].sources + assert (gate.path, gate.start, gate.stop) == ("mlp_in", 0, 9216) + assert (up.path, up.start, up.stop) == ("mlp_in", 9216, 18432) + assert specs["mlp"].swiglu and not specs["qkv"].swiglu + + +def _reference_rope(pos, dim, theta): + scale = torch.arange(0, dim, 2, dtype=torch.float64) / dim + out = torch.einsum("...n,d->...nd", pos, 1.0 / (theta**scale)) + out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) + return out.reshape(*out.shape[:-1], 2, 2).float() + + +def test_rotation_matrices_become_per_channel_cos_and_sin(): + positions = torch.arange(5, dtype=torch.float32) + ids = torch.stack([torch.zeros(5), positions, positions * 0.5, positions], dim=-1) + matrices = torch.cat([_reference_rope(ids[..., i], 32, 10000) for i in range(4)], dim=-3) + freqs = matrices.unsqueeze(0).unsqueeze(0) + runtime = types.SimpleNamespace(__dict__={}) + cos, sin = flux3._rope(runtime, freqs) + assert cos.shape == sin.shape == (5, 128) + assert torch.equal(cos[:, 0::2], cos[:, 1::2]) + assert torch.equal(sin[:, 0::2], sin[:, 1::2]) + assert torch.equal(cos[:, ::2], matrices[..., 0, 0]) + assert torch.equal(sin[:, ::2], matrices[..., 1, 0]) + assert flux3._rope(runtime, freqs)[0] is cos # cached per forward + + +def test_modulation_is_contiguous_bf16_chunks(): + class ModeModulation(nn.Module): + def __init__(self): + super().__init__() + self.lin = nn.Linear(8, 24, bias=False) + + model = nn.Module() + model.mod = ModeModulation() + flux3.install(model) + shift, scale, gate = model.mod.forward(torch.randn(5, 8)) + for chunk in (shift, scale, gate): + assert chunk.dtype == torch.bfloat16 and chunk.is_contiguous() and chunk.shape == (5, 8) + + +def test_multiple_samples_are_rejected(): + block = nn.Module() + with pytest.raises(NotImplementedError, match="one sample"): + flux3._block(block, torch.zeros(2, 4), None, [(0, 1), (1, 2)], None) + + +def _installed_modulation(): + class ModeModulation(nn.Module): + def __init__(self): + super().__init__() + self.lin = nn.Linear(8, 24, bias=False) + + model = nn.Module() + model.mod = ModeModulation() + flux3.install(model) + return model.mod + + +def test_uniform_mode_computes_one_row_expanded_over_the_tokens(): + module = _installed_modulation() + vec = torch.randn(1, 8).expand(6, 8).clone() + chunks = module.forward(vec) + projected = module.lin(torch.nn.functional.silu(vec[:1])).expand(6, 24) + reference = [c.to(torch.bfloat16) for c in projected.chunk(3, -1)] + for chunk, ref in zip(chunks, reference, strict=True): + assert chunk.shape == (6, 8) and chunk.stride(0) == 0 + assert torch.equal(chunk, ref) + # The single-stream concatenation of slices materializes one row per token. + assert torch.equal(torch.cat([chunk[0:4], chunk[4:6]]), ref) + + +def test_varying_rows_keep_one_modulation_row_per_token(): + module = _installed_modulation() + shift, _, _ = module.forward(torch.randn(5, 8)) + assert shift.is_contiguous() and shift.stride(0) == 8 + + +def test_token_index_is_cached_per_length_and_kind(): + runtime = types.SimpleNamespace(__dict__={}) + zeros = flux3._index(runtime, 4, 1, torch.device("cpu")) + tokens = flux3._index(runtime, 4, 4, torch.device("cpu")) + assert torch.equal(zeros, torch.zeros(4, dtype=torch.int32)) + assert torch.equal(tokens, torch.arange(4, dtype=torch.int32)) + assert flux3._index(runtime, 4, 1, torch.device("cpu")) is zeros + diff --git a/tests/test_fused_layout.py b/tests/test_fused_layout.py index 73595da..93fc395 100644 --- a/tests/test_fused_layout.py +++ b/tests/test_fused_layout.py @@ -23,3 +23,14 @@ def test_recognizing_fused_checkpoints_does_not_need_the_triton_runtime(): command = [sys.executable, "-c", code] result = subprocess.run(command, capture_output=True, text=True, check=True) assert result.stdout.strip() == "False" + + +def test_host_constants_move_to_the_weight_device_once(): + import torch + + group = torch.nn.Module() + group.register_buffer("codes_w4", torch.arange(16, dtype=torch.int8), persistent=False) + first = PackedGroup._on_device(group, "codes_w4", torch.device("meta")) + assert first.device.type == "meta" and group.codes_w4 is first + assert PackedGroup._on_device(group, "codes_w4", torch.device("meta")) is first + assert "codes_w4" in dict(group.named_buffers()) diff --git a/uv.lock b/uv.lock index 695e00b..08d3609 100644 --- a/uv.lock +++ b/uv.lock @@ -844,7 +844,7 @@ wheels = [ [[package]] name = "orbitquant" -version = "0.12.0" +version = "0.13.0" source = { editable = "." } dependencies = [ { name = "huggingface-hub" },