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

[CPU] optimize flash_attn_varlen_func (#15708)

Browse Source
This commit is contained in:
Ma Mingfei
2026-01-29 22:07:05 -08:00
committed by GitHub
parent 336dc4579e
commit 88f7759402
10 changed files with 1189 additions and 339 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/decode.cpp
+7 -7
View File
@@ -1049,7 +1049,7 @@ void decode_attention_kernel_impl(
int64_t k_strideH,
int64_t v_strideN,
int64_t v_strideH,
float scaling,
float sm_scale,
float logit_cap,
int64_t max_num_reqs,
int64_t max_context_len,
@@ -1103,7 +1103,7 @@ void decode_attention_kernel_impl(
/* B */ k_buffer + head_id * k_strideH,
/* C */ s_i,
/* ind */ req_to_token + req_pool_id * max_context_len + n,
/* scl */ scaling,
/* scl */ sm_scale,
/* M */ 1,
/* N */ n_size,
/* K */ head_size,
@@ -1192,7 +1192,7 @@ void decode_attention_mla_kernel_impl(
int64_t k_strideH,
int64_t v_strideN,
int64_t v_strideH,
float scaling,
float sm_scale,
float logit_cap,
int64_t max_num_reqs,
int64_t max_context_len,
@@ -1291,7 +1291,7 @@ void decode_attention_mla_kernel_impl(
/* B */ Btmp0,
/* C */ s_i);
const Vec scale_vec = Vec(scaling);
const Vec scale_vec = Vec(sm_scale);
for (int64_t h = 0; h < h_size; ++h) {
// s_i <- s_i * scale
at::vec::map<float>(
@@ -1384,7 +1384,7 @@ void decode_attention_grouped_kernel_impl(
int64_t k_strideH,
int64_t v_strideN,
int64_t v_strideH,
float scaling,
float sm_scale,
float logit_cap,
int64_t max_num_reqs,
int64_t max_context_len,
@@ -1457,7 +1457,7 @@ void decode_attention_grouped_kernel_impl(
/* B */ k_buffer + head_kv_id * k_strideH,
/* C */ s_i,
/* ind */ req_to_token + req_pool_id * max_context_len + n,
/* scl */ scaling,
/* scl */ sm_scale,
/* M */ h_size,
/* N */ n_size,
/* K */ head_size,
@@ -1474,7 +1474,7 @@ void decode_attention_grouped_kernel_impl(
BLOCK_H * BLOCK_N);
}
// update the scaling coefficients
// update the sm_scale coefficients
for (int64_t h = 0; h < h_size; ++h) {
// m_i: max value per row
float m_i = at::vec::reduce_all<float>(