diff --git a/benchmarks/host_weight_runtime/README.md b/benchmarks/host_weight_runtime/README.md index 9a834ff1f..899811151 100644 --- a/benchmarks/host_weight_runtime/README.md +++ b/benchmarks/host_weight_runtime/README.md @@ -2,7 +2,11 @@ `safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from access to one cached tensor view. It imports neither vLLM nor HWR and uses a -synthetic 16-element float32 payload. HWR itself calls `get_tensor()` when +synthetic float32 payload controlled by `--tensor-elements` (default: 64). +Tensor size can affect retention: a negative result for 16 elements does not +rule out growth for 64 elements. The selected size is recorded in the JSON +arguments and used for both file creation and the correctness check. +HWR itself calls `get_tensor()` when acquiring a lease, then exposes cached views; this probe does not measure HWR requests or repeated lease acquisition. @@ -18,14 +22,18 @@ CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ --mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json ``` -For a comparison, run each mode twice in fresh processes with the same controls: +For a size comparison, run 16 and 64 elements twice each in fresh processes, +with cached reuse as a control at both sizes: ```bash for repetition in 1 2; do - for mode in get_tensor reuse; do - CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ - benchmarks/host_weight_runtime/safetensors_retention.py \ - --mode "$mode" > "$mode-$repetition.json" + for elements in 16 64; do + for mode in get_tensor reuse; do + CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ + benchmarks/host_weight_runtime/safetensors_retention.py \ + --mode "$mode" --tensor-elements "$elements" \ + > "$mode-$elements-$repetition.json" + done done done ``` diff --git a/benchmarks/host_weight_runtime/safetensors_retention.py b/benchmarks/host_weight_runtime/safetensors_retention.py index bcdb6427a..55795e7ef 100644 --- a/benchmarks/host_weight_runtime/safetensors_retention.py +++ b/benchmarks/host_weight_runtime/safetensors_retention.py @@ -43,6 +43,9 @@ def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True) + parser.add_argument( + "--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor" + ) parser.add_argument("--iterations", type=positive_int, default=1_000_000) parser.add_argument("--sample-every", type=positive_int, default=100_000) parser.add_argument("--warmup", type=positive_int, default=1000) @@ -61,10 +64,10 @@ def main() -> None: samples = [sample("before_file", 0, 0.0)] with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory: path = Path(directory) / "synthetic.safetensors" - save_file({"weight": torch.arange(16, dtype=torch.float32)}, str(path)) + save_file({"weight": torch.arange(args.tensor_elements, dtype=torch.float32)}, str(path)) with safe_open(path, framework="pt", device="cpu") as reader: cached = reader.get_tensor("weight") - assert torch.equal(cached, torch.arange(16, dtype=torch.float32)) + assert torch.equal(cached, torch.arange(args.tensor_elements, dtype=torch.float32)) for _ in range(args.warmup): tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached del tensor