diff --git a/benchmarks/single_node/agentic/glm5.2_fp4_mi355x_atom_mtp.sh b/benchmarks/single_node/agentic/glm5.2_fp4_mi355x_atom_mtp.sh new file mode 100644 index 0000000000..b635ae3560 --- /dev/null +++ b/benchmarks/single_node/agentic/glm5.2_fp4_mi355x_atom_mtp.sh @@ -0,0 +1,183 @@ +#!/usr/bin/env bash +set -eo pipefail +set -x + +# Agentic trace replay benchmark for GLM5.2 FP4 on MI355X using ATOM MTP +# +# Required env vars: +# MODEL, MODEL_PATH, TP, CONC, KV_OFFLOADING, KV_OFFLOAD_BACKEND, +# TOTAL_CPU_DRAM_GB, RESULT_DIR, DURATION, EP_SIZE, DP_ATTENTION + +source "$(dirname "$0")/../../benchmark_lib.sh" + +check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION + +echo "MODEL=$MODEL TP=$TP CONC=$CONC KV_OFFLOADING=$KV_OFFLOADING TOTAL_CPU_DRAM_GB=$TOTAL_CPU_DRAM_GB RESULT_DIR=$RESULT_DIR DURATION=$DURATION EP_SIZE=$EP_SIZE DP_ATTENTION=$DP_ATTENTION" + +if [[ -v SLURM_JOB_ID ]]; then + echo "JOB $SLURM_JOB_ID running on $SLURMD_NODENAME" +fi + +# ROCR/HIP visibility for vLLM 0.14+ +if [[ -v ROCR_VISIBLE_DEVICES ]]; then + export HIP_VISIBLE_DEVICES="$ROCR_VISIBLE_DEVICES" +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" + export MODEL_PATH="$MODEL" +fi + +rocm-smi || true +amd-smi || true + +resolve_trace_source +install_agentic_deps + +# Require the vLLM Prometheus stream in every official result. AIPerf +# deduplicates this endpoint against its automatic localhost discovery. +export AIPERF_SERVER_METRICS_URLS="http://localhost:${PORT}/metrics" +export AIPERF_REQUIRED_SERVER_METRIC_PREFIX="atom:" + +# VRAM space check +wait_for_amd_gpu_clean + +# ---- Server config ---------------------------------------------------------- +SERVER_LOG="$RESULT_DIR/server.log" +LMCACHE_LOG="$RESULT_DIR/lmcache_server.log" +mkdir -p "$RESULT_DIR" + +SERVER_PID="" +LMCACHE_PIDS=() +cleanup_agentic_services() { + local exit_code=$? + trap - EXIT INT TERM + set +e + stop_background_process_tree "$SERVER_PID" "ATOM server" 60 + local i + for i in "${!LMCACHE_PIDS[@]}"; do + stop_background_process_tree "${LMCACHE_PIDS[$i]}" "LMCache server $i" + done + exit "$exit_code" +} +trap cleanup_agentic_services EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +OFFLOAD_ARGS=() + +case "$KV_OFFLOAD_BACKEND" in + "") + require_agentic_kv_offload_none + ;; + lmcache) + require_agentic_kv_offload_backend lmcache + + # LMCache settings + export PYTHONHASHSEED=0 + export LMCACHE_LOCAL_CPU=True + export LMCACHE_MAX_LOCAL_CPU_SIZE="$TOTAL_CPU_DRAM_GB" + export LMCACHE_CHUNK_SIZE=256 + export OFFLOAD_MIN_LOAD_TOKENS=8192 + + OFFLOAD_ARGS=( + --kv-transfer-config + "{\"kv_connector\":\"lmcache_offload\",\"kv_role\":\"offload\"}" + ) + ;; + *) + echo "Unsupported KV_OFFLOAD_BACKEND: $KV_OFFLOAD_BACKEND (expected empty or lmcache)" >&2 + exit 1 + ;; +esac + +# ---- LLM server config ---------------------------------------------------------- + +echo "Starting atom server..." +export PYTHONNOUSERSITE=1 + +# ---- ATOM env ---- +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export AITER_USE_FLYDSL_MOE_SORTING=1 + +# CUDA/HIPGRAPH settings +case "$CONC" in + 1) CUDAGRAPH_CAPTURE_SIZES='[1,2]' ;; + 2) CUDAGRAPH_CAPTURE_SIZES='[1,2,4]' ;; + 4) CUDAGRAPH_CAPTURE_SIZES='[1,2,4,8]' ;; + 8) CUDAGRAPH_CAPTURE_SIZES='[1,2,4,8,12,16]' ;; + 10) CUDAGRAPH_CAPTURE_SIZES='[1,2,4,8,12,16,20]' ;; + 12) CUDAGRAPH_CAPTURE_SIZES='[1,2,4,8,12,16,20,24]' ;; + 16) CUDAGRAPH_CAPTURE_SIZES='[1,2,4,8,12,16,20,24,28,32]' ;; + *) + echo "Unsupported CONC=$CONC" >&2 + exit 2 + ;; +esac + +# PARALLEL settings +PARALLEL_ARGS=(--tensor-parallel-size "$TP") #TP +if [ "$DP_ATTENTION" = "true" ]; then + # DPA+EP + if [ "$EP_SIZE" -gt 1 ]; then #DP+EP + PARALLEL_ARGS=(--tensor-parallel-size "$TP" --enable-dp-attention --enable-expert-parallel) + # DPA+TP + else + PARALLEL_ARGS=(--tensor-parallel-size "$TP" --enable-dp-attention ) + fi +fi + +# SPEC settings +# SIMULATE_ACC_LEN and NUM_SPEC_TOKENS reference: +# https://github.com/SemiAnalysisAI/InferenceX/blob/main/golden_al_distribution/glm5.2_mtp.yaml +SIMULATE_ACC_LEN=2.99 +NUM_SPEC_TOKENS=3 +# spec-decode-acceptance-rate = (SIMULATE_ACC_LEN - 1) / NUM_SPEC_TOKENS +SPEC_ACCEPTANCE_RATE=$(awk "BEGIN{print ($SIMULATE_ACC_LEN-1)/$NUM_SPEC_TOKENS}") +if [ "${EVAL_ONLY}" = "true" ]; then + SPEC_ARGS=( + --method mtp + --num-speculative-tokens "$NUM_SPEC_TOKENS" + ) +else + SPEC_ARGS=( + --method mtp + --num-speculative-tokens "$NUM_SPEC_TOKENS" + --spec-decode-acceptance-rate "$SPEC_ACCEPTANCE_RATE" + ) +fi +echo "SIMULATE_ACC_LEN=$SIMULATE_ACC_LEN NUM_SPEC_TOKENS=$NUM_SPEC_TOKENS SPEC_ACCEPTANCE_RATE=$SPEC_ACCEPTANCE_RATE" + +ATOM_CMD=( + python -m atom.entrypoints.openai_server + --model "$MODEL_PATH" + --host 0.0.0.0 + --server-port "$PORT" + "${PARALLEL_ARGS[@]}" + --online_quant_config '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","*.mlp.gate","*expert*"]}' + --max-num-seqs "$((2 * CONC))" + --cudagraph-capture-sizes "$CUDAGRAPH_CAPTURE_SIZES" + --max-num-batched-tokens 16384 + --kv_cache_dtype fp8 + "${SPEC_ARGS[@]}" + "${OFFLOAD_ARGS[@]}" +) +write_command "$RESULT_DIR/server_command.txt" "${ATOM_CMD[@]}" +"${ATOM_CMD[@]}" > "$SERVER_LOG" 2>&1 & +SERVER_PID=$! +echo "Server PID: $SERVER_PID" + +wait_for_server_ready --port "$PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" + +# ---- Run benchmark ---------------------------------------------------------- +if [ "${EVAL_ONLY}" = "true" ]; then + run_eval --port "$PORT" +else + build_replay_cmd "$RESULT_DIR" + REPLAY_CMD+=" --apply-chat-template" + run_agentic_replay_and_write_outputs "$RESULT_DIR" +fi diff --git a/configs/amd-master.yaml b/configs/amd-master.yaml index 552451fc7e..60c9a60d14 100644 --- a/configs/amd-master.yaml +++ b/configs/amd-master.yaml @@ -1747,3 +1747,22 @@ glm5.2-fp4-mi355x-sglang-agentic-mtp: search-space: - { tp: 4, ep: 4, kv-offloading: dram, kv-offload-backend: { name: hicache }, conc-list: [1, 2, 4, 8, 10, 12, 16], spec-decoding: mtp } - { tp: 8, ep: 8, kv-offloading: none, conc-list: [1, 2, 4], spec-decoding: mtp } + +# GLM-5.2 FP4 agentic-coding benchmark on MI355X via ATOM with MTP speculative +# decoding. TP=4, TP=8 with KV offloading to DRAM (lmcache backend) to support +# Recipe is from PR https://github.com/ROCm/ATOM/pull/1877 +glm5.2-fp4-mi355x-atom-agentic-mtp: + image: rocm/atom-dev:ubuntu24.04_py3.12_pytorch_release_2.10.0_glm52_agentic_0813 + model: amd/GLM-5.2-MXFP4 + model-prefix: glm5.2 + runner: cluster:mi355x-amds + precision: fp4 + framework: atom + multinode: false + scenarios: + agentic-coding: + - dram-utilization: 0.14 + search-space: + - { tp: 4, kv-offloading: dram, kv-offload-backend: { name: lmcache, version: "0.4.5" }, conc-list: [2, 4, 8, 10], spec-decoding: mtp } + - { tp: 8, kv-offloading: none, conc-list: [1, 2, 4], spec-decoding: mtp } + diff --git a/perf-changelog.yaml b/perf-changelog.yaml index 04ba9feb56..8b031702c1 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -5837,6 +5837,13 @@ - "Sweep TP2/TP4 GPU-resident configurations and TP4 LMCache MP DRAM-offload points around the capacity knee, with vLLM server metrics enabled." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2487 +- config-keys: + - glm5.2-fp4-mi355x-atom-agentic-mtp + description: + - "Add the MI355X GLM-5.2 FP4 ATOM AgentX MTP submission, mirroring the upstream ATOM GLM-5.2 MXFP4 + MTP recipe with ptpc_fp8 online quantization and the lmcache_offload connector for DRAM offload." + - "Sweep TP4 (conc 2/4/8/10) and TP8 (conc 1/2/4) LMCache DRAM-offload configurations, using MTP num-speculative-tokens 3 with the synthetic acceptance rate derived from the golden AL distribution (SIMULATE_ACC_LEN 2.99)." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2576 + - config-keys: - qwen3.5-fp8-mi300x-sglang-agentic-mtp scenario-type: diff --git a/utils/agentic/aggregation/backends/__init__.py b/utils/agentic/aggregation/backends/__init__.py index 9920e54a02..6858942357 100644 --- a/utils/agentic/aggregation/backends/__init__.py +++ b/utils/agentic/aggregation/backends/__init__.py @@ -4,14 +4,15 @@ from typing import Any +from .atom import AtomBackend from .base import ServerMetricsBackend from .dynamo_vllm import DynamoVllmBackend from .sglang import SglangBackend from .vllm import VllmBackend - BACKENDS: tuple[ServerMetricsBackend, ...] = ( DynamoVllmBackend(), + AtomBackend(), SglangBackend(), VllmBackend(), ) diff --git a/utils/agentic/aggregation/backends/atom.py b/utils/agentic/aggregation/backends/atom.py new file mode 100644 index 0000000000..83533c456b --- /dev/null +++ b/utils/agentic/aggregation/backends/atom.py @@ -0,0 +1,101 @@ +"""Native ATOM server metric adapter.""" + +from __future__ import annotations + +from typing import Any + +from ..aggregation_common import gauge_stat, normalize_fraction, rate, sum_stat +from .base import ServerMetricsBackend, counter_int + + +def _atom_names(suffix: str) -> list[str]: + """Accept direct worker metrics and Atomesh's colon-normalized form.""" + return [f"atom:{suffix}", f"atom_{suffix}"] + + +def _atom_counter_names(stem: str) -> list[str]: + """Accept Prometheus raw names and AIPerf's counter-family names.""" + return [*_atom_names(stem), *_atom_names(f"{stem}_total")] + + +class AtomBackend(ServerMetricsBackend): + name = "atom" + + def matches(self, metrics: dict[str, dict[str, Any]], framework: str) -> bool: + metric_names = set(metrics) + return any(name.startswith(("atom:", "atom_")) for name in metric_names) or ( + not metrics and framework.lower() == "atom" + ) + + def populate( + self, + metrics: dict[str, dict[str, Any]], + flat: dict[str, Any], + nested: dict[str, Any], + ) -> None: + prompt_total = sum_stat( + metrics, + _atom_counter_names("prompt_tokens"), + preferred_keys=("total", "sum", "max", "avg"), + ) + generation_total = sum_stat( + metrics, + _atom_counter_names("generation_tokens"), + preferred_keys=("total", "sum", "max", "avg"), + ) + flat["total_prompt_tokens"] = counter_int(prompt_total) + flat["total_generation_tokens"] = counter_int(generation_total) + + cached_tokens = sum_stat( + metrics, + _atom_counter_names("prefix_cache_cached_tokens"), + preferred_keys=("total", "sum", "max", "avg"), + ) + full_tokens = sum_stat( + metrics, + _atom_counter_names("prefix_cache_full_tokens"), + preferred_keys=("total", "sum", "max", "avg"), + ) + cache_hit_rate = rate(cached_tokens, full_tokens) + external_tokens = sum_stat( + metrics, + _atom_counter_names("lmcache_loaded_tokens"), + preferred_keys=("total", "sum", "max", "avg"), + ) + external_hit_rate = rate(external_tokens, full_tokens) + # ATOM's admitted cache counter may already include a completed + # LMCache load, so do not add external tokens a second time. + overall_hit_rate = cache_hit_rate + + flat["server_gpu_cache_hit_rate"] = cache_hit_rate + flat["server_cpu_cache_hit_rate"] = external_hit_rate + flat["server_external_cache_hit_rate"] = external_hit_rate + flat["server_overall_cache_hit_rate"] = overall_hit_rate + flat["gpu_kv_cache_usage_pct"] = normalize_fraction( + gauge_stat( + metrics, + _atom_names("kv_cache_usage_ratio"), + preferred_keys=("max", "avg", "total"), + combine="max", + ) + ) + + nested["cache"].update( + { + "gpu_cache_hit_rate": cache_hit_rate, + "cpu_cache_hit_rate": external_hit_rate, + "external_cache_hit_rate": external_hit_rate, + "overall_cache_hit_rate": overall_hit_rate, + "prefix_cache_hits": cached_tokens, + "prefix_cache_queries": full_tokens, + "external_prefix_cache_hits": external_tokens, + "external_prefix_cache_queries": full_tokens, + } + ) + nested["kv_cache"]["gpu_usage_pct"] = flat["gpu_kv_cache_usage_pct"] + nested["tokens"].update( + { + "prompt_total": flat["total_prompt_tokens"], + "generation_total": flat["total_generation_tokens"], + } + ) diff --git a/utils/generate_aiperf_plots.py b/utils/generate_aiperf_plots.py index baefa7db21..9c80701fd6 100755 --- a/utils/generate_aiperf_plots.py +++ b/utils/generate_aiperf_plots.py @@ -110,6 +110,17 @@ def metric_entry(server_metrics: dict, name: str) -> dict | None: return entry if isinstance(entry, dict) else None +def first_metric_name(server_metrics: dict, *names: str) -> str: + """Return the first metric name present in an AIPerf export.""" + metrics = server_metrics.get("metrics") or {} + return next((name for name in names if name in metrics), names[0]) + + +def has_atom_metrics(server_metrics: dict) -> bool: + metrics = server_metrics.get("metrics") or {} + return any(name.startswith("atom:") for name in metrics) + + def all_series(entry: dict | None) -> list[dict]: if entry is None: return [] @@ -129,7 +140,9 @@ def series_with_label( def timeseries_from_series( - series: dict | None, t0_ns: int | None, value_key_priority=("avg", "rate", "total", "max") + series: dict | None, + t0_ns: int | None, + value_key_priority=("avg", "rate", "total", "max"), ) -> tuple[list[float], list[float]]: """Extract (relative-time-s, value) pairs from a series' timeslices.""" if series is None or t0_ns is None: @@ -153,7 +166,9 @@ def timeseries_from_series( def aggregate_timeseries( - server_metrics: dict, name: str, t0_ns: int | None, + server_metrics: dict, + name: str, + t0_ns: int | None, *, aggregator=sum, value_key_priority=("avg", "rate", "total", "max"), @@ -205,8 +220,11 @@ def rolling_window(n: int, max_window: int = 50) -> int: def panel_kv_cache_usage(ax, server_metrics: dict, t0_ns: int | None) -> None: + gpu_metric = first_metric_name( + server_metrics, "vllm:kv_cache_usage_perc", "atom:kv_cache_usage_ratio" + ) times, values = aggregate_timeseries( - server_metrics, "vllm:kv_cache_usage_perc", t0_ns, aggregator=max + server_metrics, gpu_metric, t0_ns, aggregator=max ) cpu_times, cpu_values = aggregate_timeseries( server_metrics, "vllm:cpu_kv_cache_usage_perc", t0_ns, aggregator=max @@ -242,12 +260,14 @@ def _norm(v: float) -> float: def panel_queue_depth(ax, server_metrics: dict, t0_ns: int | None) -> None: - rt, rv = aggregate_timeseries( - server_metrics, "vllm:num_requests_running", t0_ns, aggregator=max + running_metric = first_metric_name( + server_metrics, "vllm:num_requests_running", "atom:requests_running" ) - wt, wv = aggregate_timeseries( - server_metrics, "vllm:num_requests_waiting", t0_ns, aggregator=max + waiting_metric = first_metric_name( + server_metrics, "vllm:num_requests_waiting", "atom:requests_waiting" ) + rt, rv = aggregate_timeseries(server_metrics, running_metric, t0_ns, aggregator=max) + wt, wv = aggregate_timeseries(server_metrics, waiting_metric, t0_ns, aggregator=max) if rt: win = rolling_window(len(rv)) running = rolling_average(rv, win) if win > 1 else rv @@ -298,10 +318,20 @@ def _hit_rate_intervals( def panel_prefix_cache_hit_rate(ax, server_metrics: dict, t0_ns: int | None) -> None: - gpu_t, gpu_r = _hit_rate_intervals( + hits_metric = first_metric_name( server_metrics, "vllm:prefix_cache_hits", + "atom:prefix_cache_cached_tokens", + ) + queries_metric = first_metric_name( + server_metrics, "vllm:prefix_cache_queries", + "atom:prefix_cache_full_tokens", + ) + gpu_t, gpu_r = _hit_rate_intervals( + server_metrics, + hits_metric, + queries_metric, t0_ns, ) ext_t, ext_r = _hit_rate_intervals( @@ -358,11 +388,17 @@ def panel_prefix_cache_hit_rate(ax, server_metrics: dict, t0_ns: int | None) -> def panel_throughput(ax, server_metrics: dict, t0_ns: int | None) -> None: + generation_metric = first_metric_name( + server_metrics, "vllm:generation_tokens", "atom:generation_tokens" + ) + prompt_metric = first_metric_name( + server_metrics, "vllm:prompt_tokens", "atom:prompt_tokens" + ) gen_t, gen_v = aggregate_timeseries( - server_metrics, "vllm:generation_tokens", t0_ns, value_key_priority=("rate",) + server_metrics, generation_metric, t0_ns, value_key_priority=("rate",) ) prompt_t, prompt_v = aggregate_timeseries( - server_metrics, "vllm:prompt_tokens", t0_ns, value_key_priority=("rate",) + server_metrics, prompt_metric, t0_ns, value_key_priority=("rate",) ) if gen_t and prompt_t and len(gen_t) == len(prompt_t): total = [g + p for g, p in zip(gen_v, prompt_v)] @@ -396,7 +432,9 @@ def panel_throughput(ax, server_metrics: dict, t0_ns: int | None) -> None: running += total[i] * width elapsed = t - t0 if t > t0 else 1e-9 cumulative_total.append(running / elapsed if elapsed > 0 else 0.0) - ax.plot(gen_t, cumulative_total, "red", linewidth=2, label="Total Running Avg") + ax.plot( + gen_t, cumulative_total, "red", linewidth=2, label="Total Running Avg" + ) ax.legend(fontsize=8) ax.set_xlabel("Time (s)") ax.set_ylabel("Tokens/sec") @@ -404,18 +442,27 @@ def panel_throughput(ax, server_metrics: dict, t0_ns: int | None) -> None: ax.grid(True, alpha=0.3) -def panel_kv_offload_transfer_rate( - ax, server_metrics: dict, t0_ns: int | None -) -> None: - g2c_t, g2c_v = aggregate_timeseries( +def panel_kv_offload_transfer_rate(ax, server_metrics: dict, t0_ns: int | None) -> None: + atom_metrics = has_atom_metrics(server_metrics) + gpu_to_cpu_metric = first_metric_name( server_metrics, "vllm:kv_offload_bytes_gpu_to_cpu", + "atom:lmcache_saved_tokens", + ) + cpu_to_gpu_metric = first_metric_name( + server_metrics, + "vllm:kv_offload_bytes_cpu_to_gpu", + "atom:lmcache_loaded_tokens", + ) + g2c_t, g2c_v = aggregate_timeseries( + server_metrics, + gpu_to_cpu_metric, t0_ns, value_key_priority=("rate",), ) c2g_t, c2g_v = aggregate_timeseries( server_metrics, - "vllm:kv_offload_bytes_cpu_to_gpu", + cpu_to_gpu_metric, t0_ns, value_key_priority=("rate",), ) @@ -424,36 +471,52 @@ def panel_kv_offload_transfer_rate( ) if has_data: if g2c_t: - mb = [v / 1e6 for v in g2c_v] - ax.scatter(g2c_t, mb, alpha=0.15, s=3, c="blue") - win = rolling_window(len(mb)) + scaled = [v / 1e6 for v in g2c_v] + ax.scatter(g2c_t, scaled, alpha=0.15, s=3, c="blue") + win = rolling_window(len(scaled)) if win > 1: ax.plot( g2c_t, - rolling_average(mb, win), + rolling_average(scaled, win), "b-", linewidth=1.5, label=f"GPU→CPU (avg n={win})", ) else: - ax.plot(g2c_t, mb, "b-", linewidth=1, alpha=0.8, label="GPU→CPU") + ax.plot( + g2c_t, + scaled, + "b-", + linewidth=1, + alpha=0.8, + label="GPU→CPU", + ) if c2g_t: - mb = [v / 1e6 for v in c2g_v] - ax.scatter(c2g_t, mb, alpha=0.15, s=3, c="red") - win = rolling_window(len(mb)) + scaled = [v / 1e6 for v in c2g_v] + ax.scatter(c2g_t, scaled, alpha=0.15, s=3, c="red") + win = rolling_window(len(scaled)) if win > 1: ax.plot( c2g_t, - rolling_average(mb, win), + rolling_average(scaled, win), "r-", linewidth=1.5, label=f"CPU→GPU (avg n={win})", ) else: - ax.plot(c2g_t, mb, "r-", linewidth=1, alpha=0.8, label="CPU→GPU") + ax.plot( + c2g_t, + scaled, + "r-", + linewidth=1, + alpha=0.8, + label="CPU→GPU", + ) ax.legend(fontsize=8) ax.set_xlabel("Time (s)") - ax.set_ylabel("Transfer Rate (MB/s)") + ax.set_ylabel( + "Transfer Rate (M tokens/s)" if atom_metrics else "Transfer Rate (MB/s)" + ) ax.set_title("KV Offload Transfer Rate") ax.grid(True, alpha=0.3) @@ -467,14 +530,44 @@ def _prompt_token_source_series( return timeseries_from_series(s, t0_ns, value_key_priority=("total",)) -def panel_prefill_source_breakdown( - ax, server_metrics: dict, t0_ns: int | None -) -> None: - c_t, c_v = _prompt_token_source_series(server_metrics, "local_compute", t0_ns) - h_t, h_v = _prompt_token_source_series(server_metrics, "local_cache_hit", t0_ns) - e_t, e_v = _prompt_token_source_series( - server_metrics, "external_kv_transfer", t0_ns - ) +def panel_prefill_source_breakdown(ax, server_metrics: dict, t0_ns: int | None) -> None: + if has_atom_metrics(server_metrics): + full_t, full_v = aggregate_timeseries( + server_metrics, + "atom:prefix_cache_full_tokens", + t0_ns, + value_key_priority=("total",), + ) + cached_t, cached_v = aggregate_timeseries( + server_metrics, + "atom:prefix_cache_cached_tokens", + t0_ns, + value_key_priority=("total",), + ) + e_t, e_v = aggregate_timeseries( + server_metrics, + "atom:lmcache_loaded_tokens", + t0_ns, + value_key_priority=("total",), + ) + all_times = sorted(set(full_t) | set(cached_t) | set(e_t)) + full_by_t = dict(zip(full_t, full_v)) + cached_by_t = dict(zip(cached_t, cached_v)) + ext_by_t = dict(zip(e_t, e_v)) + c_t = all_times + c_v = [ + max(0.0, full_by_t.get(t, 0.0) - cached_by_t.get(t, 0.0)) for t in all_times + ] + h_t = all_times + h_v = [ + max(0.0, cached_by_t.get(t, 0.0) - ext_by_t.get(t, 0.0)) for t in all_times + ] + else: + c_t, c_v = _prompt_token_source_series(server_metrics, "local_compute", t0_ns) + h_t, h_v = _prompt_token_source_series(server_metrics, "local_cache_hit", t0_ns) + e_t, e_v = _prompt_token_source_series( + server_metrics, "external_kv_transfer", t0_ns + ) # Align timestamps: use the union of all sample timestamps. if not (c_t or h_t or e_t): ax.set_xlabel("Time (s)") @@ -546,6 +639,7 @@ def panel_kv_offload_cumulative( title: str, color: str, t0_ns: int | None, + unit: str = "bytes", ) -> None: times, values = aggregate_timeseries( server_metrics, metric_name, t0_ns, value_key_priority=("total",) @@ -555,11 +649,15 @@ def panel_kv_offload_cumulative( running = 0.0 for v in values: running += v - cumulative.append(running / 1e9) # GB + cumulative.append(running / (1e6 if unit == "tokens" else 1e9)) ax.plot(times, cumulative, f"{color}-", linewidth=1.5) ax.fill_between(times, cumulative, alpha=0.2, color=color) ax.set_xlabel("Time (s)") - ax.set_ylabel("Cumulative Transfer (GB)") + ax.set_ylabel( + "Cumulative Transfer (M tokens)" + if unit == "tokens" + else "Cumulative Transfer (GB)" + ) ax.set_title(title) ax.grid(True, alpha=0.3) @@ -597,8 +695,11 @@ def panel_per_record_metric( def panel_preemptions(ax, server_metrics: dict, t0_ns: int | None) -> None: + metric_name = first_metric_name( + server_metrics, "vllm:num_preemptions", "atom:preemptions" + ) times, values = aggregate_timeseries( - server_metrics, "vllm:num_preemptions", t0_ns, value_key_priority=("total",) + server_metrics, metric_name, t0_ns, value_key_priority=("total",) ) if not times: ax.set_xlabel("Time (s)") @@ -708,8 +809,9 @@ def main(argv: list[str]) -> int: # Interactivity: tokens/sec from per-token latency (ms). interactivities.append(1000.0 / itl if itl and itl > 0 else 0.0) + atom_metrics = has_atom_metrics(server_metrics) fig, axes = plt.subplots(6, 2, figsize=(14, 24)) - fig.suptitle("vLLM Server Metrics During Benchmark", fontsize=14) + fig.suptitle("LLM Server Metrics During Benchmark", fontsize=14) panel_kv_cache_usage(axes[0, 0], server_metrics, t0_ns) panel_queue_depth(axes[0, 1], server_metrics, t0_ns) @@ -720,18 +822,28 @@ def main(argv: list[str]) -> int: panel_kv_offload_cumulative( axes[3, 0], server_metrics, - "vllm:kv_offload_bytes_gpu_to_cpu", + ( + "atom:lmcache_saved_tokens" + if atom_metrics + else "vllm:kv_offload_bytes_gpu_to_cpu" + ), "KV Offload: GPU → CPU (Cumulative)", "b", t0_ns, + unit="tokens" if atom_metrics else "bytes", ) panel_kv_offload_cumulative( axes[3, 1], server_metrics, - "vllm:kv_offload_bytes_cpu_to_gpu", + ( + "atom:lmcache_loaded_tokens" + if atom_metrics + else "vllm:kv_offload_bytes_cpu_to_gpu" + ), "KV Offload: CPU → GPU (Cumulative)", "r", t0_ns, + unit="tokens" if atom_metrics else "bytes", ) panel_per_record_metric( axes[4, 0],