The proliferation of 100B+ parameter Large Language Models (LLMs) with 100k+ context length support have resulted in increasing demands for on-chip memory to support large KV caches. Techniques such as StreamingLLM and SnapKV demonstrate how to control KV cache size while maintaining model accuracy. Yet, these techniques are not commonly used within industrial deployments using frameworks like vLLM or SGLang. The reason is twofold: on one hand, the static graphs and continuous batching methodology employed by these frameworks make it difficult to admit modifications to the standard multi-head attention algorithm, while on the other hand, the accuracy implications of such techniques on modern instruction-following and reasoning models are not well understood, obfuscating the need for implementing these techniques. In this paper, we explore these accuracy implications on Llama-3.1-8B-Instruct and DeepSeek-R1, and develop SnapStream, a KV cache compression method that can be deployed at scale. We demonstrate the efficacy of SnapStream in a 16-way tensor-parallel deployment of DeepSeek-671B on SambaNova SN40L accelerators running at 128k context length and up to 1832 tokens per second in a real production setting. SnapStream enables $4\times$ improved on-chip memory usage and introduces minimal accuracy degradation on LongBench-v2, AIME24 and LiveCodeBench. To the best of our knowledge, this is the first implementation of sparse KV attention techniques deployed in a production inference system with static graphs and continuous batching.
翻译:随着支持10万以上上下文长度的千亿参数大语言模型(LLMs)普及,对容纳大型KV缓存(KV cache)的片上内存需求日益增长。StreamingLLM和SnapKV等技术展示了如何在保持模型精度的同时控制KV缓存大小。然而,在使用vLLM或SGLang等框架的工业部署中,这些技术尚未普遍采用。原因有二:一方面,这些框架采用的静态图与连续批处理(continuous batching)方法,使得修改标准多头注意力算法变得困难;另一方面,此类技术对现代指令遵循与推理模型精度的影响尚不明确,模糊了实施这些技术的必要性。本文在Llama-3.1-8B-Instruct和DeepSeek-R1上探究了这些精度影响,并开发了SnapStream——一种可大规模部署的KV缓存压缩方法。我们在SambaNova SN40L加速器上以16路张量并行方式部署DeepSeek-671B,在128k上下文长度下实现了高达每秒1832个token的实际生产运行,验证了SnapStream的有效性。SnapStream实现了4倍的片上内存效率提升,且在LongBench-v2、AIME24和LiveCodeBench上引入了极小的精度损失。据我们所知,这是首个在采用静态图与连续批处理的生产推理系统中部署稀疏KV注意力技术的实现。