[Refactor] Share packed head-parallel exchanges for diffusion models

[yeahdongcn]

Summary

Extract packed head-parallel exchanges into a shared diffusion distributed module, with MAGI-2 as the first caller.

• HeadParallelLayout holds immutable input-token counts in process-group rank order and validates the local rank/group contract.
• scatter_heads_gather_tokens maps [local_tokens, heads, dim] to [total_tokens, local_heads, dim]; scatter_tokens_gather_heads reverses the exchange using the same layout.
• Keep MAGI-2's existing group selection, single-rank behavior, and optional token-count discovery in its wrappers.

The functions consume an existing caller-owned process group. They do not create/destroy groups, select EP topology, pad model heads, route expert rows, or change conventional FusedMoE semantics. Explicit layouts need no token-count collective. All ranks must agree on the split tuple, dtype and head layout before entering a collective; local validation does not replace a distributed agreement protocol.

Scope

Three files only: the shared primitive, minimal MAGI-2 wrappers, and focused tests. This is the exchange-contract portion of the MAGI-2 EP decomposition, not complete configurable EP support.

No changes to DiffusionParallelConfig, the existing TP/SP/CFG/DP/EP group lifecycle, public environment variables, MoE kernels, attention backends, sampler, compilation or dependency pins. Configurable group selection and the MAGI-2 EP adapter remain separate work. #7156 stays a reference draft.

The collective layout is retained from the native MAGI-2 implementation; its source attribution is preserved. Invalid rank/shape/count inputs are rejected, and globally empty exchanges skip communication.

Validation

Commit: 7175f8575a3ab45726050c306862e6085b313616 (one signed-off commit), based on upstream b58ff5cb8b17250b76f9cdf9b9b46385cdda4376.

• Targeted pre-commit passed, including Ruff, mypy and test-marker checks.
• 64 CPU/Gloo tests passed: layout/launch contracts; real four-rank exchanges over world, contiguous replica subgroups and noncontiguous rank subgroups; existing EP-layout tests; MAGI-2 pipeline/native tests and its four-rank TP4/TP2xSP2/SP4 tiny-model oracle.
• One two-rank MCCL test passed on two 60-SM MTT S5000 cards (driver 5.2.0-server). It checks equal, uneven, one-empty and all-empty token partitions, direct exchange, head-owner-specific transformed return, and the MAGI-2 wrappers against exact CPU values (rtol=atol=0).
• Noncontiguous input and legacy inferred-count wrapper paths are covered. Shape/rank errors are checked before collective launch.

Stack: Python 3.10; torch/torch_musa 2.11.0.post1+musa5.2.0; torchada 0.1.83; vLLM 0.28.0; vllm-musa 0.1.28; Triton 3.2.0. The exact checkout was installed editable with --no-deps --no-build-isolation into an isolated container; the MUSA entrypoint imports torchada before Torch/vLLM.

python -m pytest -q -o addopts= -m cpu \
  tests/diffusion/distributed/test_head_parallel.py \
  tests/diffusion/distributed/test_expert_parallel_layout.py \
  tests/diffusion/models/magi2/test_native_preview.py \
  tests/diffusion/models/magi2/test_pipeline_magi2.py \
  tests/diffusion/models/magi2/test_native_distributed_parity.py

For MUSA, run the new test file with -m musa from a torchada-first Python entrypoint with a __main__ guard; it spawns two workers. Tests use bounded process-group timeouts.

Not run: CUDA/NCCL hardware, multi-node groups, graph/compile integration, full-checkpoint/video E2E, or performance benchmarks. This refactor does not claim a latency improvement or complete MAGI-2 EP enablement; it remains Draft.

[yeahdongcn]

Superseded by #8511, which carries this change rebased on main together with the head-sharded EP groups and the MAGI-2 activation.