Commit 42c787e8c for llama.cpp
commit 42c787e8c191d49c01c757f4b7029d69ec0bf2c8
Author: Shobhit <sobhit.me@gmail.com>
Date: Thu Oct 8 01:53:43 2026 +0800
cuda: update uncoalesced memory reads in pool2d (#29425)
* cuda: update uncoalesced memory reads in pool2d
* cuda: Added perf and __restrict__ in pool2d
* Update POOL2D_WARP_KERNEL_MIN_WINDOW macro
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
* cuda: Fix compiler bugs
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/ggml/src/ggml-cuda/pool2d.cu b/ggml/src/ggml-cuda/pool2d.cu
index c6d51e4d6..83a49753c 100644
--- a/ggml/src/ggml-cuda/pool2d.cu
+++ b/ggml/src/ggml-cuda/pool2d.cu
@@ -50,6 +50,65 @@ static __global__ void pool2d_nchw_kernel(
o_ptr[cur_oh * ow + cur_ow] = res;
}
+template <typename Ti, typename To>
+static __global__ void pool2d_nchw_kernel_warp(
+ const int ih, const int iw, const int oh, const int ow,
+ const int kh, const int kw, const int sh, const int sw,
+ const int ph, const int pw, const int parallel_elements,
+ const Ti * __restrict__ src, To * __restrict__ dst, const enum ggml_op_pool op) {
+ const int warp_id = (threadIdx.x + blockIdx.x * blockDim.x) / WARP_SIZE;
+ const int lane = threadIdx.x % WARP_SIZE;
+ if (warp_id >= parallel_elements) {
+ return;
+ }
+
+ const int I_HW = ih * iw;
+ const int O_HW = oh * ow;
+ const int nc = warp_id / O_HW;
+ const int cur_oh = warp_id % O_HW / ow;
+ const int cur_ow = warp_id % O_HW % ow;
+ const Ti* i_ptr = src + nc * I_HW;
+
+ const int start_h = cur_oh * sh - ph;
+ const int bh = max(0, start_h);
+ const int eh = min(ih, start_h + kh);
+ const int start_w = cur_ow * sw - pw;
+ const int bw = max(0, start_w);
+ const int ew = min(iw, start_w + kw);
+
+ const int win_w = ew - bw;
+ const int win_elems = (eh - bh) * win_w;
+ const To scale = 1. / (kh * kw);
+
+ To res;
+ switch (op) {
+ case GGML_OP_POOL_AVG: res = 0; break;
+ case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
+ default: res = 0; assert(false);
+ }
+
+ for (int t = lane; t < win_elems; t += WARP_SIZE) {
+ const int i = bh + t / win_w;
+ const int j = bw + t % win_w;
+ const Ti cur = i_ptr[i * iw + j];
+ switch (op) {
+ case GGML_OP_POOL_AVG: res += cur * scale; break;
+ case GGML_OP_POOL_MAX: res = max(res, (To)cur); break;
+ default: break;
+ }
+ }
+
+#pragma unroll
+ for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
+ const To other = __shfl_xor_sync(0xFFFFFFFF, res, offset, WARP_SIZE);
+ res = (op == GGML_OP_POOL_MAX) ? max(res, other) : res + other;
+ }
+
+ if (lane == 0) {
+ dst[nc * O_HW + cur_oh * ow + cur_ow] = res;
+ }
+}
+
static void pool2d_nchw_kernel_f32_f32_cuda(
const int ih, const int iw, const int oh, const int ow,
const int kh, const int kw, const int sh, const int sw,
@@ -57,6 +116,13 @@ static void pool2d_nchw_kernel_f32_f32_cuda(
const float * src, float * dst, const enum ggml_op_pool op,
cudaStream_t stream) {
+ if (kh * kw >= CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW) {
+ const int warps_per_block = CUDA_POOL2D_BLOCK_SIZE / WARP_SIZE;
+ const int num_blocks = (parallel_elements + warps_per_block - 1) / warps_per_block;
+ pool2d_nchw_kernel_warp<<<num_blocks, CUDA_POOL2D_BLOCK_SIZE, 0, stream>>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);
+ return;
+ }
+
const int num_blocks = (parallel_elements + CUDA_POOL2D_BLOCK_SIZE - 1) / CUDA_POOL2D_BLOCK_SIZE;
dim3 block_nums(num_blocks);
pool2d_nchw_kernel<<<block_nums, CUDA_POOL2D_BLOCK_SIZE, 0, stream>>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);
diff --git a/ggml/src/ggml-cuda/pool2d.cuh b/ggml/src/ggml-cuda/pool2d.cuh
index 7841292bc..e2c7fb441 100644
--- a/ggml/src/ggml-cuda/pool2d.cuh
+++ b/ggml/src/ggml-cuda/pool2d.cuh
@@ -1,5 +1,6 @@
#include "common.cuh"
#define CUDA_POOL2D_BLOCK_SIZE 256
+#define CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW 32
void ggml_cuda_op_pool2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 89554d4d2..cfe5b73e7 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -9558,6 +9558,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
}
+ // MSA indexer block-max shape: wide 1-D window along ne0, no padding, non-divisible width
+ for (int64_t iw : {4096, 4160}) {
+ for (int blk : {32, 64}) {
+ test_cases.emplace_back(new test_pool2d(pool_type, type_input, {iw, 2, 1, 1}, blk, 1, blk, 1, 0, 0));
+ }
+ }
}
}
@@ -11652,6 +11658,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
GGML_TYPE_F32, {n_kv, 512, 64, 1}, false, {2, 1, 0, 3}));
}
+ // POOL_2D max over MiniMax-M3 indexer block scores:
+ // sc is [n_ps, n_head=4, n_tokens] and gets ggml_pool_2d(.., MAX, blk=128, 1, blk=128, 1, 0, 0).
+ for (int64_t n_ps : { 8192, 32768 }) {
+ for (int64_t n_tokens : { 1, 512 }) {
+ test_cases.emplace_back(new test_pool2d(
+ GGML_OP_POOL_MAX, GGML_TYPE_F32, {n_ps, 4, n_tokens, 1}, 128, 1, 128, 1, 0, 0));
+ }
+ }
+
// LEAKY_RELU at FFN activation width, for direct comparison with RELU
for (int64_t n_tokens : {512, 2048}) {
test_cases.emplace_back(new test_leaky_relu(GGML_TYPE_F32, { 17408, n_tokens, 1, 1 }, 0.1f));