From 70d4aabe42b8b7c058fb83c148141904c4337753 Mon Sep 17 00:00:00 2001 From: Chongchong Tian <58495634+tianchongchong@users.noreply.github.com> Date: Fri, 13 Mar 2026 14:10:55 +0800 Subject: [PATCH] Add CLI args to conveniently support tuning more models (#12922) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- benchmark/kernels/deepep/tuning_deepep.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) 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)