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

k_get_rows_float did a scalar one-element-per-thread copy and recomputed the
row-invariant work (index load, fast_div_modulo, src/dst row pointers) for
every element. Hoist that out of the per-element loop, and add a vectorized
path (k_get_rows_float_vec) that copies one int4 (16 B) per thread for the
contiguous same-type (no-cast) case.

The vectorized path is gated at compile time (is_same<src0_t, dst_t>) and at
runtime on 16-byte alignment of the base pointers and all row strides and on
ne00 % VEC == 0. Vectorizing divides the block count by VEC, so a small
single-row gather can drop below the device CU count and regress; an
occupancy gate keeps those on the block-rich scalar path.

On Strix Halo (gfx1151) the DeltaNet recurrent-state gather (ne00=524288)
drops 18.6us -> 13.0us (rocprofv3 HW timestamps), faster than the Vulkan
backend, with no regression on the small conv-state gather; total get_rows
-27%. test-backend-ops GET_ROWS passes (47/47).

Assisted-by: Claude Opus 4.8
This commit is contained in:
Piotr Wilkin (ilintar)
2026-07-21 15:53:57 +02:00
committed by GitHub
parent 76f46ad29d
commit 305ba519ab

View File

@@ -55,27 +55,51 @@ 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;
if (i00 >= ne00) {
return;
}
const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
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);
dst_row[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
}
}
}
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) {
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;
const int i01 = src1_ptr[i10*s10 + i11*s11 + i12*s12];
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];
}
}
}
template<typename grad_t, typename dst_t>
static __global__ void k_get_rows_back_float(
const grad_t * __restrict__ grad, const int32_t * __restrict__ rows, dst_t * __restrict__ dst,
@@ -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);
@@ -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,