[Bugfix] Fix Precision Mismatch in MoE Router of DeepSeek V2/V3 Models and Fused Kernels (BF16 -> FP32)

[DaizeDong]

Description

• Updated the computation of router_logits in the MoE gate to use FP32 instead of the default BF16 to enhance numerical stability.
• Ensured that activation functions (softmax and sigmoid) for router_logits operate in FP32, improving the precision of exponential calculations.
• Addressed a related TODO in vllm/_custom_ops.py.

Results

The additional computation cost is negligible as the routing operation is lightweight. However, this adjustment results in a consistent performance improvement, with Winogrande accuracy increasing from 71.27 to 71.43 for the DeepSeek-V2-Lite model.

Before Winogrande Accuracy:

After Winogrande Accuracy:

References

FP32 Computation for router_logits

• This approach is a new trick in Megatron (Issue 1421) that stabilizes MoE pretraining.
• The original implementations of DeepSeek V2 and V3 also adopt this method.

FP32 Activation for softmax and sigmoid

• Similarly, DeepSeek V2 and V3 use FP32 for activation calculations.

Potential Enhancement

Currently, the precision fix for router_logits is implemented by explicitly setting the router weights to FP32. A potentially better approach would be to retain the router weights in FP16 and adjust computation precision dynamically during execution. However, I am not fully familiar with the quant_config in ReplicatedLinear, so I believe someone with more expertise in this area could refine this further.

For now, this implementation effectively ensures stable and correct model behavior.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[jikunshang]

I feel bias should also initiate to fp32.

[DaizeDong]

Seems not necessary for a change here?
I found that the initialization method for router bias is the same in vllm and HF for DS V3, using nn.Parameter that defaults as FP32. And there is no type conversion before calculation in the original HF implementation.

[mgoin]

We should run some performance benchmarks to ensure we aren't regressing cc @tlrmchlsmth