KV cache란?

• GPT는 이전 token 생성 시 발생된 중간값인 activation (KV cache)를 캐싱하여, 이전 token값을 재계산하기 위한 GPU의 FLOPs를 절감하는 대신, KV cache를 위한 추가적인 메모리 공간이 필요함
• 또한 LLM의 context window size가 증가할수록 KV cache의 크기 또한 linear하게 증가하므로 context window size는 메모리 용량에 제한을 받음
• KV caching 은 multiple token generation step 시 decoder에서만 발생

• scaled dot-product attention 과정에서 우리는 결국 “새로운 token”만을 생성하기를 원하기 때문에, 이전 정보에 대해서는 caching을 진행한다.

 

• caching을 적용해서 진행하는 step은 다음과 같다.

key token과 value token에 대해서는 cache된 값을 가져다가 쓴다.

• KV caching을 ㅇ통해서 얻어진 matrix는 훨씬 크기가 작고, 이는 곧 훨씬 더 적은 matrix multiplication을 가능하게 할 수 있다 .
• 1000개의 token을 generation 할 때 실제 결과:

with KV caching: 11.885 +- 0.272 seconds
without KV caching: 56.197 +- 1.855 seconds

 

위 결과를 계산하는 code는 다음과 같음

import numpy as np
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)

# use_cache는 AutoModelForCausalLM이라는 모듈에 내장되어 있다 !
for use_cache in (True, False):
  times = []
  for _ in range(10):  # measuring 10 generations
    start = time.time()
    model.generate(**tokenizer("What is KV caching?", return_tensors="pt").to(device), use_cache=use_cache, max_new_tokens=1000)
    times.append(time.time() - start)
  print(f"{'with' if use_cache else 'without'} KV caching: {round(np.mean(times), 3)} +- {round(np.std(times), 3)} seconds")

 

• KV cache 관련 논문
  • EFFICIENTLY SCALING TRANSFORMER INFERENCE

 

 

- 참고자료 출처 : https://medium.com/@joaolages/kv-caching-explained-276520203249