vLLM 핵심 개발팀이 창업한 Inferact가 일련의 추론 벤치마크를 공개했다. 구글 TPU v7 Ironwood 16개로 문샷AI(Moonshot AI)의 Kimi K3를 구동한 결과 단일 스트림 디코딩 속도가 초당 709토큰에 달했다. 반면 비교군인 엔비디아 GB200 16개는 초당 452토큰에 그쳐, 약 57% 더 빨랐다.
벤치마크 보고서와 코드는 9월 23일 함께 공개됐으며, 저장소 inferact/tpu-megakernels는 아파치 2.0 라이선스로 오픈소스화됐다. Inferact는 vLLM 원년 멤버들이 세운 회사로, 앞서 a16z와 라이트스피드(Lightspeed) 주도로 시드 라운드 1억5000만 달러를 유치했으며 기업가치는 8억 달러로 평가받았다.
스펙만 보면 GB200이 우위
먼저 양쪽의 원 스펙을 보자. 보고서에 나온 수치에 따르면 TPU v7 1개는 BF16 최대 연산 성능 2.31 PFLOPS, FP8 4.61 PFLOPS를 내며 HBM 용량 206GB, 대역폭 7380 GB/s를 갖췄다. 이에 비해 GB200 1개는 각각 2.5 PFLOPS, 5 PFLOPS이고 HBM 용량 186GB, 대역폭 8000 GB/s다. 연산 성능과 대역폭 모두 GB200이 앞서며, TPU가 앞서는 부분은 메모리 용량이 20GB 더 많다는 점뿐이다.
Kimi K3는 92개 레이어로 구성된 MoE 모델로, 어텐션 부분은 Kimi 델타 어텐션과 멀티헤드 잠재 어텐션(MLA)을 혼합해 사용한다. 테스트에서는 텐서 병렬 4, 전문가(expert) 병렬 8 방식으로 16개 칩에 모델을 분할했다.
추측 디코딩을 켜지 않은 순수 속도 비교는 다음과 같다.
| 배치 크기 | TPU v7(megakernel) | GB200(vLLM) |
|---|---|---|
| 1 | 249 | 127 |
| 2 | 392 | 227 |
| 4 | 515 | 373 |
| 8 | 865 | 636 |
단위는 초당 토큰 수다. 배치 크기가 1일 때는 TPU가 거의 두 배 빠르지만, 배치 크기가 8이 되면 격차가 30% 남짓으로 좁혀진다. 딥시크(DeepSeek)가 제안한 추측 디코딩 기법 DSpark를 켜면 단일 스트림 속도는 709 대 452가 되며, 스텝당 디코딩 시간은 약 8.5밀리초다.
비결은 92개 레이어를 하나의 프로그램으로 합친 것
Inferact는 이 방식을 메가커널(megakernel)이라 부른다. 일반적인 추론 프레임워크에서는 각 레이어의 연산이 여러 개의 독립적인 커널로 나뉘어 순차적으로 실행되고, 커널 사이에는 공백이 생긴다. 또한 가중치는 해당 커널이 실행을 시작해야만 메모리에서 칩으로 옮겨지기 시작한다.
Inferact는 Pallas를 이용해 92개 레이어 전체의 디코딩 스텝을 하나의 프로그램으로 작성했다. 보고서에 따르면 메가커널은 커널 간 경계를 없애기 때문에 가중치 로딩이 그것을 사용하는 커널에 묶이지 않게 되고, 칩 내부 메모리가 허용하는 만큼 미리 앞당겨 가져올 수 있다. 단일 스트림 디코딩에서는 칩이 대부분의 시간을 데이터 대기에 쓰는데, 이 프리페치가 정확히 그 공백을 메운다. 배치 크기가 커져 연산이 차지하는 비중이 늘어날수록 이 기법의 효과는 줄어드는데, 이는 위 표의 추세와도 맞아떨어진다.
부수적인 이점은 컴파일 시간이다. 기존 XLA 컴파일은 30분 이상 걸렸지만, 메가커널은 90초가 채 걸리지 않는다. 속도 향상이 정확도 희생 없이 이뤄졌음을 보이기 위해 보고서에는 TPU에서 구동한 Kimi K3의 점수도 실렸다. GPQA-Diamond 94.4%, GSM8K 97.2%다.
연산 성능 대비 산출량을 대략 계산해보면
단일 스트림 추측 디코딩 결과로 대략 계산하면, TPU v7의 BF16 최대 성능은 GB200의 약 92%에 불과하지만 토큰 처리량은 1.57배에 달한다. 최대 연산 성능 단위당 산출량으로 환산하면 TPU 쪽이 GB200의 약 1.7배가 된다. 대역폭 기준으로 환산해도 비율은 비슷하게 약 1.7배로 나온다.
다만 이 수치에는 유보 조건이 붙는다. GB200 쪽은 vLLM이 공개한 Kimi K3 레시피를 사용했는데, 이는 메가커널 수준의 최적화를 거치지 않은 일반적인 멀티커널 방식이다. 보고서에는 동등한 최적화를 적용한 GB200 결과가 포함돼 있지 않아, 격차 중 얼마가 하드웨어에서 오고 얼마가 소프트웨어 투자에서 오는지는 현재로선 분리해낼 수 없다. 엔비디아는 이 수치에 대해 아직 응답하지 않았다.
문샷AI 입장에서 이번 벤치마크는 Kimi K3를 엔비디아가 아닌 하드웨어에서 구동할 수 있는 실질적인 경로를 제시한 셈이다. Kimi K3는 파라미터 2조8000억 개 규모의 오픈 웨이트 모델로, 이 정도 규모는 직접 구축하는 데 진입 장벽이 높다. 오픈소스화된 TPU 커널이 생기면서 클라우드에서 TPU를 빌려 K3를 구동하는 선택지가 하나 더 늘었다.
참고 출처: Inferact 기술 블로그, CocoLoop, inferact/tpu-megakernels 코드 저장소, 량쯔웨이(量子位, QbitAI); 처리량 수치는 Inferact가 진행한 16 대 16 칩 테스트 기준이며, 투자 유치 정보는 공개 보도를 바탕으로 한다.