fix: Fix AMD CI failures with HIP layernorm and PyPI connectivity (#13814)

Co-authored-by: root <root@mi300x8-005.atl1.do.cpe.ice.amd.com>
This commit is contained in:
sunxxuns
2025-11-26 22:30:37 -05:00
committed by GitHub
parent 9f340ab1fb
commit 5443db8759
2 changed files with 43 additions and 8 deletions

View File

@@ -165,8 +165,12 @@ class RMSNorm(CustomOp):
# NOTE: Remove this if aiter kernel supports discontinuous input
x = x.contiguous()
if residual is not None:
fused_add_rms_norm(x, residual, self.weight.data, self.variance_epsilon)
return x, residual
out = torch.empty_like(x)
residual_out = torch.empty_like(x)
fused_add_rms_norm(
out, x, residual_out, residual, self.weight.data, self.variance_epsilon
)
return out, residual_out
out = torch.empty_like(x)
rms_norm(out, x, self.weight.data, self.variance_epsilon)
return out