[MUSA] Use MUSA-optimized operators in piecewise CUDA graph

[popsiclexu]

Motivation

MUSA devices previously could not use piecewise CUDA graph due to multiple incompatibilities: torch.compile on MUSA fails to serialize custom device types in FX-generated code, and several operators (SiluAndMul, RMSNorm, FP8 quantization) lacked fake kernel registrations needed for torch.compile tracing. As a workaround, these operators fell back to native PyTorch implementations when piecewise CUDA graph was enabled, resulting in suboptimal performance. This PR enables piecewise CUDA graph to work correctly on MUSA devices and allows the use of MUSA-optimized operators (e.g., nn.SwishGLU).

Modifications

1. Patch FX codegen for custom devices (patch_torch.py): Add patch_fx_custom_device() that post-processes FX-generated source code, replacing device(type='musa', index=N) with torch.device('musa:N'). This fixes torch.compile serialization failure on MUSA, which was the root cause of piecewise CUDA graph crashes.

2. Register fake kernels for torch.compile tracing:

   • activation.py: Register aten::_fused_swiglu_forward fake kernel so nn.SwishGLU can be traced.
   • fp8_kernel.py: Register sgl_kernel::sgl_per_token_group_quant_8bit_v2 fake kernel so FP8 quantization can be traced.

3. Remove piecewise CUDA graph workarounds in MUSA operators (activation.py, layernorm.py): Remove the if not get_global_server_args().disable_piecewise_cuda_graph: return self.forward_native(x) guards in SiluAndMul.forward_musa and RMSNorm.forward_musa, so MUSA-optimized code paths are used even when piecewise CUDA graph is enabled.

Accuracy Tests

run server:

python3 -m sglang.launch_server \
  --model-path "/home/dist/DeepSeek-V2-Lite-Chat-FP8/ \
  --served-model-name "deepseek" \
  --mem-fraction-static 0.5 \
  --disable-overlap-schedule \
  --attention-backend fa3 \
  --cuda-graph-bs $(seq 1 16) \
  --chunked-prefill-size -1 \
  --disable-radix-cache  \
  --piecewise-cuda-graph-max-tokens 2048

python3 -m sglang.test.few_shot_gsm8k   --host http://127.0.0.1   --port 30000   --num-questions 200   --parallel 16
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [01:06<00:00,  3.00it/s]
Accuracy: 0.650
Invalid: 0.005
Latency: 66.860 s
Output throughput: 381.153 token/s

Known Limitations

Current piecewise CUDA graph does not support --moe-runner-backend=deepgemm on MUSA. We plan to add support in a follow-up PR.

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[gemini-code-assist]

Code Review

This pull request implements MUSA device support, focusing on piecewise CUDA graph execution and kernel registrations. Key updates include a torch.fx patch for MUSA device serialization, fake kernel registrations for activation and quantization layers, and the removal of redundant CUDA graph checks. Feedback highlights improvements for robustness, such as using _replace for NamedTuple objects, safely accessing torch.version attributes, and employing conditional fake registrations for kernels.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[yeahdongcn]

/tag-and-rerun-ci