vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes
Congyue Deng
cs.CV, cs.DC, cs.LG
2026-09-16
vidax ports five video generators to Cloud TPU with zero-copy JAX weights. On TPU v4-8, LTX-Video 13B is 5.2 s/step; HunyuanVideo-1.5 is 221 s/step.
Open-source video generators ship as PyTorch/CUDA reference code. Wan, Cosmos, LTX, HunyuanVideo, and CogVideoX have no production inference path on Cloud TPU, even though TPU HBM is large and relatively cheap for the tens-of-thousands-of-token sequences that spatiotemporal attention produces.
Video is not image generation with a time axis glued on. Patchifying multi-frame latents blows sequence length into the tens of thousands. Quadratic attention and per-token activations then dwarf the parameter count. Native reference resolution needs sequence parallelism, flash attention, and weight offload as baseline requirements.
vidax is a JAX/Flax inference engine plus a zero-copy PyTorch-to-JAX weight translator, written by Congyue Deng at MIT. The execution path has no PyTorch dependency.
Weights stage on the host as NumPy arrays, then a single deviceput places them already sharded. Calling jnp.array too early allocates an unsharded copy on one chip and OOMs large DiTs. Modules use explicit setup() instead of @nn.compact, so the forward pass splits into preprocess, a block loop, and postprocess. That split is what makes per-layer offload and piecewise JIT practical.
Parallelism sits on one 3-axis JAX mesh: data, Megatron-style 1D tensor parallel, and DeepSpeed-Ulysses sequence parallel. Tensor parallel shards linear layers. Sequence parallel cuts tokens along time and alltoalls into a head-sharded layout only for attention. Wan2.2's per-token AdaLN (a separate scale/shift per token) makes activations, not weights, the bottleneck, so sequence parallel is required. Inside shardmap, local output shapes and psum have to be written by hand; otherwise a replicated bias is added once per TP device.
Unmasked attention dispatches to Pallas TPU flash-attention kernels so the full QK matrix never materializes. Those custom calls skip GSPMD, so multi-chip flash-attention is wrapped in shardmap. LTX-2.5's neighborhood-attention VAE has no native TPU kernel; vidax falls back to a scan/vmap windowed primitive. 3D RoPE conventions are not interchangeable across families, so each model keeps its official formula.
When the full parameter tree leaves no room for VAE buffers (Wan2.1 14B at 720p is the running example), layers stream from host into a fixed-shape HBM buffer, ZeRO-Offload style. Transfer does not fully overlap compute, so this is an OOM escape hatch, not a default. JIT rules: keep spatial and temporal dims static so compiled graphs can be reused; leave sampling, offload sweeps, and chunked VAE decode as host Python loops (jitting the whole sampler traces every step's activations into one HLO); downcast precision on the host before device placement, so the chip never holds a dual-precision copy.
Five families ship in the first release: Wan 2.1/2.2, Cosmos Predict2.5/Cosmos3, LTX-Video/LTX-2.5, HunyuanVideo 13B and 1.5, CogVideoX 1.0/1.5. Brand names are not architecture names. Wan2.1 vs Wan2.2, Predict2.5 vs Cosmos3, LTX-Video vs LTX-2.5 each needed a separate implementation. Three sampler stacks are shared: Euler rectified flow (plus the ancestral SDE LTX-2.5 needs), UniPC, and DDIM/DPM-Solver. Most weights run in bfloat16; Wan2.1's DiT and LTX-2.5's AdaLN tables stay fp32 to match the official checkpoints.
All numbers are on one TPU v4-8 slice (4 chips), JAX 0.11.0, batch 1 with classifier-free guidance, mean of 5 runs, each model's own reference resolution, frame count, and step budget. Usable HBM is about 30.75 GB per chip.
| Model | Split | Resolution / frames | Compile / step | Peak HBM |
| Cosmos3 Nano 16B | TP=4 | 1280×704 / 93 | 64.2 s / 7.1 s | 29.5 GB |
| LTX-Video 13B | TP=4 | 1216×704 / 121 | 134.7 s / 5.2 s | 15.3 GB |
| LTX-2.5 22B | TP=4 + offload | 1216×704 / 121 | 87.7 s / 7.3 s | 16.7 GB |
| Wan2.2 A14B | TP=2 SP=2 + offload | 832×480 / 81 | 65.8 s / 43.2 s | 28.4 GB |
| Wan2.1 14B | TP=4 + offload | 1280×720 / 81 | 108.2 s / 123 s | 23.0 GB |
| HunyuanVideo-1.5 8.3B | TP=4 | 1280×720 / 121 | 410.8 s / 221 s | 30.3 GB |
| HunyuanVideo 13B | TP=4 + offload | 1280×720 / 129 | 506.7 s / 299.6 s | 18.4 GB |
| CogVideoX1.5 5B | SP=4 | 1360×768 / 81 | 306.6 s / 52.8 s | 31.5 GB |
The same JAX path produces prompt-faithful clips. Step time still spans more than an order of magnitude. LTX-family models sit at a few seconds per step. HunyuanVideo-1.5, with unwindowed joint attention and a replicated Qwen2.5-VL tower, takes 221 s/step at 720p and 410.8 s to compile. Offload configs are slower still. A chunk-size sweep on Wan2.1 14B at native 720p goes from 141.7 s/step and 15.2 GB at chunk 1 down to 111.3 s and 26.1 GB with all 40 blocks resident. CogVideoX1.5's joint attention is about 45k tokens; tensor parallel fails to compile on a v4-8 slice, so the run uses SP=4 and sits at 31.5 GB, right at the per-chip ceiling.
The porting bugs are the more reusable result. Wan2.1 14B at 720p/81 frames goes hazy if residuals stay in bf16; the reference keeps the residual stream in fp32. CogVideoX's T5 residual hits about 1e5, and unmasked padding in bf16 drifts embeddings 16-37% versus PyTorch. LTX-Video agrees with PyTorch only to about 2 decimals until jaxdefaultmatmulprecision="highest" (max diff about 3e-5). LTX-2.5 washes out without guidance rescale 0.7, and 290 AdaLN tables must stay fp32. Row-parallel psum of replicated biases under offload plus SP makes Wan output blocky. A naive LTX-2.5 VAE window attention traced 147 GB of HLO temporaries; scan plus tp=4 brought that to 14.76 GB/chip. Color-grid artifacts on Cosmos-Predict2.5 came from a wrong EDM preconditioning wrapper and an unpatchify channel swap.
This is plumbing, not a new generator. If you have TPU quota (Google TRC or GCP) and want to run the open video checkpoints, this is the public path. The weight translator, 3-axis mesh, flash-attention wrap, and per-layer offload are the reusable pieces.
For anyone porting video models, the appendix bug taxonomy is the part to read first. Mixed-precision leaves cannot be cast with one global dtype. Custom kernels sit outside GSPMD. VAE window attention can explode the HLO. Synthetic checkpoints do not surface those failures.
There is no head-to-head latency against the official PyTorch/CUDA stacks, and no FVD or VBench. Do not read the tables as "TPU is faster" or as a quality ranking.
Inference only. Training, LoRA, and full fine-tune are listed as future work. Benchmarks stop at TPU v4-8; v5e and v6e are unmeasured. Offload is explicitly a memory-fit mechanism because transfer and compute do not overlap cleanly. HunyuanVideo-1.5's 400+ s compile and 221 s/step are far from interactive serving. Qualitative clips share one prompt. Single author, funded by Google TRC.
"Per-step" is (sampling loop + VAE decode) divided by the step count. LTX-2.5 rows that switch to the diffusion VAE decoder are dominated by a one-time decode compile, so they are not comparable to DiT sampling cost.