[sgl-kernel/cpu] support w8a8 int8 model for arm cpu (#16045)

skip gpu test as this one is not related to gpu backend.
This commit is contained in:
Yibo Cai
2026-05-08 14:47:06 +08:00
committed by GitHub
parent 47e9ec11ad
commit 55d8223c2b
7 changed files with 832 additions and 16 deletions
+2 -2
View File
@@ -211,7 +211,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<at::Tensor>& scale);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
// fused moe
at::Tensor fused_experts_cpu(
at::Tensor& hidden_states,
@@ -228,6 +227,7 @@ at::Tensor fused_experts_cpu(
const std::optional<std::vector<int64_t>> block_size,
bool is_vnni);
#if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
at::Tensor shared_expert_cpu(
at::Tensor& hidden_states,
at::Tensor& w1,
@@ -546,7 +546,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 "
@@ -554,6 +553,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, 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 "