Skip to content

Commit 0711a5f

Browse files
apage43apage43
authored
metal : add norm, cpy f16->f16, alibi kernels (#1823)
1 parent fc45a81 commit 0711a5f

File tree

2 files changed

+222
-0
lines changed

2 files changed

+222
-0
lines changed

‎ggml-metal.m

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
GGML_METAL_DECL_KERNEL(get_rows_q5_k);
5858
GGML_METAL_DECL_KERNEL(get_rows_q6_k);
5959
GGML_METAL_DECL_KERNEL(rms_norm);
60+
GGML_METAL_DECL_KERNEL(norm);
6061
GGML_METAL_DECL_KERNEL(mul_mat_f16_f32);
6162
GGML_METAL_DECL_KERNEL(mul_mat_q4_0_f32);
6263
GGML_METAL_DECL_KERNEL(mul_mat_q4_1_f32);
@@ -66,8 +67,10 @@
6667
GGML_METAL_DECL_KERNEL(mul_mat_q5_k_f32);
6768
GGML_METAL_DECL_KERNEL(mul_mat_q6_k_f32);
6869
GGML_METAL_DECL_KERNEL(rope);
70+
GGML_METAL_DECL_KERNEL(alibi_f32);
6971
GGML_METAL_DECL_KERNEL(cpy_f32_f16);
7072
GGML_METAL_DECL_KERNEL(cpy_f32_f32);
73+
GGML_METAL_DECL_KERNEL(cpy_f16_f16);
7174

7275
#undef GGML_METAL_DECL_KERNEL
7376
};
@@ -162,6 +165,7 @@ @implementation GGMLMetalClass
162165
GGML_METAL_ADD_KERNEL(get_rows_q5_k);
163166
GGML_METAL_ADD_KERNEL(get_rows_q6_k);
164167
GGML_METAL_ADD_KERNEL(rms_norm);
168+
GGML_METAL_ADD_KERNEL(norm);
165169
GGML_METAL_ADD_KERNEL(mul_mat_f16_f32);
166170
GGML_METAL_ADD_KERNEL(mul_mat_q4_0_f32);
167171
GGML_METAL_ADD_KERNEL(mul_mat_q4_1_f32);
@@ -171,8 +175,10 @@ @implementation GGMLMetalClass
171175
GGML_METAL_ADD_KERNEL(mul_mat_q5_k_f32);
172176
GGML_METAL_ADD_KERNEL(mul_mat_q6_k_f32);
173177
GGML_METAL_ADD_KERNEL(rope);
178+
GGML_METAL_ADD_KERNEL(alibi_f32);
174179
GGML_METAL_ADD_KERNEL(cpy_f32_f16);
175180
GGML_METAL_ADD_KERNEL(cpy_f32_f32);
181+
GGML_METAL_ADD_KERNEL(cpy_f16_f16);
176182

177183
#undef GGML_METAL_ADD_KERNEL
178184
}
@@ -735,6 +741,65 @@ void ggml_metal_graph_compute(
735741

736742
[encoder dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
737743
} break;
744+
case GGML_OP_NORM:
745+
{
746+
if (encoder == nil) {
747+
encoder = [command_buffer computeCommandEncoder];
748+
}
749+
750+
const float eps = 1e-5f;
751+
752+
const int nth = 256;
753+
754+
[encoder setComputePipelineState:ctx->pipeline_norm];
755+
[encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
756+
[encoder setBuffer:id_dst offset:offs_dst atIndex:1];
757+
[encoder setBytes:&ne00 length:sizeof( int64_t) atIndex:2];
758+
[encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:3];
759+
[encoder setBytes:&eps length:sizeof( float) atIndex:4];
760+
[encoder setThreadgroupMemoryLength:nth*sizeof(float) atIndex:0];
761+
762+
const int64_t nrows = ggml_nrows(src0);
763+
764+
[encoder dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
765+
} break;
766+
case GGML_OP_ALIBI:
767+
{
768+
GGML_ASSERT((src0t == GGML_TYPE_F32));
769+
const int n_past = ((int32_t *) src1->data)[0];
770+
const int n_head = ((int32_t *) src1->data)[1];
771+
const float max_bias = ((float *) src1->data)[2];
772+
if (__builtin_popcount(n_head) != 1) {
773+
GGML_ASSERT(false && "only power-of-two n_head implemented");
774+
}
775+
const int n_heads_log2_floor = 1 << (int) floor(log2(n_head));
776+
const float m0 = powf(2.0f, -(max_bias) / n_heads_log2_floor);
777+
if (encoder == nil) {
778+
encoder = [command_buffer computeCommandEncoder];
779+
}
780+
[encoder setComputePipelineState:ctx->pipeline_alibi_f32];
781+
[encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
782+
[encoder setBuffer:id_dst offset:offs_dst atIndex:1];
783+
[encoder setBytes:&ne00 length:sizeof( int64_t) atIndex:2];
784+
[encoder setBytes:&ne01 length:sizeof( int64_t) atIndex:3];
785+
[encoder setBytes:&ne02 length:sizeof( int64_t) atIndex:4];
786+
[encoder setBytes:&ne03 length:sizeof( int64_t) atIndex:5];
787+
[encoder setBytes:&nb00 length:sizeof(uint64_t) atIndex:6];
788+
[encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:7];
789+
[encoder setBytes:&nb02 length:sizeof(uint64_t) atIndex:8];
790+
[encoder setBytes:&nb03 length:sizeof(uint64_t) atIndex:9];
791+
[encoder setBytes:&ne0 length:sizeof( int64_t) atIndex:10];
792+
[encoder setBytes:&ne1 length:sizeof( int64_t) atIndex:11];
793+
[encoder setBytes:&ne2 length:sizeof( int64_t) atIndex:12];
794+
[encoder setBytes:&ne3 length:sizeof( int64_t) atIndex:13];
795+
[encoder setBytes:&nb0 length:sizeof(uint64_t) atIndex:14];
796+
[encoder setBytes:&nb1 length:sizeof(uint64_t) atIndex:15];
797+
[encoder setBytes:&nb2 length:sizeof(uint64_t) atIndex:16];
798+
[encoder setBytes:&nb3 length:sizeof(uint64_t) atIndex:17];
799+
[encoder setBytes:&m0 length:sizeof( float) atIndex:18];
800+
const int nth = 32;
801+
[encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
802+
} break;
738803
case GGML_OP_ROPE:
739804
{
740805
if (encoder == nil) {
@@ -788,6 +853,14 @@ void ggml_metal_graph_compute(
788853
default: GGML_ASSERT(false && "not implemented");
789854
};
790855
} break;
856+
case GGML_TYPE_F16:
857+
{
858+
switch (dstt) {
859+
case GGML_TYPE_F16: [encoder setComputePipelineState:ctx->pipeline_cpy_f16_f16]; break;
860+
case GGML_TYPE_F32: GGML_ASSERT(false && "cpy_f16_f32 not implemented"); break;
861+
default: GGML_ASSERT(false && "not implemented");
862+
};
863+
} break;
791864
default: GGML_ASSERT(false && "not implemented");
792865
}
793866

‎ggml-metal.metal

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,6 +256,72 @@ kernel void kernel_get_rows_q4_1(
256256
(device float *) ((device char *) dst + i*nb1), ne00);
257257
}
258258

259+
kernel void kernel_norm(
260+
device const void * src0,
261+
device float * dst,
262+
constant int64_t & ne00,
263+
constant uint64_t & nb01,
264+
constant float & eps,
265+
threadgroup float * sum [[threadgroup(0)]],
266+
uint tgpig[[threadgroup_position_in_grid]],
267+
uint tpitg[[thread_position_in_threadgroup]],
268+
uint ntg[[threads_per_threadgroup]]) {
269+
device const float * x = (device const float *) ((device const char *) src0 + tgpig*nb01);
270+
// MEAN
271+
// parallel sum
272+
sum[tpitg] = 0.0f;
273+
for (int i00 = tpitg; i00 < ne00; i00 += ntg) {
274+
sum[tpitg] += x[i00];
275+
}
276+
// reduce
277+
threadgroup_barrier(mem_flags::mem_threadgroup);
278+
for (uint i = ntg/2; i > 0; i /= 2) {
279+
if (tpitg < i) {
280+
sum[tpitg] += sum[tpitg + i];
281+
}
282+
threadgroup_barrier(mem_flags::mem_threadgroup);
283+
}
284+
// broadcast
285+
if (tpitg == 0) {
286+
sum[0] /= ne00;
287+
}
288+
threadgroup_barrier(mem_flags::mem_threadgroup);
289+
const float mean = sum[0];
290+
291+
// recenter
292+
device float * y = dst + tgpig*ne00;
293+
for (int i00 = tpitg; i00 < ne00; i00 += ntg) {
294+
y[i00] = x[i00] - mean;
295+
}
296+
297+
// VARIANCE
298+
// parallel sum
299+
sum[tpitg] = 0.0f;
300+
for (int i00 = tpitg; i00 < ne00; i00 += ntg) {
301+
sum[tpitg] += y[i00] * y[i00];
302+
}
303+
// reduce
304+
threadgroup_barrier(mem_flags::mem_threadgroup);
305+
for (uint i = ntg/2; i > 0; i /= 2) {
306+
if (tpitg < i) {
307+
sum[tpitg] += sum[tpitg + i];
308+
}
309+
threadgroup_barrier(mem_flags::mem_threadgroup);
310+
}
311+
// broadcast
312+
if (tpitg == 0) {
313+
sum[0] /= ne00;
314+
}
315+
threadgroup_barrier(mem_flags::mem_threadgroup);
316+
const float variance = sum[0];
317+
318+
const float scale = 1.0f/sqrt(variance + eps);
319+
for (int i00 = tpitg; i00 < ne00; i00 += ntg) {
320+
y[i00] = y[i00] * scale;
321+
}
322+
}
323+
324+
259325
kernel void kernel_rms_norm(
260326
device const void * src0,
261327
device float * dst,
@@ -485,6 +551,48 @@ kernel void kernel_mul_mat_f16_f32(
485551
}
486552
}
487553

554+
kernel void kernel_alibi_f32(
555+
device const float * src0,
556+
device float * dst,
557+
constant int64_t & ne00,
558+
constant int64_t & ne01,
559+
constant int64_t & ne02,
560+
constant int64_t & ne03,
561+
constant uint64_t & nb00,
562+
constant uint64_t & nb01,
563+
constant uint64_t & nb02,
564+
constant uint64_t & nb03,
565+
constant int64_t & ne0,
566+
constant int64_t & ne1,
567+
constant int64_t & ne2,
568+
constant int64_t & ne3,
569+
constant uint64_t & nb0,
570+
constant uint64_t & nb1,
571+
constant uint64_t & nb2,
572+
constant uint64_t & nb3,
573+
constant float & m0,
574+
uint3 tgpig[[threadgroup_position_in_grid]],
575+
uint3 tpitg[[thread_position_in_threadgroup]],
576+
uint3 ntg[[threads_per_threadgroup]]) {
577+
const int64_t i03 = tgpig[2];
578+
const int64_t i02 = tgpig[1];
579+
const int64_t i01 = tgpig[0];
580+
581+
const int64_t n = i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00;
582+
583+
const int64_t i3 = n / (ne2*ne1*ne0);
584+
const int64_t i2 = (n - i3*ne2*ne1*ne0) / (ne1*ne0);
585+
const int64_t i1 = (n - i3*ne2*ne1*ne0 - i2*ne1*ne0) / ne0;
586+
const int64_t i0 = (n - i3*ne2*ne1*ne0 - i2*ne1*ne0 - i1*ne0);
587+
588+
device float * dst_data = (device float *) ((device char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
589+
float m_k = pow(m0, i2 + 1);
590+
for (int64_t i00 = tpitg.x; i00 < ne00; i00 += ntg.x) {
591+
device const float * src = (device float *)((device char *) src0 + i03*nb03 + i02*nb02 + i01*nb01 + i00*nb00);
592+
dst_data[i00] = src[0] + m_k * (i00 - ne00 + 1);
593+
}
594+
}
595+
488596
kernel void kernel_rope(
489597
device const void * src0,
490598
device float * dst,
@@ -540,6 +648,47 @@ kernel void kernel_rope(
540648
}
541649
}
542650

651+
kernel void kernel_cpy_f16_f16(
652+
device const half * src0,
653+
device half * dst,
654+
constant int64_t & ne00,
655+
constant int64_t & ne01,
656+
constant int64_t & ne02,
657+
constant int64_t & ne03,
658+
constant uint64_t & nb00,
659+
constant uint64_t & nb01,
660+
constant uint64_t & nb02,
661+
constant uint64_t & nb03,
662+
constant int64_t & ne0,
663+
constant int64_t & ne1,
664+
constant int64_t & ne2,
665+
constant int64_t & ne3,
666+
constant uint64_t & nb0,
667+
constant uint64_t & nb1,
668+
constant uint64_t & nb2,
669+
constant uint64_t & nb3,
670+
uint3 tgpig[[threadgroup_position_in_grid]],
671+
uint3 tpitg[[thread_position_in_threadgroup]],
672+
uint3 ntg[[threads_per_threadgroup]]) {
673+
const int64_t i03 = tgpig[2];
674+
const int64_t i02 = tgpig[1];
675+
const int64_t i01 = tgpig[0];
676+
677+
const int64_t n = i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00;
678+
679+
const int64_t i3 = n / (ne2*ne1*ne0);
680+
const int64_t i2 = (n - i3*ne2*ne1*ne0) / (ne1*ne0);
681+
const int64_t i1 = (n - i3*ne2*ne1*ne0 - i2*ne1*ne0) / ne0;
682+
const int64_t i0 = (n - i3*ne2*ne1*ne0 - i2*ne1*ne0 - i1*ne0);
683+
684+
device half * dst_data = (device half *) ((device char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
685+
686+
for (int64_t i00 = tpitg.x; i00 < ne00; i00 += ntg.x) {
687+
device const half * src = (device half *)((device char *) src0 + i03*nb03 + i02*nb02 + i01*nb01 + i00*nb00);
688+
dst_data[i00] = src[0];
689+
}
690+
}
691+
543692
kernel void kernel_cpy_f32_f16(
544693
device const float * src0,
545694
device half * dst,

0 commit comments

Comments
 (0)