Visko, Cloud TPUs에서 Orbis 라이브 월드 모델 30일 만에 배포

AI 스타트업 Visko가 Google Cloud와 협력해 인터랙티브 라이브 월드 모델 Orbis를 Cloud TPUs에 단 30일 만에 배포했다. 핵심 과제는 평균 처리량이 아니라 엄격한 실시간 시계였다. 모든 비디오 및 오디오 청크가 재생 속도보다 빠르게 렌더링되어야 세계가 멈추지 않기 때문이다. 배치 크기 1과 엄격한 순차적 의존성을 고려해 파이프라인을 최적화한 결과, 1.07x의 재생 속도를 달성했다.

멀티모달 파이프라인 설계

Orbis 파이프라인은 BF16 정밀도로 실행되는 단일 Cloud TPU slice에서 세 가지 생성 모델을 연결한다.

• 핵심 비디오 확산 모델: 기본 비디오 잠재 변수를 생성한다.
• 초해상도 업스케일러: 해상도를 18 fps에서 1664x936으로 향상시킨다.
• 오디오 생성기: 동기화된 오디오를 생성한다.

개발 속도와 성능의 균형을 맞추기 위해 스택은 세 가지 오픈소스 기둥에 의존한다.

• JAX: 고성능 배열 연산 및 XLA 컴파일용.
• torchax: 완전한 재작성 없이 PyTorch 모델을 JAX를 통해 TPU에서 실행할 수 있게 하는 상호운용성 계층.
• Pallas: TPU용 커스텀 커널 프로그래밍 언어로, 특수화된 어텐션 커널 작성에 사용된다.

실시간 스트리밍을 위한 주요 최적화

레이턴시 예산을 충족하기 위해 엔지니어링 팀은 세 가지 특정 최적화를 구현했다.

1. Ulysses 컨텍스트 병렬 처리: 배치 크기 1에서 모델 가중치 분할 대신 토큰 시퀀스를 TPU 칩 간에 분할했다. 고속 all-to-all 인터커넥트 교환으로 텐서를 재구성하여 각 칩이 어텐션 헤드의 하위 집합에 대해 전체 비디오 시퀀스를 처리하도록 했다. 이 방식은 해당 워크로드에서 텐서 병렬 처리나 링 어텐션보다 빠름을 입증했다.
2. 커스텀 Pallas Splash Attention: Google Cloud의 오픈소스 참조 라이브러리인 MaxDiffusion의 커스텀 splash attention 커널을 통합했다. 이 커널은 FlashAttention 스타일의 타일링을 적용해 연산을 온칩 VMEM에 유지하고, HBM에 전체 2차 어텐션 행렬을 물리적으로 생성하는 메모리 대역폭 병목을 피한다.
3. 선택적 JAX 재작성: torchax는 비디오 및 업스케일링 모델을 TPU로 빠르게 가져오는 데 성공했지만, 오디오 단계는 여전히 병목이었다. 오디오 모델만 torchax에서 네이티브 JAX로 재작성한 결과 5.8x의 속도 향상이 이루어져 전체 파이프라인이 재생 마감 시간 내에 들어왔다.

스타트업을 위한 실무 교훈

이 프로젝트는 초저지연 생성형 AI를 위한 몇 가지 전략을 부각했다.

• 네이티브 BF16 정밀도: 오프라인 생성과 달리, 연속 스트리밍은 정밀도를 지나치게 낮추면 품질 저하를 겪는다. BF16은 속도와 시각적 충실도 사이의 필요한 균형을 제공했다.
• 전략적 포팅: torchax를 빠른 브리지로 사용하여 PyTorch 모델을 실행한 후, 네이티브 JAX 재작성에 투자하기 전에 프로파일링을 통해 실제 병목을 식별한다.
• XLA 캐싱: 컴파일된 XLA 그래프를 캐싱하여 워커 재시작 시간을 몇 분에서 몇 초로 단축했으며, 이는 반복적 개발에 필수적이었다.