Fix Float8Tensor quantize op kernrel preference dispatch

Summary:
Previously we didn't handle kernel_preference == "fbgemm" properly for the quantize op,
this PR makes sure we dispatch to fbgemm kernels when kernel_preference is fbgemm

This doesn't have much impact on BC, the serialized checkpoints will use AUTO which is going to be dispatched
to triton op for quantize, only thing is fixing the kernel choice for fbgemm kernel preference, which
is supposed to be a developer facing API (we expect most users to just use AUTO without worrying about details)

Test Plan:
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_kernel_preference_numerical_equivalence

Reviewers:

Subscribers:

Tasks:

Tags:

stack-info: PR: #2883, branch: jerryzh168/stack/59

Author: jerryzh168

torchao/quantization/quantize_/common/kernel_preference.py

@@ -30,5 +30,9 @@ class KernelPreference(str, Enum):
     """
     FBGEMM = "fbgemm"
 
+    """Use triton quantize and quantized mm kernels (if available), requires fbgemm_gpu_genai library, if no triton kernel for the quantize op or mm kernel is available, we'll fallback to torch ops
+    """
+    TRITON = "triton"
+
 
 torch.serialization.add_safe_globals([KernelPreference])

torchao/quantization/quantize_/workflows/float8/float8_tensor.py

@@ -22,7 +22,7 @@
     preprocess_data,
     preprocess_scale,
 )
-from torchao.quantization.granularity import PerRow
+from torchao.quantization.granularity import PerRow, PerTensor
 from torchao.quantization.observer import get_block_size
 from torchao.quantization.quant_primitives import (
     _choose_scale_float8,
@@ -177,18 +177,33 @@ def to_float8(
         block_size = get_block_size(hp_tensor.shape, granularity)
         block_size = list(block_size)
 
-        # for per row quantization and kernel_preference default setting, we'll use triton kernel for best performance
+        kernel_choice = None
         if (
             kernel_preference == KernelPreference.AUTO
             and _is_fbgemm_genai_gpu_available()
-            and (
-                tuple(block_size)
-                == (1,) * (hp_tensor.ndim - 1) + (hp_tensor.shape[-1],)
+            and is_sm_at_least_90()
+            and isinstance(granularity, PerRow)
+            and float8_dtype == torch.float8_e4m3fn
+            and hp_value_lb is None
+        ) or kernel_preference == KernelPreference.TRITON:
+            # for per row quantization and kernel_preference auto setting
+            # we'll use triton quantize kernel for best performance
+            kernel_choice = "triton"
+        elif kernel_preference == KernelPreference.FBGEMM:
+            # we'll use fbgemm quantize kernel if it's explicitly chosen by user
+            assert _is_fbgemm_genai_gpu_available() and is_sm_at_least_90(), (
+                "Specified fbgemm but fbgemm_gpu_genai is not installed or hardware is not >= SM 9.0 (> H100)"
             )
-        ):
-            assert float8_dtype == torch.float8_e4m3fn, (
-                f"Only torch.float8_e4m3fn is supported, got: {float8_dtype}"
+            assert hp_value_lb is None, (
+                "hp_value_lb should not be specified if FBGEMM is explicitly chosen"
             )
+            kernel_choice = "fbgemm"
+        else:
+            # fallback quantize kernel for everything else will be torch
+            kernel_choice = "torch"
+
+        if kernel_choice == "triton":
+            assert hp_value_lb is None, f"{hp_value_lb=} is not supported"
             if hp_value_ub is not None:
                 maybe_hp_value_ub_tensor = torch.tensor(
                     hp_value_ub, dtype=torch.float, device=hp_tensor.device
@@ -202,7 +217,39 @@ def to_float8(
             for i in range(hp_tensor.ndim):
                 scale_shape.append(hp_tensor.shape[i] // block_size[i])
             scale = scale.reshape(*scale_shape)
+        elif kernel_choice == "fbgemm":
+            assert hp_value_lb is None, f"{hp_value_lb=} is not supported"
+            if hp_value_ub is not None:
+                maybe_hp_value_ub_tensor = torch.tensor(
+                    hp_value_ub, dtype=torch.float, device=hp_tensor.device
+                )
+            else:
+                maybe_hp_value_ub_tensor = None
+            # not used
+            num_tokens = torch.empty([hp_tensor.size(0)], device=hp_tensor.device)
+            if isinstance(granularity, PerRow):
+                data, scale = torch.ops.fbgemm.quantize_fp8_per_row(
+                    hp_tensor, num_tokens, scale_ub=maybe_hp_value_ub_tensor
+                )
+            else:
+                assert isinstance(granularity, PerTensor), (
+                    f"Expected per tensor, got {granularity}"
+                )
+                # TODO: use fbgemm kernel when it works
+                # current error: torch.AcceleratorError: CUDA error: an illegal memory access was encountered
+                # data, scale = torch.ops.fbgemm.quantize_fp8_per_tensor(
+                #     hp_tensor, num_tokens, scale_ub=maybe_hp_value_ub_tensor
+                # )
+                scale = _choose_scale_float8(
+                    hp_tensor,
+                    float8_dtype=float8_dtype,
+                    block_size=block_size,
+                    hp_value_lb=hp_value_lb,
+                    hp_value_ub=hp_value_ub,
+                )
+                data = _quantize_affine_float8(hp_tensor, scale, float8_dtype)
         else:
+            assert kernel_choice == "torch", f"Expected torch, got {kernel_choice}"
             scale = _choose_scale_float8(
                 hp_tensor,
                 float8_dtype=float8_dtype,
@@ -256,6 +303,9 @@ def _(func, types, args, kwargs):
                 kernel_choice = "fbgemm"
         elif weight_tensor.kernel_preference == KernelPreference.FBGEMM:
             kernel_choice = "fbgemm"
+        elif weight_tensor.kernel_preference == KernelPreference.TRITON:
+            # no triton gemm op is available, so we'll fallback to torch
+            kernel_choice = "torch"
         else:
             assert weight_tensor.kernel_preference == KernelPreference.TORCH, (
                 f"{weight_tensor.kernel_preference=} not handled"