Grouped Value Attention: Efficient KV Caching via On-Demand Key Reconstruction

The KV cache is a primary bottleneck for Transformer decoding: its memory footprint and cache-read traffic grow with sequence length. Grouped-query attention (GQA) reduces this cost by sharing key-value heads, but still stores both a key and a value at every step. We introduce Grouped Value Attention (GVA), which stores grouped values and reconstructs content keys with a learned linear map. At inference, the map can be absorbed into the query, eliminating the need to materialize content keys in the intended decode path. A small shared decoupled RoPE channel retains positional information through a separately cached positional key. For the configurations studied, this representation reduces persistent cache scalars by approximately 45-47% relative to matched GQA. At the 350M-parameter scale with 30B FineWeb-Edu tokens, the 16-dimensional positional variant reaches 44.18 average accuracy across five tasks, compared with 44.36 for GQA and 43.88 for MLA. These results demonstrate near-GQA benchmark accuracy with a more compact cache representation. To translate this compact representation into faster autoregressive inference, we have developed custom decoding kernels and are currently evaluating their end-to-end inference performance with an open-source release planned soon.

한 줄 요약

GVA는 KV 캐시 공간을 45-47% 절감하면서 GQA 수준의 성능을 유지하는 새로운 어텐션 메커니즘이다.

핵심 기여도

GVA는 `K = V M` 형태의 학습된 선형 매핑을 통해 키를 재구성하고, 이를 쿼리에 흡수하여 KV 캐시 공간을 45-47% 절감.
`16차원 positional variant`를 사용한 GVA는 350M 파라미터, 30B FineWeb-Edu 토큰에서 44.35 평균 정확도를 달성 (GQA: 44.36, MLA: 43.88).
`decoupled RoPE` 채널을 도입해 위치 정보를 별도의 캐시 없이 유지.
`custom decoding kernels`를 개발하여 추론 성능 향상을 목표로 하고 있음.

핵심 아이디어

기존 GQA는 키-벨류 헤드를 그룹화하여 KV 캐시 공간을 절감하지만, 여전히 키와 벨류를 모두 저장해야 한다. GVA는 이 문제를 해결하기 위해 벨류만 저장하고, 키는 `K = V M` 형태의 학습된 선형 매핑을 통해 재구성한다. 이 매핑은 추론 시 쿼리에 흡수되어, 별도의 키 캐시 생성이 불필요하다. 이는 `value-only content cache`를 가능하게 하며, KV 캐시 공간을 약 절반으로 줄인다.

또한, 재구성된 키에 표준 RoPE를 적용하면 흡수가 불가능하다. 이를 해결하기 위해 `small shared decoupled RoPE channel`을 도입하여, 위치 정보를 별도의 캐시 없이 유지한다. 이는 `unrotated content slice`와 `short rotated slice`를 결합한 형태로, `Td_r` 차원만 추가된다.

기술적 접근법

**GVA 메커니즘**: `K = V M` 형태의 선형 매핑을 통해 키를 벨류에서 재구성.
**decoupled RoPE**: 위치 정보를 유지하기 위해 `unrotated content slice`와 `shared rotated slice`를 결합.
**캐시 크기 비교**: GQA는 `2TGd_h`, GVA는 `TGd_h + Td_r` (Td_r: decoupled RoPE 차원).
**모델 규모**: 350M 파라미터, 30B FineWeb-Edu 토큰.
**평가 지표**: HellaSwag, WinoGrande, OpenBookQA, ARC 등 5개 태스크의 평균 정확도.

주요 결과

GVA는 350M 파라미터, 30B FineWeb-Edu 토큰에서 44.35 평균 정확도 달성 (GQA: 44.36, MLA: 43.88).
KV 캐시 공간은 GQA 대비 약 45-47% 절감.
`16-dimensional positional variant`에서 GQA와 유사한 정확도를 유지하면서도 캐시 공간을 줄임.

의의 및 한계

GVA는 KV 캐시 공간을 대폭 줄이며 GQA 수준의 성능을 유지함으로써, 메모리 효율적인 추론을 가능하게 한다. 특히, `value-only content cache`와 `decoupled RoPE`는 기존 어텐션 메커니즘과 비교해 독창적인 접근법이다.

하지만, 현재 연구는 `prefill` 단계의 성능 개선 여부는 명시되지 않았으며, `custom decoding kernels`의 실제 추론 성능은 아직 평가 중이다. 또한, `RoPE width`, `모델 규모`, `랜덤 시드`에 따른 영향은 추가 실험 필요.

실용적 활용

GVA는 메모리 효율적인 추론이 요구되는 대규모 언어 모델, 특히 자동 생성 시스템(예: 챗봇, 번역, 요약)에 적용 가능하다. `KV 캐시 절감`은 장문 처리 시 성능 향상과 에너지 효율 향상에 기여할 수 있다.