[Layers] Improve the first-token perf. in w8a8(f32 InT) by enabling flashAttn (#43)

Author: abenmao

src/layers/attention.h

@@ -372,7 +372,7 @@ class Attention {
 
         TimeLine t4("MHA");
         if (pastSeqLen == 0) {
-            if (ctx->inputSeqLen > Env::getInstance().getFlashAttnThreshold()) {
+            if (Env::getInstance().getFlashAttnEnabled<InT>(ctx->inputSeqLen)) {
                 flashAttention(ctx, query, key, value, attnSplit, presentKey, presentValue, attnMask, pastSeqLen);
             } else if constexpr ((std::is_same_v<InT, bfloat16_t> && std::is_same_v<OutT, bfloat16_t>)
 #if defined(AMX_FP16_WEIGHT_ONLY_FP16)
@@ -581,7 +581,7 @@ class Attention {
 
         TimeLine t4("MHA");
         if (seqs[0]->getStep() == 0) { // First token generation
-            if (totInSeqLen > Env::getInstance().getFlashAttnThreshold() * seqs.size()) {
+            if (Env::getInstance().getFlashAttnEnabled<InT>(totInSeqLen / seqs.size())) {
                 flashAttention(ctx, query, key, value, attnSplit, keyCaches, valueCaches, seqs);
             } else if constexpr ((std::is_same_v<InT, bfloat16_t> && std::is_same_v<OutT, bfloat16_t>)
 #if defined(AMX_FP16_WEIGHT_ONLY_FP16)

src/utils/environment.h

@@ -48,7 +48,14 @@ class Env {
     int getAMXThresholdM() { return AMXThresholdMValue; }
 
     // get FLASH_ATTN_THRESHOLD
-    int getFlashAttnThreshold() { return FlashAttnThresholdValue; }
+    template <typename T>
+    bool getFlashAttnEnabled(int inputLen) {
+        if (FlashAttnThresholdValue >= 0 &&
+            (std::is_same_v<T, float> || inputLen > FlashAttnThresholdValue))
+            return true;
+        else
+            return false;
+    }
 
     // get ENABLE_CAT_MLP
     bool getMlpCatEnabled() { return MlpCatEnabled; }
@@ -245,13 +252,12 @@ class Env {
         // > threshold to enable flash attention, default 8192
         char *flashAttnThresholdValue = getenv("FLASH_ATTN_THRESHOLD");
         if (flashAttnThresholdValue != NULL) {
-            int value = atoi(flashAttnThresholdValue);
-            if (value >= 0)
-                FlashAttnThresholdValue = value;
-            else
-                printf("[ERROR] FLASH_ATTN_THRESHOLD value need to be greater than or equal to 0.\n");
+            FlashAttnThresholdValue = atoi(flashAttnThresholdValue);
         }
-        printf("[INFO] SeqLen > FLASH_ATTN_THRESHOLD(%d) will enable FlashAttn.\n", FlashAttnThresholdValue);
+        if (FlashAttnThresholdValue < 0)
+            printf("[INFO] FlashAttn is disabled (FLASH_ATTN_THRESHOLD = %d).\n", FlashAttnThresholdValue);
+        else
+            printf("[INFO] SeqLen > FLASH_ATTN_THRESHOLD(%d) will enable FlashAttn.\n", FlashAttnThresholdValue);
     }
 
     // ENABLE_CAT_MLP