Baseten 엔지니어가 분석한 7년간의 어텐션 메커니즘 진화: GPT-2부터 Kimi K3까지, 실행 가능한 PyTorch 코드로 구현
요약
Baseten 엔지니어가 GPT-2부터 Kimi K3까지 지난 7년간의 어텐션 메커니즘 진화 과정을 PyTorch 코드로 분석했습니다. KV 캐시의 등장 배경부터 메모리 대역폭 병목 현상, 그리고 이를 해결하기 위한 선형 어텐션의 원리까지 단계별로 설명합니다.
핵심 포인트
- GPT-2에서 Kimi K3까지의 파라미터 규모는 약 22,580배 성장함
- KV 캐시는 재계산을 방지하지만 시퀀스 길이에 따라 메모리 사용량이 선형적으로 증가함
- 최신 모델의 핵심 병목은 연산량이 아닌 메모리 대역폭 문제임
- 선형 어텐션은 소프트맥스를 피처 맵으로 교체하여 메모리 효율성을 높임
지난주, @waterloo_intern이라는 이름으로 활동하는 Baseten의 추론(inference) 엔지니어가 **"22,580: GPT-2부터 Kimi K3까지, 설명하기"**라는 제목의 기술 블로그 포스트를 게시했습니다. 이 글은 며칠 만에 240만 회의 조회수를 기록했습니다.
그는 보도 자료를 작성한 것이 아닙니다. 그는 실행 가능한 PyTorch 코드를 작성했습니다. GPT-2의 어텐션 블록(attention block)부터 시작하여 모든 아키텍처(architectural) 변화를 단계별로 짚어가며, 각 반복(iteration)마다 발생하는 하나의 문제와 하나의 비용을 설명했습니다. 이는 제가 본 트랜스포머(transformer) 계보 설명 중 최고였습니다.
저는 그의 포스트를 탐독한 후, 5편의 원문 논문과 주요 주장들을 교차 검증했습니다. 전체적인 그림은 다음과 같습니다.
22,580배라는 숫자
2019년 2월, OpenAI는 124M 파라미터(parameters)를 가진 GPT-2를 출시했습니다. 7년 후, Moonshot AI는 2.8T 파라미터를 가진 Kimi K3를 오픈 소스로 공개했습니다. Kimi K3 하나 안에 22,580개의 GPT-2를 집어넣을 수 있습니다.
하지만 이것은 단순히 "연산량(compute)을 더 많이 투입한다"는 이야기가 아닙니다. 이것은 우리가 메모리(memory)를 어떻게 저장하고, 업데이트하며, 검색하는가에 대한 이야기입니다.
시작점: GPT-2
class Block(nn.Module):
def forward(self, x):
x = x + self.attn(self.ln_1(x))
...
모델이 새로운 토큰(token)을 생성할 때마다, 모델은 모든 과거 토큰에 대해 Q, K, V 프로젝션(projections)을 다시 계산한 다음, O(N²) 소프트맥스 어텐션(softmax attention)을 실행합니다. 토큰 1부터 N-1까지의 K와 V는 어떻게 될까요? 버려집니다. 토큰 N+1이 도착하면? 모든 것을 다시 계산합니다.
이것이 바로 KV 캐시(KV Cache)가 발명된 이유입니다.
KV 캐시: 재계산하지 말고 저장하라
단순한 아이디어입니다: 이미 계산된 키(keys)와 값(values)을 캐시(cache)하는 것입니다. 다음 토큰을 위해, 새로운 Q는 캐시된 K와 단 한 번의 내적(dot product)만 수행하면 됩니다.
문제는 해결되었지만, 새로운 문제가 발생했습니다. KV 캐시는 시퀀스 길이(sequence length)에 따라 **선형적(linearly)**으로 증가합니다. 1M 토큰 × d_model × 레이어(layers)를 기준으로 하면, 이는 수십 GB의 VRAM을 차지합니다. 모든 디코딩(decoding) 단계마다 HBM에서 이 모든 데이터를 읽어와야 합니다.
병목 현상(bottleneck)은 연산량이 아닙니다. 메모리 대역폭(memory bandwidth)입니다. 이것이 이후에 이어지는 모든 개선 사항을 이해하는 핵심입니다.
선형 어텐션 (Linear Attention): 고정된 크기의 메모리
O(N²D)를 O(ND²)로 압축할 수 있을까요?
아이디어는 다음과 같습니다: 소프트맥스(softmax)를 피처 맵(feature map)으로 교체하는 것입니다.
# 표준 소프트맥스 (먼저 N×N 행렬을 실체화해야 함)
attention = softmax(QKᵀ / √d) × V
...
이제 우리는 고정된 크기의 D×D 행렬인 K'^T × V'를 먼저 계산한 다음, 이를 Q'와 곱할 수 있습니다. 과거의 KV 정보가 상수 크기의 상태(state) 행렬로 "접힙니다(folded)". 캐시(Cache)는 더 이상 N에 따라 증가하지 않습니다.
비용은 무엇일까요? ELU+1은 소프트맥스 커널(softmax kernel)의 근사치입니다. 표현력(Expressiveness)이 떨어집니다. 하지만 긴 문맥(long-context) 작업에서는 이러한 트레이드오프(tradeoff)가 종종 가치가 있습니다.
참고: FlashAttention은 이렇게 작동하지 않았습니다
흔히 혼동하는 지점입니다. Ali의 포스트에서는 "2020년에는 FlashAttention이 없었다"라고 언급하지만, FlashAttention과 선형 어텐션(linear attention)은 근본적으로 다른 문제를 해결합니다.
FlashAttention (Tri Dao, NeurIPS 2022)은 어텐션 **알고리즘(algorithm)**을 변경하지 않았습니다. 대신 GPU의 **IO 패턴(IO pattern)**을 최적화했습니다. 즉, N×N 행렬을 타일링(tiling)하여 HBM에 완전히 올라가지 않도록 한 것입니다. 이는 소프트맥스를 더 빠르게 만들지만, 여전히 O(N²)의 복잡도를 가집니다.
선형 어텐션은 알고리즘 자체를 재정의했습니다. 소프트맥스를 피처 맵(feature map)으로 교체하여, 복잡도를 O(N²D)에서 O(ND²)로 낮추었습니다.
하나는 IO를 최적화하고, 다른 하나는 알고리즘을 바꿉니다. 서로 직교(Orthogonal)하는 개념입니다.
DeltaNet: 이제 메모리를 _편집_할 수 있습니다
선형 어텐션에는 치명적인 결함이 있습니다. 오직 **추가(add)**만 할 수 있을 뿐, 결코 **업데이트(update)**할 수 없다는 점입니다.
모든 새로운 토큰은 상태에 쌓이기만 합니다: S = S + K'^T × V'. 정보는 계속 늘어나기만 합니다. 이는 마치 글을 쓸 수는 있지만, 지우거나 수정할 수는 없는 공책과 같습니다.
DeltaNet (Songlin Yang et al., NeurIPS 2024)은 이 문제를 해결합니다.
핵심 아이디어는 1990년대 Schmidhuber의 "Fast Weight Programmers"로 거슬러 올라가며, 이후 Schlag et al. (ICML 2021)이 선형 어텐션을 fast weights와 연결하며 공식화되었습니다. DeltaNet의 접근 방식은 다음과 같습니다: 쓰기 전에, 현재 이 키(key) 위치에 무엇이 저장되어 있는지 읽습니다.
v_old = k @ S_old # 키 위치의 현재 값을 읽음
delta = v_new - v_old # 차이(delta)를 계산
S_new = S_old + k^T @ delta # 차이값만 기록
만약 v_new == v_old라면, delta는 0이 되어 아무것도 변하지 않습니다. 만약 완전히 다르다면, delta는 v_new와 같아지며 이는 덮어쓰기(overwrite)와 동일합니다. 그 사이의 모든 값은 부드러운 보간(interpolation)을 수행합니다.
이것이 바로 **델타 규칙 (delta rule)**입니다. 무분별한 누적 대신 정밀한 메모리 업데이트를 수행합니다.
병렬화 트릭 (Ali는 이를 이해하는 데 7시간이 걸렸다고 말했습니다)
DeltaNet의 상태 업데이트(state update)는 엄격하게 순차적입니다. 즉, 각 단계가 이전의 $S$에 의존합니다. 이로 인해 병렬화가 불가능해 보이지만, 실제로는 가능합니다.
방법은 다음과 같습니다: 시퀀스를 크기 $C$인 청크(chunk)로 나눕니다. 각 청크 내부에서는 일반적인 마스크 어텐션 (masked attention, GPU 병렬 처리)을 사용합니다. 청크 사이에는 상태 행렬(state matrix)과 하나의 행렬 곱셈 (matmul, $Q @ S$)을 사용합니다. Householder 변환(Householder transformation)은 델타 업데이트를 재매개변수화(reparameterize)하여, 청크 내의 모든 델타를 한 번에 계산할 수 있게 합니다.
복잡도(Complexity): 고정 비용 $2LD^2$ (상태 유지) + 가변 비용 $2LCD$ (청크 내 어텐션). $C$가 커질수록 가변 비용은 증가하지만 GPU 효율성은 좋아집니다. 실제로는 $C=64$ 또는 $128$이 가장 잘 작동합니다. FLOPs가 유일한 지표는 아니며, 텐서 코어(tensor core) 활용도 또한 중요합니다.
Gated DeltaNet: 이제는 '망각'할 수 있습니다
DeltaNet은 개별 키-값(key-value) 쌍을 정밀하게 편집할 수 있습니다. 하지만 **대규모 망각 (forgetting at scale)**은 어떨까요?
1M 토큰 컨텍스트에서 주제 A에 관한 모든 문서를 읽은 다음, 주제 B로 전환한다고 상상해 보십시오. 모델은 이상적으로 B를 위한 용량을 확보하기 위해 주제 A를 "망각"해야 합니다.
DeltaNet은 특정 항목을 덮어쓸 수는 있지만, 전역 메모리를 일괄적으로 감쇠(bulk-decay)시킬 수는 없습니다. Mamba-2 (Dao & Gu, ICML 2024)는 다음과 같이 할 수 있습니다:
cache = α × S_old + S_new
여기서 $\alpha$는 0과 1 사이의 게이팅 값(gating value)으로, 모든 오래된 메모리를 균일하게 감쇠시킵니다. 이것이 Mamba-2의 "상태 공간 이중성 (State Space Duality)" 이론의 핵심입니다. 즉, 소프트맥스 어텐션 (softmax attention)과 SSM은 수학적으로 동일한 것을 다르게 표현한 것입니다. Mamba-2는 게이팅을 통해 이 둘을 통합했습니다.
Gated DeltaNet (Songlin Yang et al., ICLR 2025, NVIDIA)은 DeltaNet의 델타 업데이트와 Mamba-2의 게이팅을 결합합니다:
S_new = α × S_old + k^T @ delta # 먼저 감쇠시킨 후, 정밀하게 기록
$\alpha=1$이면 순수한 DeltaNet입니다. $\alpha=0$이면 메모리 삭제입니다. 그 사이의 모든 값은 망각과 업데이트를 동시에 수행합니다.
핵심적인 메커니즘은 다음과 같습니다: 타임스텝 $x$에서 기록되고 $x+t$에서 읽히는 토큰은 $t$번의 누적 $\alpha$ 감쇠($\alpha_x \times \alpha_{x+1} \times \dots \times \alpha_{x+t}$)를 거치게 됩니다. 서로 다른 시점에 기록된 정보는 서로 다른 속도로 망각됩니다 — 최근 항목은 거의 감쇠되지 않지만, 오래된 항목은 완전히 사라질 수 있습니다.
KDA: Kimi의 비법
Kimi Linear (Moonshot AI, arXiv 2510.26692, 2025년 10월) — K3의 전신 — 는 이 아이디어를 더욱 정교화했습니다: 스칼라 게이팅 (scalar gating)에서 차원별 게이팅 (per-dimension gating)으로 진화한 것입니다.
Gated DeltaNet은 $\alpha$를 단일 스칼라로 사용합니다. 즉, 하나의 숫자가 모든 메모리 차원의 망각 속도를 제어합니다. 반면 KDA는 $\alpha$를 벡터 (vector) (또는 행렬)로 변환합니다. 각 차원이 자신의 망각률을 독립적으로 제어합니다. 어떤 개념을 유지하고 어떤 차원을 감쇠시킬지는 모델이 스스로 학습합니다.
논문의 주요 데이터:
| 지표 | Full MLA | Kimi Linear |
|---|---|---|
| KV Cache | 100% | -75% |
| ... |
이는 모든 시나리오에 걸친 공정한 비교 하에서 리니어 어텐션 (linear attention)이 풀 어텐션 (full attention)의 성능을 능가한 첫 번째 사례입니다. 긴 컨텍스트 (long context)뿐만 아니라 짧은 컨텍스트 (short context)에서도 마찬가지입니다. 그리고 이 "-75% KV 캐시"는 추론 비용 절감으로 직결됩니다.
저는 논문의 초록(abstract)을 통해 이 주장을 직접 확인했습니다: "for the first time, outperforms full attention under fair comparisons across various scenarios." (사상 처음으로, 다양한 시나리오에서의 공정한 비교 하에 풀 어텐션의 성능을 능가함) 이는 마케팅 용어가 아닌 핵심 결론입니다.
K3의 최종 하이브리드 아키텍처
K3의 기술 보고서 (arXiv 2607.24653, 2026년 7월)는 다음을 확인해 줍니다:
- 3/4의 레이어는 KDA (linear)를 사용하고, 1/4은 gated MLA (full softmax)를 사용합니다.
- MoE 라우팅 (routing): 896개의 전문가 (experts), 토큰당 16개의 활성 전문가 (~1.8%), 부하 분산을 위한 분위수 균형 (Quantile Balancing) 적용
- 어텐션 잔차 (Attention Residuals): 레이어가 특정 이전 레이어의 표현 (representations)을 "다시 돌아볼" 수 있음
- MXFP4 가중치 (weights) + MXFP8 활성화 (activations), 양자화 인식 학습 (quantization-aware training)
K3는 순수한 선형 어텐션 (linear-attention) 모델이 아닙니다. 이는 **하이브리드 시스템 (hybrid system)**입니다. KDA가 대량의 처리 (저렴하고 빠르며 고정된 상태)를 담당하는 동안, 주기적인 소프트맥스 (softmax) 레이어가 정밀한 검색 (precision retrieval)을 수행하여 선형 압축 과정에서 손실될 수 있는 세부 사항을 복구합니다. 이는 Ali가 분석한 트레이드오프 (tradeoff)와 정확히 일치합니다: 선형 어텐션 (linear attention) + 주기적인 소프트맥스 (softmax) 검색.
그렇다면 22,580배는 실제로 무엇을 의미하는가?
Ali의 결론은 제가 작성할 수 있는 그 어떤 것보다 날카롭습니다:
GPT-2부터 Kimi K3에 이르기까지, 각 세대의 핵심적인 개선 사항은 "더 많은 파라미터 (parameters)"가 아니었습니다. 그것은 메모리에 접근하는 방식, 즉 어떻게 저장하고, 어떻게 잊고, 어떻게 검색하는지를 재설계하는 것이었습니다.
표로 정리한 진화 과정:
| 단계 | 대표 모델 | 메모리 메커니즘 (Memory Mechanism) | 해결된 문제 |
|---|---|---|---|
| 1 | GPT-2 + KV Cache | 전체 캐시 (Full cache), O(N) 증가 | 재계산 (recomputation) 제거 |
| ... |
파라미터는 22,580배 성장했습니다. 하지만 만약 당신이 그것만을 본다면, 전체 이야기의 핵심을 놓친 것입니다.
이것이 당신에게 의미하는 바
만약 당신이 긴 컨텍스트 (long-context) 작업 (코드 리뷰, 문서 분석, 멀티턴 에이전트)을 수행하고 있다면, 어텐션 아키텍처 (attention architecture)는 비용과 출력 품질에 직접적인 영향을 미칩니다.
-
비용은 고정되어 있지 않습니다. 동일한 100만 토큰 (1M-token) 컨텍스트의 경우: KDA의 KV 캐시 (KV cache)는 전체 어텐션 (full attention)의 25%에 불과합니다. 메모리 대역폭 (memory bandwidth) 압박이 낮아져 지연 시간 (latency)과 처리량 (throughput)이 현저히 개선됩니다.
-
길다고 해서 반드시 비싼 것은 아닙니다. K3의 100만 토큰 컨텍스트 윈도우 (context window)는 "무차별 대입 (brute-forced)" 방식이 아닙니다. 작업의 75%가 선형 경로 (linear path)를 통해 처리됩니다.
-
하이브리드 (Hybrid)가 추세입니다. 선형 (Linear) 방식이 소프트맥스 (softmax)를 대체하지는 않을 것입니다. 대신 서로를 보완할 것입니다. 소프트맥스를 통한 정밀한 검색 (precision retrieval), 선형을 통한 대량 처리 (bulk processing). 이러한 패러다임이 확산될 것입니다.
핵심 요약 (Key Takeaways)
- 22,580배의 파라미터(parameter) 성장은 표면적인 현상입니다. 진짜 핵심은 메모리 관리(memory management)입니다: 전체 캐시(full cache) → 고정된 상태(fixed state) → 정밀한 쓰기(precise writes) → 적응형 망각(adaptive forgetting)으로의 진화입니다.
- 계보가 명확합니다: DeltaNet (NeurIPS 2024) → GatedDeltaNet (ICLR 2025) → Kimi Linear (2025년 10월) → Kimi K3 (2026년 7월).
- KDA는 단기, 장기, 그리고 강화학습 (RL) 스케일링 시나리오 전반에 걸쳐 공정한 비교에서 풀 어텐션 (full attention)을 처음으로 능가했습니다. 이는 단순한 "저렴한 대안"이 아닙니다.
- K3의 75% 선형 (linear) / 25% 소프트맥스 (softmax) 하이브리드 구조는 논문만을 위한 구성이 아닌, 공학적 최적점 (engineering optimum)입니다.
- 다음에 모델을 평가할 때: 파라미터 수(parameter count)가 아니라 어텐션 아키텍처 (attention architecture)에 대해 질문하십시오. 그것이 훨씬 더 중요합니다.
참고 문헌 (References)
AI 자동 생성 콘텐츠
본 콘텐츠는 Dev.to AI tag의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기