Reduce write cache polution to improve IAKV performance (#2457) (#2476)

Co-authored-by: Chunyuan WU <chunyuan.wu@intel.com>
intel · Jan 17, 2024 · c95eb77 · c95eb77
1 parent df2387e
commit c95eb77
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/csrc/cpu/aten/kernels/MaskedMultiHeadAttentionKrnl.cpp b/csrc/cpu/aten/kernels/MaskedMultiHeadAttentionKrnl.cpp
@@ -798,7 +798,8 @@ scale_dot_product_for_indirect_access_kv_cache(
               }
             }
           }
-          flag_access[thread_id][bi][hi] = 1;
+          if (flag_access[thread_id][bi][hi] == 0)
+            flag_access[thread_id][bi][hi] = 1;
         }
       }
     }
@@ -1102,7 +1103,8 @@ scale_dot_product_for_indirect_access_kv_cache_half(
                     flag_access[thread_id][bi][hi]);
               }
             }
-            flag_access[thread_id][bi][hi] = 1;
+            if (flag_access[thread_id][bi][hi] == 0)
+              flag_access[thread_id][bi][hi] = 1;
           }
         }
       }