From 9aee6e1509ea3361c48d22cc490e38cdb84ceca8 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 11 Feb 2026 13:25:53 +0000 Subject: [PATCH 1/9] Add performance regression test suite Replaces the manual workflow of running examples/infer.py on two branches and eyeballing TTFT/TBT numbers. Tests run a parametrized matrix (batch_size x prompt_length x decode_length), average metrics over multiple iterations after warmup, and compare against stored JSON baselines with a configurable tolerance threshold. https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/performance/__init__.py | 0 tests/performance/baselines/.gitkeep | 0 tests/performance/conftest.py | 210 +++++++++++++++++++++ tests/performance/test_perf_regression.py | 219 ++++++++++++++++++++++ 4 files changed, 429 insertions(+) create mode 100644 tests/performance/__init__.py create mode 100644 tests/performance/baselines/.gitkeep create mode 100644 tests/performance/conftest.py create mode 100644 tests/performance/test_perf_regression.py diff --git a/tests/performance/__init__.py b/tests/performance/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/performance/baselines/.gitkeep b/tests/performance/baselines/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py new file mode 100644 index 00000000..3995e388 --- /dev/null +++ b/tests/performance/conftest.py @@ -0,0 +1,210 @@ +import json +import os +import subprocess +from datetime import datetime, timezone + +import pytest +import torch.distributed as dist +from transformers import AutoTokenizer + +from yalis import ModelConfig, InferenceConfig, LLMEngine +from tests.sample_dataset import AlpacaDataset + +BASELINE_DIR = os.path.join(os.path.dirname(__file__), "baselines") +DEFAULT_BASELINE_PATH = os.path.join(BASELINE_DIR, "perf_baselines.json") + + +# ------------------------------------------------------------------ # +# CLI options # +# ------------------------------------------------------------------ # + + +def pytest_addoption(parser): + parser.addini( + "model", + "Model to use for the test", + type="string", + default="meta-llama/Llama-3.1-8B-Instruct", + ) + parser.addini( + "dtype", + "Data type to use for the test", + type="string", + default="bf16", + ) + parser.addini( + "attn_backend", + "Attention backend to use for the test", + type="string", + default="sdpa", + ) + parser.addoption( + "--perf-update-baselines", + action="store_true", + default=False, + help="Update performance baselines instead of comparing.", + ) + parser.addoption( + "--perf-tolerance", + type=float, + default=0.10, + help="Max allowed regression fraction (default: 0.10 = 10%%).", + ) + parser.addoption( + "--perf-warmup-iters", + type=int, + default=3, + help="Warmup iterations before measurement (default: 3).", + ) + parser.addoption( + "--perf-measure-iters", + type=int, + default=5, + help="Measurement iterations for averaging (default: 5).", + ) + parser.addoption( + "--perf-baseline-path", + type=str, + default=DEFAULT_BASELINE_PATH, + help="Path to the baseline JSON file.", + ) + + +# ------------------------------------------------------------------ # +# Baseline store # +# ------------------------------------------------------------------ # + + +class BaselineStore: + """Thin wrapper around a JSON file that holds perf baselines.""" + + def __init__(self, path): + self.path = path + self._data = self._load() + + # -- persistence ------------------------------------------------ # + + def _load(self): + if os.path.exists(self.path): + with open(self.path) as f: + return json.load(f) + return {"metadata": {}, "benchmarks": {}} + + def flush(self): + os.makedirs(os.path.dirname(self.path), exist_ok=True) + with open(self.path, "w") as f: + json.dump(self._data, f, indent=2) + + # -- read / write ----------------------------------------------- # + + def get(self, key): + return self._data["benchmarks"].get(key) + + def put(self, key, entry): + self._data["benchmarks"][key] = entry + + def set_metadata(self, **kwargs): + self._data["metadata"].update(kwargs) + + +# ------------------------------------------------------------------ # +# Fixtures # +# ------------------------------------------------------------------ # + + +@pytest.fixture(scope="module", autouse=True) +def cleanup_dist(): + yield + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() + + +@pytest.fixture(scope="session") +def model_id(request): + return request.config.getini("model") + + +@pytest.fixture(scope="session") +def dtype(request): + return request.config.getini("dtype").lower() + + +@pytest.fixture(scope="session") +def attn_backend(request): + return request.config.getini("attn_backend").lower() + + +@pytest.fixture(scope="module") +def perf_engine(model_id, dtype, attn_backend): + """LLMEngine configured for performance measurement.""" + model_config = ModelConfig(model_name=model_id, precision=dtype) + inference_config = InferenceConfig( + max_batch_size=8, + max_length_of_generated_sequences=2048, + top_p=0.0, + temperature=0.0, + tp_dims=None, + attention_backend=attn_backend, + use_paged_kv_caching=False, + prestore_kv_cache=True, + ) + return LLMEngine( + model_config=model_config, + inference_config=inference_config, + ) + + +@pytest.fixture(scope="session") +def tokenizer(model_id): + tok = AutoTokenizer.from_pretrained(model_id) + tok.pad_token = tok.eos_token + tok.padding_side = "left" + return tok + + +@pytest.fixture(scope="session") +def alpaca_dataset(): + return AlpacaDataset(random_seed=42) + + +@pytest.fixture(scope="session") +def baseline_store(request, model_id, dtype, attn_backend): + """Load (or create) the baseline store and flush on teardown.""" + path = request.config.getoption("--perf-baseline-path") + store = BaselineStore(path) + + update = request.config.getoption("--perf-update-baselines") + if update: + git_sha = _git_sha() + store.set_metadata( + model=model_id, + attention_backend=attn_backend, + precision=dtype, + updated_at=datetime.now(timezone.utc).isoformat(), + git_commit=git_sha, + ) + + yield store + + if update: + store.flush() + + +# ------------------------------------------------------------------ # +# Helpers # +# ------------------------------------------------------------------ # + + +def _git_sha(): + try: + return ( + subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"], + stderr=subprocess.DEVNULL, + ) + .decode() + .strip() + ) + except Exception: + return "unknown" diff --git a/tests/performance/test_perf_regression.py b/tests/performance/test_perf_regression.py new file mode 100644 index 00000000..37b91bce --- /dev/null +++ b/tests/performance/test_perf_regression.py @@ -0,0 +1,219 @@ +""" +Performance regression tests for YALIS. + +Workflow +-------- +1. On the develop (baseline) branch, generate baselines:: + + pytest tests/performance/ --perf-update-baselines + +2. On your feature branch, run the tests to check for regressions:: + + pytest tests/performance/ + + Any metric that regresses beyond the tolerance (default 10 %) + will cause the test to fail with a detailed report. + +Options +------- +--perf-tolerance FLOAT Allowed regression fraction (default 0.10). +--perf-warmup-iters INT Warmup iterations (default 3). +--perf-measure-iters INT Measurement iterations (default 5). +--perf-baseline-path PATH Path to the baselines JSON file. +""" + +import pytest +import torch.distributed as dist + +from tests.basic_correctness.utils import alpaca_prompt + +BATCH_SIZES = [1, 8] +PROMPT_LENGTHS = [128, 512] +DECODE_LENGTHS = [32, 128] + +# Metrics where *lower* is better (latencies). +_LOWER_IS_BETTER = {"ttft_ms", "tbt_ms", "e2e_ms"} +# Metrics where *higher* is better (throughput). +_HIGHER_IS_BETTER = {"throughput_tps"} + +_ALL_METRICS = [ + ("ttft_ms", "TTFT"), + ("tbt_ms", "TBT"), + ("throughput_tps", "Throughput"), + ("e2e_ms", "E2E"), +] + + +def _bench_key(batch_size, prompt_length, decode_length): + return ( + f"batch_{batch_size}" + f"_prompt_{prompt_length}" + f"_decode_{decode_length}" + ) + + +def _run_iterations(engine, prompts, decode_length, n_iters): + """Run *n_iters* generate calls and return the list of metric dicts.""" + collected = [] + for _ in range(n_iters): + _, metrics = engine.generate( + prompts, + report_throughput=False, + tokens_to_generate=decode_length, + ignore_eos=True, + ) + collected.append(metrics) + return collected + + +def _average_metrics(metrics_list): + n = len(metrics_list) + return { + "ttft_ms": sum(m["TTFT"] for m in metrics_list) / n, + "tbt_ms": sum(m["TBT"] for m in metrics_list) / n, + "throughput_tps": sum(m["Throughput"] for m in metrics_list) / n, + "e2e_ms": sum(m["E2E"] for m in metrics_list) / n, + } + + +def _check_regressions(baseline, current, tolerance): + """Return a list of (metric, baseline_val, current_val, pct) tuples + for every metric that regressed beyond *tolerance*.""" + regressions = [] + for key, label in _ALL_METRICS: + base_val = baseline[key] + curr_val = current[key] + + if base_val == 0: + continue + + if key in _LOWER_IS_BETTER: + pct = (curr_val - base_val) / base_val + regressed = pct > tolerance + else: + pct = (base_val - curr_val) / base_val + regressed = pct > tolerance + pct = -pct # show as negative when throughput drops + + if regressed: + regressions.append((label, base_val, curr_val, pct)) + + return regressions + + +def _format_report(key, current, baseline, regressions, tolerance): + """Build a human-readable report string.""" + lines = [f"Performance regression detected for [{key}]:"] + lines.append("") + lines.append( + f" {'Metric':<14} {'Baseline':>12} {'Current':>12} {'Change':>10}" + ) + lines.append(f" {'-'*50}") + + for mkey, label in _ALL_METRICS: + base_val = baseline[mkey] + curr_val = current[mkey] + if base_val != 0: + pct = (curr_val - base_val) / base_val + marker = ( + " << REGRESSION" + if any(r[0] == label for r in regressions) + else "" + ) + lines.append( + f" {label:<14} {base_val:>12.4f} {curr_val:>12.4f}" + f" {pct:>+9.1%}{marker}" + ) + else: + lines.append( + f" {label:<14} {base_val:>12.4f} {curr_val:>12.4f}" + f" N/A" + ) + + lines.append("") + lines.append(f" Tolerance: {tolerance:.0%}") + return "\n".join(lines) + + +# ------------------------------------------------------------------ # +# Tests # +# ------------------------------------------------------------------ # + + +@pytest.mark.parametrize("batch_size", BATCH_SIZES) +@pytest.mark.parametrize("prompt_length", PROMPT_LENGTHS) +@pytest.mark.parametrize("decode_length", DECODE_LENGTHS) +def test_perf_regression( + perf_engine, + tokenizer, + alpaca_dataset, + baseline_store, + batch_size, + prompt_length, + decode_length, + request, +): + config = request.config + update_mode = config.getoption("--perf-update-baselines") + tolerance = config.getoption("--perf-tolerance") + warmup_iters = config.getoption("--perf-warmup-iters") + measure_iters = config.getoption("--perf-measure-iters") + + key = _bench_key(batch_size, prompt_length, decode_length) + + # --- prepare prompts ------------------------------------------ # + prompts = alpaca_prompt( + alpaca_dataset, tokenizer, prompt_length, batch_size + ) + + # --- warmup --------------------------------------------------- # + _run_iterations(perf_engine, prompts, decode_length, warmup_iters) + + # --- measure -------------------------------------------------- # + raw = _run_iterations(perf_engine, prompts, decode_length, measure_iters) + current = _average_metrics(raw) + + # Only rank 0 performs the baseline comparison / update. + if dist.is_initialized() and dist.get_rank() != 0: + return + + # --- update mode: store and return ---------------------------- # + if update_mode: + baseline_store.put( + key, + { + "batch_size": batch_size, + "prompt_length": prompt_length, + "decode_length": decode_length, + **current, + }, + ) + print(f"\n [baseline saved] {key}") + for mkey, label in _ALL_METRICS: + print(f" {label:<14} {current[mkey]:.4f}") + return + + # --- compare mode --------------------------------------------- # + baseline = baseline_store.get(key) + if baseline is None: + pytest.skip( + f"No baseline for {key}. " + "Run with --perf-update-baselines first." + ) + + regressions = _check_regressions(baseline, current, tolerance) + + # Always print a summary so the user can eyeball the numbers. + print(f"\n [perf] {key}") + for mkey, label in _ALL_METRICS: + base_val = baseline[mkey] + curr_val = current[mkey] + pct = (curr_val - base_val) / base_val if base_val != 0 else 0 + print( + f" {label:<14} {base_val:>10.4f} -> {curr_val:>10.4f}" + f" ({pct:+.1%})" + ) + + if regressions: + report = _format_report(key, current, baseline, regressions, tolerance) + pytest.fail(report) From d9f18d386e6520d4bccfdce7dc6ca923e93b33ea Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 13 Feb 2026 18:06:26 +0000 Subject: [PATCH 2/9] Add baseline config validation and paged KV caching option to perf tests Validate that the stored baseline metadata (model, dtype, attn_backend, use_paged_kv_caching) matches the current test configuration before comparing metrics. Raises pytest.UsageError on mismatch to prevent misleading comparisons across different settings. Fields not yet tracked in older baselines are gracefully skipped. Also adds use_paged_kv_caching as an INI-configurable option (default False) wired through fixture -> InferenceConfig -> LLMEngine. https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/performance/conftest.py | 65 +++++++++++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 3 deletions(-) diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py index 3995e388..55769af6 100644 --- a/tests/performance/conftest.py +++ b/tests/performance/conftest.py @@ -38,6 +38,12 @@ def pytest_addoption(parser): type="string", default="sdpa", ) + parser.addini( + "use_paged_kv_caching", + "Enable paged KV caching (requires flash backend)", + type="bool", + default=False, + ) parser.addoption( "--perf-update-baselines", action="store_true", @@ -135,8 +141,13 @@ def attn_backend(request): return request.config.getini("attn_backend").lower() +@pytest.fixture(scope="session") +def use_paged_kv_caching(request): + return request.config.getini("use_paged_kv_caching") + + @pytest.fixture(scope="module") -def perf_engine(model_id, dtype, attn_backend): +def perf_engine(model_id, dtype, attn_backend, use_paged_kv_caching): """LLMEngine configured for performance measurement.""" model_config = ModelConfig(model_name=model_id, precision=dtype) inference_config = InferenceConfig( @@ -146,7 +157,7 @@ def perf_engine(model_id, dtype, attn_backend): temperature=0.0, tp_dims=None, attention_backend=attn_backend, - use_paged_kv_caching=False, + use_paged_kv_caching=use_paged_kv_caching, prestore_kv_cache=True, ) return LLMEngine( @@ -169,7 +180,9 @@ def alpaca_dataset(): @pytest.fixture(scope="session") -def baseline_store(request, model_id, dtype, attn_backend): +def baseline_store( + request, model_id, dtype, attn_backend, use_paged_kv_caching +): """Load (or create) the baseline store and flush on teardown.""" path = request.config.getoption("--perf-baseline-path") store = BaselineStore(path) @@ -181,9 +194,14 @@ def baseline_store(request, model_id, dtype, attn_backend): model=model_id, attention_backend=attn_backend, precision=dtype, + use_paged_kv_caching=use_paged_kv_caching, updated_at=datetime.now(timezone.utc).isoformat(), git_commit=git_sha, ) + else: + _validate_baseline_config( + store, model_id, dtype, attn_backend, use_paged_kv_caching + ) yield store @@ -196,6 +214,47 @@ def baseline_store(request, model_id, dtype, attn_backend): # ------------------------------------------------------------------ # +def _validate_baseline_config( + store, model_id, dtype, attn_backend, use_paged_kv_caching +): + """Verify the current test config matches the baseline metadata. + + Raises ``pytest.UsageError`` on mismatch so the session fails + immediately rather than producing misleading comparisons. + """ + meta = store._data.get("metadata", {}) + if not meta: + return # no baselines yet — nothing to validate + + checks = { + "model": (meta.get("model"), model_id), + "precision": (meta.get("precision"), dtype), + "attention_backend": (meta.get("attention_backend"), attn_backend), + "use_paged_kv_caching": ( + meta.get("use_paged_kv_caching"), + use_paged_kv_caching, + ), + } + + mismatches = [] + for field, (stored, current) in checks.items(): + if stored is None: + # Baseline was created before this field was tracked — skip. + continue + if stored != current: + mismatches.append( + f" {field}: baseline={stored!r}, current={current!r}" + ) + + if mismatches: + detail = "\n".join(mismatches) + raise pytest.UsageError( + f"Baseline config mismatch — the stored baselines were " + f"recorded with a different configuration:\n{detail}\n" + f"Re-run with --perf-update-baselines to regenerate." + ) + + def _git_sha(): try: return ( From e47594b0c1ad3b653cb5739edfdeb231b8d89a16 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 13 Feb 2026 18:29:10 +0000 Subject: [PATCH 3/9] Show perf comparison table in terminal summary without -s Replace print()-based output (which requires -s / --capture=no) with pytest's terminal reporter hooks so the numbers are always visible. - pytest_configure: initialise a session-wide results collector via pytest.StashKey - pytest_terminal_summary: render a formatted table at the end of the run, covering both update mode (saved values) and compare mode (baseline vs current with % change and !! markers for regressions) - perf_results fixture: gives tests write access to the collector - Tests now append structured dicts instead of printing https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/performance/conftest.py | 60 +++++++++++++++++++++++ tests/performance/test_perf_regression.py | 21 ++++---- 2 files changed, 72 insertions(+), 9 deletions(-) diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py index 55769af6..e6d2a46d 100644 --- a/tests/performance/conftest.py +++ b/tests/performance/conftest.py @@ -13,6 +13,60 @@ BASELINE_DIR = os.path.join(os.path.dirname(__file__), "baselines") DEFAULT_BASELINE_PATH = os.path.join(BASELINE_DIR, "perf_baselines.json") +_PERF_RESULTS_KEY = pytest.StashKey[list]() + + +# ------------------------------------------------------------------ # +# Hooks # +# ------------------------------------------------------------------ # + + +def pytest_configure(config): + """Initialise a session-wide list to collect perf comparison results.""" + config.stash[_PERF_RESULTS_KEY] = [] + + +def pytest_terminal_summary(terminalreporter, config): + """Print a performance comparison table after the test run.""" + results = config.stash.get(_PERF_RESULTS_KEY, []) + if not results: + return + + write = terminalreporter.write_line + update_mode = config.getoption("--perf-update-baselines", default=False) + + if update_mode: + write("") + write("=== Performance baselines saved ===", bold=True) + for entry in results: + write(f" [{entry['key']}]") + for label, value in entry["metrics"]: + write(f" {label:<14} {value:>12.4f}") + write("") + else: + tolerance = config.getoption("--perf-tolerance", default=0.10) + write("") + write("=== Performance comparison ===", bold=True) + write( + f" {'Benchmark':<40} {'Metric':<14}" + f" {'Baseline':>12} {'Current':>12} {'Change':>10}" + ) + write(f" {'-' * 92}") + for entry in results: + first = True + for label, base_val, curr_val, pct in entry["comparisons"]: + tag = entry["key"] if first else "" + marker = " !!" if abs(pct) > tolerance else "" + write( + f" {tag:<40} {label:<14}" + f" {base_val:>12.4f} {curr_val:>12.4f}" + f" {pct:>+9.1%}{marker}" + ) + first = False + write("") + write(f" Tolerance: {tolerance:.0%}") + write("") + # ------------------------------------------------------------------ # # CLI options # @@ -179,6 +233,12 @@ def alpaca_dataset(): return AlpacaDataset(random_seed=42) +@pytest.fixture(scope="session") +def perf_results(request): + """Session-wide list for collecting perf comparison data.""" + return request.config.stash[_PERF_RESULTS_KEY] + + @pytest.fixture(scope="session") def baseline_store( request, model_id, dtype, attn_backend, use_paged_kv_caching diff --git a/tests/performance/test_perf_regression.py b/tests/performance/test_perf_regression.py index 37b91bce..75d2b7e2 100644 --- a/tests/performance/test_perf_regression.py +++ b/tests/performance/test_perf_regression.py @@ -148,6 +148,7 @@ def test_perf_regression( tokenizer, alpaca_dataset, baseline_store, + perf_results, batch_size, prompt_length, decode_length, @@ -188,9 +189,14 @@ def test_perf_regression( **current, }, ) - print(f"\n [baseline saved] {key}") - for mkey, label in _ALL_METRICS: - print(f" {label:<14} {current[mkey]:.4f}") + perf_results.append( + { + "key": key, + "metrics": [ + (label, current[mkey]) for mkey, label in _ALL_METRICS + ], + } + ) return # --- compare mode --------------------------------------------- # @@ -203,16 +209,13 @@ def test_perf_regression( regressions = _check_regressions(baseline, current, tolerance) - # Always print a summary so the user can eyeball the numbers. - print(f"\n [perf] {key}") + comparisons = [] for mkey, label in _ALL_METRICS: base_val = baseline[mkey] curr_val = current[mkey] pct = (curr_val - base_val) / base_val if base_val != 0 else 0 - print( - f" {label:<14} {base_val:>10.4f} -> {curr_val:>10.4f}" - f" ({pct:+.1%})" - ) + comparisons.append((label, base_val, curr_val, pct)) + perf_results.append({"key": key, "comparisons": comparisons}) if regressions: report = _format_report(key, current, baseline, regressions, tolerance) From 140b0e56a6763203c0b1e7d3652f8f03be7d13e9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 14 Feb 2026 02:54:56 +0000 Subject: [PATCH 4/9] Validate --perf-measure-iters is at least 1 Prevents a ZeroDivisionError in _average_metrics when the option is set to 0. https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/performance/conftest.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py index e6d2a46d..1d6cb62b 100644 --- a/tests/performance/conftest.py +++ b/tests/performance/conftest.py @@ -26,6 +26,13 @@ def pytest_configure(config): config.stash[_PERF_RESULTS_KEY] = [] +def pytest_sessionstart(session): + """Validate CLI options that must be positive.""" + val = session.config.getoption("--perf-measure-iters", default=5) + if val < 1: + raise pytest.UsageError("--perf-measure-iters must be at least 1") + + def pytest_terminal_summary(terminalreporter, config): """Print a performance comparison table after the test run.""" results = config.stash.get(_PERF_RESULTS_KEY, []) @@ -120,7 +127,7 @@ def pytest_addoption(parser): "--perf-measure-iters", type=int, default=5, - help="Measurement iterations for averaging (default: 5).", + help="Measurement iterations for averaging (default: 5, min: 1).", ) parser.addoption( "--perf-baseline-path", From e54da7d8ed83aafc22d246e8b6c379f21015a965 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 14 Feb 2026 12:44:17 +0000 Subject: [PATCH 5/9] Add perf regression test runner script and tests README Add run_perf_regression_tests.sh following the same pattern as the existing correctness test scripts, setting up environment variables and launching via srun. Extra pytest flags (e.g. --perf-update-baselines) can be passed through the PERF_PYTEST_ARGS env var. Also add tests/README.md documenting the performance regression test workflow, configuration options, and available pytest flags. https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/README.md | 75 ++++++++++++++++++++++ tests/scripts/run_perf_regression_tests.sh | 38 +++++++++++ 2 files changed, 113 insertions(+) create mode 100644 tests/README.md create mode 100755 tests/scripts/run_perf_regression_tests.sh diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 00000000..cae17f4b --- /dev/null +++ b/tests/README.md @@ -0,0 +1,75 @@ +# YALIS Tests + +## Performance Regression Tests + +The performance tests live in `tests/performance/` and measure TTFT, TBT, +end-to-end latency, and throughput across a matrix of batch sizes, prompt +lengths, and decode lengths. + +### Quick start + +Performance testing is a two-step workflow: first record baselines on a +known-good branch, then compare against those baselines on a feature branch. + +#### 1. Generate baselines (on the develop / baseline branch) + +```bash +PERF_PYTEST_ARGS="--perf-update-baselines" ./tests/scripts/run_perf_regression_tests.sh +``` + +This runs every benchmark combination and writes the results to +`tests/performance/baselines/perf_baselines.json`. + +#### 2. Check for regressions (on your feature branch) + +```bash +./tests/scripts/run_perf_regression_tests.sh +``` + +Any metric that regresses beyond the tolerance (default 10 %) will cause the +test to fail with a detailed report. + +### Configuration + +The script uses `srun` to launch on a Slurm cluster. You can control the +number of GPUs with the `GPUS` environment variable: + +```bash +GPUS=4 ./tests/scripts/run_perf_regression_tests.sh +``` + +Additional pytest options can be passed through the `PERF_PYTEST_ARGS` +environment variable: + +| Option | Default | Description | +| ----------------------------- | ----------------------------------------------- | ------------------------------------------------- | +| `--perf-update-baselines` | off | Record new baselines instead of comparing. | +| `--perf-tolerance FLOAT` | `0.10` | Max allowed regression fraction (10 %). | +| `--perf-warmup-iters INT` | `3` | Warmup iterations before measurement. | +| `--perf-measure-iters INT` | `5` | Measurement iterations for averaging. | +| `--perf-baseline-path PATH` | `tests/performance/baselines/perf_baselines.json` | Path to the baselines JSON file. | + +Example — tighter tolerance with more measurement iterations: + +```bash +PERF_PYTEST_ARGS="--perf-tolerance 0.05 --perf-measure-iters 10" \ + ./tests/scripts/run_perf_regression_tests.sh +``` + +The model, precision, and attention backend are configured via pytest ini +settings. The defaults (set in `tests/performance/conftest.py`) are: + +| Setting | Default | +| ------------------------ | ------------------------------------ | +| `model` | `meta-llama/Llama-3.1-8B-Instruct` | +| `dtype` | `bf16` | +| `attn_backend` | `sdpa` | +| `use_paged_kv_caching` | `False` | + +To override these, create a `pytest.ini` (or add an `[pytest]` section to +`pyproject.toml`) with the desired values, or pass a custom `-c ` +through `PERF_PYTEST_ARGS`. + +## Correctness Tests + + diff --git a/tests/scripts/run_perf_regression_tests.sh b/tests/scripts/run_perf_regression_tests.sh new file mode 100755 index 00000000..6fb69ad4 --- /dev/null +++ b/tests/scripts/run_perf_regression_tests.sh @@ -0,0 +1,38 @@ +#!/bin/bash +NNODES=1 +GPUS_DEFAULT=1 +GPUS=${GPUS:-$GPUS_DEFAULT} + + +export MASTER_ADDR=$(hostname) +export MASTER_PORT=29500 +export WORLD_SIZE=${GPUS} + +## nccl env vars to speedup stuff +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export NCCL_NET_GDR_LEVEL=PHB +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsi +export MPICH_GPU_SUPPORT_ENABLED=0 +export CUDA_VISIBLE_DEVICES=0 + +export HF_HOME="$SCRATCH/hf_cache" +export TRANSFORMERS_HOME="$SCRATCH/hf_cache" +export HF_DATASETS_CACHE="$SCRATCH/hf_cache" +export YALIS_CACHE="/pscratch/sd/p/prajwal/SpecDec/yalis/yalis/external" +export TORCHINDUCTOR_CACHE_DIR="${SCRATCH}/.cache/torch_inductor" + +SCRIPT="tests/performance/test_perf_regression.py" + +export PYTHONPATH="$PYTHONPATH:." + +chmod +x tests/get_rank_tests.sh + +# Pass extra pytest flags via PERF_PYTEST_ARGS, e.g.: +# PERF_PYTEST_ARGS="--perf-update-baselines" ./tests/scripts/run_perf_regression_tests.sh +# PERF_PYTEST_ARGS="--perf-tolerance 0.15" ./tests/scripts/run_perf_regression_tests.sh +EXTRA_ARGS=${PERF_PYTEST_ARGS:-} + +perf_cmd="NCCL_CUMEM_ENABLE=0 srun -N $NNODES -n $GPUS -G $GPUS -c 16 --cpu-bind=cores ./tests/get_rank_tests.sh pytest $SCRIPT $EXTRA_ARGS" +echo $perf_cmd +eval $perf_cmd From 0ba22d247e9c0b142bd06123a0dfbdda985fcdfc Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 14 Feb 2026 12:54:25 +0000 Subject: [PATCH 6/9] Fix baseline store race on multi-GPU runs Only rank 0 populates benchmark entries during --perf-update-baselines, but every rank was flushing the store on teardown. On multi-GPU runs a non-zero rank could overwrite the JSON with empty benchmark data. Guard the flush so only rank 0 writes the file. https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/performance/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py index 1d6cb62b..dfaeed5b 100644 --- a/tests/performance/conftest.py +++ b/tests/performance/conftest.py @@ -272,7 +272,7 @@ def baseline_store( yield store - if update: + if update and (not dist.is_initialized() or dist.get_rank() == 0): store.flush() From ad796da48a59165545410a64b279efe9396d9e31 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 14 Feb 2026 12:57:29 +0000 Subject: [PATCH 7/9] Capture rank at fixture setup to survive dist teardown cleanup_dist (scope=module) destroys the process group before the session-scoped baseline_store fixture tears down, so checking dist.is_initialized() at teardown always returns False on every rank. Capture whether the process is rank 0 at setup time and use the saved value during teardown to ensure only rank 0 flushes baselines. https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/performance/conftest.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py index dfaeed5b..63beb673 100644 --- a/tests/performance/conftest.py +++ b/tests/performance/conftest.py @@ -255,6 +255,10 @@ def baseline_store( store = BaselineStore(path) update = request.config.getoption("--perf-update-baselines") + + # Capture rank now — dist may be torn down before session teardown. + is_rank_zero = (not dist.is_initialized()) or dist.get_rank() == 0 + if update: git_sha = _git_sha() store.set_metadata( @@ -272,7 +276,7 @@ def baseline_store( yield store - if update and (not dist.is_initialized() or dist.get_rank() == 0): + if update and is_rank_zero: store.flush() From e494c827fc7179d5624f91894d335ce98284bb47 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 14 Feb 2026 19:25:18 +0000 Subject: [PATCH 8/9] Use launcher env vars for rank instead of dist API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous approach of capturing dist.get_rank() at fixture setup still fails when the session fixture is created before init_process_group — dist.is_initialized() is false on every rank, so is_rank_zero becomes true everywhere, reintroducing the last-writer-wins race. Read RANK (set by torchrun/torch.distributed.launch) or SLURM_PROCID (set by SLURM) instead. These env vars are set by the launcher before the process starts and survive the entire process lifetime, independent of the dist init/teardown lifecycle. Falls back to 0 for single-GPU runs where neither is set. https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/performance/conftest.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py index 63beb673..c84d78d2 100644 --- a/tests/performance/conftest.py +++ b/tests/performance/conftest.py @@ -256,8 +256,11 @@ def baseline_store( update = request.config.getoption("--perf-update-baselines") - # Capture rank now — dist may be torn down before session teardown. - is_rank_zero = (not dist.is_initialized()) or dist.get_rank() == 0 + # Determine rank from env vars set by the launcher (torchrun / SLURM). + # dist may not yet be initialised (session fixture created early) or + # already torn down (cleanup_dist is module-scoped), so we must not + # rely on dist.is_initialized() / dist.get_rank(). + is_rank_zero = int(os.environ.get("RANK", os.environ.get("SLURM_PROCID", "0"))) == 0 if update: git_sha = _git_sha() From a34d951e0120418cb5e1156e71af453ddab07973 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 14 Feb 2026 19:44:03 +0000 Subject: [PATCH 9/9] Fix black formatting for rank env var lookup https://claude.ai/code/session_01W6qwSimLLJTSLXReAQ5Au2 --- tests/performance/conftest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/performance/conftest.py b/tests/performance/conftest.py index c84d78d2..f3cca2d3 100644 --- a/tests/performance/conftest.py +++ b/tests/performance/conftest.py @@ -260,7 +260,8 @@ def baseline_store( # dist may not yet be initialised (session fixture created early) or # already torn down (cleanup_dist is module-scoped), so we must not # rely on dist.is_initialized() / dist.get_rank(). - is_rank_zero = int(os.environ.get("RANK", os.environ.get("SLURM_PROCID", "0"))) == 0 + _rank = os.environ.get("RANK", os.environ.get("SLURM_PROCID", "0")) + is_rank_zero = int(_rank) == 0 if update: git_sha = _git_sha()