[2/N][Bug] Fix w4afp8 MoE NaN issue (python)

[yuhyao]

Motivation

As reported in Issue#8493, NaNs may appear in the activations during DeepSeek-R1 w4afp8 inference. You can easily reproduce this issue with the following commands.
Command to launch server:

SGL_ENABLE_JIT_DEEPGEMM=1 python -m sglang.launch_server --model-path /path/to/DeepSeek-R1-W4AFP8 --tp-size 8 --context-length 32768 --trust-remote-code --host 0.0.0.0 --port 8000 --mem-fraction-static 0.8 --cuda-graph-max-bs 256 --max-running-requests 256 --random-seed 42 --disable-radix-cache --chunked-prefill-size -1 --ep-size 8 --cuda-graph-bs 1 2 4 8 16 32 64 128 256

Command to send request:

curl -N "http:/localhost:8000/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "model",
    "stream": false,
    "max_tokens": 1280,
    "messages": [
      {
        "role": "user",
        "content": [
          {"type": "text", "text": "如何制作红烧牛肉面？"}
        ]
      }
    ],
    "n": 1
  }'

Modifications

The root cause is that the w4afp8 grouped GEMM kernel uses float16 as the output type to improve throughput. However, in some cases the results exceed the maximum value of float16, producing Infs in the output, which then propagate through RMSNorm and turn into NaNs.

To resolve this, we switch the output type from float16 to bfloat16, which provides a wider dynamic range and avoids overflow in this scenario.

Note: the sgl-kernel changes are in PR#9953.

Accuracy Tests

After fix, EP8 achieves the following:
mmlu (sglang): 87.0
AIME24 (evalscope): 80.0 (43.33 before fix, probably crash if you set bs > 1.)

Benchmarking and Profiling

As noted earlier, using bfloat16 as the output type prevents overflow, but it comes at the cost of lower throughput in the GEMM kernel epilogue. Below, we present the benchmark results.

Benchmark configuration: ISL/OSL = 1000/1000, num_prompts=256, qps=64, max_concurrency=64.

python -m sglang.bench_serving --backend sglang --host localhost --port 8000 --dataset-name random --num-prompts 256 --random-input-len 1000 --random-output-len 1000 --random-range-ratio 1 --request-rate 64 --max-concurrency 64

Launching server using EP8:

SGL_ENABLE_JIT_DEEPGEMM=1 python -m sglang.launch_server --model-path /path/to/DeepSeek-R1-W4AFP8 --tp-size 8 --context-length 32768 --trust-remote-code --host 0.0.0.0 --port 8000 --mem-fraction-static 0.8 --cuda-graph-max-bs 256 --max-running-requests 256 --random-seed 42 --disable-radix-cache --chunked-prefill-size -1 --ep-size 8 --cuda-graph-bs 1 2 4 8 16 32 64 128 256

Launching server using TP8:

SGL_ENABLE_JIT_DEEPGEMM=1 python -m sglang.launch_server --model-path /path/to/DeepSeek-R1-W4AFP8 --tp-size 8 --context-length 32768 --trust-remote-code --host 0.0.0.0 --port 8000 --mem-fraction-static 0.8 --cuda-graph-max-bs 256 --max-running-requests 256 --random-seed 42 --disable-radix-cache --chunked-prefill-size -1 --cuda-graph-bs 1 2 4 8 16 32 64 128 256

The following results were obtained on 8×H800 GPUs.

EP8 results

Before fix:

============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    64.0
Max request concurrency:                 64
Successful requests:                     256
Benchmark duration (s):                  166.11
Total input tokens:                      256000
Total generated tokens:                  256000
Total generated tokens (retokenized):    254937
Request throughput (req/s):              1.54
Input token throughput (tok/s):          1541.11
Output token throughput (tok/s):         1541.11
Total token throughput (tok/s):          3082.22
Concurrency:                             63.81
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   41403.81
Median E2E Latency (ms):                 41384.33
---------------Time to First Token----------------
Mean TTFT (ms):                          2766.29
Median TTFT (ms):                        2718.30
P99 TTFT (ms):                           3645.31
---------------Inter-Token Latency----------------
Mean ITL (ms):                           38.68
Median ITL (ms):                         37.94
P95 ITL (ms):                            38.76
P99 ITL (ms):                            39.34
Max ITL (ms):                            2835.66
==================================================

After fix:

============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    64.0
Max request concurrency:                 64
Successful requests:                     256
Benchmark duration (s):                  184.83
Total input tokens:                      256000
Total generated tokens:                  256000
Total generated tokens (retokenized):    255184
Request throughput (req/s):              1.39
Input token throughput (tok/s):          1385.07
Output token throughput (tok/s):         1385.07
Total token throughput (tok/s):          2770.14
Concurrency:                             63.83
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   46080.81
Median E2E Latency (ms):                 45965.97
---------------Time to First Token----------------
Mean TTFT (ms):                          2869.23
Median TTFT (ms):                        2852.96
P99 TTFT (ms):                           3950.73
---------------Inter-Token Latency----------------
Mean ITL (ms):                           43.26
Median ITL (ms):                         42.24
P95 ITL (ms):                            43.44
P99 ITL (ms):                            50.38
Max ITL (ms):                            3126.65
==================================================

TP8 results

Before fix:

============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    64.0
Max request concurrency:                 64
Successful requests:                     256
Benchmark duration (s):                  219.78
Total input tokens:                      256000
Total generated tokens:                  256000
Total generated tokens (retokenized):    255332
Request throughput (req/s):              1.16
Input token throughput (tok/s):          1164.78
Output token throughput (tok/s):         1164.78
Total token throughput (tok/s):          2329.56
Concurrency:                             63.70
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   54686.57
Median E2E Latency (ms):                 45596.43
---------------Time to First Token----------------
Mean TTFT (ms):                          2723.64
Median TTFT (ms):                        2729.81
P99 TTFT (ms):                           3639.44
---------------Inter-Token Latency----------------
Mean ITL (ms):                           52.02
Median ITL (ms):                         42.08
P95 ITL (ms):                            94.43
P99 ITL (ms):                            571.91
Max ITL (ms):                            2800.71
==================================================

After fix:

============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    64.0
Max request concurrency:                 64
Successful requests:                     256
Benchmark duration (s):                  203.98
Total input tokens:                      256000
Total generated tokens:                  256000
Total generated tokens (retokenized):    255363
Request throughput (req/s):              1.25
Input token throughput (tok/s):          1255.00
Output token throughput (tok/s):         1255.00
Total token throughput (tok/s):          2509.99
Concurrency:                             63.84
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   50869.33
Median E2E Latency (ms):                 50840.18
---------------Time to First Token----------------
Mean TTFT (ms):                          2976.06
Median TTFT (ms):                        2961.94
P99 TTFT (ms):                           3972.77
---------------Inter-Token Latency----------------
Mean ITL (ms):                           47.94
Median ITL (ms):                         47.17
P95 ITL (ms):                            48.16
P99 ITL (ms):                            48.85
Max ITL (ms):                            3401.90
==================================================

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.

[gemini-code-assist]

Summary of Changes

Hello @yuhyao, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical numerical stability problem affecting the DeepSeek-R1 w4afp8 Mixture-of-Experts (MoE) inference. By upgrading the precision of intermediate computations, it ensures that model activations remain within a valid numerical range, thereby preventing the propagation of NaN values and maintaining the integrity of the inference process.

Highlights

• Numerical Stability Fix: Addresses a NaN (Not-a-Number) issue observed during DeepSeek-R1 w4afp8 MoE inference, caused by float16 overflow in grouped GEMM kernel outputs.
• Data Type Transition: The core solution involves switching the output data type of the grouped GEMM kernel from float16 to bfloat16 to leverage its wider dynamic range and prevent overflow.
• Accuracy Verification: Post-fix accuracy tests show improved results for EP8, with mmlu at 87.0 and AIME24 at 80.0, confirming the resolution of the issue without negative impact on model performance.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively resolves a NaN issue in the w4afp8 MoE kernel by switching the intermediate computation and output data types from float16 to bfloat16. This change leverages bfloat16's wider dynamic range to prevent overflows that were causing Inf values and subsequent NaNs. The fix is consistently applied across the Python layer, the C++/CUDA kernel definitions, and the test files. Additionally, the test suite's reference implementation has been improved to use float32 for intermediate calculations, enhancing its robustness and reliability. The changes are correct and well-implemented.

[AniZpZ]

please provide the performace drop data

[yangsijia-serena]

Great work! I've just addressed this problem too, the same solution~

[yuhyao]

please provide the performace drop data

Just added, thanks for reviewing!

[AniZpZ]

please check the CI https://github.com/sgl-project/sglang/actions/runs/17404771879/job/49421213668?pr=9918

[AniZpZ]

we need to release a new sgl-kernel to pass the CI and get this PR merged

cc @zhyncs