ggml : add fused SINKHORN_NORM op (Sinkhorn-Knopp normalization)

[jadenmach2]

Adds a new ggml op, GGML_OP_SINKHORN_NORM (ggml_sinkhorn_norm), a fused Sinkhorn-Knopp doubly-stochastic normalization for small square matrices, and uses it in the DeepSeek-V4 mHC (hyper-connections) graph. Today that's about 137 primitive ops per layer, which on a 4x4 slice is almost entirely launch and memory overhead. The new op does it all in one kernel per layer, keeping the matrix in registers.

Profiling DeepSeek-V4-Flash (UD-Q4_K_XL, 284B) on an MI250X (gfx90a, ROCm) with rocprofv3 showed the decomposed Sinkhorn at around 30% of GPU kernel time across roughly 250k tiny dispatches (broadcast div/add ~31%, row sums ~8%, transpose ~4.5%).

Results

DeepSeek-V4-Flash, 4x MI250X GCDs, llama-bench:

variant	pp64	tg16
decomposed (baseline)	111.5	10.3
fused (this PR)	119.8	16.5
+60.9% decode, +7.4% prefill.

• ggml_sinkhorn_norm(ctx, a, n_iters, eps), operates per [n, n] slice. It's a general square-matrix op; DeepSeek-V4 is just the first consumer.
• CPU reference plus a CUDA/HIP kernel (warp-cooperative with shfl_xor reductions for power-of-two n with n*n <= warp size, generic fallback otherwise).
• test-backend-ops cases (n=4, tokens 1/7/128/1024, iters 1/2/20); CPU and CUDA match.

Test plan

• test-backend-ops -o SINKHORN_NORM passes on CPU and CUDA/HIP
• DeepSeek-V4 output matches the decomposed graph

Design notes / open questions

• Public op vs. internal helper: I exposed this as a general GGML_OP_SINKHORN_NORM rather than a DeepSeek-only helper, following the Mamba (SSM_CONV/SSM_SCAN) and RWKV (RWKV_WKV6/7, GATED_LINEAR_ATTN) precedent of keeping fused per-arch ops in the core enum. Happy to fold it into an internal helper if you'd rather keep the enum lean.
• The warp-cooperative kernel assumes n is a power of two and n*n <= warp size (true for hc=4); everything else uses the generic one-thread-per-slice fallback.
• Non-CUDA backends fall back to the CPU reference via supports_op.

AI usage disclosure

I directed the design here aka the profiling and root-causing, the decision to fuse the normalization, the math, and the kernel strategy. I used an LLM as an assistant while implementing the code (roughly a 70-30 split, me to LLM), reviewed all of it myself. This description is my own. I understand the project restricts predominantly AI-generated contributions and defer to the maintainers' judgment.

[satindergrewal]

@jadenmach2 I tested your PR #25421 on my M3 Max 128GB with model DeepSeek-V4-Flash UD-Q2_K_XL from Unsloth on master branch 8f114a9, fa=0, ub=1024.

I also added a custom Metal SINKHORN_NORM kernel and got the following results:

build	pp3072 (t/s)	tg32 (t/s)
master 8f114a9 (baseline)	103.0	4.36
+ PR #25421 (Metal = CPU fallback)	106.9 *	5.03
+ PR #25421 + Metal kernel	133.0	4.90

* measured via llama-completion; other rows are llama-bench

You can find the last changes and updates done on my fork's branch satindergrewal:deepseek-v4-flash

See if this helps you or the maintainers.

You can take the idea or code and adapt it, or I can make a PR if it helps better that way.

For some reason I got underwhelming results on decode, which I'm still investigating. Hopefully it will be resolved sometime soon too, if not today, probably by tomorrow.

I'm just trying to make DSpark work with llama.cpp and coming across many layers which can help me get there.

There's a whole lot I need to do to add DSpark support to llama.cpp properly.

Let me know if I can be of any help anywhere.

Full disclosure: I used AI to generate custom Metal kernel under my direction and verified with test-backend-ops against the CPU reference plus identical greedy output, same approach as my #25395.

Cheers,
Satinder

[satindergrewal]

re-ran under controlled conditions after finding a caching artifact in my first run - I benched right after serving a different large model, so the 90GB model was partly evicted from page cache, suppressing everything roughly 2x. Corrected table below:

build	pp3072 (t/s)	tg32 (t/s)
master 8f114a9 (baseline)	213.3 ± 1.5	8.64
+ PR #25421 (Metal = CPU fallback)	219.3 ± 1.8	10.41
+ PR #25421 + Metal kernel	219.6 ± 0.9	11.55

That means the PR alone is +20% decode on Metal even with CPU fallback, +34% with the kernel.

[satindergrewal]

Follow-up to my earlier numbers: the underwhelming decode I mentioned is resolved. It was not this PR, it was graph dispatch overhead elsewhere in the V4 path. With this PR as the base layer the full stack looked like this on M3 Max 128GB (DeepSeek-V4-Flash UD-Q2_K_XL, paired same-window runs, master 8f114a9):

build	tg32 (t/s)
master 8f114a9	8.64
+ this PR + Metal sinkhorn kernel	11.55 (+34%)
+ HC fusion (#23122 ops + Metal kernels)	16.11 (+39% more)
+ top-k gather at decode	+8-10% at 16K depth, grows with context

Cumulative at 16K depth: 8.06 to 15.1 t/s, roughly 1.9x master decode.

Cross-backend check: the same stack on CUDA (RTX Pro 6000, sm120) gave +85% decode and +10% pp over that master. The win is backend-independent, which points at graph dispatch as V4's real bottleneck rather than anything Metal-specific.

Context update: #25585 has since merged fused HC ops for CPU and CUDA, so on CUDA current master already captures most of the HC gain shown above. On Metal there are no kernels for the fused path yet, so the Metal rows still represent headroom over current master. Everything here is test-backend-ops gated and lives on satindergrewal:deepseek-v4-flash. The offer from my first comment stands: take any piece directly, or I can open PRs for the pieces separately. The two self-contained candidates are Metal kernels for the merged #25585 ops and the top-k gather-decode change.

Full disclosure, same as before: kernels are AI-generated under my direction, verified with test-backend-ops and paired benches.

[jadenmach2]

@satindergrewal Thank you for the interest. I am closing this PR as better fused ops were merged into the upstream.