Function per_tensor_quantize

convert.py:216–244

Quantize a tensor using per-tensor static scaling factor. Args: tensor: The input tensor. dtype: The data type to quantize to.

(tensor: torch.Tensor, dtype: torch.dtype)

  return tokens

def per_tensor_quantize(tensor: torch.Tensor, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
  """Quantize a tensor using per-tensor static scaling factor.
  Args:
      tensor: The input tensor.
      dtype: The data type to quantize to.
  """
  finfo = torch.finfo(dtype)
  # Calculate the scale as dtype max divided by absmax.
  # Since .abs() creates a new tensor, we use aminmax to get
  # the min and max first and then calculate the absmax.
  if tensor.numel() == 0:
    # Deal with empty tensors (triggered by empty MoE experts)
    min_val, max_val = (
      torch.tensor(-16.0, dtype=tensor.dtype),
      torch.tensor(16.0, dtype=tensor.dtype),
    )
  else:
    min_val, max_val = tensor.aminmax()
  amax = torch.maximum(min_val.abs(), max_val.abs())
  scale = finfo.max / amax.clamp(min=1e-12)
  # scale and clamp the tensor to bring it to
  # the representative range of float8 data type
  # (as default cast is unsaturated)
  qweight = (tensor * scale).clamp(min=finfo.min, max=finfo.max)
  # Return both float8 data and the inverse scale (as float),
  # as both required as inputs to torch._scaled_mm
  qweight = qweight.to(dtype)
  scale = scale.float().reciprocal()
  return qweight, scale

def per_tensor_dequantize(qweight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
  assert scale.numel() == 1