diff --git a/scripts/sweep_e4_pressured.sh b/scripts/sweep_e4_pressured.sh index 71ade30..dff4846 100755 --- a/scripts/sweep_e4_pressured.sh +++ b/scripts/sweep_e4_pressured.sh @@ -6,8 +6,10 @@ # --kvcache-migration-reject-threshold 1 (was 3) # After ONE rejection the policy migrates the session to a different # D, which in turn triggers _invoke_kvcache_seeded_router → D→P sync. -# (rely solely on reject_threshold=1 for now; mem-fraction reduction -# would need extra_server_args plumbing which is out-of-scope here) +# --decode-mem-fraction-static 0.4 +# Plumbed through cli.py → topology.decode_extra_server_args → +# launcher. Shrinks per-decode KV pool, forcing admit_direct_append +# to reject more often. # # Hypotheses (same as docs/E4_PROTOCOL_ZH.md but in a stressed regime): # H1' E4-pressured TTFT p99 ≤ E1 TTFT p99 @@ -75,6 +77,7 @@ uv run --no-sync python -m agentic_pd_hybrid.cli benchmark-live \ --kvcache-migration-reject-threshold "$REJECT_THRESHOLD" \ --kvcache-direct-max-uncached-tokens 8192 \ --kvcache-load-floor-bonus "$LOAD_FLOOR_BONUS" \ + --decode-mem-fraction-static "${DECODE_MEM_FRAC:-0.4}" \ --enable-d-to-p-sync 2>&1 | tee -a "$LOG" run_dir=$(ls -td "$OUTPUT"/kvcache-centric-*/ 2>/dev/null | head -1) diff --git a/src/agentic_pd_hybrid/cli.py b/src/agentic_pd_hybrid/cli.py index e4f1663..4788266 100644 --- a/src/agentic_pd_hybrid/cli.py +++ b/src/agentic_pd_hybrid/cli.py @@ -566,6 +566,23 @@ def main() -> None: "See docs/D_TO_P_SYNC_DESIGN_ZH.md." ), ) + benchmark.add_argument( + "--decode-mem-fraction-static", + type=float, + default=None, + help=( + "Override SGLang's --mem-fraction-static on decode workers. " + "Smaller value → smaller KV pool → admit_direct_append rejects " + "more often → reseed path fires more often. Pressure tool for " + "E4-style D→P sync experiments." + ), + ) + benchmark.add_argument( + "--prefill-mem-fraction-static", + type=float, + default=None, + help="Override --mem-fraction-static on prefill workers.", + ) benchmark.add_argument( "--sample-profile", choices=["default", "small-append"], @@ -897,11 +914,25 @@ def _topology_from_args(args: argparse.Namespace): trust_remote_code=not args.no_trust_remote_code, ib_device=args.ib_device, enable_d_to_p_sync=getattr(args, "enable_d_to_p_sync", False), - prefill_extra_server_args=("--disable-overlap-schedule",), - decode_extra_server_args=("--disable-overlap-schedule",), - direct_extra_server_args=("--enable-streaming-session",), + prefill_extra_server_args=_build_extra_server_args(args, "prefill"), + decode_extra_server_args=_build_extra_server_args(args, "decode"), + direct_extra_server_args=_build_extra_server_args(args, "direct"), ) +def _build_extra_server_args(args, role: str) -> tuple[str, ...]: + base: tuple[str, ...] + if role == "direct": + base = ("--enable-streaming-session",) + else: + base = ("--disable-overlap-schedule",) + mem_frac = getattr(args, "decode_mem_fraction_static", None) if role == "decode" else None + if mem_frac is None and role == "prefill": + mem_frac = getattr(args, "prefill_mem_fraction_static", None) + if mem_frac is not None and mem_frac > 0: + base = base + ("--mem-fraction-static", f"{mem_frac:.3f}") + return base + + if __name__ == "__main__": main()