Fix DeepSeek-V4 routing: sqrtsoftplus underflow and unfloored renorm (#34459)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>

python/sglang/kernels/jit/csrc/deepseek_v4/hash_topk.cuh

@@ -12,6 +12,9 @@
 
 namespace sglang {
 
+// Hash routing ignores scores, so an all-underflowed row would renorm 0/0 to NaN.
+constexpr float kRenormalizeSumEpsilon = 1e-20f;
+
 [[maybe_unused]]
 SGL_DEVICE float act_sqrt_softplus(float x) {
   const float softplus = fmaxf(x, 0.0f) + log1pf(expf(-fabsf(x)));
@@ -61,7 +64,8 @@ __global__ void moe_hash_topk_fused(const MoEHashTopKParams __grid_constant__ pa
     const bool is_shared = lane_id >= topk;
     const auto output_offset = warp_id * topk_fused + lane_id;
     topk_ids[output_offset] = is_shared ? num_routed_experts + lane_id - topk : expert_id;
-    topk_weights[output_offset] = is_shared ? 1.0f / routed_scaling_factor : routed_weight / routed_sum;
+    topk_weights[output_offset] =
+        is_shared ? 1.0f / routed_scaling_factor : routed_weight / (routed_sum + kRenormalizeSumEpsilon);
   }
 
   PDLTriggerSecondary<kUsePDL>();

python/sglang/kernels/jit/csrc/moe/moe_fused_gate.cuh

@@ -52,8 +52,8 @@ __device__ __forceinline__ float compute_score(float x) {
     // sigmoid(x) = 1 / (1 + exp(-x))
     return 1.0f / (1.0f + expf(-x));
   } else {
-    // sqrt(softplus(x)) = sqrt(log(1 + exp(x)))
+    // sqrt(softplus(x)); sign folded out because expf overflows above 88.7.
-    float softplus = log1pf(expf(x));
+    const float softplus = fmaxf(x, 0.0f) + log1pf(expf(-fabsf(x)));
     return sqrtf(softplus);
   }
 }

python/sglang/kernels/ops/moe/moe_fused_gate.py

@@ -149,8 +149,13 @@ def _router_triton_kernel(
         activated = tl.sigmoid(scores)
         biased = activated + bias[None, :]
     elif SCORING_FUNC == 1:
-        # sqrt(softplus(x)) = sqrt(log1p(exp(x))); guard against overflow when x is large.
+        # sqrt(softplus(x)). log(1.0 + exp(x)) rounds to 0 below -16.64 and overflows
-        sp = tl.where(scores > 20.0, scores, tl.log(1.0 + tl.exp(scores)))
+        # above 88.7; Triton has no log1p, so recover it from log via z*log(u)/(u-1).
+        z = tl.exp(-tl.abs(scores))
+        u = 1.0 + z
+        exact = u == 1.0
+        log1p_z = tl.where(exact, z, z * tl.log(u) / tl.where(exact, 1.0, u - 1.0))
+        sp = tl.maximum(scores, 0.0) + log1p_z
         activated = tl.sqrt(sp)
         biased = activated + bias[None, :]
     else:

python/sglang/srt/layers/moe/hash_topk.py

@@ -15,6 +15,7 @@ from sglang.srt.eplb.expert_location_dispatch import (
     topk_ids_logical_to_physical,
 )
 from sglang.srt.layers.moe.topk import (
+    _RENORMALIZE_SUM_EPSILON,
     StandardTopKOutput,
     TopKConfig,
     _mask_topk_ids_padded_region,
@@ -157,7 +158,10 @@ class HashTopK(nn.Module):
             topk_weights[:, :-1] = scores.gather(1, topk_ids[:, :-1])
 
             if self.score_func != "softmax":
-                topk_weights[:, :-1] /= topk_weights[:, :-1].sum(dim=-1, keepdim=True)
+                topk_weights[:, :-1] /= (
+                    topk_weights[:, :-1].sum(dim=-1, keepdim=True)
+                    + _RENORMALIZE_SUM_EPSILON
+                )
 
             topk_ids[:, -1] = torch.randint(
                 low=self.num_experts,
@@ -174,7 +178,10 @@ class HashTopK(nn.Module):
             topk_ids[:, :] = self.tid2eid[input_ids]
             topk_weights[:, :] = scores.gather(1, topk_ids[:, :])
             if self.score_func != "softmax":
-                topk_weights[:, :] /= topk_weights[:, :].sum(dim=-1, keepdim=True)
+                topk_weights[:, :] /= (
+                    topk_weights[:, :].sum(dim=-1, keepdim=True)
+                    + _RENORMALIZE_SUM_EPSILON
+                )
 
         return topk_weights, topk_ids