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
Related Papers
Training an LLM in Swift, Part 1: Taking matrix mult from Gflop/s to Tflop/s
Apple Silicon에서 Swift로 직접 행렬 곱셈 커널을 구현하며 CPU, SIMD, AMX, GPU(Metal)를 단계별로 최적화해 Gflop/s에서 Tflop/s 수준까지 성능을 높이는 과정을 상세히 설명한 글이다. 프레임워크 없이 LLM 학습의 핵심 연산을 밑바닥부터 구현하고 싶은 개발자에게 Apple Silicon의 성능 한계를 체감할 수 있는 드문 자료다.
Removing fsync from our local storage engine
FractalBits가 fsync 없이 SSD 전용 KV 스토리지 엔진을 구현해 동일 조건 대비 약 65% 높은 쓰기 성능을 달성한 설계 방법을 공유했다. fsync의 메타데이터 오버헤드를 피하기 위해 사전 할당, O_DIRECT, SSD 원자 쓰기 단위 정렬 저널을 조합한 구조가 핵심이다.
Google Chrome silently installs a 4 GB AI model on your device without consent
Google Chrome이 사용자 동의 없이 Gemini Nano 4GB 모델 파일을 자동 다운로드하고, 삭제해도 재다운로드되는 문제가 발견됐다. GDPR 위반 가능성과 수십억 대 기기에 적용될 때의 환경 비용 문제가 제기되고 있다.
How OpenAI delivers low-latency voice AI at scale
OpenAI redesigned its WebRTC stack to serve real-time voice AI to over 900 million users, detailing the design decisions and trade-offs of a relay + transceiver split architecture.
Efficient Test-Time Inference via Deterministic Exploration of Truncated Decoding Trees
Deterministic Leaf Enumeration (DLE) cuts self-consistency’s redundant sampling by deterministically exploring a tree of possible sequences, simultaneously improving math/code reasoning performance and speed.
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.