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

mtl : add diag_mask_inf kernel

Browse source
This commit is contained in:
Georgi Gerganov 2023-06-01 20:41:54 +03:00
parent 0f1c580860
commit 17a70362a6
No known key found for this signature in database
GPG key ID: 449E073F9DC10735
6 changed files with 92 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

9
examples/mtl/mtl.cpp
View file

@ -41,17 +41,14 @@ int main(int argc, char ** argv) {
// TODO: tmp to match the input used when creating the cgraph
{
const int n_ctx = 128;
const int n_past = 128;
const int n_batch = 32;
const std::vector<int> tmp(n_batch, 1); // BOS
struct ggml_tensor * input = ggml_graph_get_tensor(&gf, "embd");
memcpy(input->data, tmp.data(), tmp.size() * sizeof(int));
}
// the actual inference happens here
llama_mtl_eval(ctx_mtl, &gf);
llama_mtl_eval(ctx_mtl, &gf, tmp.data(), tmp.size(), n_past);
}
llama_mtl_free(ctx_mtl);

5
examples/mtl/mtl.h
View file

@ -20,7 +20,10 @@ void llama_mtl_free(struct ggml_mtl_context * ctx);
// return 0 on success
int llama_mtl_eval(
struct ggml_mtl_context * ctx,
struct ggml_cgraph * gf);
struct ggml_cgraph * gf,
const int * tokens,
int n_tokens,
int n_past);
#ifdef __cplusplus
}

51
examples/mtl/mtl.m
View file

@ -36,6 +36,9 @@ struct ggml_mtl_context {
id<MTLFunction> function_soft_max;
id<MTLComputePipelineState> pipeline_soft_max;
id<MTLFunction> function_diag_mask_inf;
id<MTLComputePipelineState> pipeline_diag_mask_inf;
id<MTLFunction> function_get_rows_q4_0;
id<MTLComputePipelineState> pipeline_get_rows_q4_0;
@ -150,6 +153,10 @@ struct ggml_mtl_context * llama_mtl_init(
ctx->pipeline_soft_max = [ctx->device newComputePipelineStateWithFunction:ctx->function_soft_max error:nil];
fprintf(stderr, "%s: loaded kernel_soft_max: %p\n", __func__, (void *) ctx->pipeline_soft_max);
ctx->function_diag_mask_inf = [ctx->library newFunctionWithName:@"kernel_diag_mask_inf" constantValues:constants error:nil];
ctx->pipeline_diag_mask_inf = [ctx->device newComputePipelineStateWithFunction:ctx->function_diag_mask_inf error:nil];
fprintf(stderr, "%s: loaded kernel_diag_mask_inf: %p\n", __func__, (void *) ctx->pipeline_diag_mask_inf);
ctx->function_get_rows_q4_0 = [ctx->library newFunctionWithName:@"kernel_get_rows_q4_0"];
ctx->pipeline_get_rows_q4_0 = [ctx->device newComputePipelineStateWithFunction:ctx->function_get_rows_q4_0 error:nil];
fprintf(stderr, "%s: loaded kernel_get_rows_q4_0: %p\n", __func__, (void *) ctx->pipeline_get_rows_q4_0);
@ -248,8 +255,14 @@ id<MTLBuffer> llama_mtl_get_buffer(struct ggml_mtl_context * ctx, struct ggml_te
int llama_mtl_eval(
struct ggml_mtl_context * ctx,
struct ggml_cgraph * gf) {
fprintf(stderr, "%s: evaluating\n", __func__);
struct ggml_cgraph * gf,
const int * tokens,
int n_tokens,
int n_past) {
fprintf(stderr, "%s: evaluating, n_tokens = %d, n_past = %d\n", __func__, n_tokens, n_past);
struct ggml_tensor * input = ggml_graph_get_tensor(gf, "embd");
memcpy(input->data, tokens, n_tokens * sizeof(int));
id<MTLCommandBuffer> command_buffer = [ctx->queue commandBuffer];
id<MTLComputeCommandEncoder> encoder = nil;
@ -371,6 +384,28 @@ int llama_mtl_eval(
[encoder dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
} break;
case GGML_OP_DIAG_MASK_INF:
{
if (encoder == nil) {
encoder = [command_buffer computeCommandEncoder];
}
id<MTLBuffer> id_src = llama_mtl_get_buffer(ctx, gf->nodes[i]->src0, &offs_src0);
id<MTLBuffer> id_dst = llama_mtl_get_buffer(ctx, gf->nodes[i], &offs_dst);
const int64_t ne00 = gf->nodes[i]->src0->ne[0];
const int64_t ne01 = gf->nodes[i]->src0->ne[1];
const int64_t ne02 = gf->nodes[i]->src0->ne[2];
[encoder setComputePipelineState:ctx->pipeline_diag_mask_inf];
[encoder setBuffer:id_src 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:&n_past length:sizeof(int) atIndex:4];
[encoder dispatchThreadgroups:MTLSizeMake(ne00, ne01, ne02) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
} break;
case GGML_OP_MUL_MAT:
{
id<MTLBuffer> id_src0 = llama_mtl_get_buffer(ctx, gf->nodes[i]->src0, &offs_src0);
@ -550,7 +585,7 @@ int llama_mtl_eval(
const uint64_t nb2 = gf->nodes[i]->nb[2];
const uint64_t nb3 = gf->nodes[i]->nb[3];
const int n_past = ((int32_t *) gf->nodes[i]->src1->data)[0]; // TODO: TMP !!!!!
//const int n_past = ((int32_t *) gf->nodes[i]->src1->data)[0]; // TODO: TMP !!!!!
const int n_dims = ((int32_t *) gf->nodes[i]->src1->data)[1];
const int mode = ((int32_t *) gf->nodes[i]->src1->data)[2];
@ -697,17 +732,15 @@ int llama_mtl_eval(
if (t->type == GGML_TYPE_F32) {
const const float * data = (float *) ctx->out.contents;
printf("data: ");
int n = ggml_nelements(t);
if (n > 10) {
n = 10;
}
for (int i = 0; i < n; i++) {
for (int i = 0; i < (int) t->ne[0]; i++) {
printf("%f ", data[i]);
}
printf("\n");
double sum = 0.0;
for (int i = 0; i < ggml_nelements(t); i++) {
sum += data[i];
double cur = data[i];
if (isinf(cur)) continue;
sum += cur;
}
printf("sum: %f\n", sum);
} else if (t->type == GGML_TYPE_F16) {

36
examples/mtl/mtl.metal
View file

@ -86,6 +86,42 @@ kernel void kernel_soft_max(
}
}
//const int n = ggml_nrows(src0);
//const int nc = src0->ne[0];
//const int nr = src0->ne[1];
//const int nz = n/nr;
//
//assert( dst->nb[0] == sizeof(float));
//assert(src0->nb[0] == sizeof(float));
//
//for (int k = 0; k < nz; k++) {
// for (int j = ith; j < nr; j += nth) {
// for (int i = n_past; i < nc; i++) {
// if (i > n_past + j) {
// *(float *)((char *) dst->data + k*dst->nb[2] + j*dst->nb[1] + i*dst->nb[0]) = value;
// }
// }
// }
//}
kernel void kernel_diag_mask_inf(
device const float * src0,
device float * dst,
constant int64_t & ne00,
constant int64_t & ne01,
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];
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];
}
}
kernel void kernel_get_rows_q4_0(
device const void * src0,
device const int * src1,

1
ggml.c
View file

@ -15061,7 +15061,6 @@ struct ggml_cgraph ggml_graph_import(const char * fname, struct ggml_context **
memcpy(&offs, args[2]->data, sizeof(offs));
tensor->data = ((char *) tensor->data) + offs;
printf("xxxxxx offs: %zu\n", offs);
} break;
case GGML_OP_TRANSPOSE:
{

12
llama.cpp
View file

@ -1332,14 +1332,14 @@ static bool llama_eval_internal(
// KQ_scaled shape [n_past + N, N, n_head, 1]
struct ggml_tensor * KQ_scaled = ggml_scale_inplace(ctx0, KQ, KQ_scale);
ggml_set_name(KQ_scaled, "KQ_scaled");
// TODO: TMP !!!!
if (il == 0) {
ggml_set_name(KQ_scaled, "mtl-check");
}
// KQ_masked = mask_past(KQ_scaled)
struct ggml_tensor * KQ_masked = ggml_diag_mask_inf_inplace(ctx0, KQ_scaled, n_past);
ggml_set_name(KQ_masked, "KQ_masked");
// TODO: TMP !!!!
if (il == 0) {
ggml_set_name(KQ_masked, "mtl-check");
}
// KQ = soft_max(KQ_masked)
struct ggml_tensor * KQ_soft_max = ggml_soft_max_inplace(ctx0, KQ_masked);
@ -1464,12 +1464,14 @@ static bool llama_eval_internal(
auto print_t_f32 = [&](struct ggml_tensor * t) {
float * data = (float *)t->data;
printf("data: ");
for (int i = 0; i < std::min((int) t->ne[0], 10); i++) {
for (int i = 0; i < (int) t->ne[0]; i++) {
printf("%f ", data[i]);
}
printf("\n");
double sum = 0.0;
for (int i = 0; i < ggml_nelements(t); i++) {
double cur = data[i];
if (isinf(cur)) continue;
sum += data[i];
}
printf("sum: %f\n", sum);