add memory placement(--mp) option

Author: davidz-ampere

benchmarks/utils/benchmark.py

@@ -35,6 +35,9 @@ def parse_args():
     parser.add_argument("-t", "--num_threads",
                         type=int, default=1,
                         help="number of threads to use per process")
+    parser.add_argument("--mp",
+                        type=str, default="local",
+                        help="memory placement policy, 'local','interleave' or 'none'")
     return parser.parse_args()
 
 
@@ -115,21 +118,40 @@ def main():
     logs_dir = os.path.join("/tmp", str(uuid.uuid4()))
     os.mkdir(logs_dir)
     current_subprocesses = list()
+    if args.mp == "local":
+        mem_place = "--localalloc"
+    elif args.mp == "interleave":
+        mem_place = "--interleave=all"
+    else:
+        mem_place = "none"
+
     for n in range(args.num_processes):
         logfile = f"{logs_dir}/log_{n}"
         if os.path.exists("/llm/batched-bench"):
             # command-line for v1
-            cmd = ["numactl", f"--physcpubind={gen_threads_config(args.num_threads, n)}", "--localalloc",
-                   "/llm/batched-bench", args.model, str(args.kv_cache), "2048", "512", "0", "0", "0", str(args.prompt_size), str(TOKENS),
-                   str(args.batch_size), str(args.num_threads)]
+            if mem_place == "none":
+                cmd = ["numactl", f"--physcpubind={gen_threads_config(args.num_threads, n)}",
+                       "/llm/batched-bench", args.model, str(args.kv_cache), "2048", "512", "0", "0", "0", str(args.prompt_size), str(TOKENS),
+                       str(args.batch_size), str(args.num_threads)]
+            else:
+                 cmd = ["numactl", f"--physcpubind={gen_threads_config(args.num_threads, n)}", str(mem_place),
+                       "/llm/batched-bench", args.model, str(args.kv_cache), "2048", "512", "0", "0", "0", str(args.prompt_size), str(TOKENS),
+                       str(args.batch_size), str(args.num_threads)]
         elif os.path.exists("/llm/llama-batched-bench"):
             # command-line for v2
-            cmd = ["numactl", f"--physcpubind={gen_threads_config(args.num_threads, n)}", "--localalloc",
-                   "/llm/llama-batched-bench", "-m", args.model, "-c", str(args.kv_cache), "-b", "2048", "-ub", "512", "-npp", str(args.prompt_size), "-ntg", str(TOKENS),
-                   "-npl", str(args.batch_size), "-t", str(args.num_threads), "-tb", str(args.num_threads), "-td", str(args.num_threads)]
+            if mem_place == "none":
+                cmd = ["numactl", f"--physcpubind={gen_threads_config(args.num_threads, n)}",
+                       "/llm/llama-batched-bench", "-m", args.model, "-c", str(args.kv_cache), "-b", "2048", "-ub", "512", "-npp", str(args.prompt_size), "-ntg", str(TOKENS),
+                       "-npl", str(args.batch_size), "-t", str(args.num_threads), "-tb", str(args.num_threads), "-td", str(args.num_threads)]
+            else:
+                cmd = ["numactl", f"--physcpubind={gen_threads_config(args.num_threads, n)}",str(mem_place),
+                       "/llm/llama-batched-bench", "-m", args.model, "-c", str(args.kv_cache), "-b", "2048", "-ub", "512", "-npp", str(args.prompt_size), "-ntg", str(TOKENS),
+                       "-npl", str(args.batch_size), "-t", str(args.num_threads), "-tb", str(args.num_threads), "-td", str(args.num_threads)]
+
         else:
             print("FAIL: batched-bench not found!")
             sys.exit(1)
+
         current_subprocesses.append(
             subprocess.Popen(cmd, stdout=open(logfile, 'wb'), stderr=open(logfile, 'wb')))
     start = time.time()