LLM 추론 효율성을 향샹 시키는 KV Caching, MQA,GQA

LLM 추론 효율성 극대화: KV 캐싱, MQA, GQA 심층 분석

대규모 언어 모델(LLM)은 자연어 처리 분야에서 혁신적인 발전을 이끌고 있지만, 특히 텍스트를 순차적으로 생성하는 자기회귀(autoregressive) 방식의 디코더-온리 모델에서는 추론 과정에서의 엄청난 계산량과 메모리 사용량이라는 숙제를 안고 있습니다. 마치 긴 소설을 한 문장씩 이어서 쓰면서, 매번 새로운 문장을 쓸 때마다 앞서 쓴 모든 내용을 다시 검토해야 하는 작가의 고뇌와 비슷하다고 할 수 있습니다. 이러한 비효율성을 해결하고 LLM을 더 빠르고 경제적으로 활용하기 위해 다양한 최적화 기법들이 등장했는데, 그중에서도 KV 캐싱(KV Caching), 멀티-쿼리 어텐션(Multi-Query Attention, MQA), 그리고 그룹-쿼리 어텐션(Grouped-Query Attention, GQA)은 추론 효율성을 크게 향상시키는 핵심 기술로 주목받고 있습니다.

이번 글에서는 이 세 가지 기법이 각각 어떤 원리로 작동하며, 어떻게 LLM의 추론 성능을 개선하는지 개념, 비유, 예시를 통해 자세히 알아보겠습니다.

이전 계산은 재활용! : KV 캐싱 (KV Caching)

LLM이 다음 단어를 예측할 때, 트랜스포머 아키텍처의 핵심인 '어텐션(Attention)' 메커니즘은 이전에 등장한 모든 단어(토큰)들과의 관계를 계산합니다. 새로운 단어가 생성될 때마다 이 계산을 처음부터 반복하는 것은 매우 비효율적입니다. KV 캐싱은 바로 이 반복적인 계산을 줄이기 위한 영리한 전략입니다.

개념: 자기회귀적 텍스트 생성 과정에서, 모델은 각 단계마다 새로운 토큰을 예측하기 위해 이전까지 생성된 모든 토큰에 대한 어텐션 계산을 수행합니다. 이때 어텐션 계산에 사용되는 두 가지 주요 요소가 바로 '키(Key, K)'와 '밸류(Value, V)' 벡터입니다. 중요한 점은, 일단 생성된 토큰의 K와 V 벡터는 이후 단계에서 변하지 않는다는 것입니다. KV 캐싱은 이 점에 착안하여, 이미 계산된 K와 V 벡터들을 메모리(캐시)에 저장해두고 다음 토큰을 생성할 때 재활용하는 기법입니다.

작동 방식: 요리사가 여러 손님에게 순서대로 같은 코스 요리를 내놓는 상황을 상상해 보세요. 첫 번째 손님을 위해 육수를 끓이고 채소를 다듬었다면, 다음 손님을 위해 또다시 육수를 끓이고 채소를 다듬을 필요 없이 미리 준비된 재료를 사용하는 것과 같습니다. KV 캐싱도 마찬가지로, 첫 번째 토큰 생성 시 계산된 K, V 값을 저장해두고, 다음 토큰을 생성할 때는 이 저장된 값을 불러와 현재 토큰의 Q(쿼리)와만 계산합니다.

• 첫 토큰 생성: 입력 프롬프트에 대한 K, V 벡터를 계산하고 캐시에 저장합니다.
                      예) 캐시: [K1, V1].
• 다음 토큰 생성: 새로운 토큰에 대한 K, V 벡터(예: K2, V2)만 계산하고, 이를 기존 캐시에 추가합니다.
                          예) 캐시: [K1, K2], [V1, V2]. 어텐션 계산 시 캐시된 모든 K, V와 현재 토큰의 Q를 사용합니다.
• 이 과정을 통해 매번 전체 시퀀스에 대한 K, V 계산을 반복하는 것을 피해 계산량을 크게 줄일 수 있습니다.

장점:

• 계산량 감소: 이전 토큰들에 대한 중복 계산을 피하여 추론 속도를 크게 향상시킵니다.
• 일정한 처리 시간: 첫 토큰 이후에는 거의 일정한 시간 내에 다음 토큰을 생성할 수 있게 됩니다.

단점:

• 메모리 증가: 시퀀스 길이가 길어질수록 캐시가 차지하는 메모리 양이 선형적으로 증가합니다. 예를 들어, LLaMA-2 13B 모델의 경우 토큰당 약 1MB의 캐시를 사용하여 4,000 토큰 컨텍스트에서 약 4GB의 메모리를 소모할 수 있으며, 이는 모델 크기에 버금가는 수준입니다. 이로 인해 메모리 병목 현상이 발생할 수 있습니다.

똑똑한 공유로 메모리 절약: 멀티-쿼리 어텐션 (Multi-Query Attention, MQA)

KV 캐싱은 계산량을 줄이는 데 효과적이지만, 특히 어텐션 헤드가 많은 멀티-헤드 어텐션(Multi-Head Attention, MHA) 구조에서는 캐시 자체의 메모리 부담이 커지는 문제가 있습니다. MQA는 이러한 메모리 문제를 해결하기 위해 등장했습니다.

개념: 표준 MHA에서는 각 어텐션 헤드(Head)가 자신만의 독립적인 키(K) 및 밸류(V) 프로젝션 행렬을 가집니다. 즉, 헤드 수만큼 K, V 쌍이 존재합니다. 반면, MQA에서는 모든 쿼리(Query, Q) 헤드가 단 하나의 키(K) 헤드와 단 하나의 밸류(V) 헤드를 공유합니다.

비유: 여러 명의 기자가 각자 다른 질문(Q)을 던지지만, 모두가 단 한 대의 카메라(K)와 단 한 대의 마이크(V)를 통해 답변을 받는 상황과 유사합니다. 기자마다 개별 장비를 사용할 때보다 훨씬 적은 장비로도 여러 질문에 대한 정보를 수집할 수 있는 것입니다.

장점:

• 메모리 절감: 저장해야 할 KV 캐시의 크기가 헤드 수만큼 줄어들어 메모리 사용량과 메모리 대역폭 요구량을 크게 절감할 수 있습니다. 특히 긴 시퀀스에서 효과적입니다.
• 처리량 및 지연 시간 개선: MQA 기반 효율성 기술을 사용하면 추론 시 처리량이 최대 11배 향상되고 지연 시간이 30% 감소할 수 있다는 보고도 있습니다.

단점:

• 품질 저하 가능성: 모든 Q 헤드가 동일한 K, V 정보를 공유하기 때문에, MHA에 비해 표현력이 다소 떨어져 모델의 품질이 저하될 수 있습니다.

균형 잡힌 효율성 추구: 그룹-쿼리 어텐션 (Grouped-Query Attention, GQA)

MQA가 메모리 효율성은 뛰어나지만 품질 저하의 우려가 있다면, GQA는 이 둘 사이의 균형을 맞추려는 시도입니다.

개념: GQA는 표준 MHA와 MQA 사이의 절충안입니다. 여러 개의 쿼리(Q) 헤드를 몇 개의 그룹으로 나누고, 각 그룹 내의 Q 헤드들이 하나의 키(K) 헤드와 하나의 밸류(V) 헤드를 공유하는 방식입니다. 즉, MHA처럼 모든 Q 헤드가 독립적인 K, V를 갖지도 않고, MQA처럼 모든 Q 헤드가 단 하나의 K, V를 공유하지도 않습니다.

비유: 앞선 기자 비유를 다시 사용하자면, 여러 명의 기자를 몇 개의 취재팀(그룹)으로 나눕니다. 각 취재팀은 자신들만의 카메라(K)와 마이크(V)를 공유하여 사용합니다. 모든 기자가 단 한 대의 장비를 공유하는 것(MQA)보다는 다양한 각도에서 정보를 얻을 수 있고, 각 기자마다 개별 장비를 사용하는 것(MHA)보다는 효율적인 방식입니다.

장점:

• 성능과 효율성의 균형: MQA보다는 KV 캐시 크기가 크지만, 표준 MHA보다는 훨씬 작습니다. 이를 통해 메모리 사용량을 크게 줄이면서도, MQA에서 발생할 수 있는 품질 저하를 완화하여 MQA보다 더 나은 성능 균형을 보이는 경우가 많습니다.
• MHA에 가까운 품질, MQA에 가까운 속도: GQA는 MHA에 근접하는 품질을 유지하면서 MQA와 비슷한 수준의 추론 속도를 달성하는 것을 목표로 합니다.

마무리하며

KV 캐싱, 멀티-쿼리 어텐션, 그리고 그룹-쿼리 어텐션은 대규모 언어 모델의 추론 과정에서 발생하는 계산 및 메모리 병목 현상을 완화하는 데 핵심적인 역할을 합니다. KV 캐싱은 불필요한 반복 계산을 제거하여 속도를 높이고, MQA와 GQA는 KV 캐시의 메모리 부담을 줄여 더 긴 컨텍스트를 처리하거나 더 큰 모델을 효율적으로 운영할 수 있도록 돕습니다.

---

https://arxiv.org/abs/2305.13245

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

 

https://medium.com/@joaolages/kv-caching-explained-276520203249

KV Caching Explained