From 05bea0d3c3ff7180c306223088916eee7aabe3ad Mon Sep 17 00:00:00 2001 From: cloudhan Date: Wed, 7 Jun 2023 15:04:27 +0800 Subject: [PATCH] Add new cases for non biased mha tests (#16097) 1. Add new test data GetSelfAttentionData_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias, also added non-biased data 2. Add new test data GetCrossAttentionData_DiffSequenceLengths_HeadSize8, also added non-biased data 3. Disabled the new tests for CUDA EP due to qkv is not correctly transposed. --- .../contrib_ops/attention_op_test_helper.cc | 239 ++++++++++++++++++ .../contrib_ops/attention_op_test_helper.h | 4 + .../multihead_attention_op_test.cc | 28 +- .../multihead_attention_op_test_data_gen.py | 62 ++++- 4 files changed, 323 insertions(+), 10 deletions(-) diff --git a/onnxruntime/test/contrib_ops/attention_op_test_helper.cc b/onnxruntime/test/contrib_ops/attention_op_test_helper.cc index 93f4995ae1..21cab01fab 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test_helper.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test_helper.cc @@ -4102,6 +4102,120 @@ void GetCrossAttentionData_DiffSequenceLengths(AttentionTestData& data) { data.is_static_kv = true; } +void GetCrossAttentionData_DiffSequenceLengths_HeadSize8(AttentionTestData& data) { + data.hidden_size = 16; + data.v_hidden_size = 16; + data.num_heads = 2; + data.batch_size = 1; + data.sequence_length = 2; + data.kv_sequence_length = 4; + data.mask_type = AttentionMaskType::MASK_NONE; + + data.skip_kernel_types = { + AttentionKernelType::AttentionKernel_TrtFlashAttention, + AttentionKernelType::AttentionKernel_TrtFusedCrossAttention, + AttentionKernelType::AttentionKernel_TrtFusedAttention, + AttentionKernelType::AttentionKernel_CutlassMemoryEfficientAttention, + }; + + data.query_data = { + 0.74714613f, -2.49789214f, -0.11628322f, 1.33038604f, 0.82568336f, 0.07685500f, 2.47562003f, 2.61135578f, + 1.55278158f, -1.85635769f, 0.36962336f, 0.87219834f, 0.69827259f, 0.95257485f, -0.77894646f, 1.46218395f, + 1.29534733f, 2.14051294f, 1.09895217f, 1.39164531f, -0.01471180f, -1.40148544f, -0.50825417f, 0.26134527f, + -0.70491123f, 0.63738143f, 2.13708138f, 0.05667466f, -0.44220763f, 0.85254443f, 2.00844359f, -1.23413038f}; + data.key_data = { + 1.70455408f, 0.07344571f, 0.18893155f, -1.48390186f, -0.86155319f, 0.10993601f, -0.29869685f, 0.73800445f, + 0.94670546f, -1.36712539f, -0.41328859f, 0.88237023f, 1.62447476f, 0.80396229f, -1.38206959f, 1.62546301f, + -1.61546838f, -0.56213129f, -0.23501799f, 0.89255226f, -1.95987988f, 0.85192877f, -0.06520678f, -1.32849765f, + 2.07457638f, -0.08192353f, -2.03260493f, 0.58190948f, 2.22535419f, -0.60754669f, 1.14538383f, 0.22928622f, + -0.11596665f, -0.57144678f, -0.23428933f, -0.68404931f, -1.46875453f, 1.32763886f, 0.28525546f, -0.11347114f, + 1.63199806f, -1.44967401f, -2.54707336f, 0.78083873f, -0.19109090f, 0.59508920f, 0.58886564f, 0.81380880f, + 0.54458785f, 0.21354041f, 0.03592521f, 0.27506533f, -0.87245977f, 0.65083951f, 0.32723498f, -0.48263270f, + 1.54926634f, -0.49424326f, 0.65303659f, 1.03355420f, 1.36284590f, 0.78218406f, -1.22659552f, 2.27737451f}; + data.value_data = { + 0.20429733f, -0.57036293f, 0.22116289f, 0.07601038f, 1.79898310f, 0.62182522f, 0.48815370f, -1.59284389f, + 0.33195397f, 0.34822315f, 0.54315579f, 1.06468117f, 1.34500551f, -0.09528533f, -1.30459058f, -0.07034321f, + -1.34877563f, 1.58868146f, -1.44948101f, 0.74792957f, 0.91922742f, -0.56811053f, 0.59939134f, -1.10749292f, + 1.36371183f, -0.89673072f, -0.28341034f, 0.93497890f, 1.62986696f, -0.83026254f, -0.20963377f, -2.14284325f, + -0.95242530f, 0.37379366f, 1.17815948f, -0.55676895f, 0.74420613f, 0.58715403f, -0.43127203f, 0.62706453f, + 0.50881875f, 2.14387321f, 0.85787302f, 2.32273459f, -0.04902139f, -0.04061748f, 1.55004728f, -0.25090796f, + -1.05647933f, -0.87152874f, 0.70491379f, -0.44258183f, 0.26371828f, 1.58179390f, 0.02976498f, 0.20476237f, + 1.20772719f, -0.99407929f, -0.15339416f, 0.54562038f, 1.29705775f, -0.28651321f, -0.90150839f, -1.09473300f}; + + data.bias_data = { + -0.38124341f, 0.02696526f, -0.11914945f, -0.43795273f, -0.34948170f, -0.19608477f, 0.19725692f, 0.39987487f, + 0.04772711f, -0.03419551f, -0.30606642f, 0.42656231f, -0.23178342f, -0.13692456f, -0.04889601f, 0.48739988f, + 0.27079183f, 0.42074734f, -0.40314156f, -0.43726659f, 0.27376485f, -0.38174152f, -0.43700469f, 0.38040614f, + -0.40546918f, 0.06927037f, 0.16979086f, 0.41458064f, 0.07120579f, -0.08055863f, 0.12095112f, -0.27988660f, + -0.10567203f, 0.26791072f, -0.08976898f, 0.31341976f, 0.06027532f, 0.14307594f, 0.31587386f, 0.16180152f, + 0.34785229f, 0.00531715f, -0.35168743f, -0.11641458f, 0.39196932f, 0.44535065f, 0.43545735f, 0.15593112f}; + + data.fp32_output_data = { + -0.73531479f, 0.17652693f, 0.43294340f, 0.10832195f, 1.06569219f, 0.84791648f, 0.37950000f, -0.19036117f, + 1.14446115f, -0.09822676f, -0.12345937f, 0.85349751f, 1.55650973f, 0.24238834f, -0.30006421f, -0.48439008f, + + -0.55764329f, -0.13255559f, 0.21757828f, 0.22631887f, 1.18009710f, 0.97638327f, 0.58746934f, -0.60751349f, + 1.29570508f, -0.01765576f, -0.19382195f, 1.03120971f, 1.47274101f, 0.09958147f, 0.22852209f, -0.86264789f}; + data.fp16_output_data = data.fp32_output_data; + + data.present_key_data = { + 1.97534585f, 0.49419305f, -0.21421000f, -1.92116845f, -0.58778834f, -0.27180552f, -0.73570156f, 1.11841059f, + -1.34467649f, -0.14138395f, -0.63815951f, 0.45528567f, -1.68611503f, 0.47018725f, -0.50221145f, -0.94809151f, + 0.15482518f, -0.15069944f, -0.63743091f, -1.12131596f, -1.19498968f, 0.94589734f, -0.15174922f, 0.26693499f, + 0.81537968f, 0.63428771f, -0.36721635f, -0.16220126f, -0.59869492f, 0.26909798f, -0.10976970f, -0.10222656f, + + 0.54123628f, -1.29785502f, -0.24349773f, 1.29695082f, 1.69568062f, 0.72340369f, -1.26111841f, 1.34557641f, + 1.66910720f, -0.01265316f, -1.86281407f, 0.99649012f, 2.29656005f, -0.68810534f, 1.26633501f, -0.05060038f, + 1.22652888f, -1.38040364f, -2.37728262f, 1.19541931f, -0.11988510f, 0.51453054f, 0.70981675f, 0.53392220f, + 1.14379716f, -0.42497289f, 0.82282746f, 1.44813490f, 1.43405175f, 0.70162547f, -1.10564446f, 1.99748790f}; + data.present_value_data = { + 0.09862530f, -0.30245221f, 0.13139391f, 0.38943014f, 1.85925841f, 0.76490116f, 0.80402756f, -1.43104243f, + -1.45444763f, 1.85659218f, -1.53925002f, 1.06134939f, 0.97950274f, -0.42503458f, 0.91526520f, -0.94569141f, + -1.05809736f, 0.64170438f, 1.08839047f, -0.24334919f, 0.80448145f, 0.73022997f, -0.11539817f, 0.78886604f, + -1.16215134f, -0.60361803f, 0.61514485f, -0.12916207f, 0.32399359f, 1.72486985f, 0.34563884f, 0.36656389f, + + 0.67980623f, 0.35354030f, 0.19146836f, 0.94826663f, 1.73697484f, 0.35006532f, -0.86913323f, 0.08558790f, + 1.71156406f, -0.89141357f, -0.63509774f, 0.81856430f, 2.02183628f, -0.38491189f, 0.22582358f, -1.98691213f, + 0.85667104f, 2.14919043f, 0.50618559f, 2.20632005f, 0.34294793f, 0.40473318f, 1.98550463f, -0.09497684f, + 1.55557942f, -0.98876214f, -0.50508159f, 0.42920581f, 1.68902707f, 0.15883744f, -0.46605104f, -0.93880188f}; + + data.is_static_kv = true; +} + +void GetCrossAttentionData_DiffSequenceLengths_HeadSize8_NoBias(AttentionTestData& data) { + GetCrossAttentionData_DiffSequenceLengths_HeadSize8(data); + data.bias_data.clear(); + data.fp32_output_data = { + -0.65824336f, -0.13227919f, 0.50652772f, -0.20617434f, 0.96178705f, 0.75033671f, 0.07374202f, -0.33853477f, + 0.79871416f, -0.14487229f, 0.21765450f, 0.94517386f, 1.19351113f, -0.20844331f, -0.78461820f, -0.64874089f, + + -0.51557893f, -0.43041292f, 0.32726392f, -0.11220890f, 1.03794634f, 0.89304829f, 0.25361922f, -0.68402284f, + 0.98060286f, -0.14295936f, 0.11289610f, 1.09369516f, 1.15080333f, -0.37929729f, -0.27804646f, -1.09627187f}; + data.fp16_output_data = data.fp32_output_data; + data.present_key_data = { + 1.70455408f, 0.07344571f, 0.18893155f, -1.48390186f, -0.86155319f, 0.10993601f, -0.29869685f, 0.73800445f, + -1.61546838f, -0.56213129f, -0.23501799f, 0.89255226f, -1.95987988f, 0.85192877f, -0.06520678f, -1.32849765f, + -0.11596665f, -0.57144678f, -0.23428933f, -0.68404931f, -1.46875453f, 1.32763886f, 0.28525546f, -0.11347114f, + 0.54458785f, 0.21354041f, 0.03592521f, 0.27506533f, -0.87245977f, 0.65083951f, 0.32723498f, -0.48263270f, + + 0.94670546f, -1.36712539f, -0.41328859f, 0.88237023f, 1.62447476f, 0.80396229f, -1.38206959f, 1.62546301f, + 2.07457638f, -0.08192353f, -2.03260493f, 0.58190948f, 2.22535419f, -0.60754669f, 1.14538383f, 0.22928622f, + 1.63199806f, -1.44967401f, -2.54707336f, 0.78083873f, -0.19109090f, 0.59508920f, 0.58886564f, 0.81380880f, + 1.54926634f, -0.49424326f, 0.65303659f, 1.03355420f, 1.36284590f, 0.78218406f, -1.22659552f, 2.27737451f}; + data.present_value_data = { + 0.20429733f, -0.57036293f, 0.22116289f, 0.07601038f, 1.79898310f, 0.62182522f, 0.48815370f, -1.59284389f, + -1.34877563f, 1.58868146f, -1.44948101f, 0.74792957f, 0.91922742f, -0.56811053f, 0.59939134f, -1.10749292f, + -0.95242530f, 0.37379366f, 1.17815948f, -0.55676895f, 0.74420613f, 0.58715403f, -0.43127203f, 0.62706453f, + -1.05647933f, -0.87152874f, 0.70491379f, -0.44258183f, 0.26371828f, 1.58179390f, 0.02976498f, 0.20476237f, + + 0.33195397f, 0.34822315f, 0.54315579f, 1.06468117f, 1.34500551f, -0.09528533f, -1.30459058f, -0.07034321f, + 1.36371183f, -0.89673072f, -0.28341034f, 0.93497890f, 1.62986696f, -0.83026254f, -0.20963377f, -2.14284325f, + 0.50881875f, 2.14387321f, 0.85787302f, 2.32273459f, -0.04902139f, -0.04061748f, 1.55004728f, -0.25090796f, + 1.20772719f, -0.99407929f, -0.15339416f, 0.54562038f, 1.29705775f, -0.28651321f, -0.90150839f, -1.09473300f}; + + data.is_static_kv = true; +} + void GetSelfAttentionData_WithPastAndPresent_NoMask_NoRelPosBias(AttentionTestData& data) { data.hidden_size = 8; data.v_hidden_size = 8; @@ -4334,6 +4448,131 @@ void GetSelfAttentionData_WithPastAndPresent_NoMask_NoRelPosBias(AttentionTestDa data.is_static_kv = false; } +void GetSelfAttentionData_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias(AttentionTestData& data) { + data.hidden_size = 16; + data.v_hidden_size = 16; + data.num_heads = 2; + data.batch_size = 2; + data.sequence_length = 1; + data.kv_sequence_length = 1; + data.mask_type = AttentionMaskType::MASK_NONE; + + data.skip_kernel_types = { + AttentionKernelType::AttentionKernel_TrtFlashAttention, + AttentionKernelType::AttentionKernel_TrtFusedCrossAttention, + AttentionKernelType::AttentionKernel_TrtFusedAttention, + AttentionKernelType::AttentionKernel_CutlassMemoryEfficientAttention, + }; + + data.query_data = { + 1.29534733f, 2.14051294f, 1.09895217f, 1.39164531f, -0.01471180f, -1.40148544f, -0.50825417f, 0.26134527f, + -0.70491123f, 0.63738143f, 2.13708138f, 0.05667466f, -0.44220763f, 0.85254443f, 2.00844359f, -1.23413038f, + -0.08030051f, -1.25450790f, -0.89664006f, -0.69433510f, 0.20943037f, 1.41880298f, 1.42875051f, 0.79920006f, + 1.57896936f, -1.13204634f, -0.61002654f, 0.43365243f, 0.22888106f, -0.38688308f, -0.45924744f, 0.99473029f}; + + data.key_data = { + 0.37680483f, 0.15317714f, 0.05767500f, 0.37780648f, -2.27755547f, 0.89294612f, -0.85582626f, 0.54963046f, + 1.67390800f, -1.06330085f, -2.99566054f, 0.68927419f, 1.66056263f, -0.77022851f, 0.15417719f, 0.94860524f, + -1.84928346f, -0.52135336f, 0.70491475f, 0.37400877f, 0.55338752f, 0.52915680f, 0.52876079f, -0.55780333f, + -1.49814773f, 0.18675917f, 0.31246936f, -1.32707596f, 0.42132780f, -1.69121027f, 0.20342645f, -0.34370381f}; + + data.value_data = { + 0.60890561f, -0.88021755f, 1.63002241f, 0.86171651f, 1.80559230f, 1.26110435f, -0.97890180f, -1.60215497f, + -0.79229754f, 1.07830989f, -0.85298145f, 2.76264572f, 0.01659799f, -1.49499071f, 0.85316724f, -2.56763911f, + 0.53017867f, 1.31909978f, -1.10940945f, 0.68858552f, -1.07115889f, -2.34016919f, 0.48310637f, -0.05351824f, + -0.08850761f, -0.56362265f, 0.05224326f, -2.47377181f, 0.44249821f, -0.10389519f, -0.46113095f, 2.81619215f}; + + data.bias_data = { + -0.38124341f, 0.02696526f, -0.11914945f, -0.43795273f, -0.34948170f, -0.19608477f, 0.19725692f, 0.39987487f, + 0.04772711f, -0.03419551f, -0.30606642f, 0.42656231f, -0.23178342f, -0.13692456f, -0.04889601f, 0.48739988f, + 0.27079183f, 0.42074734f, -0.40314156f, -0.43726659f, 0.27376485f, -0.38174152f, -0.43700469f, 0.38040614f, + -0.40546918f, 0.06927037f, 0.16979086f, 0.41458064f, 0.07120579f, -0.08055863f, 0.12095112f, -0.27988660f, + -0.10567203f, 0.26791072f, -0.08976898f, 0.31341976f, 0.06027532f, 0.14307594f, 0.31587386f, 0.16180152f, + 0.34785229f, 0.00531715f, -0.35168743f, -0.11641458f, 0.39196932f, 0.44535065f, 0.43545735f, 0.15593112f}; + + data.past_key_data = { + -1.00657940f, -0.46509427f, -1.65118766f, -0.17705369f, 1.71204090f, 0.53921354f, -1.67056096f, 0.42517155f, + -2.00129080f, 1.26244307f, 0.28864837f, 1.38792157f, -0.59647840f, -1.18904924f, 0.58950418f, -2.26774645f, + 1.88496518f, 0.59231639f, 0.33360308f, -1.23532701f, 0.10543400f, -1.77481365f, -0.79397631f, -0.22495472f, + -0.26800078f, -0.20456636f, 1.43141091f, 1.55566478f, -0.22702518f, 1.75312757f, -1.29037595f, -0.95538902f}; + data.past_value_data = { + 3.18056512f, 0.13370860f, -2.20253444f, 2.30826044f, 0.86762893f, -1.91499686f, 2.18277764f, 0.53384149f, + -0.43230706f, 0.49148068f, -0.29957789f, -3.56583714f, -1.46747136f, -0.40299624f, 1.78018796f, 2.84104395f, + -0.68692255f, 1.25688624f, -0.42734757f, -1.03185725f, 0.47858545f, 1.18466282f, -1.06095874f, -0.63918531f, + 1.41408277f, 0.74389833f, 0.89590931f, 1.06388271f, 1.29734015f, 0.42640167f, -0.99740052f, -2.79366398f}; + + data.fp32_output_data = { + 0.72723210f, -0.54989153f, 1.22711349f, 1.26993895f, 1.78235006f, 1.12648177f, -0.42493403f, -1.27518260f, + -0.43240935f, 0.49647018f, -0.30720428f, -3.51349354f, -1.45166361f, -0.40844491f, 1.77604592f, 2.79678369f, + 0.25752395f, 1.53741217f, -1.08321750f, 0.69643497f, -0.78710371f, -1.68901348f, 0.51954043f, -0.00401744f, + 1.11207914f, 0.40332735f, 0.58328331f, 0.10821819f, 1.17628312f, 0.40418532f, -0.74326056f, -1.28571272f}; + data.fp16_output_data = data.fp32_output_data; + + data.present_key_data = { + -1.00657940f, -0.46509427f, -1.65118766f, -0.17705369f, 1.71204090f, 0.53921354f, -1.67056096f, 0.42517155f, + 0.64759666f, 0.57392448f, -0.34546655f, -0.05946010f, -2.00379062f, 0.51120460f, -1.29283094f, 0.93003660f, + -2.00129080f, 1.26244307f, 0.28864837f, 1.38792157f, -0.59647840f, -1.18904924f, 0.58950418f, -2.26774645f, + 1.26843882f, -0.99403048f, -2.82586956f, 1.10385489f, 1.73176837f, -0.85078716f, 0.27512830f, 0.66871864f, + 1.88496518f, 0.59231639f, 0.33360308f, -1.23532701f, 0.10543400f, -1.77481365f, -0.79397631f, -0.22495472f, + -1.57849169f, -0.10060602f, 0.30177319f, -0.06325781f, 0.82715237f, 0.14741528f, 0.09175611f, -0.17739719f, + -0.26800078f, -0.20456636f, 1.43141091f, 1.55566478f, -0.22702518f, 1.75312757f, -1.29037595f, -0.95538902f, + -1.90361691f, 0.25602955f, 0.48226023f, -0.91249532f, 0.49253359f, -1.77176893f, 0.32437757f, -0.62359041f}; + data.present_value_data = { + 3.18056512f, 0.13370860f, -2.20253444f, 2.30826044f, 0.86762893f, -1.91499686f, 2.18277764f, 0.53384149f, + 0.50323355f, -0.61230683f, 1.54025340f, 1.17513633f, 1.86586761f, 1.40418029f, -0.66302794f, -1.44035339f, + -0.43230706f, 0.49148068f, -0.29957789f, -3.56583714f, -1.46747136f, -0.40299624f, 1.78018796f, 2.84104395f, + -0.44444525f, 1.08362699f, -1.20466888f, 2.64623117f, 0.40856731f, -1.04964006f, 1.28862453f, -2.41170788f, + -0.68692255f, 1.25688624f, -0.42734757f, -1.03185725f, 0.47858545f, 1.18466282f, -1.06095874f, -0.63918531f, + 0.42450663f, 1.58701050f, -1.19917846f, 1.00200534f, -1.01088357f, -2.19709325f, 0.79898024f, 0.10828328f, + 1.41408277f, 0.74389833f, 0.89590931f, 1.06388271f, 1.29734015f, 0.42640167f, -0.99740052f, -2.79366398f, + 0.25934470f, -0.55830550f, -0.29944417f, -2.59018636f, 0.83446753f, 0.34145546f, -0.02567360f, 2.97212315f}; + + data.is_static_kv = false; +} + +void GetSelfAttentionData_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias_NoBias(AttentionTestData& data) { + GetSelfAttentionData_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias(data); + data.bias_data.clear(); + data.past_key_data = { + -1.27737117f, -0.88584161f, -1.24804604f, 0.26021290f, 1.43827605f, 0.92095506f, -1.23355627f, 0.04476542f, + -1.59582162f, 1.19317269f, 0.11885749f, 0.97334087f, -0.66768420f, -1.10849059f, 0.46855307f, -1.98785996f, + 1.61417341f, 0.17156902f, 0.73674464f, -0.79806042f, -0.16833085f, -1.39307213f, -0.35697165f, -0.60536087f, + 0.13746840f, -0.27383673f, 1.26162004f, 1.14108407f, -0.29823098f, 1.83368623f, -1.41132712f, -0.67550242f}; + data.past_value_data = { + 3.28623724f, -0.13420212f, -2.11276555f, 1.99484074f, 0.80735362f, -2.05807281f, 1.86690378f, 0.37204000f, + -0.78015935f, 0.48616353f, 0.05210955f, -3.44942260f, -1.85944068f, -0.84834689f, 1.34473062f, 2.68511271f, + -0.58125055f, 0.98897558f, -0.33757859f, -1.34527707f, 0.41831014f, 1.04158688f, -1.37683260f, -0.80098683f, + 1.06623054f, 0.73858118f, 1.24759674f, 1.18029726f, 0.90537083f, -0.01894896f, -1.43285787f, -2.94959521f}; + + data.fp32_output_data = { + 0.89556247f, -0.80034304f, 1.22928894f, 0.98303795f, 1.69871271f, 0.90572613f, -0.67420667f, -1.39078152f, + -0.78021139f, 0.48869953f, 0.04823331f, -3.42281842f, -1.85140634f, -0.85111630f, 1.34262550f, 2.66261697f, + 0.34449580f, 1.26394701f, -0.98046219f, 0.34879467f, -0.82231814f, -1.77519011f, 0.17237240f, -0.17839541f, + 0.72679031f, 0.35579273f, 0.89621741f, 0.10616791f, 0.76930743f, -0.04391927f, -1.14721453f, -1.25471735f}; + data.fp16_output_data = data.fp32_output_data; + + data.present_key_data = { + -1.27737117f, -0.88584161f, -1.24804604f, 0.26021290f, 1.43827605f, 0.92095506f, -1.23355627f, 0.04476542f, + 0.37680483f, 0.15317714f, 0.05767500f, 0.37780648f, -2.27755547f, 0.89294612f, -0.85582626f, 0.54963046f, + -1.59582162f, 1.19317269f, 0.11885749f, 0.97334087f, -0.66768420f, -1.10849059f, 0.46855307f, -1.98785996f, + 1.67390800f, -1.06330085f, -2.99566054f, 0.68927419f, 1.66056263f, -0.77022851f, 0.15417719f, 0.94860524f, + 1.61417341f, 0.17156902f, 0.73674464f, -0.79806042f, -0.16833085f, -1.39307213f, -0.35697165f, -0.60536087f, + -1.84928346f, -0.52135336f, 0.70491475f, 0.37400877f, 0.55338752f, 0.52915680f, 0.52876079f, -0.55780333f, + 0.13746840f, -0.27383673f, 1.26162004f, 1.14108407f, -0.29823098f, 1.83368623f, -1.41132712f, -0.67550242f, + -1.49814773f, 0.18675917f, 0.31246936f, -1.32707596f, 0.42132780f, -1.69121027f, 0.20342645f, -0.34370381f}; + data.present_value_data = { + 3.28623724f, -0.13420212f, -2.11276555f, 1.99484074f, 0.80735362f, -2.05807281f, 1.86690378f, 0.37204000f, + 0.60890561f, -0.88021755f, 1.63002241f, 0.86171651f, 1.80559230f, 1.26110435f, -0.97890180f, -1.60215497f, + -0.78015935f, 0.48616353f, 0.05210955f, -3.44942260f, -1.85944068f, -0.84834689f, 1.34473062f, 2.68511271f, + -0.79229754f, 1.07830989f, -0.85298145f, 2.76264572f, 0.01659799f, -1.49499071f, 0.85316724f, -2.56763911f, + -0.58125055f, 0.98897558f, -0.33757859f, -1.34527707f, 0.41831014f, 1.04158688f, -1.37683260f, -0.80098683f, + 0.53017867f, 1.31909978f, -1.10940945f, 0.68858552f, -1.07115889f, -2.34016919f, 0.48310637f, -0.05351824f, + 1.06623054f, 0.73858118f, 1.24759674f, 1.18029726f, 0.90537083f, -0.01894896f, -1.43285787f, -2.94959521f, + -0.08850761f, -0.56362265f, 0.05224326f, -2.47377181f, 0.44249821f, -0.10389519f, -0.46113095f, 2.81619215f}; + + data.is_static_kv = false; +} + void GetCrossAttentionData_WithPastPassedInDirectly_NoMask(AttentionTestData& data) { data.hidden_size = 4; data.v_hidden_size = 4; diff --git a/onnxruntime/test/contrib_ops/attention_op_test_helper.h b/onnxruntime/test/contrib_ops/attention_op_test_helper.h index db9f3edfd4..3da171532c 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test_helper.h +++ b/onnxruntime/test/contrib_ops/attention_op_test_helper.h @@ -68,7 +68,11 @@ void GetCrossAttentionDataWithPast(AttentionTestData& data); void GetSelfAttentionData_WithPast_WithRelPosBias_ForT5(AttentionTestData& data); void GetCrossAttentionData_DiffSequenceLengths(AttentionTestData& data); +void GetCrossAttentionData_DiffSequenceLengths_HeadSize8(AttentionTestData& data); +void GetCrossAttentionData_DiffSequenceLengths_HeadSize8_NoBias(AttentionTestData& data); void GetSelfAttentionData_WithPastAndPresent_NoMask_NoRelPosBias(AttentionTestData& data); +void GetSelfAttentionData_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias(AttentionTestData& data); +void GetSelfAttentionData_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias_NoBias(AttentionTestData& data); void GetCrossAttentionData_WithPastPassedInDirectly_NoMask(AttentionTestData& data); void GetAttentionDataCutlassRelPosBias(AttentionTestData& data); diff --git a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc index 9f2d47cf5b..2295873b42 100644 --- a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc @@ -370,7 +370,7 @@ static void RunMultiHeadAttentionKernel( } } -static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu = false) { +static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu = false, bool disable_cuda = false) { if (data.fp32_output_data.size() > 0) { constexpr bool use_float16 = false; @@ -381,7 +381,7 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data, data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp32_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size, - data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu); + data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu, disable_cuda); } #if USE_FLASH_ATTENTION @@ -394,7 +394,7 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data, data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp32_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size, - data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu); + data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu, disable_cuda); } } #endif @@ -405,7 +405,7 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data, data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp32_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size, - data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu); + data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu, disable_cuda); } if (data.fp16_output_data.size() > 0) { @@ -417,7 +417,7 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data, data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp16_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size, - data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu); + data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu, disable_cuda); } kernel_type = AttentionKernelType::AttentionKernel_TrtFusedAttention; @@ -427,7 +427,7 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data, data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp16_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size, - data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu); + data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu, disable_cuda); } #if USE_FLASH_ATTENTION @@ -438,7 +438,7 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data, data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp16_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size, - data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu); + data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu, disable_cuda); } #endif @@ -448,7 +448,7 @@ static void RunMultiHeadAttentionTests(AttentionTestData& data, bool disable_cpu data.rel_pos_bias_data, data.past_key_data, data.past_value_data, data.present_key_data, data.present_value_data, data.key_padding_mask_data, data.mask_type, data.fp16_output_data, data.num_heads, data.batch_size, data.sequence_length, data.kv_sequence_length, data.hidden_size, - data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu); + data.v_hidden_size, kernel_type, use_float16, data.is_static_kv, disable_cpu, disable_cuda); } } @@ -559,6 +559,12 @@ TEST(MultiHeadAttentionTest, CrossAttention_DiffSequenceLengths) { AttentionTestData data; GetCrossAttentionData_DiffSequenceLengths(data); RunMultiHeadAttentionTests(data); + + GetCrossAttentionData_DiffSequenceLengths_HeadSize8(data); + RunMultiHeadAttentionTests(data, /*disable_cpu=*/false, /*disable_cuda=*/true); + + GetCrossAttentionData_DiffSequenceLengths_HeadSize8_NoBias(data); + RunMultiHeadAttentionTests(data, /*disable_cpu=*/false, /*disable_cuda=*/true); } TEST(MultiHeadAttentionTest, SelfAttention_WithPastAndPresent_NoMask_NoRelPosBias) { @@ -567,6 +573,12 @@ TEST(MultiHeadAttentionTest, SelfAttention_WithPastAndPresent_NoMask_NoRelPosBia AttentionTestData data; GetSelfAttentionData_WithPastAndPresent_NoMask_NoRelPosBias(data); RunMultiHeadAttentionTests(data); + + GetSelfAttentionData_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias(data); + RunMultiHeadAttentionTests(data, /*disable_cpu=*/false, /*disable_cuda=*/true); + + GetSelfAttentionData_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias_NoBias(data); + RunMultiHeadAttentionTests(data, /*disable_cpu=*/false, /*disable_cuda=*/true); } TEST(MultiHeadAttentionTest, CrossAttention_WithPastPassedInDirectly_NoMask) { diff --git a/onnxruntime/test/contrib_ops/multihead_attention_op_test_data_gen.py b/onnxruntime/test/contrib_ops/multihead_attention_op_test_data_gen.py index 2e65ea266c..96df6dc52b 100644 --- a/onnxruntime/test/contrib_ops/multihead_attention_op_test_data_gen.py +++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test_data_gen.py @@ -112,6 +112,8 @@ class Attention(nn.Module): # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) # if encoder bi-directional self-attention `past_key_value` is always `None` past_key_value = (key_layer, value_layer) + print("k cache", key_layer) + print("v cache", value_layer) # Take the dot product between "query" and "key" to get the raw attention scores. attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) @@ -157,6 +159,7 @@ def run_cross_attention( kv_sequence_length, key_padding_mask=None, has_bias=True, + is_decoder=False, ): seed = 123 torch.manual_seed(seed) @@ -164,7 +167,7 @@ def run_cross_attention( torch.use_deterministic_algorithms(True) device = torch.device("cuda:0") - mha = Attention(num_heads, hidden_dim, q_head_size, v_head_size, is_decoder=False).to(device).eval() + mha = Attention(num_heads, hidden_dim, q_head_size, v_head_size, is_decoder=is_decoder).to(device).eval() if key_padding_mask is not None: key_padding_mask = key_padding_mask.to(device) torch.nn.init.uniform_(mha.query.weight, -0.5, 0.5) @@ -255,6 +258,7 @@ def run_self_attention( sequence_length, key_padding_mask=None, has_bias=True, + is_decoder=False, ): seed = 123 torch.manual_seed(seed) @@ -262,7 +266,7 @@ def run_self_attention( torch.use_deterministic_algorithms(True) device = torch.device("cuda:0") - mha = Attention(num_heads, hidden_dim, q_head_size, v_head_size, is_decoder=False).to(device).eval() + mha = Attention(num_heads, hidden_dim, q_head_size, v_head_size, is_decoder=is_decoder).to(device).eval() if key_padding_mask is not None: key_padding_mask = key_padding_mask.to(device) torch.nn.init.uniform_(mha.query.weight, -0.5, 0.5) @@ -461,6 +465,54 @@ def run_self_batch2_headsize_32_packed_qkv(): ) +def run_cross_diff_seqlen_headsize_8(): + hidden_dim = 16 + q_head_size = 8 + v_head_size = 8 + num_heads = 2 + batch_size = 1 + sequence_length = 2 + kv_sequence_length = 4 + key_padding_mask = None + has_bias = True + run_cross_attention( + hidden_dim, + q_head_size, + v_head_size, + num_heads, + batch_size, + sequence_length, + kv_sequence_length, + key_padding_mask, + has_bias, + is_decoder=True, + ) + + +def run_self_past_present_headsize_8_nomask_norelposbias(): + hidden_dim = 16 + q_head_size = 8 + v_head_size = 8 + num_heads = 2 + batch_size = 2 + # In cpp side we use sequence_length = 1, we manually split the data of the first and second token. + # Then we use first token related data as past and second token related data as true input. + sequence_length = 2 + key_padding_mask = None + has_bias = True + run_self_attention( + hidden_dim, + q_head_size, + v_head_size, + num_heads, + batch_size, + sequence_length, + key_padding_mask, + has_bias, + is_decoder=True, + ) + + def create_test_data(): """ Create test data used in attention_op_test_helper.cc and multihead_attention_op_test.cc @@ -486,6 +538,12 @@ def create_test_data(): print("SelfAttention_Batch2_HeadSize32_PackedQKV") run_self_batch2_headsize_32_packed_qkv() + print("SelfAttention_WithPastAndPresent_HeadSize8_NoMask_NoRelPosBias") + run_self_past_present_headsize_8_nomask_norelposbias() + + print("CrossAttention_DiffSequenceLengths_HeadSize8") + run_cross_diff_seqlen_headsize_8() + with torch.no_grad(): create_test_data()