Skip to content

CUDA: vectorize same-type get_rows with int4 copy #25929

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pwilkin merged 1 commit into from
Jul 21, 2026
+64 −14
Merged

CUDA: vectorize same-type get_rows with int4 copy #25929

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 64 additions & 14 deletions ggml/src/ggml-cuda/getrows.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,23 +55,47 @@ static __global__ void k_get_rows_float(
dst_t * GGML_CUDA_RESTRICT dst = dst_ptr;
ggml_cuda_pdl_sync();
for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
// The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
const int i10 = blockIdx.x;
const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
const int i11 = dm.x;
const int i12 = dm.y;

const int i01 = src1[i10*s10 + i11*s11 + i12*s12];

dst_t * GGML_CUDA_RESTRICT dst_row = dst + i10*s1 + i11*s2 + i12*s3;
const src0_t * GGML_CUDA_RESTRICT src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);

for (int64_t i00 = blockIdx.y*blockDim.x + threadIdx.x; i00 < ne00; i00 += gridDim.y*blockDim.x) {
// The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
const int i10 = blockIdx.x;
const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
const int i11 = dm.x;
const int i12 = dm.y;
dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
}
}
}

if (i00 >= ne00) {
return;
}
template<typename dst_t>
static __global__ void k_get_rows_float_vec(
const dst_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
const int64_t ne00v,
const int64_t ne11, const uint3 ne12_fdv,
const size_t s1, const size_t s2, const size_t s3,
const size_t nb01, const size_t nb02, const size_t nb03,
const size_t s10, const size_t s11, const size_t s12) {

const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
ggml_cuda_pdl_lc();
ggml_cuda_pdl_sync();
for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
const int i10 = blockIdx.x;
const uint2 dm = fast_div_modulo((uint32_t)z, ne12_fdv);
const int i11 = dm.x;
const int i12 = dm.y;

dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
const src0_t * src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
const int i01 = src1_ptr[i10*s10 + i11*s11 + i12*s12];

dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
int4 * GGML_CUDA_RESTRICT dst_row = (int4 *) (dst_ptr + i10*s1 + i11*s2 + i12*s3);
const int4 * GGML_CUDA_RESTRICT src0_row = (const int4 *)((const char *) src0_ptr + i01*nb01 + i11*nb02 + i12*nb03);

for (int64_t i = blockIdx.y*blockDim.x + threadIdx.x; i < ne00v; i += gridDim.y*blockDim.x) {
dst_row[i] = src0_row[i];
}
}
}
Expand Down Expand Up @@ -148,8 +172,6 @@ static void get_rows_cuda_float(
const size_t nb1, const size_t nb2, const size_t nb3,
cudaStream_t stream) {
const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));

// strides in elements
// const size_t s0 = nb0 / sizeof(dst_t);
Expand All @@ -166,6 +188,34 @@ static void get_rows_cuda_float(
GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
const uint3 ne12_fdv = init_fastdiv_values(ne12);

if constexpr (std::is_same<src0_t, dst_t>::value) {
constexpr int VEC = 16 / sizeof(dst_t);
const int64_t ne00v = ne00 / VEC;
const int64_t vec_block_num_y = (ne00v + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
const bool enough_blocks = vec_block_num_y * ne10 * ne11 * ne12 >= 128;
const bool can_vec = VEC > 1 && enough_blocks &&
(ne00 % VEC == 0) &&
(nb01 % 16 == 0) && (nb02 % 16 == 0) && (nb03 % 16 == 0) &&
(nb1 % 16 == 0) && (nb2 % 16 == 0) && (nb3 % 16 == 0) &&
(((uintptr_t) src0_d) % 16 == 0) && (((uintptr_t) dst_d) % 16 == 0);

if (can_vec) {
const int block_num_y = vec_block_num_y;
const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
ggml_cuda_kernel_launch(k_get_rows_float_vec<dst_t>, launch_params,
(const dst_t *) src0_d, src1_d, dst_d,
ne00v, ne11, ne12_fdv,
s1, s2, s3,
nb01, nb02, nb03,
s10, s11, s12);
return;
}
}

const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
const dim3 block_nums(ne10, MIN(block_num_y, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));

const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{block_nums, block_dims, 0, stream};
ggml_cuda_kernel_launch(k_get_rows_float<src0_t, dst_t>, launch_params,
src0_d, src1_d, dst_d,
Expand Down
Loading