[Major] Fuse bias+gemm and layernorm+quantization for more efficient ViT #254
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
ys-2020
reviewed
Jan 14, 2025
| int hidden_size = input.size(-1); | ||
| int num_tokens = input.numel() / hidden_size; | ||
| dim3 grid(num_tokens); | ||
| dim3 block(std::min(hidden_size, 1024)); | ||
| dim3 block(std::min(hidden_size/2, 1024));//Prevent thread idling when the embedding size is greater than 1024 and not an integer multiple of it. |
There was a problem hiding this comment.
Can we fix this in a more elegant way to improve utilzation?
ys-2020
reviewed
Jan 14, 2025
| @@ -41,23 +40,24 @@ __inline__ __device__ Tf compute_layernorm(Tf val, float s_mean, float s_varianc | |||
| * First pass (loop) computes the mean. | |||
| * Second computes the variance via Var[x] = E[(x - E[x])²]. | |||
| * Third pass computes and writes normed_output | |||
| * | |||
| * with USE_DIFF_OF_SQUARES set to true (may be faster but less accurate): | |||
| * For better speedup, we set USE_DIFF_OF_SQUARES to true (may be faster but less accurate): | |||
There was a problem hiding this comment.
Lets keep the original template for better flexibility.
ys-2020
reviewed
Jan 14, 2025
| #pragma unroll | ||
| for (int i = 0; i < CTA_N ; i++) | ||
| { | ||
| Bias_shared[i] = __half2float(Bias[cta_offset_n+i]); |
ys-2020
reviewed
Jan 14, 2025
| @@ -592,10 +606,10 @@ void w8a8_gemm_fuse_bias_forward_cuda(torch::Tensor _in_feats, | |||
| constexpr int CTA_M = 128; | |||
| constexpr int CTA_N = 128; | |||
| constexpr int CTA_K = 64; | |||
| constexpr int WARP_M = 128; | |||
| constexpr int WARP_M = 64; | |||
There was a problem hiding this comment.
[Issue] Why couldn't we have 128 here?
ys-2020
reviewed
Jan 14, 2025
| @@ -604,7 +618,7 @@ void w8a8_gemm_fuse_bias_forward_cuda(torch::Tensor _in_feats, | |||
| constexpr int CTA_N = 64; | |||
| constexpr int CTA_K = 64; | |||
| constexpr int WARP_M = 32; | |||
| constexpr int WARP_N = 32; | |||
| constexpr int WARP_N = 16; | |||
There was a problem hiding this comment.
[Issue] Why couldn't we have 16 here?
ys-2020
changed the title
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.
@ys-2020 I have fused bias+gemm and layernorm+quantization. These optimizations achieve 1.5-1.6x and 1.5-2.0x kernel speedup respectively on RTX 4090, which lead to about 1.15x e2e speedup for ViT. I have also made some optimizations on these for Orin, but there is still room for improvement. I will complete it later.