[Feature] Add CuTeDSL Gram Newton-Schulz support for Muon

[RangiLyu]

Motivation

This PR integrates Gram Newton-Schulz into XTuner's distributed Muon optimizer

• Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon (https://dao-ailab.github.io/blog/2026/gram-newton-schulz/)
• Reference implementation (https://github.com/Dao-AILab/gram-newton-schulz)

Gram Newton-Schulz performs most iterations on the smaller square Gram matrix and uses symmetric CuTeDSL GEMM kernels to reduce the cost of Muon orthogonalization.

Changes

• Add an opt-in use_gram_newton_schulz switch to Muon and MuonConfig.
• Use the recommended safety-scaled Polar Express coefficients.
• Use FP32 normalization followed by FP16 internal computation.
• Use gram_restart_iterations=(2,) by default.
• Use CuTeDSL kernels for supported matrix shapes.
• Automatically fall back to standard Newton-Schulz for square matrices and matrices with aspect ratio no greater than 2.0.
• Add a gram-ns optional dependency group for Quack and CUTLASS DSL.

Optimizer step benchmark

Benchmark configuration:

• Model: Qwen3.5-35B-A3B
• World size: 32 GPUs
• Maximum sequence length: 4096
• Optimizer timing warmup: 20 steps
• Sampling interval: 10 steps
• Samples: 48, covering steps 30–500
• Measurement scope: complete optimizer.step()
• Reported values are averages over all collected samples

Optimizer step metric	Original XTuner Muon	Gram Newton-Schulz	Latency reduction	Speedup
Wall time, mean across 32 ranks	127.8944 ms	104.3441 ms	18.41%	1.2257×
Wall time, maximum across 32 ranks	128.8304 ms	105.0166 ms	18.48%	1.2268×
CUDA event time, maximum across 32 ranks	128.8231 ms	105.0095 ms	18.49%	1.2268×

The cross-rank maximum wall time represents the synchronous optimizer critical path. Under this configuration, Gram Newton-Schulz reduces the complete optimizer step by approximately 23.81 ms, corresponding to an 18.48% latency reduction or 1.2268× speedup.

Usage

Install the optional dependencies:

pip install -e '.[gram-ns]' --no-build-isolation

Enable it through MuonConfig:

optim_cfg = MuonConfig(
    use_gram_newton_schulz=True,
)

The Qwen3.5 example also supports:

USE_GRAM_NEWTON_SCHULZ=1