(self, original: nn.Module, kernel_fn: Callable)
| 479 | """Wraps RMSNorm-like modules to use an optimized kernel_fn.""" |
| 480 | |
| 481 | def __init__(self, original: nn.Module, kernel_fn: Callable): |
| 482 | super().__init__() |
| 483 | self.original = original |
| 484 | self.kernel_fn = kernel_fn |
| 485 | # RMSNorm typically has a 'weight' attribute |
| 486 | self.weight = getattr(original, "weight", None) |
| 487 | self.eps = getattr(original, "eps", 1e-6) |
| 488 | |
| 489 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 490 | orig_shape = x.shape |