vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes
cs.CV · cs.DC · cs.LG
Quick verdict
Ports five video-gen families to JAX/TPU with verified weights and full benchmarks; solid systems work, but lacks GPU baselines and quality metrics.
Summary
vidax is an open-source JAX/Flax inference engine plus a zero-copy PyTorch-to-JAX weight translator covering five video-generation families (Wan, Cosmos, LTX, HunyuanVideo, CogVideoX), spanning diffusion transformers, Mixture-of-Transformers, 3D VAEs, text encoders, and native samplers, with no PyTorch in the execution path. The system combines a three-axis sharding mesh (data, Megatron 1D tensor, DeepSpeed-Ulysses sequence parallelism), Pallas TPU flash-attention kernels, host-side precision handling, and per-layer weight offloading so that reference-resolution generation fits TPU v4 HBM (~30.75 GB/chip). Benchmarks on a v4-8 slice (JAX 0.11.0; mean of 5 runs) report compile time, per-step latency, and peak HBM for 16+ configurations, and Appendix A documents real checkpoint-translation bugs (fp32 residual accumulation for Wan2.1, T5 embedding drift for CogVideoX, psum bias double-counting, a scan-based VAE decoder fix). Limitations: generation quality is assessed only qualitatively (Figure 1) with no FVD/VBench or human evaluation; there is no latency comparison against PyTorch/CUDA references on comparable hardware; no run variance is reported; offloaded configurations are slow (Wan2.1 14B at 720p: ~111-141 s/step). Code, configs, and a pinned environment are released, so reproducibility looks strong.
Review complete
Assessment at a glance
5
risk flags identified
● high 0
● medium 2
● low 3
5 strengths · 5 weaknesses · 5 author questions
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- vidax runs five major open video-generation families at native reference resolutions on a TPU v4-8 slice with zero PyTorch dependency in the execution path.
- A unified three-axis mesh (data, Megatron-style 1D tensor, DeepSpeed-Ulysses sequence parallelism) plus Pallas flash attention and per-layer offloading keeps peak HBM within the ~30.75 GB/chip budget for all evaluated configurations.
- Weight translation is exact (1:1 key mapping) with bit-exact or high-correlation block-level parity against official PyTorch checkpoints.
- Per-layer weight offloading is a memory-fit mechanism rather than a throughput-neutral optimization: on Wan2.1 14B at 720p, latency improves from 141.7 to 111.3 s/step as offload chunk size grows from 1 to 40 blocks while peak HBM rises from 15.2 to 26.1 GB.
- Porting real checkpoints surfaced correctness bugs invisible in synthetic testing, including fp32 residual-accumulation requirements in Wan2.1 and psum bias double-counting under offload plus sequence parallelism.
Strengths
Broad, verified architecture coverage
Five model families and 16+ configurations (Table 5) run end-to-end through one JAX execution path, each with exact 1:1 parameter-tree mapping and block-level parity checks against official PyTorch checkpoints (Section 3.2).
Genuinely useful debugging taxonomy
Appendix A/Table 4 documents real-world porting bugs with root causes and fixes: row-parallel psum double-counting replicated biases, dense attention masks bypassing Pallas O(S) savings (fixed via SegmentIds), a 147 GB HLO temporary in LTX-2.5's VAE decoder solved by lax.scan, and T5 embedding drift of 16-37% from bf16 accumulation. Rarely documented at this specificity.
Honest, quantified offloading trade-off
Transferable systems rules
Reproducible benchmark methodology
Weaknesses
No quantitative output-quality evaluation
No cross-framework or cross-hardware baseline
Large performance anomalies left under-analyzed
Per-step metric conflates sampling and VAE decode
Single hardware point and undefined checkpoint labels
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is vidax?
vidax is an open-source JAX/Flax inference engine with a zero-copy PyTorch-to-JAX weight translator. It lets major open-source video generation models (Wan, Cosmos, LTX, HunyuanVideo, CogVideoX) run on Google Cloud TPUs with no PyTorch dependency at execution time, using tensor and sequence parallelism, flash-attention kernels, and per-layer weight offloading to fit reference resolutions in accelerator memory.
Which video models does vidax support?
The initial release covers Wan2.1/2.2, Cosmos-Predict2.5 and Cosmos3, LTX-Video and LTX-2.5, HunyuanVideo and HunyuanVideo-1.5, and CogVideoX/1.5 — including their 3D VAEs, text encoders (UMT5, T5, Qwen2.5-VL, Gemma, Llama, byT5, SigLIP, CLIP), and Euler, UniPC, DDIM, and DPM-Solver samplers reimplemented in JAX.
What hardware was used to benchmark vidax, and how fast is it?
A single TPU v4-8 slice (4 chips, about 30.75 GB usable HBM each) running JAX 0.11.0. Reported per-step times range from about 2 seconds for small models at 480p to roughly 300 seconds for HunyuanVideo 13B at 720p with weight offloading. The paper reports absolute TPU numbers only and includes no GPU comparison.
Does offloading weights to the host slow down video generation in vidax?
Yes, substantially — the paper treats it strictly as a memory-fit fallback, not a speed feature. On Wan2.1 14B at 720p, per-step latency drops from 141.7 seconds with one block per transfer to 111.3 seconds with all 40 blocks resident, while peak per-chip HBM rises from 15.2 GB to 26.1 GB, because host-to-device transfers do not fully overlap with compute.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figure 1 (video stills) and Figure 2 (latency/memory plots) cannot be inspected from extracted text; judgments about visual quality and plot trends rely on captions and in-text numbers only.
- The open-source repository (github.com/FlyingGiraffe/vidax) was not inspected; verification claims such as 'bit-exact or high-correlation parity' rest on the paper's assertions.
- Extracted text appears complete through Appendix B including Table 5 footnotes, so no truncation was detected; confidence is correspondingly high.
- Benchmark values are transcribed from dense multi-column tables; individual numbers carry a small transcription-risk despite careful reading.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.