확률적 KV 라우팅: 적응형 층별 캐시 공유

Transformer 언어모델 서빙에서 KV cache는 추론 속도를 지키는 대신 메모리 비용을 크게 키운다. 이 연구는 시간축 압축이나 eviction 대신, 레이어 깊이 방향의 중복을 줄이는 depth-wise cache sharing을 겨냥한다.

기존 cross-layer cache sharing은 처리량 저하나 첫 토큰 지연을 낳기 쉬웠다. 대신 random cross-layer attention으로 학습해 모델이 여러 캐시 공유 방식에 미리 적응하도록 만든다.

• 학습 중 각 레이어는 자신의 KV 상태 또는 바로 앞 레이어의 KV 상태를 무작위로 참조한다.
• 그 결과 배포 시에는 일부 레이어의 캐시를 제거해도 정보 손실 없이 동작할 수 있다.
• pre-training과 fine-tuning 모두에 적용할 수 있고, 여러 모델 계열에서 depth-wise cache sharing을 지원한다.
• 큰 모델과 데이터가 제한된 환경에서는 정규화처럼 작동해 성능을 유지하거나 개선하면서 캐시 메모리를 크게 줄였다.