From 4e7a40614b5c78fa614ad78c0a5c2a1c48224b29 Mon Sep 17 00:00:00 2001 From: jadenmach2 Date: Fri, 10 Jul 2026 11:17:15 -0400 Subject: [PATCH] ggml : add fused SINKHORN_NORM op (Sinkhorn-Knopp normalization) --- ggml/include/ggml.h | 7 + ggml/src/ggml-cpu/ggml-cpu.c | 5 + ggml/src/ggml-cpu/ops.cpp | 119 +++++++++++++++++ ggml/src/ggml-cpu/ops.h | 1 + ggml/src/ggml-cuda/ggml-cuda.cu | 9 ++ ggml/src/ggml-cuda/sinkhorn-norm.cu | 192 +++++++++++++++++++++++++++ ggml/src/ggml-cuda/sinkhorn-norm.cuh | 3 + ggml/src/ggml.c | 27 +++- src/models/deepseek4.cpp | 30 +---- tests/test-backend-ops.cpp | 44 ++++++ 10 files changed, 407 insertions(+), 30 deletions(-) create mode 100644 ggml/src/ggml-cuda/sinkhorn-norm.cu create mode 100644 ggml/src/ggml-cuda/sinkhorn-norm.cuh diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index ac133665d978..8ddf70c691c3 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -509,6 +509,7 @@ extern "C" { GGML_OP_RMS_NORM_BACK, GGML_OP_GROUP_NORM, GGML_OP_L2_NORM, + GGML_OP_SINKHORN_NORM, GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID, @@ -1406,6 +1407,12 @@ extern "C" { struct ggml_tensor * a, float eps); + GGML_API struct ggml_tensor * ggml_sinkhorn_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_iters, + float eps); + // a - x // b - dy GGML_API struct ggml_tensor * ggml_rms_norm_back( diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index a82842fcffc0..17f9c8a2ad7e 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1833,6 +1833,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_l2_norm(params, tensor); } break; + case GGML_OP_SINKHORN_NORM: + { + ggml_compute_forward_sinkhorn_norm(params, tensor); + } break; case GGML_OP_MUL_MAT: { ggml_compute_forward_mul_mat(params, tensor); @@ -2306,6 +2310,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_RMS_NORM: case GGML_OP_RMS_NORM_BACK: case GGML_OP_L2_NORM: + case GGML_OP_SINKHORN_NORM: case GGML_OP_GROUP_NORM: case GGML_OP_CONCAT: case GGML_OP_MUL_MAT: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index df0028cf15e3..602a325a5722 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4237,6 +4237,125 @@ void ggml_compute_forward_l2_norm( } } +// ggml_compute_forward_sinkhorn_norm + +#define GGML_SINKHORN_NORM_MAX_N 32 + +static void ggml_compute_forward_sinkhorn_norm_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_contiguous(src0)); + + const int ith = params->ith; + const int nth = params->nth; + + GGML_TENSOR_UNARY_OP_LOCALS + + const int32_t n_iters = ggml_get_op_params_i32(dst, 0); + const float eps = ggml_get_op_params_f32(dst, 1); + + const int64_t n = ne00; + GGML_ASSERT(ne01 == n); + GGML_ASSERT(n <= GGML_SINKHORN_NORM_MAX_N); + GGML_ASSERT(n_iters >= 1); + + const int64_t n_slices = ne02 * ne03; + + float m[GGML_SINKHORN_NORM_MAX_N * GGML_SINKHORN_NORM_MAX_N]; + + for (int64_t s = ith; s < n_slices; s += nth) { + const int64_t i02 = s % ne02; + const int64_t i03 = s / ne02; + + const float * x = (const float *) ((const char *) src0->data + i02*nb02 + i03*nb03); + float * y = (float *) (( char *) dst->data + i02*nb2 + i03*nb3); + + for (int64_t b = 0; b < n; ++b) { + for (int64_t a = 0; a < n; ++a) { + m[b*n + a] = x[b*n + a]; + } + } + + for (int64_t b = 0; b < n; ++b) { + float mx = -INFINITY; + for (int64_t a = 0; a < n; ++a) { + mx = fmaxf(mx, m[b*n + a]); + } + float sum = 0.0f; + for (int64_t a = 0; a < n; ++a) { + const float e = expf(m[b*n + a] - mx); + m[b*n + a] = e; + sum += e; + } + for (int64_t a = 0; a < n; ++a) { + m[b*n + a] /= sum; + } + } + + for (int64_t i = 0; i < n*n; ++i) { + m[i] += eps; + } + + auto norm_cols = [&]() { + for (int64_t a = 0; a < n; ++a) { + float r = 0.0f; + for (int64_t b = 0; b < n; ++b) { + r += m[b*n + a]; + } + r += eps; + for (int64_t b = 0; b < n; ++b) { + m[b*n + a] /= r; + } + } + }; + + auto norm_rows = [&]() { + for (int64_t b = 0; b < n; ++b) { + float c = 0.0f; + for (int64_t a = 0; a < n; ++a) { + c += m[b*n + a]; + } + c += eps; + for (int64_t a = 0; a < n; ++a) { + m[b*n + a] /= c; + } + } + }; + + norm_cols(); + for (int32_t it = 1; it < n_iters; ++it) { + norm_rows(); + norm_cols(); + } + + for (int64_t i = 0; i < n*n; ++i) { + y[i] = m[i]; + } + } +} + +void ggml_compute_forward_sinkhorn_norm( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_sinkhorn_norm_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + // ggml_compute_forward_out_prod static void ggml_compute_forward_out_prod_f32( diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index a8e18c716db7..21d935b6c5f2 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -48,6 +48,7 @@ void ggml_compute_forward_rms_norm_mul_fused(const struct ggml_compute_params * void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_sinkhorn_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 98816f885cf6..097cacac6930 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -33,6 +33,7 @@ #include "ggml-cuda/mmvf.cuh" #include "ggml-cuda/mmvq.cuh" #include "ggml-cuda/norm.cuh" +#include "ggml-cuda/sinkhorn-norm.cuh" #include "ggml-cuda/opt-step-adamw.cuh" #include "ggml-cuda/opt-step-sgd.cuh" #include "ggml-cuda/out-prod.cuh" @@ -2083,6 +2084,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_GROUP_NORM: ggml_cuda_op_group_norm(ctx, dst); break; + case GGML_OP_SINKHORN_NORM: + ggml_cuda_op_sinkhorn_norm(ctx, dst); + break; case GGML_OP_L2_NORM: ggml_cuda_op_l2_norm(ctx, dst); break; @@ -4850,6 +4854,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_SILU_BACK: return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; break; + case GGML_OP_SINKHORN_NORM: + return op->src[0]->type == GGML_TYPE_F32 && + ggml_is_contiguous(op->src[0]) && + op->src[0]->ne[0] == op->src[0]->ne[1] && + op->src[0]->ne[0] <= 8; case GGML_OP_NORM: case GGML_OP_RMS_NORM: case GGML_OP_L2_NORM: diff --git a/ggml/src/ggml-cuda/sinkhorn-norm.cu b/ggml/src/ggml-cuda/sinkhorn-norm.cu new file mode 100644 index 000000000000..fbf896138564 --- /dev/null +++ b/ggml/src/ggml-cuda/sinkhorn-norm.cu @@ -0,0 +1,192 @@ +#include "sinkhorn-norm.cuh" + +#define SINKHORN_NORM_MAX_N 8 + +template +static __global__ void sinkhorn_norm_f32( + const float * __restrict__ x, + float * __restrict__ y, + const int n, + const int n_iters, + const float eps, + const int64_t n_slices, + const int64_t slice_stride) { + + const int64_t s = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (s >= n_slices) { + return; + } + + const float * xs = x + s * slice_stride; + float * ys = y + s * slice_stride; + + float m[MAXN * MAXN]; + const int nn = n * n; + + for (int i = 0; i < nn; ++i) { + m[i] = xs[i]; + } + + for (int b = 0; b < n; ++b) { + float mx = -INFINITY; + for (int a = 0; a < n; ++a) { + mx = fmaxf(mx, m[b*n + a]); + } + float sum = 0.0f; + for (int a = 0; a < n; ++a) { + const float e = expf(m[b*n + a] - mx); + m[b*n + a] = e; + sum += e; + } + for (int a = 0; a < n; ++a) { + m[b*n + a] /= sum; + } + } + + for (int i = 0; i < nn; ++i) { + m[i] += eps; + } + + #define NORM_COLS() do { \ + for (int a = 0; a < n; ++a) { \ + float r = 0.0f; \ + for (int b = 0; b < n; ++b) r += m[b*n + a]; \ + r += eps; \ + for (int b = 0; b < n; ++b) m[b*n + a] /= r; \ + } \ + } while (0) + + #define NORM_ROWS() do { \ + for (int b = 0; b < n; ++b) { \ + float c = 0.0f; \ + for (int a = 0; a < n; ++a) c += m[b*n + a]; \ + c += eps; \ + for (int a = 0; a < n; ++a) m[b*n + a] /= c; \ + } \ + } while (0) + + NORM_COLS(); + for (int it = 1; it < n_iters; ++it) { + NORM_ROWS(); + NORM_COLS(); + } + + #undef NORM_COLS + #undef NORM_ROWS + + for (int i = 0; i < nn; ++i) { + ys[i] = m[i]; + } +} + +template +static __global__ void sinkhorn_norm_warp_f32( + const float * __restrict__ x, + float * __restrict__ y, + const int n_iters, + const float eps, + const int64_t n_slices) { + + constexpr int NN = N * N; + + const int64_t g = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + const int64_t token = g / NN; + if (token >= n_slices) { + return; + } + + const int e = threadIdx.x % NN; + + float v = x[token * NN + e]; + + float mx = v; +#pragma unroll + for (int off = 1; off < N; off <<= 1) { + mx = fmaxf(mx, __shfl_xor_sync(0xffffffff, mx, off, NN)); + } + float ex = expf(v - mx); + float sum = ex; +#pragma unroll + for (int off = 1; off < N; off <<= 1) { + sum += __shfl_xor_sync(0xffffffff, sum, off, NN); + } + v = ex / sum + eps; + + { + float r = v; +#pragma unroll + for (int off = N; off < NN; off <<= 1) { + r += __shfl_xor_sync(0xffffffff, r, off, NN); + } + v /= (r + eps); + } + + for (int it = 1; it < n_iters; ++it) { + { + float c = v; +#pragma unroll + for (int off = 1; off < N; off <<= 1) { + c += __shfl_xor_sync(0xffffffff, c, off, NN); + } + v /= (c + eps); + } + { + float r = v; +#pragma unroll + for (int off = N; off < NN; off <<= 1) { + r += __shfl_xor_sync(0xffffffff, r, off, NN); + } + v /= (r + eps); + } + } + + y[token * NN + e] = v; +} + +static bool sinkhorn_norm_is_pow2(int n) { + return n > 0 && (n & (n - 1)) == 0; +} + +void ggml_cuda_op_sinkhorn_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const float * src0_d = (const float *) src0->data; + float * dst_d = (float *) dst->data; + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src0)); + + GGML_TENSOR_UNARY_OP_LOCALS; + + const int32_t n_iters = ggml_get_op_params_i32(dst, 0); + const float eps = ggml_get_op_params_f32(dst, 1); + + const int n = (int) ne00; + GGML_ASSERT(ne01 == ne00); + GGML_ASSERT(n <= SINKHORN_NORM_MAX_N); + GGML_ASSERT(n_iters >= 1); + + const int64_t n_slices = ne02 * ne03; + const int64_t slice_stride = (int64_t) n * n; + + const int block_size = 256; + + const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + + if (sinkhorn_norm_is_pow2(n) && n*n <= warp_size) { + const int64_t total_threads = n_slices * (int64_t) (n*n); + const int64_t num_blocks = (total_threads + block_size - 1) / block_size; + switch (n) { + case 2: sinkhorn_norm_warp_f32<2><<>>(src0_d, dst_d, n_iters, eps, n_slices); break; + case 4: sinkhorn_norm_warp_f32<4><<>>(src0_d, dst_d, n_iters, eps, n_slices); break; + case 8: sinkhorn_norm_warp_f32<8><<>>(src0_d, dst_d, n_iters, eps, n_slices); break; + default: GGML_ABORT("sinkhorn_norm: unsupported warp N=%d", n); + } + return; + } + + const int64_t num_blocks = (n_slices + block_size - 1) / block_size; + sinkhorn_norm_f32<<>>( + src0_d, dst_d, n, n_iters, eps, n_slices, slice_stride); +} diff --git a/ggml/src/ggml-cuda/sinkhorn-norm.cuh b/ggml/src/ggml-cuda/sinkhorn-norm.cuh new file mode 100644 index 000000000000..d671c7b5bf0c --- /dev/null +++ b/ggml/src/ggml-cuda/sinkhorn-norm.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +void ggml_cuda_op_sinkhorn_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index de0321d9ffd9..42787cd3dd46 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1018,6 +1018,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "RMS_NORM_BACK", "GROUP_NORM", "L2_NORM", + "SINKHORN_NORM", "MUL_MAT", "MUL_MAT_ID", @@ -1096,7 +1097,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); +static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1129,6 +1130,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "rms_norm_back(x)", "group_norm(x)", "l2_norm(x)", + "sinkhorn_norm(x)", "X*Y", "X[i]*Y", @@ -1207,7 +1209,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); +static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -3245,6 +3247,27 @@ struct ggml_tensor * ggml_l2_norm_inplace( return ggml_l2_norm_impl(ctx, a, eps, true); } +// ggml_sinkhorn_norm + +struct ggml_tensor * ggml_sinkhorn_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_iters, + float eps) { + GGML_ASSERT(a->ne[0] == a->ne[1]); + GGML_ASSERT(n_iters >= 1); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, n_iters); + ggml_set_op_params_f32(result, 1, eps); + + result->op = GGML_OP_SINKHORN_NORM; + result->src[0] = a; + + return result; +} + // ggml_mul_mat static inline bool ggml_can_mul_mat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 07aa477e1e47..04f0c6282a72 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -220,34 +220,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn( int il) const { GGML_UNUSED(il); - // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference: - // row softmax over dst, one column normalization, then repeated row/column normalization. - comb = ggml_soft_max(ctx0, comb); - - ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1); - eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps); - - comb = ggml_add(ctx0, comb, eps); - - auto norm_cols = [&]() { - ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); - ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst); - col_sum = ggml_add(ctx0, col_sum, eps); - col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3); - comb = ggml_div(ctx0, comb, col_sum); - }; - - auto norm_rows = [&]() { - ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb); - row_sum = ggml_add(ctx0, row_sum, eps); - comb = ggml_div(ctx0, comb, row_sum); - }; - - norm_cols(); - for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) { - norm_rows(); - norm_cols(); - } + comb = ggml_cont(ctx0, comb); + comb = ggml_sinkhorn_norm(ctx0, comb, (int) hparams.dsv4_hc_sinkhorn_iters, hparams.dsv4_hc_eps); return comb; } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 1fae3f5176c3..88cf2d795c37 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6350,6 +6350,36 @@ struct test_l2_norm : public test_case { } }; +// GGML_OP_SINKHORN_NORM +struct test_sinkhorn_norm : public test_case { + const ggml_type type; + const int64_t n; + const int64_t n_tokens; + const int n_iters; + const float eps; + + std::string vars() override { + return VARS_TO_STR5(type, n, n_tokens, n_iters, eps); + } + + test_sinkhorn_norm(ggml_type type = GGML_TYPE_F32, + int64_t n = 4, + int64_t n_tokens = 128, + int n_iters = 20, + float eps = 1e-6f) + : type(type), n(n), n_tokens(n_tokens), n_iters(n_iters), eps(eps) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a = ggml_new_tensor_3d(ctx, type, n, n, n_tokens); + ggml_set_name(a, "a"); + + ggml_tensor * out = ggml_sinkhorn_norm(ctx, a, n_iters, eps); + ggml_set_name(out, "out"); + + return out; + } +}; + // GGML_OP_ACC struct test_acc : public test_case { const ggml_type type; @@ -8437,6 +8467,15 @@ static std::vector> make_test_cases_eval() { } } + // DeepSeek-V4 mHC fused Sinkhorn-Knopp + for (int64_t hc : {4}) { + for (int64_t nt : {1, 7, 128, 1024}) { + for (int iters : {1, 2, 20}) { + test_cases.emplace_back(new test_sinkhorn_norm(GGML_TYPE_F32, hc, nt, iters, 1e-6f)); + } + } + } + // in-place tests test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 5, 4, 3}, false, 1e-6f, true)); @@ -9403,6 +9442,11 @@ static std::vector> make_test_cases_eval() { static std::vector> make_test_cases_perf() { std::vector> test_cases; + // DeepSeek-V4 mHC fused Sinkhorn-Knopp (4x4 slices, iters=20) + for (int64_t nt : {4096, 65536, 262144}) { + test_cases.emplace_back(new test_sinkhorn_norm(GGML_TYPE_F32, 4, nt, 20, 1e-6f)); + } + // Conv2d: K=CRS=NPQ=4096 matmul performance uint32_t iwh_idx = 0; uint32_t kwh_idx = 1;