| 21 | |
| 22 | |
| 23 | class RMSNorm(torch.nn.Module): |
| 24 | def __init__(self, dim: int, eps: float = 1e-6): |
| 25 | super().__init__() |
| 26 | self.eps = eps |
| 27 | self.weight = nn.Parameter(torch.ones(dim)) |
| 28 | |
| 29 | def _norm(self, x): |
| 30 | return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| 31 | |
| 32 | def forward(self, x): |
| 33 | output = torch.empty_like(x) |
| 34 | awq_inference_engine.layernorm_forward_cuda(x, self.weight, output, self.eps) |
| 35 | return output |
| 36 | |
| 37 | |
| 38 | def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): |