JAX와 TPU v5e로 SDXL 추론 가속

Hugging Face Diffusers가 JAX를 지원하여 Cloud TPU v5e 환경에서 Stable Diffusion XL(SDXL) 추론을 가속화할 수 있게 되었다.

JAX의 JIT(Just-In-Time) 컴파일 기능을 활용해 정적 형태(static shapes)를 가진 이미지 생성 작업에 최적화된 TPU 바이너리를 생성하며, 이를 통해 매우 빠른 추론 속도를 구현한다.

또한, **XLA 컴파일러 기반의 pmap**을 통해 여러 TPU 칩에 작업을 병렬로 분산할 수 있다. 예를 들어, 8개의 칩을 사용하면 단일 칩에서 이미지를 생성할 때와 동일한 시간 내에 8개의 이미지를 생성할 수 있어 높은 처리량을 보장한다.

Cloud TPU v5e는 TPU v4 대비 비용이 절반 이하로 저렴하면서도 대규모 AI 모델의 학습 및 추론에 최적화된 성능을 제공하여, SDXL과 같은 대형 모델의 효율적인 배포를 가능하게 한다.