Hi Thomas,
Thank you for your reply. I have tried looping over the segments, it does the job but not really efficient.
In case some one also interests in this question: at the end I found a relatively more efficient way: construct a matrix to first calculate the denominator of softmax and then do the division element-wise.
def segment_softmax(logits, segment_ids):
    logits_len = len(segment_ids)
    num_segments = max(segment_ids) + 1
    logits_exp = torch.exp(-logits).unsqueeze(1) # e^{-logit} N x 1

    # calculate summation of exponential of logits value for each group
    sparse_index = torch.LongTensor(np.stack([segment_ids, np.arange(logits_len)]))
    sparse_value = torch.ones(logits_len, dtype=torch.float)
    trans_matrix_sparse = torch.sparse.FloatTensor(sparse_index, sparse_value,
                                                   torch.Size([num_segments, logits_len]))
    softmax_den = torch.sparse.mm(trans_matrix_sparse, logits_exp)

    # repeat softmax denominator to have the same length as logits
    sparse_index2 = torch.LongTensor(np.stack([np.arange(logits_len), segment_ids]))
    sparse_value2 = torch.ones(logits_len, dtype=torch.float)
    trans_matrix_sparse2 = torch.sparse.FloatTensor(sparse_index2, sparse_value2,
                                                   torch.Size([logits_len, num_segments])).to(device)
    softmax_den_repeat = torch.sparse.mm(trans_matrix_sparse2, softmax_den)

    return torch.squeeze(logits_exp/softmax_den_repeat)

Best,
Peng