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

2026-05-13

AI Summary

This work addresses the memory bottleneck in long-context inference caused by the linear growth of KV cache, which severely limits decoding efficiency. To overcome this challenge, the authors propose a joint optimization framework combining Angular-Domain Attention (ADA) and Rate-Distortion Retention (RDR). By parameterizing key vectors on a hypersphere, ADA computes attention scores directly in the angular domain, eliminating the need for dense reconstruction. RDR dynamically allocates precision and storage resources based on estimated future utility, enabling block-fusion-friendly computation. Under a fixed memory budget, the proposed method substantially reduces KV cache footprint and HBM traffic while preserving high decoding quality, achieving superior long-context inference performance in real-world deployment scenarios.

Abstract

Long-context inference is increasingly constrained by the KV cache: resident memory grows with context length, and decoding becomes limited by repeated High Bandwidth Memory (HBM) streaming rather than arithmetic. Existing methods such as eviction, windowing, quantization, and offloading reduce footprint, but often leave the critical-path bottleneck only partially addressed, especially when compressed states must still be reconstructed into dense vectors during decoding. We present Spherical KV, a long-context inference method that treats KV allocation as a rate-distortion problem grounded in attention geometry for efficient decoding. The method is built on two ideas: (i) represent directional information cheaply in the decode hot loop, and (ii) allocate retention and precision according to estimated future utility. Its first component, Angle-Domain Attention (ADA), stores keys in a spherical parameterization consisting of a scalar radius and compact angle codes, and computes attention logits directly from these codes without reconstructing dense keys. This preserves a paged, block-local, fusion-friendly decode path and directly targets HBM traffic in realistic serving settings. Its second component, Rate-Distortion Retention (RDR), jointly chooses keep/drop decisions and precision tiers per token and head under a fixed budget, producing tier-homogeneous pages with lightweight metadata and coalesced reads. Together, ADA and RDR provide a deployment-oriented mechanism for reducing KV residency while preserving decode efficiency.

Problem

• KV cache
• long-context inference
• HBM bandwidth
• rate-distortion
• attention mechanism

Innovation

• Spherical KV
• Angle-Domain Attention
• Rate-Distortion Retention
• KV Cache Compression
• Long-Context Inference