FLUX.3 Klein fused family; one-time device move of fused group constants (0.13.0)

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

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"

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()
 

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)

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
+

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,

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.