[Feature] Add SonicMoE support

[sallyjunjun]

Summary

This PR adds an optional SonicMoE expert backend for XTuner MoE models.

The feature is disabled by default. Existing MoE configs still use expert_backend="grouped_gemm" unless users explicitly opt in with expert_backend="sonicmoe", so the original MoE flow, dispatcher logic, checkpoint layout, and default training behavior remain unchanged.

Implementation

The implementation introduces SonicMoE as a backend adapter instead of changing XTuner’s parameter ownership or checkpoint format.

Main changes:

• Add expert_backend to MoEConfig, with default value "grouped_gemm".
• Add SonicMoEBackendConfig and SonicMoEBackend.
• Add a SonicMoE op wrapper around the official sonic-moe package.
• Route MoE expert execution through SonicMoE only when expert_backend="sonicmoe".
• Keep XTuner’s native expert parameters and checkpoint-compatible layout:
  • XTuner keeps [E, 2I, H] / [E, H, I].
  • The SonicMoE adapter passes permuted views to SonicMoE without materializing new parameters.
• Preserve XTuner router output contract:
  • SonicMoE consumes XTuner’s unpermuted topk_ids and topk_weights.
  • General/dropless routing preserves every native top-k assignment.
  • Optional token rounding metadata is implemented behind sonicmoe_cfg.routing_mode="token_rounding".

The initial integration intentionally validates and rejects unsupported combinations:

• ep_size must be 1.
• expert_tp_size must be 1.
• dispatcher must be None.
• float8_cfg must be None.
• Expert activation must be SwiGLU.
• SonicMoE execution is CUDA-only and BF16-only.

How To Enable

Install the optional dependency:

pip install -e ".[sonicmoe]"

Then enable SonicMoE in the MoE model config:

model = Qwen3MoE30BA3Config(
    expert_backend="sonicmoe",
)

Optional SonicMoE routing config can be passed through sonicmoe_cfg:

  from xtuner.v1.module.moe_backend import SonicMoEBackendConfig

  model = Qwen3MoE30BA3Config(
      expert_backend="sonicmoe",
      sonicmoe_cfg=SonicMoEBackendConfig(
          routing_mode="general",
      ),
  )

For token rounding mode:

  model = Qwen3MoE30BA3Config(
      expert_backend="sonicmoe",
      sonicmoe_cfg=SonicMoEBackendConfig(
          routing_mode="token_rounding",
      ),
  )

Token rounding has additional validation requirements: greedy router, softmax scores, no grouped router selection, norm_topk_prob=True, and router_scaling_factor=1.0.

Tests And Validation

Added unit coverage in tests/module/test_sonicmoe_backend.py for:

• SonicMoE config validation.
• Default backend behavior remaining unchanged.
• Decoder-layer SonicMoE routing path.
• SonicMoE call contract against XTuner router outputs.
• Empty-token behavior preserving autograd participation.
• Official SonicMoE forward/backward parity path.
• 8-GPU FSDP smoke/parity coverage for distributed training.

A standalone benchmark/profiler is also added:

tests/profiler/qwen35_sonicmoe_fsdp_benchmark.py

It compares XTuner native grouped GEMM MoE vs SonicMoE on Qwen3.5 MoE FSDP using identical model weights and identical per-rank input batches. Each measured step performs forward and backward without optimizer update, so loss comparison is directly attributable to the expert backend implementation. The profiler emits both JSON
results and a markdown report including loss deltas, throughput, per-step latency, and SonicMoE speedup ratio.

Example:

  XTUNER_DETERMINISTIC=true torchrun --nproc-per-node 8 \
    tests/profiler/qwen35_sonicmoe_fsdp_benchmark.py \
    --model-path /path/to/Qwen3.5-35B-A3B \
    --output /path/to/qwen35_sonicmoe_fsdp_benchmark.json \
    --steps 10 --warmup-steps 5 --deterministic