Add kvcache-centric profiling and admission controls

src/agentic_pd_hybrid/cli.py

@@ -7,6 +7,7 @@ from pathlib import Path
 from agentic_pd_hybrid.benchmark import BenchmarkConfig, run_live_benchmark
 from agentic_pd_hybrid.launcher import build_launch_plan
 from agentic_pd_hybrid.microbench import SmallAppendTraceConfig, write_small_append_trace
+from agentic_pd_hybrid.profile import ProfileConfig, print_profile_summary, write_profile
 from agentic_pd_hybrid.replay import ReplayConfig, replay_trace
 from agentic_pd_hybrid.sampling import SessionSampleConfig, sample_trace_sessions
 from agentic_pd_hybrid.trace_profiles import (
@@ -142,6 +143,71 @@ def main() -> None:
             "or query the decode worker on the critical path."
         ),
     )
+    replay.add_argument(
+        "--kvcache-seed-max-resident-tokens",
+        type=int,
+        default=None,
+        help=(
+            "For kvcache-centric routing, do not seed/reseed a decode session "
+            "when input+output tokens exceed this value."
+        ),
+    )
+    replay.add_argument(
+        "--kvcache-seed-max-output-tokens",
+        type=int,
+        default=None,
+        help=(
+            "For kvcache-centric routing, do not seed/reseed a decode session "
+            "when output tokens exceed this value."
+        ),
+    )
+    replay.add_argument(
+        "--kvcache-seed-min-turn-id",
+        type=int,
+        default=1,
+        help=(
+            "For kvcache-centric routing, do not seed/reseed a decode session "
+            "before this turn id."
+        ),
+    )
+    replay.add_argument(
+        "--kvcache-seed-only-multiturn-sessions",
+        action="store_true",
+        help=(
+            "Oracle ablation for kvcache-centric routing: only seed sessions "
+            "that have more than one turn in the replay trace."
+        ),
+    )
+    replay.add_argument(
+        "--kvcache-prefill-backup-policy",
+        choices=["release-after-transfer", "capacity-backup"],
+        default="release-after-transfer",
+        help=(
+            "For kvcache-centric seed/reseed, release the P-side session after "
+            "P->D transfer or keep a capacity-limited P-side backup."
+        ),
+    )
+    replay.add_argument(
+        "--kvcache-seed-max-inflight-decode",
+        type=int,
+        default=3,
+        help=(
+            "For kvcache-centric routing, skip seed/reseed when the router "
+            "shadow inflight decode load assigned to the target D exceeds this value. "
+            "Use a negative value to disable this filter."
+        ),
+    )
+    replay.add_argument(
+        "--kvcache-prefill-priority-eviction",
+        action="store_true",
+        help=(
+            "For kvcache-centric routing, mark P-side prefixes predicted to move "
+            "direct-to-D as lower priority. Requires P workers to use "
+            "--radix-eviction-policy priority."
+        ),
+    )
+    replay.add_argument("--kvcache-prefill-direct-priority", type=int, default=-100)
+    replay.add_argument("--kvcache-prefill-normal-priority", type=int, default=100)
 
     sample = subparsers.add_parser(
         "sample-sessions",
@@ -177,6 +243,17 @@ def main() -> None:
     normalize.add_argument("--output-length", type=int, default=1_000)
     normalize.add_argument("--max-requests", type=int, default=None)
 
+    profile = subparsers.add_parser(
+        "profile",
+        help="Profile a trace and optional request metrics for routing analysis",
+    )
+    profile.add_argument("--trace", type=Path, required=True)
+    profile.add_argument("--output", type=Path, default=None)
+    profile.add_argument("--metrics", type=Path, default=None)
+    profile.add_argument("--baseline-metrics", type=Path, default=None)
+    profile.add_argument("--candidate-metrics", type=Path, default=None)
+    profile.add_argument("--direct-max-uncached-tokens", type=int, default=2048)
+
     micro = subparsers.add_parser(
         "make-small-append-trace",
         help="Generate a synthetic multi-turn trace with small turn2+ appends",
@@ -223,6 +300,15 @@ def main() -> None:
     benchmark.add_argument("--time-scale", type=float, default=1.0)
     benchmark.add_argument("--concurrency-limit", type=int, default=32)
     benchmark.add_argument("--timeout-s", type=float, default=1200.0)
+    benchmark.add_argument(
+        "--request-timeout-s",
+        type=float,
+        default=None,
+        help=(
+            "Per-request replay/router timeout. If unset, --timeout-s is used. "
+            "--timeout-s still controls stack startup."
+        ),
+    )
     benchmark.add_argument(
         "--no-stream",
         action="store_true",
@@ -249,6 +335,70 @@ def main() -> None:
             "or query the decode worker on the critical path."
         ),
     )
+    benchmark.add_argument(
+        "--kvcache-seed-max-resident-tokens",
+        type=int,
+        default=None,
+        help=(
+            "For kvcache-centric routing, do not seed/reseed a decode session "
+            "when input+output tokens exceed this value."
+        ),
+    )
+    benchmark.add_argument(
+        "--kvcache-seed-max-output-tokens",
+        type=int,
+        default=None,
+        help=(
+            "For kvcache-centric routing, do not seed/reseed a decode session "
+            "when output tokens exceed this value."
+        ),
+    )
+    benchmark.add_argument(
+        "--kvcache-seed-min-turn-id",
+        type=int,
+        default=1,
+        help=(
+            "For kvcache-centric routing, do not seed/reseed a decode session "
+            "before this turn id."
+        ),
+    )
+    benchmark.add_argument(
+        "--kvcache-seed-only-multiturn-sessions",
+        action="store_true",
+        help=(
+            "Oracle ablation for kvcache-centric routing: only seed sessions "
+            "that have more than one turn in the replay trace."
+        ),
+    )
+    benchmark.add_argument(
+        "--kvcache-prefill-backup-policy",
+        choices=["release-after-transfer", "capacity-backup"],
+        default="release-after-transfer",
+        help=(
+            "For kvcache-centric seed/reseed, release the P-side session after "
+            "P->D transfer or keep a capacity-limited P-side backup."
+        ),
+    )
+    benchmark.add_argument(
+        "--kvcache-seed-max-inflight-decode",
+        type=int,
+        default=3,
+        help=(
+            "For kvcache-centric routing, skip seed/reseed when the router "
+            "shadow inflight decode load assigned to the target D exceeds this value. "
+            "Use a negative value to disable this filter."
+        ),
+    )
+    benchmark.add_argument(
+        "--kvcache-prefill-priority-eviction",
+        action="store_true",
+        help=(
+            "For kvcache-centric benchmark-live, launch P workers with priority "
+            "radix eviction and mark direct-to-D predicted prefixes lower priority."
+        ),
+    )
+    benchmark.add_argument("--kvcache-prefill-direct-priority", type=int, default=-100)
+    benchmark.add_argument("--kvcache-prefill-normal-priority", type=int, default=100)
     benchmark.add_argument(
         "--sample-profile",
         choices=["default", "small-append"],
@@ -294,6 +444,23 @@ def main() -> None:
             stream_idle_timeout_s=args.stream_idle_timeout_s,
             kvcache_direct_max_uncached_tokens=args.kvcache_direct_max_uncached_tokens,
             kvcache_admission_mode=args.kvcache_admission_mode,
+            kvcache_seed_max_resident_tokens=args.kvcache_seed_max_resident_tokens,
+            kvcache_seed_max_output_tokens=args.kvcache_seed_max_output_tokens,
+            kvcache_seed_min_turn_id=args.kvcache_seed_min_turn_id,
+            kvcache_seed_only_multiturn_sessions=(
+                args.kvcache_seed_only_multiturn_sessions
+            ),
+            kvcache_prefill_backup_policy=args.kvcache_prefill_backup_policy,
+            kvcache_seed_max_inflight_decode=(
+                None
+                if args.kvcache_seed_max_inflight_decode < 0
+                else args.kvcache_seed_max_inflight_decode
+            ),
+            kvcache_prefill_priority_eviction=(
+                args.kvcache_prefill_priority_eviction
+            ),
+            kvcache_prefill_direct_priority=args.kvcache_prefill_direct_priority,
+            kvcache_prefill_normal_priority=args.kvcache_prefill_normal_priority,
         )
         results = asyncio.run(replay_trace(config))
         print(
@@ -302,6 +469,26 @@ def main() -> None:
         )
         return
 
+    if args.command == "profile":
+        if (args.baseline_metrics is None) != (args.candidate_metrics is None):
+            raise ValueError(
+                "--baseline-metrics and --candidate-metrics must be provided together"
+            )
+        report = write_profile(
+            ProfileConfig(
+                trace_path=args.trace,
+                output_path=args.output,
+                metrics_path=args.metrics,
+                baseline_metrics_path=args.baseline_metrics,
+                candidate_metrics_path=args.candidate_metrics,
+                direct_max_uncached_tokens=args.direct_max_uncached_tokens,
+            )
+        )
+        print_profile_summary(report)
+        if args.output is not None:
+            print(f"wrote profile to {args.output}")
+        return
+
     if args.command == "sample-sessions":
         summary = sample_trace_sessions(
             SessionSampleConfig(
@@ -378,10 +565,32 @@ def main() -> None:
                 time_scale=args.time_scale,
                 concurrency_limit=args.concurrency_limit,
                 timeout_s=args.timeout_s,
+                request_timeout_s=args.request_timeout_s,
                 stream=not args.no_stream,
                 stream_idle_timeout_s=args.stream_idle_timeout_s,
                 kvcache_direct_max_uncached_tokens=args.kvcache_direct_max_uncached_tokens,
                 kvcache_admission_mode=args.kvcache_admission_mode,
+                kvcache_seed_max_resident_tokens=args.kvcache_seed_max_resident_tokens,
+                kvcache_seed_max_output_tokens=args.kvcache_seed_max_output_tokens,
+                kvcache_seed_min_turn_id=args.kvcache_seed_min_turn_id,
+                kvcache_seed_only_multiturn_sessions=(
+                    args.kvcache_seed_only_multiturn_sessions
+                ),
+                kvcache_prefill_backup_policy=args.kvcache_prefill_backup_policy,
+                kvcache_seed_max_inflight_decode=(
+                    None
+                    if args.kvcache_seed_max_inflight_decode < 0
+                    else args.kvcache_seed_max_inflight_decode
+                ),
+                kvcache_prefill_priority_eviction=(
+                    args.kvcache_prefill_priority_eviction
+                ),
+                kvcache_prefill_direct_priority=(
+                    args.kvcache_prefill_direct_priority
+                ),
+                kvcache_prefill_normal_priority=(
+                    args.kvcache_prefill_normal_priority
+                ),
                 sample_profile=args.sample_profile,
                 min_initial_input_tokens=args.min_initial_input_tokens,
                 max_initial_input_tokens=args.max_initial_input_tokens,