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
+159
View File
@@ -11278,3 +11278,162 @@ kernel void kernel_count_equal(
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];
}
}