From 70d25873246bb02335b0a107575e289a35662f96 Mon Sep 17 00:00:00 2001 From: jianan-gu Date: Thu, 4 Dec 2025 16:38:47 +0800 Subject: [PATCH] [CPU] Optimize small oc GEMM for Qwen3-next on CPU (#12446) Co-authored-by: Zheng, Beilei --- python/sglang/srt/layers/amx_utils.py | 8 +- python/sglang/srt/models/qwen2_moe.py | 39 ++- sgl-kernel/csrc/cpu/gemm.cpp | 271 +++++++++++++++++++- sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 13 + test/srt/cpu/test_gemm.py | 44 ++++ 5 files changed, 357 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/layers/amx_utils.py b/python/sglang/srt/layers/amx_utils.py index df2a05ba5..8e1209ea0 100644 --- a/python/sglang/srt/layers/amx_utils.py +++ b/python/sglang/srt/layers/amx_utils.py @@ -17,15 +17,17 @@ def amx_process_weight_after_loading(weight): # TODO: currently gemm kernel has the below requirements: -# OC % TILE_N == 0, where TILE_N = 16 -# IC % TILE_K == 0, where TILE_K = 32 +# OC: OC % TILE_N == 0 or OC < TILE_N, where TILE_N = 16 +# IC: IC % TILE_K == 0, where TILE_K = 32 def dim_is_supported(weight): TILE_N = 16 TILE_K = 32 ndim = weight.ndim OC = weight.size(1) if ndim == 3 else weight.size(0) IC = weight.size(2) if ndim == 3 else weight.size(1) - return OC % TILE_N == 0 and IC % TILE_K == 0 + is_oc_support = OC < TILE_N or OC % TILE_N == 0 + is_ic_support = IC % TILE_K == 0 + return is_oc_support and is_ic_support def _amx_process_weight_after_loading( diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 656968192..caeae9c23 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -71,11 +71,20 @@ from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import add_prefix, is_cuda, make_layers +from sglang.srt.utils import ( + add_prefix, + cpu_has_amx_support, + is_cpu, + is_cuda, + make_layers, + use_intel_amx_backend, +) logger = logging.getLogger(__name__) _is_cuda = is_cuda() +_is_cpu = is_cpu() +_is_cpu_amx_available = cpu_has_amx_support() class Qwen2MoeMLP(nn.Module): @@ -189,7 +198,16 @@ class Qwen2MoeSparseMoeBlock(nn.Module): ) else: self.shared_expert = None - self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) + if _is_cpu and _is_cpu_amx_available: + self.shared_expert_gate = ReplicatedLinear( + config.hidden_size, + 1, + bias=False, + quant_config=None, + prefix=add_prefix("shared_expert_gate", prefix), + ) + else: + self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) if get_moe_a2a_backend().is_deepep(): # TODO: we will support tp < ep in the future @@ -211,9 +229,20 @@ class Qwen2MoeSparseMoeBlock(nn.Module): if self.shared_expert is not None: shared_output = self.shared_expert(hidden_states) if self.shared_expert_gate is not None: - shared_output = ( - F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_output - ) + if use_intel_amx_backend(self.shared_expert_gate): + shared_output = torch.ops.sgl_kernel.fused_linear_sigmoid_mul( + hidden_states, + self.shared_expert_gate.weight, + self.shared_expert_gate.bias, + True, + shared_output, + ) + else: + shared_output = ( + F.sigmoid(self.shared_expert_gate(hidden_states)) + * shared_output + ) + return shared_output def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): diff --git a/sgl-kernel/csrc/cpu/gemm.cpp b/sgl-kernel/csrc/cpu/gemm.cpp index 48655b9f7..e2fdc8951 100644 --- a/sgl-kernel/csrc/cpu/gemm.cpp +++ b/sgl-kernel/csrc/cpu/gemm.cpp @@ -84,6 +84,26 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu } } +template +inline void copy_stub(float* __restrict__ out, const scalar_t* __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) { + fVec data0, data1; + bVec b_vec = bVec::loadu(input + d); + std::tie(data0, data1) = at::vec::convert_to_float(b_vec); + data0.store(out + d); + data1.store(out + d + fVec::size()); + } + for (; d < size; ++d) { + out[d] = static_cast(input[d]); + } +} + template inline void copy_add_stub( scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) { @@ -104,6 +124,40 @@ inline void copy_add_stub( } } +template +inline void scalar_sigmoid_and_mul( + scalar_t* __restrict__ out, + const float* __restrict__ input, + const float* __restrict__ bias, + const scalar_t* __restrict__ mul, + int SIZE) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + // 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()); + + // 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(m_fvec0, m_fvec1); + out_vec.store(out + d); + } +} + template struct tinygemm_kernel_nn { static inline void apply( @@ -233,6 +287,21 @@ struct brgemm { } } } + static inline void apply( + const float* __restrict__ A, + const float* __restrict__ B, + scalar_t* __restrict__ C, + float* __restrict__ Ctmp, + const float* __restrict__ bias, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc) { + constexpr int BLOCK_N = block_size_n(); + at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp); + } }; template @@ -326,6 +395,28 @@ void tinygemm_kernel( } } +template +void tinygemm_kernel( + const float* __restrict__ A, + const float* __restrict__ B, + scalar_t* __restrict__ C, + float* __restrict__ Ctmp, + const float* __restrict__ bias, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + bool brg) { + TORCH_CHECK(brg, "Expected to use fp32 brgemm for small N GEMM"); + if (brg) { + brgemm::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc); + return; + } + // TODO : add intrinsic path +} + template void weight_packed_linear_kernel_impl( scalar_t* __restrict__ out, @@ -378,6 +469,81 @@ void weight_packed_linear_kernel_impl( }); } +template +void weight_packed_linear_kernel_impl( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ mat1, + const float* __restrict__ mat2, + const float* __restrict__ bias, + const scalar_t* __restrict__ post_mul_mat, + int64_t M, + int64_t N, + int64_t K, + int64_t mat1_strideM, + int64_t out_strideM) { + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N = block_size_n(); + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + const bool use_brgemm = true; // TODO: add intrinsic path + // parallel on [MB, NB] + AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] { + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { + // for brgemm, use float32 for accumulate + alignas(64) float Atmp[BLOCK_M * K]; + alignas(64) float Ctmp[BLOCK_M * BLOCK_N]; + + loop_2d(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) { + int64_t mb_start = mb * BLOCK_M; + int64_t mb_size = std::min(M - mb_start, BLOCK_M); + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(N - nb_start, BLOCK_N); + for (int64_t m = 0; m < mb_size; ++m) { + copy_stub(Atmp + m * K, mat1 + mb_start * mat1_strideM + m * K, K); + } + tinygemm_kernel( + /* A */ Atmp, + /* B */ mat2 + nb_start * K /* nb * BLOCK_N * K */, + /* C */ out + mb_start * out_strideM + nb_start, + /* Ctmp*/ Ctmp, + /* bias*/ bias + nb_start, + /* M */ mb_size, + /* N */ nb_size, + /* K */ K, + /* lda */ mat1_strideM, + /* ldb */ nb_size, + /* ldc */ out_strideM, + /* brg */ use_brgemm); + + if (post_mul_mat != nullptr) { + for (int64_t m = 0; m < mb_size; ++m) { + scalar_sigmoid_and_mul( + out + mb_start * out_strideM + nb_start + m * out_strideM, + Ctmp + m * BLOCK_N, + bias + nb_start, + post_mul_mat + mb_start * out_strideM + m * out_strideM, + out_strideM); + } + } else { + for (int64_t m = 0; m < mb_size; ++m) { + if constexpr (has_bias) { + copy_add_stub( + out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, bias + nb_start, N); + } else { + copy_stub(out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, N); + } + } + } + }); + + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); + }); +} + } // anonymous namespace // tinygemm interface @@ -423,6 +589,12 @@ at::Tensor convert_weight_packed(at::Tensor& weight) { const int64_t ndim = weight.ndimension(); TORCH_CHECK(ndim == 2 || ndim == 3, "expect weight to be 2d or 3d, got ", ndim, "d tensor."); + + if (ndim == 2 && weight.size(0) < TILE_N) { + // for 2D weight and small OC shape, we use fma linear path, which needs transpose not pack + return weight.to(at::kFloat).t().contiguous(); + } + const auto st = weight.scalar_type(); const int64_t E = ndim == 3 ? weight.size(0) : 1; const int64_t OC = ndim == 3 ? weight.size(1) : weight.size(0); @@ -475,7 +647,7 @@ at::Tensor convert_weight_packed(at::Tensor& weight) { } // mat1 : [M, K] -// mat2 : [N, K] +// mat2 : [N, K] ([K, N] if use_fma_gemm) // bias : [N] // out : [M, N] // @@ -484,22 +656,28 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional({mat1, mat2, bias})); auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2); + bool use_fma_gemm = false; + if (packed_w.scalar_type() == at::kFloat) { + use_fma_gemm = true; + } + + int64_t M = mat1.size(0); + int64_t K = mat1.size(1); + int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0); CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1); CHECK_INPUT(mat2); - - int64_t M = mat1.size(0); - int64_t N = mat2.size(0); - int64_t K = mat2.size(1); - CHECK_EQ(mat1.size(1), K); CHECK_DIM(2, mat1); CHECK_DIM(2, mat2); + if (!use_fma_gemm) { + CHECK_EQ(mat1.size(1), K); + } + auto dispatch_type = mat1.scalar_type(); auto out = at::empty({M, N}, mat1.options()); - // strides - int64_t mat1_strideM = mat1.stride(0); int64_t out_strideM = out.stride(0); + int64_t mat1_strideM = mat1.stride(0); const bool has_bias = bias.has_value(); const float* bias_data = nullptr; @@ -508,12 +686,85 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional(); } - AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "weight_packed_linear_kernel_impl", [&] { + AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "weight_packed_linear_kernel_impl", [&] { + if (use_fma_gemm) { + weight_packed_linear_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), + bias_data, + nullptr, + M, + N, + K, + mat1_strideM, + out_strideM); + } else { + weight_packed_linear_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), + bias_data, + M, + N, + K, + mat1_strideM, + out_strideM); + } + }); + + return out; +} + +// mat1 : [M, K] +// mat2 : [K, 1] +// post_mul_mat : [M, K] +// bias : [N] +// out : [M, N] +// +at::Tensor fused_linear_sigmoid_mul( + at::Tensor& mat1, + at::Tensor& mat2, + const std::optional& bias, + bool is_vnni, + const at::Tensor& post_mul_mat) { + RECORD_FUNCTION("sgl-kernel::fused_linear_sigmoid_mul", std::vector({mat1, mat2, bias, post_mul_mat})); + + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2); + TORCH_CHECK(packed_w.scalar_type() == at::kFloat, "fused_linear_sigmoid_mul requires packed float weight") + + int64_t M = mat1.size(0); + int64_t K = mat1.size(1); + int64_t N = mat2.size(1); + + CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1); + CHECK_INPUT(mat2); + CHECK_DIM(2, mat1); + CHECK_DIM(2, mat2); + + int64_t out_strideM = post_mul_mat.size(1); + int64_t mat1_strideM = mat1.stride(0); + auto dispatch_type = mat1.scalar_type(); + auto out = at::empty({M, out_strideM}, mat1.options()); + + TORCH_CHECK( + N == 1 && out_strideM % 32 == 0, + "post_mul_mat tensor size(1) should be 32 dividable, and the mat2 OC=1 (Mx1 as linear output shape)") + + const bool has_bias = bias.has_value(); + const float* bias_data = nullptr; + if (has_bias) { + CHECK_EQ(bias.value().size(0), N); + bias_data = bias.value().data_ptr(); + } + + AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "fused_linear_sigmoid_mul", [&] { weight_packed_linear_kernel_impl( out.data_ptr(), mat1.data_ptr(), - packed_w.data_ptr(), + packed_w.data_ptr(), bias_data, + post_mul_mat.data_ptr(), M, N, K, diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 44c32659e..b39850a69 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -124,6 +124,14 @@ std::tuple per_token_quant_int8_cpu(at::Tensor& A); at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional& bias, bool is_vnni); +// gemm fusion +at::Tensor fused_linear_sigmoid_mul( + at::Tensor& mat1, + at::Tensor& mat2, + const std::optional& bias, + bool is_vnni, + const at::Tensor& post_mul_mat); + // igemm at::Tensor int8_scaled_mm_cpu( at::Tensor& mat1, @@ -355,6 +363,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("weight_packed_linear(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni) -> Tensor"); m.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear); + // gemm fusion + m.def( + "fused_linear_sigmoid_mul(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni, Tensor post_mul_mat) -> Tensor"); + m.impl("fused_linear_sigmoid_mul", torch::kCPU, &fused_linear_sigmoid_mul); + // igemm m.def( "int8_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales1, Tensor scales2, Tensor? bias, ScalarType " diff --git a/test/srt/cpu/test_gemm.py b/test/srt/cpu/test_gemm.py index b9ec1e7bf..0d3113c6e 100644 --- a/test/srt/cpu/test_gemm.py +++ b/test/srt/cpu/test_gemm.py @@ -79,6 +79,50 @@ class TestGemm(CustomTestCase): ): self._bf16_gemm(*params) + def _bf16_gemm_with_small_oc(self, M, N, K, has_bias, use_post_sigmul): + use_post_sigmul = use_post_sigmul and N == 1 + mat_mul = ( + None if not use_post_sigmul else torch.randn(M, 2 * K, dtype=torch.bfloat16) + ) + mat1 = torch.randn(M, K, dtype=torch.bfloat16) + mat2 = torch.randn(N, K, dtype=torch.bfloat16) + + ref = torch.nn.functional.linear(mat1, mat2) + if has_bias: + bias = torch.randn(N, dtype=torch.float32) + ref.add_(bias) + if use_post_sigmul: + ref = torch.nn.functional.sigmoid(ref) * mat_mul + out = torch.ops.sgl_kernel.fused_linear_sigmoid_mul( + mat1, + torch.ops.sgl_kernel.convert_weight_packed(mat2), + bias if has_bias else None, + True, + mat_mul if use_post_sigmul else None, + ) + else: + out = torch.ops.sgl_kernel.weight_packed_linear( + mat1, + torch.ops.sgl_kernel.convert_weight_packed(mat2), + bias if has_bias else None, + True, + ) + atol = rtol = precision[ref.dtype] + torch.testing.assert_close(ref, out, atol=atol, rtol=rtol) + + def test_bf16_gemm_with_small_oc(self): + for params in itertools.product( + [1, 8, 32, 1024], [12, 1], self.K, self.has_bias, [False, True] + ): + with self.subTest( + M=params[0], + N=params[1], + K=params[2], + has_bias=params[3], + use_post_sigmul=params[4], + ): + self._bf16_gemm_with_small_oc(*params) + def _int8_gemm(self, M, N, K, has_bias): dtype = torch.bfloat16 A = torch.randn((M, K), dtype=dtype) / 10