update_test (#2156)

* update_test

* update

* update

Author: amd-ruitang3

op_tests/test_moe.py

@@ -367,7 +367,7 @@ def calculateTensorsSize(*args):
     "--token",
     type=int,
     nargs="*",
-    default=None,
+    default=[128],
     help="""Token Num.
     e.g.: -m 128""",
 )
@@ -376,7 +376,7 @@ def calculateTensorsSize(*args):
     "--hidden_dim",
     type=int,
     nargs="*",
-    default=None,
+    default=[4096],
     help="""Hidden states dim.
     e.g.: -hd 4096""",
 )
@@ -385,7 +385,7 @@ def calculateTensorsSize(*args):
     "--inter_dim",
     type=int,
     nargs="*",
-    default=None,
+    default=[1024],
     help="""Intermediate dim.
     e.g.: -id 1024""",
 )
@@ -410,7 +410,7 @@ def calculateTensorsSize(*args):
 parser.add_argument(
     "-a",
     "--activation",
-    type=str,
+    type=dtypes.str2ActivationType,
     choices=[
         "silu",
         "gelu",
@@ -424,21 +424,16 @@ def calculateTensorsSize(*args):
 
 args = parser.parse_args()
 
-args.activation = {"gelu": ActivationType.Gelu, "silu": ActivationType.Silu}[
-    args.activation
-]
 
 for test in args.test:
     print(f"\nRunning test: {test}")
     if test == "test_fmoe_16_bit":
         print("test test_fmoe 16 bit")
         print("\ng1u0 no quant")
         for dtype in args.dtype:
-            for m in [128, 256] if args.token is None else args.token:
-                for hdim in (
-                    [4096, 8192] if args.hidden_dim is None else args.hidden_dim
-                ):
-                    for idim in [1024] if args.inter_dim is None else args.inter_dim:
+            for m in args.token:
+                for hdim in args.hidden_dim:
+                    for idim in args.inter_dim:
                         expert = 32 if args.expert is None else args.expert
                         topk = 5 if args.topk is None else args.topk
                         test_fmoe(
@@ -453,11 +448,9 @@ def calculateTensorsSize(*args):
                         )
     elif test == "g1u1_no_quant":
         for dtype in args.dtype:
-            for m in [128, 256] if args.token is None else args.token:
-                for hdim in (
-                    [4096, 8192] if args.hidden_dim is None else args.hidden_dim
-                ):
-                    for idim in [1024] if args.inter_dim is None else args.inter_dim:
+            for m in args.token:
+                for hdim in args.hidden_dim:
+                    for idim in args.inter_dim:
                         expert = 32 if args.expert is None else args.expert
                         topk = 5 if args.topk is None else args.topk
                         test_fmoe(
@@ -473,11 +466,9 @@ def calculateTensorsSize(*args):
                         )
     elif test == "g1u1_int8quant":
         for dtype in args.dtype:
-            for m in [128, 256] if args.token is None else args.token:
-                for hdim in (
-                    [4096, 8192] if args.hidden_dim is None else args.hidden_dim
-                ):
-                    for idim in [1024] if args.inter_dim is None else args.inter_dim:
+            for m in args.token:
+                for hdim in args.hidden_dim:
+                    for idim in args.inter_dim:
                         expert = 32 if args.expert is None else args.expert
                         topk = 5 if args.topk is None else args.topk
                         test_fmoe(
@@ -495,11 +486,9 @@ def calculateTensorsSize(*args):
 
     elif test == "g1u1_fp8quant":
         for dtype in args.dtype:
-            for m in [128, 256] if args.token is None else args.token:
-                for hdim in (
-                    [4096, 8192] if args.hidden_dim is None else args.hidden_dim
-                ):
-                    for idim in [1024] if args.inter_dim is None else args.inter_dim:
+            for m in args.token:
+                for hdim in args.hidden_dim:
+                    for idim in args.inter_dim:
                         expert = 32 if args.expert is None else args.expert
                         topk = 5 if args.topk is None else args.topk
                         test_fmoe(
@@ -518,13 +507,9 @@ def calculateTensorsSize(*args):
 
     elif test == "g1u0_int8smoothquant":
         for dtype in args.dtype:
-            for m in [128] if args.token is None else args.token:
-                for hdim in (
-                    [4096, 6144, 8192] if args.hidden_dim is None else args.hidden_dim
-                ):
-                    for idim in (
-                        [512, 1024] if args.inter_dim is None else args.inter_dim
-                    ):
+            for m in args.token:
+                for hdim in args.hidden_dim:
+                    for idim in args.inter_dim:
                         expert = 32 if args.expert is None else args.expert
                         topk = 5 if args.topk is None else args.topk
                         test_fmoe(
@@ -541,15 +526,9 @@ def calculateTensorsSize(*args):
 
     elif test == "g1u1_int8smoothquant":
         for dtype in args.dtype:
-            for m in [128] if args.token is None else args.token:
-                for hdim in (
-                    [4096, 6144, 8192] if args.hidden_dim is None else args.hidden_dim
-                ):
-                    for idim in (
-                        [512, 1024, 1280, 1536]
-                        if args.inter_dim is None
-                        else args.inter_dim
-                    ):
+            for m in args.token:
+                for hdim in args.hidden_dim:
+                    for idim in args.inter_dim:
                         expert = 32 if args.expert is None else args.expert
                         topk = 5 if args.topk is None else args.topk
                         test_fmoe(
@@ -566,13 +545,9 @@ def calculateTensorsSize(*args):
 
     elif test == "g1u1_fp8smoothquant":
         for dtype in args.dtype:
-            for m in [128] if args.token is None else args.token:
-                for hdim in (
-                    [4096, 6144, 8192] if args.hidden_dim is None else args.hidden_dim
-                ):
-                    for idim in (
-                        [512, 1024, 1280] if args.inter_dim is None else args.inter_dim
-                    ):
+            for m in args.token:
+                for hdim in args.hidden_dim:
+                    for idim in args.inter_dim:
                         expert = 32 if args.expert is None else args.expert
                         topk = 5 if args.topk is None else args.topk
                         test_fmoe(
@@ -588,13 +563,9 @@ def calculateTensorsSize(*args):
                         )
     elif test == "g1u1_int4":
         for dtype in args.dtype:
-            for m in [32, 128] if args.token is None else args.token:
-                for hdim in (
-                    [4096, 6144] if args.hidden_dim is None else args.hidden_dim
-                ):
-                    for idim in (
-                        [1024, 4096] if args.inter_dim is None else args.inter_dim
-                    ):
+            for m in args.token:
+                for hdim in args.hidden_dim:
+                    for idim in args.inter_dim:
                         expert = 8 if args.expert is None else args.expert
                         topk = 3 if args.topk is None else args.topk
                         test_fmoe(

op_tests/test_pa.py

@@ -974,9 +974,6 @@ def test_paged_attention(
 
 
 df = []
-l_num_heads = [(4, 1), (8, 1), (32, 8)]
-l_ctx_len = [7, 26, 57, 66, 109, 128, 257, 282, 4097]
-l_dtype = ["fp16", "bf16"]
 
 parser = argparse.ArgumentParser(
     formatter_class=argparse.RawTextHelpFormatter,
@@ -985,53 +982,35 @@ def test_paged_attention(
 parser.add_argument(
     "-d",
     "--dtype",
-    type=str,
-    choices=l_dtype,
-    nargs="?",
-    const=None,
-    default=None,
+    type=dtypes.str2Dtype,
+    nargs="*",
+    default=[dtypes.d_dtypes["fp16"], dtypes.d_dtypes["bf16"]],
     help="""Data type.
     e.g.: -d bf16""",
 )
-
 parser.add_argument(
     "-n",
     "--num_heads",
     type=dtypes.str2tuple,
-    choices=l_num_heads,
-    nargs="?",
-    const=None,
-    default=None,
+    nargs="*",
+    default=[(4, 1), (8, 1), (32, 8)],
     help="""Number of heads (num_query_heads, num_kv_heads)
     e.g.: -n 4,1""",
 )
-
 parser.add_argument(
     "-c",
     "--ctx_len",
     type=int,
-    choices=l_ctx_len,
-    nargs="?",
-    const=None,
-    default=None,
+    nargs="*",
+    default=[7, 26, 57, 66, 109, 128, 257, 282, 4097],
     help="""Context length.
     e.g. -c 128""",
 )
-
 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.num_heads is not None:
-    l_num_heads = [args.num_heads]
-if args.ctx_len is not None:
-    l_ctx_len = [args.ctx_len]
-
-
-for num_heads in l_num_heads:
-    for ctx_len in l_ctx_len:
-        for dtype in l_dtype:
+
+for num_heads in args.num_heads:
+    for ctx_len in args.ctx_len:
+        for dtype in args.dtype:
             ret = test_paged_attention(
                 ctx_len, 128, num_heads, 128, False, 16, dtype, "auto", 0, "cuda:0"
             )

op_tests/test_quant.py

@@ -12,6 +12,7 @@
 from aiter import get_hip_quant, get_torch_quant, get_triton_quant
 import itertools
 import argparse
+import pandas as pd
 
 torch.set_default_device("cuda")
 
@@ -64,30 +65,16 @@ def test_quant(m, n, q_type, q_dtype, h_dtype):
     return ret
 
 
-d_quant = {
-    "fp8_tensor": (aiter.QuantType.per_Tensor, dtypes.fp8),
-    "fp8_token": (aiter.QuantType.per_Token, dtypes.fp8),
-    "fp8_1x128": (aiter.QuantType.per_1x128, dtypes.fp8),
-    "i8_token": (aiter.QuantType.per_Token, dtypes.i8),
-    # 'fp4x2-1x32': (aiter.QuantType.per_1x32, dtypes.fp4x2),
-}
-list_dtype = ["fp16", "bf16"]
-l_n = [4096, 8192]
-l_m = [1, 2, 16, 32, 64, 128, 192, 256, 512, 1024, 16384, 163840]
-import pandas as pd
-
 parser = argparse.ArgumentParser(
     formatter_class=argparse.RawTextHelpFormatter,
     description="config input of test",
 )
 parser.add_argument(
     "-d",
     "--dtype",
-    type=str,
-    choices=list_dtype,
-    nargs="?",
-    const=None,
-    default=None,
+    type=dtypes.str2Dtype,
+    nargs="*",
+    default=[dtypes.d_dtypes["fp16"], dtypes.d_dtypes["bf16"]],
     help="""Data type.
     e.g.: -d bf16""",
 )
@@ -96,7 +83,7 @@ def test_quant(m, n, q_type, q_dtype, h_dtype):
     "--n",
     type=int,
     nargs="*",
-    default=None,
+    default=[4096, 8192],
     help="""N of mnk.
     e.g.: -n 1024""",
 )
@@ -105,10 +92,17 @@ def test_quant(m, n, q_type, q_dtype, h_dtype):
     "--m",
     type=int,
     nargs="*",
-    default=None,
+    default=[1, 2, 16, 32, 64, 128, 192, 256, 512, 1024, 16384, 163840],
     help="""M of mnk.
     e.g.: -m 32""",
 )
+d_quant = {
+    "fp8_tensor": (aiter.QuantType.per_Tensor, dtypes.fp8),
+    "fp8_token": (aiter.QuantType.per_Token, dtypes.fp8),
+    "fp8_1x128": (aiter.QuantType.per_1x128, dtypes.fp8),
+    "i8_token": (aiter.QuantType.per_Token, dtypes.i8),
+    # 'fp4x2-1x32': (aiter.QuantType.per_1x32, dtypes.fp4x2),
+}
 parser.add_argument(
     "-q",
     "--quant",
@@ -121,23 +115,15 @@ def test_quant(m, n, q_type, q_dtype, h_dtype):
 )
 
 args = parser.parse_args()
-if args.dtype is None:
-    list_dtype = [dtypes.d_dtypes[key] for key in list_dtype]
-else:
-    list_dtype = [dtypes.d_dtypes[args.dtype]]
 list_quant = [d_quant[key] for key in args.quant]
-if args.n is not None:
-    l_n = args.n
-if args.m is not None:
-    l_m = args.m
 
 for (
     (q_type, q_dtype),
     h_dtype,
-) in itertools.product(list_quant, list_dtype):
+) in itertools.product(list_quant, args.dtype):
     df = []
-    for n in l_n:
-        for m in l_m:
+    for n in args.n:
+        for m in args.m:
             ret = test_quant(m, n, q_type, q_dtype, h_dtype)
             df.append(ret)
     df = pd.DataFrame(df)