Skip to content

Fix Attention GQA implementation on CPU #25966

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

Merged
merged 15 commits into from
Sep 15, 2025
+129 −57
Merged

Fix Attention GQA implementation on CPU #25966

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
28de5c5
Fix Attention GQA implementation on CPU
xadupre Sep 5, 2025
f407fd7
add one more unit test
xadupre Sep 5, 2025
7ff71da
negative infinity
xadupre Sep 8, 2025
6eb8ee9
fix last unittest
xadupre Sep 9, 2025
627fe07
fix merge conflicts
xadupre Sep 10, 2025
07745c4
disable onnx test for attention
xadupre Sep 10, 2025
542b5ab
disable
xadupre Sep 10, 2025
08841de
disable
xadupre Sep 10, 2025
db6df58
fix issues
xadupre Sep 10, 2025
abb196e
disable
xadupre Sep 10, 2025
805fe69
Merge branch 'main' of https://github.com/microsoft/onnxruntime into …
xadupre Sep 11, 2025
59ded0c
disable test waiting for onnx
xadupre Sep 11, 2025
5f93b47
disable
xadupre Sep 11, 2025
5c99cf0
disable
xadupre Sep 11, 2025
18b0787
inlien
xadupre Sep 12, 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
10 changes: 6 additions & 4 deletions onnxruntime/core/providers/cpu/llm/attention.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,7 +317,8 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
}

// handling GQA
std::ptrdiff_t ki = batch_i * parameters.kv_num_heads + head_i % parameters.kv_num_heads;
std::ptrdiff_t head_ki = head_i * parameters.kv_num_heads / parameters.q_num_heads;
std::ptrdiff_t ki = batch_i * parameters.kv_num_heads + head_ki;
const T* k = K + k_input_chunk_length * ki;

if (nullptr != present_key) {
Expand Down Expand Up @@ -347,7 +348,7 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
alpha,
Q + q_input_chunk_length * parameters.q_num_heads * batch_i + head_i * parameters.head_size,
parameters.head_size * parameters.q_num_heads, // lda
transposed_k ? K + k_input_chunk_length * parameters.kv_num_heads * batch_i + head_i * parameters.head_size : k,
transposed_k ? K + k_input_chunk_length * parameters.kv_num_heads * batch_i + head_ki * parameters.head_size : k,
transposed_k ? parameters.head_size * parameters.kv_num_heads : parameters.head_size, // ldb
beta,
output,
Expand Down Expand Up @@ -555,7 +556,8 @@ void AttentionBase<T>::ComputeVxAttentionScore(T* output, // bu
// handling GQA
std::ptrdiff_t batch_i = i / num_heads;
std::ptrdiff_t head_i = i % num_heads;
std::ptrdiff_t vi = batch_i * kv_num_heads + head_i % kv_num_heads;
std::ptrdiff_t head_vi = head_i * kv_num_heads / num_heads;
std::ptrdiff_t vi = batch_i * kv_num_heads + head_vi;
const T* v = V + v_input_chunk_length * vi;

if (nullptr != present_value) {
Expand Down Expand Up @@ -585,7 +587,7 @@ void AttentionBase<T>::ComputeVxAttentionScore(T* output, // bu
1.f, // alpha
attention_probs + attention_probs_offset, // QK
total_sequence_length, // lda
transposed_v ? V + head_i * v_head_size + v_input_chunk_length * kv_num_heads * batch_i : v, // V
transposed_v ? V + head_vi * v_head_size + v_input_chunk_length * kv_num_heads * batch_i : v, // V
transposed_v ? v_head_size * kv_num_heads : v_head_size, // ldb
0.f, // beta
output + ((batch_i * sequence_length * num_heads + head_i) * v_head_size),
Expand Down
Loading
Loading