vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes
Congyue Deng
cs.CV, cs.DC, cs.LG
2026-09-16
MIT开源的JAX推理引擎vidax把五家主流视频生成模型搬上Cloud TPU：权重零拷贝导入，TPU v4-8上LTX-Video 13B单步5.2秒，HunyuanVideo-1.5单步221秒，附录记下移植时踩到的数值坑。
开源视频生成模型几乎全是 PyTorch/CUDA 参考实现。Wan、Cosmos、LTX、HunyuanVideo、CogVideoX 这些已经能用的权重，Cloud TPU 上没有现成的生产级推理路径。TPU 的 HBM 大、单价相对便宜，本来适合视频这种「时空注意力把序列拉到几万 token」的负载，JAX 这边却一直缺一层能把各家架构收进同一套并行网格的推理框架。
视频生成和 2D 图像不是同一档内存画像。多帧 latent 切成 patch 之后，序列长度上到数万，二次注意力和每 token 激活往往比参数本身更吃显存。原生参考分辨率要跑起来，序列并行、flash-attention、权重卸载是刚需。
vidax 是 MIT 的 Congyue Deng 做的 JAX/Flax 推理引擎，外加一套零拷贝的 PyTorch 到 JAX 权重翻译。执行路径不依赖 PyTorch。
权重先在 host 上以 NumPy 数组暂存，再一次性 deviceput 到已经按 mesh 切好的设备上。太早上卡、直接 jnp.array，会先在单芯分配一份未切分副本，大 DiT 直接 OOM。模块写成显式 setup()，不用 @nn.compact，前向拆成预处理、block 循环、后处理，按层卸载和分段 JIT 才好做。
并行叠在一张三轴 JAX mesh 上：数据并行、Megatron 风格的 1D 张量并行、DeepSpeed-Ulysses 序列并行。张量并行切线性层。序列并行沿时间维切开 token，只在 attention 前 alltoall 成按头切的布局。Wan2.2 的 per-token AdaLN 给每个 token 一套 scale/shift，激活内存压过参数量，序列并行变成必选项。shardmap 里要手写局部输出形状和 psum，否则 replicated 的 bias 会被加 TP 份。
注意力走 Pallas 的 TPU flash-attention，完整 QK 矩阵不再物化。自定义 kernel 不走 GSPMD 自动切分，多卡 flash-attention 必须包进 shardmap。LTX-2.5 的 neighborhood-attention VAE 没有现成 TPU kernel，退回 scan/vmap 窗口注意力。3D RoPE 各家约定不能互换，每家保留官方公式。
整棵参数树常驻会挤掉 VAE 缓冲时，典型是 Wan2.1 14B 的 720p，按层从 host 流进一块固定形状的 HBM，做法接近 ZeRO-Offload。host 到设备的传输和计算叠不齐，所以这是防 OOM 的备选，不是默认加速。JIT 三条硬规矩：时空维度保持静态，编译图才能复用；采样循环、卸载扫描、分块 VAE 解码留在 Python 宿主循环，整段采样不要套一个 jax.jit，否则所有步的中间激活被描进同一张 HLO；精度下降在 host 做完再上卡，避免设备上短暂出现双精度两份数组。
首发覆盖五家：Wan 2.1/2.2、Cosmos Predict2.5/Cosmos3、LTX-Video/LTX-2.5、HunyuanVideo 13B 与 1.5、CogVideoX 1.0/1.5。同一品牌下架构并不连续，Wan2.1 和 Wan2.2、Cosmos-Predict2.5 和 Cosmos3、LTX-Video 和 LTX-2.5 都是分开实现的。采样器三套共用：Euler 整流流（含 LTX-2.5 要的 ancestral SDE）、UniPC、DDIM/DPM-Solver。多数权重跑 bfloat16；Wan2.1 的 DiT 和 LTX-2.5 的 AdaLN 表按官方 checkpoint 留 fp32。
评测在一片 TPU v4-8（4 芯）、JAX 0.11.0，batch 1、开 classifier-free guidance，每项 5 次平均，分辨率、帧数、步数用各模型官方默认。可用 HBM 大约每芯 30.75 GB。
| 模型 | 切分 | 分辨率 / 帧 | 编译 / 单步 | 峰值 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 |
同一套 JAX 路径能出提示词对得上的视频，速度却差一个数量级以上。LTX 系单步几秒。HunyuanVideo-1.5 无窗全局注意力，旁边还要常驻一份 Qwen2.5-VL，720p 单步 221 秒，编译 410.8 秒。需要 offload 的配置更慢。Wan2.1 14B 在原生 720p 上扫 chunk size：chunk=1 时单步 141.7 秒、15.2 GB，40 层全部常驻降到 111.3 秒、26.1 GB。CogVideoX1.5 的联合注意力大约 4.5 万 token，张量并行在 v4-8 上编不过，只能 SP=4，峰值 31.5 GB，贴着每芯上限。
移植过程里踩到的数值坑比速度表更有用。Wan2.1 14B 在 720p、81 帧下残差用 bf16 会糊，必须按官方在 fp32 里累加。CogVideoX 的 T5 残差到约 10^5 量级，再对未遮罩的 padding 做注意力，bf16 下 embedding 相对 PyTorch 漂 16%–37%。LTX-Video 默认 TPU matmul 精度只对到约 2 位小数，改成 highest 最大差约 3e-5。LTX-2.5 缺 guidance rescale 0.7 会过曝，checkpoint 里 290 张 AdaLN 表必须留 fp32。序列并行加卸载时，row-parallel 的 psum 把 replicated bias 加了 TP 份，Wan 画面变块。LTX-2.5 VAE 朴素窗口注意力一度描出 147 GB 临时量，改成按行 scan 再加 tp=4，压到每芯 14.76 GB。Cosmos-Predict2.5 的彩色网格伪影来自错误的 EDM preconditioning 和 unpatchify 通道顺序。
这是基建，不是新的生成方法。手里有 TPU 配额（Google TRC 或 GCP）又想跑开源视频权重，这是目前公开的那条路。权重翻译、三轴 mesh、flash-attention 封装、按层卸载，都是可复用的工程件。
做移植的人，附录那张 bug 分类表比延迟数字值钱。混合精度不能按全局 dtype 一刀切，自定义 kernel 是 GSPMD 的边界，VAE 窗口注意力能把 HLO 撑爆。这些坑用合成权重测不出来。
论文没有和 GPU 上的官方 PyTorch 实现做同设置延迟对照，也没有 FVD、VBench。表格不能读成「TPU 更快」，也不能读成模型质量排名。
只做了推理。训练、LoRA、全参微调写在未来工作里。benchmark 停在 TPU v4-8，v5e 和 v6e 没测。offload 自己写明传输和计算叠不齐，是内存拟合手段。HunyuanVideo-1.5 编译超过 400 秒、单步 221 秒，离交互式服务还远。定性结果共用一条提示词，没有标准视频质量指标。单作者，Google TRC 赞助。
「单步延迟」把采样循环和 VAE 解码加在一起再除以步数。LTX-2.5 换成 diffusion VAE decoder 的那几行，被 decode 的一次编译主导，不能拿来跟 DiT 采样比。