cuda: fuse MoE weighted expert reduction by anujj · Pull Request #25952 · ggml-org/llama.cpp

anujj · 2026-07-21T08:08:59Z

Overview

cuda: fuse MoE weighted expert reduction

The MoE combine tail currently writes weighted expert outputs to
global memory before reducing them. The existing top-k=8 CUDA path
generally uses two physical fused kernels. This path does weighting
and ordered expert reduction in one CUDA kernel. The main benefit is
removing that intermediate global-memory traffic; launch count goes
from two to one.
Supported graphs:

unscaled: experts * router_weights
scaled: (experts * expert_scale) * router_weights

k = 2..15 is handled by one runtime-k kernel. Matching uses ops
shapes, strides, views, add-chain topology, and memory safety. It
does not depend on the model or quantization type.
Allocator integration uses add_alloc_dep from merged PR #27301 so
experts, router weights, and optional expert scales stay live until
the fused destination completes. Memory ranges are rechecked before
the fused kernel runs. There is no scratch allocation and no
device-to-device preservation copy.
Unsupported or unsafe graphs keep the existing per-op path.
Set GGML_CUDA_MOE_WEIGHTED_REDUCTION=0 to disable the fusion.
k = 16 deliberately exercises that fallback.

Additional information

Performance — prefill t/s, default-ON vs =0, same binary, mean-vs-mean:

Hardware	Model	Quant	pp2048	pp8192
RTX 5090	Qwen3.6-35B-A3B	NVFP4	+3.6%	+3.6%
RTX 5090	Qwen3.6-35B-A3B	Q4_K_M	+4.1%	+3.9%
RTX 5090	Gemma-4-26B-A4B	NVFP4	+4.3%	+4.3%
DGX Spark (GB10)	Qwen3.6-35B-A3B	NVFP4	+5.4%	+5.4%
DGX Spark (GB10)	Gemma-4-26B-A4B	Q4_K_M	+7.1%	+6.2%

Qwen3.6-35B-A3B NVFP4 tested with Single GPU, multi-GPU, CPU MoE 10, multi-GPU + CPU MoE 10 and the largest PPL diff ON/OFF was 0.0843%

Testing (both modes, same binary): test-backend-ops 13,178/13,178 pass; CTest 42/42 pass; CPU-vs-CUDA compared; WikiText-2 PPL identical 5.8857 ± 0.03661; llama-cli output unchanged. CUDA backend only; no changes to weights or numerics.

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES - paired with codex on this

ggml-gh-bot · 2026-07-21T08:13:49Z

Hi @anujj, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.
AI-generated content: This project does not accept PRs, descriptions or commit messages that are fully or predominantly AI-generated. If you have used AI to assist you in writing code, please make sure to disclose that explicitly.

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

anujj · 2026-07-21T08:59:07Z

@gaugarg-nv @ORippler @JohannesGaessler for review

ORippler

It seems to me the PR description of previous state doesn't 100% match the current state of the CUDA backend:

We already fuse Multi-MUL|ADD in the CUDA backend, launching 2 k_bin_b_cast instead of N kernels

Taking nsight compute to those kernels, they seem to be not running at speed-of-light yet.

Please comment on what this PR does differently than #17100, where we revert this exact fusion due to PPL issues (was it memory-overlap of input/outputs?)

How I think we should proceed:

We should fold the ops that do not require GEMM-completion (i.e. everything except the reduction over the experts) into the GEMM-epilogue, similar to GEMV kernel. This does not have to happen all at once. Alternatively, one can look to optimize k_bin_bcast_kernel for the above configurations.
We should look to bring remaining contraction kernel to speed-of-light (i.e. fully coalesced loads and stores, ideally only once)

@am17an @JohannesGaessler thoughts on above?

ORippler · 2026-07-21T12:53:01Z

+    float4        first        = experts[first_row * n_embd4 + col4];
+    if constexpr (has_expert_scale) {
+        const float scale = expert_scale[first_row];
+        first.x = __fmul_rn(first.x, scale);


Why use explicit __fmul_rn? If it is to prevent compiler optimizations: Given we already enable --use_fast_math for the cuda backend, we don't care about bit-wise numerical exactness.

addressed and revalidated the model correctness though the PPL after allowing normal compiler optimization

ORippler · 2026-07-21T12:53:35Z

+// __fmul_rn/__fadd_rn ordering are identical to the templated kernels, so the result is
+// bit-identical. The hot small-k cases keep their fully-unrolled kernels.


ORippler · 2026-07-21T12:54:50Z

+
+#include <climits>
+
+template <int n_expert_used, bool has_expert_scale>


have we validated templating is necessary? A ternary expression where one multiplies with 1 should be fine (unless we actually run into math pipe throttles). May need to massage compiler as in 8a6ca53

No meaningful gain was observed with the template, so I removed it and handled the scaling based on the suggestions in the review comments.

anujj · 2026-07-21T15:42:05Z

It seems to me the PR description of previous state doesn't 100% match the current state of the CUDA backend:

We already fuse Multi-MUL|ADD in the CUDA backend, launching 2 k_bin_b_cast instead of N kernels

2. Taking nsight compute to those kernels, they seem to be not running at speed-of-light yet. 3. Please comment on what this PR does differently than #17100, where we revert this exact fusion due to PPL issues (was it memory-overlap of input/outputs?)
How I think we should proceed:

We should fold the ops that do not require GEMM-completion (i.e. everything except the reduction over the experts) into the GEMM-epilogue, similar to GEMV kernel. This does not have to happen all at once. Alternatively, one can look to optimize k_bin_bcast_kernel for the above configurations.

We should look to bring remaining contraction kernel to speed-of-light (i.e. fully coalesced loads and stores, ideally only once)

@am17an @JohannesGaessler thoughts on above?

My original number came from the focused unit test, where that graph allocation didn't trigger the generic MUL/ADD fusion, so the baseline showed 8 launches. On the real Qwen model the generic fusion does kick in, and the baseline is only 2 kernels (a fused two-MUL + a fused seven-ADD) — which my Nsight Systems trace confirms. I've updated the PR description to describe it as 2 → 1 for the production workload (not 9 → 1) and to frame the main benefit as the eliminated intermediate-tensor traffic rather than the launch-count drop.

am17an

IMO this code needs to be deslopped. Please keep only relevant comments + code minimal. I think fusing the NVFP4 scale should be separate and follow in a subsequent PR.

Please also check why #16857 was reverted, there was a PPL blowup using --n-cpu-moe or multi-GPU IIRC.

am17an · 2026-07-23T08:47:57Z

+    ggml_tensor *       dst           = cgraph->nodes[output_idx];
+    const uintptr_t     experts_begin = (uintptr_t) experts->data;
+    const uintptr_t     experts_end   = experts_begin + ggml_nbytes(experts);
+    const uintptr_t     dst_begin     = (uintptr_t) dst->data;
+    const uintptr_t     dst_end       = dst_begin + ggml_nbytes(dst);
+    // The fused op preserves overlapping weights, but cannot safely overwrite expert rows while reading them.
+    if (experts_begin < dst_end && dst_begin < experts_end) {
+        return false;
+    }


this should be check via ggml_cuda_check_fusion_memory_ranges

am17an · 2026-07-23T08:49:14Z

+// Largest router top-k we fuse. The bound comes from ggml_can_fuse_subgraph(), which matches at
+// most 31 nodes (fixed int idxs[32] with GGML_ASSERT(count < 32)). Our fused subgraph spans
+// node_count = 2*k + n_muls - 1 nodes (n_muls is 1 or 2), so 2*k + 1 <= 31 gives k <= 15. Real MoE
+// routing uses a small top-k, so this covers every current model; larger k simply falls back to
+// the per-op path. Dispatch: k <= 8 -> per-k fully-unrolled kernels, 9..15 -> one runtime-k kernel
+// (see ggml_cuda_op_moe_weighted_reduction in moe-weighted-reduction.cu).


is this true? we're fusing 32 nodes here when max_experts=15?

scaled form: 2 MUL + 15 VIEW + 14 ADD = 31 nodes
unscaled form: 1 MUL + 15 VIEW + 14 ADD = 30 nodes

anujj · 2026-08-04T20:51:04Z

IMO this code needs to be deslopped. Please keep only relevant comments + code minimal. I think fusing the NVFP4 scale should be separate and follow in a subsequent PR.

I have reworked the implementation substantially to simplify the PR.

Regarding the NVFP4 scale fusion, I just wanted to confirm that I understood your suggestion correctly. Are you proposing that (experts * expert_scale) * router_weight should be implemented in a follow-up PR?

After the rework, I think the scaled path is much simpler to review as well. However, if you would still prefer to keep it separate, I'm happy to split it into a follow-up PR.

Please also check why #16857 was reverted, there was a PPL blowup using --n-cpu-moe or multi-GPU IIRC.

I checked #16857 and ran full WikiText-2 with fusion ON/OFF for single GPU, --n-cpu-moe 10, multi-GPU, and both combined on current PR; all paired PPL blowup is not observed. Also the new implementation validates overlap with ggml_cuda_check_fusion_memory_ranges, preserves aliased weights/scales in scratch, and rejects expert/output overlap

Testing results:

Qwen3.6-35B-A3B NVFP4

Configuration	Fusion OFF PPL	Fusion ON PPL	Result
Single GPU	5.8857 ± 0.03661	5.8874 ± 0.03659	PASS
Multi-GPU	5.8099 ± 0.03598	5.8090 ± 0.03595	PASS
Single GPU + `--n-cpu-moe 10`	5.8857 ± 0.03661	5.8874 ± 0.03659	PASS
Multi-GPU + `--n-cpu-moe 10`	5.8099 ± 0.03598	5.8090 ± 0.03595	PASS

Qwen3.6-35B-A3B Q4_K_M

Configuration	Fusion OFF PPL	Fusion ON PPL	Result
Single GPU	5.7009 ± 0.03509	5.7038 ± 0.03512	PASS
Multi-GPU	5.7005 ± 0.03509	5.7021 ± 0.03509	PASS

Qwen3-30B-A3B Q4_0

Configuration	Fusion OFF PPL	Fusion ON PPL	Result
Single GPU	7.6244 ± 0.05636	7.6248 ± 0.05636	PASS
Multi-GPU	7.6240 ± 0.05636	7.6233 ± 0.05634	PASS

Based on these results, enabling the fusion does not introduce a meaningful PPL regression for any of the measured configurations.

am17an · 2026-08-05T15:07:29Z

The fusion code to me now is starting to look quite bad and it was the only solution before we had support for arbitary fused ops without maintainence burden.

So my recommendation would be add this as MOE_REDUCE as well remove everything for the top-k and call it TOPK_MOE (in a separate PR), the ops can be handled quite simply using the llm_fused_op structure. Let me know if this is a good way forward @ggerganov, simply because the fusion detection is quite a complex piece of code and is brittle has led to quite subtle allocator/lifetime bugs.

ggerganov · 2026-08-05T15:14:33Z

So my recommendation would be add this as MOE_REDUCE as well remove everything for the top-k and call it TOPK_MOE

It's not clear - do you ask if we should introduce GGML_OP_MOE_REDUCE and GGML_OP_TOPK_MOE?

am17an · 2026-08-05T15:19:45Z

I'm asking if we fuse operations like this using llm_graph_fused method going forward to avoid the fusion detection code, as now it is much easy to add ops without other backends suffering.

ggerganov · 2026-08-05T15:36:55Z

The llm_graph_fused needs to have a dedicated ggml op. Generally, we should avoid adding new ops unless really necessary, so it would be better to look for ways to do this through fusing kernels in the backend.

I'll need to look more to understand better what is the problem here and why the fusion logic is so complicated.

am17an · 2026-08-05T15:42:01Z

These MoE ops will be used much more than say dsv4 ops, so it makes sense for them to have dedicated ops. If you take topk-moe for example the fusion detection is way more than the actual kernel, which doesn't let other backends implement the op properly. It also does not play well with the graph allocator as it is not fusion aware.

anujj · 2026-08-07T10:31:51Z

Thanks for the discussion, @am17an and @ggerganov . Before I rework the PR further, could we align on the preferred direction: should the MoE expert weighting and reduction become a dedicated GGML operation, such as GGML_OP_MOE_REDUCE, or should we continue with the current backend fusion-matching approach? I’m happy to implement either direction, but it would be helpful to agree on the architecture first to avoid unnecessary rework.

am17an · 2026-08-07T10:44:31Z

I support having new ops for these MoE operations:

They are bound to have less bugs
Other backends can quickly adapt to them without the fusion logic
They can be used in both prefill and decode, as there as cases when the fusion doesn't fire due to the graph allocations
They help a wide variety of models (namely MoE moels)

ggerganov · 2026-08-07T11:05:12Z

From what I understand, one of the fusions requires 16 source tensors (2x MUL + 14x ADD). This would require to bump GGML_MAX_SRC. And who's to say that there won't be a model with even more experts in the future.

It's not great that we had to introduce DSv4-specific ops. The experience is that such ops get obsolete over time (e.g. the RWKV ops). I would rather look for ways to express them with backend fusion logic and deprecate them.

So I am not convinced that it's better to introduce dedicated ggml ops in this case. The main purpose of the fusion mechanism in the backends is to avoid the operator explosion that otherwise would occur if we add new ops for each new thing that we want to fuse. IMO it's better to think along the direction to refactor the backend fusion logic, make it easy to maintain and test.

am17an · 2026-08-07T11:15:44Z

From what I understand, one of the fusions requires 16 source tensors (2x MUL + 14x ADD).

It's only that way because llama-graph does this way, the actual operations would operate on the source tensor.

    // order the views before the adds
    for (uint32_t i = 0; i < hparams.n_expert_used; ++i) {
        cur_experts[i] = ggml_view_2d(ctx0, experts, n_embd, n_tokens, experts->nb[2], i*experts->nb[1]);

        ggml_build_forward_expand(gf, cur_experts[i]);
    }

    // aggregate experts
    // note: here we explicitly use hparams.n_expert_used instead of n_expert_used
    //       to avoid potentially a large number of add nodes during warmup
    //       ref: https://github.com/ggml-org/llama.cpp/pull/14753
    ggml_tensor * moe_out = cur_experts[0];

    for (uint32_t i = 1; i < hparams.n_expert_used; ++i) {
        moe_out = ggml_add(ctx0, moe_out, cur_experts[i]);

        ggml_build_forward_expand(gf, moe_out);
    }

    if (hparams.n_expert_used == 1) {
        // avoid returning a non-contiguous tensor
        moe_out = ggml_cont(ctx0, moe_out);
    }

The experience is that such ops get obsolete over time (e.g. the RWKV ops)

Yes, though now with llm_graph_fused it does not cost much to introduce the op IMO. The most utility of the model is within the first months of it's existence, after that all code regarding that model is obsolete. So it makes sense to have the ability to move fast when required. In this case however this OP will be around for a long time just like topk-moe.

ggerganov · 2026-08-07T11:21:26Z

Which part of the graph does the topk-moe replace?

am17an · 2026-08-07T11:24:09Z

IMO it's better to think along the direction to refactor the backend fusion logic, make it easy to maintain and test.

Regarding this point specifically, we don't have a clear idea as how to fuse ops without telling the allocator that we wish to fuse these ops. This leads to various races/clobbering issues e.g. #21566. One possible implementation I proposed was #21897 (comment) which may be sufficient for fusion.

Which part of the graph does the topk-moe replace?

in the same function it replaces from probs->logits

 // add experts selection bias - introduced in DeepSeek V3
    // leave probs unbiased as it's later used to get expert weights
    ggml_tensor * selection_probs = probs;
    if (exp_probs_b != nullptr) {
        selection_probs = ggml_add(ctx0, probs, exp_probs_b);
        cb(selection_probs, "ffn_moe_probs_biased", il);
    }

    // llama4 doesn't have exp_probs_b, and sigmoid is only used after top_k
    // see: https://github.com/meta-llama/llama-models/blob/699a02993512fb36936b1b0741e13c06790bcf98/models/llama4/moe.py#L183-L198
    if (arch == LLM_ARCH_LLAMA4) {
        selection_probs = logits;
    }

    if (arch == LLM_ARCH_GROVEMOE) {
        selection_probs = ggml_sigmoid(ctx0, logits); // [n_expert, n_tokens]
        cb(selection_probs, "ffn_moe_probs_biased", il);
    }
 ...
 
     //call early so that topk-moe can be used
    ggml_build_forward_expand(gf, weights);

ggerganov · 2026-08-07T11:30:12Z

Yes, though now with llm_graph_fused it does not cost much to introduce the op IMO.

The llm_graph_fused is not a silver bullet - it still has at least 2 problems:

Increases the initialization time when checking all the fusions - can be resolved if the models graphs start to announce what fusions they want to use
Leaves untested branches of the graph once a backend implements the fused operator. This is basically unsolvable because testing all possible branches would lead to explosion of the test cases

am17an · 2026-08-07T11:38:27Z

Yes it is not perfect, we need to a way to consolidate that check from O(graph * fused_ops) to O(graph) at some point. The second point can be solved via how we solved the FA stuff on the CPU using the ref parameter to always use the reference implementation.

ggerganov · 2026-08-07T11:38:39Z

The moe reduction op is now clear, but do you know how the top-k op signature would look like?

am17an · 2026-08-07T11:40:35Z

top-k already has a signature in the CUDA backend, I imagine it would be something similar

void ggml_cuda_op_topk_moe(ggml_backend_cuda_context &     ctx,
                           const ggml_tensor *             logits,
                           ggml_tensor *                   weights,
                           ggml_tensor *                   ids,
                           const ggml_tensor *             clamp,
                           const ggml_tensor *             scale,
                           const ggml_tensor *             bias,
                           const ggml_cuda_topk_moe_args & args);

am17an · 2026-08-07T11:41:49Z

Note that I'm not advocating for top-k to refactored immediately, but say for adding it to Metal after the refactor would be much easier than adding it now.

ggerganov · 2026-08-07T11:59:30Z

I think the topk-moe dedicated op would be tricky to define because there is a lot of branching, so whatever API we come up with, it would likely break really quickly.

Regarding the expert reduction - do I understand correctly that the CUDA backend currently does not even fuse the ADDs? That would be surprising - it's quite simple to fuse those ops.

anujj · 2026-08-12T04:36:41Z

Aman, did you get chance to review this PR ?

am17an · 2026-08-12T13:20:46Z

+    const float * weights_data      = (const float *) weights->data;
+    const float * expert_scale_data = expert_scale ? (const float *) expert_scale->data : nullptr;
+    const uintptr_t weights_begin = (uintptr_t) weights->data;
+    const uintptr_t weights_end   = weights_begin + ggml_nbytes(weights);
+    const uintptr_t dst_begin     = (uintptr_t) dst->data;
+    const uintptr_t dst_end       = dst_begin + ggml_nbytes(dst);
+    ggml_cuda_pool_alloc<float> weights_copy(ctx.pool());
+    if (weights_begin < dst_end && dst_begin < weights_end) {
+        // The graph allocator may reuse weights for dst after the original MUL. Fusion reads both at once.
+        weights_data = weights_copy.alloc(ggml_nelements(weights));
+        CUDA_CHECK(cudaMemcpyAsync((void *) weights_data, weights->data, ggml_nbytes(weights),
+                                   cudaMemcpyDeviceToDevice, stream));
+    }
+
+    ggml_cuda_pool_alloc<float> expert_scale_copy(ctx.pool());
+    if (expert_scale != nullptr) {
+        const uintptr_t scale_begin = (uintptr_t) expert_scale->data;
+        const uintptr_t scale_end   = scale_begin + ggml_nbytes(expert_scale);
+        if (scale_begin < dst_end && dst_begin < scale_end) {
+            expert_scale_data = expert_scale_copy.alloc(ggml_nelements(expert_scale));
+            CUDA_CHECK(cudaMemcpyAsync((void *) expert_scale_data, expert_scale->data, ggml_nbytes(expert_scale),
+                                       cudaMemcpyDeviceToDevice, stream));
+        }
+    }


This is the bad pattern I was talking about earlier. To enable fusion in all cases we basically have to act like a graph allocator ourselves. So for now let's just not enable the fusion if there are overlapping ranges. If after this change basically no fusions happen then we should come back to this once we have a better idea to how to do fusion

Aman, I tested the strict overlap behavior you suggested on Qwen3.6-35B-A3B NVFP4, Q4_K_M, and and one other model. Without preserving overlapping inputs in scratch, all reduction-fusion sites were rejected for for all 3 models. Therefore, the current end-of-chain fusion is not practically useful without the scratch-copy behavior you objected to.

Would a cleaner direction be to fuse the immediate router-weight MUL into the preceding MUL_MAT_ID GEMM epilogue, then use the CUDA backend’s existing multi-input ADD fusion for the expert reduction?

We could initially support the simple MUL_MAT_ID -> router MUL form and consider scale_2/NVFP4 separately afterward. This would provide:

A small, fixed two-node matcher instead of matching the complete VIEW/ADD chain.

Reuse of the existing ADD-reduction fusion.

No new reduction kernel.

No fusion-specific scratch allocation or input copying.

No need for the CUDA backend to compensate for graph-allocation decisions.

Strict overlap rejection with fallback to the existing path.

A smaller and less brittle implementation that is easier to test and maintain.

Would this direction address your concerns with the current design?

Would a cleaner direction be to fuse the immediate router-weight MUL into the preceding MUL_MAT_ID GEMM epilogue, then use the CUDA backend’s existing multi-input ADD fusion for the expert reduction?

Yes - at least in the meaning that it would not require transient memory-allocations

This is the bad pattern I was talking about earlier.

I don't share this sentiment personally, so long as transient workspace memory-allocation is contained within a single "op" in the CUDA backend. ggml_cuda_pool_alloc exists for this exact reason (i.e. to allocate transient/workspace memory), and having workspace memory is common practice for GPU programming.

The problem is that it doesn't play well with --fit, also vmm is broken on HIP, we want to move away from pool allocations. That sentiment is shared by both me and @JohannesGaessler. I personally want fusion to be allocator aware so that it works across all batch sizes

The problem is that it doesn't play well with --fit, also vmm is broken on HIP, we want to move away from pool allocations. That sentiment is shared by both me and @JohannesGaessler.

Hmm. I think having some pre-sized/constant workspace allocation that can then work with --fit would be suitable middle ground (similar to how the memory claimed by CudaRT has to be considered) - We want to keep CudaX libraries (or their cross-IHV components) in there for performance, and as said having workspace memory is a common thing to do in GPU programming (see CUB/CUBLAS).

ORippler · 2026-08-14T15:54:32Z

In theory, there should be a unified fusion-detection mechanism that is used by all backends, instead of implementing adhoc logic individually.

Agree that we should look to formalize this long term. Combined with more static graphs on llama-side, this would enable graph-compilers to go beyond static-pattern-matching we currently do.

am17an · 2026-08-18T06:31:42Z

@anujj can you try #27301? The idea is to signal to the allocator to not re-use the intermediate tensors

anujj · 2026-08-18T09:59:38Z

Thanks for putting this up @am17an , i will try and update

anujj · 2026-08-18T17:35:35Z

@am17an , I tested #27301 locally with this PR, and it solves the lifetime issue cleanly. I registered experts, weights, and the optional expert_scale as allocation dependencies.

am17an · 2026-08-18T17:36:48Z

Great, we can wait for that to be merged and then rebase this on top

anskumar01 · 2026-08-31T16:55:21Z

#27301 is merged now. @anujj, is this PR updated based on that?
Cc @am17an for final review / approval

anujj · 2026-08-31T16:56:46Z

i am working on it, i will update the final patch tomorrow

The MoE combine tail currently writes weighted expert outputs to global memory before reducing them. That intermediate global-memory traffic is the main cost. The production baseline generally runs two physical fused kernels; this path runs one. This change matches the full expert-weighting plus ordered-reduction subgraph and replaces it with one weighted-reduction kernel. Supported graphs: - unscaled: experts * router_weights - scaled: (experts * expert_scale) * router_weights k = 2..15 is handled by one runtime-k kernel. Matching is structural: op sequence, shapes, strides, expert views, and the left-to-right ADD chain. The fused kernel keeps that same reduction order. Results are not claimed bit-identical; CUDA FP32 contraction can change rounding slightly. Allocator integration uses add_alloc_dep from the graph-optimizer API so experts, router weights, and optional expert scales stay live until the fused destination is written. Memory ranges are rechecked before the fused kernel runs. Unrecognized or unsafe graphs are left alone and keep the existing per-op path. Set GGML_CUDA_MOE_WEIGHTED_REDUCTION=0 to disable the fusion. test-backend-ops covers scaled/unscaled, aligned/unaligned, and representative values across k=2..15, plus a k=16 case that must stay on the per-op path.

anujj · 2026-09-01T06:30:10Z

Rebased onto #27301. Experts, router weights, and the optional expert_scale are now registered as allocation dependencies, and the scratch buffer allocation is gone.
Updated the PR description to match the current implementation, including the multi-GPU and CPU MoE results in the description

Final Spark numbers with Qwen3.6-35B-A3B NVFP4:

ISL	Fusion OFF	Fusion ON	Gain
1K	2861.09 tok/s	3048.40 tok/s	+6.55%
2K	2921.74 tok/s	3136.92 tok/s	+7.36%
4K	2907.73 tok/s	3105.44 tok/s	+6.80%
8K	2872.56 tok/s	3071.02 tok/s	+6.91%

@am17an : could you help to revie it , please

ORippler

I still think this PR needs clean-up. If we cannot do it by today EoD, we should commit to a follow-up PR similar to how we did it for #25635

ORippler · 2026-09-01T10:21:49Z

 }

-// returns whether the write (out) nodes overwrite the read nodes in operation
+// Returns whether the write (out) nodes overwrite the read nodes in operation.


unnecessary change

ORippler · 2026-09-01T11:09:15Z

+            params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.experts), match.dst);
+            params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.weights), match.dst);
+            if (match.expert_scale != nullptr) {
+                params->add_alloc_dep(
+                    params->user_data, const_cast<ggml_tensor *>(match.expert_scale), match.dst);
+            }


This path is not tested for at the moment. Beyond the scope of this PR, but I feel we should have test-backend-ops reflect the graph_optimize -> graph_compute flow

ORippler · 2026-09-01T11:10:36Z

+        return;
+    }
+
+    const int64_t token     = index / n_embd;


if n_embed and index fit in uint_32t, fastdiv should be used (didn't check this myself)

ORippler · 2026-09-01T11:12:40Z

+// The long form spans 2*k + 1 nodes. ggml_can_fuse_subgraph() accepts at most
+// 31 nodes, so k <= 15; larger values use the per-operation path.
+static constexpr int MOE_WEIGHTED_REDUCTION_MAX_EXPERTS = 15;


Feels like we should relax that artificial constraint to 32 nodes on ggml side in a follow-up PR

ORippler · 2026-09-01T11:15:24Z

+    auto is_weights = [](const ggml_tensor * tensor, const ggml_tensor * full) {
+        return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) &&
+            tensor->ne[0] == 1 && tensor->ne[1] == full->ne[1] &&
+            tensor->ne[2] == full->ne[2] && tensor->ne[3] == full->ne[3];
+    };
+    auto is_experts = [](const ggml_tensor * tensor, const ggml_tensor * full) {
+        return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) &&
+            ggml_are_same_shape(tensor, full);
+    };


Given they are used only in split_mul lambda, we can define them in there

ORippler · 2026-09-01T11:28:45Z

+static bool ggml_cuda_fusion_disabled() {
+    static const bool disabled = [] {
+        const char * env = getenv("GGML_CUDA_DISABLE_FUSION");
+        return env != nullptr && atoi(env) != 0;
+    }();
+    return disabled;
+}


Unrelated refactor that could be reverted

ORippler · 2026-09-01T11:29:03Z

+static bool ggml_cuda_use_moe_weighted_reduction() {
+    static const bool enabled = [] {
+        // Enabled by default. Unrecognized graphs use the per-operation path.
+        // Set GGML_CUDA_MOE_WEIGHTED_REDUCTION=0 to disable the fusion.
+        const char * env = getenv("GGML_CUDA_MOE_WEIGHTED_REDUCTION");
+        return env == nullptr || atoi(env) != 0;
+    }();
+    return enabled;
+}


Please remove, there is no need for a separate toggle

ORippler · 2026-09-01T11:33:50Z

+    int                 node_count   = 0;
+};
+
+static bool ggml_cuda_match_moe_weighted_reduction(


Please move the run-time memory-check out of this function (i.e. ggml_cuda_match_moe_weighted_reduction is used by both graph_optimize and try_fuse to match the fusion pattern, but try_fuse afterwards does the check via ggml_cuda_check_fusion_memory_ranges

ORippler · 2026-09-01T11:46:29Z

+        for (int64_t expert = 0; expert < n_expert_used; ++expert) {
+            views[expert] = ggml_view_2d(
+                ctx, weighted, n_embd, n_tokens, weighted->nb[2], expert * weighted->nb[1]);
+        }
+
+        ggml_tensor * out = views[0];
+        for (int64_t expert = 1; expert < n_expert_used; ++expert) {
+            out = ggml_add(ctx, out, views[expert]);
+        }


If we support interleaving view and adds in the cuda backend, we should also test for it

am17an

pass gridDim.x as the token and gridDim.y as the ceil_div(n_embd / n_threads). Also we should clean-up before merging this

am17an · 2026-09-01T15:16:02Z

+                                                  int64_t n_embd,
+                                                  int64_t n_tokens,
+                                                  int     n_expert_used) {


should be const

anujj · 2026-09-01T16:53:52Z

@ORippler and @am17an , thanks for the review comments, i have addressed the comments

ORippler

@am17an @ggerganov I feel we should add ggml_backend_graph_optimize(backend, gf) to test_case::eval and test_case::eval_perf in test-backend-ops to ensure we run the fusion code. Can be done in a follow-up PR though, if @anujj checked this PR for correctness (PPL value-diff) nvm we don't use the realloaction beahvior in test-backend-ops

ORippler · 2026-09-01T19:28:07Z

-
    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;

+    static const bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION"));


Let's extract this to a shared helper in a follow-up pr (used here and in ggml_cuda_try_fuse)

ORippler · 2026-09-01T19:48:41Z

(base) osimons@ub-osimons:~/anuj.cpp$ ./test_ops/bin/test-backend-ops -o MOE_WEIGHTED_REDUCTION
ggml_cuda_init: found 1 CUDA devices (Total VRAM: 32165 MiB):
  Device 0: NVIDIA RTX PRO 4500 Blackwell, compute capability 12.0, VMM: yes, VRAM: 32165 MiB
Testing 2 devices

Backend 1/2: CUDA0
  Device description: NVIDIA RTX PRO 4500 Blackwell
  Device memory: 32165 MB (31868 MB free)

  MOE_WEIGHTED_REDUCTION(n_embd=63,n_expert_used=2,n_tokens=17,unaligned_experts=0,with_expert_scale=0,interleaved_views_adds=0): OK
  MOE_WEIGHTED_REDUCTION(n_embd=2048,n_expert_used=8,n_tokens=128,unaligned_experts=0,with_expert_scale=0,interleaved_views_adds=0): OK
  MOE_WEIGHTED_REDUCTION(n_embd=2048,n_expert_used=8,n_tokens=128,unaligned_experts=0,with_expert_scale=1,interleaved_views_adds=0): OK
  MOE_WEIGHTED_REDUCTION(n_embd=63,n_expert_used=12,n_tokens=33,unaligned_experts=1,with_expert_scale=1,interleaved_views_adds=1): OK
  MOE_WEIGHTED_REDUCTION(n_embd=2048,n_expert_used=15,n_tokens=40,unaligned_experts=0,with_expert_scale=1,interleaved_views_adds=0): OK
  MOE_WEIGHTED_REDUCTION(n_embd=2048,n_expert_used=16,n_tokens=32,unaligned_experts=0,with_expert_scale=1,interleaved_views_adds=0): OK

test-suite passes locally, merging

evilJazz · 2026-09-04T19:35:13Z

For everyone that arrived at this exact PR / merged commit because of bisecting the significant performance regression in your MI50 / gfx906 or other AMD hardware:

Make sure to set the environment variable GGML_CUDA_DISABLE_FUSION=1 to regain performance. The variable is no longer called GGML_CUDA_MOE_WEIGHTED_REDUCTION as stated in the PR description.

Has anyone in this ticket actually tested on AMD hardware?

unsloth/GLM-4.7-GGUF:UD-IQ2_M without fusion: ~325/9 (pp8192/tg1024)
unsloth/GLM-4.7-GGUF:UD-IQ2_M with fusion: ~93/9 (pp8192/tg1024)
(this is on 9x MI50 layer split mode)

* cuda : fuse MoE weighted reduction (mul + view + add) The MoE combine tail currently writes weighted expert outputs to global memory before reducing them. That intermediate global-memory traffic is the main cost. The production baseline generally runs two physical fused kernels; this path runs one. This change matches the full expert-weighting plus ordered-reduction subgraph and replaces it with one weighted-reduction kernel. Supported graphs: - unscaled: experts * router_weights - scaled: (experts * expert_scale) * router_weights k = 2..15 is handled by one runtime-k kernel. Matching is structural: op sequence, shapes, strides, expert views, and the left-to-right ADD chain. The fused kernel keeps that same reduction order. Results are not claimed bit-identical; CUDA FP32 contraction can change rounding slightly. Allocator integration uses add_alloc_dep from the graph-optimizer API so experts, router weights, and optional expert scales stay live until the fused destination is written. Memory ranges are rechecked before the fused kernel runs. Unrecognized or unsafe graphs are left alone and keep the existing per-op path. Set GGML_CUDA_MOE_WEIGHTED_REDUCTION=0 to disable the fusion. test-backend-ops covers scaled/unscaled, aligned/unaligned, and representative values across k=2..15, plus a k=16 case that must stay on the per-op path. * Pruned the test matrix from 15 to 6 * Addressed the aman and olivers review comments

anujj requested review from a team and ggerganov as code owners July 21, 2026 08:09

github-actions Bot added testing Everything test related ggml changes relating to the ggml tensor library for machine learning CUDA Related to the CUDA backend labels Jul 21, 2026

ORippler requested changes Jul 21, 2026

View reviewed changes

am17an reviewed Jul 23, 2026

View reviewed changes

am17an mentioned this pull request Jul 27, 2026

cuda: fuse MoE expert weighting and reduction (+3% to +9% prefill, MoE models only) #26160

Closed

ggerganov self-assigned this Aug 6, 2026

am17an reviewed Aug 12, 2026

View reviewed changes

am17an mentioned this pull request Aug 18, 2026

ggml: allow passing alloc dependencies in graph_optimize #27301

Merged

anujj force-pushed the cuda-moe-weighted-reduction-upstream branch from 286a7a8 to 2e14285 Compare August 30, 2026 12:55

anujj force-pushed the cuda-moe-weighted-reduction-upstream branch from 2e14285 to eb24eb8 Compare August 31, 2026 17:55

Pruned the test matrix from 15 to 6

7274321

ORippler reviewed Sep 1, 2026

View reviewed changes

am17an requested changes Sep 1, 2026

View reviewed changes

Addressed the aman and olivers review comments

7a62274

am17an approved these changes Sep 1, 2026

View reviewed changes

ORippler approved these changes Sep 1, 2026

View reviewed changes

ORippler merged commit 3466812 into ggml-org:master Sep 1, 2026
27 of 30 checks passed

This was referenced Sep 3, 2026

cuda: shared-quantize cache, residual and elementwise-chain fusions (generic half of #338) TheTom/llama-cpp-turboquant#343

Merged

cuda: TurboQuant TQ4_1S decode optimisations TheTom/llama-cpp-turboquant#338

Closed

Patt92 mentioned this pull request Sep 3, 2026

Eval bug: CUDA/HIP: MoE models produce garbage output on RDNA3.5 (gfx1151) since #27621 #28113

Closed

ggerganov mentioned this pull request Sep 4, 2026

ggml : bump version to 0.23.0 ggml-org/ggml#1618

Merged

BrewTestBot mentioned this pull request Sep 4, 2026

ggml 0.23.0 Homebrew/homebrew-core#302393

Merged

1 task

		// __fmul_rn/__fadd_rn ordering are identical to the templated kernels, so the result is
		// bit-identical. The hot small-k cases keep their fully-unrolled kernels.


		#include <climits>

		template <int n_expert_used, bool has_expert_scale>


		ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;

		static const bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION"));

Uh oh!

Conversation

anujj commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Additional information

Requirements

Uh oh!

ggml-gh-bot Bot commented Jul 21, 2026

Uh oh!

anujj commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ORippler left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

anujj commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

am17an left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anujj commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Qwen3.6-35B-A3B NVFP4

Qwen3.6-35B-A3B Q4_K_M

Qwen3-30B-A3B Q4_0

Uh oh!

am17an commented Aug 5, 2026

Uh oh!

ggerganov commented Aug 5, 2026

Uh oh!

am17an commented Aug 5, 2026

Uh oh!

ggerganov commented Aug 5, 2026

Uh oh!

am17an commented Aug 5, 2026

Uh oh!

anujj commented Aug 7, 2026

Uh oh!

am17an commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ggerganov commented Aug 7, 2026

Uh oh!

am17an commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ggerganov commented Aug 7, 2026

Uh oh!

am17an commented Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ggerganov commented Aug 7, 2026

Uh oh!

am17an commented Aug 7, 2026

Uh oh!

ggerganov commented Aug 7, 2026

anujj commented Jul 21, 2026 •

edited

Loading

anujj commented Jul 21, 2026 •

edited

Loading

anujj commented Jul 21, 2026 •

edited

Loading

anujj commented Aug 4, 2026 •

edited

Loading

am17an commented Aug 7, 2026 •

edited

Loading

am17an commented Aug 7, 2026 •

edited

Loading

am17an commented Aug 7, 2026 •

edited

Loading

anujj commented Aug 12, 2026 •

edited

Loading

ORippler Aug 14, 2026 •

edited

Loading

ORippler Aug 14, 2026 •

edited

Loading

ORippler Aug 14, 2026 •

edited

Loading

am17an commented Aug 18, 2026 •

edited

Loading