[Refactor] Reuse FP32 RMSNorm normalization in MAGI-2

[yeahdongcn]

Summary

Use the existing PyTorch RMSNorm operation for MAGI-2's FP32 normalization, while keeping its checkpoint-specific affine operation outside the call.

• Input is still converted to FP32 before normalization.
• F.rms_norm receives no affine weight; MAGI-2 still applies its flattened, modality-major (weight + 1) scales afterward, then casts to the requested output dtype.
• Existing constructor, state-dict key/layout, modality grouping, multiple-pattern behavior and singleton-head broadcasting are retained.
• Nonstandard input shapes keep the original broadcast behavior. No per-call dtype override, new environment variable, model wiring, cache or custom kernel is added.

This is the normalization portion of the reference #7156 decomposition, without its BF16/FP16 intermediate downcasts or separate mHC overrides. Omni's standard RMSNorm parameter represents effective scale directly, so substituting it for MAGI-2's delta-scale parameter would change checkpoint semantics; only the unweighted normalization operation is reused here.

Numerical contract and investigation

This is not a bitwise-equivalence claim. An initial MUSA test used a uniform rtol=atol=1e-6 on all output dtypes: FP32 and BF16 cases passed, but 2 of 43008 FP16 values differed by one FP16 ULP.

An independent same-source diagnostic compared outputs before and after the final cast. Across nine dtype/shape cases, all pre-cast FP32 comparisons passed the original 1e-6 gate (largest absolute difference 9.54e-7); both implementations' low-precision outputs exactly matched their own FP32 outputs cast at the end. The discrepant FP16 values straddled a rounding midpoint. For example, pre-cast 1.6381838322 versus 1.63818359375 rounded to 1.638671875 versus 1.6376953125.

The final test makes the acceptance contract explicit and expands coverage to three seeds:

1. Pre-cast FP32 affine outputs must still satisfy rtol=atol=1e-6.
2. Each returned output must exactly equal the corresponding FP32 result cast to its dtype.
3. FP16/BF16 output differences must be at most one local ULP, measured as the larger spacing to the two representable neighbors of the reference value.

This replaces the inappropriate uniform FP32 tolerance on quantized outputs; it is an explicitly documented gate adjustment, not a retroactive pass for the initial run. Separate FA3 numerical gates are unchanged. Full-model acceptance remains a separate requirement.

Validation

Tested commit da47bf6772f7b91425d25045d8ea3e6e99a6fc19, one signed-off commit based on upstream d59d1664858956976a12797a8613bd4b1d803097.

• Targeted pre-commit passed, including Ruff, mypy and test markers.
• 32 CPU tests passed, 4 MUSA deselected: explicit original-formula parity, FP32 operation arguments, modality/pattern grouping with empty groups, input/state preservation, low-precision parameter offsets, gradients, state-dict reload and weight updates, shape/value edge cases and two-layer native model parity.
• A Dynamo test asserts actual FX graph creation and changed-weight behavior. Its backend runs the graph eagerly; it is not Inductor/capture validation.
• 4 MUSA tests passed under the contract above. Three seeds, widths 128/768/1536, FP32/BF16/FP16 inputs, FP32 output overrides, strided/multi-head inputs, multiple scale patterns and empty modalities are covered. The real RMSNorm operation is observed receiving FP32 input and no weight.
• Final hardware: one 48-SM MTT S5000, driver 5.2.0-server; no driver pin. Exact checkout installed editable without dependency changes.

Stack: torch/torch_musa 2.11.0.post1+musa5.2.0, torchada 0.1.83, vLLM 0.28.0, vllm-musa 0.1.28. MUSA entrypoint imports torchada first. The current dispatch table reports a composite RMSNorm registration; using this API alone does not establish that a particular kernel is fused or faster.

python -m pytest -q -o addopts= -m cpu \
  tests/diffusion/models/magi2/test_rmsnorm.py \
  tests/diffusion/models/magi2/test_native_preview.py

For MUSA, run the new test file with -m musa from a torchada-first entrypoint with one visible GPU.

Not run: CUDA hardware, GPU backward, Inductor/graph, full MUSA MAGI-2 checkpoint/video or performance. No DiT-step improvement is claimed; this remains Draft pending review and broader integration gates.

[hsliuustc0106]

This PR touches tests/diffusion/, vllm_omni/diffusion/ (2 files). Based on CODEOWNERS coverage of the changed files, the most-related reviewers appear to be:

@fhfuih @Bounty-hunter @wtomin

Could one of you take a look when you get a chance? Thanks!

[hsliuustc0106]

@yeahdongcn Final follow-up: are you planning to continue work on this PR? Please reply within 72 hours of this comment's published timestamp.

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, CI, or other maintainer action, or need more time, please let us know.

If there is no reply, new commit, or maintainer takeover after that period, we will recheck and close this PR as inactive. If you explicitly confirm that you will not continue, we will close it after checking that nobody has taken over. It can be reopened when work resumes.