diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index f3f5d4d77dd7..25a8a3ef6b1f 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -507,6 +507,47 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexe return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_mask( + ggml_metal_library_t lib, + const ggml_tensor * op, + bool fill, + bool has_score, + bool has_blocks, + bool has_extra, + bool vec) { + GGML_ASSERT(op->op == GGML_OP_TOP_K_MASK); + + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_top_k_mask%s_%s", fill ? "_fill" : "", ggml_type_name(op->type)); + + ggml_metal_pipeline_with_params res; + + ggml_metal_cv_t cv = ggml_metal_cv_init(); + + if (fill) { + snprintf(name, 256, "%s_vec=%d", base, vec ? 1 : 0); + + ggml_metal_cv_set_bool(cv, vec, FC_TOP_K_MASK + 3); + } else { + snprintf(name, 256, "%s_sc=%d_bl=%d_ex=%d", base, has_score ? 1 : 0, has_blocks ? 1 : 0, has_extra ? 1 : 0); + + ggml_metal_cv_set_bool(cv, has_score, FC_TOP_K_MASK + 0); + ggml_metal_cv_set_bool(cv, has_blocks, FC_TOP_K_MASK + 1); + ggml_metal_cv_set_bool(cv, has_extra, FC_TOP_K_MASK + 2); + } + + res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, cv); + } + + ggml_metal_cv_free(cv); + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, const ggml_tensor * op) { char name[256]; const char * base = nullptr; diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 6d77c57ffe0b..dd7cdb958382 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -126,6 +126,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_ad struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_mask (ggml_metal_library_t lib, const struct ggml_tensor * op, bool fill, bool has_score, bool has_blocks, bool has_extra, bool vec); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t nc, bool use_silu); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs, int32_t nc, bool use_silu); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index d1bce8bad85a..3cfd24f1ad31 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1834,6 +1834,15 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te ggml_is_contiguous_rows(op->src[1]) && ggml_is_contiguous_rows(op->src[2]) && (op->src[3] == NULL || ggml_is_contiguous_rows(op->src[3])); + case GGML_OP_TOP_K_MASK: + return (op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32) && + op->src[0]->type == op->type && + op->src[0]->nb[0] == ggml_type_size(op->src[0]->type) && + op->nb[0] == ggml_type_size(op->type) && + op->src[1]->type == GGML_TYPE_I32 && + (op->src[2] == NULL || op->src[2]->type == GGML_TYPE_F32) && + (op->src[3] == NULL || op->src[3]->type == GGML_TYPE_I32) && + (op->src[4] == NULL || op->src[4]->type == GGML_TYPE_I32); case GGML_OP_SSM_SCAN: return has_simdgroup_reduction; case GGML_OP_SSM_CONV: diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 9db266f44231..6e566084a144 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -125,6 +125,7 @@ #define FC_FLASH_ATTN_EXT_TENSOR 2200 #define FC_LIGHTNING_INDEXER 2200 #define FC_MUL_MV_MMA 2300 +#define FC_TOP_K_MASK 2400 // op-specific constants #define OP_FLASH_ATTN_EXT_NQPSG 8 @@ -1295,6 +1296,38 @@ typedef struct { uint64_t nbm3; } ggml_metal_kargs_lightning_indexer; +typedef struct { + int32_t n_kv; + int32_t m1; + int32_t m2; + int32_t m3; + int32_t n_top; + int32_t n_pick; + int32_t n_blk; + int32_t n_extra; + int64_t n_elems; + uint64_t nb1; + uint64_t nb2; + uint64_t nb3; + uint64_t nbm1; + uint64_t nbm2; + uint64_t nbm3; + uint64_t nbk0; + uint64_t nbk1; + uint64_t nbk2; + uint64_t nbk3; + uint64_t nbs0; + uint64_t nbs1; + uint64_t nbs2; + uint64_t nbs3; + uint64_t nbb0; + uint64_t nbb1; + uint64_t nbe0; + uint64_t nbe1; + uint64_t nbe2; + uint64_t nbe3; +} ggml_metal_kargs_top_k_mask; + typedef struct { int32_t n_tokens; int32_t n_iter; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 0cebbe5d6e1f..14712c3e8d28 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -345,6 +345,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_lightning_indexer(ctx, idx); } break; + case GGML_OP_TOP_K_MASK: + { + n_fuse = ggml_metal_op_top_k_mask(ctx, idx); + } break; case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: @@ -1424,6 +1428,124 @@ int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_top_k_mask(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + const ggml_tensor * mask = op->src[0]; + const ggml_tensor * top_k = op->src[1]; + const ggml_tensor * score = op->src[2]; + const ggml_tensor * blocks = op->src[3]; + const ggml_tensor * extra = op->src[4]; + + GGML_ASSERT(op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32); + GGML_ASSERT(mask->type == op->type); + GGML_ASSERT(top_k->type == GGML_TYPE_I32); + GGML_ASSERT(!score || score->type == GGML_TYPE_F32); + GGML_ASSERT(!blocks || blocks->type == GGML_TYPE_I32); + GGML_ASSERT(!extra || extra->type == GGML_TYPE_I32); + + const int32_t n_kv = mask->ne[0]; + const int32_t n_top = (int32_t) top_k->ne[0]; + const int32_t n_pick = blocks ? blocks->ne[1] : score ? score->ne[0] : n_kv; + const int32_t n_blk = blocks ? blocks->ne[0] : 0; + const int32_t n_extra = extra ? extra->ne[0] : 0; + + const ggml_metal_buffer_id bid = ggml_metal_get_buffer_id(op); + + // the vector fill stores 4 elements at a time + const int32_t ts = ggml_type_size(op->type); + const int64_t n_elems = ggml_nelements(op); + + const bool has_score = score != NULL; + const bool has_blocks = blocks != NULL; + const bool has_extra = extra != NULL; + + const bool vec = n_elems % 4 == 0 && bid.offs % (4*ts) == 0; + + ggml_metal_kargs_top_k_mask args = { + /*.n_kv =*/ n_kv, + /*.m1 =*/ (int32_t) mask->ne[1], + /*.m2 =*/ (int32_t) mask->ne[2], + /*.m3 =*/ (int32_t) mask->ne[3], + /*.n_top =*/ n_top, + /*.n_pick =*/ n_pick, + /*.n_blk =*/ n_blk, + /*.n_extra =*/ n_extra, + /*.n_elems =*/ n_elems, + /*.nb1 =*/ op->nb[1], + /*.nb2 =*/ op->nb[2], + /*.nb3 =*/ op->nb[3], + /*.nbm1 =*/ mask->nb[1], + /*.nbm2 =*/ mask->nb[2], + /*.nbm3 =*/ mask->nb[3], + /*.nbk0 =*/ top_k->nb[0], + /*.nbk1 =*/ top_k->nb[1], + /*.nbk2 =*/ top_k->nb[2], + /*.nbk3 =*/ top_k->nb[3], + /*.nbs0 =*/ score ? score->nb[0] : 0, + /*.nbs1 =*/ score ? score->nb[1] : 0, + /*.nbs2 =*/ score ? score->nb[2] : 0, + /*.nbs3 =*/ score ? score->nb[3] : 0, + /*.nbb0 =*/ blocks ? blocks->nb[0] : 0, + /*.nbb1 =*/ blocks ? blocks->nb[1] : 0, + /*.nbe0 =*/ extra ? extra->nb[0] : 0, + /*.nbe1 =*/ extra ? extra->nb[1] : 0, + /*.nbe2 =*/ extra ? extra->nb[2] : 0, + /*.nbe3 =*/ extra ? extra->nb[3] : 0, + }; + + // the fill covers the whole dst, the scatter overwrites the kept cells with the mask values + { + auto pipeline_fill = ggml_metal_library_get_pipeline_top_k_mask(lib, op, true, false, false, false, vec); + + const int64_t n_fill = n_elems/(vec ? 4 : 1); + const int32_t nth_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_fill); + + int nth = 32; // SIMD width + + while (nth < nth_max && nth < n_fill) { + nth *= 2; + } + + ggml_metal_encoder_set_pipeline(enc, pipeline_fill); + ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, bid, 1); + + ggml_metal_encoder_dispatch_threadgroups(enc, (n_fill + nth - 1)/nth, 1, 1, nth, 1, 1); + + ggml_metal_op_concurrency_reset(ctx); + } + + { + auto pipeline = ggml_metal_library_get_pipeline_top_k_mask(lib, op, false, has_score, has_blocks, has_extra, false); + + const int32_t n_max = n_top + n_extra; + const int32_t nth_max = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); + + int nth = 32; // SIMD width + + while (nth < nth_max && nth < n_max) { + nth *= 2; + } + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(mask), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(top_k), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(score ? score : mask), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(blocks ? blocks : mask), 4); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(extra ? extra : mask), 5); + ggml_metal_encoder_set_buffer(enc, bid, 6); + + ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], op->ne[2], op->ne[3], nth, 1, 1); + } + + return 1; +} + int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index e14a2dd5f732..0a77d8d69b7d 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -60,6 +60,7 @@ int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx); int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_top_k_mask (ggml_metal_op_t ctx, int idx); int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx); int ggml_metal_op_soft_max (ggml_metal_op_t ctx, int idx); int ggml_metal_op_ssm_conv (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/kernels/misc.metal b/ggml/src/ggml-metal/kernels/misc.metal index 90d3fbe211af..5dfd4a60c3ab 100644 --- a/ggml/src/ggml-metal/kernels/misc.metal +++ b/ggml/src/ggml-metal/kernels/misc.metal @@ -906,3 +906,99 @@ kernel void kernel_dsv4_hc_post_f32( *(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst]; } } + +constant bool FC_top_k_mask_has_score [[function_constant(FC_TOP_K_MASK + 0)]]; +constant bool FC_top_k_mask_has_blocks [[function_constant(FC_TOP_K_MASK + 1)]]; +constant bool FC_top_k_mask_has_extra [[function_constant(FC_TOP_K_MASK + 2)]]; +constant bool FC_top_k_mask_vec [[function_constant(FC_TOP_K_MASK + 3)]]; + +// the fill runs before the scatter, both are fully parallel over the whole dst and over the picks of one row +template +kernel void kernel_top_k_mask_fill( + constant ggml_metal_kargs_top_k_mask & args, + device char * dst, + uint tid[[thread_position_in_grid]], + uint ntot[[threads_per_grid]]) { + const T ninf = T(-INFINITY); + + if (FC_top_k_mask_vec) { + device T4 * dst_4 = (device T4 *) dst; + for (int64_t i = tid; i < args.n_elems/4; i += ntot) { + dst_4[i] = T4(ninf); + } + } else { + device T * dst_t = (device T *) dst; + for (int64_t i = tid; i < args.n_elems; i += ntot) { + dst_t[i] = ninf; + } + } +} + +typedef decltype(kernel_top_k_mask_fill) kernel_top_k_mask_fill_h_t; +typedef decltype(kernel_top_k_mask_fill) kernel_top_k_mask_fill_f_t; + +template [[host_name("kernel_top_k_mask_fill_f16")]] kernel kernel_top_k_mask_fill_h_t kernel_top_k_mask_fill; +template [[host_name("kernel_top_k_mask_fill_f32")]] kernel kernel_top_k_mask_fill_f_t kernel_top_k_mask_fill; + +// one threadgroup per row: copy the mask cells of the live picks and of the extra cells over the fill +template +kernel void kernel_top_k_mask( + constant ggml_metal_kargs_top_k_mask & args, + device const char * mask, + device const char * top_k, + device const char * score, + device const char * blocks, + device const char * extra, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg_max[[threads_per_threadgroup]]) { + const int i1 = tgpig.x; + const int i2 = tgpig.y; + const int i3 = tgpig.z; + + const int tiitg = tpitg.x; + const int ntg = ntg_max.x; + + device const char * mask_row = mask + (i1 % args.m1)*args.nbm1 + (i2 % args.m2)*args.nbm2 + (i3 % args.m3)*args.nbm3; + device char * dst_row = dst + i1*args.nb1 + i2*args.nb2 + i3*args.nb3; + device const char * top_k_row = top_k + i1*args.nbk1 + i2*args.nbk2 + i3*args.nbk3; + + device const T * mask_row_t = (device const T *) mask_row; + device T * dst_row_t = (device T *) dst_row; + + for (int i = tiitg; i < args.n_top; i += ntg) { + const int32_t p = *(device const int32_t *) (top_k_row + i*args.nbk0); + if (p < 0 || p >= args.n_pick) { + continue; + } + if (FC_top_k_mask_has_score && !(*(device const float *) (score + p*args.nbs0 + i1*args.nbs1 + i2*args.nbs2 + i3*args.nbs3) > -INFINITY)) { + continue; + } + if (FC_top_k_mask_has_blocks) { + FOR_UNROLL (int k = 0; k < args.n_blk; ++k) { + const int32_t c = *(device const int32_t *) (blocks + k*args.nbb0 + p*args.nbb1); + if (c >= 0 && c < args.n_kv) { + dst_row_t[c] = mask_row_t[c]; + } + } + } else if (p < args.n_kv) { + dst_row_t[p] = mask_row_t[p]; + } + } + + if (FC_top_k_mask_has_extra) { + for (int i = tiitg; i < args.n_extra; i += ntg) { + const int32_t c = *(device const int32_t *) (extra + i*args.nbe0 + i1*args.nbe1 + i2*args.nbe2 + i3*args.nbe3); + if (c >= 0 && c < args.n_kv) { + dst_row_t[c] = mask_row_t[c]; + } + } + } +} + +typedef decltype(kernel_top_k_mask) kernel_top_k_mask_h_t; +typedef decltype(kernel_top_k_mask) kernel_top_k_mask_f_t; + +template [[host_name("kernel_top_k_mask_f16")]] kernel kernel_top_k_mask_h_t kernel_top_k_mask; +template [[host_name("kernel_top_k_mask_f32")]] kernel kernel_top_k_mask_f_t kernel_top_k_mask;