File size: 8,180 Bytes
4a28d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 | """Tests for TurboQuant (quant_policy=QuantPolicy.TURBO_QUANT).
This module contains kernel-level tests for TurboQuant MSE quantization,
which is used by quant_policy=QuantPolicy.TURBO_QUANT (K=QJL4, V=2bit mixed precision).
TurboQuant is a quantization method that:
- Uses Lloyd-Max algorithm for optimal quantization
- Applies random rotation for better distribution
- Stores only L2 norms (not scales/zeros) for dequantization
"""
import math
import pytest
import torch
# Also import turbo_quant kernels for direct access when needed
from lmdeploy.pytorch.kernels.cuda.turbo_quant import (
get_hadamard_matrix,
get_lloyd_max_codebook,
)
# Import shared TurboQuant utilities to avoid duplication
from .turboquant_utils import (
_unpack_indices,
_unpack_qjl4_nibbles,
dequantize_turboquant_mse,
dequantize_turboquant_qjl4,
quant_turboquant_mse,
quant_turboquant_qjl4,
)
class TestTurboQuantMSE:
"""Verify TurboQuant MSE quantization-dequantization correctness.
These tests verify the core TurboQuant MSE algorithm used by quant_policy=QuantPolicy.TURBO_QUANT.
"""
@pytest.fixture
def head_dim(self):
yield 128
@pytest.fixture
def n_vectors(self):
yield 100
@pytest.mark.parametrize('nbits', [2])
def test_quant_dequant_roundtrip(self, head_dim, n_vectors, nbits):
"""Test quantization-dequantization roundtrip."""
torch.manual_seed(42)
x = torch.randn(n_vectors, head_dim).cuda()
# Quantize
q_x, norms = quant_turboquant_mse(x, nbits)
# Verify norms shape is correct
assert norms.shape == (n_vectors,), f'norms shape incorrect: {norms.shape}'
# Verify quantized values are in valid range
max_val = 2 ** nbits - 1
# Unpack and verify
unpacked = _unpack_indices(q_x, nbits, head_dim)
assert unpacked.max().item() <= max_val, 'quantized value exceeds range'
assert unpacked.min().item() >= 0, 'quantized value less than 0'
print(f' bits={nbits}: quant OK, norms range=[{norms.min():.3f}, {norms.max():.3f}]')
@pytest.mark.parametrize('nbits', [2])
def test_mse_within_theoretical_bound(self, head_dim, n_vectors, nbits):
"""Verify quantization-dequantization MSE is within theoretical bound
(for unit vectors)."""
torch.manual_seed(42)
x = torch.randn(n_vectors, head_dim).cuda()
# Normalize to unit sphere (theoretical bound is for unit vectors)
x = x / torch.norm(x, dim=-1, keepdim=True)
# Quantize
q_x, norms = quant_turboquant_mse(x, nbits)
# Dequantize
x_reconstructed = dequantize_turboquant_mse(q_x, norms, nbits)
# Compute MSE
mse = ((x - x_reconstructed) ** 2).mean().item()
# Theoretical bound: D_mse <= sqrt(3)*pi/2 * (1/4^bits)
theoretical_bound = math.sqrt(3) * math.pi / 2 * (1 / (4 ** nbits))
ratio = mse / theoretical_bound
print(f' bits={nbits}: MSE={mse:.6f}, theory_bound={theoretical_bound:.6f}, ratio={ratio:.3f}')
# Theoretical bound is an upper bound, actual MSE must be less
assert ratio < 1, f'MSE {mse} exceeds theoretical bound {theoretical_bound} (ratio={ratio:.3f})'
@pytest.mark.parametrize('nbits', [2])
def test_reconstruction_quality(self, head_dim, n_vectors, nbits):
"""Verify reconstruction quality (using cosine similarity for unit
vectors).
For unit vectors, cosine similarity better reflects the effect of quantization on direction.
"""
torch.manual_seed(42)
x = torch.randn(n_vectors, head_dim).cuda()
# Normalize to unit sphere
x = x / torch.norm(x, dim=-1, keepdim=True)
# Quantize
q_x, norms = quant_turboquant_mse(x, nbits)
# Dequantize
x_reconstructed = dequantize_turboquant_mse(q_x, norms, nbits)
# Compute cosine similarity (after normalization)
x_norm = x / (x.norm(dim=-1, keepdim=True) + 1e-10)
recon_norm = x_reconstructed / (x_reconstructed.norm(dim=-1, keepdim=True) + 1e-10)
cos_sim = (x_norm * recon_norm).sum(dim=-1).mean().item()
print(f' bits={nbits}: cos_sim={cos_sim:.4f}')
# Cosine similarity should be close to 1.0
# 2bit: ~0.80
assert cos_sim > 0.79, f'2bit cosine similarity {cos_sim} too low'
def test_determinism(self, head_dim):
"""Verify same input produces same output."""
torch.manual_seed(42)
x = torch.randn(10, head_dim).cuda()
# Two quantizations should produce the same result
q1, n1 = quant_turboquant_mse(x, 2)
q2, n2 = quant_turboquant_mse(x, 2)
torch.testing.assert_close(q1, q2)
torch.testing.assert_close(n1, n2)
# Two dequantizations should produce the same result
r1 = dequantize_turboquant_mse(q1, n1, 2)
r2 = dequantize_turboquant_mse(q2, n2, 2)
torch.testing.assert_close(r1, r2)
print(' determinism: OK')
class TestTurboQuantQJL4:
"""Verify 4bit TurboQuant reference with 3bit MSE + 1bit QJL."""
@pytest.fixture
def head_dim(self):
yield 128
@pytest.fixture
def n_vectors(self):
yield 100
def test_quant_dequant_roundtrip(self, head_dim, n_vectors):
torch.manual_seed(42)
x = torch.randn(n_vectors, head_dim).cuda()
q_x, meta = quant_turboquant_qjl4(x)
assert q_x.shape == (n_vectors, head_dim // 2)
assert meta.shape == (n_vectors, 2)
idx3, bit1 = _unpack_qjl4_nibbles(q_x, head_dim)
assert idx3.min().item() >= 0
assert idx3.max().item() <= 7
assert bit1.min().item() >= 0
assert bit1.max().item() <= 1
print(f' qjl4: mse_norm range=[{meta[:,0].min():.3f}, {meta[:,0].max():.3f}]')
print(f' qjl4: qjl_norm range=[{meta[:,1].min():.3f}, {meta[:,1].max():.3f}]')
def test_reconstruction_quality(self, head_dim, n_vectors):
torch.manual_seed(42)
x = torch.randn(n_vectors, head_dim).cuda()
x = x / torch.norm(x, dim=-1, keepdim=True)
q_x, meta = quant_turboquant_qjl4(x)
x_reconstructed = dequantize_turboquant_qjl4(q_x, meta)
x_norm = x / (x.norm(dim=-1, keepdim=True) + 1e-10)
recon_norm = x_reconstructed / (x_reconstructed.norm(dim=-1, keepdim=True) + 1e-10)
cos_sim = (x_norm * recon_norm).sum(dim=-1).mean().item()
mse = ((x - x_reconstructed)**2).mean().item()
print(f' qjl4: mse={mse:.6f}, cos_sim={cos_sim:.4f}')
# This is a test-side reference construction, so use a moderate threshold first.
assert cos_sim > 0.86, f'QJL4 cosine similarity {cos_sim} too low'
def test_qjl4_not_worse_than_3bit_mse(self, head_dim, n_vectors):
torch.manual_seed(42)
x = torch.randn(n_vectors, head_dim).cuda()
x = x / torch.norm(x, dim=-1, keepdim=True)
# Pure 3bit MSE baseline
Pi = get_hadamard_matrix(head_dim, device=str(x.device))
centroids3, boundaries3 = get_lloyd_max_codebook(head_dim, bits=3, device=str(x.device))
y = torch.matmul(x, Pi.T)
idx3 = torch.searchsorted(boundaries3, y.contiguous()).clamp(0, 7)
y3 = centroids3[idx3]
x3 = torch.matmul(y3, Pi)
mse_3bit = ((x - x3)**2).mean().item()
q_x, meta = quant_turboquant_qjl4(x)
x4 = dequantize_turboquant_qjl4(q_x, meta)
mse_qjl4 = ((x - x4)**2).mean().item()
print(f' 3bit_mse={mse_3bit:.6f}, qjl4={mse_qjl4:.6f}')
assert mse_qjl4 <= mse_3bit * 1.05, 'QJL4 should not be significantly worse than pure 3bit MSE'
def test_determinism(self, head_dim):
torch.manual_seed(42)
x = torch.randn(10, head_dim).cuda()
q1, m1 = quant_turboquant_qjl4(x)
q2, m2 = quant_turboquant_qjl4(x)
torch.testing.assert_close(q1, q2)
torch.testing.assert_close(m1, m2)
r1 = dequantize_turboquant_qjl4(q1, m1)
r2 = dequantize_turboquant_qjl4(q2, m2)
torch.testing.assert_close(r1, r2)
print(' qjl4 determinism: OK')
|