RaBitQCache: Rotated Binary Quantization for KVCache in Long Context LLM Inference

Abstract: Long-context LLM inference is severely bottlenecked by the massive Key-Value (KV) cache, yet existing sparse attention methods often suffer from static fixed-budget (Top-k) retrieval or rely on proxy scores that are computationally expensive and biased. To address these limitations, we propose RaBitQCache, a novel sparse attention framework that utilizes randomized rotated binary quantization and high-throughput binary-INT4 arithmetic to efficiently estimate attention weights. Our proxy score serves as an unbiased estimator with a proven error bound, enabling adaptive Top-p retrieval that dynamically adjusts the token budget based on actual attention sparsity. We further implement a hardware-aware system with asynchronous pipelining and lazy updates to mask overhead. Evaluations demonstrate that RaBitQCache significantly accelerates inference and reduces memory I/O while preserving generation quality compared to state-of-the-art baselines. Code is available at https://github.com/Sakuraaa0/RaBitQCache.git.

• The paper proposes a novel rotated binary quantization framework that yields an unbiased estimator for sparse attention with an error of O(1/√D).
• The methodology integrates a prefill phase with sign-based quantization and a decode phase using INT4 quantization paired with custom CUDA kernels.
• Empirical results demonstrate up to a 3.88x decoding speedup, minimal memory overhead, and maintained generation quality across various benchmarks.

RaBitQCache: Rotated Binary Quantization for Efficient Sparse Attention in Long-Context LLM Inference

Motivation and Background

With the escalating operational demands for ever-increasing context lengths in LLMs, efficient management of the key-value (KV) cache during inference is a significant bottleneck in both memory and computation. Existing solutions largely fall into static token selection, which cannot adapt to the input's attention structure, or dynamic methods that employ computationally expensive or inaccurate proxy scores without theoretical error guarantees. These limitations hinder both efficiency and precision, particularly when attention patterns are highly variable across layers and heads.

Methodology: Rotated Binary Quantization Framework

RaBitQCache introduces an attention sparsification scheme grounded in randomized rotated binary quantization of KV caches, drawing inspiration from the Johnson-Lindenstrauss (JL) lemma for dimensionality reduction with theoretical guarantees. The process operates in two core phases:

• Prefill Phase:
  • Keys are centered, normalized, and then rotated using a random orthogonal matrix.
  • Each rotated key vector is quantized to its sign, yielding a D-bit binary code by projecting onto the nearest vertex of the unit hypercube.
  • For each token, a scalar correction factor is precomputed to correct the estimator's bias, enabling subsequent efficient computation.
• Decode Phase:
  • Each query is similarly centered, rotated, and quantized using uniform scalar quantization to INT4.
  • The attention score proxy is computed using an efficient INT4-by-binary dot product, leveraging custom CUDA kernel implementations.
  • A theoretically unbiased estimator, with error $O(1/\sqrt{D})$, reconstructs the normalized similarity between the query and each key.

Crucially, this proxy score is not merely an order-statistic but an unbiased estimator of the true attention score, allowing adaptive Top-p token selection—the minimum set whose softmaxed attention probabilities sum to probability $p$. This contrasts with Top-k or fixed-budget approaches that are insensitive to true sparsity and often require conservative budget allocation, leading to inefficiency or quality degradation.

System and Kernel Optimizations

RaBitQCache's efficacy is determined not just by its theoretical underpinnings but also by its hardware-conscious system design:

• Asynchronous, Pipelined Prefill: Binary index construction is offloaded to a low-priority CUDA stream concurrent with dense prefill attention, fully masking the quantization latency.
• Lazy Batched Decode Updates: Recent tokens are retained in a full-precision local buffer and quantized in batches, further amortizing launch overhead and safeguarding generation quality.
• Efficient Memory Layout: The binary key index is densely packed, reducing memory bandwidth; queries are cached in shared (on-chip) memory for repeated access.
• Custom Token Selection: An optimized linear-time Top-p selection kernel enables dynamic, context-aware retrieval.
• INT4-by-Binary GEMV: The scoring operation leverages efficient vectorized and bitwise instructions for maximal GPU throughput.

Empirical Results

Rigorous evaluations demonstrate that RaBitQCache yields marked improvements in both efficiency and retrieval fidelity:

• Efficiency: Achieves up to 3.88x decoding speedup and overall 2.16x end-to-end acceleration over highly optimized full-precision attention baselines at context lengths up to 30K tokens.
• Memory Footprint: The hardware-resident index incurs minimal (≈3.5%) additional memory overhead relative to standard FP16 KV cache.
• Quality Preservation: Maintains generation scores virtually indistinguishable from the full attention baseline—even surpassing it in several tasks—while retrieving only ~17% of tokens, with recall rates near upper-bound or oracle approaches.
• Adaptive Sparsity: The Top-p mechanism dynamically adjusts retrieval budgets per query, automatically exploiting the inherent variable sparsity of attention maps without manual tuning.
• Generalization: Demonstrates robust performance across model capacities (from 7B to 70B parameters) and diverse benchmarks (LongBench, RULER, GSM8K), consistently outperforming SOTA baselines like Quest, SparQ, DoubleSparsity, and quantization/eviction-based approaches. Notably, in several settings, adaptive sparsification via RaBitQCache not only reduces compute but also serves as a mild regularizer, occasionally yielding better downstream accuracy than dense attention.

Theoretical and Practical Implications

RaBitQCache's principal innovations are twofold:

1. Unbiased Score Estimation with Provable Error Bounds: By leveraging the JL lemma for randomized projections and quantization, the framework constructs unbiased estimators of the attention score, enabling adaptive thresholding—this directly enables Top-p selection and guarantees precise control over sparsity-induced approximation error in attention.
2. Tight Algorithm-hardware Co-design: Through dense memory formats, fused kernel design, and asynchronicity, RaBitQCache translates theoretical reductions to practical wall-clock improvements on modern GPUs, achieving genuine decoupling of inference latency from context length.

Practically, this enables long-context LLM inference with dramatically reduced hardware requirements—critical for democratization of long-range language modeling and Green AI objectives. Future work could explore orthogonal integration with lossy KV cache compression methods (e.g., KIVI), more adaptive error-bound calibration for non-uniformly distributed activation spaces, and investigating the implications of estimator regularization in downstream tasks and robustness.

Conclusion

RaBitQCache establishes a new standard for sparse attention in long-context LLM inference, unifying principled randomized quantization, unbiased score estimation, and low-level hardware optimizations. It achieves substantial efficiency gains with virtually no degradation in generation quality and supports true adaptive retrieval, a capability not available in prior fixed-budget systems. This framework opens promising avenues for efficient, scalable long-context deployment and further theoretical investigations into randomized quantization for high-dimensional attention mechanisms.