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:
Georgi Gerganov
2026-08-02 21:06:02 +03:00
committed by GitHub
co-authored by Thiago Padilha
parent bb4e0e1b3f
commit fffbcbdb9d
8 changed files with 397 additions and 0 deletions
+137
View File
@@ -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);