tiny-turboquant 0.1.0__py3-none-any.whl

benchmarks/bench_ann.py

@@ -0,0 +1,188 @@
+"""
+ANN A/B benchmark: TurboQuant vs faiss PQ vs faiss IVFPQ vs RaBitQ (optional).
+
+Reports recall@k, indexing time, and queries/second for each method on the
+same dataset. The dataset can be loaded from a .npy file (most realistic)
+or generated synthetically.
+
+Usage:
+    python -m benchmarks.bench_ann --data path/to/embeddings.npy --bits 4 --k 10
+    python -m benchmarks.bench_ann --synthetic 50000 --d 768 --bits 4
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from tiny_turboquant import TurboQuantProd
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data",      type=str, default=None,
+                   help="Path to .npy of shape (N, D), float32.")
+    p.add_argument("--synthetic", type=int, default=50_000)
+    p.add_argument("--d",         type=int, default=768)
+    p.add_argument("--bits",      type=int, default=4)
+    p.add_argument("--k",         type=int, default=10)
+    p.add_argument("--n_queries", type=int, default=500)
+    p.add_argument("--device",    default="cuda" if torch.cuda.is_available() else "cpu")
+    return p.parse_args()
+
+
+def load_data(args) -> tuple[np.ndarray, np.ndarray]:
+    if args.data is not None:
+        X = np.load(args.data).astype(np.float32)
+    else:
+        rng = np.random.default_rng(0)
+        # Cluster structure mimicking real embedding stores.
+        n_clusters = max(50, args.synthetic // 500)
+        centers = rng.standard_normal((n_clusters, args.d)).astype(np.float32)
+        centers /= np.linalg.norm(centers, axis=1, keepdims=True)
+        cluster_id = rng.integers(0, n_clusters, size=args.synthetic)
+        X = centers[cluster_id] + 0.25 * rng.standard_normal((args.synthetic, args.d)).astype(np.float32)
+
+    X = X / (np.linalg.norm(X, axis=1, keepdims=True) + 1e-12)
+    rng = np.random.default_rng(1)
+    qi = rng.choice(X.shape[0], args.n_queries, replace=False)
+    return X, X[qi].copy()
+
+
+def recall_at_k(true_top: np.ndarray, est_top: np.ndarray) -> float:
+    hits = sum(len(set(t.tolist()) & set(e.tolist()))
+               for t, e in zip(true_top, est_top))
+    return hits / true_top.size
+
+
+# ---- methods ----------------------------------------------------------
+
+def turboquant_run(X, Q, k, bits, device):
+    Xt = torch.from_numpy(X).to(device)
+    Qt = torch.from_numpy(Q).to(device)
+    d = X.shape[1]
+
+    t0 = time.perf_counter()
+    q  = TurboQuantProd.build(d, bits, device=device, seed=0,
+                              dtype=torch.float32)
+    idx, signs, gamma = q.quant(Xt)
+    if device.startswith("cuda"):
+        torch.cuda.synchronize()
+    t_index = time.perf_counter() - t0
+
+    t0 = time.perf_counter()
+    Xh = q.dequant(idx, signs, gamma)
+    sims = Qt @ Xh.T
+    top = sims.topk(k, dim=-1).indices.cpu().numpy()
+    if device.startswith("cuda"):
+        torch.cuda.synchronize()
+    t_query = time.perf_counter() - t0
+    return top, t_index, t_query
+
+
+def faiss_pq_run(X, Q, k, bits):
+    try:
+        import faiss
+    except Exception as e:   # pragma: no cover
+        print(f"  (faiss unavailable: {e}) — skipping PQ")
+        return None, None, None
+    d = X.shape[1]
+    # Match bit-budget: PQ stores log2(ks) bits per sub-quantizer;
+    # M sub-vectors * log2(ks) / d total bits per coord.
+    # Use ks=2**bits and M=d so PQ's bit-budget = bits per coord.
+    M  = d
+    ks = 2 ** bits
+    if ks > 256:                          # faiss IndexPQ stores 1 byte per code
+        ks = 256
+    index = faiss.IndexPQ(d, M, int(np.log2(ks)))
+    index.metric_type = faiss.METRIC_INNER_PRODUCT
+    t0 = time.perf_counter(); index.train(X); index.add(X); t_index = time.perf_counter() - t0
+    t0 = time.perf_counter()
+    _, top = index.search(Q, k)
+    t_query = time.perf_counter() - t0
+    return top, t_index, t_query
+
+
+def faiss_ivfpq_run(X, Q, k, bits):
+    try:
+        import faiss
+    except Exception:
+        return None, None, None
+    d = X.shape[1]
+    nlist = 100
+    M = d
+    ks_bits = min(bits, 8)
+    quantizer = faiss.IndexFlatIP(d)
+    index = faiss.IndexIVFPQ(quantizer, d, nlist, M, ks_bits)
+    index.metric_type = faiss.METRIC_INNER_PRODUCT
+    t0 = time.perf_counter(); index.train(X); index.add(X); t_index = time.perf_counter() - t0
+    index.nprobe = 16
+    t0 = time.perf_counter()
+    _, top = index.search(Q, k)
+    t_query = time.perf_counter() - t0
+    return top, t_index, t_query
+
+
+def rabitq_run(X, Q, k, bits):                  # pragma: no cover
+    """Optional RaBitQ baseline; skipped silently if package missing."""
+    try:
+        import rabitqlib                          # noqa: F401
+    except Exception as e:
+        print(f"  (rabitq unavailable: {e}) — skipping RaBitQ")
+        return None, None, None
+    print("  TODO: integrate rabitqlib (left as exercise — official API churns)")
+    return None, None, None
+
+
+# ---- driver ----------------------------------------------------------
+
+def main():
+    args = parse_args()
+    X, Q = load_data(args)
+    n, d = X.shape
+
+    # Ground truth: brute-force top-k inner product
+    print(f"\nDataset: {n} vectors of dim {d}, {len(Q)} queries, k={args.k}, "
+          f"bits={args.bits}, device={args.device}\n")
+    t0 = time.perf_counter()
+    sims = Q @ X.T
+    true_top = np.argpartition(-sims, args.k, axis=1)[:, :args.k]
+    # exact top-k order
+    true_top = np.array([row[np.argsort(-sims[i, row])] for i, row in enumerate(true_top)])
+    t_brute = time.perf_counter() - t0
+    print(f"brute force baseline: {t_brute:.3f}s\n")
+
+    rows = []
+    for name, fn in (
+        ("TurboQuant",       lambda: turboquant_run(X, Q, args.k, args.bits, args.device)),
+        ("faiss PQ",         lambda: faiss_pq_run(X, Q, args.k, args.bits)),
+        ("faiss IVF-PQ",     lambda: faiss_ivfpq_run(X, Q, args.k, args.bits)),
+        ("RaBitQ (opt.)",    lambda: rabitq_run(X, Q, args.k, args.bits)),
+    ):
+        top, t_index, t_query = fn()
+        if top is None:
+            continue
+        rec = recall_at_k(true_top, top)
+        qps = len(Q) / max(t_query, 1e-9)
+        rows.append((name, t_index, t_query, qps, rec))
+
+    print(f"{'method':16s} | {'index(s)':>9} | {'query(s)':>9} | {'qps':>10} | {'recall@k':>9}")
+    print("-" * 72)
+    for name, ti, tq, qps, rec in rows:
+        print(f"{name:16s} | {ti:>9.3f} | {tq:>9.3f} | {qps:>10.0f} | {rec:>9.3f}")
+    print()
+
+    if rows:
+        tq_row = next((r for r in rows if r[0] == "TurboQuant"), None)
+        pq_row = next((r for r in rows if r[0] == "faiss PQ"),   None)
+        if tq_row and pq_row:
+            print(f"TurboQuant indexing speedup vs PQ: {pq_row[1] / tq_row[1]:.1f}×")
+            print(f"TurboQuant recall delta vs PQ    : {tq_row[4] - pq_row[4]:+.3f}")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/bench_kv_real.py

@@ -0,0 +1,192 @@
+"""
+Real-LLM benchmark: TurboQuant KV cache vs fp16 baseline.
+
+Runs a short prompt-completion task on a small open model and reports:
+    - actual packed KV cache memory
+    - per-token attention output cosine similarity
+    - generation perplexity / acceptance vs baseline
+
+Default model: HuggingFaceTB/SmolLM2-360M-Instruct (small enough for
+free Kaggle T4, big enough to be representative).
+
+Usage:
+    python -m benchmarks.bench_kv_real --model SmolLM2-360M --bits 4
+    python -m benchmarks.bench_kv_real --bits 2 --bits_outlier 3 --n_outlier 32
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import time
+
+# Quieten the noisy weight-loading tqdm bar from transformers / accelerate.
+# (When stdout is captured by a subprocess pipe, tqdm prints one line per
+# update instead of redrawing in place, which floods the log.)
+os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
+os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+os.environ.setdefault("ACCELERATE_DISABLE_RICH", "1")
+os.environ.setdefault("TQDM_DISABLE", "1")
+
+import torch
+
+from tiny_turboquant import TurboQuantKVCache
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="HuggingFaceTB/SmolLM2-360M-Instruct")
+    p.add_argument("--bits",          type=int, default=4)
+    p.add_argument("--bits_outlier",  type=int, default=None,
+                   help="If set, enable outlier-channel split with this bit-width.")
+    p.add_argument("--n_outlier",     type=int, default=32)
+    p.add_argument("--prompt", default=(
+        "The capital of France is Paris. The capital of Germany is Berlin. "
+        "The capital of Spain is Madrid. The capital of Italy is"))
+    p.add_argument("--max_new_tokens", type=int, default=64)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--dtype",  default="float16")
+    return p.parse_args()
+
+
+@torch.no_grad()
+def measure(model, tokenizer, prompt, cache, max_new_tokens):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    t0 = time.perf_counter()
+    out = model.generate(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        do_sample=False,
+        past_key_values=cache,
+    )
+    elapsed = time.perf_counter() - t0
+    text = tokenizer.decode(out[0], skip_special_tokens=True)
+    return text, elapsed, out
+
+
+@torch.no_grad()
+def measure_logit_kl(model, tokenizer, prompt, cache_factory, max_new_tokens=32):
+    """Teacher-forced KL between fp16 baseline logits and TurboQuant logits.
+
+    Both paths see identical input at every step. This isolates pure
+    quantization error. Free-running generations are still produced for
+    the "identical generation?" check, but they are NOT used for KL —
+    otherwise the metric explodes the moment paths diverge by one token.
+    """
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    # 1) Free-running baseline + TQ generations (for identical-generation check only)
+    out_ref = model.generate(
+        **inputs, max_new_tokens=max_new_tokens, do_sample=False,
+        return_dict_in_generate=True,
+    )
+    cache = cache_factory()
+    out_tq = model.generate(
+        **inputs, max_new_tokens=max_new_tokens, do_sample=False,
+        return_dict_in_generate=True,
+        past_key_values=cache,
+    )
+
+    text_ref = tokenizer.decode(out_ref.sequences[0], skip_special_tokens=True)
+    text_tq  = tokenizer.decode(out_tq.sequences[0],  skip_special_tokens=True)
+
+    # First-divergence index: how many decode steps stay greedy-identical?
+    # This is the production-relevant signal — small KL means the answer
+    # token is identical even when later filler tokens drift.
+    prompt_len_for_div = inputs["input_ids"].shape[1]
+    ref_new = out_ref.sequences[0, prompt_len_for_div:].tolist()
+    tq_new  = out_tq .sequences[0, prompt_len_for_div:].tolist()
+    n_new   = min(len(ref_new), len(tq_new))
+    first_diverge = next(
+        (i for i in range(n_new) if ref_new[i] != tq_new[i]),
+        n_new,
+    )
+
+    # 2) Teacher-forced KL: feed the SAME baseline-generated sequence into
+    #    both fp16-cache and TQ-cache, compare per-position logits.
+    seq = out_ref.sequences  # (1, prompt_len + new_tokens)
+    prompt_len = inputs["input_ids"].shape[1]
+
+    # fp16 reference: one forward pass, no cache compression.
+    ref_logits = model(seq).logits  # (1, T, V)
+
+    # TQ path: prefill + decode through TurboQuantKVCache so attention
+    # is computed against quantised K/V at every step.
+    tq_cache = cache_factory()
+    # prefill on the prompt
+    tq_logits_chunks = []
+    out = model(seq[:, :prompt_len], past_key_values=tq_cache, use_cache=True)
+    tq_logits_chunks.append(out.logits)
+    # decode one token at a time, feeding the *baseline* tokens
+    for t in range(prompt_len, seq.shape[1]):
+        out = model(seq[:, t:t + 1], past_key_values=tq_cache, use_cache=True)
+        tq_logits_chunks.append(out.logits)
+    tq_logits = torch.cat(tq_logits_chunks, dim=1)  # (1, T, V)
+
+    # Compare logits at positions [prompt_len-1 .. T-2] — the ones that
+    # predict the new tokens. Both saw identical inputs.
+    p = torch.softmax(ref_logits[:, prompt_len - 1:-1].float(), dim=-1)
+    q = torch.softmax(tq_logits [:, prompt_len - 1:-1].float(), dim=-1).clamp_min_(1e-12)
+    kls = (p * (p.clamp_min_(1e-12).log() - q.log())).sum(-1).squeeze(0).tolist()
+
+    return text_ref, text_tq, kls, cache, first_diverge, n_new
+
+
+def main() -> None:
+    args = parse_args()
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from transformers.utils import logging as hf_logging
+    import huggingface_hub.utils as hf_hub_utils
+
+    # Hard-disable every progress bar transformers / huggingface_hub know
+    # about. Env vars alone are unreliable when this script is launched as
+    # a subprocess from kaggle_run.py (pipes confuse tqdm).
+    hf_logging.disable_progress_bar()
+    hf_logging.set_verbosity_error()
+    hf_hub_utils.disable_progress_bars()
+
+    dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16,
+             "float32": torch.float32}[args.dtype]
+    print(f"Loading {args.model}  ({args.dtype} on {args.device})")
+    tok   = AutoTokenizer.from_pretrained(args.model)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model, torch_dtype=dtype
+    ).to(args.device).eval()
+
+    def factory():
+        return TurboQuantKVCache(
+            bits=args.bits,
+            bits_outlier=args.bits_outlier,
+            n_outlier=args.n_outlier,
+        )
+
+    text_ref, text_tq, kls, cache, first_diverge, n_new = measure_logit_kl(
+        model, tok, args.prompt, factory, max_new_tokens=args.max_new_tokens,
+    )
+
+    print("\n--- BASELINE (fp16, default DynamicCache) ---")
+    print(text_ref)
+    print("\n--- TURBOQUANT ---")
+    print(text_tq)
+
+    n_tokens = cache.get_seq_length(0)
+    fp16_bytes = cache.fp16_baseline_bytes()
+    actual_tq_bytes = cache.actual_memory_bytes()
+    theoretical_tq_bytes = cache.theoretical_memory_bytes()
+
+    print("\n--- METRICS ---")
+    print(f"prompt + decode tokens : {n_tokens}")
+    print(f"fp16 KV cache actual   : {fp16_bytes / 1e6:.2f} MB")
+    print(f"TurboQuant actual      : {actual_tq_bytes / 1e6:.2f} MB  "
+          f"(× compression = {fp16_bytes / max(actual_tq_bytes, 1):.2f})")
+    print(f"TurboQuant theoretical : {theoretical_tq_bytes / 1e6:.2f} MB")
+    print(f"mean per-token logit KL: {sum(kls) / len(kls):.5f}  (teacher-forced; lower is better)")
+    print(f"max  per-token logit KL: {max(kls):.5f}")
+    print(f"first divergence at    : {first_diverge}/{n_new} decode tokens "
+          f"(higher is better; {n_new}/{n_new} = exact match)")
+    print(f"identical generation?  : {text_ref == text_tq}")
+
+
+if __name__ == "__main__":
+    main()

demos/demo1_distortion_vs_theory.py

@@ -0,0 +1,75 @@
+"""
+Demo 1 — Empirical distortion matches Shannon-theoretic bounds.
+
+Story for the audience:
+  "We compress unit vectors to b bits per coordinate. The black dashed line
+   is the information-theoretic *lower* bound from Shannon source coding;
+   the orange line is what *any* possible algorithm could ever achieve
+   in the limit. TurboQuant (blue) sits within a small constant of optimum
+   for every bit-width, with zero tuning."
+
+Run:
+    python -m demos.demo1_distortion_vs_theory
+Outputs:
+    demos/out/distortion.png
+"""
+
+from __future__ import annotations
+
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from tiny_turboquant.numpy_reference import TurboQuantMSE, TurboQuantProd
+
+OUT = os.path.join(os.path.dirname(__file__), "out")
+os.makedirs(OUT, exist_ok=True)
+
+
+def main() -> None:
+    rng = np.random.default_rng(0)
+    d, n = 512, 4000
+    X = rng.standard_normal((n, d))
+    X /= np.linalg.norm(X, axis=1, keepdims=True)
+
+    bits = [1, 2, 3, 4, 5, 6]
+    mse_emp, mse_lb, mse_ub = [], [], []
+    for b in bits:
+        q = TurboQuantMSE(d, b, seed=0)
+        Xh = q.dequant(q.quant(X))
+        mse_emp.append(float(np.mean(np.sum((X - Xh) ** 2, axis=1))))
+        mse_lb.append(4.0 ** (-b))                    # Shannon lower bound
+        mse_ub.append(3 * np.pi / 2 * 4.0 ** (-b))    # paper Theorem 1
+
+    fig, ax = plt.subplots(figsize=(7, 4.5))
+    ax.plot(bits, mse_ub,  "--",  color="orange", label="Paper upper bound (3π/2·4⁻ᵇ)")
+    ax.plot(bits, mse_lb,  "--",  color="black",  label="Shannon lower bound (4⁻ᵇ)")
+    ax.plot(bits, mse_emp, "o-",  color="tab:blue", lw=2, label="TurboQuant empirical")
+    ax.set_yscale("log")
+    ax.set_xlabel("bits per coordinate")
+    ax.set_ylabel("MSE  E‖x − x̂‖²")
+    ax.set_title(f"TurboQuant-MSE vs information-theoretic bounds  (d={d}, n={n})")
+    ax.grid(True, which="both", alpha=0.3)
+    ax.legend()
+    fig.tight_layout()
+    out = os.path.join(OUT, "distortion.png")
+    fig.savefig(out, dpi=140)
+    print(f"saved {out}")
+
+    # Also show inner-product unbiasedness — the *killer* feature for ANN/RAG.
+    Y = rng.standard_normal((n, d)); Y /= np.linalg.norm(Y, axis=1, keepdims=True)
+    print("\n  bits |   bias   |  ip-MSE   |  ratio vs full-precision")
+    print("  -----+----------+-----------+--------------------------")
+    for b in (2, 3, 4):
+        qp = TurboQuantProd(d, b, seed=0)
+        idx, s, g = qp.quant(X)
+        Xh = qp.dequant(idx, s, g)
+        true = np.sum(X * Y, 1)
+        est  = np.sum(Xh * Y, 1)
+        bias = float(np.mean(est - true))
+        var  = float(np.mean((est - true) ** 2))
+        print(f"   {b}   | {bias:+.4f} | {var:.5f}   | "
+              f"{32/b:.0f}× smaller than fp32")
+
+
+if __name__ == "__main__":
+    main()

demos/demo2_ann_vs_pq.py

@@ -0,0 +1,133 @@
+"""
+Demo 2 — ANN search: TurboQuant vs Product Quantization.
+
+Story for the audience:
+  "We index 50k random embeddings and answer top-10 NN queries. TurboQuant
+   matches PQ's recall at the same bit budget while indexing in
+   ~milliseconds — PQ needs to k-means-train per subspace."
+
+Run:
+    python -m demos.demo2_ann_vs_pq
+"""
+
+from __future__ import annotations
+
+import time
+import numpy as np
+
+from tiny_turboquant.numpy_reference import TurboQuantProd
+
+
+# ----- toy product quantization baseline ----------------------------------
+
+def pq_train(X: np.ndarray, n_sub: int, ks: int, seed: int = 0):
+    """Train PQ codebooks by vectorised k-means per subspace."""
+    rng = np.random.default_rng(seed)
+    n, d = X.shape
+    sub_dim = d // n_sub
+    codebooks = np.empty((n_sub, ks, sub_dim))
+    codes = np.empty((n, n_sub), dtype=np.int32)
+
+    for m in range(n_sub):
+        sub = X[:, m * sub_dim : (m + 1) * sub_dim]               # (n, sub_dim)
+        centers = sub[rng.choice(n, ks, replace=False)].copy()    # (ks, sub_dim)
+        sub_sq = (sub * sub).sum(1, keepdims=True)                # (n, 1)
+        for _ in range(10):
+            # ||a-b||^2 = ||a||^2 + ||b||^2 - 2 a·b  (vectorised)
+            c_sq = (centers * centers).sum(1)                     # (ks,)
+            d2 = sub_sq + c_sq[None, :] - 2.0 * sub @ centers.T
+            assign = d2.argmin(1)
+            # Update centers
+            new_centers = np.zeros_like(centers)
+            counts = np.bincount(assign, minlength=ks)
+            np.add.at(new_centers, assign, sub)
+            mask = counts > 0
+            new_centers[mask] /= counts[mask, None]
+            new_centers[~mask] = centers[~mask]
+            centers = new_centers
+        codebooks[m] = centers
+        c_sq = (centers * centers).sum(1)
+        codes[:, m] = (sub_sq + c_sq[None, :] - 2.0 * sub @ centers.T).argmin(1)
+    return codebooks, codes, sub_dim
+
+
+def pq_decode(codebooks: np.ndarray, codes: np.ndarray) -> np.ndarray:
+    n_sub, ks, sub_dim = codebooks.shape
+    n = codes.shape[0]
+    out = np.empty((n, n_sub * sub_dim))
+    for m in range(n_sub):
+        out[:, m * sub_dim : (m + 1) * sub_dim] = codebooks[m, codes[:, m]]
+    return out
+
+
+# ----- evaluation ---------------------------------------------------------
+
+def recall_at_k(true_top: np.ndarray, est_top: np.ndarray) -> float:
+    hits = 0
+    for t, e in zip(true_top, est_top):
+        hits += len(set(t.tolist()) & set(e.tolist()))
+    return hits / true_top.size
+
+
+def main() -> None:
+    rng = np.random.default_rng(0)
+    n_db, n_q, d, k = 20_000, 200, 128, 10
+    bits = 4
+
+    # Use clustered data: real embeddings (BERT, OpenAI, etc.) are highly
+    # structured, *not* i.i.d. Gaussian. On i.i.d. Gaussian PQ has a built-in
+    # advantage because it gets to train on the same distribution. On
+    # clustered/structured data, TurboQuant's data-oblivious random rotation
+    # is competitive while keeping its huge indexing-time advantage.
+    n_clusters = 50
+    centers = rng.standard_normal((n_clusters, d))
+    centers /= np.linalg.norm(centers, axis=1, keepdims=True)
+    cluster_id = rng.integers(0, n_clusters, size=n_db)
+    X = centers[cluster_id] + 0.25 * rng.standard_normal((n_db, d))
+    X = X.astype(np.float64)
+    X /= np.linalg.norm(X, axis=1, keepdims=True)
+
+    q_cluster = rng.integers(0, n_clusters, size=n_q)
+    Q = centers[q_cluster] + 0.25 * rng.standard_normal((n_q, d))
+    Q = Q.astype(np.float64)
+    Q /= np.linalg.norm(Q, axis=1, keepdims=True)
+
+    # Ground truth
+    true_top = np.argsort(-Q @ X.T, axis=1)[:, :k]
+
+    # ----- TurboQuant -----
+    t0 = time.perf_counter()
+    tq = TurboQuantProd(d, bits=bits, seed=0)
+    idx, signs, gamma = tq.quant(X)
+    t_tq_index = time.perf_counter() - t0
+
+    t0 = time.perf_counter()
+    Xh = tq.dequant(idx, signs, gamma)
+    est_top_tq = np.argsort(-Q @ Xh.T, axis=1)[:, :k]
+    t_tq_query = time.perf_counter() - t0
+
+    rec_tq = recall_at_k(true_top, est_top_tq)
+
+    # ----- PQ at matched bit budget -----
+    # bits/coord = log2(ks) / sub_dim => ks=256, sub_dim=2 gives 4 bits/coord
+    n_sub, ks = d // 2, 256
+    t0 = time.perf_counter()
+    cb, codes, _ = pq_train(X, n_sub=n_sub, ks=ks, seed=0)
+    t_pq_index = time.perf_counter() - t0
+
+    t0 = time.perf_counter()
+    Xh_pq = pq_decode(cb, codes)
+    est_top_pq = np.argsort(-Q @ Xh_pq.T, axis=1)[:, :k]
+    t_pq_query = time.perf_counter() - t0
+    rec_pq = recall_at_k(true_top, est_top_pq)
+
+    print(f"\nDataset: {n_db} unit vectors in R^{d}, top-{k} ANN, bit budget = {bits}\n")
+    print(f"{'method':14s} | indexing(s) | query(s) | recall@{k}")
+    print("-" * 56)
+    print(f"{'TurboQuant':14s} | {t_tq_index:11.3f} | {t_tq_query:8.3f} | {rec_tq:.3f}")
+    print(f"{'PQ (k-means)':14s} | {t_pq_index:11.3f} | {t_pq_query:8.3f} | {rec_pq:.3f}")
+    print(f"\nIndexing speedup: {t_pq_index / t_tq_index:,.0f}×  in TurboQuant's favour")
+
+
+if __name__ == "__main__":
+    main()