Back to AI Research

AI Research

DepthWeave-KV: Token-Adaptive Cross-Layer Residual... | AI Research

Key Takeaways

  • DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression Long-context language models often struggle with the "...
  • A fused CUDA implementation jointly performs basis lookup, residual dequantization, and attention projection to reduce decode-time memory traffic.
  • DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression
  • Long-context language models often struggle with the "KV cache" bottleneck, where storing the massive amount of data required to process long documents consumes too much memory and slows down performance.
  • Instead of storing a full, independent cache for every layer, DepthWeave-KV "weaves" a shared foundation across neighboring layers.
Paper AbstractExpand

Long-context language model inference is increasingly limited by the memory bandwidth and capacity required to store key-value caches, yet existing compression methods often apply uniform budgets across layers or tokens and degrade retrieval when lexical cues and semantic states require different preservation. We introduce DepthWeave-KV, a token-adaptive cache compression method that factorizes key and value states across neighboring transformer layers using shared low-rank channel bases while retaining lightweight token-specific residuals where attention behavior is sensitive. DepthWeave-KV combines cross-depth residual factorization with a token-conditional depth router that allocates higher reconstruction rank to instruction-bearing and retrieval-critical tokens, and uses calibration-free online error tracking from attention-output probes to adapt compression during generation without retraining the base model. A fused CUDA implementation jointly performs basis lookup, residual dequantization, and attention projection to reduce decode-time memory traffic. Across LongBench, Needle-in-a-Haystack, L-Eval, and long-form QA and summarization benchmarks, DepthWeave-KV achieves near-full-cache task quality with substantially lower memory use, improving average score and retrieval accuracy over prior compressed caches while reaching 8.3x KV memory reduction and 72.8 tokens per second at 64K context.

DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression
Long-context language models often struggle with the "KV cache" bottleneck, where storing the massive amount of data required to process long documents consumes too much memory and slows down performance. While many existing methods try to compress this cache by deleting tokens or sharing data across layers, they often do so uniformly, which can lead to "brittle" performance where the model loses critical information needed for retrieval or instruction following. DepthWeave-KV addresses this by using a smarter, adaptive approach that treats different tokens and layers based on their actual importance, allowing the model to maintain high accuracy while significantly reducing memory usage.

How the Approach Works

Instead of storing a full, independent cache for every layer, DepthWeave-KV "weaves" a shared foundation across neighboring layers. It uses a compact set of shared low-rank bases to represent the common information found in adjacent transformer layers.
To ensure accuracy, the system adds "token-specific residuals"—small, lightweight pieces of data that act as corrections. A "token-conditional depth router" decides how much of this residual data to keep for each token. For example, a standard word in a long document might be represented mostly by the shared base, while an instruction-bearing or retrieval-critical token is assigned a higher "rank," meaning it gets a more detailed, high-fidelity residual to ensure the model doesn't lose track of it.

Online Adaptation Without Retraining

A key innovation of DepthWeave-KV is its ability to adjust compression on the fly without needing to retrain the base model. It uses "attention-output probes" during the generation process. Every few steps, the system briefly checks the difference between the compressed cache and a high-fidelity reference. If the error exceeds a certain threshold, the system automatically increases the residual rank for that specific window of tokens. This feedback loop allows the model to stay efficient during simple tasks but become more precise when the prompt requires complex retrieval.

Performance and Efficiency

DepthWeave-KV was tested across a wide range of benchmarks, including Needle-in-a-Haystack, LongBench, and various summarization tasks. The results show that it achieves near-full-cache quality while reducing memory usage by 8.3x.
Beyond memory savings, the method improves system speed. By using a custom-fused CUDA implementation, the system performs basis lookup, residual dequantization, and attention projection in a single step. This reduces the amount of data moving between memory and the processor, resulting in a decode speed of 72.8 tokens per second at a 64K context length—a significant improvement over previous compression techniques.

Why It Matters

The primary takeaway from this research is that not all cache data is created equal. By moving away from uniform compression and toward a system that dynamically allocates resources based on token importance and online error tracking, DepthWeave-KV avoids the common pitfalls of previous methods. It provides a robust solution for long-context inference that balances the need for extreme memory efficiency with the requirement for high-quality, reliable model responses.

Comments (0)

No comments yet

Be the first to share your thoughts!