vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes

2026-09-15

AI Summary

为了解决Cloud TPU缺乏生产就绪的视频生成模型问题，提出vidax框架，使用JAX/Flax实现并提供PyTorch到JAX的权重转换器。

Abstract

Open-source video generative models ship almost exclusively as PyTorch/CUDA reference implementations. This leaves Cloud TPU pods without a production-ready inference path, despite offering large, cost-effective accelerator memory pools ideal for long-sequence spatiotemporal attention. We present vidax, an open-source JAX/Flax inference engine and zero-copy PyTorch-to-JAX weight translator for modern video generation architectures. vidax covers a diverse set of spatiotemporal models --- including Diffusion Transformers, omnimodal Mixture-of-Transformers, 3D VAEs, text encoders, and native samplers --- with zero PyTorch dependency in the execution path. The framework unifies 1D tensor parallelism with DeepSpeed-Ulysses sequence parallelism on a single JAX sharding mesh, integrates TPU flash-attention kernels, and implements per-layer weight offloading to support reference resolutions that exceed single-device memory. We benchmark compile times, latency, and peak memory utilization on TPU v4-8 hardware, and document real-world numerical bugs surfaced during checkpoint translation. vidax is released open-source as a baseline for JAX and TPU video generation research.

Problem

• video generative models
• Cloud TPU
• inference path
• spatiotemporal attention
• accelerator memory

Innovation

• JAX/Flax inference engine
• zero-copy weight translation
• spatiotemporal models