turboquant-gpu 0.1.0__tar.gz

turboquant_gpu-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Anirudh Bharadwaj Vangara
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

turboquant_gpu-0.1.0/PKG-INFO

@@ -0,0 +1,96 @@
+Metadata-Version: 2.4
+Name: turboquant-gpu
+Version: 0.1.0
+Summary: TurboQuant KV cache compression for LLM inference — cuTile GPU kernels
+Author: Anirudh Bharadwaj Vangara
+License-Expression: MIT
+Keywords: quantization,kv-cache,llm,inference,cutile,cuda,gpu,attention,blackwell,hopper,h100,b200
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: scipy
+Provides-Extra: gpu
+Requires-Dist: cuda-tile; extra == "gpu"
+Dynamic: license-file
+
+# turboquant-gpu
+
+**5.02x KV cache compression for LLM inference** — GPU-accelerated cuTile kernels with PyTorch fallback.
+
+```
+pip install turboquant-gpu
+```
+
+## quick start
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from turboquant_gpu import TurboQuantEngine
+import torch
+
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-1.5B", torch_dtype=torch.float16, device_map="cuda")
+tok   = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B")
+
+engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cuda")
+result = engine.generate(model, tok, "The key to efficient LLM inference is")
+
+print(result["text"])
+print(f"{result['tokens']} tokens | {result['stats']['ratio']:.1f}x compression")
+```
+
+## how it works
+
+Implements the [TurboQuant](https://arxiv.org/abs/2501.09747) algorithm:
+
+1. **normalize + rotate** — random orthogonal rotation (Pi) makes coordinates near-Gaussian
+2. **Lloyd-Max quantize** — optimal scalar quantization against N(0, 1/d)
+3. **QJL bias correction** — 1-bit sign sketch of the residual for unbiased key scores
+
+At 3-bit (2-bit MSE + 1-bit QJL) this gives ~5x compression with negligible quality loss.
+
+## step-by-step api
+
+```python
+engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cuda")
+
+# after model prefill:
+compressed = engine.compress_kv_cache(out.past_key_values)
+cache      = engine.build_cache(compressed)
+stats      = engine.compression_stats(out.past_key_values)
+
+# or just do it all in one call:
+result = engine.generate(model, tokenizer, "your prompt here")
+```
+
+## gpu support
+
+Written in [cuTile](https://docs.nvidia.com/cuda/cutile-python/) for cross-architecture portability.
+Falls back to PyTorch if cuTile or a compatible driver isn't available.
+
+| GPU family | Architecture | Status |
+|------------|-------------|--------|
+| A100       | Ampere      | supported (PyTorch fallback) |
+| H100       | Hopper      | supported |
+| B200/B300  | Blackwell   | supported + swizzle fast path |
+
+## kernels
+
+| kernel | what it does |
+|--------|-------------|
+| `compress_keys` | normalize → rotate(Pi^T) → Lloyd-Max quantize → QJL signs |
+| `compress_values` | normalize → rotate(Pi^T) → Lloyd-Max quantize |
+| `decompress_values` | dequantize → un-rotate(Pi) → scale by norms |
+| `attention_scores` | asymmetric dot product with QJL correction |
+| `fused_attention` | scores + online softmax + V accumulation |
+
+## license
+
+MIT

turboquant_gpu-0.1.0/README.md

@@ -0,0 +1,73 @@
+# turboquant-gpu
+
+**5.02x KV cache compression for LLM inference** — GPU-accelerated cuTile kernels with PyTorch fallback.
+
+```
+pip install turboquant-gpu
+```
+
+## quick start
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from turboquant_gpu import TurboQuantEngine
+import torch
+
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-1.5B", torch_dtype=torch.float16, device_map="cuda")
+tok   = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B")
+
+engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cuda")
+result = engine.generate(model, tok, "The key to efficient LLM inference is")
+
+print(result["text"])
+print(f"{result['tokens']} tokens | {result['stats']['ratio']:.1f}x compression")
+```
+
+## how it works
+
+Implements the [TurboQuant](https://arxiv.org/abs/2501.09747) algorithm:
+
+1. **normalize + rotate** — random orthogonal rotation (Pi) makes coordinates near-Gaussian
+2. **Lloyd-Max quantize** — optimal scalar quantization against N(0, 1/d)
+3. **QJL bias correction** — 1-bit sign sketch of the residual for unbiased key scores
+
+At 3-bit (2-bit MSE + 1-bit QJL) this gives ~5x compression with negligible quality loss.
+
+## step-by-step api
+
+```python
+engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cuda")
+
+# after model prefill:
+compressed = engine.compress_kv_cache(out.past_key_values)
+cache      = engine.build_cache(compressed)
+stats      = engine.compression_stats(out.past_key_values)
+
+# or just do it all in one call:
+result = engine.generate(model, tokenizer, "your prompt here")
+```
+
+## gpu support
+
+Written in [cuTile](https://docs.nvidia.com/cuda/cutile-python/) for cross-architecture portability.
+Falls back to PyTorch if cuTile or a compatible driver isn't available.
+
+| GPU family | Architecture | Status |
+|------------|-------------|--------|
+| A100       | Ampere      | supported (PyTorch fallback) |
+| H100       | Hopper      | supported |
+| B200/B300  | Blackwell   | supported + swizzle fast path |
+
+## kernels
+
+| kernel | what it does |
+|--------|-------------|
+| `compress_keys` | normalize → rotate(Pi^T) → Lloyd-Max quantize → QJL signs |
+| `compress_values` | normalize → rotate(Pi^T) → Lloyd-Max quantize |
+| `decompress_values` | dequantize → un-rotate(Pi) → scale by norms |
+| `attention_scores` | asymmetric dot product with QJL correction |
+| `fused_attention` | scores + online softmax + V accumulation |
+
+## license
+
+MIT

turboquant_gpu-0.1.0/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "turboquant-gpu"
+version = "0.1.0"
+description = "TurboQuant KV cache compression for LLM inference — cuTile GPU kernels"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.10"
+authors = [
+    { name = "Anirudh Bharadwaj Vangara" },
+]
+keywords = [
+    "quantization", "kv-cache", "llm", "inference",
+    "cutile", "cuda", "gpu", "attention",
+    "blackwell", "hopper", "h100", "b200",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+dependencies = [
+    "torch",
+    "scipy",
+]
+
+[project.optional-dependencies]
+gpu = ["cuda-tile"]
+
+[tool.setuptools.packages.find]
+include = ["turboquant_gpu*"]

turboquant_gpu-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

turboquant_gpu-0.1.0/tests/test_attention.py

@@ -0,0 +1,153 @@
+"""
+Test attention kernel correctness against PyTorch reference and ground truth.
+
+Key properties verified:
+  - Scores match PyTorch reference within FP16 tolerance
+  - Inner product estimator is unbiased (mean error ≈ 0)
+  - Correct 1/√d scaling
+  - Sweep over multiple bit widths and sequence lengths
+"""
+
+import math
+import sys
+import os
+import torch
+import pytest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from turboquant_gpu import TurboQuantEngine
+
+
+@pytest.mark.parametrize("seq_q,seq_k", [(1, 64), (1, 256), (1, 1024), (16, 512)])
+@pytest.mark.parametrize("total_bits", [3, 4])
+def test_scores_shape(seq_q, seq_k, total_bits):
+    engine = TurboQuantEngine(head_dim=128, total_bits=total_bits, device="cpu")
+
+    Q = torch.randn(seq_q, 128).half()
+    K = torch.randn(seq_k, 128).half()
+    compressed_k = engine.compress_keys_pytorch(K)
+    scores = engine.attention_scores_pytorch(Q, compressed_k)
+
+    assert scores.shape == (seq_q, seq_k)
+
+
+@pytest.mark.parametrize("total_bits", [3, 4])
+def test_unbiasedness(total_bits):
+    """
+    The asymmetric estimator should be unbiased: E[estimated_ip] ≈ true_ip.
+    We test by averaging over many random pairs.
+    """
+    d = 128
+    n = 2000
+    engine = TurboQuantEngine(head_dim=d, total_bits=total_bits, seed=42, device="cpu")
+
+    K = torch.randn(n, d)
+    K = K / torch.norm(K, dim=-1, keepdim=True)
+    Q = torch.randn(n, d)
+    Q = Q / torch.norm(Q, dim=-1, keepdim=True)
+
+    true_ip = (Q * K).sum(dim=-1)
+
+    compressed_k = engine.compress_keys_pytorch(K.half())
+    scores_unscaled = engine.attention_scores_pytorch(Q.half(), compressed_k)
+
+    estimated_per_pair = torch.diag(scores_unscaled) / engine.scale
+
+    bias = (estimated_per_pair - true_ip).mean().abs().item()
+    assert bias < 0.05, f"Bias too large: {bias:.4f} (should be ~0)"
+
+
+@pytest.mark.parametrize("total_bits", [2, 3, 4])
+def test_correlation_with_true_scores(total_bits):
+    """Estimated scores should correlate well with true Q·K^T scores."""
+    d = 128
+    seq_q, seq_k = 1, 512
+    engine = TurboQuantEngine(head_dim=d, total_bits=total_bits, seed=42, device="cpu")
+
+    Q = torch.randn(seq_q, d).half()
+    K = torch.randn(seq_k, d).half()
+
+    true_scores = (Q.float() @ K.float().T) * engine.scale
+    compressed_k = engine.compress_keys_pytorch(K)
+    estimated_scores = engine.attention_scores_pytorch(Q, compressed_k)
+
+    corr = torch.corrcoef(
+        torch.stack([true_scores.flatten(), estimated_scores.flatten()])
+    )[0, 1].item()
+
+    min_corr = {2: 0.45, 3: 0.75, 4: 0.90}[total_bits]
+    assert corr > min_corr, (
+        f"bits={total_bits}: correlation={corr:.4f} < {min_corr}"
+    )
+
+
+def test_scaling_correct():
+    """Scores should be scaled by 1/√d."""
+    d = 128
+    engine = TurboQuantEngine(head_dim=d, total_bits=3, device="cpu")
+
+    Q = torch.randn(1, d).half()
+    K = torch.randn(64, d).half()
+
+    compressed_k = engine.compress_keys_pytorch(K)
+    scores = engine.attention_scores_pytorch(Q, compressed_k)
+
+    raw_ip = Q.float() @ compressed_k["k_mse"].float().T
+    assert scores.abs().mean().item() < raw_ip.abs().mean().item(), (
+        "Scaled scores should be smaller than raw inner products"
+    )
+
+
+def test_single_decode_token():
+    """Typical decode scenario: seq_q=1, seq_k=large."""
+    engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cpu")
+    Q = torch.randn(1, 128).half()
+    K = torch.randn(2048, 128).half()
+
+    compressed_k = engine.compress_keys_pytorch(K)
+    scores = engine.attention_scores_pytorch(Q, compressed_k)
+
+    assert scores.shape == (1, 2048)
+    assert not torch.isnan(scores).any()
+    assert not torch.isinf(scores).any()
+
+
+def test_needle_in_haystack():
+    """Can we still find the most-attended key after compression?"""
+    d = 128
+    seq_k = 1024
+    engine = TurboQuantEngine(head_dim=d, total_bits=3, seed=42, device="cpu")
+
+    K = torch.randn(seq_k, d)
+    K = K / torch.norm(K, dim=-1, keepdim=True)
+
+    needle_pos = seq_k // 3
+    Q = K[needle_pos].unsqueeze(0)
+
+    true_scores = Q @ K.T
+    true_top1 = true_scores.argmax(dim=-1).item()
+
+    compressed_k = engine.compress_keys_pytorch(K.half())
+    estimated_scores = engine.attention_scores_pytorch(Q.half(), compressed_k)
+    estimated_top1 = estimated_scores.argmax(dim=-1).item()
+
+    assert estimated_top1 == true_top1 or abs(estimated_top1 - needle_pos) < 5, (
+        f"Needle at {needle_pos}, true top1={true_top1}, estimated top1={estimated_top1}"
+    )
+
+
+@pytest.mark.parametrize("total_bits", [3, 4])
+def test_bits_sweep_scores_reasonable(total_bits):
+    """Higher bits should produce lower MSE in scores."""
+    d = 128
+    engine = TurboQuantEngine(head_dim=d, total_bits=total_bits, device="cpu")
+
+    Q = torch.randn(4, d).half()
+    K = torch.randn(256, d).half()
+
+    true_scores = (Q.float() @ K.float().T) * engine.scale
+    compressed_k = engine.compress_keys_pytorch(K)
+    estimated_scores = engine.attention_scores_pytorch(Q, compressed_k)
+
+    score_mse = ((true_scores - estimated_scores) ** 2).mean().item()
+    assert score_mse < 0.5, f"Score MSE too high: {score_mse:.6f}"

turboquant_gpu-0.1.0/tests/test_codebook.py

@@ -0,0 +1,92 @@
+"""
+Test Lloyd-Max codebook correctness.
+
+Verifies:
+  - Centroid symmetry around zero
+  - Correct number of levels per bit-width
+  - Centroids are monotonically sorted
+  - Distortion is within theoretical upper bound
+  - Quantize → dequantize round-trip consistency
+"""
+
+import math
+import sys
+import os
+import torch
+import pytest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from turboquant_gpu.codebook import LloydMaxCodebook, solve_lloyd_max
+
+
+@pytest.mark.parametrize("d", [64, 128, 256])
+@pytest.mark.parametrize("bits", [1, 2, 3, 4])
+def test_correct_num_levels(d, bits):
+    cb = LloydMaxCodebook(d, bits)
+    assert cb.n_levels == 2**bits
+    assert cb.centroids.shape == (2**bits,)
+    assert cb.boundaries.shape == (2**bits - 1,)
+
+
+@pytest.mark.parametrize("d", [64, 128, 256])
+@pytest.mark.parametrize("bits", [2, 3, 4])
+def test_symmetry(d, bits):
+    cb = LloydMaxCodebook(d, bits)
+    assert cb.centroids.sum().abs().item() < 1e-4, (
+        f"Centroids should be symmetric: sum={cb.centroids.sum().item()}"
+    )
+
+
+@pytest.mark.parametrize("bits", [1, 2, 3, 4])
+def test_sorted(bits):
+    cb = LloydMaxCodebook(128, bits)
+    for i in range(len(cb.centroids) - 1):
+        assert cb.centroids[i] < cb.centroids[i + 1]
+    for i in range(len(cb.boundaries) - 1):
+        assert cb.boundaries[i] < cb.boundaries[i + 1]
+
+
+@pytest.mark.parametrize("bits", [1, 2, 3, 4])
+def test_boundaries_between_centroids(bits):
+    cb = LloydMaxCodebook(128, bits)
+    for i in range(len(cb.boundaries)):
+        assert cb.centroids[i] < cb.boundaries[i] < cb.centroids[i + 1]
+
+
+@pytest.mark.parametrize("bits", [1, 2, 3, 4])
+def test_distortion_within_paper_bound(bits):
+    """MSE distortion per vector <= sqrt(3) * pi/2 * (1/4^b) for unit vectors."""
+    d = 128
+    cb = LloydMaxCodebook(d, bits)
+    sigma = 1.0 / math.sqrt(d)
+
+    n_samples = 5000
+    x = torch.randn(n_samples, d)
+    x = x / torch.norm(x, dim=-1, keepdim=True)
+
+    rotated = x  # with identity rotation, still approximately N(0, 1/d)
+    indices = cb.quantize(rotated)
+    reconstructed = cb.dequantize(indices)
+    mse = ((rotated - reconstructed) ** 2).sum(dim=-1).mean().item()
+
+    upper_bound = math.sqrt(3) * math.pi / 2 * (1.0 / (4**bits))
+    assert mse < upper_bound * 2.0, (
+        f"bits={bits}: MSE={mse:.6f} exceeds 2× paper bound {upper_bound:.6f}"
+    )
+
+
+def test_roundtrip_identity():
+    """Quantize → dequantize should map each centroid exactly to itself."""
+    cb = LloydMaxCodebook(128, 3)
+    indices = torch.arange(cb.n_levels)
+    reconstructed = cb.dequantize(indices)
+    assert torch.allclose(reconstructed, cb.centroids, atol=1e-6)
+
+
+def test_quantize_nearest_centroid():
+    """Values close to a centroid should map to that centroid's index."""
+    cb = LloydMaxCodebook(128, 2)
+    for i, c in enumerate(cb.centroids):
+        test_val = c + 1e-6
+        idx = cb.quantize(test_val.unsqueeze(0))
+        assert idx.item() == i, f"Expected index {i}, got {idx.item()}"

turboquant_gpu-0.1.0/tests/test_compress.py

@@ -0,0 +1,133 @@
+"""
+Test compression kernel correctness against PyTorch reference.
+
+Compares cuTile kernel output (via TurboQuantEngine) against the pure PyTorch
+reference path for:
+  - MSE indices match
+  - QJL signs match
+  - Norms are within FP16 tolerance
+  - Various sequence lengths and bit widths
+"""
+
+import sys
+import os
+import torch
+import pytest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from turboquant_gpu import TurboQuantEngine
+
+
+def _make_random_keys(seq_k: int, head_dim: int = 128) -> torch.Tensor:
+    K = torch.randn(seq_k, head_dim)
+    return K.half()
+
+
+@pytest.mark.parametrize("seq_k", [64, 128, 256, 512])
+@pytest.mark.parametrize("total_bits", [3, 4])
+def test_compress_keys_shapes(seq_k, total_bits):
+    engine = TurboQuantEngine(head_dim=128, total_bits=total_bits, device="cpu")
+    K = _make_random_keys(seq_k)
+    compressed = engine.compress_keys_pytorch(K)
+
+    assert compressed["indices"].shape == (seq_k, 128)
+    assert compressed["k_mse"].shape == (seq_k, 128)
+    assert compressed["qjl_signs"].shape == (seq_k, 128)
+    assert compressed["vec_norms"].shape == (seq_k,)
+    assert compressed["residual_norms"].shape == (seq_k,)
+
+
+@pytest.mark.parametrize("total_bits", [3, 4])
+def test_indices_in_valid_range(total_bits):
+    engine = TurboQuantEngine(head_dim=128, total_bits=total_bits, device="cpu")
+    K = _make_random_keys(256)
+    compressed = engine.compress_keys_pytorch(K)
+
+    mse_bits = max(total_bits - 1, 1)
+    max_idx = (1 << mse_bits) - 1
+    assert compressed["indices"].max().item() <= max_idx
+    assert compressed["indices"].min().item() >= 0
+
+
+def test_signs_are_pm1():
+    engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cpu")
+    K = _make_random_keys(256)
+    compressed = engine.compress_keys_pytorch(K)
+
+    unique_vals = set(compressed["qjl_signs"].unique().tolist())
+    assert unique_vals.issubset({-1, 1}), f"Signs should be +/-1, got {unique_vals}"
+
+
+def test_norms_positive():
+    engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cpu")
+    K = _make_random_keys(256)
+    compressed = engine.compress_keys_pytorch(K)
+
+    assert (compressed["vec_norms"] >= 0).all()
+    assert (compressed["residual_norms"] >= 0).all()
+
+
+def test_compress_values_shapes():
+    engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cpu")
+    V = _make_random_keys(256)
+    compressed = engine.compress_values_pytorch(V)
+
+    assert compressed["indices"].shape == (256, 128)
+    assert compressed["vec_norms"].shape == (256,)
+
+
+def test_value_indices_in_valid_range():
+    engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cpu")
+    V = _make_random_keys(256)
+    compressed = engine.compress_values_pytorch(V)
+
+    max_idx = (1 << 3) - 1  # values use all bits (no QJL)
+    assert compressed["indices"].max().item() <= max_idx
+
+
+def test_compress_matches_reference():
+    """Compare our engine against the original cutiledump compressors.py implementation."""
+    engine = TurboQuantEngine(head_dim=128, total_bits=3, seed=42, device="cpu")
+
+    K = torch.randn(64, 128).half()
+    compressed = engine.compress_keys_pytorch(K)
+
+    K_f = K.float()
+    vec_norms = torch.norm(K_f, dim=-1, keepdim=True)
+    K_normed = K_f / (vec_norms + 1e-8)
+    rotated = K_normed @ engine.PiT.float()
+    centroids = engine.key_codebook.centroids
+    diffs = rotated.unsqueeze(-1) - centroids
+    expected_indices = diffs.abs().argmin(dim=-1).to(torch.uint8)
+
+    assert torch.equal(compressed["indices"], expected_indices)
+
+
+@pytest.mark.parametrize("total_bits", [3, 4])
+def test_k_mse_reconstruction_quality(total_bits):
+    """k_mse should be a reasonable approximation of K."""
+    engine = TurboQuantEngine(head_dim=128, total_bits=total_bits, device="cpu")
+    K = torch.randn(512, 128).half()
+    compressed = engine.compress_keys_pytorch(K)
+
+    cos_sim = torch.nn.functional.cosine_similarity(
+        K.float(), compressed["k_mse"].float(), dim=-1
+    )
+    assert cos_sim.mean().item() > 0.85, (
+        f"Mean cosine similarity too low: {cos_sim.mean().item():.4f}"
+    )
+
+
+def test_edge_case_zero_vector():
+    engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cpu")
+    K = torch.zeros(1, 128).half()
+    compressed = engine.compress_keys_pytorch(K)
+    assert not torch.isnan(compressed["k_mse"]).any()
+    assert not torch.isinf(compressed["k_mse"]).any()
+
+
+def test_edge_case_large_values():
+    engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cpu")
+    K = torch.randn(4, 128).half() * 100
+    compressed = engine.compress_keys_pytorch(K)
+    assert not torch.isnan(compressed["k_mse"]).any()