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

Co-authored-by: Wu, Chunyuan <chunyuan.wu@intel.com>

test/srt/cpu/test_extend.py

@@ -118,8 +118,8 @@ class TestExtendAttention(CustomTestCase):
             v_extend[extend_start:extend_end] = v_buffer[
                 extend_start_in_buffer:extend_end_in_buffer
             ]
-            q_extend[extend_start:extend_end] = torch.randn(
-                (b_seq_len_extend[i], H_Q, D), dtype=dtype
+            q_extend[extend_start:extend_end] = (
+                torch.randn((b_seq_len_extend[i], H_Q, D), dtype=dtype) * 20
             )
 
         # q_extend, k_extend, v_extend, k_buffer and v_buffer supports non-contiguous tensors

test/srt/cpu/test_flash_attn.py

@@ -47,6 +47,37 @@ def flash_attn_varlen_ref(
     return out.transpose(1, 2).squeeze(0)
 
 
+# faster version ref kernel for non varlen case
+def flash_attn_non_varlen_ref(
+    q,
+    k,
+    v,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    is_causal,
+    enable_gqa,
+):
+    cu_q = cu_seqlens_q.tolist()
+    cu_k = cu_seqlens_k.tolist()
+    batch = len(cu_k) - 1
+
+    B_T, H, D = q.shape
+    T = B_T // batch
+
+    # [T, H, D] -> [1, H, T, D]
+    q, k, v = [x.reshape(batch, T, H, D).transpose(1, 2) for x in [q, k, v]]
+
+    out = F.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        is_causal=is_causal,
+        enable_gqa=enable_gqa,
+    )
+    # [B, H, T, D] -> [B * T, H, D]
+    return out.transpose(1, 2).reshape(batch * T, H, D)
+
+
 class TestFlashAttn(CustomTestCase):
 
     @parametrize(
@@ -110,6 +141,69 @@ class TestFlashAttn(CustomTestCase):
         atol = rtol = precision[dtype]
         torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol)
 
+    # test with large size to capture overflow issue
+    @parametrize(
+        batch=[4097],
+        max_seqlen_q=[4097],
+        max_seqlen_k=[4097],
+        num_heads=[4],
+        num_heads_kv=[4],
+        head_dim=[32],
+        head_dim_v=[32],
+        is_causal=[False],
+    )
+    def test_flash_attn_large_size(
+        self,
+        batch,
+        max_seqlen_q,
+        max_seqlen_k,
+        num_heads,
+        num_heads_kv,
+        head_dim,
+        head_dim_v,
+        is_causal,
+    ):
+        dtype = torch.bfloat16
+
+        # test the non varlen case
+        seqlens_q = torch.full((batch,), max_seqlen_q, dtype=torch.int32)
+        seqlens_k = torch.full((batch,), max_seqlen_k, dtype=torch.int32)
+
+        cu_seqlens_q = torch.zeros((batch + 1,), dtype=torch.int32)
+        cu_seqlens_k = torch.zeros((batch + 1,), dtype=torch.int32)
+        cu_seqlens_q[1:] = torch.cumsum(seqlens_q, 0)
+        cu_seqlens_k[1:] = torch.cumsum(seqlens_k, 0)
+
+        sum_seqlen_q = seqlens_q.sum().item()
+        sum_seqlen_k = seqlens_k.sum().item()
+        q = torch.randn(sum_seqlen_q, num_heads, head_dim).to(dtype)
+        k = torch.randn(sum_seqlen_k, num_heads_kv, head_dim).to(dtype)
+        v = torch.randn(sum_seqlen_k, num_heads_kv, head_dim_v).to(dtype)
+
+        out_ref = flash_attn_non_varlen_ref(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            is_causal=is_causal,
+            enable_gqa=num_heads != num_heads_kv,
+        )
+
+        out = flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            seqlens_q.max().item(),
+            seqlens_k.max().item(),
+            is_causal,
+        )
+
+        atol = rtol = precision[dtype]
+        torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol)
+
 
 if __name__ == "__main__":
     unittest.main()