(num_bits, symmetric, x, group_dim, group_size, quant_type)
| 44 | |
| 45 | |
| 46 | def _quantize(num_bits, symmetric, x, group_dim, group_size, quant_type): |
| 47 | x_float = x.float() |
| 48 | x_reshaped = x_float.reshape( |
| 49 | x.shape[:group_dim] |
| 50 | + (x.shape[group_dim] // group_size, group_size) |
| 51 | + x.shape[group_dim + 1 :] |
| 52 | ) |
| 53 | if symmetric: |
| 54 | signed_max = float(2 ** (num_bits - 1)) - 1 |
| 55 | offset = signed_max if quant_type is flow.uint8 else 0.0 |
| 56 | scale_float = ( |
| 57 | x_reshaped.abs().max(dim=group_dim + 1, keepdim=True).values / signed_max |
| 58 | ) |
| 59 | quantized = ( |
| 60 | flow.round(x_reshaped / scale_float + offset) |
| 61 | .reshape(x.shape) |
| 62 | .to(quant_type) |
| 63 | ) |
| 64 | if num_bits == 4: |
| 65 | quantized = _pack_int8_to_int4(quantized) |
| 66 | return (quantized, scale_float.squeeze(group_dim + 1).to(x.dtype), None) |
| 67 | else: |
| 68 | unsigned_max = float(2 ** num_bits) - 1 |
| 69 | mn = x_reshaped.min(dim=group_dim + 1, keepdim=True).values |
| 70 | mx = x_reshaped.max(dim=group_dim + 1, keepdim=True).values |
| 71 | scale_float = (mx - mn) / unsigned_max |
| 72 | quantized = ( |
| 73 | flow.round((x_reshaped - mn) / scale_float).reshape(x.shape).to(flow.uint8) |
| 74 | ) |
| 75 | if num_bits == 4: |
| 76 | quantized = _pack_int8_to_int4(quantized) |
| 77 | return ( |
| 78 | quantized, |
| 79 | scale_float.squeeze(group_dim + 1).to(x.dtype), |
| 80 | mn.squeeze(group_dim + 1).to(x.dtype), |
| 81 | ) |
| 82 | |
| 83 | |
| 84 | def _dequantize_ref(num_bits, symmetric, quantized, scale, zero, group_dim, group_size): |
no test coverage detected