구글, 대형 언어 모델의 KV 캐시 옮기는 오픈소스 도구 TPU Sync 공개

구글은 대형 언어 모델의 KV 캐시를 TPU 메모리와 서버 메모리, 다른 TPU 노드 사이에서 옮겨 관리하는 오픈소스 도구 TPU Sync를 공개한다고 밝혔습니다. Amin Vahdat 구글 AI·인프라 부문 수석부사장 겸 최고기술책임자는 11월 9일 공식 블로그에서 이 도구를 TPU 클러스터의 모델 상태 계층으로 소개했습니다.

TPU는 인공지능 계산을 위해 설계한 가속기이고, 고대역폭 메모리(HBM)는 칩에서 계산에 필요한 데이터를 빠르게 공급합니다. 대형 언어 모델은 입력 문맥을 읽고 답을 만들 때 앞서 처리한 토큰의 계산 결과를 KV 캐시에 저장해 활용합니다. 문맥이 길어질수록 캐시도 커집니다. 구글은 에이전트가 백만 토큰 규모의 문맥을 다루거나 강화학습 작업을 이어 가는 상황에서 메모리가 병목이 된다고 설명했습니다.

TPU Sync는 HBM과 서버의 호스트 메모리(DRAM), 공유 저장장치, 데이터센터 네트워크를 하나의 메모리 계층으로 연결합니다. 구글은 큰 KV 캐시를 칩 밖으로 밀리초 단위로 옮기고, 입력 문맥을 처리하는 프리필 노드에서 다음 토큰을 생성하는 디코드 노드로 문맥을 전달한다고 밝혔습니다. 필요한 캐시를 다른 메모리 계층이나 노드로 옮겨, 긴 문맥을 처리할 때 가속기 메모리에 집중되는 부담을 나누는 방식입니다.

공개 저장소는 캐시 이동 과정을 더 구체적으로 설명합니다. HBM 용량이 차면 최근에 덜 사용한 캐시 블록을 호스트 DRAM으로 보내고, 다시 필요할 때 비동기식 DMA 전송으로 불러옵니다. 저장소에는 호스트 메모리에 캐시를 옮기는 예제와 여러 TPU 가상 머신 사이에서 문맥을 전달하는 예제가 있습니다. TPU Sync는 JAX와 PyTorch 기반 작업 아래에서 데이터 이동과 캐시 관리를 맡는 C++ 라이브러리로 소개됩니다.

이 도구는 추론뿐 아니라 강화학습에도 쓰입니다. 구글은 강화학습 가중치 동기화를 지원한다고 밝혔고, 공개 저장소는 학습기와 샘플러 사이에서 가중치와 활성값을 옮기는 기능을 설명합니다. 학습기가 모델을 갱신하고 샘플러가 이를 실행해 결과를 모으는 과정에서 두 작업 사이의 데이터 이동을 맡습니다.

구글은 TPU Sync가 KV 캐시 전송 지연과 첫 토큰 응답시간(TTFT)을 90% 넘게 줄인다고 밝혔습니다. 다만 블로그의 해당 설명에는 이 수치의 측정 모델·작업·비교 기준이 제시되지 않았습니다.

공개 저장소의 지원 장비 목록에는 현재 Cloud TPU v6e와 v7x가 적혀 있습니다. 이날 함께 발표한 8세대 TPU 8tx와 8i Max는 목록에 없습니다. 저장소는 API가 활발히 바뀌는 개발 단계라고 안내하며, 첫 안정판이 나오기 전까지 PyPI에서 개발용 nightly 패키지만 제공한다고 설명합니다.