Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-device.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
1 change: 1 addition & 0 deletions ggml/src/ggml-metal/ggml-metal-device.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
36 changes: 36 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-device.m
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
43 changes: 43 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
137 changes: 137 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-ops.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);

Expand Down
1 change: 1 addition & 0 deletions ggml/src/ggml-metal/ggml-metal-ops.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Loading
Loading