Skip to content

[CPU] MoE Kernel #25958

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 11 commits into from
Sep 10, 2025
Merged

[CPU] MoE Kernel #25958

merged 11 commits into from
Sep 10, 2025

Conversation

@apsonawane
Copy link
Contributor

@apsonawane apsonawane commented Sep 4, 2025 •
edited
Loading

CPU MoE Kernel

name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 1, seq_len: 16, max_diff: 2.682209014892578e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 1, seq_len: 32, max_diff: 2.980232238769531e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 2, seq_len: 16, max_diff: 2.980232238769531e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 2, seq_len: 32, max_diff: 4.172325134277344e-07
.MoE CPU kernel time: 15.721677541732786 ms
.
----------------------------------------------------------------------
Ran 5 tests in 30.217s

github-advanced-security[bot]
github-advanced-security bot found potential problems Sep 4, 2025
View reviewed changes
@apsonawane apsonawane marked this pull request as ready for review September 8, 2025 17:43
@apsonawane apsonawane requested a review from tianleiwu September 8, 2025 17:43
tianleiwu
tianleiwu reviewed Sep 9, 2025
View reviewed changes
tianleiwu
tianleiwu approved these changes Sep 10, 2025
View reviewed changes
Copy link
Contributor

@tianleiwu tianleiwu left a comment

Choose a reason for hiding this comment

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

The logic is correct, the performance can be improved later like memory allocation out of loop etc.

@apsonawane
Copy link
Contributor Author

apsonawane commented Sep 10, 2025

Will add the optimizations immediately after the QMoE kernel work is done

@apsonawane apsonawane merged commit 930e640 into main Sep 10, 2025
92 of 94 checks passed
@apsonawane apsonawane deleted the asonawane/moe_cpu branch September 10, 2025 23:54
adrianlizarraga pushed a commit that referenced this pull request Sep 24, 2025
@apsonawane @adrianlizarraga
[CPU] MoE Kernel (#25958)
18e066d 
CPU MoE Kernel
```
name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 1, seq_len: 16, max_diff: 2.682209014892578e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 1, seq_len: 32, max_diff: 2.980232238769531e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 2, seq_len: 16, max_diff: 2.980232238769531e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 2, seq_len: 32, max_diff: 4.172325134277344e-07
.MoE CPU kernel time: 15.721677541732786 ms
.
----------------------------------------------------------------------
Ran 5 tests in 30.217s
```
@adrianlizarraga adrianlizarraga mentioned this pull request Sep 24, 2025
Closed
adrianlizarraga pushed a commit that referenced this pull request Sep 26, 2025
@apsonawane @adrianlizarraga
[CPU] MoE Kernel (#25958)
55aa524 
CPU MoE Kernel
```
name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 1, seq_len: 16, max_diff: 2.682209014892578e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 1, seq_len: 32, max_diff: 2.980232238769531e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 2, seq_len: 16, max_diff: 2.980232238769531e-07
.name: SwigluMoEBlock, quant_bits: 0, dtype: FP32, batch: 2, seq_len: 32, max_diff: 4.172325134277344e-07
.MoE CPU kernel time: 15.721677541732786 ms
.
----------------------------------------------------------------------
Ran 5 tests in 30.217s
```
@adrianlizarraga adrianlizarraga mentioned this pull request Sep 26, 2025
Merged
snnn pushed a commit that referenced this pull request Sep 27, 2025
@adrianlizarraga @xieofxie
@fs-eire @apsonawane @yuslepukhin @edgchen1 @quic-tirupath @quic-ashwshan
ORT 1.23.1 cherrypick 2 (#26182)
d9b2048 
### Description
Adds the following commits to the `rel-1.23.1` branch for ORT 1.23.1:


- add session_id_ to LogEvaluationStart/Stop, LogSessionCreationStart
  - main merge date: July 31, 1:05am
  - pr: #25590
  - commit: e753643
- [build] fix WebAssembly build on macOS/arm64
  - main merge date: Aug 5, 8:07am
  - pr: #25653
  - commit: 53f152b
- [CPU] MoE Kernel (#25958)
  - main merge date: Sept 10, 4:54pm
  - pr: #25958
  - commit: 930e640
- [CPU] Block-wise QMoE kernel for CPU
  - main merge date: Sept 15, 8:32am
  - pr: #26009
  - commit: 5d17734
- [C#] Implement missing APIs
  - main merge date: Sept 24, 10:50am
  - pr: #26101
  - commit: 35dcab5
- Regenerate test model with ONNX IR < 12
  - main merge date: Sept 24, 2:50pm
  - pr: #26149
  - commit: 88f2652
- [CPU] Fix compilation errors because of unused variables
  - main merge date: Sept 25, 1:21pm
  - pr: #26147
  - commit: 42fcd71
- [EP ABI] Check if nodes specified in GetCapability() have already been
assigned
  - main merge date: Sept 26, 1:24am
  - pr: #26156
  - commit: 67d3ba0
- [QNN EP] Add dynamic option to set HTP performance mode
  - main merge date: Sept 26, 11:55am
  - pr: #26135
  - commit: 6cc40fd

---------

Co-authored-by: xieofxie <[email protected]>
Co-authored-by: hualxie <[email protected]>
Co-authored-by: Yulong Wang <[email protected]>
Co-authored-by: Akshay Sonawane <[email protected]>
Co-authored-by: Dmitri Smirnov <[email protected]>
Co-authored-by: Edward Chen <[email protected]>
Co-authored-by: quic-tirupath <[email protected]>
Co-authored-by: quic-ashwshan <[email protected]>
@snnn
Copy link
Member

snnn commented Sep 27, 2025

This PR has been cherry-picked into the rel-1.23.1 branch in PR #26182. Removing the release:1.23.1 label.

TedThemistokleous added a commit to ROCm/onnxruntime that referenced this pull request Oct 17, 2025
@TedThemistokleous @adrianlizarraga
@yuslepukhin @edgchen1 @chilo-ms @xieofxie @fs-eire @apsonawane @quic-tirupath @quic-ashwshan
V1.23.1 merge (#181)
8d330d8 
* ORT 1.23.1 cherrypick 1 [REDO] (microsoft#26140)

### Description
Cherry-pick the following PRs into the ORT 1.23.1 branch:

- Fix Attention GQA implementation on CPU
- **MANUAL MERGE**: see
microsoft#26057
  - main merge date: Sept 15, 11:33am
  - pr: microsoft#25966
  - commit: d530b29
- Address edge GetMemInfo edge cases
  - main merge date: Sept 16, 10:32am
  - pr: microsoft#26021
  - commit: d251f3a
- Implement new Python APIs
  - main merge date: Sept 17, 11:44am
  - pr: microsoft#25999
  - commit: abc63e8
- MemcpyFromHost and MemcpyToHost support for plugin EPs
- **MERGE CONFLICT** on file
onnxruntime/test/optimizer/transpose_optimizer_test.cc. Conflicts with
microsoft#25689
  - main merge date: Sept 23, 10:42am
  - pr: microsoft#26088
  - commit: 4545732
- [TRT RTX EP] Fix bug for generating the correct subgraph in
GetCapability microsoft#26132
  - main merge date: Sept 23, 8:54pm
  - pr: microsoft#26132
  - commit: 72e56e7


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

---------

Co-authored-by: Dmitri Smirnov <[email protected]>
Co-authored-by: Edward Chen <[email protected]>
Co-authored-by: Chi Lo <[email protected]>

* ORT 1.23.1 cherrypick 2 (microsoft#26182)

### Description
Adds the following commits to the `rel-1.23.1` branch for ORT 1.23.1:


- add session_id_ to LogEvaluationStart/Stop, LogSessionCreationStart
  - main merge date: July 31, 1:05am
  - pr: microsoft#25590
  - commit: e753643
- [build] fix WebAssembly build on macOS/arm64
  - main merge date: Aug 5, 8:07am
  - pr: microsoft#25653
  - commit: 53f152b
- [CPU] MoE Kernel (microsoft#25958)
  - main merge date: Sept 10, 4:54pm
  - pr: microsoft#25958
  - commit: 930e640
- [CPU] Block-wise QMoE kernel for CPU
  - main merge date: Sept 15, 8:32am
  - pr: microsoft#26009
  - commit: 5d17734
- [C#] Implement missing APIs
  - main merge date: Sept 24, 10:50am
  - pr: microsoft#26101
  - commit: 35dcab5
- Regenerate test model with ONNX IR < 12
  - main merge date: Sept 24, 2:50pm
  - pr: microsoft#26149
  - commit: 88f2652
- [CPU] Fix compilation errors because of unused variables
  - main merge date: Sept 25, 1:21pm
  - pr: microsoft#26147
  - commit: 42fcd71
- [EP ABI] Check if nodes specified in GetCapability() have already been
assigned
  - main merge date: Sept 26, 1:24am
  - pr: microsoft#26156
  - commit: 67d3ba0
- [QNN EP] Add dynamic option to set HTP performance mode
  - main merge date: Sept 26, 11:55am
  - pr: microsoft#26135
  - commit: 6cc40fd

---------

Co-authored-by: xieofxie <[email protected]>
Co-authored-by: hualxie <[email protected]>
Co-authored-by: Yulong Wang <[email protected]>
Co-authored-by: Akshay Sonawane <[email protected]>
Co-authored-by: Dmitri Smirnov <[email protected]>
Co-authored-by: Edward Chen <[email protected]>
Co-authored-by: quic-tirupath <[email protected]>
Co-authored-by: quic-ashwshan <[email protected]>

---------

Co-authored-by: Adrian Lizarraga <[email protected]>
Co-authored-by: Dmitri Smirnov <[email protected]>
Co-authored-by: Edward Chen <[email protected]>
Co-authored-by: Chi Lo <[email protected]>
Co-authored-by: xieofxie <[email protected]>
Co-authored-by: hualxie <[email protected]>
Co-authored-by: Yulong Wang <[email protected]>
Co-authored-by: Akshay Sonawane <[email protected]>
Co-authored-by: quic-tirupath <[email protected]>
Co-authored-by: quic-ashwshan <[email protected]>
@apsonawane apsonawane added the cherry-picked Cherry-picked for a cherrypicks branch label Oct 17, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tianleiwu tianleiwu tianleiwu approved these changes

Assignees

No one assigned

Labels

cherry-picked Cherry-picked for a cherrypicks branch

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@apsonawane @snnn @tianleiwu @adrianlizarraga