Skip to content

cuda: opt-in MoE epilogue fusions (GGML_CUDA_MOE_EP_FUSION) - #423

Merged
bri-prism merged 1 commit into
prismfrom
cuda/moe-epilogue-fusion
Oct 31, 2026
Merged

bri-prism merged 1 commit into
prismfrom
cuda/moe-epilogue-fusion

Conversation

@bri-prism

Copy link
Copy Markdown
Collaborator

What

  • Opt-in fusions for the MoE layer epilogue, selected by the bit mask GGML_CUDA_MOE_EP_FUSION (default 0, off):
    • bit 1: fold the per-expert scale gather (REPEAT + GET_ROWS of the down-expert scales by the selected ids) into the weighted expert reduction, which reads scale_table[ids] directly;
    • bit 2: fuse the router RMS_NORM -> SCALE -> MUL(weight) into one kernel.
  • GGML_CUDA_MOE_EP_VERIFY=1|2 compares the bit-1 fusion against the existing fused reduction or the fully unfused nodes.

Why

  • MoE decode runs these as several small serial kernels per layer.

How

  • Bit 1 reads the same float values, so the reduction is unchanged. Bit 2 follows the reduction path the unfused RMS_NORM takes (scalar, or float4 with one float4 per thread) and does not fuse rows wider than that path covers, so it stays bit-identical.

Testing

  • Greedy output is byte-identical with the switch off, the switch on and the base build, on two MoE models (Q2_0 experts and rotated PQ2_0 experts).
  • GGML_CUDA_MOE_EP_VERIFY 1 and 2: 0 mismatches in decode and in 512-token batches.
  • test-backend-ops (MUL_MAT_ID, MOE_WEIGHTED_REDUCE, RMS_NORM, TOP_K, GET_ROWS, ADD, MUL, SOFT_MAX, ARGSORT) identical to the base.
  • H200 tg128, switch 3 vs 0, alternating runs: +3.7% (rotated PQ2_0 experts) and +4.2% (Q2_0 experts). Switch off is within noise of the base.

Bit 1: fold the per-expert scale gather (REPEAT + GET_ROWS of the down-expert scales
by the selected ids) into the weighted reduction kernel, which reads scale_table[ids]
directly (same float values) for every expert slot.

Bit 2: fuse the router RMS_NORM -> SCALE -> MUL(weight) into one kernel with the same
reduction and expression order as the unfused rms_norm kernel, so the result is
bit-identical; matched with a subgraph check.

GGML_CUDA_MOE_EP_FUSION is a bit mask (1, 2, or 3 for both; default 0: off).
GGML_CUDA_MOE_EP_VERIFY=1 compares the bit-1 fusion against the existing fused
reduction, =2 against the fully unfused MUL/MUL/VIEW/ADD nodes (default 0: off).

The fused router norm kernel follows the reduction path the unfused RMS_NORM takes
(scalar, or float4 with one float4 per thread) and is not used for rows wider than
that float4 path covers, so it stays bit-identical to the unfused nodes.
@bri-prism
bri-prism merged commit 75033ed into prism Oct 31, 2026
8 of 16 checks passed
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