ggml-cuda: overlap the MoE shared expert on a separate stream#36
Open
roberteg16 wants to merge 6 commits into
Open
ggml-cuda: overlap the MoE shared expert on a separate stream#36roberteg16 wants to merge 6 commits into
roberteg16 wants to merge 6 commits into
Conversation
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>
roberteg16
changed the title
Add a shared-expert detector to 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 decode (single token). The overlap only helps when the routed branch leaves the GPU underutilized for the shared expert to run alongside it: that holds in decode (batch 1, latency/occupancy-bound) but not in prefill, where the routed matmuls already saturate the GPU and 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>
The branch scratch buffer made the node interleaving obsolete, so the concurrent branches already run in their original graph order. The "restore original order" pass in ggml_cuda_graph_evaluate_and_capture is therefore a no-op, and the ggml_cuda_concurrent_event::original_order it consumed is now unused. Remove the reorder pass and original_order, keeping only the is_valid() check that clears the events when a region cannot be safely overlapped. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
roberteg16
force-pushed
the
rogarcia.cuda-multistream-shexp
branch
from
e11e7b5 to
6277fbe
Compare
Move the diamond-detection heuristics (attention QKV fan-out and the MoE shared-expert diamond) out of ggml_backend_cuda_graph_optimize into a header-only function ggml_cuda_detect_concurrent_regions() that operates only on the public ggml graph API and holds no CUDA state. The backend now turns the detected regions into ggml_cuda_concurrent_event objects and lays their branch tensors into the scratch buffer. Behavior is unchanged, and the detection heuristics can now be exercised in isolation. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Built only for CUDA/HIP/MUSA: - detection: drive ggml_cuda_detect_concurrent_regions() with synthetic diamond graphs and assert which regions are and are not detected - attention 3-way fan-out, MoE shared-expert diamond, and negatives (fan-out of 2 or 4, multi-row/prefill, wrong fork name, wrong MoE operand names). This also pins the tensor-name contract the detector relies on. - end-to-end: run the diamonds through ggml_backend_sched with GGML_CUDA_GRAPH_OPT=1 and check the output against an analytic reference; a log hook confirms both concurrency paths engaged. Skipped when no GPU is available. Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
|
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this changes
A decode-only optimization for MoE models that have a gated shared expert, plus a supporting refactor and tests.
Optimization
ggml_cuda_op_mul_mathardcoded 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 usectx.stream()on the non-split path and a per-(device,stream) cuBLAS handle.ggml_backend_cuda_graph_optimizedetects thejoin = add(ffn_moe_out, ffn_shexp)diamond by callback names. 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).Gated to decode (single token) and byte-identical to sequential. Requires CUDA graphs and a single GPU. Enabled by default on RDNA3.5 (the architecture it is tuned for) and off elsewhere;
GGML_CUDA_GRAPH_OPT=1/=0forces it on/off on any architecture.Cleanup & refactor
ggml_cuda_graph_evaluate_and_capture(and theoriginal_orderit consumed) is now a no-op and is removed; only theis_valid()fallback that disables an unsafe region is kept.ggml_backend_cuda_graph_optimizeinto a header-onlyggml_cuda_detect_concurrent_regions()that operates only on the public ggml graph API and holds no CUDA state. Behavior unchanged; the heuristics can now be exercised in isolation.Default
GGML_CUDA_GRAPH_OPTis unset, the optimization now defaults on for RDNA3.5 (the architecture it is tuned for) and off elsewhere; the env var still forces it on/off on any architecture.Tests
ggml_cuda_detect_concurrent_regions()with synthetic diamond graphs and assert which regions are and are not detected: attention 3-way fan-out, MoE shared-expert diamond, and negatives (fan-out of 2 or 4, multi-row/prefill, wrong fork name, wrong MoE operand names). Also pins the tensor-name contract the detector relies on.ggml_backend_schedwithGGML_CUDA_GRAPH_OPT=1and check the output against an analytic referenceMeasured results
Command:
llama-bench -hf unsloth/Qwen3.6-35B-A3B-GGUF:Q4_K_M -ngl 999 -n 128,1024 -r 1