Skip to content

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

Merged
justinchuby merged 16 commits into
mainfrom
justinchu/attn-keys
Jan 12, 2026
Merged

Allow present_key to be empty when past_key is provided in Attention#26303
justinchuby merged 16 commits into
mainfrom
justinchu/attn-keys

Conversation

@justinchuby

@justinchuby justinchuby commented Oct 14, 2025

Copy link
Copy Markdown
Contributor

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

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

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.

Comment thread onnxruntime/core/providers/cpu/llm/attention.cc Outdated
@justinchuby
justinchuby marked this pull request as draft October 14, 2025 18:08
Comment thread onnxruntime/core/providers/cpu/llm/attention.cc Outdated
@justinchuby
justinchuby marked this pull request as ready for review October 14, 2025 18:33
Comment thread onnxruntime/core/providers/cpu/llm/attention.cc
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@titaiwangms
titaiwangms self-requested a review October 28, 2025 23:05
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@justinchuby

justinchuby commented Dec 19, 2025

Copy link
Copy Markdown
Contributor Author

For some reason present key and present values are not the same as k and v when there is no past:

2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 1.2701495885848999, which exceeds tolerance, where
2: cur_expected[i] evaluates to 2.0174980163574219,
2: cur_actual[i] evaluates to 0.74734842777252197, and
2: tolerance evaluates to 0.00023174978559836745.
2: i:0
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 1.2701497077941895, which exceeds tolerance, where
2: cur_expected[i] evaluates to 3.0174980163574219,
2: cur_actual[i] evaluates to 1.7473483085632324, and
2: tolerance evaluates to 0.00033174979034811258.
2: i:1
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.73147225379943848, which exceeds tolerance, where
2: cur_expected[i] evaluates to 2,
2: cur_actual[i] evaluates to 2.7314722537994385, and
2: tolerance evaluates to 0.00022999999055173248.
2: i:2
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.73147225379943848, which exceeds tolerance, where
2: cur_expected[i] evaluates to 3,
2: cur_actual[i] evaluates to 3.7314722537994385, and
2: tolerance evaluates to 0.00032999998074956238.
2: i:3
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.73406648635864258, which exceeds tolerance, where
2: cur_expected[i] evaluates to 1.4550299644470215,
2: cur_actual[i] evaluates to 0.72096347808837891, and
2: tolerance evaluates to 0.00017550299526192248.
2: i:4
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.73406648635864258, which exceeds tolerance, where
2: cur_expected[i] evaluates to 2.4550299644470215,
2: cur_actual[i] evaluates to 1.7209634780883789, and
2: tolerance evaluates to 0.00027550300001166761.
2: i:5
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 1.3152304887771606, which exceeds tolerance, where
2: cur_expected[i] evaluates to 1.4400719404220581,
2: cur_actual[i] evaluates to 2.7553024291992188, and
2: tolerance evaluates to 0.00017400718934368342.
2: i:6
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 1.3152303695678711, which exceeds tolerance, where
2: cur_expected[i] evaluates to 2.4400720596313477,
2: cur_actual[i] evaluates to 3.7553024291992188, and
2: tolerance evaluates to 0.00027400720864534378.
2: i:7
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.40000000596046448, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.40000000596046448,
2: cur_actual[i] evaluates to 0.80000001192092896, and
2: tolerance evaluates to 4.9999998736893758e-05.
2: i:2
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.30000001192092896, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.60000002384185791,
2: cur_actual[i] evaluates to 0.30000001192092896, and
2: tolerance evaluates to 7.0000001869630069e-05.
2: i:3
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.40000000596046448, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.80000001192092896,
2: cur_actual[i] evaluates to 0.40000000596046448, and
2: tolerance evaluates to 9.0000001364387572e-05.
2: i:4
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.30000001192092896, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.30000001192092896,
2: cur_actual[i] evaluates to 0.60000002384185791, and
2: tolerance evaluates to 3.9999998989515007e-05.
2: i:5
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 2.5, which exceeds tolerance, where
2: cur_expected[i] evaluates to 3,
2: cur_actual[i] evaluates to 0.5, and
2: tolerance evaluates to 0.00030999997397884727.
2: i:2
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 2.5, which exceeds tolerance, where
2: cur_expected[i] evaluates to 4,
2: cur_actual[i] evaluates to 1.5, and
2: tolerance evaluates to 0.0004099999787285924.
2: i:3
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 2.5, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.5,
2: cur_actual[i] evaluates to 3, and
2: tolerance evaluates to 5.999999848427251e-05.
2: i:4
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 2.5, which exceeds tolerance, where
2: cur_expected[i] evaluates to 1.5,
2: cur_actual[i] evaluates to 4, and
2: tolerance evaluates to 0.00015999999595806003.
2: i:5
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider

@justinchuby

Copy link
Copy Markdown
Contributor Author
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.40000000596046448, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.40000000596046448,
2: cur_actual[i] evaluates to 0.80000001192092896, and
2: tolerance evaluates to 4.9999998736893758e-05.
2: i:2
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.30000001192092896, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.60000002384185791,
2: cur_actual[i] evaluates to 0.30000001192092896, and
2: tolerance evaluates to 7.0000001869630069e-05.
2: i:3
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.40000000596046448, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.80000001192092896,
2: cur_actual[i] evaluates to 0.40000000596046448, and
2: tolerance evaluates to 9.0000001364387572e-05.
2: i:4
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 0.30000001192092896, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.30000001192092896,
2: cur_actual[i] evaluates to 0.60000002384185791, and
2: tolerance evaluates to 3.9999998989515007e-05.
2: i:5
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 2.5, which exceeds tolerance, where
2: cur_expected[i] evaluates to 3,
2: cur_actual[i] evaluates to 0.5, and
2: tolerance evaluates to 0.00030999997397884727.
2: i:2
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 2.5, which exceeds tolerance, where
2: cur_expected[i] evaluates to 4,
2: cur_actual[i] evaluates to 1.5, and
2: tolerance evaluates to 0.0004099999787285924.
2: i:3
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 2.5, which exceeds tolerance, where
2: cur_expected[i] evaluates to 0.5,
2: cur_actual[i] evaluates to 3, and
2: tolerance evaluates to 5.999999848427251e-05.
2: i:4
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(432): error: The difference between cur_expected[i] and cur_actual[i] is 2.5, which exceeds tolerance, where
2: cur_expected[i] evaluates to 1.5,
2: cur_actual[i] evaluates to 4, and
2: tolerance evaluates to 0.00015999999595806003.
2: i:5
2: Google Test trace:
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\checkers.cc(610): provider type: CPUExecutionProvider
2: E:\_work\onnxruntime\onnxruntime\onnxruntime\test\unittest_util\base_tester.cc(877): registered execution providers: CPUExecutionProvider
2: 
2: [  FAILED  ] AttentionTest.AttentionNoPastWithPresentOutput (0 ms)

@justinchuby justinchuby added this to the 1.24.0 milestone Jan 7, 2026
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
@justinchuby
justinchuby enabled auto-merge (squash) January 10, 2026 00:05
@titaiwangms

Copy link
Copy Markdown
Contributor

Sorry, I wasn't following this. Curious if this is going to change any assumption we have with MultiHeadAttention? I am using MultiHeadAttention as CUDA kernel for Attention. I wonder whether we enforced this because of contrib op cuda implementation?

cc @tianleiwu

@justinchuby

Copy link
Copy Markdown
Contributor Author

Shouldn't be affecting as this is the cpu implementation. The cuda kernel can have its own constraints, but in general this is already supported as far as I understand. We just had redundant checks.

@justinchuby
justinchuby merged commit 1d149f5 into main Jan 12, 2026
91 of 104 checks passed
@justinchuby
justinchuby deleted the justinchu/attn-keys branch January 12, 2026 08:17
alex-spacemit pushed a commit to spacemit-com/onnxruntime that referenced this pull request Jan 20, 2026
…icrosoft#26303)

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.

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
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.

5 participants