(weight: torch.Tensor,
quant_algo: str,
plugin: bool = True)
| 68 | |
| 69 | |
| 70 | def weight_only_quantize(weight: torch.Tensor, |
| 71 | quant_algo: str, |
| 72 | plugin: bool = True): |
| 73 | assert quant_algo in [QuantAlgo.W4A16, QuantAlgo.W8A16 |
| 74 | ], f'unsupported quant algo: {quant_algo}' |
| 75 | if quant_algo == QuantAlgo.W4A16: |
| 76 | assert plugin, 'W4A16 is only supported with plugin' |
| 77 | if weight.dim() > 2: |
| 78 | v = weight.transpose(-1, -2) |
| 79 | else: |
| 80 | v = weight.t() |
| 81 | t = torch.quint4x2 if quant_algo == QuantAlgo.W4A16 else torch.int8 |
| 82 | processed_torch_weights, torch_weight_scales = \ |
| 83 | torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( |
| 84 | v.contiguous(), t) |
| 85 | if plugin: |
| 86 | return processed_torch_weights, torch_weight_scales |
| 87 | else: |
| 88 | return v, torch_weight_scales |
| 89 | |
| 90 | |
| 91 | def get_weight(params: Dict[str, torch.Tensor], prefix: str, |
no test coverage detected