serve/pd_server: EPLB off by default, same --eplb-interval / --eplb-after-prefill flags

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK
This commit is contained in:
Jiejing 2026-10-22 17:17:44 +00:00
parent 4159328eb5
commit 94d032c2f0

View File

@ -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(