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

[CPU] Add GPT-OSS model optimization for CPU (#16775)

Browse Source 
Co-authored-by: mingfeima <mingfei.ma@intel.com>
Co-authored-by: jianan-gu <jianan.gu@intel.com>
This commit is contained in:
blzheng
2026-05-29 16:05:26 +08:00
committed by GitHub
co-authored by mingfeima jianan-gu
parent 5601b7139d
commit 3ecf2c76ad
35 changed files with 2000 additions and 530 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/decode.cpp
+51 -14
View File
@@ -981,16 +981,20 @@ template <typename scalar_t>
void decode_accumulate_kv_splits(
scalar_t* __restrict__ output,
float* __restrict__ attn_logits,
const scalar_t* __restrict__ sinks_ptr,
int64_t batches,
int64_t num_heads,
int64_t head_size_v,
int64_t num_kv_splits,
int64_t l_stride1,
int64_t l_stride2) {
int64_t l_stride2,
bool has_sink) {
using Vec = at::vec::Vectorized<float>;
// parallel on [batches, num_heads]
at::parallel_for(0, batches * num_heads, 0, [&](int64_t begin, int64_t end) {
int64_t bi{0}, ni{0};
data_index_init(begin, bi, batches, ni, num_heads);
// NB: here we use logits[b][h][0] as acc, since
// for the first kv split (kv_id == 0):
// m_delta = std::exp(-inf) = 0
@@ -1022,8 +1026,12 @@ void decode_accumulate_kv_splits(
s_prime = s_prime * m_delta + e_logic;
m_prime = m_i;
}
if (has_sink) {
s_prime += std::exp(sinks_ptr[ni] - m_prime);
}
copy_stub<scalar_t>(output + i * head_size_v, acc, 1 / s_prime, head_size_v);
// move to the next index
data_index_step(bi, batches, ni, num_heads);
}
});
}
@@ -1039,6 +1047,7 @@ void decode_attention_kernel_impl(
const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
const scalar_t* __restrict__ sinks,
int64_t batches,
int64_t num_heads,
int64_t head_size,
@@ -1055,8 +1064,10 @@ void decode_attention_kernel_impl(
int64_t max_num_reqs,
int64_t max_context_len,
int64_t max_total_num_tokens,
int64_t sliding_window_size,
bool is_cross_attn,
bool has_encoder_lens) {
bool has_encoder_lens,
bool has_sink) {
using Vec = at::vec::Vectorized<float>;
// strides
@@ -1083,6 +1094,10 @@ void decode_attention_kernel_impl(
int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs];
int64_t req_pool_id = req_pool_indices[bs];
int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
if (sliding_window_size > 0 && seq_len_kv > sliding_window_size) {
kv_offset = seq_len_kv - sliding_window_size;
seq_len_kv = sliding_window_size;
}
TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
@@ -1173,7 +1188,7 @@ void decode_attention_kernel_impl(
});
decode_accumulate_kv_splits(
output, attn_logits, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2);
output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink);
} // MHA
template <typename scalar_t, typename index_t, int64_t BLOCK_N>
@@ -1187,6 +1202,7 @@ void decode_attention_mla_kernel_impl(
const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens,
scalar_t* __restrict__ buffer,
const scalar_t* __restrict__ sinks,
int64_t batches,
int64_t num_heads,
int64_t head_size,
@@ -1203,7 +1219,8 @@ void decode_attention_mla_kernel_impl(
int64_t max_num_reqs,
int64_t max_context_len,
int64_t max_total_num_tokens,
int64_t buffer_size_per_thread) {
int64_t buffer_size_per_thread,
bool has_sink) {
using Vec = at::vec::Vectorized<float>;
// block length for heads
@@ -1369,7 +1386,7 @@ void decode_attention_mla_kernel_impl(
});
decode_accumulate_kv_splits(
output, attn_logits, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2);
output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink);
} // MLA
template <typename scalar_t, typename index_t, int64_t BLOCK_N>
@@ -1383,6 +1400,7 @@ void decode_attention_grouped_kernel_impl(
const int64_t* __restrict__ req_pool_indices,
const int64_t* __restrict__ seq_lens,
const int64_t* __restrict__ encoder_lens,
const scalar_t* __restrict__ sinks,
int64_t batches,
int64_t num_heads,
int64_t num_heads_kv,
@@ -1400,8 +1418,10 @@ void decode_attention_grouped_kernel_impl(
int64_t max_num_reqs,
int64_t max_context_len,
int64_t max_total_num_tokens,
int64_t sliding_window_size,
bool is_cross_attn,
bool has_encoder_lens) {
bool has_encoder_lens,
bool has_sink) {
using Vec = at::vec::Vectorized<float>;
// block length for heads
@@ -1447,7 +1467,10 @@ void decode_attention_grouped_kernel_impl(
int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
if (sliding_window_size > 0 && seq_len_kv > sliding_window_size) {
kv_offset = seq_len_kv - sliding_window_size;
seq_len_kv = sliding_window_size;
}
const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits);
const int64_t kv_start = kv_id * SPLIT_SIZE;
const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv);
@@ -1545,7 +1568,7 @@ void decode_attention_grouped_kernel_impl(
});
decode_accumulate_kv_splits(
output, attn_logits, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2);
output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink);
} // GQA/MQA
} // anonymous namespace
@@ -1559,7 +1582,7 @@ void decode_attention_grouped_kernel_impl(
// req_pool_indices: [num_seqs] int64
// seq_lens: [num_seqs] int64
// encoder_lens: [num_seqs] int64 or None
//
// sinks: [num_heads] or None
void decode_attention_cpu(
at::Tensor& query,
at::Tensor& k_buffer,
@@ -1575,7 +1598,9 @@ void decode_attention_cpu(
double sm_scale,
double logit_cap,
bool is_cross_attn,
std::optional<at::Tensor> encoder_lens) {
int64_t sliding_window_size,
std::optional<at::Tensor> encoder_lens,
std::optional<at::Tensor> sinks) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer);
@@ -1642,6 +1667,10 @@ void decode_attention_cpu(
encoder_lens_t = encoder_lens.value();
CHECK_EQ(encoder_lens_t.size(0), num_seqs);
}
bool has_sink = sinks.has_value();
at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, query.options());
CHECK_DIM(1, sinks_tensor);
CHECK_EQ(sinks_tensor.size(0), num_heads);
AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "decode_attention_kernel", [&] {
AT_DISPATCH_INDEX_TYPES(index_dtype, "decode_attention_indices", [&] {
if (key.has_value()) {
@@ -1693,6 +1722,7 @@ void decode_attention_cpu(
req_pool_indices.data_ptr<int64_t>(),
seq_lens.data_ptr<int64_t>(),
encoder_lens_t.data_ptr<int64_t>(),
sinks_tensor.data_ptr<scalar_t>(),
num_seqs,
num_heads,
head_size,
@@ -1709,8 +1739,10 @@ void decode_attention_cpu(
max_num_reqs,
max_context_len,
max_total_num_tokens,
sliding_window_size,
is_cross_attn,
has_encoder_lens);
has_encoder_lens,
has_sink);
} else if (is_mla) {
// MLA
decode_attention_mla_kernel_impl<scalar_t, index_t, BLOCK_N>(
@@ -1723,6 +1755,7 @@ void decode_attention_cpu(
req_pool_indices.data_ptr<int64_t>(),
seq_lens.data_ptr<int64_t>(),
buffer.data_ptr<scalar_t>(),
sinks_tensor.data_ptr<scalar_t>(),
num_seqs,
num_heads,
head_size,
@@ -1739,7 +1772,8 @@ void decode_attention_cpu(
max_num_reqs,
max_context_len,
max_total_num_tokens,
size_per_thread);
size_per_thread,
has_sink);
} else {
// GQA/MQA
decode_attention_grouped_kernel_impl<scalar_t, index_t, BLOCK_N>(
@@ -1752,6 +1786,7 @@ void decode_attention_cpu(
req_pool_indices.data_ptr<int64_t>(),
seq_lens.data_ptr<int64_t>(),
encoder_lens_t.data_ptr<int64_t>(),
sinks_tensor.data_ptr<scalar_t>(),
num_seqs,
num_heads,
num_heads_kv,
@@ -1769,8 +1804,10 @@ void decode_attention_cpu(
max_num_reqs,
max_context_len,
max_total_num_tokens,
sliding_window_size,
is_cross_attn,
has_encoder_lens);
has_encoder_lens,
has_sink);
}
});
});