JAXBench: Google Cloud TPU를 위한 AI 생성 커널 최적화 벤치마크
JAXBench: Benchmarking Autonomous TPU Kernel Optimization
JAXBench는 Google Cloud TPUs에서 AI 생성 커널 최적화를 위한 TPU 네이티브 벤치마크 스위트로, 50개의 JAX 작업을 포함하고 있다. 17개의 생산 ML 연산자는 Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, AlphaFold2와 같은 공개 MaxText 라이브러리에서 추출되었으며, 33개의 연산자는 KernelBench에서 검증된 후 새로운 문제 크기로 설정되었다. JAXBench는 1.28배의 기하 평균 속도 향상을 달성하며, Autocomp의 빔 검색 파이프라인은 XLA에 비해 1.36배의 기하 평균 속도 향상을 보인다. 또한, 8개의 수동 조정 커널에서 Autocomp는 XLA에 비해 1.60배의 기하 평균 속도 향상을 기록하였다.
원문 출처
arXiv cs.AI (인공지능)