MAGI-2 预览版在 MUSA 上通过 vLLM-Omni 实现每步 565 毫秒

摘要

MAGI-2 预览版是 Sand AI 的 1140 亿参数音视频生成模型，现可通过 torchada 兼容层在摩尔线程 MUSA GPU 上运行，每去噪步 565 毫秒（8 × MTT S5000，272p，静态编译），比首个可用版本（每步 15.6 秒）提升了 27.6 倍。加速来自一系列针对性 PR，这些 PR 消除了主机同步、融合了 MoE 内核、重写了 mHC 残差，并优化了专家并行分片。

1. MAGI-2 预览版作为独特推理负载的特点

  • 单流 DiT 主干将文本、视频潜变量和音频潜变量 token 打包到一个自注意力序列中；只有第一层和最后四层具有模态特定权重。
  • 无 adaLN – 时间步直接编码在每个 token 的前 64 个通道中，允许视频和音频时间步独立。
  • 多头 MoE（第 2-37 层）将每个 3072 维 token 拆分为十二个 256 维头；每个头路由到其 256 个专家中的前 6 个（SwiGLU7），外加一个共享专家和三个按模态共享的专家。路由专家约占参数的 95%。
  • mHC 超连接提供四流残差（总宽度 12288），具有通过 20 次 Sinkhorn 迭代计算的双随机 4 × 4 混合矩阵。
  • 采样运行 100 个未蒸馏步骤，每个步骤有两个 CFG 分支，一次性对完整的 10 秒视频进行去噪。
指标 值
总参数 1140 亿（每个 token 约 60 亿活跃）
主干 40 层 DiT，隐藏层 3072，24 头 × 128
MoE 12 头 × 256 专家，每头 top-6
残差 mHC，4 流，20 次 Sinkhorn 迭代
输出 10 秒，125 帧，12.5 fps，44.1 kHz 立体声
支持的分辨率 272p（448 × 256）和 540p（896 × 512）

与其他开源视频模型相比，MAGI-2 预览版是唯一一个在单个自注意力流中结合了每 token MoE 路由、mHC 残差和联合音视频生成的模型。

2. 使用 torchada 在 MUSA 上运行 CUDA 优先模型

  • torchada 拦截 CUDA 风格的导入（torch.cuda、设备字符串、NCCL、CUDA Graph），并将其映射到 MUSA 等效项（torch.musa、MCCL、MUSAGraph）。映射刻意保持 真实（torch.cuda.is_available() 保持 False），因此平台检查必须使用 torch.version.musa。
  • 代码复用：eager PyTorch、torch.compile/Inductor 以及上游 Triton 内核（融合 BF16 MoE、SwiGLU7、mHC）在 MUSA 上无需修改即可运行。只有启动配置和少数平台特定路径不同。
  • 关键适配：
    • torch.isfinite 替换为 abs()<inf 表达式以适配 MUSA（torchada #120）。
    • Int32 路由计数（scatter_add_）在 MUSA 上替换 int64，每次调用节省约 39 毫秒（PR #8496）。
    • GEMM 操作数布局重新排序为流主序，消除了每步约 80 次复制（PR #8510）。
    • Triton-MUSA 缺少 BF16 原子操作，因此始终选择融合 BF16 MoE 路径。

“torchada 使 MUSA 张量的 is_cuda 为 true，因此 CUDA 风格的代码可以正常运行。” – torchada 文档。

3. 优化时间线（每步延迟）

阶段 每步时间（毫秒） 主要变化
初始启动（9 月） 15,600 8 × MTT S5000 上的基线
主机同步移除 8,800 .item() → .tolist()，减少 all-gather 元数据
MoE 融合 GEMM（MATE） 2,390 分组 BF16 GEMM 替换逐专家 Python 循环
融合 W13 + SwiGLU7 Triton 954 单个内核用于 gate/up + 激活
Int32 路由计数 756 PR #8496
K 维分块（BLOCK_SIZE_K=64） 734 PR #8506
mHC 逐元素重写 734 → 734（基线）→ 733.6 PR #8510（FP32 收缩）
Sinkhorn 单内核 –6.6 PR #8594
W13 累加器拆分 590 – 591 PR #8855（内核提速 38–40%）
完整 PR 栈（10 月 19 日） 565 应用上述所有 PR，静态编译

27.6 倍整体加速（15.6 秒 → 565 毫秒）是在相同硬件和工作负载（272p，100 步，T2VA 提示，SP8×CFG1+EP4）上测量的。每个项目都针对其自身基线进行评估；数字不是跨行累积的。

4. 详细技术收益

4.1 主机端同步

  • 逐专家 .item()：替换为批量 .tolist()，每步节省约 3.4 秒。
  • All-gather token 计数元数据：从每步 39 次调用减少到 3 次，节省约 61 毫秒。
  • torch.isfinite：在 MUSA 上重新实现为 abs()<inf，每步减少约 1.8 秒。

4.2 MoE 内核融合

  1. MATE 分组 GEMM 将 12 × 256 维专家 GEMM 折叠为单个 BF16 GEMM（21.9 毫秒 → 0.69 毫秒）。
  2. 融合 W13 + SwiGLU7 Triton 内核 结合了 gate、up 投影和激活，将 MoE 瓶颈从约 1.5 秒减少到 <1 秒。
  3. 权重布局预交错 消除了每次调用的转置，将步时间缩短约 7.5%。
  4. Int32 路由直方图 从路由阶段削减约 1.2 秒。
  5. K 维分块（BLOCK_SIZE_K=64）带来约 2% 的适度提升。
  6. 确定性路由构建（PR #8592）增加了约 12 毫秒。

4.3 mHC 残差重写

  • 将小 K bmm 替换为 FP32 逐元素收缩和流主序布局，每步减少约 250 毫秒（约 25%）。
  • 将 20 次 Sinkhorn 迭代合并为单个 Triton 内核（-6.6 毫秒）。

4.4 注意力

  • 从 Torch 实现切换到 MATE FlashAttention-3（FA3），带有学习到的 LSE 汇点，在 SP4×CFG2 配置上实现了 16.8% 的减少。
  • 缓存 RoPE、模态分组和汇点-LSE 元数据（-6 毫秒）。

4.5 专家并行分片

  • 头分片 EP（EP4）与序列并行解耦，提供了最快的配置（565 毫秒）。EP8 较慢（633 毫秒），因为更多分片增加了通信开销。
  • 移除 Ulysses 发送/接收周围的缓冲区复制，额外节省了 10.6 毫秒（PR #8595）。

4.6 编译

  • 区域 torch.compile 相比 eager 执行将每步时间减少了约 16%。
  • 在动态编译下固定非 token 维度，将 1.4 毫秒的内核压缩到 55 微秒（PR #8627）。
  • MUSA 上的动态编译需要 Sinkhorn 修复（#8594）才能产生正确输出。

5. 实测性能结果

5.1 PR 树端到端计时（静态 vs 动态）

并行配置 最终 PR 轮次前（静态） 静态编译（毫秒/步） 动态编译（毫秒/步）
SP8×CFG1+EP4 714 / 715 665.7 / 669.9 678.3 / 678.9 / 683.7
SP4×CFG2+EP4 732.6 690.8 –
SP4×CFG2 783.1 729.9 –

565 毫秒数字（第 4 节）来自 2026 年 10 月 19 日测量的 最新 树；它与上述行不直接可比，因为它使用不同的代码库和会话。

5.2 W13 累加器拆分 A/B 测试（PR #8855）

树 静态编译（毫秒/步） 输出哈希（音频 / 视频）
对照（无拆分） 690.8 / 691.1 32f1b4fca4 / 44b9747e9e
+ #8855（拆分） 591.3 / 590.4 相同

14.5% 的减少是位相同的；相同的 SHA-256 哈希确认数值等价。

5.3 最新树上的专家并行度扫描

配置 平均每步时间（毫秒） 相对于 EP4
SP8×CFG1，EP 关闭（旧版） 697.8 +23.5%
SP8×CFG1+EP4 565.0 基线
SP8×CFG1+EP8 633.4 +12.1%

所有三种配置产生相同的音频和视频哈希，确认速度差异纯粹是通信/计算权衡。

6. 示例输出和质量说明

  • 基准提示（金毛犬，272p，SP8×CFG1+EP4）在每步 565 毫秒内生成了 10 秒视频；片段可在博客资源中获取。
  • 官方 Sand AI 提示（咖啡店、动漫、黏土动画）在 T2VA 和 I2VA 模式下运行。所有样本均符合文本描述，包含非零音频，咖啡店片段包含可听的英语对话。
  • 540p 预览以约 3.6 秒/步（静态）和约 3.9 秒/步（动态）运行；这些数字与 272p 计时不直接可比，因为它们使用不同的代码树和编译模式。
  • 尚未进行定量多种子质量评估；视觉和听觉保真度已手动验证。

7. 上游状态和兼容性

7.1 PR 概况（截至 2026 年 10 月 28 日）

PR 功能 状态 实测每步影响
#7206 融合 BF16 MoE（W13） 已合并 –
#7243 可配置 Dynamo 重编译限制 已合并 –
#7252 融合 SwiGLU7 已合并 –
#7261 融合 mHC 后残差 已合并 –
#8496 int32 路由计数 已合并 1,952 → 756 毫秒
#8497 W13 单副本 已合并 –
#8498 MUSA 支持（FA3 varlen + sink） 已合并 –
#8506 MUSA 上 MoE GEMM 的 BLOCK_SIZE_K=64 已合并 754.5 → 733.9 毫秒
#8510 mHC FP32 收缩和流主序布局 已合并 984.7 → 733.6 毫秒
#8511 头分片 EP 独立于 SP 开放 864.6 → 717.5 毫秒
#8592 位精确路由构建 已合并 –12 毫秒
#8593 RoPE 和汇点-LSE 元数据缓存 开放 –6 毫秒
#8594 单内核 Sinkhorn 开放 –6.6 毫秒
#8595 移除 Ulysses 缓冲区复制 开放 –10.6 毫秒
#8627 固定注意力头布局（动态） 已合并 1,429 微秒 → 55 微秒
#8855 拆分 W13 累加器（MUSA） 开放 690.8 → 590.4 毫秒（静态）
#8915 1080p 精炼器（块窗口注意力） 草稿 –

所有列出的 PR 要么已合并，要么正在审查中；本文中的性能数字反映了每个 PR 测量时的代码状态。

7.2 兼容性边界

  • 完全基准测试：MUSA + SP8×CFG1+EP4 在 272p 下（第 5.1-5.3 节）。
  • 部分支持：540p 预览（仅样本）、I2VA 模式（仅样本）、1080p 精炼器（草稿实现，未计时）。
  • 未验证：逐层卸载、量化、Cache-DiT 和其他内存节省技术。

8. 结论和未来工作

  • 加速总结：通过系统消除主机同步、融合 MoE 内核、重写 mHC 残差和改进专家并行分片，从 15.6 秒基线降至每步 565 毫秒（27.6 倍）。
  • 剩余瓶颈：
    1. MoE W13 GEMM 仍受内存带宽限制；进一步分块（128 宽 K）或自定义分组内核可能再节省约 2–3 毫秒。
    2. 通信-计算重叠：EP4 每步产生 60–90 毫秒的集合通信，未与计算重叠。
    3. 主机空闲时间：FlowUniPC 主机逻辑（约 19 毫秒）和内核内启动间隙（约 14 毫秒）仍然存在。
    4. 动态编译开销：多个内核中的 int64 大小算术；转换为 int32 可恢复约 85% 的性能损失。
  • 下一步：
    • 在 MUSA 上启用图重放（目前比区域编译慢）。
    • 完成 1080p 精炼器集成（块窗口注意力、内存规划）。
    • 扩展 CI 以覆盖 540p 基准测试和仅包含已合并 PR 的主分支。
    • 探索量化和 cache-DiT 技术以进一步降低延迟。

9. 致谢

作者感谢 Sand AI 开源 MAGI-2 预览版并提供官方提示，感谢摩尔线程 MUSA、torchada 和 MATE 团队提供 GPU 栈，以及原始博客文章中列出的众多贡献者实现和审查 PR。

10. 可复现性

  • 容器镜像：registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview（摘要 sha256:2785b4626d1da1e0765a4cb515b9afc5349a22c6417e14c81091add0c1a9b902）。
  • 命令（在 8 × MTT S5000 上运行）：
    docker pull registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview
docker run -d --name magi2 -v /ipfs:/ipfs -v "$PWD/magi2-out":/work \
  --entrypoint sleep registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview infinity
docker exec magi2 bash -lc '
  cd /opt/vllm-omni-src &&
  python /opt/magi2/instrument-steps.py &&
  VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONUNBUFFERED=1 \
  python /opt/magi2/magi2_e2e_gate.py \
    --ulysses-degree 8 --cfg-parallel-size 1 \
    --enable-expert-parallel --expert-parallel-size 4 \
    --resolution 272p --height 256 --width 448 \
    --seconds 10 --fps 12.5 --num-frames 125 --num-inference-steps 100 \
    --no-pin-cpu-memory --metrics-dir /work/steps --save-outputs /work/outputs'
  • 环境变量用于确定性 Inductor 行为，已在容器内设置，如源文章附录 A.5 所述。

所有数字、表格和引用均直接取自 2026 年 10 月 12 日的 vLLM 博客文章。没有虚构任何额外数据。

Sources