Skip to content

Fix CPU MoE input corruption from in-place buffer reuse - #33885

Open
Tianlei Wu (tianleiwu) with Copilot wants to merge 4 commits into
mainfrom
copilot/fix-moe-wrong-results-cpu
Open

Tianlei Wu (tianleiwu) with Copilot wants to merge 4 commits into
mainfrom
copilot/fix-moe-wrong-results-cpu

Conversation

Copilot AI commented Nov 4, 2026 •

Copy link
Copy Markdown
Contributor

Description

  • Buffer preservation: Copy tokens when input/output buffers alias, independently of routing-weight normalization.
  • Regression coverage: Cover graph inputs, Add and Reshape(Add) producers, intermediate/exposed MoE outputs, and both normalization modes.

Motivation and Context

CPU MoE can return incorrect results in Add → MoE → Relu when routing normalization is disabled. Its in-place allocation contract permits output to reuse token storage, but clearing that output destroys tokens before expert computation reads them.

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

Copilot AI and others added 2 commits November 4, 2026 19:57
Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>
Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>
Copilot AI changed the title [WIP] Fix MoE output results on CPU with graph inputs Fix CPU MoE input corruption from in-place buffer reuse Nov 4, 2026
@tianleiwu
Tianlei Wu (tianleiwu) marked this pull request as ready for review November 6, 2026 17:37
Copilot AI balanced review requested due to automatic review settings November 6, 2026 17:37

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Approval recommended

No active review issues were identified for the CPU MoE buffer-aliasing fix.

No open issues 🎉

What changed in this PR

Fixes CPU MoE input corruption when its output reuses token storage, addressing #33866 without tying buffer preservation to routing-weight normalization.

Changes:

  • Copies tokens before clearing the output when input and output buffers alias.
  • Adds float32 and float16 regression tests covering graph inputs, Add and Reshape(Add) producers, intermediate and exposed MoE outputs, both normalization modes, and repeated runs.
File Description
onnxruntime/​test/​contrib_ops/​moe_test.cc Adds regression coverage for CPU MoE buffer reuse across graph configurations.
onnxruntime/​contrib_ops/​cpu/​moe/​moe_cpu.cc Preserves tokens whenever input and output buffers alias.

Review effort: Balanced

Bring the SentencePiece 0.2.2 pin compatible with CI's protobuf 6.x runtime and refresh current-main validation.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CPU MoE gives wrong results when its input is computed in the graph and its output feeds another node

3 participants