sampling: dual pivot rejection sampling algorithm to improve top-p/top-k sampling efficiency

[yzh119]

In our previous sampling algorithms, the rejection sampling algorithm for top-p/top-k sampling is not guaranteed to stop within given number of rounds, and the API would return an success array indicating whether the sampling is successful or not. If not, serving engines will fall back to naive sorting-base sampling algorithm.

This PR improves the rejection sampling algorithm, instead of relying on a single pivot, we propose to use dual pivot, and the bound [low, high] is guaranteed to shrink by half each round. After n rounds, the gap between low and high will be within 2^-n.

Design doc: https://docs.google.com/document/d/1rhdgOM5VawSMAK6jjapFS02-1neGYd8dNazhtmZg7fA/edit?usp=sharing

Breaking Changes

• This PR removes the success return value of all sampling API, which is not compatible with earlier design.
• Instead of passing uniform tensor, we changed to interface to accept torch.Generator (optional, https://pytorch.org/docs/stable/generated/torch.Generator.html), to align with the behavior of torch.
• C++ API and TVM interface would break in this PR, let's fix the behavior later.

Co-authored-by: Shanli Xing [email protected]

[yzh119]

@MasterJH5574 we need to figure out how to get a torch.Generator equivalent in tvm.

We obtain philox seed/offset from the generator:

  uint64_t philox_seed, philox_offset;
  auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
      gen_, at::cuda::detail::getDefaultCUDAGenerator());
  std::lock_guard<std::mutex> lock(gen->mutex_);
  at::PhiloxCudaState rng_engine_inputs = gen->philox_cuda_state(increment_size);
  philox_seed = rng_engine_inputs.seed_.val;
  philox_offset = rng_engine_inputs.offset_.val;

Definition of PhiloxCudaState:

https://github.com/pytorch/pytorch/blob/85467ed063d284fa21a2f1d2adfec8fda544923d/aten/src/ATen/cuda/detail/PhiloxCudaStateRaw.cuh