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) <noreply@anthropic.com>

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()

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()

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."""

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."""

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.

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

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)."""