Function blockwise_dequantize

convert.py:250–260

(qweight: torch.Tensor, scale: torch.Tensor, block_size: torch.Tensor)

  return qweight.to(torch.float32) * scale

def blockwise_dequantize(qweight: torch.Tensor, scale: torch.Tensor, block_size: torch.Tensor) -> torch.Tensor:
  assert qweight.ndim == scale.ndim and scale.ndim == block_size.numel() and scale.ndim == 2
  assert torch.all((torch.tensor(list(qweight.shape)) / block_size).ceil() == torch.tensor(list(scale.shape)))
  out = torch.empty_like(qweight, dtype=torch.float32)
  for i in range(scale.shape[0]):
    for j in range(scale.shape[1]):
      block_size_i = block_size[0]
      block_size_j = block_size[1]
      qw_block = qweight[i*block_size_i:(i+1)*block_size_i, j*block_size_j:(j+1)*block_size_j]
      out[i*block_size_i:(i+1)*block_size_i, j*block_size_j:(j+1)*block_size_j] = per_tensor_dequantize(qw_block, scale[i, j])
  return out

def blockwise_quantize(weight: torch.Tensor, block_size: torch.Tensor, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
  assert weight.ndim == block_size.numel() and weight.ndim == 2