From c022ea5c27fb848490c526c04441a3db33179e7b Mon Sep 17 00:00:00 2001 From: vivienfanghua Date: Mon, 7 Jul 2025 20:20:56 +0800 Subject: [PATCH] draft change rejectionsampling topk=40 --- custom_ops/gpu_ops/sample_kernels/sampling.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/custom_ops/gpu_ops/sample_kernels/sampling.cuh b/custom_ops/gpu_ops/sample_kernels/sampling.cuh index eb5f6f1b84..c4764c00c5 100644 --- a/custom_ops/gpu_ops/sample_kernels/sampling.cuh +++ b/custom_ops/gpu_ops/sample_kernels/sampling.cuh @@ -287,7 +287,7 @@ __global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* output, flo curandStatePhilox4_32_10_t state; curand_init(philox_seed, bx, philox_offset, &state); const uint32_t row_idx = bx; - const uint32_t k = top_p_arr[row_idx] == 0 ? 1 : 20; + const uint32_t k = top_p_arr[row_idx] == 0 ? 1 : 40; const float p = top_p_arr[row_idx] == 0 ? 1e-6 : top_p_arr[row_idx]; extern __shared__ __align__(