[CPU] fix bug in AVX512 implementation of flash_attn_softmax (#20220)

Co-authored-by: Wu, Chunyuan <chunyuan.wu@intel.com>
2026-03-18 22:18:47 -07:00
parent 687d9eb66f
commit c2b01bd2fc
4 changed files with 107 additions and 6 deletions
@@ -190,6 +190,7 @@ struct flash_attn_softmax<at::BFloat16, BLOCK_M, BLOCK_N> {

      // m_i: max value per row
      float m_i = _mm512_reduce_max_ps(vmax);
+      m_i = std::max(m_i, m_prime[m]);
      vmax = _mm512_set1_ps(m_i);

      // m_delta <- exp(m' - m_i)