feat: deterministic sampling (#417)

Our previous sampling kernels relies on cub's BlockScan, which is not
deterministic as reported in NVIDIA/cub#454.

This PR implements the deterministic BlockScan using Belloch scan
algorithm, which is slower than cub but guarantees determinism.

Author: yzh119

Diff for: include/flashinfer/sampling.cuh

@@ -54,26 +54,29 @@ void merge_state_in_place(torch::Tensor v, torch::Tensor s, torch::Tensor v_othe
 
 std::vector<torch::Tensor> merge_states(torch::Tensor v, torch::Tensor s);
 
-torch::Tensor sampling_from_probs(torch::Tensor probs, torch::Tensor uniform_samples);
+torch::Tensor sampling_from_probs(torch::Tensor probs, torch::Tensor uniform_samples,
+                                  bool deterministic);
 
 std::vector<torch::Tensor> top_p_sampling_from_probs(torch::Tensor probs,
-                                                     torch::Tensor uniform_samples, double top_p);
+                                                     torch::Tensor uniform_samples, double top_p,
+                                                     bool deterministic);
 
 std::vector<torch::Tensor> top_k_sampling_from_probs(torch::Tensor probs,
                                                      torch::Tensor uniform_samples,
-                                                     unsigned int top_k);
+                                                     unsigned int top_k, bool deterministic);
 
 std::vector<torch::Tensor> top_k_top_p_sampling_from_probs(torch::Tensor probs,
                                                            torch::Tensor uniform_samples,
-                                                           torch::Tensor top_k,
-                                                           torch::Tensor top_p);
+                                                           torch::Tensor top_k, torch::Tensor top_p,
+                                                           bool deterministic);
 
 torch::Tensor top_p_renorm_prob(torch::Tensor probs, double top_p, double eps);
 
 torch::Tensor top_k_renorm_prob(torch::Tensor probs, unsigned int top_k, double eps);
 
 torch::Tensor chain_speculative_sampling(torch::Tensor draft_probs, torch::Tensor draft_token_ids,
-                                         torch::Tensor uniform_samples, torch::Tensor target_probs);
+                                         torch::Tensor uniform_samples, torch::Tensor target_probs,
+                                         bool deterministic);
 
 torch::Tensor rmsnorm(torch::Tensor x, torch::Tensor w, double eps);
 

Diff for: python/csrc/flashinfer_ops.h

@@ -20,7 +20,8 @@
 
 using namespace flashinfer;
 
-torch::Tensor sampling_from_probs(torch::Tensor probs, torch::Tensor uniform_samples) {
+torch::Tensor sampling_from_probs(torch::Tensor probs, torch::Tensor uniform_samples,
+                                  bool deterministic) {
   CHECK_INPUT(probs);
   CHECK_INPUT(uniform_samples);
   auto device = probs.device();
@@ -36,16 +37,18 @@ torch::Tensor sampling_from_probs(torch::Tensor probs, torch::Tensor uniform_sam
   cudaStream_t torch_current_stream = c10::cuda::getCurrentCUDAStream(device.index());
   auto samples = torch::empty({batch_size}, torch::dtype(torch::kInt32).device(device));
 
-  cudaError_t status = sampling::SamplingFromProb(
-      static_cast<float*>(probs.data_ptr()), static_cast<float*>(uniform_samples.data_ptr()),
-      static_cast<int*>(samples.data_ptr()), batch_size, vocab_size, torch_current_stream);
+  cudaError_t status = sampling::SamplingFromProb(static_cast<float*>(probs.data_ptr()),
+                                                  static_cast<float*>(uniform_samples.data_ptr()),
+                                                  static_cast<int*>(samples.data_ptr()), batch_size,
+                                                  vocab_size, deterministic, torch_current_stream);
   TORCH_CHECK(status == cudaSuccess, "SamplingFromProbs failed with error code " +
                                          std::string(cudaGetErrorString(status)));
   return samples;
 }
 
 std::vector<torch::Tensor> top_p_sampling_from_probs(torch::Tensor probs,
-                                                     torch::Tensor uniform_samples, double top_p) {
+                                                     torch::Tensor uniform_samples, double top_p,
+                                                     bool deterministic) {
   CHECK_INPUT(probs);
   CHECK_INPUT(uniform_samples);
   auto device = probs.device();
@@ -66,7 +69,7 @@ std::vector<torch::Tensor> top_p_sampling_from_probs(torch::Tensor probs,
   cudaError_t status = sampling::TopPSamplingFromProb<float, int>(
       static_cast<float*>(probs.data_ptr()), static_cast<float*>(uniform_samples.data_ptr()),
       static_cast<int*>(samples.data_ptr()), static_cast<bool*>(success.data_ptr()), top_p,
-      batch_size, vocab_size, max_top_p_rounds, torch_current_stream);
+      batch_size, vocab_size, max_top_p_rounds, deterministic, torch_current_stream);
   TORCH_CHECK(status == cudaSuccess, "TopPSamplingFromProbs failed with error code " +
                                          std::string(cudaGetErrorString(status)));
 
@@ -75,7 +78,7 @@ std::vector<torch::Tensor> top_p_sampling_from_probs(torch::Tensor probs,
 
 std::vector<torch::Tensor> top_k_sampling_from_probs(torch::Tensor probs,
                                                      torch::Tensor uniform_samples,
-                                                     unsigned int top_k) {
+                                                     unsigned int top_k, bool deterministic) {
   CHECK_INPUT(probs);
   CHECK_INPUT(uniform_samples);
   auto device = probs.device();
@@ -96,7 +99,7 @@ std::vector<torch::Tensor> top_k_sampling_from_probs(torch::Tensor probs,
   cudaError_t status = sampling::TopKSamplingFromProb<float, int>(
       static_cast<float*>(probs.data_ptr()), static_cast<float*>(uniform_samples.data_ptr()),
       static_cast<int*>(samples.data_ptr()), static_cast<bool*>(success.data_ptr()), top_k,
-      batch_size, vocab_size, max_top_k_rounds, torch_current_stream);
+      batch_size, vocab_size, max_top_k_rounds, deterministic, torch_current_stream);
   TORCH_CHECK(status == cudaSuccess, "TopKSamplingFromProbs failed with error code " +
                                          std::string(cudaGetErrorString(status)));
 
@@ -105,8 +108,8 @@ std::vector<torch::Tensor> top_k_sampling_from_probs(torch::Tensor probs,
 
 std::vector<torch::Tensor> top_k_top_p_sampling_from_probs(torch::Tensor probs,
                                                            torch::Tensor uniform_samples,
-                                                           torch::Tensor top_k,
-                                                           torch::Tensor top_p) {
+                                                           torch::Tensor top_k, torch::Tensor top_p,
+                                                           bool deterministic) {
   CHECK_INPUT(probs);
   CHECK_INPUT(uniform_samples);
   CHECK_INPUT(top_k);
@@ -138,7 +141,7 @@ std::vector<torch::Tensor> top_k_top_p_sampling_from_probs(torch::Tensor probs,
       static_cast<float*>(probs.data_ptr()), static_cast<float*>(uniform_samples.data_ptr()),
       static_cast<int*>(top_k.data_ptr()), static_cast<float*>(top_p.data_ptr()),
       static_cast<int*>(samples.data_ptr()), static_cast<bool*>(success.data_ptr()), batch_size,
-      vocab_size, max_rounds, torch_current_stream);
+      vocab_size, max_rounds, deterministic, torch_current_stream);
   TORCH_CHECK(status == cudaSuccess, "TopKTopPSamplingFromProbs failed with error code " +
                                          std::string(cudaGetErrorString(status)));
 
@@ -187,8 +190,8 @@ torch::Tensor top_k_renorm_prob(torch::Tensor probs, unsigned int top_k, double
 }
 
 torch::Tensor chain_speculative_sampling(torch::Tensor draft_probs, torch::Tensor draft_token_ids,
-                                         torch::Tensor uniform_samples,
-                                         torch::Tensor target_probs) {
+                                         torch::Tensor uniform_samples, torch::Tensor target_probs,
+                                         bool deterministic) {
   CHECK_INPUT(draft_probs);
   CHECK_INPUT(draft_token_ids);
   CHECK_INPUT(uniform_samples);
@@ -224,7 +227,7 @@ torch::Tensor chain_speculative_sampling(torch::Tensor draft_probs, torch::Tenso
       static_cast<float*>(draft_probs.data_ptr()), static_cast<int*>(draft_token_ids.data_ptr()),
       static_cast<float*>(uniform_samples.data_ptr()), static_cast<float*>(target_probs.data_ptr()),
       static_cast<int*>(output_token_ids.data_ptr()), batch_size, num_speculate_tokens, vocab_size,
-      torch_current_stream);
+      deterministic, torch_current_stream);
 
   TORCH_CHECK(status == cudaSuccess, "ChainSpeculativeSampling failed with error code " +
                                          std::string(cudaGetErrorString(status)));

Diff for: python/csrc/sampling.cu

@@ -32,7 +32,7 @@
 
 
 def sampling_from_probs(
-    probs: torch.Tensor, uniform_samples: torch.Tensor
+    probs: torch.Tensor, uniform_samples: torch.Tensor, deterministic: bool = True
 ) -> torch.Tensor:
     r"""Fused GPU kernel for category sampling from probabilities.
 
@@ -43,6 +43,8 @@ def sampling_from_probs(
     uniform_samples: torch.Tensor
         The uniform samples used as needle for sampling, shape ``(batch_size,)``.
         Expected to be uniformly distributed in ``[0, 1)``.
+    deterministic: bool
+        Whether to use deterministic kernel implementation, default is ``True``.
 
     Returns
     -------
@@ -73,11 +75,14 @@ def sampling_from_probs(
     -----
     This function expects float32 inputs, and the output is int32.
     """
-    return _kernels.sampling_from_probs(probs, uniform_samples)
+    return _kernels.sampling_from_probs(probs, uniform_samples, deterministic)
 
 
 def top_p_sampling_from_probs(
-    probs: torch.Tensor, uniform_samples: torch.Tensor, top_p: float
+    probs: torch.Tensor,
+    uniform_samples: torch.Tensor,
+    top_p: float,
+    deterministic: bool = True,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     r"""Fused GPU kernel for top-p sampling (nucleus sampling) from probabilities,
     this operator implements GPU-based rejection sampling without explicit sorting.
@@ -95,6 +100,8 @@ def top_p_sampling_from_probs(
         Expected to be uniformly distributed in ``[0, 1)``.
     top_p: float
         The threshold for top-p sampling.
+    deterministic: bool
+        Whether to use deterministic kernel implementation, default is ``True``.
 
     Returns
     -------
@@ -134,11 +141,16 @@ def top_p_sampling_from_probs(
     We encourage users to set ``max_top_p_rounds`` to a reasonable value, e.g., 32. The actual
     implementation usually use much fewer rounds for rejection sampling because of early stopping.
     """
-    return _kernels.top_p_sampling_from_probs(probs, uniform_samples, top_p)
+    return _kernels.top_p_sampling_from_probs(
+        probs, uniform_samples, top_p, deterministic
+    )
 
 
 def top_k_sampling_from_probs(
-    probs: torch.Tensor, uniform_samples: torch.Tensor, top_k: int
+    probs: torch.Tensor,
+    uniform_samples: torch.Tensor,
+    top_k: int,
+    deterministic: bool = True,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     r"""Fused GPU kernel for top-k sampling from probabilities,
     this operator implements GPU-based rejection sampling without explicit sorting.
@@ -156,6 +168,8 @@ def top_k_sampling_from_probs(
         Expected to be uniformly distributed in ``[0, 1)``.
     top_k: int
         The k in "top-k".
+    deterministic: bool
+        Whether to use deterministic kernel implementation, default is ``True``.
 
     Returns
     -------
@@ -195,14 +209,17 @@ def top_k_sampling_from_probs(
     We encourage users to set ``max_top_k_rounds`` to a reasonable value, e.g., 32. The actual
     implementation usually use much fewer rounds for rejection sampling because of early stopping.
     """
-    return _kernels.top_k_sampling_from_probs(probs, uniform_samples, top_k)
+    return _kernels.top_k_sampling_from_probs(
+        probs, uniform_samples, top_k, deterministic
+    )
 
 
 def top_k_top_p_sampling_from_probs(
     probs: torch.Tensor,
     uniform_samples: torch.Tensor,
     top_k: torch.Tensor,
     top_p: torch.Tensor,
+    deterministic: bool = True,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     r"""Fused GPU kernel for joint top-k and top-p sampling from probabilities,
 
@@ -223,6 +240,8 @@ def top_k_top_p_sampling_from_probs(
         The k in "top-k" for each request, shape ``(batch_size,)``.
     top_p: torch.Tensor
         The threshold for top-p sampling for each request, shape ``(batch_size,)``.
+    deterministic: bool
+        Whether to use deterministic kernel implementation, default is ``True``.
 
     Returns
     -------
@@ -264,7 +283,7 @@ def top_k_top_p_sampling_from_probs(
     implementation usually use much fewer rounds for rejection sampling because of early stopping.
     """
     return _kernels.top_k_top_p_sampling_from_probs(
-        probs, uniform_samples, top_k, top_p
+        probs, uniform_samples, top_k, top_p, deterministic
     )
 
 
@@ -328,6 +347,7 @@ def chain_speculative_sampling(
     draft_token_ids,
     uniform_samples,
     target_probs,
+    deterministic: bool = True,
 ) -> torch.Tensor:
     r"""Fused-GPU kernel for speculative sampling for sequence generation (proposed in
     paper `Accelerating Large Language Model Decoding with Speculative Sampling <https://arxiv.org/pdf/2302.01318>`_),
@@ -349,6 +369,8 @@ def chain_speculative_sampling(
         Compared to input :attr:`draft_probs`, the target model's probability has an additional
         slot at the end because the target model will generate one more token than the draft model.
         Shape: ``(batch_size, num_speculate_tokens + 1, vocab_size)``
+    deterministic: bool
+        Whether to use deterministic kernel implementation, default is ``True``.
 
     Returns
     -------
@@ -361,5 +383,5 @@ def chain_speculative_sampling(
         Shape: (batch_size, num_specutate_tokens + 1)
     """
     return _kernels.chain_speculative_sampling(
-        draft_probs, draft_token_ids, uniform_samples, target_probs
+        draft_probs, draft_token_ids, uniform_samples, target_probs, deterministic
     )

Diff for: python/flashinfer/sampling.py

@@ -26,6 +26,7 @@ template <typename T>
 void bench_sampling_with_probability(nvbench::state& state) {
   size_t batch_size = state.get_int64("batch_size");
   size_t vocab_size = state.get_int64("vocab_size");
+  bool deterministic = state.get_int64("determinisic");
 
   std::vector<T> probs_h(batch_size * vocab_size);
   std::vector<T> uniform_samples_h(batch_size);
@@ -55,7 +56,7 @@ void bench_sampling_with_probability(nvbench::state& state) {
     cudaError_t status = sampling::SamplingFromProb<T>(
         thrust::raw_pointer_cast(probs_d.data()),
         thrust::raw_pointer_cast(uniform_samples_d.data()),
-        thrust::raw_pointer_cast(output_d.data()), batch_size, vocab_size);
+        thrust::raw_pointer_cast(output_d.data()), batch_size, vocab_size, deterministic);
     timer.stop();
     if (status != cudaSuccess) {
       state.skip("CUDA error: " + std::string(cudaGetErrorString(status)));
@@ -67,6 +68,7 @@ template <typename T>
 void bench_top_p_sampling_with_probability(nvbench::state& state) {
   size_t batch_size = state.get_int64("batch_size");
   size_t vocab_size = state.get_int64("vocab_size");
+  bool deterministic = state.get_int64("determinisic");
   double p = state.get_float64("p");
   constexpr uint32_t max_top_p_rounds = 32;
 
@@ -100,7 +102,7 @@ void bench_top_p_sampling_with_probability(nvbench::state& state) {
         thrust::raw_pointer_cast(probs_d.data()),
         thrust::raw_pointer_cast(uniform_samples_d.data()),
         thrust::raw_pointer_cast(output_d.data()), thrust::raw_pointer_cast(success_d.data()), p,
-        batch_size, vocab_size, max_top_p_rounds);
+        batch_size, vocab_size, max_top_p_rounds, deterministic);
     timer.stop();
     if (status != cudaSuccess) {
       state.skip("CUDA error: " + std::string(cudaGetErrorString(status)));
@@ -113,6 +115,7 @@ void bench_top_k_sampling_with_probability(nvbench::state& state) {
   size_t batch_size = state.get_int64("batch_size");
   size_t vocab_size = state.get_int64("vocab_size");
   size_t k = state.get_int64("k");
+  bool deterministic = state.get_int64("determinisic");
   constexpr uint32_t max_top_k_rounds = 32;
 
   std::vector<T> probs_h(batch_size * vocab_size);
@@ -145,7 +148,7 @@ void bench_top_k_sampling_with_probability(nvbench::state& state) {
         thrust::raw_pointer_cast(probs_d.data()),
         thrust::raw_pointer_cast(uniform_samples_d.data()),
         thrust::raw_pointer_cast(output_d.data()), thrust::raw_pointer_cast(success_d.data()), k,
-        batch_size, vocab_size, max_top_k_rounds);
+        batch_size, vocab_size, max_top_k_rounds, deterministic);
     timer.stop();
     if (status != cudaSuccess) {
       state.skip("CUDA error: " + std::string(cudaGetErrorString(status)));
@@ -157,18 +160,21 @@ auto bench_sampling_with_probability_f32 = bench_sampling_with_probability<float
 NVBENCH_BENCH(bench_sampling_with_probability_f32)
     .set_name("bench_sampling_with_probability_f32")
     .add_int64_axis("batch_size", {16, 32, 128, 512, 2048})
-    .add_int64_axis("vocab_size", {32000, 32001, 32002, 128000, 256000});
+    .add_int64_axis("vocab_size", {32000, 32001, 32002, 128000, 256000})
+    .add_int64_axis("determinisic", {0, 1});
 
 auto bench_top_p_sampling_with_probability_f32 = bench_top_p_sampling_with_probability<float>;
 NVBENCH_BENCH(bench_top_p_sampling_with_probability_f32)
     .set_name("bench_top_p_sampling_with_probability_f32")
     .add_int64_axis("batch_size", {16, 32, 128, 512, 2048})
     .add_int64_axis("vocab_size", {32000, 32001, 32002, 128000, 256000})
-    .add_float64_axis("p", {0.1, 0.5, 0.9, 1.0});
+    .add_float64_axis("p", {0.1, 0.5, 0.9, 1.0})
+    .add_int64_axis("determinisic", {0, 1});
 
 auto bench_top_k_sampling_with_probability_f32 = bench_top_k_sampling_with_probability<float>;
 NVBENCH_BENCH(bench_top_k_sampling_with_probability_f32)
     .set_name("bench_top_k_sampling_with_probability_f32")
     .add_int64_axis("batch_size", {16, 32, 128, 512, 2048})
     .add_int64_axis("vocab_size", {32000, 32001, 32002, 128000, 256000})
-    .add_int64_axis("k", {16, 32, 128, 1024});
+    .add_int64_axis("k", {16, 32, 128, 1024})
+    .add_int64_axis("determinisic", {0, 1});