월~금 오전 7시, 일요일 오전 8시 — AI 뉴스와 용어를 보내드립니다메일로 받아보기

METAL LAB

LLM이 답을 만들 때 쓰는 '기억 캐시(KV cache)'를 GPU 밖에 두고도, 다음에 필요한 부분만 미리 예측해서 당겨오는 방식으로 처리량을 2배 가까이 늘렸다

arXiv:2608.080972026-08-07

OasisKV: Scaling In-Decode KV Cache Beyond HBM with Lookahead Sparse Prefetching

LLM이 답을 만들 때 쓰는 '기억 캐시(KV cache)'를 GPU 밖에 두고도, 다음에 필요한 부분만 미리 예측해서 당겨오는 방식으로 처리량을 2배 가까이 늘렸다

대형 언어모델이 문장을 이어 생성할 때마다 이전 토큰들의 정보(KV 캐시)를 GPU의 고속 메모리(HBM)에 계속 들고 있어야 해서, 문맥이 길어지면 GPU 메모리가 금방 꽉 찬다. OasisKV는 스펙큘레이티브 디코딩(미리 몇 토큰을 추측해보는 기법)에서 나오는 '미리보기 토큰'을 이용해 다음 단계에 진짜 필요한 KV 블록을 미리 정확하게 예측하고, 이를 CPU 메모리나 원격 메모리에서 GPU로 백그라운드로 미리 가져와 둔다. vLLM 위에 구현해 검증한 결과, 정확도 손실을 0.7점 이내로 유지하면서 처리량을 최대 2배 가까이 끌어올렸다.

METAL LAB 해설 도표

OasisKV의 미리보기 기반 KV 프리페치 구조

증거 상태측정 결과가 보고됨

  1. 문제 상황긴 문맥을 다루는 디코딩에서 KV 캐시가 GPU 고속메모리(HBM) 용량을 다 차지해 배치 크기와 처리량이 제한된다
  2. 미리보기 예측스펙큘레이티브 디코딩의 draft 토큰을 이용해 다음 단계에 필요한 KV 블록을 별도 학습 없이 예측한다(top-K 예측)
  3. 선택과 전송 제한예측된 블록과 현재 GPU에 있는 블록을 비교해 필요한 블록만 골라 단계당 전송량을 제한(capped eviction)해 PCIe 대역폭 안에서 처리한다
  4. 백그라운드 비동기 파이프라인예측-선택-전송 단계를 레이어별로 CUDA 스트림에서 겹쳐 실행해 GPU 연산과 동시에 CPU/원격 메모리에서 KV를 가져온다
  5. 결과vLLM 기반 구현에서 정확도 손실 0.7점 이내로 처리량을 최대 2.1배까지, PD 분리 서빙에서는 약 2배까지 향상시켰다
METAL LAB이 원문을 바탕으로 재구성한 해설 도표이며, 논문 저자의 원문 figure가 아닙니다.

무엇을 했나

  1. 문제 제기: 긴 문맥·긴 추론을 다루는 에이전트형 작업이 늘면서 KV 캐시가 GPU 고속메모리(HBM) 용량을 가장 많이 잡아먹어 배치 크기와 처리량을 제한하는 병목이 되었다.
  2. 방법: 스펙큘레이티브 디코딩에서 이미 만들어지는 '미리보기(draft) 토큰'을 그대로 활용해, 다음 디코딩 단계에서 실제로 중요할 KV 블록을 별도의 예측 모델 학습 없이 미리 알아낸다. 이를 백그라운드 파이프라인(예측→선택→전송)으로 CPU/원격 메모리에서 GPU로 겹쳐서 가져온다.
  3. 구현: vLLM 기반으로 프로토타입을 만들어 페이지 단위 KV 관리, 헤드별 매핑, 단계별 전송량 제한(capped eviction) 등을 추가해 실제 서빙 엔진에서 동작하게 했다.
  4. 결과: 2,048토큰 KV 예산에서 전체 어텐션(모든 KV를 다 보는 방식) 대비 정확도 차이를 0.7점 이내로 유지하면서, 추론 작업에서 vLLM 대비 1.69배, 멀티 GPU 장문맥 서빙에서 최대 2.1배 처리량 향상을 얻었다.
  5. 확장: prefill과 decode를 분리한 서빙 구조에서도 전체 KV를 다 옮기는 방식보다 6.5~9.7배 적은 KV만 전송하면서 처리량은 약 2배, 디코드 노드의 호스트 메모리 사용은 2.2~2.6배 적게 유지했다.
Figure 1. KV-cache memory roofline. Decode throughput is capped by Token-KV intensity × KV-cache bandwidth (three-tier slope-1 roofs) and by the dense/sparse throughput ceilings. Higher Token-KV intensity means more KV reuse (e.g., GQA (1), MLA (5)) or more KV sparsity (e.g., NSA (42), DSA (4)). It is defined as the tokens generated per decode pass per KV-cache byte transferred over HBM or off-GPU IOs. Two throughput ceilings are determined by the maximum batch size and minimum time per output token (TPOT) bounded by compute and weight-loading time.
Figure 1. KV-cache memory roofline. Decode throughput is capped by Token-KV intensity × KV-cache bandwidth (three-tier slope-1 roofs) and by the dense/sparse throughput ceilings. Higher Token-KV intensity means more KV reuse (e.g., GQA (1), MLA (5)) or more KV sparsity (e.g., NSA (42), DSA (4)). It is defined as the tokens generated per decode pass per KV-cache byte transferred over HBM or off-GPU IOs. Two throughput ceilings are determined by the maximum batch size and minimum time per output token (TPOT) bounded by compute and weight-loading time.
Figure 2. The per-output-token decode latency (TPOT) breakdown of on-demand KV retrieval over PCIe (Qwen3-32B, one H100 SXM-HBM3 GPU, context length 32K, BF16). All latencies are computed from a roofline model, assuming 10% of the attended KV is fetched from CPU DRAM each step. The percentage above each bar denotes the KV retrieval overhead.
Figure 2. The per-output-token decode latency (TPOT) breakdown of on-demand KV retrieval over PCIe (Qwen3-32B, one H100 SXM-HBM3 GPU, context length 32K, BF16). All latencies are computed from a roofline model, assuming 10% of the attended KV is fetched from CPU DRAM each step. The percentage above each bar denotes the KV retrieval overhead.
Table 1. Accuracy under the same 2,048-token KV budget. Each retrieval method is read against the full-attention anchor of its own stack: Δ is the difference from that stack’s Full column (Quest, FreeKV vs. Full (HF); Ours vs. Full (vLLM)), never across stacks.
HuggingFace Transformers stackvLLM stack
Dataset / SubsetMetricFullQuestΔFreeKVΔFullOursΔ
Long input — Llama-3.1-8B-Instruct, LongBench v2
Overall29.6229.42−0.2029.03−0.5930.2329.62−0.61
Short34.4434.440.0035.00+0.5635.5633.89−1.67
Medium28.3728.84+0.4726.51−1.8627.4427.440.00
Long24.0722.22−1.8524.070.0026.8526.850.00
Long input — Qwen3-8B, LongBench v2
Overall32.2131.61−0.6031.01−1.2033.6033.20−0.40
Short36.6737.22+0.5536.670.0039.4438.33−1.11
Medium29.3026.98−2.3227.44−1.8629.3030.70+1.40
Long30.5631.48+0.9228.70−1.8632.4129.63−2.78
Long output — Qwen3-8B, reasoning
Overallpass@k81.4178.86−2.5677.91−3.5078.8478.18−0.66
avg@k69.4866.65−2.8366.85−2.6367.6367.28−0.35
AIME24pass@886.6786.670.0080.00−6.6783.3383.330.00
avg@877.5075.42−2.0872.08−5.4279.1776.67−2.50
AIME25pass@883.3376.67−6.6680.00−3.3380.0080.000.00
avg@870.8364.17−6.6668.75−2.0865.4267.97+2.55
GPQA-Diamondpass@474.2473.23−1.0173.74−0.5073.2071.21−1.99
avg@460.1060.36+0.2659.72−0.3858.3057.20−1.10
Figure 3. Decode pipelines of dense attention, sparse attention, KV retrieval, and KV prefetching.
Figure 3. Decode pipelines of dense attention, sparse attention, KV retrieval, and KV prefetching.
Figure 4. Top-20 KV-block prediction accuracy across model layers using the previous token vs. a lookahead token.
Figure 4. Top-20 KV-block prediction accuracy across model layers using the previous token vs. a lookahead token.
Table 2. Fetch-cap ablation on Qwen3-8B (AIME24, LRU eviction). The cap bounds the blocks fetched per step by a fetch ratio. Bold marks our default operating point. Dense attention scores 76.04 / 90.00.
Fetch RatioFetchBWTPSAIME24
(GB/step)(GB/s)(tok/s)avg@32pass@32
0.010.305.02,17874.9090.00
0.020.609.82,06675.1090.00
0.051.4923.82,08375.9486.67
0.102.8731.41,42176.7786.67
0.204.3434.01,03577.4093.33
Fetch all5.0533.582476.4686.67
Figure 5. Overview of the OasisKV architecture.
Figure 5. Overview of the OasisKV architecture.
Figure 6. Per-layer agreement between the top-K set predicted by the propagated draft query and the exact set of the true next-token query. The profile uses Qwen3-8B on GSM8K with Tengyunw/qwen3_8b_eagle3 as the EAGLE-3 draft model.
Figure 6. Per-layer agreement between the top-K set predicted by the propagated draft query and the exact set of the true next-token query. The profile uses Qwen3-8B on GSM8K with Tengyunw/qwen3_8b_eagle3 as the EAGLE-3 draft model.

실제로 확인된 결과

  • 2,048토큰 KV 예산 조건에서 전체 어텐션과 비교해 정확도 차이가 0.7점 이내로 유지됐다.
  • 추론 워크로드에서 vLLM 대비 1.69배 처리량 향상을 얻었으며 이때 정확도 손실은 0.1점이었다.
  • 멀티 GPU 장문맥 서빙에서 최대 2.1배 처리량 향상을 얻었다.
  • prefill-decode 분리 서빙에서 전체 KV 전송 대비 6.5~9.7배 적은 KV로 요청을 처리하면서 약 2배의 처리량을 냈고, 디코드 노드 호스트 메모리 사용은 2.2~2.6배 적었다.
  • Qwen3-8B 기준 KV 단계당 전송량 제한(fetch cap)을 0.05로 설정했을 때 정확도는 dense 방식과 0.1점 차이(75.94 대비 76.04)를 보이면서 전체 전송 대비 2.5배 높은 처리량(2,083 tok/s)을 기록했다.
Figure 7. Look-ahead attention. Left: the CPU full KV cache maps to two GPU caches — compressed keys via per-block min/max pooling and the sparse KV working set via block-wise sparsification. Right: the attention kernel processes the normal and draft queries together over the sparse KV; the draft query then scans the compressed keys to predict the next step’s top-K blocks, which are prefetched from the CPU cache.
Figure 7. Look-ahead attention. Left: the CPU full KV cache maps to two GPU caches — compressed keys via per-block min/max pooling and the sparse KV working set via block-wise sparsification. Right: the attention kernel processes the normal and draft queries together over the sparse KV; the draft query then scans the compressed keys to predict the next step’s top-K blocks, which are prefetched from the CPU cache.
Figure 8. The asynchronous prefetch pipeline across two decoding steps. Red arrows trace one layer’s chain: the draft query at step t drives top-K prediction, KV selection, and KV transfer before that layer’s attention at step t+1.
Figure 8. The asynchronous prefetch pipeline across two decoding steps. Red arrows trace one layer’s chain: the draft query at step t drives top-K prediction, KV selection, and KV transfer before that layer’s attention at step t+1.

어디에 쓸 수 있나

  • 긴 문맥이나 긴 추론 과정을 다루는 챗봇·코딩 에이전트·웹 사용 에이전트 서비스에서 GPU 메모리 제약을 완화하고 동시 처리 요청 수를 늘리는 인프라 최적화
  • 여러 GPU에 걸쳐 서비스되는 장문맥 LLM 서빙 시스템의 처리량 개선
  • prefill과 decode 단계를 분리해 운영하는 대규모 서빙 클러스터에서 네트워크 전송량과 디코드 노드 메모리 사용을 줄이는 설계
Figure 9. Head-wise mapping between the bounded GPU working set and the CPU full-KV cache. The original page tables retain their logical-to-physical translations. An additional table maps each GPU logical block to a CPU logical block for every KV head.
Figure 9. Head-wise mapping between the bounded GPU working set and the CPU full-KV cache. The original page tables retain their logical-to-physical translations. An additional table maps each GPU logical block to a CPU logical block for every KV head.
Figure 10. Remote partial fetching: a partial transfer at admission (top) and a network fetch at decode (bottom). Numbered steps are described in Sections 4.4.1 and 4.4.2.
Figure 10. Remote partial fetching: a partial transfer at admission (top) and a network fetch at decode (bottom). Numbered steps are described in Sections 4.4.1 and 4.4.2.

한계와 남은 검증

  • 평가는 Qwen3-8B, Qwen3-32B, Qwen3-235B 등 특정 모델과 AIME24/25, GPQA-Diamond, LongBench v2, GSM8K 같은 특정 벤치마크에 한정되어 다른 모델·작업으로의 일반화는 추가 검증이 필요하다.
  • 프로토타입은 아직 프리픽스 캐싱(prefix caching)을 지원하지 않으며, 이를 반영한 TTFT 개선 효과는 실측이 아닌 분석적 모델링 결과로만 제시됐다.
  • 성능은 특정 하드웨어(H100 GPU, PCIe 대역폭 등) 조건에서 측정된 것으로, 다른 인터커넥트·메모리 계층 구성에서는 결과가 달라질 수 있다.
  • 예측이 항상 정확하지는 않아 놓친 블록에 대한 보정 전송(decode-time miss)이 발생하며, 이는 네트워크 트래픽 절감 효과를 일부 상쇄한다.
Figure 11. Synthetic decode sweep over max concurrency at 16K and 32K context, on Qwen3-8B (single H100) and Qwen3-235B (TP8, eight H100s). We compare OasisKV against dense attention on unmodified vLLM (labeled vLLM) and three hierarchical-KV baselines: ShadowKV, InfiniGen, and FreeKV. Rows: decode throughput (TPS), running batch, and per-token latency (TPOT); each request generates 2,048 tokens. Arrows mark the OasisKV/vLLM ratio at max concurrency 256.
Figure 11. Synthetic decode sweep over max concurrency at 16K and 32K context, on Qwen3-8B (single H100) and Qwen3-235B (TP8, eight H100s). We compare OasisKV against dense attention on unmodified vLLM (labeled vLLM) and three hierarchical-KV baselines: ShadowKV, InfiniGen, and FreeKV. Rows: decode throughput (TPS), running batch, and per-token latency (TPOT); each request generates 2,048 tokens. Arrows mark the OasisKV/vLLM ratio at max concurrency 256.
Figure 12. Disaggregated serving over the offered request rate, Qwen3-8B, at 24K (left) and 32K (right) context with 2,048 output tokens. Top: decode throughput. Bottom: average and peak host memory usage in the decode node.
Figure 12. Disaggregated serving over the offered request rate, Qwen3-8B, at 24K (left) and 32K (right) context with 2,048 output tokens. Top: decode throughput. Bottom: average and peak host memory usage in the decode node.

왜 중요한가

GPU의 고속 메모리는 비싸고 한정돼 있어 긴 문맥을 다루는 서비스일수록 배치 크기와 처리량이 급격히 줄어드는데, 이 방법은 정확도를 거의 유지한 채 같은 하드웨어로 더 많은 요청을 동시에 처리할 길을 보여준다. 특히 여러 GPU에 걸친 서빙이나 prefill-decode를 분리한 대규모 서빙 환경에서 메모리 비용을 크게 줄일 수 있다는 점에서 실제 서비스 운영에 직결된다.

이 논문의 용어

  • KV 캐시 · 언어모델이 이전에 처리한 토큰들의 정보를 저장해두는 캐시로, 다음 토큰을 생성할 때마다 다시 참조된다
  • HBM · GPU에 탑재된 고속이지만 용량이 제한적인 메모리
  • 스펙큘레이티브 디코딩 · 느린 정식 모델 대신 작은 모델로 몇 토큰을 미리 추측해 생성 속도를 높이는 기법
  • prefill-decode 분리 서빙 · 프롬프트를 처리하는 단계와 토큰을 하나씩 생성하는 단계를 서로 다른 서버(노드)에서 처리하는 서빙 방식
  • capped eviction · 한 단계에서 새로 불러올 KV 블록 수를 제한해 전송량이 너무 커지지 않도록 하는 정책

저자 · Can Xiao

arXiv에서 원문 보기

최신 논문

논문 전체 보기 →

METAL LAB 최신 기사

그림 출처: Can Xiao et al., arXiv:2608.08097, arxiv-nonexclusive