TurboQuant-PyTorch: When 99.5% Attention Fidelity Still Breaks Your LLM

Hook

What if your compression algorithm achieved 99.5% attention score similarity but still generated complete garbage? TurboQuant-PyTorch demonstrates exactly this paradox—and why the paper's headline innovation actually makes things worse.

Context

Large language models face a brutal memory bottleneck: the key-value cache. Every token in your context window stores two vectors (a key and a value) for every attention layer, and at 7B parameters with 32 layers, that's roughly 1GB of VRAM per 1000 tokens. Serve a 32k context window and you're looking at 32GB just for the cache, before you even load the model weights.

Google's TurboQuant paper (ICLR 2026) promised a solution: 5x compression at 3-bit quantization while maintaining 99.5% attention fidelity through clever use of randomized orthogonal transforms and residual correction. The tonbistudio/turboquant-pytorch repository set out to implement this from scratch in PyTorch—and in doing so, uncovered something fascinating. The paper's key innovation, the QJL (Quantized Johnson-Lindenstrauss) residual correction stage, doesn't just fail to help: it actively hurts performance. This implementation's journey from attempting faithful reproduction to discovering what actually works offers more value than most successful replications ever could.

Technical Insight

TurboQuant's core architecture uses a two-stage vector quantization pipeline. First, it multiplies each KV vector by a Hadamard transform matrix—a fast, deterministic orthogonal transformation that normalizes the distribution of vector coefficients. This normalization is crucial because it makes quantization errors more uniform across dimensions. Then it applies Lloyd-Max optimal scalar quantization to each coordinate independently, storing only the quantized indices plus a single scalar norm for reconstruction.

Here's what the basic quantization flow looks like in practice:

def quantize_kv(kv_tensor, num_bits=3):
    # kv_tensor shape: [batch, num_heads, seq_len, head_dim]
    
    # Step 1: Apply Hadamard transform for normalization
    hadamard_matrix = generate_hadamard(kv_tensor.shape[-1])
    transformed = torch.matmul(kv_tensor, hadamard_matrix)
    
    # Step 2: Extract and normalize by L2 norm
    norms = torch.norm(transformed, dim=-1, keepdim=True)
    normalized = transformed / (norms + 1e-8)
    
    # Step 3: Scalar quantization per coordinate
    num_levels = 2 ** num_bits
    quantization_levels = torch.linspace(-1, 1, num_levels)
    
    # Find nearest quantization level for each element
    distances = torch.abs(normalized.unsqueeze(-1) - quantization_levels)
    indices = torch.argmin(distances, dim=-1)
    
    # Step 4: Bit-pack indices for true compression
    packed = pack_bits(indices, num_bits)  # Custom kernel
    
    return packed, norms, quantization_levels

The V3 implementation introduces three critical improvements over the paper. First, it removes the QJL residual correction entirely. The paper's approach added carefully crafted unbiased noise to bound quantization error, but this noise gets exponentially amplified through the softmax operation in attention. Six independent implementations confirmed this failure mode—the theoretical guarantees don't survive contact with softmax's exponential sensitivity.

Second, asymmetric bit allocation. The repository discovered that keys and values play fundamentally different roles in attention. Keys determine which tokens attend to which (precision-critical), while values get averaged together (errors cancel out). Allocating 4 bits to keys and 2 bits to values at a 3-bit average budget outperforms uniform 3-bit allocation:

def asymmetric_quantize(keys, values, k_bits=4, v_bits=2):
    # Keys need precision for attention pattern accuracy
    quantized_keys, k_norms, k_levels = quantize_kv(keys, num_bits=k_bits)
    
    # Values can tolerate more aggressive compression
    quantized_values, v_norms, v_levels = quantize_kv(values, num_bits=v_bits)
    
    # Average compression: (4 + 2) / 2 = 3 bits, but better quality
    return quantized_keys, quantized_values, (k_norms, k_levels), (v_norms, v_levels)

Third, layer-adaptive precision protection. Not all layers tolerate compression equally. The implementation adds configurable precision scaling for early and late layers, which handle initial token embeddings and final output logits respectively:

layer_bits = {
    0: 6,  # First layer: minimal compression
    1: 5,
    # ... middle layers use configured bits ...
    30: 5,
    31: 6  # Last layer: minimal compression
}

The repository also implements true bit-packed storage, which many "compression" implementations skip. Without packing, you might store 3-bit quantized values in 32-bit floats, actually increasing memory usage. The V2 implementation fell into this trap, ending up 38% larger than uncompressed caches. V3 uses custom CUDA kernels to pack quantization indices into contiguous bit arrays, achieving actual 5x compression ratios in memory footprint.

But here's the brutal reality the benchmarks reveal: attention score similarity is a lying metric. At 5x compression (K3/V3), you get 99.6% attention score similarity but completely broken text generation. The model outputs gibberish. Reliable generation requires backing down to approximately 2x compression (K6/V4 with 128-token windows), far below the paper's advertised performance. The repository's honest benchmarking—comparing 18 different prompts and documenting that only modest compression works—provides more practical value than the paper's optimistic claims.

Gotcha

This implementation has three significant limitations that potential users need to understand upfront. First, there's a massive gap between theoretical metrics and practical performance. The repository achieves 99.5%+ attention score similarity at high compression ratios, which sounds impressive until you actually generate text and get nonsense. The only compression level that reliably produces coherent output is around 2x (K6/V4), not the advertised 5x. This isn't a bug in the implementation—it's a fundamental insight that attention score similarity doesn't capture the cascading error effects through multiple transformer layers during autoregressive generation.

Second, portability is severely limited. The implementation requires CUDA-capable NVIDIA GPUs and has only been tested on Windows with RTX hardware. The bit-packing kernels and Hadamard transforms are CUDA-specific, and there's no CPU fallback or support for other accelerators like AMD ROCm or Apple Metal. If you're not running Windows with an NVIDIA GPU, you'll need to rewrite significant portions.

Third, the repository's development history reveals concerning testing gaps. The README openly acknowledges a critical bug in earlier validation where setting residual_window=0 accidentally disabled compression entirely, invalidating initial claims of 18/18 perfect generation scores. While the transparency is commendable, it raises questions about how many other edge cases might not have been thoroughly tested. The honest acknowledgment of failures (both the QJL approach and the validation bug) is actually the repository's greatest strength—but it also means you're working with an educational implementation, not battle-tested production code.

Verdict

Use if: You're researching KV cache compression techniques and want a clean, well-documented reference implementation with honest benchmarks about what doesn't work. The repository's value lies in its transparency about failures—both the theoretical QJL approach and practical validation challenges—which is rare and valuable in academic reproductions. It's also useful if you're learning about quantization techniques and want to understand the gap between attention metrics and generation quality. Skip if: You need production-ready compression for serving LLMs at scale. The practical 2x compression ratio doesn't justify the CUDA complexity and platform restrictions when alternatives like FlashAttention-2's kernel optimizations or vLLM's prefix caching deliver comparable or better performance with wider hardware support. Also skip if you need the advertised 5x compression—this implementation conclusively demonstrates those numbers don't hold for real text generation, regardless of how good the attention similarity scores look.