diff --git a/op_tests/test_gemm_a4w4.py b/op_tests/test_gemm_a4w4.py index 51b34d1cfb..aeaca56d35 100644 --- a/op_tests/test_gemm_a4w4.py +++ b/op_tests/test_gemm_a4w4.py @@ -155,83 +155,6 @@ def test_gemm(dtype, M, N, K): return ret -l_dtype = ["bf16"] -l_mnk = [ - # pure_compute - (256, 2048, 8192), - (2048, 8192, 8192), - (16384, 16384, 16384), - (32768, 106496, 16384), - (32768, 16384, 53248), - (32768, 18432, 16384), - (32768, 16384, 16384), - (128, 106496, 16384), - (128, 16384, 53248), - (128, 18432, 16384), - (128, 16384, 16384), - (64, 106496, 16384), - (64, 16384, 53248), - (64, 18432, 16384), - (64, 16384, 16384), - (64, 106496, 16384), - (32, 106496, 16384), - (32, 16384, 53248), - (32, 18432, 16384), - (32, 16384, 16384), - # qkv_proj - (1, 1280, 8192), - (64, 1280, 8192), - (127, 1280, 8192), - (129, 1280, 8192), - (65, 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), - # attn_out - (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), - (16384, 8192, 1024), - # tune - (1552, 8192, 8192), - (1664, 8192, 8192), - (1792, 8192, 8192), - (1920, 8192, 8192), - (3072, 8192, 8192), - (1552, 10240, 8192), - (1664, 10240, 8192), - (1792, 10240, 8192), - (1920, 10240, 8192), - (3072, 10240, 8192), - (1552, 57344, 8192), - (1664, 57344, 8192), - (1792, 57344, 8192), - (1920, 57344, 8192), - (3072, 57344, 8192), - (1552, 8192, 28672), - (1664, 8192, 28672), - (1792, 8192, 28672), - (1920, 8192, 28672), - (3072, 8192, 28672), -] - parser = argparse.ArgumentParser( formatter_class=argparse.RawTextHelpFormatter, description="config input of test", @@ -239,11 +162,11 @@ def test_gemm(dtype, M, N, K): parser.add_argument( "-d", "--dtype", - type=str, - choices=l_dtype, - nargs="?", - const=None, - default=None, + type=dtypes.str2Dtype, + nargs="*", + choices=[dtypes.d_dtypes["bf16"]], + metavar="{bf16}", + default=[dtypes.d_dtypes["bf16"]], help="""Data type. e.g.: -d bf16""", ) @@ -251,24 +174,91 @@ def test_gemm(dtype, M, N, K): "-mnk", "--shape", type=dtypes.str2tuple, - nargs="?", - const=None, - default=None, + nargs="*", + default=[ + # pure_compute + (256, 2048, 8192), + (2048, 8192, 8192), + (16384, 16384, 16384), + (32768, 106496, 16384), + (32768, 16384, 53248), + (32768, 18432, 16384), + (32768, 16384, 16384), + (128, 106496, 16384), + (128, 16384, 53248), + (128, 18432, 16384), + (128, 16384, 16384), + (64, 106496, 16384), + (64, 16384, 53248), + (64, 18432, 16384), + (64, 16384, 16384), + (64, 106496, 16384), + (32, 106496, 16384), + (32, 16384, 53248), + (32, 18432, 16384), + (32, 16384, 16384), + # qkv_proj + (1, 1280, 8192), + (64, 1280, 8192), + (127, 1280, 8192), + (129, 1280, 8192), + (65, 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), + # attn_out + (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), + (16384, 8192, 1024), + # tune + (1552, 8192, 8192), + (1664, 8192, 8192), + (1792, 8192, 8192), + (1920, 8192, 8192), + (3072, 8192, 8192), + (1552, 10240, 8192), + (1664, 10240, 8192), + (1792, 10240, 8192), + (1920, 10240, 8192), + (3072, 10240, 8192), + (1552, 57344, 8192), + (1664, 57344, 8192), + (1792, 57344, 8192), + (1920, 57344, 8192), + (3072, 57344, 8192), + (1552, 8192, 28672), + (1664, 8192, 28672), + (1792, 8192, 28672), + (1920, 8192, 28672), + (3072, 8192, 28672), + ], help="""Shape of mnk. e.g. -mnk 1280,8192,1024""", ) args = parser.parse_args() -if args.dtype is None: - l_dtype = [dtypes.d_dtypes[key] for key in l_dtype] -else: - l_dtype = [dtypes.d_dtypes[args.dtype]] -if args.shape is not None: - l_mnk = [args.shape] df = [] -for dtype in l_dtype: - for m, n, k in l_mnk: +for dtype in args.dtype: + for m, n, k in args.shape: ret = test_gemm(dtype, m, n, k) df.append(ret) df = pd.DataFrame(df) diff --git a/op_tests/test_gemm_a8w8_blockscale.py b/op_tests/test_gemm_a8w8_blockscale.py index 0ee7b2f12e..1c4e0bdd77 100755 --- a/op_tests/test_gemm_a8w8_blockscale.py +++ b/op_tests/test_gemm_a8w8_blockscale.py @@ -138,54 +138,19 @@ def run_asm(x, weight, x_scale, w_scale, dtype=dtypes.bf16, kernel_name=None): parser.add_argument( "-d", "--dtype", - type=str, - choices=["bf16"], - nargs="?", - const=None, - default=None, + type=dtypes.str2Dtype, + choices=[dtypes.d_dtypes["bf16"]], + nargs="*", + default=[dtypes.d_dtypes["bf16"]], + metavar="{bf16}", help="""Data type. e.g.: -d bf16""", ) parser.add_argument( "-m", type=int, - nargs="?", - const=None, - default=None, - help="""M of mnk. - e.g.: -m 32""", -) -parser.add_argument( - "-nk", - type=dtypes.str2tuple, - nargs="?", - const=None, - default=None, - help="""N&K of mnk. - e.g.: -nk 4096,512""", -) -parser.add_argument( - "--ck_preshuffle", - nargs="?", - default=[True, False], - help="weight ck_preshuffle or not", -) - -args = parser.parse_args() -if args.dtype is None: - l_dtype = [dtypes.d_dtypes[key] for key in ["bf16"]] -else: - l_dtype = [dtypes.d_dtypes[args.dtype]] -if args.m is not None: - l_m = [args.m] -if args.nk is not None: - l_nk = [args.nk] -l_preshuffle: List[bool] = args.ck_preshuffle - -df = [] -for dtype in [dtypes.bf16]: - # deepseek-r1 - for m in [ + nargs="*", + choices=[ 1, 2, 4, @@ -213,15 +178,79 @@ def run_asm(x, weight, x_scale, w_scale, dtype=dtypes.bf16, kernel_name=None): 6144, 8192, 10240, - ]: - for n, k in [ - (24576, 1536), - # (32768, 512), - # (7168, 16384), - # (36864, 7168), - ]: - ret = test_gemm(dtype, m, n, k, ck_preshuffle=False) - df.append(ret) + ], + default=[ + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 96, + 128, + 160, + 192, + 224, + 256, + 288, + 320, + 352, + 384, + 416, + 448, + 480, + 512, + 1024, + 2048, + 4096, + 6144, + 8192, + 10240, + ], + help="""M of mnk. + e.g.: -m 32""", +) +parser.add_argument( + "-nk", + type=dtypes.str2tuple, + nargs="*", + choices=[ + (24576, 1536), + # (32768, 512), + # (7168, 16384), + # (36864, 7168), + ], + default=[ + (24576, 1536), + # (32768, 512), + # (7168, 16384), + # (36864, 7168), + ], + help="""N&K of mnk. + e.g.: -nk 24576,1536""", +) +parser.add_argument( + "--ck_preshuffle", + type=dtypes.str2bool, + nargs="*", + default=[True, False], + help="""weight ck_preshuffle or not. + e.g.: --ck_preshuffle True + or --ck_preshuffle False + """, +) + +args = parser.parse_args() + +df = [] +for dtype in args.dtype: + # deepseek-r1 + for m in args.m: + for n, k in args.nk: + for ck_p in args.ck_preshuffle: + ret = test_gemm(dtype, m, n, k, ck_preshuffle=ck_p) + df.append(ret) df = pd.DataFrame(df) # Configure pandas to show all columns without truncation diff --git a/op_tests/test_gemm_a8w8_blockscale_mi350.py b/op_tests/test_gemm_a8w8_blockscale_mi350.py deleted file mode 100755 index 511f1c2837..0000000000 --- a/op_tests/test_gemm_a8w8_blockscale_mi350.py +++ /dev/null @@ -1,214 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. - -import torch -import torch.nn.functional as F -import aiter -from aiter import dtypes -from aiter.test_common import checkAllclose, perftest, benchmark -from einops import rearrange -from einops import repeat as eirp -from aiter.ops.shuffle import shuffle_weight -import pandas as pd -import argparse - -block_shape = (128, 128) - - -@perftest(num_iters=5) -def run_torch(x, weight, x_scale, w_scale, dtype=dtypes.bf16): - block_shape_n, block_shape_k = block_shape - m, k = x.shape - n = weight.shape[0] - scale_n = (n + block_shape_n - 1) // block_shape_n - scale_k = (k + block_shape_k - 1) // block_shape_k - x = x.to(x_scale.dtype).view( - m, k // block_shape[1], block_shape[1] - ) * x_scale.unsqueeze(-1) - x = x.view(m, k) - - w_scale = rearrange( - w_scale.view(-1, 1) - .repeat(1, block_shape_n * block_shape_k) - .view(scale_n, scale_k, block_shape_n, block_shape_k), - "num_blk_n num_blk_k blk_n blk_k -> (num_blk_n blk_n) (num_blk_k blk_k)", - ) - w_scale = w_scale[:n, :k] - weight = weight.to(w_scale.dtype) * w_scale - - out = F.linear(x.to(dtypes.fp32), weight.to(dtypes.fp32)) - return out.to(dtype) - - -@perftest() -def run_gemm_ck(x, weight, x_scale, w_scale, dtype=dtypes.bf16): - return aiter.gemm_a8w8_blockscale(x, weight, x_scale, w_scale, dtype) - - -@benchmark() -def test_gemm(dtype, m, n, k): - dim = (m, n, k) - block_shape_n, block_shape_k = block_shape - scale_n = (n + block_shape_n - 1) // block_shape_n - scale_k = (k + block_shape_k - 1) // block_shape_k - x = (torch.rand((m, k), dtype=dtypes.fp16, device="cuda") / 10).to(dtypes.fp8) - weight = (torch.rand((n, k), dtype=dtypes.fp16, device="cuda") / 10).to(dtypes.fp8) - x_scale = torch.rand([m, scale_k], dtype=dtypes.fp32, device="cuda") - w_scale = torch.rand([scale_n, scale_k], dtype=dtypes.fp32, device="cuda") - - a, avg_a = run_torch(x, weight, x_scale, w_scale, dtype) - b, avg_b = run_gemm_ck(x, weight, x_scale, w_scale, dtype) - - msg = f"[perf] dim: {str(dim):<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) - - return {"us": avg_b} - - -@perftest(num_iters=5) -def run_torch2(x, weight, x_scale, w_scale, dtype=dtypes.bf16): - block_shape_n, block_shape_k = block_shape - m, k = x.shape - n = weight.shape[0] - - x_scale_ = eirp(x_scale, "m k -> m (k repeat)", repeat=block_shape_k) - x_scale_ = x_scale_[:m, :k] - - w_scale_ = eirp(w_scale, "n k -> (n repeat) k", repeat=block_shape_n) - w_scale_ = eirp(w_scale_, "n k -> n (k repeat)", repeat=block_shape_k) - w_scale_ = w_scale_[:n, :k] - - x_ = x.to(x_scale.dtype) * x_scale_ - weight_ = weight.to(w_scale.dtype) * w_scale_ - - out = F.linear(x_.to(dtypes.fp32), weight_.to(dtypes.fp32)) - return out.to(dtype) - - -@perftest() -def run_asm(x, weight, x_scale, w_scale, dtype=dtypes.bf16): - # return aiter.flatmm_a8w8_blockscale_ASM(x, weight, x_scale, w_scale, dtype) - return aiter.gemm_a8w8_blockscale(x, weight, x_scale, w_scale, dtype, True) - - -@benchmark() -def test_gemm_asm_gfx950(dtype, m, n, k): - from aiter.jit.utils.chip_info import get_gfx - - if get_gfx() not in ["gfx950"]: - return - dim = (m, n, k) - block_shape_n, block_shape_k = block_shape - print("block_shape", block_shape) - scale_m = m - scale_n = (n + block_shape_n - 1) // block_shape_n - scale_k = (k + block_shape_k - 1) // block_shape_k - - x = (torch.rand((m, k), dtype=dtypes.fp16, device="cuda") / 10).to(dtypes.fp8) - weight = (torch.rand((n, k), dtype=dtypes.fp16, device="cuda") / 10).to(dtypes.fp8) - x_scale = torch.rand([scale_k, scale_m], dtype=dtypes.fp32, device="cuda") - w_scale = torch.rand([scale_k, scale_n], dtype=dtypes.fp32, device="cuda") - x_scale_trans = torch.transpose(x_scale, 0, 1) - w_scale_trans = torch.transpose(w_scale, 0, 1).contiguous() - flat_weight = shuffle_weight(weight, layout=(16, 16)) - a, avg_a = run_torch2(x, weight, x_scale_trans, w_scale_trans, dtypes.bf16) - b, avg_b = run_asm(x, flat_weight, x_scale, w_scale_trans, dtype) - a = a.to(dtypes.fp32) - b = b.to(dtypes.fp32) - tflops = 2 * m * n * k / (avg_b) / 1e6 - msg = f"[perf] dim: {str(dim):<20} dtype: {dtype}, torch avg: {avg_a:<8.2f} us, asm avg: {avg_b:<8.2f} us, uplift: {avg_a/avg_b -1:<5.1%}, tflops:{tflops:<8.2f}" - checkAllclose(a, b, msg="a,b: " + msg, rtol=1e-2, atol=1e-2) - - -l_dtype = ["bf16"] -l_m = [16, 32, 64, 128, 256, 512, 1024, 1536, 2048, 4096, 8192, 16384, 20480] -l_nk = [ - (1536, 7168), - (3072, 1536), - (576, 7168), - (7168, 256), - (7168, 2048), - (4608, 7168), - (7168, 2304), - (512, 7168), - (4096, 512), -] - -parser = argparse.ArgumentParser( - formatter_class=argparse.RawTextHelpFormatter, - description="config input of test", -) -parser.add_argument( - "-d", - "--dtype", - type=str, - choices=l_dtype, - nargs="?", - const=None, - default=None, - help="""Data type. - e.g.: -d bf16""", -) -parser.add_argument( - "-m", - type=int, - nargs="?", - const=None, - default=None, - help="""M of mnk. - e.g.: -m 32""", -) -parser.add_argument( - "-nk", - type=dtypes.str2tuple, - nargs="?", - const=None, - default=None, - help="""N&K of mnk. - e.g.: -nk 4096,512""", -) - -args = parser.parse_args() -if args.dtype is None: - l_dtype = [dtypes.d_dtypes[key] for key in l_dtype] -else: - l_dtype = [dtypes.d_dtypes[args.dtype]] -if args.m is not None: - l_m = [args.m] -if args.nk is not None: - l_nk = [args.nk] - -for dtype in [dtypes.bf16]: - # deepseek-r1 - for m in [ - 16, - 17, - 31, - 33, - 127, - 129, - 32, - 64, - 128, - 256, - 512, - 1024, - 1536, - 2048, - 4096, - 8192, - 16384, - 20480, - ]: - for n, k in [ - (1536, 7168), - (3072, 1536), - (7168, 2048), - (4608, 7168), - (7168, 2304), - (512, 7168), - (4096, 512), - ][1:2]: - continue # Disabled now due to coredump issue - test_gemm_asm_gfx950(dtype, m, n, k) - break