长上下文语言模型推理越来越受到存储键值缓存,所需的内存带宽和容量的限制,但现有的压缩方法通常跨层或令牌应用统一的预算,并且当词汇提示和语义状态需要不同的保存时,会降低检索性能。我们引入了 DepthWeave-KV, 一种令牌自适应缓存压缩方法,该方法使用共享的低秩通道基来分解相邻变压器层的键和值状态,同时保留注意力行为敏感的轻量级特定于令牌的残差。 DepthWeave-KV 将交叉深度残差分解与令牌条件深度路由器相结合,为指令承载和检索关键令牌, 分配更高的重建等级,并使用来自注意力输出探针的免校准在线错误跟踪来适应生成过程中的压缩,而无需重新训练基本模型。融合 CUDA 实现联合执行基础查找, 残差反量化, 和注意力投影,以减少解码时内存流量。在 LongBench, Needle-in-a-Haystack, L-Eval, 以及长格式 QA 和摘要基准测试中,, DepthWeave-KV 实现了接近满缓存的任务质量,同时显着降低了内存使用, 与之前的压缩缓存相比,提高了平均得分和检索精度,同时在 64K 上下文中实现了 8.3 倍的 KV 内存减少和每秒 72.8 个令牌。

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.

科目:人工智能(cs.AI)

Subjects: Artificial Intelligence (cs.AI)