DASC: Decay-Aware State Compression for Hybrid Linear-Attention Serving
作者: Yanqi Yu, Pingwei Sun, Jianchao Tan, Tao Zhang, Yuchen Xie, Xunliang Cai, Yao Liu
分类: cs.LG, cs.AI
发布日期: 2026-08-31
💡 一句话要点
提出DASC以解决混合线性注意力模型的状态压缩问题
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 混合线性注意力 状态压缩 模型优化 推理效率 深度学习
📋 核心要点
- 现有混合线性注意力模型在状态更新时面临缓存管理复杂化的问题,导致内存压力和性能下降。
- 论文提出的DASC方法通过分析模型权重推导保留视域,选择长视域状态单元并优化状态检查点布局。
- 实验结果显示,DASC在固定状态检查点内存预算下,平均首次令牌时间减少42.6%,输入吞吐量提升68.4%。
📝 摘要(中文)
混合线性注意力架构最近已扩展到大型开放权重模型,提供与全注意力相竞争的质量,同时显著减少键/值(KV)缓存增长。然而,其就地递归状态更新使得缓存管理复杂化:前缀重用需要状态检查点与全注意力KV并存,而在全注意力中存储状态检查点则增加了内存压力,导致更多的驱逐和重复预填。通过分析Gated DeltaNet(GDN)和Kimi Delta Attention(KDA)的衰减结构,发现不同头和通道在不同时间尺度上保留前缀信息,称为“保留视域”。基于这一观察,提出了“衰减感知状态压缩”(DASC),从模型权重中推导保留视域,选择长视域状态单元,并将其打包成不规则状态检查点布局。DASC在重用时要么用零填充省略的单元,要么从有界后缀刷新它们,带来额外的计算成本。实验表明,DASC在压缩KDA递归状态检查点的同时,保持接近全缓存的性能。
🔬 方法详解
问题定义:当前混合线性注意力模型在进行状态更新时,缓存管理变得复杂,导致内存压力增加和性能下降,尤其是在需要前缀重用的情况下。
核心思路:DASC通过分析模型的衰减结构,识别不同头和通道的保留视域,从而选择合适的状态单元进行压缩,优化缓存使用效率。
技术框架:DASC的整体架构包括状态单元的选择、压缩布局的设计以及与张量并行推理引擎的集成,确保高效的状态管理和计算。
关键创新:DASC的主要创新在于其衰减感知的状态压缩方法,通过识别和利用保留视域,显著提高了状态检查点的压缩效率,与传统方法相比具有本质区别。
关键设计:DASC在参数设置上采用了长视域状态单元的选择策略,并在压缩布局中实现了不规则状态检查点设计,同时在重用时采用零填充或后缀刷新策略,平衡了计算成本与性能。
🖼️ 关键图片
📊 实验亮点
实验结果表明,DASC在压缩KDA递归状态检查点时,能够实现2.63倍的压缩比,同时在固定状态检查点内存预算下,平均首次令牌时间减少42.6%,输入吞吐量提升68.4%。
🎯 应用场景
DASC方法在自然语言处理、图像识别等领域具有广泛的应用潜力,尤其是在需要高效状态管理和快速推理的场景中。其优化的状态压缩技术可以有效提升模型的推理速度和资源利用率,推动大规模模型的实际应用。
📄 摘要(原文)
Hybrid linear-attention architectures have recently scaled to large open-weight models, offering quality competitive with full attention while substantially reducing key/value (KV) cache growth. However, their in-place recurrent-state updates complicate cache management: prefix reuse requires state checkpoints alongside full-attention KV, while storing state checkpoints in full increases memory pressure, leading to more evictions and repeated prefill. By analyzing the decay structure of Gated DeltaNet (GDN) and Kimi Delta Attention (KDA), we find that different heads and channels retain prefix information over markedly different timescales, which we term \emph{retention horizons}. This variation suggests substantial compression potential in persistent state checkpoints. Building on this observation, we introduce \emph{Decay-Aware State Compression} (DASC), which derives retention horizons from model weights, selects long-horizon state units, and packs them into a ragged state checkpoint layout. To integrate efficiently with tensor-parallel inference engines, DASC furtherly balances compressed state checkpoints across TP ranks. On reuse, DASC either zero-fills omitted units or refreshes them from a bounded suffix with additional compute cost. Across retrieval and end-to-end reasoning benchmarks on Kimi-Linear, conservative DASC configurations remain close to full caching while compressing KDA recurrent state checkpoints by $2.63\times$. Under fixed state checkpoint memory budgets, the resulting capacity gains reduce mean Time to First Token (TTFT) by 42.6\% and improve input throughput by 68.4\%. At larger compression ratio, suffix refresh recovers much of the accuracy lost to more aggressive omission, at the cost of additional replay computation. Qwen with GDN exhibits a similar quality--efficiency trend, showing that DASC extends from channel-wise KDA to head-wise GDN.