如何通过KV稀疏实现对vLLM的1.5倍加速

696 0 0

文章摘要

【关键词】 KV稀疏、推理加速、分层稀疏、性能测试、Tensor Parallelism

在过去一年中，KV稀疏算法的研究取得了显著进展，但在实际应用中，如何将这些算法与现有的大型语言模型（LLM）框架如vLLM相融合，以提高推理性能，仍是一个挑战。vLLM框架采用了如Continuous Batching、FlashAttention和PagedAttention等优化特性，但与大多数稀疏算法不兼容或性能不佳。为了解决这一问题，研究者们基于KV稀疏原理对vLLM框架进行了修改，实现了1.5倍的推理加速。

LLM的Massive Activations特性表明，少数激活值远高于其他激活，这为通过KV稀疏方法提升推理性能提供了可能。Llama2和Llama3模型的分析数据显示，Massive Activations出现在特定的特征维度和token类型中。Llama3模型的不同层展现出不同的激活特性，因此需要分层稀疏策略。

在实现KV稀疏时，面临的挑战包括显卡容量、计算能力和IO限制。通过分层稀疏，可以减少KV开销，从而加速推理过程。具体来说，通过淘汰策略删除低分KV，保留高分和近距离KV，以节约内存并降低计算量和IO开销。

性能测试表明，KV稀疏在保证首token时间（TTFT）可用的基础上，能显著提升vLLM的吞吐量。在较大Batch Size场景下，KV稀疏版本在高并发度下仍能保持性能稳定。模型性能测试显示，精度损失控制在3%以内，长文本场景下的精度损失约为10%。

关键技术包括分层稀疏和Tensor Parallelism。vLLM采用Continuous Batching调度策略，需要解决不同层的内存管理问题。通过调整底层结构，实现了不同层的不同稀疏程度。此外，Attention计算单元的改造，如FlashAttention和PagedAttention，也是实现KV稀疏的关键。

总结来说，基于vLLM 0.5.1改造的PPIO Sparse0.5.1支持Llama3-8B和Llama3-70B等模型，主要在消费级显卡上部署。尽管在工程和算法上仍有优化空间，但PPIO Sparse0.5.1已经在推理性能和模型性能方面取得了显著成果。