DepthWeave-KV: Token-Adaptive Compression for Long-Context LLMs

A new token-adaptive cache compression method, DepthWeave-KV, has been developed to tackle the memory and bandwidth constraints of Key-Value (KV) caches in long-context language model inference. Unlike existing methods that apply uniform compression, DepthWeave-KV intelligently factorizes key and value states across adjacent transformer layers using shared low-rank channel bases, while preserving lightweight, token-specific residuals for critical attention behaviors. The system incorporates a token-conditional depth router that dynamically assigns higher reconstruction rank to important tokens, such as those bearing instructions or critical for retrieval. This adaptive approach, combined with calibration-free online error tracking from attention-output probes, allows the compression policy to adjust during generation without requiring retraining of the base model. A fused CUDA implementation further optimizes performance by jointly handling basis lookup, residual dequantization, and attention projection, thereby reducing decode-time memory traffic. Benchmarks across various long-context tasks, including LongBench and Needle-in-a-Haystack, show that DepthWeave-KV achieves near-full-cache task quality with an 8.3x KV memory reduction and improved throughput, reaching 72.8 tokens per second at 64K context.