class Linear4bit(nn.Module):
    def init(self, ):
        
def forward(self, x):
    return my_4bit_kernel(x, self.weight, self.bias)

This way, Transformers models can be easily quantized by replacing some instances of nn.Linear with a target class.

The quantization method should be serializable. You can save the quantized weights locally or push them to the Hub.
Make sure the package that contains the quantization kernels/primitive is stable (no frequent breaking changes).