MB5 PD reuse-centric ablation: tooling, data, Fig 1-3

Three-axis controlled ablation of PD-colo vs PD-disagg on synthetic regular
traces (closed-loop, controlled reuse via REPLAY_NO_REALIZED_PREFIX) on the
clean stack (e13391e gated off).

  Axis 1 (Fig 1) -- reuse 6%->94% at N=8, in8192/out256
  Axis 2 (Fig 2) -- shape in2048/out2048 -> in32768/out64 at N=8, reuse~70%
  Axis 3 (Fig 3) -- concurrency N=8/16/32/64 at reuse~71%, in8192/out256

Findings:
  * APC parity colo=PD at every reuse (5.5/22/44/66/77/82%) -- contamination
    fix validated.
  * PD edge erodes 1.57x->1.10x with reuse; prefill GPUs strand 26%->9%.
  * Shape: PD-best peaks mid-sweep (1.34x at in8192/out512); wrong PD ratio
    catastrophic at prefill extreme (in32768/out64 pd2 = 378/400, p99 432s).
  * Concurrency: PD wins N<=32 (1.23-1.29x), TIPS at N=64 -- pd2/pd4
    crater (APC 71%->1.4%, TPS -30%) while colo scales cleanly.

Infrastructure:
  * replayer: --max-inflight-sessions, --inter-turn-think, --no-realized-prefix
    (env-defaulted via REPLAY_MAX_INFLIGHT, REPLAY_INTER_TURN_THINK_S,
    REPLAY_NO_REALIZED_PREFIX).
  * mb5_run.sh: writes bench_config.json + gpu_util.csv + run_window.json +
    instance_apc.txt + metrics.jsonl for bench_report/fig_agg ingest.
  * fig_agg.py: per-arm GPU role split + producer-side APC; --json mode.
  * gpu_util_report.py: companion per-GPU util report from gpu_util.csv.
  * partial_summary.py: stats from in-flight replay_metrics.jsonl
    (works before metrics.summary.json exists).

Data: analysis/mb5_pd_ablation/fig{1,2,3}.json (24 + 20 + 16 rows).
Figures: figs/mb5_pd_ablation/fig{1_reuse,2_shape,3_concurrency}_axis.png.

replayer/__main__.py

@@ -30,12 +30,23 @@ def main() -> None:
                    default=float(_env_think) if _env_think else None,
                    help="Closed-loop think-time (s) after each turn completes; "
                         "ignore absolute trace schedule. Env: REPLAY_INTER_TURN_THINK_S")
+    p.add_argument("--no-realized-prefix",
+                   action="store_true",
+                   default=bool(os.environ.get("REPLAY_NO_REALIZED_PREFIX")),
+                   help="Controlled-reuse mode: prompt = hash-built tokens only "
+                        "(reuse set by hash_ids). Env: REPLAY_NO_REALIZED_PREFIX")
     p.add_argument("--dispatch-mode", choices=["tracets", "thinktime"],
                    default=os.environ.get("REPLAY_DISPATCH_MODE", "tracets"),
                    help="tracets (Mode 1): absolute trace ts = max(prev_finished, ts). "
                         "thinktime (Mode 2): turn-k at prev_finished + "
                         "time_to_parent_chat. Env: REPLAY_DISPATCH_MODE")
     p.add_argument("--request-timeout", type=float, default=600.0)
+    _env_maxdur = os.environ.get("REPLAY_MAX_DURATION")
+    p.add_argument("--max-duration", type=float,
+                   default=float(_env_maxdur) if _env_maxdur else None,
+                   help="Overall wall-clock deadline (s): cancel in-flight + write "
+                        "summary (un-run turns counted as failures) to bound a "
+                        "collapsed config's drain. Env: REPLAY_MAX_DURATION")
     p.add_argument("--request-limit", type=int, default=None,
                    help="Limit number of requests to replay")
     p.add_argument("-v", "--verbose", action="store_true")
@@ -56,7 +67,9 @@ def main() -> None:
         request_limit=args.request_limit,
         max_inflight_sessions=args.max_inflight_sessions,
         inter_turn_think_s=args.inter_turn_think,
+        no_realized_prefix=args.no_realized_prefix,
         dispatch_mode=args.dispatch_mode,
+        max_duration_s=args.max_duration,
     )
 
     results = asyncio.run(replay_trace(config))

replayer/replay.py

@@ -66,6 +66,13 @@ class ReplayConfig:
     # max_inflight_sessions=N this is a stable N-user closed-loop (no open-loop
     # runaway), so it removes the "immediate retrigger under load" artifact.
     inter_turn_think_s: float | None = None
+    # Controlled-reuse mode: skip _apply_realized_prefix so each turn's prompt is
+    # exactly the hash-built tokens. Then prefix-cache reuse is governed solely by
+    # the generated hash_ids (shared prefix blocks hit, fresh delta blocks miss) —
+    # required for the reuse-fraction sweep, where realized-prefix would otherwise
+    # force every fixed-length turn to ≈ the prior turn (≈100% reuse regardless).
+    # Keep OFF (realized-prefix ON) for the real agentic trace.
+    no_realized_prefix: bool = False
     # Dispatch timing for intra-session turns:
     #  "tracets"   (Mode 1): fire at absolute trace timestamp -> effectively
     #              max(prev_finished, trace_ts); collapses think-time to 0 when
@@ -73,6 +80,25 @@ class ReplayConfig:
     #  "thinktime" (Mode 2): turn-1 at trace arrival; turn-k at
     #              prev_finished + time_to_parent_chat (real production gap).
     dispatch_mode: str = "tracets"
+    # Overall wall-clock deadline for the whole replay (seconds). When exceeded,
+    # stop awaiting in-flight sessions, cancel them, and write the summary over
+    # whatever completed — un-run turns are counted as failures so completion%
+    # stays honest (request_count == full trace). None = no deadline (default,
+    # original behavior unchanged). Used to bound the slow drain of a collapsed
+    # config in a sweep. Env: REPLAY_MAX_DURATION.
+    max_duration_s: float | None = None
+
+
+def _skipped_metric() -> "RequestMetrics":
+    """Placeholder failure row for a turn never run due to a max_duration cutoff.
+    Only its error (non-None) matters: it counts toward request/error totals but
+    is excluded from latency/ttft/tpot percentiles (successes only)."""
+    return RequestMetrics(
+        request_id="deadline_skipped", session_id="", turn_id=-1,
+        trace_timestamp_s=0.0, input_length=0, output_length=0,
+        request_type="skipped", effective_input_length=None, cached_tokens=0,
+        latency_s=None, ttft_s=None, tpot_s=None, error="deadline_skipped",
+    )
 
 
 def _build_prompt_token_ids(req: TraceRequest) -> list[int]:
@@ -318,10 +344,9 @@ async def _run_session(
                 if elapsed < target_wall:
                     await asyncio.sleep(target_wall - elapsed)
 
-            token_ids = _apply_realized_prefix(
-                _build_prompt_token_ids(req),
-                realized_context,
-            )
+            token_ids = _build_prompt_token_ids(req)
+            if not config.no_realized_prefix:
+                token_ids = _apply_realized_prefix(token_ids, realized_context)
             result = await _dispatch_request(
                 client=client, config=config, req=req,
                 prompt_token_ids=token_ids, sem=request_sem,
@@ -410,25 +435,44 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]:
             trust_env=False,
             limits=limits,
         ) as client:
+            states = [_SessionState(session_id=sid, turns=turns)
+                      for sid, turns in sessions]
             tasks = [
                 asyncio.create_task(_run_session(
-                    state=_SessionState(session_id=sid, turns=turns),
-                    config=config, client=client,
+                    state=st, config=config, client=client,
                     request_sem=request_sem,
                     earliest_ts=earliest_ts, sweep_start=sweep_start,
                     sink=sink,
                     session_sem=session_sem,
                 ))
-                for sid, turns in sessions
+                for st in states
             ]
-            all_results = await asyncio.gather(*tasks)
+            if config.max_duration_s and config.max_duration_s > 0:
+                _done, pending = await asyncio.wait(
+                    tasks, timeout=config.max_duration_s)
+                if pending:
+                    logger.warning(
+                        "max_duration %.0fs reached: cancelling %d in-flight "
+                        "session(s); un-run turns counted as failures",
+                        config.max_duration_s, len(pending))
+                    for t in pending:
+                        t.cancel()
+                    await asyncio.gather(*pending, return_exceptions=True)
+            else:
+                await asyncio.gather(*tasks)
     finally:
         sink.close()
 
     sweep_elapsed = time.perf_counter() - sweep_start
     post_metrics = await _snapshot_prefix_cache_metrics(config.endpoint_url)
 
-    flat = [m for group in all_results for m in group]
+    # Build from the session states (identical to the gather return in the
+    # uncapped path) so partially-completed (cancelled) sessions still contribute
+    # their finished turns; pad un-run turns as failures so request_count == trace.
+    flat = [m for st in states for m in st.metrics]
+    missing = n_requests - len(flat)
+    if missing > 0:
+        flat.extend(_skipped_metric() for _ in range(missing))
     summary_path = config.output_path.with_suffix(".summary.json")
     write_summary_json(summary_path, flat)