feat(connector): support tensor-parallel KV cache

benchmarks/bench_load.py

@@ -67,6 +67,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
     add_dataset_cli_args(parser)
     parser.add_argument("--max-samples", type=int, default=20)
     parser.add_argument("--block-size", type=int, default=BLOCK_TOKENS)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
     parser.add_argument("--max-context-tokens", type=int, default=0)
     parser.add_argument("--no-dedup-context", action="store_true")
     parser.add_argument("--max-inflight", type=int, default=32)
@@ -94,8 +95,13 @@ async def main_async(args: argparse.Namespace) -> None:
             f"--block-size={args.block_size} does not match manifest "
             f"block_size={manifest.block_size}"
         )
-    model = args.model if args.prepare_only else manifest.model
-    store_dir = args.store_dir if args.prepare_only else manifest.store_dir
+    if args.prepare_only:
+        model = args.model
+        store_dir = args.store_dir
+    else:
+        assert manifest is not None
+        model = manifest.model
+        store_dir = manifest.store_dir
     if model is None:
         raise ValueError("--model is required with --prepare-only")
     if store_dir is None:
@@ -141,9 +147,14 @@ async def main_async(args: argparse.Namespace) -> None:
     )
     token_counts = count_prompt_payload_tokens(tokenizer, prompts)
     total_blocks, max_prompt_blocks = workload_blocks(token_counts, args.block_size)
-    slot_size = slot_size_for_block_tokens(args.block_size)
     sizing = None
+    slot_size = None
     if args.prepare_only:
+        slot_size = slot_size_for_block_tokens(
+            model,
+            args.block_size,
+            args.tensor_parallel_size,
+        )
         capacity_limits = _capacity_limits(args, store_dir)
         sizing = derive_benchmark_sizing(
             total_blocks=total_blocks,
@@ -177,6 +188,7 @@ async def main_async(args: argparse.Namespace) -> None:
         block_size=args.block_size,
         evict=args.evict,
         sizing=sizing,
+        slot_size=slot_size,
     )
     if args.prepare_only:
         output = {"config": common_config}
@@ -392,6 +404,7 @@ def _common_config_for_run(
     evict: bool,
     block_size: int,
     sizing: BenchmarkSizing | None,
+    slot_size: int | None = None,
 ) -> dict[str, Any]:
     """Build benchmark config without re-inferring capacities during load.
 
@@ -409,6 +422,7 @@ def _common_config_for_run(
         block_size: vLLM KV block size in tokens.
         evict: Whether eviction sizing was requested.
         sizing: Prepare-time sizing, required for prepare-only invocations.
+        slot_size: Model-derived aggregate KV slot bytes for prepare invocations.
 
     Returns:
         JSON-serializable benchmark config.
@@ -436,6 +450,8 @@ def _common_config_for_run(
     if prepare_only:
         if sizing is None:
             raise ValueError("prepare-only config requires sizing")
+        if slot_size is None:
+            raise ValueError("prepare-only config requires slot size")
         return {
             "dataset": dataset,
             "num_samples": num_samples,
@@ -445,6 +461,7 @@ def _common_config_for_run(
             "total_blocks": total_blocks,
             "max_prompt_blocks": max_prompt_blocks,
             "block_size": block_size,
+            "slot_size_bytes": slot_size,
             "derived_l1_size_bytes": sizing.daser_l1_bytes,
             "derived_l1_size": format_capacity(sizing.daser_l1_bytes),
             "derived_l2_size_bytes": sizing.daser_l2_bytes,

benchmarks/bench_staging_restore.py

@@ -19,6 +19,7 @@
     sys.path.insert(0, str(ROOT))
 
 # First Party
+from benchmarks.utils.constants import BLOCK_TOKENS  # noqa: E402
 from daser.connector.staging import copy_staging_to_kv_cache  # noqa: E402
 from daser.ops.rope_apply import clear_rope_apply_cache  # noqa: E402
 
@@ -65,7 +66,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--device", default="cuda:0")
     parser.add_argument("--blocks", type=int, default=64)
     parser.add_argument("--layers", type=int, default=36)
-    parser.add_argument("--block-tokens", type=int, default=16)
+    parser.add_argument("--block-tokens", type=int, default=BLOCK_TOKENS)
     parser.add_argument("--heads", type=int, default=8)
     parser.add_argument("--head-dim", type=int, default=128)
     parser.add_argument("--rotary-dim", type=int, default=128)

benchmarks/bench_start_servers.py

@@ -11,6 +11,7 @@
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
 
+from benchmarks.utils.constants import BLOCK_TOKENS
 from benchmarks.utils.servers import ServerManager
 from benchmarks.utils.sizing import parse_size_bytes
 from benchmarks.utils.system import apply_gpu_selection
@@ -36,8 +37,10 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
     parser.add_argument("--gpu-util", type=float, default=0.85)
     parser.add_argument("--max-num-seqs", type=int, default=32)
     parser.add_argument("--max-num-batched-tokens", type=int, default=0)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
+    parser.add_argument("--trust-remote-code", action="store_true")
     parser.add_argument("--max-model-len", type=int, default=0)
-    parser.add_argument("--block-size", type=int, default=16)
+    parser.add_argument("--block-size", type=int, default=BLOCK_TOKENS)
     parser.add_argument("--l1-size", type=parse_size_bytes, default="256gib")
     parser.add_argument("--l2-size", type=parse_size_bytes, default="300gib")
     parser.add_argument(
@@ -82,6 +85,8 @@ async def main_async(args: argparse.Namespace) -> None:
         startup_timeout=args.startup_timeout,
         max_model_len=args.max_model_len if args.max_model_len > 0 else None,
         skip_l2=args.skip_l2,
+        tensor_parallel_size=args.tensor_parallel_size,
+        trust_remote_code=args.trust_remote_code,
     )
     manifest = await manager.start()
     print(f"manifest={args.store_dir}/manifest.json")

benchmarks/run_bench.py

@@ -18,6 +18,7 @@
 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
 
 from benchmarks.utils import vllm_bench
+from benchmarks.utils.constants import BLOCK_TOKENS
 from benchmarks.utils.datasets import add_dataset_cli_args
 from benchmarks.utils.servers import BenchmarkManifest, stop_from_pid_file
 
@@ -99,7 +100,9 @@ class RunBenchArgs:
     gpu_util: float = 0.85
     max_num_seqs: int = 32
     max_num_batched_tokens: int = 0
-    block_size: int = 16
+    tensor_parallel_size: int = 1
+    trust_remote_code: bool = False
+    block_size: int = BLOCK_TOKENS
     max_inflight: int = 32
     gen_max_tokens: int = 128
     max_context_tokens: int = 0
@@ -168,7 +171,9 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs:
     parser.add_argument("--gpu-util", type=float, default=0.85)
     parser.add_argument("--max-num-seqs", type=int, default=32)
     parser.add_argument("--max-num-batched-tokens", type=int, default=0)
-    parser.add_argument("--block-size", type=int, default=16)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument("--block-size", type=int, default=BLOCK_TOKENS)
     parser.add_argument("--max-inflight", type=int, default=32)
     parser.add_argument("--gen-max-tokens", type=int, default=128)
     parser.add_argument("--max-context-tokens", type=int, default=0)
@@ -206,6 +211,8 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs:
         gpu_util=args.gpu_util,
         max_num_seqs=args.max_num_seqs,
         max_num_batched_tokens=args.max_num_batched_tokens,
+        tensor_parallel_size=args.tensor_parallel_size,
+        trust_remote_code=args.trust_remote_code,
         block_size=args.block_size,
         max_inflight=args.max_inflight,
         gen_max_tokens=args.gen_max_tokens,
@@ -375,6 +382,7 @@ def _validate_run_args(args: RunBenchArgs) -> None:
         "max_num_seqs": args.max_num_seqs,
         "max_inflight": args.max_inflight,
         "gen_max_tokens": args.gen_max_tokens,
+        "tensor_parallel_size": args.tensor_parallel_size,
     }
     for name, value in positive_ints.items():
         if value <= 0:
@@ -481,6 +489,8 @@ def _prepare_command(
         str(args.max_inflight),
         "--block-size",
         str(args.block_size),
+        "--tensor-parallel-size",
+        str(args.tensor_parallel_size),
         "--gen-max-tokens",
         str(args.gen_max_tokens),
         "--max-context-tokens",
@@ -519,6 +529,8 @@ def _start_command(
         str(args.max_num_seqs),
         "--max-num-batched-tokens",
         str(args.max_num_batched_tokens),
+        "--tensor-parallel-size",
+        str(args.tensor_parallel_size),
         "--block-size",
         str(args.block_size),
         "--l1-size",
@@ -528,6 +540,8 @@ def _start_command(
     ]
     if backend_run.backend == "daser":
         command.extend(["--cache-reuse-mode", backend_run.reuse_mode])
+    if args.trust_remote_code:
+        command.append("--trust-remote-code")
     if not args.evict and backend_run.backend in ("daser", "lmcache"):
         command.append("--skip-l2")
     return command

benchmarks/utils/constants.py

@@ -3,12 +3,10 @@
 
 from __future__ import annotations
 
+from daser.config import model_geometry_from_path
+
 BYTES_PER_GIB: int = 1024**3
-BLOCK_TOKENS: int = 16
-NUM_KV_HEADS: int = 8
-HEAD_DIM: int = 128
-NUM_LAYERS: int = 36
-DTYPE_BYTES: int = 2
+BLOCK_TOKENS: int = 128
 
 COMPARISON_IOURING_MEM = "iouring-mem-vs-lmcache-local-ssd-mem"
 
@@ -18,19 +16,25 @@
 DEFAULT_IMDB_QUESTION: str = "Summarize the sentiment of this review."
 
 
-def slot_size_for_block_tokens(block_tokens: int) -> int:
-    """Return bytes required for one model KV block.
+def slot_size_for_block_tokens(
+    model_path: str,
+    block_tokens: int,
+    tensor_parallel_size: int = 1,
+) -> int:
+    """Return bytes required for one model KV block across all TP ranks.
 
     Args:
+        model_path: HuggingFace model directory containing ``config.json``.
         block_tokens: Number of tokens in a vLLM KV block.
+        tensor_parallel_size: Number of vLLM tensor-parallel ranks.
 
     Returns:
-        Slot size in bytes for the benchmark model geometry.
+        Aggregate slot size derived from the model geometry.
 
     Thread-safety:
-        Pure calculation over constants.
+        Reads the model config without mutating shared state.
     """
-    return NUM_KV_HEADS * HEAD_DIM * 2 * NUM_LAYERS * block_tokens * DTYPE_BYTES
-
-
-SLOT_SIZE: int = slot_size_for_block_tokens(BLOCK_TOKENS)
+    return model_geometry_from_path(model_path).slot_size_for_block_tokens(
+        block_tokens,
+        tensor_parallel_size,
+    )

benchmarks/utils/servers.py

@@ -103,7 +103,7 @@ def read(cls, path: str | Path) -> "BenchmarkManifest":
             for name, endpoint in payload["endpoints"].items()
         }
         payload["endpoints"] = endpoints
-        payload.setdefault("block_size", BLOCK_TOKENS)
+        payload.setdefault("block_size", 16)  # Legacy manifests predate this field.
         return cls(**payload)
 
 
@@ -130,6 +130,8 @@ def __init__(
         startup_timeout: float = 240.0,
         max_model_len: int | None = None,
         skip_l2: bool = False,
+        tensor_parallel_size: int = 1,
+        trust_remote_code: bool = False,
     ) -> None:
         """Initialize the service manager.
 
@@ -152,7 +154,11 @@ def __init__(
             startup_timeout: Health-check timeout.
             max_model_len: Optional vLLM max model length.
             skip_l2: Disable L2 persistence/adapters for L1-only no-evict runs.
+            tensor_parallel_size: vLLM tensor-parallel rank count.
+            trust_remote_code: allow model/tokenizer repository Python code.
         """
+        if tensor_parallel_size <= 0:
+            raise ValueError("tensor_parallel_size must be positive")
         self.run_id = run_id
         self.backend = backend
         self.model = model
@@ -171,6 +177,8 @@ def __init__(
         self.startup_timeout = startup_timeout
         self.max_model_len = max_model_len
         self.skip_l2 = skip_l2
+        self.tensor_parallel_size = tensor_parallel_size
+        self.trust_remote_code = trust_remote_code
         self.log_dir = self.store_dir / "logs"
         self.pid_file = self.store_dir / "pids.json"
         self.socket_path = self.store_dir / "daser.sock"
@@ -397,6 +405,10 @@ def _daser_server_command(self) -> list[str]:
             cmd.append("--skip-l2")
         else:
             cmd.extend(["--l2-size", str(self.l2_size_bytes)])
+        if self.tensor_parallel_size > 1:
+            cmd.extend(["--tensor-parallel-size", str(self.tensor_parallel_size)])
+        if self.trust_remote_code:
+            cmd.append("--trust-remote-code")
         return cmd
 
     async def _start_vllm(
@@ -447,6 +459,10 @@ def vllm_command(
             cmd.extend(["--max-model-len", str(self.max_model_len)])
         if self.max_num_batched_tokens is not None and self.max_num_batched_tokens > 0:
             cmd.extend(["--max-num-batched-tokens", str(self.max_num_batched_tokens)])
+        if self.tensor_parallel_size > 1:
+            cmd.extend(["--tensor-parallel-size", str(self.tensor_parallel_size)])
+        if self.trust_remote_code:
+            cmd.append("--trust-remote-code")
         if kv_transfer_config is not None:
             cmd.extend(["--kv-transfer-config", json.dumps(kv_transfer_config)])
         return cmd

benchmarks/utils/vllm_bench.py

@@ -182,7 +182,11 @@ def prepare_config(args: RunBenchArgs, run_root: Path) -> dict[str, Any]:
     prompt_tokens = _max_prompt_tokens(args)
     max_prompt_blocks = max(1, math.ceil(prompt_tokens / args.block_size))
     total_blocks = args.bench_num_prompts * max_prompt_blocks
-    slot_size = slot_size_for_block_tokens(args.block_size)
+    slot_size = slot_size_for_block_tokens(
+        args.model,
+        args.block_size,
+        args.tensor_parallel_size,
+    )
     sizing = derive_benchmark_sizing(
         total_blocks=total_blocks,
         max_prompt_blocks=max_prompt_blocks,
@@ -210,6 +214,7 @@ def prepare_config(args: RunBenchArgs, run_root: Path) -> dict[str, Any]:
         "max_prompt_blocks": max_prompt_blocks,
         "max_prompt_tokens": prompt_tokens,
         "block_size": args.block_size,
+        "slot_size_bytes": slot_size,
         "bench_num_prompts": args.bench_num_prompts,
         "bench_input_len": args.bench_input_len,
         "bench_suffix_input_len": bench_suffix_input_len(args),
@@ -434,6 +439,7 @@ def _bench_command(
         "/v1/completions",
         "--model",
         args.model,
+        *(["--trust-remote-code"] if args.trust_remote_code else []),
         "--dataset-name",
         "random",
         "--num-prompts",