Optimize cub::DeviceSegmented[Radix]Sort for small number of segments

Currently, `cub::DeviceSegmentedRadixSort` launches `num_segments` blocks and each block works on one segment. This approach does not have good performance when the number of segments is small: https://github.com/pytorch/pytorch/issues/63456. For small number of segments, there should be a separate code path for it.