From d01df34f35b75c8bf085487107e6a527506a4d7a Mon Sep 17 00:00:00 2001 From: yukun-hh Date: Fri, 2 Oct 2026 15:50:44 +0800 Subject: [PATCH] =?UTF-8?q?topp=E5=92=8Ctopk=E7=AE=97=E5=AD=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- 3(topp)/1.cu | 111 ++++++++++++++++++++++++++++++++++++++ 3(topp)/2.cu | 150 +++++++++++++++++++++++++++++++++++++++++++++++++++ 3(topp)/3.cu | 0 4(topk)/1.cu | 134 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 395 insertions(+) create mode 100644 3(topp)/1.cu create mode 100644 3(topp)/2.cu create mode 100644 3(topp)/3.cu create mode 100644 4(topk)/1.cu diff --git a/3(topp)/1.cu b/3(topp)/1.cu new file mode 100644 index 0000000..9d969c4 --- /dev/null +++ b/3(topp)/1.cu @@ -0,0 +1,111 @@ +#include +#include +#include +namespace cg = cooperative_groups; + +__device__ __inline__ uint64_t make_key(float prob, int64_t col) { + return (((uint64_t)__float_as_uint(prob)) << 32) | + (uint64_t)(0xffffffffu - (uint32_t)col); +} + +__device__ __inline__ float block_reduce_sum(float val) { + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block); + + // 第一步：warp 内归约 + val = cg::reduce(tile, val, cg::plus()); + + __shared__ float warp_sums[32]; + const int warp_id = threadIdx.x >> 5; + const int lane_id = threadIdx.x & 31; + const int num_warps = blockDim.x >> 5; + + if (lane_id == 0) warp_sums[warp_id] = val; + __syncthreads(); + + // 第二步：第一个 warp 做二次归约 + if (warp_id == 0) { + val = (lane_id < num_warps) ? warp_sums[lane_id] : 0.0f; + val = cg::reduce(tile, val, cg::plus()); + if (lane_id == 0) warp_sums[0] = val; + } + __syncthreads(); + + return warp_sums[0]; // 广播给所有线程 +} + +// ============================== 主 kernel ============================== +__global__ void topp(const float* __restrict__ probs, + const float* __restrict__ top_p, + float* __restrict__ renorm_probs, + int64_t batch_size, + int64_t num_classes) { + int tid = threadIdx.x; + + __shared__ uint64_t mask_share; + __shared__ uint64_t current_mask_share; + __shared__ float accept_sum; + + if (tid == 0) { + mask_share = 0ull; + current_mask_share = 0ull; + accept_sum = 0.0f; + } + __syncthreads(); + + const int64_t row_offset = (int64_t)blockIdx.x * num_classes; + const float row_top_p = top_p[blockIdx.x]; + + // ---------- 1. 从高位到低位寻找分界点 ---------- + for (int bit = 63; bit >= 0; --bit) { + const uint64_t current_mask = current_mask_share; + const uint64_t mask = mask_share; + const uint64_t now_flag = 1ull << bit; + + float local_sum = 0.0f; + for (int64_t col = tid; col < num_classes; col += blockDim.x) { + float p = probs[row_offset + col]; + uint64_t key = make_key(p, col); + if (((key & mask) == current_mask) && ((key & now_flag) != 0)) { + local_sum += p; + } + } + const float total_sum = block_reduce_sum(local_sum); + + if (tid == 0) { + if (total_sum > 0.0f && accept_sum + total_sum <= row_top_p) { + accept_sum += total_sum; // 包含该位 + } else if (total_sum > 0.0f) { + current_mask_share |= now_flag; // 分界点在此位 + } + mask_share |= now_flag; + } + __syncthreads(); + } + + // ---------- 2. 重归一化 ---------- + const uint64_t threshold = current_mask_share; + + float final_sum = 0.0f; + for (int64_t col = tid; col < num_classes; col += blockDim.x) { + float p = probs[row_offset + col]; + uint64_t key = make_key(p, col); + if (key >= threshold) final_sum += p; + } + final_sum = block_reduce_sum(final_sum); + + const float inv = (final_sum > 0.0f) ? (1.0f / final_sum) : 0.0f; + for (int64_t col = tid; col < num_classes; col += blockDim.x) { + float p = probs[row_offset + col]; + uint64_t key = make_key(p, col); + renorm_probs[row_offset + col] = (key >= threshold) ? (p * inv) : 0.0f; + } +} + +extern "C" void run_kernel(const float* probs, + const float* top_p, + float* renorm_probs, + int64_t batch_size, + int64_t num_classes) { + topp<<>>(probs, top_p, renorm_probs, batch_size, num_classes); +} \ No newline at end of file diff --git a/3(topp)/2.cu b/3(topp)/2.cu new file mode 100644 index 0000000..c79a6e4 --- /dev/null +++ b/3(topp)/2.cu @@ -0,0 +1,150 @@ +//float4 对齐版本 +#include +#include +#include +namespace cg = cooperative_groups; + +__device__ __inline__ uint64_t make_key(float prob, int64_t col) { + return (((uint64_t)__float_as_uint(prob)) << 32) | + (uint64_t)(0xffffffffu - (uint32_t)col); +} + +__device__ __inline__ bool accept_key(float p, int64_t col, + uint64_t mask, + uint64_t current_mask, + uint64_t now_flag) { + uint64_t key = make_key(p, col); + return ((key & mask) == current_mask) && ((key & now_flag) != 0); +} + +__device__ __inline__ bool keep_key(float p, int64_t col, uint64_t threshold) { + return make_key(p, col) >= threshold; +} + +__device__ __inline__ float block_reduce_sum(float val) { + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block); + + val = cg::reduce(tile, val, cg::plus()); + + __shared__ float warp_sums[32]; + const int warp_id = threadIdx.x >> 5; + const int lane_id = threadIdx.x & 31; + const int num_warps = blockDim.x >> 5; + + if (lane_id == 0) warp_sums[warp_id] = val; + __syncthreads(); + + if (warp_id == 0) { + val = (lane_id < num_warps) ? warp_sums[lane_id] : 0.0f; + val = cg::reduce(tile, val, cg::plus()); + if (lane_id == 0) warp_sums[0] = val; + } + __syncthreads(); + + return warp_sums[0]; +} + +__global__ void topp_float4(const float* __restrict__ probs, + const float* __restrict__ top_p, + float* __restrict__ renorm_probs, + int64_t batch_size, + int64_t num_classes) { + int tid = threadIdx.x; + + __shared__ uint64_t mask_share; + __shared__ uint64_t current_mask_share; + __shared__ float accept_sum; + + if (tid == 0) { + mask_share = 0ull; + current_mask_share = 0ull; + accept_sum = 0.0f; + } + __syncthreads(); + + const int64_t row_offset = (int64_t)blockIdx.x * num_classes; + const float row_top_p = top_p[blockIdx.x]; + const float* row_probs = probs + row_offset; + + // 只有行首 16B 对齐时才启用 float4 + const bool aligned = ((reinterpret_cast(row_probs) & 15) == 0); + const int64_t vec_end = aligned ? (num_classes & ~3LL) : 0; + const int64_t stride4 = (int64_t)blockDim.x * 4; + + // ---------- 1. 从高位到低位寻找分界点 ---------- + for (int bit = 63; bit >= 0; --bit) { + const uint64_t current_mask = current_mask_share; + const uint64_t mask = mask_share; + const uint64_t now_flag = 1ull << bit; + + float local_sum = 0.0f; + + // float4 向量化部分 + for (int64_t col4 = (int64_t)tid * 4; col4 < vec_end; col4 += stride4) { + float4 v = *reinterpret_cast(row_probs + col4); + if (accept_key(v.x, col4 + 0, mask, current_mask, now_flag)) local_sum += v.x; + if (accept_key(v.y, col4 + 1, mask, current_mask, now_flag)) local_sum += v.y; + if (accept_key(v.z, col4 + 2, mask, current_mask, now_flag)) local_sum += v.z; + if (accept_key(v.w, col4 + 3, mask, current_mask, now_flag)) local_sum += v.w; + } + + // 尾部标量处理 + for (int64_t col = vec_end + tid; col < num_classes; col += blockDim.x) { + float p = row_probs[col]; + if (accept_key(p, col, mask, current_mask, now_flag)) { + local_sum += p; + } + } + + const float total_sum = block_reduce_sum(local_sum); + + if (tid == 0) { + if (total_sum > 0.0f && accept_sum + total_sum <= row_top_p) { + accept_sum += total_sum; + } else if (total_sum > 0.0f) { + current_mask_share |= now_flag; + } + mask_share |= now_flag; + } + __syncthreads(); + } + + // ---------- 2. 重归一化 ---------- + const uint64_t threshold = current_mask_share; + + float final_sum = 0.0f; + + // 向量化求 final_sum + for (int64_t col4 = (int64_t)tid * 4; col4 < vec_end; col4 += stride4) { + float4 v = *reinterpret_cast(row_probs + col4); + if (keep_key(v.x, col4 + 0, threshold)) final_sum += v.x; + if (keep_key(v.y, col4 + 1, threshold)) final_sum += v.y; + if (keep_key(v.z, col4 + 2, threshold)) final_sum += v.z; + if (keep_key(v.w, col4 + 3, threshold)) final_sum += v.w; + } + + for (int64_t col = vec_end + tid; col < num_classes; col += blockDim.x) { + float p = row_probs[col]; + if (keep_key(p, col, threshold)) final_sum += p; + } + + final_sum = block_reduce_sum(final_sum); + const float inv = (final_sum > 0.0f) ? (1.0f / final_sum) : 0.0f; + + // 写回用标量，避免处理写对齐 + for (int64_t col = tid; col < num_classes; col += blockDim.x) { + float p = row_probs[col]; + renorm_probs[row_offset + col] = + keep_key(p, col, threshold) ? (p * inv) : 0.0f; + } +} + +extern "C" void run_kernel(const float* probs, + const float* top_p, + float* renorm_probs, + int64_t batch_size, + int64_t num_classes) { + topp_float4<<>>(probs, top_p, renorm_probs, + batch_size, num_classes); +} \ No newline at end of file diff --git a/3(topp)/3.cu b/3(topp)/3.cu new file mode 100644 index 0000000..e69de29 diff --git a/4(topk)/1.cu b/4(topk)/1.cu new file mode 100644 index 0000000..82d158c --- /dev/null +++ b/4(topk)/1.cu @@ -0,0 +1,134 @@ +#include +#include +#include +namespace cg = cooperative_groups; + +__device__ __inline__ uint64_t make_key(float prob, int64_t col) { + return (((uint64_t)__float_as_uint(prob)) << 32) | + (uint64_t)(0xffffffffu - (uint32_t)col); +} + +__device__ __inline__ int block_reduce_sum(int val) { + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block); + + // 第一步：warp 内归约 + val = cg::reduce(tile, val, cg::plus()); + + __shared__ int warp_sums[32]; + const int warp_id = threadIdx.x >> 5; + const int lane_id = threadIdx.x & 31; + const int num_warps = blockDim.x >> 5; + + if (lane_id == 0) warp_sums[warp_id] = val; + __syncthreads(); + + // 第二步：第一个 warp 做二次归约 + if (warp_id == 0) { + val = (lane_id < num_warps) ? warp_sums[lane_id] : 0; + val = cg::reduce(tile, val, cg::plus()); + if (lane_id == 0) warp_sums[0] = val; + } + __syncthreads(); + + return warp_sums[0]; // 广播给所有线程 +} +__device__ __inline__ float block_reduce_sum_float(float val) { + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block); + + val = cg::reduce(tile, val, cg::plus()); + + __shared__ float warp_sums[32]; + const int warp_id = threadIdx.x >> 5; + const int lane_id = threadIdx.x & 31; + const int num_warps = blockDim.x >> 5; + + if (lane_id == 0) warp_sums[warp_id] = val; + __syncthreads(); + + if (warp_id == 0) { + val = (lane_id < num_warps) ? warp_sums[lane_id] : 0.0f; + val = cg::reduce(tile, val, cg::plus()); + if (lane_id == 0) warp_sums[0] = val; + } + __syncthreads(); + + return warp_sums[0]; +} + +// ============================== 主 kernel ============================== +__global__ void topk(const float* __restrict__ probs, + const int32_t* __restrict__ top_k, + float* __restrict__ renorm_probs, + int64_t batch_size, + int64_t num_classes) { + int tid = threadIdx.x; + + __shared__ uint64_t mask_share; + __shared__ uint64_t current_mask_share; + __shared__ int32_t accept_cnt; + + if (tid == 0) { + mask_share = 0ull; + current_mask_share = 0ull; + accept_cnt = 0; + } + __syncthreads(); + + const int64_t row_offset = (int64_t)blockIdx.x * num_classes; + const int row_top_k = top_k[blockIdx.x]; + + // ---------- 1. 从高位到低位寻找分界点 ---------- + for (int bit = 63; bit >= 0; --bit) { + const uint64_t current_mask = current_mask_share; + const uint64_t mask = mask_share; + const uint64_t now_flag = 1ull << bit; + + int local_cnt = 0; + for (int64_t col = tid; col < num_classes; col += blockDim.x) { + float p = probs[row_offset + col]; + uint64_t key = make_key(p, col); + if (((key & mask) == current_mask) && ((key & now_flag) != 0)) { + local_cnt += 1; + } + } + const int total_cnt = block_reduce_sum(local_cnt); + + if (tid == 0) { + if (accept_cnt + total_cnt < row_top_k) { + accept_cnt += total_cnt; // 包含该位 + } else { + current_mask_share |= now_flag; // 分界点在此位 + } + mask_share |= now_flag; + } + __syncthreads(); + } + + // ---------- 2. 重归一化 ---------- + const uint64_t threshold = current_mask_share; + + float final_sum = 0.0f; + for (int64_t col = tid; col < num_classes; col += blockDim.x) { + float p = probs[row_offset + col]; + uint64_t key = make_key(p, col); + if (key >= threshold) final_sum += p; + } + final_sum = block_reduce_sum_float(final_sum); + + const float inv = (final_sum > 0.0f) ? (1.0f / final_sum) : 0.0f; + for (int64_t col = tid; col < num_classes; col += blockDim.x) { + float p = probs[row_offset + col]; + uint64_t key = make_key(p, col); + renorm_probs[row_offset + col] = (key >= threshold) ? (p * inv) : 0.0f; + } +} + +extern "C" void run_kernel(const float* probs, + const int32_t* top_k, + float* renorm_probs, + int64_t batch_size, + int64_t num_classes) { + topk<<>>(probs, top_k, renorm_probs, batch_size, num_classes); +} \ No newline at end of file