feat(ggml): in-graph collective node via a caller-registered callback

[maikzz32]

Multi-device inference that partitions work across processes has to combine partial sums between the kernels that produce them and the kernels that consume them. Doing that outside the graph means a host round trip per combination: download, reduce, upload, synchronize. For a 43-layer MoE model that is 43 round trips per forward.

ggml_cluster_allreduce(ctx, a, fn, user) adds a node whose value is a summed element-wise across the caller's process group. The CUDA/HIP backend runs it on the same stream as the surrounding kernels, so it is ordered after the producers and before the consumers with no host synchronization.

ggml gains no dependency

The node carries a callback and a user pointer in op_params; the caller decides how the sum is produced (NCCL, RCCL, MPI, a test double):

typedef void (*ggml_cluster_allreduce_fn)(void * user, void * data, size_t n, void * stream);

No new library, no new build option, no new backend entry point. It is a sub-op of the existing GGML_OP_MOE_FUSED family, so supports_op needs no change.

Graph capture

A graph containing such a node is not captured by default, and the reason is measured rather than assumed.

On gfx1151 with RCCL 2.30.4, capture is bit-exact — a 128-token greedy completion stays byte-identical — and replaying a 5445-node graph that holds 43 collectives is slower than launching it eagerly:

	eager	captured
autoregressive decode	21.5 tok/s	19.15 tok/s
speculative decode (q=4)	29.45 tok/s	26.65 tok/s

Two nodes over RoCE v2, DeepSeek V4 Flash, medians of 3 runs, same binary and clocks, GGML_CUDA_GRAPH_STATS=1 confirming 23 replays per 25 forwards (so this is not capture churn). GGML_CUDA_COLLECTIVE_GRAPH_CAPTURE=1 allows it for anyone whose runtime behaves differently.

What the node itself is worth

Measured by returning early from the callback — wrong output, exact timing — the whole 43-collective budget is 2.7 ms of a 48 ms decode step on that hardware. The node's value is not the collective's speed but that it removes the per-layer host round trip and lets a rank keep using the fused whole-model graph; without it the same model falls back to 43 host-driven per-layer graphs.

Hardware note: all numbers are gfx1151 (Radeon 8060S, Ryzen AI Max 395, ROCm 10), not the sm_86+ hardware in CONTRIBUTING. Happy to have them re-run on a 3090 or the 395 box if that is useful.

[cubic-dev-ai]

No issues found across 4 files

_{Re-trigger cubic}

[davide221]

Pushed onto your branch: merged main, and added test_ggml_cluster_allreduce. It uses a stand-in two-rank callback to check ordering between producer and consumer, the stream and element count the callback receives, and that the callback runs on every compute. It's wired into the NVIDIA and ROCm CI jobs. It passes on R9700 (gfx1201) and Strix (gfx1151). With GGML_CUDA_COLLECTIVE_GRAPH_CAPTURE=1 it aborts, because the stand-in synchronizes. That shows the default-off capture guard is what keeps the node eager.

We plan to run two boxes over ConnectX, so this is wanted. Would you open a PR with the runtime that drives it (rank bootstrap, RCCL communicator, which layers get reduced)?