LLM의 Context Length 늘리기 (2) - KV Cache 최적화 기법들

 

LLM의 Context Length 늘리기 (2) - KV Cache 최적화 기법들

 

 

물론 무료 Perfromance 아닙니다

 

그래 YaRN으로 모델이 긴 시퀀스 볼 수 있다 치자 ! 근데 메모리 부족한 건 어떻게 해결할 건데 ? 

학습 서버는 그래도 좋은 GPU를 사용할 수 있지만
AI 서비스가 나갈 서버의 GPU를 H100로 할 순 없으니 ...  

서빙용 GPU (가령 A100 40G)에서 이 작디작은 메모리를 어떻게 아껴가며 쓸지 생각해봐야 하는 부분이다 

 

지난 주 올렸던 글에서 (https://asidefine.tistory.com/341) 봤듯이
추론 시에는 1) 모델, 2) KV Cache가 메모리에서 가장 큰 파이를 차지하게 되고
이 중 Input Length와 관련된 부분은 KV Cache다 

 

KV Cache가 VRAM에서 차지하는 양은 배치 × 시퀀스 × 히든차원 × 레이어 × 2 × 정밀도다. 

배치와 시퀀스를 생각해보면
많은 사용자 요청이 있을수록, 또 긴 input을 받을수록 
메모리 사용량이 매우 빠르게 증가함을 예상할 수 있다

 

KV Cache 관련해서는 아래의 글들에서도 이미 정리해두었기에 내 환경(vLLM + A100 40G)에서 할 수 있는 방법을 아래에 정리하고자 한다 

• https://asidefine.tistory.com/320
• https://asidefine.tistory.com/346

 

 

1. vLLM Docs에서 최적화 방법 샅샅이 뒤지기 

 

 

 

2. LMCache 

 

요 LMCache는

Github Trend 구경하다가 발견한건데 타이밍 좋았다

 

vLLM과 함께 쓰면 TTFT도 줄어들고, 이미 vLLM에도 docs 잘 정리되어있길래 (https://docs.vllm.ai/en/latest/examples/others/lmcache.html) 바로 적용해볼 수 있을 거 같다 

 

원리는 아래와 같디 

기존의 LLM 최적화 도구로는 Hugging Face의 Accelerate, DeepSpeed, vLLM 자체의 텐서 캐싱 구조 등이 있습니다. 그러나 이들은 대부분 같은 인스턴스 내에서만 캐시를 유지하거나, 접두사 캐싱에만 초점이 맞춰져 있어 재사용 가능한 비접두사 문맥을 공유하지 못하는 한계가 있습니다.

LMCache는 비접두사 캐시까지도 다른 인스턴스와 공유 가능하도록 설계되어, 분산 환경에서의 캐시 활용을 극대화합니다. 이는 CPU나 디스크에 저장된 캐시까지도 다시 불러올 수 있기 때문에 기존 방법들보다 훨씬 범용적이고 확장성이 뛰어납니다.

 

 

 

Reference. 

https://arxiv.org/abs/2412.19442

A Survey on Large Language Model Acceleration based on KV Cache Management

https://discuss.pytorch.kr/t/lmcache-llm/7179

LMCache: LLM 서빙 효율성을 높여주는 캐시 시스템

https://github.com/LMCache/LMCache

GitHub - LMCache/LMCache: Supercharge Your LLM with the Fastest KV Cache Layer