[CPU] Add GPT-OSS model optimization for CPU (#16775)
Co-authored-by: mingfeima <mingfei.ma@intel.com> Co-authored-by: jianan-gu <jianan.gu@intel.com>
This commit is contained in:
co-authored by
mingfeima
jianan-gu
parent
5601b7139d
commit
3ecf2c76ad
@@ -161,9 +161,10 @@ void flash_attn_kernel_impl(
|
||||
}
|
||||
}
|
||||
|
||||
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
|
||||
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale);
|
||||
|
||||
for (int row = 0; row < m_size; ++row) {
|
||||
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
|
||||
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row);
|
||||
}
|
||||
// get value and pack
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ Btmp,
|
||||
@@ -344,8 +345,10 @@ void flash_attn_varlen_kernel_impl(
|
||||
}
|
||||
}
|
||||
|
||||
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
|
||||
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale);
|
||||
for (int row = 0; row < m_size; ++row) {
|
||||
flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply(
|
||||
s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row);
|
||||
}
|
||||
|
||||
// get value and pack
|
||||
pack_vnni2<scalar_t>(
|
||||
|
||||
Reference in New Issue
Block a user