Skip to content

cuda: MoE prefill — fused SwiGLU epilogue and fused weighted reduction - #107

Open
bri-prism wants to merge 2 commits into
prismfrom
perf/moe-prefill-cda
Open

cuda: MoE prefill — fused SwiGLU epilogue and fused weighted reduction#107
bri-prism wants to merge 2 commits into
prismfrom
perf/moe-prefill-cda

Conversation

@bri-prism

@bri-prism bri-prism commented Aug 3, 2026

Copy link
Copy Markdown

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.cu and moe-reduce.cuh. Disable with GGML_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:

  • Each change was measured alone and then composed, with interleaved arms and multiple rounds per arm, reporting every run rather than only means.
  • Marginal contribution was measured, not inferred. This matters: a third candidate change looked strong in isolation and turned out to contribute nothing once the SwiGLU fusion was present, because the fused path bypasses the cache that change populates. It is excluded for that reason. See below.
  • Every timing run was gated on an idle GPU, no competing CPU load, and a low load average. Host CPU load perturbs launch overhead, which is exactly what these changes reduce, so an ungated run on a shared machine produces meaningless results.
  • Correctness was re-verified on the exact binaries that produced the timings rather than carried over from development.

Correctness

  • test-backend-ops test -b CUDA0 -o MUL_MAT and -o MUL_MAT_ID: all backends pass, zero failures, on every configuration built and timed.
  • Both commits produce bit-identical output against the unpatched build.

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_helper directly 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 through ggml_cuda_get_mm_ids_map instead 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-format wants to reorder the whole pre-existing include block in ggml-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.
  • New .cu and .cuh files are picked up by the existing CMake glob, so no build file changes were needed.

Test plan

  • Builds clean, CUDA sm_90, Release.
  • test-backend-ops MUL_MAT and MUL_MAT_ID, zero failures on every configuration timed.
  • Bit-identical output verified for both commits.
  • Isolated and composed A/B with interleaved arms, gated on an idle machine.
  • Correctness re-verified on the exact binaries used for timing.

@bri-prism
bri-prism requested a review from khosravipasha August 3, 2026 06:57
@bri-prism
bri-prism force-pushed the perf/moe-prefill-cda branch from 807c154 to 75e3f84 Compare August 3, 2026 07:14
@bri-prism bri-prism changed the title cuda: MoE prefill — fused SwiGLU epilogue, fused weighted reduction, prefill top-k routing cuda: MoE prefill — fused SwiGLU epilogue and fused weighted reduction Aug 3, 2026
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
bri-prism force-pushed the perf/moe-prefill-cda branch from 75e3f84 to 0678955 Compare August 3, 2026 17:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant