CUDA 算子笔记 4
§ 参考资料 共 2 条
Decaying Causal Attention
输入 $\boldsymbol Q\in\mathbb R^{M\times d}$、$\boldsymbol K\in\mathbb R^{M\times d}$、$\boldsymbol V\in\mathbb R^{M\times d}$ 和 $\gamma\in(0,1]$，输出：
$$ \boldsymbol O_n=\sum_{m=0}^{n}\gamma^{n-m}\cdot\frac{\boldsymbol Q_n\cdot\boldsymbol K_m}{\sqrt d}\cdot\boldsymbol V_m $$其中，$\boldsymbol O_i$、$\boldsymbol Q_i$、$\boldsymbol K_i$ 和 $\boldsymbol V_i$ 分别表示输出和输入矩阵的第 $i$ 行
例子：
Input:
Q (2x4): [[1.0, 1.0, 0.0, 0.0],
[1.0, 1.0, 0.0, 0.0]]
K (2x4): [[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0]]
V (2x4): [[4.0, 8.0, 12.0, 16.0],
[4.0, 8.0, 12.0, 16.0]]
Output:
O (2x4): [[2.0, 4.0, 6.0, 8.0],
[3.0, 6.0, 9.0, 12.0]]
数据范围：$1\le M\le 8192$，$1\le d\le 256$
最简单的做法
其实这就是去掉了 Softmat 的另一种形式的更简单的 Attention，改造一下 Softmat Attention 的代码即可
Softmax Attention 参考 CUDA 算子笔记 2
代码：
#include <cuda_runtime.h>
#define CEIL(x, y) (((x) + (y) - 1) / (y))
template <int BX, int BY, int KS, int TM, int TN, bool Transpose, bool Norm,
bool Decay>
__global__ void multiple_kernel(const float* __restrict__ A,
const float* __restrict__ B,
float* __restrict__ output, int M, int N, int K,
float scale = 1.0f, float log_gamma = 0.0f) {
constexpr int BM = BY * TM;
constexpr int BN = BX * TN;
constexpr int T = BX * BY;
__shared__ float tile_A[BM][KS + 1];
__shared__ float tile_B[KS][BN + 1];
const int tid = threadIdx.y * BX + threadIdx.x;
const int block_row = blockIdx.y * BM;
const int block_col = blockIdx.x * BN;
const int thread_row = threadIdx.y * TM;
const int thread_col = threadIdx.x * TN;
float acc[TM][TN] = {};
#pragma unroll
for (int k0 = 0; k0 < K; k0 += KS) {
#pragma unroll
for (int i = 0; i < BM * KS; i += T) {
const int local_row = (tid + i) / KS;
const int local_col = (tid + i) % KS;
const int global_row = block_row + local_row;
const int global_col = k0 + local_col;
if (global_row < M && global_col < K) {
tile_A[local_row][local_col] = A[global_row * K + global_col];
} else {
tile_A[local_row][local_col] = 0.0f;
}
}
if constexpr (Transpose) {
#pragma unroll
for (int i = 0; i < KS * BN; i += T) {
const int local_row = (tid + i) / KS;
const int local_col = (tid + i) % KS;
const int global_row = block_col + local_row;
const int global_col = k0 + local_col;
if (global_row < N && global_col < K) {
tile_B[local_col][local_row] = B[global_row * K + global_col];
} else {
tile_B[local_col][local_row] = 0.0f;
}
}
} else {
#pragma unroll
for (int i = 0; i < KS * BN; i += T) {
const int local_row = (tid + i) / BN;
const int local_col = (tid + i) % BN;
const int global_row = k0 + local_row;
const int global_col = block_col + local_col;
if (global_row < K && global_col < N) {
tile_B[local_row][local_col] = B[global_row * N + global_col];
} else {
tile_B[local_row][local_col] = 0.0f;
}
}
}
__syncthreads();
#pragma unroll
for (int k = 0; k < KS; ++k) {
float a_reg[TM];
float b_reg[TN];
#pragma unroll
for (int i = 0; i < TM; ++i) {
a_reg[i] = tile_A[thread_row + i][k];
}
#pragma unroll
for (int j = 0; j < TN; ++j) {
b_reg[j] = tile_B[k][thread_col + j];
}
#pragma unroll
for (int i = 0; i < TM; ++i) {
#pragma unroll
for (int j = 0; j < TN; ++j) {
acc[i][j] = fmaf(a_reg[i], b_reg[j], acc[i][j]);
}
}
}
__syncthreads();
}
#pragma unroll
for (int i = 0; i < TM; ++i) {
const int row = block_row + thread_row + i;
#pragma unroll
for (int j = 0; j < TN; ++j) {
const int col = block_col + thread_col + j;
if (row < M && col < N) {
float value = acc[i][j];
if constexpr (Norm) {
value *= scale;
}
if constexpr (Decay) {
if (col <= row) {
value *= __expf((row - col) * log_gamma);
} else {
value = 0.0f;
}
}
output[row * N + col] = value;
}
}
}
}
// Q, K, V, output are device pointers
extern "C" void solve(const float* Q, const float* K, const float* V,
float* output, int seq_len, int d_model, float gamma) {
const int M = seq_len, N = seq_len, d = d_model;
float* attention;
cudaMalloc(&attention, M * N * sizeof(float));
{
constexpr int BX = 16, BY = 16;
constexpr int TM = 4, TN = 4;
constexpr int BM = BY * TM, BN = BX * TN;
constexpr int KS = 32;
const dim3 block_dim(BX, BY);
const dim3 grid_dim(CEIL(N, BN), CEIL(M, BM));
multiple_kernel<BX, BY, KS, TM, TN, true, true, true>
<<<grid_dim, block_dim>>>(Q, K, attention, M, N, d, 1.0f / sqrtf(d),
logf(gamma));
}
{
constexpr int BX = 16, BY = 16;
constexpr int TM = 4, TN = 4;
constexpr int BM = BY * TM, BN = BX * TN;
constexpr int KS = 32;
const dim3 block_dim(BX, BY);
const dim3 grid_dim(CEIL(d, BN), CEIL(M, BM));
multiple_kernel<BX, BY, KS, TM, TN, false, false, false>
<<<grid_dim, block_dim>>>(attention, V, output, M, d, N);
}
cudaFree(attention);
}
基于递推的方法
相比于 Softmax Attention，这个 Attention 方法去掉了 Softmax，所以可以做这样的变换：
$$ \begin{aligned} \boldsymbol O_{n,j}&=\sum_{m=0}^{n}\gamma^{n-m}\cdot\frac{\boldsymbol Q_n\cdot\boldsymbol K_m}{\sqrt d}\cdot\boldsymbol V_{m,j} \\ &=\frac{1}{\sqrt d}\sum_{m=0}^{n}\gamma^{n-m}\left(\sum_{i=0}^d\boldsymbol Q_{n,i}\boldsymbol K_{m,i}\right)\boldsymbol V_{m,j} \\ &=\frac{1}{\sqrt d}\sum_{i=0}^{d}\boldsymbol Q_{n,i}\left(\sum_{m=0}^{n}\gamma^{n-m}\boldsymbol K_{m,i}\boldsymbol V_{m,j}\right) \end{aligned} $$即调换内外两个求和符号
然后定义：
$$ \boldsymbol S_n[i,j]=\sum_{m=0}^n\gamma^{n-m}\boldsymbol K_{m,i}\boldsymbol V_{m,j} $$写成矩阵形式：
$$ \boldsymbol S_n=\sum_{m=0}^n\gamma^{n-m}\boldsymbol K_m^T\boldsymbol V_m $$有递推公式：
$$ \boldsymbol S_n=\gamma\boldsymbol S_{n-1}+\boldsymbol K_n^T\boldsymbol V_n $$那么：
$$ \boldsymbol O_{n,j}=\frac{1}{\sqrt d}\sum_{i=0}^{d}\boldsymbol Q_{n,i}\boldsymbol S_n[i,j] $$相当于 $\boldsymbol Q$ 的第 $n$ 行与 $\boldsymbol S_n$ 的第 $j$ 列相乘，即一个向量乘上一个矩阵：
$$ \boldsymbol O_n=\frac{\boldsymbol Q_n\boldsymbol S_n}{\sqrt d} $$在这里，$\boldsymbol Q_n\in\mathbb R^{1\times d}$、$\boldsymbol S_n\in\mathbb R^{d\times d}$，根本不需要中间的 $M\times M$ 的 Attention，直接：
S = 0
for n:
S = gamma * S + outer(K[n], V[n])
O[n] = Q[n] @ S / sqrt(d)
就可以了，这样把 $O(s^2d)$ 的 Attention 改写成一个 $d\times d$ 状态的递推
这个递推可以算是 Prefix Scan，代码：
#include <cuda_runtime.h>
#define CEIL(x, y) (((x) + (y) - 1) / (y))
constexpr int k_max_d = 256;
constexpr int k_threads_per_warp = 32;
// 每个 d*d 的 state 由多个 warp 共同完成，每个 warp 完成 4
// 列（一次完整的矩阵乘法运算需要第二个矩阵的完整一列，所以这里一个 warp
// 算完整一列）
// 每个 thread 分到 8*4
constexpr int k_cols_per_warp = 4;
constexpr int k_rows_per_thread = k_max_d / k_threads_per_warp;
// 每个 block 有 4 个 warp
constexpr int k_warps_per_block = 4;
constexpr int k_threads_per_block = k_threads_per_warp * k_warps_per_block;
constexpr int k_cols_per_block = k_warps_per_block * k_cols_per_warp;
// 每个 block 推进 256 个 state step
constexpr int k_chs_per_block = 256;
__device__ __forceinline__ float warp_reduce(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
val += __shfl_down_sync(0xffffffffu, val, offset);
}
return val;
}
__global__ void decay_atten_k1(const float* __restrict__ Q,
const float* __restrict__ K,
const float* __restrict__ V,
float* __restrict__ output,
float* __restrict__ d_block_states, int M, int D,
float gamma, float scale) {
const int warp = threadIdx.x >> 5;
const int lane = threadIdx.x & 31;
const int ch_base = blockIdx.y * k_chs_per_block;
const int col_base = blockIdx.x * k_cols_per_block + warp * k_cols_per_warp;
float reg[k_rows_per_thread][k_cols_per_warp];
#pragma unroll
for (int u = 0; u < k_rows_per_thread; ++u) {
#pragma unroll
for (int j = 0; j < k_cols_per_warp; ++j) {
reg[u][j] = 0.0f;
}
}
for (int t = 0; t < k_chs_per_block; ++t) {
const int ch = t + ch_base;
if (ch < M) {
float vv[k_cols_per_warp];
float dot[k_cols_per_warp];
#pragma unroll
for (int j = 0; j < k_cols_per_warp; ++j) {
const int col = col_base + j;
vv[j] = col < D ? V[ch * D + col] : 0.0f;
dot[j] = 0.0f;
}
#pragma unroll
for (int i = 0; i < k_rows_per_thread; ++i) {
const int row = lane + i * k_threads_per_warp;
if (row < D) {
float kv = K[ch * D + row];
float qv = Q[ch * D + row];
#pragma unroll
for (int j = 0; j < k_cols_per_warp; ++j) {
const int col = col_base + j;
if (col < D) {
reg[i][j] = fmaf(gamma, reg[i][j], kv * vv[j]);
dot[j] = fmaf(qv, reg[i][j], dot[j]);
}
}
}
}
#pragma unroll
for (int j = 0; j < k_cols_per_warp; j++) {
dot[j] = warp_reduce(dot[j]);
}
if (lane == 0) {
#pragma unroll
for (int j = 0; j < k_cols_per_warp; j++) {
int col = col_base + j;
if (col < D) {
output[ch * D + col] = dot[j] * scale;
}
}
}
}
}
#pragma unroll
for (int i = 0; i < k_rows_per_thread; ++i) {
const int row = lane + i * k_threads_per_warp;
if (row < D) {
#pragma unroll
for (int j = 0; j < k_cols_per_warp; ++j) {
const int col = col_base + j;
if (col < D) {
d_block_states[(blockIdx.y * D + col) * D + row] = reg[i][j];
}
}
}
}
}
__global__ void decay_atten_k2(float* __restrict__ d_block_states,
int num_chunks, int M, int D, float gamma_chunk,
float gamma) {
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
const int num_items = D * D;
if (tid >= num_items) {
return;
}
float carry = 0.0f;
for (int i = 0; i < num_chunks; ++i) {
const int idx = i * num_items + tid;
const float val = d_block_states[idx];
d_block_states[idx] = carry;
const int len = min(k_chs_per_block, M - i * k_chs_per_block);
const float decay =
(len == k_chs_per_block) ? gamma_chunk : __powf(gamma, len);
carry = fmaf(decay, carry, val);
}
}
__global__ void decay_atten_k3(const float* __restrict__ Q,
float* __restrict__ output,
const float* __restrict__ d_block_states, int M,
int D, float gamma, float scale) {
const int warp = threadIdx.x >> 5;
const int lane = threadIdx.x & 31;
const int chunk = blockIdx.y + 1;
const int ch_base = chunk * k_chs_per_block;
const int col_base = blockIdx.x * k_cols_per_block + warp * k_cols_per_warp;
float reg[k_rows_per_thread][k_cols_per_warp];
#pragma unroll
for (int i = 0; i < k_rows_per_thread; ++i) {
const int row = lane + i * k_threads_per_warp;
#pragma unroll
for (int j = 0; j < k_cols_per_warp; ++j) {
const int col = col_base + j;
if (row < D && col < D) {
reg[i][j] = d_block_states[(chunk * D + col) * D + row];
} else {
reg[i][j] = 0.0f;
}
}
}
float decay = 1.0f;
for (int t = 0; t < k_chs_per_block; ++t) {
const int ch = t + ch_base;
if (ch < M) {
decay *= gamma;
float dot[k_cols_per_warp] = {};
#pragma unroll
for (int u = 0; u < k_rows_per_thread; ++u) {
const int row = lane + u * k_threads_per_warp;
if (row < D) {
const float qv = Q[ch * D + row];
#pragma unroll
for (int j = 0; j < k_cols_per_warp; ++j) {
const int col = col_base + j;
if (col < D) {
dot[j] = fmaf(qv, reg[u][j], dot[j]);
}
}
}
}
#pragma unroll
for (int j = 0; j < k_cols_per_warp; ++j) {
dot[j] = warp_reduce(dot[j]);
}
if (lane == 0) {
#pragma unroll
for (int j = 0; j < k_cols_per_warp; ++j) {
const int col = col_base + j;
if (col < D) {
const int idx = ch * D + col;
output[idx] = fmaf(dot[j], scale * decay, output[idx]);
}
}
}
}
}
}
// Q, K, V, output are device pointers
extern "C" void solve(const float* Q, const float* K, const float* V,
float* output, int seq_len, int d_model, float gamma) {
const int M = seq_len;
const int D = d_model;
const int num_chunks = CEIL(M, k_chs_per_block);
float* d_block_states = nullptr;
cudaMalloc(&d_block_states, num_chunks * D * D * sizeof(float));
{
constexpr int block_dim = k_threads_per_block;
const dim3 grid_dim(CEIL(D, k_cols_per_block), num_chunks);
decay_atten_k1<<<grid_dim, block_dim>>>(Q, K, V, output, d_block_states, M,
D, gamma, 1.0f / sqrtf(D));
}
{
constexpr int block_dim = 256;
const int grid_dim = CEIL(D * D, block_dim);
const float gamma_chunk = powf(gamma, k_chs_per_block);
decay_atten_k2<<<grid_dim, block_dim>>>(d_block_states, num_chunks, M, D,
gamma_chunk, gamma);
}
{
constexpr int block_dim = k_threads_per_block;
const dim3 grid_dim(CEIL(D, k_cols_per_block), num_chunks - 1);
decay_atten_k3<<<grid_dim, block_dim>>>(Q, output, d_block_states, M, D,
gamma, 1.0f / sqrtf(D));
}
cudaFree(d_block_states);
}
把 $d\times d$ 的 State 分到若干个 Warp 里面，每个 Warp 计算 State 的一个完整列（这里矩阵乘法的第二个矩阵需要完整的列，所以要这样，否则计算一个完整的值需要跨 Warp 甚至跨 Block 求和，这里直接 Warp Reduce 即可）
- Kernel 1 每 256 步计算一个 State，并且计算 Local 结果
- Kernel 2 对 Chunk State 计算前缀和
- Kernel 3 计算 Global 结果
Parallel Merge
输入两个非下降序列 A 和 B，长度分别为 M 和 N，把它们合并为一个非下降序列 C
例子：
Input:
A = [-1.0, 1.0, 3.0], M = 3
B = [2.0, 5.0], N = 2
Output:
C = [-1.0, 1.0, 2.0, 3.0, 5.0]
基于 Merge Path 的方法
CPU 上的 Merge 是这样的：
i = 0
j = 0
t = 0
while i < M and j < N:
if A[i] <= B[j]:
C[t++] = A[i++]
else:
C[t++] = B[j++]
现在要把它搬到 GPU 上，假设让每个线程处理 VT 个元素，只要我们得到了初始的 i 和 j 就可以采用上面的算法来搞了
如果规定 A[i] == B[j] 时取 A[i]。就转换成了这样的问题：已经完成了 diag 个数的处理，此时满足 i + j = diag，且 A[i-1] <= B[j]、B[j-1] < A[i] 的 i 是多少
首先，这个值是唯一的。如果有两个分界点 (i1, j1) 和 (i2, j2)（i1 < i2）都满足 i + j = diag 的条件，那么有 j1 > j2。所以 B[j1-1] < A[i1] <= A[i2-1] <= B[j2]，与 B[j1-1] >= B[j2] 矛盾
然后，它有一个单调的 Predicate P(i) = (B[j-1] >= A[i])，随着 i 的增大会逐渐由 true 变成 false，我们要找的是第一个 P(i) = false 的点
所以只需要一个二分就可以找到这个分界点，这个方法叫做 Merge Path
代码：
#include <cuda_runtime.h>
// 最终输出完前 diag 个元素时，A 指针的位置（相等时优先取 A）
// 即在 i + j = diag 条件下找到满足 A[i-1] <= B[j] & B[j-1] < A[i] 的点
//
// 这个点唯一：
// 如果有两个分界点 (i1,j1) 和 (i2,j2) 都满足条件（i1 < i2），那么 j1 > j2
// 所以 B[j1-1] < A[i1] <= A[i2-1] <= B[j2]，与 B[j1-1] >= B[j2] 矛盾
//
// 满足单调性：
// P(i) = (B[j-1] >= A[i]) 随着 i 的增大会逐渐由 true 变成 false，分界点是第一个
// P(i) = false 的点
__device__ __forceinline__ int merge_path(const float* __restrict__ A,
const float* __restrict__ B, int M,
int N, int diag) {
int l = max(0, diag - N);
int r = min(diag, M);
while (l < r) {
const int i = (l + r) >> 1;
const int j = diag - i;
if (B[j - 1] >= A[i]) {
l = i + 1;
} else {
r = i;
}
}
return l;
}
template <int VT>
__global__ void merge_kernel(const float* __restrict__ A,
const float* __restrict__ B, float* __restrict__ C,
int M, int N) {
const int diag = (blockIdx.x * blockDim.x + threadIdx.x) * VT;
if (diag >= M + N) {
return;
}
int i = merge_path(A, B, M, N, diag);
int j = diag - i;
#pragma unroll
for (int k = 0; k < VT; ++k) {
const int out = diag + k;
if (out >= M + N) {
break;
}
if ((j >= N) || (i < M && A[i] <= B[j])) {
C[out] = A[i];
++i;
} else {
C[out] = B[j];
++j;
}
}
}
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int M, int N) {
constexpr int VT = 8;
constexpr int block_dim = 256;
const int grid_dim = (M + N + block_dim * VT - 1) / (block_dim * VT);
merge_kernel<VT><<<grid_dim, block_dim>>>(A, B, C, M, N);
}
二级 Merge Path
整体思路是：
先用一级 Merge Path 找到整个 Block 对应的输入区间，从而确定该 Block 可能需要访问的 A、B 范围，并将这些数据从 Global Memory 搬到 Shared Memory
然后，每个线程在这个较小的 Shared Memory 区间内再执行一次 Merge Path，找到自己的起始位置，最后完成局部 merge 并写入输出
其中，一级 Merge Path 可以由一个 Warp 的 32 个线程协作完成，相当于进行一次 32 路并行查找，以减少 Global Memory 上串行二分带来的延迟
代码：
#include <cuda_runtime.h>
__device__ __forceinline__ int merge_path(const float* __restrict__ A,
const float* __restrict__ B, int M,
int N, int diag) {
int l = max(0, diag - N);
int r = min(diag, M);
while (l < r) {
const int i = (l + r) >> 1;
const int j = diag - i;
if (B[j - 1] >= A[i]) {
l = i + 1;
} else {
r = i;
}
}
return l;
}
// 一个 warp 的 32 个 lane 合作找同一个分界点
__device__ __forceinline__ int merge_path_warp(const float* __restrict__ A,
const float* __restrict__ B,
int M, int N, int diag) {
const int lane = threadIdx.x & 31;
int l = max(0, diag - N);
int r = min(diag, M);
while (l < r) {
const int step = (r - l + 31) >> 5;
const int i = l + lane * step;
const int j = diag - i;
const bool pred = (i < r) && (B[j - 1] >= A[i]);
const unsigned int mask = __ballot_sync(0xffffffffu, pred);
if (mask == 0u) {
return l;
}
const int last_true_lane = 31 - __clz(mask);
const int last_true_i = l + last_true_lane * step;
// 边界在最后一个 true 的区间内
l = last_true_i + 1;
r = min(r, last_true_i + step);
}
return l;
}
template <int BX, int VT>
__global__ void merge_kernel(const float* __restrict__ A,
const float* __restrict__ B, float* __restrict__ C,
int M, int N) {
constexpr int BN = BX * VT;
const int block_diag_start = blockIdx.x * BN;
const int block_diag_end = min(block_diag_start + BN, M + N);
const int block_count = block_diag_end - block_diag_start;
if (block_count <= 0) {
return;
}
// warp 0 搜 block start
// warp 1 搜 block end
__shared__ int s_block_a_start, s_block_a_end;
if (threadIdx.x < 64) {
const int warp = threadIdx.x >> 5;
const int lane = threadIdx.x & 31;
const int diag = (warp == 0) ? block_diag_start : block_diag_end;
const int i = merge_path_warp(A, B, M, N, diag);
if (threadIdx.x == 0) {
s_block_a_start = i;
} else if (threadIdx.x == 32) {
s_block_a_end = i;
}
}
__syncthreads();
const int block_a_start = s_block_a_start;
const int block_a_end = s_block_a_end;
const int block_a_count = block_a_end - block_a_start;
const int block_b_start = block_diag_start - block_a_start;
const int block_b_end = block_diag_end - block_a_end;
const int block_b_count = block_b_end - block_b_start;
__shared__ float s_output[BN];
float* s_a = s_output;
float* s_b = s_output + block_a_count;
// 把当前 block 真正需要的 A/B 数据搬到 shared memory
for (int x = threadIdx.x; x < block_a_count; x += BX) {
s_a[x] = A[block_a_start + x];
}
for (int y = threadIdx.x; y < block_b_count; y += BX) {
s_b[y] = B[block_b_start + y];
}
__syncthreads();
// 每个 thread 在 shared-memory tile 内找自己的 partition
const int thread_diag = threadIdx.x * VT;
if (thread_diag >= block_count) {
return;
}
int i = merge_path(s_a, s_b, block_a_count, block_b_count, thread_diag);
int j = thread_diag - i;
#pragma unroll
for (int k = 0; k < VT; ++k) {
const int out = thread_diag + k;
if (out >= block_count) {
break;
}
if ((j >= block_b_count) || (i < block_a_count && s_a[i] <= s_b[j])) {
C[block_diag_start + out] = s_a[i++];
} else {
C[block_diag_start + out] = s_b[j++];
}
}
}
// A, B, C are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* C, int M, int N) {
constexpr int BX = 256;
constexpr int VT = 8;
constexpr int block_dim = BX;
const int grid_dim = (M + N + block_dim * VT - 1) / (block_dim * VT);
merge_kernel<BX, VT><<<grid_dim, block_dim>>>(A, B, C, M, N);
}
Top K Selection
输入一个数组，输出最大的 k 个数，其中 k 很小
例子：
Input:
input = [1.0, 5.0, 3.0, 2.0, 4.0]
N = 5
k = 3
Output:
output = [5.0, 4.0, 3.0]
基于双调排序的方法
最简单的方法就是先排序，然后拿到前 k 个，但是因为 k 很小，所以没有必要
由于 k 很小，可以在 Block 中排序，每个 Block 都拿到 k 个局部最大值，然后一直合并
代码：
#include <cuda_runtime.h>
template <int BX, int TN>
__device__ __forceinline__ void bitonic_sort_desc(float* __restrict__ data) {
constexpr int BN = BX * TN;
#pragma unroll
for (int k = 2; k <= BN; k <<= 1) {
#pragma unroll
for (int j = k >> 1; j > 0; j >>= 1) {
#pragma unroll
for (int b = 0; b < BN; b += BX) {
const int i = threadIdx.x + b;
const int ixj = i ^ j;
if (ixj > i) {
const float x = data[i];
const float y = data[ixj];
const bool descending = ((i & k) == 0);
if (descending) {
if (x < y) {
data[i] = y;
data[ixj] = x;
}
} else {
if (x > y) {
data[i] = y;
data[ixj] = x;
}
}
}
}
__syncthreads();
}
}
}
template <int BX, int TN>
__global__ void topk_kernel(const float* __restrict__ input,
float* __restrict__ output, int N, int k) {
constexpr int BN = BX * TN;
__shared__ float s_data[BN];
#pragma unroll
for (int b = 0; b < BN; b += BX) {
const int i = threadIdx.x + b;
const int idx = blockIdx.x * BN + i;
if (idx < N) {
s_data[i] = input[idx];
} else {
s_data[i] = -INFINITY;
}
}
__syncthreads();
bitonic_sort_desc<BX, TN>(s_data);
for (int i = threadIdx.x; i < k; i += BX) {
output[blockIdx.x * k + i] = s_data[i];
}
}
// input, output are device pointers
extern "C" void solve(const float* input, float* output, int N, int k) {
constexpr int BX = 256;
constexpr int TN = 4;
constexpr int BN = BX * TN;
// 不支持
if (k > BN) {
return;
}
if (N <= BN) {
topk_kernel<BX, TN><<<1, BX>>>(input, output, N, k);
return;
}
const int max_blocks = (N + BN - 1) / BN;
float* buf1;
float* buf2;
cudaMalloc(&buf1, max_blocks * k * sizeof(float));
cudaMalloc(&buf2, max_blocks * k * sizeof(float));
const float* src = input;
float* dst = buf1;
int current_n = N;
while (current_n > BN) {
const int grid_dim = (current_n + BN - 1) / BN;
topk_kernel<BX, TN><<<grid_dim, BX>>>(src, dst, current_n, k);
current_n = grid_dim * k;
src = dst;
if (dst == buf1) {
dst = buf2;
} else {
dst = buf1;
}
}
topk_kernel<BX, TN><<<1, BX>>>(src, output, current_n, k);
cudaFree(buf1);
cudaFree(buf2);
}
这里用到了 Bitonic Sort（双调排序）
双调序列（Bitonic Sequence）：一个序列如果满足“先单调上升（下降）、后单调下降（上升）”，就称为双调序列，例如：
1 4 6 8 7 5 3 2
↑ ↑ ↑ ↑
↓ ↓ ↓ ↓
双调归并（Bitonic Merge）：将一个双调序列变成升序（降序）的方法。对于上面的这个双调序列，现在希望把它变成升序
长度为 8，那么把序列分成前后两半，对距离为 n/2 = 4 的元素进行比较，然后因为目标是升序，小的放前面，大的放后面：
1 4 3 2 | 7 5 6 8
此时，前半部分中的所有元素都不会大于后半部分中的对应高位部分，并且两个子序列仍然具有双调性质。于是我们可以继续分别处理两个长度为 4 的序列。然后一直递归这样做，最终得到：
1 2 3 4 5 6 7 8
这就是 Bitonic Merge
但是普通数组一开始并不是双调序列，那么从底部开始向上处理，先排序然后再合并
完整递归算法如下：
def bitonic_merge(a, low, n, dir):
if n > 1:
k = n / 2
for i = low ... low+k-1:
if dir == ASC:
if a[i] > a[i+k]:
swap(a[i], a[i+k])
else:
if a[i] < a[i+k]:
swap(a[i], a[i+k])
bitonic_merge(a, low, k, dir)
bitonic_merge(a, low+k, k, dir)
def bitonic_sort(a, low, n, dir):
if n > 1:
k = n / 2
bitonic_sort(a, low, k, ASC) # 前一半升序
bitonic_sort(a, low+k, k, DESC) # 后一半降序
bitonic_merge(a, low, n, dir) # 得到双调序列后归并
这种排序算法非常适合并行计算，因为每一层的比较完全没有数据依赖，代码一般这样：
for (k = 2; k <= N; k *= 2) {
for (j = k / 2; j > 0; j /= 2) {
parallel_for (i = 0 ... N-1) {
ixj = i ^ j; // i ^ k 表示当前 i 应该和哪个元素比较
if (ixj > i) {
if ((i & k) == 0) { // i & k 决定当前这一组执行升序还是降序
if (a[i] > a[ixj]) {
swap(a[i], a[ixj]);
}
} else {
if (a[i] < a[ixj]) {
swap(a[i], a[ixj]);
}
}
}
}
}
}
改造 Radix Sort
改造一下 Radix Sort：
N 个 float
↓
把 float 转成可以按 unsigned int 比较的 key
↓
从最高 8 bit 开始统计 256 个桶
↓
找出第 k 大元素落在哪个桶
↓
只保留这个桶对应的前缀
↓
继续检查接下来的 8 bit
↓
重复 4 次，共 32 bit
↓
得到精确的第 k 大元素 threshold
↓
重新扫描数组，收集 >= threshold 的 k 个元素
↓
对这 k 个元素做 GPU Bitonic Sort
↓
降序 Top-K
代码：
#include <cuda_runtime.h>
#include <algorithm>
#define CEIL(x, y) (((x) + (y) - 1) / (y))
constexpr int k_bits = 32;
constexpr int k_radix_bits = 8;
constexpr int k_radix_digits = 1 << k_radix_bits;
constexpr int k_radix_mask = k_radix_digits - 1;
constexpr int k_radix_passes = k_bits / k_radix_bits;
constexpr int k_items_per_thread = 12;
constexpr int k_threads_per_block = 256;
constexpr int k_items_per_block = k_threads_per_block * k_items_per_thread;
constexpr int k_hist_parts = 4;
constexpr int k_bin_count = k_radix_passes * k_radix_digits;
__device__ __forceinline__ unsigned int float_to_key(float x) {
unsigned int u = __float_as_uint(x);
return u ^ ((u & 0x80000000u) ? 0xffffffffu : 0x80000000u);
}
__device__ __forceinline__ float key_to_float(unsigned int key) {
unsigned int u = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key;
return __uint_as_float(u);
}
struct SelectState {
unsigned int prefix;
unsigned int mask;
int rank;
int greater;
};
__global__ void init_select_state_kernel(SelectState* state, int k) {
if (threadIdx.x == 0) {
state->prefix = 0;
state->mask = 0;
state->rank = k;
state->greater = 0;
}
}
__global__ void select_histogram_kernel(const float* __restrict__ input, int n,
int shift,
const SelectState* __restrict__ state,
unsigned int* __restrict__ bins) {
__shared__ unsigned int hist[k_radix_digits * k_hist_parts];
__shared__ unsigned int s_prefix;
__shared__ unsigned int s_mask;
const int part = (threadIdx.x & 31) % k_hist_parts;
for (int i = threadIdx.x; i < k_radix_digits * k_hist_parts;
i += k_threads_per_block) {
hist[i] = 0;
}
if (threadIdx.x == 0) {
s_prefix = state->prefix;
s_mask = state->mask;
}
__syncthreads();
for (int base = blockIdx.x * k_items_per_block; base < n;
base += gridDim.x * k_items_per_block) {
#pragma unroll
for (int item = 0; item < k_items_per_thread; ++item) {
const int idx = base + item * k_threads_per_block + threadIdx.x;
if (idx < n) {
const unsigned int key = float_to_key(input[idx]);
if ((key & s_mask) == s_prefix) {
const unsigned int digit = (key >> shift) & k_radix_mask;
atomicAdd(&hist[digit * k_hist_parts + part], 1u);
}
}
}
}
__syncthreads();
for (int digit = threadIdx.x; digit < k_radix_digits;
digit += k_threads_per_block) {
unsigned int count = 0;
#pragma unroll
for (int part_id = 0; part_id < k_hist_parts; ++part_id) {
count += hist[digit * k_hist_parts + part_id];
}
if (count != 0) {
atomicAdd(&bins[digit], count);
}
}
}
__global__ void select_digit_kernel(unsigned int* __restrict__ bins, int shift,
SelectState* __restrict__ state) {
if (threadIdx.x == 0) {
int rank = state->rank;
int greater = 0;
int chosen = 0;
for (int digit = k_radix_digits - 1; digit >= 0; --digit) {
const int count = bins[digit];
if (rank > count) {
rank -= count;
greater += count;
} else {
chosen = digit;
break;
}
}
state->rank = rank;
state->greater += greater;
state->prefix |= static_cast<unsigned int>(chosen) << shift;
state->mask |= static_cast<unsigned int>(k_radix_mask) << shift;
}
__syncthreads();
if (threadIdx.x < k_radix_digits) {
bins[threadIdx.x] = 0;
}
}
__device__ __forceinline__ unsigned int warp_inclusive_scan(unsigned int val) {
const int lane = threadIdx.x & 31;
#pragma unroll
for (int offset = 1; offset < 32; offset <<= 1) {
const unsigned int other = __shfl_up_sync(0xffffffffu, val, offset);
if (lane >= offset) {
val += other;
}
}
return val;
}
__device__ __forceinline__ unsigned int warp_exclusive_scan(unsigned int val) {
return warp_inclusive_scan(val) - val;
}
template <int k_block_dim>
__device__ __forceinline__ int block_inclusive_scan(int val) {
const int warp = threadIdx.x >> 5;
const int lane = threadIdx.x & 31;
const int k_num_warps = k_block_dim >> 5;
__shared__ int warp_offsets[k_num_warps];
const int prefix = warp_inclusive_scan(val);
if (lane == 31) {
warp_offsets[warp] = prefix;
}
__syncthreads();
if (warp == 0) {
const int warp_sum = lane < k_num_warps ? warp_offsets[lane] : 0;
const int warp_offset = warp_exclusive_scan(warp_sum);
if (lane < k_num_warps) {
warp_offsets[lane] = warp_offset;
}
}
__syncthreads();
return prefix + warp_offsets[warp];
}
template <int k_block_dim>
__device__ __forceinline__ unsigned int block_exclusive_scan(unsigned int val) {
return block_inclusive_scan<k_block_dim>(val) - val;
}
__global__ void gather_topk_kernel(const float* __restrict__ input,
unsigned int* __restrict__ output, int n,
int k, const SelectState* __restrict__ state,
unsigned int* __restrict__ ctrs) {
__shared__ unsigned int s_threshold;
__shared__ unsigned int s_greater;
__shared__ unsigned int s_gt_count;
__shared__ unsigned int s_eq_count;
__shared__ unsigned int s_gt_base;
__shared__ unsigned int s_eq_base;
if (threadIdx.x == 0) {
s_threshold = state->prefix;
s_greater = state->greater;
}
__syncthreads();
for (int base = blockIdx.x * k_items_per_block; base < n;
base += gridDim.x * k_items_per_block) {
unsigned int keys[k_items_per_thread];
unsigned int gt_count = 0;
unsigned int eq_count = 0;
#pragma unroll
for (int item = 0; item < k_items_per_thread; ++item) {
const int idx = base + item * k_threads_per_block + threadIdx.x;
const unsigned int key = idx < n ? float_to_key(input[idx]) : 0u;
keys[item] = key;
if (idx < n) {
gt_count += key > s_threshold;
eq_count += key == s_threshold;
}
}
const unsigned int gt_offset =
block_exclusive_scan<k_threads_per_block>(gt_count);
const unsigned int eq_offset =
block_exclusive_scan<k_threads_per_block>(eq_count);
if (threadIdx.x == k_threads_per_block - 1) {
s_gt_count = gt_offset + gt_count;
s_eq_count = eq_offset + eq_count;
}
__syncthreads();
if (threadIdx.x == 0) {
s_gt_base = atomicAdd(ctrs, s_gt_count);
s_eq_base = atomicAdd(ctrs + 1, s_eq_count);
}
__syncthreads();
unsigned int gt_seen = 0;
unsigned int eq_seen = 0;
#pragma unroll
for (int item = 0; item < k_items_per_thread; ++item) {
const int idx = base + item * k_threads_per_block + threadIdx.x;
if (idx < n) {
const unsigned int key = keys[item];
if (key > s_threshold) {
const unsigned int pos = s_gt_base + gt_offset + gt_seen;
if (pos < s_greater) {
output[pos] = key;
}
++gt_seen;
} else if (key == s_threshold) {
const unsigned int pos = s_greater + s_eq_base + eq_offset + eq_seen;
if (pos < static_cast<unsigned int>(k)) {
output[pos] = key;
}
++eq_seen;
}
}
}
__syncthreads();
}
}
__global__ void small_sort_desc_kernel(float* output, int k, int sort_size) {
__shared__ unsigned int values[1024];
const int tid = threadIdx.x;
if (tid < sort_size) {
values[tid] = tid < k ? __float_as_uint(output[tid]) : 0u;
}
__syncthreads();
for (int width = 2; width <= sort_size; width <<= 1) {
for (int stride = width >> 1; stride > 0; stride >>= 1) {
const int other = tid ^ stride;
if (other > tid && other < sort_size) {
const unsigned int lhs = values[tid];
const unsigned int rhs = values[other];
const bool ascending = (tid & width) == 0;
if ((lhs > rhs) == ascending) {
values[tid] = rhs;
values[other] = lhs;
}
}
__syncthreads();
}
}
if (tid < k) {
output[tid] = key_to_float(values[sort_size - 1 - tid]);
}
}
extern "C" void solve(const float* input, float* output, int N, int k) {
const int num_blocks = CEIL(N, k_items_per_block);
int device = 0;
int sm_count = 0;
int active_blocks_per_sm = 0;
cudaGetDevice(&device);
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device);
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&active_blocks_per_sm, select_histogram_kernel, k_threads_per_block, 0);
const int histogram_blocks =
std::min(sm_count * active_blocks_per_sm, num_blocks);
unsigned int* d_bins;
unsigned int* d_ctrs;
SelectState* d_state;
const int bins_bytes = k_radix_digits * sizeof(unsigned int);
const int ctrs_bytes = 2 * sizeof(unsigned int);
cudaMalloc(&d_bins, bins_bytes);
cudaMalloc(&d_ctrs, ctrs_bytes);
cudaMalloc(&d_state, sizeof(SelectState));
cudaMemsetAsync(d_bins, 0, bins_bytes);
init_select_state_kernel<<<1, 1>>>(d_state, k);
#pragma unroll
for (int pass = 0; pass < k_radix_passes; ++pass) {
const int shift = k_bits - (pass + 1) * k_radix_bits;
select_histogram_kernel<<<histogram_blocks, k_threads_per_block>>>(
input, N, shift, d_state, d_bins);
select_digit_kernel<<<1, k_radix_digits>>>(d_bins, shift, d_state);
}
cudaMemsetAsync(d_ctrs, 0, ctrs_bytes);
gather_topk_kernel<<<histogram_blocks, k_threads_per_block>>>(
input, reinterpret_cast<unsigned int*>(output), N, k, d_state, d_ctrs);
int sort_size = 1;
while (sort_size < k) {
sort_size <<= 1;
}
small_sort_desc_kernel<<<1, sort_size>>>(output, k, sort_size);
cudaDeviceSynchronize();
cudaFree(d_state);
cudaFree(d_ctrs);
cudaFree(d_bins);
}