PyTorch로 구현한 GPT-2 소형 모델과 KV 캐시 성능 비교
GPT-2 KV cache from scratch: correctness, latency, and memory measurements [P]
저자는 PyTorch로 GPT-2 소형 모델(124M)과 KV 캐시를 구현하고, 수학적 절감 효과를 측정한 지연 시간과 메모리 사용량을 비교했다. 실험은 NVIDIA A40에서 진행되었으며, 캐싱을 사용했을 때 1,000개의 컨텍스트 길이에서 FLOP 추정치는 약 1,000배 작아졌고, 토큰당 시간은 29.69ms에서 4.95ms로 줄어들었다. 1,016개의 토큰을 생성하는 데 uncached는 14.66초, cached는 5.03초가 소요되었다. 메모리 사용량은 72 KiB/token으로 측정되었다.
원문 출처
r/MachineLearning