diff --git a/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu b/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu index f674e010c..7f719e52e 100644 --- a/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu +++ b/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu @@ -126,6 +126,9 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token( if (expert_id >= 0 && expert_id < NUM_EXPERTS) { output_ptr[row_idx * topk + k] = shared_original_scores[expert_id]; indices_ptr[row_idx * topk + k] = expert_id; + } else { + output_ptr[row_idx * topk + k] = 0.0f; + indices_ptr[row_idx * topk + k] = 0; } } @@ -219,11 +222,16 @@ __global__ void kimi_k2_moe_fused_gate_kernel( } } - if (lane_id == 0 && max_expert != -1) { + if (lane_id == 0) { int64_t output_idx = row_idx * topk + k; - output_ptr[output_idx] = warp_original_scores[max_expert]; - indices_ptr[output_idx] = max_expert; - warp_scores[max_expert] = -FLT_MAX; + if (max_expert != -1) { + output_ptr[output_idx] = warp_original_scores[max_expert]; + indices_ptr[output_idx] = max_expert; + warp_scores[max_expert] = -FLT_MAX; + } else { + output_ptr[output_idx] = 0.0f; + indices_ptr[output_idx] = 0; + } } __syncwarp();