Windowed-MTP: 백만 토큰 컨텍스트에서 Full-Context Draft-KV 비용 제거하기
요약
백만 토큰 이상의 긴 컨텍스트 환경에서 Multi-Token-Prediction(MTP) 초안 생성 시 발생하는 KV 캐시 읽기 비용 문제를 해결하기 위한 Windowed-MTP 방식을 제안합니다. 슬라이딩 윈도우와 어텐션 싱크를 적용하여 학습 없이도 초안 생성 비용을 획기적으로 절감하면서 타겟 모델의 출력 품질은 그대로 유지합니다.
핵심 포인트
- MTP 초안 생성 시 컨텍스트 길이에 비례해 증가하는 KV 캐시 읽기 비용 문제 지적
- Windowed-MTP는 별도의 학습이 필요 없는 training-free 및 lossless 방식
- 1M 컨텍스트 기준 KV 엔트리의 약 99%를 제거하여 비용 절감
- Qwen 및 Mamba2-hybrid 모델 테스트 결과 디코딩 비용을 최대 44% 절감
- 타겟 모델의 검증 프로세스를 유지하여 출력 분포의 품질 저하 방지
추측적 디코딩 (Speculative decoding)은 저렴한 초안 (draft) 모델이 토큰을 제안하면 타겟 모델이 이를 병렬로 검증함으로써 자기회귀 생성 (autoregressive generation)을 가속화합니다. 최첨단 모델 (Frontier models)들은 초안 생성 비용이 무시할 수 있을 정도로 저렴하다는 가정하에, 내장된 Multi-Token-Prediction (MTP/NEXTN) 초안 헤드를 점점 더 많이 탑재하고 있습니다. 하지만 백만 토큰 컨텍스트 (million-token context) 환경에서는 이 가정이 깨집니다. MTP 초안 헤드는 일반적으로 매 초안 단계마다 전체 KV 캐시 (KV cache)에 대해 전체 어텐션 (full attention)을 실행하므로, 읽기 비용이 컨텍스트 길이에 따라 선형적으로 증가하며 초안 생성 비용의 대부분을 차지하게 됩니다. 이는 추측적 디코딩이 가장 가치 있는 지점에서 발생합니다. 이러한 효과는 초안 길이 (draft length)가 길어질수록 심화되며 (깊은 네이티브 초안은 순효과를 음수로 만들어 추측을 하지 않을 때보다 더 느려질 수 있음), 검증 비용은 저렴하지만 초안의 전체 어텐션 읽기 비용은 그대로 노출되는 하이브리드/선형 어텐션 (hybrid/linear-attention) 타겟 모델에서 더욱 두드러집니다.
우리는 초안의 어텐션에만 StreamingLLM 스타일의 슬라이딩 윈도우 (sliding window)와 어텐션 싱크 (attention sink)를 적용하는 Windowed-MTP 방식을 제안하며, 전체 어텐션 검증 (full-attention verification)은 그대로 유지합니다. 이 방식은 별도의 학습이 필요 없는 (training-free), 즉시 적용 가능한 (drop-in), 그리고 구조적으로 손실이 없는 (lossless) 방식입니다. 전체 어텐션 타겟 모델이 여전히 수락된 모든 토큰을 결정하므로, 윈도우 방식은 제안되는 토큰이 무엇인지에만 영향을 줄 뿐, 수락되는 토큰에는 결코 영향을 주지 않습니다. 이를 통해 초안의 KV 작업 집합 (working set)을 상수로 제한하며, 1M 컨텍스트에서 KV 엔트리의 약 99%를 제거합니다.
SGLang 환경의 단일 GPU에서 1M 컨텍스트를 대상으로 세 가지 아키텍처 제품군 (Qwen GDN-MoE 35B/122B 및 Mamba2-hybrid NoPE 120B)에 대해 테스트한 결과, 윈도우 방식은 기존의 네이티브 MTP 초안 대비 디코딩 단계당 비용을 +28%에서 +44%까지 절감하였으며, 이 차이는 컨텍스트가 길어질수록 더 커지는 입력 불변적 (input-invariant) 마진을 보여주었습니다. 토큰당 지연 시간 (per-token latency)은 이 비용을 수락 길이 (acceptance length)로 나눈 값이므로, 수락 조건이 동일할 때 엔드 투 엔드 (end-to-end) 디코딩 지연 시간은 동일한 수준으로 개선되며, 윈도우 방식이 수락률을 높이는 경우에는 더 큰 개선을 보입니다. 동시에 타겟 모델의 검증된 출력 분포는 그대로 유지합니다. 마지막으로, 읽히지 않은 초안 KV (1M 컨텍스트에서 전체 KV의 7.7-11%)는 수락률이나 품질 저하 없이 컴팩트한 링 버퍼 (ring buffer)를 통해 재사용됩니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.CL (NLP)의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기