Oliver Simons
36f0132464
* CUDA: Refactor and expose two_stage_warp_reduce_* function * Use `two_stage_warp_reduce` also in softmax kernel, move smem out of it Moving smem out of `__device__` function to `__global__` function allows for explicit smem reuse, as either compiler or cuda rt seem to not free it afterwards (`cudaFuncSetAttribute` fails when not accounting for it once for each call to two_stage_warp_reduce) * Update ggml/src/ggml-cuda/common.cuh Co-authored-by: Aman Gupta <amangupta052@gmail.com> * Use two_stage_warp_reduce in group_norm_f32 * Use two_stage_warp_reduce in rms_norm_f32 * Fix smem calculation which expects bytes * Make `two_stage_warp_reduce` accept all values warp_reduce accepts Also integrate it into norm_f32 function * Use two_stage_warp_reduce in l2_norm_f32 * Use type traits for block reduction for better legibility Also adresss other requests by @am17an such as variable renaming * Make norm tests cover all cuda paths * Mark columns % WARP_SIZE !=0 as supported for RMS_NORM_BACK Unit-tests passed locally, let's see if they pass in the CI as well * Use `enum class` for `block_reduce_method` This is more type-safe than plain enum * Rename variables as suggested in code review by @am17an * Rename two_stage_warp_reduce -> block_reduce * Fix trailing whitespace in common.cuh * Make condition of static_assert type-dependent This delays evaluation until the template is actually instantiated. Otherwise, some compilers may evaluate the assert when parsing the template, resulting in build errors as observed here: https://github.com/ggml-org/llama.cpp/actions/runs/20960323123/job/60235530068?pr=18785 * Inline definitions --------- Co-authored-by: Aman Gupta <amangupta052@gmail.com>
40 lines
1.1 KiB
Plaintext
40 lines
1.1 KiB
Plaintext
#include "common.cuh"
|
|
|
|
// Row reduction kernel template - compute sum (norm=false) or mean (norm=true)
|
|
template <bool norm>
|
|
static __global__ void reduce_rows_f32(const float * __restrict__ x, float * __restrict__ dst, const int ncols) {
|
|
const int row = blockIdx.x;
|
|
const int col = threadIdx.x;
|
|
|
|
float sum = 0.0f;
|
|
const int num_unroll = 8;
|
|
float temp[num_unroll];
|
|
float sum_temp[num_unroll] = { 0.0f };
|
|
for (int i = col; i < ncols;) {
|
|
for (int j = 0; j < num_unroll; ++j) {
|
|
if (i < ncols) {
|
|
temp[j] = x[row * ncols + i];
|
|
} else {
|
|
temp[j] = 0;
|
|
}
|
|
i += blockDim.x;
|
|
}
|
|
for (int j = 0; j < num_unroll; ++j) {
|
|
sum_temp[j] += temp[j];
|
|
}
|
|
}
|
|
for (int j = 0; j < num_unroll; ++j) {
|
|
sum += sum_temp[j];
|
|
}
|
|
|
|
// sum up partial sums
|
|
__shared__ float shared_vals[32];
|
|
sum = block_reduce<block_reduce_method::SUM>(sum, shared_vals);
|
|
|
|
if (col != 0) {
|
|
return;
|
|
}
|
|
|
|
dst[row] = norm ? sum / ncols : sum;
|
|
}
|