SwiftKV：让大模型推理更快、吞吐量更大！

SwiftKV：让大模型推理更快、吞吐量更大！

原创 格知致能 寒武纪人工智能 2025年01月21日 07:20 北京

    最近，有关模型性能优化方面又有新的成果方法了。研究公司Snowflake公布了一款名为“SwiftKV”的AI模型调校技术，并源三款利用“SwiftKV”技术进行调校的 Llama 3.1模型。笔者连夜阅读了论文，发现这项技术创新的角度比较有意思，是针对提示词处理的，值得一读。

1

性能瓶颈：提示词太长

    大语言模型 (LLM) 正在迅速成为企业应用程序和产品、检索增强生成 (RAG)、摘要和自主工作流程的核心。 但推理的成本和速度决定了它们的实用性。 因此，提高 LLM 推理的聚合吞吐量和降低延迟已成为一个越来越重要的关注点，各种努力从多个角度解决了这个问题。

    论文作者发现，通常情况下，提示词的长度远大于生成内容的长度，平均约为生成内容的10倍（的确如此，大多数开发者都是把文档一股脑全都丢给大模型）。这种特性导致预填充（prefill）阶段的计算开销巨大，增加了响应延迟和推理成本。为适应这种大模型输入文本远大于输出本文的应用场景，作者提出了SwiftKV。

2

SwiftKV 是什么？

    SwiftKV 是一种新型的模型优化技术，专为提高大模型在企业场景中的推理效率而设计。它的核心思想非常简单：

    在处理输入时，跳过不必要的计算，并减少内存占用，直接利用较早层的输出来填充后续层的缓存。SwiftKV 还通过合并相邻层的缓存来减少内存占用。这意味着，原本需要为每一层单独分配的内存现在可以共享，从而支持更大的批量处理，进一步提高了推理效率。

    SwiftKV 的厉害之处在于，它能够在几乎不影响模型性能的前提下，显著降低推理成本和延迟。

3

SwiftKV 技术原理

    3.1 主要思路和方法：

    SingleInputKV：减少预填充计算

    核心思想：在 Transformer 模型中，较深层的隐藏状态（hidden states）变化较小，因此可以利用较早层的输出来为后续层填充 KV 缓存，从而跳过后续层的计算。

    实现方式：通过重写模型架构，使得从某一层（例如第 L/2 层）开始，后续所有层的 KV 缓存直接由该层的输出计算得到，而无需逐层计算。

    效果：减少了约 50% 的预填充计算量，同时保持了较高的生成质量。

    AcrossKV：减少 KV 缓存内存占用

    核心思想：将相邻层的 KV 缓存合并为一个共享缓存，从而减少内存占用。

    实现方式：通过将多个连续层的 KV 缓存合并为一个，减少冗余的 KV 缓存，支持更大的批量大小，从而提高吞吐量。

    效果：在 2-way 和 4-way 合并的情况下，KV 缓存内存占用分别减少了 25% 和 37.5%，并且与量化技术结合时，可以进一步减少内存占用。

    知识恢复（Knowledge Recovery）

    核心思想：由于 SingleInputKV 和 AcrossKV 改变了模型架构，需要通过知识蒸馏（distillation）从原始模型中恢复模型性能。

    实现方式：仅对受影响的层（如 Q、K、V 投影矩阵）进行微调，使用原始模型的输出作为蒸馏目标，从而在少量数据上快速恢复模型性能。

    效果：通过轻量级蒸馏，模型在多种任务上的质量仅略有下降。

    优化推理实现

    核心思想：将 SwiftKV 集成到高效的推理框架（如 vLLM）中，进一步优化计算和内存管理。

    实现方式：通过融合 KV 缓存计算、减少内存分配等方式，实现更高的推理效率。

    效果：在实际推理中，SwiftKV 实现了高达 2 倍的吞吐量提升和 60% 的输出延迟降低。

4

实验结果    

    在 Llama-3.1-8B 和 Llama-3.1-70B 模型上，SwiftKV 在多种任务（如 ARC-Challenge、Winogrande、MMLU 等）上表现出色，平均质量下降不到 1 个百分点，同时显著减少了预填充计算量和 KV 缓存内存占用。

    在批量推理场景下，SwiftKV 的吞吐量比基线模型高出 1.2 到 2 倍。

    在交互式推理场景下，SwiftKV 的首次输出延迟（TTFT）和每输出延迟（TPOT）分别降低了 50% 和 60%。

5

总结

    这篇论文提出了一种针对 LLM 推理的高效优化方法，特别适用于输入长度远大于输出长度的场景。通过减少预填充计算和 KV 缓存内存占用，显著提高了模型的吞吐量和响应速度。同时通过轻量级蒸馏，成功地在少量数据上恢复了模型性能，保持了较高的生成质量。

    作者提供了开源代码和预训练模型，方便社区使用和进一步研究。

    格知致能点评：针对特定的应用场景，对模型进行特定的优化，并且还不影响模型的性能，这是一个技术活，也可能是做大模型垂直行业落地的企业的价值和核心竞争力吧。

    论文地址：https://arxiv.org/abs/2410.03960