Abstract
Long-context language model inference is increasingly limited by the memory bandwidth and capacity required to store key-value caches. Existing compression methods often apply uniform budgets across layers or tokens, degrading 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, using 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.
Blogger's Review: The innovation of DepthWeave-KV lies in its adaptive cache compression strategy, which enhances long-context processing efficiency through intricate inter-layer factorization and token-specific residual handling. This approach not only optimizes memory usage but also maintains task performance, paving the way for new possibilities in future language model inference.