-
Notifications
You must be signed in to change notification settings - Fork 3.5k
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.
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
base: main
Are you sure you want to change the base?
Conversation
Signed-off-by: Justin Chu <[email protected]>
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.
Pull Request Overview
This PR relaxes the validation logic for past/present key tensor requirements in the Attention operator. Previously, the code enforced that past_key
and present_key
must both be null or both non-null. The new logic allows present_key
to be empty when past_key
is provided, accommodating ONNX models where dead code elimination may remove unused outputs.
Key changes:
- Replaced strict bidirectional enforcement with unidirectional validation
- Removed implementation limitation comments that are no longer applicable
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
// 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) { |
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.
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
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.
So, the deleted comment above pretty much says the same thing
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.
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.
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.
I believe line 493-497 handles it?
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
The original check enforces both the present_key and the past_key must be present. But with IO-binding there may be an issue: The past_key can be nullptr even when present_key is allocated. In reality, the kernel should just do the computation when it has the data, or when the output is requested.