OasisKV: 先読みスパースプリフェッチによるHBMを超えたデコード時KVキャッシュのスケーリング

要旨

大規模言語モデル（LLM）の推論サービングは、計算量よりもメモリによって制約されることが増えている。長コンテキストや長文推論のワークロードが普及するにつれ、キー値（KV）キャッシュがLLMのトークン生成、すなわちデコード中において、メモリフットプリントとメモリトラフィックの両方を支配するようになっている。特に、HBM容量は希少かつ高コストなリソースとなり、推論のバッチサイズとシステムスループットを大きく制限している。本論文は、LLMデコード中に完全なKVキャッシュの格納をHBMから分離することでHBM容量の圧迫を軽減する、メモリ中心のLLM推論システム設計であるOasisKVを提案する。デコード時のアテンションは自然にスパースであるため、OasisKVはアテンション計算のために最も関連性の高いトークンのKVエントリのみをHBMに保持する。我々は、投機的デコード（SD）によって生成される先読みトークンを用いることで、将来重要となるトークンを事前に正確に予測できることを観察した。OasisKVは、重要なKVブロックを特定するための効率的なアテンションのバックグラウンドパイプラインを採用する。それらのブロックは、より大容量のメモリ階層（例えば、ホストメモリやリモートメモリ）からプリフェッチされ、次のデコードステップで使用される前にHBMにステージングされる。我々はOasisKVをvLLMに基づいて実装した。先読み予測は十分に正確であり、2,048トークンのKV予算の下で、フルアテンションとの精度差を0.7ポイント以内に抑える。これによりOasisKVはスパース性をスループット向上に変換できる：推論ワークロードでは精度損失0.1ポイントで高密度（dense）vLLMの1.69倍、マルチGPUの長コンテキストサービングでは最大2.1倍のスループットを達成する。プリフィル・デコード分離の下では、OasisKVは各リクエストの受け入れに際して、完全なKV転送と比較して6.5〜9.7倍少ないKVで済み、デコードノードのホストメモリを2.2〜2.6倍少なく保持しながら、高密度スループットの約2倍を達成する。

English

Large language model (LLM) inference serving is increasingly constrained by memory rather than compute. As long-context and long-form reasoning workloads become more prevalent, the key-value (KV) cache dominates both memory footprint and memory traffic during LLM token generation, i.e., decode. In particular, HBM capacity has become a scarce and costly resource that heavily limits inference batch size and system throughput. This paper presents OasisKV, a memory-centric LLM inference system design that alleviates HBM capacity pressure by decoupling full KV-cache storage from HBM during LLM decoding. Because decode-time attention is naturally sparse, OasisKV keeps only the KV entries of the most relevant tokens in HBMs for attention computation. We observe that future important tokens can be predicted accurately in advance using lookahead tokens drafted by speculative decoding (SD). OasisKV employs an efficient attention background pipeline to identify important KV blocks. They are then prefetched from higher-capacity memory tiers (e.g., host or remote memory) and staged in HBMs before being used in the next decode step. We implement OasisKV based on vLLM. The lookahead prediction is accurate enough to keep accuracy within 0.7 points of full attention under a 2,048-token KV budget. This lets OasisKV turn sparsity into throughput gain: 1.69times over dense vLLM on the reasoning workload at 0.1 points of accuracy loss, and up to 2.1times on multi-GPU long-context serving. Under prefill--decode disaggregation, OasisKV reaches about 2times dense throughput while admitting each request with 6.5--9.7times less KV and holding 2.2-2.6 less decode-node host memory than full KV transfer.