Grouped Value Attention, 키 온디맨드 재구성으로 KV Cache 압축

들어가며

긴 컨텍스트를 생성하는 Transformer에서는 모델 파라미터보다 KV Cache가 실제 추론 환경의 제약이 되기도 합니다. 토큰을 하나 생성할 때마다 새로운 Key와 Value가 캐시에 추가되고, 다음 디코딩 단계에서는 점점 길어지는 과거 캐시를 읽어야 하기 때문입니다. Grouped-query attention(GQA)은 여러 Query 헤드가 Key-Value 헤드를 공유하도록 해 비용을 낮췄지만, 각 위치의 Key와 Value를 모두 저장하는 구조까지 없애지는 못했습니다.

논문 「Grouped Value Attention: Efficient KV Caching via On-Demand Key Reconstruction」은 헤드 수를 줄이는 대신 캐시에 저장하는 표현 자체를 바꾸는 접근을 제시합니다.

핵심 내용

• 그룹화한 Value를 저장합니다. GVA는 Value를 그룹 단위로 구성하고, 기존 방식처럼 모든 위치의 완전한 Content Key를 지속적으로 저장하지 않습니다.
• Content Key를 필요할 때 재구성합니다. 학습된 선형 변환이 Value 표현과 Content Key 표현을 연결합니다. 추론 과정에서는 이 변환을 Query 변환에 흡수할 수 있어, 설계된 디코딩 경로에서 Content Key를 실제로 물질화하지 않아도 됩니다.
• 위치 정보는 별도 채널로 유지합니다. 내용 표현을 압축하면 위치 정보가 손실될 수 있으므로, 작고 공유되는 분리형 RoPE 채널을 사용하고 별도의 Positional Key를 캐시합니다. 실험에는 16차원 위치 채널 설정도 포함됐습니다.
• 지속 캐시를 줄입니다. 논문에서 평가한 구성 기준으로 GVA는 동일 조건의 GQA보다 지속 캐시 스칼라 수를 약 45~47% 줄였습니다. 다만 저장량 감소가 곧바로 지연 시간이나 처리량 향상을 의미하는 것은 아닙니다.
• GQA에 가까운 정확도를 보였습니다. 350M 파라미터 규모에서 30B FineWeb-Edu 토큰으로 학습한 실험의 경우, 16차원 위치 변형은 5개 작업 평균 정확도 44.18을 기록했습니다. GQA는 44.36, MLA는 43.88이었습니다.

의미와 한계

GVA의 핵심은 KV Cache에서 내용 정보와 위치 정보를 분리하고, 키 재구성의 일부를 Query 측 계산으로 옮기는 데 있습니다. 긴 시퀀스, 다중 요청 배치, GPU 메모리가 제한된 서비스 환경에서는 캐시 용량과 읽기 트래픽을 줄여 더 큰 배치나 긴 컨텍스트를 수용할 여지를 만들 수 있습니다.

그러나 캐시가 작아졌다고 해서 생성이 자동으로 빨라지는 것은 아닙니다. 재구성에 필요한 연산량, 메모리 접근 패턴, 커널 융합, 배치 크기, 하드웨어 활용률이 최종 성능을 결정합니다. 저자들은 맞춤형 디코딩 커널을 개발했으며 현재 엔드투엔드 추론 성능을 평가하고 있다고 밝혔습니다. 오픈소스 공개도 예정돼 있지만, 제공된 소재만으로는 구체적인 속도 향상을 확인할 수 없습니다.

후속 평가에서는 더 큰 모델과 더 긴 컨텍스트에서도 정확도와 메모리 절감 효과가 유지되는지, 그리고 커널 구현이 줄어든 캐시 트래픽을 실제 지연 시간 감소로 연결하는지가 중요합니다. 따라서 GVA는 가능성 있는 KV Cache 표현 설계로 볼 수 있지만, 이미 검증된 범용 추론 가속 기법으로 단정하기에는 추가 시스템 실험이 필요합니다.

출처: Hugging Face Daily Papers