Fix attention UT issue.

intel · Jun 4, 2024 · 3cfa650 · 3cfa650
1 parent 0faa9de
commit 3cfa650
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/tests/ut/layers_attention_test.cpp b/tests/ut/layers_attention_test.cpp
@@ -87,14 +87,17 @@ void test_AttentionLLaMA(void) {
     int nextTokenNum = 1;
 
     compareAttentionLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum,
-            kvHeadNum, maxPositions, maxPosEmbed, hiddenSize, qkvProj, qkvProj + qSize, qkvProj + kvSize, oProj);
+            kvHeadNum, maxPositions, maxPosEmbed, hiddenSize, qkvProj, qkvProj + qSize, qkvProj + qSize + kvSize,
+            oProj);
     pastSeqLen += inputSeqLen;
     currentSeqLen = nextTokenNum;
     compareAttentionLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum,
-            kvHeadNum, maxPositions, maxPosEmbed, hiddenSize, qkvProj, qkvProj + qSize, qkvProj + kvSize, oProj);
+            kvHeadNum, maxPositions, maxPosEmbed, hiddenSize, qkvProj, qkvProj + qSize, qkvProj + qSize + kvSize,
+            oProj);
     pastSeqLen += nextTokenNum;
     compareAttentionLLaMA<T>(step++, batchSize, inputSeqLen, pastSeqLen, currentSeqLen, attHeadDim, attHeadNum,
-            kvHeadNum, maxPositions, maxPosEmbed, hiddenSize, qkvProj, qkvProj + qSize, qkvProj + kvSize, oProj);
+            kvHeadNum, maxPositions, maxPosEmbed, hiddenSize, qkvProj, qkvProj + qSize, qkvProj + qSize + kvSize,
+            oProj);
 
     free(qkvProj);
     free(oProj);