From 94d032c2f0dfd7dd95099b5f253ae86d757574eb Mon Sep 17 00:00:00 2001 From: Jiejing <113606+kzjeef@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:17:44 +0000 Subject: [PATCH] serve/pd_server: EPLB off by default, same --eplb-interval / --eplb-after-prefill flags Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK --- serve/pd_server.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/serve/pd_server.py b/serve/pd_server.py index d94672e..94bd4cb 100644 --- a/serve/pd_server.py +++ b/serve/pd_server.py @@ -71,6 +71,8 @@ def parse_args() -> argparse.Namespace: p.add_argument("--hicache-state-slots", type=int, default=4096) p.add_argument("--routing", default="affinity", choices=["affinity", "least-loaded"]) p.add_argument("--prefill-any-interval", type=int, default=2) + p.add_argument("--eplb-interval", type=int, default=0, help="see openai_server.py; 0 (default) disables EPLB") + p.add_argument("--eplb-after-prefill", type=int, default=0, choices=[0, 1], help="see openai_server.py") return p.parse_args() @@ -103,7 +105,10 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1" os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512") # identical state layout on PE and CE os.environ.setdefault("DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL", str(args.prefill_any_interval)) - setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=True) + setup_common_env(model_name=args.model, sampler_force_topk=1024, set_override_experts=True, enable_eplb=False) + if args.eplb_interval > 0: + os.environ["DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL"] = str(args.eplb_interval) + os.environ["DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE"] = str(args.eplb_after_prefill) engine_type = EngineType.PROMPT if args.role == "pe" else EngineType.COMPLETION llm = create_llm(