Skip to content

Reduce allocation overhead in quantized sdpa #15610

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.

Sign up for GitHub

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

Open
wants to merge 3 commits into
base: gh/kimishpatel/202/base
Choose a base branch
from
Open

Reduce allocation overhead in quantized sdpa #15610

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ae61ab4
Reduce allocation overhead in quantized sdpa
kimishpatel Nov 5, 2025
99902b8
Update on "Reduce allocation overhead in quantized sdpa"
kimishpatel Nov 6, 2025
72292a7
Update on "Reduce allocation overhead in quantized sdpa"
kimishpatel Nov 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 23 additions & 7 deletions extension/llm/custom_ops/op_sdpa_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
std::vector<float> dequantized_v_data(v_data.m * v_data.n);
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),
dequantized_v_data.data(),
buf_qdq_ptr,
-128,
127,
1,
Expand All @@ -237,7 +237,7 @@ void dequant_and_gemm(
m,
k,
static_cast<float>(1),
dequantized_v_data.data(),
buf_qdq_ptr,
v_data.n,
qk_data,
qk_stride_m,
Expand All @@ -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) {
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) {
Expand All @@ -273,7 +274,8 @@ void _qk_at_v_gemm(
v_stride_n,
o_data,
o_stride_m,
beta);
beta,
buf_qdq_ptr);
} else {
// For smaller batch sizes, use quantized gemm
int a_stride_m_tmp, b_stride_n_tmp;
Expand Down Expand Up @@ -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,
// 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));
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>();
Expand All @@ -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;
Expand Down Expand Up @@ -1053,7 +1068,8 @@ void cpu_flash_attention(
vStrideN,
dst_data,
headSize,
n == 0 ? static_cast<accum_t>(0) : static_cast<accum_t>(1));
n == 0 ? static_cast<accum_t>(0) : static_cast<accum_t>(1),
buf_qdq_ptr);
}
// dst <- dst / sum[row]
// reorder MHA output with strides
Expand Down
Loading