Skip to content
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

feat(gpu): speed up packing KS for levels==1 #1875

Open
wants to merge 18 commits into
base: main
Choose a base branch
from

Conversation

andrei-stoian-zama
Copy link
Contributor

@andrei-stoian-zama andrei-stoian-zama commented Dec 13, 2024

Optimize packing ks for level count == 1

  • Adds an optimized GEMM kernel
  • Adds an alternative packing keyswitch implementation using the gemm kernel
  • Implements a fast path for packing keyswitch for levels==1 using the gemm kernel
  • Adds the fast path to the integer compression GPU code and adds a test with custom parameters that use the fast path

Achieves ~4x speedup for packing keyswitch

@cla-bot cla-bot bot added the cla-signed label Dec 13, 2024
@andrei-stoian-zama andrei-stoian-zama changed the title Feat/as generalize gemm pks all params feat(gpu): speed up packing KS for levels==1 Dec 13, 2024
Copy link
Contributor

@agnesLeroy agnesLeroy left a comment

Choose a reason for hiding this comment

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

Hei @andrei-stoian-zama! It looks good to me overall, but I think it would be good that @pdroalves take a look since he was the one to implement the packing keyswitch on GPU.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants