cuda: MoE prefill — fused SwiGLU epilogue and fused weighted reduction - #107
Open
bri-prism wants to merge 2 commits into
Open
cuda: MoE prefill — fused SwiGLU epilogue and fused weighted reduction#107bri-prism wants to merge 2 commits into
bri-prism wants to merge 2 commits into
Conversation
bri-prism
force-pushed
the
perf/moe-prefill-cda
branch
from
August 3, 2026 07:14
807c154 to
75e3f84
Compare
The merged gate/up result was written to global memory and read back by a separate SwiGLU op. This writes the activated result directly from the matmul, removing a round trip through the intermediate tensor and one dispatch per layer. Output is bit-identical. Prefill improves at every sequence length measured; decode is unaffected. Disable with GGML_CUDA_MMQ_GLU_FUSION_DISABLE.
The down projection wrote one slice per selected expert, which was then scaled by the router weight and reduced by an add chain. This applies the weight and accumulates across the selected experts before the final store, removing a full-size intermediate tensor and the round trip the add chain needed. Output is bit-identical. Prefill improves; decode is unaffected. Disable with GGML_CUDA_MOE_REDUCE_DISABLE.
bri-prism
force-pushed
the
perf/moe-prefill-cda
branch
from
August 3, 2026 17:40
75e3f84 to
0678955
Compare
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
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
Two independent CUDA optimizations for the routed-expert MoE prefill path, one per commit so they can be reviewed and reverted separately.
1. SwiGLU epilogue fused into the routed MMQ kernel. The merged gate/up output was written to global memory and read back by a separate SwiGLU op. This writes the activated result directly from the matmul, removing a full round trip through the intermediate tensor and one dispatch per layer. Disable with
GGML_CUDA_MMQ_GLU_FUSION_DISABLE.2. Router weighting and expert reduction folded into one kernel. The down projection wrote one slice per selected expert, which was then scaled by the router weight and reduced. This applies the weight and accumulates across the selected experts before the final store. Adds
moe-reduce.cuandmoe-reduce.cuh. Disable withGGML_CUDA_MOE_REDUCE_DISABLE.Both changes are bit-identical to the current output.
Validation
Measured internally on Hopper. Prefill improves at every sequence length tested, with a small decode cost. Numbers are in the internal notes rather than here.
Methodology, since it is the part worth reviewing:
Correctness
test-backend-ops test -b CUDA0 -o MUL_MATand-o MUL_MAT_ID: all backends pass, zero failures, on every configuration built and timed.Two changes deliberately not included
Expert-map reuse. Sharing the routed-expert grouping map between the gate/up and down matmuls looked like one of the stronger candidates in isolation. Its marginal contribution once the SwiGLU fusion lands is zero. The fused GLU path calls
ggml_cuda_launch_mm_ids_helperdirectly and so bypasses the map cache; instrumented, the cache shows full reuse on its own and zero reuse once the fusion is present. It is obsoleted rather than broken, and would be worth revisiting if the fused path went throughggml_cuda_get_mm_ids_mapinstead of launching the helper itself.Fused top-k routing extended to prefill. This was in an earlier revision of this PR and is dropped because it fails an accuracy gate. It is the only one of the three that is not bit-identical, and against a healthy reference with a same-path control the deviation was roughly an order of magnitude larger than other numerics-changing paths we have accepted, with a meaningful fraction of tokens changing their argmax. Perplexity was unchanged, so this is a behavioural difference rather than a quality regression, but the signature points at router selection changing rather than arithmetic reassociation: where expert logits are near-tied, a small floating point difference selects a different expert, and a different expert means different weights. It contributed the smallest gain of the three, so the trade was not worth it as a default.
Notes for review
clang-formatwants to reorder the whole pre-existing include block inggml-cuda.cu, triggered by the single include this branch adds. I left it alone to keep the diff reviewable and avoid conflicts with future upstream merges. The added include is placed alphabetically..cuand.cuhfiles are picked up by the existing CMake glob, so no build file changes were needed.Test plan
test-backend-opsMUL_MAT and MUL_MAT_ID, zero failures on every configuration timed.