diff --git a/benchmarks/single_node/agentic/qwen3small_bf16_h100.sh b/benchmarks/single_node/agentic/qwen3small_bf16_h100.sh new file mode 100755 index 0000000000..86fbcb60a0 --- /dev/null +++ b/benchmarks/single_node/agentic/qwen3small_bf16_h100.sh @@ -0,0 +1,210 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +# Small AgentX integration lane for vLLM cache-source Prometheus metrics. +# This is intentionally a plumbing/observability smoke, not a publishable +# performance submission. + +source "$(dirname "$0")/../../benchmark_lib.sh" + +check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION PORT EVAL_ONLY + +if [[ "$TP" != "1" ]]; then + echo "Error: qwen3small cache-source smoke supports TP=1 only" >&2 + exit 1 +fi + +if [[ -n "${MODEL_PATH:-}" ]]; then + if [[ ! -d "$MODEL_PATH" || -z "$(ls -A "$MODEL_PATH" 2>/dev/null)" ]]; then + hf download "$MODEL" --local-dir "$MODEL_PATH" + fi +else + hf download "$MODEL" + MODEL_PATH="$MODEL" +fi + +install_agentic_deps +# The canonical AgentX corpus deliberately contains only long-context traces; +# its 256k-capped variant has no trajectories that fit Qwen3-0.6B's native +# 40,960-token window. This plumbing smoke therefore uses a checked-in Weka +# trajectory with the same growing-prefix shape. AgentX marks the result as an +# unsafe/non-submission run because the fixture is local and intentionally tiny. +export TRACE_SOURCE_FLAG="--input-file /workspace/utils/agentic/fixtures/vllm_cache_source_weka --custom-dataset-type weka_trace" +export AIPERF_UNSAFE_OVERRIDE=true + +SERVER_LOG="$RESULT_DIR/server.log" +mkdir -p "$RESULT_DIR" + +SERVER_PID="" +cleanup_services() { + local exit_code=$? + trap - EXIT INT TERM + set +e + stop_background_process_tree "$SERVER_PID" "vLLM server" 60 + exit "$exit_code" +} +trap cleanup_services EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +OFFLOAD_ARGS=() +GPU_MEMORY_UTILIZATION=0.80 +EXPECTED_CACHE_SOURCE=device +if [[ "$KV_OFFLOADING" == "none" ]]; then + require_agentic_kv_offload_none +elif require_agentic_kv_offload_backend vllm-simple; then + # Keep only enough GPU KV for one native-length request. Multiple AgentX + # lanes then evict one another's prefixes into the CPU tier and exercise + # physical reload attribution, rather than reporting device hits only. + GPU_MEMORY_UTILIZATION=0.10 + EXPECTED_CACHE_SOURCE=cpu + CPU_OFFLOAD_BYTES=$((TOTAL_CPU_DRAM_GB * 1000 * 1000 * 1000)) + OFFLOAD_CONFIG=$(printf \ + '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"kv_offload_backend":"cpu","cpu_bytes_to_use":%d,"lazy_offload":false}}' \ + "$CPU_OFFLOAD_BYTES") + OFFLOAD_ARGS=(--kv-transfer-config "$OFFLOAD_CONFIG") +else + echo "Error: unsupported KV offload backend: ${KV_OFFLOAD_BACKEND:-unset}" >&2 + exit 1 +fi + +# Match the benchmark request URL's hostname so AIPerf de-duplicates its +# auto-discovered endpoint and this explicit metrics endpoint. +export AIPERF_SERVER_METRICS_URLS="http://localhost:${PORT}/metrics" +export AIPERF_REQUIRED_SERVER_METRIC_PREFIX="vllm:" +export PYTHONNOUSERSITE=1 +export VLLM_ENABLE_CUDA_COMPATIBILITY=1 +# Qwen3-0.6B's native context is 40,960 tokens. Use the same limit for vLLM +# and AgentX trace selection so an oversized warmup request cannot reach CUDA. +export MAX_MODEL_LEN=40960 + +VLLM_CMD=( + vllm serve "$MODEL_PATH" + --served-model-name "$MODEL" + --host 0.0.0.0 + --port "$PORT" + --tensor-parallel-size 1 + --max-model-len "$MAX_MODEL_LEN" + --max-num-seqs 8 + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" + --enable-prefix-caching + "${OFFLOAD_ARGS[@]}" +) +printf '%q ' "${VLLM_CMD[@]}" | tee "$RESULT_DIR/vllm_command.txt" +printf '\n' | tee -a "$RESULT_DIR/vllm_command.txt" +"${VLLM_CMD[@]}" >"$SERVER_LOG" 2>&1 & +SERVER_PID=$! + +wait_for_server_ready --port "$PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" + +# Scrapers establish a counter baseline from their first observation. Ensure +# every built-in physical tier exists at zero before warmup traffic, otherwise +# tokens served before a newly labelled series first appears are lost from the +# exported delta. +python3 - "http://127.0.0.1:${PORT}" <<'PY' +import math +import sys + +from utils.validate_vllm_cache_source_metrics import snapshot + +builtins = {"device", "cpu", "disk", "mixed", "external"} +observed = snapshot(sys.argv[1]).cached_by_source +if set(observed) != builtins: + raise SystemExit( + f"startup cache-source labels differ: expected {sorted(builtins)}, " + f"observed {observed}" + ) +nonzero = {source: value for source, value in observed.items() if not math.isclose(value, 0)} +if nonzero: + raise SystemExit(f"startup cache-source counters are not zero: {nonzero}") +print(f"validated startup cache-source series: {observed}") +PY + +if [[ "$EVAL_ONLY" == "true" ]]; then + run_eval --port "$PORT" +else + build_replay_cmd "$RESULT_DIR" + run_agentic_replay_and_write_outputs "$RESULT_DIR" + + python3 - "$RESULT_DIR/aiperf_artifacts/server_metrics_export.json" \ + "$EXPECTED_CACHE_SOURCE" "$AIPERF_SERVER_METRICS_URLS" <<'PY' +import json +import math +import sys + +path = sys.argv[1] +expected_source = sys.argv[2] +expected_endpoint = sys.argv[3] +with open(path) as file: + metrics = json.load(file).get("metrics", {}) + +def series_totals(name, label=None): + entry = metrics.get(name) + if not isinstance(entry, dict): + raise SystemExit(f"missing {name} export") + totals = {} + for series in entry.get("series", []): + value = series.get("stats", {}).get("total") + if value is None: + continue + key = series.get("labels", {}).get(label) if label else "total" + if key is not None: + totals[key] = totals.get(key, 0.0) + float(value) + return totals + + +cached_total = sum(series_totals("vllm:prompt_tokens_cached").values()) +physical = series_totals("vllm:prompt_tokens_cached_by_source", "source") +logical = series_totals("vllm:prompt_tokens_by_source", "source") + +physical_entry = metrics["vllm:prompt_tokens_cached_by_source"] +endpoints = { + series.get("endpoint_url") + for series in physical_entry.get("series", []) + if series.get("endpoint_url") is not None +} +if endpoints != {expected_endpoint}: + raise SystemExit( + f"cache-source export contains duplicate or unexpected endpoints: {endpoints}" + ) + +builtins = {"device", "cpu", "disk", "mixed", "external"} +if set(physical) != builtins: + raise SystemExit( + f"exported cache-source labels differ: expected {sorted(builtins)}, " + f"observed {physical}" + ) +unexpected_positive = { + source: value + for source, value in physical.items() + if source not in {"device", "cpu"} and not math.isclose(value, 0) +} +if unexpected_positive: + raise SystemExit(f"unexpected positive cache-source totals: {unexpected_positive}") +if cached_total <= 0: + raise SystemExit(f"no cached prompt tokens were exported: {cached_total}") +if physical.get(expected_source, 0) <= 0: + raise SystemExit( + f"expected positive {expected_source!r} cached-token samples: {physical}" + ) +physical_total = sum(physical.values()) +logical_total = sum( + logical.get(source, 0) + for source in ("local_cache_hit", "external_kv_transfer") +) +for description, observed in ( + ("physical cache-source", physical_total), + ("logical cache-source", logical_total), +): + if not math.isclose(observed, cached_total, rel_tol=0, abs_tol=0.5): + raise SystemExit( + f"{description} total does not conserve cached tokens: " + f"observed={observed}, cached={cached_total}" + ) +print( + "validated AgentX cached-token conservation: " + f"cached={cached_total}, physical={physical}, logical={logical}" +) +PY +fi diff --git a/configs/nvidia-master.yaml b/configs/nvidia-master.yaml index 438a3e9da0..053f4d2a23 100644 --- a/configs/nvidia-master.yaml +++ b/configs/nvidia-master.yaml @@ -7180,6 +7180,26 @@ dsv4-fp4-gb300-dynamo-vllm-agentic: ep: 8 dp-attn: true +# Short-lived integration lane for the vLLM physical cache-source metric PR. +# It deliberately uses a small model and one GPU: this validates AgentX metric +# collection/aggregation for GPU-resident and native DRAM-offload sources, not +# model performance. Remove the temporary image once the vLLM PR has a durable +# published build. +qwen3small-bf16-h100-vllm-agentic-cache-source: + image: ttl.sh/cquil11-vllm-tier-dfb25dbef4-20260828:24h + model: Qwen/Qwen3-0.6B + model-prefix: qwen3small + runner: cluster:h100-dgxc + precision: bf16 + framework: vllm + multinode: false + scenarios: + agentic-coding: + - dram-utilization: 0.05 + search-space: + - { tp: 1, kv-offloading: none, conc-list: [1] } + - { tp: 1, kv-offloading: dram, kv-offload-backend: { name: vllm-simple }, conc-list: [4] } + qwen3.5-fp8-h100-sglang-agentic: image: lmsysorg/sglang:v0.5.12-cu130 model: Qwen/Qwen3.5-397B-A17B-FP8 diff --git a/perf-changelog.yaml b/perf-changelog.yaml index ab92e7921c..5053238a2e 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -6576,3 +6576,14 @@ - "Recipes sourced from srt-slurm (recipes/trtllm/qwen3.5-fp4/inferencex/gb300/{mtp,stp})." - "Runner: launch_gb300-nv.sh bumped from NVIDIA/srt-slurm@v1.0.29 to v1.0.72 for the dynamo-trt+qwen3.5+fp4 path." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2730 + +- config-keys: + - qwen3small-bf16-h100-vllm-agentic-cache-source + scenario-type: + - agentic-coding + description: + - "Add a two-point H100 AgentX integration smoke for vLLM physical cached-token source metrics: GPU-resident KV and 13 GB native CPU KV offload at TP1/concurrency 1." + - "Ingest vllm:prompt_tokens_cached_by_source into server_metrics.cache.cached_tokens_by_source and aggregate labels across metrics endpoints and engine ranks." + - "Fail the smoke when the source metric is missing or non-positive, or when the device/CPU test exposes an unexpected source label." + - "Add reusable H100/H200 validators for device, CPU, native disk, generic external connector, and real NIXL producer/consumer paths." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2766 diff --git a/utils/agentic/aggregation/backends/vllm.py b/utils/agentic/aggregation/backends/vllm.py index 18ed30ece8..03a389f520 100644 --- a/utils/agentic/aggregation/backends/vllm.py +++ b/utils/agentic/aggregation/backends/vllm.py @@ -67,6 +67,13 @@ def populate( flat["server_external_cache_hit_rate"] = external_rate flat["server_cpu_cache_hit_rate"] = external_rate + cached_by_tier_source = sum_by_label( + metrics, + "vllm:prompt_tokens_cached_by_source", + "source", + preferred_keys=("total", "sum", "max", "avg"), + ) + prompt_by_source = sum_by_label( metrics, "vllm:prompt_tokens_by_source", @@ -142,6 +149,7 @@ def populate( "prefix_cache_queries": prefix_queries, "external_prefix_cache_hits": external_hits, "external_prefix_cache_queries": external_queries, + "cached_tokens_by_source": cached_by_tier_source, } ) nested["kv_cache"].update( @@ -258,6 +266,7 @@ def _vllm_sources(metrics: dict[str, dict[str, Any]]) -> list[dict[str, Any]]: "vllm:prefix_cache_queries", "vllm:kv_cache_usage_perc", "vllm:prompt_tokens_by_source", + "vllm:prompt_tokens_cached_by_source", ): for series in metric_series(metrics, metric_name): source_ids.add(_source_id(series)) diff --git a/utils/agentic/aggregation/test_process_agentic_result.py b/utils/agentic/aggregation/test_process_agentic_result.py index a567e9aced..63d1c08408 100644 --- a/utils/agentic/aggregation/test_process_agentic_result.py +++ b/utils/agentic/aggregation/test_process_agentic_result.py @@ -1180,6 +1180,51 @@ def test_processor_aggregates_across_multiple_series(tmp_path: Path): assert agg["server_metrics"]["cache"]["gpu_cache_hit_rate"] == pytest.approx(0.3) +def test_processor_surfaces_vllm_cached_tokens_by_physical_tier(tmp_path: Path): + result_dir = _write_fixture(tmp_path) + artifact = result_dir / "aiperf_artifacts" + server_metrics = { + "metrics": { + "vllm:prompt_tokens_cached_by_source": { + "type": "counter", + "series": [ + { + "endpoint_url": "http://prefill-0:8000/metrics", + "labels": {"source": "device", "engine": "0"}, + "stats": {"total": 400.0}, + }, + { + "endpoint_url": "http://prefill-0:8000/metrics", + "labels": {"source": "cpu", "engine": "0"}, + "stats": {"total": 100.0}, + }, + { + "endpoint_url": "http://prefill-1:8000/metrics", + "labels": {"source": "device", "engine": "1"}, + "stats": {"total": 250.0}, + }, + { + "endpoint_url": "http://prefill-1:8000/metrics", + "labels": {"source": "disk", "engine": "1"}, + "stats": {"total": 75.0}, + }, + ], + } + } + } + with open(artifact / "server_metrics_export.json", "w") as f: + json.dump(server_metrics, f) + + agg = _run_processor(result_dir, tmp_path / "out") + + assert agg["server_metrics"]["cache"]["cached_tokens_by_source"] == { + "device": 650.0, + "cpu": 100.0, + "disk": 75.0, + } + assert len(agg["server_metrics"]["sources"]) == 2 + + def test_processor_surfaces_vllm_kv_offload_transfer_stats(tmp_path: Path): result_dir = _write_fixture(tmp_path) artifact = result_dir / "aiperf_artifacts" diff --git a/utils/agentic/fixtures/vllm_cache_source_weka/growing_prefix.json b/utils/agentic/fixtures/vllm_cache_source_weka/growing_prefix.json new file mode 100644 index 0000000000..333deac4be --- /dev/null +++ b/utils/agentic/fixtures/vllm_cache_source_weka/growing_prefix.json @@ -0,0 +1,19 @@ +{ + "id": "vllm_cache_source_growing_prefix", + "models": ["Qwen/Qwen3-0.6B"], + "block_size": 1024, + "hash_id_scope": "local", + "requests": [ + {"t": 0.0, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 8192, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.1, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 11264, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.2, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 14336, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.3, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 17408, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.4, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 20480, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.5, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 23552, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.6, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 26624, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.7, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 29696, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.8, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 32768, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 0.9, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 35840, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0}, + {"t": 1.0, "type": "n", "model": "Qwen/Qwen3-0.6B", "in": 38912, "out": 16, "hash_ids": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38], "input_types": ["text"], "output_types": ["text"], "stop": "end_turn", "api_time": 0.1, "think_time": 0.0} + ] +} diff --git a/utils/run_vllm_cache_source_validation.sh b/utils/run_vllm_cache_source_validation.sh new file mode 100755 index 0000000000..01d737bd79 --- /dev/null +++ b/utils/run_vllm_cache_source_validation.sh @@ -0,0 +1,139 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Run the cache-source metric validator against a one-GPU vLLM Docker server. +# Required: IMAGE. Optional: MODE=device|cpu|disk|external, HOST_HF_CACHE, +# DISK_ROOT, PORT. + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +IMAGE=${IMAGE:?Set IMAGE to the vLLM image under test} +MODE=${MODE:-device} +PORT=${PORT:-18000} +MODEL=${MODEL:-Qwen/Qwen3.5-0.8B} +PROMPT_MODE=${PROMPT_MODE:-token_ids} +HOST_HF_CACHE=${HOST_HF_CACHE:-/models/gharunners/hf-hub-cache} +DISK_ROOT=${DISK_ROOT:-/tmp} +RUN_KEY=${SLURM_JOB_ID:-$$} +VISIBLE_GPU_LIST=${CUDA_VISIBLE_DEVICES:-0} +GPU_DEVICE=${DOCKER_GPU_DEVICE:-${VISIBLE_GPU_LIST%%,*}} +GPU_DEVICE=${GPU_DEVICE:-0} +CONTAINER="vllm-cache-source-${MODE}-${RUN_KEY}" +SERVER_LOG="${TMPDIR:-/tmp}/${CONTAINER}.log" +DISK_DIR="" + +cleanup() { + local exit_code=$? + trap - EXIT INT TERM + if docker inspect "$CONTAINER" >/dev/null 2>&1; then + docker logs "$CONTAINER" >"$SERVER_LOG" 2>&1 || true + fi + docker rm -f "$CONTAINER" >/dev/null 2>&1 || true + if [[ -n "$DISK_DIR" && -d "$DISK_DIR" ]]; then + docker run --rm --volume "$DISK_DIR:/cleanup" --entrypoint /bin/sh \ + "$IMAGE" -c 'rm -rf /cleanup/* /cleanup/.[!.]* /cleanup/..?*' \ + >/dev/null 2>&1 || true + rmdir -- "$DISK_DIR" 2>/dev/null || true + fi + exit "$exit_code" +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +if [[ ! -d "$HOST_HF_CACHE" ]]; then + echo "Hugging Face cache does not exist: $HOST_HF_CACHE" >&2 + exit 1 +fi + +if ! docker image inspect "$IMAGE" >/dev/null 2>&1; then + docker pull "$IMAGE" +fi + +DOCKER_ARGS=( + --detach + --name "$CONTAINER" + --gpus "device=$GPU_DEVICE" + --network host + --ipc host + --env VLLM_ENABLE_CUDA_COMPATIBILITY=1 + --env VLLM_SERVER_DEV_MODE=1 + --volume "$HOST_HF_CACHE:/root/.cache/huggingface/hub" +) +SERVER_ARGS=( + "$MODEL" + --host 127.0.0.1 + --port "$PORT" + --language-model-only + --enable-prefix-caching + --enforce-eager + --max-model-len 4096 + --max-num-seqs 4 + --kv-cache-memory-bytes 1073741824 +) +EXPECTED_TIER=() + +case "$MODE" in + device) + ;; + cpu) + SERVER_ARGS+=( + --kv-transfer-config + '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"kv_offload_backend":"cpu","cpu_bytes_to_use":536870912,"lazy_offload":false}}' + ) + EXPECTED_TIER=(--expected-tier cpu) + ;; + disk) + DISK_DIR=$(mktemp -d "$DISK_ROOT/vllm-cache-source.XXXXXX") + DOCKER_ARGS+=(--volume "$DISK_DIR:/kv-offload") + SERVER_ARGS+=( + --kv-transfer-config + '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"kv_offload_backend":"disk","disk_path":"/kv-offload/cache.bin","disk_capacity_bytes":536870912,"disk_buffer_slots":4,"lazy_offload":false}}' + ) + EXPECTED_TIER=(--expected-tier disk) + ;; + external) + DISK_DIR=$(mktemp -d "$DISK_ROOT/vllm-cache-source.XXXXXX") + DOCKER_ARGS+=(--volume "$DISK_DIR:/external-kv") + SERVER_ARGS+=( + --kv-transfer-config + '{"kv_connector":"ExampleConnector","kv_role":"kv_both","kv_connector_extra_config":{"shared_storage_path":"/external-kv"}}' + ) + EXPECTED_TIER=(--expected-tier external) + ;; + *) + echo "Unsupported MODE=$MODE; expected device, cpu, disk, or external" >&2 + exit 1 + ;; +esac + +docker run "${DOCKER_ARGS[@]}" "$IMAGE" "${SERVER_ARGS[@]}" >/dev/null + +ready=false +for _ in $(seq 1 180); do + if curl --fail --silent "http://127.0.0.1:${PORT}/health" >/dev/null; then + ready=true + break + fi + if ! docker inspect --format '{{.State.Running}}' "$CONTAINER" 2>/dev/null | grep -qx true; then + echo "vLLM container stopped before becoming ready" >&2 + docker logs "$CONTAINER" >&2 || true + exit 1 + fi + sleep 2 +done +if [[ "$ready" != true ]]; then + echo "vLLM server did not become ready" >&2 + docker logs "$CONTAINER" >&2 || true + exit 1 +fi + +python3 "$SCRIPT_DIR/validate_vllm_cache_source_metrics.py" \ + --base-url "http://127.0.0.1:${PORT}" \ + --prompt-mode "$PROMPT_MODE" \ + "${EXPECTED_TIER[@]}" + +curl --fail --silent "http://127.0.0.1:${PORT}/metrics" \ + | grep -E '^vllm:prompt_tokens_(cached|cached_by_source|by_source)' \ + | tee "${SERVER_LOG%.log}.metrics" +docker logs "$CONTAINER" >"$SERVER_LOG" 2>&1 +echo "server log: $SERVER_LOG" diff --git a/utils/run_vllm_cache_source_validation_enroot.sh b/utils/run_vllm_cache_source_validation_enroot.sh new file mode 100755 index 0000000000..6a511a4391 --- /dev/null +++ b/utils/run_vllm_cache_source_validation_enroot.sh @@ -0,0 +1,106 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Run inside a Slurm/Pyxis or Enroot container on clusters without Docker. + +MODE=${MODE:-device} +PORT=${PORT:-18000} +MODEL=${MODEL:-Qwen/Qwen3-0.6B} +PROMPT_MODE=${PROMPT_MODE:-token_ids} +VALIDATION_DIR=${VALIDATION_DIR:-/validation} +RESULT_DIR=${RESULT_DIR:-$VALIDATION_DIR/results} +RUN_KEY=${SLURM_JOB_ID:-$$} +DISK_ROOT=${DISK_ROOT:-/mnt/numa0/enroot/runtime/user-${UID}} +SERVER_LOG="$RESULT_DIR/${MODE}-${RUN_KEY}.server.log" +METRICS_LOG="$RESULT_DIR/${MODE}-${RUN_KEY}.metrics" +DISK_DIR="" +SERVER_PID="" +mkdir -p "$RESULT_DIR" + +cleanup() { + local exit_code=$? + trap - EXIT INT TERM + if [[ -n "$SERVER_PID" ]]; then + kill "$SERVER_PID" >/dev/null 2>&1 || true + wait "$SERVER_PID" >/dev/null 2>&1 || true + fi + if [[ -n "$DISK_DIR" && -d "$DISK_DIR" ]]; then + rm -rf -- "$DISK_DIR" + fi + exit "$exit_code" +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +SERVER_ARGS=( + serve "$MODEL" + --host 127.0.0.1 + --port "$PORT" + --language-model-only + --enable-prefix-caching + --enforce-eager + --max-model-len 4096 + --max-num-seqs 4 + --kv-cache-memory-bytes 1073741824 +) +EXPECTED_TIER=() + +case "$MODE" in + device) + ;; + cpu) + SERVER_ARGS+=( + --kv-transfer-config + '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"kv_offload_backend":"cpu","cpu_bytes_to_use":536870912,"lazy_offload":false}}' + ) + EXPECTED_TIER=(--expected-tier cpu) + ;; + disk) + DISK_DIR=$(mktemp -d "$DISK_ROOT/vllm-cache-source.XXXXXX") + SERVER_ARGS+=( + --kv-transfer-config + "{\"kv_connector\":\"SimpleCPUOffloadConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"kv_offload_backend\":\"disk\",\"disk_path\":\"$DISK_DIR/cache.bin\",\"disk_capacity_bytes\":536870912,\"disk_buffer_slots\":4,\"lazy_offload\":false}}" + ) + EXPECTED_TIER=(--expected-tier disk) + ;; + *) + echo "Unsupported MODE=$MODE; expected device, cpu, or disk" >&2 + exit 1 + ;; +esac + +export VLLM_ENABLE_CUDA_COMPATIBILITY=1 +export VLLM_SERVER_DEV_MODE=1 +vllm "${SERVER_ARGS[@]}" >"$SERVER_LOG" 2>&1 & +SERVER_PID=$! + +ready=false +for _ in $(seq 1 180); do + if curl --fail --silent "http://127.0.0.1:${PORT}/health" >/dev/null; then + ready=true + break + fi + if ! kill -0 "$SERVER_PID" >/dev/null 2>&1; then + echo "vLLM server stopped before becoming ready" >&2 + tail -n 200 "$SERVER_LOG" >&2 || true + exit 1 + fi + sleep 2 +done +if [[ "$ready" != true ]]; then + echo "vLLM server did not become ready" >&2 + tail -n 200 "$SERVER_LOG" >&2 || true + exit 1 +fi + +python3 "$VALIDATION_DIR/validate_vllm_cache_source_metrics.py" \ + --base-url "http://127.0.0.1:${PORT}" \ + --prompt-mode "$PROMPT_MODE" \ + "${EXPECTED_TIER[@]}" + +curl --fail --silent "http://127.0.0.1:${PORT}/metrics" \ + | grep -E '^vllm:prompt_tokens_(cached|cached_by_source|by_source)' \ + | tee "$METRICS_LOG" +echo "server log: $SERVER_LOG" +echo "metrics: $METRICS_LOG" diff --git a/utils/run_vllm_nixl_source_validation.sh b/utils/run_vllm_nixl_source_validation.sh new file mode 100755 index 0000000000..942c862ed0 --- /dev/null +++ b/utils/run_vllm_nixl_source_validation.sh @@ -0,0 +1,147 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Run inside a two-GPU vLLM container. The upstream toy proxy must be mounted +# at NIXL_INTEGRATION_DIR and the cache-source validator at VALIDATION_DIR. + +MODEL=${MODEL:-Qwen/Qwen3-0.6B} +VALIDATION_DIR=${VALIDATION_DIR:-/validation} +NIXL_INTEGRATION_DIR=${NIXL_INTEGRATION_DIR:-/nixl-integration} +RESULT_DIR=${RESULT_DIR:-$VALIDATION_DIR/results} +RUN_KEY=${RUN_KEY:-${SLURM_JOB_ID:-$$}} +PREFILL_PORT=${PREFILL_PORT:-18100} +DECODE_PORT=${DECODE_PORT:-18200} +PROXY_PORT=${PROXY_PORT:-18192} +PREFILL_INTERNAL_PORT=${PREFILL_INTERNAL_PORT:-28100} +DECODE_INTERNAL_PORT=${DECODE_INTERNAL_PORT:-28200} +PREFILL_SIDE_PORT=${PREFILL_SIDE_PORT:-15559} +DECODE_SIDE_PORT=${DECODE_SIDE_PORT:-15659} +PREFILL_LOG="$RESULT_DIR/nixl-${RUN_KEY}.prefill.log" +DECODE_LOG="$RESULT_DIR/nixl-${RUN_KEY}.decode.log" +PROXY_LOG="$RESULT_DIR/nixl-${RUN_KEY}.proxy.log" +METRICS_LOG="$RESULT_DIR/nixl-${RUN_KEY}.metrics" +PIDS=() +LAST_SERVER_PID="" +mkdir -p "$RESULT_DIR" + +cleanup() { + local exit_code=$? + trap - EXIT INT TERM + if ((${#PIDS[@]})); then + kill "${PIDS[@]}" >/dev/null 2>&1 || true + wait "${PIDS[@]}" >/dev/null 2>&1 || true + fi + exit "$exit_code" +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +start_server() { + local gpu=$1 + local port=$2 + local internal_port=$3 + local side_port=$4 + local role=$5 + local log=$6 + CUDA_VISIBLE_DEVICES="$gpu" \ + VLLM_PORT="$internal_port" \ + VLLM_NIXL_SIDE_CHANNEL_PORT="$side_port" \ + VLLM_ENABLE_CUDA_COMPATIBILITY=1 \ + VLLM_SERVER_DEV_MODE=1 \ + VLLM_SSM_CONV_STATE_LAYOUT=DS \ + UCX_NET_DEVICES=all \ + vllm serve "$MODEL" \ + --host 127.0.0.1 \ + --port "$port" \ + --language-model-only \ + --block-size 128 \ + --gpu-memory-utilization 0.2 \ + --enforce-eager \ + --max-model-len 4096 \ + --kv-transfer-config \ + "{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"$role\"}" \ + >"$log" 2>&1 & + LAST_SERVER_PID=$! + PIDS+=("$LAST_SERVER_PID") +} + +wait_for_health() { + local port=$1 + local log=$2 + local pid=$3 + for _ in $(seq 1 300); do + if curl --fail --silent "http://127.0.0.1:${port}/health" >/dev/null; then + return + fi + if ! kill -0 "$pid" >/dev/null 2>&1; then + echo "server on port $port stopped before becoming ready" >&2 + tail -n 200 "$log" >&2 || true + return 1 + fi + sleep 2 + done + echo "server on port $port did not become ready" >&2 + tail -n 200 "$log" >&2 || true + return 1 +} + +start_server 0 "$PREFILL_PORT" "$PREFILL_INTERNAL_PORT" "$PREFILL_SIDE_PORT" \ + kv_producer "$PREFILL_LOG" +PREFILL_PID=$LAST_SERVER_PID +start_server 1 "$DECODE_PORT" "$DECODE_INTERNAL_PORT" "$DECODE_SIDE_PORT" \ + kv_consumer "$DECODE_LOG" +DECODE_PID=$LAST_SERVER_PID +wait_for_health "$PREFILL_PORT" "$PREFILL_LOG" "$PREFILL_PID" +wait_for_health "$DECODE_PORT" "$DECODE_LOG" "$DECODE_PID" + +python3 "$NIXL_INTEGRATION_DIR/toy_proxy_server.py" \ + --host 127.0.0.1 \ + --port "$PROXY_PORT" \ + --prefiller-hosts 127.0.0.1 \ + --prefiller-ports "$PREFILL_PORT" \ + --decoder-hosts 127.0.0.1 \ + --decoder-ports "$DECODE_PORT" \ + >"$PROXY_LOG" 2>&1 & +PIDS+=("$!") + +for _ in $(seq 1 60); do + if curl --fail --silent "http://127.0.0.1:${PROXY_PORT}/healthcheck" >/dev/null; then + break + fi + sleep 1 +done +curl --fail --silent "http://127.0.0.1:${PROXY_PORT}/healthcheck" >/dev/null + +PYTHONPATH="$VALIDATION_DIR${PYTHONPATH:+:$PYTHONPATH}" python3 - \ + "http://127.0.0.1:${PROXY_PORT}" \ + "http://127.0.0.1:${DECODE_PORT}" \ + "$MODEL" <<'PY' +import sys + +from validate_vllm_cache_source_metrics import ( + completion, + snapshot, + validate_delta, + wait_for_accounting, +) + +proxy_url, decoder_url, model = sys.argv[1:] +prompt = [100] * 3264 +before = snapshot(decoder_url) +completion(proxy_url, model, prompt) +after = wait_for_accounting(decoder_url, before, len(prompt), timeout=30) +delta = after.delta(before) +validate_delta("nixl_external_transfer", delta, {"external"}) +if delta.prompt_by_source.get("external_kv_transfer", 0) <= 0: + raise RuntimeError(f"NIXL request had no external KV transfer: {delta}") +PY + +curl --fail --silent "http://127.0.0.1:${DECODE_PORT}/metrics" \ + | grep -E '^vllm:prompt_tokens_(cached|cached_by_source|by_source)' \ + | tee "$METRICS_LOG" +echo "NIXL source validation passed" +echo "prefill log: $PREFILL_LOG" +echo "decode log: $DECODE_LOG" +echo "proxy log: $PROXY_LOG" +echo "metrics: $METRICS_LOG" diff --git a/utils/test_validate_vllm_cache_source_metrics.py b/utils/test_validate_vllm_cache_source_metrics.py new file mode 100644 index 0000000000..d99b2d279f --- /dev/null +++ b/utils/test_validate_vllm_cache_source_metrics.py @@ -0,0 +1,39 @@ +import pytest + +from validate_vllm_cache_source_metrics import ValidationError, _parse_labels + + +def test_parse_labels() -> None: + assert _parse_labels(None) == {} + assert _parse_labels('engine="0", source="cpu"') == { + "engine": "0", + "source": "cpu", + } + + +def test_parse_labels_unescapes_prometheus_values() -> None: + assert _parse_labels(r'value="line\nquoted\"slash\\"') == { + "value": 'line\nquoted"slash\\' + } + + +@pytest.mark.parametrize( + "labels", + [ + '1bad="value"', + 'key="unterminated', + 'key="bad\\tvalue"', + 'key="value",', + 'key="value" unexpected', + ], +) +def test_parse_labels_rejects_malformed_input(labels: str) -> None: + with pytest.raises(ValidationError): + _parse_labels(labels) + + +def test_parse_labels_handles_long_escape_sequences_linearly() -> None: + escaped_backslashes = r"\\" * 100_000 + assert _parse_labels(f'value="{escaped_backslashes}"') == { + "value": "\\" * 100_000 + } diff --git a/utils/validate_vllm_cache_source_metrics.py b/utils/validate_vllm_cache_source_metrics.py new file mode 100755 index 0000000000..85c0b04bed --- /dev/null +++ b/utils/validate_vllm_cache_source_metrics.py @@ -0,0 +1,428 @@ +#!/usr/bin/env python3 +"""Exercise and validate vLLM cached-prompt token source metrics. + +The target server must run with prefix caching and ``VLLM_SERVER_DEV_MODE=1``. +For an offload tier, use an eager SimpleCPUOffloadConnector so resetting the +local prefix cache leaves a copy available in the connector-managed cache. +""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import sys +import time +import urllib.error +import urllib.request +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + + +_SAMPLE_RE = re.compile( + r'^(?P[^\s{]+)(?:\{(?P.*)\})?\s+(?P[-+0-9.eE]+)$' +) +_CACHED_TOTAL = "vllm:prompt_tokens_cached_total" +_CACHED_BY_SOURCE = "vllm:prompt_tokens_cached_by_source_total" +_PROMPT_BY_SOURCE = "vllm:prompt_tokens_by_source_total" + + +class ValidationError(RuntimeError): + pass + + +@dataclass(frozen=True) +class Snapshot: + cached_total: float + cached_by_source: dict[str, float] + prompt_by_source: dict[str, float] + + def delta(self, previous: "Snapshot") -> "Snapshot": + return Snapshot( + cached_total=self.cached_total - previous.cached_total, + cached_by_source=_dict_delta(self.cached_by_source, previous.cached_by_source), + prompt_by_source=_dict_delta(self.prompt_by_source, previous.prompt_by_source), + ) + + +def _dict_delta(current: Mapping[str, float], previous: Mapping[str, float]) -> dict[str, float]: + return { + key: current.get(key, 0.0) - previous.get(key, 0.0) + for key in current.keys() | previous.keys() + if not math.isclose(current.get(key, 0.0), previous.get(key, 0.0)) + } + + +def _request( + base_url: str, + path: str, + *, + body: dict[str, Any] | None = None, + timeout: float = 120.0, +) -> bytes: + data = json.dumps(body).encode() if body is not None else None + request = urllib.request.Request( + f"{base_url.rstrip('/')}{path}", + data=data, + headers={"Content-Type": "application/json"} if data is not None else {}, + method="POST" if data is not None or path.startswith("/reset_") else "GET", + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + return response.read() + except urllib.error.HTTPError as error: + detail = error.read().decode(errors="replace") + raise ValidationError(f"{request.method} {path} failed: HTTP {error.code}: {detail}") from error + + +def _json_request( + base_url: str, + path: str, + *, + body: dict[str, Any] | None = None, + timeout: float = 120.0, +) -> dict[str, Any]: + value = json.loads(_request(base_url, path, body=body, timeout=timeout)) + if not isinstance(value, dict): + raise ValidationError(f"{path} returned a non-object JSON value") + return value + + +def _parse_labels(raw_labels: str | None) -> dict[str, str]: + if raw_labels is None: + return {} + + labels: dict[str, str] = {} + cursor = 0 + length = len(raw_labels) + + def fail(message: str) -> ValidationError: + return ValidationError( + f"invalid Prometheus labels at offset {cursor}: {message}: {raw_labels!r}" + ) + + while cursor < length: + while cursor < length and raw_labels[cursor] in " \t": + cursor += 1 + key_start = cursor + if cursor >= length or not (raw_labels[cursor].isalpha() or raw_labels[cursor] == "_"): + raise fail("expected label name") + cursor += 1 + while cursor < length and ( + raw_labels[cursor].isalnum() or raw_labels[cursor] == "_" + ): + cursor += 1 + key = raw_labels[key_start:cursor] + + if cursor >= length or raw_labels[cursor] != "=": + raise fail("expected '='") + cursor += 1 + if cursor >= length or raw_labels[cursor] != '"': + raise fail("expected opening quote") + cursor += 1 + + value: list[str] = [] + while cursor < length and raw_labels[cursor] != '"': + character = raw_labels[cursor] + cursor += 1 + if character != "\\": + value.append(character) + continue + if cursor >= length: + raise fail("unterminated escape") + escaped = raw_labels[cursor] + cursor += 1 + if escaped == "n": + value.append("\n") + elif escaped in {'"', "\\"}: + value.append(escaped) + else: + raise fail(f"unsupported escape {escaped!r}") + + if cursor >= length: + raise fail("unterminated quoted value") + cursor += 1 + labels[key] = "".join(value) + + while cursor < length and raw_labels[cursor] in " \t": + cursor += 1 + if cursor == length: + break + if raw_labels[cursor] != ",": + raise fail("expected ','") + cursor += 1 + if cursor == length: + raise fail("trailing ','") + + return labels + + +def snapshot(base_url: str) -> Snapshot: + values: dict[str, list[tuple[dict[str, str], float]]] = {} + text = _request(base_url, "/metrics", timeout=30.0).decode() + for line in text.splitlines(): + if not line or line.startswith("#"): + continue + match = _SAMPLE_RE.match(line) + if match is None: + continue + name = match.group("name") + if name not in {_CACHED_TOTAL, _CACHED_BY_SOURCE, _PROMPT_BY_SOURCE}: + continue + values.setdefault(name, []).append( + (_parse_labels(match.group("labels")), float(match.group("value"))) + ) + + def total(name: str) -> float: + return sum(value for _, value in values.get(name, [])) + + def by_source(name: str) -> dict[str, float]: + result: dict[str, float] = {} + for labels, value in values.get(name, []): + source = labels.get("source") + if source is not None: + result[source] = result.get(source, 0.0) + value + return result + + return Snapshot( + cached_total=total(_CACHED_TOTAL), + cached_by_source=by_source(_CACHED_BY_SOURCE), + prompt_by_source=by_source(_PROMPT_BY_SOURCE), + ) + + +def completion(base_url: str, model: str, token_ids: list[int]) -> dict[str, Any]: + response = _json_request( + base_url, + "/v1/completions", + body={ + "model": model, + "prompt": token_ids, + "max_tokens": 1, + "temperature": 0, + "ignore_eos": True, + }, + timeout=300.0, + ) + usage = response.get("usage") + if not isinstance(usage, dict) or usage.get("prompt_tokens") != len(token_ids): + raise ValidationError( + f"completion usage did not report the exact {len(token_ids)}-token prompt: {usage!r}" + ) + return response + + +def tokenized_text_prompt(base_url: str, model: str, minimum_tokens: int) -> list[int]: + seed = ( + "A cache validation sentence with ordinary words and punctuation. " + * (minimum_tokens // 4 + 1) + ) + response = _json_request( + base_url, + "/tokenize", + body={"model": model, "prompt": seed}, + timeout=120.0, + ) + tokens = response.get("tokens") + if not isinstance(tokens, list) or not all(isinstance(token, int) for token in tokens): + raise ValidationError(f"/tokenize returned invalid tokens: {type(tokens).__name__}") + if len(tokens) < minimum_tokens: + raise ValidationError( + f"tokenizer returned only {len(tokens)} tokens; need {minimum_tokens}" + ) + return tokens + + +def wait_for_accounting( + base_url: str, + previous: Snapshot, + prompt_tokens: int, + timeout: float = 15.0, +) -> Snapshot: + deadline = time.monotonic() + timeout + latest = snapshot(base_url) + while time.monotonic() < deadline: + logical_delta = sum(latest.delta(previous).prompt_by_source.values()) + if logical_delta >= prompt_tokens: + return latest + time.sleep(0.1) + latest = snapshot(base_url) + raise ValidationError( + "metrics logger did not account for the completed request within " + f"{timeout:g}s; observed logical delta {latest.delta(previous).prompt_by_source}" + ) + + +def reset_local_cache(base_url: str, attempts: int = 30) -> None: + for _ in range(attempts): + response = _json_request(base_url, "/reset_prefix_cache", body={}) + if response.get("success") is True: + return + time.sleep(1) + raise ValidationError("local prefix-cache reset never succeeded") + + +def drain_async_transfers(base_url: str, model: str, previous: Snapshot) -> Snapshot: + """Give connector workers time, then run one scheduler step to reap them.""" + time.sleep(2) + completion(base_url, model, [102]) + wait_for_accounting(base_url, previous, 1) + time.sleep(1) + return snapshot(base_url) + + +def _assert_close(actual: float, expected: float, message: str) -> None: + if not math.isclose(actual, expected, rel_tol=0, abs_tol=0.001): + raise ValidationError(f"{message}: expected {expected}, observed {actual}") + + +def validate_delta(name: str, delta: Snapshot, expected_sources: set[str]) -> None: + physical_total = sum(delta.cached_by_source.values()) + logical_total = sum(delta.prompt_by_source.values()) + _assert_close( + physical_total, + delta.cached_total, + f"{name}: physical source sum must equal cached-token counter", + ) + if logical_total: + cached_logical = ( + delta.prompt_by_source.get("local_cache_hit", 0.0) + + delta.prompt_by_source.get("external_kv_transfer", 0.0) + ) + _assert_close( + cached_logical, + delta.cached_total, + f"{name}: logical cache-hit sum must equal cached-token counter", + ) + observed_sources = { + source for source, value in delta.cached_by_source.items() if value > 0 + } + if observed_sources != expected_sources: + raise ValidationError( + f"{name}: expected positive sources {sorted(expected_sources)}, " + f"observed {delta.cached_by_source}" + ) + print(json.dumps({"case": name, "delta": delta.__dict__}, sort_keys=True)) + + +def run( + base_url: str, + expected_tier: str | None, + short_tokens: int, + long_tokens: int, + prompt_mode: str, +) -> None: + models = _json_request(base_url, "/v1/models").get("data") + if not isinstance(models, list) or not models or not isinstance(models[0], dict): + raise ValidationError("/v1/models returned no model") + model = models[0].get("id") + if not isinstance(model, str): + raise ValidationError("/v1/models returned an invalid model id") + + # The short prompt is a block-aligned prefix of the long prompt. Direct + # token IDs give the most deterministic probe; tokenizer mode provides a + # realistic-input cross-check for hybrid or multimodal model families. + if prompt_mode == "text": + token_ids = tokenized_text_prompt(base_url, model, long_tokens) + long_prompt = token_ids[:long_tokens] + short_prompt = long_prompt[:short_tokens] + else: + short_prompt = [100] * short_tokens + long_prompt = short_prompt + [101] * (long_tokens - short_tokens) + + baseline = snapshot(base_url) + completion(base_url, model, long_prompt) + after_cold = wait_for_accounting(base_url, baseline, len(long_prompt)) + validate_delta("cold_miss", after_cold.delta(baseline), set()) + + completion(base_url, model, long_prompt) + after_device = wait_for_accounting(base_url, after_cold, len(long_prompt)) + validate_delta("exact_device_hit", after_device.delta(after_cold), {"device"}) + + if expected_tier is None: + reset_local_cache(base_url) + completion(base_url, model, long_prompt) + after_reset = wait_for_accounting(base_url, after_device, len(long_prompt)) + validate_delta("cold_after_local_reset", after_reset.delta(after_device), set()) + return + + # Keep the connector-managed cache populated while dropping local block + # references. The first reload must come entirely from the second tier. + if expected_tier in {"cpu", "disk"}: + after_device = drain_async_transfers(base_url, model, after_device) + reset_local_cache(base_url) + completion(base_url, model, long_prompt) + after_tier = wait_for_accounting(base_url, after_device, len(long_prompt)) + validate_delta("exact_second_tier_hit", after_tier.delta(after_device), {expected_tier}) + + # The reference ExampleConnector addresses whole-request objects rather + # than individual prefix blocks. It validates generic KV-transfer source + # attribution, but cannot construct a partial external-prefix hit. + if expected_tier == "external": + return + + # Restore only the short prefix to the device, while the connector still + # retains the full long prompt. A long request should consume both tiers. + if expected_tier in {"cpu", "disk"}: + after_tier = drain_async_transfers(base_url, model, after_tier) + reset_local_cache(base_url) + completion(base_url, model, short_prompt) + after_short_tier = wait_for_accounting(base_url, after_tier, len(short_prompt)) + validate_delta("short_second_tier_hit", after_short_tier.delta(after_tier), {expected_tier}) + + completion(base_url, model, long_prompt) + after_mixed = wait_for_accounting(base_url, after_short_tier, len(long_prompt)) + validate_delta( + "mixed_device_and_second_tier_hit", + after_mixed.delta(after_short_tier), + {"device", expected_tier}, + ) + + allowed_sources = {"device", expected_tier} + final_sources = {source for source, value in after_mixed.cached_by_source.items() if value} + if not final_sources <= allowed_sources: + raise ValidationError( + f"unexpected source-label cardinality: {sorted(final_sources - allowed_sources)}" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--base-url", default="http://127.0.0.1:8000") + parser.add_argument( + "--expected-tier", + choices=("cpu", "disk", "external"), + help="Omit for a device-only server.", + ) + # 1088 and 3264 are multiples of both the common 16-token block and the + # 544-token hybrid page used by Qwen3.5-0.8B. + parser.add_argument("--short-tokens", type=int, default=1088) + parser.add_argument("--long-tokens", type=int, default=3264) + parser.add_argument("--prompt-mode", choices=("token_ids", "text"), default="token_ids") + args = parser.parse_args() + if args.short_tokens <= 0 or args.long_tokens <= args.short_tokens: + parser.error("require 0 < --short-tokens < --long-tokens") + return args + + +def main() -> int: + args = parse_args() + try: + run( + args.base_url, + args.expected_tier, + args.short_tokens, + args.long_tokens, + args.prompt_mode, + ) + except ValidationError as error: + print(f"validation failed: {error}", file=sys.stderr) + return 1 + print("validation passed") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())