TPU Sync 공개: KV 캐시 전송 병목과 서빙 성능을 구분하는 법

AI 생성 개념 이미지. 실제 TPU 장비나 성능 측정 화면이 아니다.

모델을 더 빠른 가속기에 올렸는데도 응답이 느리다면, 연산뿐 아니라 데이터를 옮기는 시간을 봐야 한다. 특히 긴 문맥을 처리하는 서빙에서는 KV 캐시를 어디에 보관하고 언제 가져오는지가 중요하다.

Google이 2026년 11월 9일 공개한 TPU Sync는 이 데이터 이동을 다루는 계층이다. PyTorch 모델을 TPU에서 실행하는 TorchTPU와 달리, KV 캐시와 가중치를 메모리·호스트 사이에서 옮기는 문제에 초점을 둔다. 기존 TorchTPU 실행 방식 소개와 구분해서 볼 만한 새 변화다.

무엇이 달라졌나

공식 발표에 따르면 TPU Sync는 가속기의 HBM, 호스트 DRAM, 공유 스토리지와 네트워크를 연결하는 모델 상태 전송 계층이다. 텐서의 원래 배치를 유지하는 DMA 전송, 연산과 전송의 중첩, XLA 그래프 밖에서 동작하는 네이티브 C++ 경로를 제공한다.

적용 대상은 추론 중 KV 캐시 이동만이 아니다. 강화학습에서 학습기가 갱신한 가중치를 샘플러에 전달하는 흐름도 포함한다. 발표에는 높은 전송 처리량이 제시되지만, 특정 장비와 소프트웨어 조합에서 Google이 측정한 결과다. 이를 서비스 응답 속도의 동일한 배수 개선으로 읽으면 안 된다. Google 공식 발표

KV 캐시가 얼마나 커지는지 먼저 계산한다

KV 캐시는 이미 처리한 토큰의 attention key와 value를 저장해 재사용하는 데이터다. 다음은 모든 계층에서 같은 구조를 사용하는 단순한 모델의 비압축 캐시 용량을 계산하는 예다. 특정 제품이나 실제 모델의 측정값이 아니다.

이 산술 코드를 로컬 Python에서 실행한 결과는 4.00 GiB였다. TPU 실행이나 전송 성능 실험을 수행한 것은 아니다.

여기서 헤드 수는 query head가 아니라 KV head 수다. 실제 사용량에는 패딩, 블록 할당, 캐시 공유, 양자화, 샤딩, 계층별 attention 구조가 영향을 준다. 따라서 이 계산은 용량 감각을 잡는 출발점이지 배포용 메모리 견적의 최종값이 아니다.

이 정도의 데이터를 저장할 공간이 있다는 것과, 다음 토큰을 내보내기 전에 제때 가져올 수 있다는 것은 다른 문제다. 캐시를 외부로 내보내 HBM 여유가 늘어도 전송 대기 때문에 지연이 커질 수 있다. 이는 오프로딩을 검토할 때 함께 확인해야 할 설계상 절충이다.

도입 전에는 버전부터 고정한다

공개 저장소는 Apache 2.0 라이선스이며 TPU v7x와 v6e 검증을 명시한다. API가 변화 중이므로 main을 무조건 따라가기보다 태그 또는 저장소의 lkg.version이 가리키는 검증된 커밋을 사용하도록 권고한다.

특히 JAX, libtpu, XLA 계열 의존성을 임의로 섞지 않아야 한다. 저장소는 호환되지 않는 조합이 빌드를 통과하고도 런타임에 실패할 수 있다고 경고한다. 도입 시점의 호환표를 다시 확인하고, 모델 서버까지 포함한 버전 묶음을 기록하는 편이 안전하다.

호스트 DRAM 캐시와 컨테이너 재시작 시 재사용하는 공유 메모리 캐시도 구분해야 한다. /dev/shm을 활용한다는 설명을 물리 서버 재부팅 후에도 유지되는 영구 저장소라는 뜻으로 확대하면 안 된다. 공식 저장소와 호환성 안내

전송 벤치마크와 사용자 지연을 나눠 검증한다

다음은 공식 성능 수치가 아니라, 도입 효과를 판단하기 위한 검증 제안이다.

1. 기준선을 고정한다. 같은 모델, 프롬프트 길이 분포, 출력 길이, 동시 요청 수와 하드웨어에서 기존 설정을 측정한다.

2. 캐시 상태를 나눈다. 캐시가 없는 첫 요청과 재사용 가능한 요청을 별도로 기록한다. 두 결과를 섞으면 무엇이 개선됐는지 알기 어렵다.

3. 사용자 지연을 잰다. 첫 토큰까지 시간과 토큰 간 지연을 각각 p50·p95로 비교한다. 전송 처리량만으로 통과 판정을 내리지 않는다.

4. 자원 비용을 함께 본다. HBM 감소량뿐 아니라 호스트 RAM, 네트워크 트래픽, 동시 요청 처리량을 기록한다.

5. 실패 경로를 확인한다. 캐시 미스, 용량 부족, 프로세스 재시작 때 요청이 정상 복구되는지 별도 테스트한다.

예를 들어 전송량이 늘었어도 첫 토큰 지연이 개선되지 않았다면, 해당 워크로드의 병목이 전송이 아니거나 전송이 임계 경로에서 충분히 겹쳐지지 않았을 수 있다. 이는 측정 후 확인할 가설이지 TPU Sync 자체의 결함을 단정하는 근거는 아니다.

지금 필요한 판단

TPU Sync는 모든 개발자가 즉시 설치할 범용 최적화 옵션이라기보다, TPU 기반 추론과 학습 시스템에서 모델 상태 이동을 다루는 도구다. TPU 환경이 없다면 발표 숫자를 재현하려고 유료 자원을 먼저 만들 필요는 없다.

우선 자신의 서비스가 연산, 메모리 용량, 상태 이동 중 어디에 묶여 있는지 분리하자. 새로운 전송 계층의 가치는 가장 큰 GB/s 숫자가 아니라, 실제 요청의 지연과 비용을 얼마나 바꾸는지로 판단하는 것이 맞다.

검증 범위: 공식 발표·공개 저장소 확인, 용량 계산 코드 로컬 실행. TPU 하드웨어 실행, 호환성 실험, 성능 비교는 수행하지 않았다.