Skip to content

CUDA: Optimize PAD_REFLECT_1D #15957

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
merged 11 commits into from
Sep 18, 2025
Merged

CUDA: Optimize PAD_REFLECT_1D #15957

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
1e29faf
CUDA: Optimize PAD_REFLECT_1D
bugparty Sep 12, 2025
9494833
use fast_div to improve performance
bugparty Sep 15, 2025
8583552
Apply suggestion from @JohannesGaessler
bugparty Sep 15, 2025
a5ef1d0
Apply suggestion from @JohannesGaessler
bugparty Sep 15, 2025
b3cf133
optimize
bugparty Sep 15, 2025
d73ba84
use a concise expression to further speedup the cuda kernel
bugparty Sep 15, 2025
e280cb8
add comment for rel_i0
bugparty Sep 15, 2025
188ce93
Merge branch 'ggml-org:master' into PAD_REFLECT_1D_expriment
bugparty Sep 17, 2025
4286ea7
Merge branch 'ggml-org:master' into PAD_REFLECT_1D_expriment
bugparty Sep 17, 2025
dd6789b
Merge branch 'ggml-org:master' into PAD_REFLECT_1D_expriment
bugparty Sep 17, 2025
aa12620
Merge branch 'ggml-org:master' into PAD_REFLECT_1D_expriment
bugparty Sep 18, 2025
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
8 changes: 8 additions & 0 deletions ggml/src/ggml-cuda/common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -636,6 +636,14 @@ static __device__ __forceinline__ uint32_t fastmodulo(uint32_t n, const uint3 fa
return n - fastdiv(n, fastdiv_values) * fastdiv_values.z;
}

// Calculate both division and modulo at once, returns <n/divisor, n%divisor>
static __device__ __forceinline__ uint2 fast_div_modulo(uint32_t n, const uint3 fastdiv_values) {
// expects fastdiv_values to contain <mp, L, divisor> in <x, y, z> (see init_fastdiv_values)
const uint32_t div_val = fastdiv(n, fastdiv_values);
const uint32_t mod_val = n - div_val * fastdiv_values.z;
return make_uint2(div_val, mod_val);
}

typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, float2 & v);

static __device__ __forceinline__ float get_alibi_slope(
Expand Down
50 changes: 30 additions & 20 deletions ggml/src/ggml-cuda/pad_reflect_1d.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ static __global__ void pad_reflect_1d_kernel_f32(
const int64_t ne0,
const int64_t ne00,
const int64_t ne01,
const uint3 ne01_packed,
const int64_t ne02,
const int64_t ne03,
const int64_t nb00,
Expand All @@ -21,32 +22,34 @@ static __global__ void pad_reflect_1d_kernel_f32(

const int64_t i3 = blockIdx.z;
const int64_t i2 = blockIdx.y;
const int64_t i1 = blockIdx.x;

if (i1 >= ne01 || i2 >= ne02 || i3 >= ne03) {
const uint2 div_mod_packed = fast_div_modulo(blockIdx.x, ne01_packed);
const int64_t tile1 = div_mod_packed.y; // i1
const int64_t tile0 = div_mod_packed.x; // nth i0 tile
const int64_t i1 = tile1;
const int64_t i0 = threadIdx.x + tile0 * blockDim.x;
if (i0 >= ne0 || i1 >= ne01 || i2 >= ne02 || i3 >= ne03) {
return;
}

const char * src0_ptr = (const char *)src0 + i3*nb03 + i2*nb02 + i1*nb01;
char * dst_ptr = (char *)dst + i3*nb3 + i2*nb2 + i1*nb1;

for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) {
float value;
float value;
const int64_t j = i0 - p0;

if (i0 < p0) {
// Left padding - reflect
value = *(const float *)(src0_ptr + (p0 - i0) * nb00);
} else if (i0 < ne0 - p1) {
// Middle - copy
value = *(const float *)(src0_ptr + (i0 - p0) * nb00);
} else {
// Right padding - reflect
int64_t src_idx = (ne0 - p1 - p0) - (p1 + 1 - (ne0 - i0)) - 1;
value = *(const float *)(src0_ptr + src_idx * nb00);
}

*(float *)(dst_ptr + i0 * nb0) = value;
if (j < 0) {// i0<p0
// Left padding - reflect
value = *(const float *)(src0_ptr - j * nb00);
} else if ( j < ne00 ) { //i0 < ne0 - p1
// Middle - copy
value = *(const float *)(src0_ptr + j * nb00);
} else {
// Right padding - reflect
const int64_t src_idx = (ne0 - p1 - p0) - (p1 + 1 - (ne0 - i0)) - 1;
value = *(const float *)(src0_ptr + src_idx * nb00);
}
*(float *)(dst_ptr + i0 * nb0) = value;
}

void ggml_cuda_op_pad_reflect_1d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
Expand All @@ -62,19 +65,26 @@ void ggml_cuda_op_pad_reflect_1d(ggml_backend_cuda_context & ctx, ggml_tensor *

const int64_t ne00 = src0->ne[0];
const int64_t ne01 = src0->ne[1];
const uint3 ne01_packed = init_fastdiv_values(ne01);
const int64_t ne02 = src0->ne[2];
const int64_t ne03 = src0->ne[3];

const int64_t ne0 = dst->ne[0];

// sanity: padded length matches
GGML_ASSERT(ne0 == ne00 + p0 + p1);

const dim3 block_dims(CUDA_PAD_REFLECT_1D_BLOCK_SIZE, 1, 1);
const dim3 grid_dims(ne01, ne02, ne03);
constexpr int64_t bx = CUDA_PAD_REFLECT_1D_BLOCK_SIZE; // threads per block (x)
const int64_t tiles0 = (ne0 + bx - 1) / bx; // number of tiles along i0
// grid.x covers i1 and all tiles of i0: [ne01 * tiles0]
// grid.y covers i2: [ne02]
// grid.z covers i3: [ne03]
const dim3 grid_dims((unsigned)(ne01 * tiles0), (unsigned)ne02, (unsigned)ne03);
const dim3 block_dims((unsigned)bx, 1, 1);

pad_reflect_1d_kernel_f32<<<grid_dims, block_dims, 0, stream>>>(
src0->data, dst->data,
ne0, ne00, ne01, ne02, ne03,
ne0, ne00, ne01, ne01_packed, ne02, ne03,
src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3],
p0, p1
Expand Down
7 changes: 7 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6495,6 +6495,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_pad());
test_cases.emplace_back(new test_pad_ext());
test_cases.emplace_back(new test_pad_reflect_1d());
test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 384, 4, 1}));
test_cases.emplace_back(new test_roll());
test_cases.emplace_back(new test_arange());
test_cases.emplace_back(new test_timestep_embedding());
Expand Down Expand Up @@ -6633,6 +6634,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {1024, 10, 1, 1}));
test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {32000, 512, 1, 1}));

test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {512, 34, 2, 1}));
test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 80, 1, 1}));
test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 384, 1, 1}));
test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 80, 4, 1}));
test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 384, 4, 1}));

test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 16416, 1, 128, {8, 1}, {4, 1}, {0, 2, 1, 3}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 128, 1, 16416, {8, 1}, {4, 1}, {0, 1, 2, 3}, true));

Expand Down