chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Wu, Chunyuan <chunyuan.wu@intel.com>
This commit is contained in:
blzheng
2026-03-19 13:18:47 +08:00
committed by GitHub
parent 687d9eb66f
commit c2b01bd2fc
4 changed files with 107 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
sgl-kernel/csrc/cpu/flash_attn.cpp
View File

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

1
sgl-kernel/csrc/cpu/flash_attn.h
View File

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