Archive

1. Introduction: 왜 KV Cache 공간 최적화가 필요한가?

자기회귀(Autoregressive) 디코딩 과정에서 이전 토큰들의 Key, Value 벡터를 저장하는 KV Cache는 연산의 중복을 막아주지만, 시퀀스 길이 $N$ 에 비례하여 메모리 요구량이 선형적으로 증가( $O(N)$ )하는 치명적인 단점이 있습니다.

특히 SRAM(수 MB)과 DRAM(수십 GB) 대역폭이 극도로 제한된 Edge 하드웨어(Jetson, MPU) 에서는 이 KV Cache가 SRAM 용량을 초과하여 느린 DRAM으로 밀려나는 순간(Spill-over), 추론 속도가 급감하는 Memory Wall에 직면하게 됩니다. 본 문서는 모든 토큰을 저장하지 않고도 모델의 성능을 유지하는 Token-level Sparsity(토큰 수준 희소성) 전략의 수학적 원리와 하드웨어적 이점을 분석합니다.

2. 하드웨어 관점의 문제 정의: KV Cache Memory Footprint

구체적으로 한 토큰을 생성할 때마다 누적되는 KV Cache의 물리적 크기는 다음과 같이 계산됩니다.

$\text{KV\_Cache\_Size (Bytes)} = 2 \times N \times L \times H \times P$

$N$ : 현재까지 생성된 시퀀스 길이 (Context Length)
$L$ : 트랜스포머 레이어 수 (Number of Layers)
$H$ : 히든 차원 크기 (Hidden Dimension)
$P$ : 파라미터 정밀도 (FP16/BF16의 경우 2 Bytes)
$2$ : Key와 Value 두 개의 텐서

Example: LLaMA-2 7B 모델(FP16) 기준, $4096$ 토큰을 처리할 때 KV Cache는 약 2GB를 차지합니다. Edge 디바이스의 제한된 캐시(SRAM)에는 절대 담을 수 없는 크기이며, 필연적으로 엄청난 DRAM I/O 병목을 유발합니다.

3. The Phenomenon: Attention Score의 멱법칙(Power-law) 분포

모든 토큰이 다음 토큰 생성에 동일한 기여를 할까요? 최신 연구(H2O, StreamingLLM 등)에 따르면, Attention 점수 행렬은 극도의 Sparsity(희소성) 를 보입니다.

표준 Attention의 Softmax 수식을 살펴보면:

$\text{Attention}(x_q, x_{k, \le q}) = \sum_{j \le q} \frac{\exp(x_q \cdot x_{k_j} / \sqrt{d})}{\sum_{i \le q} \exp(x_q \cdot x_{k_i} / \sqrt{d})} x_{v_j}$

실제 디코딩 과정에서 위 Softmax 확률값의 90% 이상은 전체 토큰 중 단 5% 미만의 특정 토큰들(Heavy-Hitters) 에 집중됩니다. 즉, 나머지 95%의 토큰은 확률값이 $0$ 에 수렴하여 결과 벡터 $V$ 에 거의 영향을 주지 못합니다. 이 잉여 토큰들의 KV 벡터를 메모리에서 과감히 삭제(Eviction)하는 것이 Token-level Sparsity의 핵심입니다.

4. 핵심 유지 전략: 무엇을 남기고 무엇을 버릴 것인가?

메모리 사용량을 $O(N)$ 에서 $O(1)$ 로(고정된 크기로) 제한하기 위해, 알고리즘은 다음 두 가지 토큰 그룹만 SRAM(KV Cache)에 유지합니다.

4.1. Attention Sink (초기 토큰)

문맥의 의미와 무관하게, 항상 첫 1~4개의 초기 토큰(예: <s> 보스 토큰)은 막대한 Attention 점수를 가져갑니다. 이를 Attention Sink 현상이라고 합니다.

수학적 이유: Softmax 함수는 확률의 합이 1이 되어야 합니다. 특정 토큰에 주목할 필요가 없는 레이어에서도, 모델은 불필요한 정보 갱신을 막기 위해 잉여 Attention 점수를 '안전한 피난처(Sink)'인 첫 번째 토큰에 쏟아붓도록 학습됩니다.
전략 (StreamingLLM): 이 초기 4개의 토큰을 KV Cache에서 지우면 Softmax의 분모(Denominator)가 붕괴되어 모델이 완전히 망가집니다.[하단 이미지의 (b)와 (d)참고] 따라서 초기 토큰은 영구적으로 보존합니다.

attention_sink

4.2. Local Window (최신 토큰)와 Heavy-Hitters (핵심 토큰)

Local Window: 언어의 본질상 바로 직전에 생성된 토큰들( $L$ 개)이 다음 단어 예측에 가장 큰 영향을 미칩니다(Recency Bias).
Heavy-Hitters (H2O 알고리즘): 누적 Attention 점수를 평가하여 가장 점수가 높은 상위 $K$ 개의 토큰을 추적합니다.

$\text{Score}_j = \sum_{i=j}^{N} \text{Attention}(i, j)$

새로운 토큰이 생성되어 할당된 캐시 용량( $K + L$ )을 초과하면, 이 $\text{Score}$ 가 가장 낮은 토큰의 KV 텐서를 메모리에서 방출(Eviction)합니다.

H2O

5. Edge 하드웨어 환경에서의 파괴적 시너지 (System Impact)

위의 알고리즘을 하드웨어(SRAM/DRAM) 레벨로 내려다보면 엄청난 이점이 발생합니다.

SRAM 상주(Residency) 보장: KV Cache의 최대 크기를 Edge NPU/GPU의 SRAM 용량(예: 2~4MB) 이내로 엄격하게 고정( $K+L \ll N$ )할 수 있습니다.
DRAM I/O 우회: 매 토큰 생성 시 거대한 KV Cache 전체를 DRAM에서 SRAM으로 끌고 오는 행위(Memory Fetch)가 완전히 사라집니다. 연산기는 오직 가중치(Weights)만 DRAM에서 읽어오면 되므로 Memory Bandwidth 병목이 절반 이하로 줄어듭니다.

6. 결론 및 Trade-off

Token-level Sparsity는 "용량이 작은 데이터(KV Cache)는 초고속 SRAM에 가두고, 용량이 큰 가중치만 DRAM에서 가져온다"는 하드웨어 메모리 계층 최적화의 정석을 알고리즘적으로 구현한 훌륭한 사례입니다.

한계점 (Limitations):

버려진 토큰에 포함된 세밀한 팩트(Fact)를 모델이 나중에 회상(Recall)해야 할 경우 환각(Hallucination)이 발생할 수 있습니다.
매 스텝마다 Attention Score를 계산하여 토큰을 정렬하고 Eviction하는 동적 메모리 관리(Pointer Management) 오버헤드가 발생합니다.