minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add CLI args to conveniently support tuning more models (#12922)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Chongchong Tian
2026-03-12 23:10:55 -07:00
committed by GitHub
co-authored by gemini-code-assist[bot]
parent 63ecdcbb18
commit 70d4aabe42
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/deepep/tuning_deepep.py
+8 -4
View File
@@ -40,11 +40,11 @@ def test_main(
): ):
# Settings # Settings
num_tokens, hidden, num_topk_groups, num_topk, num_experts = ( num_tokens, hidden, num_topk_groups, num_topk, num_experts = (
4096, args.num_tokens,
7168, args.hidden,
min(num_nodes, 4), min(num_nodes, 4),
8, args.num_topk,
(256 // num_ranks) * num_ranks, (args.num_experts // num_ranks) * num_ranks,
) )
assert num_experts % num_ranks == 0 and num_local_ranks == 8 assert num_experts % num_ranks == 0 and num_local_ranks == 8
if local_rank == 0: if local_rank == 0:
@@ -462,6 +462,10 @@ def test_loop(local_rank: int, num_local_ranks: int, args):
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--num-sms", type=int, default=24) 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("--output-path", type=str, default="deepep_tuned.json")
parser.add_argument("--nnodes", type=int, default=1) parser.add_argument("--nnodes", type=int, default=1)
parser.add_argument("--node-rank", type=int, default=0) parser.add_argument("--node-rank", type=int, default=0)