refactor model benchmarks (#704)

Added couple fixes that were noticed when benchmarking FA triton compared to torch:

output initialization before int8. Otherwise the output will be int8.
changed the 'sl' to 'sq' (marking sequence length) in rmsnorm and softmax aswell, to be more continuous between kernels.
removed max_ctx_len as its not a well defined model parameter. N_CTX_Q is rather picked from args.sq which has a default value 4096.
for clarity I also print the D_HEAD in the output because different models can have different values for this.
converting thd and bshd layouts to torch compatible

Co-authored-by: Tianxing Wu <[email protected]>

Author: juuso-oskari

python/perf-kernels/flash-attention.py

@@ -1879,10 +1879,10 @@ def model_benchmark_configs(args):
     for model_name, config in configs.items():
         HQ = config["num_attention_heads"]
         HK = HQ if config["num_key_value_heads"] is None else config["num_key_value_heads"]
-        max_ctx_len = config["max_ctx_len"]
-        N_CTX_Q = args.sq if args.sq else max_ctx_len
-        N_CTX_K = args.sk if args.sk else max_ctx_len
-        fa_configs.append((model_name, batch_size, HQ, HK, N_CTX_Q, N_CTX_K))
+        N_CTX_Q = args.sq if args.sq else 4096
+        N_CTX_K = args.sk if args.sk else N_CTX_Q
+        HEAD_DIM = config["hidden_size"] // HQ
+        fa_configs.append((model_name, batch_size, HQ, HK, N_CTX_Q, N_CTX_K, HEAD_DIM))
 
     return fa_configs
 
@@ -1902,6 +1902,7 @@ def run_benchmark(custom, args):
     varlen = args.layout == 'thd'
     configs = []
     plot_name = f'fused-attention-{mode}-d{head_size}-layout{args.layout}'
+    extra_args = {'D_HEAD': head_size, 'dtype': dtype, 'causal': causal, 'mode': mode}
     if custom:
         x_vals_list = [(args.b, args.hq, hk, args.sq, sk)]
     else:
@@ -1912,16 +1913,16 @@ def run_benchmark(custom, args):
 
         if args.model:
             x_vals_list = model_benchmark_configs(args)
-            x_names = ['model', 'BATCH', 'HQ', 'HK', 'N_CTX_Q', 'N_CTX_K']
+            x_names = ['model', 'BATCH', 'HQ', 'HK', 'N_CTX_Q', 'N_CTX_K', 'D_HEAD']
             plot_name = f'fused-attention-{mode}-layout{args.layout}'
+            extra_args = {'dtype': dtype, 'causal': causal, 'mode': mode}
 
     print_time = args.return_time
     line_vals = ['triton', 'torch']  # 'Time (ms)' if print_time else 'TFLOPS'
     configs.append(
         triton.testing.Benchmark(x_names=x_names, x_vals=x_vals_list, line_arg='provider', line_vals=line_vals,
-                                 line_names=line_vals, styles=[('red', '-'),
-                                                               ('green', '-')], ylabel='ms', plot_name=plot_name,
-                                 args={'D_HEAD': head_size, 'dtype': dtype, 'causal': causal, 'mode': mode}))
+                                 line_names=line_vals, styles=[('green', '-'), ('red', '-')],
+                                 ylabel='Time (ms)' if print_time else 'TFLOPS', plot_name=plot_name, args=extra_args))
 
     @triton.testing.perf_report(configs)
     def bench_flash_attention(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, causal, mode, provider, device="cuda",
@@ -1956,26 +1957,35 @@ def bench_flash_attention(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, causal
             flops_per_matmul = 2.0 * BATCH * HQ * N_CTX_Q * N_CTX_K * D_HEAD
         if causal:
             input_metadata.need_causal()
-        if int8:
-            q, k, v = quantize_input(q, k, v, input_metadata, quantize_p=quantize_p, int8_kv=int8_kv)
 
-        input_metadata.set_persistent(args.persistent)
-        o = torch.empty_like(q)
-        fn = lambda: attention(q, k, v, o, input_metadata)
-        if mode == 'bwd':
-            o, _ = fn()
-            do = torch.randn_like(o)
-            fn = lambda: o.backward(do, retain_graph=True)
-
-        if "torch" in provider:
-            if HQ != HK:
-                k = k.view(k.shape[0], k.shape[1], -1, k.shape[2],
-                           k.shape[3]).expand(-1, -1, HQ // HK, -1, -1).reshape(k.shape[0], -1, k.shape[2], k.shape[3])
-                v = v.view(v.shape[0], v.shape[1], -1, v.shape[2],
-                           v.shape[3]).expand(-1, -1, HQ // HK, -1, -1).reshape(v.shape[0], -1, v.shape[2], v.shape[3])
-
-            fn = lambda: torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0,
-                                                                          is_causal=causal, scale=None)
+        if "triton" in provider:
+            o = torch.empty_like(q)
+            if int8:
+                q, k, v = quantize_input(q, k, v, input_metadata, quantize_p=quantize_p, int8_kv=int8_kv)
+            input_metadata.set_persistent(args.persistent)
+            fn = lambda: attention(q, k, v, o, input_metadata)
+            if mode == 'bwd':
+                o, _ = fn()
+                do = torch.randn_like(o)
+                fn = lambda: o.backward(do, retain_graph=True)
+
+        elif "torch" in provider and args.layout in ["thd", "bhsd", "bshd"]:
+            # torch requires the layout to be (b (optional),...,h,s,d)
+            if args.layout in ["thd", "bshd"]:
+                q = q.transpose(-3, -2)
+                k = k.transpose(-3, -2)
+                v = v.transpose(-3, -2)
+            # check if GQA
+            HQ = q.shape[-3]
+            HK = k.shape[-3]
+            if HQ != HK:  # TODO: sdpa(..., enable_gqa=True work) should work
+                k = k.repeat_interleave(q.size(-3) // k.size(-3), -3)
+                v = v.repeat_interleave(q.size(-3) // v.size(-3), -3)
+
+            fn = lambda: torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=0.0, is_causal=causal, scale=input_metadata.sm_scale)
+        else:
+            assert False, f"Unknown provider {provider} in flash-attention."
 
         ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
         total_flops = 2 * flops_per_matmul
@@ -1984,9 +1994,9 @@ def bench_flash_attention(BATCH, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, causal
             seqlen_q = N_CTX_Q
             seqlen_k = N_CTX_K
             if seqlen_q > seqlen_k:
-                total_flops *= seqlen_k / (2 * seqlen_q)
+                total_flops *= (seqlen_k / (2 * seqlen_q))
             else:
-                total_flops *= 1 - seqlen_q / (2 * seqlen_k)
+                total_flops *= (1 - seqlen_q / (2 * seqlen_k))
         if mode == "bwd":
             total_flops *= 2.5  # 2.0(bwd) + 0.5(recompute)
         if print_time:
@@ -2014,8 +2024,9 @@ def parse_args():
     parser.add_argument('-model_configs', type=str, default="model_configs.json", help="Model config json file.")
 
     available_models = get_available_models(model_families=["llama3"])  # Dynamically load model names
-    model_help = ("Model name to benchmark. Select from: [" + ", ".join(available_models) +
-                  "]. Use 'all' to benchmark all models or leave blank for the default benchmark script.")
+    model_help = (
+        "Model name to benchmark. Select from: [" + ", ".join(available_models) +
+        "]. Use 'all' to benchmark all models. Not providing runs the default benchmark script with custom configs.")
     parser.add_argument('-model', type=str, default=None, help=model_help)
     parser.add_argument("-b", type=int, default=0)
     parser.add_argument("-hq", type=int, default=0)

python/perf-kernels/gemm.py

@@ -315,14 +315,12 @@ def parse_args():
     parser.add_argument('-model_configs', type=str, default="model_configs.json", help="Model config json file.")
 
     available_models = get_available_models(model_families=["llama3"])  # Dynamically load model names
-    model_help = ("Model name to benchmark. Select from: [" + ", ".join(available_models) +
-                  "]. Use 'all' to benchmark all models or leave blank for the default benchmark script.")
+    model_help = (
+        "Model name to benchmark. Select from: [" + ", ".join(available_models) +
+        "]. Use 'all' to benchmark all models. Not providing runs the default benchmark script with custom configs.")
     parser.add_argument('-model', type=str, default=None, help=model_help)
-    parser.add_argument('-b', type=int, default=0,
-                        help="Batch size used together with model. Defaults to 1 if not provided.")
-    parser.add_argument(
-        '-sl', type=int, default=0,
-        help="Sequence length used together with model. Defaults to max_seq_len from model config if not provided.")
+    parser.add_argument('-b', type=int, default=0, help="Batch size used together with model.")
+    parser.add_argument('-sq', type=int, default=0, help="Sequence length used together with model.")
 
     parser.add_argument("-v", action='store_true', default=False, help="Print out the best tuning config")
     parser.add_argument("-M", type=int, default=0)
@@ -348,7 +346,7 @@ def main():
         batch_size = args.b if args.b else 1
 
         for model_name, config in configs.items():
-            seq_len = args.sl if args.sl else config["max_ctx_len"]
+            seq_len = args.sq if args.sq else 4096
             M, N, K = batch_size * seq_len, config["hidden_size"], config["intermediate_size"]
             mnk_list.append((model_name, M, N, K))
 

python/perf-kernels/model_configs.json

@@ -4,23 +4,20 @@
       "num_attention_heads": 32,
       "num_key_value_heads": 8,
       "hidden_size": 4096,
-      "max_ctx_len": 8192,
       "intermediate_size": 14336,
       "vocab_size": 128256
     },
     "70B": {
       "num_attention_heads": 64,
       "num_key_value_heads": 8,
       "hidden_size": 8192,
-      "max_ctx_len": 8192,
       "intermediate_size": 28672,
       "vocab_size": 128256
     },
     "405B": {
       "num_attention_heads": 128,
       "num_key_value_heads": 8,
       "hidden_size": 16384,
-      "max_ctx_len": 8192,
       "intermediate_size": 53248,
       "vocab_size": 128256
     }
@@ -40,5 +37,6 @@
       "num_key_value_heads": 8,
       "vocab_size": 32000
     }
+
   }
 }

python/perf-kernels/rmsnorm.py

@@ -223,7 +223,7 @@ def model_benchmark_configs(args):
     batch_size = args.b if args.b else 1
 
     for model_name, config in configs.items():
-        seq_len = args.sl if args.sl else config["max_ctx_len"]
+        seq_len = args.sq if args.sq else 4096
         x_vals_list.append((model_name, batch_size * seq_len, config["hidden_size"]))
 
     return x_vals_list
@@ -309,14 +309,12 @@ def parse_args():
     parser.add_argument('-model_configs', type=str, default="model_configs.json", help="Model config json file.")
 
     available_models = get_available_models(model_families=["llama3"])  # Dynamically load model names
-    model_help = ("Model name to benchmark. Select from: [" + ", ".join(available_models) +
-                  "]. Use 'all' to benchmark all models or leave blank for the default benchmark script.")
+    model_help = (
+        "Model name to benchmark. Select from: [" + ", ".join(available_models) +
+        "]. Use 'all' to benchmark all models. Not providing runs the default benchmark script with custom configs.")
     parser.add_argument('-model', type=str, default=None, help=model_help)
-    parser.add_argument('-b', type=int, default=0,
-                        help="Batch size used together with model. Defaults to 1 if not provided.")
-    parser.add_argument(
-        '-sl', type=int, default=0,
-        help="Sequence length used together with model. Defaults to max_seq_len from model config if not provided.")
+    parser.add_argument('-b', type=int, default=0, help="Batch size used together with model.")
+    parser.add_argument('-sq', type=int, default=0, help="Sequence length used together with model.")
     parser.add_argument('-M', "--M_start", default="1", type=int)
     parser.add_argument('-Ms', "--M_step", default="2", type=int)  #This is multiplicative step
     parser.add_argument('-Me', "--M_end", default="512", type=int)

python/perf-kernels/softmax.py

@@ -142,7 +142,7 @@ def model_benchmark_configs(args):
     batch_size = args.b if args.b else 1
 
     for model_name, config in configs.items():
-        seq_len = args.sl if args.sl else config["max_ctx_len"]
+        seq_len = args.sq if args.sq else 4096
         x_vals_list.append((model_name, batch_size * seq_len, config["vocab_size"]))
 
     return x_vals_list
@@ -217,14 +217,12 @@ def parse_args():
     parser.add_argument('-model_configs', type=str, default="model_configs.json", help="Model config json file.")
 
     available_models = get_available_models(model_families=["llama3"])  # Dynamically load model names
-    model_help = ("Model name to benchmark. Select from: [" + ", ".join(available_models) +
-                  "]. Use 'all' to benchmark all models or leave blank for the default benchmark script.")
+    model_help = (
+        "Model name to benchmark. Select from: [" + ", ".join(available_models) +
+        "]. Use 'all' to benchmark all models. Not providing runs the default benchmark script with custom configs.")
     parser.add_argument('-model', type=str, default=None, help=model_help)
-    parser.add_argument('-b', type=int, default=0,
-                        help="Batch size used together with model. Defaults to 1 if not provided.")
-    parser.add_argument(
-        '-sl', type=int, default=0,
-        help="Sequence length used together with model. Defaults to max_seq_len from model config if not provided.")
+    parser.add_argument('-b', type=int, default=0, help="Batch size used together with model.")
+    parser.add_argument('-sq', type=int, default=0, help="Sequence length used together with model.")
     parser.add_argument('-M', "--M_start", default="1", type=int)
     parser.add_argument('-Ms', "--M_step", default="2", type=int)
     parser.add_argument('-Me', "--M_end", default="512", type=int)