Add external_key/external_value inputs to GroupQueryAttention for KV-shared layers #28242
+138
−21
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -98,6 +98,14 @@ Status PrepareQKV( | |||||||||||||||||||||||||||||||||||
| q_out = nullptr; | ||||||||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||
| // present_key/present_value are required for the CUDA path since flash attention | ||||||||||||||||||||||||||||||||||||
| // and memory-efficient attention read directly from the present KV buffers. | ||||||||||||||||||||||||||||||||||||
| // The CPU path supports optional present outputs for KV-shared layers. | ||||||||||||||||||||||||||||||||||||
| if (data.present_key == nullptr || data.present_value == nullptr) { | ||||||||||||||||||||||||||||||||||||
| return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, | ||||||||||||||||||||||||||||||||||||
| "present_key and present_value outputs are required for the CUDA GroupQueryAttention kernel."); | ||||||||||||||||||||||||||||||||||||
|
Comment on lines
+101
to
+106
There was a problem hiding this comment. The schema/docs now declare
Suggested change
|
||||||||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||
| U* k = reinterpret_cast<U*>(data.present_key); | ||||||||||||||||||||||||||||||||||||
| U* v = reinterpret_cast<U*>(data.present_value); | ||||||||||||||||||||||||||||||||||||
| int max_cache_length = parameters.seqlen_present_kv_cache; | ||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -307,5 +307,108 @@ TEST(GroupQueryAttentionTest, SeqlensKWrongLength) { | |
| {}, nullptr, &execution_providers); | ||
| } | ||
|
|
||
| // ============================================================================ | ||
| // Optional present_key/present_value output tests | ||
| // ============================================================================ | ||
|
|
||
| // Helper for tests with optional present outputs. | ||
| // When omit_present=true, present_key and present_value outputs are not connected. | ||
| static void RunGQAOptionalPresentTest( | ||
| int batch_size, | ||
| int sequence_length, | ||
| int total_seq_len, | ||
| bool omit_present, | ||
| OpTester::ExpectResult expect, | ||
| const std::string& expected_message) { | ||
| constexpr int num_heads = 2; | ||
| constexpr int kv_num_heads = 1; | ||
| constexpr int head_size = 8; | ||
| constexpr int hidden_size = num_heads * head_size; | ||
| constexpr int kv_hidden_size = kv_num_heads * head_size; | ||
|
|
||
| OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); | ||
| tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(num_heads)); | ||
| tester.AddAttribute<int64_t>("kv_num_heads", static_cast<int64_t>(kv_num_heads)); | ||
|
|
||
| std::vector<float> query_data(batch_size * sequence_length * hidden_size, 1.0f); | ||
| tester.AddInput<float>("query", {batch_size, sequence_length, hidden_size}, query_data); | ||
|
|
||
| std::vector<float> key_data(batch_size * sequence_length * kv_hidden_size, 0.5f); | ||
| tester.AddInput<float>("key", {batch_size, sequence_length, kv_hidden_size}, key_data); | ||
|
|
||
| std::vector<float> value_data(batch_size * sequence_length * kv_hidden_size, 0.5f); | ||
| tester.AddInput<float>("value", {batch_size, sequence_length, kv_hidden_size}, value_data); | ||
|
|
||
| tester.AddOptionalInputEdge<float>(); // past_key | ||
| tester.AddOptionalInputEdge<float>(); // past_value | ||
|
|
||
| std::vector<int32_t> seqlens_k_data(batch_size, static_cast<int32_t>(total_seq_len - 1)); | ||
| tester.AddInput<int32_t>("seqlens_k", {batch_size}, seqlens_k_data); | ||
| tester.AddInput<int32_t>("total_sequence_length", {1}, {static_cast<int32_t>(total_seq_len)}); | ||
|
|
||
| tester.AddOptionalInputEdge<float>(); // cos_cache | ||
| tester.AddOptionalInputEdge<float>(); // sin_cache | ||
| tester.AddOptionalInputEdge<int64_t>(); // position_ids | ||
| tester.AddOptionalInputEdge<float>(); // attention_bias | ||
| tester.AddOptionalInputEdge<float>(); // head_sink | ||
|
|
||
| // Output 0: output (always required) | ||
| tester.AddOutput<float>("output", {batch_size, sequence_length, hidden_size}, | ||
| std::vector<float>(batch_size * sequence_length * hidden_size, 0.0f)); | ||
|
|
||
| if (omit_present) { | ||
| // Omit present_key and present_value — they are optional | ||
| tester.AddOptionalOutputEdge<float>(); // present_key | ||
| tester.AddOptionalOutputEdge<float>(); // present_value | ||
| } else { | ||
| int present_seq_len = total_seq_len; | ||
| tester.AddOutput<float>("present_key", {batch_size, kv_num_heads, present_seq_len, head_size}, | ||
| std::vector<float>(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); | ||
| tester.AddOutput<float>("present_value", {batch_size, kv_num_heads, present_seq_len, head_size}, | ||
| std::vector<float>(batch_size * kv_num_heads * present_seq_len * head_size, 0.0f)); | ||
| } | ||
|
|
||
| if (expect == OpTester::ExpectResult::kExpectSuccess) { | ||
| tester.SetOutputTolerance(1e6f); | ||
| } | ||
|
Comment on lines
+345
to
+373
There was a problem hiding this comment. These tests effectively only verify the kernel doesn’t fail: the expected |
||
|
|
||
| std::vector<std::unique_ptr<IExecutionProvider>> execution_providers; | ||
| execution_providers.push_back(DefaultCpuExecutionProvider()); | ||
| tester.Run(expect, expected_message, {}, nullptr, &execution_providers); | ||
| } | ||
|
|
||
| // Baseline: GQA with present outputs connected works as before | ||
| TEST(GroupQueryAttentionTest, OptionalPresent_WithPresent) { | ||
| RunGQAOptionalPresentTest( | ||
| /*batch_size=*/1, | ||
| /*sequence_length=*/4, | ||
| /*total_seq_len=*/4, | ||
| /*omit_present=*/false, | ||
| OpTester::ExpectResult::kExpectSuccess, | ||
| ""); | ||
| } | ||
|
|
||
| // KV-shared layer scenario: present outputs omitted, attention uses K,V directly | ||
| TEST(GroupQueryAttentionTest, OptionalPresent_WithoutPresent) { | ||
| RunGQAOptionalPresentTest( | ||
| /*batch_size=*/1, | ||
| /*sequence_length=*/4, | ||
| /*total_seq_len=*/4, | ||
| /*omit_present=*/true, | ||
| OpTester::ExpectResult::kExpectSuccess, | ||
| ""); | ||
| } | ||
|
|
||
| // Batched: present outputs omitted with batch_size > 1 | ||
| TEST(GroupQueryAttentionTest, OptionalPresent_Batched) { | ||
| RunGQAOptionalPresentTest( | ||
| /*batch_size=*/2, | ||
| /*sequence_length=*/3, | ||
| /*total_seq_len=*/3, | ||
| /*omit_present=*/true, | ||
| OpTester::ExpectResult::kExpectSuccess, | ||
| ""); | ||
| } | ||
|
|
||
| } // namespace test | ||
| } // namespace onnxruntime | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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.
Docs correctly mark
present_key/present_valueas optional, but (per CUDA implementation) omitting them is not currently supported by the CUDA EP. Add a brief note here (mandatory if CUDA limitation remains) stating that some execution providers (CUDA) may require these outputs to be connected, so users aren’t surprised by provider-specific INVALID_ARGUMENT errors.