From 4d5917e744578184237288c62c54a2035f69c72f Mon Sep 17 00:00:00 2001 From: Jun Liu Date: Fri, 24 Jul 2026 20:45:44 +0900 Subject: [PATCH] Add DeepSeek-reference 1e-20 epsilon to top-k renormalization to prevent 0/0 NaN (#31017) Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- python/sglang/srt/layers/moe/topk.py | 53 +++++-- .../moe/test_topk_renormalize_degenerate.py | 141 ++++++++++++++++++ 2 files changed, 182 insertions(+), 12 deletions(-) create mode 100644 test/registered/moe/test_topk_renormalize_degenerate.py diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index accbc5d17..bd166b363 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -131,6 +131,17 @@ _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_musa = is_musa() +# Epsilon added to the top-k weight sum before renormalization, matching the +# DeepSeek reference gate (modeling_deepseek.py: `topk_weight.sum(...) + 1e-20`) +# and flashinfer's trtllm routing kernels (mSumEpsilon). With sigmoid scoring +# plus a selection bias, a token whose selected experts all have deeply negative +# router logits can have every gathered sigmoid weight underflow to exactly +# zero; a bare division then yields 0/0 = NaN and poisons the token's output +# row. For healthy tokens the sum is >= sigmoid(logit_max) >> 1e-20, so results +# are unchanged. The renormalization is performed in float32 (the reference gate +# computes the whole gate in fp32); the epsilon underflows to zero in float16. +_RENORMALIZE_SUM_EPSILON = 1e-20 + # Experimental: skip the HIP padded-token routing-weight masking entirely. # Padded (CUDA-graph) rows are discarded downstream and the MoE combine is # per-token, so zeroing their weights is in principle unnecessary. Gated off by @@ -696,7 +707,13 @@ def fused_topk_torch_native( topk_weights, topk_ids = torch.topk(topk_weights, topk, dim=-1) if renormalize: - topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + # fp32 like the reference gate (the epsilon is not representable in + # fp16); the sum dtype and the division's type promotion upcast inside + # the existing kernels, so no extra cast launch is needed + topk_weights = topk_weights / ( + topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32) + + _RENORMALIZE_SUM_EPSILON + ) return topk_weights, topk_ids @@ -983,12 +1000,15 @@ def grouped_topk_gpu( ) if renormalize: + # fp32 like the reference gate (the epsilon is not representable in + # fp16); the sum dtype and the division's type promotion upcast inside + # the existing kernels, so no extra cast launch is needed topk_weights_sum = ( - topk_weights.sum(dim=-1, keepdim=True) + topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32) if num_fused_shared_experts == 0 - else topk_weights[:, :-1].sum(dim=-1, keepdim=True) + else topk_weights[:, :-1].sum(dim=-1, keepdim=True, dtype=torch.float32) ) - topk_weights = topk_weights / topk_weights_sum + topk_weights = topk_weights / (topk_weights_sum + _RENORMALIZE_SUM_EPSILON) if apply_routed_scaling_factor_on_output: topk_weights *= routed_scaling_factor @@ -1109,8 +1129,11 @@ def kimi_k2_biased_topk_impl( topk_weights = scores.gather(1, topk_ids) if renormalize: - topk_weights_sum = topk_weights.sum(dim=-1, keepdim=True) - topk_weights = topk_weights / topk_weights_sum + # fp32 like the reference gate (the epsilon is not representable in + # fp16); the sum dtype and the division's type promotion upcast inside + # the existing kernels, so no extra cast launch is needed + topk_weights_sum = topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32) + topk_weights = topk_weights / (topk_weights_sum + _RENORMALIZE_SUM_EPSILON) if apply_routed_scaling_factor_on_output: topk_weights *= routed_scaling_factor @@ -1165,12 +1188,15 @@ def biased_topk_impl( ) if renormalize: + # fp32 like the reference gate (the epsilon is not representable in + # fp16); the sum dtype and the division's type promotion upcast inside + # the existing kernels, so no extra cast launch is needed topk_weights_sum = ( - topk_weights.sum(dim=-1, keepdim=True) + topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32) if num_fused_shared_experts == 0 - else topk_weights[:, :-1].sum(dim=-1, keepdim=True) + else topk_weights[:, :-1].sum(dim=-1, keepdim=True, dtype=torch.float32) ) - topk_weights = topk_weights / topk_weights_sum + topk_weights = topk_weights / (topk_weights_sum + _RENORMALIZE_SUM_EPSILON) if apply_routed_scaling_factor_on_output: topk_weights *= routed_scaling_factor @@ -1294,12 +1320,15 @@ def biased_grouped_topk_impl( ) if renormalize: + # fp32 like the reference gate (the epsilon is not representable in + # fp16); the sum dtype and the division's type promotion upcast inside + # the existing kernels, so no extra cast launch is needed topk_weights_sum = ( - topk_weights.sum(dim=-1, keepdim=True) + topk_weights.sum(dim=-1, keepdim=True, dtype=torch.float32) if num_fused_shared_experts == 0 - else topk_weights[:, :-1].sum(dim=-1, keepdim=True) + else topk_weights[:, :-1].sum(dim=-1, keepdim=True, dtype=torch.float32) ) - topk_weights = topk_weights / topk_weights_sum + topk_weights = topk_weights / (topk_weights_sum + _RENORMALIZE_SUM_EPSILON) if apply_routed_scaling_factor_on_output: topk_weights *= routed_scaling_factor diff --git a/test/registered/moe/test_topk_renormalize_degenerate.py b/test/registered/moe/test_topk_renormalize_degenerate.py new file mode 100644 index 000000000..b5ff48ff6 --- /dev/null +++ b/test/registered/moe/test_topk_renormalize_degenerate.py @@ -0,0 +1,141 @@ +"""Regression test: top-k renormalization must not emit NaN for degenerate tokens. + +With sigmoid scoring plus a selection bias (DeepSeek noaux_tc-style gates), +experts are selected by `sigmoid(logits) + bias` but weighted by the raw +sigmoid. A token whose router logits are all deeply negative (< ~-88) has +every gathered sigmoid weight underflow to exactly 0.0, so a bare +`weights / weights.sum()` renormalization computes 0/0 = NaN and poisons the +token's whole output row (observed in production as '!'-spam / NaN logits, +see sgl-project/sglang#30989). The DeepSeek reference gate guards this with +`sum + 1e-20` (modeling_deepseek.py), as does flashinfer's trtllm routing +(flashinfer-ai/flashinfer#3803); these torch implementations must match. +""" + +import unittest + +import torch + +from sglang.srt.layers.moe.topk import ( + biased_grouped_topk_impl, + biased_topk_impl, + fused_topk_torch_native, + grouped_topk_gpu, + kimi_k2_biased_topk_impl, +) +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-large") +register_amd_ci(est_time=60, stage="stage-b", runner_config="1-gpu-small-amd") + +torch.manual_seed(1234) + +NUM_EXPERTS = 256 +TOPK = 8 +HIDDEN = 64 + + +@unittest.skipUnless(torch.cuda.is_available(), "needs a GPU") +class TestTopkRenormalizeDegenerate(CustomTestCase): + DEVICE = "cuda" + + def _inputs(self, num_tokens=4, degenerate_rows=(0,), dtype=torch.float32): + """Rows in `degenerate_rows` get all-very-negative logits so every + selected expert's sigmoid weight underflows to exactly zero. In + float16 the sigmoid already underflows below logits of about -18.""" + hidden = torch.randn(num_tokens, HIDDEN, device=self.DEVICE, dtype=dtype) + logits = torch.randn(num_tokens, NUM_EXPERTS, device=self.DEVICE) + for r in degenerate_rows: + logits[r] = -100.0 + torch.rand(NUM_EXPERTS, device=self.DEVICE) + bias = 11.2 + torch.rand(NUM_EXPERTS, device=self.DEVICE) * 0.1 + return hidden, logits.to(dtype), bias.to(dtype) + + def _check(self, topk_weights, name): + self.assertFalse( + torch.isnan(topk_weights).any().item(), + f"{name}: renormalized top-k weights contain NaN", + ) + self.assertTrue( + torch.isfinite(topk_weights).all().item(), + f"{name}: renormalized top-k weights are not finite", + ) + # healthy rows (>=1) must still renormalize to 1 + row_sums = topk_weights.float().sum(dim=-1) + self.assertTrue( + torch.allclose(row_sums[1:], torch.ones_like(row_sums[1:]), atol=1e-3), + f"{name}: healthy rows no longer sum to 1: {row_sums.tolist()}", + ) + + def test_biased_topk_impl(self): + hidden, logits, bias = self._inputs() + weights, _ = biased_topk_impl(hidden, logits, bias, topk=TOPK, renormalize=True) + self._check(weights, "biased_topk_impl") + + def test_biased_topk_impl_fp16(self): + # the 1e-20 epsilon underflows to zero in float16; the renormalization + # must run in float32 for the guard to hold (fp16 sigmoid already + # underflows below logits of about -18) + hidden, logits, bias = self._inputs(dtype=torch.float16) + weights, _ = biased_topk_impl(hidden, logits, bias, topk=TOPK, renormalize=True) + self._check(weights, "biased_topk_impl[fp16]") + + def test_fused_topk_torch_native_sigmoid_bias_fp16(self): + hidden, logits, bias = self._inputs(dtype=torch.float16) + weights, _ = fused_topk_torch_native( + hidden, + logits, + topk=TOPK, + renormalize=True, + correction_bias=bias, + scoring_func="sigmoid", + ) + self._check(weights, "fused_topk_torch_native[fp16]") + + def test_biased_grouped_topk_impl(self): + hidden, logits, bias = self._inputs() + weights, _ = biased_grouped_topk_impl( + hidden, + logits, + bias, + topk=TOPK, + renormalize=True, + num_expert_group=8, + topk_group=4, + ) + self._check(weights, "biased_grouped_topk_impl") + + def test_kimi_k2_biased_topk_impl(self): + hidden, logits, bias = self._inputs() + weights, _ = kimi_k2_biased_topk_impl( + hidden, logits, bias, topk=TOPK, renormalize=True + ) + self._check(weights, "kimi_k2_biased_topk_impl") + + def test_fused_topk_torch_native_sigmoid_bias(self): + hidden, logits, bias = self._inputs() + weights, _ = fused_topk_torch_native( + hidden, + logits, + topk=TOPK, + renormalize=True, + correction_bias=bias, + scoring_func="sigmoid", + ) + self._check(weights, "fused_topk_torch_native") + + def test_grouped_topk_gpu_sigmoid(self): + hidden, logits, _ = self._inputs() + weights, _ = grouped_topk_gpu( + hidden, + logits, + topk=TOPK, + renormalize=True, + num_expert_group=8, + topk_group=4, + scoring_func="sigmoid", + ) + self._check(weights, "grouped_topk_gpu") + + +if __name__ == "__main__": + unittest.main()