PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Congyue Deng

arXiv:2609.18077v1cs.CVcs.DCcs.LG

TL;DR

Open-source video models largely lack a production-ready Cloud TPU inference path, despite video attention's substantial memory and compute demands. vidax provides a JAX/Flax engine that combines parallelism, TPU attention kernels, and weight offloading. It verifies five model families and runs reference-resolution generation within device memory constraints.

  • Problem

    Long spatiotemporal token sequences make video generation memory- and compute-intensive, while Cloud TPU users lack a unified JAX inference framework for these models.

    Source: §1 · p. 1

  • Method

    vidax is an open-source JAX/Flax inference engine and weight translator that combines tensor and sequence parallelism, TPU flash attention, and per-layer offloading.

    Source: §1 · p. 2, §2 · p. 2

  • Results

    All five model families are verified end-to-end, and vidax runs reference-resolution video generation within device memory constraints.

    Source: §3.2 · p. 4, §4 · p. 5, §5 · p. 6

  • Takeaways & Limitations

    The porting process surfaced numerical and structural edge cases in official checkpoints, and vidax is released as a baseline for JAX and TPU video-generation research.

    Source: §5 · p. 6

    Benchmarks focus on TPU v4-8; the authors state that the same execution path carries over to TPU v5e and v6e.

    Source: §4 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Video generation's long spatiotemporal token sequences make attention and activation memory demanding, while Cloud TPU users lack a unified JAX inference framework for these models. vidax addresses this gap with an open-source JAX/Flax engine and weight translator.

Motivation and gap

  • Video patch volumes create tens of thousands of tokens, making quadratic attention and per-token activations major memory and compute demands.The paper argues that native-resolution execution therefore needs sequence parallelism, optimized attention kernels, and memory offloading.Source: §1 · p. 1
  • Cloud TPU users lack a unified JAX inference framework for video models whose sequence lengths demand multidimensional parallelism.Porting also requires addressing checkpoint-loading memory staging, compilation-boundary execution safety, and composing tensor with sequence parallelism.Source: §1 · p. 1

Contributions

  • vidax provides an open-source JAX/Flax inference engine and weight translator, with no PyTorch dependency in execution.Its contributions include broad model support, a three-axis sharding mesh, TPU flash attention, per-layer offloading, and TPU v4-8 evaluation.Source: §1 · p. 2

2 System Architecture & Design

vidax combines host-resident weight staging, tensor and sequence parallelism, TPU attention kernels, and optional per-layer offloading to execute video models under accelerator memory constraints. Static shapes and host-side loops are used to control compilation and memory behavior.

Host-resident weight translation

  • Host-side NumPy staging lets a single device_put allocate checkpoint parameters once, already sharded, avoiding unsharded single-device allocation spikes.The model separates preprocessing, block-loop, and postprocessing stages using explicit setup() hooks.Source: §2 · p. 2

Spatiotemporal tensor and sequence sharding

  • A three-axis JAX mesh combines tensor and DeepSpeed-Ulysses sequence parallelism, with token partitioning addressing activation-heavy models that tensor parallelism alone cannot handle.Sequence parallelism uses all_to_all transposes to head-sharded views for attention; local shapes and psum reductions are manually defined when composing both forms.Source: §2 · p. 2

3D RoPE, 3D VAEs, and flash attention

  • Pallas flash attention avoids materializing quadratic query-key matrices, while explicit shard_map wrappers handle its multi-device calls.Models without native TPU kernels can use a scan/vmap windowed-attention primitive; 3D VAE encoding and decoding run as JIT-compiled chunks.Source: §2 · p. 2

Per-layer weight offloading

  • Per-layer host-to-device offloading streams parameters into a fixed HBM buffer to prevent out-of-memory errors at reference resolutions.A reused JIT-compiled block function mutates buffers in place; transfer latency does not fully overlap with compute, so offloading is an optional fallback.Source: §2 · p. 2

JIT compilation and memory ergonomics

  • Static spatial and temporal dimensions enable graph reuse, while host Python loops around JIT functions avoid tracing whole sampling or decoding passes into one HLO graph.Host-side precision downcasting avoids temporary dual-precision arrays on accelerator chips.Source: §2 · pp. 2–3

3 Supported Models

vidax supports five video-model families with distinct architectures, verifies checkpoint conversion and end-to-end behavior, and provides shared samplers for different diffusion and flow-matching formulations.

3.1 Architecture coverage

  • The five supported families—Wan, Cosmos, LTX, HunyuanVideo, and CogVideoX—span DiTs, omnimodal models, varied conditioning towers, and distinct VAE designs.Examples include Cosmos3's Mixture-of-Transformers, LTX-2.5's optional neighborhood-attention VAE decoder, and CogVideoX's 3D causal VAE.Source: §3.1 · p. 3

3.2 Weight import, precision, and verification

  • All five families are verified end-to-end, with exact parameter-tree matches and bit-exact or high-correlation block-level parity checks against PyTorch.Wan2.1 retains fp32 DiT weights and LTX-2.5 retains specific fp32 AdaLN tables to prevent measurable quality degradation.Source: §3.2 · pp. 3–4

3.3 Diffusion and Flow Matching Schedulers

  • Three shared sampler families cover rectified-flow Euler, UniPC, and DDIM/DPM-Solver, supporting the different formulations used across models.UniPC achieves comparable sampling quality in fewer steps, with 35 steps for Cosmos versus approximately 50 for Wan.Source: §3.3 · p. 4

4 Empirical Evaluation & Benchmarks

The evaluation benchmarks five model families on a TPU v4-8 slice at reference settings and examines generation quality, latency, and memory use. Offloading can make demanding configurations fit, but its chunk size trades HBM use against per-step latency.

Experimental setup

  • Benchmarks cover all five supported model families on a single TPU v4-8 slice, using reference resolutions, frame counts, and sampling steps; each metric averages five end-to-end runs.Runs use batch size 1 with classifier-free guidance enabled, equivalent to a 2× internal batch; the compilation cache is cleared before each run.Source: §4 · p. 4

Qualitative results

  • Across five text-conditioning tower families, three scheduler formulations, and seven VAE designs, configurations produce temporally coherent, prompt-faithful videos through one JAX execution path.Figure 1 shows representative model pairs generated using a shared standardized prompt at each model’s reference settings.Source: §4 · p. 4, Figure 1 · p. 5

Denoising latency and memory

  • Offloaded configurations generally have significantly higher per-step latency because host-to-device parameter transfers do not fully overlap with compute; LTX-2.5 22B instead offloads to free intermediate activations.LTX-2.5 22B weights fit within TP = 4 HBM.Source: §4 · p. 5, Figure 2 · p. 6

Memory optimization analysis

  • Peak HBM stays near the usable TPU v4 budget of ∼30.75 GB per chip; without the sharding and offloading mechanisms, native reference resolutions trigger out-of-memory errors.CogVideoX1.5 and native-720p Wan2.1 I2V sit at the budget ceiling; Table 3 lists primary constraints and mitigations.Source: §4 · p. 5, Table 3 · p. 6

Denoising latency and memory

  • 141.7 s at offload chunk size 1 falls to 111.3 s with all 40 blocks resident, while peak HBM rises from 15.2 GB to 26.1 GB.The sweep is for Wan2.1 14B at native 720p; the reported latency decreases monotonically as chunk size increases.Source: §4 · p. 6, Figure 2 · p. 6

5 Conclusions & Future Work

vidax brings video-generation models to Cloud TPUs through a JAX/Flax inference engine and weight translator, combining parallelism, flash attention, and offloading. Future work targets hardware generalization, fine-tuning, and lower-latency kernels.

Conclusions

  • vidax runs major open-source video-generation models on Cloud TPUs with zero PyTorch dependency during execution.The engine combines tensor and sequence parallelism, TPU flash-attention kernels, and per-layer weight offloading to run reference-resolution generation within device memory constraints.Source: §5 · p. 6

Future Work

  • Future work includes extending benchmarks and performance validation beyond TPU v4 slices to TPU v5e and v6e.Source: §5 · p. 6
  • The existing sharded execution graph and offloading infrastructure may be extended to post-training workloads, including LoRA and full-parameter fine-tuning.Source: §5 · p. 6
  • Future kernel work targets custom Pallas and Mosaic kernels tailored to the TPU memory hierarchy to reduce latency overheads.Source: §5 · p. 6

A Debugging Notes and System Verification

Validation on pre-trained checkpoints exposed correctness edge cases and performance bottlenecks not revealed by synthetic testing or unsharded execution. The section describes diagnostic practices and a boundary between automated sharding and custom-kernel communication.

System verification

  • Real-checkpoint validation exposed correctness edge cases and performance bottlenecks that synthetic testing or unsharded execution missed.Table 4 categorizes issues by numerical precision, parallel execution, architecture parity, and memory graph boundaries.Source: A Debugging Notes and System… · p. 8

Diagnostic Methodologies

  • Encoding a real photo, adding slight noise, then denoising once and decoding isolated a preconditioning bug from general weight or attention faults.The probe was used when outputs contained unstructured noise or grid artifacts, which made standard forward passes difficult to diagnose.Source: A Debugging Notes and System… · p. 8
  • Numerical parity required tensor-level precision audits, preserving fp32 for LTX-2.5 AdaLN tables and Wan2.1 residual accumulators.Source: A Debugging Notes and System… · p. 8

Systems Insights on Declarative Sharding

B Full benchmark table

This section reports performance and memory evaluations across supported DiT configurations on TPU v4-8, alongside measurement conventions and execution settings.

B Full benchmark table

  • Each benchmark row averages five end-to-end runs at batch size 1 with classifier-free guidance, after clearing XLA's compilation cache and using bf16 inputs and outputs.Resolutions, frame counts, and step counts follow each model family's reference specifications; prompts and conditioning images are standardized for comparability.Source: B Full benchmark table · p. 8
  • Compile time is a one-time cost per execution signature, per-step latency includes sampling plus VAE decoding, and peak HBM is measured per chip.LTX-2.5 diffusion-VAE-decoder per-step values are dominated by VAE decode compilation and execution, so they are not directly comparable across VAE variants; usable TPU v4 memory is about 30.75 GB per chip.Source: B Full benchmark table · pp. 8, 10
  • Table 5 reports performance and memory across evaluated TPU v4-8 configurations; Table 4 classifies correctness and performance bugs from cross-framework porting.Source: B Full benchmark table · p. 8, Table 4 · p. 9, Table 5 · p. 10
  • Tensor and sequence parallelism plus per-layer offloading manage memory; offloading can preserve VAE buffers or fit a 93-frame context, while Wan2.2 A14B also shards activations.Wan2.2 5B instead uses TP = 4; offloading chunk size trades transfer-compute overlap against peak memory.Source: B Full benchmark table · p. 10
  • LTX uses TP = 4 over 121 frames, and LTX-2.5 22B uses chunk-8 offloading to free inter-block activations; native-720p Wan2.1 requires fp32 parameters.Wan2.1's smaller 480p/1.3B rows use bf16 safely; HunyuanVideo-1.5's 720p execution has 221.0 s/step under unwindowed global attention, and its 13B configuration uses chunks 20/40.Source: B Full benchmark table · pp. 10–11
  • CogVideoX1.5 5B uses SP = 4 because plain tensor parallelism fails TPU v4-8 compilation; its 5B I2V variants use fixed sibling-T2V resolutions.CogVideoX 2B uses TP = 2 because its 30 attention heads are indivisible by 4; other I2V models dynamically set spatial dimensions from aspect ratio and token-area limits.Source: B Full benchmark table · p. 11

Related papers

Loading 2609.18077v1…