diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/hash_topk.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/hash_topk.cuh index 11f91bbe8..6557802e1 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/hash_topk.cuh +++ b/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(); diff --git a/python/sglang/kernels/jit/csrc/moe/moe_fused_gate.cuh b/python/sglang/kernels/jit/csrc/moe/moe_fused_gate.cuh index 99a9eb5ea..7ea1a83e9 100644 --- a/python/sglang/kernels/jit/csrc/moe/moe_fused_gate.cuh +++ b/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))) - float softplus = log1pf(expf(x)); + // sqrt(softplus(x)); sign folded out because expf overflows above 88.7. + const float softplus = fmaxf(x, 0.0f) + log1pf(expf(-fabsf(x))); return sqrtf(softplus); } } diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index b60ecce2e..624c331b2 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/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. - sp = tl.where(scores > 20.0, scores, tl.log(1.0 + tl.exp(scores))) + # sqrt(softplus(x)). log(1.0 + exp(x)) rounds to 0 below -16.64 and overflows + # 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: diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py index 72f707211..b7befbae9 100644 --- a/python/sglang/srt/layers/moe/hash_topk.py +++ b/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