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

[CPU] improve silu performance by replacing fp32 div with rcp14 (#31304)

Browse Source
This commit is contained in:
Ma Mingfei
2026-07-17 12:02:37 +08:00
committed by GitHub
parent 27a52d2530
commit 486a56be56
15 changed files with 509 additions and 557 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/activation.cpp
+6 -13
View File
@@ -25,13 +25,8 @@ void act_and_mul_kernel_impl(
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= dim - kVecSize; d += kVecSize) {
bVec x_bvec = bVec::loadu(input_ptr + d);
fVec x_fvec0, x_fvec1;
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
bVec y_bvec = bVec::loadu(input_other_ptr + d);
fVec y_fvec0, y_fvec1;
std::tie(y_fvec0, y_fvec1) = at::vec::convert_to_float(y_bvec);
auto [x_fvec0, x_fvec1] = load_float_vec2(input_ptr + d);
auto [y_fvec0, y_fvec1] = load_float_vec2(input_other_ptr + d);
x_fvec0 = vf(x_fvec0);
x_fvec1 = vf(x_fvec1);
@@ -39,8 +34,7 @@ void act_and_mul_kernel_impl(
x_fvec0 = x_fvec0 * y_fvec0;
x_fvec1 = x_fvec1 * y_fvec1;
x_bvec = convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1);
x_bvec.store(output_ptr + d);
convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1).store(output_ptr + d);
}
#pragma GCC unroll 4
for (; d < dim; ++d) {
@@ -69,7 +63,6 @@ void fused_sigmoid_mul_kernel_impl(
using fVec = at::vec::Vectorized<float>;
constexpr int64_t kVecSize = bVec::size();
const fVec one = fVec(1.f);
at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
for (int64_t i = begin; i < end; ++i) {
const scalar_t* __restrict__ i_ptr = input + i * dim;
@@ -86,8 +79,8 @@ void fused_sigmoid_mul_kernel_impl(
for (; d <= head_dim - kVecSize; d += kVecSize) {
auto [x_fvec0, x_fvec1] = load_float_vec2(attn_ptr + d);
auto [g_fvec0, g_fvec1] = load_float_vec2(gate_ptr + d);
x_fvec0 = x_fvec0 / (one + g_fvec0.neg().exp_u20());
x_fvec1 = x_fvec1 / (one + g_fvec1.neg().exp_u20());
x_fvec0 = x_fvec0 * fast_sigmoid(g_fvec0);
x_fvec1 = x_fvec1 * fast_sigmoid(g_fvec1);
convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1).store(out_ptr + d);
}
#pragma GCC unroll 4
@@ -121,7 +114,7 @@ at::Tensor silu_and_mul_cpu(at::Tensor& input) {
num_tokens,
d,
[](float x) { return x / (1.f + std::exp(-x)); },
[](Vec x) { return x / (Vec(1.f) + x.neg().exp_u20()); });
[](Vec x) { return fast_silu(x); });
});
return out;
}
sgl-kernel/csrc/cpu/conv3d.cpp
+4 -5
View File
@@ -59,13 +59,12 @@ inline void copy_add_stub(
constexpr int kVecSize = bVec::size();
for (int64_t d = 0; d < N; d += kVecSize) {
fVec bias0, bias1;
bVec bias_vec = bVec::loadu(bias + d);
std::tie(bias0, bias1) = at::vec::convert_to_float(bias_vec);
auto [bias0, bias1] = load_float_vec2(bias + d);
for (int64_t m = 0; m < M; ++m) {
fVec data0 = fVec::loadu(Ctmp + m * N + d) + bias0;
fVec data1 = fVec::loadu(Ctmp + m * N + d + fVec::size()) + bias1;
auto [data0, data1] = load_float_vec2(Ctmp + m * N + d);
data0 = data0 + bias0;
data1 = data1 + bias1;
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
out_vec.store(C + m * ldc + d);
}
sgl-kernel/csrc/cpu/decode.cpp
+4 -4
View File
@@ -150,8 +150,9 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc,
int64_t d = 0;
#pragma GCC unroll 4
for (; d <= size - kVecSize; d += kVecSize) {
fVec a_fvec0 = fVec::loadu(acc + d) * s_fvec;
fVec a_fvec1 = fVec::loadu(acc + d + fVec::size()) * s_fvec;
auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d);
a_fvec0 = a_fvec0 * s_fvec;
a_fvec1 = a_fvec1 * s_fvec;
bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
out_bvec.store(out + d);
}
@@ -186,8 +187,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
constexpr int col = i % COLS;
// for COLS = 2, 4 use 512bit store
if constexpr (col % 2 == 0) {
fVec a_fvec0 = fVec::loadu(input + col * 16);
fVec a_fvec1 = fVec::loadu(input + col * 16 + 16);
auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16);
bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
out_bvec.store(out + col * 16);
}
sgl-kernel/csrc/cpu/flash_attn.h
+4 -4
View File
@@ -29,8 +29,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
constexpr int col = i % COLS;
// for COLS = 2, 4 use 512bit store
if constexpr (col % 2 == 0) {
fVec a_fvec0 = fVec::loadu(input + col * 16);
fVec a_fvec1 = fVec::loadu(input + col * 16 + 16);
auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16);
bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
out_bvec.store(out + col * 16);
}
@@ -47,8 +46,9 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc,
int d = 0;
#pragma GCC unroll 4
for (; d <= size - kVecSize; d += kVecSize) {
fVec a_fvec0 = fVec::loadu(acc + d) * s_fvec;
fVec a_fvec1 = fVec::loadu(acc + d + fVec::size()) * s_fvec;
auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d);
a_fvec0 = a_fvec0 * s_fvec;
a_fvec1 = a_fvec1 * s_fvec;
bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
out_bvec.store(out + d);
}
sgl-kernel/csrc/cpu/gemm.cpp
+8 -17
View File
@@ -111,8 +111,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data0 = fVec::loadu(input + d);
fVec data1 = fVec::loadu(input + d + fVec::size());
auto [data0, data1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
out_vec.store(out + d);
}
@@ -130,9 +129,7 @@ inline void copy_stub(float* __restrict__ out, const scalar_t* __restrict__ inpu
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data0, data1;
bVec b_vec = bVec::loadu(input + d);
std::tie(data0, data1) = at::vec::convert_to_float(b_vec);
auto [data0, data1] = load_float_vec2(input + d);
data0.store(out + d);
data1.store(out + d + fVec::size());
}
@@ -151,9 +148,9 @@ inline void copy_add_stub(
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
auto [data0, data1] = load_float_vec2(input + d);
auto [bias0, bias1] = load_float_vec2(bias + d);
bVec out_vec = convert_from_float_ext<scalar_t>(data0 + bias0, data1 + bias1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
@@ -171,7 +168,6 @@ inline void scalar_sigmoid_and_mul(
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
// scalar sigmoid
const fVec one = fVec(1.f);
fVec X;
if constexpr (has_bias) {
assert(bias != nullptr);
@@ -179,18 +175,13 @@ inline void scalar_sigmoid_and_mul(
} else {
X = fVec(input[0]);
}
X = one / (one + X.neg().exp_u20());
X = fast_sigmoid(X);
// vec mul
constexpr int kVecSize = bVec::size();
for (int d = 0; d < SIZE; d += kVecSize) {
bVec m_bvec = bVec::loadu(mul + d);
fVec m_fvec0, m_fvec1;
std::tie(m_fvec0, m_fvec1) = at::vec::convert_to_float(m_bvec);
m_fvec0 = m_fvec0 * X;
m_fvec1 = m_fvec1 * X;
bVec out_vec = convert_from_float_ext<scalar_t>(m_fvec0, m_fvec1);
auto [m_fvec0, m_fvec1] = load_float_vec2(mul + d);
bVec out_vec = convert_from_float_ext<scalar_t>(m_fvec0 * X, m_fvec1 * X);
out_vec.store(out + d);
}
}
sgl-kernel/csrc/cpu/gemm_fp8.cpp
+6 -8
View File
@@ -13,8 +13,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data0 = fVec::loadu(input + d);
fVec data1 = fVec::loadu(input + d + fVec::size());
auto [data0, data1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
out_vec.store(out + d);
}
@@ -33,9 +32,9 @@ inline void copy_add_stub(
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
auto [data0, data1] = load_float_vec2(input + d);
auto [bias0, bias1] = load_float_vec2(bias + d);
bVec out_vec = convert_from_float_ext<scalar_t>(data0 + bias0, data1 + bias1);
out_vec.store(out + d);
}
for (; d < size; ++d) {
@@ -52,9 +51,8 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
int d;
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
fVec data0 = fVec::loadu(input + d) * vscale;
fVec data1 = fVec::loadu(input + d + fVec::size()) * vscale;
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
auto [data0, data1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(data0 * vscale, data1 * vscale);
out_vec.store(out + d);
}
for (; d < size; ++d) {
sgl-kernel/csrc/cpu/mamba/conv.cpp
+2 -3
View File
@@ -177,14 +177,13 @@ struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
using fVec = at::vec::Vectorized<float>;
using bVec = at::vec::Vectorized<at::BFloat16>;
const fVec one = fVec(1.f);
auto storec = [&](auto i, int64_t m) {
constexpr int col = i;
fVec x0 = fVec(vc[col * 2 + 0]);
fVec x1 = fVec(vc[col * 2 + 1]);
if constexpr (has_silu) {
x0 = x0 / (one + x0.neg().exp_u20());
x1 = x1 / (one + x1.neg().exp_u20());
x0 = fast_silu(x0);
x1 = fast_silu(x1);
}
bVec out_vec = convert_from_float_ext<at::BFloat16>(x0, x1);
out_vec.store(C + m * lda + col * 32);
sgl-kernel/csrc/cpu/mamba/fla.cpp
+17 -38
View File
@@ -1146,14 +1146,10 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
int64_t d;
#pragma GCC unroll 4
for (d = 0; d <= head_dim - VecSize; d += VecSize) {
bVec q_bvec = bVec::loadu(q_ptr + q_offset + d);
fVec q_fvec0, q_fvec1;
std::tie(q_fvec0, q_fvec1) = at::vec::convert_to_float(q_bvec);
auto [q_fvec0, q_fvec1] = load_float_vec2(q_ptr + q_offset + d);
sum_q_fvec += q_fvec0 * q_fvec0;
sum_q_fvec += q_fvec1 * q_fvec1;
bVec k_bvec = bVec::loadu(k_ptr + k_offset + d);
fVec k_fvec0, k_fvec1;
std::tie(k_fvec0, k_fvec1) = at::vec::convert_to_float(k_bvec);
auto [k_fvec0, k_fvec1] = load_float_vec2(k_ptr + k_offset + d);
sum_k_fvec += k_fvec0 * k_fvec0;
sum_k_fvec += k_fvec1 * k_fvec1;
}
@@ -1200,14 +1196,11 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
fVec kv_mem_vec1 = fVec(float(0));
for (int di = 0; di < head_dim; ++di) {
fVec k_val_vec = fVec(k_ptr[k_offset + di] * k_scale);
fVec state_vec0 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi);
fVec state_vec1 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
auto [state_vec0, state_vec1] = load_float_vec2(state_ptr + state_offset + di * v_head_dim + dvi);
kv_mem_vec0 = kv_mem_vec0 + state_vec0 * g_val_exp_vec * k_val_vec;
kv_mem_vec1 = kv_mem_vec1 + state_vec1 * g_val_exp_vec * k_val_vec;
}
bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi);
fVec v_vec0, v_vec1;
std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec);
auto [v_vec0, v_vec1] = load_float_vec2(v_ptr + v_offset + dvi);
fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec;
fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec;
fVec o_vec0 = fVec(float(0));
@@ -1215,8 +1208,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
for (int di = 0; di < head_dim; ++di) {
fVec q_vec = fVec(q_ptr[q_offset + di] * q_scale);
fVec k_vec = fVec(k_ptr[k_offset + di] * k_scale);
fVec state_vec0 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi);
fVec state_vec1 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
auto [state_vec0, state_vec1] = load_float_vec2(state_ptr + state_offset + di * v_head_dim + dvi);
state_vec0 = state_vec0 * g_val_exp_vec + k_vec * dt_vec0;
state_vec1 = state_vec1 * g_val_exp_vec + k_vec * dt_vec1;
o_vec0 = o_vec0 + state_vec0 * q_vec * scale_vec;
@@ -1265,25 +1257,19 @@ void fused_gdn_gating_kernel_impl(
constexpr int vec_size = bVec::size();
constexpr int fvec_size = fVec::size();
const fVec neg_one(-1.0f);
const fVec one(1.0f);
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t i = begin; i < end; ++i) {
int64_t j = 0;
for (; j < num_heads - (num_heads % vec_size); j += vec_size) {
fVec A_log_vec0 = fVec::loadu(A_log + j);
fVec A_log_vec1 = fVec::loadu(A_log + j + fvec_size);
bVec dt_bias_vec = bVec::loadu(dt_bias + j);
bVec a_bvec = bVec::loadu(a + i * num_heads + j);
bVec b_bvec = bVec::loadu(b + i * num_heads + j);
fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1;
std::tie(a0, a1) = at::vec::convert_to_float(a_bvec);
std::tie(b0, b1) = at::vec::convert_to_float(b_bvec);
std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec);
auto [A_log_vec0, A_log_vec1] = load_float_vec2(A_log + j);
auto [dt_bias_vec0, dt_bias_vec1] = load_float_vec2(dt_bias + j);
auto [a0, a1] = load_float_vec2(a + i * num_heads + j);
auto [b0, b1] = load_float_vec2(b + i * num_heads + j);
fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0);
fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1);
fVec beta0 = one / (one + (neg_one * b0).exp_u20());
fVec beta1 = one / (one + (neg_one * b1).exp_u20());
fVec beta0 = fast_sigmoid(b0);
fVec beta1 = fast_sigmoid(b1);
g0.store(out + i * num_heads + j);
g1.store(out + i * num_heads + j + fvec_size);
@@ -1313,26 +1299,19 @@ void fused_gdn_gating_kernel_impl(
constexpr int vec_size = bVec::size();
constexpr int fvec_size = fVec::size();
const fVec neg_one(-1.0f);
const fVec one(1.0f);
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t i = begin; i < end; ++i) {
int64_t j = 0;
for (; j < num_heads - (num_heads % vec_size); j += vec_size) {
bVec A_log_bvec = bVec::loadu(A_log + j);
fVec A_log_vec0, A_log_vec1;
std::tie(A_log_vec0, A_log_vec1) = at::vec::convert_to_float(A_log_bvec);
bVec dt_bias_vec = bVec::loadu(dt_bias + j);
bVec a_bvec = bVec::loadu(a + i * num_heads + j);
bVec b_bvec = bVec::loadu(b + i * num_heads + j);
fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1;
std::tie(a0, a1) = at::vec::convert_to_float(a_bvec);
std::tie(b0, b1) = at::vec::convert_to_float(b_bvec);
std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec);
auto [A_log_vec0, A_log_vec1] = load_float_vec2(A_log + j);
auto [dt_bias_vec0, dt_bias_vec1] = load_float_vec2(dt_bias + j);
auto [a0, a1] = load_float_vec2(a + i * num_heads + j);
auto [b0, b1] = load_float_vec2(b + i * num_heads + j);
fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0);
fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1);
fVec beta0 = one / (one + (neg_one * b0).exp_u20());
fVec beta1 = one / (one + (neg_one * b1).exp_u20());
fVec beta0 = fast_sigmoid(b0);
fVec beta1 = fast_sigmoid(b1);
g0.store(out + i * num_heads + j);
g1.store(out + i * num_heads + j + fvec_size);
sgl-kernel/csrc/cpu/moe.cpp
+17 -107
View File
@@ -118,96 +118,6 @@ int moe_align_block_size(
return num_tokens_post_pad;
}
// silu : shape leading dimension
// input0 [m_size, BLOCK_N] BLOCK_N
// input1 [m_size, BLOCK_N] BLOCK_N
// output [M * topk, N] N
template <typename scalar_t, int BLOCK_N>
inline void silu_and_mul(
scalar_t* __restrict__ output,
const float* __restrict__ input0, // x: x0, x1
const float* __restrict__ input1, // y: y0, y1
int64_t m_size,
int64_t N) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
// no remainder
for (int64_t m = 0; m < m_size; ++m) {
scalar_t* __restrict__ out = output + m * N;
const float* __restrict__ x = input0 + m * BLOCK_N;
const float* __restrict__ y = input1 + m * BLOCK_N;
for (int64_t d = 0; d < BLOCK_N; d += bVec::size()) {
fVec x0 = fVec::loadu(x + d);
fVec x1 = fVec::loadu(x + d + fVec::size());
fVec y0 = fVec::loadu(y + d);
fVec y1 = fVec::loadu(y + d + fVec::size());
// silu
x0 = x0 / (one + x0.neg().exp_u20());
x1 = x1 / (one + x1.neg().exp_u20());
// mul
x0 = x0 * y0;
x1 = x1 * y1;
// convert
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(out + d);
}
}
}
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(
@@ -284,23 +194,17 @@ struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
Unroll<ROWS * COLS>{}(compute, k);
}
using Vec = at::vec::Vectorized<float>;
const Vec one = Vec(1.f);
auto storec = [&](auto i) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
// for COLS = 2, 4 use 512bit store
if constexpr (col % 2 == 0) {
Vec x0 = vc0[row * COLS + col + 0];
Vec x1 = vc0[row * COLS + col + 1];
Vec y0 = vc1[row * COLS + col + 0];
Vec y1 = vc1[row * COLS + col + 1];
// silu
x0 = x0 / (one + x0.neg().exp_u20());
x1 = x1 / (one + x1.neg().exp_u20());
// mul
x0 = x0 * y0;
x1 = x1 * y1;
__m512 x0 = vc0[row * COLS + col + 0];
__m512 x1 = vc0[row * COLS + col + 1];
__m512 y0 = vc1[row * COLS + col + 0];
__m512 y1 = vc1[row * COLS + col + 1];
x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0);
x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1);
_mm512_storeu_si512(
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
@@ -638,11 +542,15 @@ void fused_experts_kernel_impl(
// 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);
for (int64_t m = 0; m < m_size; ++m) {
silu_and_mul_stub(ic1 + (offset + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_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);
for (int64_t m = 0; m < m_size; ++m) {
scalar_t* __restrict__ ic1_row = ic1 + (offset + m) * N;
clamp_sigmoid_and_mul_stub(ic1_row + nb * BLOCK_N / 2, C0 + m * BLOCK_N, BLOCK_N / 2, alpha, limit);
clamp_sigmoid_and_mul_stub(ic1_row + N / 2 + nb * BLOCK_N / 2, C1 + m * BLOCK_N, BLOCK_N / 2, alpha, limit);
}
}
});
@@ -811,7 +719,9 @@ void shared_expert_kernel_impl(
/* C */ C1);
// 1.d silu and mul
silu_and_mul<scalar_t, BLOCK_N>(ic1 + mb * BLOCK_M * N + nb * BLOCK_N, C0, C1, m_size, N);
for (int64_t m = 0; m < m_size; ++m) {
silu_and_mul_stub(ic1 + (mb * BLOCK_M + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N);
}
} else {
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
tinygemm_kernel(
sgl-kernel/csrc/cpu/moe.h
+41 -82
View File
@@ -59,9 +59,7 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict_
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
x0 = x0 * weight_vec;
x1 = x1 * weight_vec;
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
out_vec.store(out + d);
}
for (; d < size; ++d) {
@@ -86,10 +84,7 @@ inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ in
fVec sum_fvec0 = fVec(0.f);
fVec sum_fvec1 = fVec(0.f);
for (int t = 0; t < topk; ++t) {
bVec x_bvec = bVec::loadu(input + t * K + d);
fVec x_fvec0, x_fvec1;
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
auto [x_fvec0, x_fvec1] = load_float_vec2(input + t * K + d);
sum_fvec0 += x_fvec0;
sum_fvec1 += x_fvec1;
}
@@ -132,11 +127,7 @@ inline void add_mul_stub(
#pragma GCC unroll 4
for (d = 0; d <= size - kVecSize; d += kVecSize) {
auto [x0, x1] = load_float_vec2(input + d);
bVec y_bvec = bVec::loadu(input2 + d);
fVec y0, y1;
std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
auto [y0, y1] = load_float_vec2(input2 + d);
x0 = x0 + y0 * s_vec;
x1 = x1 + y1 * s_vec;
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
@@ -147,31 +138,52 @@ inline void add_mul_stub(
}
}
template <typename scalar_t>
template <typename scalar_t, typename input_t>
inline void silu_and_mul_stub(
scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) {
scalar_t* __restrict__ out, const input_t* __restrict__ input, const input_t* __restrict__ input2, int64_t size) {
static_assert(
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
"silu_and_mul_stub only supports input_t == float or input_t == scalar_t");
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
bVec x = bVec::loadu(input + d);
fVec x0, x1;
std::tie(x0, x1) = at::vec::convert_to_float(x);
bVec y = bVec::loadu(input2 + d);
fVec y0, y1;
std::tie(y0, y1) = at::vec::convert_to_float(y);
x0 = x0 / (one + x0.neg().exp_u20());
x1 = x1 / (one + x1.neg().exp_u20());
x0 = x0 * y0;
x1 = x1 * y1;
auto [x0, x1] = load_float_vec2(input + d);
auto [y0, y1] = load_float_vec2(input2 + d);
x0 = fast_silu(x0) * y0;
x1 = fast_silu(x1) * y1;
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
out_vec.store(out + d);
}
}
template <typename scalar_t, typename input_t>
inline void clamp_sigmoid_and_mul_stub(
scalar_t* __restrict__ out, const input_t* __restrict__ input, int64_t size, const float alpha, const float limit) {
static_assert(
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
"clamp_sigmoid_and_mul_stub only supports input_t == float or input_t == scalar_t");
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
const fVec limit_v = fVec(limit);
const fVec nlimit_v = fVec(-limit);
const fVec alpha_v = fVec(alpha);
#pragma GCC unroll 4
for (int64_t d = 0; d < 2 * size; d += bVec::size()) {
auto [x0_, y0_] = load_float_vec2(input + d);
auto [x0, y0] = at::vec::deinterleave2<float>(x0_, y0_);
x0 = at::vec::minimum(x0, limit_v);
y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
x0 = fast_sigmoid_glu(x0, alpha_v) * (y0 + one);
store_from_float_ext(out + d / 2, x0);
}
}
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>;
@@ -181,9 +193,8 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
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);
auto [x0, x1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
out_vec.store(out + d);
}
for (; d < size; ++d) {
@@ -217,63 +228,11 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict
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);
auto [x0, x1] = load_float_vec2(input + d);
bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
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
+2 -2
View File
@@ -126,8 +126,8 @@ void fused_experts_fp_kernel_impl(
} 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);
clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N / 2, alpha, limit);
clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N / 2, alpha, limit);
}
});
}
sgl-kernel/csrc/cpu/moe_int8.cpp
+11 -19
View File
@@ -48,16 +48,14 @@ inline void silu_and_mul(
vc1[col] = _mm512_mul_ps(_mm512_mul_ps(vc1[col], vas), vbs1[col]);
};
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
const fVec one = fVec(1.f);
auto silu_and_mul = [&](auto col) {
fVec x = fVec(vc0[col]);
fVec y = fVec(vc1[col]);
x = x / (one + x.neg().exp_u20());
vc0[col] = x * y;
__m512 x = vc0[col];
__m512 y = vc1[col];
vc0[col] = _mm512_mul_ps(_mm512_rcp14_silu_ps(x), y);
};
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
auto storec = [&](auto col, int64_t m) {
if constexpr (col % 2 == 0) {
fVec x0 = fVec(vc0[col + 0]);
@@ -224,23 +222,17 @@ struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
};
Unroll<ROWS * COLS>{}(scalec);
using Vec = at::vec::Vectorized<float>;
const Vec one = Vec(1.f);
auto storec = [&](auto i) {
constexpr int row = i / COLS;
constexpr int col = i % COLS;
// for COLS = 2, 4 use 512bit store
if constexpr (col % 2 == 0) {
Vec x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]);
Vec x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
Vec y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
Vec y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
// silu
x0 = x0 / (one + x0.neg().exp_u20());
x1 = x1 / (one + x1.neg().exp_u20());
// mul
x0 = x0 * y0;
x1 = x1 * y1;
__m512 x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]);
__m512 x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
__m512 y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
__m512 y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0);
x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1);
_mm512_storeu_si512(
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
sgl-kernel/csrc/cpu/qkv_proj.cpp
+3 -9
View File
@@ -240,9 +240,7 @@ inline float reduce(const scalar_t* __restrict__ x, int64_t size) {
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
bVec x_bvec = bVec::loadu(x + d);
fVec x_fvec0, x_fvec1;
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
auto [x_fvec0, x_fvec1] = load_float_vec2(x + d);
sum_fvec += x_fvec0 * x_fvec0;
sum_fvec += x_fvec1 * x_fvec1;
}
@@ -259,12 +257,8 @@ inline void map2(scalar_t* y, const scalar_t* x, const scalar_t* __restrict__ w,
// no remainder
#pragma GCC unroll 4
for (int64_t d = 0; d < size; d += bVec::size()) {
bVec x_bvec = bVec::loadu(x + d);
fVec x_fvec0, x_fvec1;
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
bVec w_bvec = bVec::loadu(w + d);
fVec w_fvec0, w_fvec1;
std::tie(w_fvec0, w_fvec1) = at::vec::convert_to_float(w_bvec);
auto [x_fvec0, x_fvec1] = load_float_vec2(x + d);
auto [w_fvec0, w_fvec1] = load_float_vec2(w + d);
x_fvec0 = x_fvec0 * scale_fvec * w_fvec0;
x_fvec1 = x_fvec1 * scale_fvec * w_fvec1;
bVec out_bvec = convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1);
sgl-kernel/csrc/cpu/vec.h
+57 -6
View File
@@ -17,11 +17,11 @@ inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, c
}
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()];
inline void store_from_float_ext(scalar_t* out, const Vectorized<float>& a) {
float out_buffer[Vectorized<float>::size()];
a.store(out_buffer);
for (int i = 0; i < 16; i++) {
out[i] = (scalar_t)out_buffer[i];
for (int i = 0; i < Vectorized<float>::size(); ++i) {
out[i] = static_cast<scalar_t>(out_buffer[i]);
}
}
@@ -55,8 +55,14 @@ convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorize
}
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))));
inline void store_from_float_ext<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
_mm256_storeu_si256(reinterpret_cast<__m256i*>(out), (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
}
template <>
inline void store_from_float_ext<at::Half>(at::Half* out, const Vectorized<float>& a) {
_mm256_storeu_si256(
reinterpret_cast<__m256i*>(out), _mm512_cvtps_ph(__m512(a), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
}
#define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
@@ -557,6 +563,51 @@ inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 val
// final interpretation to float
return _mm512_castsi512_ps(casted_integer);
}
// sigmoid(x) = 1 / (1 + exp(-x)); avoid vdivps via rcp14
inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_ps(__m512 x) {
__m512 minus_x = _mm512_xor_ps(_mm512_set1_ps(-0.f), x);
__m512 denom = _mm512_add_ps(_mm512_exp_u20_ps(minus_x), _mm512_set1_ps(1.f));
return _mm512_rcp14_ps(denom);
}
// SiLU(x) = x * sigmoid(x)
inline __attribute__((always_inline)) __m512 _mm512_rcp14_silu_ps(__m512 x) {
return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(x));
}
// x * sigmoid(x * alpha) for clamped SwiGLU
inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_glu_ps(__m512 x, __m512 alpha) {
__m512 xa = _mm512_mul_ps(x, alpha);
return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(xa));
}
#endif
inline at::vec::Vectorized<float> fast_sigmoid(const at::vec::Vectorized<float>& x) {
#if defined(CPU_CAPABILITY_AVX512)
return at::vec::Vectorized<float>(_mm512_rcp14_sigmoid_ps(x));
#else
const auto one = at::vec::Vectorized<float>(1.f);
return one / (one + x.neg().exp_u20());
#endif
}
inline at::vec::Vectorized<float> fast_silu(const at::vec::Vectorized<float>& x) {
#if defined(CPU_CAPABILITY_AVX512)
return at::vec::Vectorized<float>(_mm512_rcp14_silu_ps(x));
#else
return x * fast_sigmoid(x);
#endif
}
inline at::vec::Vectorized<float>
fast_sigmoid_glu(const at::vec::Vectorized<float>& x, const at::vec::Vectorized<float>& alpha) {
#if defined(CPU_CAPABILITY_AVX512)
return at::vec::Vectorized<float>(_mm512_rcp14_sigmoid_glu_ps(x, alpha));
#else
return x * fast_sigmoid(x * alpha);
#endif
}
} // anonymous namespace
test/registered/cpu/test_moe.py
+326 -239
View File
@@ -1,5 +1,6 @@
import math
import unittest
import pytest
# TODO: use interface in cpu.py
import torch
@@ -10,6 +11,11 @@ kernel = torch.ops.sgl_kernel
torch.manual_seed(1183)
dtype = torch.bfloat16
prepack = True
alpha = 1.702
limit = 7.0
from utils import (
BLOCK_K,
BLOCK_N,
@@ -18,7 +24,6 @@ from utils import (
fp8_max,
fp8_min,
native_fp8_fused_moe,
parametrize,
precision,
scaled_weight,
torch_naive_fused_moe,
@@ -28,114 +33,199 @@ from utils import (
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
def fused_moe(a, w1, w2, score, topk, renormalize, prepack):
G = 1
topk_group = 1
B, D = a.shape
topk_weights = torch.empty(B, topk, dtype=torch.float32)
topk_ids = torch.empty(B, topk, dtype=torch.int32)
topk_weights, topk_ids = kernel.grouped_topk_cpu(
a, score, topk, renormalize, G, topk_group, 0, None, None
def run_fused_experts(
a,
w1,
w2,
topk_weight,
topk_ids,
*,
quant=CPUQuantMethod.UNQUANT,
w1_scale=None,
w2_scale=None,
w1_zp=None,
w2_zp=None,
block_size=None,
w1_bias=None,
w2_bias=None,
alpha=None,
limit=None,
is_vnni=True,
inplace=False,
):
return kernel.fused_experts_cpu(
a,
w1,
w2,
topk_weight,
topk_ids.to(torch.int32),
inplace,
quant,
w1_scale,
w2_scale,
w1_zp,
w2_zp,
block_size,
w1_bias,
w2_bias,
alpha,
limit,
is_vnni,
)
def make_routing(m, e, topk, dtype, renormalize=False, score=None, return_score=False):
if score is None:
score = torch.randn((m, e), dtype=dtype)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
if renormalize:
topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True)
if return_score:
return topk_weight, topk_ids, score
return topk_weight, topk_ids
def make_bf16_weights(e, out_dim, in_dim, with_bias=False):
weight = torch.randn((e, out_dim, in_dim), dtype=dtype) / 10
if not with_bias:
return weight
bias = torch.randn((e, out_dim), dtype=torch.float32) / 10
return weight, bias
def make_int8_weights(e, out_dim, in_dim, int8_max=127, int8_min=-128):
weight_fp32 = (torch.rand((e, out_dim, in_dim), dtype=torch.float32) - 0.5) * 2
weight = (weight_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
return weight
def make_fp8_weights(e, out_dim, in_dim):
weight_fp32 = torch.randn(e, out_dim, in_dim)
weight = (
(weight_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
)
weight_s = (
torch.rand(e, math.ceil(out_dim / BLOCK_N), math.ceil(in_dim / BLOCK_K))
* factor_for_scale
)
weight_scaled = scaled_weight(weight, weight_s)
return weight, weight_s, weight_scaled
def make_mxfp4_weights(e, out_dim, in_dim, dtype, with_bias=False):
weight_bf16 = torch.randn((e, out_dim, in_dim), dtype=dtype) / 10
weight_q, weight_s = MXFP4QuantizeUtil.quantize(weight_bf16)
weight_s = weight_s.reshape(e, out_dim, in_dim // 32)
weight_dq = MXFP4QuantizeUtil.dequantize(weight_q, dtype, weight_s)
weight_packed = kernel.convert_weight_packed(weight_q)
weight_s_packed = kernel.convert_scale_packed(weight_s)
if not with_bias:
return weight_dq, weight_packed, weight_s_packed
bias = torch.randn((e, out_dim), dtype=torch.float32) / 10
return weight_dq, bias, weight_packed, weight_s_packed
class TestFusedExperts:
@pytest.mark.parametrize("m", [2, 114])
@pytest.mark.parametrize("n", [32])
@pytest.mark.parametrize("k", [32])
@pytest.mark.parametrize("e", [4])
@pytest.mark.parametrize("topk", [2])
@pytest.mark.parametrize("renormalize", [False, True])
def test_bf16_moe(self, m, n, k, e, topk, renormalize):
a = torch.randn((m, k), dtype=dtype) / 10
w1 = make_bf16_weights(e, 2 * n, k)
w2 = make_bf16_weights(e, k, n)
topk_weights, topk_ids, score = make_routing(
m,
e,
topk,
dtype=dtype,
renormalize=renormalize,
return_score=True,
)
torch_output = torch_naive_fused_moe(a, w1, w2, score, topk, renormalize)
packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1
packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2
inplace = True
return kernel.fused_experts_cpu(
fused_output = run_fused_experts(
a,
packed_w1,
packed_w2,
topk_weights,
topk_ids,
inplace,
CPUQuantMethod.UNQUANT,
None,
None,
None,
None,
None,
None,
None,
None,
None,
prepack,
quant=CPUQuantMethod.UNQUANT,
is_vnni=prepack,
inplace=True,
)
class TestFusedExperts(CustomTestCase):
@parametrize(m=[2, 114], n=[32], k=[32], e=[4], topk=[2], renormalize=[False, True])
def test_bf16_moe(self, m, n, k, e, topk, renormalize):
dtype = torch.bfloat16
prepack = True
a = torch.randn((m, k), device="cpu", dtype=dtype) / 10
w1 = torch.randn((e, 2 * n, k), device="cpu", dtype=dtype) / 10
w2 = torch.randn((e, k, n), device="cpu", dtype=dtype) / 10
score = torch.randn((m, e), device="cpu", dtype=dtype)
torch_output = torch_naive_fused_moe(a, w1, w2, score, topk, renormalize)
fused_output = fused_moe(a, w1, w2, score, topk, renormalize, prepack)
atol = rtol = precision[torch_output.dtype]
torch.testing.assert_close(torch_output, fused_output, atol=atol, rtol=rtol)
@parametrize(
m=[1, 32], n=[128, 64], k=[128, 64], e=[4], topk=[2], renormalize=[False]
)
@pytest.mark.parametrize("m", [1, 32])
@pytest.mark.parametrize("n", [128, 64])
@pytest.mark.parametrize("k", [128, 64])
@pytest.mark.parametrize("e", [4])
@pytest.mark.parametrize("topk", [2])
@pytest.mark.parametrize("renormalize", [False])
def test_bf16_moe_bias(self, m, n, k, e, topk, renormalize):
dtype = torch.bfloat16
a = torch.randn((m, k), device="cpu", dtype=dtype) / 10
w1 = torch.randn((e, 2 * n, k), device="cpu", dtype=dtype) / 10
w1_b = torch.randn((e, 2 * n), device="cpu", dtype=torch.float) / 10
w2 = torch.randn((e, k, n), device="cpu", dtype=dtype) / 10
w2_b = torch.randn((e, k), device="cpu", dtype=torch.float) / 10
score = torch.randn((m, e), device="cpu", dtype=dtype)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
alpha = 1.702
limit = 7.0
a = torch.randn((m, k), dtype=dtype) / 10
w1, w1_b = make_bf16_weights(e, 2 * n, k, with_bias=True)
w2, w2_b = make_bf16_weights(e, k, n, with_bias=True)
topk_weight, topk_ids = make_routing(
m, e, topk, dtype=dtype, renormalize=renormalize
)
torch_output = torch_naive_fused_moe_gptoss(
a, w1, w2, w1_b, w2_b, topk_weight, topk_ids, renormalize, alpha, limit, e
a,
w1,
w2,
w1_b,
w2_b,
topk_weight,
topk_ids,
renormalize,
alpha,
limit,
e,
)
packed_w1 = kernel.convert_weight_packed(w1)
packed_w2 = kernel.convert_weight_packed(w2)
fused_output = torch.ops.sgl_kernel.fused_experts_cpu(
fused_output = run_fused_experts(
a,
packed_w1,
packed_w2,
topk_weight,
topk_ids.to(torch.int),
False, # inplace # See [Note] inplace should be False in fused_experts.
CPUQuantMethod.UNQUANT,
None, # w1_scale
None, # w2_scale
None, # w1_zp
None, # w2_zp
None, # block_size
w1_b,
w2_b,
alpha,
limit,
True, # is_vnni
topk_ids,
quant=CPUQuantMethod.UNQUANT,
w1_bias=w1_b,
w2_bias=w2_b,
alpha=alpha,
limit=limit,
is_vnni=True,
inplace=False,
)
atol = rtol = precision[torch_output.dtype]
torch.testing.assert_close(torch_output, fused_output, atol=atol, rtol=rtol)
@parametrize(M=[1, 39], N=[128], K=[256], E=[8], topk=[3])
@pytest.mark.parametrize("M", [1, 39])
@pytest.mark.parametrize("N", [128])
@pytest.mark.parametrize("K", [256])
@pytest.mark.parametrize("E", [8])
@pytest.mark.parametrize("topk", [3])
def test_int8_moe(self, M, N, K, E, topk):
dtype = torch.bfloat16
prepack = True
# Initialize int8 quantization parameters
int8_factor_for_scale = 1e-2
int8_max = 127
@@ -146,20 +236,15 @@ class TestFusedExperts(CustomTestCase):
a = torch.randn((M, K), dtype=dtype) / math.sqrt(K)
# Generate int8 weights
w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
w1 = make_int8_weights(E, 2 * N, K, int8_max=int8_max, int8_min=int8_min)
w2 = make_int8_weights(E, K, N, int8_max=int8_max, int8_min=int8_min)
# Generate scale for each column (per-column quantization)
w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * int8_factor_for_scale
w2_s = torch.rand(E, K, device=w2_fp32.device) * int8_factor_for_scale
w1_s = torch.rand(E, 2 * N) * int8_factor_for_scale
w2_s = torch.rand(E, K) * int8_factor_for_scale
# Calculate routing
score = torch.randn((M, E), dtype=dtype)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
topk_weight, topk_ids = make_routing(M, E, topk, dtype=dtype)
ref_out = torch_w8a8_per_column_fused_moe(
a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, topk
@@ -168,24 +253,17 @@ class TestFusedExperts(CustomTestCase):
inplace = True
packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1
packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2
out = kernel.fused_experts_cpu(
out = run_fused_experts(
a,
packed_w1,
packed_w2,
topk_weight,
topk_ids.to(torch.int32),
inplace,
CPUQuantMethod.INT8_W8A8,
w1_s,
w2_s,
None,
None,
None,
None,
None,
None,
None,
prepack,
topk_ids,
quant=CPUQuantMethod.INT8_W8A8,
w1_scale=w1_s,
w2_scale=w2_s,
is_vnni=prepack,
inplace=inplace,
)
atol = rtol = precision[ref_out.dtype]
@@ -194,33 +272,18 @@ class TestFusedExperts(CustomTestCase):
atol = rtol = 0.02
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
@parametrize(M=[2, 121], N=[352, 512], K=[256, 320], E=[8], topk=[4])
@pytest.mark.parametrize("M", [2, 121])
@pytest.mark.parametrize("N", [352, 512])
@pytest.mark.parametrize("K", [256, 320])
@pytest.mark.parametrize("E", [8])
@pytest.mark.parametrize("topk", [4])
def test_fp8_moe(self, M, N, K, E, topk):
dtype = torch.bfloat16
a = torch.randn(M, K, dtype=dtype) / math.sqrt(K)
w1_fp32 = torch.randn(E, 2 * N, K)
w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
w1, w1s, w1_scaled = make_fp8_weights(E, 2 * N, K)
w2, w2s, w2_scaled = make_fp8_weights(E, K, N)
w2_fp32 = torch.randn(E, K, N)
w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
w1s = (
torch.randn(E, math.ceil(2 * N / BLOCK_N), math.ceil(K / BLOCK_K))
* factor_for_scale
)
w2s = (
torch.randn(E, math.ceil(K / BLOCK_N), math.ceil(N / BLOCK_K))
* factor_for_scale
)
w1_scaled = scaled_weight(w1, w1s)
w2_scaled = scaled_weight(w2, w2s)
score = torch.randn((M, E), dtype=dtype)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
topk_weight, topk_ids = make_routing(M, E, topk, dtype=dtype)
w1 = kernel.convert_weight_packed(w1)
w2 = kernel.convert_weight_packed(w2)
@@ -228,102 +291,132 @@ class TestFusedExperts(CustomTestCase):
ref_out = native_fp8_fused_moe(
a, w1_scaled, w2_scaled, topk_weight, topk_ids, topk
)
out = kernel.fused_experts_cpu(
out = run_fused_experts(
a,
w1,
w2,
topk_weight,
topk_ids.to(torch.int32),
False,
CPUQuantMethod.FP8_W8A16,
w1s,
w2s,
None,
None,
[BLOCK_N, BLOCK_K],
None,
None,
None,
None,
True,
topk_ids,
quant=CPUQuantMethod.FP8_W8A16,
w1_scale=w1s,
w2_scale=w2s,
block_size=[BLOCK_N, BLOCK_K],
is_vnni=True,
inplace=False,
)
atol = rtol = precision[dtype]
torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
@parametrize(M=[2, 121], N=[352, 512], K=[256, 320], E=[8], topk=[4])
def test_mxfp4_moe(self, M, N, K, E, topk):
dtype = torch.bfloat16
@pytest.mark.parametrize("m", [1, 32])
@pytest.mark.parametrize("n", [128, 64])
@pytest.mark.parametrize("k", [128, 64])
@pytest.mark.parametrize("e", [4])
@pytest.mark.parametrize("topk", [2])
@pytest.mark.parametrize("renormalize", [False])
def test_fp8_moe_bias(self, m, n, k, e, topk, renormalize):
a = torch.randn((m, k), dtype=dtype) / 10
w1, w1s, w1_scaled = make_fp8_weights(e, 2 * n, k)
w2, w2s, w2_scaled = make_fp8_weights(e, k, n)
w1_b = torch.randn((e, 2 * n), dtype=torch.float32) / 10
w2_b = torch.randn((e, k), dtype=torch.float32) / 10
w1_scaled = w1_scaled.to(dtype)
w2_scaled = w2_scaled.to(dtype)
topk_weight, topk_ids = make_routing(
m, e, topk, dtype=dtype, renormalize=renormalize
)
ref_out = torch_naive_fused_moe_gptoss(
a,
w1_scaled,
w2_scaled,
w1_b,
w2_b,
topk_weight,
topk_ids,
renormalize,
alpha,
limit,
e,
)
w1 = kernel.convert_weight_packed(w1)
w2 = kernel.convert_weight_packed(w2)
out = run_fused_experts(
a,
w1,
w2,
topk_weight,
topk_ids,
quant=CPUQuantMethod.FP8_W8A16,
w1_scale=w1s,
w2_scale=w2s,
block_size=[BLOCK_N, BLOCK_K],
w1_bias=w1_b,
w2_bias=w2_b,
alpha=alpha,
limit=limit,
is_vnni=True,
inplace=False,
)
atol = rtol = precision[dtype]
torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
@pytest.mark.parametrize("M", [2, 121])
@pytest.mark.parametrize("N", [352, 512])
@pytest.mark.parametrize("K", [256, 320])
@pytest.mark.parametrize("E", [8])
@pytest.mark.parametrize("topk", [4])
def test_mxfp4_moe(self, M, N, K, E, topk):
a = torch.randn(M, K, dtype=dtype) / 10
w1_bf16 = torch.randn((E, 2 * N, K), dtype=dtype) / 10
w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
w1s = w1s.reshape(E, 2 * N, K // 32)
w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
w1dq, w1_packed, w1s_packed = make_mxfp4_weights(E, 2 * N, K, dtype=dtype)
w2dq, w2_packed, w2s_packed = make_mxfp4_weights(E, K, N, dtype=dtype)
w2_bf16 = torch.randn((E, K, N), dtype=dtype) / 10
w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
w2s = w2s.reshape(E, K, N // 32)
w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
score = torch.randn((M, E), dtype=dtype)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
w1 = kernel.convert_weight_packed(w1q)
w2 = kernel.convert_weight_packed(w2q)
w1s = kernel.convert_scale_packed(w1s)
w2s = kernel.convert_scale_packed(w2s)
topk_weight, topk_ids = make_routing(M, E, topk, dtype=dtype)
ref_out = native_fp8_fused_moe(
a, w1dq.float(), w2dq.float(), topk_weight, topk_ids, topk
)
out = kernel.fused_experts_cpu(
out = run_fused_experts(
a,
w1,
w2,
w1_packed,
w2_packed,
topk_weight,
topk_ids.to(torch.int32),
False,
CPUQuantMethod.MXFP4,
w1s,
w2s,
None,
None,
None,
None,
None,
None,
None,
True,
topk_ids,
quant=CPUQuantMethod.MXFP4,
w1_scale=w1s_packed,
w2_scale=w2s_packed,
is_vnni=True,
inplace=False,
)
atol = rtol = precision[dtype]
torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
@parametrize(
m=[1, 32], n=[128, 64], k=[128, 64], e=[4], topk=[2], renormalize=[False]
)
@pytest.mark.parametrize("m", [1, 32])
@pytest.mark.parametrize("n", [128, 64])
@pytest.mark.parametrize("k", [128, 64])
@pytest.mark.parametrize("e", [4])
@pytest.mark.parametrize("topk", [2])
@pytest.mark.parametrize("renormalize", [False])
def test_mxfp4_moe_bias(self, m, n, k, e, topk, renormalize):
dtype = torch.bfloat16
a = torch.randn((m, k), device="cpu", dtype=dtype) / 10
w1_bf16 = torch.randn((e, 2 * n, k), device="cpu", dtype=dtype) / 10
w1q, w1s = MXFP4QuantizeUtil.quantize(w1_bf16)
w1s = w1s.reshape(e, 2 * n, k // 32)
w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s)
w1_b = torch.randn((e, 2 * n), device="cpu", dtype=torch.float32) / 10
w2_bf16 = torch.randn((e, k, n), device="cpu", dtype=dtype) / 10
w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16)
w2s = w2s.reshape(e, k, n // 32)
w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s)
w2_b = torch.randn((e, k), device="cpu", dtype=torch.float32) / 10
score = torch.randn((m, e), device="cpu", dtype=dtype)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
alpha = 1.702
limit = 7.0
a = torch.randn((m, k), dtype=dtype) / 10
w1dq, w1_b, w1_packed, w1s_packed = make_mxfp4_weights(
e, 2 * n, k, dtype=dtype, with_bias=True
)
w2dq, w2_b, w2_packed, w2s_packed = make_mxfp4_weights(
e, k, n, dtype=dtype, with_bias=True
)
topk_weight, topk_ids = make_routing(
m, e, topk, dtype=dtype, renormalize=renormalize
)
torch_output = torch_naive_fused_moe_gptoss(
a,
w1dq,
@@ -338,48 +431,46 @@ class TestFusedExperts(CustomTestCase):
e,
)
w1 = kernel.convert_weight_packed(w1q)
w2 = kernel.convert_weight_packed(w2q)
w1s = kernel.convert_scale_packed(w1s)
w2s = kernel.convert_scale_packed(w2s)
fused_output = torch.ops.sgl_kernel.fused_experts_cpu(
fused_output = run_fused_experts(
a,
w1,
w2,
w1_packed,
w2_packed,
topk_weight,
topk_ids.to(torch.int32),
False, # inplace # See [Note] inplace should be False in fused_experts.
CPUQuantMethod.MXFP4, # use_mxfp4
w1s, # w1_scale
w2s, # w2_scale
None, # w1_zp
None, # w2_zp
None, # block_size
w1_b,
w2_b,
alpha,
limit,
True, # is_vnni
topk_ids,
quant=CPUQuantMethod.MXFP4,
w1_scale=w1s_packed,
w2_scale=w2s_packed,
w1_bias=w1_b,
w2_bias=w2_b,
alpha=alpha,
limit=limit,
is_vnni=True,
inplace=False,
)
atol = rtol = precision[torch_output.dtype]
torch.testing.assert_close(torch_output, fused_output, atol=atol, rtol=rtol)
@parametrize(M=[1, 6], N=[512], K=[256], E=[8], topk=[4])
@pytest.mark.parametrize("M", [1, 6])
@pytest.mark.parametrize("N", [512])
@pytest.mark.parametrize("K", [256])
@pytest.mark.parametrize("E", [8])
@pytest.mark.parametrize("topk", [4])
def test_int4_moe(self, M, N, K, E, topk, group_size=128):
dtype = torch.bfloat16
a = torch.rand(M, K, dtype=dtype) / math.sqrt(K)
awq_w13_weight = torch.randint(-127, 128, (E, K, 2 * N // 8)).to(torch.int)
awq_w13_zero = torch.randint(0, 10, (E, K // group_size, 2 * N // 8)).to(
torch.int
)
awq_w13_scales = torch.rand(E, int(K // group_size), 2 * N).to(torch.bfloat16)
awq_w13_scales = (
torch.rand(E, int(K // group_size), 2 * N) * factor_for_scale
).to(torch.bfloat16)
awq_w2_weight = torch.randint(-127, 128, (E, N, K // 8)).to(torch.int)
awq_w2_zero = torch.randint(0, 10, (E, N // group_size, K // 8)).to(torch.int)
awq_w2_scales = torch.rand(E, int(N // group_size), K).to(torch.bfloat16)
awq_w2_scales = (torch.rand(E, int(N // group_size), K) * factor_for_scale).to(
torch.bfloat16
)
bf16_w13_weight = []
bf16_w2_weight = []
for i in range(E):
@@ -399,8 +490,7 @@ class TestFusedExperts(CustomTestCase):
ref_out = torch_naive_fused_moe(
a, bf16_w13_weight, bf16_w2_weight, score, topk, False
)
score = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weight, topk_ids = torch.topk(score, topk)
topk_weight, topk_ids = make_routing(M, E, topk, dtype=dtype, score=score)
awq_w13_weight_pack, awq_w13_zero_pack, awq_w13_scales_pack = (
torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
awq_w13_weight, awq_w13_zero, awq_w13_scales, 0
@@ -412,24 +502,19 @@ class TestFusedExperts(CustomTestCase):
)
)
out = kernel.fused_experts_cpu(
out = run_fused_experts(
a,
awq_w13_weight_pack,
awq_w2_weight_pack,
topk_weight,
topk_ids.to(torch.int32),
False,
CPUQuantMethod.INT4_W4A8,
awq_w13_scales_pack,
awq_w2_scales_pack,
awq_w13_zero_pack,
awq_w2_zero_pack,
None,
None,
None,
None,
None,
True,
topk_ids,
quant=CPUQuantMethod.INT4_W4A8,
w1_scale=awq_w13_scales_pack,
w2_scale=awq_w2_scales_pack,
w1_zp=awq_w13_zero_pack,
w2_zp=awq_w2_zero_pack,
is_vnni=True,
inplace=False,
)
atol = rtol = precision[dtype]
@@ -437,4 +522,6 @@ class TestFusedExperts(CustomTestCase):
if __name__ == "__main__":
unittest.main()
import sys
sys.exit(pytest.main([__file__]))