diff --git a/benchmark/kernels/deepep/tuning_deepep.py b/benchmark/kernels/deepep/tuning_deepep.py
index db08a8f14..191819d2c 100644
--- a/benchmark/kernels/deepep/tuning_deepep.py
+++ b/benchmark/kernels/deepep/tuning_deepep.py
@@ -40,11 +40,11 @@ def test_main(
 ):
     # Settings
     num_tokens, hidden, num_topk_groups, num_topk, num_experts = (
-        4096,
-        7168,
+        args.num_tokens,
+        args.hidden,
         min(num_nodes, 4),
-        8,
-        (256 // num_ranks) * num_ranks,
+        args.num_topk,
+        (args.num_experts // num_ranks) * num_ranks,
     )
     assert num_experts % num_ranks == 0 and num_local_ranks == 8
     if local_rank == 0:
@@ -462,6 +462,10 @@ def test_loop(local_rank: int, num_local_ranks: int, args):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--num-sms", type=int, default=24)
+    parser.add_argument("--num-tokens", type=int, default=4096)
+    parser.add_argument("--hidden", type=int, default=7168)
+    parser.add_argument("--num-topk", type=int, default=8)
+    parser.add_argument("--num-experts", type=int, default=256)
     parser.add_argument("--output-path", type=str, default="deepep_tuned.json")
     parser.add_argument("--nnodes", type=int, default=1)
     parser.add_argument("--node-rank", type=int, default=0)