TriForce: 계층적 Speculative Decoding으로
긴 문맥(128K) 생성을 손실 없이 가속하기

문제: 긴 문맥에서는 "매 토큰마다 거대한 KV cache 전체를 읽는 것" 자체가 병목

한 줄 요약 — LLM이 토큰을 1개 생성할 때마다 모델 weights 전체 + 지금까지 쌓인 KV cache 전체를 GPU 메모리에서 읽어와야 하는데, 문맥이 128K처럼 길어지면 KV cache가 모델보다 커져서(78GB > 14GB) 생성이 극도로 느려진다. 기존 해법인 KV cache 압축(eviction)은 속도를 얻는 대신 생성 품질을 잃는다.

"key-value (KV) cache … has emerged as a critical bottleneck" — Abstract

"they suffer from degradation in generation quality" — Abstract (기존 KV cache 압축 기법에 대해)

구체적 수치 Fig. 1 : Llama2-7B-128K 기준, full KV cache는 78GB / 순전파 56ms, 4K만 골라낸 partial cache는 16GB / 14ms, 초소형 draft 모델(Llama-68M)은 150MB / 0.45ms. 즉 "매번 78GB를 읽는 것"이 문제의 정체다.

기존 접근법: 무엇을 줄였고, 그 대가는 무엇이었나

TriForce: "3%짜리 나 자신"을 중간 보스로 세운 2단 계층 speculation

핵심 아이디어 한 줄 — 병목이 2개(① 모델 weights, ② KV cache)이므로, draft도 2단으로 쌓는다. 초소형 모델(150MB)이 먼저 초안을 쓰고 → 같은 target 모델이 KV cache의 3%(4K)만 검색해 붙인 상태로 1차 검증하고 → 마지막에 full cache로 최종 검증한다. 최종 검증이 표준 speculative decoding과 동일하므로 출력 분포는 수학적으로 무손실이다 §4.2 Correctness.

먼저 알아야 할 용어 3개 (중학생 버전)

설계의 근거가 된 3가지 관찰 §3

관찰 ① 병목이 문맥 길이에 따라 바뀐다 Fig. 2b, 2c — 짧은 문맥에선 모델 weights가, 긴 문맥에선 KV cache가 지배적 병목. 그래서 각 병목마다 다른 draft를 배정하는 계층 구조가 필요하다.

관찰 ② Attention은 매우 희소(sparse)하다 §3.1, Fig. 3a — 120K 문맥에서 단 4K 토큰으로 attention 점수의 96.8%를 복원 가능. 즉 cache의 3%만 있어도 "거의 같은 모델"처럼 행동한다 → 자기 자신을 draft로 쓰는 self-speculation의 근거.

관찰 ③ 문맥 참조에는 지역성(contextual locality)이 있다 §3.2, Fig. 3c — 지금 토큰이 참조하는 옛 문맥과, 바로 다음 토큰들이 참조하는 옛 문맥이 거의 같다. 따라서 한 번 골라둔 4K 캐시를 여러 디코딩 스텝에 재사용할 수 있어 캐시 구축 비용이 분할상환(amortize)된다.

구성 요소 ① Retrieval-based drafting: 버리지 말고, 찾아 쓰자 §4.1

full cache를 어차피 (최종 검증용으로) 보관하므로, draft cache는 버리는(eviction) 방식이 아니라 "골라오는(retrieval)" 방식으로 만들 수 있다. KV cache를 chunk(작은 덩어리)로 쪼개고, 현재 query와 각 chunk의 평균 key 사이 attention 점수를 계산해 점수 높은 chunk들만 4K 예산에 채운다. 시간(recency)이 아니라 관련성(relevance) 기준 선택이다.

"evicted tokens are permanently discarded" — §4.1 (StreamingLLM/H2O가 needle retrieval에 부적합한 이유)

구성 요소 ② Hierarchical speculation: 병목 2개 → draft 2단 §4.2, §4.3

retrieval cache로 KV 병목은 풀었지만, draft가 7B weights 전체를 여전히 로드하므로(14ms) 이번엔 weights가 병목이 된다. 그래서 그 아래에 Llama-68M + StreamingLLM cache(0.45ms)를 한 단 더 깐다.

구성 요소 ③ Retrieval cache의 생애주기 §4.3

prefix의 마지막 토큰으로 중요도 순 캐시를 최초 구축 → 이후에는 새 토큰의 KV가 중요도 최하위 항목을 덮어쓰며(overwrite) 유지 → 수락률의 이동평균이 임계값 아래로 떨어지거나 정해진 stride마다 전체 재구축(rebuild). 관찰 ③(지역성)이 있기에 재구축을 드물게 해도 된다.

숫자로 손으로 돌려보기 — 정상 케이스 [논문에 없음: Fig.1 수치로 직접 계산한 예시]

γ₁=2, γ₂=6, Fig. 1의 계층별 비용(0.45 / 14 / 56ms)을 그대로 사용한 1 iteration:

합계 ≈ 100.7ms에 토큰 8개 = ~12.6ms/token. vanilla는 56ms/token이므로 산술상 ~4.4×. 실제 측정치 2.31× Table 2는 여기에 retrieval 계산·캐시 갱신·커널 오버헤드가 더해진 값이다. (거절이 잦아질수록 이 수치는 빠르게 나빠진다 — s₈의 붕괴 케이스 참고)

모든 Figure/Table은 각각 하나의 주장을 판다

실험: 어떤 설계 지점이 어떤 숫자를 만들었나

Table 2 — On-chip (A100, 122K prompt, 256 생성) §5.1

Table 3 — Offloading (소비자 GPU) §5.1

Table 4 — Batching (A100) §5.1

Fig. 5 — Scalability, Table 5 — 타 기법 비교, Fig. 6 — Ablation §5.1~5.3

Fig. 5: 문맥 32K→128K로 갈수록 TriForce의 speedup은 우상향하고(이론 상한 13.1×), naive policy는 상한이 1.87×에 묶인다. 이유는 두 가지 — 문맥이 길수록 ① draft(4K 고정)와 target(선형 증가)의 지연 격차가 벌어지고 ② acceptance는 유지되기 때문 §5.1 Analysis.

Fig. 6a: 예산 4K에서 acceptance-오버헤드 균형이 최적(그 이상은 drafting 지연만 증가). Fig. 6b: chunk가 너무 작으면 개별 토큰에 과적합, 너무 크면 고득점 토큰이 저득점에 희석 → 중간이 최적. Fig. 6c: Sequoia 같은 tree 기반 speculation과 결합하면 이론 생성 토큰 수가 추가 상승 — TriForce는 그들과 직교(orthogonal)함을 어필.

결론

§6 논문 명시 결론: attention sparsity와 contextual locality라는 두 관찰로부터, KV cache와 model weights라는 이중 병목을 계층적 speculation으로 완화했다. A100에서 최대 2.31×, 2×RTX4090 offloading에서 7.78×(0.108s/token)를 달성했고, 이는 장문맥 모델 serving의 실용화 가능성을 보여준다. 별도의 future work 절은 없으며, tree-based speculation과의 결합(Fig. 6c)이 사실상의 확장 방향으로 제시된다.

Appendix에서 놓치기 쉬운 것들

구현 디테일 Appendix A — 캐시 예산이 고정이라 CUDA graphs로 커널 런칭 오버헤드를 제거할 수 있었고(speculation처럼 작은 실행이 많은 구조에서 중요), FlashAttention 사용. 특히 §3에서는 설명 편의상 첫 2개 layer를 full cache로 뒀지만, 실제 시스템은 모든 layer를 sparse로 유지한다 — acceptance가 약간 떨어지는 대신 drafting 비용이 상수로 고정되는 쪽을 택했다.

더 긴 입력 Appendix B.1 — L40 1장 + offloading으로 LWM-Text-256K에서 11.81×, 512K에서 12.10×. Fig. 5의 "길수록 유리" 주장을 실측으로 연장.

더 긴 출력 Appendix B.2 — 출력 256~2048 토큰에서 speedup이 2.28~2.34×로 안정. 캐시가 주기적으로 재구축되며 문맥 진화를 따라가기 때문.

γ 최적값 Appendix B.3 — γ₁=2, γ₂=6이 최적. 68M draft와 7B+retrieval 사이의 격차 때문에 1단계 speculation의 acceptance가 낮아 γ₁을 길게 못 가져간다는 고백이 여기 숨어 있다.

Limitations 정리

FAQ

비판적으로 읽기: 이 시스템은 언제 무너지는가

TriForce의 모든 이득은 하나의 가정 위에 서 있다: "다음 수백 토큰이 참조할 옛 문맥 ≈ 방금 참조한 옛 문맥"(contextual locality). 이 가정이 깨지는 순간, 계층 전체가 연쇄적으로 비효율이 된다.

수치 워크스루 — 붕괴 케이스 [논문에 없음: 동일 비용 모델로 구성한 사고실험]

합계 ≈ 145ms + rebuild 비용에 3토큰 → ~50ms+/token으로 vanilla(56ms)와 사실상 동급. 정상 케이스(12.6ms)와 4배 차이. 수락률이 시스템의 손익분기점을 직접 결정함을 보여준다.

논문이 다루지 않은 층위

① 대규모 serving 층위 — 평가는 batch ≤ 10, 그마저 speedup이 단조 감소(Table 4). vLLM류 continuous batching과의 결합, 요청 간 캐시 관리 등은 미답. ② 메모리 층위 — "무손실"의 대가로 78GB full cache 보관이 전제라, 애초에 그 메모리가 없는 환경(순수 메모리 부족 문제)은 이 논문의 해법 범위 밖이며 그때 필요한 것이 오히려 ZeRO-Inference류다. ③ 워크로드 층위 — needle retrieval은 "retrieval cache에 유리한" 과제이기도 하다. 지역성과 희소성이 동시에 약한 과제(예: 문서 전체를 골고루 훑는 요약)의 acceptance는 보고되지 않았다. ④ 비교 공정성 — ZeRO-Inference 대비 4.86×는 상대 시스템이 CPU-attention 구현이라는 조건이 붙는다(각주 1).

유효 적용 범위