diff --git a/sgl-kernel/csrc/cpu/aarch64/moe.cpp b/sgl-kernel/csrc/cpu/aarch64/moe.cpp index 3e0774376..171481fd3 100644 --- a/sgl-kernel/csrc/cpu/aarch64/moe.cpp +++ b/sgl-kernel/csrc/cpu/aarch64/moe.cpp @@ -238,6 +238,10 @@ at::Tensor fused_experts_cpu( const std::optional& /*w13_zero*/, const std::optional& /*w2_zero*/, const std::optional> block_size, + const std::optional& /*w1_bias*/, + const std::optional& /*w2_bias*/, + const std::optional& /*alpha*/, + const std::optional& /*limit*/, bool /*is_vnni*/) { const auto st = hidden_states.scalar_type(); CHECK_INPUT(hidden_states); diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 4885bba8d..3a52d9938 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -220,8 +220,6 @@ at::Tensor fused_linear_sigmoid_mul( // bmm void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional& scale); -#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS) - // fused moe at::Tensor fused_experts_cpu( at::Tensor& hidden_states, @@ -241,6 +239,8 @@ at::Tensor fused_experts_cpu( const std::optional& alpha, const std::optional& limit, bool is_vnni); + +#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS) at::Tensor shared_expert_cpu( at::Tensor& hidden_states, at::Tensor& w1, @@ -322,23 +322,6 @@ at::Tensor causal_conv1d_update_cpu( const std::optional& conv_state_indices, int64_t pad_slot_id, bool is_vnni); -#else - -// fused moe -at::Tensor fused_experts_cpu( - at::Tensor& hidden_states, - at::Tensor& w1, - at::Tensor& w2, - at::Tensor& topk_weights, - at::Tensor& topk_ids, - bool inplace, - int64_t moe_comp_method, - const std::optional& w1_scale, - const std::optional& w2_scale, - const std::optional& w1_zero, - const std::optional& w2_zero, - const std::optional> block_size, - bool is_vnni); #endif // conv3d fast path for patch embedding @@ -588,7 +571,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, Tensor? scale) -> ()"); m.impl("bmm_cpu", torch::kCPU, &bmm_cpu); -#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS) // moe m.def( "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor topk_weights, Tensor topk_ids, bool " @@ -596,6 +578,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? " "limit, bool is_vnni) -> Tensor"); m.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu); + +#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS) // weight absorption m.def( "qkv_proj_with_rope(Tensor hidden_states, Tensor q_a_proj_weight, Tensor q_b_proj_weight, Tensor " @@ -634,13 +618,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor weight, Tensor? bias, bool silu_activation," "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, bool is_vnni) -> Tensor"); m.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); -#else - // moe - m.def( - "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor topk_weights, Tensor topk_ids, bool " - "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, " - "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> Tensor"); - m.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu); #endif // conv3d fast path for patch embedding diff --git a/test/registered/cpu/arm64/test_moe.py b/test/registered/cpu/arm64/test_moe.py index 9d6c3c9e6..e700fc918 100644 --- a/test/registered/cpu/arm64/test_moe.py +++ b/test/registered/cpu/arm64/test_moe.py @@ -86,6 +86,10 @@ class TestFusedExpertsInt8(CustomTestCase): None, None, None, + None, + None, + None, + None, prepack, )