Multi-head Latent Attention과 KV 캐시 압축

KV 캐시를 줄이는 길은 헤드를 지우는 쪽과 차원을 접는 쪽으로 갈리고, MLA는 뒤쪽을 택합니다.

헤드를 줄이는 압축이 부딪히는 벽

디코딩은 토큰을 하나씩 만들면서 앞선 토큰 전부의 Key와 Value를 다시 읽습니다. 그래서 KV 캐시는 요청이 끝날 때까지 GPU 메모리에 남아 있어야 합니다. 크기는 배치와 컨텍스트 길이에 선형으로 늘어나, 컨텍스트 128K에 배치 32 정도면 캐시만 수십 GB를 차지합니다.

토큰 하나가 차지하는 캐시 크기는 $2 L n_h d_h$입니다. $L$은 레이어 수, $n_h$는 어텐션 헤드 수, $d_h$는 헤드 차원입니다. 앞의 2는 Key와 Value 두 벌을 뜻합니다. 이 식에서 손댈 수 있는 항은 헤드 수와 헤드 차원, 둘뿐입니다.

GQA(Grouped-Query Attention)와 MQA(Multi-Query Attention)는 헤드 수 쪽을 건드립니다. 여러 쿼리 헤드가 Key와 Value 한 벌을 공유하게 만들어 $n_h$를 그룹 수 $n_g$로 낮춥니다. 캐시는 비례해서 줄지만, 한 그룹 안에서는 헤드마다 다른 특징을 잡을 수 없습니다. 압축률과 품질이 같은 손잡이에 묶여 있는 구조입니다.

MLA(Multi-head Latent Attention)는 남은 축인 차원 쪽을 접습니다. 헤드 수는 그대로 두고, 캐시에 넣는 벡터의 차원을 낮춥니다. 표에서 MLA 행에만 2가 붙지 않은 이유는 Key와 Value를 각각 저장하지 않기 때문입니다.

잠재 벡터 하나만 캐시한다

MLA는 DeepSeek-V2(arXiv:2405.04434, 2024)가 도입한 어텐션 변형입니다. Key와 Value를 그대로 보관하는 대신, 입력 은닉 상태를 저차원 잠재 벡터 $c$로 눌러 그것만 캐시합니다. Key와 Value는 필요할 때 $c$에서 되살립니다.

프리필에서 다운프로젝션 행렬 $W_d$가 은닉 상태를 $d_c$차원으로 압축합니다. 어텐션 계산에 쓰이는 Key와 Value는 업프로젝션 행렬 $W_u$가 $c$를 펼쳐서 만듭니다. 두 행렬은 학습으로 얻은 저랭크 구조이므로, 압축과 복원이 고정된 규칙을 따릅니다.

• 캐시에 들어가는 것은 잠재 벡터 $c$와 위치 키 $k_r$ 두 조각뿐입니다.
• 헤드마다 다른 것은 업프로젝션 행렬이지 캐시가 아니므로, 헤드 수가 캐시 크기에 곱해지지 않습니다.
• DeepSeek-V2 설정은 $d_c = 512$, $d_r = 64$이며, 토큰당 캐시가 MHA 대비 93.3% 줄었습니다.
• 위치 정보를 별도 스트림으로 뺀 이유는 다음 절에서 다룹니다.

압축률의 출처는 차원 축소 한 번입니다. 헤드별로 흩어져 있던 Key와 Value가 하나의 잠재 벡터로 모이면서, 캐시 크기가 $2 n_h d_h$에서 $d_c + d_r$로 내려갑니다.

흡수 트릭과 분리형 RoPE

잠재 벡터만 캐시해도, 스텝마다 Key와 Value를 복원하면 아낀 메모리를 계산 비용으로 되갚게 됩니다. 흡수 트릭(absorption trick)이 이 되갚음을 없앱니다. 어텐션 점수는 쿼리와 키의 내적인데, 키가 $k_j = W_u c_j$이므로 점수는 다음과 같이 다시 쓸 수 있습니다.

$q_t^\top k_j = q_t^\top W_u c_j = (W_u^\top q_t)^\top c_j$

$W_u$는 학습이 끝나면 고정 행렬입니다. 그래서 쿼리 투영 행렬에 미리 곱해 둘 수 있고, 추론 중에는 캐시에서 읽은 $c_j$와 흡수된 쿼리를 바로 내적하면 됩니다. Key를 메모리에 물질화(materialize)하는 단계 자체가 사라집니다. Value 쪽 업프로젝션도 같은 방식으로 출력 투영 행렬에 흡수됩니다.

이 트릭은 추론에서만 성립합니다. 학습 중에는 Key와 Value를 명시적으로 만들어 역전파해야 하므로 흡수를 쓸 수 없고, MLA의 학습 메모리는 MHA와 비슷합니다. 흡수 후 쿼리 측 행렬이 커지면 프리필 단계 연산량이 늘어날 수 있다는 점도 함께 봐야 합니다.

RoPE(Rotary Position Embedding)는 여기에 정면으로 걸립니다. RoPE는 토큰 위치 $j$에 따라 달라지는 회전 행렬 $R_j$를 키에 곱하므로, 키가 $R_j W_u c_j$가 됩니다. $R_j$가 위치마다 바뀌니 $W_u$만 떼어 쿼리 쪽으로 옮길 수 없습니다.

MLA는 키를 두 스트림으로 나눠 이 충돌을 피합니다.

• NoPE 스트림: 위치와 무관한 의미 정보. 잠재 벡터 $c$로 압축해 캐시하고, 여기서만 흡수를 적용합니다.
• RoPE 스트림: 위치 정보만 담은 $d_r$차원 키. 회전을 적용한 상태로 작은 별도 캐시에 둡니다.

최종 점수는 두 스트림의 내적을 더해서 만듭니다. 위치 정보를 잠재 벡터에 통합했다면 위치마다 다른 잠재 벡터를 만들어야 해서 압축 이점이 사라지고, 잠재 공간의 표현이 위치에 종속돼 학습이 불안정해집니다. 분리형 RoPE는 위치 의존성을 작은 캐시로 격리해 나머지 전부를 압축 대상으로 남깁니다.

같은 캐시 예산에서 갈리는 표현력

MLA와 GQA를 나란히 놓을 때 던져야 할 질문은 캐시 크기가 아니라 같은 캐시로 무엇을 표현할 수 있느냐입니다. TransMLA(arXiv:2502.07864, 2025)는 Theorem 1로 답을 내놓습니다. 같은 KV 캐시 예산에서 MLA의 표현력은 GQA보다 엄격히 큽니다. 모든 GQA는 동등한 MLA로 변환할 수 있지만, 역방향은 일반적으로 불가능합니다.

구조를 보면 이유가 드러납니다. GQA는 헤드끼리 Key와 Value를 그대로 나눠 쓰므로, 한 그룹에 묶인 헤드들은 같은 캐시를 같은 방식으로 읽습니다. MLA는 잠재 공간에서 학습한 저랭크 구조를 헤드마다 다른 업프로젝션으로 펼칩니다. 공유되는 정보와 헤드별 정보가 같은 캐시 안에 함께 담깁니다.

실측도 이 방향입니다. 같은 논문은 Qwen2.5 7B와 14B를 GQA에서 MLA로 변환했을 때 수학·코딩 태스크 정확도가 올랐고, 파라미터 증가는 1.3% 미만이었다고 보고합니다. 캐시 예산은 늘어나지 않았습니다.

대가는 계산 쪽에 있습니다. MLA는 학습 중 업프로젝션 연산 비용이 GQA보다 조금 높고, 앞서 본 대로 흡수된 쿼리 행렬이 커지면 프리필 연산량도 늘어납니다. 메모리를 줄인 만큼을 연산으로 지불하는 교환이며, 디코드가 메모리 대역폭에 묶여 있는 서빙 환경에서 이 교환이 유리해집니다.

학습된 MHA 모델을 옮기는 길

MLA의 이점을 쓰려고 모델을 처음부터 다시 학습할 필요는 없습니다. MHA2MLA(arXiv:2502.14837, 2025)는 이미 학습된 MHA 가중치를 SVD로 분해해 MLA 구조로 옮깁니다. Key와 Value의 투영 행렬을 이어 붙여 한 번에 분해하는 SVDjoint 방식이 권장되며, 둘을 따로 분해하는 방식보다 0.91~1.38% 성능이 앞섭니다.

Key와 Value를 함께 분해하는 편이 나은 이유는 두 행렬이 같은 잠재 벡터를 공유해야 하기 때문입니다. 개별 분해는 각자에게 최적인 부분공간을 찾지만, 공유 캐시에서 되살릴 때는 그 최적이 유지되지 않습니다.

양자화를 얹으면 압축률이 더 올라갑니다. 같은 연구는 Int4HQQ 양자화를 함께 적용해 KV 캐시를 92.19%까지 줄였습니다. DeepSeek-V3 기준으로 토큰당 70 KB이며, GQA 모델의 192328 KB보다 2.74.7배 효율적이라고 보고합니다.

잠재 벡터를 양자화하면 그 오차가 업프로젝션을 지나며 Key와 Value 전체로 증폭됩니다. 헤드별 스케일 보정과 교정 데이터셋 기반 캘리브레이션이 필요하고, 압축률이 높은 만큼 다운스트림 태스크별 검증도 함께 가야 합니다.

분리형 서빙에서 줄어드는 전송량

프리필과 디코드를 다른 GPU 풀에 두는 분리형(disaggregated) 서빙에서는 KV 캐시가 네트워크 페이로드가 됩니다. 캐시 크기가 곧 전송량이므로, 어떤 어텐션 변형을 쓰느냐가 대역폭 요구로 직결됩니다.

규모 감각은 KVServe(arXiv:2605.13734, SIGCOMM 2026)의 측정으로 잡을 수 있습니다. Qwen3-235B에서 32K 토큰 요청 하나가 수 GB의 KV 캐시를 만들고, 이를 제때 옮기려면 2.1 Tbps의 송출 대역폭이 필요합니다. 일반적인 클라우드의 클러스터 간 대역폭은 100 Gbps 이하입니다. 같은 실험에서 KV 통신 시간이 작업 완료 시간(Job Completion Time, JCT)의 최대 60%를 차지했습니다.

이 격차에 대한 대응은 두 갈래입니다. 만들어진 캐시를 사후 압축하는 쪽과, 애초에 작게 만드는 쪽입니다. KVServe는 앞쪽으로, 대역폭·지연 SLO·정확도 요건을 보고 변환·양자화·엔트로피 코딩 조합을 요청마다 골라 TTFT를 최대 32.8배, JCT를 최대 9.13배 줄였습니다. MLA는 뒤쪽이며, 압축이 모델 구조 안에 들어가 있으므로 전송 경로에 추가 지연을 넣지 않습니다.

다만 MLA는 분리형 서빙에 새로운 선택지를 하나 만듭니다. 업프로젝션을 어디서 수행할지가 열려 있기 때문입니다. 프리필 노드에서 Key와 Value를 복원해 보내면 전송량이 MHA 수준으로 되돌아가고, 디코드 노드에서 복원하면 전송은 작지만 디코드 쪽 연산이 늘어납니다. TPLA(arXiv:2508.15881, 2025)는 MLA와 텐서 병렬 분리를 결합하면 KV 전송 비용이 MHA 기반 분리 서빙 대비 크게 줄어든다고 보고합니다. 같은 연구는 그 자리에 지연 대 대역폭 교환이 새로 생긴다는 점도 지적합니다.

정리

MLA는 헤드 수를 깎는 대신 캐시에 넣는 벡터의 차원을 접어 KV 캐시를 줄입니다. 업프로젝션 행렬이 학습 후 고정이라는 성질을 이용해 쿼리 쪽으로 흡수하므로, 추론 중에는 Key와 Value를 복원하지 않고 잠재 벡터만 읽습니다. 위치에 따라 변하는 RoPE는 이 흡수와 맞지 않아 작은 별도 스트림으로 분리했고, 덕분에 나머지 전부가 압축 대상으로 남았습니다. DeepSeek-V2 기준 토큰당 캐시가 MHA 대비 93.3% 줄었고, TransMLA는 같은 캐시 예산에서 MLA의 표현력이 GQA보다 엄격히 크다는 것을 증명했습니다. 분리형 서빙에서는 이 압축이 그대로 전송량 감소로 이어지지만, 업프로젝션을 프리필과 디코드 중 어디에서 수행할지에 따라 지연과 대역폭 사이의 새 교환을 결정해야 합니다.