From 86bcbbe59ace150aea9aa62ccc5bf21250a26de7 Mon Sep 17 00:00:00 2001 From: TheTom Date: Mon, 30 Mar 2026 10:57:38 -0500 Subject: [PATCH] Add norm correction to Python prototype quantization path PolarQuant.dequantize() now re-normalizes reconstructed rotated coordinates to unit norm before inverse rotation, matching the production llama.cpp grp_norm / recon_norm correction. On Qwen2.5-3B, this drops 4-bit PPL from 19.39 to 8.47 (near-lossless) and 3-bit from 1369 to 70.6. No regression on TinyLlama or Mistral. Includes benchmark scripts for TQ vs IsoQuant/RotorQuant comparison and norm correction A/B testing. Co-Authored-By: tturney@psyguard.ai Co-Authored-By: Claude Opus 4.6 (1M context) --- benchmarks/benchmark_norm_correction.py | 288 ++++++++++++++++++++++ benchmarks/benchmark_ppl_tq_vs_rq.py | 313 ++++++++++++++++++++++++ tests/test_polar_quant.py | 20 ++ tests/test_turboquant.py | 18 ++ turboquant/kv_cache.py | 17 +- turboquant/polar_quant.py | 15 +- turboquant/turboquant.py | 12 +- 7 files changed, 674 insertions(+), 9 deletions(-) create mode 100644 benchmarks/benchmark_norm_correction.py create mode 100644 benchmarks/benchmark_ppl_tq_vs_rq.py diff --git a/benchmarks/benchmark_norm_correction.py b/benchmarks/benchmark_norm_correction.py new file mode 100644 index 0000000..409c9e8 --- /dev/null +++ b/benchmarks/benchmark_norm_correction.py @@ -0,0 +1,288 @@ +#!/usr/bin/env python3 +"""Norm correction experiment: does re-normalizing after dequant fix TQ on Qwen2.5? + +Hypothesis: The Python TQ prototype lacks the norm correction that the production +C/Metal implementation has (grp_norm / recon_norm). This causes catastrophic PPL +on Qwen2.5 at 3-bit. IsoQuant accidentally avoids the issue because block-local +rotation produces smaller norm distortion. + +The fix: after dequantizing rotated coordinates y_hat, re-normalize to unit norm +before applying the inverse rotation. Since the original rotated vector y has +||y|| = 1 (orthogonal rotation of a unit vector), this restores the correct magnitude. + +This script tests the shared PolarQuant / TurboQuant implementation with +norm_correction enabled vs disabled on the same model, same data. Compares +against IQ Full as reference. + +Usage: + python3 benchmarks/benchmark_norm_correction.py + PPL_MODEL='TinyLlama/TinyLlama-1.1B-Chat-v1.0' python3 benchmarks/benchmark_norm_correction.py +""" + +import sys +import os +import time +import math +import numpy as np +import torch +from torch import nn + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from turboquant.turboquant import TurboQuantMSE +from turboquant.rotorquant_numpy import IsoQuantMSENp +from transformers import AutoModelForCausalLM, AutoTokenizer + +WIKITEXT_LOCAL = os.environ.get( + "WIKITEXT_PATH", + "/Users/tom/local_llms/llama.cpp/wikitext-2-raw/wiki.test.raw" +) + + +# --------------------------------------------------------------------------- +# Quantize-dequantize helper (same as main benchmark) +# --------------------------------------------------------------------------- + +def quant_dequant_keys(key_states: torch.Tensor, quantizer_factory) -> torch.Tensor: + B, H, S, D = key_states.shape + device = key_states.device + dtype = key_states.dtype + + out = torch.empty_like(key_states) + for b in range(B): + for h in range(H): + k_np = key_states[b, h].float().cpu().numpy() + q = quantizer_factory(D, seed=h) + indices, norms = q.quantize(k_np) + k_hat = q.dequantize(indices, norms) + out[b, h] = torch.from_numpy(k_hat).to(dtype=dtype, device=device) + return out + + +def patch_model_attention(model, quant_factory): + hooks = [] + head_dim = model.config.hidden_size // model.config.num_attention_heads + num_kv_heads = model.config.num_key_value_heads + + for layer in model.model.layers: + attn = layer.self_attn + original_k_proj = attn.k_proj + + class QuantizedKProj(nn.Module): + def __init__(self, orig, qf, hd, nkv): + super().__init__() + self.orig = orig + self.qf = qf + self.hd = hd + self.nkv = nkv + + def forward(self, x): + k = self.orig(x) + B, S, _ = k.shape + k_reshaped = k.view(B, S, self.nkv, self.hd) + k_reshaped = k_reshaped.permute(0, 2, 1, 3) + k_quant = quant_dequant_keys(k_reshaped, self.qf) + k_quant = k_quant.permute(0, 2, 1, 3).reshape(B, S, -1) + return k_quant + + attn.k_proj = QuantizedKProj(original_k_proj, quant_factory, head_dim, num_kv_heads) + hooks.append((attn, original_k_proj)) + + return hooks + + +def unpatch_model(hooks): + for attn, orig in hooks: + attn.k_proj = orig + + +@torch.no_grad() +def evaluate_ppl(model, tokenizer, encodings, device, max_length=1024, stride=512): + seq_len = encodings.size(1) + nlls = [] + n_tokens = 0 + prev_end = 0 + + for begin in range(0, seq_len, stride): + end = min(begin + max_length, seq_len) + input_ids = encodings[:, begin:end].to(device) + + target_len = end - prev_end if begin > 0 else end - begin + target_ids = input_ids.clone() + if begin > 0: + target_ids[:, :-target_len] = -100 + else: + target_ids[:, 0] = -100 + target_len -= 1 + + outputs = model(input_ids, labels=target_ids) + neg_log_likelihood = outputs.loss * target_len + nlls.append(neg_log_likelihood.item()) + n_tokens += target_len + + prev_end = end + if end == seq_len: + break + + return math.exp(sum(nlls) / n_tokens), n_tokens + + +@torch.no_grad() +def measure_k_mse(model, tokenizer, encodings, device, quant_factory, + max_chunks=8, chunk_size=256): + total_mse = 0.0 + total_count = 0 + seq_len = encodings.size(1) + + for i in range(min(max_chunks, seq_len // chunk_size)): + begin = i * chunk_size + end = begin + chunk_size + input_ids = encodings[:, begin:end].to(device) + + outputs = model(input_ids, output_attentions=False, use_cache=True) + past_kv = outputs.past_key_values + + for layer_idx, kv in enumerate(past_kv): + k = kv[0] + k_quant = quant_dequant_keys(k, quant_factory) + mse = ((k.float().cpu() - k_quant.float().cpu()) ** 2).mean().item() + total_mse += mse + total_count += 1 + + return total_mse / total_count if total_count > 0 else float('inf') + + +# --------------------------------------------------------------------------- +# Factories +# --------------------------------------------------------------------------- + +def make_tq_factory(bits): + def factory(d, seed=42): + return TurboQuantMSE(d=d, bit_width=bits, seed=seed, norm_correction=False) + return factory + +def make_tqnc_factory(bits): + def factory(d, seed=42): + return TurboQuantMSE(d=d, bit_width=bits, seed=seed, norm_correction=True) + return factory + +def make_iq_factory(bits): + def factory(d, seed=42): + return IsoQuantMSENp(d=d, bit_width=bits, seed=seed, mode='full') + return factory + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + model_name = os.environ.get("PPL_MODEL", "Qwen/Qwen2.5-3B") + device = "mps" if torch.backends.mps.is_available() else "cpu" + print(f"Device: {device}") + print(f"Model: {model_name}") + print(f"Experiment: Norm correction effect on TQ") + + print("\nLoading model...") + t0 = time.time() + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + model_name, torch_dtype=torch.float16, trust_remote_code=True, + ).to(device).eval() + print(f" Loaded in {time.time() - t0:.1f}s") + + wikitext_path = WIKITEXT_LOCAL + if not os.path.exists(wikitext_path): + # Try M2 path + wikitext_path = os.path.expanduser("~/dev/turbo_test/wikitext-2-raw/wiki.test.raw") + print(f"Loading wikitext-2 from {wikitext_path}...") + with open(wikitext_path, "r") as f: + text = f.read() + encodings = tokenizer(text, return_tensors="pt")["input_ids"] + print(f" {encodings.size(1)} tokens") + + max_tokens = 4096 + if encodings.size(1) > max_tokens: + encodings = encodings[:, :max_tokens] + print(f" Trimmed to {max_tokens} tokens") + + # Focused config: only TQ, TQ+NC, and IQ Full as reference + configs = [ + ("TQ 3-bit", make_tq_factory(3)), + ("TQ+NC 3-bit", make_tqnc_factory(3)), + ("TQ 4-bit", make_tq_factory(4)), + ("TQ+NC 4-bit", make_tqnc_factory(4)), + ("IQ Full 3-bit", make_iq_factory(3)), + ("IQ Full 4-bit", make_iq_factory(4)), + ] + + # Phase 1: MSE + print("\n" + "=" * 70) + print("PHASE 1: K-cache MSE") + print("=" * 70) + + mse_results = {} + for name, factory in configs: + print(f" {name}...", end=" ", flush=True) + t0 = time.time() + mse = measure_k_mse(model, tokenizer, encodings, device, factory) + elapsed = time.time() - t0 + mse_results[name] = mse + print(f"MSE={mse:.6f} ({elapsed:.1f}s)") + + # Phase 2: PPL + print("\n" + "=" * 70) + print("PHASE 2: Perplexity") + print("=" * 70) + + ppl_results = {} + + print(" fp16 baseline...", end=" ", flush=True) + t0 = time.time() + ppl, n_tok = evaluate_ppl(model, tokenizer, encodings, device) + elapsed = time.time() - t0 + ppl_results["fp16"] = ppl + print(f"PPL={ppl:.2f} ({n_tok} tokens, {elapsed:.1f}s)") + + for name, factory in configs: + print(f" {name}...", end=" ", flush=True) + t0 = time.time() + hooks = patch_model_attention(model, factory) + try: + ppl, n_tok = evaluate_ppl(model, tokenizer, encodings, device) + finally: + unpatch_model(hooks) + elapsed = time.time() - t0 + ppl_results[name] = ppl + print(f"PPL={ppl:.2f} ({n_tok} tokens, {elapsed:.1f}s)") + + # Summary + print("\n" + "=" * 70) + print("SUMMARY — Norm Correction Experiment") + print("=" * 70) + + baseline = ppl_results.get("fp16", 0) + all_names = ["fp16"] + [n for n, _ in configs] + + print(f"\n {'Method':<16s} {'PPL':>8s} {'K MSE':>12s} {'vs fp16':>8s}") + print(f" {'─'*16} {'─'*8} {'─'*12} {'─'*8}") + + for name in all_names: + ppl_val = ppl_results.get(name, float('nan')) + mse_val = mse_results.get(name, 0.0) + delta = ppl_val - baseline if name != "fp16" else 0.0 + mse_str = f"{mse_val:.6f}" if mse_val > 0 else "—" + delta_str = f"+{delta:.2f}" if name != "fp16" else "—" + print(f" {name:<16s} {ppl_val:>8.2f} {mse_str:>12s} {delta_str:>8s}") + + # Before/after comparison + print(f"\n BEFORE vs AFTER norm correction:") + for bits in [3, 4]: + old = ppl_results.get(f"TQ {bits}-bit", float('nan')) + new = ppl_results.get(f"TQ+NC {bits}-bit", float('nan')) + iq = ppl_results.get(f"IQ Full {bits}-bit", float('nan')) + print(f" {bits}-bit: TQ {old:.2f} → TQ+NC {new:.2f} (IQ Full: {iq:.2f})") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_ppl_tq_vs_rq.py b/benchmarks/benchmark_ppl_tq_vs_rq.py new file mode 100644 index 0000000..f5d6702 --- /dev/null +++ b/benchmarks/benchmark_ppl_tq_vs_rq.py @@ -0,0 +1,313 @@ +#!/usr/bin/env python3 +"""Head-to-head perplexity comparison: TurboQuant vs RotorQuant on real model K tensors. + +Approach: +1. Load Qwen2.5-0.5B on MPS +2. Load wikitext-2-test +3. Run forward passes, monkey-patching the KV cache to quantize-dequantize K tensors +4. Measure PPL for: fp16 baseline, TQ 3-bit, RQ 3-bit, TQ 4-bit, RQ 4-bit + +Usage: python3 benchmarks/benchmark_ppl_tq_vs_rq.py +""" + +import sys +import os +import time +import math +import numpy as np +import torch +from torch import nn +from functools import partial + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from turboquant.turboquant import TurboQuantMSE +from turboquant.rotorquant_numpy import RotorQuantMSENp, IsoQuantMSENp +from transformers import AutoModelForCausalLM, AutoTokenizer +# datasets import optional — we prefer local wikitext file +WIKITEXT_LOCAL = "/Users/tom/local_llms/llama.cpp/wikitext-2-raw/wiki.test.raw" + + +# --------------------------------------------------------------------------- +# Quantize-dequantize helper +# --------------------------------------------------------------------------- + +def quant_dequant_keys(key_states: torch.Tensor, quantizer_factory) -> torch.Tensor: + """Quantize-dequantize K tensor through a given quantizer. + + key_states: (batch, n_heads, seq_len, head_dim) float16/float32 + quantizer_factory: callable(head_dim, seed) -> quantizer with .quantize/.dequantize + """ + B, H, S, D = key_states.shape + device = key_states.device + dtype = key_states.dtype + + # Process each head independently (quantizers are per-head-dim) + out = torch.empty_like(key_states) + for b in range(B): + for h in range(H): + # (seq_len, head_dim) -> numpy + k_np = key_states[b, h].float().cpu().numpy() # (S, D) + + q = quantizer_factory(D, seed=h) # seed per head for diversity + indices, norms = q.quantize(k_np) + k_hat = q.dequantize(indices, norms) + + out[b, h] = torch.from_numpy(k_hat).to(dtype=dtype, device=device) + return out + + +def make_tq_factory(bits): + def factory(d, seed=42): + return TurboQuantMSE(d=d, bit_width=bits, seed=seed) + return factory + + +def make_rq_factory(bits): + def factory(d, seed=42): + return RotorQuantMSENp(d=d, bit_width=bits, seed=seed) + return factory + + +def make_iq_factory(bits, mode='full'): + def factory(d, seed=42): + return IsoQuantMSENp(d=d, bit_width=bits, seed=seed, mode=mode) + return factory + + +# --------------------------------------------------------------------------- +# Monkey-patch approach: wrap the model's attention to intercept K cache +# --------------------------------------------------------------------------- + +def patch_model_attention(model, quant_factory): + """Patch all attention layers to quantize K before caching. + + Works with Qwen2 architecture (Qwen2Attention / Qwen2SdpaAttention). + """ + hooks = [] + head_dim = model.config.hidden_size // model.config.num_attention_heads + num_kv_heads = model.config.num_key_value_heads + + for layer in model.model.layers: + attn = layer.self_attn + original_k_proj = attn.k_proj + + class QuantizedKProj(nn.Module): + def __init__(self, orig, qf, hd, nkv): + super().__init__() + self.orig = orig + self.qf = qf + self.hd = hd + self.nkv = nkv + + def forward(self, x): + k = self.orig(x) + B, S, _ = k.shape + k_reshaped = k.view(B, S, self.nkv, self.hd) + k_reshaped = k_reshaped.permute(0, 2, 1, 3) # (B, H, S, D) + k_quant = quant_dequant_keys(k_reshaped, self.qf) + k_quant = k_quant.permute(0, 2, 1, 3).reshape(B, S, -1) + return k_quant + + attn.k_proj = QuantizedKProj(original_k_proj, quant_factory, head_dim, num_kv_heads) + hooks.append((attn, original_k_proj)) + + return hooks + + +def unpatch_model(hooks): + for attn, orig in hooks: + attn.k_proj = orig + + +# --------------------------------------------------------------------------- +# Perplexity evaluation (sliding window) +# --------------------------------------------------------------------------- + +@torch.no_grad() +def evaluate_ppl(model, tokenizer, encodings, device, max_length=1024, stride=512): + """Sliding window perplexity on tokenized text.""" + seq_len = encodings.size(1) + nlls = [] + n_tokens = 0 + + prev_end = 0 + for begin in range(0, seq_len, stride): + end = min(begin + max_length, seq_len) + input_ids = encodings[:, begin:end].to(device) + + target_len = end - prev_end if begin > 0 else end - begin + target_ids = input_ids.clone() + # Mask out tokens we've already scored + if begin > 0: + target_ids[:, :-target_len] = -100 + else: + # First window: skip first token (no context) + target_ids[:, 0] = -100 + target_len -= 1 + + outputs = model(input_ids, labels=target_ids) + # Loss is averaged over non-ignored tokens + neg_log_likelihood = outputs.loss * target_len + nlls.append(neg_log_likelihood.item()) + n_tokens += target_len + + prev_end = end + if end == seq_len: + break + + ppl = math.exp(sum(nlls) / n_tokens) + return ppl, n_tokens + + +# --------------------------------------------------------------------------- +# Also measure MSE on real K tensors (simpler, more robust) +# --------------------------------------------------------------------------- + +@torch.no_grad() +def measure_k_mse(model, tokenizer, encodings, device, quant_factory, + max_chunks=10, chunk_size=256): + """Run forward passes and measure MSE of quantized K vs original K.""" + total_mse = 0.0 + total_count = 0 + seq_len = encodings.size(1) + + for i in range(min(max_chunks, seq_len // chunk_size)): + begin = i * chunk_size + end = begin + chunk_size + input_ids = encodings[:, begin:end].to(device) + + outputs = model(input_ids, output_attentions=False, use_cache=True) + past_kv = outputs.past_key_values + + for layer_idx, kv in enumerate(past_kv): + k = kv[0] # (batch, n_heads, seq_len, head_dim) + k_quant = quant_dequant_keys(k, quant_factory) + mse = ((k.float().cpu() - k_quant.float().cpu()) ** 2).mean().item() + total_mse += mse + total_count += 1 + + return total_mse / total_count if total_count > 0 else float('inf') + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + model_name = os.environ.get("PPL_MODEL", "Qwen/Qwen2.5-3B") + device = "mps" if torch.backends.mps.is_available() else "cpu" + print(f"Device: {device}") + print(f"Model: {model_name}") + + # Load model + print("\nLoading model...") + t0 = time.time() + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + model_name, + torch_dtype=torch.float16, + trust_remote_code=True, + ).to(device).eval() + print(f" Loaded in {time.time() - t0:.1f}s") + + # Load wikitext-2 + print(f"Loading wikitext-2 from {WIKITEXT_LOCAL}...") + with open(WIKITEXT_LOCAL, "r") as f: + text = f.read() + encodings = tokenizer(text, return_tensors="pt")["input_ids"] + print(f" {encodings.size(1)} tokens") + + # Trim to manageable size for PPL eval + max_tokens = 4096 # keep it fast + if encodings.size(1) > max_tokens: + encodings = encodings[:, :max_tokens] + print(f" Trimmed to {max_tokens} tokens for speed") + + # ------------------------------------------------------------------ + # Phase 1: MSE on real K tensors (fast, always works) + # ------------------------------------------------------------------ + print("\n" + "=" * 70) + print("PHASE 1: MSE on Real Model K Tensors") + print("=" * 70) + + configs = [ + ("TQ 3-bit", make_tq_factory(3)), + ("RQ 3-bit", make_rq_factory(3)), + ("IQ-F 3-bit", make_iq_factory(3, mode='fast')), + ("IQ 3-bit", make_iq_factory(3, mode='full')), + ("TQ 4-bit", make_tq_factory(4)), + ("RQ 4-bit", make_rq_factory(4)), + ("IQ-F 4-bit", make_iq_factory(4, mode='fast')), + ("IQ 4-bit", make_iq_factory(4, mode='full')), + ] + + mse_results = {} + for name, factory in configs: + print(f" {name}...", end=" ", flush=True) + t0 = time.time() + mse = measure_k_mse(model, tokenizer, encodings, device, factory, + max_chunks=8, chunk_size=256) + elapsed = time.time() - t0 + mse_results[name] = mse + print(f"MSE={mse:.6f} ({elapsed:.1f}s)") + + print(f"\n {'Method':<14s} {'K-cache MSE':>12s}") + print(f" {'─'*14} {'─'*12}") + for name, mse in mse_results.items(): + print(f" {name:<14s} {mse:>12.6f}") + + # ------------------------------------------------------------------ + # Phase 2: Full perplexity with monkey-patched K projection + # ------------------------------------------------------------------ + print("\n" + "=" * 70) + print("PHASE 2: Perplexity (monkey-patched K quantization)") + print("=" * 70) + + ppl_results = {} + + # Baseline + print(" fp16 baseline...", end=" ", flush=True) + t0 = time.time() + ppl, n_tok = evaluate_ppl(model, tokenizer, encodings, device, + max_length=1024, stride=512) + elapsed = time.time() - t0 + ppl_results["fp16"] = ppl + print(f"PPL={ppl:.2f} ({n_tok} tokens, {elapsed:.1f}s)") + + # Quantized variants + for name, factory in configs: + print(f" {name}...", end=" ", flush=True) + t0 = time.time() + hooks = patch_model_attention(model, factory) + try: + ppl, n_tok = evaluate_ppl(model, tokenizer, encodings, device, + max_length=1024, stride=512) + finally: + unpatch_model(hooks) + elapsed = time.time() - t0 + ppl_results[name] = ppl + print(f"PPL={ppl:.2f} ({n_tok} tokens, {elapsed:.1f}s)") + + # ------------------------------------------------------------------ + # Summary table + # ------------------------------------------------------------------ + print("\n" + "=" * 70) + print("SUMMARY") + print("=" * 70) + print(f"\n {'Method':<14s} {'PPL':>8s} {'K MSE':>12s} {'PPL Δ':>8s}") + print(f" {'─'*14} {'─'*8} {'─'*12} {'─'*8}") + + baseline_ppl = ppl_results.get("fp16", 0) + for name in ["fp16", "TQ 3-bit", "RQ 3-bit", "IQ-F 3-bit", "IQ 3-bit", + "TQ 4-bit", "RQ 4-bit", "IQ-F 4-bit", "IQ 4-bit"]: + ppl_val = ppl_results.get(name, float('nan')) + mse_val = mse_results.get(name, 0.0) + delta = ppl_val - baseline_ppl if name != "fp16" else 0.0 + mse_str = f"{mse_val:.6f}" if mse_val > 0 else "—" + delta_str = f"+{delta:.2f}" if name != "fp16" else "—" + print(f" {name:<14s} {ppl_val:>8.2f} {mse_str:>12s} {delta_str:>8s}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_polar_quant.py b/tests/test_polar_quant.py index b69caff..dfc1845 100644 --- a/tests/test_polar_quant.py +++ b/tests/test_polar_quant.py @@ -108,6 +108,26 @@ class TestPolarQuantRoundTrip: assert indices.min() >= 0 assert indices.max() < (1 << 3) + def test_norm_correction_preserves_unit_norm_better(self): + """Norm correction should bring reconstructed unit vectors closer to norm 1.""" + d = 128 + rng = np.random.default_rng(11) + x = rng.standard_normal(d) + x = x / np.linalg.norm(x) + + pq_uncorrected = PolarQuant(d=d, bit_width=3, seed=42, norm_correction=False) + pq_corrected = PolarQuant(d=d, bit_width=3, seed=42, norm_correction=True) + + idx_u, norm_u = pq_uncorrected.quantize(x) + idx_c, norm_c = pq_corrected.quantize(x) + + x_hat_u = pq_uncorrected.dequantize(idx_u, norm_u) + x_hat_c = pq_corrected.dequantize(idx_c, norm_c) + + err_u = abs(np.linalg.norm(x_hat_u) - 1.0) + err_c = abs(np.linalg.norm(x_hat_c) - 1.0) + assert err_c < err_u + class TestPolarQuantResidual: """Test quantize_and_residual method.""" diff --git a/tests/test_turboquant.py b/tests/test_turboquant.py index b49dc0d..2ab098b 100644 --- a/tests/test_turboquant.py +++ b/tests/test_turboquant.py @@ -147,6 +147,24 @@ class TestTurboQuantMSE: # because all 3 bits go to MSE (no QJL stage) assert mse < 0.1, f"MSE-only 3-bit MSE {mse:.4f} too high" + def test_norm_correction_can_be_disabled(self): + d = 64 + x = np.random.default_rng(2).standard_normal(d) + + tq_old = TurboQuantMSE(d=d, bit_width=3, seed=42, norm_correction=False) + tq_new = TurboQuantMSE(d=d, bit_width=3, seed=42, norm_correction=True) + + idx_old, norm_old = tq_old.quantize(x) + idx_new, norm_new = tq_new.quantize(x) + + np.testing.assert_array_equal(idx_old, idx_new) + np.testing.assert_allclose(norm_old, norm_new) + + x_hat_old = tq_old.dequantize(idx_old, norm_old) + x_hat_new = tq_new.dequantize(idx_new, norm_new) + + assert not np.allclose(x_hat_old, x_hat_new) + class TestCompressedSizeBits: """Test compressed_size_bits method.""" diff --git a/turboquant/kv_cache.py b/turboquant/kv_cache.py index 1a6ec1f..80c61f9 100644 --- a/turboquant/kv_cache.py +++ b/turboquant/kv_cache.py @@ -52,7 +52,14 @@ class KVCacheCompressor: compressor.compress_token(k_vec, v_vec, layer=0, head=0) """ - def __init__(self, head_dim: int, k_bits: int = 3, v_bits: int = 3, seed: int = 42): + def __init__( + self, + head_dim: int, + k_bits: int = 3, + v_bits: int = 3, + seed: int = 42, + norm_correction: bool = True, + ): """ Args: head_dim: Dimension of each attention head vector. @@ -65,10 +72,14 @@ class KVCacheCompressor: self.v_bits = v_bits # K cache uses full TurboQuant (inner product preservation) - self.k_quantizer = TurboQuant(head_dim, bit_width=k_bits, seed=seed) + self.k_quantizer = TurboQuant( + head_dim, bit_width=k_bits, seed=seed, norm_correction=norm_correction, + ) # V cache uses MSE-only PolarQuant (value reconstruction) - self.v_quantizer = TurboQuantMSE(head_dim, bit_width=v_bits, seed=seed + 500) + self.v_quantizer = TurboQuantMSE( + head_dim, bit_width=v_bits, seed=seed + 500, norm_correction=norm_correction, + ) def compress(self, k_cache: np.ndarray, v_cache: np.ndarray) -> CompressedKVCache: """Compress full KV cache tensors. diff --git a/turboquant/polar_quant.py b/turboquant/polar_quant.py index 97bacb9..73f9c00 100644 --- a/turboquant/polar_quant.py +++ b/turboquant/polar_quant.py @@ -29,10 +29,11 @@ class PolarQuant: x_hat = pq.dequantize(indices, norms) # reconstructed """ - def __init__(self, d: int, bit_width: int, seed: int = 42): + def __init__(self, d: int, bit_width: int, seed: int = 42, norm_correction: bool = True): self.d = d self.bit_width = bit_width self.n_centroids = 1 << bit_width + self.norm_correction = norm_correction rng = np.random.default_rng(seed) self.rotation = random_rotation_dense(d, rng) @@ -84,8 +85,18 @@ class PolarQuant: indices = indices[np.newaxis, :] norms = np.array([norms]) - # Look up centroids → unit-norm reconstruction + # Look up centroids in the rotated domain. y_hat = self.centroids[indices] + + # Production TurboQuant rescales the reconstructed rotated vector by + # grp_norm / recon_norm before the inverse rotation. In the normalized + # Python prototype, grp_norm = 1, so this simplifies to re-normalizing + # y_hat back to unit norm. + if self.norm_correction: + y_hat_norms = np.linalg.norm(y_hat, axis=1, keepdims=True) + y_hat_norms = np.where(y_hat_norms > 1e-10, y_hat_norms, 1.0) + y_hat = y_hat / y_hat_norms + x_hat_unit = (self.rotation.T @ y_hat.T).T # Rescale by original norms diff --git a/turboquant/turboquant.py b/turboquant/turboquant.py index 51bc5cf..2a3bc88 100644 --- a/turboquant/turboquant.py +++ b/turboquant/turboquant.py @@ -40,7 +40,7 @@ class TurboQuant: tq.dequantize(tq.quantize(y))) """ - def __init__(self, d: int, bit_width: int, seed: int = 42): + def __init__(self, d: int, bit_width: int, seed: int = 42, norm_correction: bool = True): """ Args: d: Vector dimension. @@ -55,7 +55,9 @@ class TurboQuant: self.bit_width = bit_width # Stage 1: PolarQuant at (b-1) bits - self.polar_quant = PolarQuant(d, bit_width=bit_width - 1, seed=seed) + self.polar_quant = PolarQuant( + d, bit_width=bit_width - 1, seed=seed, norm_correction=norm_correction, + ) # Stage 2: QJL for residual (uses different seed) self.qjl = QJL(d, seed=seed + 1000) @@ -133,10 +135,12 @@ class TurboQuantMSE: Simpler, slightly less storage overhead (no QJL signs needed). """ - def __init__(self, d: int, bit_width: int, seed: int = 42): + def __init__(self, d: int, bit_width: int, seed: int = 42, norm_correction: bool = True): self.d = d self.bit_width = bit_width - self.polar_quant = PolarQuant(d, bit_width=bit_width, seed=seed) + self.polar_quant = PolarQuant( + d, bit_width=bit_width, seed=seed, norm_correction=norm_correction, + ) def quantize(self, x: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """Returns (indices, norms)."""