[CPU] layernorm & fused add-layernorm kernels (#14074)

2025-12-12 08:58:23 +08:00
parent 9975acf50f
commit d6bd2d1126
5 changed files with 258 additions and 10 deletions
--- a/python/sglang/srt/layers/layernorm.py
+++ b/python/sglang/srt/layers/layernorm.py
@@ -361,7 +361,12 @@ class LayerNorm(CustomOp):
        self,
        x: torch.Tensor,
    ) -> torch.Tensor:
-        return self.forward_native(x)
+        if _is_cpu_amx_available:
+            return torch.ops.sgl_kernel.layernorm_cpu(
+                x, self.weight.data, self.variance_epsilon
+            )
+        else:
+            return self.forward_native(x)


 class GemmaRMSNorm(CustomOp):