aiter/op_tests/test_batched_gemm_a8w8.py at main · lucas-santos-amd/aiter · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

import argparse

import torch
import torch.nn.functional as F

import aiter
from aiter import dtypes
from aiter.test_common import checkAllclose, perftest


@perftest(num_iters=5)
def run_torch(x, weight, x_scale, w_scale, bias=None, dtype=dtypes.bf16):
    B = x.size(0)
    M = x.size(1)
    N = weight.size(1)
    out = torch.empty(B, M, N, dtype=dtypes.bf16, device="cuda")
    for b in range(B):
        b_x = F.linear(x[b, :, :].to(dtypes.fp32), weight[b, :, :].to(dtypes.fp32))
        b_scale = torch.matmul(x_scale[b, :, :], w_scale[b, :, :])
        b_out = torch.mul(b_x, b_scale)
        if bias is not None:
            b_out = b_out.to(bias[b, :, :]) + bias[b, :, :]
        out[b, :, :] = b_out
    return out.to(dtype)


@perftest()
def run_gemm_ck(x, weight, x_scale, w_scale, bias=None, dtype=dtypes.bf16):
    return aiter.batched_gemm_a8w8_CK(x, weight, x_scale, w_scale, bias)


def test_gemm(dtype, b, m, n, k):
    dim = (b, m, n, k)
    x = torch.randint(-20, 20, (b, m, k), dtype=dtypes.i8).cuda()
    weight = torch.randint(-20, 20, (b, n, k), dtype=dtypes.i8).cuda()
    x_scale = torch.rand([b, m, 1], dtype=dtypes.fp32).cuda() + 1e-6
    w_scale = torch.rand([b, 1, n], dtype=dtypes.fp32).cuda() + 1e-6

    a, avg_a = run_torch(x, weight, x_scale, w_scale, None, dtype)
    b, avg_b = run_gemm_ck(x, weight, x_scale, w_scale, None, dtype)
    msg = f"[perf] dim: {dim!s:<20} dtype: {dtype}, torch avg: {avg_a:<8.2f} us, ck avg: {avg_b:<8.2f} us, uplift: {avg_a/avg_b-1:<5.1%}"
    checkAllclose(
        a, b, msg="a,b: " + msg, rtol=1e-2, atol=0.01, catastrophic_check=True
    )


parser = argparse.ArgumentParser(
    formatter_class=argparse.RawTextHelpFormatter,
    description="config input of test",
)
parser.add_argument(
    "-d",
    "--dtype",
    type=dtypes.str2Dtype,
    choices=[dtypes.d_dtypes["bf16"]],
    nargs="*",
    default="bf16,",
    metavar="{bf16}",
    help="""Data type.
    e.g.: -d bf16""",
)
parser.add_argument(
    "-b",
    "--batch",
    type=int,
    choices=[16],
    nargs="*",
    default=[16],
    help="""Batch size.
    e.g.: -b 16""",
)
parser.add_argument(
    "-s",
    "--mnk",
    type=dtypes.str2tuple,
    nargs="*",
    default=[
        (1, 1280, 8192),
        (32, 1280, 8192),
        (64, 1280, 8192),
        (128, 1280, 8192),
        (192, 1280, 8192),
        (256, 1280, 8192),
        (320, 1280, 8192),
        (512, 1280, 8192),
        (1024, 1280, 8192),
        (2048, 1280, 8192),
        (4096, 1280, 8192),
        (8192, 1280, 8192),
        (1, 8192, 1024),
        (32, 8192, 1024),
        (64, 8192, 1024),
        (128, 8192, 1024),
        (192, 8192, 1024),
        (256, 8192, 1024),
        (320, 8192, 1024),
        (512, 8192, 1024),
        (1024, 8192, 1024),
        (2048, 8192, 1024),
        (4096, 8192, 1024),
        (8192, 8192, 1024),
    ],
    help="""Shape of mnk.
    e.g.:   -s 1024,8192,1024
            --mnk 1024,8192,1024""",
)

args = parser.parse_args()


for dtype in args.dtype:
    for b in args.batch:
        for m, n, k in args.mnk:
            test_gemm(dtype, b, m, n, k)