vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

metal: faster diagonal infinity

Browse source 
Although, to me it looks like one should simply
fuse scale + diagnonal infinity + soft_max on the
KQtensor.
This commit is contained in:
Iwan Kawrakow 2023-09-08 12:06:23 +02:00
parent 43ca76976d
commit fa5a989104
2 changed files with 24 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

8
ggml-metal.m
View file

@ -815,15 +815,17 @@ void ggml_metal_graph_compute(
case GGML_OP_DIAG_MASK_INF:
{
const int n_past = ((int32_t *)(dst->op_params))[0];
const int64_t n00x01 = ne00*ne01;
assert((n00x01*ne02)%8 == 0);
[encoder setComputePipelineState:ctx->pipeline_diag_mask_inf];
[encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
[encoder setBuffer:id_dst offset:offs_dst atIndex:1];
[encoder setBytes:&ne00 length:sizeof(ne00) atIndex:2];
[encoder setBytes:&ne01 length:sizeof(ne01) atIndex:3];
[encoder setBytes:&ne00 length:sizeof(ne00) atIndex:2];
[encoder setBytes:&n00x01 length:sizeof(n00x01) atIndex:3];
[encoder setBytes:&n_past length:sizeof(int) atIndex:4];
[encoder dispatchThreadgroups:MTLSizeMake(ne00, ne01, ne02) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
[encoder dispatchThreadgroups:MTLSizeMake(n00x01*ne02/8, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
} break;
case GGML_OP_MUL_MAT:
{

29
ggml-metal.metal
View file

@ -183,20 +183,29 @@ kernel void kernel_soft_max_4(
}
kernel void kernel_diag_mask_inf(
device const float * src0,
device float * dst,
device const float4 * src0,
device float4 * dst,
constant int64_t & ne00,
constant int64_t & ne01,
constant int64_t & n00x01,
constant int & n_past,
uint3 tpig[[thread_position_in_grid]]) {
const int64_t i02 = tpig[2];
const int64_t i01 = tpig[1];
const int64_t i00 = tpig[0];
const int64_t i = 2*tpig[0];
if (i00 > n_past + i01) {
dst[i02*ne01*ne00 + i01*ne00 + i00] = -INFINITY;
} else {
dst[i02*ne01*ne00 + i01*ne00 + i00] = src0[i02*ne01*ne00 + i01*ne00 + i00];
dst[i+0] = src0[i+0];
dst[i+1] = src0[i+1];
int64_t i4 = 4*i;
const int64_t i02 = i4/n00x01;
i4 -= i02*n00x01;
const int64_t i01 = i4/ne00;
const int64_t i00 = i4 - i01*ne00;
for (int k = 3; k >= 0; --k) {
if (i00 + 4 + k <= n_past + i01) {
break;
}
dst[i+1][k] = -INFINITY;
if (i00 + k > n_past + i01) {
dst[i][k] = -INFINITY;
}
}
}