ggml: add fused sinkhorn op (eps + output-layout params) (#2115)
* ggml: add fused sinkhorn op (eps + output-layout params); use it for openPangu mHC * openpangu: call ggml_sinkhorn directly from mhc_post --------- Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
This commit is contained in:
Joel Farthing
GitHub
parent
0a1dd13c95
commit
97bc869552
|
|
@ -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 *);
|
||||
|
|
|
|||
|
|
@ -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 "ggml-cuda/indexer_topk.cuh"
|
||||
|
||||
|
|
@ -4127,6 +4128,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;
|
||||
|
|
@ -5038,6 +5042,11 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
|
|||
case GGML_OP_DELTA_NET:
|
||||
case GGML_OP_INDEXER_TOPK:
|
||||
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;
|
||||
|
|
|
|||
|
|
@ -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 <int S>
|
||||
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><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
|
||||
case 2: k_sinkhorn<2><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
|
||||
case 3: k_sinkhorn<3><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
|
||||
case 4: k_sinkhorn<4><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
|
||||
case 5: k_sinkhorn<5><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
|
||||
case 6: k_sinkhorn<6><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
|
||||
case 7: k_sinkhorn<7><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
|
||||
case 8: k_sinkhorn<8><<<grid, block, 0, stream>>>(x, y, T, iters, eps, transposed, src0->nb[1]); break;
|
||||
default: GGML_ABORT("sinkhorn: unsupported S");
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
#include "common.cuh"
|
||||
|
||||
void ggml_cuda_op_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
|
||||
115
ggml/src/ggml.c
115
ggml/src/ggml.c
|
|
@ -4333,9 +4333,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",
|
||||
|
|
@ -4455,10 +4456,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");
|
||||
|
||||
|
|
@ -10139,6 +10141,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(
|
||||
|
|
@ -23086,6 +23114,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(
|
||||
|
|
@ -24828,6 +24933,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)) {
|
||||
|
|
@ -25896,6 +26005,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
|
||||
}
|
||||
|
|
@ -26633,6 +26743,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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -876,7 +843,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));
|
||||
|
||||
|
||||
|
|
@ -1076,7 +1042,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);
|
||||
|
|
@ -1101,7 +1067,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);
|
||||
|
|
|
|||
Loading…
Reference in New Issue