chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add fused_rmsnorm_gated_cpu kernel for CPU to support Qwen3-Next (#11577)

Browse Source
This commit is contained in:
YanbingJiang
2025-11-21 01:33:31 +08:00
committed by GitHub
parent 4528cb7d40
commit acde21d8d5
3 changed files with 164 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+5
View File
@@ -33,6 +33,9 @@ at::Tensor l2norm_cpu(at::Tensor& input, double eps);
// rmsnorm
at::Tensor rmsnorm_cpu(at::Tensor& input, at::Tensor& weight, double eps);
// qwen3_next_rmsnorm_gated
at::Tensor fused_rmsnorm_gated_cpu(at::Tensor& input, at::Tensor& weight, at::Tensor& gate, double eps);
// fused_add_rmsnorm
void fused_add_rmsnorm_cpu(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps);
@@ -247,6 +250,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.impl("rmsnorm_cpu", torch::kCPU, &rmsnorm_cpu);
m.def("l2norm_cpu(Tensor input, float eps) -> Tensor");
m.impl("l2norm_cpu", torch::kCPU, &l2norm_cpu);
m.def("fused_rmsnorm_gated_cpu(Tensor input, Tensor weight, Tensor gate, float eps) -> Tensor");
m.impl("fused_rmsnorm_gated_cpu", torch::kCPU, &fused_rmsnorm_gated_cpu);
m.def("fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor residual, Tensor weight, float eps) -> ()");
m.impl("fused_add_rmsnorm_cpu", torch::kCPU, &fused_add_rmsnorm_cpu);