AoiZora: Diffusion Transformer 추론을 위한 토폴로지 인식 자동 병렬 최적화

비디오 확산 (Video diffusion) 모델은 빠르게 핵심적인 생성 서비스 워크로드로 성장했지만, 각 클립을 생성하기 위해서는 거대한 시공간 잠재 변수 (spatio-temporal latents)에 대해 수많은 디노이징 (denoising) 반복 과정이 필요하며, 이는 단일 장치에서 저지연 추론 (low-latency inference)을 달성하기 어렵게 만듭니다. 따라서 디노이징 단계는 일반적으로 여러 가속기에 분산되어 수행되며, TPU 서브 슬라이스 (sub-slices)는 이를 수행하기 위한 매력적이고 실용적인 패브릭 (fabric)이 되었습니다. 그러나 현재의 자동 병렬 (auto-parallel) 시스템은 거의 전적으로 논리적 장치 메쉬 (logical device meshes) 상에서만 탐색을 수행하며, 선택된 샤딩 (sharding)이 물리적 TPU 상호 연결 (interconnect)에 실제로 어떻게 배치되는지는 무시합니다. 이러한 간과로 인해 토폴로지 (topology)에 따라 달라질 수 있는 상당한 성능 향상의 기회를 놓치게 됩니다. 우리는 TPU 서브 슬라이스에서의 저지연 비디오 확산 추론을 위해 구축된 컴파일러 매개 토폴로지 플래너인 AoiZora를 통해 이 격차를 해결합니다. AoiZora의 핵심 원칙은 컴파일 흐름의 서로 다른 지점들을 활용하여 논리적 샤딩과 물리적 배치를 다시 연결하는 것입니다. AoiZora는 먼저 비용이 적게 드는 사전 컴파일 IR (pre-compilation IRs)에서 취약한 샤딩 후보들을 제거한 다음, 살아남은 후보들만 컴파일하고 컴파일된 HLO와 토폴로지 인식 통신 모델을 사용하여 이들의 물리적 배치를 정렬합니다. 최종 선정된 계획은 일반적인 컴파일러 경로를 따라 구현되므로, 모델 코드, 컴파일러 로워링 (lowering), 집합 커널 (collective kernels) 및 네트워크 라우팅 (network routing)을 완전히 온전하게 유지합니다. TPU v5e 서브 슬라이스에서 AoiZora는 기존 솔루션 대비 Wan 2.1의 1단계 디노이징 지연 시간을 최대 1.42배까지 단축합니다.