Efficient LLM Inference with KCache

🔥🔥[KCache] EFFICIENT LLM INFERENCE WITH KCACHE(@Qiaozhi He, Zhihua Wu)

https://arxiv.org/pdf/2404.18057

---

TL; DR

이 논문은 "KCache"라는 새로운 기술에 대해 설명합니다. 이 기술은 대형 언어 모델(LLM)의 추론 과정에서 메모리 병목 현상을 완화하고 시스템 처리량을 40% 향상시키면서도 정확성을 유지하는 방법을 제시합니다. KCache는 기존의 KV Cache를 대체하여, 메모리 사용량을 줄이면서 추론 속도를 개선합니다. 이를 위해 K 캐시는 HBM(High Bandwidth Memory)에 유지하고, V 캐시는 CPU 메모리에 저장하여 필요한 정보만 동적으로 선택하고 GPU로 복사합니다.

또한, 이 논문은 LLM의 두 가지 주요 단계인 prefill 단계와 decode 단계에서 KCache의 구현 방법을 설명합니다. Prefill 단계에서는 입력 프롬프트를 받아 병렬로 계산을 수행하고, decode 단계에서는 생성된 토큰을 기반으로 하나씩 출력 토큰을 생성합니다. 이 과정에서 KCache는 Attention 계산의 결과를 바탕으로 중요한 V 캐시 데이터만을 선택적으로 GPU로 전송하여 처리량을 최적화합니다.

논문은 KCache가 여러 LLM 모델에서 메모리 효율성을 개선하고, 높은 처리량을 달성하면서도 추론 정확도를 유지하거나 향상시킬 수 있음을 실험을 통해 보여줍니다. 결과적으로, KCach는 추론 과정에서의 메모리 사용과 GPU 부하를 효과적으로 관리하면서 높은 성능을 제공하는 유연하고 확장 가능한 솔루션을 제공합니다.

 

---

Instruction

LLM 추론의 메모리 사용량은 주로 세 부분으로 구성됩니다: Weight, Activations, 그리고 KV Cache입니다. 예를 들어, LLaMA2-7B 모델의 경우, 가중치는 fp16에서 약 14GB의 메모리를 차지합니다. 배치 크기가 8이고 시퀀스 길이가 32 × 1024일 때, KV 캐시는 약 128GB의 메모리를 차지하며, layer-wise memory sharing 전략을 사용하면 activation은 단지 2GB의 메모리만을 차지합니다. 배치 크기와 시퀀스 길이가 증가함에 따라, KV 캐시의 메모리 사용량은 선형적으로 증가합니다.

KV Cache에 대한 내용을 더 자세히 알아보시고 싶으시다면, 이전에 블로그에 작성한 내용을 참고하시면 좋습니다

 

LLM 인퍼런스 훑어보기 (2) - KV Cache

 

최근의 LLM들은 추론 과정에서 KV Cache에서 병목 현상이 발생합니다. 이러한 병목 현상을 완화하기 위해 여러 최적화 방법이 제안되고 있습니다. KV 캐시를 압축하기 위한 방안들로써 바이트의 관점에서 압축하는 양자화 알고리즘, 시퀀스 길이(s) 관점에서 압축하는 컨텍스트 윈도우 압축 알고리즘이 제안되었습니다. Adaptive computation 연산 알고리즘은 레이어(l) 관점에서 계산을 줄이기 위해 early exit decoding을 제안하였으며, 멀티 헤드 어텐션(MHA)의 구조를 개선하여 추론을 가속화하는 방법도 제안되었습니다. K 캐시와 V 캐시 관점에서, 추론 중에 CPU로 오프로딩하고 GPU로 다시 로딩하는 것은 GPU 메모리에 대한 압박을 완화할 수 있지만, 현재 호스트-디바이스(H2D) 및 디바이스-호스트(D2H) 대역폭이 추론의 새로운 병목 지점이 됩니다.

 

Method

아래는 KCache를 표현하는 그림입니다.

 

 

Prefill 단계 동안에는 각 레이어의 계산 결과가 고대역폭 메모리(HBM)로 push됩니다. 그 후, V 캐시의 일부가 CPU로 비동기적으로 복사되며, 이 과정에서 해당 V 캐시가 차지하던 GPU 메모리가 해제됩니다. Decode 단계에서는 K state가 KV Cache로 push되고 pull됩니다. 그러나 KCache는 TopN의 attention score를 계산하고, TopN 결과의 인덱스에 기반하여, 해당 V Cache를 CPU에서 HBM으로 실시간으로 가져와 후속 계산을 완료합니다.

 

 

장문의 Context에서 사용자는 일반적으로 긴 시퀀스에 기반한 여러 라운드의 질문을 합니다. 각 질문은 긴 맥락의 다른 세그먼트에 초점을 맞출 수 있습니다. 각 라운드에서 결과의 정확성을 극대화하기 위해, KCache에서는 KV state의 축소나 압축을 피함으로써 모델의 upper bound를 보장합니다. 그러나, 추론 중에 KV state를 CPU 메모리로 오프로딩하고 GPU로 다시 로딩하는 것은 추론의 E2E 시간을 크게 증가시킬 것입니다. 따라서 모델의 효과성과 추론 지연 시간 사이의 균형을 맞추기 위해서는 필요한 정보만 HBM으로 다시 로드하는 방법을 찾아야 하며, 이는 어떤 정보가 중요한지 결정하는 모듈이 필요함을 의미합니다. 다행히도, Attention 메커니즘에서 Key와 Value 쌍의 의미를 고려할 때, Key는 Query와의 관련성을 계산하는 데 사용되고 Value는 Key와 관련된 실제 정보를 나타내므로, K 캐시와 V 캐시의 일부를 CPU 메모리로 오프로드하는 것이 KCache에 영감을 줬습니다.

 

Analysis of KCache Performance

Prefill 페이즈에서는 V Cache의 일부는 CPU 메모리에 비동기적으로 복사되어야 합니다. 이러한 상황에서 저자들은 각 레이어에 대한 computation time이 이전 레이어의 데이터를 카피하는 시간과 오버랩될 것으로 바랬습니다. 각 트랜스포머 블록에 대해서 $ 2bsd $ bytes의 데이터가 D2H로 전송되어야하고, $ 22bsd^2 + 4bs^2d $ FLops가 연산되어야 합니다.

NVIDIA A100 (80GB) GPU를 사용했을 때, LLaMA2-7B는 $ d = 4096 $이고, 이는 아래와 같은 부등식을 만족합니다.

 

Decode 페이즈에서는 MHA 모듈은 전형적인 memory-boud 작업입니다. Arithmetic Intensity는 부동 소수점 연산(FLOPs)과 입출력 바이트의 비율로 정의됩니다. 이는 디코딩 중 MHA 모듈의 계산 시간이 메모리 접근량에 크게 의존함을 나타냅니다. 특히, 디코드 단계에서 MHA 모듈의 성능은 hidden size와 시퀀스 길이와는 독립적이며 오로지 batch size에만 영향을 받습니다. 이 관찰은 다음과 같은 기대를 이끌어냅니다: 제안된 KCache MHA 모듈의 계산 시간과 데이터 전송 시간은 기존의 KV 캐시 MHA 구현보다 적을 수 있습니다.

NVIDIA A100 (80GB) GPU를 사용했을 때, 메모리 대역폭은 2039GB/s GPU Memory Bandwidth and 32GB/s H2D Bandwidth이고, KCache는 이에 따라 A100에서 $s/N > 64$인 경우에는 성능이 떨어지지 않습니다.

 

 

Analysis of KCache Accuracy

방법론에 따르면 Value 텐서가 CPU 메모리에 비동기적으로 복사되는 prefill 페이즈에서는 추론 성능과 정확도에 영향을 주지 않습니다. Decode 페이즈에서는 H2D 데이터의 양을 줄이는 것이 필연적입니다. 만약 어텐션 스코어의 결과가 충분히 sparse하다면, 최종 결과에 미치는 영향은 무시할 수 있을 정도로 작아집니다.

 

 

Experiments

 

Accuracy

Accuracy 측면에서, KCache는 기본적으로 손실 없이 정확성을 유지했으며, 여러 데이터 세트와 모델에서 더 나은 성능을 달성했습니다. L은 특히 모델 성능이 충분히 강력한 경우 모델 정확도에 상대적으로 작은 영향을 미칩니다. 모델 성능이 상대적으로 약한 경우에는 L을 더 크게 설정하는 것이 좋습니다.

 

 

N이 클수록 정확도가 높아지며, N = 128일 때 KCache는 기준과 동일하거나 더 높은 정확도를 유지했습니다. 저자들은 TopN이 소프트맥스를 정규화하고 노이즈 정보를 추가로 필터링했다고 믿습니다. 세 가지 데이터 세트에 대한 실험을 통해 컨텍스트 길이가 약 2K 이하인 경우 N을 64 또는 128로 설정해도 정확도에 큰 영향을 미치지 않는 것으로 확인되었습니다.

 

 

Performance

Performance 측면에서, 64GB 메모리, 1TB GPU 메모리 대역폭 및 180TFLOPS를 갖춘 GPU에서 실험을 수행했습니다. LLaMA2-7B 모델을 평가한 결과, 표 3은 실험 결과를 보여줍니다. 전반적으로 S >> N 일 때 KCache가 성능 이점을 보여줍니다. 동시에 KCache는 15K 컨텍스트 길이와 N = 128 설정에서 추론 처리량을 40% 이상 향상시켰습니다.