OasisKV: 룩어헤드 희소 프리페칭으로 HBM 너머로 디코딩 KV 캐시 확장

초록

대규모 언어 모델(LLM) 추론 서빙은 점점 더 연산보다 메모리에 의해 제약을 받고 있다. 장문 컨텍스트(long-context) 및 장문 추론(long-form reasoning) 워크로드가 보편화됨에 따라, 키-값(KV) 캐시는 LLM 토큰 생성, 즉 디코딩 과정에서 메모리 사용량과 메모리 트래픽을 모두 지배하게 되었다. 특히 HBM 용량은 추론 배치 크기와 시스템 처리량을 크게 제한하는 희소하고 비용이 많이 드는 자원이 되었다. 본 논문은 LLM 디코딩 중 전체 KV 캐시 저장소를 HBM에서 분리함으로써 HBM 용량 압력을 완화하는 메모리 중심의 LLM 추론 시스템 설계인 OasisKV를 제시한다. 디코딩 시점의 어텐션은 자연적으로 희소하기 때문에, OasisKV는 어텐션 계산을 위해 가장 관련성 높은 토큰의 KV 항목만 HBM에 유지한다. 우리는 투기적 디코딩(speculative decoding, SD)이 제안한 선행 토큰(lookahead tokens)을 사용하여 미래에 중요한 토큰을 사전에 정확하게 예측할 수 있음을 관찰한다. OasisKV는 중요한 KV 블록을 식별하기 위해 효율적인 어텐션 백그라운드 파이프라인을 사용한다. 그런 다음 이 블록들은 더 큰 용량의 메모리 계층(예: 호스트 또는 원격 메모리)에서 프리페치되어 다음 디코드 단계에서 사용되기 전에 HBM에 배치된다. 우리는 vLLM을 기반으로 OasisKV를 구현한다. 선행 예측은 2,048토큰 KV 예산 하에서 전체 어텐션 대비 정확도를 0.7포인트 이내로 유지할 만큼 정확하다. 이를 통해 OasisKV는 희소성을 처리량 향상으로 전환한다: 추론 워크로드에서는 정확도 손실 0.1포인트로 밀집(dense) vLLM 대비 1.69배, 다중 GPU 장문 컨텍스트 서빙에서는 최대 2.1배의 처리량을 달성한다. 프리필-디코드 분리(prefill-decode disaggregation) 환경에서 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.