diff --git a/examples/mtl/mtl.cpp b/examples/mtl/mtl.cpp index e15a1b02e..40e8fbcee 100644 --- a/examples/mtl/mtl.cpp +++ b/examples/mtl/mtl.cpp @@ -41,18 +41,15 @@ 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 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, tmp.data(), tmp.size(), n_past); } - // the actual inference happens here - llama_mtl_eval(ctx_mtl, &gf); - llama_mtl_free(ctx_mtl); ggml_free(ctx_work); diff --git a/examples/mtl/mtl.h b/examples/mtl/mtl.h index a40d57111..a6a336eaa 100644 --- a/examples/mtl/mtl.h +++ b/examples/mtl/mtl.h @@ -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 } diff --git a/examples/mtl/mtl.m b/examples/mtl/mtl.m index 8f55f8467..06d8961ee 100644 --- a/examples/mtl/mtl.m +++ b/examples/mtl/mtl.m @@ -36,6 +36,9 @@ struct ggml_mtl_context { id function_soft_max; id pipeline_soft_max; + id function_diag_mask_inf; + id pipeline_diag_mask_inf; + id function_get_rows_q4_0; id 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 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 command_buffer = [ctx->queue commandBuffer]; id 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 id_src = llama_mtl_get_buffer(ctx, gf->nodes[i]->src0, &offs_src0); + id 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 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) { diff --git a/examples/mtl/mtl.metal b/examples/mtl/mtl.metal index b132be15e..ef2b690c1 100644 --- a/examples/mtl/mtl.metal +++ b/examples/mtl/mtl.metal @@ -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, diff --git a/ggml.c b/ggml.c index 114136122..fc3bdcf6b 100644 --- a/ggml.c +++ b/ggml.c @@ -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: { diff --git a/llama.cpp b/llama.cpp index 28d489016..ff4268ed6 100644 --- a/llama.cpp +++ b/llama.cpp @@ -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);