멀티 헤드 잠재 어텐션

Multi-head latent attention, MLA; 멀티 헤드 잠재 어텐션
키와 값을 저차원 잠재 벡터 하나로 함께 압축해 KV 캐시에는 그 벡터만 저장하고, 추론 때는 복원 행렬을 쿼리·출력 투영에 흡수해 압축된 상태로 어텐션을 계산하는 방식

멀티 헤드 잠재 어텐션(Multi-head latent attention, MLA)은 딥시크가 2024년 5월 DeepSeek-V2에서 도입한 어텐션 구조다.^[1] 멀티 헤드 어텐션(MHA)은 토큰마다 헤드별 키·값을 모두 KV 캐시에 저장하고, 그룹 쿼리 어텐션(GQA, MQA)은 키·값 헤드 수를 줄여 캐시를 줄인다. MLA는 헤드 수를 줄이지 않고, 키와 값을 함께 만들어 낼 수 있는 저차원 잠재 벡터 $c^{K V}$ 하나만 저장한다. 그리고 추론할 때 키·값을 실제로 복원하지 않도록 복원 행렬을 다른 투영 행렬에 흡수(absorb)한다.

DeepSeek-V2 설정에서 MLA의 토큰·층당 캐시는 576개 원소로, 헤드 128개짜리 MHA(32,768개)의 약 1.8%이자 그룹 2.25개짜리 GQA와 같은 크기다. 그런데도 딥시크의 비교 실험에서는 MHA보다 성능이 좋았다.^[1] DeepSeek-V3·R1이 같은 구조를 이어 썼고, 문샷 AI의 키미 K2 등 다른 회사 모델도 채택했다. 이 문서는 딥시크 문서의 MLA 절보다 한 단계 들어가, 압축·흡수·분리 RoPE의 수식과 구현을 다룬다.

동기

긴 문맥 서비스에서 GPU 메모리를 가장 많이 차지하는 것은 KV 캐시이고, 디코딩 속도는 이 캐시를 읽는 메모리 대역폭에 묶인다. GQA·MQA는 캐시를 줄이지만 품질을 잃는다. 딥시크가 같은 조건(1.33조 토큰)으로 학습한 약 70억 매개변수 밀집 모델 비교에서 MMLU는 MHA 45.2, GQA(그룹 8) 41.2, MQA 37.9였다.^[1] 딥시크의 문제의식은 "캐시는 MQA 수준으로 작게, 표현력은 MHA 이상으로"였고, 그 답이 헤드를 줄이는 대신 키·값의 정보를 저계수(low-rank)로 압축하는 것이었다.

구조

기호: 모델 차원 $d$ , 헤드 수 $n_{h}$ , 헤드 차원 $d_{h}$ , 층 수 $l$ , 토큰 $t$ 의 어텐션 입력 $h_{t} \in ℝ^{d}$ .

키·값 공동 압축

$c_{t}^{K V} = W^{D K V} h_{t}, k_{t}^{C} = W^{U K} c_{t}^{K V}, v_{t}^{C} = W^{U V} c_{t}^{K V}$

$c_{t}^{K V} \in ℝ^{d_{c}}$ ( $d_{c} ≪ n_{h} d_{h}$ )가 압축 잠재 벡터, $W^{D K V} \in ℝ^{d_{c} \times d}$ 가 압축(down) 투영, $W^{U K}, W^{U V} \in ℝ^{n_{h} d_{h} \times d_{c}}$ 가 복원(up) 투영이다. 캐시에는 $c_{t}^{K V}$ 만 넣는다. 키와 값이 같은 잠재 벡터에서 나오므로 "공동"(joint) 압축이다.^[1]

학습 때 활성값 메모리를 줄이려고 쿼리도 같은 방식으로 압축한다( $c_{t}^{Q} = W^{D Q} h_{t}$ , $q_{t}^{C} = W^{U Q} c_{t}^{Q}$ ). 쿼리는 캐시에 들어가지 않으므로 캐시 크기와는 무관하다.^[1]

행렬 흡수

헤드 $i$ 의 점수는 $q_{t, i}^{⊤} k_{j, i} = (W_{i}^{U Q} c_{t}^{Q})^{⊤} (W_{i}^{U K} c_{j}^{K V}) = (c_{t}^{Q})^{⊤} (W_{i}^{U Q ⊤} W_{i}^{U K}) c_{j}^{K V}$ 이다. 괄호 안 행렬은 위치와 상관없는 상수이므로 미리 곱해 둘 수 있다. 곧 쿼리를 잠재 공간으로 옮긴 뒤 캐시된 $c_{j}^{K V}$ 와 바로 내적하면 되고, 키 $k_{j}$ 를 복원할 필요가 없다. 값 쪽도 마찬가지로 $W^{U V}$ 를 출력 투영 $W^{O}$ 에 흡수해, 어텐션 가중치로 $c_{j}^{K V}$ 를 가중합한 다음 한 번에 투영한다.^[1]

흡수하면 모든 쿼리 헤드가 같은 $c^{K V}$ 를 키이자 값으로 보게 된다. 구조상 MQA와 같은 모양이며, 딥시크는 이를 MLA의 "MQA 모드", 키·값을 복원해 헤드별로 계산하는 쪽을 "MHA 모드"라 부른다. DeepSeek-V3.1-Terminus는 학습과 프리필에 MHA 모드, 디코딩에 MQA 모드를 썼다.^[2] 긴 프롬프트를 한 번에 처리할 때는 헤드별 차원이 작은 MHA 모드가 연산량이 적고, 토큰 하나씩 생성할 때는 캐시를 읽는 양이 적은 MQA 모드가 유리하기 때문이다.

분리 RoPE

RoPE를 키에 그대로 적용하면 흡수가 깨진다. 키가 $R_{j} W^{U K} c_{j}^{K V}$ 가 되면 점수 사이에 위치마다 다른 회전 행렬이 끼어들어 $W^{U Q ⊤} R_{j - t} W^{U K}$ 를 상수로 미리 곱할 수 없다. 그러면 생성할 때마다 앞선 모든 토큰의 키를 다시 계산해야 한다.^[1]

딥시크는 위치 정보를 별도 차원으로 떼어 냈다(decoupled RoPE). 헤드마다 RoPE 전용 쿼리 $q_{t, i}^{R} \in ℝ^{d_{h}^{R}}$ 를 만들고, 키 쪽은 모든 헤드가 공유하는 RoPE 전용 키 $k_{t}^{R}$ 하나를 둔다.

$q_{t}^{R} = RoPE (W^{Q R} c_{t}^{Q}), k_{t}^{R} = RoPE (W^{K R} h_{t})$

$q_{t, i} = [q_{t, i}^{C}; q_{t, i}^{R}], k_{t, i} = [k_{t, i}^{C}; k_{t}^{R}]$

$o_{t, i} = \sum_{j \leq t} {softmax}_{j} (\frac{q_{t, i}^{⊤} k_{j, i}}{\sqrt{d_{h} + d_{h}^{R}}}) v_{j, i}^{C}, u_{t} = W^{O} [o_{t, 1}; \dots; o_{t, n_{h}}]$

점수는 내용 항 $q^{C ⊤} k^{C}$ (흡수 가능)와 위치 항 $q^{R ⊤} k^{R}$ 의 합이 된다. 캐시에는 $c_{t}^{K V}$ 와 $k_{t}^{R}$ 를 함께 넣으므로 토큰·층당 $d_{c} + d_{h}^{R}$ 개 원소가 든다.^[1]

DeepSeek-V2·V3의 설정

기호	의미	V2·V3 값
$n_{h}$	어텐션 헤드 수	128
$d_{h}$	헤드 차원	128
$d_{c}$	키·값 압축 차원	512 ( $= 4 d_{h}$ )
${d_{c}}^{'}$	쿼리 압축 차원	1536
$d_{h}^{R}$	분리 RoPE 헤드 차원	64 ( $= d_{h} / 2$ )
층 수		V2 60, V3 61

V2와 V3 기술 보고서 모두 같은 MLA 차원을 썼다.^[1]^[3] 실제 모델에는 압축 잠재 벡터 뒤에 RMSNorm을 하나 더 두고 폭이 좁아지는 지점에 스케일 계수를 곱해 학습을 안정시켰다. 총 157억 매개변수의 DeepSeek-V2-Lite는 쿼리를 압축하지 않는다.^[1] 긴 문맥 확장(YaRN)은 위치 정보를 담는 공유 키 $k^{R}$ 에만 적용했다.^[1]

캐시 크기 비교

방식	토큰당 KV 캐시(원소 수)	V2 설정 대입(층당)	딥시크가 평가한 성능
MHA	$2 n_{h} d_{h} l$	32,768	강함
GQA (그룹 $n_{g}$ )	$2 n_{g} d_{h} l$	256 × $n_{g}$	보통
MQA	$2 d_{h} l$	256	약함
MLA	$(d_{c} + d_{h}^{R}) l \approx \frac{9}{2} d_{h} l$	576	더 강함

표의 성능 칸은 V2 논문의 정성 평가를 옮긴 것이다.^[1] 층 수까지 곱하면 V3(61층)는 토큰당 35,136개 원소, BF16 기준 약 69KB다. 비교하자면 GQA를 쓰는 라마 3 70B(80층, 키·값 헤드 8개, 헤드 차원 128)는 토큰당 163,840개 원소, 약 320KB다.

같은 조건 비교 실험은 다음과 같다. 두 쌍의 MoE 모델을 어텐션만 바꿔 학습했다.^[1]

모델	토큰당 KV 캐시	BBH	MMLU	C-Eval	CMMLU
소형 MoE(총 158억) + MHA	110.6K	37.9	48.7	51.6	52.3
소형 MoE(총 157억) + MLA	15.6K	39.0	50.0	50.9	53.4
대형 MoE(총 2,508억) + MHA	860.2K	46.6	57.5	57.9	60.7
대형 MoE(총 2,474억) + MLA	34.6K	50.7	59.0	59.2	62.5

MLA의 캐시는 MHA의 14%(소형), 4%(대형)였다. 흔히 인용되는 "KV 캐시 93.3% 감소, 최대 생성 처리량 5.76배"는 DeepSeek-V2를 이전 밀집 모델 DeepSeek 67B와 비교한 수치다. 이 비교에는 MLA뿐 아니라 매개변수 FP8 변환과 KV 캐시 원소당 평균 6비트 양자화가 함께 들어가 있으므로, MLA 하나의 효과로 읽으면 안 된다.^[1]

구현

아래는 정규화와 스케일 보정을 뺀 단순화 구현이다. 같은 가중치로 MHA 모드(forward_naive)와 흡수한 MQA 모드(forward_absorbed)를 계산하고, float64에서 두 출력이 10^-15 수준까지 일치하는 것을 확인했다(PyTorch 2.5).

import torch
import torch.nn as nn

def rope(x, pos, base=10000.0):              # 절반 분할 방식 RoPE, x: (..., T, d_r)
    d = x.shape[-1]
    inv = 1.0 / base ** (torch.arange(0, d, 2, dtype=torch.float32) / d)
    ang = torch.cat([pos[:, None].float() * inv] * 2, -1)
    x1, x2 = x.chunk(2, -1)
    return x * ang.cos() + torch.cat([-x2, x1], -1) * ang.sin()

class MLA(nn.Module):
    def __init__(self, d=1024, n_h=8, d_h=64, d_c=128, d_cq=256, d_r=32):
        super().__init__()
        self.n_h, self.d_h, self.d_r = n_h, d_h, d_r
        self.w_dq  = nn.Linear(d, d_cq, bias=False)            # W^DQ
        self.w_uq  = nn.Linear(d_cq, n_h * d_h, bias=False)    # W^UQ
        self.w_qr  = nn.Linear(d_cq, n_h * d_r, bias=False)    # W^QR
        self.w_dkv = nn.Linear(d, d_c, bias=False)             # W^DKV -> c^KV (캐시)
        self.w_kr  = nn.Linear(d, d_r, bias=False)             # W^KR  -> k^R  (캐시, 헤드 공유)
        self.w_uk  = nn.Linear(d_c, n_h * d_h, bias=False)     # W^UK
        self.w_uv  = nn.Linear(d_c, n_h * d_h, bias=False)     # W^UV
        self.w_o   = nn.Linear(n_h * d_h, d, bias=False)       # W^O

    def _common(self, h, pos):
        B, T, _ = h.shape
        c_kv, k_r = self.w_dkv(h), rope(self.w_kr(h), pos)    # 토큰당 d_c + d_r 개만 캐시
        c_q = self.w_dq(h)
        q_c = self.w_uq(c_q).view(B, T, self.n_h, self.d_h).transpose(1, 2)
        q_r = rope(self.w_qr(c_q).view(B, T, self.n_h, self.d_r).transpose(1, 2), pos)
        mask = torch.ones(T, T, dtype=torch.bool).triu(1)
        return B, T, c_kv, k_r, q_c, q_r, mask

    def forward_naive(self, h, pos):           # MHA 모드: 키·값을 복원
        B, T, c_kv, k_r, q_c, q_r, mask = self._common(h, pos)
        k_c = self.w_uk(c_kv).view(B, T, self.n_h, self.d_h).transpose(1, 2)
        v   = self.w_uv(c_kv).view(B, T, self.n_h, self.d_h).transpose(1, 2)
        s = q_c @ k_c.transpose(-1, -2) + q_r @ k_r[:, None].transpose(-1, -2)
        s = (s / (self.d_h + self.d_r) ** 0.5).masked_fill(mask, float("-inf"))
        o = s.softmax(-1) @ v
        return self.w_o(o.transpose(1, 2).reshape(B, T, -1))

    def forward_absorbed(self, h, pos):        # MQA 모드: W^UK, W^UV 흡수
        B, T, c_kv, k_r, q_c, q_r, mask = self._common(h, pos)
        W_uk = self.w_uk.weight.view(self.n_h, self.d_h, -1)   # (n_h, d_h, d_c)
        W_uv = self.w_uv.weight.view(self.n_h, self.d_h, -1)
        q_lat = q_c @ W_uk                                     # 쿼리를 잠재 공간으로 (B, n_h, T, d_c)
        s = q_lat @ c_kv[:, None].transpose(-1, -2) + q_r @ k_r[:, None].transpose(-1, -2)
        s = (s / (self.d_h + self.d_r) ** 0.5).masked_fill(mask, float("-inf"))
        o_lat = s.softmax(-1) @ c_kv[:, None]                  # 모든 헤드가 같은 c^KV를 값으로 사용
        o = o_lat @ W_uv.transpose(-1, -2)                     # (B, n_h, T, d_h)
        return self.w_o(o.transpose(1, 2).reshape(B, T, -1))

m = MLA().double()
h, pos = torch.randn(2, 12, 1024, dtype=torch.float64), torch.arange(12)
print((m.forward_naive(h, pos) - m.forward_absorbed(h, pos)).abs().max())   # ~1e-16

흡수 모드에서 키 쪽 차원은 $d_{c} + d_{h}^{R}$ (V2·V3에서 576), 값 쪽 차원은 $d_{c}$ (512)이고 헤드 128개가 이를 공유한다. 일반 GQA·MHA 커널은 이런 모양(키와 값 차원이 다르고 헤드 차원이 큼)을 가정하지 않아서 전용 커널이 필요하다. 딥시크는 2025년 2월 MLA 디코딩 커널 라이브러리 FlashMLA를 공개했다.^[4]

채택과 후속 연구

딥시크 모델: V2(2024년 5월), V3(2024년 12월)와 이를 바탕으로 한 R1이 MLA를 썼다.^[3] V3.2-Exp는 MLA 위에 희소 어텐션 DSA를 얹었는데, 커널 효율을 위해 하나의 키·값 항목을 여러 쿼리가 공유해야 해서 MLA의 MQA 모드를 기반으로 구현했다.^[2]
키미 K2: 문샷 AI는 1조 매개변수 MoE 모델 키미 K2(2025년 7월)에 DeepSeek-V3와 비슷한 구조로 MLA를 썼다. 학습 안정화에 흔히 쓰는 QK-Norm은 추론 때 키 행렬이 다 만들어지지 않는 MLA에는 적용할 수 없어서, 대신 가중치를 잘라 내는 QK-Clip을 고안했다고 밝혔다.^[5] 같은 회사의 Kimi Linear(2025년 10월)는 선형 어텐션 층과 MLA 층을 섞은 하이브리드다.^[6]
기존 모델 변환: 이미 학습된 MHA·GQA 모델을 MLA로 바꾸는 연구가 2025년 잇달아 나왔다. TransMLA는 GQA 모델을 MLA로 변환해 LLaMA-2-7B의 KV 캐시를 93% 압축하고 8K 문맥에서 10.6배 추론 속도 향상을 보고했다.^[7] MHA2MLA는 기여가 작은 차원에서 RoPE를 빼는 부분 RoPE와 키·값 공동 SVD 근사를 써서, 원 데이터의 0.3~0.6%만으로 성능을 회복하고 Llama2-7B의 KV 캐시를 92.19% 줄이면서 LongBench 점수 하락을 0.5%로 막았다고 보고했다.^[8]

한계

복잡도: 투영 행렬이 MHA의 4개에서 8개 안팎으로 늘고, 분리 RoPE, 잠재 벡터 정규화, 두 계산 모드, 전용 커널까지 갖춰야 한다. RoPE 길이 확장도 위치 담당 차원에만 따로 적용해야 한다.
텐서 병렬: 잠재 벡터 $c^{K V}$ 는 모든 헤드가 공유하므로, 헤드를 여러 GPU에 나누는 텐서 병렬에서는 각 GPU가 잠재 캐시 전체를 가져야 한다. GQA처럼 키·값 헤드를 장치별로 나눠 캐시를 쪼개는 이득은 없다.
연산량: 흡수 모드는 캐시 읽기를 줄이는 대신 헤드마다 512차원 잠재 공간에서 내적하므로 토큰당 연산은 늘어난다. 디코딩처럼 메모리에 묶인 상황에서 이득이 크고, 연산에 묶인 프리필에서는 MHA 모드가 낫다.
검증 범위: "MHA보다 성능이 좋다"는 결과는 딥시크가 자사 MoE 모델 두 규모에서 비교한 실험에서 나온 것이다.

같이 보기

각주

↑ ^1.00 ^1.01 ^1.02 ^1.03 ^1.04 ^1.05 ^1.06 ^1.07 ^1.08 ^1.09 ^1.10 ^1.11 ^1.12 ^1.13 DeepSeek-AI, 〈DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model〉, 2024. arXiv:2405.04434
↑ ^2.0 ^2.1 DeepSeek-AI, 〈DeepSeek-V3.2-Exp: Boosting Long-Context Efficiency with DeepSeek Sparse Attention〉, 2025. GitHub PDF
↑ ^3.0 ^3.1 DeepSeek-AI, 〈DeepSeek-V3 Technical Report〉, 2024. arXiv:2412.19437
↑ deepseek-ai/FlashMLA, GitHub. 저장소 생성일 2025-02-21.
↑ Kimi Team, 〈Kimi K2: Open Agentic Intelligence〉, 2025. arXiv:2507.20534
↑ Kimi Team, 〈Kimi Linear: An Expressive, Efficient Attention Architecture〉, 2025. arXiv:2510.26692
↑ Fanxu Meng 외, 〈TransMLA: Multi-Head Latent Attention Is All You Need〉, 2025. arXiv:2502.07864
↑ Tao Ji 외, 〈Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs〉, 2025. arXiv:2502.14837

[dsv2-1] 1.00 ^1.01 ^1.02 ^1.03 ^1.04 ^1.05 ^1.06 ^1.07 ^1.08 ^1.09 ^1.10 ^1.11 ^1.12 ^1.13 DeepSeek-AI, 〈DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model〉, 2024. arXiv:2405.04434

[dsv32-2] 2.0 ^2.1 DeepSeek-AI, 〈DeepSeek-V3.2-Exp: Boosting Long-Context Efficiency with DeepSeek Sparse Attention〉, 2025. GitHub PDF

[dsv3-3] 3.0 ^3.1 DeepSeek-AI, 〈DeepSeek-V3 Technical Report〉, 2024. arXiv:2412.19437

[flashmla-4] seek-ai/FlashMLA, GitHub. 저장소 생성일 2025-02-21.

[kimik2-5] Kimi Team, 〈Kimi K2: Open Agentic Intelligence〉, 2025. arXiv:2507.20534

[kimilinear-6] Kimi Team, 〈Kimi Linear: An Expressive, Efficient Attention Architecture〉, 2025. arXiv:2510.26692

[transmla-7] Fanxu Meng 외, 〈TransMLA: Multi-Head Latent Attention Is All You Need〉, 2025. arXiv:2502.07864

[mha2mla-8] Tao Ji 외, 〈Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs〉, 2025. arXiv:2502.14837

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]