| import torch |
| import torch.nn as nn |
| import quiptools_cuda |
| from lib.utils import dtype_from_str, get_hadK |
| from lib import codebook |
| from .quantized_linear import QuantizedLinear |
| import time |
|
|
|
|
| class FusedQuantizedLinear(QuantizedLinear): |
|
|
| def __init__(self, fuse_dim, fuse_sizes, *QL_args, **QL_kwargs): |
| super(FusedQuantizedLinear, self).__init__(*QL_args, **QL_kwargs) |
| self.fuse_dim = fuse_dim |
| self.fuse_sizes = fuse_sizes |
| self.register_buffer('fuse_scales', torch.ones(len(self.fuse_sizes))) |
| self.n = len(self.fuse_sizes) |
|
|
| def forward(self, input): |
| fused_output = super(FusedQuantizedLinear, self).forward(input) |
| split_outputs = torch.split(fused_output, self.fuse_sizes, self.fuse_dim) |
| return tuple(split_outputs[i] * self.fuse_scales[i] for i in range(self.n)) |
|
|
