Skip to content

Commit 1d149f5

Browse files
justinchubyjustinchuby
authored
Allow present_key to be empty when past_key is provided in Attention (microsoft#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>
1 parent e0d5b58 commit 1d149f5

2 files changed

Lines changed: 115 additions & 4 deletions

File tree

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -204,8 +204,8 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
204204
// However, if present_key is not requested, we should avoid allocated more memory than needed but that mean
205205
// allocating one buffer per thread. That's why the implementation is not done.
206206
// The user should define a model with a present_key even if not used if past_key is not null.
207-
ORT_ENFORCE((past_key == nullptr) == (present_key == nullptr),
208-
"The implementation only supports past_key and present_key both null or both not null.");
207+
ORT_ENFORCE(!((past_key != nullptr) && (present_key == nullptr)),
208+
"The implementation does not support past_key provided and present_key being null.");
209209
const size_t past_chunk_length = static_cast<size_t>(parameters.past_sequence_length) * parameters.head_size; // P x H
210210
const size_t q_input_chunk_length = static_cast<size_t>(parameters.q_sequence_length) * parameters.head_size; // S x H
211211
const size_t k_input_chunk_length = static_cast<size_t>(parameters.kv_sequence_length) * parameters.head_size; // L x H
@@ -529,8 +529,8 @@ void AttentionBase<T>::ComputeVxAttentionScore(T* output, // bu
529529
T* present_value, // present value only (if not using present state)
530530
bool transpose_output, // whether to transpose the output (0, 2, 1, 3)
531531
ThreadPool* tp) const {
532-
ORT_ENFORCE((past_value == nullptr) == (present_value == nullptr),
533-
"The implementation only supports past_value and present_value both null or both not null.");
532+
ORT_ENFORCE(!((past_value != nullptr) && (present_value == nullptr)),
533+
"The implementation does not support past_value provided and present_value being null.");
534534
const ptrdiff_t past_chunk_length = SafeInt<ptrdiff_t>(past_sequence_length) * v_head_size; // P x H_v
535535
const ptrdiff_t v_input_chunk_length = SafeInt<ptrdiff_t>(kv_sequence_length) * v_head_size; // L x H_v
536536
const ptrdiff_t present_chunk_length = past_chunk_length + v_input_chunk_length; // T x H_v

‎onnxruntime/test/providers/cpu/llm/attention_op_test.cc‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1267,5 +1267,116 @@ TEST(AttentionTest, TestAttention4DWithPastAndPresentQkMatmulBias4DMaskCausal) {
12671267
);
12681268
}
12691269

1270+
// Test for attention when past_key and past_value are nullptr.
1271+
// This test verifies the updated logic handles the case correctly when no past state is provided.
1272+
TEST(AttentionTest, AttentionNoPastKeyValue) {
1273+
int batch_size = 1; // Q.shape[0]
1274+
int q_num_heads = 2; // Q.shape[1]
1275+
int q_sequence_length = 2; // Q.shape[2]
1276+
int head_size = 3; // Q.shape[3]
1277+
int kv_sequence_length = 2; // K.shape[2] and V.shape[2]
1278+
int kv_num_heads = 2; // K.shape[1] and V.shape[1]
1279+
int v_head_size = 3; // V.shape[3]
1280+
int past_sequence_length = 0; // No past state
1281+
1282+
std::vector<float> q = {
1283+
0.548814f, 0.715189f, 0.602763f,
1284+
0.423655f, 0.645894f, 0.437587f,
1285+
0.963663f, 0.383442f, 0.791725f,
1286+
0.568045f, 0.925597f, 0.071036f};
1287+
1288+
std::vector<float> k = {
1289+
0.186193f, 0.944372f, 0.739551f,
1290+
0.227415f, 0.254356f, 0.058029f,
1291+
0.311796f, 0.696343f, 0.377752f,
1292+
0.024679f, 0.067250f, 0.679393f};
1293+
1294+
std::vector<float> v = {
1295+
0.070870f, 0.292794f, 0.152355f,
1296+
0.131289f, 0.604118f, 0.382808f,
1297+
0.967795f, 0.546885f, 0.274824f,
1298+
0.896761f, 0.406733f, 0.552078f};
1299+
1300+
ASSERT_EQ(q.size(), batch_size * q_num_heads * q_sequence_length * head_size);
1301+
ASSERT_EQ(k.size(), batch_size * kv_num_heads * kv_sequence_length * head_size);
1302+
ASSERT_EQ(v.size(), batch_size * kv_num_heads * kv_sequence_length * v_head_size);
1303+
1304+
std::vector<float> expected_y = {
1305+
0.477184f, 0.407899f, 0.207834f,
1306+
0.521223f, 0.503569f, 0.469034f,
1307+
0.485670f, 0.410303f, 0.208993f,
1308+
0.487087f, 0.512371f, 0.461486f};
1309+
1310+
// Test with no past_key or past_value (empty vectors)
1311+
// The empty vectors will cause AddOptionalInputEdge to be called, resulting in nullptr tensors
1312+
RunTest3D(batch_size, q_num_heads, q_sequence_length, head_size, kv_sequence_length, kv_num_heads, v_head_size, past_sequence_length,
1313+
q, k, v, std::vector<float>(), std::initializer_list<bool>(), std::vector<float>(), std::vector<float>(),
1314+
-1, -1, std::numeric_limits<float>::quiet_NaN(), std::numeric_limits<float>::quiet_NaN(), -1, TensorType::kFloat,
1315+
expected_y, std::vector<float>(), std::vector<float>(), std::vector<float>(),
1316+
false, true, true // disable_cpu, disable_cuda, disable_dml
1317+
);
1318+
}
1319+
1320+
// Test for attention with present_key/present_value outputs but no past state
1321+
// This ensures the ConcatStateChunk logic works correctly when past_key/past_value are nullptr
1322+
TEST(AttentionTest, AttentionNoPastWithPresentOutput) {
1323+
int batch_size = 1; // Q.shape[0]
1324+
int q_num_heads = 2; // Q.shape[1]
1325+
int q_sequence_length = 2; // Q.shape[2]
1326+
int head_size = 2; // Q.shape[3]
1327+
int kv_sequence_length = 2; // K.shape[2] and V.shape[2]
1328+
int kv_num_heads = 2; // K.shape[1] and V.shape[1]
1329+
int v_head_size = 2; // V.shape[3]
1330+
int past_sequence_length = 0; // No past state
1331+
1332+
std::vector<float> q = {
1333+
0.5f, 0.8f,
1334+
0.3f, 0.9f,
1335+
0.7f, 0.4f,
1336+
0.6f, 0.2f};
1337+
1338+
std::vector<float> k = {
1339+
0.1f, 0.7f,
1340+
0.4f, 0.6f,
1341+
0.8f, 0.3f,
1342+
0.2f, 0.9f};
1343+
1344+
std::vector<float> v = {
1345+
1.0f, 2.0f,
1346+
3.0f, 4.0f,
1347+
0.5f, 1.5f,
1348+
2.5f, 3.5f};
1349+
1350+
// The output key/value has 4d shapes. So the data is transposed
1351+
std::vector<float> expected_present_key = {
1352+
0.1f, 0.7f,
1353+
0.8f, 0.3f,
1354+
0.4f, 0.6f,
1355+
0.2f, 0.9f};
1356+
1357+
std::vector<float> expected_present_value = {
1358+
1.0f, 2.0f,
1359+
0.5f, 1.5f,
1360+
3.0f, 4.0f,
1361+
2.5f, 3.5f};
1362+
1363+
ASSERT_EQ(q.size(), batch_size * q_num_heads * q_sequence_length * head_size);
1364+
ASSERT_EQ(k.size(), batch_size * kv_num_heads * kv_sequence_length * head_size);
1365+
ASSERT_EQ(v.size(), batch_size * kv_num_heads * kv_sequence_length * v_head_size);
1366+
1367+
std::vector<float> expected_y = {
1368+
0.747348f, 1.747348f,
1369+
2.731472f, 3.731472f,
1370+
0.720963f, 1.720963f,
1371+
2.755302f, 3.755302f};
1372+
1373+
RunTest3D(batch_size, q_num_heads, q_sequence_length, head_size, kv_sequence_length, kv_num_heads, v_head_size, past_sequence_length,
1374+
q, k, v, std::vector<float>(), std::initializer_list<bool>(), std::vector<float>(), std::vector<float>(),
1375+
-1, -1, std::numeric_limits<float>::quiet_NaN(), std::numeric_limits<float>::quiet_NaN(), -1, TensorType::kFloat,
1376+
expected_y, expected_present_key, expected_present_value, std::vector<float>(),
1377+
false, true, true // disable_cpu, disable_cuda, disable_dml
1378+
);
1379+
}
1380+
12701381
} // namespace test
12711382
} // namespace onnxruntime

0 commit comments

Comments
 (0)