SnapStream: Efficient Long Sequence Decoding on Dataflow Accelerators

TL;DR

This paper introduces SnapStream, a novel KV cache compression technique enabling efficient long sequence decoding on dataflow accelerators with minimal accuracy loss, demonstrated on large language models in production.

Contribution

Develops SnapStream, the first sparse KV attention method deployed in production inference systems with static graphs and continuous batching, improving memory efficiency and scalability.

Findings

SnapStream achieves 4x on-chip memory savings.

Minimal accuracy degradation on LongBench-v2, AIME24, and LiveCodeBench.

Demonstrated in a 16-way tensor-parallel deployment on SambaNova accelerators.

Abstract

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.