From f6d6dc01eaebf63a1d864ce1f6ea1ac1c33dff29 Mon Sep 17 00:00:00 2001
From: Claude Code Agent <VtimothyeveeE@europe.com>
Date: Wed, 13 May 2026 10:43:26 +0800
Subject: [PATCH] feat(cli): per-role --mem-fraction-static + use in
 E4-pressured
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

E4-v1 / v2 / pressured-v1 all failed to fire admission rejections in
this workload because the default 0.6 mem-fraction-static gives
288K-token kv_pool per decoder, more than enough to absorb the
50-session trace even at concurrency=32.

This commit adds:
  --decode-mem-fraction-static  (overrides per-decode SGLang arg)
  --prefill-mem-fraction-static (symmetric for completeness)

Plumbed via topology.{decode,prefill}_extra_server_args. The
pressured sweep now uses --decode-mem-fraction-static 0.4 which
shrinks decoder kv_pool to ~192K tokens — should force enough
admission rejections to actually exercise the D→P snapshot path.
---
 scripts/sweep_e4_pressured.sh |  7 +++++--
 src/agentic_pd_hybrid/cli.py  | 37 ++++++++++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 5 deletions(-)

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