Fix grid overflow in get_rows on CUDA (#2367)

This commit is contained in:
Kawrakow 2026-08-27 15:40:41 +02:00 committed by GitHub
parent 6b261d771e
commit 78ce50c163
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 6 deletions

View File

@ -17,8 +17,8 @@ static __global__ void k_get_rows(
/*size_t nb00,*/ size_t nb01, size_t nb02, size_t nb03,
size_t s10, size_t s11, size_t s12/*, size_t s13*/) {
const int i00 = (blockIdx.x*blockDim.x + threadIdx.x)*2;
const int i10 = blockDim.y*blockIdx.y + threadIdx.y;
const int i00 = (blockIdx.y*blockDim.x + threadIdx.x)*2;
const int i10 = blockIdx.x*blockDim.y + threadIdx.y;
const int i11 = (blockIdx.z*blockDim.z + threadIdx.z)/ne12;
const int i12 = (blockIdx.z*blockDim.z + threadIdx.z)%ne12;
@ -57,8 +57,8 @@ static __global__ void k_get_rows_float(
/*size_t nb00,*/ size_t nb01, size_t nb02, size_t nb03,
size_t s10, size_t s11, size_t s12/*, size_t s13*/) {
const int i00 = blockIdx.x*blockDim.x + threadIdx.x;
const int i10 = blockDim.y*blockIdx.y + threadIdx.y;
const int i00 = blockIdx.y*blockDim.x + threadIdx.x;
const int i10 = blockIdx.x*blockDim.y + threadIdx.y;
const int i11 = (blockIdx.z*blockDim.z + threadIdx.z)/ne12;
const int i12 = (blockIdx.z*blockDim.z + threadIdx.z)%ne12;
@ -81,7 +81,8 @@ static void get_rows_cuda(const ggml_tensor * src0, const ggml_tensor * src1, gg
const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
const int block_num_x = (ne00 + 2*CUDA_GET_ROWS_BLOCK_SIZE - 1) / (2*CUDA_GET_ROWS_BLOCK_SIZE);
const dim3 block_nums(block_num_x, ne10, ne11*ne12);
GGML_ASSERT(ne11*ne12 < 65536);
const dim3 block_nums(ne10, block_num_x, ne11*ne12);
// strides in elements
//const size_t s0 = nb0 / ggml_element_size(dst);
@ -115,7 +116,8 @@ static void get_rows_cuda_float(const ggml_tensor * src0, const ggml_tensor * sr
const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
const int block_num_x = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
const dim3 block_nums(block_num_x, ne10, ne11*ne12);
GGML_ASSERT(ne11*ne12 < 65536);
const dim3 block_nums(ne10, block_num_x, ne11*ne12);
// strides in elements
//const size_t s0 = nb0 / ggml_element_size(dst);