Skip to content

GQA Memory Efficient Kernel #17920

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Tianlei Wu (tianleiwu) merged 20 commits into from
Nov 2, 2023
Merged

GQA Memory Efficient Kernel #17920

Changes from 1 commit
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
ef82d4d
squash merge
aciddelgado Oct 12, 2023
54f2526
add unit test and fix build
aciddelgado Oct 12, 2023
e2e1157
undo work in attention_impl file
aciddelgado Oct 17, 2023
415440b
reduce tests and change default behavior for past-kv is nullptr
aciddelgado Oct 19, 2023
0e84d2d
test compatibility w/ no cuda
aciddelgado Oct 19, 2023
9ba6963
exclude from amd
aciddelgado Oct 19, 2023
6573133
fix test script
aciddelgado Oct 20, 2023
a87f211
make kernels more efficient and make present output required
aciddelgado Oct 24, 2023
16bda28
Merge branch 'main' into aciddelgado/gqa_memeff_v2
aciddelgado Oct 24, 2023
08f553d
merge main and memeff changes
aciddelgado Oct 24, 2023
eb12522
address comments
aciddelgado Oct 25, 2023
6c6aead
update ContribOperators.md
aciddelgado Oct 25, 2023
db307f3
clarify input and output formats memory efficient attention
aciddelgado Oct 27, 2023
e7a50ee
max sequence length for memory efficient attention
aciddelgado Oct 27, 2023
bc1cf0a
clang and fix test file
aciddelgado Oct 27, 2023
44ca857
undo clang on unrelated files
aciddelgado Oct 27, 2023
f08495f
Merge branch 'main' into aciddelgado/gqa_memeff_v2
aciddelgado Oct 30, 2023
660c8fd
check value and key inputs
aciddelgado Oct 31, 2023
b9c4d15
key and value dont check for nullptr since they are required
aciddelgado Oct 31, 2023
fd0ecc3
lint
aciddelgado Nov 1, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 47 additions & 3 deletions onnxruntime/contrib_ops/cuda/bert/group_query_attention_helper.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,9 +48,7 @@ Status CheckInputs(const Tensor* query,
int head_size = static_cast<int>(q_hidden_size) / num_heads;

int kv_sequence_length = static_cast<int>(key_dims[1]);
int kv_hidden_size = (key_dims.size() == 3)
? static_cast<int>(key_dims[2])
: (kv_num_heads * static_cast<int>(key_dims[3]));
int kv_hidden_size = static_cast<int>(key_dims[2]);

int max_sequence_length = 0;
if (past_key != nullptr && past_value != nullptr) {
Expand Down Expand Up @@ -134,6 +132,52 @@ Status CheckInputs(const Tensor* query,
"Input 'past_key' and 'past_value' shall be both present or both absent");
}

if (key != nullptr) {
const auto& key_dims = key->Shape().GetDims();
if (key_dims.size() != 3) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'key' is expected to have 3 dimensions, got ",
key_dims.size());
}
if (query_dims[0] != key_dims[0]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'query' and 'key' shall have same dim 0 (batch size)");
}

if (num_heads % kv_num_heads != 0) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"num_heads must be a multiple of kv_num_heads. Got num_heads % kv_num_heads == ",
num_heads % kv_num_heads);
}
} else {
Comment thread
tianleiwu marked this conversation as resolved.
Outdated
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Missing key tensor.");
}

if (value != nullptr) {
const auto& value_dims = value->Shape().GetDims();
if (value_dims.size() != 3) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'value' is expected to have 3 dimensions, got ",
value_dims.size());
}

if (query_dims[0] != value_dims[0]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'query' and 'value' shall have same dim 0 (batch_size)");
}

if (static_cast<int64_t>(kv_sequence_length) != value_dims[1]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 'key' and 'value' shall have the same dim 1 (kv_sequence_length)");
}

if (value_dims[2] != kv_hidden_size) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'value' is expected to have same hidden size as key.");
}
} else {
Comment thread
tianleiwu marked this conversation as resolved.
Outdated
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Missing value tensor.");
}

Comment thread
tianleiwu marked this conversation as resolved.
// When kv-cache, we take past_seq_len as an argument... otherwise we use sequence length of past kv directly.
int32_t past_sequence_length = 0;
int present_sequence_length = kv_sequence_length;
Expand Down