TPU가 Kimi를 57% 빠르게 실행한 이유는 추론 소프트웨어

들어가며

대규모 모델 추론에서 Google TPU가 Nvidia GPU를 앞설 수 있는지는 피크 연산 성능만으로 결정되지 않습니다. Inferact가 공개한 테스트에 따르면 16개의 TPU v7 Ironwood로 Kimi K3를 실행했을 때 초당 709토큰을 기록했습니다. 같은 수의 GB200은 초당 452토큰으로, TPU가 57% 빠른 결과입니다. 이 사례의 핵심은 칩 자체보다 하드웨어에 맞춘 소프트웨어 설계입니다.

테스트의 핵심

• 두 시스템은 같은 Kimi K3 모델과 vLLM 추론 엔진을 사용했고, 주요 차이는 가속기와 저수준 Kernel 구현이었습니다.
• TPU 구성에는 DSpark 추측 디코딩이 적용됐습니다. 작은 모델이 후보 토큰을 먼저 제안하고 큰 모델이 이를 묶어서 검증하는 방식이며, 평균 수용 길이는 6으로 보고됐습니다.
• 추측 디코딩을 끈 Batch Size 1에서는 TPU가 초당 249토큰, GB200이 127토큰을 기록했습니다. Batch Size 8에서는 각각 865토큰과 636토큰이었습니다.
• Qwen 3.8 27B 테스트에서는 4개의 TPU가 초당 1515토큰, 4개의 GB200이 695토큰을 기록했습니다.
• Inferact는 TPU에서 실행한 Kimi K3의 GPQA-Diamond 점수가 94.4%, GSM8K 점수가 97.2%로 GPU 결과와 같았다고 설명했습니다.

다만 위 수치는 Inferact가 제시한 벤치마크입니다. 구현의 가능성을 보여주지만 모든 모델과 서비스 조건에서의 보편적인 우위를 뜻하지는 않습니다.

Megakernel이 바꾸는 것

자기회귀 추론에서는 계산량보다 데이터 이동이 병목이 되는 경우가 많습니다. 기존 방식은 추론을 수백 개의 작은 Kernel로 나누기 때문에 Kernel 사이의 실행과 전달 과정에서 메모리 대역폭이 비게 됩니다. Inferact의 Megakernel은 Kimi K3의 여러 계층 연산을 하나의 Pallas 프로그램으로 합쳤습니다.

이 구조에서는 현재 계층을 계산하는 동안 다음 계층의 가중치를 미리 가져올 수 있습니다. TPU의 VMEM은 소프트웨어가 직접 관리하므로 어떤 데이터를 적재하고 유지하며 해제할지 세밀하게 결정할 수 있습니다. Inferact는 XLA의 자동 최적화에만 의존하지 않고 모델 전체의 데이터 이동을 직접 구성했다고 설명합니다. 컴파일 시간도 30분 이상에서 90초 미만으로 줄었다고 밝혔습니다.

의미와 한계

이번 사례는 추론 경쟁의 기준이 피크 연산 성능에서 데이터 흐름과 메모리 활용 효율로 이동하고 있음을 보여줍니다. TPU에는 새로운 칩뿐 아니라 오픈소스 추론 프레임워크와의 폭넓은 통합이 중요해질 수 있습니다. Nvidia의 CUDA Graphs와 PDL 역시 실행 간격을 줄이려는 같은 방향의 기술입니다.

그러나 현재 Megakernel은 Kimi K3와 같은 특정 구조에 맞춘 구현입니다. 다른 모델로 확장하려면 추가적인 적응 작업이 필요합니다. 더 많은 아키텍처를 지원하고 vLLM 생태계에 안정적으로 통합되며, 실제 지연시간·비용·동시성 조건에서도 우위를 유지할 수 있는지가 향후 평가의 핵심입니다.

출처: 퀀텀비트