diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index c3758cd50cf..ddec3c431ab 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -3,38 +3,46 @@ template static __global__ void norm_f32( - const float * x, float * dst, const int ncols, const int64_t stride_row, const int64_t stride_channel, - const int64_t stride_sample, const float eps) { - const int nrows = gridDim.x; - const int nchannels = gridDim.y; + const float * x, float * dst, const int ncols, const int nchannels, const int nsamples, + const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps) { + const int nrows = gridDim.x; + const int row = blockIdx.x; + const int tid = threadIdx.x; - const int row = blockIdx.x; - const int channel = blockIdx.y; - const int sample = blockIdx.z; - const int tid = threadIdx.x; + extern __shared__ float2 s_sum2[]; - x += sample*stride_sample + channel*stride_channel + row*stride_row; - dst += ((sample*nchannels + channel)*nrows + row)*ncols; + ggml_cuda_pdl_sync(); - float2 mean_var = make_float2(0.0f, 0.0f); + // grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples + for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) { + for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) { + const float * xc = x + sample*stride_sample + channel*stride_channel + row*stride_row; + float * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols; - ggml_cuda_pdl_sync(); - for (int col = tid; col < ncols; col += block_size) { - const float xi = x[col]; - mean_var.x += xi; - mean_var.y += xi * xi; - } + float2 mean_var = make_float2(0.0f, 0.0f); - // sum up partial sums - extern __shared__ float2 s_sum2[]; - mean_var = block_reduce(mean_var, s_sum2); + for (int col = tid; col < ncols; col += block_size) { + const float xi = xc[col]; + mean_var.x += xi; + mean_var.y += xi * xi; + } - const float mean = mean_var.x / ncols; - const float var = mean_var.y / ncols - mean * mean; - const float inv_std = rsqrtf(var + eps); + // sum up partial sums + mean_var = block_reduce(mean_var, s_sum2); - for (int col = tid; col < ncols; col += block_size) { - dst[col] = (x[col] - mean) * inv_std; + const float mean = mean_var.x / ncols; + const float var = mean_var.y / ncols - mean * mean; + const float inv_std = rsqrtf(var + eps); + + for (int col = tid; col < ncols; col += block_size) { + dstc[col] = (xc[col] - mean) * inv_std; + } + + if constexpr (block_size > WARP_SIZE) { + // sync is needed as we reuse s_sum2 across block_reduce invocations, see #26385 + __syncthreads(); + } + } } } @@ -77,6 +85,8 @@ template static __global__ void rms_norm_f32(const float * x, float * dst, const int ncols, + const int nchannels, + const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, @@ -98,58 +108,68 @@ static __global__ void rms_norm_f32(const float * x, const uint3 add_nchannels_packed = make_uint3(0, 0, 0), const uint3 add_nsamples_packed = make_uint3(0, 0, 0)) { ggml_cuda_pdl_lc(); - const int nrows = gridDim.x; - const int nchannels = gridDim.y; - - const int row = blockIdx.x; - const int channel = blockIdx.y; - const int sample = blockIdx.z; - const int tid = threadIdx.x; + const int nrows = gridDim.x; + const int row = blockIdx.x; + const int tid = threadIdx.x; static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying"); - x += sample*stride_sample + channel*stride_channel + row*stride_row; - dst += ((sample*nchannels + channel)*nrows + row)*ncols; - - if constexpr (do_multiply) { - const uint32_t mul_row = fastmodulo(row, mul_nrows_packed); - const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed); - const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed); - mul += mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row; - } - - if constexpr (do_add) { - const int add_row = fastmodulo(row, add_nrows_packed); - const int add_channel = fastmodulo(channel, add_nchannels_packed); - const int add_sample = fastmodulo(sample, add_nsamples_packed); - add += add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row; - } - - float tmp = 0.0f; // partial sum for thread in warp - - ggml_cuda_pdl_sync(); - for (int col = tid; col < ncols; col += block_size) { - const float xi = x[col]; - tmp += xi * xi; - } - - // sum up partial sums extern __shared__ float s_sum[]; - tmp = block_reduce(tmp, s_sum); - const float mean = tmp / ncols; - const float scale = rsqrtf(mean + eps); + ggml_cuda_pdl_sync(); - for (int col = tid; col < ncols; col += block_size) { - if constexpr (do_multiply && do_add) { - const int mul_col = fastmodulo(col, mul_ncols_packed); - const int add_col = fastmodulo(col, add_ncols_packed); - dst[col] = scale * x[col] * mul[mul_col] + add[add_col]; - } else if constexpr (do_multiply) { - const int mul_col = fastmodulo(col, mul_ncols_packed); - dst[col] = scale * x[col] * mul[mul_col]; - } else { - dst[col] = scale * x[col]; + // grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples + for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) { + for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) { + const float * xc = x + sample*stride_sample + channel*stride_channel + row*stride_row; + float * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols; + + const float * mulc = nullptr; + if constexpr (do_multiply) { + const uint32_t mul_row = fastmodulo(row, mul_nrows_packed); + const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed); + const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed); + mulc = mul + mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row; + } + + const float * addc = nullptr; + if constexpr (do_add) { + const int add_row = fastmodulo(row, add_nrows_packed); + const int add_channel = fastmodulo(channel, add_nchannels_packed); + const int add_sample = fastmodulo(sample, add_nsamples_packed); + addc = add + add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row; + } + + float tmp = 0.0f; // partial sum for thread in warp + + for (int col = tid; col < ncols; col += block_size) { + const float xi = xc[col]; + tmp += xi * xi; + } + + // sum up partial sums + tmp = block_reduce(tmp, s_sum); + + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + + for (int col = tid; col < ncols; col += block_size) { + if constexpr (do_multiply && do_add) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + const int add_col = fastmodulo(col, add_ncols_packed); + dstc[col] = scale * xc[col] * mulc[mul_col] + addc[add_col]; + } else if constexpr (do_multiply) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + dstc[col] = scale * xc[col] * mulc[mul_col]; + } else { + dstc[col] = scale * xc[col]; + } + } + + if constexpr (block_size > WARP_SIZE) { + // sync is needed as we reuse s_sum across block_reduce invocations, see #26385 + __syncthreads(); + } } } } @@ -243,50 +263,58 @@ static __global__ void rms_norm_back_f32( template static __global__ void l2_norm_f32( - const float * x, float * dst, const int ncols, const int64_t stride_row, const int64_t stride_channel, - const int64_t stride_sample, const float eps) { - const int nrows = gridDim.x; - const int nchannels = gridDim.y; + const float * x, float * dst, const int ncols, const int nchannels, const int nsamples, + const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps) { + const int nrows = gridDim.x; + const int row = blockIdx.x; + const int tid = threadIdx.x; - const int row = blockIdx.x; - const int channel = blockIdx.y; - const int sample = blockIdx.z; - const int tid = threadIdx.x; + extern __shared__ float s_sum[]; - x += sample*stride_sample + channel*stride_channel + row*stride_row; - dst += ((sample*nchannels + channel)*nrows + row)*ncols; + ggml_cuda_pdl_lc(); + ggml_cuda_pdl_sync(); - float tmp = 0.0f; // partial sum for thread in warp + // grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples + for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) { + for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) { + const float * xc = x + sample*stride_sample + channel*stride_channel + row*stride_row; + float * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols; - ggml_cuda_pdl_sync(); - for (int col = tid; col < ncols; col += block_size) { - const float xi = x[col]; - tmp += xi * xi; - } + float tmp = 0.0f; // partial sum for thread in warp - // sum up partial sums - extern __shared__ float s_sum[]; - tmp = block_reduce(tmp, s_sum); - ggml_cuda_pdl_lc(); + for (int col = tid; col < ncols; col += block_size) { + const float xi = xc[col]; + tmp += xi * xi; + } - // from https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html - const float scale = rsqrtf(fmaxf(tmp, eps * eps)); + // sum up partial sums + tmp = block_reduce(tmp, s_sum); - for (int col = tid; col < ncols; col += block_size) { - dst[col] = scale * x[col]; + // from https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html + const float scale = rsqrtf(fmaxf(tmp, eps * eps)); + + for (int col = tid; col < ncols; col += block_size) { + dstc[col] = scale * xc[col]; + } + + if constexpr (block_size > WARP_SIZE) { + // sync is needed as we reuse s_sum across block_reduce invocations, see #26385 + __syncthreads(); + } + } } } static void norm_f32_cuda( const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { - const dim3 blocks_num(nrows, nchannels, nsamples); + const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX)); if (ncols < 1024) { const dim3 block_dims(WARP_SIZE, 1, 1); - norm_f32<<>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); + norm_f32<<>>(x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps); } else { const dim3 block_dims(1024, 1, 1); - norm_f32<1024><< WARP_SIZE ? 32 * sizeof(float2): 0, stream>>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps); + norm_f32<1024><< WARP_SIZE ? 32 * sizeof(float2): 0, stream>>>(x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps); } } @@ -304,19 +332,19 @@ static void group_norm_f32_cuda( static void rms_norm_f32_cuda( const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { - const dim3 blocks_num(nrows, nchannels, nsamples); + const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX)); if (ncols < 1024) { const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<256, false>, launch_params, - x, dst, ncols, stride_row, stride_channel, stride_sample, eps, + x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); } else { const dim3 block_dims(1024, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<1024, false>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps, + ggml_cuda_kernel_launch(rms_norm_f32<1024, false>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); @@ -350,7 +378,7 @@ static void rms_norm_mul_f32_cuda(const float * x, const uint32_t add_nsamples, const float eps, cudaStream_t stream) { - const dim3 blocks_num(nrows, nchannels, nsamples); + const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX)); if (mul == nullptr) { rms_norm_f32_cuda(x, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream); return; @@ -364,7 +392,7 @@ static void rms_norm_mul_f32_cuda(const float * x, const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<256, true>, launch_params, - x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); @@ -372,7 +400,7 @@ static void rms_norm_mul_f32_cuda(const float * x, const dim3 block_dims(1024, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<1024, true>, launch_params, - x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); @@ -391,7 +419,7 @@ static void rms_norm_mul_f32_cuda(const float * x, const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims,block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<256, true, true>, launch_params, - x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed); @@ -399,7 +427,7 @@ static void rms_norm_mul_f32_cuda(const float * x, const dim3 block_dims(1024, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<1024, true, true>, launch_params, - x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed); @@ -420,15 +448,15 @@ static void rms_norm_back_f32_cuda(const float * grad, const float * xf, float * static void l2_norm_f32_cuda( const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { - const dim3 blocks_num(nrows, nchannels, nsamples); + const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX)); if (ncols < 1024) { const dim3 block_dims(WARP_SIZE, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, 0, stream}; - ggml_cuda_kernel_launch(l2_norm_f32, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps); + ggml_cuda_kernel_launch(l2_norm_f32, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps); } else { const dim3 block_dims(1024, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(l2_norm_f32<1024>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps); + ggml_cuda_kernel_launch(l2_norm_f32<1024>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps); } } diff --git a/ggml/src/ggml-cuda/rope.cu b/ggml/src/ggml-cuda/rope.cu index e546fb6553c..b4fca778fa6 100644 --- a/ggml/src/ggml-cuda/rope.cu +++ b/ggml/src/ggml-cuda/rope.cu @@ -709,7 +709,7 @@ void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rope // one block per row: block_reduce gives the norm scale, then each thread applies mul and rope to the elements it owns template static __global__ void rms_norm_mul_rope_f32( - const float * x, D * dst, const int ncols, + const float * x, D * dst, const int ncols, const int nchannels, const int nsamples, const int64_t s01, const int64_t s02, const int64_t s03, const int64_t s1, const int64_t s2, const int64_t s3, const float eps, @@ -724,66 +724,76 @@ static __global__ void rms_norm_mul_rope_f32( const int64_t * row_indices, const int set_rows_stride, const bool is_neox) { ggml_cuda_pdl_lc(); - const int row = blockIdx.x; - const int channel = blockIdx.y; - const int sample = blockIdx.z; - const int tid = threadIdx.x; - - x += sample*s03 + channel*s02 + row*s01; - - const uint32_t mul_row = fastmodulo(row, mul_nrows_packed); - const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed); - const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed); - mul += mul_sample*mul_s03 + mul_channel*mul_s02 + mul_row*mul_s01; - - float tmp = 0.0f; - - ggml_cuda_pdl_sync(); - for (int col = tid; col < ncols; col += block_size) { - const float xi = x[col]; - tmp += xi * xi; - } + const int row = blockIdx.x; + const int tid = threadIdx.x; extern __shared__ float s_sum[]; - tmp = block_reduce(tmp, s_sum); - - const float scale = rsqrtf(tmp/ncols + eps); - - int64_t idst = sample*s3 + channel*s2 + row*s1; - if (set_rows_stride != 0) { - idst = row*s1 + row_indices[channel]*set_rows_stride; - } - dst += idst; - - for (int i0 = 2*tid; i0 < ncols; i0 += 2*block_size) { - int ix0; - int ix1; - if (is_neox && i0 < n_dims) { - ix0 = i0/2; - ix1 = i0/2 + n_dims/2; - } else { - ix0 = i0 + 0; - ix1 = i0 + 1; - } - const float x0 = scale * x[ix0] * mul[fastmodulo(ix0, mul_ncols_packed)]; - const float x1 = scale * x[ix1] * mul[fastmodulo(ix1, mul_ncols_packed)]; + ggml_cuda_pdl_sync(); - if (i0 >= n_dims) { - dst[ix0] = ggml_cuda_cast(x0); - dst[ix1] = ggml_cuda_cast(x1); - continue; + // grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples + for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) { + for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) { + const float * xc = x + sample*s03 + channel*s02 + row*s01; + + const uint32_t mul_row = fastmodulo(row, mul_nrows_packed); + const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed); + const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed); + const float * mulc = mul + mul_sample*mul_s03 + mul_channel*mul_s02 + mul_row*mul_s01; + + float tmp = 0.0f; + + for (int col = tid; col < ncols; col += block_size) { + const float xi = xc[col]; + tmp += xi * xi; + } + + tmp = block_reduce(tmp, s_sum); + + const float scale = rsqrtf(tmp/ncols + eps); + + int64_t idst = sample*s3 + channel*s2 + row*s1; + if (set_rows_stride != 0) { + idst = row*s1 + row_indices[channel]*set_rows_stride; + } + D * dstc = dst + idst; + + for (int i0 = 2*tid; i0 < ncols; i0 += 2*block_size) { + int ix0; + int ix1; + if (is_neox && i0 < n_dims) { + ix0 = i0/2; + ix1 = i0/2 + n_dims/2; + } else { + ix0 = i0 + 0; + ix1 = i0 + 1; + } + + const float x0 = scale * xc[ix0] * mulc[fastmodulo(ix0, mul_ncols_packed)]; + const float x1 = scale * xc[ix1] * mulc[fastmodulo(ix1, mul_ncols_packed)]; + + if (i0 >= n_dims) { + dstc[ix0] = ggml_cuda_cast(x0); + dstc[ix1] = ggml_cuda_cast(x1); + continue; + } + + const float theta_base = pos[channel]*powf(theta_scale, i0/2.0f); + const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + + float cos_theta; + float sin_theta; + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + + dstc[ix0] = ggml_cuda_cast(x0*cos_theta - x1*sin_theta); + dstc[ix1] = ggml_cuda_cast(x0*sin_theta + x1*cos_theta); + } + + if constexpr (block_size > WARP_SIZE) { + // sync is needed as we reuse s_sum across block_reduce invocations, see #26385 + __syncthreads(); + } } - - const float theta_base = pos[channel]*powf(theta_scale, i0/2.0f); - const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; - - float cos_theta; - float sin_theta; - rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); - - dst[ix0] = ggml_cuda_cast(x0*cos_theta - x1*sin_theta); - dst[ix1] = ggml_cuda_cast(x0*sin_theta + x1*cos_theta); } } @@ -806,7 +816,7 @@ static void rms_norm_mul_rope_cuda( const bool is_neox, cudaStream_t stream) { GGML_ASSERT(ncols % 2 == 0); - const dim3 blocks_num(nrows, nchannels, nsamples); + const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX)); const float theta_scale = powf(freq_base, -2.0f/n_dims); @@ -820,13 +830,13 @@ static void rms_norm_mul_rope_cuda( const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream}; if (freq_factors == nullptr) { ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, false, D>, launch_params, - x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03, + x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, is_neox); } else { ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, true, D>, launch_params, - x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03, + x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, is_neox); @@ -836,13 +846,13 @@ static void rms_norm_mul_rope_cuda( const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream}; if (freq_factors == nullptr) { ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, false, D>, launch_params, - x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03, + x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, is_neox); } else { ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, true, D>, launch_params, - x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03, + x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, is_neox); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index d61b379282d..77657397b8c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2656,6 +2656,11 @@ struct test_rms_norm_mul_rope : public test_case { bool run_whole_graph() override { return true; } + double max_nmse_err() override { + // large positions amplify the difference between the GPU and CPU trig functions + return ne[2] > 8192 ? 1e-5 : test_case::max_nmse_err(); + } + std::string vars() override { return VARS_TO_STR6(ne, eps, multi_add, set_rows, broadcast, mode); } @@ -9190,6 +9195,12 @@ static std::vector> make_test_cases_eval() { // in-place tests test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 5, 4, 3}, false, 1e-6f, true)); + // shapes below exceed the CUDA gridDim.y/gridDim.z limit of 65535 (#27901) + test_cases.emplace_back(new test_norm (GGML_TYPE_F32, {4, 1, 65536, 1}, false, 1e-6f)); + test_cases.emplace_back(new test_rms_norm (GGML_TYPE_F32, {4, 1, 65536, 1}, false, 1e-6f, false)); + test_cases.emplace_back(new test_rms_norm (GGML_TYPE_F32, {4, 1, 1, 65536}, false, 1e-6f, false)); + test_cases.emplace_back(new test_l2_norm (GGML_TYPE_F32, {4, 1, 65536, 1}, 1e-12f, false, false)); + test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, {4, 1, 65536, 1}, 1e-6f, false, false)); for (float eps : { 0.0f, 1e-6f, 1e-4f, 1e-1f, 1.0f }) { for (uint32_t n : { 64, 1025 }) { @@ -9225,6 +9236,9 @@ static std::vector> make_test_cases_eval() { } } } + // ne[2] > 65535 exceeds the CUDA gridDim.y limit (#27901) + test_cases.emplace_back(new test_rms_norm_mul_rope({4, 1, 65536, 1}, 1e-6f, false, false, false, GGML_ROPE_TYPE_NORMAL)); + test_cases.emplace_back(new test_rms_norm_mul_rope({4, 1, 65536, 1}, 1e-6f, false, true, false, GGML_ROPE_TYPE_NEOX)); for (int64_t d_conv : {3, 4, 9}) { for (int64_t d_inner: {1024, 1536, 2048}) { test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv, d_inner, 1, 1}, {d_conv, d_inner, 1, 1}));