Merge branch 'master' into w

Signed-off-by: YY4994 <15141404368@163.com>

Author: YY4994

benchmark/test_binary_pointwise_perf.py

@@ -72,9 +72,7 @@ def get_tflops(self, op, *args, **kwargs):
         for name, op, dtype in [
             # Arithmetic operations
             ("add", torch.add, FLOAT_DTYPES + COMPLEX_DTYPES),
-            ("atan2", torch.atan2, FLOAT_DTYPES),
-            ("copysign", torch.copysign, FLOAT_DTYPES),
-            ("div", torch.div, FLOAT_DTYPES),
+            ("div", torch.div, FLOAT_DTYPES + COMPLEX_DTYPES),
             ("mul", torch.mul, FLOAT_DTYPES + COMPLEX_DTYPES),
             ("sub", torch.sub, FLOAT_DTYPES + COMPLEX_DTYPES),
             ("pow", torch.pow, FLOAT_DTYPES),

benchmark/test_blas_perf_parallel.py

@@ -16,9 +16,24 @@
 from benchmark.attri_util import BenchmarkMetrics, BenchmarkResult, OperationAttribute
 from benchmark.conftest import Config, emit_record_logger
 
+try:
+    from vllm.utils.deep_gemm import (
+        fp8_gemm_nt,
+        is_deep_gemm_supported,
+        transform_sf_into_required_layout,
+    )
+
+    DEEPGEMM_AVAILABLE = is_deep_gemm_supported()
+except Exception:
+    fp8_gemm_nt = None
+    transform_sf_into_required_layout = None
+    DEEPGEMM_AVAILABLE = False
+
 PARALLEL_WORKER_ENV = "FLAGGEMS_BENCH_PARALLEL_WORKER"
 PARALLEL_RESULT_FILE_ENV = "FLAGGEMS_BENCH_RESULT_FILE"
 torch_device_object = flag_gems.runtime.backend.gen_torch_device_object()
+DEEPGEMM_N_MULTIPLE = 64
+DEEPGEMM_K_MULTIPLE = 128
 
 
 def _parallel_device_is_available():
@@ -197,6 +212,7 @@ def estimate_shape_cost(shape):
                 "bmm",
                 "baddbmm",
                 "w8a8_block_fp8_matmul",
+                "w8a8_block_fp8_matmul_deepgemm",
             }:
                 normalized_shape = shape
                 if len(shape) == 3:
@@ -209,7 +225,11 @@ def estimate_shape_cost(shape):
                 if normalized_shape is None:
                     return 1
 
-                if self.op_name in {"mm", "bmm", "w8a8_block_fp8_matmul"}:
+                if self.op_name in {
+                    "mm",
+                    "bmm",
+                    "w8a8_block_fp8_matmul",
+                }:
                     return m * n * k * 2
                 return m * n * (2 * k + 1)
 
@@ -482,7 +502,7 @@ class ParallelAddrBenchmark(ParallelBenchmarkMixin, blas_perf.AddrBenchmark):
 class ParallelW8A8BlockFP8MatmulBenchmark(
     ParallelBenchmarkMixin, blas_perf.W8A8BlockFP8MatmulBenchmark
 ):
-    SHAPE_CONFIG_KEYS = ("mm", "BlasBenchmark")
+    SHAPE_CONFIG_KEYS = ("w8a8_block_fp8_matmul", "BlasBenchmark")
 
     def set_more_shapes(self):
         if os.environ.get(PARALLEL_WORKER_ENV):
@@ -537,6 +557,118 @@ def set_shapes(self, shape_file_path=None):
         self.shape_desc = "M, N, K"
 
 
+def _deepgemm_block_scaled_mm(A, B, As_dg, Bs_dg, output):
+    fp8_gemm_nt((A, As_dg), (B, Bs_dg), output)
+    return output
+
+
+class ParallelW8A8BlockFP8DeepGemmBenchmark(ParallelW8A8BlockFP8MatmulBenchmark):
+    def __init__(self, *args, output_dtype=torch.bfloat16, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.output_dtype = output_dtype
+
+    def set_shapes(self, shape_file_path=None):
+        super().set_shapes(shape_file_path)
+        self.shapes = [
+            (m, n, k)
+            for m, n, k in self.shapes
+            if n % DEEPGEMM_N_MULTIPLE == 0 and k % DEEPGEMM_K_MULTIPLE == 0
+        ]
+
+    def get_input_iter(self, cur_dtype):
+        fp8_dtype = blas_perf.get_w8a8_block_fp8_dtype()
+        if fp8_dtype is None:
+            raise RuntimeError(
+                "DeepGEMM benchmark requires CUDA device with FP8 support"
+            )
+
+        block_n, block_k = self.block_size
+        recipe = (1, 128, 128)
+
+        for m, n, k in self.shapes:
+            num_k_groups = (k + block_k - 1) // block_k
+            num_n_groups = (n + block_n - 1) // block_n
+
+            A = blas_perf.rand_fp8_tensor((m, k), self.device, fp8_dtype).contiguous()
+            B = blas_perf.rand_fp8_tensor((n, k), self.device, fp8_dtype).contiguous()
+            As = (
+                0.01
+                * torch.rand((m, num_k_groups), dtype=torch.float32, device=self.device)
+                + 0.005
+            ).contiguous()
+            Bs = (
+                0.01
+                * torch.rand(
+                    (num_n_groups, num_k_groups),
+                    dtype=torch.float32,
+                    device=self.device,
+                )
+                + 0.005
+            ).contiguous()
+
+            As_dg = transform_sf_into_required_layout(
+                sf=As.unsqueeze(0),
+                mn=m,
+                k=k,
+                recipe=recipe,
+                num_groups=1,
+                is_sfa=True,
+            ).squeeze(0)
+            Bs_dg = transform_sf_into_required_layout(
+                sf=Bs.unsqueeze(0),
+                mn=n,
+                k=k,
+                recipe=recipe,
+                num_groups=1,
+                is_sfa=False,
+            ).squeeze(0)
+            output = torch.empty((m, n), dtype=self.output_dtype, device=self.device)
+
+            yield (
+                A,
+                B,
+                As_dg,
+                Bs_dg,
+                output,
+            ), (
+                A,
+                B,
+                As,
+                Bs,
+                self.block_size[:],
+                self.output_dtype,
+            )
+
+    def _build_metric_from_input(self, input_item):
+        dg_input, gems_input = input_item
+        metric = BenchmarkMetrics()
+
+        dg_args, dg_kwargs = self.unpack_to_args_kwargs(dg_input)
+        gems_args, gems_kwargs = self.unpack_to_args_kwargs(gems_input)
+        metric.shape_detail = self.record_shapes(*gems_args, **gems_kwargs)
+
+        if "latency_base" in self.to_bench_metrics:
+            metric.latency_base = self.get_latency(self.torch_op, *dg_args, **dg_kwargs)
+        if "latency" in self.to_bench_metrics:
+            metric.latency = self.get_latency(self.gems_op, *gems_args, **gems_kwargs)
+        if "speedup" in self.to_bench_metrics:
+            metric.speedup = metric.latency_base / metric.latency
+        if "tflops" in self.to_bench_metrics:
+            metric.tflops = (
+                self.get_tflops(self.torch_op, *dg_args, **dg_kwargs)
+                / metric.latency
+                / 1e12
+                * 1e3
+            )
+        return metric
+
+    def get_tflops(self, op, *args, **kwargs):
+        A, B = args[0], args[1]
+        m, k = A.shape
+        n = B.shape[0]
+        return 2 * m * n * k
+
+
 @pytest.mark.parametrize(
     "op_name, torch_op, input_fn, bench_cls",
     [
@@ -604,6 +736,25 @@ def test_perf_w8a8_block_fp8_matmul():
     bench.run()
 
 
+@pytest.mark.w8a8_block_fp8_matmul_deepgemm
+def test_perf_w8a8_block_fp8_matmul_deepgemm():
+    if not DEEPGEMM_AVAILABLE:
+        pytest.skip("DeepGEMM is not available on this platform")
+    if blas_perf.get_w8a8_block_fp8_dtype() is None:
+        pytest.skip(
+            "w8a8_block_fp8_matmul benchmark requires CUDA device with FP8 support"
+        )
+
+    bench = ParallelW8A8BlockFP8DeepGemmBenchmark(
+        op_name="w8a8_block_fp8_matmul_deepgemm",
+        torch_op=_deepgemm_block_scaled_mm,
+        dtypes=["fp8"],
+        output_dtype=torch.bfloat16,
+    )
+    bench.set_gems(flag_gems.w8a8_block_fp8_matmul)
+    bench.run()
+
+
 @pytest.mark.parametrize(
     "op_name, torch_op, input_fn",
     [

benchmark/test_select_and_slice_perf.py

@@ -78,7 +78,6 @@ def index_select_gbps(bench_fn_args, latency):
     return io_amount * 1e-9 / (latency * 1e-3)
 
 
-@pytest.mark.index_select
 @pytest.mark.parametrize(
     "op_name, torch_op, input_fn, gbps_fn, dtypes",
     [
@@ -128,7 +127,6 @@ def test_perf_masked_scatter_inplace():
     bench.run()
 
 
-@pytest.mark.masked_select
 @pytest.mark.parametrize(
     "op_name, torch_op, input_fn, gbps_fn, dtypes",
     [
@@ -142,7 +140,7 @@ def test_perf_masked_scatter_inplace():
         ),
     ],
 )
-def test_perf_masked_select(op_name, torch_op, input_fn, gbps_fn, dtypes):
+def test_masked_select(op_name, torch_op, input_fn, gbps_fn, dtypes):
     bench = TensorSelectBenchmark(
         input_fn=input_fn,
         op_name=op_name,
@@ -193,8 +191,8 @@ def inner(shape, dtype, device):
     return inner
 
 
-@pytest.mark.scatter
-def test_perf_scatter():
+@pytest.mark.scatter_src
+def test_scatter_src():
     bench = TensorSelectBenchmark(
         op_name="scatter.src",
         torch_op=torch.scatter,
@@ -205,8 +203,8 @@ def test_perf_scatter():
     bench.run()
 
 
-@pytest.mark.scatter
-def test_perf_scatter_add():
+@pytest.mark.scatter_reduce
+def test_scatter_reduce_add():
     bench = TensorSelectBenchmark(
         op_name="scatter.reduce",
         torch_op=torch.scatter,
@@ -217,29 +215,8 @@ def test_perf_scatter_add():
     bench.run()
 
 
-@pytest.mark.scatter_add_
-def test_perf_scatter_add_():
-    def scatter_input_fn(shape, dtype, device):
-        input_gen = gather_input_fn(shape, dtype, device)
-        inp, dim, index = next(input_gen)
-        src_shape = list(size + 16 for size in index.shape)
-        src = torch.randn(src_shape, dtype=dtype, device=device)
-
-        yield inp, dim, index, src
-
-    bench = TensorSelectBenchmark(
-        op_name="scatter_add_",
-        torch_op=torch.Tensor.scatter_add_,
-        input_fn=scatter_input_fn,
-        get_gbps=gather_scatter_gbps,
-        dtypes=FLOAT_DTYPES,
-    )
-    bench.run()
-
-
-@pytest.mark.scatter_multiply
-@pytest.mark.scatter
-def test_perf_scatter_multiply():
+@pytest.mark.scatter_reduce
+def test_scatter_reduce_multiply():
     bench = TensorSelectBenchmark(
         op_name="scatter.reduce",
         torch_op=torch.scatter,
@@ -250,8 +227,8 @@ def test_perf_scatter_multiply():
     bench.run()
 
 
-@pytest.mark.scatter_
-def test_perf_scatter_inplace():
+@pytest.mark.scatter_src_
+def test_scatter_src_inplace():
     bench = TensorSelectBenchmark(
         op_name="scatter_.src",
         torch_op=torch.Tensor.scatter_,
@@ -263,8 +240,8 @@ def test_perf_scatter_inplace():
     bench.run()
 
 
-@pytest.mark.scatter_
-def test_perf_scatter_add_inplace():
+@pytest.mark.scatter_reduce_
+def test_scatter_reduce_add_inplace():
     bench = TensorSelectBenchmark(
         op_name="scatter_.reduce",
         torch_op=torch.Tensor.scatter_,
@@ -276,8 +253,8 @@ def test_perf_scatter_add_inplace():
     bench.run()
 
 
-@pytest.mark.scatter_
-def test_perf_scatter_multiply_inplace():
+@pytest.mark.scatter_reduce_
+def test_scatter_reduce_multiply_inplace():
     bench = TensorSelectBenchmark(
         op_name="scatter_.reduce",
         torch_op=torch.Tensor.scatter_,
@@ -289,6 +266,26 @@ def test_perf_scatter_multiply_inplace():
     bench.run()
 
 
+@pytest.mark.scatter_add_
+def test_scatter_add_inplace():
+    def scatter_input_fn(shape, dtype, device):
+        input_gen = gather_input_fn(shape, dtype, device)
+        inp, dim, index = next(input_gen)
+        src_shape = list(size + 16 for size in index.shape)
+        src = torch.randn(src_shape, dtype=dtype, device=device)
+
+        yield inp, dim, index, src
+
+    bench = TensorSelectBenchmark(
+        op_name="scatter_add_",
+        torch_op=torch.Tensor.scatter_add_,
+        input_fn=scatter_input_fn,
+        get_gbps=gather_scatter_gbps,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()
+
+
 def gather_input_fn(shape, dtype, device):
     inp = torch.randn(shape, dtype=dtype, device=device)
 

benchmark/test_special_perf.py

@@ -1511,8 +1511,7 @@ def _functional_sym_constrain_range_for_size_input_fn(shape, cur_dtype, device):
 
 
 @pytest.mark.functional_sym_constrain_range_for_size
-@pytest.mark.performance
-def test_perf_functional_sym_constrain_range_for_size():
+def test_functional_sym_constrain_range_for_size():
     bench = GenericBenchmark(
         op_name="functional_sym_constrain_range_for_size",
         torch_op=torch.ops.aten._functional_sym_constrain_range_for_size,

benchmark/test_unary_pointwise_perf.py

@@ -418,8 +418,7 @@ def set_gems(self, gems_op):
     flag_gems.vendor_name in UNSUPPORTED_VENDORS, reason="Vendor not supported"
 )
 @pytest.mark.apply_repetition_penalties
-@pytest.mark.performance
-def test_perf_repetition_penalty():
+def test_apply_repetition_penalties():
     vllm_ops = pytest.importorskip("vllm._custom_ops")
 
     bench = RepetitionPenaltyBenchmark(

benchmark/test_vllm_perf.py

@@ -236,7 +236,6 @@ def get_input_iter(self, dtype):
     reason="requires vLLM and NVIDIA Hopper architecture",
 )
 @pytest.mark.cutlass_scaled_mm
-@pytest.mark.performance
 def test_cutlass_scaled_mm_benchmark():
     bench = CutlassScaledMMBenchmark()
     bench.run()
@@ -1403,11 +1402,7 @@ def make_input_flashmla(param: Flashmla_Sparse_Test_Param):
 
 
 @pytest.mark.flashmla_sparse
-@pytest.mark.performance
 @pytest.mark.skipif(not HAS_VLLM_FLASHMLA_SPARSE, reason="vllm not installed")
 def test_perf_flashmla_sparse_gems_vs_vllm():
-    """
-    Benchmark FlagGems flash_mla_sparse_fwd vs vLLM flash_mla_sparse_fwd.
-    """
     bench = FlashmlaSparseBenchmark()
     bench.run()