Skip to content

Conversation

justinchuby
Copy link
Contributor

@justinchuby justinchuby commented Oct 14, 2025

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.

@justinchuby justinchuby requested a review from Copilot October 14, 2025 17:02
Copy link
Contributor

@Copilot Copilot AI left a 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.

@justinchuby justinchuby marked this pull request as draft October 14, 2025 18:08
// 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?

Signed-off-by: Justin Chu <[email protected]>
@justinchuby justinchuby marked this pull request as ready for review October 14, 2025 18:33
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Signed-off-by: Justin Chu <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants