如何通过KV稀疏实现对vLLM的1.5倍加速
如何通过KV稀疏实现对vLLM的1.5倍加速
近一年来，关于KV稀疏的论文层出不穷，但如何将这些学术成果应用到实际框架中却面临诸多挑战。本文将介绍如何基于LLM的Massive Activations特性，对vLLM框架进行基于KV稀疏的改造，最终实现1.5倍的推理加速。
LLM的Massive Activations特性
在LLM中，存在一个显著的特性：少数激活值明显高于其他激活值，有时甚至高出100,000倍以上。这意味着少数token在模型推理中起着至关重要的作用。这种现象在Llama2和Llama3中都有体现。
Activation Magnitudes (-axis) in LLaMA2-7B. and axes are sequence and feature dimensions. For this specific model, we observe that activations with massive magnitudes appear in two fixed feature dimensions (1415, 2533), and two types of tokens—the starting token, and the first period (.) or newline token (\n).
Llama3与Llama2有所不同，其在不同层表现出不同的特性：在lower layers呈现均匀分布特性，在middle layers呈现localized attention特性，只有在upper layers才呈现出Massive Activations特性。这也是为什么llama3等模型需要使用分层稀疏的原因。
Attention patterns of retrieval-augmented generation across layers in LlaMa (Touvron et al., 2023a;b) reveal that in the lower layers, the model exhibits a broad-spectrum mode of attention, distributing attention scores uniformly across all content. In the middle layers, attention becomes more localized within each document, indicating refined information aggregation (dotted red triangular shapes in layers 6 and 10). This culminates in the upper layers, where “massive attention” focuses on a few key tokens (concentrated attention bars after layer 18), efficiently extracting essential information for answers.
KV稀疏的实现原理
在大模型推理过程中，显卡的容量、计算能力和IO往往是推理的三大瓶颈。为了避免重复计算，通常会缓存KV，这会导致推理过程中KV Cache占用大量显存。同时，计算量和IO开销在Prefill和Decode两个阶段也是重要瓶颈。
基于LLM的Massive Activations特性，我们对vLLM进行了分层稀疏改造，即不同层的稀疏程度不同，从而大大降低KV开销，实现推理加速。
如上图所示，如果我们在推理的过程中，我们对模型的不同层分别进行 KV 稀疏，即通过淘汰策略将打分较低的 KV 进行删除，同时保留打分较高与距离较近的 KV，从而节约内存并同时降低计算量与 IO 开销，最终实现推理加速。
性能测试结果
推理性能
在实际应用中，输入/输出长度为4000/500是最常见的长度，同时RTX4090显卡也是业界应用最广泛的显卡之一。我们在此基础上进行分批次的性能对比测试。对照组是vLLM0.6.1.p2，实验组是PPIO Sparse0.5.1（vLLM 0.5.1的kv稀疏改造版本），两者进行多轮的性能对比测试，参考的主要指标为TTFT（首token指标，影响用户体验）和Throughput（吞吐量，实际的推理速度）。
最终的测试结果显示，通过KV稀疏，在保证TTFT可用的基础上（P50在1秒之内），能将vLLM的吞吐量提升约1.58倍。
如上表所示，在较大 Batch Size 的场景下，vLLM0.6.1.p2 在并发度为 10 的情况下已经到达极限，而 PPIO Sparse0.5.1 在并发度为 20 的情况下依旧能保持 TTFT 性能稳定，从一定程度上保证了 KV 稀疏在实际生产中的性能稳定性。
模型性能
由于KV稀疏是一种有损的压缩算法，因此模型性能评测同样至关重要。我们基于Llama3-8B进行了mmlu等性能评测，发现精度损失基本上在3%以内。
我们通过 mmlu、humanities 进行模型性能测试
针对实际业务中的长文本场景，我们进行了输入长度为7k-30k的QA任务评测，压缩比控制在10倍以上，结果显示整体精度损失在10%左右。
针对实际业务当中的长文本场景，我们进行了输入长度为 7k-30k 的 QA 任务评测，压缩比控制在 10 倍以上，其结果显示整体精度损失在 10% 左右。
关键技术实现
分层稀疏与Tensor Parallelism
vLLM本身采用的是Iteration-Level Schedule调度策略，即Continuous Batching。基于Continuous Batching的改造，我们遇到了一系列挑战。主要需要解决的问题是不同层的内存管理问题。我们对vLLM的底层结构进行了调整，使其同时支持Full Cache和Sliding Window策略，以满足不同层选择不同缓存策略的需求。
Completing seven sequences using continuous batching. Left shows the batch after a single iteration, right shows the batch after several iterations. Once a sequence emits an end-of-sequence token, we insert a new sequence in its place (i.e. sequences S5, S6, and S7). This achieves higher GPU utilization since the GPU does not wait for all sequences to complete before starting a new one.
整体上可以分成三个阶段的调整：
- 服务启动阶段：在服务启动过程中，通过剩余空间计算对三层分配的空间进行初始化，同时对Metadata结构进行分层初始化。
- Step准备阶段：在vLLM的step初始化过程中，对不同层进行不同的block table管理策略，控制KV的更新机制。
- 稀疏化计算阶段：在模型执行后，通过计算得到的稀疏化打分执行稀疏操作，重点考虑Tensor Parallelism、CUDA graphs、GQA等技术的兼容性。
在实际工程实现中，Tensor Parallelism需要重点考虑。在单机多卡场景中，需要通过合理的切分机制将KV稀疏操作限制在每块GPU内部，避免GPU卡之间的通信。幸运的是，Attention单元采用的是基于ColumnLinear的切分方式，实现相对容易。
以上，主要是针对 vllm 支持分层稀疏在框架层上的修改，接下来重点介绍一下底层 CUDA 的优化。
Attention改造
底层CUDA方面的改造主要集中在Attention计算单元的改造，特别是FlashAttention和PagedAttention的改造。这里重点介绍FlashAttention的改造。
回顾一下Attention的计算公式：
如上图所示，softmax 自身的计算过程是对 QK 的计算，同时是一种 2-pass 算法（循环 2 次），然而最终的计算目标 O 却可以通过 FlashAttention 算法，通过 1-pass（循环 1 次）进行实现。FlashAttention 的实现逻辑可以参考下面关于 FlashAttention2 论文的截图，简而言之，即通过 Q 以分块遍历的方式对 KV 进行分块计算，同时逐步更新 O/P/rowmax 等数据，直到循环结束，再让 O 除以ℓ，即可实现 1-pass 的 FlashAttention 计算。
对于实现KV稀疏来说，需要重点注意的是，FlashAttention的计算过程中已经计算出softmax所需的rowmax/ℓ（全局的最大化和累加值），但FlashAttention并未返回这两者的结果。
有了 FlashAttention 的基础理解，我们再讨论一下 PYRAMIDKV 等论文中关于稀疏化的打分公式，如下图所示：
上述公式的 A 为 softmax 打分，S 为邻近 a 个 softmax 打分的求和，通过上面的内容可知，softmax 本身是一种 2-pass 算法，无法直接融合到 FlashAttention 当中，也说明稀疏化打分同样无法 1-pass 实现。
基于上述分析，我们对FlashAttention进行了如下改进：
- FlashAttention增加了rowmax/ℓ的输出，由于这些值是最大值或累加值，数据量不大，对FlashAttention本身影响不大，但通过复用rowmax/ℓ可以避免一次1-pass计算，提升整体计算效率。
- 增加Online Softmax步骤，在FlashAttention的基础上，根据p=e((kq)-row_max)/ℓ公式计算出softmax打分（1-pass逻辑），再进行行求和，即可算出稀疏化打分S。
在Online Softmax的代码实现中，需要注意以下两个细节：
- 调整循环顺序，因为S是基于k维度的打分求和，对比attention中的qk，需要在计算Online Softmax时调整为kq，通过这种方式避免了各个block之间的无效通信。
- q的起止位置，因为存在causal编码以及邻近a求和的逻辑，不需要k遍历所有q。举例说明，例如输入的prompt长度为5000，邻近a的长度为256，常规计算需要50005000的计算量，而通过记录起止位置只需要5000256的计算量，从而提升计算效率。
PagedAttention所做的工作与FlashAttention类似，目标都是求最终的S，这里不再详细说明。
总结
我们基于vLLM 0.5.1改造的PPIO Sparse0.5.1，当前主要支持Llama3-8B、Llama3-70B等模型，推理模式上支持CUDA graphs模式，运行环境主要在RTX 4090等消费级显卡上进行部署。当前无论在工程上，还是在算法上，尚有巨大的优化空间。
对比最新的vLLM-0.6.2版本，以下特性我们还不支持：
- Chunked-prefill
- Prefix caching
- FlashInfer and other non-FlashAttention attention backends
参考论文
[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
[2] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference
[3] SnapKV: LLM Knows What You are Looking for Before Generation
[4] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
[5] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference
[6] MiniCache: KV Cache Compression in Depth Dimension for Large Language Models
[7] Layer-Condensed KV Cache for Efficient Inference of Large Language Models
[8] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding
[9] CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion
[10] KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation