[RFC] Native MAGI-2 Preview architecture and distributed deployment

Context

PR #5918 adds a native vLLM-Omni implementation of sand-ai/MAGI-2-preview. The change is substantially larger than 500 non-test LOC, so this issue records the architecture, ownership boundaries, rollout phases, and explicit scope exceptions for maintainer review.

The runtime does not import SandAI's inference package, alter sys.path, start the reference process manager, or require a reference checkout. SandAI's Apache-2.0 implementation at f68a0f9bbccbea56e0177a9bd912abc3b18ffe61 and model revision 2dea51b64db47ee5b4402d36fd90829a0c58913b are accuracy references only.

Native architecture

The in-tree graph owns:

• Qwen3.5-27B text conditioning and structured-prompt normalization;
• optional aspect-preserving I2VA conditioning through the Wan2.2 VAE;
• packed conditional/unconditional video, audio, text, and image-reference tokens;
• the 40-layer Preview DiT with four-stream mHC residual mixing;
• learned-sink, softcapped, variable-length attention;
• the 12-head / 256-expert MagiMoE and native SwiGLU7 Triton path;
• independent FlowUniPC video/audio updates;
• TurboVAE video decode and Oobleck audio decode/resample.

The released checkpoint hierarchy is preserved and strict loading rejects missing, unexpected, or shape-invalid tensors.

Attention ownership

MAGI-2 needs packed CFG isolation, learned sink logits corrected through FlashAttention LSE, softcap, uneven Ulysses token splits, and model-specific Q/K FP32 normalization/RoPE order. PR #5918 routes this kernel through the shared vllm_omni.diffusion.attention.layer.Attention abstraction using an explicit custom_attention extension point. The shared layer retains compile/dispatch ownership; the native MAGI kernel owns the specialized packed computation and communication. Generic KV-cache quantization is disabled for this kernel.

Parallelism and DLO

Supported four-device layouts:

• resident TP4, TP2SP2, and SP4;
• DLO DP4 and DP2SP2 with AllGather;
• rank-local DLO SP4 without AllGather.

TP owns matrix and MoE-head shards; SP owns packed token/head exchange. For DLO DP2SP2, native local transforms are applied before orthogonal DP sharding so equal-SP-coordinate DP peers reconstruct identical tensors. SP-only DLO streams non-interchangeable rank-local MoE-head weights and therefore rejects AllGather.

The reference-aligned default remains resident SP4. Compatible eight-worker factorizations pass topology validation, but the official eight-Hopper runtime was unavailable on this host.

Cache-DiT

Cache-DiT wraps only transformer.block.layers, the repeated denoising stack. Pre/post adapters, packed CFG preparation, attention sinks, and decoders still run normally. MAGI packs both CFG branches into one forward, so the adapter declares has_separate_cfg=False.

A custom-enabler result carries the exact nested BlockAdapter target for refresh, statistics, and teardown. This composes with DLO because the offloader continues to own layer placement while Cache-DiT decides whether middle layers execute.

Validation includes:

• a three-layer CPU test that forces a cache hit and proves only the front block reruns;
• a four-L20X, rank-local DLO SP4 shared-example smoke at 272p/4 steps;
• a valid 125-frame, 12.5 fps, stereo 44.1 kHz output contract.

Cache-DiT is approximate; resident SP4 without cache remains the fidelity baseline.

Rollout / review phases

1. Native checkpoint and component parity.
2. Packed attention, mHC, MoE, and decoder parity.
3. TP4 / TP2SP2 / SP4 distributed parity.
4. DLO DP4 / DP2SP2 / rank-local SP4.
5. Full 540p, 100-step T2VA and I2VA qualification on four local devices.
6. Shared Attention integration, Cache-DiT support, strict docs/precheck cleanup.
7. Follow-up scope: Preview 1080p refiner and official eight-Hopper runtime qualification.

Validation and scope exceptions

The final targeted CPU suite passes 309 tests; Ruff check/format, strict MkDocs, DCO, and git diff --check pass. Full 540p/100-step T2VA and I2VA completed through shared examples with rank-local DLO SP4.

No benchmark file and no model-specific example script are added. Performance evidence remains tied to existing shared scripts and is explicitly qualification data, not a benchmark suite.