Optimize MX4 padding to minimize need for tuning (pytorch#3040)

Summary:
X-link: facebookresearch/FBGEMM#137

Pull Request resolved: pytorch#3040

D61447274 introduced a very cool way of doing 2D indexing over input tensors during MX4 quantization, however, it is fairly reliant on tuning configurations to get good performance. It turns out the use case for MX4 has highly dynamic shapes, so we spend a huge amount of time tuning those shapes.

After deep meditation I realized there's a much simpler indexing scheme we can use, which is similar to the 1D accesses we used previously but adds shifts for padding.

With this approach we should get the best of both worlds; support for padding rows not divisible by group size and minimizing tuning while maintaining good performance.

After further experimentation, we can actually remove tuning entirely and just use a reasonably large `GROUP_LOAD`. This gives good performance across all shapes and removes any chance of overhead. Empirically, `GROUP_LOAD=64` seems to be the sweet spot.

Differential Revision: D61816830

Author: jwfromm

fbgemm_gpu/fbgemm_gpu/triton/common.py

@@ -8,6 +8,26 @@
 # pyre-unsafe
 from enum import IntEnum
 
+import torch
+
+
+# LUTS need to be allocated ahead of time and copied to GPU to avoid expensive copies later.
+if torch.version.cuda:
+    lut_device = "cuda"
+else:
+    lut_device = "cpu"
+
+E2M1_LUT = torch.tensor(
+    [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6],
+    dtype=torch.float32,
+    device=lut_device,
+)
+E3M0_LUT = torch.tensor(
+    [0, 0.25, 0.5, 1, 2, 4, 8, 16, -0, -0.25, -0.5, -1, -2, -4, -8, -16],
+    dtype=torch.float32,
+    device=lut_device,
+)
+
 
 class RoundingMode(IntEnum):
     """Rounding options for quantization."""
@@ -47,26 +67,9 @@ def get_mx4_lookup_table(ebits, mbits):
         The lookup table for the specified mx4 format.
     """
     if ebits == 2 and mbits == 1:
-        return [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6]
+        return E2M1_LUT
     elif ebits == 3 and mbits == 0:
-        return [
-            0,
-            0.25,
-            0.5,
-            1,
-            2,
-            4,
-            8,
-            16,
-            -0,
-            -0.25,
-            -0.5,
-            -1,
-            -2,
-            -4,
-            -8,
-            -16,
-        ]
+        return E3M0_LUT
     else:
         raise NotImplementedError(
             f"MX4 with ebits={ebits} and mbits={mbits} not supported."