Skip to content

[CPU] Merge MoE ±bias tests via bias=[False, True] #94

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mingfeima wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

[CPU] Merge MoE ±bias tests via bias=[False, True] #94

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f9876d6
replace fp32 div with rcp14 and simplify code
mingfeima Jul 15, 2026
27d7a48
Fix test_int4_moe: keep shared routing for ref vs kernel
cursoragent Jul 16, 2026
d1109b2
Fix FP8 MoE bias test: keep float32 scaled weights for ref
cursoragent Jul 16, 2026
92fcf42
fix error
mingfeima Jul 16, 2026
8c2ac79
Refactor CPU MoE tests: merge ±bias cases via bias=[False, True]
cursoragent Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 6 additions & 13 deletions sgl-kernel/csrc/cpu/activation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,22 +25,16 @@ 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);

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) {
Expand Down Expand Up @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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;
}
Expand Down
9 changes: 4 additions & 5 deletions sgl-kernel/csrc/cpu/conv3d.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
8 changes: 4 additions & 4 deletions sgl-kernel/csrc/cpu/decode.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down Expand Up @@ -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);
}
Expand Down
8 changes: 4 additions & 4 deletions sgl-kernel/csrc/cpu/flash_attn.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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);
}
Expand Down
25 changes: 8 additions & 17 deletions sgl-kernel/csrc/cpu/gemm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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());
}
Expand All @@ -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) {
Expand All @@ -171,26 +168,20 @@ 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);
X = fVec(input[0] + bias[0]);
} 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);
}
}
Expand Down
14 changes: 6 additions & 8 deletions sgl-kernel/csrc/cpu/gemm_fp8.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand All @@ -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) {
Expand All @@ -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) {
Expand Down
5 changes: 2 additions & 3 deletions sgl-kernel/csrc/cpu/mamba/conv.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
53 changes: 17 additions & 36 deletions sgl-kernel/csrc/cpu/mamba/fla.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -1200,23 +1196,19 @@ 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));
fVec o_vec1 = fVec(float(0));
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;
Expand Down Expand Up @@ -1270,20 +1262,15 @@ void fused_gdn_gating_kernel_impl(
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);
Expand Down Expand Up @@ -1318,21 +1305,15 @@ void fused_gdn_gating_kernel_impl(
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);
Expand Down
Loading
Loading