[CPU] expand the interface of shared_expert without scaling factor (#22933)

merge since this is CPU only change on sgl-kernel.
This commit is contained in:
Ma Mingfei
2026-04-21 20:03:39 +08:00
committed by GitHub
parent 48daa831ea
commit 929e00eeab
9 changed files with 313 additions and 623 deletions
+3 -3
View File
@@ -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<at::Tensor>& fused_experts_out,
const std::optional<double> 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);