mirror of
https://github.com/TheTom/turboquant_plus.git
synced 2026-10-22 02:53:13 -05:00
dccd96bc71
CRITICAL BUG (found by roast review): PolarQuant codebook was calibrated for unit-norm vectors only. Real KV cache tensors have norms ~10-50, causing every coordinate to saturate to extreme centroids. Fix (per paper page 5): extract L2 norms before quantization, normalize to unit vectors, quantize, then rescale on dequantization. Norms stored as float32 metadata alongside indices. API change: quantize() now returns (indices, norms), dequantize() takes (indices, norms). All downstream code updated (TurboQuant, KVCache, tests). Added test_non_unit_norm_vectors at scales [0.01, 1.0, 10.0, 50.0, 100.0]. Zero vector now correctly reconstructs to zero (norm=0). 128 tests, 98% coverage, all green. Co-Authored-By: tturney@psyguard.ai Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
126 lines
4.3 KiB
Python
126 lines
4.3 KiB
Python
"""Tests for PolarQuant (Algorithm 1)."""
|
||
|
||
import numpy as np
|
||
import pytest
|
||
|
||
from turboquant.polar_quant import PolarQuant
|
||
|
||
|
||
class TestPolarQuantRoundTrip:
|
||
"""Quantize → dequantize should produce bounded MSE."""
|
||
|
||
@pytest.mark.parametrize("bit_width", [1, 2, 3])
|
||
@pytest.mark.parametrize("d", [64, 128, 256])
|
||
def test_round_trip_mse_within_bounds(self, bit_width, d):
|
||
"""MSE distortion should be within paper's bounds (Table 2)."""
|
||
expected_mse = {1: 0.36, 2: 0.117, 3: 0.03}
|
||
|
||
pq = PolarQuant(d=d, bit_width=bit_width, seed=42)
|
||
rng = np.random.default_rng(99)
|
||
|
||
n_samples = 500
|
||
mse_total = 0.0
|
||
for _ in range(n_samples):
|
||
x = rng.standard_normal(d)
|
||
x = x / np.linalg.norm(x)
|
||
|
||
indices, norms = pq.quantize(x)
|
||
x_hat = pq.dequantize(indices, norms)
|
||
mse_total += np.mean((x - x_hat) ** 2)
|
||
|
||
avg_mse = mse_total / n_samples
|
||
assert avg_mse < expected_mse[bit_width] * 2.0, (
|
||
f"MSE {avg_mse:.4f} exceeds 2× paper bound {expected_mse[bit_width]} "
|
||
f"at d={d}, b={bit_width}"
|
||
)
|
||
|
||
def test_zero_vector(self):
|
||
"""Zero vector should reconstruct to zero (norm=0 → rescale to 0)."""
|
||
pq = PolarQuant(d=128, bit_width=2, seed=42)
|
||
x = np.zeros(128)
|
||
indices, norms = pq.quantize(x)
|
||
x_hat = pq.dequantize(indices, norms)
|
||
np.testing.assert_allclose(x_hat, 0.0, atol=1e-15)
|
||
|
||
def test_non_unit_norm_vectors(self):
|
||
"""CRITICAL: PolarQuant must work on non-unit-norm vectors (real KV cache).
|
||
|
||
Roast review found this was the biggest bug — codebook is calibrated for
|
||
unit norm, but real KV tensors have norms ~10-50.
|
||
"""
|
||
d = 128
|
||
pq = PolarQuant(d=d, bit_width=3, seed=42)
|
||
rng = np.random.default_rng(42)
|
||
|
||
for scale in [0.01, 1.0, 10.0, 50.0, 100.0]:
|
||
x = rng.standard_normal(d) * scale
|
||
|
||
indices, norms = pq.quantize(x)
|
||
x_hat = pq.dequantize(indices, norms)
|
||
|
||
# Relative MSE should be bounded regardless of scale
|
||
# (same as unit-norm MSE since we extract/rescale norms)
|
||
norm_sq_per_d = np.linalg.norm(x) ** 2 / d
|
||
if norm_sq_per_d > 1e-10:
|
||
relative_mse = np.mean((x - x_hat) ** 2) / norm_sq_per_d
|
||
assert relative_mse < 0.5, (
|
||
f"Relative MSE {relative_mse:.4f} too high at scale={scale}"
|
||
)
|
||
|
||
def test_deterministic(self):
|
||
"""Same seed, same input → same output."""
|
||
d = 128
|
||
x = np.random.default_rng(1).standard_normal(d)
|
||
|
||
pq1 = PolarQuant(d=d, bit_width=2, seed=42)
|
||
pq2 = PolarQuant(d=d, bit_width=2, seed=42)
|
||
|
||
idx1, n1 = pq1.quantize(x)
|
||
idx2, n2 = pq2.quantize(x)
|
||
np.testing.assert_array_equal(idx1, idx2)
|
||
assert n1 == n2
|
||
|
||
def test_batch_matches_single(self):
|
||
"""Batch quantization should match single-vector quantization."""
|
||
d = 128
|
||
pq = PolarQuant(d=d, bit_width=2, seed=42)
|
||
rng = np.random.default_rng(7)
|
||
|
||
X = rng.standard_normal((10, d))
|
||
|
||
batch_indices, batch_norms = pq.quantize(X)
|
||
batch_recon = pq.dequantize(batch_indices, batch_norms)
|
||
|
||
for i in range(10):
|
||
single_idx, single_norm = pq.quantize(X[i])
|
||
single_recon = pq.dequantize(single_idx, single_norm)
|
||
np.testing.assert_array_equal(batch_indices[i], single_idx)
|
||
np.testing.assert_allclose(batch_recon[i], single_recon, atol=1e-12)
|
||
|
||
def test_indices_in_range(self):
|
||
"""All indices should be in [0, 2^bit_width)."""
|
||
d = 256
|
||
pq = PolarQuant(d=d, bit_width=3, seed=42)
|
||
rng = np.random.default_rng(5)
|
||
x = rng.standard_normal(d)
|
||
|
||
indices, _ = pq.quantize(x)
|
||
assert indices.min() >= 0
|
||
assert indices.max() < (1 << 3)
|
||
|
||
|
||
class TestPolarQuantResidual:
|
||
"""Test quantize_and_residual method."""
|
||
|
||
def test_residual_identity(self):
|
||
"""residual = x - dequantize(quantize(x))."""
|
||
d = 128
|
||
pq = PolarQuant(d=d, bit_width=2, seed=42)
|
||
rng = np.random.default_rng(3)
|
||
x = rng.standard_normal(d)
|
||
|
||
indices, norms, residual = pq.quantize_and_residual(x)
|
||
x_hat = pq.dequantize(indices, norms)
|
||
|
||
np.testing.assert_allclose(residual, x - x_hat, atol=1e-12)
|