diff --git a/benchmarks/bench_load.py b/benchmarks/bench_load.py index 066d386..3b3d1ed 100644 --- a/benchmarks/bench_load.py +++ b/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, diff --git a/benchmarks/bench_staging_restore.py b/benchmarks/bench_staging_restore.py index cefc521..990e118 100644 --- a/benchmarks/bench_staging_restore.py +++ b/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) diff --git a/benchmarks/bench_start_servers.py b/benchmarks/bench_start_servers.py index 7033425..9f5539c 100644 --- a/benchmarks/bench_start_servers.py +++ b/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") diff --git a/benchmarks/run_bench.py b/benchmarks/run_bench.py index f7d5eff..0828aea 100644 --- a/benchmarks/run_bench.py +++ b/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 diff --git a/benchmarks/utils/constants.py b/benchmarks/utils/constants.py index 7b117c3..a8a8a2c 100644 --- a/benchmarks/utils/constants.py +++ b/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, + ) diff --git a/benchmarks/utils/servers.py b/benchmarks/utils/servers.py index 6e7df2f..d30af4f 100644 --- a/benchmarks/utils/servers.py +++ b/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 diff --git a/benchmarks/utils/vllm_bench.py b/benchmarks/utils/vllm_bench.py index 44b460f..73f4005 100644 --- a/benchmarks/utils/vllm_bench.py +++ b/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", diff --git a/daser/config.py b/daser/config.py index 5ff2331..bb6f9aa 100644 --- a/daser/config.py +++ b/daser/config.py @@ -31,17 +31,35 @@ class ModelGeometry: num_layers: int dtype_bytes: int - def slot_size_for_block_tokens(self, block_tokens: int) -> int: - """Return bytes required for one vLLM KV block across all layers. + def local_slot_size_for_block_tokens( + self, block_tokens: int, tensor_parallel_size: int = 1 + ) -> int: + """Return bytes stored by one TP rank for one vLLM KV block. Args: block_tokens: Number of tokens in each vLLM KV block. + tensor_parallel_size: Number of vLLM tensor-parallel ranks. Returns: - Bytes required for one KV slot with ``block_tokens`` tokens. + Bytes required by one rank for one KV slot. + + Raises: + ValueError: if ``tensor_parallel_size`` is not positive. """ + if tensor_parallel_size <= 0: + raise ValueError("tensor_parallel_size must be positive") + if self.num_kv_heads >= tensor_parallel_size: + compatible = self.num_kv_heads % tensor_parallel_size == 0 + else: + compatible = tensor_parallel_size % self.num_kv_heads == 0 + if not compatible: + raise ValueError( + "tensor_parallel_size must divide num_kv_heads or be divisible " + "by num_kv_heads" + ) + local_kv_heads = max(1, self.num_kv_heads // tensor_parallel_size) return ( - self.num_kv_heads + local_kv_heads * self.head_dim * 2 # K and V * self.num_layers @@ -49,6 +67,23 @@ def slot_size_for_block_tokens(self, block_tokens: int) -> int: * self.dtype_bytes ) + def slot_size_for_block_tokens( + self, block_tokens: int, tensor_parallel_size: int = 1 + ) -> int: + """Return bytes required for one vLLM KV block across all layers. + + Args: + block_tokens: Number of tokens in each vLLM KV block. + tensor_parallel_size: Number of vLLM tensor-parallel ranks. + + Returns: + Aggregate bytes required for one KV slot across all TP ranks. + """ + return ( + self.local_slot_size_for_block_tokens(block_tokens, tensor_parallel_size) + * tensor_parallel_size + ) + def _dtype_bytes(dtype: object) -> int: """Return storage bytes for a HuggingFace dtype string. @@ -60,6 +95,8 @@ def _dtype_bytes(dtype: object) -> int: Number of bytes per element. """ dtype_str = str(dtype or "").lower() + if "float8" in dtype_str or "fp8" in dtype_str: + raise ValueError("FP8 KV cache geometry is not supported") if "float32" in dtype_str or "fp32" in dtype_str: return 4 return 2 @@ -88,7 +125,12 @@ def model_geometry_from_path(model_path: str) -> ModelGeometry: or payload.get("multi_query_group_num") or num_attention_heads ) - num_layers = int(payload.get("num_hidden_layers") or payload.get("n_layer") or 0) + num_layers = int( + payload.get("num_hidden_layers") + or payload.get("num_layers") + or payload.get("n_layer") + or 0 + ) head_dim = int(payload.get("head_dim") or payload.get("kv_channels") or 0) if head_dim == 0 and hidden_size > 0 and num_attention_heads > 0: head_dim = hidden_size // num_attention_heads @@ -139,6 +181,8 @@ class DaserConfig: l1_size_bytes: memory-tier capacity for iouring mode. skip_l2: when True, keep all transfer bytes in volatile L1 memory and do not allocate or persist SSD-tier store files. + tensor_parallel_size: vLLM tensor-parallel rank count used by the + physical KV store layout. """ model_path: str = "" @@ -153,6 +197,7 @@ class DaserConfig: transfer_mode: str = "iouring" l1_size_bytes: int = DEFAULT_IOURING_L1_BYTES skip_l2: bool = False + tensor_parallel_size: int = 1 @property def store_path(self) -> str: @@ -167,7 +212,10 @@ def index_path(self) -> str: @property def model_id(self) -> str: """Model identifier used for cache isolation.""" - return self.vllm_model_id or self.model_path + model_id = self.vllm_model_id or self.model_path + if self.tensor_parallel_size == 1: + return model_id + return f"{model_id}::tp{self.tensor_parallel_size}" @property def total_slots(self) -> int: @@ -191,9 +239,19 @@ def resolved_slot_size(self) -> int: Slot size in bytes. """ return model_geometry_from_path(self.model_path).slot_size_for_block_tokens( - self.block_tokens + self.block_tokens, self.tensor_parallel_size ) + def resolved_local_slot_size(self) -> int: + """Return bytes stored by one TP rank for one logical KV slot. + + Returns: + Per-rank KV slot bytes derived from model and TP configuration. + """ + return model_geometry_from_path( + self.model_path + ).local_slot_size_for_block_tokens(self.block_tokens, self.tensor_parallel_size) + def runtime_config(self) -> dict[str, object]: """Return connector runtime config owned by DaseR server. @@ -204,6 +262,9 @@ def runtime_config(self) -> dict[str, object]: "socket_path": self.ipc_socket_path, "store_path": "" if self.skip_l2 else self.store_path, "slot_size": self.resolved_slot_size(), + "local_slot_size": self.resolved_local_slot_size(), + "tensor_parallel_size": self.tensor_parallel_size, + "rank_stride_bytes": self.total_slots * self.resolved_local_slot_size(), "block_tokens": self.block_tokens, "model_id": self.model_id, "cache_reuse_mode": self.cache_reuse_mode, diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 1e2527b..32e00dd 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -7,6 +7,7 @@ # Third Party import torch +from vllm.distributed import get_tensor_model_parallel_rank from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, KVConnectorRole, @@ -95,6 +96,16 @@ def __init__( self._socket_path: str = extra.get("socket_path", "/tmp/daser.sock") self._store_path: str = "" self._slot_size: int = 0 + parallel_config = getattr(vllm_config, "parallel_config", None) + self._tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1) or 1) + self._tp_rank = ( + get_tensor_model_parallel_rank() + if role == KVConnectorRole.WORKER and self._tp_size > 1 + else 0 + ) + self._server_tp_size = 1 + self._local_slot_size = 0 + self._rank_stride_bytes = 0 self._block_tokens: int = 16 self._model_id: str = "default" self._skip_l2: bool = bool(extra.get("skip_l2", False)) @@ -211,6 +222,12 @@ def _refresh_runtime_config(self) -> None: self._store_path = str(config.get("store_path", self._store_path)) self._slot_size = int(config.get("slot_size", self._slot_size)) + self._server_tp_size = int( + config.get("tensor_parallel_size", self._server_tp_size) + ) + self._rank_stride_bytes = int( + config.get("rank_stride_bytes", self._rank_stride_bytes) + ) self._block_tokens = int(config.get("block_tokens", self._block_tokens)) self._model_id = str(config.get("model_id", self._model_id)) self._set_cache_reuse_strategy(str(config["cache_reuse_mode"])) diff --git a/daser/connector/ipc_client.py b/daser/connector/ipc_client.py index 15bd404..f024c10 100644 --- a/daser/connector/ipc_client.py +++ b/daser/connector/ipc_client.py @@ -237,13 +237,22 @@ def alloc_chunks( raise RuntimeError("[IPC] invalid alloc_chunks response") return [dict(alloc) for alloc in allocations] - def commit_chunk(self, chunk_key: str) -> None: + def commit_chunk(self, chunk_key: str, tp_rank: int = 0, tp_size: int = 1) -> None: """Mark a chunk as committed (GDS write complete). Args: chunk_key: xxh3_128 hex of the chunk's token IDs. + tp_rank: tensor-parallel rank whose shard was stored. + tp_size: total tensor-parallel ranks required for publication. """ - self.call({"op": "commit_chunk", "chunk_key": chunk_key}) + self.call( + { + "op": "commit_chunk", + "chunk_key": chunk_key, + "tp_rank": tp_rank, + "tp_size": tp_size, + } + ) def transfer_drain(self) -> None: """Wait for server-owned transfer-layer background work. @@ -381,21 +390,43 @@ async def close(self) -> None: async with self._lock: await self._reset() - async def commit_chunk(self, chunk_key: str) -> None: + async def commit_chunk( + self, chunk_key: str, tp_rank: int = 0, tp_size: int = 1 + ) -> None: """Async: mark a chunk as committed. Args: chunk_key: xxh3_128 hex of the token IDs. + tp_rank: tensor-parallel rank whose shard was stored. + tp_size: total tensor-parallel ranks required for publication. """ - await self.call({"op": "commit_chunk", "chunk_key": chunk_key}) + await self.call( + { + "op": "commit_chunk", + "chunk_key": chunk_key, + "tp_rank": tp_rank, + "tp_size": tp_size, + } + ) - async def commit_chunks(self, chunk_keys: list[str]) -> None: + async def commit_chunks( + self, chunk_keys: list[str], tp_rank: int = 0, tp_size: int = 1 + ) -> None: """Async: mark multiple chunks as committed in one RPC. Args: chunk_keys: xxh3_128 hex chunk keys. + tp_rank: tensor-parallel rank whose shards were stored. + tp_size: total tensor-parallel ranks required for publication. """ - await self.call({"op": "commit_chunks", "chunk_keys": chunk_keys}) + await self.call( + { + "op": "commit_chunks", + "chunk_keys": chunk_keys, + "tp_rank": tp_rank, + "tp_size": tp_size, + } + ) async def transfer_drain(self) -> None: """Async: wait for server-owned transfer-layer background work. @@ -597,6 +628,7 @@ async def register_load_staging_cuda( async def transfer_load_registered_cuda( self, buffer_index: int, + producer_pid: int, nbytes: int, spans: list[dict[str, int]], ) -> dict[str, Any]: @@ -605,6 +637,7 @@ async def transfer_load_registered_cuda( Args: buffer_index: Worker-local fixed staging buffer index registered through ``register_load_staging_cuda``. + producer_pid: Process ID that registered the staging buffer. nbytes: logical bytes to write for this transfer. spans: byte spans containing target_offset, nbytes, and file_offset. @@ -620,6 +653,7 @@ async def transfer_load_registered_cuda( "op": "transfer_load", "payload": { "load_staging_buffer_index": int(buffer_index), + "producer_pid": int(producer_pid), "nbytes": int(nbytes), }, "spans": spans, diff --git a/daser/connector/scheduler.py b/daser/connector/scheduler.py index f8ef8f9..6064340 100644 --- a/daser/connector/scheduler.py +++ b/daser/connector/scheduler.py @@ -930,7 +930,8 @@ def update_connector_output(self, connector_output: Any) -> None: self._req_tokens.pop(req_id, None) self._discard_pending_request(req_id) for req_id in getattr(connector_output, "finished_recving", None) or (): - self._req_tokens.pop(req_id, None) + if req_id not in self._pending_alloc: + self._req_tokens.pop(req_id, None) for pending_req_id in list(self._pending_loads): if _matches_request_or_store_id(pending_req_id, req_id): self._pending_loads.pop(pending_req_id, None) diff --git a/daser/connector/worker.py b/daser/connector/worker.py index 78143e1..35b2375 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -5,7 +5,7 @@ # Standard import asyncio from collections import deque -from dataclasses import dataclass +from dataclasses import dataclass, replace import os import queue import threading @@ -90,6 +90,76 @@ _LoadBatch = tuple[int, list[dict[str, int]], list[Any]] +def _rank_lane_offset( + start_slot: int, + local_slot_size: int, + rank_stride_bytes: int, + tp_rank: int, +) -> int: + """Return the physical offset for one rank-local logical slot range. + + Args: + start_slot: First server-owned logical slot. + local_slot_size: Bytes stored per logical slot by one TP rank. + rank_stride_bytes: Byte distance between adjacent rank lanes. + tp_rank: Current vLLM tensor-parallel rank. + + Returns: + Physical store offset for ``start_slot`` in ``tp_rank``'s lane. + + Async/thread-safety: + Pure arithmetic used on the worker thread before IPC submission. + """ + return tp_rank * rank_stride_bytes + start_slot * local_slot_size + + +def _local_slot_bytes(connector: Any) -> int: + """Return per-rank slot bytes, falling back for TP=1 test probes.""" + local_slot_size = int(getattr(connector, "_local_slot_size", 0)) + return local_slot_size or int(connector._slot_size) # noqa: SLF001 + + +def _validate_tp_layout( + local_slot_size: int, + storage_slot_size: int, + tp_size: int, + server_tp_size: int, + tp_rank: int, + rank_stride_bytes: int = 0, +) -> None: + """Validate worker KV geometry against the server-owned TP layout. + + Args: + local_slot_size: Slot bytes measured from the worker KV tensor. + storage_slot_size: Aggregate slot bytes reported by the server. + tp_size: vLLM worker tensor-parallel size. + server_tp_size: Tensor-parallel size reported by the server. + tp_rank: Current vLLM tensor-parallel rank. + rank_stride_bytes: Byte distance between server-owned rank lanes. + + Raises: + ValueError: if rank counts or slot geometry do not match. + + Async/thread-safety: + Pure startup validation called before request traffic. + """ + if tp_size <= 0 or not 0 <= tp_rank < tp_size: + raise ValueError(f"invalid TP rank {tp_rank} for size {tp_size}") + if not storage_slot_size: + return + if server_tp_size != tp_size: + raise ValueError( + f"vLLM TP size {tp_size} does not match DaseR TP size {server_tp_size}" + ) + if local_slot_size * tp_size != storage_slot_size: + raise ValueError( + "worker KV slot geometry does not match DaseR storage layout: " + f"local={local_slot_size} tp={tp_size} storage={storage_slot_size}" + ) + if tp_size > 1 and rank_stride_bytes <= 0: + raise ValueError("DaseR runtime config is missing TP rank stride") + + def _store_staging_pool_depth(buffer_bytes: int, pending_limit_bytes: int) -> int: """Return fixed store staging pool depth for the configured byte budget. @@ -618,20 +688,32 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: self._pending_store_staging_limit_bytes, ) - if self._slot_size == 0 and self._layer_names and sample is not None: + if self._layer_names and sample is not None: num_blocks = sample.shape[1] if sample.dim() >= 2 else 1 layer_size = sample.nbytes // num_blocks - self._slot_size = layer_size * len(self._layer_names) - logger.info( - "[CONNECTOR] computed slot_size=%d from %d layers", + local_slot_size = layer_size * len(self._layer_names) + tp_size = int(getattr(self, "_tp_size", 1)) + _validate_tp_layout( + local_slot_size, self._slot_size, + tp_size, + int(getattr(self, "_server_tp_size", tp_size)), + int(getattr(self, "_tp_rank", 0)), + int(getattr(self, "_rank_stride_bytes", 0)), + ) + self._local_slot_size = local_slot_size + if self._slot_size == 0: + self._slot_size = local_slot_size * tp_size + logger.info( + "[CONNECTOR] registered local_slot_size=%d from %d layers", + self._local_slot_size, len(self._layer_names), ) if sample is not None: self._store_staging_bytes = max( self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, - self._slot_size, + self._local_slot_size, ) self._store_staging_pool = StoreCudaStagingPool( device=sample.device, @@ -712,16 +794,27 @@ def register_cross_layers_kv_cache( self._pending_store_staging_limit_bytes, ) = _derive_store_staging_limits(kv_cache.device) layer_size = kv_cache[0, 0].nbytes + local_slot_size = layer_size * len(self._layer_names) + tp_size = int(getattr(self, "_tp_size", 1)) + _validate_tp_layout( + local_slot_size, + self._slot_size, + tp_size, + int(getattr(self, "_server_tp_size", tp_size)), + int(getattr(self, "_tp_rank", 0)), + int(getattr(self, "_rank_stride_bytes", 0)), + ) + self._local_slot_size = local_slot_size if self._slot_size == 0: - self._slot_size = layer_size * len(self._layer_names) - logger.info( - "[CONNECTOR] computed cross-layer slot_size=%d from %d layers", - self._slot_size, - len(self._layer_names), - ) + self._slot_size = local_slot_size * tp_size + logger.info( + "[CONNECTOR] registered cross-layer local_slot_size=%d from %d layers", + self._local_slot_size, + len(self._layer_names), + ) self._store_staging_bytes = max( self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, - self._slot_size, + self._local_slot_size, ) self._store_staging_pool = StoreCudaStagingPool( device=kv_cache.device, @@ -976,6 +1069,8 @@ async def _run_load_request_dispatcher(self, sample_tensor: torch.Tensor) -> Non to coalesce requests; it submits each queued request immediately when both an in-flight slot and a fixed staging buffer are available. """ + if sample_tensor.device.type == "cuda": + torch.cuda.set_device(sample_tensor.device) load_queue = self._ensure_load_request_queue(sample_tensor) load_staging_pool = self._ensure_load_staging_pool(sample_tensor) dispatcher = LoadRequestDispatcher( @@ -1075,7 +1170,7 @@ def _ensure_load_staging_pool( return pool buffer_bytes = max( self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES, - self._slot_size, + _local_slot_bytes(self), ) pool = FixedCudaStagingPool( device=sample_tensor.device, @@ -1119,6 +1214,7 @@ def _submit_load_batch( if getattr(self, "_load_staging_registered", False): transfer_coro = self._transfer_load_registered_cuda( buffer_index=buffer_index, + producer_pid=os.getpid(), nbytes=total_bytes, spans=spans, ) @@ -1171,9 +1267,18 @@ def _submit_request_load_for_dispatcher( next segment only after the previous segment has been restored. """ load_staging_pool = self._ensure_load_staging_pool(sample_tensor) + spec = replace( + item.spec, + file_offset=_rank_lane_offset( + item.spec.start_slot, + _local_slot_bytes(self), + int(getattr(self, "_rank_stride_bytes", 0)), + int(getattr(self, "_tp_rank", 0)), + ), + ) load_batches = _build_load_read_batches( - {item.spec_id: item.spec}, - self._slot_size, + {item.spec_id: spec}, + _local_slot_bytes(self), max_batch_bytes=load_staging_pool.buffer_bytes, include_req_ids=True, ) @@ -1289,7 +1394,7 @@ def _consume_loaded_batch( kv_caches=self._kv_caches, layer_names=self._layer_names, block_ids=run.block_ids, - slot_size=self._slot_size, + slot_size=_local_slot_bytes(self), load_key_scale=self._load_key_scale, load_value_scale=self._load_value_scale, pos_offset=run.pos_offset, @@ -1562,6 +1667,15 @@ def _ensure_transfer_ready(self) -> bool: ) return False + _validate_tp_layout( + _local_slot_bytes(self), + self._slot_size, + int(getattr(self, "_tp_size", 1)), + int(getattr(self, "_server_tp_size", 1)), + int(getattr(self, "_tp_rank", 0)), + int(getattr(self, "_rank_stride_bytes", 0)), + ) + self._transfer_ready = True logger.info("[CONNECTOR] server transfer mode=%s", self._transfer_mode) return True @@ -1756,23 +1870,21 @@ def _wait_for_store_staging_release(self, nbytes: int) -> None: Async/thread-safety: Called from the worker thread when the fixed store staging pool is exhausted. It first applies byte-budget backpressure, then waits - for one oldest store future if no lease was released yet. + for oldest store futures until a lease is released. """ pool = getattr(self, "_store_staging_pool", None) - before = pool.available if pool is not None else 0 self._wait_for_save_staging_capacity(nbytes) - if pool is None or pool.available > before or not self._save_futures: - return - record = self._save_futures.pop(0) - try: - record.future.result(timeout=120.0) - finally: - self._pending_save_staging_bytes = max( - 0, - self._pending_save_staging_bytes - record.staging_bytes, - ) - record.release() - self._reap_save_futures(block=False) + while pool is not None and pool.available == 0 and self._save_futures: + record = self._save_futures.pop(0) + try: + record.future.result(timeout=120.0) + finally: + self._pending_save_staging_bytes = max( + 0, + self._pending_save_staging_bytes - record.staging_bytes, + ) + record.release() + self._reap_save_futures(block=False) def _acquire_staging( self, @@ -1841,7 +1953,8 @@ def _stage_store_batch( return None if not block_ids or not spans: return None - nbytes = len(block_ids) * self._slot_size + local_slot_size = _local_slot_bytes(self) + nbytes = len(block_ids) * local_slot_size self._wait_for_save_staging_capacity(nbytes) staging_lease = self._acquire_staging(nbytes, sample_tensor.device) staging = staging_lease.view @@ -1857,7 +1970,7 @@ def _stage_store_batch( kv_cache=cross_layer_kv_cache, block_ids=block_ids, num_layers=num_layers, - slot_size=self._slot_size, + slot_size=local_slot_size, block_index=block_index, ) else: @@ -1868,7 +1981,7 @@ def _stage_store_batch( layer_idx=self._layer_idx_map[layer_name], block_ids=block_ids, num_layers=num_layers, - slot_size=self._slot_size, + slot_size=local_slot_size, block_index=block_index, ) return StagedStoreBatch( @@ -1893,9 +2006,21 @@ def _submit_finished_save(self, save: _DeferredFinishedSave) -> Any | None: still holding the finished request's KV blocks. """ batch_futures = [] + reqs_to_store = { + req_id: replace( + spec, + file_offset=_rank_lane_offset( + spec.start_slot, + _local_slot_bytes(self), + int(getattr(self, "_rank_stride_bytes", 0)), + int(getattr(self, "_tp_rank", 0)), + ), + ) + for req_id, spec in save.reqs_to_store.items() + } batches = _build_staging_store_batches( - save.reqs_to_store, - self._slot_size, + reqs_to_store, + _local_slot_bytes(self), max_batch_bytes=(self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES), ) for block_ids, spans in batches: @@ -1955,6 +2080,8 @@ async def _write_cuda_buffer( Returns: Chunk keys accepted by the server for this buffer. """ + if buffer.device.type == "cuda": + torch.cuda.set_device(buffer.device) if ready_event is not None: ready_event.synchronize() else: @@ -1995,7 +2122,11 @@ async def _commit_stored_keys( requested = set(commit_keys) candidate_keys = [key for key in stored_keys if key in requested] keys_to_commit = list(dict.fromkeys(candidate_keys)) - await self._ipc_store_async.commit_chunks(keys_to_commit) + await self._ipc_store_async.commit_chunks( + keys_to_commit, + tp_rank=int(getattr(self, "_tp_rank", 0)), + tp_size=int(getattr(self, "_tp_size", 1)), + ) async def _transfer_load_cuda(self, **kwargs: Any) -> dict[str, Any]: """Load through the dedicated worker load IPC client. diff --git a/daser/server/__main__.py b/daser/server/__main__.py index 82e3fb2..4351603 100644 --- a/daser/server/__main__.py +++ b/daser/server/__main__.py @@ -169,6 +169,11 @@ def _parse_args() -> argparse.Namespace: default=None, help="HuggingFace model path used by vLLM and tokenizer loading", ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow custom model repository code when loading the tokenizer.", + ) parser.add_argument( "--store-dir", required=True, @@ -219,6 +224,12 @@ def _parse_args() -> argparse.Namespace: default=BLOCK_TOKENS, help="vLLM KV block size in tokens. Must match vLLM --block-size.", ) + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="vLLM tensor parallel size. Must match vLLM --tensor-parallel-size.", + ) return parser.parse_args() @@ -322,6 +333,7 @@ def _build_daser_config(args: argparse.Namespace) -> DaserConfig: transfer_mode=transfer_mode, l1_size_bytes=l1_size, skip_l2=skip_l2, + tensor_parallel_size=int(args.tensor_parallel_size), ) slot_size = cfg.resolved_slot_size() if cfg.total_store_bytes <= 0 or cfg.total_slots <= 0: @@ -356,6 +368,7 @@ def _build_http_config(args: argparse.Namespace) -> HTTPServerConfig: cache_reuse_mode=args.cache_reuse_mode, align_document_chunks=args.cache_reuse_mode == CACHE_REUSE_CHUNK, transfer_mode=args.transfer_mode, + trust_remote_code=bool(args.trust_remote_code), ) diff --git a/daser/server/chunk_lifecycle.py b/daser/server/chunk_lifecycle.py index 9d905fc..10f70d4 100644 --- a/daser/server/chunk_lifecycle.py +++ b/daser/server/chunk_lifecycle.py @@ -24,6 +24,8 @@ def __init__(self) -> None: self._write_owners: set[str] = set() self._evicted: set[str] = set() self._commit_waiters: dict[str, set[asyncio.Future[None]]] = {} + self._commit_shards: dict[str, tuple[int, set[int]]] = {} + self._publishing: set[str] = set() def is_committed(self, chunk_key: str) -> bool: """Return whether ``chunk_key`` is committed and visible to lookup.""" @@ -51,22 +53,69 @@ def mark_committed(self, chunk_key: str) -> None: """Mark ``chunk_key`` committed, claim ownership, and wake waiters.""" self._committed.add(chunk_key) self._write_owners.add(chunk_key) + self._commit_shards.pop(chunk_key, None) + self._publishing.discard(chunk_key) self._notify_commit_waiters(chunk_key) + def record_commit_shard(self, chunk_key: str, tp_rank: int, tp_size: int) -> bool: + """Record one TP rank and return whether this call should publish. + + Args: + chunk_key: Cache key for the pending chunk. + tp_rank: Tensor-parallel rank that completed its store. + tp_size: Required number of tensor-parallel ranks. + + Returns: + True exactly once when all distinct ranks have arrived. + + Raises: + ValueError: if rank metadata is invalid or changes mid-commit. + + Async/thread-safety: + Runs on the server event loop. ``_publishing`` prevents another + request from publishing while the retrieval insert awaits. + """ + if tp_size <= 0 or not 0 <= tp_rank < tp_size: + raise ValueError(f"invalid TP rank {tp_rank} for size {tp_size}") + if chunk_key in self._committed or chunk_key in self._publishing: + return False + expected_size, ranks = self._commit_shards.setdefault( + chunk_key, (tp_size, set()) + ) + if expected_size != tp_size: + raise ValueError( + f"inconsistent TP size for chunk: {tp_size} != {expected_size}" + ) + ranks.add(tp_rank) + if len(ranks) != tp_size: + return False + self._publishing.add(chunk_key) + return True + + def abort_publish(self, chunk_key: str) -> None: + """Allow a failed retrieval-index publish to be retried.""" + self._publishing.discard(chunk_key) + def mark_evicted(self, chunk_key: str) -> None: """Drop committed/owner state for ``chunk_key`` and record eviction.""" self._committed.discard(chunk_key) self._write_owners.discard(chunk_key) + self._commit_shards.pop(chunk_key, None) + self._publishing.discard(chunk_key) self._evicted.add(chunk_key) def discard_owner(self, chunk_key: str) -> None: """Release a store-writer claim without committing or evicting.""" self._write_owners.discard(chunk_key) + self._commit_shards.pop(chunk_key, None) + self._publishing.discard(chunk_key) def discard(self, chunk_key: str) -> None: """Drop committed and write-owner state without recording eviction.""" self._committed.discard(chunk_key) self._write_owners.discard(chunk_key) + self._commit_shards.pop(chunk_key, None) + self._publishing.discard(chunk_key) async def wait_for_committed( self, diff --git a/daser/server/core.py b/daser/server/core.py index c3c3096..3a4e6a0 100644 --- a/daser/server/core.py +++ b/daser/server/core.py @@ -422,11 +422,15 @@ async def match_and_alloc( ), ) - async def commit_chunk(self, chunk_key: str) -> None: + async def commit_chunk( + self, chunk_key: str, tp_rank: int = 0, tp_size: int = 1 + ) -> None: """Mark a chunk as committed and visible to lookup. Args: chunk_key: cache key for the chunk. + tp_rank: tensor-parallel rank whose shard finished storing. + tp_size: total tensor-parallel ranks required before publication. Raises: ValueError: if the chunk was not allocated. @@ -449,15 +453,26 @@ async def commit_chunk(self, chunk_key: str) -> None: ) return raise ValueError(f"chunk_key not found: {chunk_key}") - await self._ri.insert(meta) - self._lifecycle.mark_committed(chunk_key) self._commit_requests += 1 + if not self._lifecycle.record_commit_shard(chunk_key, tp_rank, tp_size): + return + try: + await self._ri.insert(meta) + except BaseException: + self._lifecycle.abort_publish(chunk_key) + raise + self._lifecycle.mark_committed(chunk_key) self._metrics.counter( "daser_cache_committed_chunks_total", "Chunks committed and published to the retrieval index.", ).inc() self._record_capacity_metrics() - logger.debug("[CORE] commit_chunk key=%s", chunk_key[:8]) + logger.debug( + "[CORE] commit_chunk key=%s tp_rank=%d tp_size=%d", + chunk_key[:8], + tp_rank, + tp_size, + ) def is_chunk_committed(self, chunk_key: str) -> bool: """Return whether a chunk key has been committed. @@ -521,6 +536,22 @@ async def wait_for_committed_chunks( """ await self._lifecycle.wait_for_committed(chunk_keys, timeout_s) + async def wait_for_pending_chunks(self, timeout_s: float) -> None: + """Wait for currently allocated store writers to commit. + + Args: + timeout_s: maximum seconds to wait. + + Raises: + TimeoutError: if any pending chunk remains uncommitted at timeout. + + Async/thread-safety: + Snapshots lifecycle state and waits without blocking the server + event loop. Writers registered after the snapshot are not included. + """ + pending = list(self._lifecycle.write_owners - self._lifecycle.committed) + await self._lifecycle.wait_for_committed(pending, timeout_s) + async def commit_stats(self) -> dict[str, int]: """Return connector commit counters for benchmark synchronization. diff --git a/daser/server/http/app.py b/daser/server/http/app.py index 45a41cd..e121fcf 100644 --- a/daser/server/http/app.py +++ b/daser/server/http/app.py @@ -45,6 +45,7 @@ class HTTPServerConfig: cache_reuse_mode: cache reuse strategy selected by the DaseR server. align_document_chunks: when True, insert padding tokens before each document so document chunks begin on vLLM block boundaries. + trust_remote_code: allow tokenizer repository Python code. """ vllm_base_url: str @@ -59,6 +60,7 @@ class HTTPServerConfig: cache_reuse_mode: str = DEFAULT_CACHE_REUSE_MODE align_document_chunks: bool = False transfer_mode: str = "iouring" + trust_remote_code: bool = False class UploadRequest(BaseModel): @@ -534,7 +536,9 @@ def build_http_app( # Third Party from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer) + tokenizer = AutoTokenizer.from_pretrained( + cfg.tokenizer, trust_remote_code=cfg.trust_remote_code + ) if vllm is None: vllm = VLLMClient(base_url=cfg.vllm_base_url, model=cfg.model) pad_token = _resolve_pad_token(tokenizer) @@ -622,9 +626,16 @@ async def metrics_endpoint() -> Response: @app.post("/drain") async def drain_endpoint() -> dict[str, bool]: - """Wait for server-owned transfer background work to finish.""" + """Wait for transfer work and allocated stores to commit.""" if drain_transfer is not None: await drain_transfer() + try: + await core.wait_for_pending_chunks(timeout_s=_DOCUMENT_STORE_SYNC_TIMEOUT_S) + except TimeoutError as exc: + raise HTTPException( + status_code=504, + detail="DaseR drain timed out waiting for pending stores", + ) from exc return {"ok": True} @app.post("/documents", status_code=201) diff --git a/daser/server/ipc/server.py b/daser/server/ipc/server.py index 95102db..3571e62 100644 --- a/daser/server/ipc/server.py +++ b/daser/server/ipc/server.py @@ -126,7 +126,7 @@ def __init__( self._cuda_ipc_cache: OrderedDict[ tuple[int, int, int, int | None], "_CachedCudaArray" ] = OrderedDict() - self._load_staging_buffers: dict[int, _CachedCudaArray] = {} + self._load_staging_buffers: dict[tuple[int, int], _CachedCudaArray] = {} self._op_handlers: dict[ str, Callable[[dict[str, Any]], Awaitable[dict[str, Any]]] ] = { @@ -351,13 +351,21 @@ async def _op_match_and_alloc(self, msg: dict[str, Any]) -> dict[str, Any]: async def _op_commit_chunk(self, msg: dict[str, Any]) -> dict[str, Any]: """Handle a single ``commit_chunk`` request.""" - await self._core.commit_chunk(msg["chunk_key"]) + await self._core.commit_chunk( + msg["chunk_key"], + tp_rank=int(msg.get("tp_rank", 0)), + tp_size=int(msg.get("tp_size", 1)), + ) return {"ok": True} async def _op_commit_chunks(self, msg: dict[str, Any]) -> dict[str, Any]: """Handle a batched ``commit_chunks`` request.""" for chunk_key in msg.get("chunk_keys", []): - await self._core.commit_chunk(chunk_key) + await self._core.commit_chunk( + chunk_key, + tp_rank=int(msg.get("tp_rank", 0)), + tp_size=int(msg.get("tp_size", 1)), + ) return {"ok": True} async def _op_commit_stats(self, msg: dict[str, Any]) -> dict[str, Any]: @@ -565,13 +573,15 @@ async def _register_load_staging(self, msg: dict[str, Any]) -> dict[str, Any]: """ payload = msg.get("payload", {}) buffer_index = int(payload["buffer_index"]) + producer_pid = int(payload["producer_pid"]) + buffer_key = (producer_pid, buffer_index) opened = self._open_cuda_ipc_payload( payload=payload, nbytes_key="allocation_bytes", cache_mapping=False, ) - previous = self._load_staging_buffers.get(buffer_index) - self._load_staging_buffers[buffer_index] = opened + previous = self._load_staging_buffers.get(buffer_key) + self._load_staging_buffers[buffer_key] = opened if previous is not None: previous.close() return {"ok": True} @@ -704,11 +714,14 @@ def _payload_buffer(self, payload: dict[str, Any]) -> Any: return bytearray(payload["data"]) if "load_staging_buffer_index" in payload: buffer_index = int(payload["load_staging_buffer_index"]) + producer_pid = int(payload["producer_pid"]) + buffer_key = (producer_pid, buffer_index) try: - return self._load_staging_buffers[buffer_index] + return self._load_staging_buffers[buffer_key] except KeyError as exc: raise ValueError( - f"unknown load staging buffer index: {buffer_index}" + "unknown load staging buffer: " + f"producer_pid={producer_pid} index={buffer_index}" ) from exc if "cuda_ipc_handle" in payload: return self._open_cuda_ipc_payload(payload=payload, nbytes_key="nbytes") @@ -819,9 +832,11 @@ def __getitem__(self, item: Any) -> Any: def synchronize(self) -> None: """Synchronize CUDA writes issued through the opened array.""" + import cupy from cupy.cuda import runtime - runtime.streamSynchronize(0) + with cupy.cuda.Device(int(self._opened.array.device.id)): + runtime.streamSynchronize(0) def close(self) -> None: """Close the CUDA IPC handle.""" diff --git a/daser/transfer/cuda_ipc.py b/daser/transfer/cuda_ipc.py index 0560325..26868b0 100644 --- a/daser/transfer/cuda_ipc.py +++ b/daser/transfer/cuda_ipc.py @@ -25,10 +25,12 @@ class CudaIPCBuffer: def close(self) -> None: """Close the CUDA IPC memory handle.""" + import cupy # Third Party from cupy.cuda import runtime # Third Party if self.owns_handle: - runtime.ipcCloseMemHandle(self.ptr) + with cupy.cuda.Device(int(self.array.device.id)): + runtime.ipcCloseMemHandle(self.ptr) def open_cuda_ipc_buffer( diff --git a/daser/transfer/iouring/copy_ops.py b/daser/transfer/iouring/copy_ops.py index 7c86187..519c6ac 100644 --- a/daser/transfer/iouring/copy_ops.py +++ b/daser/transfer/iouring/copy_ops.py @@ -97,14 +97,16 @@ def copy_src_to_pinned( nbytes: number of bytes to copy. """ if hasattr(src, "data") and getattr(src.data, "ptr", None) is not None: + import cupy from cupy.cuda import runtime - runtime.memcpy( - pinned.ptr_at(target_offset), - int(src.data.ptr), - nbytes, - runtime.memcpyDeviceToHost, - ) + with cupy.cuda.Device(int(src.device.id)): + runtime.memcpy( + pinned.ptr_at(target_offset), + int(src.data.ptr), + nbytes, + runtime.memcpyDeviceToHost, + ) return pinned.view()[target_offset : target_offset + nbytes] = memoryview(src).cast("B")[ :nbytes @@ -145,6 +147,7 @@ def copy_grouped_to_cuda_dst(dst: Any, chunks: list[CopyChunk]) -> None: dst: CuPy ndarray destination. chunks: grouped copy chunks from the L1 tier. """ + import cupy from cupy.cuda import runtime ordered = sorted(chunks, key=lambda item: item[0]) @@ -168,13 +171,14 @@ def copy_grouped_to_cuda_dst(dst: Any, chunks: list[CopyChunk]) -> None: dst_ptr = cuda_array_ptr(target) if dst_ptr is None: raise TypeError("grouped CUDA copy target lost CUDA array interface") - runtime.memcpyAsync( - dst_ptr, - source_ptr, - nbytes, - runtime.memcpyHostToDevice, - 0, - ) + with cupy.cuda.Device(int(target.device.id)): + runtime.memcpyAsync( + dst_ptr, + source_ptr, + nbytes, + runtime.memcpyHostToDevice, + 0, + ) def copy_grouped_to_dst(dst: Any, chunks: list[CopyChunk]) -> None: diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index d347c37..ea18b19 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -73,7 +73,9 @@ _InflightRequestLoad, _load_staging_pool_depth, _PendingLoad, + _rank_lane_offset, _RequestLoadFuture, + _SaveFuture, ) BLOCK_TOKENS = 4 @@ -345,12 +347,16 @@ class _CommitProbe(WorkerConnectorMixin): def __init__(self) -> None: self.committed: list[list[str]] = [] + self.commit_metadata: list[tuple[int, int]] = [] self._ipc_async = self self._ipc_store_async = self - async def commit_chunks(self, chunk_keys: list[str]) -> None: + async def commit_chunks( + self, chunk_keys: list[str], tp_rank: int = 0, tp_size: int = 1 + ) -> None: """Record chunk keys submitted to the async IPC client.""" self.committed.append(list(chunk_keys)) + self.commit_metadata.append((tp_rank, tp_size)) async def commit_stored_keys( self, stored_keys: list[str], commit_keys: list[str] @@ -2306,6 +2312,46 @@ def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): assert calls == [] +def test_store_staging_wait_skips_future_without_lease() -> None: + """Store backpressure waits until a future returns a staging lease.""" + + class _Future: + def __init__(self, name: str) -> None: + self.name = name + + def done(self) -> bool: + return False + + def result(self, timeout: float) -> None: + assert timeout > 0 + completed.append(self.name) + + class Probe(WorkerConnectorMixin): + def __init__(self) -> None: + self._store_staging_pool = StoreCudaStagingPool( + device=torch.device("cpu"), + buffer_bytes=16, + depth=1, + ) + lease = self._store_staging_pool.acquire(8) + self._pending_store_staging_limit_bytes = 64 + self._pending_save_staging_bytes = 8 + self._save_futures = [ + _SaveFuture(_Future("commit"), 0, None), + _SaveFuture(_Future("store"), 8, lease), + ] + + def wait_for_staging(self) -> None: + self._wait_for_store_staging_release(8) + + completed: list[str] = [] + probe = Probe() + probe.wait_for_staging() + + assert completed == ["commit", "store"] + assert probe._store_staging_pool.available == 1 # noqa: SLF001 + + def test_stage_store_batch_records_ready_event_without_synchronizing(monkeypatch): from daser.connector import worker @@ -3180,6 +3226,42 @@ def test_finished_recving_keeps_pending_suffix_stores() -> None: assert not connector.has_req_tokens("req") +def test_finished_recving_keeps_tokens_until_suffix_allocation() -> None: + """Async prefix loads retain tokens needed to key later suffix blocks.""" + tokens = list(range(12)) + key0, key1, key2 = rolling_keys(tokens, BLOCK_TOKENS) + connector = _AllocatingSchedulerProbe() + connector.use_prefix_reuse_strategy() + connector.seed_tokens("req", tokens) + connector._pending_alloc["req"] = PendingStore( # noqa: SLF001 + chunk_key="", + token_count=12, + start_slot_index=1, + rolling_key=key0, + rolling_slot_index=1, + block_ids=[10], + ) + + connector.update_connector_output_for_test(finished_recving={"req"}) + assert connector.has_req_tokens("req") + + class Cached: + req_ids = ["req"] + new_block_ids = [([11, 12],)] + resumed_req_ids: set[str] = set() + + class Output: + scheduled_cached_reqs = Cached() + + connector.record_cached_blocks(Output()) + + pending_alloc, pending_stores = connector.pending_state + assert pending_alloc == {} + assert connector.alloc_calls == [("batch", 2, "m")] + assert pending_stores["req:store:1"]["chunk_key"] == key1 + assert pending_stores["req:store:2"]["chunk_key"] == key2 + + def test_record_cached_store_blocks_appends_resumed_incremental_blocks(): """Resumed chunked-prefill steps may report only newly allocated blocks.""" connector = _AllocatingSchedulerProbe() @@ -3714,6 +3796,20 @@ async def test_commit_empty_stored_keys_does_not_publish_requested_chunks(): await connector.commit_stored_keys([], ["stale-key"]) assert connector.committed == [[]] + assert connector.commit_metadata == [(0, 1)] + + +def test_tensor_parallel_rank_lanes_are_contiguous_and_disjoint() -> None: + """Each TP rank maps a logical slot run into its own contiguous lane.""" + local_slot_size = 32 + rank_stride = 10 * local_slot_size + + rank_0 = _rank_lane_offset(3, local_slot_size, rank_stride, tp_rank=0) + rank_1 = _rank_lane_offset(3, local_slot_size, rank_stride, tp_rank=1) + + assert rank_0 == 3 * local_slot_size + assert rank_1 == rank_stride + 3 * local_slot_size + assert rank_0 + 2 * local_slot_size <= rank_1 @pytest.mark.asyncio diff --git a/tests/connector/test_ipc_client.py b/tests/connector/test_ipc_client.py index 1922677..a542eef 100644 --- a/tests/connector/test_ipc_client.py +++ b/tests/connector/test_ipc_client.py @@ -398,6 +398,7 @@ async def fake_call(self, payload: dict) -> dict: ) response = await client.transfer_load_registered_cuda( buffer_index=1, + producer_pid=43, nbytes=16, spans=[{"target_offset": 0, "nbytes": 16, "file_offset": 0}], ) @@ -421,6 +422,7 @@ async def fake_call(self, payload: dict) -> dict: "op": "transfer_load", "payload": { "load_staging_buffer_index": 1, + "producer_pid": 43, "nbytes": 16, }, "spans": [{"target_offset": 0, "nbytes": 16, "file_offset": 0}], diff --git a/tests/server/test_core.py b/tests/server/test_core.py index 3b9ef33..0d0f52c 100644 --- a/tests/server/test_core.py +++ b/tests/server/test_core.py @@ -190,6 +190,43 @@ async def test_wait_for_committed_chunks_times_out() -> None: await core.wait_for_committed_chunks(["missing"], timeout_s=0.001) +@pytest.mark.asyncio +async def test_wait_for_pending_chunks_waits_for_tp_commit_quorum() -> None: + core = make_core() + await core.wait_for_pending_chunks(timeout_s=0.001) + + tokens = [1, 2, 3, 4] + key = first_rolling_key(tokens) + await core.alloc_chunk(key, token_count=len(tokens), model_id="m") + await core.commit_chunk(key, tp_rank=0, tp_size=2) + + waiter = asyncio.create_task(core.wait_for_pending_chunks(timeout_s=1.0)) + await asyncio.sleep(0) + assert waiter.done() is False + + await core.commit_chunk(key, tp_rank=1, tp_size=2) + await waiter + + +@pytest.mark.asyncio +async def test_tensor_parallel_commit_waits_for_all_distinct_ranks() -> None: + """A chunk becomes visible only after every declared TP rank commits.""" + core = make_core() + tokens = [1, 2, 3, 4] + key = first_rolling_key(tokens) + await core.alloc_chunk(key, token_count=len(tokens), model_id="m") + + await core.commit_chunk(key, tp_rank=0, tp_size=2) + await core.commit_chunk(key, tp_rank=0, tp_size=2) + assert await core.lookup(tokens, "m") == [] + + with pytest.raises(ValueError, match="inconsistent TP size"): + await core.commit_chunk(key, tp_rank=1, tp_size=3) + + await core.commit_chunk(key, tp_rank=1, tp_size=2) + assert len(await core.lookup(tokens, "m")) == 1 + + @pytest.mark.asyncio async def test_restored_orphan_committed_chunk_can_be_reused(tmp_path) -> None: tokens = [1, 2, 3, 4] diff --git a/tests/server/test_http_server.py b/tests/server/test_http_server.py index f96e917..454fe84 100644 --- a/tests/server/test_http_server.py +++ b/tests/server/test_http_server.py @@ -179,13 +179,18 @@ def _make_client( def test_drain_endpoint_waits_for_core_transfer_work() -> None: """POST /drain exposes a benchmark-safe transfer drain primitive.""" - calls = 0 + calls: list[str] = [] async def drain_transfer() -> None: - nonlocal calls - calls += 1 + calls.append("transfer") core = make_core() + + async def wait_for_pending_chunks(timeout_s: float) -> None: + assert timeout_s > 0 + calls.append("commit") + + core.wait_for_pending_chunks = wait_for_pending_chunks # type: ignore[method-assign] fake_vllm = FakeVLLMClient(commit_core=core) app = build_http_app( HTTPServerConfig( @@ -205,7 +210,36 @@ async def drain_transfer() -> None: assert response.status_code == 200 assert response.json() == {"ok": True} - assert calls == 1 + assert calls == ["transfer", "commit"] + + +def test_drain_endpoint_reports_pending_store_timeout() -> None: + """POST /drain reports a bounded failure when stores do not commit.""" + core = make_core() + + async def wait_for_pending_chunks(timeout_s: float) -> None: + del timeout_s + raise TimeoutError + + core.wait_for_pending_chunks = wait_for_pending_chunks # type: ignore[method-assign] + app = build_http_app( + HTTPServerConfig( + vllm_base_url="http://vllm", + model="m", + tokenizer="fake", + block_tokens=4, + ), + core, + tokenizer=FakeTokenizer(), + vllm=FakeVLLMClient(commit_core=core), + ) + + response = TestClient(app).post("/drain") + + assert response.status_code == 504 + assert response.json()["detail"] == ( + "DaseR drain timed out waiting for pending stores" + ) def _ids(text: str) -> list[int]: diff --git a/tests/server/test_ipc_server.py b/tests/server/test_ipc_server.py index e783ecf..f217b37 100644 --- a/tests/server/test_ipc_server.py +++ b/tests/server/test_ipc_server.py @@ -566,15 +566,15 @@ def fake_ensure_transfer(_server: IPCServer) -> FakeTransfer: @pytest.mark.asyncio -async def test_registered_load_staging_reuses_preopened_cuda_mapping( +async def test_registered_load_staging_is_scoped_by_producer( tmp_path, monkeypatch: pytest.MonkeyPatch, ) -> None: - """Registered load staging buffers are opened once and reused by index.""" + """Equal worker-local indexes resolve to each producer's CUDA mapping.""" class FakeOpened: - def __init__(self) -> None: - self.array = bytearray(1024) + def __init__(self, marker: int) -> None: + self.array = bytearray([marker]) * 1024 self.closed = 0 def close(self) -> None: @@ -586,7 +586,7 @@ def close(self) -> None: def fake_open_cuda_ipc_buffer(**kwargs: Any) -> FakeOpened: open_calls.append(kwargs) - opened = FakeOpened() + opened = FakeOpened(len(opened_buffers) + 1) opened_buffers.append(opened) return opened @@ -626,43 +626,51 @@ def fake_ensure_transfer(_server: IPCServer) -> FakeTransfer: await server.start() try: - registered = await _send_recv( - str(tmp_path / "test.sock"), - { - "op": "register_load_staging", - "payload": { - "buffer_index": 1, - "cuda_ipc_handle": b"h" * 64, - "allocation_bytes": 1024, - "device_id": 0, - "device_ptr": 123456, - "allocation_base_ptr": 122880, - "allocation_offset": 576, - "producer_pid": 42, - }, - }, - ) - loaded = await _send_recv( - str(tmp_path / "test.sock"), - { - "op": "transfer_load", - "payload": { - "load_staging_buffer_index": 1, - "nbytes": 128, - }, - "spans": [{"target_offset": 0, "nbytes": 128, "file_offset": 0}], - }, - ) + registered = [] + loaded = [] + for producer_pid in (42, 43): + registered.append( + await _send_recv( + str(tmp_path / "test.sock"), + { + "op": "register_load_staging", + "payload": { + "buffer_index": 1, + "cuda_ipc_handle": b"h" * 64, + "allocation_bytes": 1024, + "device_id": 0, + "device_ptr": 123456, + "allocation_base_ptr": 122880, + "allocation_offset": 576, + "producer_pid": producer_pid, + }, + }, + ) + ) + for producer_pid in (42, 43): + loaded.append( + await _send_recv( + str(tmp_path / "test.sock"), + { + "op": "transfer_load", + "payload": { + "load_staging_buffer_index": 1, + "producer_pid": producer_pid, + "nbytes": 128, + }, + "spans": [ + {"target_offset": 0, "nbytes": 128, "file_offset": 0} + ], + }, + ) + ) finally: await server.stop() - assert registered == {"ok": True} - assert loaded["ok"] is True - assert loaded["bytes"] == 128 - assert len(opened_buffers) == 1 - assert len(loaded_buffers) == 1 - assert loaded_buffers[0][0:128] == opened_buffers[0].array[0:128] - assert open_calls == [ + assert registered == [{"ok": True}, {"ok": True}] + assert [response["bytes"] for response in loaded] == [128, 128] + assert [bytes(buffer[0:1]) for buffer in loaded_buffers] == [b"\x01", b"\x02"] + assert open_calls == 2 * [ { "handle": b"h" * 64, "nbytes": 1024, @@ -671,7 +679,7 @@ def fake_ensure_transfer(_server: IPCServer) -> FakeTransfer: "allocation_offset": 576, } ] - assert opened_buffers[0].closed == 1 + assert [opened.closed for opened in opened_buffers] == [1, 1] @pytest.mark.asyncio diff --git a/tests/server/test_main_cli.py b/tests/server/test_main_cli.py index 3cad6bf..03a67b1 100644 --- a/tests/server/test_main_cli.py +++ b/tests/server/test_main_cli.py @@ -89,6 +89,7 @@ def test_documented_flags_populate_config(tmp_path: Path) -> None: "iouring", "--l1-size", "1gb", + "--trust-remote-code", ] ) cfg = _build_daser_config(args) @@ -111,6 +112,7 @@ def test_documented_flags_populate_config(tmp_path: Path) -> None: assert http_cfg.model == str(model_path) assert http_cfg.tokenizer == str(model_path) assert http_cfg.align_document_chunks is True + assert http_cfg.trust_remote_code is True assert args.cache_reuse_mode == "chunk" runtime = cfg.runtime_config() assert runtime["transfer_mode"] == "iouring" diff --git a/tests/test_config.py b/tests/test_config.py index b792e88..16556c0 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4,6 +4,8 @@ import json from pathlib import Path +import pytest + # First Party from daser.config import ( BLOCK_TOKENS, @@ -131,7 +133,7 @@ def test_runtime_config_reuses_server_parameters(tmp_path: Path) -> None: "hidden_size": 512, "num_attention_heads": 8, "num_key_value_heads": 8, - "num_hidden_layers": 4, + "num_layers": 4, }, ) cfg = DaserConfig( @@ -147,6 +149,9 @@ def test_runtime_config_reuses_server_parameters(tmp_path: Path) -> None: "socket_path": "/tmp/custom.sock", "store_path": str(store_dir / "daser.store"), "slot_size": 8 * 64 * 2 * 4 * BLOCK_TOKENS * 2, + "local_slot_size": 8 * 64 * 2 * 4 * BLOCK_TOKENS * 2, + "tensor_parallel_size": 1, + "rank_stride_bytes": 512 * 64 * 2 * 4 * BLOCK_TOKENS * 2, "block_tokens": BLOCK_TOKENS, "model_id": str(model_path), "cache_reuse_mode": "chunk", @@ -232,3 +237,68 @@ def test_model_id_can_differ_from_model_path(tmp_path: Path) -> None: assert cfg.model_id == "served-model" assert cfg.runtime_config()["model_id"] == "served-model" + + +def test_tensor_parallel_runtime_uses_rank_lanes(tmp_path: Path) -> None: + """TP geometry follows vLLM per-rank KV-head replication semantics.""" + model_path = tmp_path / "model" + _write_model_config( + model_path, + { + "hidden_size": 512, + "num_attention_heads": 8, + "num_key_value_heads": 1, + "num_hidden_layers": 4, + }, + ) + local_slot_size = 64 * 2 * 4 * BLOCK_TOKENS * 2 + cfg = DaserConfig( + model_path=str(model_path), + store_dir=str(tmp_path / "store"), + total_store_bytes=local_slot_size * 2 * 8, + tensor_parallel_size=2, + ) + + runtime = cfg.runtime_config() + + assert cfg.resolved_local_slot_size() == local_slot_size + assert cfg.resolved_slot_size() == local_slot_size * 2 + assert cfg.total_slots == 8 + assert cfg.model_id == f"{model_path}::tp2" + assert runtime["tensor_parallel_size"] == 2 + assert runtime["rank_stride_bytes"] == local_slot_size * 8 + + +def test_model_geometry_rejects_incompatible_tp_head_count(tmp_path: Path) -> None: + """TP geometry follows vLLM's KV-head divisibility requirement.""" + model_path = tmp_path / "model" + _write_model_config( + model_path, + { + "hidden_size": 768, + "num_attention_heads": 6, + "num_key_value_heads": 6, + "num_hidden_layers": 2, + }, + ) + + with pytest.raises(ValueError, match="tensor_parallel_size"): + model_geometry_from_path(str(model_path)).slot_size_for_block_tokens(16, 4) + + +def test_model_geometry_rejects_fp8_without_runtime_kv_spec(tmp_path: Path) -> None: + """FP8 layouts fail before DaseR silently sizes them as BF16.""" + model_path = tmp_path / "model" + _write_model_config( + model_path, + { + "hidden_size": 512, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "num_hidden_layers": 2, + "torch_dtype": "float8_e4m3fn", + }, + ) + + with pytest.raises(ValueError, match="FP8"): + model_geometry_from_path(str(model_path)) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index d673822..de02ecc 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -88,7 +88,7 @@ def _set(_: str, __: Any) -> None: if importlib.util.find_spec("vllm") is None: base = _ensure_module("vllm.distributed.kv_transfer.kv_connector.v1.base") _ensure_module("vllm") - _ensure_module("vllm.distributed") + distributed = _ensure_module("vllm.distributed") _ensure_module("vllm.distributed.kv_transfer") _ensure_module("vllm.distributed.kv_transfer.kv_connector") _ensure_module("vllm.distributed.kv_transfer.kv_connector.v1") @@ -154,9 +154,15 @@ def clear_connector_metadata(self) -> None: """ self._connector_metadata = None + def get_tensor_model_parallel_rank() -> int: + """Return the single-rank value used by CPU-only unit tests.""" + + return 0 + base.KVConnectorBase_V1 = KVConnectorBase_V1 base.KVConnectorMetadata = KVConnectorMetadata base.KVConnectorRole = KVConnectorRole + distributed.get_tensor_model_parallel_rank = get_tensor_model_parallel_rank _install_optional_runtime_stubs() diff --git a/tests/unit/test_benchmark_unified_utils.py b/tests/unit/test_benchmark_unified_utils.py index e44ac79..4218bfc 100644 --- a/tests/unit/test_benchmark_unified_utils.py +++ b/tests/unit/test_benchmark_unified_utils.py @@ -36,7 +36,11 @@ from benchmarks.utils import vllm_bench # First Party -from benchmarks.utils.constants import BYTES_PER_GIB, COMPARISON_IOURING_MEM +from benchmarks.utils.constants import ( + BYTES_PER_GIB, + COMPARISON_IOURING_MEM, + slot_size_for_block_tokens, +) from benchmarks.utils.datasets import BenchmarkSample, ImdbDataset, LongBenchDataset from benchmarks.utils.loadgen import ( PhaseResult, @@ -107,6 +111,23 @@ def apply_chat_template( ) +def _write_model_config( + path: Path, + payload: dict[str, object] | None = None, +) -> str: + if payload is None: + payload = { + "hidden_size": 4096, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "num_hidden_layers": 36, + "torch_dtype": "bfloat16", + } + path.mkdir() + (path / "config.json").write_text(json.dumps(payload), encoding="utf-8") + return str(path) + + def test_longbench_dataset_loads_samples_with_answers(tmp_path: Path) -> None: """LongBench loader preserves context, question, and answers.""" data = tmp_path / "2wikimqa.jsonl" @@ -1989,6 +2010,7 @@ def fake_run_command(command: list[str]) -> None: datasets="triviaqa", max_samples=1, block_size=128, + tensor_parallel_size=2, ) ) @@ -2031,6 +2053,10 @@ def fake_run_command(command: list[str]) -> None: assert block_size_commands for command in block_size_commands: assert command[command.index("--block-size") + 1] == "128" + prepare_command = next( + command for command in commands if "--prepare-only" in command + ) + assert prepare_command[prepare_command.index("--tensor-parallel-size") + 1] == "2" start_commands = [ command for command in commands @@ -2107,9 +2133,10 @@ def test_run_benchmark_validates_direct_args_before_creating_run_dir( def test_vllm_bench_prepare_config_uses_synthetic_lengths(tmp_path: Path) -> None: """Synthetic vLLM bench sizing uses configured input length and block size.""" + model = _write_model_config(tmp_path / "model") args = RunBenchArgs( backend="baseline,lmcache,daser-prefix", - model="/models/qwen", + model=model, store_dir=str(tmp_path), load_generator="vllm-bench", block_size=128, @@ -2136,9 +2163,10 @@ def test_vllm_bench_prefix_prepare_config_derives_prefix_ratio( tmp_path: Path, ) -> None: """Prefix benchmark sizing treats bench_input_len as total input length.""" + model = _write_model_config(tmp_path / "model") args = RunBenchArgs( backend="baseline,lmcache,daser-prefix", - model="/models/qwen", + model=model, store_dir=str(tmp_path), load_generator="vllm-bench-prefix", block_size=128, @@ -2160,6 +2188,25 @@ def test_vllm_bench_prefix_prepare_config_derives_prefix_ratio( assert config["total_blocks"] == 80 +def test_benchmark_slot_size_uses_glm_geometry_and_tp_replication( + tmp_path: Path, +) -> None: + """Benchmark sizing shares DaseR's model-aware TP geometry.""" + model = _write_model_config( + tmp_path / "glm", + { + "hidden_size": 4096, + "num_attention_heads": 32, + "multi_query_group_num": 4, + "kv_channels": 128, + "num_hidden_layers": 40, + "torch_dtype": "bfloat16", + }, + ) + + assert slot_size_for_block_tokens(model, 16, 2) == 1_310_720 + + def test_vllm_bench_command_uses_random_dataset(tmp_path: Path) -> None: """vLLM bench commands target completions with deterministic random load.""" args = RunBenchArgs( @@ -2176,6 +2223,7 @@ def test_vllm_bench_command_uses_random_dataset(tmp_path: Path) -> None: bench_burstiness=2.0, bench_random_prefix_len=256, bench_random_range_ratio=0.25, + trust_remote_code=True, ) raw_path = tmp_path / "raw.json" @@ -2189,6 +2237,7 @@ def test_vllm_bench_command_uses_random_dataset(tmp_path: Path) -> None: assert command[command.index("--backend") + 1] == "openai" assert command[command.index("--base-url") + 1] == "http://127.0.0.1:8001" assert command[command.index("--endpoint") + 1] == "/v1/completions" + assert "--trust-remote-code" in command assert command[command.index("--dataset-name") + 1] == "random" assert command[command.index("--num-prompts") + 1] == "12" assert command[command.index("--input-len") + 1] == "4096" @@ -2423,6 +2472,7 @@ def test_run_bench_vllm_bench_entrypoint_runs_openai_rows( capsys, ) -> None: """vLLM bench mode starts OpenAI-compatible rows and writes a summary.""" + model = _write_model_config(tmp_path / "model") run_root = tmp_path / "run_20260102_030405" commands: list[list[str]] = [] lmcache_waits: list[str] = [] @@ -2552,7 +2602,7 @@ async def fake_collect_phase_metrics( RunBenchArgs( backend="baseline,lmcache,daser-prefix", load_generator="vllm-bench", - model="/models/qwen", + model=model, store_dir=str(tmp_path), block_size=128, bench_num_prompts=3, @@ -2598,6 +2648,7 @@ def test_run_bench_vllm_prefix_entrypoint_runs_single_phase_against_baseline( capsys, ) -> None: """Prefix bench mode runs one phase per row and compares with baseline.""" + model = _write_model_config(tmp_path / "model") run_root = tmp_path / "run_20260102_030405" commands: list[list[str]] = [] result_names: list[str] = [] @@ -2708,7 +2759,7 @@ async def fake_collect_phase_metrics( RunBenchArgs( backend="baseline,lmcache,daser-prefix", load_generator="vllm-bench-prefix", - model="/models/qwen", + model=model, store_dir=str(tmp_path), block_size=128, bench_num_prompts=3, @@ -2880,6 +2931,31 @@ def test_server_commands_propagate_custom_block_size(tmp_path: Path) -> None: assert manager.manifest().block_size == 128 +def test_server_commands_propagate_tensor_parallel_size(tmp_path: Path) -> None: + """TP size reaches both vLLM workers and the DaseR storage layout.""" + manager = ServerManager( + run_id="run1", + backend="daser", + model="/models/glm", + store_dir=tmp_path, + gpu_id="0,2", + gpu_util=0.85, + max_num_seqs=1, + l1_size_bytes=1024**3, + l2_size_bytes=2 * 1024**3, + tensor_parallel_size=2, + trust_remote_code=True, + ) + + vllm_command = manager.vllm_command(None) + daser_command = manager._daser_server_command() # noqa: SLF001 + + assert vllm_command[vllm_command.index("--tensor-parallel-size") + 1] == "2" + assert daser_command[daser_command.index("--tensor-parallel-size") + 1] == "2" + assert "--trust-remote-code" in vllm_command + assert "--trust-remote-code" in daser_command + + def test_lmcache_metrics_use_http_server_endpoint() -> None: """LMCache MP metrics are exposed by the HTTP server, not port 9090.""" manifest = BenchmarkManifest(