[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:
@@ -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 "
|
||||
|
||||
Reference in New Issue
Block a user