MAGI-2 Preview 在 vLLM-Omni 上跑通摩尔线程 MUSA：单步从 15.6 秒降到 565 毫秒

torchada 首先让 MAGI-2 Preview 在摩尔线程 MUSA 上运行起来；随后沿执行路径逐项进行优化。第一个可运行的版本每个去噪步骤需要 15.6 秒；应用下文所有 PR 后，最新测量结果为 565 毫秒，而仅 #8855 一项在自身配对的 A/B 测试中就贡献了 14.5%。

MAGI-2 Preview 是 Sand AI 的音视频生成模型：总参数量约 114B，每个 token 约 6B 激活参数，一次生成可产出 10 秒视频并带同步立体声音频，输入可以是文本（T2VA）或图像（I2VA）。其骨干是单流 DiT，将文本、视频和音频 token 打包进一个序列；其 MoE 按“头”路由，残差使用 mHC 超连接。这种组合使其推理成本特征与常见视频 DiT 不同。

vLLM-Omni 已原生支持 MAGI-2 Preview。在摩尔线程 MTT S5000 上，该实现只需少量平台相关改动即可运行：torchada 将 CUDA 风格的 PyTorch 调用映射到 MUSA，模型中以 CUDA 编写的 Triton 内核也可原样运行。随后我们沿执行路径进行优化，并将每项改动作为独立 PR 提交。

在 8× MTT S5000 上，MAGI-2 Preview 在 272p、100 步下，启用静态编译并应用本文所有 PR 后，每步运行时间为 565 毫秒（第 5.3 节）。单个数字说明不了什么，因此这些结论采用配对测量：一次会话内 A/B 测得对照组为 691 毫秒，应用最后一项改动后为 590–591 毫秒，提升 14.5% 且输出逐位一致（第 5.2 节）；第 5.1 节报告该轮 PR 代码树在静态编译下为 665–670 毫秒、动态编译下为 678–684 毫秒。这些小节使用不同的代码树和会话，彼此之间不作比较。静态编译除 #8510 的 mHC 收缩外保持输出逐位一致，该改动会改变求和顺序。第 2.1 节定义了“每步”的含义。

1. 为什么 MAGI-2 Preview 是一个不同的推理问题

MAGI-2 Preview 只有一个 DiT 骨干，文本、视频和音频共享同一个自注意力序列。图 1 展示了该阶段，以及其下方 Sand AI 为 1080p 输出添加的 refiner：

text prompt (+ first-frame image) -> Qwen3.5-27B text encoder / Wan2.2 VAE encoder
                                  -> single-stream DiT × 100 steps (two CFG branches)
                                  -> TurboVAE (video) + Stable Audio Open (audio) -> 540p output
 
1080p tier (in review as #8915, not measured here):
  Preview video latent -> upsample + renoise -> 6.7B dense refiner × 5 steps -> TurboVAE -> 1920×1088
  audio: decoded once in the Preview stage, reused unchanged

图 1：MAGI-2 推理流水线。Preview 阶段（上方条带）是本文唯一测量的阶段：其 40 层中，两端各 4 层为按模态分别设置权重的稠密层，中间 36 层为多头 MoE，并由 mHC 混合残差流。Sand AI 的 1080p 层级（下方条带）在解码前取 Preview 的 latent，对其进行上采样并重新加噪，用单独 checkpoint 的 6.7B 稠密 refiner 配合块窗口注意力运行 5 个引导步，然后解码；音频只解码一次并复用。该 refiner 在 vLLM-Omni 中对应 #8915。来源：Sand AI 博客、模型卡以及 vllm_omni/diffusion/models/magi2/。

• 一个序列，三种模态：视频 latent、音频 latent 和文本特征按“视频 → 音频 → 文本”的顺序打包进一个变长序列，不使用交叉注意力。两端各 4 层（第 0、1、38、39 层）为每种模态各有一套独立权重。
• 无 adaLN：时间步以 64 维正弦编码直接写入每个 token 的前 64 个通道，因此视频和音频可以各自拥有自己的时间步。
• 多头 MoE（第 2–37 层）：每个 3072 维 token 被拆分为十二个 256 维的“头”，每个头在自己的 256 个专家中选择 top-6。每个专家是一个 256→1280→256 的 SwiGLU7，此外还有 1 个共享专家和 3 个按模态划分的共享专家。路由专家占全部参数的 95%。
• mHC 超连接：残差有 4 条流，总共 12288 宽。每个分支先用一个 FP32 [12288×24] 投影计算混合系数。它们之间的 4×4 残差混合矩阵通过 20 次 Sinkhorn 迭代归一化为双随机矩阵（mHC 论文）。
• 注意力：24 头全多头注意力，head_dim 为 128，双向；每个头有一个可学习的 FP32 注意力汇，输出乘以一个 sigmoid 门。
• 采样：100 步，DiT 未蒸馏。每一步运行两个 CFG 分支，完整的 10 秒一次性去噪，而不是像 MAGI-1 那样按自回归分块。

在下表的开源视频模型中，只有 MAGI-2 Preview 将逐 token MoE 路由、mHC 超连接以及音频和文本结合在同一个自注意力序列中：

对于推理而言，这些差异归结为四个具体问题：

1. 按头分片的 MoE：路由发生在 256 维子向量上，因此专家 GEMM 很小（K=256、N=2560 和 K=1280、N=256）但数量众多，专家并行只能沿“头”维度分片。
2. mHC 是大量小型 FP32 操作：残差流携带的字节数是普通残差的 4 倍，每个分支还要运行 20 次 Sinkhorn 迭代，因此 mHC 受限于内存流量和内核启动次数。
3. 文本也要经过完整模型：文本 token 在每一步都经过注意力和 MoE。两个 CFG 分支被打包在一起，因此注意力必须支持汇和变长拆分。
4. 100 个未蒸馏步骤：几乎全部端到端时间都花在 DiT 上，因此每步耗时是主要的优化杠杆。

2. 测量约定与证据边界

本文分别报告三类证据：

两轮使用不同的代码树，因此本文中任何单项改动的收益都不跨越该边界。唯一横跨两轮的数字是第 4 节中从 15.6 s 到 566 ms 的总体数字：它衡量的是整体工作，而非两个代码状态之间的对比，图中标注了基线重置的位置。

2.1 固定控制项

本文中的并行配置记法：SP 是 Ulysses 序列并行；CFG 是 CFG 并行，将两个 CFG 分支放在不同 GPU 上（配置名中的 CFG1 和 CFG2 表示两个分支的放置方式，而非 5.0 和 7.0 的引导尺度）；EP 是 MAGI-2 的按头分片专家并行（见第 4.5 节）。

比较仅在单行内进行：同一棵树、同一并行配置、同一编译模式，每次优化前后的版本交替运行。每个性能数字都配有正确性检查——在静态编译下，前后输出哈希完全一致。例外是 #8510 的第一个提交，它改变了求和顺序，改为对照 FP64 CPU 参考结果进行检查（第 4.3 节）。

3. 在 MUSA 上运行 CUDA 优先的模型并使用 torchada

MUSA 是摩尔线程的 GPU 编程环境，其接口设计与 CUDA 对齐。在 PyTorch 侧，torch_musa 提供 musa 设备。在此基础上，torchada 将 CUDA 风格的调用映射到 MUSA，使为 CUDA 编写的代码可以照常运行。图 2 展示了 MAGI-2 在 MUSA 上的软件栈。

图 2：MAGI-2 在 MUSA 上的软件栈。torchada 在导入时起作用，将 CUDA 风格的调用（torch.cuda、"cuda" 设备、NCCL/MCCL、CUDA Graph/MUSAGraph、Triton 扩展）映射到 torch_musa 及其下层库上。来源：torchada、torch_musa、MATE、vllm-musa、#8498。

3.1 直接复用的部分

导入 torchada 后，vLLM-Omni 中 MAGI-2 的大部分代码可在 MUSA 上原样运行：

• Eager PyTorch 代码：CUDA 风格的张量、设备和流操作可直接工作。
• torch.compile / Inductor：MUSA 平台支持 Inductor，区域性的 torch.compile 相比 eager 将每步时间缩短约 16%。
• CUDA 编写的 Triton 内核：上游的融合 BF16 MoE（#7206）、SwiGLU7（#7252）和 mHC（#7261）内核可在 MUSA Triton 3.2 上原样运行。每个平台选择两项内容：启动配置（MUSA 上为 BLOCK_SIZE_K=64、num_warps=16、num_stages=1；CUDA 上为 32 / 4 / 3），以及 MoE 路径，因为 MUSA 始终采用无需原子操作的融合 BF16 路径（第 3.3 节）。
• 分布式：Ulysses all-to-all 和按头分片的专家并行均通过 ProcessGroupMCCL 运行，比直接加载 libmccl 的 PyNccl 路径更快（780.7 对 823.7 ms/步）。

3.2 torchada 如何映射 CUDA 风格的 API

import torchada  # import before other torch-ecosystem imports
import torch
 
torch.cuda.is_available()          # False: left unpatched so platform detection stays truthful
torch.version.musa                 # 50200; platform checks use this
x = torch.randn(4, device="cuda")  # the tensor lands on musa:0
x.is_cuda                          # True: intentional, so CUDA-style code paths run as usual

• torch.cuda 指向 torch.musa，amp、stream、graph、memory 和 profiler 等子模块也随之映射。
• "cuda" / "cuda:N" 设备字符串、torch.device("cuda")、.cuda() 和 .to("cuda") 均落到 MUSA 上。
• NCCL 进程组映射到 MCCL，CUDA Graph 映射到 MUSAGraph，profiler 的 CUDA 活动映射到 MUSA。
• Triton 的 tl.extra.cuda 映射到 tl.extra.musa；CUDAExtension 源码在构建时自动移植到 MUSA。
• torch.cuda.is_available() 和 torch.version.cuda 有意保持不变；平台检查使用 torch.version.musa 或框架的平台接口。

3.3 适配说明

在兼容层之上仍有几项需要处理。MUSA 特有的选择要么在 torchada 中一次性解决，要么放在平台检查之后；CUDA 路径保持不变。路由计数（#8496）和动态形状固定（#8627）是所有平台的通用代码；它们不改变输出，其 CUDA 开销未测量：

另外，在审查第三方 PR #7545（融合 mHC post_residual，仍处于打开状态）时，我们建议在 MUSA 上用 tldevice.fma_rn(a, 1.0, b) 替换 PTX 内联汇编 add.rn.f32。两者结果相同，且在 CUDA 上生成的 PTX 不变。

torchada 也在过程中得到改进：FX eager / aot_eager 后端的设备构建（#124）以及 torch 2.6 及更高版本的 cpp_extension 路径接口（#121）均已合并。

这遵循 vLLM 的平台设计：仅限 CUDA 的内核由平台对象控制，而非张量的设备类型。

3.4 首次运行

vLLM-Omni 在九月初合并了对 MAGI-2 的原生支持（#5918）。导入 torchada 后，上游代码在 vLLM 0.28 + vllm-musa 上仅需少量按平台的改动即可运行：设备字符串和加速器检查改为平台检查，两个仅限 CUDA 的快速路径（vLLM 自带的 FlashAttention 扩展和需要 BF16 原子操作的 Triton MoE）在接入 MATE 和 MUSA Triton 之前使用 PyTorch 参考实现。首批数据来自 4 步诊断运行：

SP8 剖析显示，大部分时间并非用于计算。设备内核总计 27.7 s，其中一个 MCCL 广播内核占 16.3 s，主要是在等待其他 rank。BF16 tensor-core GEMM 约耗时 1.0 s；路由排序、bincount 和标量同步合计约 6.8 s。按专家逐个执行的 Python 循环会为每个专家执行一次主机同步。因此顺序是：先移除主机同步，再用真正的 GPU 内核替换 MoE 和注意力。

最初几次运行产生的几乎全是黑帧。这不是内核 bug：在只有 1 或 4 个去噪步的情况下，VAE 输出的均值约为 −1.1，落在被截断的范围内。要检查视觉语义，需要运行全部 100 步。

4. 沿执行路径的优化

图 3：从首次可运行版本到本文所测技术栈的每步耗时（对数刻度）：从 15.6 s 降至 565 ms，快了 27.6 倍。第一个点是 4 步诊断运行的稳态，其余是 100 步运行；除最后一个点外，每个点只应用一项改动，最后一个点是完整技术栈。该序列重新建立了一次基线，即把改动重新应用到干净代码树上，因此只有最后几个点描述的是所测技术栈。所有点：8× MTT S5000；最后一个点是 10 月 19 日的完整 PR 栈（第 5.3 节）。第 5.2 节的 591 ms 来自另一次会话，未在图中绘制。

每步耗时从首次可运行版本的 15.6 s 降至本文所测技术栈的 565 ms，整体快了 27.6 倍。该路径分两轮进行：第一轮在 SP4×CFG2 下达到 770 ms，随后重新实现为上游 PR，其在干净代码树上的基线为 4.0 s。图中标出了这次重新建立基线，只有最后一个点描述的是第 5 节所测的技术栈。下面每个条目都是相对其自身基线测量的，因此各条目之间的时间并不逐项递减。

4.1 移除主机同步

• 按专家 .item()：用单个 .tolist() 替代它，SP8 下的每步耗时从 12.2 s 降至 8.8 s。之后路由对齐完全移到设备上，采用固定容量分配且不回读到主机，在 SP8×CFG1+EP4 下又节省约 99 ms。
• 序列拆分元数据：每一层都用 all_gather 交换各 rank 的 token 计数，然后再读回。从调用方向下传递这些数据，可将每步的 all_gather 调用从 39 次减少到 3 次，并节省约 61 ms。
• torch.isfinite：在 torch_musa 2.11 上，注意力中的有限性检查会等待设备队列排空，每次约 38.6 ms，每步 41 次。torchada 改为对 MUSA 浮点张量计算等价的 abs() < inf（#120），结果相同。在该轮中，每步耗时从 2,894 ms 降至 1,137 ms。

4.2 MoE：从 Python 循环到一个融合内核

1. MATE 分组 GEMM：用 MATE 的 ragged 分组 BF16 GEMM 替换按专家循环，使单个投影从 21.9 ms 降至 0.69 ms，每步耗时从 14.4 s 降至 2.39 s。
2. 融合 W13 + SwiGLU7 Triton 内核：一个内核将路由的 gate/up GEMM 与 SwiGLU7 激活合并，针对 MUSA 调整了分块参数，每步耗时首次降至 1 s 以下（1,464 → 954 ms，基线来自 20 步运行）。该内核后来成为 #7206 和 #7252。
3. 预布局权重：在加载时将 W13 以 kernel 所需的交错布局存储，省去了每次调用时的转置拷贝。每步时间从 841 ms 降至 778 ms（−7.5%），输出不变。在 main 上对应 #8497：它只保留一份 W13 副本，这也避免了在加载时额外存储一份权重副本。
4. int32 计数（#8496）：路由直方图的 scatter_add_ 从 int64 改为 int32。在 MUSA 上，int32 是更快的路径（隔离测试中为 0.071 ms 对 0.883 ms；模型中 int64 每次调用约 39 ms）。其余不变的情况下，每步时间从 1,952 ms 降至 756 ms，输出逐位相同。该改动适用于所有平台；其在 CUDA 上的效果未测量。
5. K 维分块（#8506）：在 MUSA 上，两个 MoE GEMM 的 BLOCK_SIZE_K 从 32 改为 64，CUDA 不变。每步时间从 754.5 ms 降至 733.9 ms，输出逐位相同。
6. 路由构建（#8592）：路由通过排序加 searchsorted 确定性地构建，共享专家的 SwiGLU7 按投影分别计算，总计约 −12 ms。

4.3 mHC：将小型 GEMM 重写为逐元素收缩

每个 mHC 分支都有一组小型 FP32 运算：[12288×24] 投影、4×4 流混合，以及 20 次 Sinkhorn 迭代。在 MUSA 上，逐元素收缩比小 K 的 bmm 更适合这些运算，而流主序布局还能省去一次操作数拷贝。

• 9 月，我们首先集成了融合的 mHC 混合与 Sinkhorn kernel。每步时间从 2,040 ms 降至 1,795 ms（−12%）。
• #8510 的第一个提交在 MUSA 上用 FP32 逐元素或按流收缩替代小 K 的 bmm，用于流混合和投影。每步时间从 984.7 ms 降至 733.6 ms（−25.5%）。求和顺序改变，因此输出不再逐位相同；针对 FP64 CPU 参考的数值测试覆盖了正确性。
• #8510 的第二个提交让 mHC norm 直接以“流主序”布局写出输出。这每步节省 80 次拷贝，约 −24 ms，输出逐位相同。
• #8594 将 20 次 Sinkhorn 迭代写成一个 Triton kernel，并通过 triton_op + wrap_triton 接入编译路径：−6.6 ms。

4.4 注意力

• 9 月，注意力从分块的 Torch 实现迁移到 MATE FA3。在同一代码树上，SP4×CFG2 下的 100 步端到端时间从 2,271 s 降至 1,889 s（−16.8%）。
• #8498 调用 FA3 varlen 获取 LSE，然后以 FP32 将学习到的 attention sink 作为 LSE 修正加入。
• #8593 缓存每请求的 RoPE、模态分组和 sink-LSE 元数据，总计约 −6 ms。
• 注意力 kernel 本身已接近峰值：在 profile 中，FA3 达到 8 卡配置 500 TF/s BF16 峰值的 82–87%。

4.5 按头分片的专家并行

MAGI-2 的 MoE 按头路由，因此专家并行也按头分片：使用 EP4 时，每块 GPU 拥有 3 个头的全部专家。Sand AI 的博客指出：“跨设备通信不会随每个头的 Top-k 增长。”EP8 需要将 12 个头填充到 16，这会留下 2 块 GPU 没有真实头。

• #8511 将 EP 并行度与 SP 并行度解耦。在 SP8×CFG1 下，EP=SP=8 时每步耗时 864.6 ms，EP4 时为 717.5 ms，而在整个 SP 组之间交换 head-EP token 还能再节省 12.9 ms。在最新的代码树上，并行度是实测而非假设的：显式标志关闭时为 697.8 ms（head 分片仍跟随 SP 组），EP4 时为 565.0 ms，EP8 时为 633.4 ms，输出全部逐字节一致（第 5.3 节）。
• #8595 移除了 Ulysses 发送和接收周围的缓冲区拷贝：−10.6 ms。
• 在同样的 8 块 GPU 上，SP8×CFG1+EP4 比 SP4×CFG2 更快。四路 Ulysses 只用到 GPU 之间 7 条链路中的 3 条（约 15 ms）。两个 CFG 分支的 token 数不同（15,015 和 14,601），因此把它们放在不同的 GPU 上运行会让彼此互相等待（约 10 ms）。

4.6 编译

• 上游的区域编译（#7174）会编译重复的 Transformer 块；在 MUSA 上启用它需要做一些适配。在 10 月初的代码树上，与 eager 模式相比，区域编译将每步耗时降低了约 16%。
• vLLM-Omni 默认使用 dynamic=True，因此每个维度都是符号化的，attention 输出门控 kernel 对每个元素做 64 位索引运算。通过 torch._check（#8627）固定非 token 维度后，该 kernel 从 1,429 µs（外加 212 µs 的拷贝）降至 55 µs，并去掉了这次拷贝。#8595 改变了这一布局，因此包含它的代码树使用该 PR 自带的修复（2,628 → 39 µs）。未报告端到端的动态编译差异：基线早于 #8594，而后者的动态输出是无效的。
• MUSA 上的动态编译需要 #8594 才能正确工作。 MUSA Triton 工具链（Triton 3.2.0、mcc 5.2.0）会错误编译动态形状 _moe_output 区域中的一个融合逐点 kernel：第一个之后的每个 token 都是错的，视频解码出来是噪声。静态形状区域可以正确编译，而 #8594 将最后一次 Sinkhorn 除法保留在该 kernel 之外，从而使动态输出与静态输出逐位一致。这里的静态结果不依赖它；#8915 使 MUSA transformer 区域默认以静态形状编译。

5. 实测结果

5.1 PR 代码树的端到端结果

这些结果是在第 2.1 节的条件下测得的（ms/步；多次运行分别列出）：

“最后一轮优化之前”指的是不含 #8592–#8595、不含 #8510 的第二个提交、不含 #8511 的第四个提交的代码树。其动态编译列使用来自 #8595 的等效修复，并针对该布局做了适配（第 4.6 节）。该代码树上最快的配置是 SP8×CFG1+EP4，静态约 668 ms，动态约 680 ms。该列之所以有意义，只是因为代码树包含 #8594；更早的代码树没有动态数值，其动态输出是无效的（第 4.6 节）。

• 这些数字需要四个仍未合并的 PR（#8511、#8593–#8595）。仅包含已合并 PR 的 main 尚未在 MUSA 上完整测量。
• 对于读者今天会运行的代码树（main 加 #8595），静态编译测得每步 670.2 ms，输出完全一致。上面的 665.7 / 669.9 运行早于该修复，保留用于对比。
• 10 月 19 日，在应用了本文所有 PR 的情况下，相同工作负载测得每步 565 ms（第 5.3 节）。那是在不同会话中测量的另一棵树，因此不声称与上述各行存在差值。
• 仅在同一行和同一编译模式内进行比较。不同的并行配置在通信量和负载均衡上有所不同。

5.2 拆分 W13 累加器（#8855）

图 4：SP8×CFG1+EP4 的每步时间分解，来自稍早的代码树（每步总计约 778 ms）。#8510（已合并）移除了 bmm 输入拷贝，#8594（审查中）将每次调用启动的 42 个 Sinkhorn 内核融合为一个。#8855（第 5.2 节）拆分 W13 累加器，使该内核耗时减少约 40%。来源：作者使用 torch profiler 进行的测量。

W13 是融合的 gate/up GEMM，是一步中最大的 MoE GEMM——在图 4 中约占三分之一。在 MUSA 启动配置（BM 128、BN 128、BK 64、16 warps、1 stage）下，已发布的内核将交错的 gate 和 up 行累加在一个 FP32 共享内存 tile 中，SwiGLU7 尾声从该 tile 读取，这限制了能在 MP 上常驻的程序数量。#8855 改为将 gate 和 up 保留在两个半宽累加器 tile 中，因此尾声在寄存器上运行且无溢出；路由、掩码和填充块处理保持不变。仅该内核在全部 13 次捕获调用中就快了 38.8–40.3%，从 50.2 → 82.9 TF/s，而一个每步约值 260 ms 的内核减少 40% 约为 100 ms——与端到端测得的 99.5 ms 规模相同。

在 8× MTT S5000 上以 SP8×CFG1+EP4、272p、100 步、静态编译、在同一会话内交替运行 control 和 candidate 的端到端结果：

两侧各启动两次，且全部四次运行共享同一个输出哈希，因此 14.5% 的减少是逐位相同的。此 A/B 是单一会话，因此其绝对数字仅在这些行内比较：此处的 control 为 690.8 ms，而第 5.1 节在不同代码树上测得 665.7 ms，两者不互相换算。该代码树早于 #8592–#8595，后者的收益是在更早的代码树上增量测得的；它们在此处是否叠加尚未作为一对重新测量，不过携带全部这些改动的代码树在第 5.3 节中进行了测量。

该片段使用基准的提示词和种子，在第 6 节描述的示例设置下生成，以展示画面和声音。它不是计时运行的输出。

5.3 最新的代码树：EP 度数的价值

10 月 19 日，整套 PR——#8511、#8592–#8595、#8855、#8915 的 refiner 分支以及每步 instrumentation——在第 2.1 节的条件下，在一棵此前从未执行过的代码树上进行了测量。refiner 分支在代码树中，但 272p 不会执行它。按第 2.1 节定义的每步，272p、100 步、静态编译的时间为：

关闭 --enable-expert-parallel 后，MoE 头分片并不会关闭：当 TP 大于 1 时，get_magi2_ep_group() 回退到 TP 组；当 TP 等于 1 时，回退到 Ulysses 组，因此在 SP8 下，这一实验组仍将头库分成八份。这是头并行与 Ulysses 重叠的传统布局，而不是未分片的对照组。三个实验组仅在 --enable-expert-parallel 和 --expert-parallel-size 上不同；运行横幅记录了最终解析出的配置（ulysses_degree、cfg_parallel_size、enable_expert_parallel、expert_parallel_size）。

EP4 配置的三次运行——第一次在一棵从未运行过的代码树上，紧接着重复一次，以及本次对比中的一次——分别得到 565.0、566.3 和 565.0 ms，因此运行间波动为 1.3 ms（0.2%）。三种并行配置产生相同的音频（3188e16c3487fc50…，sha256）和相同的图像（61e6da48517e0d0f…），因此它们之间的差异并非数值取舍。

本次扫描在最新代码树上为第 4.5 节的 EP4 与 EP8 对比增加了关闭该标志的对照组，并区分了两种效应。在相同的八路分片下，显式头 EP 路径比传统 SP 耦合布局更快：633.4 对 697.8 ms。而在同一路径下，分片越少越好：EP4 比 EP8 快 68.4 ms。在本次扫描之前，EP4 恰好是每个最佳数字所使用的配置——864.6 对 717.5 ms 这一对是八分片对四分片，因此这两个分片数都是已知的，而关闭标志的实验组和固定 SP 度下的 EP8 都尚未测量。在此规模下，EP8 并非更快的配置；#8915 采用它是为了内存，而非速度。

有两点需要注意。无 EP 和 EP8 配置各只测量了一次，EP4 测量了三次，因此只有 EP4 有运行间波动的证据。而且这并非针对第 5.1 节的配对 A/B 对比：它是在更晚的会话中测量的更晚的代码树，不声称 565 ms 与 670 ms 之间存在差值。其间的变更是第 4 节的各个 PR，每个 PR 都带有自己的配对测量。

6. 生成的样本与质量边界

以下所有视频均使用 8× MTT S5000 上的基准设置生成：10 秒、125 帧、12.5 fps、100 步、种子 42、默认 CFG（视频 5.0 / 音频 7.0）以及默认负面提示词。模型将音频与画面一起生成，没有后期配音。每个提示词只生成一次，未进行挑选。有一项设置与第 5 节的计时代码树不同：样本在带有 #8594 的代码树上使用动态编译，因此输出有效（第 4.6 节），并且是在单独的会话中生成的。它们仅作为视觉和音频样本。

vLLM-Omni 实现了 MAGI-2 的 Preview 阶段；Sand AI 用于 1080p 输出的 refiner 正在以 #8915 进行审查，因此这些帧无法直接与 Sand AI 网站上的 1080p 样本进行比较。

6.1 Sand AI 官方提示词

提示词来自 Sand AI 仓库中的 assets/demo_samples.json：两个长提示词，每个分别作为 T2VA 和 I2VA 各运行一次，外加一个带对话的 T2VA 提示词。与 Sand AI 的入口脚本一样，每个提示词文件在去除首尾空白后原样传给模型。每个提示词约 9,000–11,000 个字符，分别在独立章节中描述画面、镜头和声音设计。

T2VA 使用 SP8×CFG1+EP4。I2VA 使用 SP4×CFG2：用于这些样本的离线图生视频脚本无法在该树上单独设置 EP 程度，因此 MoE 头以 EP=SP=4 分片，每块 GPU 同样拥有 3 个头。

每个片段为 10 秒。I2VA 运行以同一提示词附带的静帧为条件（动画为 sample_000.jpeg，黏土动画为 sample_001.jpeg）；咖啡店提示词未附带静帧，因此没有 I2VA 运行。T2VA 运行仅使用提示词。

6.2 540p

SP8×CFG1+EP4 有足够内存运行 540p。在生成这些样本的会话中，动态编译下每步耗时约 3.9 秒（同一提示词在 272p 下约 0.9 秒）。540p 尚未在第 5 节的计时树上进行基准测试；之后的一次 100 步静态运行测得每步约 3.6 秒，因此两个数字仅供参考：它们在树和编译模式上均不同，彼此不可比较。

6.3 质量检查

7. 上游状态与兼容性边界

7.1 上游合并状态

状态截至 2026 年 10 月 28 日。在 torchada 侧，#120、#121、#124 和 #125 已合并。vLLM-Omni 中的 MUSA 进展在 RFC #2347 中跟踪，涵盖平台集成、来自 MATE 的 FA3、适配模型、CI 计划以及 MAGI-2 PR。

7.2 兼容性边界

8. 结论与未来工作

借助 torchada，MUSA 上的 MAGI-2 Preview 复用了 vLLM-Omni 的现有实现。CUDA 风格的 PyTorch 代码、Inductor 编译路径和分布式通信均直接复用。Triton 内核仅在每平台启动配置和路径选择上有所不同，平台特定处理集中在少数几处。在此基础上，沿执行路径的优化将 9 月的 bring-up 从每步 15.6 s 降至 770 ms；应用本文所有 PR 后，最新测量为每步 565 ms（第 5.3 节），最后一轮在 10 月代码树上实现了 14.5% 的步级降低，且与其自身对照逐位一致（第 5.2 节）。第 4 节的 27.6× 描述的是整体工作，而非两个代码状态之间的对比。这些更改已作为 PR 上游提交；部分已合并，其余正在审查中。

以下前三项来自图 4；动态编译在隔离测试中测量，其余为平台和功能覆盖工作：

• MoE W13 GEMM：融合的 SwiGLU7 尾声使用了大量共享内存，限制了每个 MP（MUSA 的多处理器，对应 CUDA 中的 SM）上同时驻留的块数。#8855 的做法是：gate 和 up 在单独的 tile 中累加，尾声在寄存器上运行，内核速度提升 38.8–40.3%（第 5.2 节）。剩下的是 GEMM 本身，仍受延迟限制，约为 83 TF/s。根据 profile 评估了三个逐位精确的选项，但此处未采用：128 宽的 K tile（2–3 ms/步）、排序中间结果（0–7 ms/步），以及通过 MATE 专门构建的分组内核（0–140 ms/步，工作量最大）。
• 通信-计算重叠：使用 EP4 时，每步 60–90 ms 的集合通信尚未与计算重叠。
• 步间空闲时间：FlowUniPC 在步间的主机计算约需 19 ms，前向传播内部的启动间隙约需 14 ms。
• 动态编译：其他内核仍执行 int64 尺寸运算；在隔离测试中，将尺寸以 int32 传入可挽回约 85% 的该开销。
• MUSA 上的图重放：MUSAGraph 已能正确捕获 MAGI-2 的区域图。9 月时，启用图重放的每步时间为 3,601 ms（未启用为 972 ms）；图重放调优留待后续工作。
• CI 与覆盖率：自托管的 8× MTT S5000 runner 和 MUSA CI 流水线是 RFC #2347 计划的一部分。540p 基准测试以及仅用已合并 PR 测量 main 也是后续工作。
• 1080p 精修器：Sand AI 的默认 1080p 输出有两个阶段。Preview DiT 首先生成 540p，然后一个独立的 6.7B 稠密精修器（30 层、块窗口局部注意力、5 步）将视频提升至 1920×1088、25 fps，并复用 Preview 阶段的音频。vLLM-Omni 目前已实现 Preview 阶段，精修器阶段正在 #8915 中评审。在 MUSA 上，精修器需要块窗口局部注意力的实现，以及同时容纳 Preview 和精修器权重的内存方案。本文优化的是 Preview DiT；测量 540p Preview 阶段是迈向 1080p 的第一步。

几条同样适用于下一个模型的经验：

1. 先找主机同步，再看内核。逐专家 .item()、逐层 all_gather 加回读，以及有限性检查都很容易修复。.item() 和有限性检查的修复是收益最大的几项之一。
2. 针对分片后的局部形状进行优化。在 EP4 下，每块 GPU 只持有 3 个头的专家，GEMM 的 K 仅为 256。分块和融合的选择必须在该形状下测量。
3. 在每个性能数字旁放一个正确性检查，并确认该检查确实生效。
4. 在平台对象上而非张量的设备类型上对仅 CUDA 内核做门控。torchada 有意映射张量语义，使 CUDA 代码照常运行。
5. 把可以放在兼容层里的修复放进兼容层。isfinite 修复位于 torchada 中，因此其他基于 torchada 的框架也能受益。

致谢

我们感谢 Sand AI 开源 MAGI-2 Preview，并在其仓库中提供官方示例提示词。

MAGI-2 的 MUSA 支持与优化、本文中的 vLLM-Omni PR 以及 torchada 修复均由 @yeahdongcn（R0CKSTAR）完成。

我们感谢 @hsliuustc0106 在 vLLM-Omni 中实现原生 MAGI-2 Preview 支持，并评审和合并本文中的 MAGI-2 PR。我们感谢 @cr-gao 实现区域编译，@IDEA-V 在 #7175 中融合 MAGI-2 的多头 MoE 路由构建，以及 @Dong1017 在 H200 上独立验证 MAGI-2 内核和 MUSA 相关 PR（#7261 也在 Ascend 上验证）。我们还感谢 @RuixiangMa、@SamitHuang、@wtomin、@lishunyang12、@alex-jw-brooks 和 @daii-0818 的评审与讨论；以及 @gcanlin 和 @lishunyang12 对 MUSA RFC #2347 的支持。

我们感谢摩尔线程 MATE、torch_musa 和 torchada 团队提供 FlashAttention-3 与分组 GEMM、PyTorch 后端以及 torchada CUDA 兼容层。

附录 A. 复现

A.1 计时边界

A.2 容器镜像

整套栈以单个镜像形式发布，因此读者无需自行组装源码树：

该镜像固定的代码树比表格中的更新：5.1–5.3 节测量的是它们所引用的代码树，而该镜像是同一方法最后一次运行所用的代码树。其中没有任何数字与那些表格进行比较。权重不在镜像中，且此基础镜像上无需为 vLLM 提供兼容层。

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

MUSA_VISIBLE_DEVICES 用于选择容器内的显卡，与任何 MUSA 容器中的做法相同。

A.3 计时运行

gate 是入口点，它直接接受并行和形状标志，包括 --compile-static。容器内分两步——先插桩，再运行：

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'

gate 按 2.1 节定义的指标进行聚合——每步取最慢的 rank，丢弃第 0 步和第 1 步，其余取平均——在 --metrics-dir 下为每个 rank 和步骤写入一条 JSONL 记录，并在 --save-outputs 下写入一个包含各数组 sha256 的 manifest.json。在 8× MTT S5000 上，此工作负载测得每步 566.8 ms，重复测量为 565.7 ms。

说明：

• 插桩以未应用状态发布。 早期的每步计时来自 Magi2PreviewSampler.denoise_step，而采样器重构移除了该方法，因此此栈的每步序列来自一个由 instrument-steps.py 在容器内应用到 sampler_magi2.py 的补丁。没有它，gate 会报告阶段总计和 no step records。它就地编辑已安装的源码，是幂等的，并在计时区域之外进行同步。
• 编译。 #8915 使 MUSA transformer 区域即使在 diffusion_compile_dynamic 为 true 时也能以静态形状编译（4.6 节），因此 --compile-static 明确说明了这一点，并且在 MUSA 上不做任何更改。
• 专家并行。 在 8× 80 GB 上，--expert-parallel-size 必须至少为 4：为 2 时，加载会内存不足（最满的卡上请求 960 MiB，而空闲 723 MiB）。MAGI2_EP_SIZE 在此栈中无效；该环境变量不会创建任何组。关闭该标志并不会关闭头分片——当 TP=1 时，MoE 头轴回退到 Ulysses 组（5.3 节）。
• 此处需要 --no-pin-cpu-memory。 使用默认值时，固定内存暂存器会在约 60 GiB 权重已驻留后引发看似 MUSA OOM 的错误。
• 哈希是更强的检查。 5.3 节的并行配置产生相同的音频和图像数组，因此复现哈希的运行复现的是计算，而不仅仅是计时。

A.4 其他分辨率和样本

# 540p (896×512)
docker exec magi2 bash -lc 'cd /opt/vllm-omni-src && python /opt/magi2/magi2_e2e_gate.py \
  --ulysses-degree 8 --cfg-parallel-size 1 --enable-expert-parallel --expert-parallel-size 4 \
  --resolution 540p --height 512 --width 896 \
  --seconds 10 --fps 12.5 --num-frames 125 --num-inference-steps 100 \
  --no-pin-cpu-memory --save-outputs /work/540p'
 
# 1080p, with the refiner
docker exec magi2 bash -lc 'cd /opt/vllm-omni-src && python /opt/magi2/magi2_e2e_gate.py \
  --ulysses-degree 8 --cfg-parallel-size 1 --enable-expert-parallel --expert-parallel-size 4 \
  --resolution 1080p --height 1088 --width 1920 \
  --seconds 10 --fps 25 --num-frames 249 --refiner-num-inference-steps 5 \
  --vae-patch-parallel-size 8 --no-pin-cpu-memory --save-outputs /work/1080p'

在 1080p 下，--vae-patch-parallel-size 8 不是可选项：使用默认值 1 时，运行会在解码期间无回溯地终止（Stage-0 inline diffusion engine is dead），且该单卡故障的原因尚未确定。分辨率检查仅接受带有精确 1920x1088 配对的 1080p。样本提示和图像经过同一个 gate：T2VA 使用 --prompt 或 --prompt-file，I2VA 使用 --image 和 --prompt。提示和图像文件必须在容器内可见。

A.5 位级一致性检查

部分 Inductor 选择（reduction 配置、GEMM padding）是在编译时通过基准测试做出的，因此同一份代码在全新的编译缓存下可能产生不同的输出哈希。TORCHINDUCTOR_DETERMINISTIC=1 在这里没有影响：MAGI-2 流水线在构造期间调用 torch.use_deterministic_algorithms(False)，在首次编译前就将其关闭。本文中的测量使用以下设置：

export TORCHINDUCTOR_FORCE_FILTER_REDUCTION_CONFIGS=1  # keep only one config per reduction
export TORCHINDUCTOR_CACHE_DIR=$(mktemp -d)            # use a fresh cache for every run
export TORCHINDUCTOR_DETERMINISTIC=1                   # set, but has no effect, as explained above

在这些设置下，GEMM 是否被 padding（pad_mm）仍由基准测试决定。使用这些设置时，相同的哈希证明输出相同；不同的哈希并不能证明数值发生了变化，需要进一步检查。仅保留一种 reduction 配置本身没有可测量的性能损失（697.8 ms，对照组为 700–703 ms）。

参考文献

• vLLM-Omni
• MAGI-2 Preview 博客、模型卡、SandAI-org/MAGI-2-preview
• mHC 论文（arXiv 2512.24880）
• vLLM-Omni 原生 MAGI-2 支持 #5918、区域编译 #7174、融合多头 MoE 路由构建 #7175
• 原生 MAGI-2 Preview 架构 RFC #5949
• MUSA 支持路线图 RFC #2347
• torchada
• torch_musa
• MATE（MUSA AI Tensor Engine）
• vllm-musa