문제: KV 캐시 윈도우를 밀면 왜 망가지는가
긴 대화나 무한 스트리밍 생성에서는 KV 캐시가 계속 커진다. 메모리를 아끼려고 오래된 토큰을 버리는 슬라이딩 윈도우를 쓰면, 일정 시점부터 출력이 갑자기 반복되거나 의미 없는 토큰을 뱉으며 붕괴한다. 원인은 단순한 정보 손실이 아니다. 소프트맥스 어텐션은 확률의 합이 1이어야 하므로, 문맥상 어디에도 집중할 필요가 없을 때 남는 확률 질량을 어딘가에 "버려야" 한다. 학습 과정에서 모델은 이 잉여 질량을 초반 몇 개 토큰(주로 첫 토큰)에 몰아두는 습관을 들인다.
어텐션 싱크의 정체
이 초반 토큰들을 어텐션 싱크(attention sink)라고 부른다. 내용상 의미는 거의 없지만, 어텐션 분포의 배출구 역할을 한다. 슬라이딩 윈도우가 이 싱크 토큰을 캐시에서 밀어내는 순간, 쏟아낼 곳을 잃은 확률 질량이 윈도우 내부 토큰으로 재분배되면서 분포가 왜곡되고 생성이 무너진다. 핵심 통찰은 "최근 토큰만 유지"가 아니라 "초반 싱크 몇 개 + 최근 윈도우"를 함께 유지하면 캐시 크기를 고정한 채로 안정성을 되찾는다는 것이다.
해결: 싱크 + 슬라이딩 윈도우
구현은 KV 캐시를 두 영역으로 나눈다. 맨 앞 N개(보통 4개) 싱크 토큰은 절대 버리지 않고, 그 뒤로는 최근 W개만 유지한다. 한 가지 주의가 필요한데, 회전 위치 임베딩(RoPE)을 쓸 때는 원래 절대 위치가 아니라 캐시 내부 상대 위치로 다시 매겨야 한다. 그렇지 않으면 싱크와 최근 토큰 사이의 위치 간격이 벌어져 오히려 성능이 떨어진다.
def evict_kv_cache(keys, values, sink=4, window=1020):
# keys/values: [n_layers, n_heads, seq_len, head_dim]
total = keys.shape[2]
if total <= sink + window:
return keys, values
# 앞쪽 sink개는 항상 보존, 나머지는 최근 window개만 유지
keep = list(range(sink)) + list(range(total - window, total))
idx = torch.tensor(keep, device=keys.device)
return keys.index_select(2, idx), values.index_select(2, idx)
def build_position_ids(sink, window, device):
# 캐시 내부의 연속된 상대 위치로 재부여 (RoPE 사용 시 필수)
length = sink + window
return torch.arange(length, device=device).unsqueeze(0)
서빙 프레임워크 설정
대부분의 추론 서버는 이 기능을 옵션으로 제공한다. 직접 커널을 짜기 전에 프레임워크 플래그부터 확인하는 편이 빠르다. 아래는 싱크 토큰 수와 윈도우 길이를 노출하는 설정 예시다.
engine:
model: my-org/llm-7b
max_model_len: 8192
# 어텐션 싱크 관련 설정
attention_sink:
enabled: true
sink_tokens: 4 # 보존할 초반 토큰 수
recent_window: 4092 # 유지할 최근 토큰 수
# 캐시 총량 = sink_tokens + recent_window 로 고정
kv_cache:
block_size: 16
reuse_position_ids: true
방식 비교
| 전략 | KV 메모리 | 장문 안정성 | 비고 |
|---|---|---|---|
| 전체 캐시 | 선형 증가 | 높음 | 길이 제한에서 OOM |
| 슬라이딩 윈도우만 | 고정 | 낮음(붕괴) | 싱크 유실 |
| 싱크 + 윈도우 | 고정 | 높음 | 권장 |
주의점
몇 가지를 놓치기 쉽다. 첫째, 이 기법은 캐시 크기를 늘리는 것이 아니라 고정하는 것이다. 윈도우 밖으로 밀려난 실제 문맥은 사라지므로, 초반에 등장한 사실을 뒤에서 정확히 인용해야 하는 작업에는 부적합하다. 유효 문맥 길이를 늘리는 방법이 아니라, 무한 생성 시 붕괴를 막는 방법이다. 둘째, 싱크 개수는 4개 안팎이면 충분하고 무작정 늘려도 이득이 없다. 셋째, RoPE 위치 재부여를 빼먹으면 안정성 이득이 거의 사라지므로 반드시 검증한다. 넷째, 프롬프트 캐시나 배치 프리필과 함께 쓸 때 싱크 영역이 공유 접두부와 겹치는지 확인해야 인덱싱 버그를 피할 수 있다.
검증 방법
도입 후에는 실제 붕괴 지점을 넘겨 생성해보며 확인한다. 윈도우 길이의 2~3배까지 토큰을 뽑았을 때 퍼플렉시티가 급등하지 않고 평탄하게 유지되는지, 반복 루프에 빠지지 않는지를 본다. 슬라이딩 윈도우만 켠 대조군과 나란히 돌려 붕괴 시점 차이를 로그로 남겨두면, 설정 회귀를 조기에 잡을 수 있다.