[Misc] Make retention probe tensor size configurable

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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
 ```

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