Commit 31953bcf7 for llama.cpp
commit 31953bcf74ba489a8bbb0c936c5fcf24fcf9485c
Author: anujj <ajalota@nvidia.com>
Date: Fri Nov 13 19:11:17 2026 +0530
ggml : add BF16 support to RMS_NORM (CPU, CUDA) (#32135)
* ggml : add BF16 support to RMS_NORM (CPU, CUDA)
* cuda : use one helper for the F32/BF16 RMS_NORM launches
diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp
index 0bfb96009..7e692f1b5 100644
--- a/ggml/src/ggml-cann/ggml-cann.cpp
+++ b/ggml/src/ggml-cann/ggml-cann.cpp
@@ -2601,6 +2601,9 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
}
case GGML_OP_SUM:
return ggml_is_contiguous_rows(op->src[0]);
+ case GGML_OP_RMS_NORM:
+ // the cached gamma of ones does not track its type, so F32 and BF16 can not share it
+ return op->src[0]->type != GGML_TYPE_BF16;
case GGML_OP_L2_NORM:
case GGML_OP_CROSS_ENTROPY_LOSS:
case GGML_OP_DUP:
@@ -2617,7 +2620,6 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
case GGML_OP_SUB:
case GGML_OP_MUL:
case GGML_OP_DIV:
- case GGML_OP_RMS_NORM:
case GGML_OP_SQR:
case GGML_OP_SQRT:
case GGML_OP_CLAMP:
diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp
index eed5d69e3..7c2542741 100644
--- a/ggml/src/ggml-cpu/ops.cpp
+++ b/ggml/src/ggml-cpu/ops.cpp
@@ -4144,6 +4144,50 @@ static void ggml_compute_forward_rms_norm_f32(
}
}
+static void ggml_compute_forward_rms_norm_bf16(
+ const ggml_compute_params * params,
+ ggml_tensor * dst) {
+
+ const ggml_tensor * src0 = dst->src[0];
+
+ GGML_ASSERT(ggml_are_same_shape(src0, dst));
+ GGML_ASSERT(src0->nb[0] == sizeof(ggml_bf16_t));
+
+ const int ith = params->ith;
+ const int nth = params->nth;
+
+ GGML_TENSOR_UNARY_OP_LOCALS
+
+ float eps;
+ memcpy(&eps, dst->op_params, sizeof(float));
+ GGML_ASSERT(eps >= 0.0f);
+
+ for (int64_t i03 = 0; i03 < ne03; i03++) {
+ for (int64_t i02 = 0; i02 < ne02; i02++) {
+ for (int64_t i01 = ith; i01 < ne01; i01 += nth) {
+ const ggml_bf16_t * x = (ggml_bf16_t *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03);
+ ggml_bf16_t * y = (ggml_bf16_t *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3);
+
+ ggml_float sum = 0.0;
+ for (int64_t i00 = 0; i00 < ne00; i00++) {
+ const float v = GGML_BF16_TO_FP32(x[i00]);
+ sum += (ggml_float)(v*v);
+ }
+
+ const float mean = sum/ne00;
+ const float scale = 1.0f/sqrtf(mean + eps);
+
+ // if you hit this, likely you got an inf somewhere earlier
+ assert(scale > 0.0f);
+
+ for (int64_t i00 = 0; i00 < ne00; i00++) {
+ y[i00] = GGML_FP32_TO_BF16(GGML_BF16_TO_FP32(x[i00])*scale);
+ }
+ }
+ }
+ }
+}
+
void ggml_compute_forward_rms_norm(
const ggml_compute_params * params,
ggml_tensor * dst) {
@@ -4155,6 +4199,10 @@ void ggml_compute_forward_rms_norm(
{
ggml_compute_forward_rms_norm_f32<GGML_RMS_NORM_FUSE_OP_NONE>(params, dst);
} break;
+ case GGML_TYPE_BF16:
+ {
+ ggml_compute_forward_rms_norm_bf16(params, dst);
+ } break;
default:
{
GGML_ABORT("fatal error");
diff --git a/ggml/src/ggml-cpu/spacemit/ime.cpp b/ggml/src/ggml-cpu/spacemit/ime.cpp
index dc51ad0aa..0794c115d 100644
--- a/ggml/src/ggml-cpu/spacemit/ime.cpp
+++ b/ggml/src/ggml-cpu/spacemit/ime.cpp
@@ -995,7 +995,8 @@ class tensor_traits_common : public tensor_traits_base {
spacemit_kernels::rvv::forward_rms_norm_f32(params, op);
return true;
default:
- GGML_ABORT("fatal error");
+ ggml_compute_forward_rms_norm(params, op);
+ return true;
}
case GGML_OP_ADD:
switch (op->src[0]->type) {
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index fed14ddfc..540c53614 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -4165,22 +4165,25 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph,
add = cgraph->nodes[node_idx+2];
}
- GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32);
- GGML_ASSERT(rms_norm->type == GGML_TYPE_F32);
-
- //rms norm only supports F32
- if (mul->src[0]->type != GGML_TYPE_F32 ||
- mul->src[1]->type != GGML_TYPE_F32 ||
- mul->type != GGML_TYPE_F32) {
+ // the mul operand is a norm weight, which model files store as F32;
+ // the add operand is an F32 bias or a residual with the type of x
+ const ggml_type type = rms_norm->type;
+ if ((type != GGML_TYPE_F32 && type != GGML_TYPE_BF16) || rms_norm->src[0]->type != type) {
return false;
}
- if (add && (add->src[0]->type != GGML_TYPE_F32 ||
- add->src[1]->type != GGML_TYPE_F32 ||
- add->type != GGML_TYPE_F32) ) {
+ const ggml_tensor * mul_src = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0];
+ if (mul_src->type != GGML_TYPE_F32 || mul->type != type) {
return false;
}
+ if (add) {
+ const ggml_tensor * add_src = add->src[0] == mul ? add->src[1] : add->src[0];
+ if ((add_src->type != GGML_TYPE_F32 && add_src->type != type) || add->type != type) {
+ return false;
+ }
+ }
+
//if rms norm is the B operand, then we don't handle broadcast
if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) {
return false;
@@ -4202,8 +4205,9 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph,
const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
const ggml_tensor * scale = cgraph->nodes[node_idx+1];
- GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32);
- GGML_ASSERT(rms_norm->type == GGML_TYPE_F32);
+ if (rms_norm->src[0]->type != GGML_TYPE_F32 || rms_norm->type != GGML_TYPE_F32) {
+ return false;
+ }
float bias;
memcpy(&bias, (const float *) scale->op_params + 1, sizeof(float));
@@ -6524,9 +6528,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
case GGML_OP_SILU_BACK:
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_NORM:
- case GGML_OP_RMS_NORM:
case GGML_OP_L2_NORM:
return ggml_is_contiguous_rows(op->src[0]);
+ case GGML_OP_RMS_NORM:
+ return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16) && op->type == op->src[0]->type &&
+ ggml_is_contiguous_rows(op->src[0]);
case GGML_OP_RMS_NORM_BACK:
return ggml_is_contiguous(op->src[0]);
case GGML_OP_NONE:
diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu
index 3f0ea6f2a..d6471bcc4 100644
--- a/ggml/src/ggml-cuda/norm.cu
+++ b/ggml/src/ggml-cuda/norm.cu
@@ -1,4 +1,5 @@
#include "norm.cuh"
+#include "convert.cuh"
#include <cstdint>
template <int block_size>
@@ -81,9 +82,35 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr
}
}
-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,
+// 4 consecutive values of x as float4, x + 4*i is aligned to 4 values
+template <typename T>
+static __device__ __forceinline__ float4 rms_norm_load4(const T * x, const int i) {
+ if constexpr (std::is_same_v<T, nv_bfloat16>) {
+ __align__(8) nv_bfloat162 v[2];
+ ggml_cuda_memcpy_1<sizeof(v)>(v, x + 4*i);
+ const float2 lo = ggml_cuda_cast<float2>(v[0]);
+ const float2 hi = ggml_cuda_cast<float2>(v[1]);
+ return make_float4(lo.x, lo.y, hi.x, hi.y);
+ } else {
+ return ((const float4 *) x)[i];
+ }
+}
+
+template <typename T>
+static __device__ __forceinline__ void rms_norm_store4(T * dst, const int i, const float4 r) {
+ if constexpr (std::is_same_v<T, nv_bfloat16>) {
+ const __align__(8) nv_bfloat16 v[4] = {
+ ggml_cuda_cast<nv_bfloat16>(r.x), ggml_cuda_cast<nv_bfloat16>(r.y), ggml_cuda_cast<nv_bfloat16>(r.z), ggml_cuda_cast<nv_bfloat16>(r.w) };
+ ggml_cuda_memcpy_1<sizeof(v)>(dst + 4*i, v);
+ } else {
+ ((float4 *) dst)[i] = r;
+ }
+}
+
+// x and dst are float or nv_bfloat16, mul is float, add is float or the type of x
+template <typename T, int block_size, int vpt, bool do_vec, bool do_multiply = false, bool do_add = false, bool do_scale = false, typename T_add = float>
+static __global__ void rms_norm_f32(const T * x,
+ T * dst,
const int ncols,
const int nchannels,
const int nsamples,
@@ -99,7 +126,7 @@ static __global__ void rms_norm_f32(const float * x,
const uint3 mul_nrows_packed = make_uint3(0, 0, 0),
const uint3 mul_nchannels_packed = make_uint3(0, 0, 0),
const uint3 mul_nsamples_packed = make_uint3(0, 0, 0),
- const float * add = nullptr,
+ const T_add * add = nullptr,
const int64_t add_stride_row = 0,
const int64_t add_stride_channel = 0,
const int64_t add_stride_sample = 0,
@@ -123,8 +150,8 @@ static __global__ void rms_norm_f32(const float * x,
// 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 T * xc = x + sample*stride_sample + channel*stride_channel + row*stride_row;
+ T * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols;
[[maybe_unused]] const float * mulc = nullptr;
if constexpr (do_multiply) {
@@ -134,7 +161,7 @@ static __global__ void rms_norm_f32(const float * x,
mulc = mul + mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row;
}
- [[maybe_unused]] const float * addc = nullptr;
+ [[maybe_unused]] const T_add * addc = nullptr;
if constexpr (do_add) {
const int add_row = fastmodulo(row, add_nrows_packed);
const int add_channel = fastmodulo(channel, add_nchannels_packed);
@@ -144,10 +171,7 @@ static __global__ void rms_norm_f32(const float * x,
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;
@@ -156,7 +180,7 @@ static __global__ void rms_norm_f32(const float * x,
for (int k = 0; k < vpt; k++) {
const int i = tid + k*block_size;
if (i < ncols4) {
- xs[k] = x4[i];
+ xs[k] = rms_norm_load4(xc, 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;
}
}
@@ -176,7 +200,7 @@ static __global__ void rms_norm_f32(const float * x,
r.w *= m.w;
}
if constexpr (do_add) {
- const float4 a = add4[i];
+ const float4 a = rms_norm_load4(addc, i);
r.x += a.x;
r.y += a.y;
r.z += a.z;
@@ -188,14 +212,14 @@ static __global__ void rms_norm_f32(const float * x,
r.z *= scale_out;
r.w *= scale_out;
}
- dst4[i] = r;
+ rms_norm_store4(dstc, i, r);
}
}
} else {
float tmp = 0.0f; // partial sum for thread in warp
for (int col = tid; col < ncols; col += block_size) {
- const float xi = xc[col];
+ const float xi = ggml_cuda_cast<float>(xc[col]);
tmp += xi * xi;
}
@@ -206,17 +230,18 @@ static __global__ void rms_norm_f32(const float * x,
const float scale = rsqrtf(mean + eps);
for (int col = tid; col < ncols; col += block_size) {
+ const float xi = ggml_cuda_cast<float>(xc[col]);
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];
+ dstc[col] = ggml_cuda_cast<T>(scale * xi * mulc[mul_col] + ggml_cuda_cast<float>(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];
+ dstc[col] = ggml_cuda_cast<T>(scale * xi * mulc[mul_col]);
} else if constexpr (do_scale) {
- dstc[col] = scale_out * (scale * xc[col]);
+ dstc[col] = ggml_cuda_cast<T>(scale_out * (scale * xi));
} else {
- dstc[col] = scale * xc[col];
+ dstc[col] = ggml_cuda_cast<T>(scale * xi);
}
}
}
@@ -387,7 +412,7 @@ static void group_norm_f32_cuda(
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, \
+ ggml_cuda_kernel_launch(rms_norm_f32<T, BLOCK_SIZE, VPT, VEC, do_multiply, do_add, do_scale, T_add>, 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, \
@@ -395,10 +420,10 @@ static void group_norm_f32_cuda(
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>
+template <bool do_multiply, bool do_add, bool do_scale, typename T_add = float, typename T>
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 ggml_tensor * x_t, T * 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,
@@ -408,20 +433,20 @@ static void rms_norm_f32_dispatch(
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 T * x = (const T *) 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;
+ const T_add * add = add_t ? (const T_add *) 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);
+ bool vec = ncols % 4 == 0 && ncols > 0 && ((uintptr_t) dst % (4*sizeof(T))) == 0;
+ vec = vec && ggml_cuda_is_aligned(x_t, 4*sizeof(T));
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);
+ && ggml_cuda_is_aligned(add_t, 4*sizeof(T_add));
}
const int n4 = ncols / 4;
@@ -449,9 +474,9 @@ static void rms_norm_f32_dispatch(
}
#undef RMS_NORM_LAUNCH
-template <bool do_scale = false>
+template <bool do_scale = false, typename T>
static void rms_norm_f32_cuda(
- const ggml_tensor * x_t, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
+ const ggml_tensor * x_t, T * 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));
@@ -462,10 +487,11 @@ static void rms_norm_f32_cuda(
stream, scale_out);
}
+template <typename T>
static void rms_norm_mul_f32_cuda(const ggml_tensor * x_t,
const ggml_tensor * mul_t,
const ggml_tensor * add_t,
- float * dst,
+ T * dst,
const int ncols,
const int nrows,
const int nchannels,
@@ -511,12 +537,32 @@ static void rms_norm_mul_f32_cuda(const ggml_tensor * x_t,
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);
- 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);
+ // the add operand is F32 or has the type of x (e.g. a BF16 residual)
+ if (add_t->type == GGML_TYPE_F32) {
+ rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/true, /*do_scale*/false, float>(
+ 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);
+ } else {
+ rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/true, /*do_scale*/false, T>(
+ 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);
+ }
+ }
+}
+
+// calls launch with the data of dst as float * or nv_bfloat16 *, the output types of the RMS_NORM kernels
+template <typename F>
+static void rms_norm_launch_typed(ggml_tensor * dst, const F & launch) {
+ if (dst->type == GGML_TYPE_BF16) {
+ launch((nv_bfloat16 *) dst->data);
+ } else {
+ launch((float *) dst->data);
}
}
@@ -590,11 +636,10 @@ 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];
- float * dst_d = (float *) dst->data;
cudaStream_t stream = ctx.stream();
- GGML_ASSERT(src0->type == GGML_TYPE_F32);
- GGML_ASSERT( dst->type == GGML_TYPE_F32);
+ GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16);
+ GGML_ASSERT( dst->type == src0->type);
GGML_TENSOR_UNARY_OP_LOCALS;
@@ -608,7 +653,9 @@ 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, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream);
+ rms_norm_launch_typed(dst, [&](auto * dst_d) {
+ 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) {
@@ -653,12 +700,12 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor *
GGML_ASSERT(false);
}
- float * dst_d = (float *) mul_tensor->data;
cudaStream_t stream = ctx.stream();
- GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32);
- GGML_ASSERT(dst->type == GGML_TYPE_F32);
- GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32);
+ GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32 || rms_norm_src->type == GGML_TYPE_BF16);
+ GGML_ASSERT(dst->type == rms_norm_src->type);
+ GGML_ASSERT(mul_tensor->type == rms_norm_src->type);
+ GGML_ASSERT(mul_src->type == GGML_TYPE_F32);
GGML_ASSERT(eps >= 0.0f);
const int64_t ne00 = rms_norm_src->ne[0];
@@ -683,14 +730,16 @@ 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(rms_norm_src, mul_src, nullptr, dst_d,
- ne00, ne01, ne02, ne03,
- /*s00*/ s01, s02, s03,
- /*mul_s00*/ mul_s01, mul_s02, mul_s03,
- mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
- /*add_s00*/ 0, 0, 0,
- 0, 0, 0, 0,
- eps, stream);
+ rms_norm_launch_typed(mul_tensor, [&](auto * 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,
+ mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
+ /*add_s00*/ 0, 0, 0,
+ 0, 0, 0, 0,
+ eps, stream);
+ });
}
void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
@@ -722,13 +771,14 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx,
GGML_ASSERT(false);
}
- float * dst_d = (float *) add_tensor->data;
cudaStream_t stream = ctx.stream();
- GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32);
- GGML_ASSERT(dst->type == GGML_TYPE_F32);
- GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32);
- GGML_ASSERT(add_tensor->type == GGML_TYPE_F32);
+ GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32 || rms_norm_src->type == GGML_TYPE_BF16);
+ GGML_ASSERT(dst->type == rms_norm_src->type);
+ GGML_ASSERT(mul_tensor->type == rms_norm_src->type);
+ GGML_ASSERT(add_tensor->type == rms_norm_src->type);
+ GGML_ASSERT(mul_src->type == GGML_TYPE_F32);
+ GGML_ASSERT(add_src->type == GGML_TYPE_F32 || add_src->type == rms_norm_src->type);
GGML_ASSERT(eps >= 0.0f);
const int64_t ne00 = rms_norm_src->ne[0];
@@ -764,14 +814,16 @@ 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(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,
- mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
- /*add_s00*/ add_s01, add_s02, add_s03,
- add_ncols, add_nrows, add_nchannels, add_nsamples,
- eps, stream);
+ rms_norm_launch_typed(add_tensor, [&](auto * 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,
+ mul_ncols, mul_nrows, mul_nchannels, mul_nsamples,
+ /*add_s00*/ add_s01, add_s02, add_s03,
+ add_ncols, add_nrows, add_nchannels, add_nsamples,
+ eps, stream);
+ });
}
void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index 0a07a76bf..0ea4a821a 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -1708,8 +1708,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
case GGML_OP_ARGMAX:
return has_simdgroup_reduction;
case GGML_OP_NORM:
- case GGML_OP_RMS_NORM:
return has_simdgroup_reduction && (ggml_is_contiguous_rows(op->src[0]));
+ case GGML_OP_RMS_NORM:
+ return has_simdgroup_reduction && op->src[0]->type == GGML_TYPE_F32 && (ggml_is_contiguous_rows(op->src[0]));
case GGML_OP_ROPE:
case GGML_OP_ROPE_BACK:
return ggml_rope_get_truncate(op);
diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp
index 7fa1a8e35..a60929abc 100644
--- a/ggml/src/ggml-opencl/ggml-opencl.cpp
+++ b/ggml/src/ggml-opencl/ggml-opencl.cpp
@@ -11966,7 +11966,7 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
case GGML_OP_NORM:
return true;
case GGML_OP_RMS_NORM:
- return op->ne[0] % 4 == 0 && ggml_is_contiguous_rows(op->src[0]);
+ return op->src[0]->type == GGML_TYPE_F32 && op->ne[0] % 4 == 0 && ggml_is_contiguous_rows(op->src[0]);
case GGML_OP_L2_NORM:
return ggml_is_contiguous_rows(op->src[0]);
case GGML_OP_REPEAT:
diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp
index 3aff66724..810a30a2e 100644
--- a/ggml/src/ggml-openvino/ggml-openvino.cpp
+++ b/ggml/src/ggml-openvino/ggml-openvino.cpp
@@ -1455,6 +1455,12 @@ static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) {
}
break;
}
+ case GGML_OP_RMS_NORM: {
+ if (op->type == GGML_TYPE_BF16) {
+ return {false, "RMS_NORM with BF16 type is not supported"};
+ }
+ break;
+ }
case GGML_OP_ADD_ID: {
// Keep support aligned with the CPU backend implementation, which only handles f32 inputs/output and i32 ids.
if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32 ||
diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp
index fdd526425..3bcbad6f8 100644
--- a/ggml/src/ggml-sycl/ggml-sycl.cpp
+++ b/ggml/src/ggml-sycl/ggml-sycl.cpp
@@ -7598,8 +7598,9 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
case GGML_OP_NORM:
case GGML_OP_L2_NORM:
case GGML_OP_GROUP_NORM:
- case GGML_OP_RMS_NORM:
return true;
+ case GGML_OP_RMS_NORM:
+ return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
case GGML_OP_RMS_NORM_BACK:
return ggml_is_contiguous(op->src[0]);
case GGML_OP_SCALE:
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 785d6fa69..bb5aec365 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -16313,11 +16313,11 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
case GGML_OP_VIEW:
case GGML_OP_PERMUTE:
case GGML_OP_TRANSPOSE:
- case GGML_OP_RMS_NORM:
return true;
case GGML_OP_GROUP_NORM:
return ggml_is_contiguous(op->src[0]);
case GGML_OP_NORM:
+ case GGML_OP_RMS_NORM:
case GGML_OP_L2_NORM:
return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
case GGML_OP_ADD:
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 3abdff5fd..b73e838c5 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -2975,6 +2975,7 @@ struct test_rms_norm_mul_rope : public test_case {
int mode;
const bool yarn;
const bool truncate;
+ const ggml_type type; // type of a, the weights are F32
std::string op_desc(ggml_tensor * t) override {
GGML_UNUSED(t);
@@ -2984,18 +2985,22 @@ struct test_rms_norm_mul_rope : public test_case {
bool run_whole_graph() override { return true; }
std::string vars() override {
- return VARS_TO_STR11(ne, eps, multi_add, mul, rope, set_rows, broadcast, mode, set_rows_type, yarn, truncate);
+ std::string s = VARS_TO_STR11(ne, eps, multi_add, mul, rope, set_rows, broadcast, mode, set_rows_type, yarn, truncate);
+ if (type != GGML_TYPE_F32) {
+ s += "," + VAR_TO_STR(type);
+ }
+ return s;
}
test_rms_norm_mul_rope(std::array<int64_t, 4> ne, float eps = 1e-6f, bool multi_add = false,
bool set_rows = false, bool broadcast = false, int mode = GGML_ROPE_TYPE_NORMAL,
bool mul = true, bool rope = true, ggml_type set_rows_type = GGML_TYPE_F16,
- bool yarn = false, bool truncate = true)
+ bool yarn = false, bool truncate = true, ggml_type type = GGML_TYPE_F32)
: ne(ne), eps(eps), multi_add(multi_add), mul(mul), rope(rope), set_rows(set_rows), broadcast(broadcast),
- set_rows_type(set_rows_type), mode(mode), yarn(yarn), truncate(truncate) {}
+ set_rows_type(set_rows_type), mode(mode), yarn(yarn), truncate(truncate), type(type) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
- ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], ne[3]);
+ ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]);
ggml_tensor * b = nullptr;
ggml_tensor * c = nullptr;
@@ -3074,6 +3079,10 @@ struct test_rms_norm_mul_rope : public test_case {
}
double max_nmse_err() override {
+ if (type == GGML_TYPE_BF16) {
+ // the fused kernel rounds to BF16 once, the CPU rounds after each op
+ return 1e-4;
+ }
if (ne[0] == 8192) {
return 5e-6;
}
@@ -4036,6 +4045,11 @@ struct test_rms_norm : public test_case {
}
}
+ double max_nmse_err() override {
+ // the row scale can differ by a few F32 ulps between backends, so a few outputs round to the neighbouring BF16 value
+ return type == GGML_TYPE_BF16 ? 1e-5 : test_case::max_nmse_err();
+ }
+
float grad_eps() override {
return 1.0f;
}
@@ -4090,6 +4104,8 @@ struct test_rms_norm_mul_add : public test_case {
const bool post_mul;
const bool alias_rms_input;
const bool weight_broadcast;
+ const ggml_type type_w; // type of b and d, GGML_TYPE_COUNT -> same as a
+ const ggml_type type_add; // type of c, GGML_TYPE_COUNT -> same as b
std::string op_desc(ggml_tensor * t) override {
GGML_UNUSED(t);
@@ -4099,22 +4115,33 @@ struct test_rms_norm_mul_add : public test_case {
bool run_whole_graph() override { return true; }
std::string vars() override {
- return VARS_TO_STR8(type, ne, eps, broadcast, multi_add, post_mul, alias_rms_input, weight_broadcast);
+ std::string s = VARS_TO_STR8(type, ne, eps, broadcast, multi_add, post_mul, alias_rms_input, weight_broadcast);
+ if (type_w != GGML_TYPE_COUNT) {
+ s += "," + VAR_TO_STR(type_w);
+ }
+ if (type_add != GGML_TYPE_COUNT) {
+ s += "," + VAR_TO_STR(type_add);
+ }
+ return s;
}
test_rms_norm_mul_add(ggml_type type = GGML_TYPE_F32,
std::array<int64_t, 4> ne = {64, 5, 4, 3},
float eps = 1e-6f, bool broadcast = false, bool multi_add = false, bool post_mul = false,
- bool alias_rms_input = false, bool weight_broadcast = false)
+ bool alias_rms_input = false, bool weight_broadcast = false, ggml_type type_w = GGML_TYPE_COUNT,
+ ggml_type type_add = GGML_TYPE_COUNT)
: type(type), ne(ne), eps(eps), broadcast(broadcast), multi_add(multi_add), post_mul(post_mul),
- alias_rms_input(alias_rms_input), weight_broadcast(weight_broadcast) {}
+ alias_rms_input(alias_rms_input), weight_broadcast(weight_broadcast), type_w(type_w), type_add(type_add) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
std::array<int64_t, 4> broadcast_dims = {ne[0]*2, ne[1]*3, ne[2]*3, ne[3]*4};
+ const ggml_type tw = type_w == GGML_TYPE_COUNT ? type : type_w;
+ const ggml_type ta = type_add == GGML_TYPE_COUNT ? tw : type_add;
+
ggml_tensor * a = ggml_new_tensor(ctx, type, 4, broadcast ? broadcast_dims.data() : ne.data());
- ggml_tensor * b = weight_broadcast ? ggml_new_tensor_1d(ctx, type, ne[0]) : ggml_new_tensor(ctx, type, 4, ne.data());
- ggml_tensor * c = ggml_new_tensor(ctx, type, 4, ne.data());
+ ggml_tensor * b = weight_broadcast ? ggml_new_tensor_1d(ctx, tw, ne[0]) : ggml_new_tensor(ctx, tw, 4, ne.data());
+ ggml_tensor * c = ggml_new_tensor(ctx, ta, 4, ne.data());
ggml_tensor * d = nullptr;
ggml_set_param(a);
@@ -4127,7 +4154,7 @@ struct test_rms_norm_mul_add : public test_case {
// Use a, b and c early, so we don't end up with an OP_NONE between rms_norm and mul
a = ggml_add(ctx, ggml_add(ctx, a, b), c);
if (post_mul) {
- d = ggml_new_tensor_1d(ctx, type, 1);
+ d = ggml_new_tensor_1d(ctx, tw, 1);
ggml_set_param(d);
ggml_set_name(d, "d");
a = ggml_add(ctx, a, d);
@@ -4151,6 +4178,11 @@ struct test_rms_norm_mul_add : public test_case {
}
}
+ double max_nmse_err() override {
+ // fused ops round to BF16 once, the CPU rounds after each op
+ return type == GGML_TYPE_BF16 ? 1e-4 : test_case::max_nmse_err();
+ }
+
float grad_eps() override {
return 1.0f;
}
@@ -11508,6 +11540,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
for (bool v : { false, true }) {
test_cases.emplace_back(new test_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, v, eps));
test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, v, eps));
+ test_cases.emplace_back(new test_rms_norm(GGML_TYPE_BF16, { n, 5, 4, 3 }, v, eps));
}
test_cases.emplace_back(new test_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, false, eps, true));
test_cases.emplace_back(new test_norm_scale(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false, 1.5f));
@@ -11526,12 +11559,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
for (bool v : { false, true }) {
test_cases.emplace_back(new test_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, v, eps));
test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, v, eps));
+ test_cases.emplace_back(new test_rms_norm(GGML_TYPE_BF16, { n, 5, 4, 3 }, v, eps));
}
}
}
// in-place tests
test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 5, 4, 3}, false, 1e-6f, true));
+ test_cases.emplace_back(new test_rms_norm(GGML_TYPE_BF16, {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));
@@ -11545,6 +11580,19 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, {n, 3, 2, 1}, 1e-6f, false, false, false, false, true));
}
+ // BF16 x with F32 weights, n > 1 because the NMSE of a single BF16 output is too noisy
+ // the widths cover every CUDA block size, with and without 4-value loads
+ for (uint32_t n : {127, 128, 512, 1024, 2048, 2816, 4096, 8192, 16385}) {
+ test_cases.emplace_back(new test_rms_norm(GGML_TYPE_BF16, {n, 3, 2, 1}));
+ test_cases.emplace_back(new test_rms_norm_mul_rope({n, 3, 2, 1}, 1e-6f, false, false, true, GGML_ROPE_TYPE_NORMAL, true, false, GGML_TYPE_F16, false, true, GGML_TYPE_BF16));
+ test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_BF16, {n, 3, 2, 1}, 1e-6f, false, false, false, false, true, GGML_TYPE_F32));
+ // BF16 residual added after the norm (post-norm models, e.g. Gemma)
+ test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_BF16, {n, 3, 2, 1}, 1e-6f, false, false, false, false, true, GGML_TYPE_F32, GGML_TYPE_BF16));
+ for (bool multi_add : {false, true}) {
+ test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_BF16, {n, 1, 1, 1}, 1e-6f, false, multi_add, false, false, false, GGML_TYPE_F32));
+ }
+ }
+
for (ggml_type set_rows_type : { GGML_TYPE_F32, GGML_TYPE_F16 }) {
test_cases.emplace_back(new test_rms_norm_mul_rope({ 256, 1, 1, 1 }, 1e-6f, false, true, false, GGML_ROPE_TYPE_NORMAL, false, false, set_rows_type));
test_cases.emplace_back(new test_rms_norm_mul_rope({ 128, 4, 3, 1 }, 1e-6f, false, true, false, GGML_ROPE_TYPE_NORMAL, false, false, set_rows_type));
@@ -11559,6 +11607,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
for (uint32_t n : { 64, 1025 }) {
test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false));
test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, true));
+ test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_BF16, { n, 5, 4, 3 }, eps, false, false, false, false, false, GGML_TYPE_F32));
+ test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_BF16, { n, 5, 4, 3 }, eps, true, false, false, false, false, GGML_TYPE_F32));
+ test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_BF16, { n, 5, 4, 3 }, eps, true, false, false, false, false, GGML_TYPE_F32, GGML_TYPE_BF16));
test_cases.emplace_back(new test_norm_mul_add(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false));
test_cases.emplace_back(new test_norm_mul_add(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, true));
test_cases.emplace_back(new test_add_rms_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false));
@@ -11586,6 +11637,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 256, 4, 3, 2 }, 1e-6f, false, false, true, false, true));
test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 1536, 1, 1, 1 }, 1e-6f, false, false, false, true));
test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_F32, { 256, 4, 1, 1 }, 1e-6f, false, false, false, true));
+ test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_BF16, { 1536, 1, 1, 1 }, 1e-6f, false, false, true, false, false, GGML_TYPE_F32));
+ // the ADD operand is the BF16 input itself, added in place
+ test_cases.emplace_back(new test_rms_norm_mul_add(GGML_TYPE_BF16, { 1536, 1, 1, 1 }, 1e-6f, false, false, false, true, false, GGML_TYPE_F32));
test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 7, 2}));
test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 7, 2}, 1e-6f, false, true));
@@ -13492,8 +13546,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
// RMS_NORM and fused variants at decode and prefill shapes
for (int64_t n_tokens : {1, 512}) {
- test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {4096, n_tokens, 1, 1}));
- test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {128, 32*n_tokens, 1, 1}));
+ for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_BF16}) {
+ test_cases.emplace_back(new test_rms_norm(type, {4096, n_tokens, 1, 1}));
+ test_cases.emplace_back(new test_rms_norm(type, {128, 32*n_tokens, 1, 1}));
+ }
test_cases.emplace_back(new test_rms_norm_mul_rope({4096, n_tokens, 1, 1}, 1e-6f, false, false, true, GGML_ROPE_TYPE_NORMAL, true, false));
test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, n_tokens, 1}, 1e-6f, false, false, true, GGML_ROPE_TYPE_NEOX));
test_cases.emplace_back(new test_rms_norm_mul_rope({128, 8, n_tokens, 1}, 1e-6f, false, true, true, GGML_ROPE_TYPE_NEOX));