[CPU] support LayerNorm with 3D shape (#15075)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Zaili Wang
2026-03-19 13:15:24 +08:00
committed by GitHub
parent dc6aa26ce9
commit 2f4babe32b
4 changed files with 196 additions and 54 deletions

View File

@@ -421,8 +421,9 @@ class LayerNorm(MultiPlatformOp):
x: torch.Tensor,
) -> torch.Tensor:
if _is_cpu_amx_available:
bias_data = self.bias.data if self.use_bias else None
return torch.ops.sgl_kernel.layernorm_cpu(
x, self.weight.data, self.variance_epsilon
x, self.weight.data, bias_data, self.variance_epsilon
)
else:
return self.forward_native(x)