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/common.h
+7
View File
@@ -352,6 +352,13 @@ inline int get_cache_blocks<at::Float8_e4m3fn>(int chunk_size) {
return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
}
template <>
inline int get_cache_blocks<uint8_t>(int chunk_size) {
// mxfp4 uses bf16 as accumulate type
int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
}
// 2d sequential loop in range : [mb0, mb1), [nb0, nb1)
template <typename T, typename func_t>
inline void loop_2d(int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1, int64_t chunk_size, const func_t& f) {
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);
}
});
});
sgl-kernel/csrc/cpu/extend.cpp
+51 -13
View File
@@ -26,6 +26,7 @@ void extend_attention_kernel_impl(
const index_t* __restrict__ extend_seq_lens,
const index_t* __restrict__ extend_start_loc,
const void* __restrict__ buffer,
const scalar_t* __restrict__ sinks,
int batches,
int num_heads,
int num_heads_kv,
@@ -47,9 +48,11 @@ void extend_attention_kernel_impl(
int max_total_num_tokens,
int max_len_extend,
int buffer_size_per_thread,
int64_t sliding_window_size,
bool is_prefix_skipped,
bool is_cross_attn,
bool has_encoder_lens) {
bool has_encoder_lens,
bool has_sink) {
// strides
const int o_strideM = num_heads * head_size_v;
const int o_strideH = head_size_v;
@@ -69,18 +72,21 @@ void extend_attention_kernel_impl(
data_index_init(begin, bs, batches, head_id, num_heads, mb, MB);
int tid = at::get_thread_num();
// s_i and s_delta: [BLOCK_M, BLOCK_N]
// s_i: [BLOCK_M, BLOCK_N]
float* __restrict__ s_i = reinterpret_cast<float*>((char*)(buffer) + tid * buffer_size_per_thread);
scalar_t* __restrict__ s_delta = reinterpret_cast<scalar_t*>(s_i);
// v_prime: [BLOCK_M, head_size_v]
float* __restrict__ v_prime = s_i + BLOCK_M * BLOCK_N;
// s_delta: [BLOCK_M, BLOCK_N]
scalar_t* __restrict__ s_delta = reinterpret_cast<scalar_t*>(v_prime + BLOCK_M * head_size_v);
// Btmp: [BLOCK_N, max(head_size, head_size_v)]
scalar_t* __restrict__ Btmp = reinterpret_cast<scalar_t*>(v_prime + BLOCK_M * head_size_v);
scalar_t* __restrict__ Btmp = reinterpret_cast<scalar_t*>(s_delta + BLOCK_M * BLOCK_N);
// init Btmp just once for each thread to prevent NaN
fill_stub(Btmp, 0.f, BLOCK_N * ldb_tmp);
fill_stub(s_delta, 0.f, BLOCK_M * BLOCK_N);
alignas(64) float s_prime[BLOCK_M];
alignas(64) float m_prime[BLOCK_M];
@@ -151,11 +157,20 @@ void extend_attention_kernel_impl(
/* B */ Btmp,
/* C */ s_i);
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) {
if (sliding_window_size > 0) {
int last_col = seq_len_prefix + row + m - sliding_window_size + 1;
if (last_col >= n + n_size) {
continue;
}
fill_stub(s_i + row * BLOCK_N, -std::numeric_limits<float>::infinity(), last_col - n);
}
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, index_t>(
pack_vnni2<scalar_t>(
/* dst */ Btmp,
/* src */ v_buffer + head_kv_id * v_strideH,
/* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
@@ -233,8 +248,17 @@ void extend_attention_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) {
if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size - 1 &&
row + m + 1 < n + sliding_window_size + n_size) {
fill_stub(
s_i + row * BLOCK_N, -std::numeric_limits<float>::infinity(), row + m - n - sliding_window_size + 1);
} else if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size) {
continue;
}
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>(
@@ -261,6 +285,9 @@ void extend_attention_kernel_impl(
}
scalar_t* __restrict__ out_ptr = o_extend + (seq_extend_start_loc + m) * o_strideM + head_id * o_strideH;
for (int row = 0; row < m_size; ++row) {
if (has_sink) {
s_prime[row] += std::exp(sinks[head_id] - m_prime[row]);
}
float s = 1 / s_prime[row];
copy_stub<scalar_t>(out_ptr + row * o_strideM, v_prime + row * head_size_v, s, head_size_v);
}
@@ -280,6 +307,7 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
const int size_per_thread =
/* s_i */ BLOCK_M * BLOCK_N * sizeof(float) +
/* v_prime */ BLOCK_M * head_size_v * sizeof(float) +
/* s_delta */ BLOCK_M * BLOCK_N * sizeof(uint16_t) +
/* Btmp */ BLOCK_N * std::max(head_size, head_size_v) * sizeof(uint16_t);
buffer.resize_({num_threads, size_per_thread});
@@ -304,6 +332,7 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
extend_seq_lens.data_ptr<index_t>(), \
extend_start_loc.data_ptr<index_t>(), \
buffer.data_ptr(), \
sinks_tensor.data_ptr<scalar_t>(), \
num_seqs, \
num_heads, \
num_heads_kv, \
@@ -325,9 +354,11 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
max_total_num_tokens, \
max_len_extend, \
sz, \
sliding_window_size, \
is_prefix_skipped, \
is_cross_attn, \
has_encoder_lens); \
has_encoder_lens, \
has_sink); \
} while (0)
// q_extend, k_extend, v_extend, o_extend: contiguous tensors
@@ -344,8 +375,8 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
// seq_lens: [num_seqs] int64
// extend_seq_lens: [num_seqs]
// extend_start_loc: [num_seqs]
// encoder_lens: [num_seqs] int64
//
// encoder_lens: [num_seqs] int64 or None
// sinks: [num_heads] or None
void extend_attention_cpu(
at::Tensor& q_extend,
const std::optional<at::Tensor>& k_extend_opt,
@@ -362,7 +393,9 @@ void extend_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) {
if (!is_cross_attn) {
TORCH_CHECK(
k_extend_opt.has_value() && v_extend_opt.has_value(),
@@ -443,6 +476,11 @@ void extend_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}, q_extend.options());
CHECK_DIM(1, sinks_tensor);
CHECK_EQ(sinks_tensor.size(0), num_heads);
AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] {
AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] {
if (max_len_extend <= 256) {
sgl-kernel/csrc/cpu/flash_attn.cpp
+8 -5
View File
@@ -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>(
sgl-kernel/csrc/cpu/flash_attn.h
+84 -86
View File
@@ -100,45 +100,43 @@ struct flash_attn_softmax {
int n_size,
int padded_n_size,
int head_size_v,
const float sm_scale) {
const float sm_scale,
int row) {
using Vec = at::vec::Vectorized<float>;
const Vec scale_vec = Vec(sm_scale);
float* s_delta = s_i;
for (int row = 0; row < m_size; ++row) {
// s_i <- s_i * scale
at::vec::map<float>(
[scale_vec](Vec x) { return x * scale_vec; }, s_i + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
// s_i <- s_i * scale
at::vec::map<float>([scale_vec](Vec x) { return x * scale_vec; }, s_i + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
// m_i: max value per row
float m_i = at::vec::reduce_all<float>(
[](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + row * BLOCK_N, n_size);
m_i = std::max(m_i, m_prime[row]);
// m_i: max value per row
float m_i =
at::vec::reduce_all<float>([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + row * BLOCK_N, n_size);
m_i = std::max(m_i, m_prime[row]);
// m_delta <- exp(m' - m_i)
float m_delta = std::exp(m_prime[row] - m_i);
// m_delta <- exp(m' - m_i)
float m_delta = std::exp(m_prime[row] - m_i);
// s_delta <- exp(s_i - m_i)
at::vec::map<float>(
[m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
// s_delta <- exp(s_i - m_i)
at::vec::map<float>(
[m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
// s' <- s' * m_delta + sum(s_delta)
s_prime[row] *= m_delta;
s_prime[row] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + row * BLOCK_N, n_size);
// s' <- s' * m_delta + sum(s_delta)
s_prime[row] *= m_delta;
s_prime[row] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + row * BLOCK_N, n_size);
m_prime[row] = m_i;
m_prime[row] = m_i;
// v' <- v' * m_delta
at::vec::map<float>(
[m_delta](Vec x) { return x * Vec(m_delta); },
v_prime + row * head_size_v,
v_prime + row * head_size_v,
head_size_v);
// v' <- v' * m_delta
at::vec::map<float>(
[m_delta](Vec x) { return x * Vec(m_delta); },
v_prime + row * head_size_v,
v_prime + row * head_size_v,
head_size_v);
// Keep s_delta row-major for the following brgemm(P @ V), and only
// convert the columns that brgemm will consume.
fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size);
copy_stub<scalar_t>(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N, 1.f, padded_n_size);
}
// Keep s_delta row-major for the following brgemm(P @ V), and only
// convert the columns that brgemm will consume.
fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size);
copy_stub<scalar_t>(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N, 1.f, padded_n_size);
}
};
@@ -155,7 +153,8 @@ struct flash_attn_softmax<at::BFloat16, BLOCK_M, BLOCK_N> {
int n_size,
int padded_n_size,
int head_size_v,
const float sm_scale) {
const float sm_scale,
int row) {
float* s_delta = s_i;
const __m512 vscale = _mm512_set1_ps(sm_scale);
@@ -175,72 +174,71 @@ struct flash_attn_softmax<at::BFloat16, BLOCK_M, BLOCK_N> {
const __m512 vneg_inf = _mm512_set1_ps(NEG_INF);
for (int m = 0; m < m_size; ++m) {
vmax = vneg_inf;
int m = row;
vmax = vneg_inf;
// s_i <- s_i * scale
int n = 0;
for (; n <= n_size - 16; n += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
vmax = _mm512_max_ps(va, vmax);
}
if (n_remainder > 0) {
va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
vmax = _mm512_max_ps(va, vmax);
}
// s_i <- s_i * scale
int n = 0;
for (; n <= n_size - 16; n += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
vmax = _mm512_max_ps(va, vmax);
}
if (n_remainder > 0) {
va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
vmax = _mm512_max_ps(va, vmax);
}
// 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_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)
float m_delta = std::exp(m_prime[m] - m_i);
// m_delta <- exp(m' - m_i)
float m_delta = std::exp(m_prime[m] - m_i);
// s_delta <- exp(s_i - m_i)
vsum = _mm512_setzero_ps();
for (n = 0; n <= n_size - 16; n += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
vsum = _mm512_add_ps(vsum, va);
// s_delta <- exp(s_i - m_i)
vsum = _mm512_setzero_ps();
for (n = 0; n <= n_size - 16; n += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
vsum = _mm512_add_ps(vsum, va);
vb = (__m256i)(_mm512_cvtneps_pbh(va));
_mm256_storeu_si256(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vb);
}
if (n_remainder > 0) {
va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
vsum = _mm512_add_ps(vsum, va);
vb = (__m256i)(_mm512_cvtneps_pbh(va));
_mm256_storeu_si256(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vb);
}
if (n_remainder > 0) {
va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
vsum = _mm512_add_ps(vsum, va);
vb = (__m256i)(_mm512_cvtneps_pbh(va));
_mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vmask, vb);
}
vb = (__m256i)(_mm512_cvtneps_pbh(va));
_mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vmask, vb);
}
// s' <- s' * m_delta + sum(s_delta)
s_prime[m] *= m_delta;
s_prime[m] += _mm512_reduce_add_ps(vsum);
// s' <- s' * m_delta + sum(s_delta)
s_prime[m] *= m_delta;
s_prime[m] += _mm512_reduce_add_ps(vsum);
m_prime[m] = m_i;
m_prime[m] = m_i;
// pad s_delta with 0, pad_size range from [0, 32)
int pad_size = padded_n_size - n_size;
if (pad_size > 0) {
const __m512i vzero = _mm512_setzero_si512();
__mmask32 vmask2 = (1ULL << pad_size) - 1;
_mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(s_delta2 + m * BLOCK_N + n_size), vmask2, vzero);
}
// pad s_delta with 0, pad_size range from [0, 32)
int pad_size = padded_n_size - n_size;
if (pad_size > 0) {
const __m512i vzero = _mm512_setzero_si512();
__mmask32 vmask2 = (1ULL << pad_size) - 1;
_mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(s_delta2 + m * BLOCK_N + n_size), vmask2, vzero);
}
// v' <- v' * m_delta
vmdelta = _mm512_set1_ps(m_delta);
int k = 0;
for (; k <= head_size_v - 16; k += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(v_prime + m * head_size_v + k), vmdelta);
_mm512_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), va);
}
if (v_remainder > 0) {
va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask1, v_prime + m * head_size_v + k), vmdelta);
_mm512_mask_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), vmask1, va);
}
// v' <- v' * m_delta
vmdelta = _mm512_set1_ps(m_delta);
int k = 0;
for (; k <= head_size_v - 16; k += 16) {
va = _mm512_mul_ps(_mm512_loadu_ps(v_prime + m * head_size_v + k), vmdelta);
_mm512_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), va);
}
if (v_remainder > 0) {
va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask1, v_prime + m * head_size_v + k), vmdelta);
_mm512_mask_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), vmask1, va);
}
}
};
sgl-kernel/csrc/cpu/gemm.h
+53 -9
View File
@@ -67,7 +67,13 @@ inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
return use_int8_w8a8 ? K + sizeof(int32_t) : K;
}
enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3 };
enum class CPUActMethod : int {
silu_and_mul = 0,
swiglu = 1,
gelu_and_mul = 2,
};
enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };
constexpr bool operator==(CPUQuantMethod a, int64_t b) {
return static_cast<int64_t>(a) == b;
@@ -87,6 +93,17 @@ constexpr bool operator==(int64_t a, CPUQuantAlgo b) {
return a == static_cast<int64_t>(b);
}
inline int64_t get_row_size(CPUQuantMethod quant, int64_t K) {
switch (quant) {
case CPUQuantMethod::INT8_W8A8:
return K + sizeof(int32_t);
case CPUQuantMethod::MXFP4:
return K >> 1;
default:
return K;
}
}
inline int64_t get_4bit_block_k_size(int64_t group_size) {
return group_size > 128 ? 128 : group_size;
}
@@ -124,9 +141,9 @@ void fused_experts_int8_kernel_impl(
int64_t topk,
int64_t num_tokens_post_pad);
// moe implementations for fp8 w8a16
template <typename scalar_t>
void fused_experts_fp8_kernel_impl(
// moe implementations for fp8 w8a16 and mxfp4
template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
void fused_experts_fp_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic0,
scalar_t* __restrict__ ic1,
@@ -135,10 +152,12 @@ void fused_experts_fp8_kernel_impl(
scalar_t* __restrict__ B_tmp,
float* __restrict__ C_tmp,
const scalar_t* __restrict__ input,
const at::Float8_e4m3fn* __restrict__ packed_w1,
const at::Float8_e4m3fn* __restrict__ packed_w2,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
const packed_t* __restrict__ packed_w1,
const packed_t* __restrict__ packed_w2,
const float* __restrict__ w1_bias,
const float* __restrict__ w2_bias,
const param_t* __restrict__ w1s,
const param_t* __restrict__ w2s,
int64_t block_size_N,
int64_t block_size_K,
const float* __restrict__ topk_weights,
@@ -150,7 +169,11 @@ void fused_experts_fp8_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad);
int64_t num_tokens_post_pad,
float alpha,
float limit,
CPUActMethod act_func,
bool with_bias);
// shared expert implementation for int8 w8a8
template <typename scalar_t>
@@ -261,6 +284,7 @@ void tinygemm_kernel(
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const float* __restrict__ Bbias,
const float* __restrict__ scale,
int64_t M,
int64_t N,
@@ -289,6 +313,26 @@ void tinygemm_kernel(
int64_t ldc,
bool brg);
// mxfp4
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const uint8_t* __restrict__ B,
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const float* __restrict__ Bbias,
const uint8_t* __restrict__ scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg,
int64_t block_size_K,
bool do_unpack = true);
template <typename scalar_t>
void tinygemm_kernel(
scalar_t* C,
sgl-kernel/csrc/cpu/gemm_fp8.cpp
+83
View File
@@ -62,6 +62,23 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
}
}
template <>
inline void
copy_add_stub(float* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = fVec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data = fVec::loadu(input + d) + fVec::loadu(bias + d);
data.store(out + d);
}
for (; d < size; ++d) {
out[d] = input[d] + bias[d];
}
}
inline void unpack_B(
at::BFloat16* __restrict__ Btmp,
const at::Float8_e4m3fn* __restrict__ packed_B,
@@ -913,6 +930,7 @@ void tinygemm_kernel(
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const float* __restrict__ Bbias,
const float* __restrict__ scale,
int64_t M,
int64_t N,
@@ -923,6 +941,11 @@ void tinygemm_kernel(
bool brg,
int64_t block_size_K,
bool do_unpack) {
if (Bbias != nullptr) {
tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
return;
}
tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
}
@@ -952,6 +975,7 @@ void tinygemm_kernel(
scalar_t* __restrict__ C,
scalar_t* __restrict__ Btmp,
float* __restrict__ Ctmp,
const float* __restrict__ Bbias,
const uint8_t* __restrict__ scale,
int64_t M,
int64_t N,
@@ -962,10 +986,68 @@ void tinygemm_kernel(
bool brg,
int64_t block_size_K,
bool do_unpack) {
if (Bbias != nullptr) {
tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
return;
}
tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
}
// tinygemm interface
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const at::Float8_e4m3fn* __restrict__ B,
float* __restrict__ C,
scalar_t* __restrict__ Btmp,
const float* __restrict__ Bbias,
const float* __restrict__ scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg,
int64_t block_size_K,
bool do_unpack) {
if (Bbias != nullptr) {
tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
return;
}
tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
}
template <typename scalar_t>
void tinygemm_kernel(
const scalar_t* __restrict__ A,
const uint8_t* __restrict__ B,
float* __restrict__ C,
scalar_t* __restrict__ Btmp,
const float* __restrict__ Bbias,
const uint8_t* __restrict__ scale,
int64_t M,
int64_t N,
int64_t K,
int64_t lda,
int64_t ldb,
int64_t ldc,
bool brg,
int64_t block_size_K,
bool do_unpack) {
if (Bbias != nullptr) {
tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
return;
}
tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
}
#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE_A, TYPE_B, TYPE_S) \
template void tinygemm_kernel<TYPE_A>( \
const TYPE_A* __restrict__ A, \
@@ -973,6 +1055,7 @@ void tinygemm_kernel(
TYPE_A* __restrict__ C, \
TYPE_A* __restrict__ Btmp, \
float* __restrict__ Ctmp, \
const float* __restrict__ Bbias, \
const TYPE_S* __restrict__ scale, \
int64_t M, \
int64_t N, \
sgl-kernel/csrc/cpu/moe.cpp
+217 -37
View File
@@ -158,6 +158,56 @@ inline void silu_and_mul(
}
}
template <typename scalar_t, int BLOCK_N>
inline void clamp_sigmoid_and_mul(
scalar_t* __restrict__ output,
const float* __restrict__ input0,
int64_t m_size,
int64_t N,
const float alpha,
const float limit,
int64_t offset) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
const fVec zero = fVec(0.f);
const fVec limit_v = fVec(limit);
const fVec nlimit_v = fVec(-limit);
const fVec alpha_v = fVec(alpha);
// no remainder
for (int64_t m = 0; m < m_size; ++m) {
scalar_t* __restrict__ out = output + m * N;
const float* __restrict__ cur_ptr = input0 + m * BLOCK_N;
for (int64_t d = 0; d < BLOCK_N; d += bVec::size()) {
float tmp_glu0[fVec::size()]; // 16
float tmp_linear0[fVec::size()]; // 16
// interleaved: x[2i] = glu, x[2i+1] = linear
for (int j = 0; j < fVec::size(); ++j) {
// x0 [0,2,..30]
tmp_glu0[j] = cur_ptr[d + j * 2];
// y0 [1,3,...31]
tmp_linear0[j] = cur_ptr[d + j * 2 + 1];
}
fVec x0 = fVec::loadu(tmp_glu0);
fVec y0 = fVec::loadu(tmp_linear0);
// clamp
x0 = at::vec::minimum(x0, limit_v);
y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
// x * sigmoid(x * alpha)
x0 = x0 / (one + (x0 * alpha_v).neg().exp_u20());
// (y + 1) * x
y0 = y0 + one;
x0 = x0 * y0;
// // convert
convert_from_float_and_store<scalar_t>(out + d / 2 + offset, x0);
}
}
}
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
struct tinygemm_kernel_nn2 {
static inline void apply(
@@ -450,6 +500,8 @@ void fused_experts_kernel_impl(
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ packed_w1,
const scalar_t* __restrict__ packed_w2,
const float* __restrict__ w1_bias,
const float* __restrict__ w2_bias,
const float* __restrict__ topk_weights,
const int32_t* __restrict__ sorted_ids,
const int32_t* __restrict__ expert_ids,
@@ -459,7 +511,11 @@ void fused_experts_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad) {
int64_t num_tokens_post_pad,
float alpha,
float limit,
CPUActMethod act_func,
bool with_bias) {
// handle 2 tiles per block
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
@@ -494,6 +550,8 @@ void fused_experts_kernel_impl(
int32_t expert_id = expert_ids[mb];
const scalar_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n;
const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;
const float* __restrict__ B0_bias = w1_bias + expert_id * 2 * N + nb_upper * BLOCK_N;
const float* __restrict__ B1_bias = w1_bias + expert_id * 2 * N + nb_lower * BLOCK_N;
int64_t m_size = offsets[mb + 1] - offsets[mb];
@@ -533,23 +591,58 @@ void fused_experts_kernel_impl(
/* B */ B1,
/* C */ C1);
// 1.d silu and mul
const int64_t offset = offsets[mb];
silu_and_mul<scalar_t, BLOCK_N>(ic1 + offset * N + nb * BLOCK_N, C0, C1, m_size, N);
} else {
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
const int64_t offset = offsets[mb];
tinygemm_kernel(
/* A */ A,
/* B0 */ B0,
/* B1 */ B1,
/* C */ ic1 + offset * N + nb * BLOCK_N,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ N);
if (act_func == CPUActMethod::swiglu) {
tinygemm_kernel(
/* A */ A,
/* B */ B0,
/* C */ C0,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ BLOCK_N);
tinygemm_kernel(
/* A */ A,
/* B */ B1,
/* C */ C1,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ BLOCK_N);
} else {
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
tinygemm_kernel(
/* A */ A,
/* B0 */ B0,
/* B1 */ B1,
/* C */ ic1 + offset * N + nb * BLOCK_N,
/* M */ m_size,
/* N */ n_size,
/* K */ K,
/* lda */ K,
/* ldb */ n_size,
/* ldc */ N);
}
}
if (with_bias) {
for (int64_t m = 0; m < m_size; ++m) {
add_bias_stub(C0 + m * BLOCK_N, B0_bias, n_size);
add_bias_stub(C1 + m * BLOCK_N, B1_bias, n_size);
}
}
// 1.d silu and mul
const int64_t offset = offsets[mb];
if (act_func == CPUActMethod::silu_and_mul && use_brgemm) {
silu_and_mul<scalar_t, BLOCK_N>(ic1 + offset * N + nb * BLOCK_N, C0, C1, m_size, N);
} else if (act_func == CPUActMethod::swiglu) {
clamp_sigmoid_and_mul<scalar_t, BLOCK_N>(ic1 + offset * N, C0, m_size, N, alpha, limit, 0 + nb * BLOCK_N / 2);
clamp_sigmoid_and_mul<scalar_t, BLOCK_N>(
ic1 + offset * N, C1, m_size, N, alpha, limit, N / 2 + nb * BLOCK_N / 2);
}
});
@@ -586,6 +679,7 @@ void fused_experts_kernel_impl(
// B shape [IC, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const scalar_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
const float* __restrict__ B_bias = w2_bias + expert_id * OC + nb * BLOCK_N;
// 2.a gemm: C = A @ B
if (use_brgemm) {
@@ -613,6 +707,11 @@ void fused_experts_kernel_impl(
/* ldc */ BLOCK_N);
}
if (with_bias) {
for (int64_t m = 0; m < m_size; ++m) {
add_bias_stub(C + m * BLOCK_N, B_bias, n_size);
}
}
// 2.b copy from C to ic2 in original order
// and also mul topk_weights in float32
for (int64_t m = 0; m < m_size; ++m) {
@@ -804,21 +903,38 @@ void shared_expert_kernel_impl(
} // anonymous namespace
// common checks
template <CPUQuantMethod quant>
static inline void check_moe_scales(
bool use_int8_w8a8,
bool use_fp8_w8a16,
const std::optional<at::Tensor>& w1_scale,
const std::optional<at::Tensor>& w2_scale,
const std::optional<std::vector<int64_t>> block_size) {
if (use_int8_w8a8) {
if constexpr (quant == CPUQuantMethod::INT8_W8A8) {
TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for int8 w8a8.");
TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for int8 w8a8.");
}
if (use_fp8_w8a16) {
} else if constexpr (quant == CPUQuantMethod::FP8_W8A16) {
TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a16.");
TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for fp8 w8a16.");
TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a16.");
TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2.");
} else if constexpr (quant == CPUQuantMethod::MXFP4) {
TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for mxfp4.");
TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4.");
TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8.");
TORCH_CHECK(w2_scale.value().scalar_type() == at::kByte, "expect w2_scale to be uint8.");
}
}
static inline void check_moe_scales(
int64_t moe_comp_method,
const std::optional<at::Tensor>& w1_scale,
const std::optional<at::Tensor>& w2_scale,
const std::optional<std::vector<int64_t>> block_size) {
if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
check_moe_scales<CPUQuantMethod::INT8_W8A8>(w1_scale, w2_scale, block_size);
} else if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
check_moe_scales<CPUQuantMethod::FP8_W8A16>(w1_scale, w2_scale, block_size);
} else if (moe_comp_method == CPUQuantMethod::MXFP4) {
check_moe_scales<CPUQuantMethod::MXFP4>(w1_scale, w2_scale, block_size);
}
}
@@ -834,8 +950,8 @@ static inline void check_moe_scales(
TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K))
// hidden_states: [M, K]
// w1: [E, 2N, K]
// w2: [E, K, N]
// w1: [E, 2N, K] or [E, 2N, K / 2] for uint8
// w2: [E, K, N] or [E, K, N / 2] for uint8
// topk_weights: [M, topk]
// topk_ids: [M, topk] (int32_t)
//
@@ -853,6 +969,10 @@ at::Tensor fused_experts_cpu(
const std::optional<at::Tensor>& w1_zero,
const std::optional<at::Tensor>& w2_zero,
const std::optional<std::vector<int64_t>> block_size,
const std::optional<at::Tensor>& w1_bias,
const std::optional<at::Tensor>& w2_bias,
const std::optional<double>& alpha,
const std::optional<double>& limit,
bool is_vnni) {
auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
@@ -895,8 +1015,8 @@ at::Tensor fused_experts_cpu(
int64_t topk = topk_weights_.size(1);
// we use int32_t compensation for int8 w8a8
int64_t packed_K = get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8);
int64_t packed_N = get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8);
int64_t packed_K = get_row_size(static_cast<CPUQuantMethod>(moe_comp_method), K);
int64_t packed_N = get_row_size(static_cast<CPUQuantMethod>(moe_comp_method), N);
// check weight shapes
CHECK_EQ(w2.size(0), E);
@@ -906,12 +1026,7 @@ at::Tensor fused_experts_cpu(
CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
}
// check scales
check_moe_scales(
moe_comp_method == CPUQuantMethod::INT8_W8A8,
moe_comp_method == CPUQuantMethod::FP8_W8A16,
w1_scale,
w2_scale,
block_size);
check_moe_scales(moe_comp_method, w1_scale, w2_scale, block_size);
at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
@@ -963,7 +1078,7 @@ at::Tensor fused_experts_cpu(
// 5. Aq_tmp : [M, K] or [M * topk, N]
// 6. As_tmp : [M * topk]
//
// for fp8 w8a16:
// for fp8 w8a16 and mxfp4:
// 7. intermediate_cache0 : [M * topk, 2N]
// 8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
//
@@ -976,7 +1091,7 @@ at::Tensor fused_experts_cpu(
if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
}
if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4) {
buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2;
}
if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
@@ -1029,9 +1144,11 @@ at::Tensor fused_experts_cpu(
float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
bool with_bias = w1_bias.has_value();
auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
CHECK_MOE_SCALES_FP8(1, 2);
fused_experts_fp8_kernel_impl(
fused_experts_fp_kernel_impl<scalar_t, at::Float8_e4m3fn, float, false>(
out_hidden_states.data_ptr<scalar_t>(),
intermediate_cache0,
intermediate_cache1,
@@ -1042,6 +1159,8 @@ at::Tensor fused_experts_cpu(
hidden_states.data_ptr<scalar_t>(),
packed_w1.data_ptr<at::Float8_e4m3fn>(),
packed_w2.data_ptr<at::Float8_e4m3fn>(),
with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
w1s.data_ptr<float>(),
w2s.data_ptr<float>(),
block_size_N,
@@ -1055,7 +1174,56 @@ at::Tensor fused_experts_cpu(
K,
E,
topk,
num_tokens_post_pad);
num_tokens_post_pad,
alpha.has_value() ? float(alpha.value()) : 0,
limit.has_value() ? float(limit.value()) : 0,
act_func,
with_bias);
} else if (moe_comp_method == CPUQuantMethod::MXFP4) {
scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K));
float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
bool with_bias = w1_bias.has_value();
auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
// mxfp4 supports only group size of 32 (2^5)
constexpr int64_t group_size = 32;
auto w1s = w1_scale.value();
auto w2s = w2_scale.value();
TORCH_CHECK(w1s.numel(), E * 2 * N * K >> 5);
TORCH_CHECK(w2s.numel(), E * K * N >> 5);
fused_experts_fp_kernel_impl<scalar_t, uint8_t, uint8_t, true>(
out_hidden_states.data_ptr<scalar_t>(),
intermediate_cache0,
intermediate_cache1,
intermediate_cache2,
A_tmp,
B_tmp,
C_tmp,
hidden_states.data_ptr<scalar_t>(),
packed_w1.data_ptr<uint8_t>(),
packed_w2.data_ptr<uint8_t>(),
with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
w1s.data_ptr<uint8_t>(),
w2s.data_ptr<uint8_t>(),
/*block_size_N*/ 1,
/*block_size_K*/ group_size,
topk_weights_.data_ptr<float>(),
sorted_ids,
expert_ids,
offsets,
M,
N,
K,
E,
topk,
num_tokens_post_pad,
alpha.has_value() ? float(alpha.value()) : 0,
limit.has_value() ? float(limit.value()) : 0,
act_func,
with_bias);
} else if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
@@ -1101,6 +1269,8 @@ at::Tensor fused_experts_cpu(
} else {
scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K;
float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
bool with_bias = w1_bias.has_value();
auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
fused_experts_kernel_impl<scalar_t>(
out_hidden_states.data_ptr<scalar_t>(),
@@ -1111,6 +1281,8 @@ at::Tensor fused_experts_cpu(
hidden_states.data_ptr<scalar_t>(),
packed_w1.data_ptr<scalar_t>(),
packed_w2.data_ptr<scalar_t>(),
with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
topk_weights_.data_ptr<float>(),
sorted_ids,
expert_ids,
@@ -1120,7 +1292,11 @@ at::Tensor fused_experts_cpu(
K,
E,
topk,
num_tokens_post_pad);
num_tokens_post_pad,
alpha.has_value() ? float(alpha.value()) : 0,
limit.has_value() ? float(limit.value()) : 0,
act_func,
with_bias);
}
});
return out_hidden_states;
@@ -1183,7 +1359,11 @@ at::Tensor shared_expert_cpu(
CHECK_EQ(packed_w2.size(1), packed_N);
// check scales
check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size);
if (use_int8_w8a8) {
check_moe_scales<CPUQuantMethod::INT8_W8A8>(w1_scale, w2_scale, block_size);
} else if (use_fp8_w8a16) {
check_moe_scales<CPUQuantMethod::FP8_W8A16>(w1_scale, w2_scale, block_size);
}
at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
sgl-kernel/csrc/cpu/moe.h
+106
View File
@@ -171,3 +171,109 @@ inline void silu_and_mul_stub(
out_vec.store(out + d);
}
}
template <typename scalar_t>
inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
const fVec weight_vec = fVec(weight);
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data0 = fVec::loadu(input + d) * weight_vec;
fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] * weight);
}
}
// input = input + input2
inline void add_bias_stub(float* __restrict__ input, const float* __restrict__ input2, int64_t size) {
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = fVec::size();
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec x_fvec = fVec::loadu(input + d);
fVec y_fvec = fVec::loadu(input2 + d);
x_fvec = x_fvec + y_fvec;
x_fvec.store(input + d);
}
for (; d < size; ++d) {
input[d] = input[d] + input2[d];
}
}
template <typename scalar_t>
inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
const fVec weight_vec = fVec(weight);
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
bVec x = bVec::loadu(input + d);
fVec x0, x1;
std::tie(x0, x1) = at::vec::convert_to_float(x);
x0 = x0 * weight_vec;
x1 = x1 * weight_vec;
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
out[d] = static_cast<scalar_t>(input[d] * weight);
}
}
template <typename scalar_t>
inline void clamp_sigmoid_and_mul_stub(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
int64_t size,
const float alpha,
const float limit) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
const fVec zero = fVec(0.f);
const fVec limit_v = fVec(limit);
const fVec nlimit_v = fVec(-limit);
const fVec alpha_v = fVec(alpha);
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
bVec x = bVec::loadu(input + d);
fVec x0_, y0_;
std::tie(x0_, y0_) = at::vec::convert_to_float(x);
float tmp_buffer[fVec::size() * 2]; // 32
float tmp_glu[fVec::size()]; // 16
float tmp_linear[fVec::size()]; // 16
x0_.store(tmp_buffer);
y0_.store(tmp_buffer + fVec::size());
// interleaved: x[2i] = glu, x[2i+1] = linear
for (int j = 0; j < fVec::size(); ++j) {
// x0 [0,2,..30]
tmp_glu[j] = tmp_buffer[j * 2];
// y0 [1,3,...31]
tmp_linear[j] = tmp_buffer[j * 2 + 1];
}
fVec x0 = fVec::loadu(tmp_glu);
fVec y0 = fVec::loadu(tmp_linear);
// clamp
x0 = at::vec::minimum(x0, limit_v);
y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
// x * sigmoid(x * alpha)
x0 = x0 / (one + (x0 * alpha_v).neg().exp_u20());
// (y + 1) * x
y0 = y0 + one;
x0 = x0 * y0;
convert_from_float_and_store<scalar_t>(out + d / 2, x0);
}
}
sgl-kernel/csrc/cpu/moe_fp8.cpp
+91 -54
View File
@@ -2,8 +2,8 @@
#include "gemm.h"
#include "moe.h"
template <typename scalar_t>
void fused_experts_fp8_kernel_impl(
template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
void fused_experts_fp_kernel_impl(
scalar_t* __restrict__ output,
scalar_t* __restrict__ ic0,
scalar_t* __restrict__ ic1,
@@ -12,10 +12,12 @@ void fused_experts_fp8_kernel_impl(
scalar_t* __restrict__ B_tmp,
float* __restrict__ C_tmp,
const scalar_t* __restrict__ input,
const at::Float8_e4m3fn* __restrict__ packed_w1,
const at::Float8_e4m3fn* __restrict__ packed_w2,
const float* __restrict__ w1s,
const float* __restrict__ w2s,
const packed_t* __restrict__ packed_w1,
const packed_t* __restrict__ packed_w2,
const float* __restrict__ w1_bias,
const float* __restrict__ w2_bias,
const param_t* __restrict__ w1s,
const param_t* __restrict__ w2s,
int64_t block_size_N,
int64_t block_size_K,
const float* __restrict__ topk_weights,
@@ -27,7 +29,11 @@ void fused_experts_fp8_kernel_impl(
int64_t K,
int64_t E,
int64_t topk,
int64_t num_tokens_post_pad) {
int64_t num_tokens_post_pad,
float alpha,
float limit,
CPUActMethod act_func,
bool with_bias) {
constexpr int64_t BLOCK_M = block_size_m();
constexpr int64_t BLOCK_N = block_size_n();
@@ -37,12 +43,20 @@ void fused_experts_fp8_kernel_impl(
int64_t scale_size_N = div_up(2 * N, block_size_N);
int64_t scale_size_K = div_up(K, block_size_K);
int64_t blocks_n_per_group = block_size_N / BLOCK_N;
std::function<int64_t(int64_t)> scale_offset_per_block;
if constexpr (is_mxfp4) {
scale_offset_per_block = [&](int64_t a) { return a * BLOCK_N; };
} else {
scale_offset_per_block = [&](int64_t a) { return a / blocks_n_per_group; };
}
const int64_t stride_e = 2 * N * K;
const int64_t stride_n = K;
const int64_t packed_K = get_row_size<packed_t>(K);
const int64_t stride_e = 2 * N * packed_K;
const int64_t stride_n = packed_K;
int64_t avg_M = std::max(int64_t(1), M * topk / E);
const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(avg_M);
const bool use_brgemm = can_use_brgemm<packed_t>(avg_M);
int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
@@ -52,14 +66,15 @@ void fused_experts_fp8_kernel_impl(
int tid = get_thread_num();
scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
// B shape [K, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
const float* __restrict__ Bs =
w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
const packed_t* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
const param_t* __restrict__ Bs =
w1s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
const float* __restrict__ B_bias = with_bias ? w1_bias + expert_id * 2 * N + nb * BLOCK_N : nullptr;
// do unpacking for the first row or a new expert
int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
@@ -83,6 +98,7 @@ void fused_experts_fp8_kernel_impl(
/* C */ ic0 + offset * 2 * N + nb * BLOCK_N,
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
/* Bbias */ B_bias,
/* scale */ Bs,
/* M */ m_size,
/* N */ n_size,
@@ -101,12 +117,20 @@ void fused_experts_fp8_kernel_impl(
});
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
if (act_func == CPUActMethod::silu_and_mul) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
}
});
} else if (act_func == CPUActMethod::swiglu) {
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
for (int64_t m = begin; m < end; ++m) {
clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N, alpha, limit);
clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N, alpha, limit);
}
});
}
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
// w2 : [E, K, N] as [E, OC, IC]
const int64_t OC = K; // rename K as OC
@@ -115,15 +139,16 @@ void fused_experts_fp8_kernel_impl(
const int64_t NB2 = div_up(OC, BLOCK_N);
scale_size_N = div_up(K, block_size_N);
scale_size_K = div_up(N, block_size_K);
const int64_t stride_e2 = OC * IC;
const int64_t stride_oc = IC;
const int64_t packed_IC = get_row_size<packed_t>(IC);
const int64_t stride_e2 = OC * packed_IC;
const int64_t stride_oc = packed_IC;
// parallel on [MB2, NB2]
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
int tid = get_thread_num();
alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
int64_t m_size = offsets[mb + 1] - offsets[mb];
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
@@ -134,9 +159,10 @@ void fused_experts_fp8_kernel_impl(
// B shape [IC, n_size] in vnni format
int32_t expert_id = expert_ids[mb];
const at::Float8_e4m3fn* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
const float* __restrict__ Bs =
w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
const packed_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
const param_t* __restrict__ Bs =
w2s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
const float* __restrict__ B_bias = with_bias ? w2_bias + expert_id * OC + nb * BLOCK_N : nullptr;
// do unpacking for the first row or a new expert
int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
@@ -148,6 +174,7 @@ void fused_experts_fp8_kernel_impl(
/* C */ C,
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
/* Bbias */ B_bias,
/* scale */ Bs,
/* M */ m_size,
/* N */ n_size,
@@ -182,35 +209,43 @@ void fused_experts_fp8_kernel_impl(
});
}
#define INSTANTIATE_MOE_FP8_TEMPLATE(TYPE) \
template void fused_experts_fp8_kernel_impl<TYPE>( \
TYPE* __restrict__ output, \
TYPE* __restrict__ ic0, \
TYPE* __restrict__ ic1, \
TYPE* __restrict__ ic2, \
TYPE* __restrict__ A_tmp, \
TYPE* __restrict__ B_tmp, \
float* __restrict__ C_tmp, \
const TYPE* __restrict__ input, \
const at::Float8_e4m3fn* __restrict__ packed_w1, \
const at::Float8_e4m3fn* __restrict__ packed_w2, \
const float* __restrict__ w1s, \
const float* __restrict__ w2s, \
int64_t block_size_N, \
int64_t block_size_K, \
const float* __restrict__ topk_weights, \
const int32_t* __restrict__ sorted_ids, \
const int32_t* __restrict__ expert_ids, \
const int32_t* __restrict__ offsets, \
int64_t M, \
int64_t N, \
int64_t K, \
int64_t E, \
int64_t topk, \
int64_t num_tokens_post_pad)
#define INSTANTIATE_MOE_FP_TEMPLATE(TYPE1, TYPE2, TYPE3, IS_MXFP4) \
template void fused_experts_fp_kernel_impl<TYPE1, TYPE2, TYPE3, IS_MXFP4>( \
TYPE1* __restrict__ output, \
TYPE1* __restrict__ ic0, \
TYPE1* __restrict__ ic1, \
TYPE1* __restrict__ ic2, \
TYPE1* __restrict__ A_tmp, \
TYPE1* __restrict__ B_tmp, \
float* __restrict__ C_tmp, \
const TYPE1* __restrict__ input, \
const TYPE2* __restrict__ packed_w1, \
const TYPE2* __restrict__ packed_w2, \
const float* __restrict__ w1_bias, \
const float* __restrict__ w2_bias, \
const TYPE3* __restrict__ w1s, \
const TYPE3* __restrict__ w2s, \
int64_t block_size_N, \
int64_t block_size_K, \
const float* __restrict__ topk_weights, \
const int32_t* __restrict__ sorted_ids, \
const int32_t* __restrict__ expert_ids, \
const int32_t* __restrict__ offsets, \
int64_t M, \
int64_t N, \
int64_t K, \
int64_t E, \
int64_t topk, \
int64_t num_tokens_post_pad, \
float alpha, \
float limit, \
CPUActMethod act_func, \
bool with_bias)
INSTANTIATE_MOE_FP8_TEMPLATE(at::BFloat16);
INSTANTIATE_MOE_FP8_TEMPLATE(at::Half);
INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false);
INSTANTIATE_MOE_FP_TEMPLATE(at::Half, at::Float8_e4m3fn, float, false);
INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, uint8_t, uint8_t, true);
INSTANTIATE_MOE_FP_TEMPLATE(at::Half, uint8_t, uint8_t, true);
template <typename scalar_t>
void shared_expert_fp8_kernel_impl(
@@ -261,6 +296,7 @@ void shared_expert_fp8_kernel_impl(
/* C */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
/* Bbias */ nullptr,
/* scale */ w1s + (nb / blocks_n_per_group) * scale_size_K,
/* M */ m_size,
/* N */ n_size,
@@ -312,6 +348,7 @@ void shared_expert_fp8_kernel_impl(
/* C */ C,
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
/* Bbias */ nullptr,
/* scale */ w2s + (nb / blocks_n_per_group) * scale_size_K,
/* M */ m_size,
/* N */ n_size,
sgl-kernel/csrc/cpu/qkv_proj.cpp
+1
View File
@@ -210,6 +210,7 @@ void segment_gemm_kernel_impl(
/* C */ C + mb_start * ldc + local_nb_start,
/* Btmp*/ Btmp + tid * BLOCK_N * K,
/* Ctmp*/ Ctmp,
/*Bbias*/ nullptr,
/* Bs */ Bs + (new_nb / blocks_n_per_group) * scale_size_K,
/* M */ mb_size,
/* N */ nb_size,
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+43 -10
View File
@@ -100,7 +100,9 @@ void decode_attention_cpu(
double sm_scale,
double logit_cap,
bool is_cross_attn,
std::optional<at::Tensor> encoder_lens);
int64_t slidling_window_size,
std::optional<at::Tensor> encoder_lens,
std::optional<at::Tensor> sinks);
void extend_attention_cpu(
at::Tensor& q_extend,
@@ -118,7 +120,9 @@ void extend_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);
// flash attention
at::Tensor flash_attn_varlen_func(
@@ -215,6 +219,8 @@ at::Tensor fused_linear_sigmoid_mul(
// bmm
void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// fused moe
at::Tensor fused_experts_cpu(
at::Tensor& hidden_states,
@@ -229,9 +235,11 @@ at::Tensor fused_experts_cpu(
const std::optional<at::Tensor>& w1_zero,
const std::optional<at::Tensor>& w2_zero,
const std::optional<std::vector<int64_t>> block_size,
const std::optional<at::Tensor>& w1_bias,
const std::optional<at::Tensor>& w2_bias,
const std::optional<double>& alpha,
const std::optional<double>& limit,
bool is_vnni);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
at::Tensor shared_expert_cpu(
at::Tensor& hidden_states,
at::Tensor& w1,
@@ -313,6 +321,23 @@ at::Tensor causal_conv1d_update_cpu(
const std::optional<at::Tensor>& conv_state_indices,
int64_t pad_slot_id,
bool is_vnni);
#else
// fused moe
at::Tensor fused_experts_cpu(
at::Tensor& hidden_states,
at::Tensor& w1,
at::Tensor& w2,
at::Tensor& topk_weights,
at::Tensor& topk_ids,
bool inplace,
int64_t moe_comp_method,
const std::optional<at::Tensor>& w1_scale,
const std::optional<at::Tensor>& w2_scale,
const std::optional<at::Tensor>& w1_zero,
const std::optional<at::Tensor>& w2_zero,
const std::optional<std::vector<int64_t>> block_size,
bool is_vnni);
#endif
// conv3d fast path for patch embedding
@@ -478,15 +503,16 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"decode_attention_cpu(Tensor query, Tensor k_cache, Tensor v_cahce, Tensor(a!) output, Tensor? key, Tensor? "
"value, "
"Tensor loc, Tensor attn_logits, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, float sm_scale, "
"float logit_cap, bool is_cross_attn, Tensor? encoder_lens) -> ()");
"float logit_cap, bool is_cross_attn, int sliding_window_size, Tensor? encoder_lens, Tensor? sinks) -> ()");
m.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu);
// extend
m.def(
"extend_attention_cpu(Tensor q_extend, Tensor? k_extend, Tensor? v_extend, Tensor(a!) o_extend, Tensor k_buffer, "
"Tensor v_buffer, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor extend_seq_lens, Tensor "
"extend_start_loc, int max_len_extend, float sm_scale, float logit_cap, bool is_cross_attn, Tensor? "
"encoder_lens) -> ()");
"extend_start_loc, int max_len_extend, float sm_scale, float logit_cap, bool is_cross_attn, int "
"sliding_window_size, Tensor? "
"encoder_lens, Tensor? sinks) -> ()");
m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
// flash attn
@@ -561,14 +587,14 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, Tensor? scale) -> ()");
m.impl("bmm_cpu", torch::kCPU, &bmm_cpu);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// moe
m.def(
"fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor topk_weights, Tensor topk_ids, bool "
"inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
"Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> Tensor");
"Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? "
"limit, bool is_vnni) -> Tensor");
m.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// weight absorption
m.def(
"qkv_proj_with_rope(Tensor hidden_states, Tensor q_a_proj_weight, Tensor q_b_proj_weight, Tensor "
@@ -607,6 +633,13 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, bool is_vnni) -> Tensor");
m.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#else
// moe
m.def(
"fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor topk_weights, Tensor topk_ids, bool "
"inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
"Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> Tensor");
m.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
#endif
// conv3d fast path for patch embedding
sgl-kernel/csrc/cpu/vec.h
+14
View File
@@ -16,6 +16,15 @@ inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, c
return at::vec::convert_from_float<scalar_t>(a, b);
}
template <typename scalar_t>
inline void convert_from_float_and_store(scalar_t* out, const Vectorized<float>& a) {
float out_buffer[at::vec::Vectorized<float>::size()];
a.store(out_buffer);
for (int i = 0; i < 16; i++) {
out[i] = (scalar_t)out_buffer[i];
}
}
// allow f16, bf16
template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) {
@@ -45,6 +54,11 @@ convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorize
return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
}
template <>
inline void convert_from_float_and_store<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
_mm256_storeu_si256((__m256i*)out, (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
}
#define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
#define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a)