metal: implement DeepSeek V4 hyper-connections (#26459)
- Implement GGML_OP_DSV4_HC_COMB, GGML_OP_DSV4_HC_PRE, and GGML_OP_DSV4_HC_POST with SIMDgroup register and shuffle optimized kernels. - Add Metal dispatch and support plumbing and test the production Sinkhorn iteration count and embedding width. Assisted-by: Codex Co-authored-by: Thiago Padilha <thiago@padilha.cc>
This commit is contained in:
co-authored by
Thiago Padilha
parent
bb4e0e1b3f
commit
fffbcbdb9d
@@ -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);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user