diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 0f0d08ec25..40b2c7ca3e 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -705,6 +705,7 @@ extern "C" { GGML_OP_FUSED_RMS_RMS_ADD, GGML_OP_BLEND, GGML_OP_INDEXER_TOPK, + GGML_OP_SINKHORN, GGML_OP_COUNT, }; @@ -2578,6 +2579,22 @@ extern "C" { enum ggml_unary_op op, int n_top_k); + // Sinkhorn normalization of a flat [S*S, T] batch of S x S matrices into + // doubly-stochastic form: softmax over columns, then column normalization, + // then (n_iters - 1) rounds of row + column normalization (ends on columns). + // The flat input is row-major (column index fastest). eps, when non-zero, is + // added to the softmax output and to every normalization sum before dividing. + // With output_transposed the result is [S, S, T] with ne0 = row, ne1 = column + // (ready for out[c] = sum_r m[r,c] * residual[r] consumers); otherwise the + // bare input layout (ne0 = column) is kept. + GGML_API struct ggml_tensor * ggml_sinkhorn( + struct ggml_context * ctx, + struct ggml_tensor * a, + int S, + int n_iters, + float eps, + bool output_transposed); + // custom operators typedef void (*ggml_unary_op_f32_t) (const int, float *, const float *); diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index e006e661cf..242c549d3d 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -56,6 +56,7 @@ #include "ggml-cuda/reduce.cuh" #include "ggml-cuda/tri.cuh" #include "ggml-cuda/delta-net.cuh" +#include "ggml-cuda/sinkhorn.cuh" #include "ggml-cuda/blend.cuh" #include @@ -4123,6 +4124,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_DELTA_NET: ggml_cuda_op_delta_net(ctx, dst); break; + case GGML_OP_SINKHORN: + ggml_cuda_op_sinkhorn(ctx, dst); + break; case GGML_OP_FLASH_ATTN_EXT: ggml_cuda_flash_attn_ext(ctx, dst); break; @@ -5029,6 +5033,11 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons op->src[3]->ne[0] == op->src[0]->ne[2]; case GGML_OP_DELTA_NET: return true; + case GGML_OP_SINKHORN: { + const int sink_s = op->op_params[0]; + return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && + sink_s >= 1 && sink_s <= 8 && op->src[0]->ne[0] == (int64_t) sink_s*sink_s; + } case GGML_OP_FLASH_ATTN_EXT: #if defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__) return (op->src[0]->ne[0] == 64 && op->src[1]->type == GGML_TYPE_F16) || op->src[0]->ne[0] == 128; diff --git a/ggml/src/ggml-cuda/sinkhorn.cu b/ggml/src/ggml-cuda/sinkhorn.cu new file mode 100644 index 0000000000..f02a8a440c --- /dev/null +++ b/ggml/src/ggml-cuda/sinkhorn.cu @@ -0,0 +1,103 @@ +#include "common.cuh" +#include "sinkhorn.cuh" + +// Sinkhorn normalization of T independent S x S matrices (S <= 8, so a matrix is +// at most 64 floats). One thread per token: the whole matrix lives in a thread-local +// array and the 6-node-per-iteration graph chain collapses into this single kernel. +// Semantics match the reference: softmax over columns, column normalization, +// then (iters - 1) rounds of row + column normalization (ends on columns). +// Input is the flat [S*S, T] row-major tensor (column index fastest); output is +// [S, S, T] with ne0 = row, i.e. transposed on write. + +template +static __global__ void k_sinkhorn(const float * __restrict__ x, float * __restrict__ dst, + const int64_t T, const int iters, const float eps, + const int transposed, const int64_t nb1) { + const int64_t t = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + if (t >= T) { + return; + } + + const float * xt = (const float *)((const char *) x + t*nb1); + float m[S*S]; + + #pragma unroll + for (int r = 0; r < S; ++r) { + float mx = xt[r*S]; + for (int c = 1; c < S; ++c) mx = fmaxf(mx, xt[r*S + c]); + float sum = 0.0f; + for (int c = 0; c < S; ++c) { m[r*S + c] = expf(xt[r*S + c] - mx); sum += m[r*S + c]; } + for (int c = 0; c < S; ++c) m[r*S + c] = m[r*S + c]/sum + eps; + } + #pragma unroll + for (int c = 0; c < S; ++c) { + float sum = eps; + for (int r = 0; r < S; ++r) sum += m[r*S + c]; + for (int r = 0; r < S; ++r) m[r*S + c] /= sum; + } + for (int i = 0; i < iters - 1; ++i) { + #pragma unroll + for (int r = 0; r < S; ++r) { + float sum = eps; + for (int c = 0; c < S; ++c) sum += m[r*S + c]; + for (int c = 0; c < S; ++c) m[r*S + c] /= sum; + } + #pragma unroll + for (int c = 0; c < S; ++c) { + float sum = eps; + for (int r = 0; r < S; ++r) sum += m[r*S + c]; + for (int r = 0; r < S; ++r) m[r*S + c] /= sum; + } + } + + float * yt = dst + t*S*S; + if (transposed) { + #pragma unroll + for (int c = 0; c < S; ++c) { + for (int r = 0; r < S; ++r) yt[c*S + r] = m[r*S + c]; + } + } else { + #pragma unroll + for (int k = 0; k < S*S; ++k) yt[k] = m[k]; + } +} + +void ggml_cuda_op_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + + const int S = dst->op_params[0]; + const int iters = dst->op_params[1]; + float eps; + memcpy(&eps, &dst->op_params[2], sizeof(float)); + const int transposed = dst->op_params[3]; + const int64_t T = src0->ne[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(S >= 1 && S <= 8); + GGML_ASSERT(src0->ne[0] == (int64_t) S * S); + GGML_ASSERT(ggml_is_contiguous(dst)); + + if (T == 0) { + return; + } + + const int block = 256; + const int64_t grid = (T + block - 1)/block; + cudaStream_t stream = ctx.stream(); + + const float * x = (const float *) src0->data; + float * y = (float *) dst->data; + + switch (S) { + case 1: k_sinkhorn<1><<>>(x, y, T, iters, eps, transposed, src0->nb[1]); break; + case 2: k_sinkhorn<2><<>>(x, y, T, iters, eps, transposed, src0->nb[1]); break; + case 3: k_sinkhorn<3><<>>(x, y, T, iters, eps, transposed, src0->nb[1]); break; + case 4: k_sinkhorn<4><<>>(x, y, T, iters, eps, transposed, src0->nb[1]); break; + case 5: k_sinkhorn<5><<>>(x, y, T, iters, eps, transposed, src0->nb[1]); break; + case 6: k_sinkhorn<6><<>>(x, y, T, iters, eps, transposed, src0->nb[1]); break; + case 7: k_sinkhorn<7><<>>(x, y, T, iters, eps, transposed, src0->nb[1]); break; + case 8: k_sinkhorn<8><<>>(x, y, T, iters, eps, transposed, src0->nb[1]); break; + default: GGML_ABORT("sinkhorn: unsupported S"); + } +} diff --git a/ggml/src/ggml-cuda/sinkhorn.cuh b/ggml/src/ggml-cuda/sinkhorn.cuh new file mode 100644 index 0000000000..c9f86186c6 --- /dev/null +++ b/ggml/src/ggml-cuda/sinkhorn.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +void ggml_cuda_op_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index c105ea37f8..6941ec35fe 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -4323,9 +4323,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "FUSED_RMS_RMS_ADD", "BLEND", "INDEXER_TOPK", + "SINKHORN", }; -static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); +static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -4445,10 +4446,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "rms(x1)+rms(x2)", "blend(a,b,c)", "indexer_topk(k, q, w, mask)", + "sinkhorn(x)", }; -static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); +static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -10129,6 +10131,32 @@ struct ggml_tensor * ggml_indexer_topk( } +struct ggml_tensor * ggml_sinkhorn( + struct ggml_context * ctx, + struct ggml_tensor * a, + int S, + int n_iters, + float eps, + bool output_transposed) { + GGML_ASSERT(eps >= 0.0f); + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(S >= 1 && S <= 8); + GGML_ASSERT(a->ne[0] == (int64_t) S * S); + GGML_ASSERT(a->ne[2] == 1 && a->ne[3] == 1); + GGML_ASSERT(n_iters >= 1); + GGML_ASSERT(ggml_is_contiguous(a)); + + struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, S, a->ne[1]); + result->op = GGML_OP_SINKHORN; + result->op_params[0] = S; + result->op_params[1] = n_iters; + memcpy(&result->op_params[2], &eps, sizeof(float)); + result->op_params[3] = output_transposed ? 1 : 0; + result->src[0] = a; + + return result; +} + // ggml_fill static struct ggml_tensor * ggml_fill_impl( @@ -23017,6 +23045,83 @@ static void ggml_compute_forward_delta_net( } } +// ggml_compute_forward_sinkhorn + +static void ggml_compute_forward_sinkhorn_f32( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * src0 = dst->src[0]; + + const int S = dst->op_params[0]; + const int iters = dst->op_params[1]; + float eps; + memcpy(&eps, &dst->op_params[2], sizeof(float)); + const int transposed = dst->op_params[3]; + const int64_t T = src0->ne[1]; + + GGML_ASSERT(S >= 1 && S <= 8); + GGML_ASSERT(src0->ne[0] == (int64_t) S * S); + GGML_ASSERT(iters >= 1); + + // one token is S*S floats (16 at S=4): parallelize over tokens only + const int64_t t0 = (T * params->ith ) / params->nth; + const int64_t t1 = (T * (params->ith+1)) / params->nth; + + float m[64]; + + for (int64_t t = t0; t < t1; ++t) { + const float * x = (const float *)((const char *)src0->data + t*src0->nb[1]); + float * y = (float *)(( char *)dst->data + t*dst->nb[2]); + + // softmax over columns c for each row r; flat input is row-major (c fastest) + for (int r = 0; r < S; ++r) { + float mx = x[r*S]; + for (int c = 1; c < S; ++c) mx = MAX(mx, x[r*S + c]); + float sum = 0.0f; + for (int c = 0; c < S; ++c) { m[r*S + c] = expf(x[r*S + c] - mx); sum += m[r*S + c]; } + for (int c = 0; c < S; ++c) m[r*S + c] = m[r*S + c]/sum + eps; + } + // column normalization first, then (iters - 1) rounds of row + column: ends on columns + for (int c = 0; c < S; ++c) { + float sum = eps; + for (int r = 0; r < S; ++r) sum += m[r*S + c]; + for (int r = 0; r < S; ++r) m[r*S + c] /= sum; + } + for (int i = 0; i < iters - 1; ++i) { + for (int r = 0; r < S; ++r) { + float sum = eps; + for (int c = 0; c < S; ++c) sum += m[r*S + c]; + for (int c = 0; c < S; ++c) m[r*S + c] /= sum; + } + for (int c = 0; c < S; ++c) { + float sum = eps; + for (int r = 0; r < S; ++r) sum += m[r*S + c]; + for (int r = 0; r < S; ++r) m[r*S + c] /= sum; + } + } + if (transposed) { + // dst is [row, col, T] (ne0 = row): transpose on write + for (int c = 0; c < S; ++c) { + for (int r = 0; r < S; ++r) y[c*S + r] = m[r*S + c]; + } + } else { + for (int k = 0; k < S*S; ++k) y[k] = m[k]; + } + } +} + +static void ggml_compute_forward_sinkhorn( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: + ggml_compute_forward_sinkhorn_f32(params, dst); + break; + default: + GGML_ABORT("fatal error"); + } +} + // ggml_compute_forward_win_part static void ggml_compute_forward_win_part_f32( @@ -24755,6 +24860,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml { ggml_compute_forward_delta_net(params, tensor); } break; + case GGML_OP_SINKHORN: + { + ggml_compute_forward_sinkhorn(params, tensor); + } break; case GGML_OP_INDEXER_TOPK: { if (!iqk_indexer_topk(tensor, params->wdata, (barrier_t)ggml_barrier, (void *)params->shared, params->ith, params->nth)) { @@ -25823,6 +25932,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor case GGML_OP_SOLVE_TRI: case GGML_OP_DELTA_NET: case GGML_OP_INDEXER_TOPK: + case GGML_OP_SINKHORN: { GGML_ABORT("fatal error"); // TODO: not implemented } @@ -26559,6 +26669,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_SOLVE_TRI: case GGML_OP_DELTA_NET: case GGML_OP_INDEXER_TOPK: + case GGML_OP_SINKHORN: { n_tasks = n_threads; } break; diff --git a/src/graphs/build_openpangu.cpp b/src/graphs/build_openpangu.cpp index b0c6faac47..91b40fb3f2 100644 --- a/src/graphs/build_openpangu.cpp +++ b/src/graphs/build_openpangu.cpp @@ -222,39 +222,6 @@ static ggml_tensor * openpangu_causal_conv(ggml_context * ctx, ggml_cgraph * gf, return out; } -// --- mHC Sinkhorn: h_res [S*S, T] -> doubly-stochastic per token, 20 iters (ends on col norm) --- -static ggml_tensor * openpangu_sinkhorn(ggml_context * ctx, ggml_tensor * h_res_flat, - int64_t S, int64_t T, int iters, float hc_eps) { - // The flat h_res is torch [r,c] row-major (c fastest), so a bare reshape gives ne0=col. - // Transpose once so ne0=row(S), ne1=col(S): every axis op below then matches the - // reference _mhc_sinkhorn_naive (softmax over col, first norm over row, end on col-sum=1) - // and mhc_post's out[c] = sum_r m[r,c]*residual[r]. - (void) hc_eps; // softmax outputs are strictly positive, so the eps is numerically inert here - ggml_tensor * m = ggml_reshape_3d(ctx, h_res_flat, S, S, T); // ne0=col (bare reshape) - // ref softmaxes h_res over columns; a bare reshape already has ne0=col, so soft_max - // (over ne0) hits the column axis directly -- no pre-permute round-trip needed. - m = ggml_soft_max(ctx, m); // softmax over col - m = ggml_cont(ctx, ggml_permute(ctx, m, 1, 0, 2, 3)); // transpose once -> [row, col, T] - - auto col_norm = [&](ggml_tensor * a) { - ggml_tensor * col_sum = ggml_sum_rows(ctx, a); // sums ne0(row) -> [1, col, T] - return ggml_div(ctx, a, col_sum); // broadcast [1,col,T] over rows - }; - auto row_norm = [&](ggml_tensor * a) { - ggml_tensor * ap = ggml_cont(ctx, ggml_permute(ctx, a, 1, 0, 2, 3)); // [col,row,T] - ggml_tensor * row_sum = ggml_sum_rows(ctx, ap); // [1, row, T] - ggml_tensor * out = ggml_div(ctx, ap, row_sum); - return ggml_cont(ctx, ggml_permute(ctx, out, 1, 0, 2, 3)); // back [row,col,T] - }; - - m = col_norm(m); - for (int i = 0; i < iters - 1; ++i) { - m = row_norm(m); - m = col_norm(m); - } - return m; // [row(S), col(S), T] -} - // Attention sublayer body, shared by the base layers and the NextN/MTP head. // x_normed = input-layernormed hidden [n_embd, T]; returns post-o_proj output [n_embd, T]. // conv_state is the recurrent MoME state slot. seq_qnext is the [1, T] sequence-id input @@ -852,7 +819,6 @@ ggml_cgraph * llm_build_context::build_openpangu() { const int64_t n_embd_head_k = hparams.n_embd_head_k(0); // 192 const int64_t S = hparams.mhc_num_stream; // 4 const int sink_iters = (int) hparams.mhc_recur_norm; // 20 - const float hc_eps = 1e-6f; const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k)); @@ -1052,7 +1018,7 @@ ggml_cgraph * llm_build_context::build_openpangu() { // cont is required: the CUDA broadcast-mul path misreads strided views (h_pre is a // row-slice of mixes), while CPU handles the strides — token 0 right, tokens 1+ garbage h_pre = ggml_add(ctx0, ggml_mul(ctx0, ggml_cont(ctx0, h_pre), a_pre), b_pre); // broadcast scalar + [S] - h_pre = ggml_sigmoid(ctx0, h_pre); // [S,T] (+hc_eps omitted, inert) + h_pre = ggml_sigmoid(ctx0, h_pre); // [S,T] (+eps omitted, inert) // combine: x[h,t] = sum_s h_pre[s,t] * R[h,s,t] ggml_tensor * hpre3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, h_pre), 1, S, n_tokens); @@ -1077,7 +1043,7 @@ ggml_cgraph * llm_build_context::build_openpangu() { h_post = ggml_scale(ctx0, ggml_sigmoid(ctx0, h_post), 2.0f); // 2*sigmoid, [S,T] ggml_tensor * m = ggml_add(ctx0, ggml_mul(ctx0, h_res, a_res), b_res); // [S*S,T] - m = openpangu_sinkhorn(ctx0, m, S, n_tokens, sink_iters, hc_eps); // [row S, col S, T] + m = ggml_sinkhorn(ctx0, m, (int) S, sink_iters, 0.0f, /*output_transposed=*/true); // [row S, col S, T] // term1: h_post[s,t]*y[h,t] -> [H,S,T] ggml_tensor * y3 = ggml_reshape_3d(ctx0, y, n_embd, 1, n_tokens);