[MUSA] Allow MAGI-2 pipeline device handling on MUSA

[yeahdongcn]

Summary

Allow MAGI-2 to use MUSA through the existing Omni platform abstraction.

• Select CUDA/MUSA devices through the platform helper; include MUSA in request synchronization and opt-in peak-memory instrumentation.
• Keep the vLLM-bundled CUDA FlashAttention resolver CUDA-only, even when a compatibility tensor reports is_cuda=True.
• For MUSA BF16 MoE, select the retained Triton kernel's existing routed-output mode to avoid its unsupported BF16 atomic_add epilogue. CUDA/ROCm, MUSA FP16/FP32 and explicit MAGI2_DETERMINISTIC=1 dispatch remain unchanged.

This is entry compatibility, not optimized full-resolution MUSA support. Attention still uses the dense Torch reference, which is only qualified by bounded tests here.

Scope

Independent of #7249 and #7206: no FA3 adapter or new BF16 MoE kernel is copied or enabled. No MoE kernel arithmetic, SwiGLU7, mHC, EP, sampler-step optimization, compilation, launch-tile or dependency-pin changes. #7156 remains the historical reference draft.

The BF16 MoE adjustment reuses deterministic=True to choose tl.store followed by the existing scatter. It is a non-atomic Triton-output selection, not a new whole-pipeline determinism guarantee. Its full-model buffer/memory/performance impact remains unmeasured.

Exact-head validation

Tested and pushed commit: 96a02eaa6948a6a26e3e8c49f9b61d61af359eac, a signed-off follow-up to f9101b2a3.

• 82 CPU tests passed. The new 36-case matrix covers CUDA/ROCm/MUSA, CPU-like dispatch, BF16/FP16/FP32 and deterministic unset/0/1. Before the fix it reproduced two MUSA BF16 mode-selection failures.
• 6 MUSA tests passed, zero skipped, on one 48-SM MTT S5000, driver 3.3.8-server (no driver pin).
• Two new cases run real default MoE routing and the retained Triton kernel with the environment override unset: D=256, I=1280, top_k=6, tokens/heads 2/3 and 129/1. Constant clipped expert inputs yield the independently derived nonzero output 105/64; both cases match it bitwise and confirm non-atomic output selection.
• The four retained MUSA checks cover device selection before checkpoint loading, seed repeatability, eager tiny-transformer CPU/MUSA parity (rtol=2e-4, atol=2e-5), and synchronization/memory instrumentation with a stub runtime.
• Targeted pre-commit passed, including Ruff, mypy, pytest markers and CUDA API checks. Container/test exit codes are 0.

Stack: Python 3.10.12; 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; MATE 0.2.6; flash_attn_3 0.2.6+musa.

Image: registry.mthreads.com/mcconline/inference/vllm-omni@sha256:7f4f2cd83a88979025ef92968334f6e075eb24d45d02e0c30dbb222f7780d58a.

Installation was skipped for this Python-only source update. The exact Git archive and runner were checksum-verified; PYTHONPATH and the resolved import path identify the candidate source. No dependency/native build was performed. On the MUSA stack, import torchada first:

import torchada
import pytest

raise SystemExit(pytest.main([
    "-vv", "-s", "-o", "addopts=", "-p", "no:cacheprovider",
    "tests/diffusion/models/magi2/test_musa_pipeline.py",
    "-m", "musa",
]))

CPU validation additionally includes test_pipeline_magi2.py and test_native_preview.py, selecting cpu.

Not run: CUDA hardware regression, multi-rank MUSA collectives, full-checkpoint/video accuracy, optimized FA3, compiled/graph execution, full-model memory sizing and performance benchmarks. This remains Draft.

[yeahdongcn]

Superseded by #8498, which carries the MUSA pipeline device handling. The BF16 non-atomic MoE change is dropped there because main has routed MUSA BF16 through the fused MoE path since #7206.