(x, scale, eps)
| 101 | |
| 102 | @flags.compile_wrap |
| 103 | def rms_norm(x, scale, eps): |
| 104 | dtype = reduce(torch.promote_types, (x.dtype, scale.dtype, torch.float32)) |
| 105 | mean_sq = torch.mean(x.to(dtype)**2, dim=-1, keepdim=True) |
| 106 | scale = scale.to(dtype) * torch.rsqrt(mean_sq + eps) |
| 107 | return x * scale.to(x.dtype) |
| 108 | |
| 109 | |
| 110 | @flags.compile_wrap |