fix post_residual_addition more generally (#17286)

2026-01-23 15:43:37 -08:00
parent 628ab5d57b
commit ad05782160
2 changed files with 22 additions and 0 deletions
@@ -50,6 +50,8 @@ def npu_wrapper_rmsnorm_forward(func):
        if not x.is_contiguous():
            x = x.contiguous()
        if residual is not None:
+            if post_residual_addition is not None:
+                residual = residual + post_residual_addition
            out, residual_out = add_rmsnorm_bias(
                x,
                residual,