[논문리뷰] OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching

링크: 논문 PDF로 바로 열기

메타데이터

저자: Can Xiao, Sukmin Cho, Junbong We, Zhixiong Niu, Jianyi Cheng, Yiren Zhao, Youngjin Kwon, Yongqiang Xiong, Rui Ma, Junyi Liu

---

## 1. Key Terms & Definitions (핵심 용어 및 정의)

• KV Cache: LLM의 디코딩 과정에서 이전 토큰들의 정보를 저장하여 중복 계산을 방지하는 메모리 구조.
• Speculative Decoding (SD): 작은 모델(Draft Model)을 사용해 미래의 토큰을 예측하고 이를 기반으로 추론 속도를 가속화하는 기법.
• PD (Prefill-Decode) Disaggregation: LLM의 연산 단계인 Prefill과 Decode를 분리하여 서로 다른 노드에서 처리함으로써 자원 활용 효율을 극대화하는 방식.
• HBM (High-Bandwidth Memory): GPU 내부에 위치하여 매우 높은 대역폭을 제공하지만 용량이 제한적인 메모리.

## 2. Motivation & Problem Statement (연구 배경 및 문제 정의) 본 논문은 대규모 언어 모델(LLM)의 long-context 및 agentic workloads가 증가함에 따라 발생하는 HBM의 용량 병목 현상을 해결하고자 한다 [Figure 1]. 기존 연구들은 Sparse Attention을 통해 연산 효율을 높이거나 Hierarchical Memory 시스템을 통해 KV Cache를 CPU DRAM 등으로 오프로딩하지만, 여전히 디코딩 과정에서 발생하는 빈번한 메모리 전송이 Critical Path를 점유하여 시스템 처리량을 제한한다 [Figure 3]. 특히, PD Disaggregation 환경에서는 원격 노드로부터의 데이터 전송 대역폭이 제한적이어서, 디코딩 성능 향상을 위해서는 미래에 필요한 KV Cache 블록을 미리 예측하여 비동기적으로 프리페칭(Prefetching)하는 시스템 설계가 필수적이다.

Figure 1 — KV 캐시 메모리 루프라인 모델

## 3. Method & Key Results (제안 방법론 및 핵심 결과) 본 논문은 Speculative Decoding에서 생성된 미래 토큰 예측을 활용하여 필요한 KV Cache 블록을 사전에 파악하고 전송하는 OasisKV를 제안한다 [Figure 5]. OasisKV는 Look-ahead Attention을 도입하여 foreground에서는 디코딩을 수행하고, background에서는 draft token을 이용해 다음 단계의 중요 KV 블록을 선택하여 프리페칭한다 [Figure 7]. 특히, compressed key summaries(min/max pooling)를 사용하여 비거주 블록의 중요도를 낮은 오버헤드로 예측하며, eviction 전략을 통해 HBM 내 Sparse working set을 효율적으로 관리한다. 실험 결과, OasisKV는 vanilla vLLM 대비 reasoning 워크로드에서 1.69x, multi-GPU long-context serving에서 최대 2.1x의 처리량 향상을 달성하였다. 또한, PD disaggregation 환경에서는 전체 KV 전송 방식 대비 2.12.3x의 높은 처리량을 기록하면서도, 6.59.7x 적은 KV cache 메모리와 2.2~2.6x 적은 호스트 메모리 점유율을 보였다. [Figure 6]은 draft query 기반의 블록 예측이 실제 다음 단계와 높은 정확도로 일치함을 입증한다.

Figure 5 — OasisKV 시스템 전체 아키텍처

Figure 7 — Look-ahead Attention 작동 원리

## 4. Conclusion & Impact (결론 및 시사점) 본 논문은 OasisKV를 통해 LLM 추론 시 HBM의 용량 제한을 극복하고 효율적인 다계층 메모리 활용을 가능하게 했다. OasisKV는 Speculative Decoding의 예측 능력을 시스템 최적화와 결합하여 고성능 인프라 환경에서 LLM 추론의 처리량을 획기적으로 향상시켰다. 이 연구는 고가의 GPU HBM 자원을 효율적으로 관리하고자 하는 대규모 AI 서비스 운영자들에게 실질적인 가이드라인을 제공하며, 향후 더 긴 문맥 처리가 요구되는 차세대 에이전트형 모델 인프라 구축의 핵심적인 기술적 토대가 될 것으로 기대된다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.