Add FP8 rowwise kernel (cutlass + triton)

Summary: As the title. The benchmark could help us quickly prototype the kernel level improvement

Reviewed By: xuzhao9

Differential Revision: D58344346

fbshipit-source-id: 597b4cb0c88457bad17fb1b551787b03f0cd02b3

torchbenchmark/operators/fp8_gemm_rowwise/__init__.py

@@ -0,0 +1 @@
+from .operator import Operator

torchbenchmark/operators/fp8_gemm_rowwise/operator.py

@@ -0,0 +1,187 @@
+import argparse
+import os
+from typing import Any, Callable, Generator, List, Optional, Tuple
+
+import fbgemm_gpu.experimental.gen_ai # noqa: F401
+
+import torch
+import triton
+
+from torchbenchmark.util.triton_op import (
+ BenchmarkOperator,
+ BenchmarkOperatorMetrics,
+ register_benchmark,
+ register_metric,
+ register_x_val,
+)
+
+
+def parse_args(args: List[str]) -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description="TorchBench Addmm operator Benchmark")
+ parser.add_argument("--m", type=int)
+ parser.add_argument("--n", type=int)
+ parser.add_argument("--k", type=int)
+ args = parser.parse_args(args)
+ return args
+
+
+try:
+ from fbgemm_gpu.experimental.gemm.triton_gemm.fp8_gemm import (
+ matmul_fp8_row as triton_fp8_row,
+ )
+
+ HAS_TRITON = True
+except ImportError:
+ HAS_TRITON = False
+
+
+try:
+ cutlass_fp8_row = torch.ops.fbgemm.f8f8bf16_rowwise
+ HAS_CUTLASS = True
+except ImportError:
+ HAS_CUTLASS = False
+
+BUILDIN_SHAPES = [
+ (1, 2304, 2048),
+ (1, 8192, 16384),
+ (4, 4096, 2304),
+ (4, 13312, 2048),
+ (8, 2304, 2304),
+ (8, 8192, 6656),
+ (16, 4096, 6656),
+ (16, 13312, 13312),
+ (32, 2304, 16384),
+ (32, 8192, 13312),
+ (64, 4096, 2048),
+ (64, 13312, 2048),
+ (128, 2304, 6656),
+ (128, 8192, 2304),
+ (2048, 8192, 2048),
+ (2048, 13312, 6656),
+ (4096, 2304, 13312),
+ (4096, 13312, 2304),
+ (16384, 4096, 16384),
+ (16384, 8192, 13312),
+]
+
+E4M3_MAX_POS: float = torch.finfo(torch.float8_e4m3fn).max
+EPS: float = 1e-12
+FP16_MAX_POS: float = torch.finfo(torch.float16).max
+
+
+def fp8_row_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ # Quantize an input tensor and return the fp8 tensor and its inverse scale.
+ x_row_max = torch.max(torch.abs(x), dim=1).values
+ scale = E4M3_MAX_POS / torch.clamp(x_row_max, EPS)
+ if x.dtype is torch.float16:
+ scale = torch.clamp(scale, max=FP16_MAX_POS)
+ xq = torch.clamp(x * scale[:, None], min=-1 * E4M3_MAX_POS, max=E4M3_MAX_POS).to(
+ torch.float8_e4m3fn
+ )
+ return xq, scale.reciprocal().to(torch.float32)
+
+
+class Operator(BenchmarkOperator):
+ DEFAULT_METRICS = ["tflops", "speedup", "accuracy"]
+ DEFAULT_PRECISION = "fp32"
+
+ def __init__(self, mode: str, device: str, extra_args: Optional[List[str]] = None):
+ super().__init__(mode=mode, device=device, extra_args=extra_args)
+ addmm_args = parse_args(self.extra_args)
+ if addmm_args.m and addmm_args.n and addmm_args.k:
+ self.shapes = [(addmm_args.m, addmm_args.n, addmm_args.k)]
+ else:
+ self.shapes = BUILDIN_SHAPES
+
+ @register_benchmark(enabled=HAS_TRITON, baseline=True)
+ def _triton(self, xq, wq, x_scale, w_scale) -> Callable:
+ return lambda: triton_fp8_row(xq, wq, x_scale, w_scale)
+
+ @register_benchmark(enabled=HAS_CUTLASS)
+ def _cutlass(self, xq, wq, x_scale, w_scale) -> Callable:
+ return lambda: cutlass_fp8_row(xq, wq, x_scale, w_scale)
+
+ # TODO: add cublas rowwise FP8 kernel
+ # @register_benchmark(baseline=True)
+ # def _torch(self, xq, wq, x_scale, w_scale) -> Callable:
+ # def _cublass(xq, wq, x_scale, w_scale):
+ # output, _ = torch._scaled_mm(
+ # xq, wq.T, use_fast_accum=True, out_dtype=torch.bfloat16
+ # )
+ # return output * x_scale[:, None] * w_scale[None, :]
+
+ # return lambda: _cublass(xq, wq, x_scale, w_scale)
+
+ @register_metric()
+ def tflops(
+ self, fn_name: str, example_inputs: Any, metrics: BenchmarkOperatorMetrics
+ ) -> List[float]:
+ xq, wq, _, _ = example_inputs
+ m, k = xq.size()
+ n, k = wq.size()
+ flops = m * k * 2 * n
+ return [flops / x / 1e12 * 1e3 for x in metrics.latency]
+
+ @register_x_val(label="(M, N, K)")
+ def get_x_val(self, example_inputs) -> Tuple[int, int, int]:
+ xq, wq, _, _ = example_inputs
+ m, k = xq.size()
+ n, k = wq.size()
+ return (m, n, k)
+
+ def get_input_iter(self) -> Generator:
+ for shape in self.shapes:
+ m, n, k = shape
+ x = torch.randn(
+ (m, k), device=self.device, dtype=torch.bfloat16
+ ).requires_grad_(False)
+ w = torch.randn(
+ (n, k), device=self.device, dtype=torch.bfloat16
+ ).requires_grad_(False)
+ xq, x_scale = fp8_row_quantize(x)
+ wq, w_scale = fp8_row_quantize(w)
+ yield xq, wq, x_scale, w_scale
+
+ def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool:
+ output = fn()
+ baseline_output = baseline_fn()
+ accuracy = True
+ try:
+ torch.testing.assert_close(output, baseline_output, atol=1e-2, rtol=0.5)
+ except Exception:
+ accuracy = False
+ finally:
+ return accuracy
+
+ def plot(self):
+ @triton.testing.perf_report(
+ triton.testing.Benchmark(
+ x_names=["density"], # argument names to use as an x-axis for the plot
+ x_vals=self.output.x_vals, # different possible values for `x_name`
+ line_arg="provider", # argument name whose value corresponds to a different line in the plot
+ line_vals=[
+ "_torch",
+ "_triton",
+ "_cutlass",
+ ], # possible values for `line_arg``
+ line_names=[
+ "Torch",
+ "Triton",
+ "Cutlass",
+ ], # label name for the lines
+ styles=[("blue", "-"), ("green", "-"), ("yellow", "-")], # line styles
+ ylabel="tflops", # label name for the y-axis
+ plot_name="gemm-performance", # name for the plot. Used also as a file name for saving the plot.
+ args={}, # values for function arguments not in `x_names` and `y_name`
+ )
+ )
+ def _plot(density, provider):
+ tflops = self.output.get_y_vals(density, provider, "tflops")
+ return tflops
+
+ save_path = "/tmp/test_fp8_gemm_rowwise"
+
+ if not os.path.exists(save_path):
+ os.mkdir(save_path)
+
+ _plot.run(show_plots=True, print_data=True, save_path=save_path)