Update fused_multihead_attention_v2.h

fix missed unrollList condition
jacobkahn · Dec 8, 2022 · 20f6181 · 20f6181
1 parent 46e1f4a
commit 20f6181
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/3rdparty/trt_fused_multihead_attention/fused_multihead_attention_v2.h b/3rdparty/trt_fused_multihead_attention/fused_multihead_attention_v2.h
@@ -5473,7 +5473,7 @@ class FusedMultiHeadAttentionXMMAKernelV2:
             };
             for (unsigned int i = 0u; i < sizeof(unrollList) / sizeof(unrollList[0]); ++i) {
                 if (mSM == unrollList[i].mSM && mDataType == unrollList[i].mDataType && params.s == unrollList[i].mS
-                    && params.b <= unrollList[i].mMaxBatch) {
+                    && params.d == unrollList[i].mD && params.b <= unrollList[i].mMaxBatch) {
                     forceUnroll = true;
                     break;
                 }