vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

CUDA op getrows fails for long sequences

Browse source 
T5 embeddings have a square input pos tensor which quickly exceeds the 65k limit of getrows

Implemented only for _float, need other implementations
This commit is contained in:
Milot Mirdita 2025-01-11 15:26:23 +09:00
parent 2739a71e4b
commit 6acdb265fc
1 changed files with 17 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

17
ggml/src/ggml-cuda/getrows.cu
View file

@ -118,13 +118,28 @@ static void get_rows_cuda_float(const ggml_tensor * src0, const ggml_tensor * sr
const size_t s12 = nb12 / ggml_element_size(src1);
//const size_t s13 = nb13 / ggml_element_size(src1);
// CUDA grids typically cannot exceed 65535 in .y dimension
static const int64_t MAX_GRID_Y = 65535;
for (int64_t startY = 0; startY < ne10; startY += MAX_GRID_Y) {
int64_t chunkSizeY = std::min<int64_t>(MAX_GRID_Y, ne10 - startY);
// launch chunkSizeY blocks in the Y-dimension instead of the entire ne10
dim3 block_nums(block_num_x, chunkSizeY, ne11 * ne12);
// pointer offset: since the kernel calculates src1[i10*s10 + ...] and dst[i10*s1 + ...],
// we shift src1_dd/dst_dd by (startY * stride) so that when i10=0 in the kernel,
// it corresponds to row "startY" in the overall space.
const int32_t * src1_dd_offset = src1_dd + startY * s10;
float * dst_dd_offset = dst_dd + startY * s1;
k_get_rows_float<<<block_nums, block_dims, 0, stream>>>(
src0_dd, src1_dd, dst_dd,
src0_dd, src1_dd_offset, dst_dd_offset,
ne00, /*ne01, ne02, ne03,*/
/*ne10, ne11,*/ ne12, /*ne13,*/
/* s0,*/ s1, s2, s3,
/* nb00,*/ nb01, nb02, nb03,
s10, s11, s12/*, s13*/);
}
GGML_UNUSED(dst);
}