Skip to content

Add int8 LUT CPU ops #2026

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.

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

Merged
merged 1 commit into from
Apr 28, 2025
Merged

Add int8 LUT CPU ops #2026

merged 1 commit into from
Apr 28, 2025

Conversation

metascroy
Copy link
Contributor

Adds op support for 1-4 bit LUT kernels (int8 multiple of FP32 scale).

Copy link

pytorch-bot bot commented Apr 7, 2025

🔗 Helpful Links

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

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

❌ 1 New Failure

As of commit 665dab3 with merge base df532f0 (image):

NEW FAILURE - The following job has failed:

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

@metascroy metascroy requested a review from digantdesai April 7, 2025 15:49
@facebook-github-bot facebook-github-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 Apr 7, 2025
@metascroy metascroy requested a review from Jack-Khuu April 7, 2025 15:49
Copy link
Contributor

@digantdesai digantdesai left a comment

Choose a reason for hiding this comment

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

Looks good, left some comments around the UK api

@@ -126,7 +126,7 @@ void pack_weights(
bias);
}

template <int weight_nbit, int nr, int kr, int sr>
template <int weight_nbit, int nr_, int kr_, int sr_>
Copy link
Contributor

Choose a reason for hiding this comment

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

Nit: suffix _ has different meaning too i.e. class private variables, what about capitalized?

Comment on lines 145 to +146
pack_weights_fn_type pack_weights{nullptr};
pack_weights_with_lut_fn_type pack_weights_with_lut{nullptr};
Copy link
Contributor

Choose a reason for hiding this comment

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

should we have a different uKernel config for LUT? or even worse a union :)

Suggested change
pack_weights_fn_type pack_weights{nullptr};
pack_weights_with_lut_fn_type pack_weights_with_lut{nullptr};
union pack_weights_fn_t {
pack_weights_fn_type tiled{nullptr};
pack_weights_with_lut_fn_type tiled_with_lut{nullptr};
}

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I prefer separate ones for now

int weight_nbit,
bool has_weight_zeros,
bool has_bias,
packed_weights_size_fn_type packed_weights_with_lut_size,
Copy link
Contributor

Choose a reason for hiding this comment

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

function overload?

weight_nbit
);
constexpr bool has_lut = true;
int preferred_alignment = 16;
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: push this down closer to the kernel?

@metascroy metascroy added the topic: not user facing Use this tag if you don't want this PR to show up in release notes label Apr 28, 2025
@metascroy metascroy merged commit 337a057 into main Apr 28, 2025
18 of 20 checks passed
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: not user facing Use this tag if you don't want this PR to show up in release notes
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants