Replace matching modules in the model. Returns number of modules replaced.
(self, repl: KernelReplacement)
| 555 | self._applied.clear() |
| 556 | |
| 557 | def _apply_replacement(self, repl: KernelReplacement) -> int: |
| 558 | """ |
| 559 | Replace matching modules in the model. Returns number of modules replaced. |
| 560 | """ |
| 561 | count = 0 |
| 562 | |
| 563 | if repl.kernel_type == "matmul": |
| 564 | count = self._replace_linear_modules(repl) |
| 565 | elif repl.kernel_type == "layernorm": |
| 566 | count = self._replace_layernorm_modules(repl) |
| 567 | elif repl.kernel_type == "rmsnorm": |
| 568 | count = self._replace_rmsnorm_modules(repl) |
| 569 | else: |
| 570 | print(f" NOTE: No replacement strategy for kernel type '{repl.kernel_type}'. " |
| 571 | f"Skipping. (Supported: matmul, layernorm, rmsnorm)") |
| 572 | |
| 573 | return count |
| 574 | |
| 575 | def _replace_linear_modules(self, repl: KernelReplacement) -> int: |
| 576 | """Replace all nn.Linear modules with optimized matmul wrapper.""" |
no test coverage detected