(self, x: torch.Tensor)
| 112 | |
| 113 | # ------------------------------------------------------------------ |
| 114 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 115 | # 1. Internal RMSNorm |
| 116 | x_normed = self.norm(x) |
| 117 | |
| 118 | # 2. Quantise activations (STE, keeps gradients) |
| 119 | x_q, act_scale = _quantize_activations_int8(x_normed) |
| 120 | |
| 121 | # 3. Quantise weights (STE, keeps gradients through round) |
| 122 | w_q, w_scale = _quantize_weights_ternary(self.weight) |
| 123 | |
| 124 | # 4. Linear using the (still fp) quantised values |
| 125 | out = F.linear(x_q, w_q, self.bias) |
| 126 | |
| 127 | # 5. Rescale: broadcast act_scale [..., 1] and scalar w_scale |
| 128 | out = out * act_scale * w_scale |
| 129 | |
| 130 | return out |
| 131 | |
| 132 | # ------------------------------------------------------------------ |
| 133 | @classmethod |
nothing calls this directly
no test coverage detected