From 929e00eeab0e0f2d5537a9019984941a4f8f7071 Mon Sep 17 00:00:00 2001 From: Ma Mingfei Date: Tue, 21 Apr 2026 20:03:39 +0800 Subject: [PATCH] [CPU] expand the interface of shared_expert without scaling factor (#22933) merge since this is CPU only change on sgl-kernel. --- sgl-kernel/csrc/cpu/moe.cpp | 144 +++------------ sgl-kernel/csrc/cpu/moe.h | 173 ++++++++++++++++++ sgl-kernel/csrc/cpu/moe_fp8.cpp | 142 +-------------- sgl-kernel/csrc/cpu/moe_int4.cpp | 186 ++------------------ sgl-kernel/csrc/cpu/moe_int8.cpp | 114 +----------- sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 6 +- test/srt/cpu/test_moe.py | 35 +--- test/srt/cpu/test_shared_expert.py | 114 ++++++------ test/srt/cpu/utils.py | 22 ++- 9 files changed, 313 insertions(+), 623 deletions(-) create mode 100644 sgl-kernel/csrc/cpu/moe.h diff --git a/sgl-kernel/csrc/cpu/moe.cpp b/sgl-kernel/csrc/cpu/moe.cpp index eaf700bf9..35215d140 100644 --- a/sgl-kernel/csrc/cpu/moe.cpp +++ b/sgl-kernel/csrc/cpu/moe.cpp @@ -1,6 +1,7 @@ +#include "moe.h" + #include "common.h" #include "gemm.h" -#include "vec.h" namespace { @@ -25,112 +26,6 @@ namespace { // 3. abstract at::native::cpublas::brgemm with WoQ gemm (M = 1 & M != 1) // -template -inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) { - using Vec = at::vec::Vectorized; - const Vec data_vec(val); - at::vec::map([data_vec](Vec out) { return out = data_vec; }, out, out, size); -} - -template -inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) { - using Vec = at::vec::Vectorized; -// no remainder -#pragma GCC unroll 4 - for (int64_t d = 0; d < size; d += Vec::size()) { - Vec data = Vec::loadu(input + d); - data.store(out + d); - } -} - -template -inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - 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(data0, data1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d] * weight); - } -} - -// acc from [topk, K] to [K] -template -inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - if (topk == 1) { - // do copy for topk = 1 - copy_stub(out, input, K); - } else { - // do sum for topk != 1 - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= K - kVecSize; d += kVecSize) { - 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); - - sum_fvec0 += x_fvec0; - sum_fvec1 += x_fvec1; - } - bVec out_bvec = convert_from_float_ext(sum_fvec0, sum_fvec1); - out_bvec.store(out + d); - } - for (; d < K; ++d) { - float sum_val = 0.f; - for (int t = 0; t < topk; ++t) { - sum_val += static_cast(input[t * K + d]); - } - out[d] = static_cast(sum_val); - } - } -} - -// out = input + input2 * scale -template -inline void add_mul_stub( - scalar_t* __restrict__ out, - const float* __restrict__ input, - const scalar_t* __restrict__ input2, - float scale, - int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - const fVec s_vec = fVec(scale); - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec x0 = fVec::loadu(input + d); - fVec x1 = fVec::loadu(input + d + fVec::size()); - - bVec y_bvec = bVec::loadu(input2 + d); - fVec y0, y1; - std::tie(y0, y1) = at::vec::convert_to_float(y_bvec); - - x0 = x0 + y0 * s_vec; - x1 = x1 + y1 * s_vec; - bVec out_vec = convert_from_float_ext(x0, x1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d] + float(input2[d]) * scale); - } -} - template int moe_align_block_size( int32_t* __restrict__ sorted_ids, @@ -765,6 +660,8 @@ void shared_expert_kernel_impl( const bool use_brgemm = can_use_brgemm(M); + const bool apply_scaling_factor = fused_experts_out != nullptr; + // here we only parallel on half of 2N to fuse silu_and_mul with gemm parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { // get local pointers @@ -888,9 +785,11 @@ void shared_expert_kernel_impl( // 2.b copy from C to output and add fused_experts_out scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N; - const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N; + const scalar_t* __restrict__ fused_out = + apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr; for (int64_t m = 0; m < m_size; ++m) { - add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size); + const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr; + add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size); } }); @@ -1235,8 +1134,8 @@ at::Tensor shared_expert_cpu( at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2, - at::Tensor& fused_experts_out, - double routed_scaling_factor, + const std::optional& fused_experts_out, + const std::optional routed_scaling_factor, bool inplace, bool use_int8_w8a8, bool use_fp8_w8a16, @@ -1252,15 +1151,22 @@ at::Tensor shared_expert_cpu( constexpr int64_t BLOCK_M = block_size_m(); constexpr int64_t BLOCK_N = block_size_n(); + double routed_scaling_factor_value = 0; + if (routed_scaling_factor.has_value()) { + TORCH_CHECK(fused_experts_out.has_value(), "shared_expert_cpu: expect fused_experts_out."); + const auto fused_experts_out_tensor = fused_experts_out.value(); + routed_scaling_factor_value = routed_scaling_factor.value(); + CHECK_INPUT(fused_experts_out_tensor); + CHECK_EQ(hidden_states.sizes(), fused_experts_out_tensor.sizes()); + } + const auto st = hidden_states.scalar_type(); CHECK_INPUT(hidden_states); - CHECK_INPUT(fused_experts_out); CHECK_INPUT(w1); CHECK_INPUT(w2); CHECK_DIM(2, hidden_states); CHECK_DIM(2, w1); CHECK_DIM(2, w2); - CHECK_EQ(hidden_states.sizes(), fused_experts_out.sizes()); CHECK_EQ(hidden_states.scalar_type(), st); int64_t M = hidden_states.size(0); @@ -1328,8 +1234,8 @@ at::Tensor shared_expert_cpu( packed_w2.data_ptr(), w1s.data_ptr(), w2s.data_ptr(), - fused_experts_out.data_ptr(), - routed_scaling_factor, + conditional_data_ptr(fused_experts_out), + routed_scaling_factor_value, M, N, K); @@ -1351,8 +1257,8 @@ at::Tensor shared_expert_cpu( w2s.data_ptr(), block_size_N, block_size_K, - fused_experts_out.data_ptr(), - routed_scaling_factor, + conditional_data_ptr(fused_experts_out), + routed_scaling_factor_value, M, N, K); @@ -1364,8 +1270,8 @@ at::Tensor shared_expert_cpu( hidden_states.data_ptr(), packed_w1.data_ptr(), packed_w2.data_ptr(), - fused_experts_out.data_ptr(), - routed_scaling_factor, + conditional_data_ptr(fused_experts_out), + routed_scaling_factor_value, M, N, K); diff --git a/sgl-kernel/csrc/cpu/moe.h b/sgl-kernel/csrc/cpu/moe.h new file mode 100644 index 000000000..f9d8afca9 --- /dev/null +++ b/sgl-kernel/csrc/cpu/moe.h @@ -0,0 +1,173 @@ +#pragma once +#include "vec.h" + +template +inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) { + using Vec = at::vec::Vectorized; + const Vec data_vec(val); + at::vec::map([data_vec](Vec out) { return out = data_vec; }, out, out, size); +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) { + using Vec = at::vec::Vectorized; + constexpr int kVecSize = Vec::size(); + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= size - kVecSize; d += kVecSize) { + Vec data = Vec::loadu(input + d); + data.store(out + d); + } + for (; d < size; ++d) { + out[d] = input[d]; + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= size - kVecSize; d += kVecSize) { + auto [x0, x1] = load_float_vec2(input + d); + bVec out_vec = convert_from_float_ext(x0, x1); + out_vec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d]); + } +} + +template <> +inline void copy_stub(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) { + // size might be 64x + 32 + std::memcpy(out, input, size * sizeof(uint8_t)); +} + +template +inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict__ input, float weight, int64_t size) { + static_assert( + std::is_same_v || std::is_same_v, + "copy_mul_stub only supports input_t == float or input_t == scalar_t"); + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + 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) { + auto [x0, x1] = load_float_vec2(input + d); + x0 = x0 * weight_vec; + x1 = x1 * weight_vec; + bVec out_vec = convert_from_float_ext(x0, x1); + out_vec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d] * weight); + } +} + +// acc from [topk, K] to [K] +template +inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + if (topk == 1) { + // do copy for topk = 1 + copy_stub(out, input, K); + } else { + // do sum for topk != 1 + int64_t d; +#pragma GCC unroll 4 + for (d = 0; d <= K - kVecSize; d += kVecSize) { + 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); + + sum_fvec0 += x_fvec0; + sum_fvec1 += x_fvec1; + } + bVec out_bvec = convert_from_float_ext(sum_fvec0, sum_fvec1); + out_bvec.store(out + d); + } + for (; d < K; ++d) { + float sum_val = 0.f; + for (int t = 0; t < topk; ++t) { + sum_val += static_cast(input[t * K + d]); + } + out[d] = static_cast(sum_val); + } + } +} + +// out = input + input2 * scale +template +inline void add_mul_stub( + scalar_t* __restrict__ out, + const input_t* __restrict__ input, + const scalar_t* __restrict__ input2, + float scale, + int64_t size) { + static_assert( + std::is_same_v || std::is_same_v, + "add_mul_stub only supports input_t == float or input_t == scalar_t"); + + // out = input (without scale factor) + if (input2 == nullptr) { + copy_stub(out, input, size); + return; + } + + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec s_vec = fVec(scale); + int64_t d; +#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); + + x0 = x0 + y0 * s_vec; + x1 = x1 + y1 * s_vec; + bVec out_vec = convert_from_float_ext(x0, x1); + out_vec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d] + float(input2[d]) * scale); + } +} + +template +inline void silu_and_mul_stub( + scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + 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; + bVec out_vec = convert_from_float_ext(x0, x1); + out_vec.store(out + d); + } +} diff --git a/sgl-kernel/csrc/cpu/moe_fp8.cpp b/sgl-kernel/csrc/cpu/moe_fp8.cpp index 281c00897..fb476b78f 100644 --- a/sgl-kernel/csrc/cpu/moe_fp8.cpp +++ b/sgl-kernel/csrc/cpu/moe_fp8.cpp @@ -1,139 +1,6 @@ #include "common.h" #include "gemm.h" -#include "vec.h" - -namespace { - -template -inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) { - using Vec = at::vec::Vectorized; -// no remainder -#pragma GCC unroll 4 - for (int64_t d = 0; d < size; d += Vec::size()) { - Vec data = Vec::loadu(input + d); - data.store(out + d); - } -} - -template -inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - 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(x0, x1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d] * weight); - } -} - -// acc from [topk, K] to [K] -template -inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - if (topk == 1) { - // do copy for topk = 1 - copy_stub(out, input, K); - } else { - // do sum for topk != 1 - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= K - kVecSize; d += kVecSize) { - 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); - - sum_fvec0 += x_fvec0; - sum_fvec1 += x_fvec1; - } - bVec out_bvec = convert_from_float_ext(sum_fvec0, sum_fvec1); - out_bvec.store(out + d); - } - for (; d < K; ++d) { - float sum_val = 0.f; - for (int t = 0; t < topk; ++t) { - sum_val += static_cast(input[t * K + d]); - } - out[d] = static_cast(sum_val); - } - } -} - -// out = input + input2 * scale -template -inline void add_mul_stub( - scalar_t* __restrict__ out, - const scalar_t* __restrict__ input, - const scalar_t* __restrict__ input2, - float scale, - int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - const fVec s_vec = fVec(scale); - - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input + d); - fVec x0, x1; - std::tie(x0, x1) = at::vec::convert_to_float(x_bvec); - - bVec y_bvec = bVec::loadu(input2 + d); - fVec y0, y1; - std::tie(y0, y1) = at::vec::convert_to_float(y_bvec); - - x0 = x0 + y0 * s_vec; - x1 = x1 + y1 * s_vec; - bVec out_vec = convert_from_float_ext(x0, x1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d] + float(input2[d]) * scale); - } -} - -template -inline void silu_and_mul_stub( - scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - 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; - bVec out_vec = convert_from_float_ext(x0, x1); - out_vec.store(out + d); - } -} - -} // anonymous namespace +#include "moe.h" template void fused_experts_fp8_kernel_impl( @@ -372,6 +239,7 @@ void shared_expert_fp8_kernel_impl( int64_t blocks_n_per_group = block_size_N / BLOCK_N; const bool use_brgemm = can_use_brgemm(M); + const bool apply_scaling_factor = fused_experts_out != nullptr; int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N); @@ -455,9 +323,11 @@ void shared_expert_fp8_kernel_impl( // 2.b copy from C to output and add fused_experts_out scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N; - const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N; + const scalar_t* __restrict__ fused_out = + apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr; for (int64_t m = 0; m < m_size; ++m) { - add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size); + const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr; + add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size); } }); }); diff --git a/sgl-kernel/csrc/cpu/moe_int4.cpp b/sgl-kernel/csrc/cpu/moe_int4.cpp index 10956769b..76403abf7 100644 --- a/sgl-kernel/csrc/cpu/moe_int4.cpp +++ b/sgl-kernel/csrc/cpu/moe_int4.cpp @@ -1,185 +1,19 @@ #include "common.h" #include "gemm.h" -#include "vec.h" -namespace { +#include "moe.h" -template -inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) { - using Vec = at::vec::Vectorized; -// no remainder -#pragma GCC unroll 4 - for (int64_t d = 0; d < size; d += Vec::size()) { - Vec data = Vec::loadu(input + d); - data.store(out + d); - } -} - -template -inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - 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); - bVec out_vec = convert_from_float_ext(x0, x1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d]); - } -} - -template -inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - 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(data0, data1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d] * weight); - } -} - -// acc from [topk, K] to [K] -template -inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - if (topk == 1) { - // do copy for topk = 1 - copy_stub(out, input, K); - } else { - // do sum for topk != 1 - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= K - kVecSize; d += kVecSize) { - 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); - - sum_fvec0 += x_fvec0; - sum_fvec1 += x_fvec1; - } - bVec out_bvec = convert_from_float_ext(sum_fvec0, sum_fvec1); - out_bvec.store(out + d); - } - for (; d < K; ++d) { - float sum_val = 0.f; - for (int t = 0; t < topk; ++t) { - sum_val += static_cast(input[t * K + d]); - } - out[d] = static_cast(sum_val); - } - } -} - -// out = input + input2 * scale -template -inline void add_mul_stub( - scalar_t* __restrict__ out, - const scalar_t* __restrict__ input, - const scalar_t* __restrict__ input2, - float scale, - int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - const fVec s_vec = fVec(scale); - - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= size - kVecSize; d += kVecSize) { - bVec x_bvec = bVec::loadu(input + d); - fVec x0, x1; - std::tie(x0, x1) = at::vec::convert_to_float(x_bvec); - - bVec y_bvec = bVec::loadu(input2 + d); - fVec y0, y1; - std::tie(y0, y1) = at::vec::convert_to_float(y_bvec); - - x0 = x0 + y0 * s_vec; - x1 = x1 + y1 * s_vec; - bVec out_vec = convert_from_float_ext(x0, x1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d] + float(input2[d]) * scale); - } -} - -template -inline void silu_and_mul_stub( - scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - 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; - bVec out_vec = convert_from_float_ext(x0, x1); - out_vec.store(out + d); - } -} - -} // anonymous namespace - -// TODO: stride access template inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m, int64_t ldn) { - if (bias_ptr) { - for (int i = 0; i < m; ++i) { - int j = 0; -#if defined(CPU_CAPABILITY_AVX512) + using Vec = at::vec::Vectorized; + constexpr int kVecSize = Vec::size(); + static_assert(N % kVecSize == 0, "copy_bias requires N to be a multiple of Vectorized::size()"); + const bool has_bias = bias_ptr != nullptr; + const Vec zero_vec(0.f); + for (int i = 0; i < m; ++i) { #pragma GCC unroll 2 - for (; j < N; j += 16) { - __m512 bias_vec = _mm512_loadu_ps(bias_ptr + j); - _mm512_storeu_ps(y_buf + i * ldn + j, bias_vec); - } -#endif - for (; j < N; ++j) { - y_buf[i * ldn + j] = bias_ptr[j]; - } - } - } else { // initialize to zero - for (int i = 0; i < m; ++i) { - int j = 0; -#if defined(CPU_CAPABILITY_AVX512) -#pragma GCC unroll 2 - for (; j < N; j += 16) { - __m512 zero_vec = _mm512_setzero_ps(); - _mm512_storeu_ps(y_buf + i * ldn + j, zero_vec); - } -#endif - for (; j < N; ++j) { - y_buf[i * ldn + j] = 0; - } + for (int j = 0; j < N; j += kVecSize) { + Vec vec = has_bias ? Vec::loadu(bias_ptr + j) : zero_vec; + vec.store(y_buf + i * ldn + j); } } } diff --git a/sgl-kernel/csrc/cpu/moe_int8.cpp b/sgl-kernel/csrc/cpu/moe_int8.cpp index 8fbac902f..95edd4562 100644 --- a/sgl-kernel/csrc/cpu/moe_int8.cpp +++ b/sgl-kernel/csrc/cpu/moe_int8.cpp @@ -1,114 +1,9 @@ #include "common.h" #include "gemm.h" -#include "vec.h" +#include "moe.h" namespace { -template -inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) { - using Vec = at::vec::Vectorized; -// no remainder -#pragma GCC unroll 4 - for (int64_t d = 0; d < size; d += Vec::size()) { - Vec data = Vec::loadu(input + d); - data.store(out + d); - } -} - -template <> -inline void copy_stub(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) { - // size might be 64x + 32 - std::memcpy(out, input, size * sizeof(uint8_t)); -} - -template -inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - 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(data0, data1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d] * weight); - } -} - -// acc from [topk, K] to [K] -template -inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - if (topk == 1) { - // do copy for topk = 1 - copy_stub(out, input, K); - } else { - // do sum for topk != 1 - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= K - kVecSize; d += kVecSize) { - 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); - - sum_fvec0 += x_fvec0; - sum_fvec1 += x_fvec1; - } - bVec out_bvec = convert_from_float_ext(sum_fvec0, sum_fvec1); - out_bvec.store(out + d); - } - for (; d < K; ++d) { - float sum_val = 0.f; - for (int t = 0; t < topk; ++t) { - sum_val += static_cast(input[t * K + d]); - } - out[d] = static_cast(sum_val); - } - } -} - -// out = input + input2 * scale -template -inline void add_mul_stub( - scalar_t* __restrict__ out, - const float* __restrict__ input, - const scalar_t* __restrict__ input2, - float scale, - int64_t size) { - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int kVecSize = bVec::size(); - const fVec s_vec = fVec(scale); - int64_t d; -#pragma GCC unroll 4 - for (d = 0; d <= size - kVecSize; d += kVecSize) { - fVec x0 = fVec::loadu(input + d); - fVec x1 = fVec::loadu(input + d + fVec::size()); - - bVec y_bvec = bVec::loadu(input2 + d); - fVec y0, y1; - std::tie(y0, y1) = at::vec::convert_to_float(y_bvec); - - x0 = x0 + y0 * s_vec; - x1 = x1 + y1 * s_vec; - bVec out_vec = convert_from_float_ext(x0, x1); - out_vec.store(out + d); - } - for (; d < size; ++d) { - out[d] = static_cast(input[d] + float(input2[d]) * scale); - } -} - template inline void silu_and_mul( scalar_t* __restrict__ C, @@ -885,6 +780,7 @@ void shared_expert_int8_kernel_impl( const int64_t stride_n = packed_K; const bool use_brgemm = can_use_brgemm(M); + const bool apply_scaling_factor = fused_experts_out != nullptr; // here we only parallel on half of 2N to fuse silu_and_mul with gemm parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { @@ -1034,9 +930,11 @@ void shared_expert_int8_kernel_impl( // 2.b copy from C to output and add fused_experts_out scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N; - const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N; + const scalar_t* __restrict__ fused_out = + apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr; for (int64_t m = 0; m < m_size; ++m) { - add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size); + const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr; + add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size); } }); diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index e5f279866..197dcbd86 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -226,8 +226,8 @@ at::Tensor shared_expert_cpu( at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2, - at::Tensor& fused_experts_out, - double routed_scaling_factor, + const std::optional& fused_experts_out, + const std::optional routed_scaling_factor, bool inplace, bool use_int8_w8a8, bool use_fp8_w8a16, @@ -554,7 +554,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { // shared expert m.def( - "shared_expert_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor fused_experts_out, float " + "shared_expert_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor? fused_experts_out, float? " "routed_scaling_factor, bool inplace, bool use_int8_w8a8, bool use_fp8_w8a16, Tensor? w1_scale, Tensor? " "w2_scale, int[]? block_size, bool is_vnni) -> Tensor"); m.impl("shared_expert_cpu", torch::kCPU, &shared_expert_cpu); diff --git a/test/srt/cpu/test_moe.py b/test/srt/cpu/test_moe.py index a1f6bf288..10cc0fd9d 100644 --- a/test/srt/cpu/test_moe.py +++ b/test/srt/cpu/test_moe.py @@ -300,35 +300,16 @@ class TestFusedExperts(CustomTestCase): ) score = torch.softmax(score, dim=-1, dtype=torch.float32) topk_weight, topk_ids = torch.topk(score, topk) - awq_w13_weight_pack = [] - awq_w13_zero_pack = [] - awq_w13_scales_pack = [] - awq_w2_weight_pack = [] - awq_w2_zero_pack = [] - awq_w2_scales_pack = [] - for i in range(E): - packed_weight_13_i, packed_zero_13_i, packed_scales_13_i = ( - torch.ops.sgl_kernel.convert_weight_packed_scale_zp( - awq_w13_weight[i], awq_w13_zero[i], awq_w13_scales[i] - ) + 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 ) - awq_w13_weight_pack.append(packed_weight_13_i) - awq_w13_zero_pack.append(packed_zero_13_i) - awq_w13_scales_pack.append(packed_scales_13_i) - packed_weight_2_i, packed_zero_2_i, packed_scales_2_i = ( - torch.ops.sgl_kernel.convert_weight_packed_scale_zp( - awq_w2_weight[i], awq_w2_zero[i], awq_w2_scales[i] - ) + ) + awq_w2_weight_pack, awq_w2_zero_pack, awq_w2_scales_pack = ( + torch.ops.sgl_kernel.convert_weight_packed_scale_zp( + awq_w2_weight, awq_w2_zero, awq_w2_scales ) - awq_w2_weight_pack.append(packed_weight_2_i) - awq_w2_zero_pack.append(packed_zero_2_i) - awq_w2_scales_pack.append(packed_scales_2_i) - awq_w13_weight_pack = torch.stack(awq_w13_weight_pack).detach() - awq_w13_zero_pack = torch.stack(awq_w13_zero_pack).detach() - awq_w13_scales_pack = torch.stack(awq_w13_scales_pack).detach() - awq_w2_weight_pack = torch.stack(awq_w2_weight_pack).detach() - awq_w2_zero_pack = torch.stack(awq_w2_zero_pack).detach() - awq_w2_scales_pack = torch.stack(awq_w2_scales_pack).detach() + ) out = kernel.fused_experts_cpu( a, diff --git a/test/srt/cpu/test_shared_expert.py b/test/srt/cpu/test_shared_expert.py index f5a8f3632..563916131 100644 --- a/test/srt/cpu/test_shared_expert.py +++ b/test/srt/cpu/test_shared_expert.py @@ -2,12 +2,10 @@ import itertools import math import unittest -# TODO: use interface in cpu.py import torch from utils import ( BLOCK_K, BLOCK_N, - SiluAndMul, factor_for_scale, fp8_max, fp8_min, @@ -18,7 +16,6 @@ from utils import ( torch_w8a8_per_column_moe, ) -from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.test.test_utils import CustomTestCase torch.manual_seed(1234) @@ -29,37 +26,41 @@ class TestSharedExpert(CustomTestCase): N = [32, 32 * 4] K = [32, 32 * 2] routed_scaling_factor = [16] + apply_scaling_factor = [True, False] M_fp8 = [2, 12] N_fp8 = [512] K_fp8 = [256] - def _bf16_shared_expert(self, m, n, k, routed_scaling_factor): + def _bf16_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor): dtype = torch.bfloat16 - prepack = True hidden_states = torch.randn(m, k, dtype=dtype) / k w1 = torch.randn(2 * n, k, dtype=dtype) w2 = torch.randn(k, n, dtype=dtype) - fused_output = torch.randn(m, k, dtype=dtype) / k + fused_output = ( + torch.randn(m, k, dtype=dtype) / k if apply_scaling_factor else None + ) + routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None # fused moe mutates content in hs hidden_states2 = hidden_states.clone() # bfloat16 ref = torch_naive_moe( - hidden_states.float(), - w1.float(), - w2.float(), - fused_output.float(), - routed_scaling_factor, - ).to(dtype=dtype) - res = torch.ops.sgl_kernel.shared_expert_cpu( hidden_states, w1, w2, fused_output, routed_scaling_factor, + output_dtype=dtype, + ) + out = torch.ops.sgl_kernel.shared_expert_cpu( + hidden_states2, + w1, + w2, + fused_output, + routed_scaling_factor, True, False, False, @@ -70,7 +71,7 @@ class TestSharedExpert(CustomTestCase): ) atol = rtol = precision[ref.dtype] - torch.testing.assert_close(ref, res, atol=atol, rtol=rtol) + torch.testing.assert_close(ref, out, atol=atol, rtol=rtol) def test_bf16_shared_expert(self): for params in itertools.product( @@ -78,39 +79,43 @@ class TestSharedExpert(CustomTestCase): self.N, self.K, self.routed_scaling_factor, + self.apply_scaling_factor, ): with self.subTest( m=params[0], n=params[1], k=params[2], routed_scaling_factor=params[3], + apply_scaling_factor=params[4], ): self._bf16_shared_expert(*params) - def _int8_shared_expert(self, m, n, k, routed_scaling_factor): + def _int8_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor): dtype = torch.bfloat16 - prepack = True hidden_states = torch.randn(m, k, dtype=dtype) / k w1 = torch.randn(2 * n, k, dtype=dtype) w2 = torch.randn(k, n, dtype=dtype) - fused_output = torch.randn(m, k, dtype=dtype) / k + fused_output = ( + torch.randn(m, k, dtype=dtype) / k if apply_scaling_factor else None + ) + routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None # fused moe mutates content in hs hidden_states2 = hidden_states.clone() w1_q, w1_s = per_token_quant_int8(w1) w2_q, w2_s = per_token_quant_int8(w2) - ref2 = torch_w8a8_per_column_moe( - hidden_states2.float(), + ref = torch_w8a8_per_column_moe( + hidden_states, w1_q, w2_q, w1_s, w2_s, - fused_output.float(), + fused_output, routed_scaling_factor, - ).to(dtype=dtype) - res2 = torch.ops.sgl_kernel.shared_expert_cpu( + ) + out = torch.ops.sgl_kernel.shared_expert_cpu( hidden_states2, w1_q, w2_q, @@ -125,8 +130,8 @@ class TestSharedExpert(CustomTestCase): False, ) - atol = rtol = precision[ref2.dtype] - torch.testing.assert_close(ref2, res2, atol=atol, rtol=rtol) + atol = rtol = precision[ref.dtype] + torch.testing.assert_close(ref, out, atol=atol, rtol=rtol) def test_int8_shared_expert(self): for params in itertools.product( @@ -134,57 +139,64 @@ class TestSharedExpert(CustomTestCase): self.N, self.K, self.routed_scaling_factor, + self.apply_scaling_factor, ): with self.subTest( m=params[0], n=params[1], k=params[2], routed_scaling_factor=params[3], + apply_scaling_factor=params[4], ): self._int8_shared_expert(*params) - def _fp8_shared_expert(self, M, N, K, routed_scaling_factor): - set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) - + def _fp8_shared_expert(self, m, n, k, routed_scaling_factor, apply_scaling_factor): dtype = torch.bfloat16 - prepack = True - a = torch.randn(M, K, dtype=dtype) / math.sqrt(K) + hidden_states = torch.randn(m, k, dtype=dtype) / math.sqrt(k) - w1_fp32 = torch.randn(1, 2 * N, K) + w1_fp32 = torch.randn(1, 2 * n, k) w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) - w2_fp32 = torch.randn(1, K, N) + w2_fp32 = torch.randn(1, k, n) w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) - w1s = torch.randn(1, 2 * N // BLOCK_N, K // BLOCK_K) * factor_for_scale - w2s = torch.randn(1, K // BLOCK_N, N // BLOCK_K) * factor_for_scale + w1s = torch.randn(1, 2 * n // BLOCK_N, k // BLOCK_K) * factor_for_scale + w2s = torch.randn(1, k // BLOCK_N, n // BLOCK_K) * factor_for_scale - w1_scaled = scaled_weight(w1, w1s).view(2 * N, K) - w2_scaled = scaled_weight(w2, w2s).view(K, N) + w1_scaled = scaled_weight(w1, w1s).view(2 * n, k) + w2_scaled = scaled_weight(w2, w2s).view(k, n) # change back to 2D w1, w2 = w1.squeeze(0), w2.squeeze(0) w1s, w2s = w1s.squeeze(0), w2s.squeeze(0) w1_scaled, w2_scaled = w1_scaled.squeeze(0), w2_scaled.squeeze(0) - fused_out = torch.randn(M, K, dtype=dtype) / math.sqrt(K) - a2 = a.clone() + fused_output = ( + torch.randn(m, k, dtype=dtype) / math.sqrt(k) + if apply_scaling_factor + else None + ) + routed_scaling_factor = routed_scaling_factor if apply_scaling_factor else None + hidden_states2 = hidden_states.clone() - # ref - ic0 = torch.matmul(a.float(), w1_scaled.transpose(0, 1)) - ic1 = SiluAndMul(ic0) - shared_out = torch.matmul(ic1, w2_scaled.transpose(0, 1)) - ref_out = shared_out + fused_out.float() * routed_scaling_factor - ref_out = ref_out.to(dtype=dtype) + # ref with bfloat16 + ref = torch_naive_moe( + hidden_states, + w1_scaled, + w2_scaled, + fused_output, + routed_scaling_factor, + output_dtype=dtype, + ) w1 = torch.ops.sgl_kernel.convert_weight_packed(w1) # [2N, K] w2 = torch.ops.sgl_kernel.convert_weight_packed(w2) # [K, N] out = torch.ops.sgl_kernel.shared_expert_cpu( - a2, + hidden_states2, w1, w2, - fused_out, + fused_output, routed_scaling_factor, True, False, @@ -195,8 +207,8 @@ class TestSharedExpert(CustomTestCase): True, ) - atol = rtol = precision[ref_out.dtype] - torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol) + atol = rtol = precision[ref.dtype] + torch.testing.assert_close(ref, out, atol=atol, rtol=rtol) def test_fp8_shared_expert(self): for params in itertools.product( @@ -204,12 +216,14 @@ class TestSharedExpert(CustomTestCase): self.N_fp8, self.K_fp8, self.routed_scaling_factor, + self.apply_scaling_factor, ): with self.subTest( - M=params[0], - N=params[1], - K=params[2], + m=params[0], + n=params[1], + k=params[2], routed_scaling_factor=params[3], + apply_scaling_factor=params[4], ): self._fp8_shared_expert(*params) diff --git a/test/srt/cpu/utils.py b/test/srt/cpu/utils.py index f90967ed7..eccdccbb1 100644 --- a/test/srt/cpu/utils.py +++ b/test/srt/cpu/utils.py @@ -126,16 +126,28 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, bias, output_dtype=torch.bfloat16 return C.reshape(origin_C_shape).to(output_dtype) -def torch_naive_moe(a, w1, w2, b, routed_scaling_factor): +def torch_naive_moe(a, w1, w2, b, routed_scaling_factor, output_dtype=torch.bfloat16): + + a = a.to(torch.float32) + w1 = w1.to(torch.float32) + w2 = w2.to(torch.float32) + b = b.to(torch.float32) if b is not None else None ic1 = torch.matmul(a, w1.transpose(0, 1)) ic2 = SiluAndMul(ic1) ic3 = torch.matmul(ic2, w2.transpose(0, 1)) - return ic3 + b * routed_scaling_factor + out = ic3 if b is None else ic3 + b * routed_scaling_factor + + return out.to(output_dtype) -def torch_w8a8_per_column_moe(a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor): +def torch_w8a8_per_column_moe( + a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor, output_dtype=torch.bfloat16 +): + + a = a.to(torch.float32) + b = b.to(torch.float32) if b is not None else None # Perform per-token quantization a_q, a_s = per_token_quant_int8(a) @@ -150,7 +162,9 @@ def torch_w8a8_per_column_moe(a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_facto a1_q, w2_q, a1_s, w2_s, bias=None, output_dtype=torch.float32 ) - return ic3 + b * routed_scaling_factor + out = ic3 if b is None else ic3 + b * routed_scaling_factor + + return out.to(output_dtype) def scaled_weight(weight, scales):