Core API

turboquant.core.codebook

Lloyd-Max codebook computation for the exact Beta distribution arising from random rotation of unit-norm vectors.

Math background (paper Section 3.1, Lemma 1): After multiplying a unit-norm vector x in S^{d-1} by a Haar-random rotation matrix Pi, each coordinate y_j = (Pi x)_j follows:

  f(t) = Gamma(d/2) / (sqrt(pi) * Gamma((d-1)/2))
         * (1 - t^2)^((d-3)/2), t in [-1, 1]

which is a scaled-shifted Beta distribution. For large d this converges to N(0, 1/d), but we use the EXACT integral at every d so the codebook is numerically precise for all head dimensions such as 64, 96, 128, 256.

Lloyd-Max algorithm: 1. Initialise 2^bits centroids at quantile midpoints of f. 2. Alternate: a. Boundaries = midpoints between consecutive centroids. b. Centroids = conditional means E[y | boundary[i] < y < boundary[i+1]]. 3. Stop when cost delta < tol.

Caching: Solved codebooks are stored under /core/codebooks/ as JSON files named "d{D}_b{B}.json". They are loaded on first use and kept in RAM.

Both the exact Beta and the Gaussian approximation are exposed. Use exact Beta for production; Gaussian is kept as a diagnostic path.

beta_pdf

beta_pdf(t: ndarray, d: int) -> np.ndarray

PDF of a single coordinate of a uniformly random point on S^{d-1}.

f(t) = Gamma(d/2) / (sqrt(pi) * Gamma((d-1)/2))
       * (1 - t^2)^((d-3)/2)

Args: t : array of evaluation points in (-1, 1). d : dimension of the embedding space, d >= 3.

Returns: Array of probability density values, same shape as t.

Source code in turboquant/core/codebook.py

def beta_pdf(t: np.ndarray, d: int) -> np.ndarray:
    """
    PDF of a single coordinate of a uniformly random point on S^{d-1}.

        f(t) = Gamma(d/2) / (sqrt(pi) * Gamma((d-1)/2))
               * (1 - t^2)^((d-3)/2)

    Args:
        t  : array of evaluation points in (-1, 1).
        d  : dimension of the embedding space, d >= 3.

    Returns:
        Array of probability density values, same shape as t.
    """
    if d < 3:
        raise ValueError(f"d must be >= 3, got {d}")
    log_const = (
        special.gammaln(d / 2.0)
        - 0.5 * math.log(math.pi)
        - special.gammaln((d - 1) / 2.0)
    )
    exponent = (d - 3) / 2.0
    t_clipped = np.clip(t, -1.0 + 1e-14, 1.0 - 1e-14)
    return np.exp(log_const + exponent * np.log(1.0 - t_clipped ** 2))

gaussian_approx_pdf

gaussian_approx_pdf(t: ndarray, d: int) -> np.ndarray

Gaussian approximation N(0, 1/d) of the Beta PDF (diagnostic path only). Exact Beta should be preferred for production codebooks.

Source code in turboquant/core/codebook.py

def gaussian_approx_pdf(t: np.ndarray, d: int) -> np.ndarray:
    """
    Gaussian approximation N(0, 1/d) of the Beta PDF (diagnostic path only).
    Exact Beta should be preferred for production codebooks.
    """
    sigma2 = 1.0 / d
    return np.exp(-0.5 * t ** 2 / sigma2) / math.sqrt(2.0 * math.pi * sigma2)

compute_lloyd_max

compute_lloyd_max(d: int, bits: int, *, use_exact: bool = True, max_iter: int = 300, tol: float = 1e-13) -> dict

Solve the 1-D Lloyd-Max problem for the Beta (or Gaussian) distribution.

Args: d : head dimension. bits : bits per coordinate (1-8). use_exact : if True use exact Beta; if False use Gaussian approximation. max_iter : maximum Lloyd-Max iterations. tol : convergence threshold on absolute cost delta.

Returns dict with keys: centroids : list of 2^bits float64 centroids (sorted). boundaries : list of 2^bits + 1 boundaries (includes -1 and +1). mse_per_dim : scalar, MSE per coordinate dimension. d, bits, use_exact.

Source code in turboquant/core/codebook.py

def compute_lloyd_max(
    d: int,
    bits: int,
    *,
    use_exact: bool = True,
    max_iter: int = 300,
    tol: float = 1e-13,
) -> dict:
    """
    Solve the 1-D Lloyd-Max problem for the Beta (or Gaussian) distribution.

    Args:
        d         : head dimension.
        bits      : bits per coordinate (1-8).
        use_exact : if True use exact Beta; if False use Gaussian approximation.
        max_iter  : maximum Lloyd-Max iterations.
        tol       : convergence threshold on absolute cost delta.

    Returns dict with keys:
        centroids   : list of 2^bits float64 centroids (sorted).
        boundaries  : list of 2^bits + 1 boundaries (includes -1 and +1).
        mse_per_dim : scalar, MSE per coordinate dimension.
        d, bits, use_exact.
    """
    if bits < 1 or bits > 8:
        raise ValueError(f"bits must be in [1, 8], got {bits}")
    n = 2 ** bits
    mean_fn = _conditional_mean if use_exact else (
        lambda lo, hi, _d: _conditional_mean_gaussian(lo, hi, _d)
    )

    centroids = _initial_centroids(d, n, use_exact)
    prev_cost = float("inf")

    for _ in range(max_iter):
        # Step A: boundaries = midpoints between consecutive centroids
        bnd = np.empty(n + 1)
        bnd[0] = -1.0
        bnd[-1] = 1.0
        bnd[1:-1] = (centroids[:-1] + centroids[1:]) / 2.0

        # Step B: centroids = conditional means
        new_c = np.array([mean_fn(bnd[i], bnd[i + 1], d) for i in range(n)])

        # Recompute cost with the EXACT Beta regardless of approximation mode
        cost = _mse_cost(new_c, bnd, d)
        centroids = new_c

        if abs(prev_cost - cost) < tol:
            break
        prev_cost = cost

    # Last boundaries
    bnd = np.empty(n + 1)
    bnd[0] = -1.0
    bnd[-1] = 1.0
    bnd[1:-1] = (centroids[:-1] + centroids[1:]) / 2.0

    return {
        "centroids": centroids.tolist(),
        "boundaries": bnd.tolist(),
        "mse_per_dim": float(cost),
        "d": d,
        "bits": bits,
        "use_exact": use_exact,
    }

get_codebook

get_codebook(d: int, bits: int, *, use_exact: bool = True) -> dict

Return a Lloyd-Max codebook, loading from disk or computing on first call.

Args: d : head dimension. bits : bits per coordinate. use_exact : use exact Beta distribution (True) or Gaussian approx (False).

Returns: dict with 'centroids', 'boundaries', 'mse_per_dim', 'd', 'bits'.

Source code in turboquant/core/codebook.py

def get_codebook(
    d: int,
    bits: int,
    *,
    use_exact: bool = True,
) -> dict:
    """
    Return a Lloyd-Max codebook, loading from disk or computing on first call.

    Args:
        d         : head dimension.
        bits      : bits per coordinate.
        use_exact : use exact Beta distribution (True) or Gaussian approx (False).

    Returns:
        dict with 'centroids', 'boundaries', 'mse_per_dim', 'd', 'bits'.
    """
    key = (d, bits, use_exact)
    if key in _CB_RAM:
        return _CB_RAM[key]

    path = _cache_path(d, bits, use_exact)
    if os.path.exists(path):
        with open(path) as fh:
            cb = json.load(fh)
        _CB_RAM[key] = cb
        return cb

    tag = "exact Beta" if use_exact else "Gaussian approx"
    logger.info("computing Lloyd-Max codebook d=%s bits=%s (%s)", d, bits, tag)
    cb = compute_lloyd_max(d, bits, use_exact=use_exact)
    with open(path, "w") as fh:
        json.dump(cb, fh, indent=2)
    logger.info("Lloyd-Max MSE per dim: %.6e", cb["mse_per_dim"])
    _CB_RAM[key] = cb
    return cb

get_codebook_tensors

get_codebook_tensors(d: int, bits: int, device: device, dtype: dtype = torch.float32, *, use_exact: bool = True) -> Tuple[torch.Tensor, torch.Tensor]

Return (centroids, decision_boundaries) as GPU tensors.

centroids : (2^bits,) centroid values decision_boundaries: (2^bits-1,) interior boundaries for searchsorted

The decision boundaries are the interior ones, i.e., boundaries[1:-1], which is what torch.searchsorted needs to map each coordinate to an index.

Source code in turboquant/core/codebook.py

def get_codebook_tensors(
    d: int,
    bits: int,
    device: torch.device,
    dtype: torch.dtype = torch.float32,
    *,
    use_exact: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Return (centroids, decision_boundaries) as GPU tensors.

    centroids          : (2^bits,) centroid values
    decision_boundaries: (2^bits-1,) interior boundaries for searchsorted

    The decision boundaries are the interior ones, i.e., boundaries[1:-1],
    which is what torch.searchsorted needs to map each coordinate to an index.
    """
    cb = get_codebook(d, bits, use_exact=use_exact)
    centroids = torch.tensor(cb["centroids"], device=device, dtype=dtype)
    all_bnd = torch.tensor(cb["boundaries"], device=device, dtype=dtype)
    decision_bnd = all_bnd[1:-1].contiguous()
    return centroids, decision_bnd

Rotation helpers for TurboQuant core codecs.

RotationState dataclass

RotationState(mode: RotationMode, d_orig: int, seed: int, matrix: Optional[Tensor] = None, signs1: Optional[Tensor] = None, signs2: Optional[Tensor] = None, d_pad: int = 0)

Parameters needed to apply and invert a layer rotation.

build_rotation

build_rotation(d: int, mode: RotationMode, device: device, dtype: dtype = torch.float32, seed: int = 42) -> RotationState

Create a rotation state for a layer/head group.

Source code in turboquant/core/rotation.py

def build_rotation(
    d: int,
    mode: RotationMode,
    device: torch.device,
    dtype: torch.dtype = torch.float32,
    seed: int = 42,
) -> RotationState:
    """Create a rotation state for a layer/head group."""
    if mode == RotationMode.QR:
        mat = _generate_qr_matrix(d, device, dtype, seed)
        return RotationState(mode=mode, d_orig=d, seed=seed, matrix=mat, d_pad=d)
    else:  # RHT
        s1, s2, d_pad = _generate_rht_signs(d, device, dtype, seed)
        return RotationState(mode=mode, d_orig=d, seed=seed, signs1=s1, signs2=s2, d_pad=d_pad)

derive_transform_seed

derive_transform_seed(base_seed: int, layer_idx: int, head_idx: int = 0) -> int

Derive a deterministic transform seed for layer/head metadata.

Source code in turboquant/core/rotation.py

def derive_transform_seed(base_seed: int, layer_idx: int, head_idx: int = 0) -> int:
    """Derive a deterministic transform seed for layer/head metadata."""
    return int(base_seed + 1_000_003 * layer_idx + 97_003 * head_idx)

rotation_from_spec

rotation_from_spec(spec: TransformSpec, device: device, dtype: dtype = torch.float32) -> RotationState

Rebuild a rotation state from serialized metadata.

Source code in turboquant/core/rotation.py

def rotation_from_spec(
    spec: TransformSpec,
    device: torch.device,
    dtype: torch.dtype = torch.float32,
) -> RotationState:
    """Rebuild a rotation state from serialized metadata."""
    if spec.kind == "qr_rotation":
        mode = RotationMode.QR
    elif spec.kind == "rht":
        mode = RotationMode.RHT
    else:
        raise ValueError(f"Unsupported rotation spec kind: {spec.kind}")
    return build_rotation(spec.dim, mode, device=device, dtype=dtype, seed=spec.seed)

rotate_forward

rotate_forward(x: Tensor, state: RotationState) -> torch.Tensor

Apply the forward rotation to row-vector inputs.

Source code in turboquant/core/rotation.py

def rotate_forward(x: torch.Tensor, state: RotationState) -> torch.Tensor:
    """Apply the forward rotation to row-vector inputs."""
    x_f = x.float()

    if state.mode == RotationMode.QR:
        return torch.matmul(x_f, state.matrix.T)

    d_orig = state.d_orig
    d_pad  = state.d_pad
    if d_pad > d_orig:
        pad = d_pad - d_orig
        x_f = F.pad(x_f, (0, pad))
    x_f = x_f * state.signs1
    x_f = _walsh_hadamard_transform(x_f)
    x_f = x_f * state.signs2
    return x_f[..., :d_orig]

rotate_backward

rotate_backward(y: Tensor, state: RotationState) -> torch.Tensor

Apply the inverse rotation to row-vector inputs.

Source code in turboquant/core/rotation.py

def rotate_backward(y: torch.Tensor, state: RotationState) -> torch.Tensor:
    """Apply the inverse rotation to row-vector inputs."""
    y_f = y.float()

    if state.mode == RotationMode.QR:
        return torch.matmul(y_f, state.matrix)

    d_orig = state.d_orig
    d_pad  = state.d_pad
    if d_pad > d_orig:
        pad = d_pad - d_orig
        y_f = F.pad(y_f, (0, pad))
    y_f = y_f * state.signs2
    y_f = _walsh_hadamard_transform(y_f)
    y_f = y_f * state.signs1
    return y_f[..., :d_orig]

build_qjl_matrix

build_qjl_matrix(d: int, device: device, dtype: dtype = torch.float32, seed: int = 12345) -> torch.Tensor

Generate a QJL sketch matrix with i.i.d. standard normal entries.

Source code in turboquant/core/rotation.py

def build_qjl_matrix(
    d: int,
    device: torch.device,
    dtype: torch.dtype = torch.float32,
    seed: int = 12345,
) -> torch.Tensor:
    """Generate a QJL sketch matrix with i.i.d. standard normal entries."""
    gen = torch.Generator(device="cpu")
    gen.manual_seed(seed)
    S = torch.randn(d, d, generator=gen, dtype=torch.float32)
    return S.to(device=device, dtype=dtype)

TurboQuant MSE quantizer, Algorithm 1 from the paper.

TorbuquantMSE

TorbuquantMSE(dim: int, bits: int, rotation: RotationState, device: device, *, use_exact: bool = True, norm_correction: bool = False)

Bases: Module

Algorithm 1 TurboQuant quantizer for reconstruction MSE.

Source code in turboquant/core/mse.py

def __init__(
    self,
    dim: int,
    bits: int,
    rotation: RotationState,
    device: torch.device,
    *,
    use_exact: bool = True,
    norm_correction: bool = False,
):
    super().__init__()
    self.dim = dim
    self.bits = bits
    self.rotation = rotation
    self.use_exact = use_exact
    self.norm_correction = norm_correction

    centroids, decision_bnd = get_codebook_tensors(
        dim, bits, device, dtype=torch.float32, use_exact=use_exact
    )
    # Register as buffers so state_dict / .to() work correctly
    self.register_buffer("centroids", centroids)          # (2^bits,)
    self.register_buffer("decision_bnd", decision_bnd)   # (2^bits - 1,)

quantize

quantize(x: Tensor) -> MSEData

Quantize a batch of vectors.

Source code in turboquant/core/mse.py

def quantize(self, x: torch.Tensor) -> MSEData:
    """Quantize a batch of vectors."""
    norms = x.norm(dim=-1)

    x_unit = x / (norms.unsqueeze(-1).clamp(min=1e-12))

    y = rotate_forward(x_unit, self.rotation)

    y_c = y.contiguous()
    indices = torch.searchsorted(
        self.decision_bnd, y_c.reshape(-1, self.dim)
    ).reshape(y_c.shape)

    packed = pack_indices(indices, self.bits)

    return MSEData(
        indices=packed,
        norms=norms.half(),
        bits=self.bits,
        dim=self.dim,
    )

dequantize

dequantize(q: MSEData) -> torch.Tensor

Reconstruct vectors from MSEData.

Source code in turboquant/core/mse.py

def dequantize(self, q: MSEData) -> torch.Tensor:
    """Reconstruct vectors from MSEData."""
    indices = unpack_indices(q.indices, q.bits, q.dim)

    y_hat = self.centroids[indices]

    x_hat = rotate_backward(y_hat, self.rotation)

    if self.norm_correction:
        x_hat = x_hat / x_hat.norm(dim=-1, keepdim=True).clamp(min=1e-12)
    x_hat = x_hat * q.norms.float().unsqueeze(-1)

    return x_hat

rotate_query

rotate_query(query: Tensor) -> torch.Tensor

Apply the quantizer rotation to query vectors.

Source code in turboquant/core/mse.py

def rotate_query(self, query: torch.Tensor) -> torch.Tensor:
    """Apply the quantizer rotation to query vectors."""
    if query.shape[-1] != self.dim:
        raise ValueError(f"query dim mismatch: expected {self.dim}, got {query.shape[-1]}")
    return rotate_forward(query, self.rotation)

score_rotated

score_rotated(query_rotated: Tensor, q: MSEData, *, scale: float = 1.0) -> torch.Tensor

Compute inner-product scores from rotated queries and packed indices.

query_rotated may be shaped (..., dim) or (..., query_tokens, dim). The packed data must be shaped (..., kv_tokens, packed_dim) with matching leading dimensions. The result has shape (..., kv_tokens) or (..., query_tokens, kv_tokens).

Source code in turboquant/core/mse.py

def score_rotated(self, query_rotated: torch.Tensor, q: MSEData, *, scale: float = 1.0) -> torch.Tensor:
    """Compute inner-product scores from rotated queries and packed indices.

    ``query_rotated`` may be shaped ``(..., dim)`` or ``(..., query_tokens,
    dim)``.  The packed data must be shaped ``(..., kv_tokens, packed_dim)``
    with matching leading dimensions.  The result has shape
    ``(..., kv_tokens)`` or ``(..., query_tokens, kv_tokens)``.
    """
    if q.dim != self.dim:
        raise ValueError(f"quantized dim mismatch: expected {self.dim}, got {q.dim}")
    if q.bits != self.bits:
        raise ValueError(f"quantized bit width mismatch: expected {self.bits}, got {q.bits}")
    if query_rotated.shape[-1] != self.dim:
        raise ValueError(f"query dim mismatch: expected {self.dim}, got {query_rotated.shape[-1]}")

    indices = unpack_indices(q.indices.to(query_rotated.device), q.bits, q.dim)
    values = self.centroids.to(query_rotated.device)[indices.long()].float()
    norms = q.norms.to(query_rotated.device).float()
    if self.norm_correction:
        unit_norm = values.norm(dim=-1).clamp(min=1e-12)
        norms = norms / unit_norm

    query_f = query_rotated.float()
    if query_f.ndim == values.ndim - 1:
        if query_f.shape[:-1] != values.shape[:-2]:
            raise ValueError("query and quantized prefixes do not match")
        scores = (query_f.unsqueeze(-2) * values).sum(dim=-1) * norms
        return scores * scale
    if query_f.ndim == values.ndim:
        if query_f.shape[:-2] != values.shape[:-2]:
            raise ValueError("query and quantized prefixes do not match")
        scores = torch.einsum("...qd,...nd->...qn", query_f, values)
        return scores * norms.unsqueeze(-2) * scale
    raise ValueError("query must be shaped (..., dim) or (..., query_tokens, dim)")

score

score(query: Tensor, q: MSEData, *, scale: float = 1.0) -> torch.Tensor

Compute inner-product scores from unrotated queries and packed indices.

Source code in turboquant/core/mse.py

def score(self, query: torch.Tensor, q: MSEData, *, scale: float = 1.0) -> torch.Tensor:
    """Compute inner-product scores from unrotated queries and packed indices."""
    return self.score_rotated(self.rotate_query(query), q, scale=scale)

pack_indices

pack_indices(indices: Tensor, bits: int) -> torch.Tensor

Bit-pack integer indices into uint8 bytes.

Source code in turboquant/core/mse.py

def pack_indices(indices: torch.Tensor, bits: int) -> torch.Tensor:
    """Bit-pack integer indices into uint8 bytes."""
    _packing_params(bits)
    return pack_unsigned(indices, bits)

unpack_indices

unpack_indices(packed: Tensor, bits: int, d: int) -> torch.Tensor

Unpack uint8 bytes back to integer indices.

Source code in turboquant/core/mse.py

def unpack_indices(packed: torch.Tensor, bits: int, d: int) -> torch.Tensor:
    """Unpack uint8 bytes back to integer indices."""
    _packing_params(bits)
    return unpack_unsigned(packed, bits, d)

turboquant.core.qjl

Quantized Johnson-Lindenstrauss (QJL) transform, with 1-bit residual correction.

Paper Section 2.2 (Definition 1 / Lemma 2): Given residual r = x - x_mse (x in S^{d-1}, x_mse from Algorithm 1), the QJL quantization is:

  z = sign(S * r) in {-1, +1}^d
  ||r||_2 stored as a scalar

Dequantized QJL contribution: x_qjl = sqrt(pi/2) / d * ||r|| * S^T * z

The combined estimator x_mse + x_qjl is unbiased: E[] = for any y in R^d

and has inner-product variance bounded by (pi/2) / d * ||y||^2 * D_mse(b-1).

Storage: Signs are packed 8 per byte (1 bit per coordinate) into uint8. One float32 scalar per vector stores ||r||.

Score computation (used during decode attention): Given pre-sketched query q_sk = q @ S^T: qjl_score = qjl_scale * ||r|| * = qjl_scale * ||r|| * sum_j q_sk[j] * z[j]

where qjl_scale = sqrt(pi/2) / d.

TorbuquantQJL

TorbuquantQJL(dim: int, S: Tensor, device: device)

Bases: Module

QJL residual quantizer (1 bit per coordinate of the residual).

Args: dim : head dimension d. S : (d, d) float32 Gaussian random matrix generated once per layer by build_qjl_matrix() in rotation.py. device : torch device.

Source code in turboquant/core/qjl.py

def __init__(self, dim: int, S: torch.Tensor, device: torch.device):
    super().__init__()
    self.dim = dim
    self.qjl_scale = math.sqrt(math.pi / 2.0) / dim
    self.register_buffer("S", S.to(device=device, dtype=torch.float32))

quantize

quantize(residual: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Quantize residual r to 1-bit signs.

Args: residual : (..., d) float, r = x - x_mse.

Returns: signs_packed : (..., ceil(d/8)) uint8. residual_norms : (...,) float32.

Source code in turboquant/core/qjl.py

def quantize(
    self, residual: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Quantize residual r to 1-bit signs.

    Args:
        residual : (..., d) float, r = x - x_mse.

    Returns:
        signs_packed   : (..., ceil(d/8))  uint8.
        residual_norms : (...,)             float32.
    """
    residual_norms = residual.float().norm(dim=-1)         # (...,)
    projected = torch.matmul(residual.float(), self.S.T)   # (..., d)
    signs_packed = pack_signs(projected)
    return signs_packed, residual_norms

dequantize

dequantize(signs_packed: Tensor, residual_norms: Tensor) -> torch.Tensor

Reconstruct QJL contribution: x_qjl = scale * ||r|| * S^T * z

Returns: x_qjl : (..., d) float32.

Source code in turboquant/core/qjl.py

def dequantize(
    self,
    signs_packed: torch.Tensor,
    residual_norms: torch.Tensor,
) -> torch.Tensor:
    """
    Reconstruct QJL contribution: x_qjl = scale * ||r|| * S^T * z

    Returns:
        x_qjl : (..., d) float32.
    """
    z = unpack_signs(signs_packed, self.dim)                 # (..., d)
    x_qjl = torch.matmul(z, self.S)
    x_qjl = x_qjl * (self.qjl_scale * residual_norms.float().unsqueeze(-1))
    return x_qjl

score

score(q_sketched: Tensor, signs_packed: Tensor, residual_norms: Tensor) -> torch.Tensor

Compute QJL attention score contribution without dequantizing.

Instead of forming x_qjl and then dotting with query, use: qjl_score[n] = scale * ||r[n]|| * where q_sk = q @ S^T is precomputed once per decode step.

Args: q_sketched : (BH, d) float32 pre-sketched query q @ S^T. signs_packed : (BH, N, ceil(d/8)) uint8. residual_norms : (BH, N) float32.

Returns: scores : (BH, N) float32.

Source code in turboquant/core/qjl.py

def score(
    self,
    q_sketched: torch.Tensor,
    signs_packed: torch.Tensor,
    residual_norms: torch.Tensor,
) -> torch.Tensor:
    """
    Compute QJL attention score contribution without dequantizing.

    Instead of forming x_qjl and then dotting with query, use:
        qjl_score[n] = scale * ||r[n]|| * <q_sk, z[n]>
    where q_sk = q @ S^T is precomputed once per decode step.

    Args:
        q_sketched     : (BH, d) float32 pre-sketched query q @ S^T.
        signs_packed   : (BH, N, ceil(d/8))  uint8.
        residual_norms : (BH, N)   float32.

    Returns:
        scores : (BH, N)  float32.
    """
    BH, N, _ = signs_packed.shape
    z = unpack_signs(signs_packed.reshape(BH * N, -1), self.dim)  # (BH*N, d)
    z = z.reshape(BH, N, self.dim)                                # (BH, N, d)
    # <q_sk, z[n]> for each token n: (BH, N)
    dot = torch.bmm(z, q_sketched.unsqueeze(-1)).squeeze(-1)      # (BH, N)
    return self.qjl_scale * residual_norms * dot

sketch_query

sketch_query(query: Tensor) -> torch.Tensor

Pre-sketch query: q_sk = q @ S^T.

Args: query : (BH, d) float query vectors.

Returns: q_sketched : (BH, d) float32.

Source code in turboquant/core/qjl.py

def sketch_query(self, query: torch.Tensor) -> torch.Tensor:
    """
    Pre-sketch query: q_sk = q @ S^T.

    Args:
        query : (BH, d) float query vectors.

    Returns:
        q_sketched : (BH, d)  float32.
    """
    return torch.matmul(query.float(), self.S.T)

pack_signs

pack_signs(projected: Tensor) -> torch.Tensor

Convert real projections to packed sign bits (8 signs per byte).

Args: projected : (..., d) float values from S * r.

Returns: packed : (..., ceil(d/8)) uint8.

Source code in turboquant/core/qjl.py

def pack_signs(projected: torch.Tensor) -> torch.Tensor:
    """
    Convert real projections to packed sign bits (8 signs per byte).

    Args:
        projected : (..., d) float values from S * r.

    Returns:
        packed : (..., ceil(d/8))  uint8.
    """
    signs = (projected > 0).to(torch.uint8)   # 1 = positive, 0 = negative
    *batch, d = signs.shape
    pad = (-d) % 8
    if pad:
        signs = F.pad(signs, (0, pad), value=0)
    s = signs.reshape(*batch, -1, 8)
    powers = torch.tensor([1, 2, 4, 8, 16, 32, 64, 128],
                          dtype=torch.uint8, device=projected.device)
    return (s * powers).sum(dim=-1, dtype=torch.uint8)  # (..., d//8)

unpack_signs

unpack_signs(packed: Tensor, d: int) -> torch.Tensor

Unpack sign bits back to float {-1, +1}.

Args: packed : (..., ceil(d/8)) uint8. d : original number of coordinates.

Returns: signs : (..., d) float32 with values in {-1, +1}.

Source code in turboquant/core/qjl.py

def unpack_signs(packed: torch.Tensor, d: int) -> torch.Tensor:
    """
    Unpack sign bits back to float {-1, +1}.

    Args:
        packed : (..., ceil(d/8))  uint8.
        d      : original number of coordinates.

    Returns:
        signs : (..., d)  float32 with values in {-1, +1}.
    """
    *batch, _ = packed.shape
    powers = torch.tensor([1, 2, 4, 8, 16, 32, 64, 128],
                          dtype=torch.uint8, device=packed.device)
    bits = ((packed.unsqueeze(-1) & powers) > 0).reshape(*batch, -1)[..., :d]
    return bits.float() * 2.0 - 1.0

TurboQuant MSE plus QJL residual codec, Algorithm 2 from the paper.

TorbuquantProd

TorbuquantProd(dim: int, bits: int, rotation: RotationState, device: device, *, qjl_seed: int = 12345, use_exact: bool = True, qjl_for_attention: bool = False)

Bases: Module

Algorithm 2: unbiased inner-product TurboQuant quantizer.

Args: dim : head dimension d. bits : total bits per coordinate (must be >= 2). MSE uses (bits-1), QJL uses 1. rotation : RotationState for the MSE stage (shared with layer). device : torch device. qjl_seed : RNG seed for the QJL Gaussian matrix S. use_exact : use exact Beta codebook for MSE stage.

Source code in turboquant/core/polar.py

def __init__(
    self,
    dim: int,
    bits: int,
    rotation: RotationState,
    device: torch.device,
    *,
    qjl_seed: int = 12345,
    use_exact: bool = True,
    qjl_for_attention: bool = False,
):
    super().__init__()
    if bits < 2:
        raise ValueError("TorbuquantProd requires bits >= 2")
    self.dim = dim
    self.bits = bits
    self.device = device
    self.qjl_for_attention = qjl_for_attention

    # Stage 1: MSE at (bits-1) bits
    self.mse = TorbuquantMSE(
        dim=dim,
        bits=bits - 1,
        rotation=rotation,
        device=device,
        use_exact=use_exact,
    )

    # Stage 2: QJL on the residual
    S = build_qjl_matrix(dim, device=device, dtype=torch.float32, seed=qjl_seed)
    self.qjl = TorbuquantQJL(dim=dim, S=S, device=device)

quantize

quantize(x: Tensor) -> ProdData

Two-stage quantization.

Source code in turboquant/core/polar.py

def quantize(self, x: torch.Tensor) -> ProdData:
    """Two-stage quantization."""
    mse_q = self.mse.quantize(x)
    x_hat = self.mse.dequantize(mse_q)

    residual = x.float() - x_hat

    signs_packed, res_norms = self.qjl.quantize(residual)

    return ProdData(
        mse_indices=mse_q.indices,
        qjl_signs=signs_packed,
        residual_norms=res_norms.float(),
        norms=mse_q.norms,
        mse_bits=mse_q.bits,
        dim=self.dim,
    )

dequantize

dequantize(q: ProdData) -> torch.Tensor

Reconstruct with MSE plus QJL contribution.

Source code in turboquant/core/polar.py

def dequantize(self, q: ProdData) -> torch.Tensor:
    """Reconstruct with MSE plus QJL contribution."""
    mse_data = MSEData(
        indices=q.mse_indices,
        norms=q.norms,
        bits=q.mse_bits,
        dim=q.dim,
    )
    x_mse = self.mse.dequantize(mse_data)
    x_qjl = self.qjl.dequantize(q.qjl_signs, q.residual_norms)
    return x_mse + x_qjl

attention_score

attention_score(query: Tensor, q: ProdData, *, include_qjl: bool | None = None) -> torch.Tensor

Compute raw scores from compressed keys.

Source code in turboquant/core/polar.py

def attention_score(
    self,
    query: torch.Tensor,
    q: ProdData,
    *,
    include_qjl: bool | None = None,
) -> torch.Tensor:
    """Compute raw scores from compressed keys."""
    if query.dim() == 3:
        query = query.squeeze(1)

    BH, d = query.shape
    if include_qjl is None:
        include_qjl = self.qjl_for_attention

    q_rot = rotate_forward(query, self.mse.rotation)

    from turboquant.core.mse import unpack_indices
    mse_idx = unpack_indices(q.mse_indices, q.mse_bits, d)

    centroids = self.mse.centroids
    c_vals = centroids[mse_idx]

    mse_scores = (c_vals * q_rot.unsqueeze(1)).sum(-1)
    mse_scores = mse_scores * q.norms.float()

    if not include_qjl:
        return mse_scores

    q_sk = self.qjl.sketch_query(query)
    qjl_scores = self.qjl.score(q_sk, q.qjl_signs, q.residual_norms)

    return mse_scores + qjl_scores

forward

forward(x: Tensor) -> torch.Tensor

Quantize and dequantize.

Source code in turboquant/core/polar.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Quantize and dequantize."""
    return self.dequantize(self.quantize(x))

Channel-split TurboQuant codecs for fractional bit targets.

TorbuquantChannelMSE

TorbuquantChannelMSE(*, dim: int, target_bits: float, device: device, seed: int = 42, calibration_scores: Tensor | None = None, rotation_mode: RotationMode = 'rht', use_exact: bool = True, norm_correction: bool = True)

Bases: Module

MSE codec that assigns one extra bit to selected channels.

Source code in turboquant/core/outlier.py

def __init__(
    self,
    *,
    dim: int,
    target_bits: float,
    device: torch.device,
    seed: int = 42,
    calibration_scores: torch.Tensor | None = None,
    rotation_mode: RotationMode = "rht",
    use_exact: bool = True,
    norm_correction: bool = True,
):
    super().__init__()
    high_idx, low_idx, high_bits, low_bits = split_channel_indices(
        dim,
        target_bits,
        calibration_scores=calibration_scores,
        device=device,
    )
    self.dim = dim
    self.target_bits = float(target_bits)
    self.high_bits = high_bits
    self.low_bits = low_bits
    self.register_buffer("high_index", high_idx)
    self.register_buffer("low_index", low_idx)

    self.high_codec: TorbuquantMSE | None = None
    self.low_codec: TorbuquantMSE | None = None
    if int(high_idx.numel()) > 0:
        high_dim = int(high_idx.numel())
        high_rotation = build_rotation(
            high_dim,
            mode=rotation_mode,
            device=device,
            dtype=torch.float32,
            seed=derive_transform_seed(seed, 10_001, high_dim),
        )
        self.high_codec = TorbuquantMSE(
            dim=high_dim,
            bits=high_bits,
            rotation=high_rotation,
            device=device,
            use_exact=use_exact,
            norm_correction=norm_correction,
        )
    if int(low_idx.numel()) > 0:
        low_dim = int(low_idx.numel())
        low_rotation = build_rotation(
            low_dim,
            mode=rotation_mode,
            device=device,
            dtype=torch.float32,
            seed=derive_transform_seed(seed, 10_002, low_dim),
        )
        self.low_codec = TorbuquantMSE(
            dim=low_dim,
            bits=low_bits,
            rotation=low_rotation,
            device=device,
            use_exact=use_exact,
            norm_correction=norm_correction,
        )

storage_bits_per_vector

storage_bits_per_vector() -> int

Return packed index and norm bits per vector for this codec.

Source code in turboquant/core/outlier.py

def storage_bits_per_vector(self) -> int:
    """Return packed index and norm bits per vector for this codec."""
    index_bits = int(self.high_index.numel()) * self.high_bits + int(self.low_index.numel()) * self.low_bits
    norm_count = int(self.high_index.numel() > 0) + int(self.low_index.numel() > 0)
    return index_bits + 16 * norm_count

split_channel_indices

split_channel_indices(dim: int, target_bits: float, *, calibration_scores: Tensor | None = None, device: device | None = None) -> tuple[torch.Tensor, torch.Tensor, int, int]

Return high-bit and low-bit channel indices for a fractional bit target.

Source code in turboquant/core/outlier.py

def split_channel_indices(
    dim: int,
    target_bits: float,
    *,
    calibration_scores: torch.Tensor | None = None,
    device: torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor, int, int]:
    """Return high-bit and low-bit channel indices for a fractional bit target."""
    if dim <= 0:
        raise ValueError("dim must be positive")
    if target_bits < 1.0 or target_bits > 8.0:
        raise ValueError("target_bits must be in [1, 8]")

    low_bits = int(math.floor(target_bits))
    high_bits = int(math.ceil(target_bits))
    frac = target_bits - low_bits
    n_high = int(round(dim * frac))
    if low_bits == high_bits:
        n_high = 0
    if high_bits > 8:
        raise ValueError("high-bit channel width exceeds 8")

    dev = device if device is not None else (
        calibration_scores.device if calibration_scores is not None else torch.device("cpu")
    )
    all_idx = torch.arange(dim, device=dev, dtype=torch.long)
    if n_high == 0:
        return all_idx[:0], all_idx, high_bits, low_bits
    if n_high == dim:
        return all_idx, all_idx[:0], high_bits, low_bits

    if calibration_scores is not None:
        if calibration_scores.shape != (dim,):
            raise ValueError(f"calibration_scores must have shape ({dim},)")
        high_idx = torch.topk(calibration_scores.to(dev).float(), k=n_high, largest=True).indices.sort().values
    else:
        high_idx = all_idx[:n_high]

    mask = torch.ones(dim, device=dev, dtype=torch.bool)
    mask[high_idx] = False
    low_idx = all_idx[mask]
    return high_idx, low_idx, high_bits, low_bits

channel_energy_scores

channel_energy_scores(x: Tensor) -> torch.Tensor

Compute mean squared magnitude per channel from calibration tensors.

Source code in turboquant/core/outlier.py

def channel_energy_scores(x: torch.Tensor) -> torch.Tensor:
    """Compute mean squared magnitude per channel from calibration tensors."""
    if x.ndim < 2:
        raise ValueError("calibration tensor must include a channel dimension")
    return x.float().reshape(-1, x.shape[-1]).pow(2).mean(dim=0)

Shared data structures for quantized tensors and codec metadata.

TransformSpec dataclass

TransformSpec(kind: Literal['qr_rotation', 'rht', 'qjl'], dim: int, seed: int, dtype: str = 'float32', device_type: str = 'cpu', pad_dim: int | None = None)

Serializable rotation or sketch transform metadata.

CodebookSpec dataclass

CodebookSpec(dim: int, bits: int, distribution: Literal['beta_sphere', 'gaussian_approx'] = 'beta_sphere', cache_key: str | None = None)

Serializable Lloyd-Max codebook metadata.

QuantizedTensor

Bases: NamedTuple

Common tensor payload plus byte-level format metadata.

QuantizedKeys

Bases: NamedTuple

Key payload with transform and codebook specs.

QuantizedValues

Bases: NamedTuple

Value payload with scale/zero metadata.

MSEData

Bases: NamedTuple

Output of TurboQuantMSE.quantize().

indices : (BH, N, Pk) uint8 norms : (BH, N) fp16 bits : int dim : int

ProdData

Bases: NamedTuple

Output of TurboQuantProd.quantize().

mse_indices : (BH, N, Pk) uint8 qjl_signs : (BH, N, Ps) uint8 residual_norms : (BH, N) fp32 norms : (BH, N) fp16 mse_bits : int dim : int

ChannelSplitData

Bases: NamedTuple

Output of channel-split MSE quantization.

high : MSEData for selected high-bit channels, or None low : MSEData for remaining channels, or None high_index : selected high-bit channel indices low_index : remaining channel indices dim : original vector dimension target_bits : average bit target before norm metadata

ValueData

Bases: NamedTuple

Output of quantize_values().

data : (BH, N, Pv) uint8 scales : (BH, N, G) fp16 zeros : (BH, N, G) fp16 bits : int dim : int group_size : int