diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 16e98eb51..ec3ac5140 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -477,6 +477,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max(ggml_me return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, ggml_op op) { + const char * name = nullptr; + + switch (op) { + case GGML_OP_DSV4_HC_COMB: name = "kernel_dsv4_hc_comb_f32"; break; + case GGML_OP_DSV4_HC_PRE: name = "kernel_dsv4_hc_pre_f32"; break; + case GGML_OP_DSV4_HC_POST: name = "kernel_dsv4_hc_post_f32"; break; + default: GGML_ABORT("fatal error"); + } + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv(ggml_metal_library_t lib, const ggml_tensor * op) { GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 91b841b67..2d712f0bf 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -124,6 +124,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_bl struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add (ggml_metal_library_t lib, const struct ggml_tensor * op); 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_dsv4_hc (ggml_metal_library_t lib, enum ggml_op op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op); 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); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 7d2a68685..6bc60a0db 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1299,6 +1299,42 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return false; } return has_simdgroup_mm; // TODO: over-restricted for vec-kernels + case GGML_OP_DSV4_HC_COMB: + return has_simdgroup_reduction && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == 24 && + op->src[1]->ne[0] >= 3 && + op->src[2]->ne[0] == 24 && + ggml_is_contiguous_rows(op->src[0]) && + ggml_is_contiguous_rows(op->src[1]) && + ggml_is_contiguous_rows(op->src[2]); + case GGML_OP_DSV4_HC_PRE: + return has_simdgroup_reduction && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + op->src[0]->ne[1] == 4 && + op->src[1]->ne[0] == 4 && + ggml_is_contiguous_rows(op->src[0]) && + ggml_is_contiguous_rows(op->src[1]); + case GGML_OP_DSV4_HC_POST: + return has_simdgroup_reduction && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && + op->src[3]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + op->src[1]->ne[1] == 4 && + op->src[2]->ne[0] == 4 && + op->src[3]->ne[0] == 4 && + op->src[3]->ne[1] == 4 && + ggml_is_contiguous_rows(op->src[0]) && + ggml_is_contiguous_rows(op->src[1]) && + ggml_is_contiguous_rows(op->src[2]) && + ggml_is_contiguous_rows(op->src[3]); case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: return has_simdgroup_reduction; diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 9f350aad5..0fa0c2d6a 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -1171,6 +1171,49 @@ typedef struct { int64_t val; } ggml_metal_kargs_memset; +typedef struct { + int32_t n_tokens; + int32_t n_iter; + uint64_t nb_m0; + uint64_t nb_m1; + uint64_t nb_s0; + uint64_t nb_b0; + uint64_t nb_d0; + uint64_t nb_d1; + uint64_t nb_d2; + float eps; +} ggml_metal_kargs_dsv4_hc_comb; + +typedef struct { + int32_t n_embd; + int32_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb_d0; + uint64_t nb_d1; +} ggml_metal_kargs_dsv4_hc_pre; + +typedef struct { + int32_t n_embd; + int32_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_r0; + uint64_t nb_r1; + uint64_t nb_r2; + uint64_t nb_p0; + uint64_t nb_p1; + uint64_t nb_c0; + uint64_t nb_c1; + uint64_t nb_c2; + uint64_t nb_d0; + uint64_t nb_d1; + uint64_t nb_d2; +} ggml_metal_kargs_dsv4_hc_post; + typedef struct { int32_t ne00; int32_t ne01; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 76626a451..49f4ee37b 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -316,6 +316,12 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_cumsum(ctx, idx); } break; + case GGML_OP_DSV4_HC_COMB: + case GGML_OP_DSV4_HC_PRE: + case GGML_OP_DSV4_HC_POST: + { + n_fuse = ggml_metal_op_dsv4_hc(ctx, idx); + } break; case GGML_OP_SOFT_MAX: { n_fuse = ggml_metal_op_soft_max(ctx, idx); @@ -1297,6 +1303,137 @@ int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_encoder_t enc = ctx->enc; + auto pipeline = ggml_metal_library_get_pipeline_dsv4_hc(ctx->lib, op->op); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + + switch (op->op) { + case GGML_OP_DSV4_HC_COMB: + { + const ggml_tensor * mixes = op->src[0]; + const ggml_tensor * scale = op->src[1]; + const ggml_tensor * base = op->src[2]; + + GGML_ASSERT(mixes->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + GGML_ASSERT(mixes->ne[0] == 24); + GGML_ASSERT(op->ne[0] == 4 && op->ne[1] == 4); + + ggml_metal_kargs_dsv4_hc_comb args = { + /*.n_tokens =*/ (int32_t) mixes->ne[1], + /*.n_iter =*/ ggml_get_op_params_i32(op, 1), + /*.nb_m0 =*/ mixes->nb[0], + /*.nb_m1 =*/ mixes->nb[1], + /*.nb_s0 =*/ scale->nb[0], + /*.nb_b0 =*/ base->nb[0], + /*.nb_d0 =*/ op->nb[0], + /*.nb_d1 =*/ op->nb[1], + /*.nb_d2 =*/ op->nb[2], + /*.eps =*/ ggml_get_op_params_f32(op, 0), + }; + + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(mixes), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(scale), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(base), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); + + // One SIMDgroup owns one 4x4 Sinkhorn matrix. Packing up to four + // independent tokens per threadgroup keeps both decode and prompt + // dispatches compact without any threadgroup-memory synchronization. + const int nsg = std::min(4, args.n_tokens); + ggml_metal_encoder_dispatch_threadgroups( + enc, (args.n_tokens + nsg - 1)/nsg, 1, 1, 32, nsg, 1); + } break; + case GGML_OP_DSV4_HC_PRE: + { + const ggml_tensor * x = op->src[0]; + const ggml_tensor * weights = op->src[1]; + + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + GGML_ASSERT(x->ne[1] == 4); + + ggml_metal_kargs_dsv4_hc_pre args = { + /*.n_embd =*/ (int32_t) x->ne[0], + /*.n_tokens =*/ (int32_t) x->ne[2], + /*.nb_x0 =*/ x->nb[0], + /*.nb_x1 =*/ x->nb[1], + /*.nb_x2 =*/ x->nb[2], + /*.nb_w0 =*/ weights->nb[0], + /*.nb_w1 =*/ weights->nb[1], + /*.nb_d0 =*/ op->nb[0], + /*.nb_d1 =*/ op->nb[1], + }; + + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weights), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3); + + const int n_tiles = (args.n_embd + 31)/32; + const int nsg = std::min(4, n_tiles); + ggml_metal_encoder_dispatch_threadgroups( + enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1); + } break; + case GGML_OP_DSV4_HC_POST: + { + const ggml_tensor * x = op->src[0]; + const ggml_tensor * residual = op->src[1]; + const ggml_tensor * post = op->src[2]; + const ggml_tensor * comb = op->src[3]; + + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(residual->type == GGML_TYPE_F32); + GGML_ASSERT(post->type == GGML_TYPE_F32); + GGML_ASSERT(comb->type == GGML_TYPE_F32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + GGML_ASSERT(residual->ne[1] == 4); + + ggml_metal_kargs_dsv4_hc_post args = { + /*.n_embd =*/ (int32_t) x->ne[0], + /*.n_tokens =*/ (int32_t) x->ne[1], + /*.nb_x0 =*/ x->nb[0], + /*.nb_x1 =*/ x->nb[1], + /*.nb_r0 =*/ residual->nb[0], + /*.nb_r1 =*/ residual->nb[1], + /*.nb_r2 =*/ residual->nb[2], + /*.nb_p0 =*/ post->nb[0], + /*.nb_p1 =*/ post->nb[1], + /*.nb_c0 =*/ comb->nb[0], + /*.nb_c1 =*/ comb->nb[1], + /*.nb_c2 =*/ comb->nb[2], + /*.nb_d0 =*/ op->nb[0], + /*.nb_d1 =*/ op->nb[1], + /*.nb_d2 =*/ op->nb[2], + }; + + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(residual), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(post), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comb), 4); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 5); + + const int n_tiles = (args.n_embd + 31)/32; + const int nsg = std::min(4, n_tiles); + ggml_metal_encoder_dispatch_threadgroups( + enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1); + } break; + default: + GGML_ABORT("fatal error"); + } + + return 1; +} + int ggml_metal_op_soft_max(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 2783ecb8b..bd20bdfbc 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -54,6 +54,7 @@ int ggml_metal_op_cumsum (ggml_metal_op_t ctx, int idx); 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_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); int ggml_metal_op_ssm_scan (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index f14ee0792..3c8c0ecd4 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -11278,3 +11278,162 @@ kernel void kernel_count_equal( typedef decltype(kernel_count_equal) kernel_count_equal_t; template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal; + +kernel void kernel_dsv4_hc_comb_f32( + constant ggml_metal_kargs_dsv4_hc_comb & args, + device const char * mixes, + device const char * scale, + device const char * base, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort hc = 4; + constexpr ushort comb_offset = 2*hc; + + const int it = tgpig.x*ntg.y + sgitg; + if (it >= args.n_tokens) { + return; + } + + float scale_lane = 0.0f; + if (tiisg == 0) { + scale_lane = *(device const float *) (scale + 2*args.nb_s0); + } + const float scale_comb = simd_shuffle(scale_lane, 0); + + float v = 0.0f; + if (tiisg < hc*hc) { + v = *(device const float *) (mixes + (comb_offset + tiisg)*args.nb_m0 + it*args.nb_m1)*scale_comb + + *(device const float *) (base + (comb_offset + tiisg)*args.nb_b0); + } + + // Softmax across destinations (the four contiguous lanes for each source). + float vmax = max(v, simd_shuffle_xor(v, 1)); + vmax = max(vmax, simd_shuffle_xor(vmax, 2)); + v = exp(v - vmax); + + float sum = v + simd_shuffle_xor(v, 1); + sum += simd_shuffle_xor(sum, 2); + v = v/sum + args.eps; + + // Normalize columns: equal destination indices are four lanes apart. + sum = v + simd_shuffle_xor(v, 4); + sum += simd_shuffle_xor(sum, 8); + v /= sum + args.eps; + + for (int i = 1; i < args.n_iter; ++i) { + sum = v + simd_shuffle_xor(v, 1); + sum += simd_shuffle_xor(sum, 2); + v /= sum + args.eps; + + sum = v + simd_shuffle_xor(v, 4); + sum += simd_shuffle_xor(sum, 8); + v /= sum + args.eps; + } + + if (tiisg < hc*hc) { + const ushort idst = tiisg & 3; + const ushort isrc = tiisg >> 2; + *(device float *) (dst + idst*args.nb_d0 + isrc*args.nb_d1 + it*args.nb_d2) = v; + } +} + +kernel void kernel_dsv4_hc_pre_f32( + constant ggml_metal_kargs_dsv4_hc_pre & args, + device const char * x, + device const char * weights, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort hc = 4; + + const int it = tgpig.y; + const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg; + + float weight_lane = 0.0f; + if (tiisg < hc) { + weight_lane = *(device const float *) (weights + tiisg*args.nb_w0 + it*args.nb_w1); + } + + float w[hc]; + FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) { + w[ih] = simd_shuffle(weight_lane, ih); + } + + if (i0 >= args.n_embd) { + return; + } + + device const char * xb = x + i0*args.nb_x0 + it*args.nb_x2; + float result = 0.0f; + FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) { + result = fma(*(device const float *) (xb + ih*args.nb_x1), w[ih], result); + } + + *(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = result; +} + +kernel void kernel_dsv4_hc_post_f32( + constant ggml_metal_kargs_dsv4_hc_post & args, + device const char * x, + device const char * residual, + device const char * post, + device const char * comb, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort hc = 4; + + const int it = tgpig.y; + const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg; + + float coeff_lane = 0.0f; + if (tiisg < hc) { + coeff_lane = *(device const float *) (post + tiisg*args.nb_p0 + it*args.nb_p1); + } else if (tiisg < hc + hc*hc) { + const ushort idx = tiisg - hc; + const ushort idst = idx & 3; + const ushort isrc = idx >> 2; + coeff_lane = *(device const float *) (comb + idst*args.nb_c0 + isrc*args.nb_c1 + it*args.nb_c2); + } + + float post_reg[hc]; + float comb_reg[hc][hc]; + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + post_reg[idst] = simd_shuffle(coeff_lane, idst); + } + FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) { + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + comb_reg[isrc][idst] = simd_shuffle(coeff_lane, hc + idst + hc*isrc); + } + } + + if (i0 >= args.n_embd) { + return; + } + + const float xv = *(device const float *) (x + i0*args.nb_x0 + it*args.nb_x1); + float result[hc]; + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + result[idst] = xv*post_reg[idst]; + } + + device const char * rb = residual + i0*args.nb_r0 + it*args.nb_r2; + FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) { + const float rv = *(device const float *) (rb + isrc*args.nb_r1); + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + result[idst] = fma(rv, comb_reg[isrc][idst], result[idst]); + } + } + + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + *(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst]; + } +} + diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 8f3fb45da..32468bf22 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8069,6 +8069,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_comb(1, 1)); test_cases.emplace_back(new test_dsv4_hc_comb(17, 4)); test_cases.emplace_back(new test_dsv4_hc_comb(257, 8)); + test_cases.emplace_back(new test_dsv4_hc_comb(17, 20)); test_cases.emplace_back(new test_dsv4_hc_pre(1, 1)); test_cases.emplace_back(new test_dsv4_hc_pre(31, 17)); @@ -8078,6 +8079,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_post(1, 1)); test_cases.emplace_back(new test_dsv4_hc_post(31, 17)); test_cases.emplace_back(new test_dsv4_hc_post(128, 257)); + test_cases.emplace_back(new test_dsv4_hc_post(4096, 21)); // glu ops for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {