diff --git a/serve/README.md b/serve/README.md index 9f204eb..1f30b79 100644 --- a/serve/README.md +++ b/serve/README.md @@ -30,6 +30,10 @@ Knobs: `--hicache-size-gb`, `--hicache-state-slots`, `--routing affinity|least-l conversation's turns on one rank, keyed by the first user message; required for hits since the cache is per rank). The cache requires decoder bounded replay (`--yoco-keep-tokens > 0`, the default 140). +EPLB expert rearrangement is off by default (`--eplb-interval 0`). Every rearrangement stalls all ranks for about +0.3 s, and with `--eplb-after-prefill 1` it fires after nearly every step under multi-turn traffic; enable it with +`--eplb-interval 30` only for decode-heavy workloads whose expert load is skewed. + ## How a hit works A new prompt's page hashes are matched from the start; the match is capped so that at least `keep_tokens` diff --git a/serve/openai_server.py b/serve/openai_server.py index 1e41bf2..e117e1b 100644 --- a/serve/openai_server.py +++ b/serve/openai_server.py @@ -67,6 +67,22 @@ def parse_args() -> argparse.Namespace: help="DEEPINFER_PREFILL_ANY_ATTEMPT_INTERVAL: a rank with a waiting prompt gets a prefill step at most this " "many decode steps later even if other ranks are idle (engine default 100 => ~1.4 s TTFT floor)", ) + p.add_argument( + "--eplb-interval", + type=int, + default=0, + help="DEEPINFER_EXPERT_REARRANGEMENT_STEP_INTERVAL: rearrange experts every N steps; 0 (default) disables EPLB. " + "Each rearrangement stalls every rank (~0.3 s on B200/GB300), which under online traffic costs far more than " + "the balance it buys", + ) + p.add_argument( + "--eplb-after-prefill", + type=int, + default=0, + choices=[0, 1], + help="DEEPINFER_REARRANGE_AFTER_PREFILL_IN_NORMAL_ENGINE: also rearrange after every prefill step (only with " + "--eplb-interval > 0); with agentic traffic nearly every step batch contains a prefill", + ) return p.parse_args() @@ -92,7 +108,10 @@ def worker_main(args: argparse.Namespace, dp_rank: int, in_q: mp.Queue, out_q: m # 128-boundaries behind the end of a prompt (see HostPrefixCache.state_lookback_blocks). os.environ.setdefault("DEEPINFER_PREFILL_STORE_EXTRA_SIZE", "512") 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) llm = create_llm( args.model,