"""OrbitQuant inference primitives — RPBH rotation and codebook dequantization. Implements the Randomized Permuted Block-Hadamard (RPBH) rotation and Lloyd-Max codebook lookup for OrbitQuant (arXiv:2607.02461) inference. The RPBH rotation is identical to the converter's implementation to ensure exact cancellation: (W @ Π^T) @ (x @ Π^T)^T = W @ x^T. """ from __future__ import annotations import math import torch import torch.nn.functional as F from ._quant_utils import unpack_int4, unpack_int8 __all__ = [ "rpbh_rotate", "codebook_dequant", "codebook_dequant_int4", "codebook_dequant_int8", "adaln_rtn_dequant", "parse_orbitquant_metadata", "load_orbitquant_globals", "clear_orbitquant_globals", "make_rpbh_params", "largest_power_of_two_dividing", "fast_walsh_hadamard", "get_quant_scheme_from_metadata", "infer_quant_scheme_from_scale", "infer_orbitquant_metadata_from_state_dict", "rpbh_rotate_fast", "codebook_dequant_fast", # RPBH uniform-grid (non-OrbitQuant) support "parse_rpbh_metadata", "load_rpbh_globals", "infer_rpbh_metadata_from_state_dict", ] # ── Constants ── ORBITQUANT_NORM_EPS = 1e-10 # paper §4.3 # ── Global tensor cache ── # Structure: # _orbitquant_globals[model_path]["rpbh"][d] = { # "perm": Tensor, "signs": Tensor, "h": int, "d_padded": int # } # _orbitquant_globals[model_path]["codebooks"][(d, bits)] = Tensor [2^b] _orbitquant_globals: dict[str, dict] = {} def clear_orbitquant_globals(model_path: str | None = None): """Free cached RPBH params and codebooks for a model (or all). Called when: - A model is unloaded (from _reset_aimdo_state or loader) - A new model is loaded (clear previous model's cache) - Tests need a clean slate """ if model_path is None: _orbitquant_globals.clear() else: _orbitquant_globals.pop(model_path, None) # ── RPBH rotation ── def largest_power_of_two_dividing(n: int) -> int: """Largest power of 2 dividing n. Identical to converter.""" if n <= 0: raise ValueError(f"n must be positive, got {n}") return n & (-n) def fast_walsh_hadamard(x: torch.Tensor) -> torch.Tensor: """Normalized FWHT on last dim. Identical to converter. Uses step-based contiguous slicing. The last dimension must be a power of 2. Normalized so H @ H^T = I. """ h = x.shape[-1] result = x.clone() step = 1 while step < h: leading = result.shape[:-1] blocks = result.reshape(*leading, -1, 2 * step) left = blocks[..., :step].clone() right = blocks[..., step:].clone() blocks[..., :step] = left + right blocks[..., step:] = left - right result = blocks.reshape(*leading, h) step *= 2 return result / math.sqrt(h) def make_rpbh_params( d: int, seed: int = 42, block_size: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, int]: """Generate RPBH rotation parameters for dimension d. Identical to the converter's parameter generation: 1. Determine block size h (largest power of 2 dividing d, or explicit) 2. Pad d to next multiple of h → d_padded 3. Generate random permutation of [0, d_padded) 4. Generate random Rademacher signs (±1) of length d_padded Returns (perm, signs, h) where: - perm: LongTensor [d_padded] — permutation indices - signs: FloatTensor [d_padded] — ±1 values - h: int — Hadamard block size """ h = block_size if block_size is not None else largest_power_of_two_dividing(d) d_padded = math.ceil(d / h) * h gen = torch.Generator() gen.manual_seed(seed) # Permutation: random shuffle of [0, d_padded) perm = torch.randperm(d_padded, generator=gen) # Rademacher signs: ±1 with equal probability signs_raw = torch.randint(0, 2, (d_padded,), generator=gen).float() signs = signs_raw * 2.0 - 1.0 # map {0, 1} → {-1, +1} return perm, signs, h def rpbh_rotate( x: torch.Tensor, perm: torch.Tensor, signs: torch.Tensor, h: int, ) -> torch.Tensor: """Apply RPBH rotation: x ↦ x @ Π_d^T. Identical to converter's implementation: 1. Permutation gather (P_π^T acts first — CRITICAL for mass-balancing) 2. Block FWHT 3. Rademacher sign flip Computes in at least float32 precision for numerical accuracy over log2(h) butterfly stages: fp16/bf16 inputs are upcast to float32, float32/float64 inputs are left at their native dtype (so float64 callers keep full precision). Returns same dtype as input. Cancellation: rpbh_rotate(W) @ rpbh_rotate(x).T = W @ x.T """ orig_d = x.shape[-1] d_padded = perm.shape[0] orig_dtype = x.dtype # Upcast low-precision inputs to float32 for the butterfly stages. # float32 and float64 are already precise enough — casting float64 # down to float32 would silently lose ~9 decimal digits. compute_dtype = torch.float32 if orig_dtype in (torch.float16, torch.bfloat16) else orig_dtype x = x.to(compute_dtype) if orig_d < d_padded: x = F.pad(x, (0, d_padded - orig_d)) # 1. Permutation gather (P_π^T acts first on row vector) x_perm = x[..., perm.long()] # 2. Block fast Walsh–Hadamard transform k = d_padded // h leading = x_perm.shape[:-1] x_blocks = x_perm.reshape(*leading, k, h) x_hadamard = fast_walsh_hadamard(x_blocks) x_hadamard = x_hadamard.reshape(*leading, d_padded) # 3. Rademacher sign flip x_rot = x_hadamard * signs.to(x_hadamard.dtype) return x_rot.to(orig_dtype) # ── Codebook dequantization ── def codebook_dequant( qdata: torch.Tensor, norms: torch.Tensor, codebook: torch.Tensor, d: int, d_padded: int, bits: int, ) -> torch.Tensor: """Dequantize OrbitQuant codebook weights to float. 1. Unpack indices from packed format 2. Offset from signed int8 range back to [0, 2^b-1] 3. Lookup codebook centroids 4. Multiply by per-row L2 norms Returns: [out, d_padded] float tensor. NOTE: Does NOT trim to d. At inference, both weight and activation are in the d_padded-dimensional RPBH-rotated basis. Trimming would cause a dimension mismatch in the matmul. The converter's §3 trimming is for offline reconstruction only. """ if bits == 4: return codebook_dequant_int4(qdata, norms, codebook, d, d_padded) else: return codebook_dequant_int8(qdata, norms, codebook, d, d_padded) def codebook_dequant_int4( packed: torch.Tensor, norms: torch.Tensor, codebook: torch.Tensor, d: int, d_padded: int, ) -> torch.Tensor: """Dequantize packed INT4 codebook indices to float weights. Returns: [out, d_padded] float tensor. """ # Unpack uint8 → int8 [-8, 7] unpacked = unpack_int4(packed, d_padded) # [out, d_padded] int8 # Offset to [0, 15] indices = (unpacked.to(torch.int16) + 8).to(torch.long) # Codebook lookup direction = codebook[indices] # [out, d_padded] float # Multiply by per-row L2 norms (reshape to [out, 1] so the broadcast # scales each row, regardless of whether norms is 1-D [out] or 2-D [out, 1]). norms_2d = norms.float().reshape(-1, 1) return (direction * norms_2d).to(direction.dtype) def codebook_dequant_int8( indices_int8: torch.Tensor, norms: torch.Tensor, codebook: torch.Tensor, d: int, d_padded: int, ) -> torch.Tensor: """Dequantize INT8 codebook indices to float weights. Returns: [out, d_padded] float tensor. """ # Offset int8 [-128, 127] → [0, 255] indices = (indices_int8.to(torch.int16) + 128).to(torch.long) # Codebook lookup direction = codebook[indices] # [out, d_padded] float # Multiply by per-row L2 norms (reshape to [out, 1] so the broadcast # scales each row, regardless of whether norms is 1-D [out] or 2-D [out, 1]). norms_2d = norms.float().reshape(-1, 1) return (direction * norms_2d).to(direction.dtype) # ── AdaLN RTN dequantization ── def adaln_rtn_dequant( packed: torch.Tensor, scales: torch.Tensor, d: int, bits: int, ) -> torch.Tensor: """Dequantize AdaLN RTN weights (standard scale-based, group_size=64). AdaLN uses per-group RTN scales with shape [out, d/64]. This is different from the codebook path which uses per-row norms. Returns: [out, d] float tensor. """ if bits == 4: unpacked = unpack_int4(packed, d) # [out, d] int8 [-8, 7] else: unpacked = unpack_int8(packed, d) # [out, d] int8 [-128, 127] # scales shape: [out, d/64] float16 # unpacked shape: [out, d] int8 # Need to reshape for group-wise multiplication out_features = unpacked.shape[0] num_groups = scales.shape[1] if d % num_groups != 0: raise ValueError( f"adaln_rtn_dequant: d ({d}) must be divisible by the number of " f"scale groups ({num_groups}), but {d} % {num_groups} != 0. " f"This indicates a malformed state dict (the converter guarantees " f"divisibility)." ) group_size = d // num_groups unpacked_grouped = unpacked.reshape(out_features, num_groups, group_size) scales_3d = scales.float().unsqueeze(2) # [out, num_groups, 1] dequant_grouped = unpacked_grouped.float() * scales_3d return dequant_grouped.reshape(out_features, d) # ── Metadata parsing ── def parse_orbitquant_metadata(metadata: dict) -> dict | None: """Extract OrbitQuant config from safetensors metadata. Returns None if not an OrbitQuant model. Raises ValueError naming the offending key if a field is present but not a valid integer. """ def _meta_int(key: str, default: int | None = None, allow_none: bool = False) -> int | None: if key not in metadata or metadata[key] is None: if allow_none: return None if default is not None: return default raw = metadata[key] else: raw = metadata[key] try: return int(raw) except (TypeError, ValueError) as exc: raise ValueError( f"OrbitQuant metadata field {key!r} is not a valid integer: " f"got {raw!r}" ) from exc if not metadata: return None if metadata.get("int_crush.format_version") != "4": return None if metadata.get("int_crush.orbitquant") != "true": return None dims_str = metadata.get("int_crush.orbitquant_dims", "") dims = [int(d) for d in dims_str.split(",") if d.strip()] d_padded_str = metadata.get("int_crush.orbitquant_d_padded", "") d_padded_map = {} for pair in d_padded_str.split(","): if ":" in pair: k, v = pair.split(":") d_padded_map[int(k)] = int(v) block_size = _meta_int("int_crush.orbitquant_block_size", allow_none=True) return { "rpbh_seed": _meta_int("int_crush.rpbh_seed", default=42), "codebook_bits": _meta_int("int_crush.codebook_bits", default=4), "act_bits": _meta_int("int_crush.act_bits", allow_none=True), "dims": dims, "d_padded_map": d_padded_map, "block_size": block_size, "normal_approx": metadata.get("int_crush.orbitquant_normal_approx") == "true", "format_version": metadata.get("int_crush.format_version", "4"), } def load_orbitquant_globals( state_dict: dict, metadata: dict, model_path: str, ) -> dict: """Load and cache all global OrbitQuant tensors from state_dict. Pops from state_dict: - _rpbh_perm_{d} for each d - _rpbh_signs_{d} for each d - _orbitquant_codebook_{d}_{bits} for each (d, bits) Returns the globals dict for this model. """ if model_path in _orbitquant_globals: return _orbitquant_globals[model_path] oq_meta = parse_orbitquant_metadata(metadata) if oq_meta is None: raise ValueError("Not an OrbitQuant model — metadata missing") globals_entry: dict = {"rpbh": {}, "codebooks": {}} for d in oq_meta["dims"]: d_padded = oq_meta["d_padded_map"].get(d, d) h = largest_power_of_two_dividing(d) # Use explicit block_size override if present if oq_meta["block_size"] is not None: h = oq_meta["block_size"] perm_key = f"_rpbh_perm_{d}" signs_key = f"_rpbh_signs_{d}" perm = state_dict.pop(perm_key, None) signs = state_dict.pop(signs_key, None) if perm is None or signs is None: raise ValueError( f"OrbitQuant: missing RPBH tensors for d={d}. " f"Expected {perm_key} and {signs_key} in state_dict." ) globals_entry["rpbh"][d] = { "perm": perm, "signs": signs, "h": h, "d_padded": d_padded, } # Load codebooks (weight + optionally activation) bits = oq_meta["codebook_bits"] act_bits = oq_meta.get("act_bits") or bits for d in oq_meta["dims"]: for b in {bits, act_bits}: cb_key = f"_orbitquant_codebook_{d}_{b}" cb = state_dict.pop(cb_key, None) if cb is not None: globals_entry["codebooks"][(d, b)] = cb _orbitquant_globals[model_path] = globals_entry return globals_entry def infer_orbitquant_metadata_from_state_dict(state_dict: dict) -> dict | None: """Infer OrbitQuant metadata from global tensor keys in state_dict. Used as a fallback when safetensors metadata is unavailable. Scans for _rpbh_perm_{d} and _orbitquant_codebook_{d}_{bits} keys to reconstruct the metadata dict needed by load_orbitquant_globals. Returns None if no OrbitQuant globals are found. """ import re dims = set() bits_values = set() d_padded_map = {} for key in state_dict: m = re.match(r"^_rpbh_perm_(\d+)$", key) if m: d = int(m.group(1)) dims.add(d) # d_padded = perm.shape[0] d_padded_map[d] = state_dict[key].shape[0] continue m = re.match(r"^_orbitquant_codebook_(\d+)_(\d+)$", key) if m: d = int(m.group(1)) b = int(m.group(2)) dims.add(d) bits_values.add(b) if not dims: return None codebook_bits = min(bits_values) if bits_values else 4 return { "int_crush.format_version": "4", "int_crush.orbitquant": "true", "int_crush.rpbh_seed": "42", "int_crush.codebook_bits": str(codebook_bits), "int_crush.orbitquant_dims": ",".join(str(d) for d in sorted(dims)), "int_crush.orbitquant_d_padded": ",".join( f"{d}:{d_padded_map.get(d, d)}" for d in sorted(dims) ), } def get_quant_scheme_from_metadata(metadata: dict, prefix: str) -> str | None: """Read per-layer quant_scheme from safetensors metadata. The converter stores keys like: "int_crush.layers.0.q_proj.weight.quant_scheme": "orbitquant_codebook" ComfyUI's _load_from_state_dict passes prefixes with a model-level prefix (e.g. "model.diffusion_model.layers.0.q_proj.") that the converter does NOT include in its metadata keys. We try progressively shorter prefixes to handle this mismatch. Args: metadata: The safetensors file metadata dict. prefix: The module prefix from _load_from_state_dict, e.g. "model.diffusion_model.layers.0.q_proj." (includes trailing dot). Returns: "orbitquant_codebook", "adaln_rtn", or None if not found. """ if not metadata: return None param_name = prefix.rstrip(".") # remove trailing dot if present if not param_name.endswith(".weight"): param_name = param_name + ".weight" # Try progressively shorter prefixes to strip model-level path components # that the converter doesn't include in metadata keys. # e.g. "model.diffusion_model.single_blocks.0.linear1.weight" → # try "model.diffusion_model.single_blocks.0.linear1.weight" # try "diffusion_model.single_blocks.0.linear1.weight" # try "single_blocks.0.linear1.weight" ← matches converter key parts = param_name.split(".") for i in range(len(parts)): candidate = ".".join(parts[i:]) key = f"int_crush.{candidate}.quant_scheme" value = metadata.get(key) if value in ("orbitquant_codebook", "adaln_rtn"): return value return None def infer_quant_scheme_from_scale(scale: torch.Tensor) -> str: """Infer quant_scheme from scale tensor shape. - [out, 1] → orbitquant_codebook (per-row L2 norms) - [out, N] where N > 1 → adaln_rtn (per-group RTN scales) Fallback when per-layer metadata is unavailable. """ if scale.dim() == 2 and scale.shape[1] == 1: return "orbitquant_codebook" elif scale.dim() == 2 and scale.shape[1] > 1: return "adaln_rtn" else: # 1D scale → per-row (orbitquant_codebook) return "orbitquant_codebook" # ── RPBH uniform-grid (non-OrbitQuant) metadata parsing ──────────────────── def _meta_int_rpbh(metadata: dict, key: str, default: int) -> int: """Parse an int metadata field, raising ValueError naming the key on malformed values.""" raw = metadata.get(key, default) try: return int(raw) except (TypeError, ValueError) as exc: raise ValueError( f"RPBH metadata field {key!r} is not a valid integer: got {raw!r}" ) from exc def parse_rpbh_metadata(metadata: dict) -> dict | None: """Extract RPBH uniform-grid config from safetensors metadata. Returns None if not an RPBH uniform-grid model (format version 1/3 with ``int_crush.rpbh == "true"``), or for OrbitQuant models (format version 4, handled by :func:`parse_orbitquant_metadata`). """ if not metadata: return None # OrbitQuant models (format 4) have their own RPBH handling. if metadata.get("int_crush.format_version") == "4": return None if metadata.get("int_crush.rpbh") != "true": return None dims_str = metadata.get("int_crush.rpbh_dims", "") dims = [int(d) for d in dims_str.split(",") if d.strip()] d_padded_str = metadata.get("int_crush.rpbh_d_padded", "") d_padded_map: dict[int, int] = {} for pair in d_padded_str.split(","): if ":" in pair: k, v = pair.split(":") d_padded_map[int(k)] = int(v) block_size = _meta_int_rpbh(metadata, "int_crush.rpbh_block_size", default=0) # rot_size > 0 means uniform h for all dims (explicit --rot-size). # rot_size == 0 means per-dim auto-detect (largest_power_of_two_dividing). rot_size = _meta_int_rpbh(metadata, "int_crush.rot_size", default=0) return { "rpbh_seed": _meta_int_rpbh(metadata, "int_crush.rpbh_seed", default=42), "block_size": block_size, "rot_size": rot_size, "dims": dims, "d_padded_map": d_padded_map, } def load_rpbh_globals( state_dict: dict, metadata: dict, model_path: str, ) -> dict: """Load and cache RPBH global tensors (perm + signs) from *state_dict*. Pops from *state_dict*: - ``_rpbh_perm_{d}`` for each *d* - ``_rpbh_signs_{d}`` for each *d* No codebooks — this is the uniform-grid path. Returns the globals dict: ``{"rpbh": {d: {"perm", "signs", "h", "d_padded"}}, "codebooks": {}}`` """ if model_path in _orbitquant_globals: return _orbitquant_globals[model_path] rpbh_meta = parse_rpbh_metadata(metadata) if rpbh_meta is None: raise ValueError("Not an RPBH uniform-grid model — metadata missing") globals_entry: dict = {"rpbh": {}, "codebooks": {}} for d in rpbh_meta["dims"]: d_padded = rpbh_meta["d_padded_map"].get(d, d) # Block size priority: # 1. rpbh_block_size from metadata (set by calibration node) # 2. rot_size (legacy: explicit --rot-size in converter) # 3. Auto-detect per dimension if rpbh_meta["block_size"] > 0: h = rpbh_meta["block_size"] elif rpbh_meta["rot_size"] > 0: h = rpbh_meta["rot_size"] else: h = largest_power_of_two_dividing(d) perm_key = f"_rpbh_perm_{d}" signs_key = f"_rpbh_signs_{d}" perm = state_dict.pop(perm_key, None) signs = state_dict.pop(signs_key, None) if perm is None or signs is None: raise ValueError( f"RPBH: missing RPBH tensors for d={d}. " f"Expected {perm_key} and {signs_key} in state_dict." ) globals_entry["rpbh"][d] = { "perm": perm, "signs": signs, "h": h, "d_padded": d_padded, } _orbitquant_globals[model_path] = globals_entry return globals_entry def infer_rpbh_metadata_from_state_dict(state_dict: dict) -> dict | None: """Infer RPBH uniform-grid metadata from global tensor keys in *state_dict*. Used as a fallback when safetensors metadata is unavailable. Scans for ``_rpbh_perm_{d}`` keys **without** ``_orbitquant_codebook_*`` keys — if codebook keys are present this is OrbitQuant, not RPBH-uniform. Returns ``None`` if no RPBH globals are found. """ import re dims: set[int] = set() d_padded_map: dict[int, int] = {} for key in state_dict: m = re.match(r"^_rpbh_perm_(\d+)$", key) if m: d = int(m.group(1)) dims.add(d) d_padded_map[d] = state_dict[key].shape[0] if not dims: return None # Distinguish from OrbitQuant: if codebook keys present, this is OrbitQuant. has_codebook = any( re.match(r"^_orbitquant_codebook_\d+_\d+$", key) for key in state_dict ) if has_codebook: return None return { "int_crush.rpbh": "true", "int_crush.rpbh_seed": "42", "int_crush.rpbh_block_size": "0", "int_crush.rpbh_dims": ",".join(str(d) for d in sorted(dims)), "int_crush.rpbh_d_padded": ",".join( f"{d}:{d_padded_map.get(d, d)}" for d in sorted(dims) ), } # ── Fast dispatchers (Triton when available, else PyTorch reference) ────────── # # These wrap the reference rpbh_rotate / codebook_dequant with a Triton fast # path. The reference functions above remain the numerical ground truth and # the CPU / non-Triton fallback. Callers that need to force the PyTorch path # (e.g. _intcrush_use_pytorch) should call the reference functions directly. def rpbh_rotate_fast( x: torch.Tensor, perm: torch.Tensor, signs: torch.Tensor, h: int, d_padded: int | None = None, ) -> torch.Tensor: """RPBH rotation via Triton when available, else PyTorch reference. If the Triton kernel raises a CUDA OOM, we clear the caching allocator and retry once with the PyTorch reference path. """ from . import _triton_runtime as _tr if _tr.TRITON_RPBH_ROTATE and x.is_cuda: try: return _tr.rpbh_rotate_triton(x, perm, signs, h, d_padded) except RuntimeError as exc: if "out of memory" in str(exc).lower(): torch.cuda.empty_cache() return rpbh_rotate(x, perm, signs, h) raise return rpbh_rotate(x, perm, signs, h) def codebook_dequant_fast( qdata: torch.Tensor, norms: torch.Tensor, codebook: torch.Tensor, d: int, d_padded: int, bits: int, ) -> torch.Tensor: """Codebook dequant via Triton when available, else PyTorch reference. Returns float32 when using the PyTorch path (matches the reference codebook_dequant exactly), or the caller-controlled dtype of the Triton output buffer when using Triton (default float16 for inference). If the Triton kernel raises a CUDA OOM, we clear the caching allocator and retry once with the PyTorch reference path. """ from . import _triton_runtime as _tr if _tr.TRITON_CODEBOOK_DEQUANT and qdata.is_cuda: try: return _tr.codebook_dequant_triton(qdata, norms, codebook, d_padded, bits) except RuntimeError as exc: if "out of memory" in str(exc).lower(): torch.cuda.empty_cache() return codebook_dequant(qdata, norms, codebook, d, d_padded, bits) raise return codebook_dequant(qdata, norms, codebook, d, d_padded, bits)