[PD] Bound cached-prefix DCP transfers by pack capacity (#40376)

benchmark/disaggregation/bench_cached_prefix.py

@@ -0,0 +1,113 @@
+"""Measure PD TTFT with a shared cached prefix and unique request suffixes.
+
+Start a PD router and workers before running this client. Both worker caches are
+flushed after warmup; the shared prefix is then warmed on prefill. Decode radix
+caching should be disabled to exercise transfer of the entire missing KV range.
+Saves per-request responses and client-observed TTFT to the requested JSON file.
+"""
+
+import argparse
+import asyncio
+import json
+import random
+import time
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+
+
+async def main(args):
+    rng = random.Random(35762)
+    prefix = [rng.randrange(1000, 30000) for _ in range(args.prefix)]
+    prompts = [
+        prefix + [rng.randrange(1000, 30000) for _ in range(args.unique)]
+        for _ in range(args.requests)
+    ]
+    timeout = aiohttp.ClientTimeout(total=1800)
+    async with aiohttp.ClientSession(timeout=timeout) as session:
+
+        async def generate(ids, stream=True):
+            start = time.perf_counter()
+            first = None
+            result = None
+            body = {
+                "input_ids": ids,
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": args.output,
+                    "ignore_eos": True,
+                },
+                "stream": stream,
+            }
+            async with session.post(args.url + "/generate", json=body) as response:
+                response.raise_for_status()
+                if stream:
+                    async for line in response.content:
+                        if not line.startswith(b"data:"):
+                            continue
+                        payload = line[5:].strip()
+                        if payload == b"[DONE]":
+                            continue
+                        result = json.loads(payload)
+                        if first is None:
+                            first = time.perf_counter() - start
+                else:
+                    result = await response.json()
+            elapsed = time.perf_counter() - start
+            if result is None or "error" in result:
+                raise RuntimeError(result)
+            return {"ttft_s": first, "elapsed_s": elapsed, "response": result}
+
+        # Warm both cold-prefill and cached-prefix/concurrent shapes. Clear both
+        # sides afterwards so measured requests never reuse a unique suffix.
+        await generate(prompts[0])
+        await asyncio.gather(*(generate(p) for p in prompts[1 : args.concurrency + 1]))
+        for url in args.workers:
+            async with session.post(url + "/flush_cache", params={"timeout": 30}) as r:
+                r.raise_for_status()
+        if prefix:
+            await generate(prefix)
+        sem = asyncio.Semaphore(args.concurrency)
+
+        async def run(ids):
+            async with sem:
+                return await generate(ids)
+
+        start = time.perf_counter()
+        records = await asyncio.gather(*(run(ids) for ids in prompts))
+        elapsed = time.perf_counter() - start
+    ttfts = [r["ttft_s"] for r in records]
+    summary = {
+        "args": vars(args),
+        "elapsed_s": elapsed,
+        "requests_per_s": len(records) / elapsed,
+        "ttft_mean_ms": float(np.mean(ttfts) * 1000),
+        "ttft_p50_ms": float(np.percentile(ttfts, 50) * 1000),
+        "ttft_p99_ms": float(np.percentile(ttfts, 99) * 1000),
+        "records": records,
+    }
+    Path(args.result).write_text(json.dumps(summary, indent=2))
+    print(json.dumps({k: v for k, v in summary.items() if k != "records"}), flush=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", default="http://127.0.0.1:30000")
+    parser.add_argument(
+        "--workers",
+        nargs="+",
+        default=["http://127.0.0.1:30001", "http://127.0.0.1:30002"],
+    )
+    parser.add_argument("--prefix", type=int, default=83264)
+    parser.add_argument("--unique", type=int, default=6720)
+    parser.add_argument("--output", type=int, default=1)
+    parser.add_argument("--requests", type=int, default=16)
+    parser.add_argument("--concurrency", type=int, default=1)
+    parser.add_argument("--result", required=True)
+    args = parser.parse_args()
+    if args.prefix < 0 or args.unique <= 0:
+        parser.error("--prefix must be nonnegative and --unique must be positive")
+    if min(args.output, args.requests, args.concurrency) <= 0:
+        parser.error("--output, --requests and --concurrency must be positive")
+    asyncio.run(main(args))