Skip to content

Support int8 output for scaled_embedding_bag#3231

Merged
Xia-Weiwen merged 13 commits into
pytorch:mainfrom
LevelDownRefine:wengshiy/int8_scaled_embedding_bag
Nov 6, 2025
Merged

Support int8 output for scaled_embedding_bag#3231
Xia-Weiwen merged 13 commits into
pytorch:mainfrom
LevelDownRefine:wengshiy/int8_scaled_embedding_bag

Conversation

@LevelDownRefine

Copy link
Copy Markdown
Contributor

No description provided.

@pytorch-bot

pytorch-bot Bot commented Oct 23, 2025

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3231

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 1 Pending

As of commit d168675 with merge base 7d91f11 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@LevelDownRefine LevelDownRefine marked this pull request as draft October 23, 2025 04:24
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 23, 2025
@LevelDownRefine LevelDownRefine changed the title Wengshiy/int8 scaled embedding bag Support int8 output for scaled embedding bag Oct 23, 2025
@LevelDownRefine LevelDownRefine changed the title Support int8 output for scaled embedding bag Support int8 output for scaled_embedding_bag Oct 23, 2025
@LevelDownRefine

Copy link
Copy Markdown
Contributor Author

@pytorchbot label "topic: new feature"

@pytorch-bot pytorch-bot Bot added the topic: new feature Use this tag if this PR adds a new feature label Oct 24, 2025
@LevelDownRefine

Copy link
Copy Markdown
Contributor Author

@Xia-Weiwen @mingfeima Could you help review this PR?

@Xia-Weiwen Xia-Weiwen left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM. Some minor changes are needed.

Comment thread test/test_ops.py Outdated
Comment thread test/test_ops.py
Comment thread torchao/csrc/cpu/aten_kernels/scaled_embedding_bag.cpp Outdated
Comment thread torchao/csrc/cpu/aten_kernels/scaled_embedding_bag.cpp Outdated
Comment thread torchao/csrc/cpu/aten_kernels/scaled_embedding_bag.cpp Outdated
Comment thread torchao/csrc/cpu/aten_kernels/scaled_embedding_bag.cpp Outdated
Comment thread torchao/csrc/cpu/aten_kernels/scaled_embedding_bag.cpp Outdated
@Xia-Weiwen

Copy link
Copy Markdown
Collaborator

CC @mingfeima for review. Thanks.

LevelDownRefine and others added 3 commits October 24, 2025 09:52
Co-authored-by: Xia Weiwen <xia.weiwen@hotmail.com>
Co-authored-by: Xia Weiwen <xia.weiwen@hotmail.com>
@LevelDownRefine LevelDownRefine marked this pull request as ready for review October 28, 2025 02:11
@LevelDownRefine

Copy link
Copy Markdown
Contributor Author

@jerryzh168 Could you help review this pr?

@Xia-Weiwen Xia-Weiwen requested a review from jerryzh168 October 28, 2025 07:45
@LevelDownRefine

Copy link
Copy Markdown
Contributor Author

@jerryzh168 Could you help review this pr?

@Xia-Weiwen Xia-Weiwen merged commit 9008f84 into pytorch:main Nov 6, 2025
18 checks passed
namgyu-youn pushed a commit to namgyu-youn/ao that referenced this pull request Nov 21, 2025
* re-enable scaled_embedding_bag

* only support fp32 out_dtype

* support int8 output

* support int8 fallback path

* fix lint

* fix lint and clang

* refine scale

Co-authored-by: Xia Weiwen <xia.weiwen@hotmail.com>

* refine code

Co-authored-by: Xia Weiwen <xia.weiwen@hotmail.com>

* refine code

---------

Co-authored-by: Xia Weiwen <xia.weiwen@hotmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: new feature Use this tag if this PR adds a new feature

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants