[Perf][Diffusion] MAGI-2: reuse SP MoE sequence-size metadata

[NancyFyong]

Purpose

Refs #7085 (M3). Two parts:

1. Remove the repeated MoE sequence-size exchange on the SP path. Magi2MultiHeadMoE.forward
   gathered the per-rank sequence length with a scalar torch.distributed.all_gather plus a host
   .item() before every ep_dispatch. MAGI-2 Preview runs MoE in 36 of its 40 layers, so one
   transformer forward paid that metadata round-trip 36 times per denoising step, even though the
   transformer already knows the Ulysses split for the request. This passes the existing
   request-scoped cp_split_sizes down to the MoE layer, so the production SP path performs zero
   per-layer metadata collectives. Shard ownership, routing and expert compute are unchanged.

2. Add a component profiler (benchmarks/diffusion/magi2_moe_communication.py) covering the
   areas M3 names — sequence-size exchange, head dispatch/undispatch, TP reduction, routed expert
   compute, and the routed/shared branch split — plus a probe that overlaps the shared branch with
   the dispatch collective, so the value of compute/communication overlap is measured rather than
   assumed.

What this PR claims, and what it does not. It removes 36 redundant collectives per transformer
forward, with bit-exact decoded output and unchanged peak memory. The saving is reproducible at the
layer level (0.23 ms/layer, about 8.2 ms/forward at SP4), and a 31-pair interleaved end-to-end study
resolves it at roughly +0.4% on the 272p/100-step T2VA workload, with the mean paired delta
landing exactly on the value predicted by the component profile. It is a small effect: the honest
summary is "removes provably redundant work, output-preserving, no memory cost, with a measurable but
sub-percent latency benefit", not a headline speedup.

Scope notes:

• The collective fallback is kept for direct callers that do not have the transformer's split, and
  tensor-valued cp_split_sizes still takes the legacy path.
• This is one item inside M3. Head dispatch/undispatch is now profiled but not optimized; TP
  reductions, shared-expert execution restructuring, and compute/communication overlap are
  profiled and, based on the measurements below, not worth pursuing on this hardware class.

Test Plan

vLLM Version: 0.28.0

vLLM-Omni Commit: 3237096

Item	Value
Accelerator	4x NVIDIA H20, 97,871 MiB each, NV18 bonded NVLink all pairs
Driver / CUDA	535.161.08 / 12.9
PyTorch	2.13.0+cu129
Python	3.12
Checkpoint	sand-ai/MAGI-2-preview revision 2dea51b64db47ee5b4402d36fd90829a0c58913b
Layout	Resident SP4, TP=1, BF16, no Cache-DiT (the reference-aligned baseline)
Kernel selected	FlashAttention 3
Workload	T2VA and I2VA 272p, 448x256, 125 frames, 12.5 fps, 100 steps, seed 42
Baseline / candidate	32370968 / 92a13868

# 1. MAGI-2 CPU suite
pytest tests/diffusion/models/magi2 -m "core_model and cpu" --run-level core_model

# 2. Four-rank TP4 / TP2SP2 / SP4 oracle parity, extended to assert no layout performs
#    redundant per-layer sequence-size collectives
pytest tests/diffusion/models/magi2/test_native_distributed_parity.py \
  -k test_four_rank_tp_sp_layouts_match_single_rank_oracle

# 3. Real four-GPU NCCL regression on the MoE module
pytest tests/diffusion/distributed/test_magi2_moe_metadata.py \
  -m "core_model and cuda and L4 and cards_4" --run-level core_model

# 4. Real two-GPU NCCL regression on the production transformer, uneven SP2 split
pytest tests/diffusion/distributed/test_magi2_moe_metadata.py \
  -m "core_model and cuda and L4 and cards_2" --run-level core_model

# 5. Component profile at the released SP4 layout
torchrun --standalone --nnodes=1 --nproc-per-node=4 \
  benchmarks/diffusion/magi2_moe_communication.py --tokens 7500 --iterations 30

# 6. End-to-end equality, deterministic mode, one run per arm
MAGI2_DETERMINISTIC=1 python examples/offline_inference/text_to_video/text_to_video.py \
  --model "$MAGI2_CKPT_ROOT" --model-class-name Magi2Pipeline --seed 42 \
  --prompt "A red fox walks through fresh snow while wind moves the pine branches." \
  --height 256 --width 448 --num-frames 125 --num-inference-steps 100 --fps 12.5 \
  --tensor-parallel-size 1 --ulysses-degree 4 \
  --extra-body '{"seconds":10,"resolution":"272p"}' --output t2va.mp4

# then compare decoded streams, not container bytes:
ffmpeg -v error -i t2va.mp4 -map 0:v -f rawvideo - | sha256sum
ffmpeg -v error -i t2va.mp4 -map 0:a -f s16le  - | sha256sum

# 7. End-to-end timing: 31 interleaved pairs, non-deterministic mode, each pair
#    running base then candidate back to back (same command as 6 without
#    MAGI2_DETERMINISTIC), then a paired analysis of the 31 deltas.

Test 4 runs the same ranks and weights twice — once on the production path and once forced back
onto the per-layer collective — and requires exact output equality plus one baseline collective per
MoE layer. Mutation check: reverting the propagation in Magi2TransformerLayer.forward to
sequence_split_sizes = None must turn test 4 red.

Test Result

Check	Result
tests/diffusion/models/magi2, core_model and cpu	46 passed
Four-rank TP4/TP2SP2/SP4 oracle parity (Gloo)	1 passed
Four-GPU NCCL MoE metadata regression	1 passed, max_error=0, scalar_all_gathers=0
Two-GPU NCCL SP2 transformer regression	1 passed
Mutation check	failed as required: still performs per-layer sequence-size collectives (2 scalar all_gather calls)
Pre-commit on changed files	passed (ruff, ruff-format, mypy-3.10, SPDX, forbidden imports, torch.cuda, CI marks, typos)

End-to-end numerical equality — exact

MAGI2_DETERMINISTIC=1, seed 42, full 100-step trajectory, released checkpoint, resident SP4.
Decoded streams, not container bytes, because the MP4 container is not byte-reproducible:

T2VA:

Arm	decoded video sha256	decoded audio sha256
base 32370968	62cc40426983d942	f4555c01a192f343
candidate 92a13868	62cc40426983d942	f4555c01a192f343

Both streams are identical, so the change is output-preserving on a real trajectory, not only on the
synthetic module tests. Frame count, geometry, frame rate and the 44.1 kHz stereo audio contract are
preserved by construction since the streams match bit for bit.

Note for anyone repeating this: default (non-deterministic) mode does not reproduce byte-identical
output across runs even with a fixed seed, so equality must be checked with MAGI2_DETERMINISTIC=1.

I2VA, same geometry and parallelism, first frame supplied as the conditioning image:

Arm	decoded video sha256	decoded audio sha256
base 32370968	570bf6a41d6dd9d3	1d21956548514ffb
candidate 92a13868	570bf6a41d6dd9d3	1d21956548514ffb

Both conditioning modes required by the acceptance criteria are therefore covered, and both are
bit-exact on the full 100-step trajectory. Each decoded output was verified to match the MAGI-2 output
contract before comparison: 448x256, 125 frames, 44.1 kHz stereo, 10.000 s.

End-to-end latency — small but resolved

31 interleaved pairs, each pair running base then candidate back to back so machine drift affects both
arms of a pair almost equally. 62 runs total, zero failures, engine-reported generation time
(checkpoint load excluded).

Arm	median	mean	sd	min	max
base	236.58 s	236.72 s	1.38 s	234.56 s	241.03 s
candidate	235.30 s	235.90 s	2.06 s	233.49 s	243.15 s

Paired delta (base minus candidate), n=31:

Statistic	Value
mean	+0.822 s (+0.347%)
median	+1.075 s (+0.454%)
10% trimmed mean	+0.983 s (+0.415%)
sd / se	2.489 s / 0.447 s
95% CI of the mean (t)	[-0.054, +1.698] s
95% CI of the mean (bootstrap, 200k)	[-0.073, +1.653] s
Wilcoxon signed-rank	z=2.25, two-sided p=0.024
sign test	23/31 pairs faster, two-sided p=0.011
predicted by the SP4 component profile	+0.822 s — ratio 1.00x

How to read this honestly: the two mean-based confidence intervals marginally include zero, because the
delta distribution carries two symmetric Tukey outliers (-6.65 s and +5.70 s) that inflate the
variance. The rank-based tests, which are the appropriate ones under those outliers, are significant
(p=0.011 and p=0.024), the median and trimmed mean are both larger than the mean, and the mean lands
exactly on the independently predicted 8.2 ms/forward x 100 steps. The consistent reading is a real
effect of roughly 0.35-0.45%, bounded above by about 0.7% at 95%.

Peak worker memory was 66,830.00 MiB (65.26 GiB) in all 62 runs of both arms, so there is no memory
regression and no memory improvement.

Transparency notes:

• A first A/B that ran all base runs before all candidate runs produced a 4.01 s median delta, 4.9x the
  predicted effect. That was an ordering artifact; it was discarded and re-run interleaved. It is
  recorded here so those numbers are not later quoted as a speedup.
• A 3-pair interleaved run could not separate the effect from the noise floor. The 31-pair study above
  supersedes it. The required sample size was estimated from that pilot's paired sd before running.

Component profile at the released SP4 layout

Released Preview geometry (40 layers, hidden 3072, 12 MoE heads, head dim 256, 256 experts, top-k 6,
expert/shared/modality intermediate 1280), 7,500 local tokens per rank, BF16, 5 warmup + 30 measured
iterations, rank-max median. Min/max stayed within about 1% of the median for every row.

Component	SP4 ms/layer	share of layer
whole layer	22.862	100.0%
local expert compute	16.209	70.9%
shared-expert branch	3.521	15.4%
split_linear	1.049	4.6%
merge_linear	1.047	4.6%
ep_undispatch	0.292	1.3%
ep_dispatch	0.271	1.2%
sequence-size exchange (removed here)	0.228	1.0%

Secondary layouts, same workload: SP2 removes 0.180 ms/layer; TP2 removes about 0 because the TP
path returns before the exchange, and its merge_linear reduction costs 0.202 ms/layer.

Two conclusions:

• The exchange removed here is 29% of all MoE-layer communication at SP4 (0.228 of 0.792 ms), and
  it grows with world size (0.180 ms at SP2 to 0.228 ms at SP4), for a change that adds no new state
  and no new collective.
• Communication is not the MAGI-2 MoE bottleneck: all SP token collectives together are 2.5% of
  the layer, while local expert compute is 70.9%. That points at M2, not further M3 communication work.

Negative result: shared/routed overlap is not worth it

The shared-expert branch depends only on normalized, not on the routed result, so it can legally run
while the dispatch collective is in flight. Measured with an async all_to_all_single and the shared
branch in the gap:

Variant	SP4 ms/layer	SP2 ms/layer
branches serialized (today)	20.199	20.356
branches overlapped (probe)	20.192	20.353

The difference is +0.03% at SP4 and +0.01% at SP2, inside run-to-run spread, and it repeated across
runs. The dispatch is only 0.27 ms and the two branches contend for the same SMs, so there is nothing
meaningful to hide. Recommendation: do not restructure the layer for overlap on this hardware
class. Reported so the next contributor does not repeat the experiment.

Compatibility

Combination	Status
Resident SP4, BF16, no Cache-DiT	Validated end to end on the released checkpoint
SP2	Covered by the two-GPU GPU regression
TP4, TP2SP2	Covered by four-rank parity; unchanged code path
Direct Magi2MultiHeadMoE callers without split metadata	Collective fallback retained
Tensor-valued cp_split_sizes	Falls back to the legacy collective
I2VA	Bit-exact on the released checkpoint, deterministic mode
CFG parallelism, DLO, HSDP, Cache-DiT, regional compile, quantization	Not exercised

Note for #7112: its regional-compile refactor calls self.mlp.moe_mlp(routed) directly, so whichever
lands second needs to forward the split metadata through that call.

[yeahdongcn]

Same change measured on MTT S5000 (MUSA), MAGI-2 272p/125 frames/100 steps at SP4xCFG2: passing cp_split_sizes into the MoE layer saves about 46 ms per step (6%) with static-shape compile and about 40 ms with the default dynamic compile, far more than the +0.4% you measured on H20, because each scalar all_gather plus .item() blocks the MUSA launch queue. A one-time check against the gathered sizes never saw a mismatch. The branch now conflicts with main in modeling_magi2.py; once it is rebased we can rerun the S5000 numbers on your head.

[hsliuustc0106]

@NancyFyong Are you planning to continue work on this PR? Please reply by 2026-10-20 04:15:52 UTC.

A yes or a brief plan is enough; code changes do not need to be finished within these three days. If you are waiting on a review or other maintainer action, please let us know.

If there is no reply, new commit, or maintainer takeover by then, we will recheck the PR's status before deciding whether to close it as inactive. It can be reopened when work resumes.