구글 연구진이 AI 기반 TPU(텐서 처리 장치) 커널 최적화의 발전을 가속화하기 위한 새로운 벤치마크 스위트 'JAXBench'를 공개했습니다. GPU(그래픽 처리 장치) 분야에서는 이미 자율적인 커널 성능 최적화를 위한 엄격한 벤치마크가 존재했지만, TPU 환경에서는 이와 동등한 표준이 없었습니다. JAXBench는 구글 클라우드 TPU에서 AI가 생성한 커널 최적화를 평가하기 위한 TPU 네이티브 벤치마크로, 이 격차를 해소하고자 합니다.
JAXBench는 총 50개의 JAX 워크로드로 구성되어 있으며, 이는 실제 관련성이 높고 최적화의 여지가 충분한 작업들입니다. 연구진은 Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, AlphaFold2와 같은 공개 MaxText 라이브러리의 아키텍처에서 17개의 프로덕션 ML 연산자(operator)를 추출했습니다. 또한, KernelBench에서 33개의 연산자를 가져와 정확성을 검증하고 TPU v6e MXU 활용도를 높이는 새로운 문제 크기로 설정했습니다. 특히, 17개의 프로덕션 연산자 중 8개는 공개 Tokamax 라이브러리의 수동 최적화된 Pallas 커널(kernel)과 블록 크기 튜닝을 통해 전문가 수준의 상한선(upper-bound baseline)을 설정하여 비교 기준으로 삼았습니다.
연구진은 JAXBench에서 후보 Pallas 커널을 생성하기 위해 네 가지 피드백 기반 방법을 평가했습니다. 그 결과, Gemini 3 Flash 모델을 사용했을 때 모델 규모보다는 특정 대상에 대한 문맥(context)이 Pallas와 같은 문서화가 부족한 DSL(도메인 특화 언어)에 더 중요하다는 것을 발견했습니다. 선별된 TPU 문서를 조건으로 부여했을 때 샘플당 정확도가 5.8%에서 37.3%로 크게 향상되었고, 50개 벤치마크 중 48개를 1.28배의 기하 평균 속도 향상으로 해결했습니다. 정확성이 확보된 후에는 Autocomp의 빔 서치(beam-search) 파이프라인이 XLA(Accelerated Linear Algebra) 대비 1.36배의 기하 평균 속도 향상을 달성하며 상당한 성능 개선을 보였습니다. 수동 튜닝된 8개 커널에서는 Autocomp가 XLA 대비 1.60배의 기하 평균 속도 향상을 기록하며 Tokamax의 상한선(2.08배) 대부분을 따라잡았지만, 특수화된 페이지드 및 래그드 어텐션(paged and ragged attention) 연산자에서는 여전히 뒤처졌습니다. 이는 고품질 TPU 커널 최적화가 여전히 도전적인 과제임을 시사하며, 연구팀은 오픈 소스 기여를 지원하기 위해 JAXBench 벤치마크, 평가 하네스(harness) 및 기준 결과를 공개했습니다.
이러한 벤치마크의 등장은 AI 모델 개발 및 배포의 효율성을 크게 높일 잠재력을 가지고 있습니다. 특히 구글 클라우드 TPU를 활용하는 개발자나 기업에게는 AI가 자동으로 최적화된 커널을 생성하고 평가할 수 있는 도구가 제공됨으로써, 수동 최적화에 드는 시간과 비용을 절감하고 더 높은 성능을 달성할 수 있게 됩니다. 이는 궁극적으로 대규모 언어 모델(LLM)과 같은 복잡한 AI 워크로드의 학습 및 추론 속도를 향상시켜, AI 기술의 상용화와 확산에 중요한 기여를 할 것입니다. JAXBench는 TPU 생태계의 성숙도를 높이고, AI 기반 최적화 연구의 새로운 방향을 제시하는 중요한 이정표가 될 것으로 보입니다.