Tensor-parallel: Fix delayed AllReduce on Gemma-4 MoE#22129
Merged
Conversation
Skip forward past nodes that don't consume the current one, and allow a chain of MULs.
github-actions
Bot
added
the
ggml
Contributor
|
What's special about the MUL mode, can't it skip over more generally looking at the src pointers? |
am17an
approved these changes
Apr 20, 2026
This comment was marked as spam.
This comment was marked as spam.
Contributor
Author
@JohannesGaessler 's original code was trying to match the common pattern found in MOE models. I have tried to extend it to cover Gemma. This check can be extended to include other ops that will work with |
JohannesGaessler
approved these changes
Apr 20, 2026
Comment on lines
+1707
to
1723
| // Chain of MULs with MIRRORED src[1] | ||
| while (true) { | ||
| skip_unrelated(); | ||
| if (id + 1 >= cgraph->n_nodes) { | ||
| return idr; | ||
| } | ||
| ggml_tensor * next = cgraph->nodes[id+1]; | ||
| if (next->op == GGML_OP_MUL && next->src[0] == node && | ||
| ggml_backend_meta_get_split_state(next->src[1], false).axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { | ||
| node = next; | ||
| id++; | ||
| idr = id; | ||
| n_used = ggml_node_get_use_count(cgraph, id); | ||
| } else { | ||
| break; | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
This code is effectively a roundabout way to check condition 2 that I outlined previously. The PR is fine like this, but I'll maybe refactor and simplify this in the future when I touch this part of the code again.
JohannesGaessler
approved these changes
Apr 20, 2026
ArberSephirotheca
pushed a commit
to ArberSephirotheca/llama.cpp
that referenced
this pull request
Apr 21, 2026
* Fix delayed AllReduce on Gemma-4 MoE Skip forward past nodes that don't consume the current one, and allow a chain of MULs. * Check for all sources before skipping nodes * Address review comments
arthw
pushed a commit
to arthw/llama.cpp
that referenced
this pull request
Apr 23, 2026
* Fix delayed AllReduce on Gemma-4 MoE Skip forward past nodes that don't consume the current one, and allow a chain of MULs. * Check for all sources before skipping nodes * Address review comments
rsenthilkumar6
pushed a commit
to rsenthilkumar6/llama.cpp
that referenced
this pull request
May 1, 2026
* Fix delayed AllReduce on Gemma-4 MoE Skip forward past nodes that don't consume the current one, and allow a chain of MULs. * Check for all sources before skipping nodes * Address review comments
ljubomirj
pushed a commit
to ljubomirj/llama.cpp
that referenced
this pull request
May 6, 2026
* Fix delayed AllReduce on Gemma-4 MoE Skip forward past nodes that don't consume the current one, and allow a chain of MULs. * Check for all sources before skipping nodes * Address review comments
my-other-github-account
pushed a commit
to my-other-github-account/llama.cpp
that referenced
this pull request
May 15, 2026
* Fix delayed AllReduce on Gemma-4 MoE Skip forward past nodes that don't consume the current one, and allow a chain of MULs. * Check for all sources before skipping nodes * Address review comments
my-other-github-account
pushed a commit
to my-other-github-account/llama.cpp
that referenced
this pull request
May 15, 2026
* Fix delayed AllReduce on Gemma-4 MoE Skip forward past nodes that don't consume the current one, and allow a chain of MULs. * Check for all sources before skipping nodes * Address review comments
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
4 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.
Skip forward past nodes that don't consume the current node, and allow a chain of MULs.
When
down_exps_sis set, build_moe_ffn pulls the scale tensor in via reshape/repeat/get_rows. Topological sort places those betweenmul_mat_idand the MUL that consumes it, so the existing nodes[id+1] check never sees an ADD_ID or MUL and fails.The scale MUL is followed by a second MUL; the old code only accepted one.
Performance on 2x 5090:
Requirements