Repository: intel/intel-xpu-backend-for-triton
Pull request #3604: Add ScaledDot to micro benchmark
Author: login: leonling-ll
Created: 2025-03-04T02:43:47Z | Merged: 2025-03-04T11:06:51Z

Per https://github.com/intel/intel-xpu-backend-for-triton/pull/3538#pullrequestreview-2654848212, split micro benchmark part from #3538 .

The kernel is extracted from UT. 
The reason of add it to micro benchmark is this Op implemtation changed much recently. We have performance concern because there are many layout conversions used in lowering.


Diff:
diff --git a/benchmarks/micro_benchmarks/core_ops/__init__.py b/benchmarks/micro_benchmarks/core_ops/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/benchmarks/micro_benchmarks/core_ops/dot_scaled.py b/benchmarks/micro_benchmarks/core_ops/dot_scaled.py
new file mode 100644
index 0000000000..65d87694ae
--- /dev/null
+++ b/benchmarks/micro_benchmarks/core_ops/dot_scaled.py
@@ -0,0 +1,155 @@
+import torch
+import triton
+import triton.language as tl
+
+import triton_kernels_benchmark as benchmark_suit
+
+
+@triton.jit
+def dot_scale_kernel(a_base, stride_a0, stride_a1, a_scale, b_base, stride_b0, stride_b1, b_scale, out,
+                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, type_a: tl.constexpr,
+                     type_b: tl.constexpr):
+    DIV_FACTOR_A: tl.constexpr = 2 if type_a == 'e2m1' else 1
+    DIV_FACTOR_B: tl.constexpr = 2 if type_b == 'e2m1' else 1
+    PACKED_BLOCK_K_A: tl.constexpr = BLOCK_K // DIV_FACTOR_A
+    PACKED_BLOCK_K_B: tl.constexpr = BLOCK_K // DIV_FACTOR_B
+    a_ptr = a_base + tl.arange(0, BLOCK_M)[:, None] * stride_a0 + tl.arange(0, PACKED_BLOCK_K_A)[None, :] * stride_a1
+    b_ptr = b_base + tl.arange(0, PACKED_BLOCK_K_B)[:, None] * stride_b0 + tl.arange(0, BLOCK_N)[None, :] * stride_b1
+
+    a = tl.load(a_ptr)
+    b = tl.load(b_ptr)
+    SCALE_BLOCK_K: tl.constexpr = BLOCK_K // 32
+    if a_scale is not None:
+        scale_a_ptr = a_scale + tl.arange(0, BLOCK_M)[:, None] * SCALE_BLOCK_K + tl.arange(0, SCALE_BLOCK_K)[None, :]
+        a_scale = tl.load(scale_a_ptr)
+    if b_scale is not None:
+        scale_b_ptr = b_scale + tl.arange(0, BLOCK_N)[:, None] * SCALE_BLOCK_K + tl.arange(0, SCALE_BLOCK_K)[None, :]
+        b_scale = tl.load(scale_b_ptr)
+    c = tl.dot_scaled(a, a_scale, type_a, b, b_scale, type_b)
+    out_ptr = out + \
+        tl.arange(0, BLOCK_M)[:, None] * BLOCK_N + \
+        tl.arange(0, BLOCK_N)[None, :]
+    tl.store(out_ptr, c.to(tl.bfloat16))
+
+
+def dot_scaled(M, N, K, x, y, z, scale_x, scale_y, type_a, type_b, num_warps):
+    kernel_kwargs = {'num_warps': num_warps}
+    dot_scale_kernel[(1, )](x, *x.stride(), scale_x, y, *y.stride(), scale_y, z, M, N, K, type_a, type_b,
+                            **kernel_kwargs)
+
+
+# Benchmark Performance
+@benchmark_suit.perf_report(
+    benchmark_suit.Benchmark(
+        # argument names to use as an x-axis for the plot
+        x_names=['M', 'K', 'N', 'col_a', 'col_b', 'rhs_scale', 'mxfp_type', 'normal_type'],
+        x_vals=[(M, N, K, col_a, col_b, rhs_scale, mxfp_type, normal_type)
+                for M, N, K in [(128, 128, 128)]
+                for col_a, col_b in [(True, True), (False, False)]
+                for rhs_scale in [True, False]
+                for mxfp_type in ['e2m1', 'e4m3']
+                for normal_type in ['bf16']],
+        line_arg='provider',
+        # argument name whose value corresponds to a different line in the plot
+        # possible values for `line_arg``
+        line_vals=['triton'],
+        # label name for the lines
+        line_names=['Triton'],
+        # line styles
+        styles=[('green', '-'), ('green', '--'), ('blue', '-'), ('blue', '--')],
+        ylabel=['GB/s', 'TFlops'],  # label name for the y-axis
+        plot_name='scaled-dot',
+        # name for the plot. Used also as a file name for saving the plot.
+        args={},
+    ))
+def benchmark(M, N, K, col_a, col_b, rhs_scale, mxfp_type, normal_type, provider):
+
+    device = 'xpu'
+    num_warps = 4
+    quantiles = [0.5, 0.0, 1.0]
+
+    comp_dtype = torch.float16 if normal_type == 'fp16' else torch.bfloat16
+    # The max exponent we use to initialize data in the x/y and associated scale tensor to avoid
+    # overflow when scaling.
+    comp_dtype_max_exp = 6 if normal_type == 'fp16' else 15
+
+    torch.manual_seed(0)
+
+    def make_arg(shape, ty, col_major=False):
+        if col_major:
+            shape = shape[:-2] + (shape[-1], shape[-2])
+        if ty in ['fp16', 'bf16']:
+            ret = torch.randn(shape, dtype=comp_dtype, device=device)
+            # Clamp to avoid relative error issues
+            ret.clamp_(-2**comp_dtype_max_exp, 2**comp_dtype_max_exp - 1)
+        else:
+            ret = torch.randint(256, shape, dtype=torch.uint8, device=device)
+        if col_major:
+            ret = ret.mT
+        return ret
+
+    type_a = normal_type if rhs_scale else mxfp_type
+    type_b = mxfp_type if rhs_scale else normal_type
+
+    DIV_FACTOR_A = 2 if type_a == 'e2m1' else 1
+    DIV_FACTOR_B = 2 if type_b == 'e2m1' else 1
+    x = make_arg((M, K // DIV_FACTOR_A), type_a, col_major=col_a)
+    y = make_arg((K // DIV_FACTOR_B, N), type_b, col_major=col_b)
+
+    min_scale, max_scale = (0, 142) if comp_dtype == torch.bfloat16 else (124, 131)
+    scale_x = torch.randint(min_scale, max_scale + 1, (M, K // 32), dtype=torch.uint8, device=device)
+    scale_y = torch.randint(min_scale, max_scale + 1, (N, K // 32), dtype=torch.uint8, device=device)
+
+    def make_finite(x, dtype):
+        # e5m2 has too many non-finite values when sampled uniformly (1 / 32) and
+        # Fp8E5M2_to_Bf16 doesn't preserve NaNs (fixme)
+        if dtype not in ('e5m2', 'e4m3'):
+            return x
+        if dtype == 'e5m2' and comp_dtype == torch.float16:
+            x = x & 0xB
+        mask = 0x7C if dtype == 'e5m2' else 0x7F
+        finite = torch.arange(x.numel(), device=device, dtype=torch.uint8).reshape_as(x) % mask
+        x_finite = torch.where(x & mask == mask, finite | (0x80 & x), x)
+        x.copy_(x_finite)
+        return x
+
+    x = make_finite(x, type_a)
+    y = make_finite(y, type_b)
+    z = x.new_empty((M, N), dtype=comp_dtype)
+    if rhs_scale:
+        scale_x = None
+    else:
+        scale_y = None
+
+    if provider == 'triton':
+        triton_fn = lambda: dot_scaled(M, N, K, x, y, z, scale_x, scale_y, type_a, type_b, num_warps)
+
+        _, min_ms, max_ms, mean_ms, cv = benchmark_suit.do_bench(triton_fn, n_warmup=10, n_repeat=10,
+                                                                 quantiles=quantiles)
+    else:
+        raise NotImplementedError(f'Unsupported provider {provider}')
+
+    def tflops(ms):
+        scale_ops = N * K if rhs_scale else M * K
+        return (2 * M * N * K + scale_ops) * (1e-12) / (ms * 1e-3)
+
+    def gbps(ms):
+
+        def size_x(m, n, ty):
+            if ty in ['e2m1']:
+                return m * n // 2
+            if ty in ['e4m3', 'e5m2']:
+                return m * n
+            if ty in ['fp16', 'bf16']:
+                return m * n * 2
+            raise NotImplementedError(f'Unsupported type {ty} for scaledot operand')
+
+        tensor_size = size_x(M, K, type_a) + size_x(K, N, type_b)
+        scale_size = (M * K // 32) if rhs_scale else (N * K // 32)
+        return (tensor_size + scale_size + 4.0 * (M * N)) * (1e-9) / (ms * 1e-3)
+
+    return (gbps(mean_ms), gbps(max_ms), gbps(min_ms)), (tflops(mean_ms), tflops(max_ms), tflops(min_ms)), cv
+
+
+if __name__ == '__main__':
+    benchmark.run(show_plots=False, print_data=True)
diff --git a/benchmarks/micro_benchmarks/run_benchmarks.py b/benchmarks/micro_benchmarks/run_benchmarks.py
index 69c4706c90..4f3aad8004 100644
--- a/benchmarks/micro_benchmarks/run_benchmarks.py
+++ b/benchmarks/micro_benchmarks/run_benchmarks.py
@@ -1,6 +1,7 @@
 import argparse
 
 from conversion import float_conversion
+from core_ops import dot_scaled
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
@@ -12,3 +13,4 @@
     )
     args = parser.parse_args()
     float_conversion.benchmark.run(print_data=True, save_path=args.reports)
+    dot_scaled.benchmark.run(print_data=True, save_path=args.reports)


Reviews:
user: login: chengjunlu
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T02:46:05Z
user: login: whitneywhtsang
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T02:51:37Z

Review comments:
diff_hunk: @@ -0,0 +1,157 @@
+import torch
+import triton
+import triton.language as tl
+
+import triton_kernels_benchmark as benchmark_suit
+
+
+@triton.jit
+def dot_scale_kernel(a_base, stride_a0, stride_a1, a_scale, b_base, stride_b0, stride_b1, b_scale, out,
+                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, type_a: tl.constexpr,
+                     type_b: tl.constexpr):
+    DIV_FACTOR_A: tl.constexpr = 2 if type_a == 'e2m1' else 1
+    DIV_FACTOR_B: tl.constexpr = 2 if type_b == 'e2m1' else 1
+    PACKED_BLOCK_K_A: tl.constexpr = BLOCK_K // DIV_FACTOR_A
+    PACKED_BLOCK_K_B: tl.constexpr = BLOCK_K // DIV_FACTOR_B
+    a_ptr = a_base + tl.arange(0, BLOCK_M)[:, None] * stride_a0 + tl.arange(0, PACKED_BLOCK_K_A)[None, :] * stride_a1
+    b_ptr = b_base + tl.arange(0, PACKED_BLOCK_K_B)[:, None] * stride_b0 + tl.arange(0, BLOCK_N)[None, :] * stride_b1
+
+    a = tl.load(a_ptr)
+    b = tl.load(b_ptr)
+    SCALE_BLOCK_K: tl.constexpr = BLOCK_K // 32
+    if a_scale is not None:
+        scale_a_ptr = a_scale + tl.arange(0, BLOCK_M)[:, None] * SCALE_BLOCK_K + tl.arange(0, SCALE_BLOCK_K)[None, :]
+        a_scale = tl.load(scale_a_ptr)
+    if b_scale is not None:
+        scale_b_ptr = b_scale + tl.arange(0, BLOCK_N)[:, None] * SCALE_BLOCK_K + tl.arange(0, SCALE_BLOCK_K)[None, :]
+        b_scale = tl.load(scale_b_ptr)
+    c = tl.dot_scaled(a, a_scale, type_a, b, b_scale, type_b)
+    out_ptr = out + \
+        tl.arange(0, BLOCK_M)[:, None] * BLOCK_N + \
+        tl.arange(0, BLOCK_N)[None, :]
+    tl.store(out_ptr, c.to(tl.bfloat16))
+
+
+def dot_scaled(M, N, K, x, y, z, scale_x, scale_y, type_a, type_b, num_warps):
+    kernel_kwargs = {'num_warps': num_warps}
+    dot_scale_kernel[(1, )](x, *x.stride(), scale_x, y, *y.stride(), scale_y, z, M, N, K, type_a, type_b,
+                            **kernel_kwargs)
+
+
+# Benchmark Performance
+
+
path: benchmarks/micro_benchmarks/core_ops/dot_scaled.py
user: login: whitneywhtsang
body: ```suggestion
```
created_at: 2025-03-04T02:49:47Z
author_association: CONTRIBUTOR

PR comments:
user: login: alexbaden
created_at: 2025-03-04T02:49:38Z
body: What is the runtime for the benchmark job before and after?
author_association: CONTRIBUTOR
user: login: leonling-ll
created_at: 2025-03-04T02:57:06Z
body: > What is the runtime for the benchmark job before and after?

There's not much change.

**Scaled_dot with Fp4ToFp**(After):
```
    M    K    N   col_a  col_b  rhs_scale   mxfp_type  normal_type  Triton-GB/s    Triton-TFlops
0  128  128  128   True   True       True      e2m1        bf16     0.582261         0.022912
1  128  128  128   True   True       True      e4m3        bf16     0.737327         0.026950
2  128  128  128   True   True      False      e2m1        bf16     0.675641         0.026586
3  128  128  128   True   True      False      e4m3        bf16     0.950809         0.034753
4  128  128  128  False  False       True      e2m1        bf16     0.676581         0.026623
5  128  128  128  False  False       True      e4m3        bf16     0.924559         0.033794
6  128  128  128  False  False      False      e2m1        bf16     0.625778         0.024624
7  128  128  128  False  False      False      e4m3        bf16     0.730594         0.026704
```

**Scaled_dot with UpcastMXFP**(Before):
```
    M    K    N   col_a  col_b  rhs_scale   mxfp_type  normal_type  Triton-GB/s    Triton-TFlops
0  128  128  128   True   True       True      e2m1        bf16     0.580115         0.022827
1  128  128  128   True   True       True      e4m3        bf16     0.736573         0.026923
2  128  128  128   True   True      False      e2m1        bf16     0.675726         0.026589
3  128  128  128   True   True      False      e4m3        bf16     0.953642         0.034857
4  128  128  128  False  False       True      e2m1        bf16     0.674874         0.026556
5  128  128  128  False  False       True      e4m3        bf16     0.936129         0.034217
6  128  128  128  False  False      False      e2m1        bf16     0.630572         0.024813
7  128  128  128  False  False      False      e4m3        bf16     0.736290         0.026912
```
author_association: CONTRIBUTOR
user: login: alexbaden
created_at: 2025-03-04T03:19:45Z
body: Thx - but I was curious how long (in ms/seconds) the benchmark takes to run as written. 
author_association: CONTRIBUTOR
user: login: leonling-ll
created_at: 2025-03-04T03:29:34Z
body: > Thx - but I was curious how long (in ms/seconds) the benchmark takes to run as written.

Just rerun it with the new implementation, the average time in ms of above cases:
```
0. avg time 0.1937 ms
1. avg time 0.1599 ms
2. avg time 0.1549 ms
3. avg time 0.1269 ms
4. avg time 0.1664 ms
5. avg time 0.1295 ms
6. avg time 0.1788 ms
7. avg time 0.1495 ms
```
author_association: CONTRIBUTOR