feat(gpu): optimize packing keyswitch on gpu #1908
Open
+44
−283
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Performance improvement on H100s is also quite impressive:
|
pdroalves
requested changes
Jan 6, 2025
| @@ -31,8 +31,7 @@ __host__ inline bool can_use_pks_fast_path(uint32_t lwe_dimension, | |||
| uint32_t polynomial_size, | |||
| uint32_t level_count, | |||
| uint32_t glwe_dimension) { | |||
| // TODO: Generalize to level_count > 1 by transposing the KSK | |||
| return level_count == 1; | |||
| return true; | |||
There was a problem hiding this comment.
I think now we can get rid of this function, now?
|
|
||
| Torus a_i = lwe_in[read_val_idx]; | ||
|
|
||
| Torus state = init_decomposer_state(a_i, base_log, level_count); | ||
|
|
||
| Torus mod_b_mask = (1ll << base_log) - 1ll; | ||
| lwe_out[write_val_idx] = decompose_one<Torus>(state, mod_b_mask, base_log); | ||
| __syncthreads(); |
There was a problem hiding this comment.
Can you replace __syncthreads() by synchronize_threads_in_block()? This is obviously not wrong but we are trying to encapsulate some CUDA intrinsics.
|
|
||
| Torus a_i = lwe_in[read_val_idx]; | ||
|
|
||
| Torus state = init_decomposer_state(a_i, base_log, level_count); | ||
|
|
||
| Torus mod_b_mask = (1ll << base_log) - 1ll; | ||
| lwe_out[write_val_idx] = decompose_one<Torus>(state, mod_b_mask, base_log); | ||
| __syncthreads(); | ||
| lwe_out[write_state_idx] = state; | ||
| __syncthreads(); |
|
|
||
| Torus mod_b_mask = (1ll << base_log) - 1ll; | ||
|
|
||
| buffer_in[val_idx] = decompose_one<Torus>(state, mod_b_mask, base_log); | ||
| __syncthreads(); | ||
| buffer_in[state_idx] = state; | ||
| __syncthreads(); |
andrei-stoian-zama
force-pushed
the
feat/as_optimize_pks_all_levels
branch
from
9589a09 to
e120ed8
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
GPU optimization for packing keyswitch for any level count using GEMM.
Bench using
make bench_integer_compression_gpuon GTX 4060 mobile