Skip to content

Allow present_key to be empty when past_key is provided in Attention #26303

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 5 commits into
base: main
Choose a base branch
from
Open

Allow present_key to be empty when past_key is provided in Attention #26303

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c06496e
Allow present_key to be empty when past_key is provided in Attention
justinchuby Oct 14, 2025
f6b7bff
Remove check
justinchuby Oct 14, 2025
348887a
Updated
justinchuby Oct 14, 2025
9df401e
try this
justinchuby Oct 14, 2025
f03d6d1
past_value
justinchuby Oct 14, 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
11 changes: 4 additions & 7 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 @@ -199,13 +199,10 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
T* output_qk, // Q*K output
ThreadPool* tp,
AllocatorPtr allocator) const {
// The case past_key != nullptr and present_key == nullptr is not supported.
// We use the fact present_key is requested to avoid any extra allocation.
// However, if present_key is not requested, we should avoid allocated more memory than needed but that mean
// allocating one buffer per thread. That's why the implementation is not done.
// The user should define a model with a present_key even if not used if past_key is not null.
ORT_ENFORCE((past_key == nullptr) == (present_key == nullptr),
"The implementation only supports past_key and present_key both null or both not null.");
if (present_key != nullptr) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There must be logic down below that concatenates past key to generate present key which will be impacted: because with this generalization you will still need to construct the concatenated present key which is used in the attention computation, but there will be no present-key-buffer to store this concatenated value; it will implicitly use the nullptr for this buffer or some other error will happen

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, the deleted comment above pretty much says the same thing

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I updated the check to match the (recovered) comment so that it does not prevent the case when past_key is null, seen with io binding.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe line 493-497 handles it?

ORT_ENFORCE(past_key != nullptr, "past_key must be provided when present_key is requested.");
}

const size_t past_chunk_length = static_cast<size_t>(parameters.past_sequence_length) * parameters.head_size; // P x H
const size_t q_input_chunk_length = static_cast<size_t>(parameters.q_sequence_length) * parameters.head_size; // S x H
const size_t k_input_chunk_length = static_cast<size_t>(parameters.kv_sequence_length) * parameters.head_size; // L x H
Expand Down
Loading