Long Context LLM — 두 개의 완전히 다른 문제

“Long context inference가 느리다”는 말은 하나의 문제처럼 들린다. 실제로는 완전히 성격이 다른 두 병목이 공존한다. 이 구분을 이해하지 못하면, StreamingLLM·YaRN·Ring Attention이 각각 왜 다른 기법인지도, 왜 셋 다 필요한지도 설명할 수 없다. 어떤 병목이 어디서 발생하는가?

Prefill vs Decode — 정반대의 계산 특성

Prefill(프롬프트 처리)과 Decode(토큰 생성)는 계산 프로파일이 정반대다.

Prefill에서는 길이 $L$의 프롬프트 안에서 모든 Query가 모든 Key/Value와 attend한다. Attention 행렬의 크기는 $L \times L$이므로 계산량은 다음과 같다.

$$\text{Prefill FLOPs} \approx L^2 \times d_{\text{total}}$$

LLaMA-70B($d_{\text{total}}=8192$), $L=128\text{K}$를 대입하면 약 $1.3 \times 10^{15}$ FLOPs. A100 기준 이론적 하한이 4초를 넘는다. Prefill의 병목은 compute다.

Decode에서는 매 스텝마다 새로운 Query 1개가 누적된 $L$개의 Key/Value와 attend한다. 계산량 자체는 $O(L)$로 작지만, 모든 이전 토큰의 K/V를 GPU 메모리에 유지해야 한다.

$$M_{\text{KV}} = B \times 2 \times L \times d_{\text{total}} \times \text{dtype\_bytes}$$

$L=128\text{K}$, 배치 8, BF16이면 약 26 GB. A100 80GB에서 배치 2–3이 한계에 도달한다. Decode의 병목은 memory다.

StreamingLLM — Decode 메모리를 상수로

Decode 메모리를 줄이는 가장 단순한 접근은 슬라이딩 윈도우다 — 최근 $W$개 토큰만 KV cache에 유지하면 메모리가 $O(W)$로 고정된다. 그런데 Xiao et al.(2024)의 관찰에 따르면, 이것만으로는 perplexity가 10배 이상 급등한다.

이유는 Softmax normalization에 있다. 초기 토큰들은 실제 내용과 무관하게 모든 Query에 대해 비정상적으로 높은 attention weight를 받는 “attention sink” 역할을 한다. 이 토큰들이 윈도우 밖으로 밀려나면 Softmax의 분모(partition function)가 불안정해진다.

$$K_{\text{eff}} = K_{[1:M]} \cup K_{[t-W:t]}, \quad |K_{\text{eff}}| = M + W = \text{const}$$

처음 $M$개(4–8개)를 sink로 고정하고 최근 $W$개를 유지하면 KV cache가 상수가 된다. $t = 100\text{K}$에서도 메모리는 $(M+W) \times d$뿐이다.

YaRN — 학습 범위 밖의 Position을 보간

LLaMA처럼 4K context로 학습된 모델을 128K로 확장하려면 RoPE의 position 범위를 늘려야 한다. 단순 외삽(extrapolation)은 $m > 4\text{K}$에서 high-frequency 성분이 aliasing을 일으켜 perplexity가 약 32배 폭등한다.

Chen et al.(2023)의 Position Interpolation은 모든 position을 균일하게 스케일한다.

$$m' = m \times \frac{L_{\text{train}}}{L_{\text{infer}}}$$

이 방식은 low-frequency(장거리) 성분에는 적합하지만, high-frequency(세밀한 위치 구분)를 과도하게 압축해 정보를 손실한다.

YaRN(Peng 2024)은 dimension별로 다른 보간 계수를 적용한다.

$$\theta_i^{\text{YaRN}} = \lambda(\theta_i) \cdot \theta_i, \quad \lambda(\theta) = \begin{cases} 1 & \text{high-freq} \\ \text{ramp}(\cdot) & \text{transition} \\ s & \text{low-freq} \end{cases}$$

High-frequency dimension은 $\lambda=1$로 학습 당시 그대로 유지하고, low-frequency만 스케일 $s = L_{\text{train}}/L_{\text{infer}}$로 내린다. 결과적으로 4K → 128K 확장 시 perplexity 손실이 2% 미만이며, fine-tuning 없이 즉시 적용 가능하다.

Ring Attention — Prefill을 분산으로

YaRN으로 128K까지 확장해도, $L=1\text{M}$에서의 Prefill은 단일 GPU로 불가능하다. Attention 행렬 $QK^T \in \mathbb{R}^{L \times L}$이 BF16 기준 2 TB를 요구하기 때문이다.

Ring Attention(Liu 2024)은 $N$개 GPU에 sequence를 분할한다.

$N=16$이면 2 TB → 8 GB 수준으로, A100에서 현실적이다. 통신량은 라운드당 $(L/N) \times d$ bytes로 computation과 pipelining하면 대부분의 overhead를 숨길 수 있다.

Ring Attention은 단일 GPU inference에서는 의미가 없다. 대규모 분산 학습 클러스터에서 1M context 모델을 훈련·추론할 때 사용하는 기법이다.

정리

• Long context 병목은 Prefill(compute) 과 Decode(memory) 두 개로, 각각 다른 해결책이 필요하다.
• StreamingLLM은 attention sink + 슬라이딩 윈도우로 Decode KV 메모리를 $O(1)$ 로 만든다. 단, 장거리 정보는 잃는다.
• YaRN은 dimension별 RoPE 보간으로 fine-tuning 없이 context를 128K까지 확장한다.
• Ring Attention은 sequence를 GPU간 분산해 Prefill의 $O(L^2/N^2)$ 메모리를 가능하게 하며, 1M context 모델의 토대가 된다.

세 기법은 경쟁이 아니라 서로 다른 축을 다룬다. 실제 1M context 시스템에서는 셋 모두 동시에 쓰인다.