Commit 623c16393 for llama.cpp
commit 623c16393c7d53a37df8bb69c6ea8ed6e5171f62
Author: uvos <carl@uvos.xyz>
Date: Fri Oct 16 09:07:34 2026 +0200
CUDA: HIP: rms_norm: add variants that load float4 at a time and store x in registers (#29720)
* CUDA: HIP: rms_norm: load float4 at a time and store x in registers
* Add template paramter comments
* use ggml_cuda_is_aligned for src
* nits
* also use ggml_cuda_is_aligned for src in fused variants
* fix nits
diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu
index d804b933b..3f0ea6f2a 100644
--- a/ggml/src/ggml-cuda/norm.cu
+++ b/ggml/src/ggml-cuda/norm.cu
@@ -81,7 +81,7 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr
}
}
-template <int block_size, bool do_multiply = false, bool do_add = false, bool do_scale = false>
+template <int block_size, int vpt, bool do_vec, bool do_multiply = false, bool do_add = false, bool do_scale = false>
static __global__ void rms_norm_f32(const float * x,
float * dst,
const int ncols,
@@ -142,31 +142,82 @@ static __global__ void rms_norm_f32(const float * x,
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<block_reduce_method::SUM, block_size>(tmp, s_sum);
+ if constexpr (do_vec) {
+ const int ncols4 = ncols / 4;
+ const float4 * x4 = (const float4 *) xc;
+ const float4 * mul4 = (const float4 *) mulc;
+ const float4 * add4 = (const float4 *) addc;
+ float4 * dst4 = (float4 *) dstc;
+
+ float4 xs[vpt];
+ float tmp = 0.0f;
+
+#pragma unroll
+ for (int k = 0; k < vpt; k++) {
+ const int i = tid + k*block_size;
+ if (i < ncols4) {
+ xs[k] = x4[i];
+ tmp += xs[k].x*xs[k].x + xs[k].y*xs[k].y + xs[k].z*xs[k].z + xs[k].w*xs[k].w;
+ }
+ }
+ tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
+ const float scale = rsqrtf(tmp / ncols + eps);
+
+#pragma unroll
+ for (int k = 0; k < vpt; k++) {
+ const int i = tid + k*block_size;
+ if (i < ncols4) {
+ float4 r = make_float4(scale*xs[k].x, scale*xs[k].y, scale*xs[k].z, scale*xs[k].w);
+ if constexpr (do_multiply) {
+ const float4 m = mul4[i];
+ r.x *= m.x;
+ r.y *= m.y;
+ r.z *= m.z;
+ r.w *= m.w;
+ }
+ if constexpr (do_add) {
+ const float4 a = add4[i];
+ r.x += a.x;
+ r.y += a.y;
+ r.z += a.z;
+ r.w += a.w;
+ }
+ if constexpr (do_scale) {
+ r.x *= scale_out;
+ r.y *= scale_out;
+ r.z *= scale_out;
+ r.w *= scale_out;
+ }
+ dst4[i] = r;
+ }
+ }
+ } else {
+ float tmp = 0.0f; // partial sum for thread in warp
- const float mean = tmp / ncols;
- const float scale = rsqrtf(mean + eps);
+ for (int col = tid; col < ncols; col += block_size) {
+ const float xi = xc[col];
+ tmp += xi * xi;
+ }
- 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 if constexpr (do_scale) {
- dstc[col] = scale_out * (scale * xc[col]);
- } else {
- dstc[col] = scale * xc[col];
+ // sum up partial sums
+ tmp = block_reduce<block_reduce_method::SUM, block_size>(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 if constexpr (do_scale) {
+ dstc[col] = scale_out * (scale * xc[col]);
+ } else {
+ dstc[col] = scale * xc[col];
+ }
}
}
@@ -332,34 +383,89 @@ static void group_norm_f32_cuda(
}
}
+#define RMS_NORM_LAUNCH(BLOCK_SIZE, VPT, VEC) \
+ do { \
+ const dim3 block_dims(BLOCK_SIZE, 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<BLOCK_SIZE, VPT, VEC, do_multiply, do_add, do_scale>, launch_params, \
+ 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, scale_out); \
+ } while (0)
+
+template <bool do_multiply, bool do_add, bool do_scale>
+static void rms_norm_f32_dispatch(
+ const dim3 blocks_num,
+ const ggml_tensor * x_t, 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 ggml_tensor * mul_t, const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample,
+ const uint3 mul_ncols_packed, const uint3 mul_nrows_packed, const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed,
+ const ggml_tensor * add_t, const int64_t add_stride_row, const int64_t add_stride_channel, const int64_t add_stride_sample,
+ const uint3 add_ncols_packed, const uint3 add_nrows_packed, const uint3 add_nchannels_packed, const uint3 add_nsamples_packed,
+ cudaStream_t stream, const float scale_out = 1.0f) {
+ static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying");
+ static_assert(!do_scale || !do_multiply, "fusing scale is not supported with multiplying");
+
+ const float * x = (const float *) x_t->data;
+ const float * mul = mul_t ? (const float *) mul_t->data : nullptr;
+ const float * add = add_t ? (const float *) add_t->data : nullptr;
+
+ // dst is written at row * ncols, i.e. always at a multiple of ncols, so only its base needs checking here
+ bool vec = ncols % 4 == 0 && ncols > 0 && ((uintptr_t) dst & 15) == 0;
+ vec = vec && ggml_cuda_is_aligned(x_t, 16);
+ if constexpr (do_multiply) {
+ vec = vec && mul_ncols_packed.z == (uint32_t) ncols
+ && ggml_cuda_is_aligned(mul_t, 16);
+ }
+ if constexpr (do_add) {
+ vec = vec && add_ncols_packed.z == (uint32_t) ncols
+ && ggml_cuda_is_aligned(add_t, 16);
+ }
+
+ const int n4 = ncols / 4;
+ if (vec) {
+ if (n4 <= 32) {
+ RMS_NORM_LAUNCH(32, 1, true);
+ } else if (n4 <= 128) {
+ RMS_NORM_LAUNCH(128, 1, true);
+ } else if (n4 <= 256) {
+ RMS_NORM_LAUNCH(256, 1, true);
+ } else if (n4 <= 512) {
+ RMS_NORM_LAUNCH(512, 1, true);
+ } else if (n4 <= 1024) {
+ RMS_NORM_LAUNCH(1024, 1, true);
+ } else if (n4 <= 4096) {
+ RMS_NORM_LAUNCH(1024, 4, true);
+ } else {
+ RMS_NORM_LAUNCH(1024, 1, false);
+ }
+ } else if (ncols < 1024) {
+ RMS_NORM_LAUNCH(256, 1, false);
+ } else {
+ RMS_NORM_LAUNCH(1024, 1, false);
+ }
+}
+#undef RMS_NORM_LAUNCH
+
template <bool do_scale = false>
static void rms_norm_f32_cuda(
- const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
+ const ggml_tensor * x_t, 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 float scale_out = 1.0f) {
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, false, do_scale>, launch_params,
- x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
- // underlying cudaLaunchKernelEx does not support default params
+ rms_norm_f32_dispatch</*do_multiply*/false, /*do_add*/false, do_scale>(
+ blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
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), scale_out);
- } 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, false, do_scale>, 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), scale_out);
- }
+ stream, scale_out);
}
-static void rms_norm_mul_f32_cuda(const float * x,
- const float * mul,
- const float * add,
- float * dst,
+static void rms_norm_mul_f32_cuda(const ggml_tensor * x_t,
+ const ggml_tensor * mul_t,
+ const ggml_tensor * add_t,
+ float * dst,
const int ncols,
const int nrows,
const int nchannels,
@@ -384,59 +490,33 @@ static void rms_norm_mul_f32_cuda(const float * x,
const float eps,
cudaStream_t stream) {
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);
+ if (mul_t == nullptr) {
+ rms_norm_f32_cuda(x_t, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream);
return;
}
- if (add == nullptr) {
- const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols);
- const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows);
- const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
- const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples);
- if (ncols < 1024) {
- 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, 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), 1.0f);
- } 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, true>, launch_params,
- 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), 1.0f);
- }
+ const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols);
+ const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows);
+ const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
+ const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples);
+ const uint3 zero_packed = make_uint3(0, 0, 0);
+
+ if (add_t == nullptr) {
+ rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/false, /*do_scale*/false>(
+ blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
+ mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample,
+ mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+ nullptr, 0, 0, 0, zero_packed, zero_packed, zero_packed, zero_packed, stream);
} else {
- const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols);
- const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows);
- const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
- const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples);
-
const uint3 add_ncols_packed = init_fastdiv_values(add_ncols);
const uint3 add_nrows_packed = init_fastdiv_values(add_nrows);
const uint3 add_nchannels_packed = init_fastdiv_values(add_nchannels);
const uint3 add_nsamples_packed = init_fastdiv_values(add_nsamples);
- if (ncols < 1024) {
- 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, 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, 1.0f);
- } 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, true, true>, launch_params,
- 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, 1.0f);
- }
+ rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/true, /*do_scale*/false>(
+ blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
+ mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample,
+ mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+ add_t, add_stride_row, add_stride_channel, add_stride_sample,
+ add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, stream);
}
}
@@ -510,7 +590,6 @@ void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
- const float * src0_d = (const float *) src0->data;
float * dst_d = (float *) dst->data;
cudaStream_t stream = ctx.stream();
@@ -529,12 +608,11 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const int64_t s02 = nb02 / ts0;
const int64_t s03 = nb03 / ts0;
- rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream);
+ rms_norm_f32_cuda(src0, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream);
}
void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor) {
const ggml_tensor * src0 = dst->src[0];
- const float * src0_d = (const float *) src0->data;
float * dst_d = (float *) scale_tensor->data;
cudaStream_t stream = ctx.stream();
@@ -556,7 +634,7 @@ void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_ten
const int64_t s02 = nb02 / ts0;
const int64_t s03 = nb03 / ts0;
- rms_norm_f32_cuda<true>(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream, scale);
+ rms_norm_f32_cuda<true>(src0, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream, scale);
}
void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) {
@@ -565,15 +643,11 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor *
memcpy(&eps, dst->op_params, sizeof(float));
- const float * src0_d = (const float *) rms_norm_src->data;
- const float * mul_d = nullptr;
const ggml_tensor * mul_src = nullptr;
if (mul_tensor->src[0] == dst) {
- mul_d = (float *) mul_tensor->src[1]->data;
mul_src = mul_tensor->src[1];
} else if(mul_tensor->src[1] == dst) {
- mul_d = (float *) mul_tensor->src[0]->data;
mul_src = mul_tensor->src[0];
} else {
GGML_ASSERT(false);
@@ -609,7 +683,7 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor *
const int mul_nchannels = mul_src->ne[2];
const int mul_nsamples = mul_src->ne[3];
- rms_norm_mul_f32_cuda(src0_d, mul_d, nullptr, dst_d,
+ rms_norm_mul_f32_cuda(rms_norm_src, mul_src, nullptr, dst_d,
ne00, ne01, ne02, ne03,
/*s00*/ s01, s02, s03,
/*mul_s00*/ mul_s01, mul_s02, mul_s03,
@@ -628,28 +702,21 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
memcpy(&eps, dst->op_params, sizeof(float));
- const float * src0_d = (const float *) rms_norm_src->data;
- const float * mul_d = nullptr;
const ggml_tensor * mul_src = nullptr;
if (mul_tensor->src[0] == dst) {
- mul_d = (float *) mul_tensor->src[1]->data;
mul_src = mul_tensor->src[1];
} else if (mul_tensor->src[1] == dst) {
- mul_d = (float *) mul_tensor->src[0]->data;
mul_src = mul_tensor->src[0];
} else {
GGML_ASSERT(false);
}
- const float * add_d = nullptr;
const ggml_tensor * add_src = nullptr;
if (add_tensor->src[0] == mul_tensor) {
- add_d = (float *) add_tensor->src[1]->data;
add_src = add_tensor->src[1];
} else if (add_tensor->src[1] == mul_tensor) {
- add_d = (float *) add_tensor->src[0]->data;
add_src = add_tensor->src[0];
} else {
GGML_ASSERT(false);
@@ -697,7 +764,7 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
const int add_nchannels = add_src->ne[2];
const int add_nsamples = add_src->ne[3];
- rms_norm_mul_f32_cuda(src0_d, mul_d,add_d,dst_d,
+ rms_norm_mul_f32_cuda(rms_norm_src, mul_src,add_src,dst_d,
ne00,ne01, ne02, ne03,
/*s00*/ s01, s02, s03,
/*mul_s00*/ mul_s01, mul_s02, mul_s03,