saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-06-28 03:20:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Disable unit tests for decoder masked multihead attention on CC 5.2 or lower GPUs (#15114)

Browse source
This commit is contained in:
Hariharan Seshadri 2023-03-20 15:09:49 -07:00 committed by GitHub
parent 226a691e05
commit 0ace27fdf7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
onnxruntime/test/contrib_ops/decoder_masked_multihead_attention_op_test.cc
View file

@ -635,6 +635,11 @@ std::vector<MLFloat16> Softmax_QK_Transpose_V(MLFloat16* softmax_qk_transpose_ma
return output;
}
TEST(DecoderMaskedMultiheadAttentionTest, Test_fp32) {
// The kernel is only supported on CC 5.3 or higher GPUs
if (NeedSkipIfCudaArchLowerThan(530)) {
return;
}
// Vary batch size
for (int batch_size = 1; batch_size <= 5; batch_size += 2) {
// Vary kv_lengths
@ -742,6 +747,11 @@ TEST(DecoderMaskedMultiheadAttentionTest, Test_fp32) {
}
TEST(DecoderMaskedMultiheadAttentionTest, Test_fp16) {
// The kernel is only supported on CC 5.3 or higher GPUs
if (NeedSkipIfCudaArchLowerThan(530)) {
return;
}
// Vary batch size
for (int batch_size = 1; batch_size <= 5; batch_size += 2) {
// Vary kv_lengths