以高吞吐量服务 Transformer 语言模型需要缓存键值对(KV),以避免自回归生成过程中的冗余计算。KV 缓存的显存占用巨大,严重影响服务成本。本文提出减少这些显存需求的方法。虽然近期工作主要通过沿时间轴进行压缩和驱逐来减少 KV 缓存,但我们认为深度维度提供了一个正交且稳健的优化途径。尽管先前研究表明,为每一层保留完整缓存是冗余的,但实现跨层缓存共享在实践中仍是一个挑战;现有方法通常会降低吞吐量或增加首 token 生成时间。在本文中,我们证明丢弃某一层的缓存可以在不损失信息的情况下实现高效优化。我们提出一种简单的训练方法:随机跨层注意力。在训练过程中,各层随机选择关注自身的 KV 状态还是前一层的 KV 状态。这种随机过程使模型能够适应各种深度维度的缓存共享策略,从而确保在部署时对未知硬件约束的灵活性。我们的评估表明,在预训练或微调期间应用此方案,可以为各种模型族实现深度维度的缓存共享。此外,对于数据受限场景下的较大模型,该方法表现出类似正则化的效果,通常能在显著减少缓存显存占用的同时保持或提升性能。
相关阅读与更新。
EpiCache:面向资源受限环境下长程对话的片段式 KV 缓存管理
现代大语言模型(LLM)将上下文长度扩展至数百万 token,能够基于长程对话历史生成连贯、个性化的回复。然而,键值对(KV)缓存随扩展的对话历史线性增长,导致模型的显存占用迅速超出设备限制。虽然近期 KV 缓存压缩方法试图减少显存使用,但大多数方法在处理完整个…
KV-Runahead:通过并行键值缓存生成实现可扩展的因果大语言模型推理
大语言模型(LLM)推理包含两个阶段:提示词阶段(或称预填充阶段),用于输出第一个模型 token;以及扩展阶段(或称解码阶段),用于生成后续模型 token。在这项工作中,我们提出了一种高效的并行化方案——KV-Runahead,用于加速提示词阶段。关键观察在于,由于键值缓存(KV-cache)的存在,扩展阶段生成模型 token 的速度快于提示词阶段。因此,KV-Runahead……
