Reduce allocation overhead in quantized sdpa #15610
There are no files selected for viewing
This file contains hidden or 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
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -213,13 +213,13 @@ void dequant_and_gemm( | |||||
| const int64_t v_stride_n, | ||||||
| float* o_data, | ||||||
| const int64_t o_stride_m, | ||||||
| const float beta, | ||||||
| float* buf_qdq_ptr) { | ||||||
| dequantize_per_channel_optimized( | ||||||
| static_cast<const int8_t*>(v_data.data), | ||||||
| static_cast<const float*>(v_data.scales), | ||||||
| static_cast<const int8_t*>(v_data.zero_points), | ||||||
| buf_qdq_ptr, | ||||||
| -128, | ||||||
| 127, | ||||||
| 1, | ||||||
|
|
@@ -237,7 +237,7 @@ void dequant_and_gemm( | |||||
| m, | ||||||
| k, | ||||||
| static_cast<float>(1), | ||||||
| buf_qdq_ptr, | ||||||
| v_data.n, | ||||||
| qk_data, | ||||||
| qk_stride_m, | ||||||
|
|
@@ -257,7 +257,8 @@ void _qk_at_v_gemm( | |||||
| const int64_t v_stride_n, | ||||||
| accum_t* o_data, | ||||||
| const int64_t o_stride_m, | ||||||
| const accum_t beta, | ||||||
| accum_t* buf_qdq_ptr) { | ||||||
| if (v_data.dtype == ScalarType::Char) { | ||||||
| if constexpr (std::is_same<accum_t, float>::value) { | ||||||
| if (m > 4) { | ||||||
|
|
@@ -273,7 +274,8 @@ void _qk_at_v_gemm( | |||||
| v_stride_n, | ||||||
| o_data, | ||||||
| o_stride_m, | ||||||
| beta, | ||||||
| buf_qdq_ptr); | ||||||
| } else { | ||||||
| // For smaller batch sizes, use quantized gemm | ||||||
| int a_stride_m_tmp, b_stride_n_tmp; | ||||||
|
|
@@ -773,6 +775,17 @@ void cpu_flash_attention( | |||||
| // at::Tensor buf_reduced = at::empty( | ||||||
| // {num_thread, qSplitSize, is_reduced_type ? kvSplitSize : 0}, | ||||||
| // query.options()); | ||||||
| int64_t size_per_thread_qdq_vec = qSplitSize * kvSplitSize * headSize; | ||||||
| // Lets align size_per_thread_qdq_vec to 64 bytes, for coalesced cache reads, | ||||||
|
There was a problem hiding this comment. The comment says "align to 64 bytes" but Either:
Suggested change
|
||||||
| // by padding with right number of per thread elements | ||||||
| constexpr int64_t kAlignment = 32; | ||||||
| size_per_thread_qdq_vec = | ||||||
| (size_per_thread_qdq_vec + kAlignment - 1) & (-(kAlignment - 1)); | ||||||
|
There was a problem hiding this comment. The alignment calculation is incorrect. The formula For aligning to a power-of-2 boundary, the correct formula is: (size_per_thread_qdq_vec + kAlignment - 1) & (-kAlignment)or equivalently: (size_per_thread_qdq_vec + kAlignment - 1) & ~(kAlignment - 1)The current code uses
Suggested change
Contributor
There was a problem hiding this comment. @copilot are you sure? Please double check again |
||||||
| int64_t size_per_thread_qdq_bytes = size_per_thread_qdq_vec * sizeof(accum_t); | ||||||
| int64_t size_qdq_bytes = size_per_thread_qdq_bytes * num_thread; | ||||||
| std::vector<char> scratch_for_quant_dequant_vec(size_qdq_bytes); | ||||||
| accum_t* scratch_for_quant_dequant = | ||||||
| reinterpret_cast<accum_t*>(scratch_for_quant_dequant_vec.data()); | ||||||
|
|
||||||
| // Data ptrs | ||||||
| const scalar_t* q_data = query.const_data_ptr<scalar_t>(); | ||||||
|
|
@@ -797,6 +810,8 @@ void cpu_flash_attention( | |||||
| scalar_t* qk_reduced_data = is_reduced_type | ||||||
| ? buf_reduced_data + ompIdx * qSplitSize * kvSplitSize | ||||||
| : nullptr; | ||||||
| accum_t* buf_qdq_ptr = | ||||||
| scratch_for_quant_dequant + ompIdx * size_per_thread_qdq_vec; | ||||||
|
|
||||||
| for (int64_t z = begin; z < end; z++) { | ||||||
| int64_t m = k * qSplitSize; | ||||||
|
|
@@ -1053,7 +1068,8 @@ void cpu_flash_attention( | |||||
| vStrideN, | ||||||
| dst_data, | ||||||
| headSize, | ||||||
| n == 0 ? static_cast<accum_t>(0) : static_cast<accum_t>(1), | ||||||
| buf_qdq_ptr); | ||||||
| } | ||||||
| // dst <- dst / sum[row] | ||||||
| // reorder MHA output with strides | ||||||
|
|
||||||
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The buffer size calculation appears to be larger than necessary. The dequantize operation needs
kvBlockSize * headSizeelements (at mostkvSplitSize * headSize), but this allocatesqSplitSize * kvSplitSize * headSize. The extraqSplitSizefactor seems unnecessary and wastes memory per thread.Consider changing to:
int64_t size_per_thread_qdq_vec = kvSplitSize * headSize;