\documentclass{article}

% NeurIPS/ICML workshop style
\usepackage[final]{neurips_2024}

\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{hyperref}
\usepackage{url}
\usepackage{booktabs}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{graphicx}
\usepackage{microtype}
\usepackage{algorithm}
\usepackage{algorithmic}
\usepackage{xcolor}

\title{TurboQuant Pro: Open-Source TurboQuant Implementation \\
for LLM KV Cache Compression on Consumer GPUs}

\author{
  Andrew H.\ Bond \\
  Department of Computer Engineering \\
  San Jose State University \\
  \texttt{andrew.bond@sjsu.edu}
}

\begin{document}

\maketitle

\begin{abstract}
We present TurboQuant Pro, the first open-source implementation of the
TurboQuant algorithm (Zandieh et al., ICLR 2026) for compressing
the key-value (KV) cache in large language model inference.  Our
implementation achieves 5.12$\times$ compression at 0.978 mean cosine
similarity using 3-bit quantization with bit-packing, and includes a
novel streaming two-tier cache with L1 (hot, uncompressed) and L2
(cold, bit-packed) storage for autoregressive generation.  The system
runs on consumer Volta-class GPUs via CuPy CUDA kernels and falls back
to vectorized NumPy on CPU.  We benchmark compression quality,
throughput, and memory savings across multiple model configurations
including Llama~3.1, Gemma~4, and Mistral, and provide integration
examples for llama.cpp and vLLM.  The package is released under the
MIT license at \url{https://github.com/ahb-sjsu/turboquant-pro}.
\end{abstract}

% ------------------------------------------------------------------ %
\section{Introduction}
\label{sec:intro}

The key-value (KV) cache is a critical memory bottleneck in autoregressive
LLM inference.  For a model with $L$ layers, $H$ KV heads, head dimension
$d$, and context length $S$, the KV cache consumes $2LHSd$ elements,
typically stored in \texttt{float16}.  For Gemma~4~27B at 8K context, this
is approximately 1.13\,GB---often exceeding the available VRAM on consumer
GPUs and limiting the maximum context window.

Zandieh et al.~\cite{turboquant} introduced TurboQuant (building on PolarQuant~\cite{polarquant} and QJL~\cite{qjl}), a
theoretically grounded approach that achieves sub-linear memory scaling by:
(1)~applying a random orthogonal rotation to map KV vectors onto the unit
hypersphere where coordinates are approximately i.i.d.\ Gaussian,
(2)~applying an optimal Lloyd-Max scalar quantizer to each coordinate
independently, and (3)~storing $b$-bit indices plus a per-vector L2 norm.

Despite the theoretical impact of this work, no open-source implementation
existed prior to ours.  We contribute:
\begin{itemize}
  \item A production-quality Python implementation with CuPy GPU kernels
    and NumPy CPU fallback, supporting 2-, 3-, and 4-bit quantization with
    bit-packing.
  \item A streaming two-tier KV cache (L1 hot / L2 cold) designed for
    autoregressive generation, which transparently compresses older tokens.
  \item Comprehensive benchmarks across model configurations, bit widths,
    and context lengths.
  \item Integration examples for llama.cpp, vLLM, and HuggingFace
    Transformers.
\end{itemize}

% ------------------------------------------------------------------ %
\section{Method}
\label{sec:method}

We follow the algorithm of Zandieh et al.~\cite{turboquant}
with the following concrete implementation choices.

\paragraph{Random Rotation.}
For head dimensions $d \leq 4096$, we generate a random orthogonal matrix
$\Pi \in \mathbb{R}^{d \times d}$ via QR factorization of a Gaussian
matrix.  For $d > 4096$, we use a structured rotation (random sign flip +
random permutation) for $O(d)$ memory and application cost.  The rotation
is fixed at initialization and shared across all KV vectors.

\paragraph{Lloyd-Max Quantization.}
After rotation, each coordinate is approximately distributed as
$\mathcal{N}(0, 1/\sqrt{d})$.  We apply precomputed Lloyd-Max centroids
for the standard normal distribution, scaled by $1/\sqrt{d}$.  The
centroids are:
\begin{align}
  \text{2-bit:} &\quad \{-1.510,\; -0.453,\; 0.453,\; 1.510\} \\
  \text{3-bit:} &\quad \{-1.748,\; -1.050,\; \ldots,\; 1.050,\; 1.748\} \\
  \text{4-bit:} &\quad \{-2.401,\; -1.844,\; \ldots,\; 1.844,\; 2.401\}
\end{align}
Quantization is performed via \texttt{searchsorted} on the midpoints
between adjacent centroids.

\paragraph{Bit-Packing.}
We pack $b$-bit indices into bytes: 8 indices $\times$ 3 bits = 3 bytes
for 3-bit, 4 $\times$ 2 = 1 byte for 2-bit, and 2 $\times$ 4 = 1 byte
for 4-bit.  This is the key step that achieves $>5\times$ compression
over fp16 (vs.\ $\sim$2$\times$ with naive uint8 storage).

\paragraph{Decompression.}
Decompression unpacks indices, looks up centroid values, applies the
inverse rotation, and rescales by the stored L2 norm.

\paragraph{Streaming Tiered Cache.}
For autoregressive generation, we implement a two-tier architecture:
\begin{itemize}
  \item \textbf{L1 (hot window)}: The most recent $W$ tokens are stored
    uncompressed for zero-latency attention.
  \item \textbf{L2 (cold storage)}: When the hot window exceeds $W$
    tokens, the oldest half is compressed with bit-packing and moved to
    cold storage.  Cold entries are decompressed on demand.
\end{itemize}
This amortizes compression cost and keeps the attention-critical recent
tokens at full precision.

% ------------------------------------------------------------------ %
\section{Implementation}
\label{sec:impl}

\paragraph{Dual Backend.}
The core \texttt{TurboQuantKV} class accepts a \texttt{use\_gpu} flag.
When True and CuPy is available, all operations (rotation via
\texttt{einsum}, quantization via \texttt{searchsorted}, bit-packing via
CUDA RawKernels) run on GPU.  Otherwise, equivalent NumPy operations
provide a CPU fallback.

\paragraph{CUDA Kernels.}
We provide six CuPy RawKernels for pack/unpack at 2, 3, and 4 bits.
The 3-bit kernel is the most complex: each CUDA thread reads 8 indices,
packs them into a 24-bit value split across 3 bytes, and writes the
result.  All kernels are Volta-compatible (compute 7.0+) and are lazily
compiled on first use.

\paragraph{Memory Layout.}
The \texttt{CompressedKV} dataclass stores:
\begin{itemize}
  \item \texttt{indices}: flat \texttt{uint8} array of bit-packed indices
  \item \texttt{norms}: \texttt{float32} array of shape $(B, H, S)$
  \item Metadata: \texttt{bits}, \texttt{original\_dtype}, \texttt{shape},
    \texttt{n\_values}
\end{itemize}

% ------------------------------------------------------------------ %
\section{Experiments}
\label{sec:experiments}

All experiments use NumPy on CPU with random Gaussian tensors matching
real model dimensions.  We report mean cosine similarity and MSE between
original and reconstructed tensors.

\begin{table}[h]
\centering
\caption{Compression quality for head\_dim=256, n\_heads=16, seq\_len=2048.}
\label{tab:quality}
\begin{tabular}{@{}cccc@{}}
\toprule
Bits & Compression Ratio & Cosine Similarity & MSE \\
\midrule
2 & 7.5$\times$ & 0.926 & 0.001178 \\
3 & 5.1$\times$ & 0.978 & 0.000349 \\
4 & 3.9$\times$ & 0.995 & 0.000082 \\
\bottomrule
\end{tabular}
\end{table}

\begin{table}[h]
\centering
\caption{KV cache memory estimates for popular models (3-bit, packed, 8K context).}
\label{tab:memory}
\begin{tabular}{@{}lcccc@{}}
\toprule
Model & Original (GB) & Compressed (GB) & Saved (GB) & Ratio \\
\midrule
Llama 3.1 8B   & 0.500 & 0.098 & 0.402 & 5.1$\times$ \\
Llama 3.1 70B  & 1.250 & 0.244 & 1.006 & 5.1$\times$ \\
Gemma 4 27B    & 1.125 & 0.220 & 0.905 & 5.1$\times$ \\
Mistral 7B     & 2.000 & 0.391 & 1.609 & 5.1$\times$ \\
\bottomrule
\end{tabular}
\end{table}

\paragraph{Quality vs.\ Bit Width.}
Table~\ref{tab:quality} shows that 3-bit quantization provides an
excellent quality--compression trade-off: 5.1$\times$ compression with
only 2.2\% cosine similarity loss.  4-bit provides near-lossless
reconstruction (0.995 cosine similarity) at 3.9$\times$ compression.
2-bit is suitable for applications where memory is extremely constrained
and moderate quality loss is acceptable.

\paragraph{Scaling with Context Length.}
The compression ratio is determined by the bit width and head dimension,
not the context length.  Memory savings scale linearly with sequence
length: at 32K context, Llama~3.1~8B saves 1.6\,GB of VRAM.

\paragraph{Streaming Cache Overhead.}
The L1/L2 tiered cache adds minimal overhead.  With a hot window of 512
tokens and 4096 total tokens, the effective compression ratio is
approximately 4.2$\times$ (reduced from 5.1$\times$ by the uncompressed
hot window), and the per-token append latency is dominated by the
periodic batch compression of flushed tokens.

% ------------------------------------------------------------------ %
\section{Integration}
\label{sec:integration}

We provide three integration pathways:

\paragraph{llama.cpp / llama-cpp-python.}
The Python wrapper pattern intercepts KV tensors after each forward pass
and stores them in a \texttt{TurboQuantKVCache}.  For production
deployment, the CUDA kernels can be compiled into a C extension callable
from the llama.cpp C++ codebase.

\paragraph{vLLM.}
TurboQuant Pro can compress cold KV pages in vLLM's PagedAttention,
reducing GPU memory pressure and enabling larger batch sizes.

\paragraph{HuggingFace Transformers.}
By subclassing the attention module, users can override the KV cache
update to store compressed tensors.

% ------------------------------------------------------------------ %
\section{Conclusion}
\label{sec:conclusion}

TurboQuant Pro provides the first publicly available implementation of
Zandieh et al.'s TurboQuant algorithm for KV cache compression.
Our implementation achieves 5.1$\times$ memory reduction at 0.978 cosine
similarity with 3-bit quantization, runs on consumer Volta-class GPUs,
and includes a streaming tiered cache for autoregressive generation.

The novelty of this work lies in the engineering: a production-ready
implementation with CUDA kernels, the streaming two-tier cache design,
and integration examples for major inference frameworks.  The algorithm
itself is due entirely to Zandieh et al.

The package is available under the MIT license at
\url{https://github.com/ahb-sjsu/turboquant-pro}.

\paragraph{Limitations.}
Our benchmarks use synthetic Gaussian data; quality on real model
activations (which may not be perfectly Gaussian after rotation) should
be validated per model.  The current implementation does not include
per-channel or per-head adaptive bit widths, which could improve the
quality--compression trade-off.

\bibliography{references}
\bibliographystyle{plain}

\begin{thebibliography}{1}

\bibitem{turboquant}
Amir Zandieh, Majid Daliri, Majid Hadian, and Vahab Mirrokni.
\newblock TurboQuant: Online vector quantization with near-optimal distortion rate.
\newblock In \emph{International Conference on Learning Representations
  (ICLR)}, 2026. arXiv:2504.19874.

\bibitem{polarquant}
Insu Han, Praneeth Kacham, Amin Karbasi, Vahab Mirrokni, and Amir Zandieh.
\newblock PolarQuant: Quantizing KV caches with polar transformation.
\newblock arXiv:2502.02617, 2025.

\bibitem{qjl}
Amir Zandieh, Majid Daliri, and Insu Han.
\newblock QJL: 1-bit quantized JL transform for KV cache quantization with zero
  overhead.
\newblock In \emph{AAAI Conference on Artificial Intelligence}, 2025.
  arXiv:2406.03482.

\bibitem{lloyd1982least}
Stuart Lloyd.
\newblock Least squares quantization in PCM.
\newblock \emph{IEEE Transactions on Information Theory},
  28(2):129--137, 1982.

\bibitem{max1960quantizing}
Joel Max.
\newblock Quantizing for minimum distortion.
\newblock \emph{IRE Transactions on Information Theory},
  6(1):7--12, 1960.

\end{thebibliography}

\end{document}
