topp和topk算子

3(topp)/1.cu

@@ -0,0 +1,111 @@
+#include <stdint.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+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<float>());
+
+    __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<float>());
+        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<<<batch_size, 1024>>>(probs, top_p, renorm_probs, batch_size, num_classes);
+}