KV-Fold: 긴 컨텍스트 추론을 위한 One-Step KV-Cache Recurrence
KV-Fold: One-Step KV-Cache Recurrence for Long-Context Inference
TL;DR Highlight
모델 수정 없이 KV 캐시를 청크 간 누산기로 쓰면 128K 토큰까지 100% 정확도로 정보를 검색할 수 있다.
Who Should Read
긴 문서 처리나 대용량 코드베이스 분석에서 컨텍스트 한계를 맞닥뜨린 백엔드/MLOps 개발자. 모델을 재학습하지 않고 긴 컨텍스트 추론을 구현하고 싶은 경우.
Core Mechanics
- 긴 시퀀스를 청크로 나눠서 각 청크 처리 시 이전 청크들의 KV 캐시를 prefix로 붙여서 attention하는 방식 — 함수형 프로그래밍의 foldl과 동일한 구조.
- 모델 파라미터 변경, 추가 학습, 특수 메모리 토큰 없이 순수 inference 프로토콜만으로 동작함.
- 청크 경계마다 오류가 쌓일 것 같지만, 실제로는 초반 몇 스텝에서만 drift가 생기고 이후엔 flat plateau로 안정화됨 — 오류 누적이 아님.
- 이 안정화 현상은 수치 정밀도 10,000배 변경(bf16→fp32)에도 drift가 고작 2.8%만 줄어드는 걸 보면, 수치 오차가 아니라 구조적 특성임.
- Llama-3.1-8B, Qwen2.5-7B-Instruct, OLMoE-1B-7B 세 가지 다른 아키텍처 모두에서 동일한 plateau 패턴이 나타나 특정 모델에 국한된 현상이 아님.
- StreamingLLM은 메모리는 아끼지만 캐시 밖으로 나간 토큰은 검색 불가 — KV-Fold는 메모리를 더 쓰는 대신 임의 거리의 정보를 정확히 검색 가능.
Evidence
- Needle-in-a-haystack 벤치마크에서 16K~128K 토큰, chain depth 최대 511까지 152번 시도 전부 100% exact-match 달성 (Llama-3.1-8B-Instruct 기준).
- 128K 토큰 처리 시 단일 full-attention forward는 attention score 행렬만 ~1TB로 A100 40GB에서 불가능 — KV-Fold는 peak 35.57GB, 171초로 완료.
- KV-Fold NLL 2.46 vs StreamingLLM NLL 2.66 (T=128K) — StreamingLLM은 d=31 이상에서 retrieval 0%, KV-Fold는 d=511까지 유지.
- Qwen2.5-7B에서 depth 15~60 구간 총 drift 변화량이 -0.0003 nats로 사실상 0 — 깊은 chain에서도 오류 누적 없음.
How to Apply
- 긴 문서를 256~1024 토큰 청크로 나누고, 각 청크 forward pass 시 이전 청크들의 KV 캐시를 prefix로 concat해서 넘기면 됨 — 모델 코드 변경 없이 inference 루프만 수정.
- 128K 이상 코드베이스 분석이나 긴 로그 파일 처리 시, StreamingLLM 대신 KV-Fold를 쓰면 오래된 위치의 정보도 정확히 검색 가능 — 단, 메모리는 0.13 KB/token씩 선형 증가함을 감안해야 함.
- int8 양자화와 조합하면 128K depth 511에서 93% retrieval 유지 가능 — GPU 메모리가 부족할 때 KV 캐시를 int8로 저장해서 운용하는 방식으로 절충 가능.
Code Example
# KV-Fold 핵심 루프 (HuggingFace transformers 스타일 pseudocode)
import torch
def kv_fold_inference(model, token_chunks, device):
"""
token_chunks: List[Tensor], 각 청크는 shape [1, chunk_size]
"""
past_key_values = None # 누산기 (accumulator)
for chunk_idx, chunk in enumerate(token_chunks):
chunk = chunk.to(device)
with torch.no_grad():
outputs = model(
input_ids=chunk,
past_key_values=past_key_values, # 이전 청크들의 KV 캐시를 prefix로
use_cache=True,
return_dict=True
)
# KV 캐시를 누적해서 다음 청크로 전달
past_key_values = outputs.past_key_values
# 마지막 청크면 logits 사용
if chunk_idx == len(token_chunks) - 1:
logits = outputs.logits
return logits, past_key_values
# 사용 예시
# tokenizer로 긴 시퀀스를 chunk_size=256으로 분할
chunk_size = 256
long_text = "...매우 긴 문서..."
tokens = tokenizer(long_text, return_tensors='pt')['input_ids'][0]
chunks = [tokens[i:i+chunk_size].unsqueeze(0) for i in range(0, len(tokens), chunk_size)]
logits, final_kv = kv_fold_inference(model, chunks, device='cuda')Terminology
관련 논문
Swift로 LLM 학습시키기 Part 1: 행렬 곱셈을 Gflop/s에서 Tflop/s로 끌어올리기
Apple Silicon에서 Swift로 직접 행렬 곱셈 커널을 구현하며 CPU, SIMD, AMX, GPU(Metal)를 단계별로 최적화해 Gflop/s에서 Tflop/s 수준까지 성능을 높이는 과정을 상세히 설명한 글이다. 프레임워크 없이 LLM 학습의 핵심 연산을 밑바닥부터 구현하고 싶은 개발자에게 Apple Silicon의 성능 한계를 체감할 수 있는 드문 자료다.
fsync 없이 로컬 스토리지 엔진을 crash-consistent하게 만든 방법
FractalBits가 fsync 없이 SSD 전용 KV 스토리지 엔진을 구현해 동일 조건 대비 약 65% 높은 쓰기 성능을 달성한 설계 방법을 공유했다. fsync의 메타데이터 오버헤드를 피하기 위해 사전 할당, O_DIRECT, SSD 원자 쓰기 단위 정렬 저널을 조합한 구조가 핵심이다.
Google Chrome, 사용자 동의 없이 4GB AI 모델(Gemini Nano)을 몰래 설치
Google Chrome이 사용자 동의 없이 Gemini Nano 4GB 모델 파일을 자동 다운로드하고, 삭제해도 재다운로드되는 문제가 발견됐다. GDPR 위반 가능성과 수십억 대 기기에 적용될 때의 환경 비용 문제가 제기되고 있다.
OpenAI가 대규모 저지연 Voice AI를 제공하는 방법
OpenAI가 9억 명 이상의 사용자에게 실시간 음성 AI를 제공하기 위해 WebRTC 스택을 어떻게 재설계했는지 설명하는 글로, relay + transceiver 분리 아키텍처의 설계 결정과 trade-off를 상세히 다룬다.
Truncated Decoding Tree의 결정론적 탐색을 통한 효율적인 Test-Time Inference
Self-consistency의 중복 샘플링 낭비를 없애는 결정론적 트리 탐색 디코딩 기법 DLE로 수학/코드 추론 성능과 속도를 동시에 개선
GoModel – Go로 작성된 오픈소스 AI Gateway
OpenAI, Anthropic, Gemini 등 여러 AI 프로바이더를 하나의 OpenAI 호환 API로 묶어주는 Go 기반 오픈소스 AI 게이트웨이로, LiteLLM의 컴파일 언어 대안이다.
Related Resources
Original Abstract (Expand)
We introduce KV-Fold, a simple, training-free long-context inference protocol that treats the key-value (KV) cache as the accumulator in a left fold over sequence chunks. At each step, the model processes the next chunk conditioned on the accumulated cache, appends the newly produced keys and values, and passes the enlarged cache forward; the same one-step update is applied repeatedly, analogous to foldl in functional programming. Building on the KV cache concatenation primitive introduced for latent multi-agent communication, we repurpose it as a chunk-to-chunk recurrence for long-context inference. When processing chunk t, the model attends to the KV cache carried from earlier chunks as a prefix, reusing its internal state across segments without modifying or retraining the model. Despite its simplicity, the induced recurrence is stable: per-step drift rises briefly and then saturates into a flat plateau that persists across deep chains. This plateau is insensitive to a 10,000x change in numerical precision, robust across chunk sizes, and consistent across model families. At the task level, KV-Fold preserves exact information over long distances. On a needle-in-a-haystack benchmark, it achieves 100% exact-match retrieval across 152 trials spanning contexts from 16K to 128K tokens and chain depths up to 511 on Llama-3.1-8B, while remaining within the memory limits of a single 40GB GPU. Compared to streaming methods, which trade fidelity for bounded memory, KV-Fold maintains long-range retrieval while operating as a sequence of tractable forward passes. Overall, our results show that frozen pretrained transformers already support a stable form of KV-cache recurrence, providing a practical route to long-context inference without architectural changes or training.