Revert "Use exp2 for mx scaling" (#1813)

Revert "Use exp2 for mx scaling (#1530)"

This reverts commit 890e0ac.

Author: jainapurva

torchao/prototype/mx_formats/custom_cast.py

@@ -12,13 +12,20 @@
     _f32_to_floatx_unpacked,
     _floatx_unpacked_to_f32,
 )
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_4
+
+# TODO(future): if needed, make the below work on previous PyTorch versions,
+# just need to hunt down the previous location of `libdevice`. An assert
+# at the callsite prevents usage of this on unsupported versions.
+if TORCH_VERSION_AT_LEAST_2_4 and has_triton():
+    from torch._inductor.runtime.triton_helpers import libdevice
+
 from torchao.prototype.mx_formats.constants import (
     E8M0_EXPONENT_BIAS,
     E8M0_EXPONENT_NAN_VAL,
     F4_E2M1_EXP_BIAS,
     F32_EXP_BIAS,
 )
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_4
 
 
 def get_bits(x: torch.Tensor) -> str:
@@ -287,8 +294,8 @@ def triton_f4_to_scaled_bf16_kernel(
         s = tl.load(s_ptr + offsets_s, mask=mask_s)
 
         # create the scale in bf16
-        # S is already biased by 127, so we just have to shift it to align w/ bf16
-        s_fp = (s.to(tl.uint16) << 7).to(tl.bfloat16, bitcast=True)
+        s_offset = s.to(tl.int16) - e8m0_exponent_bias
+        s_fp = libdevice.pow(2.0, s_offset).to(tl.bfloat16)
         s_fp = tl.where(s != e8m0_exponent_nan_val, s_fp, float("nan"))
 
         # multiply output by scale

torchao/prototype/mx_formats/mx_tensor.py

@@ -175,7 +175,10 @@ def to_mx(
 
     # For now, calculate the scale in floating point.
     # TODO(future) audit if there is a need to bit shift exponents instead.
-    scale_fp = torch.exp2(scale_e8m0_unbiased).to(torch.float32)
+    scale_fp = torch.pow(
+        torch.full(max_abs.size(), 2.0, device=scale_e8m0_biased.device),
+        scale_e8m0_unbiased,
+    )
 
     # Today, 2**-127 returns 0 in compile+inductor+triton because it is in the
     # float32 denormal range. For now, manually adjust the fp scale. This is
@@ -230,10 +233,14 @@ def to_mx(
 
 
 def get_fp_scale(scale_e8m0):
+    s_offset = scale_e8m0.to(torch.int16) - E8M0_EXPONENT_BIAS
+    # TODO(later): it would be nice if there was a way to do the 2^x operation
+    # in PyTorch without creating a tensor of twos
+    two = torch.full(s_offset.size(), 2.0, device=scale_e8m0.device)
+    # pow(two, s_offset) can be out of range of floating point formats.
     # TODO(later): handle this for float16 if we decide to support float16
     # scales.
-    s_offset = scale_e8m0.to(torch.int16) - E8M0_EXPONENT_BIAS
-    s_fp = torch.exp2(s_offset)
+    s_fp = torch.pow(two, s_offset)
 
     # If a block exponent was 255, set values of that block to NaN
     s_fp = torch.where(scale_e8m0 != E8M0_EXPONENT_NAN_VAL, s_fp, float("nan"))