| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Lloyd-Max optimal scalar quantizer for TurboQuant. |
| 4 | |
| 5 | After rotating a d-dimensional unit vector by a random orthogonal matrix, |
| 6 | each coordinate approximately follows N(0, 1/d) for d >= 64. |
| 7 | We solve the Lloyd-Max conditions to find optimal centroids. |
| 8 | |
| 9 | Based on: turboquant-pytorch/lloyd_max.py (Zandieh et al.) |
| 10 | """ |
| 11 | |
| 12 | import math |
| 13 | from functools import lru_cache |
| 14 | |
| 15 | import torch |
| 16 | |
| 17 | |
| 18 | def _gaussian_pdf(x: float, sigma2: float) -> float: |
| 19 | return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(-x * x / (2 * sigma2)) |
| 20 | |
| 21 | |
| 22 | def _trapz(f, a: float, b: float, n: int = 200) -> float: |
| 23 | """Trapezoidal numerical integration (replaces scipy.integrate.quad).""" |
| 24 | h = (b - a) / n |
| 25 | result = 0.5 * (f(a) + f(b)) |
| 26 | for i in range(1, n): |
| 27 | result += f(a + i * h) |
| 28 | return result * h |
| 29 | |
| 30 | |
| 31 | def solve_lloyd_max( |
| 32 | d: int, |
| 33 | bits: int, |
| 34 | max_iter: int = 200, |
| 35 | tol: float = 1e-10, |
| 36 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 37 | """Solve Lloyd-Max optimal quantizer for N(0, 1/d) distribution. |
| 38 | |
| 39 | Args: |
| 40 | d: Vector dimension (determines variance = 1/d). |
| 41 | bits: Number of quantization bits. |
| 42 | max_iter: Maximum Lloyd-Max iterations. |
| 43 | tol: Convergence tolerance. |
| 44 | |
| 45 | Returns: |
| 46 | centroids: Sorted tensor of 2^bits optimal centroids. |
| 47 | boundaries: Sorted tensor of 2^bits - 1 decision boundaries. |
| 48 | |
| 49 | """ |
| 50 | n_levels = 2**bits |
| 51 | sigma2 = 1.0 / d |
| 52 | sigma = math.sqrt(sigma2) |
| 53 | |
| 54 | def pdf(x): |
| 55 | return _gaussian_pdf(x, sigma2) |
| 56 | |
| 57 | lo, hi = -3.5 * sigma, 3.5 * sigma |
| 58 | centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)] |
| 59 | |
| 60 | for _ in range(max_iter): |
| 61 | boundaries = [ |
| 62 | (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1) |
| 63 | ] |
| 64 | edges = [lo * 3] + boundaries + [hi * 3] |
| 65 | new_centroids = [] |
| 66 | for i in range(n_levels): |
| 67 | a, b = edges[i], edges[i + 1] |
| 68 | num = _trapz(lambda x: x * pdf(x), a, b) |
| 69 | den = _trapz(pdf, a, b) |
| 70 | new_centroids.append(num / den if den > 1e-15 else centroids[i]) |
| 71 | |
| 72 | if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol: |
| 73 | break |
| 74 | centroids = new_centroids |
| 75 | |
| 76 | boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)] |
| 77 | return ( |
| 78 | torch.tensor(centroids, dtype=torch.float32), |
| 79 | torch.tensor(boundaries, dtype=torch.float32), |
| 80 | ) |
| 81 | |
| 82 | |
| 83 | @lru_cache(maxsize=32) |
| 84 | def get_centroids(d: int, bits: int) -> torch.Tensor: |
| 85 | """Get precomputed Lloyd-Max centroids (cached).""" |
| 86 | centroids, _ = solve_lloyd_max(d, bits) |
| 87 | return centroids |