Repository navigation
[DO NOT MERGE] Copy of #2739 - #2767
Draft
Xiaoyu Z (xiaoyu-work) wants to merge 3 commits into
Draft
Xiaoyu Z (xiaoyu-work) wants to merge 3 commits into
Xiaoyu Z (xiaoyu-work) wants to merge 3 commits into
Conversation
Static-unroll MoE blocks keep all expert weights in one [E, K, N] initializer and select one expert per MatMul with Gather(W_all, k, axis=0), optionally followed by Squeeze(axes=[0]). OnnxKQuantQuantization only quantizes MatMuls whose weight is an initializer, so every per-expert MatMul stayed in full precision. The new surgeon replaces gathered weights that use a constant index with 2-D initializers holding W_all[k], so a following weight-only quantizer picks them up. Gather/Squeeze nodes and stacked initializers left without consumers are removed. MatMuls with a runtime expert index are unchanged. Fixes #2489
Refactor conditions to improve clarity and handle graph inputs. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
…dMatMulWeights - Pick new initializer names from the whole value namespace (initializers, graph inputs, graph outputs and node outputs) instead of only initializers, so an existing value named like W_all_0 can't collide with the new weight. - Skip stacked weights, expert indices and Squeeze axes that are graph inputs. Those are overridable defaults, so materializing a slice would freeze them and inference-time overrides would stop having an effect.
This branch has not been deployed
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.
Describe your changes
Independent copy of #2739: Add UnstackGatheredMatMulWeights graph surgeon for static-unroll MoE.
Copies
Iyanuoluwa007/Olive:unstack-gathered-matmul-weightsto the same branch name inmicrosoft/Olive, preserving the original commits. The original PR remains unchanged.Validation: the copied branch head matches the source PR head. Tests and
lintrunner -awere not run for this copy-only operation; the existing code has not been revalidated.Checklist before requesting a review
lintrunner -a(Optional) Issue link
Source PR: #2739.