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

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

sgl-kernel/csrc/cpu/flash_attn.cpp

@@ -97,8 +97,11 @@ void flash_attn_kernel_impl(
     alignas(64) float m_prime[BLOCK_M];
 
     for (int i = begin; i < end; ++i) {
-      int seq_q_start_loc = bs * seqlen_q;
-      int seq_k_start_loc = bs * seqlen_k;
+      // [Note] use int64_t to avoid overflow
+      // For large inputs, for example bs = 4096, seqlen_q = 4097, m = 0, q_strideM = 128:
+      // The index calculated below: (seq_q_start_loc + m) * q_strideM = 4096 * 4097 * 128 will overflow int
+      int64_t seq_q_start_loc = bs * seqlen_q;
+      int64_t seq_k_start_loc = bs * seqlen_k;
 
       // offset and size in MB
       int m = mb * BLOCK_M;
@@ -272,8 +275,11 @@ void flash_attn_varlen_kernel_impl(
 
     for (int i = begin; i < end; ++i) {
       int32_t bs = indices[mb * 2 + 0];
-      int32_t seq_q_start_loc = cu_seqlens_q[bs];
-      int32_t seq_k_start_loc = cu_seqlens_k[bs];
+
+      // See [Note] use int64_t to avoid overflow
+      int64_t seq_q_start_loc = cu_seqlens_q[bs];
+      int64_t seq_k_start_loc = cu_seqlens_k[bs];
+
       int32_t seqlen_q = cu_seqlens_q[bs + 1] - cu_seqlens_q[bs];
 
       // offset and size in MB

sgl-kernel/csrc/cpu/flash_attn.h

@@ -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)

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()