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);
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);
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;
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;
{
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);
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);
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);
typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;
template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;
+
+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];
+ }
+}
+
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));
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}) {
overview / pkg / ggml / sources / llama.cpp / commitdiff