Skip to content

Commit cd57f51

Browse files
lhezlhez
committed
opencl: fix SOFT_MAX
* Add fp16 variant
1 parent 2923cc9 commit cd57f51

File tree

2 files changed

+147
-2
lines changed

2 files changed

+147
-2
lines changed

‎ggml/src/ggml-opencl/ggml-opencl.cpp‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,7 @@ struct ggml_backend_opencl_context {
143143
cl_kernel kernel_rms_norm;
144144
cl_kernel kernel_diag_mask_inf, kernel_diag_mask_inf_8;
145145
cl_kernel kernel_soft_max, kernel_soft_max_4;
146+
cl_kernel kernel_soft_max_f16, kernel_soft_max_4_f16;
146147
cl_kernel kernel_get_rows_f32, kernel_get_rows_f16, kernel_get_rows_q4_0;
147148
cl_kernel kernel_rope_norm_f32, kernel_rope_norm_f16, kernel_rope_neox_f32, kernel_rope_neox_f16;
148149
cl_kernel kernel_cpy_f16_f16, kernel_cpy_f16_f32, kernel_cpy_f32_f16, kernel_cpy_f32_f32;
@@ -614,6 +615,8 @@ static ggml_backend_opencl_context * ggml_cl2_init(ggml_backend_dev_t dev) {
614615
CL_CHECK((backend_ctx->kernel_diag_mask_inf_8 = clCreateKernel(backend_ctx->program, "kernel_diag_mask_inf_8", &err), err));
615616
CL_CHECK((backend_ctx->kernel_soft_max = clCreateKernel(backend_ctx->program, "kernel_soft_max", &err), err));
616617
CL_CHECK((backend_ctx->kernel_soft_max_4 = clCreateKernel(backend_ctx->program, "kernel_soft_max_4", &err), err));
618+
CL_CHECK((backend_ctx->kernel_soft_max_f16 = clCreateKernel(backend_ctx->program, "kernel_soft_max_f16", &err), err));
619+
CL_CHECK((backend_ctx->kernel_soft_max_4_f16 = clCreateKernel(backend_ctx->program, "kernel_soft_max_4_f16", &err), err));
617620
CL_CHECK((backend_ctx->kernel_rope_norm_f32 = clCreateKernel(backend_ctx->program, "kernel_rope_norm_f32", &err), err));
618621
CL_CHECK((backend_ctx->kernel_rope_norm_f16 = clCreateKernel(backend_ctx->program, "kernel_rope_norm_f16", &err), err));
619622
CL_CHECK((backend_ctx->kernel_rope_neox_f32 = clCreateKernel(backend_ctx->program, "kernel_rope_neox_f32", &err), err));
@@ -3674,6 +3677,8 @@ static void ggml_cl_soft_max(ggml_backend_t backend, const ggml_tensor * src0, c
36743677
const float m0 = powf(2.0f, -(max_bias ) / n_head_log2);
36753678
const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2);
36763679

3680+
const bool use_f16 = (src1 && src1->type == GGML_TYPE_F16);
3681+
36773682
// Local size must be wave size. Each workgroup is a wave, working on a row,
36783683
// where a row corresponds to leading dimension.
36793684
int nth = MIN(32, ne00);
@@ -3691,9 +3696,17 @@ static void ggml_cl_soft_max(ggml_backend_t backend, const ggml_tensor * src0, c
36913696
cl_kernel kernel;
36923697

36933698
if (ne00%4 == 0) {
3694-
kernel = backend_ctx->kernel_soft_max_4;
3699+
if (use_f16) {
3700+
kernel = backend_ctx->kernel_soft_max_4_f16;
3701+
} else {
3702+
kernel = backend_ctx->kernel_soft_max_4;
3703+
}
36953704
} else {
3696-
kernel = backend_ctx->kernel_soft_max;
3705+
if (use_f16) {
3706+
kernel = backend_ctx->kernel_soft_max_f16;
3707+
} else {
3708+
kernel = backend_ctx->kernel_soft_max;
3709+
}
36973710
}
36983711

36993712
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra0->data_device));

‎ggml/src/ggml-opencl/kernels/ggml-opencl.cl‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -811,6 +811,138 @@ kernel void kernel_soft_max_4(
811811
}
812812
}
813813

814+
kernel void kernel_soft_max_f16(
815+
global float * src0,
816+
ulong offset0,
817+
global half * src1,
818+
ulong offset1,
819+
global float * dst,
820+
ulong offsetd,
821+
int ne00,
822+
int ne01,
823+
int ne02,
824+
float scale,
825+
float max_bias,
826+
float m0,
827+
float m1,
828+
int n_head_log2
829+
) {
830+
src0 = (global float *)((global char *)src0 + offset0);
831+
src1 = (global half *)((global char *)src1 + offset1);
832+
dst = (global float *)((global char *)dst + offsetd);
833+
834+
int i03 = get_group_id(2);
835+
int i02 = get_group_id(1);
836+
int i01 = get_group_id(0);
837+
838+
global float * psrc0 = src0 + i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00;
839+
global half * pmask = (global char *)src1 != (global char *)src0 ? src1 + i01*ne00 : 0;
840+
global float * pdst = dst + i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00;
841+
842+
float slope = 1.0f;
843+
844+
// ALiBi
845+
if (max_bias > 0.0f) {
846+
int h = i02;
847+
848+
float base = h < n_head_log2 ? m0 : m1;
849+
int exp = h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1;
850+
851+
slope = pow(base, exp);
852+
}
853+
854+
// parallel max
855+
float lmax = -INFINITY;
856+
for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
857+
lmax = fmax(lmax, psrc0[i00]*scale + (pmask ? slope*pmask[i00] : 0.0f));
858+
}
859+
float max = sub_group_reduce_max(lmax);
860+
861+
// parallel sum
862+
float lsum = 0.0f;
863+
for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
864+
float exp_psrc0 = exp((psrc0[i00]*scale + (pmask ? slope*pmask[i00] : 0.0f)) - max);
865+
lsum += exp_psrc0;
866+
// Remember the result of exp here. exp is expensive, so we really do not
867+
// wish to compute it twice.
868+
pdst[i00] = exp_psrc0;
869+
}
870+
871+
const float sum = sub_group_reduce_add(lsum);
872+
873+
for (int i00 = get_local_id(0); i00 < ne00; i00 += get_local_size(0)) {
874+
pdst[i00] /= sum;
875+
}
876+
}
877+
878+
#ifdef ADRENO_GPU
879+
REQD_SUBGROUP_SIZE_64
880+
#endif
881+
kernel void kernel_soft_max_4_f16(
882+
global float * src0,
883+
ulong offset0,
884+
global half * src1,
885+
ulong offset1,
886+
global float * dst,
887+
ulong offsetd,
888+
int ne00,
889+
int ne01,
890+
int ne02,
891+
float scale,
892+
float max_bias,
893+
float m0,
894+
float m1,
895+
int n_head_log2
896+
) {
897+
src0 = (global float *)((global char *)src0 + offset0);
898+
src1 = (global half *)((global char *)src1 + offset1);
899+
dst = (global float *)((global char *)dst + offsetd);
900+
901+
int i03 = get_group_id(2);
902+
int i02 = get_group_id(1);
903+
int i01 = get_group_id(0);
904+
905+
global float4 * psrc4 = (global float4 *)(src0 + i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00);
906+
global half4 * pmask = (global char *)src1 != (global char *)src0 ? (global half4 *)(src1 + i01*ne00) : 0;
907+
global float4 * pdst4 = (global float4 *)(dst + i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00);
908+
909+
float slope = 1.0f;
910+
911+
// ALiBi
912+
if (max_bias > 0.0f) {
913+
int h = i02;
914+
915+
float base = h < n_head_log2 ? m0 : m1;
916+
int exp = h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1;
917+
918+
slope = pow(base, exp);
919+
}
920+
921+
// parallel max
922+
float4 lmax4 = -INFINITY;
923+
for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
924+
lmax4 = fmax(lmax4, psrc4[i00]*scale + slope*(pmask ? convert_float4(pmask[i00]) : 0.0f));
925+
}
926+
float lmax = fmax(fmax(lmax4.s0, lmax4.s1), fmax(lmax4.s2, lmax4.s3));
927+
928+
const float max = sub_group_reduce_max(lmax);
929+
930+
// parallel sum
931+
float4 lsum4 = 0.0f;
932+
for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
933+
const float4 exp_psrc4 = exp((psrc4[i00]*scale + slope*(pmask ? convert_float4(pmask[i00]) : 0.0f)) - max);
934+
lsum4 += exp_psrc4;
935+
pdst4[i00] = exp_psrc4;
936+
}
937+
float lsum = lsum4.s0 + lsum4.s1 + lsum4.s2 + lsum4.s3;
938+
939+
const float sum = sub_group_reduce_add(lsum);
940+
941+
for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
942+
pdst4[i00] /= sum;
943+
}
944+
}
945+
814946
//------------------------------------------------------------------------------
815947
// kernel_rope
816948
//------------------------------------------------------------------------------

0 commit comments

Comments
 (0)