💬 NLP

SPHERICAL KV: Angle-Domain Attention and Rate-Distortion Retention for Efficient Long-Context Inference

Spherical KV is an efficient long-context inference method that addresses HBM bandwidth bottlenecks by combining Angle-Domain Attention, which computes attention logits directly from compact spherical key representations without dense reconstruction, and Rate-Distortion Retention, which optimally allocates token retention and precision under a fixed memory budget.

Original authors: Anay Chauhan, Gurucharan Marthi Krishna Kumar, Arion Das, Amit Dhanda, Vinija Jain, Aman Chadha, Amitava Das

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Anay Chauhan, Gurucharan Marthi Krishna Kumar, Arion Das, Amit Dhanda, Vinija Jain, Aman Chadha, Amitava Das

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to read a massive, 100,000-page book to answer a single question. To do this, your brain (the AI) needs to keep a "cheat sheet" of everything it has read so far so it doesn't forget the beginning of the story while reading the end.

In AI terms, this cheat sheet is called the KV Cache.

The problem, as this paper explains, is that as the book gets longer, this cheat sheet gets so huge that it doesn't fit in your brain's short-term memory (the GPU's fast memory). Even if it fits, your brain spends more time flipping through the pages of this massive cheat sheet than it does actually thinking about the answer. The bottleneck isn't your thinking speed; it's the traffic jam of moving pages back and forth.

Existing solutions try to fix this by either:

Throwing away pages: Deleting old parts of the story (which risks forgetting important details).
Writing in tiny handwriting: Compressing the text (which usually requires re-writing it in big letters just to read it, wasting time).

Spherical KV is a new method that changes how the cheat sheet is written and how it is read, without needing to rewrite the text in big letters.

Here is how it works, using simple analogies:

1. The "Direction vs. Distance" Trick (Angle-Domain Attention)

Imagine you are giving someone directions to a coffee shop.

Old Way: You write down the exact coordinates (Latitude, Longitude, Altitude) for every single step of the journey. To find the shop, the reader has to look up the full 3D coordinates, do a complex math calculation, and then figure out the direction.
Spherical KV Way: You realize that for finding the shop, the distance (how far away it is) and the direction (which way to look) are the only things that matter.
- You write down the distance as a simple number (e.g., "5 miles").
- You write down the direction as a compact code (e.g., "North-North-East").
- The Magic: When the reader needs to find the shop, they don't need to reconstruct the full 3D map. They can look at the "North-North-East" code and the "5 miles" number and instantly know the answer. They skip the heavy math of rebuilding the full map.

In the paper, this is called Angle-Domain Attention. It stores keys (the "directions" of the text) as a radius and an angle. The AI can calculate the "relevance" of a word directly from these codes without ever rebuilding the full, heavy data. This saves a massive amount of time moving data around.

2. The "Smart Budget" (Rate-Distortion Retention)

Imagine you have a limited amount of space in your backpack for your trip. You can't carry everything.

Old Way: You might just throw away the oldest items, or you might compress everything equally (making everything slightly fuzzy).
Spherical KV Way: You act like a smart travel agent. You look at your items and ask: "How important is this?"
- Critical items: "The map to the hospital" or "The list of allergies." You keep these in high-quality, uncompressed detail.
- Important items: "The name of the hotel." You keep these, but maybe in a slightly smaller font.
- Low-value items: "The color of the sky on Tuesday." You might throw these away entirely or compress them heavily.

This is called Rate-Distortion Retention. It doesn't just decide what to keep; it decides how clearly to keep it. It spends its "memory budget" on the parts of the story that matter most for the current question, ensuring the AI doesn't get confused by fuzzy details when it needs to be precise.

The Result

By combining these two ideas, Spherical KV creates a system where:

The cheat sheet takes up less space (because it's smarter about what to keep).
The AI reads the cheat sheet faster (because it doesn't have to rebuild the full text to understand it).

The Paper's Findings:
When the authors tested this on long stories (up to 128,000 words), they found that:

The AI could generate text 1.5 to 1.7 times faster than before.
It used 24% to 42% less memory to do the same job.
Crucially, the quality of the answers didn't drop. The AI didn't start hallucinating or forgetting the plot; it just got much more efficient at managing its memory.

In Summary:
Think of Spherical KV as upgrading from a library where you have to walk to the back, pull out a giant encyclopedia, copy the whole page, and then read it, to a library where you just get a tiny, smart index card that tells you exactly what you need, instantly. It solves the "traffic jam" of long conversations by making the memory smaller and the reading process smarter.

Technical Summary: Spherical KV

1. Problem Statement

Long-context inference in Large Language Models (LLMs) is increasingly constrained by the Key-Value (KV) cache rather than arithmetic throughput. As context length ( $T$ ) grows, the resident memory footprint scales linearly with $T$ , batch size ( $B$ ), layers ( $L$ ), and heads ( $H$ ). In production serving environments utilizing paged memory (e.g., PagedAttention), the dominant bottleneck shifts from Floating Point Operations (FLOPs) to High Bandwidth Memory (HBM) traffic. Specifically, decoding becomes limited by the repeated streaming of historical KV pages from device memory to compute units.

Existing solutions attempt to mitigate this via three levers:

Retention/Eviction: Storing fewer tokens (e.g., StreamingLLM, H2O).
Quantization/Compression: Storing fewer bits (e.g., KVQuant, KIVI).
Offloading: Storing KV elsewhere (e.g., CPU/Storage).

However, these methods often fail to fully address the critical-path bottleneck. Many compression techniques require reconstructing dense vectors from compressed representations before computing attention logits (a "decompress-then-dot-product" path). This reconstruction reintroduces HBM traffic and prevents kernel fusion, negating theoretical memory savings in realistic, paged serving substrates. Furthermore, simple eviction policies often lack the granularity to preserve critical information (e.g., retrieved evidence) while discarding low-utility states.

2. Methodology: Spherical KV

The paper proposes Spherical KV, a long-context inference method that treats KV allocation as a rate–distortion problem grounded in attention geometry. It consists of two orthogonal components designed to work within a paged, ragged serving substrate: Angle-Domain Attention (ADA) and Rate–Distortion Retention (RDR).

2.1 Angle-Domain Attention (ADA)

ADA replaces the standard "decompress-then-dot-product" primitive with a compute-native approach that operates directly on compact spherical codes.

Spherical Parameterization: Keys are decomposed into a scalar radius ( $r = \|k\|$ ) and directional angles ( $\phi = \Phi(\hat{k})$ ).
Compressed-Domain Similarity: Instead of caching dense vectors, ADA stores quantized radii ( $\tilde{r}$ ) and packed angular codes ( $c_\theta$ ). During the decode hot loop, the attention kernel computes logits directly from these codes:
$\hat{\ell}(q, k) \approx \|q\| \cdot \tilde{r} \cdot f_b(c_\theta^q, c_\theta^k)$
where $f_b$ is a fused kernel operation computing cosine similarity from the angular codes.
No Reconstruction Invariant: The critical design constraint is that the decode kernel never materializes a dense key vector ( $\tilde{k} \in \mathbb{R}^d$ ) in HBM. This ensures that reductions in stored KV bytes translate directly into reduced HBM traffic per token.

2.2 Rate–Distortion Retention (RDR)

RDR is a control mechanism that jointly optimizes which states to retain and at what precision (tier) to store them, subject to a strict memory budget.

Joint Optimization: Unlike methods that treat eviction and quantization separately, RDR solves a discrete rate–distortion allocation problem. For each token/head/layer state $i$ , it selects a binary residency decision ( $z_i \in \{0, 1\}$ ) and a precision tier ( $b_i \in \mathcal{B}$ ).
Distortion Surrogate: The controller uses a mechanistic distortion proxy derived from ADA's error analysis. It estimates that logit error is dominated by angular approximation errors ( $\epsilon_\theta$ ) scaled by the query and key magnitudes.
Protection Constraints: To prevent catastrophic failure on critical information (e.g., retrieved evidence, safety spans, outliers), RDR enforces hard protection constraints ( $I_{prot}(i)=1 \implies z_i=1, b_i=b_{max}$ ).
Tier-Homogeneous Pages: The controller groups retained states into pages where all items share the same precision tier. This enables single-pass, header-guided streaming reads by the decode kernel, avoiding irregular gathers.

2.3 System Integration

Spherical KV operates on a standard paged/ragged KV substrate (compatible with vLLM-style layouts). The pipeline maps dense keys to spherical codes, applies RDR to filter and tier them, and stores the result in a compressed format. The decode kernel consumes these codes directly, maintaining a "no-dense-reconstruction" invariant.

3. Key Contributions

Angle-Domain Attention (ADA): A kernel-level primitive that computes attention logits directly from compact spherical codes, eliminating the HBM traffic penalty associated with dense key reconstruction.
Rate–Distortion Retention (RDR): A unified controller that jointly allocates residency and precision tiers under a fixed budget, prioritizing high-utility states and protecting brittle ones.
Deployment-Oriented Evaluation: The method is evaluated not on toy contiguous caches but on a realistic paged/ragged serving substrate, measuring end-to-end throughput, HBM traffic, and behavioral stability.
Pareto Frontier Shift: The paper demonstrates that Spherical KV shifts the memory–quality–throughput frontier, achieving higher throughput at lower memory footprints compared to dense KV and strong baselines.

4. Experimental Results

The authors evaluated Spherical KV on Llama-3.1-8B-Instruct, Qwen2.5-14B-Instruct, and gpt-oss-20b across context lengths of 8K, 32K, and 128K. Workloads included language modeling (PG-19), retrieval-heavy QA (LongBench: HotpotQA, 2WikiMultiHopQA), and agentic tool rollouts (AgentBench, ToolBench).

Throughput and Memory Efficiency: Under matched quality constraints (iso-quality), Spherical KV improved decode throughput by 1.55×–1.72× while reducing resident KV bytes per token by 24–42%.
Context Scaling: The gains were most pronounced in the 128K regime, where KV residency and HBM traffic dominate system performance.
Kernel-Realized Efficiency: The method successfully reduced HBM bytes per generated token, confirming that the memory savings were not offset by reconstruction overheads.
Stability: The system demonstrated smooth degradation under tighter budgets and preserved reasoning trajectories and tool-call sequences in agentic workloads, avoiding the abrupt failures seen in hard-eviction baselines.

5. Significance and Claims

The paper claims that Spherical KV addresses the fundamental bottleneck of long-context inference: the conflict between memory capacity and memory bandwidth. By treating KV compression as a compute-native problem (where the compressed representation is directly executable by the kernel) rather than just a storage problem, the method achieves a genuine shift in the operating frontier.

The authors emphasize that:

Compression must be compute-native: Merely storing fewer bits is insufficient if the decode path requires dense reconstruction.
Memory must be budgeted where it matters: Rate–distortion control allows the system to spend precision on high-impact states while saving on low-utility ones.
Evaluation must be deployment-faithful: Gains must survive realistic paged layouts, metadata overheads, and kernel constraints to be considered real.

Spherical KV is presented as a mechanism for reducing KV residency and HBM traffic while preserving decode efficiency and behavioral stability, specifically targeting the memory-bound regimes of modern LLM serving.