Grouped Query Attention (GQA)

그룹화 쿼리 어텐션 (GQA)은 여러 쿼리 헤드가 더 적은 수의 키 및 값 헤드를 공유하는 어텐션 설계입니다. 키와 값에 필요한 메모리 및 데이터 이동을 줄이면서 멀티 헤드 쿼리의 다양한 관점을 유지합니다. GQA는 모델이 한 번에 하나의 토큰을 생성하고 저장된 어텐션 상태를 반복해서 읽는 자기회귀 추론에서 특히 유용합니다.

GQA 작동 방식#

각 토큰은 어텐션 메커니즘 내에서 세 가지 표현으로 프로젝션됩니다:

• 쿼리: 현재 토큰이 찾고 있는 정보를 나타냅니다.
• 키: 사용 가능한 각 토큰이 나타내는 내용을 설명합니다.
• 값: 쿼리가 키와 일치할 때 검색되는 정보를 포함합니다.

기존 셀프 어텐션에서는 모델이 서로 다른 관계를 병렬로 학습할 수 있도록 이러한 프로젝션을 헤드로 나눕니다. GQA는 여러 쿼리 헤드를 유지하면서 그중 여러 개를 공유 키 및 값 헤드에 할당합니다. 예를 들어 어텐션 레이어에서 쿼리 헤드 8개와 키-값 헤드 2개를 사용할 수 있습니다. 그러면 각 키-값 헤드가 쿼리 헤드 4개를 담당합니다.

어텐션 계산 자체는 여전히 스케일 조정된 내적 어텐션입니다. 변경되는 부분은 헤드 배열과 생성, 저장 및 읽는 키-값 데이터의 양입니다. 따라서 PyTorch 스케일 조정 내적 어텐션과 같은 프레임워크에서는 GQA를 활성화할 때 쿼리 헤드 수가 키-값 헤드 수로 나누어떨어져야 합니다. (docs.pytorch.org)

GQA와 멀티 헤드 및 멀티 쿼리 어텐션 비교#

GQA는 멀티 헤드 어텐션과 멀티 쿼리 어텐션의 중간에 해당합니다:

• 멀티 헤드 어텐션: 모든 쿼리 헤드가 자체 키 및 값 헤드를 가집니다. 헤드 독립성이 가장 높지만 키-값 메모리 요구량도 가장 큽니다.
• 그룹화 쿼리 어텐션: 여러 쿼리 헤드가 각 키-값 헤드를 공유합니다. 표현 용량과 메모리 효율 사이의 균형을 제공합니다.
• 멀티 쿼리 어텐션: 모든 쿼리 헤드가 하나의 키 헤드와 하나의 값 헤드를 공유합니다. 키-값 저장 공간은 최소화하지만 키-값 다양성은 낮아집니다.

레이어에 쿼리 헤드가 32개 있다면 멀티 헤드 어텐션에서도 키-값 헤드 32개를 사용할 수 있고, GQA에서는 8개를 사용할 수 있으며, 멀티 쿼리 어텐션에서는 1개를 사용합니다. GQA는 더 넓은 Transformer 아키텍처를 대체하는 것이 아니라, Transformer 어텐션 레이어에서 사용할 수 있는 하나의 구성입니다. NVIDIA의 멀티 헤드, 멀티 쿼리 및 그룹화 쿼리 어텐션 문서에서는 이러한 변형이 주로 쿼리 헤드를 담당하는 키-값 헤드의 수에 따라 달라진다고 설명합니다. (nvidia.github.io)

GQA가 중요한 이유#

자기회귀 생성 중에 모델은 이전 키와 값을 KV 캐시에 저장합니다. 표준 멀티 헤드 어텐션에서는 모든 어텐션 헤드가 별도의 캐시 키와 값을 생성합니다. 따라서 긴 시퀀스, 큰 배치, 많은 모델 레이어는 상당한 GPU 메모리를 사용할 수 있습니다.

GQA는 키-값 헤드 수를 줄이므로 캐시도 해당 헤드 감소 비율에 맞춰 작아집니다. 이를 통해 다음 항목을 개선할 수 있습니다:

• 메모리 용량: 더 긴 프롬프트나 더 많은 동시 요청을 사용 가능한 메모리에 수용할 수 있습니다.
• 생성 처리량: 생성되는 각 토큰에 대해 읽어야 하는 키-값 데이터가 줄어듭니다.
• 추론 지연 시간: 메모리 트래픽이 감소하여 토큰 간 처리 시간이 단축될 수 있습니다.
• 컨텍스트 윈도우 확장성: 긴 시퀀스를 더 실용적으로 서비스할 수 있습니다.

이러한 이점은 하드웨어, 시퀀스 길이, 배치 크기 및 커널 지원 여부에 따라 달라집니다. 프로덕션 런타임에서는 여전히 효율적인 캐시 할당이 필요합니다. 예를 들어 NVIDIA의 TensorRT-LLM KV 캐시 시스템은 GQA 지원을 블록 기반 캐싱, 재사용 및 오프로딩과 결합합니다. (nvidia.github.io)

실제 적용 사례#

긴 컨텍스트 어시스턴트: 코딩 어시스턴트는 답변을 생성하기 전에 수천 개의 소스 코드 토큰을 처리할 수 있습니다. GQA는 해당 기록과 연결된 캐시 키-값 상태를 줄여 서버가 GPU 메모리를 비례적으로 늘리지 않고도 더 긴 파일이나 더 많은 동시 사용자를 처리할 수 있도록 합니다.

멀티모달 이미지 및 비디오 어시스턴트: 비전-언어 모델은 이미지 패치나 비디오 프레임을 긴 토큰 시퀀스로 표현할 수 있습니다. GQA는 이러한 시각 토큰이 언어 디코더에 입력된 후 어텐션 캐시 부담을 줄여 추가 이미지, 프레임 또는 요청을 위한 메모리를 더 많이 확보할 가능성이 있습니다. 이는 표준 Vision Transformer와는 다릅니다. 표준 Vision Transformer에서는 자기회귀 KV 캐싱 없이 모든 이미지 패치를 병렬로 처리할 수 있습니다.

실제 구현 및 절충점#

PyTorch는 enable_gqa=True을 통해 GQA를 제공합니다. 다음 문서화된 CUDA 예제에서는 쿼리 헤드 8개와 공유 키-값 헤드 2개를 사용합니다:

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

assert torch.cuda.is_available(), "A CUDA device is required."

query = torch.randn(1, 8, 16, 32, device="cuda")
key = torch.randn(1, 2, 16, 32, device="cuda")
value = torch.randn(1, 2, 16, 32, device="cuda")

# Four query heads share each key-value head.
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        enable_gqa=True,
    )

print(output.shape)

출력에는 쿼리 헤드 8개가 유지되는 반면, 키와 값에는 헤드 2개만 사용됩니다. PyTorch 어텐션 백엔드 선택기를 통해 작업을 수행할 지원 커널을 제어할 수 있으며, PyTorch Transformer 구성 요소 가이드에서는 더 폭넓은 구현 맥락을 제공합니다. (docs.pytorch.org)

GQA는 모델을 설계하거나 조정할 때 선택해야 하며, 기존 체크포인트에 대해 일반적으로 추론 시간에 전환하는 옵션이 아닙니다. 개발자는 프레임워크 지원, 헤드 나눗셈 가능 여부, 수치적 동작, 메모리 사용량 및 작업 품질을 확인해야 합니다. 컴퓨터 비전 배포에서는 Ultralytics TensorRT 통합 및 Ultralytics 벤치마크 모드와 같은 보완적인 워크플로를 통해 아키텍처 및 런타임 최적화가 대상 하드웨어에서 의미 있는 개선을 제공하는지 측정할 수 있습니다.