MAGI-2 Preview on MUSA achieves 565 ms per step with vLLM‑Omni

TL;DR

MAGI-2 Preview, Sand AI’s 114 B audio‑video generation model, now runs on Moore Threads MUSA GPUs via the torchada compatibility layer at 565 ms per denoising step (8 × MTT S5000, 272p, static compilation), a 27.6× improvement over the first working version (15.6 s per step). The speed‑up comes from a sequence of targeted PRs that remove host synchronizations, fuse MoE kernels, rewrite mHC residuals, and optimize expert‑parallel sharding.

1. What makes MAGI-2 Preview a unique inference workload

Single‑stream DiT backbone packs text, video‑latent, and audio‑latent tokens into one self‑attention sequence; only the first and last four layers have modality‑specific weights.
No adaLN – timestep is encoded directly in the first 64 channels of each token, allowing independent video and audio timesteps.
Multi‑head MoE (layers 2‑37) splits each 3072‑dim token into twelve 256‑dim heads; each head routes to its top‑6 of 256 experts (SwiGLU7), plus one shared expert and three per‑modality shared experts. Routed experts account for ~95 % of parameters.
mHC hyper‑connections provide a four‑stream residual (total width 12288) with a doubly‑stochastic 4 × 4 mixing matrix computed via 20 Sinkhorn iterations.
Sampling runs 100 undistilled steps, each with two CFG branches, denoising the full 10‑second video in one pass.

Metric	Value
Total parameters	114 B (≈6 B active per token)
Backbone	40‑layer DiT, hidden 3072, 24 heads × 128
MoE	12 heads × 256 experts, top‑6 per head
Residual	mHC, 4 streams, 20 Sinkhorn iterations
Output	10 s, 125 frames, 12.5 fps, 44.1 kHz stereo
Supported resolutions	272p (448 × 256) and 540p (896 × 512)

Compared to other open video models, MAGI-2 Preview is the only one that combines per‑token MoE routing, mHC residuals, and joint audio‑video generation in a single self‑attention stream.

2. Running a CUDA‑first model on MUSA with torchada

torchada intercepts CUDA‑style imports (torch.cuda, device strings, NCCL, CUDA Graph) and maps them to the MUSA equivalents (torch.musa, MCCL, MUSAGraph). The mapping is deliberately truthful (torch.cuda.is_available() stays False) so that platform checks must use torch.version.musa.
Code reuse: eager PyTorch, torch.compile/Inductor, and the upstream Triton kernels (fused BF16 MoE, SwiGLU7, mHC) run unchanged on MUSA. Only launch configurations and a few platform‑specific paths differ.
Key adaptations:
- torch.isfinite replaced with an abs()<inf expression for MUSA (torchada #120).
- Int32 routing counts (scatter_add_) replace int64 on MUSA, yielding a ~39 ms per‑call gain (PR #8496).
- GEMM operand layout is reordered to stream‑major, eliminating ~80 copies per step (PR #8510).
- Triton‑MUSA lacks BF16 atomics, so the fused BF16 MoE path is always selected.

"torchada makes is_cuda true for MUSA tensors, so CUDA‑style code runs as usual." – torchada documentation.

3. Optimization timeline (per‑step latency)

Stage	Per‑step time (ms)	Main change
Initial bring‑up (Sept)	15,600	Baseline on 8 × MTT S5000
Host‑sync removal	8,800	`.item()` → `.tolist()`, reduced all‑gather metadata
MoE fused GEMM (MATE)	2,390	Grouped BF16 GEMM replaces per‑expert Python loop
Fused W13 + SwiGLU7 Triton	954	Single kernel for gate/up + activation
Int32 routing counts	756	PR #8496
K‑dim tiling (BLOCK_SIZE_K=64)	734	PR #8506
mHC elementwise rewrite	734 → 734 (baseline) → 733.6	PR #8510 (FP32 contractions)
Sinkhorn single‑kernel	–6.6	PR #8594
W13 accumulator split	590 – 591	PR #8855 (38–40 % faster kernel)
Full PR stack (Oct 19)	565	All above PRs applied, static compilation

The 27.6× overall speed‑up (15.6 s → 565 ms) is measured on the same hardware and workload (272p, 100 steps, T2VA prompt, SP8×CFG1+EP4). Each bullet is evaluated against its own baseline; the numbers are not cumulative across rows.

4. Detailed technical gains

4.1 Host‑side synchronizations

Per‑expert .item(): replaced by a bulk .tolist(), saving ~3.4 s per step.
All‑gather token‑count metadata: reduced from 39 to 3 calls per step, saving ~61 ms.
torch.isfinite: reimplemented as abs()<inf on MUSA, cutting ~1.8 s per step.

4.2 MoE kernel fusion

MATE grouped GEMM collapsed 12 × 256‑dim expert GEMMs into a single BF16 GEMM (21.9 ms → 0.69 ms).
Fused W13 + SwiGLU7 Triton kernel combined gate, up‑projection, and activation, reducing the MoE bottleneck from ~1.5 s to <1 s.
Weight layout pre‑interleaving removed a per‑call transpose, shaving ~7.5 % off the step time.
Int32 routing histogram cut ~1.2 s from the routing phase.
K‑dim tiling (BLOCK_SIZE_K=64) gave a modest ~2 % gain.
Deterministic route construction (PR #8592) added ~12 ms.

4.3 mHC residual rewrite

Replaced small‑K bmm with FP32 elementwise contractions and a stream‑major layout, cutting ~250 ms per step (≈25 %).
Consolidated the 20 Sinkhorn iterations into a single Triton kernel (‑6.6 ms).

4.4 Attention

Switched from Torch implementation to MATE FlashAttention‑3 (FA3) with learned LSE sinks, yielding a 16.8 % reduction on the SP4×CFG2 configuration.
Cached RoPE, modality grouping, and sink‑LSE metadata (‑6 ms).

4.5 Expert‑parallel sharding

Head‑sharded EP (EP4) decoupled from sequence parallelism, delivering the fastest configuration (565 ms). EP8 was slower (633 ms) because more shards increase communication overhead.
Removing buffer copies around Ulysses send/receive saved an additional 10.6 ms (PR #8595).

4.6 Compilation

Regional torch.compile reduced per‑step time by ~16 % compared with eager execution.
Pinning non‑token dimensions under dynamic compilation collapsed a 1.4 ms kernel to 55 µs (PR #8627).
Dynamic compilation on MUSA required the Sinkhorn fix (#8594) to produce correct outputs.

5. Measured performance results

5.1 PR‑tree end‑to‑end timings (static vs. dynamic)

Parallel config	Before final PR round (static)	Static compilation (ms/step)	Dynamic compilation (ms/step)
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	–

The 565 ms figure (Section 4) comes from the newest tree measured on 19 Oct 2026; it is not directly comparable to the rows above because it uses a different code base and session.

5.2 A/B test of the W13 accumulator split (PR #8855)

Tree	Static compilation (ms/step)	Output hash (audio / video)
Control (no split)	690.8 / 691.1	`32f1b4fca4` / `44b9747e9e`
+ #8855 (split)	591.3 / 590.4	identical

The 14.5 % reduction is bit‑identical; the same SHA‑256 hashes confirm numerical equivalence.

5.3 Expert‑parallel degree sweep on the latest tree

Config	Mean per‑step time (ms)	Relative to EP4
SP8×CFG1, EP off (legacy)	697.8	+23.5 %
SP8×CFG1+EP4	565.0	baseline
SP8×CFG1+EP8	633.4	+12.1 %

All three configurations produce identical audio and video hashes, confirming that the speed differences are purely communication/computation trade‑offs.

6. Sample outputs and quality notes

Benchmark prompt (golden retriever, 272p, SP8×CFG1+EP4) generated a 10‑second video in 565 ms per step; the clip is available in the blog assets.
Official Sand AI prompts (coffee shop, anime, clay animation) were run in both T2VA and I2VA modes. All samples match the textual description, contain non‑zero audio, and the coffee‑shop clip includes audible English dialogue.
540p preview runs at ~3.6 s per step (static) and ~3.9 s (dynamic); these numbers are not directly comparable to the 272p timing because they use a different code tree and compilation mode.
No quantitative multi‑seed quality evaluation has been performed yet; visual and auditory fidelity were verified manually.

7. Upstream status and compatibility

7.1 PR landscape (as of 28 Oct 2026)

PR	Feature	Status	Measured per‑step impact
#7206	Fused BF16 MoE (W13)	Merged	–
#7243	Configurable Dynamo recompile limit	Merged	–
#7252	Fused SwiGLU7	Merged	–
#7261	Fused mHC post‑residual	Merged	–
#8496	int32 routing counts	Merged	1,952 → 756 ms
#8497	Single copy of W13	Merged	–
#8498	MUSA support (FA3 varlen + sink)	Merged	–
#8506	BLOCK_SIZE_K=64 for MoE GEMMs on MUSA	Merged	754.5 → 733.9 ms
#8510	mHC FP32 contractions & stream‑major layout	Merged	984.7 → 733.6 ms
#8511	Head‑sharded EP independent of SP	Open	864.6 → 717.5 ms
#8592	Bit‑exact route build	Merged	–12 ms
#8593	RoPE & sink‑LSE metadata cache	Open	–6 ms
#8594	Single‑kernel Sinkhorn	Open	–6.6 ms
#8595	Remove Ulysses buffer copies	Open	–10.6 ms
#8627	Pin attention‑head layout (dynamic)	Merged	1,429 µs → 55 µs
#8855	Split W13 accumulator (MUSA)	Open	690.8 → 590.4 ms (static)
#8915	1080p refiner (block‑window attention)	Draft	–

All listed PRs are either merged or under review; the performance numbers in this post reflect the state of the code at the time each PR was measured.

7.2 Compatibility boundaries

Fully benchmarked: MUSA + SP8×CFG1+EP4 at 272p (Sections 5.1‑5.3).
Partial support: 540p preview (samples only), I2VA mode (samples only), 1080p refiner (draft implementation, not timed).
Unvalidated: Layerwise offload, quantization, Cache‑DiT, and other memory‑saving techniques.

8. Conclusions and future work

Speed‑up summary: From a 15.6 s baseline to 565 ms per step (27.6×) by systematically eliminating host synchronizations, fusing MoE kernels, rewriting mHC residuals, and improving expert‑parallel sharding.
Remaining bottlenecks:
1. MoE W13 GEMM still limited by memory bandwidth; further tiling (128‑wide K) or a custom grouped kernel could shave another ~2–3 ms.
2. Communication‑compute overlap: EP4 incurs 60–90 ms of collective communication per step that is not overlapped with compute.
3. Idle host time: FlowUniPC host logic (~~19 ms) and intra‑kernel launch gaps (~~14 ms) remain.
4. Dynamic compilation overhead: int64 size arithmetic in several kernels; converting to int32 recovers ~85 % of the lost performance.
Next steps:
- Enable graph replay on MUSA (currently slower than regional compilation).
- Complete the 1080p refiner integration (block‑window attention, memory planning).
- Extend CI to cover 540p benchmarking and the main branch with only merged PRs.
- Explore quantization and cache‑DiT techniques for further latency reduction.

9. Acknowledgments

The authors thank Sand AI for open‑sourcing MAGI‑2 Preview and providing the official prompts, the Moore Threads MUSA, torchada, and MATE teams for the GPU stack, and the many contributors listed in the original blog post for implementing and reviewing the PRs.

10. Reproducibility

Container image: registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview (digest sha256:2785b4626d1da1e0765a4cb515b9afc5349a22c6417e14c81091add0c1a9b902).

Command (run on 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'

Environment variables for deterministic Inductor behavior are set inside the container as described in Appendix A.5 of the source post.

All numbers, tables, and quotes are taken directly from the vLLM blog post dated 12 Oct 2026. No additional data were invented.

MAGI-2 Preview on MUSA achieves 565 ms per step with vLLM‑Omni