733 lines
24 KiB
Python
733 lines
24 KiB
Python
"""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)
|