ggml-cuda: overlap the MoE shared expert on a separate stream by roberteg16 · Pull Request #36 · ROCm/llama.cpp

roberteg16 · 2026-07-06T10:39:36Z

What this changes

A small-batch (decode) optimization for MoE models that have a gated shared expert, building on the stream-based concurrency from ggml-org#16991.

honor the current stream in multi-stream concurrency - ggml_cuda_op_mul_mat hardcoded stream 0, so ops assigned to an aux stream still ran on the main stream; and the single per-device cuBLAS handle was shared across streams (its workspace corrupts concurrent GEMMs). Fixed to use ctx.stream() on the non-split path and a per-(device,stream) cuBLAS handle.
isolate branch scratch - replace the fragile node-interleaving (which could let ggml-alloc recycle a tensor read across the region, corrupting a concurrent read) with a dedicated compute buffer for the concurrent branch nodes: pre-assigned per-node offsets, reused across layers so the branches are structurally disjoint. This is what makes the concurrency correct; it covers both the existing attention QKV branches (from CUDA: add stream-based concurrency ggml-org/llama.cpp#16991) and the shared expert below.
overlap the shared expert - ggml_backend_cuda_graph_optimize detects the join = add(ffn_moe_out, ffn_shexp) diamond by callback names (in a helper ggml_cuda_detect_shared_expert_concurrency). The routed experts stay on the main stream; only the shared expert forks onto one aux stream, joined at the add (one fork/join, reusing the branch scratch).
enable by default on RDNA3.5 - when GGML_CUDA_GRAPH_OPT is unset, the optimization defaults on for RDNA3.5 (the architecture it is tuned for) and off elsewhere; GGML_CUDA_GRAPH_OPT=1/=0 forces it on/off on any architecture.

Gated to small batches - up to MMVQ_MAX_BATCH_SIZE tokens, i.e. the memory-bound mat-vec regime where the routed branch leaves the GPU underutilized; beyond that the routed matmuls saturate the GPU and the overlap only adds contention. Output is byte-identical to sequential. Requires CUDA graphs and a single GPU.

Measured results

Command: llama-bench -hf unsloth/Qwen3.6-35B-A3B-GGUF:Q4_K_M -ngl 999 -n 128,1024 -r 1

test	GGML_CUDA_GRAPH_OPT=0 (t/s)	GGML_CUDA_GRAPH_OPT=1 (t/s)	Δ absolute (t/s)	Δ relative
tg128	53.33	57.28	+3.95	+7.41%
tg512	53.72	57.55	+3.83	+7.13%
tg1024	53.62	57.63	+4.01	+7.48%
tg2048	53.48	57.41	+3.93	+7.35%
tg4096	53.23	57.14	+3.91	+7.35%

Two correctness fixes for the GGML_CUDA_GRAPH_OPT fork/join concurrency so that work assigned to a non-default stream actually runs there: - ggml_cuda_op_mul_mat (the legacy tiled dispatcher) hardcoded ctx.stream(id, 0) for the src0 setup and ctx.stream(id, is) for the per-column compute loop. For a single GPU both resolve to stream 0 regardless of curr_stream_no, so any op on this path ran on the main stream while its aux-stream consumers only waited the fork event and read stale data. Use ctx.stream() for the non-split case; the multi-GPU split path keeps its per-device stream 0 for peer synchronization. This is a no-op when concurrency is off (curr_stream_no == 0). - The cuBLAS handle was shared across streams (one per device). The handle carries a workspace that concurrent GEMMs on different streams corrupt, so give each (device, stream) its own handle. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

The attention QKV concurrency (GGML_CUDA_GRAPH_OPT) produced incorrect results because it relied on interleaving the branch nodes in the graph so that ggml-alloc would keep them non-overlapping. ggml-alloc sees the interleaved order while the executor restores the original order, and is_valid() only checks the branches against each other, not against tensors read across the region (e.g. the fork output every branch consumes concurrently). ggml-alloc could therefore recycle such a tensor's memory for a branch scratch, corrupting a concurrent read. Replace the interleave with a dedicated compute buffer for the concurrent branch nodes. Their data/buffer is pre-assigned into one CUDA buffer, reused across layers (which run sequentially) with distinct per-node offsets within a region, so the branches are structurally disjoint from each other and from the rest of the graph. is_valid() then accepts the event without any graph reordering. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

mgehre-amd · 2026-07-07T16:25:20Z

I don't understand why this changes needs so much code changes. Is there a more minimal version? Especially because large diffs will make it hard to upstream.

roberteg16 · 2026-07-08T08:20:17Z

I don't understand why this changes needs so much code changes. Is there a more minimal version? Especially because large diffs will make it hard to upstream.

Most of the changes come from the fact that I wanted to add a test. llama.cpp has no test right now for the functionality I fixed and it was broken before. I can drop it and this will make many changes disappear. Perhaps we could reintroduce it back when we do the upstream

Add a shared-expert detector, ggml_cuda_detect_shared_expert_concurrency, called from ggml_backend_cuda_graph_optimize after the attention pass, under GGML_CUDA_GRAPH_OPT. It matches the diamond join = ggml_add(ffn_moe_out, ffn_shexp*) by callback names and finds the FFN-input norm that forks into both branches. The routed experts stay on the main stream and only the shared expert forks onto a single aux stream, joined at the add. Keeping the large routed branch on the main stream (region nodes not in the stream map default to the main stream) avoids migrating it and needs just one fork/join. The shared-expert branch reuses the concurrent-branch scratch buffer so it is disjoint from the routed branch without any graph reordering. Gated to the mat-vec regime (up to MMVQ_MAX_BATCH_SIZE tokens). The overlap only helps when the routed branch leaves the GPU underutilized for the shared expert to run alongside it: that holds while the token count keeps the matmuls in the memory/occupancy-bound mat-vec path, but not once it grows past that and the routed matmuls saturate the GPU (prefill), where the overlap only adds contention. Output is byte-identical to sequential. Requires GGML_CUDA_GRAPH_OPT, CUDA graphs and a single GPU, so it is off by default. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

GGML_CUDA_GRAPH_OPT still forces the optimization on ("1") or off ("0"), but when it is unset the default is now on for RDNA3.5 - the architecture where the decode-time shared-expert overlap has been tuned - and off everywhere else. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

roberteg16 and others added 2 commits July 5, 2026 02:46

roberteg16 changed the title ~~Rogarcia.cuda multistream shexp~~ ggml-cuda: overlap the MoE shared expert on a separate stream Jul 6, 2026

roberteg16 force-pushed the rogarcia.cuda-multistream-shexp branch from e11e7b5 to 6277fbe Compare July 7, 2026 14:11

roberteg16 marked this pull request as ready for review July 7, 2026 15:53

roberteg16 requested review from liangliangchang and mgehre-amd July 7, 2026 15:54

roberteg16 force-pushed the rogarcia.cuda-multistream-shexp branch from f194596 to 1a52792 Compare July 8, 2026 09:26