nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml-cuda: fix shared memory sizing in OpenAI argsort kernel (#1512)

This commit is contained in:
Alec Koumjian 2026-03-26 03:00:00 -04:00 committed by GitHub
parent 095bd3d3b2
commit aa7fdb3259
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
ggml/src/ggml-cuda/argsort.cu
View File

@ -407,7 +407,10 @@ static void argsort_openai_f32_f32_i32_cuda(const float * x, float * weights, in
const dim3 block_dims(ncols_pad, 1, 1);
const dim3 block_nums(nrows, 1, 1);
const size_t shared_mem = (ncols_pad + ncols_pad > WARP_SIZE ? WARP_SIZE : 0) * sizeof(int);
// k_openai_f32_f32_i32 always stores dst_row[ncols_pad] in dynamic shared memory.
// When blockDim > WARP_SIZE it also uses a per-warp float scratch segment (s_sum).
const int n_warps = (ncols_pad + WARP_SIZE - 1) / WARP_SIZE;
const size_t shared_mem = ncols_pad * sizeof(int) + (ncols_pad > WARP_SIZE ? n_warps * sizeof(float) : 0);
// FIXME: this limit could be raised by ~2-4x on Ampere or newer
GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb);