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

justinchuby · 2025-10-14T17:02:10Z

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.

Copilot

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.}

onnxruntime/core/providers/cpu/llm/attention.cc

Signed-off-by: Justin Chu <[email protected]>

justinchuby · 2025-12-19T00:23:17Z

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 · 2025-12-19T19:41:18Z

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)

Signed-off-by: Justin Chu <[email protected]>

titaiwangms · 2026-01-10T00:12:20Z

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 · 2026-01-10T00:21:57Z

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 requested a review from Copilot October 14, 2025 17:02

Copilot AI reviewed Oct 14, 2025

View reviewed changes

onnxruntime/core/providers/cpu/llm/attention.cc Outdated Show resolved Hide resolved

justinchuby marked this pull request as draft October 14, 2025 18:08

gramalingam reviewed Oct 14, 2025

View reviewed changes

onnxruntime/core/providers/cpu/llm/attention.cc Outdated Show resolved Hide resolved

justinchuby marked this pull request as ready for review October 14, 2025 18:33

justinchuby requested a review from gramalingam October 14, 2025 21:40

xadupre reviewed Oct 27, 2025

View reviewed changes

onnxruntime/core/providers/cpu/llm/attention.cc Show resolved Hide resolved

justinchuby added 5 commits October 28, 2025 10:06

Allow present_key to be empty when past_key is provided in Attention

5a26381

Signed-off-by: Justin Chu <[email protected]>

Remove check

f2e5fbe

Signed-off-by: Justin Chu <[email protected]>

Updated

c26b933

Signed-off-by: Justin Chu <[email protected]>

try this

5477c2e

Signed-off-by: Justin Chu <[email protected]>

past_value

a648379

Signed-off-by: Justin Chu <[email protected]>

justinchuby force-pushed the justinchu/attn-keys branch from f03d6d1 to a648379 Compare October 28, 2025 17:06

titaiwangms self-requested a review October 28, 2025 23:05

justinchuby added 8 commits November 5, 2025 13:52

tests

54ceee6

Signed-off-by: Justin Chu <[email protected]>

ttest

b20bb20

Signed-off-by: Justin Chu <[email protected]>

test

346ded1

Signed-off-by: Justin Chu <[email protected]>

format

cc75dbb

Signed-off-by: Justin Chu <[email protected]>

update

6ed81f1

Signed-off-by: Justin Chu <[email protected]>

Merge branch 'main' into justinchu/attn-keys

12d236b

Update tests

4a20f6f

Signed-off-by: Justin Chu <[email protected]>

update tests

6ce22ca

Signed-off-by: Justin Chu <[email protected]>

justinchuby added this to the 1.24.0 milestone Jan 7, 2026

justinchuby added 2 commits January 9, 2026 13:06

Merge branch 'main' into justinchu/attn-keys

e0ad50a

test

46f44e0

Signed-off-by: Justin Chu <[email protected]>

justinchuby force-pushed the justinchu/attn-keys branch from d859a63 to 46f44e0 Compare January 9, 2026 22:15

test

472fea4

Signed-off-by: Justin Chu <[email protected]>

justinchuby enabled auto-merge (squash) January 10, 2026 00:05

justinchuby added the release:1.24.0 label Jan 10, 2026

titaiwangms approved these changes Jan 10, 2026

View reviewed changes

justinchuby merged commit 1d149f5 into main Jan 12, 2026
91 of 104 checks passed

justinchuby deleted the justinchu/attn-keys branch January 12, 2026 08:17

justinchuby removed the release:1.24.0 label Jan 12, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

justinchuby commented Oct 14, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Dec 19, 2025 •

edited

Loading

Uh oh!

justinchuby commented Dec 19, 2025

Uh oh!

titaiwangms commented Jan 10, 2026

Uh oh!

justinchuby commented Jan 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

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

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

Conversation

justinchuby commented Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Dec 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

justinchuby commented Dec 19, 2025

Uh oh!

titaiwangms commented Jan 10, 2026

Uh oh!

justinchuby commented Jan 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

justinchuby commented Oct 14, 2025 •

edited

Loading

justinchuby commented Dec 19, 2025 •

edited

Loading