하드웨어-적응형 대형 언어 모델 디코딩을 위한 Group-Query Latent Attention(GQLA)

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

대형 언어 모델의 autoregressive 디코딩은 KV 캐시 트래픽에 의해 병목이 발생한다. MLA는 KV를 잠재 공간으로 압축하나 단일 MQA-absorb 경로에 묶여 하드웨어 비대칭과 head-axis 텐서 병렬성 제약을 초래한다. GQLA는 같은 가중치를 유지하면서 업프로젝션을 그룹으로 인덱싱해 두 개의 해석 가능한 디코딩 경로를 노출하고, 런타임에서 대상 하드웨어에 맞춰 경로를 선택하도록 한다.

왜 중요한가

대형 언어 모델의 autoregressive 디코딩은 KV 캐시 트래픽에 의해 병목이 발생한다. MLA는 KV를 잠재 공간으로 압축하나 단일 MQA-absorb 경로에 묶여 하드웨어 비대칭과 head-axis 텐서 병렬성 제약을 초래한다. GQLA는 같은 가중치를 유지하면서 업프로젝션을 그룹으로 인덱싱해 두 개의 해석 가능한 디코딩 경로를 노출하고, 런타임에서 대상 하드웨어에 맞춰 경로를 선택하도록 한다.

핵심 기여

Two algebraically equivalent decoding paths from a single weight set

GQLA 가중치는 두 가지 디코딩 경로(MQA-absorb 경로, GQA 경로)를 동일 가중치로 노출한다. 이로써 H100의 MQA-absorb 경로와 H20의 GQA 경로를 모두 충족하는 roofline에 맞춘 디코딩이 가능해지며 재학습 없이 하드웨어에 맞춘 실행이 가능하다.

Group-indexed up-projections enabling 8-way tensor parallelism

업프로젝션을 query-head가 아닌 group 단위로 인덱싱함으로써 8-way tensor-parallelism을 가능하게 한다. 이로써 GQA 경로에서 그룹-별 캐시를 확장하고, TP 재현성을 보존한다.

TransGQLA: 변환 기반의 두 경로 변환 파이프라인

GQA 체크포인트와 MLA 체크포인트를 최소 변경으로 GQLA로 변환하는 두 가지 경로를 제시한다. 3.2.1은 head-merging 구성의 단일 변경으로 GQA-path를 유지하고, 3.2.2는 MLA latent로부터 그룹-별 PCA를 통해 요인을 복구한다.

Sparse GQLA 확장

fine-grained sparse attention에서 GQA 경로를 유지하고, H20-급 하드웨어에서 MMA 타일과의 매핑을 통해 성능을 보존한다. 이로써 스파스 어텐션에서도 GQA 경로의 이점을 활용한다.

Roofline 분석에 의한 하드웨어-적합성 검증

H100과 H20에서 동일한 가중치로 두 경로가 로프라인(roofline) 상의 두 지점을 모두 달성함을 분석으로 확인한다. GQLA의 GQA 경로는 H20의 ridge를, MQA-absorb 경로는 H100의 ridge를 각각 달성한다.

핵심 아이디어 이해하기

단계 1: Transformer의 KV 캐시는 디코딩 시점까지 누적되며, KV 캐시 트래픽이 성능 병목의 주된 원인이다. 단계 2: MLA는 KV를 저랭크 잠재 공간으로 압축하되 두 개의 실행 경로를 동일 가중치로 노출하지 못해 하드웨어 의존성과 텐서 병렬성의 한계를 가진다. 단계 3: GQLA는 up-projections를 group 단위로 인덱싱하고, MQA-absorb 경로를 유지하는 동시에 GQA 경로를 캐시 확장으로 구현한다. 결과적으로 같은 가중치 세트로 H100의 MQA-absorb와 H20의 GQA 경로를 모두 지원하고, 8-way Tensor Parallelism을 가능하게 한다. 또한 TransGQLA를 통해 기존 GQA/MLA 체크포인트를 추가 학습 없이 변환할 수 있다.

방법론

입력 xt에 대해 WQ, WKV, WKU, WVV 등의 가중치를 사용하여 c^KV_t를 rkv 차원으로 압축하고, q^C_t, k^C_t, q^R_t, k^R_t를 각각 분리된 RoPE 기반으로 구성한다. 두 경로는 동일한 c^KV_t를 소비하되, (1) GQA 경로는 gKV 그룹과 per-group KC/V 캐시를 사용하고, (2) MQA-absorb 경로는 WUK, WUV를 재배치해 잠재 캐시를 직접 Q/O에 흡수한다. RoPE를 그룹별로 정렬하는 RoRoPE를 적용해 그룹 간 RoPE 기반의 각도 차이를 보정하고, FreqFold로 RoPE 주파수 대역별 PCA를 수행한다. K 및 V의 NoPE와 V를 균일한 Frobenius 노름으로 재스케일링하고, 이를 WUK, WUV에 흡수하여 rkv로 차원을 축소한다. MLA에서 GQLA로의 이전은 per-group PCA를 통해 각 그룹의 K/V를 독립적으로 분해하고, U, V, O의 변환을 재활용하는 방식으로 수행된다.

주요 결과

주요 벤치마크에서의 성능 및 효율 변화가 보고된다. 예시로 LLaMA-3-8B에서 0-token 변환 시 GQA→GQLA는 ∼9.7포인트 Avg. 감소를 보였고(GLM-4.7의 0-token 변환은 ∼4.7 Avg. 포인트 감소), 같은 Latent에서의 재정렬은 30B 토큰에서 0.5포인트 이내로 회복 가능하다고 제시된다. 벤치마크별 자세한 수치는 표 3에 제시되며, H100에서 sq=1의 MQA-absorb는 2.82 µs/step, sq=2는 4.61 µs/step로 나타난다. H20의 경우 GQA 경로의 (g, sq) 조합에서 ridge를 확보하고, 221K tok/s 수준으로 3.4배 이상의 개선을 보인다. 4.2.1~4.2.3의 Roofline 분석에 따르면 두 경로가 각각 H100/H20에서 이론적 정점에 근접한다. Sparse GQLA는 m=16 MMA 타일에 정확히 맞춰 Tensor Core의 효율성을 유지한다.

기술 상세

전체 아키텍처: Group-Query Latent Attention(GQLA)은 입력 토큰 임베딩 x_t에 대해 WQ, WKV를 이용해 c^KV_t를 rkv 차원의 잠재 벡터로 압축하고, WDQ, WUQ를 통해 q^C_t, q^R_t를 구성한다. RoPE는 decoupled RoPE 방식으로 q^R_t, k^R_t에 적용되어 각 그룹 간 RoPE 기준을 맞춘다. 두 경로는 (a) GQA 경로: g KV 그룹의 KC/V를 확장 캐시로 사용하고, (b) MQA-absorb 경로: WUK, WUV를 재구성해 latent를 직접 Q/O에 흡수한다. RoRoPE를 통해 각 그룹의 RoPE 베이스를 하나의 공통 베이스로 정렬하고, FreqFold로 RoPE 주파수 대역별 PCA를 수행해 NoPE 부분(c^KV_t)과 RoPE 부분(k^R_t)을 분리한다. K/V의 합성 인자 NoPE와 V의 인자를 균등하게 정규화하고 WUK, WUV에 반영한다. TransGQLA는 (i) GQA-preserving merge를 통한 변환(머지 단계의 head-merging에서 group-indexed 업프로젝션 유지)과 (ii) MLA latent로부터 per-group PCA를 통한 요인 재구성으로 구현된다. 배포 시 두 경로는 동일한 가중치를 사용하므로, Absorb 경로와 PCA 경로 중 하나를 선택하는 캐시-압축/확장 한 번의 작업으로 구현 가능하다.

실무 활용

TransGQLA는 사전학습된 GQA/MLA 체크포인트를 추가 훈련 없이 GQLA로 변환할 수 있어 하드웨어에 따른 디코딩 경로를 동적으로 선택하는 실용적 구현을 가능하게 한다.

H100/H20 급 하드웨어에서 KV 캐시-지향 디코딩의 roofline 최적화를 달성하는 프로덕션 파이프라인에 적용
Sparse attention 시에도 GQA 경로를 유지하여 Tensor Core MMA 타일과의 호환성 확보
두 경로를 동일 가중치로 유지하므로 재훈련 없이 다양한 하드웨어로 이식

코드 공개 여부: 공개

코드 저장소 보기

키워드

GQLA (Group-Query Latent Attention)MLA (Multi-head Latent Attention)GQA (Generalized Multi-Query Attention)RoRoPEFreqFoldPCAKV cacheH100H20