diff --git a/docs/api/python/sampling.rst b/docs/api/python/sampling.rst index 49ee278a..63df15f6 100644 --- a/docs/api/python/sampling.rst +++ b/docs/api/python/sampling.rst @@ -13,6 +13,7 @@ Kernels for LLM sampling. sampling_from_probs top_p_sampling_from_probs top_k_sampling_from_probs + min_p_sampling_from_probs top_k_top_p_sampling_from_probs top_p_renorm_prob top_k_renorm_prob diff --git a/include/flashinfer/sampling.cuh b/include/flashinfer/sampling.cuh index c7a8104f..2d631a48 100644 --- a/include/flashinfer/sampling.cuh +++ b/include/flashinfer/sampling.cuh @@ -85,6 +85,7 @@ struct SamplingTempStorage { union { T value; Pair pair; + T max_p; } block_aggregate; } data; }; @@ -447,6 +448,112 @@ __global__ void TopPSamplingFromProbKernel(DType* probs, DType* uniform_samples, } } +template +__global__ void MinPSamplingFromProbKernel(DType* probs, DType* uniform_samples, DType* min_p, + IdType* output, bool* success, uint32_t d, + uint32_t max_min_p_rounds) { + const uint32_t batch_size = gridDim.x; + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + DType p = min_p[bx]; + + extern __shared__ __align__( + alignof(SamplingTempStorage)) + uint8_t smem_sampling[]; + auto& temp_storage = reinterpret_cast< + SamplingTempStorage&>(smem_sampling); + + vec_t probs_vec; + DType aggregate; + DType q = DType(1); + DType pivot = DType(0); + + DType max_p = 0; + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(DType(0)); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.load(probs + bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + DType probs_[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_[j] = probs_vec[j]; + } + max_p = max(max_p, BlockReduce(temp_storage.block_prim.reduce) + .Reduce(probs_, cub::Max())); + __syncthreads(); + } + if (tx == 0) { + temp_storage.data.block_aggregate.max_p = max_p; + } + __syncthreads(); + DType scaled_p = temp_storage.data.block_aggregate.max_p * p; + + IdType sampled_id; + for (uint32_t round = 0; round < max_min_p_rounds; ++round) { + temp_storage.data.sampled_id = d - 1; + __syncthreads(); + DType u = uniform_samples[round * batch_size + bx] * q; + aggregate = DType(0); + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(DType(0)); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.load(probs + bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + DeviceSamplingFromProb(i, d, pivot, u, probs_vec, aggregate, + &temp_storage); + if (aggregate > u) { + break; + } + } + __syncthreads(); + sampled_id = temp_storage.data.sampled_id; + pivot = max(pivot, probs[bx * d + sampled_id]); + if (pivot >= scaled_p) { + break; + } + + DType aggregate_gt_pivot = DType(0); + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(DType(0)); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.load(probs + bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + DType probs_gt_pivot[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot[j] = (probs_vec[j] > pivot) ? probs_vec[j] : DType(0); + } + + aggregate_gt_pivot += BlockReduce(temp_storage.block_prim.reduce) + .Sum(probs_gt_pivot); + if (tx == 0) { + temp_storage.data.block_aggregate.value = aggregate_gt_pivot; + } + __syncthreads(); + } + q = temp_storage.data.block_aggregate.value; + } + __syncthreads(); + if (tx == 0) { + if (pivot < scaled_p) { + // failed to sample within MAX_ROUNDS + if (success != nullptr) { + success[bx] = false; + } + } else { + output[bx] = sampled_id; + if (success != nullptr) { + success[bx] = true; + } + } + } +} + template @@ -636,6 +743,30 @@ cudaError_t TopPSamplingFromProb(T* probs, T* uniform_samples, IdType* output, b return cudaSuccess; } +template +cudaError_t MinPSamplingFromProb(T* probs, T* uniform_samples, T* min_p, IdType* output, + bool* success, uint32_t batch_size, uint32_t d, + uint32_t max_rounds, bool deterministic, cudaStream_t stream = 0) { + constexpr uint32_t BLOCK_THREADS = 1024; + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + const uint32_t smem_size = sizeof(SamplingTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&probs, &uniform_samples, &min_p, &output, &success, &d, &max_rounds}; + + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = MinPSamplingFromProbKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + return cudaSuccess; +} + template cudaError_t TopKTopPSamplingFromProb(T* probs, T* uniform_samples, IdType* top_k, T* top_p, IdType* output, bool* success, uint32_t batch_size, uint32_t d, diff --git a/python/csrc/flashinfer_ops.cu b/python/csrc/flashinfer_ops.cu index 51dbb2b8..16800f13 100644 --- a/python/csrc/flashinfer_ops.cu +++ b/python/csrc/flashinfer_ops.cu @@ -26,6 +26,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("sampling_from_probs", &sampling_from_probs, "Sample from probabilities"); m.def("top_k_sampling_from_probs", &top_k_sampling_from_probs, "Top-k sampling from probabilities"); + m.def("min_p_sampling_from_probs", &min_p_sampling_from_probs, + "Min-p sampling from probabilities"); m.def("top_p_sampling_from_probs", &top_p_sampling_from_probs, "Top-p sampling from probabilities"); m.def("top_k_top_p_sampling_from_probs", &top_k_top_p_sampling_from_probs, diff --git a/python/csrc/flashinfer_ops.h b/python/csrc/flashinfer_ops.h index e3edff64..0eee5495 100644 --- a/python/csrc/flashinfer_ops.h +++ b/python/csrc/flashinfer_ops.h @@ -46,6 +46,10 @@ std::vector top_k_sampling_from_probs(torch::Tensor probs, torch::Tensor uniform_samples, unsigned int top_k, bool deterministic); +std::vector min_p_sampling_from_probs(torch::Tensor probs, + torch::Tensor uniform_samples, + torch::Tensor min_p, bool deterministic); + std::vector top_k_top_p_sampling_from_probs(torch::Tensor probs, torch::Tensor uniform_samples, torch::Tensor top_k, torch::Tensor top_p, diff --git a/python/csrc/sampling.cu b/python/csrc/sampling.cu index 62828966..2be782cb 100644 --- a/python/csrc/sampling.cu +++ b/python/csrc/sampling.cu @@ -106,6 +106,42 @@ std::vector top_k_sampling_from_probs(torch::Tensor probs, return {samples, success}; } +std::vector min_p_sampling_from_probs(torch::Tensor probs, + torch::Tensor uniform_samples, + torch::Tensor min_p, bool deterministic) { + CHECK_INPUT(probs); + CHECK_INPUT(uniform_samples); + CHECK_INPUT(min_p); + auto device = probs.device(); + CHECK_EQ(uniform_samples.device(), device); + CHECK_EQ(min_p.device(), device); + CHECK_DIM(2, probs); // probs: (batch_size, vocab_size) + CHECK_DIM(2, uniform_samples); // uniform_samples: (max_rounds, batch_size) + CHECK_DIM(1, min_p); // min_p: (batch_size,) + unsigned int batch_size = probs.size(0); + unsigned int vocab_size = probs.size(1); + unsigned int max_rounds = uniform_samples.size(0); + CHECK_EQ(uniform_samples.size(1), batch_size); + CHECK_EQ(min_p.size(0), batch_size); + probs = probs.to(torch::kFloat32); + uniform_samples = uniform_samples.to(torch::kFloat32); + min_p = min_p.to(torch::kFloat32); + + cudaStream_t torch_current_stream = c10::cuda::getCurrentCUDAStream(device.index()); + auto samples = torch::empty({batch_size}, torch::dtype(torch::kInt32).device(device)); + auto success = torch::empty({batch_size}, torch::dtype(torch::kBool).device(device)); + + cudaError_t status = sampling::MinPSamplingFromProb( + static_cast(probs.data_ptr()), static_cast(uniform_samples.data_ptr()), + static_cast(min_p.data_ptr()), static_cast(samples.data_ptr()), + static_cast(success.data_ptr()), batch_size, vocab_size, max_rounds, deterministic, + torch_current_stream); + TORCH_CHECK(status == cudaSuccess, "MinPSamplingFromProb failed with error code " + + std::string(cudaGetErrorString(status))); + + return {samples, success}; +} + std::vector top_k_top_p_sampling_from_probs(torch::Tensor probs, torch::Tensor uniform_samples, torch::Tensor top_k, torch::Tensor top_p, diff --git a/python/flashinfer/sampling.py b/python/flashinfer/sampling.py index 528367a9..8d54ad7c 100644 --- a/python/flashinfer/sampling.py +++ b/python/flashinfer/sampling.py @@ -214,6 +214,76 @@ def top_k_sampling_from_probs( ) +def min_p_sampling_from_probs( + probs: torch.Tensor, + uniform_samples: torch.Tensor, + min_p: torch.Tensor, + deterministic: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + r"""Fused GPU kernel for `min_p sampling `_ from probabilities, + + this operator implements GPU-based rejection sampling without explicit sorting. + + The multiple rounds of rejection sampling are implemented in a single CUDA kernel, + which is more efficient than the naive implementation that launches a series of kernels. + + Parameters + ---------- + probs: torch.Tensor + Probabilities, shape ``(batch_size, num_classes)``. + uniform_samples: torch.Tensor + The uniform samples used as needle for sampling, shape ``(max_top_k_rounds, batch_size,)``, + where the first dimension is the maximum number of rounds for rejection sampling. + Expected to be uniformly distributed in ``[0, 1)``. + min_p: torch.Tensor + The :math:`p_{\text{base}}` in min_p sampling for each request, shape ``(batch_size,)``. + deterministic: bool + Whether to use deterministic kernel implementation, default is ``True``. + + Returns + ------- + samples: torch.Tensor + Sampled categories, shape ``(batch_size,)``. + success: torch.Tensor + Whether the sampling is successful within ``max_top_k_rounds`` rounds, + shape ``(batch_size,)``. + + Examples + -------- + + >>> import torch + >>> import flashinfer + >>> torch.manual_seed(42) + + >>> batch_size = 4 + >>> vocab_size = 5 + >>> max_rounds = 3 + >>> min_p = torch.full((batch_size,), 0.05).to(0) + >>> pre_norm_prob = torch.rand(batch_size, vocab_size).to(0) + >>> norm_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) + >>> norm_prob + tensor([[0.2499, 0.2592, 0.1085, 0.2718, 0.1106], + [0.2205, 0.0942, 0.2912, 0.3452, 0.0489], + [0.2522, 0.1602, 0.2346, 0.1532, 0.2000], + [0.1543, 0.3182, 0.2062, 0.0958, 0.2255]], device='cuda:0') + >>> uniform_samples = torch.rand(max_rounds, batch_size).to(0) + >>> samples, success = flashinfer.sampling.min_p_sampling_from_probs(norm_prob, uniform_samples, min_p) + >>> samples + tensor([1, 2, 1, 4], device='cuda:0', dtype=torch.int32) + >>> success + tensor([True, True, True, True], device='cuda:0') + + Notes + ----- + This function expects float32 inputs, and the output is int32. + We encourage users to set ``max_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.min_p_sampling_from_probs( + probs, uniform_samples, min_p, deterministic + ) + + def top_k_top_p_sampling_from_probs( probs: torch.Tensor, uniform_samples: torch.Tensor, diff --git a/python/tests/test_sampling.py b/python/tests/test_sampling.py index 53fd3c40..26c2395f 100644 --- a/python/tests/test_sampling.py +++ b/python/tests/test_sampling.py @@ -95,6 +95,40 @@ def test_top_k_sampling(batch_size, vocab_size, k): ] +@pytest.mark.parametrize("batch_size", [1, 19, 99, 989]) +@pytest.mark.parametrize("vocab_size", [111, 500, 32000, 128256]) +@pytest.mark.parametrize("p", [0.05, 0.1, 0.2, 0.7, 1]) +def test_min_p_sampling(batch_size, vocab_size, p): + torch.manual_seed(42) + max_min_p_trails = 32 + pre_norm_prob = torch.rand(batch_size, vocab_size).to(0) + normalized_prob = pre_norm_prob / pre_norm_prob.sum(dim=-1, keepdim=True) + sorted_prob, indices = torch.sort(normalized_prob, descending=False) + # scale min-p + top_probs = sorted_prob[:, -1].unsqueeze(-1) + scaled_p = p * top_probs + # min-p mask + mask = torch.zeros(batch_size, vocab_size, dtype=torch.int32).to(0) + mask.scatter_add_(1, indices, (sorted_prob >= scaled_p).int()) + + uniform_samples = torch.empty(max_min_p_trails, batch_size, dtype=torch.float32).to( + 0 + ) + min_p_tensor = torch.full((batch_size,), p).to(0) + + num_trails = 1000 + for _ in range(num_trails): + uniform_samples.uniform_() + samples, success = flashinfer.sampling.min_p_sampling_from_probs( + normalized_prob, uniform_samples, min_p_tensor + ) + assert torch.all(success) + assert torch.all(samples < vocab_size) and torch.all(samples >= 0) + assert torch.all(mask[torch.arange(batch_size), samples] == 1), normalized_prob[ + torch.arange(batch_size), samples + ] + + @pytest.mark.parametrize("batch_size", [1, 19, 99, 989]) @pytest.mark.parametrize("vocab_size", [111, 500, 32000, 128256]) @pytest.mark.parametrize("p", [0.1, 0.5])