From 1803f5704a4cf7521fd17c9c68a0cd831a9d32ee Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Wed, 19 Aug 2026 23:10:31 +0000 Subject: [PATCH] [Common] Add benchmarkable unit test infrastructure A test returns a Case (setup/evaluate/reference/verify) instead of asserting inline, so one test serves as both a correctness test and a benchmark. Default pytest runs evaluate and reference once each and lets verify assert. --nvte-benchmark times them over repeated runs for median metrics, gating each point on a one-time correctness check first, and writes JSON/JSONL/CSV reports that transformer_engine.common.testing.compare diffs against a baseline to flag regressions. @benchmark(argnames, values) marks a test Case-bearing and declares the values one axis takes when benchmarking. Those values are substituted into the test's existing pytest.mark.parametrize rather than adding an axis, so undeclared axes keep their correctness values and correctness-mode node IDs are unchanged. @benchmark.skip and @benchmark.skipif mark a Case-bearing test that is never benchmarked, which is how a correctness-only Case is written. On a class the declaration is shared by every test method, and each method's return value decides: a Case is run by the harness, None is left to pytest. The plugin autoloads as the nvte-benchmark pytest11 entry point, and its options carry an --nvte- prefix to stay out of pytest-benchmark's namespace. Standard QA suites run default pytest and so never benchmark. This lays the foundation and converts two tests as narrow examples; porting the rest of the suite and migrating the standalone benchmark scripts follow separately. Added: transformer_engine/common/testing/__init__.py transformer_engine/common/testing/case.py Case contract, axis renderer transformer_engine/common/testing/decorator.py the benchmark decorator transformer_engine/common/testing/declaration.py axis declarations transformer_engine/common/testing/plugin.py pytest hooks, mode selection transformer_engine/common/testing/runner.py Case execution and records transformer_engine/common/testing/timing.py wall-clock sampler transformer_engine/common/testing/device.py cuda-python device access transformer_engine/common/testing/artifacts.py report writers transformer_engine/common/testing/compare.py baseline comparison CLI docs/examples/benchmarkable_tests.rst Modified: setup.py nvte-benchmark entry point, cuda-python dep pyproject.toml a returning test is an error, not a warning tests/jax/pytest.ini the same, for the JAX suites' -c runs docs/index.rst toctree entry tests/pytorch/test_fused_rope.py converted in place as an example tests/jax/test_softmax.py converted in place as an example Signed-off-by: Alp Dener --- docs/examples/benchmarkable_tests.rst | 112 +++++ docs/index.rst | 1 + pyproject.toml | 5 + setup.py | 7 +- tests/jax/pytest.ini | 3 + tests/jax/test_softmax.py | 45 +- tests/pytorch/test_fused_rope.py | 153 ++++--- transformer_engine/common/testing/__init__.py | 17 + .../common/testing/artifacts.py | 358 ++++++++++++++++ transformer_engine/common/testing/case.py | 104 +++++ transformer_engine/common/testing/compare.py | 267 ++++++++++++ .../common/testing/declaration.py | 48 +++ .../common/testing/decorator.py | 85 ++++ transformer_engine/common/testing/device.py | 190 +++++++++ transformer_engine/common/testing/plugin.py | 393 ++++++++++++++++++ transformer_engine/common/testing/runner.py | 166 ++++++++ transformer_engine/common/testing/timing.py | 56 +++ 17 files changed, 1939 insertions(+), 71 deletions(-) create mode 100644 docs/examples/benchmarkable_tests.rst create mode 100644 transformer_engine/common/testing/__init__.py create mode 100644 transformer_engine/common/testing/artifacts.py create mode 100644 transformer_engine/common/testing/case.py create mode 100644 transformer_engine/common/testing/compare.py create mode 100644 transformer_engine/common/testing/declaration.py create mode 100644 transformer_engine/common/testing/decorator.py create mode 100644 transformer_engine/common/testing/device.py create mode 100644 transformer_engine/common/testing/plugin.py create mode 100644 transformer_engine/common/testing/runner.py create mode 100644 transformer_engine/common/testing/timing.py diff --git a/docs/examples/benchmarkable_tests.rst b/docs/examples/benchmarkable_tests.rst new file mode 100644 index 0000000000..825f49ce7b --- /dev/null +++ b/docs/examples/benchmarkable_tests.rst @@ -0,0 +1,112 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Benchmarkable Tests +=================== + +A benchmarkable test is an ordinary pytest test that returns a ``Case`` instead of asserting, so +one definition of setup, evaluation, reference and verification serves both correctness testing +and benchmarking. Running pytest normally checks correctness; adding ``--nvte-benchmark`` times +the same code instead. + +Writing the test +---------------- + +Build a ``Case`` from four callables and return it: + +.. code-block:: python + + from transformer_engine.common.testing import Case, benchmark + + def test_something(shape, dtype): + def setup(): + return make_inputs(shape, dtype) # deterministic + + def evaluate(state): + return te_implementation(state) # the Transformer Engine path + + def reference(state): + return naive_implementation(state) # what it should agree with + + def verify(actual, expected): + torch.testing.assert_close(actual, expected, **dtype_tols(dtype)) + + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify) + +``setup`` must be deterministic, because benchmark mode calls it again for each timed variant. +``verify`` is required whenever ``reference`` is set; there is no default comparator, so build one +on ``tests/pytorch/utils.py::dtype_tols`` or ``tests/jax/utils.py::assert_allclose``. Raise +``CaseSkip`` from ``setup`` when a backend or architecture is unavailable and the test is skipped. + +Optional fields: ``reset(state)`` runs between timed samples for cases that mutate their state, +``time_reference=False`` records only the Transformer Engine path, and ``bytes_moved`` / ``flops`` +add ``bandwidth_GBps`` and ``tflops`` to the recorded numbers. + +Marking it for benchmarking +--------------------------- + +``@benchmark(argnames, values)`` gives an axis the values it should take when benchmarking. It +does not create an axis: the values replace those of an existing ``pytest.mark.parametrize`` with +the same argnames, so correctness parametrization is untouched. Coupled argnames are written +exactly as parametrized (``"m,n,k"``). + +Abridged from ``tests/pytorch/test_fused_rope.py``: + +.. code-block:: python + + @benchmark("dtype", [torch.bfloat16]) + @benchmark("seq_length", [8192]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("seq_length", [2048, 4096]) + def test_fused_rope(dtype, seq_length): + ... + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify) + +An axis you do not declare keeps its full correctness values, so declare enough of them to keep +the benchmark matrix small -- a single benchmark shape usually means pinning most axes to one +value each. + +``@benchmark`` also applies to a class, where it covers every test method: + +.. code-block:: python + + @benchmark("b,s_q,s_kv,h", [(8, 2048, 2048, 16)]) + @pytest.mark.parametrize("b, s_q, s_kv, h", [...]) + class TestSoftmaxPrimitives: + @staticmethod + def test_forward(b, s_q, s_kv, h, dtype): + ... + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify) + + @staticmethod + @benchmark.skip(reason="returns no Case") + def test_backward(b, s_q, s_kv, h, dtype): + ... + +Use ``@benchmark.skip`` or ``@benchmark.skipif(condition)`` for a test that returns a ``Case`` but +should not be benchmarked -- a correctness-only test written in this style, or one whose benchmark +you are temporarily disabling. + +Running benchmarks +------------------ + +.. code-block:: shell + + python3 -m pytest tests/pytorch/test_fused_rope.py --nvte-benchmark \ + --nvte-benchmark-report-dir /tmp/te-bench + +``--nvte-benchmark`` selects benchmark mode and deselects everything else. Each point is checked +for correctness once before it is timed, so a benchmark run also verifies the shapes it measures. + +Options, with defaults: + +* ``--nvte-benchmark-iterations`` (20) -- minimum timed samples per variant. +* ``--nvte-benchmark-warmup`` (5) -- untimed calls before sampling. +* ``--nvte-benchmark-inner-iterations`` (1) -- calls per timed sample. Raise it for kernels short + enough that host launch latency dominates. +* ``--nvte-benchmark-min-run-time`` (0.0) -- keep sampling until this many seconds have elapsed. +* ``--nvte-benchmark-no-reference`` (off) -- skip timing the reference variant. +* ``--nvte-benchmark-report-dir`` (unset) -- where to write the JSON, JSONL and CSV reports. + Without it, the collected numbers are discarded with a warning. diff --git a/docs/index.rst b/docs/index.rst index fcd15a7a11..aeac2e8282 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -59,6 +59,7 @@ Transformer Engine documentation examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb examples/onnx/onnx_export.ipynb examples/te_jax_integration.rst + examples/benchmarkable_tests.rst examples/op_fuser/op_fuser.rst examples/gemm_profiling/gemm_profiling.rst diff --git a/pyproject.toml b/pyproject.toml index 2c9f224c14..1ecedfa2d1 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,3 +7,8 @@ requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "nin # Use legacy backend to import local packages in setup.py build-backend = "setuptools.build_meta:__legacy__" + +[tool.pytest.ini_options] +# Benchmarkable tests return a Case that only the nvte-benchmark plugin runs; without +# it pytest discards the Case and the test passes having verified nothing. +filterwarnings = ["error::pytest.PytestReturnNotNoneWarning"] diff --git a/setup.py b/setup.py index 2a8a9d7688..02af4736de 100644 --- a/setup.py +++ b/setup.py @@ -124,7 +124,7 @@ def setup_requirements() -> Tuple[List[str], List[str]]: "importlib-metadata>=1.0", "packaging", ] - test_reqs: List[str] = ["pytest>=8.2.1"] + test_reqs: List[str] = ["pytest>=8.2.1", "cuda-python>=12.0"] # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): @@ -403,6 +403,11 @@ def git_check_submodules() -> None: ], ), extras_require=extras_require, + entry_points={ + "pytest11": [ + "nvte-benchmark = transformer_engine.common.testing.plugin", + ], + }, description="Transformer acceleration library", long_description=long_description, long_description_content_type="text/x-rst", diff --git a/tests/jax/pytest.ini b/tests/jax/pytest.ini index 490671a631..73cb71088e 100644 --- a/tests/jax/pytest.ini +++ b/tests/jax/pytest.ini @@ -27,3 +27,6 @@ filterwarnings= ignore:Scan loop is disabled for fused ring attention.*:UserWarning ignore:jax.extend.ffi.register_ffi_target is deprecated ignore:jax.extend.ffi.ffi_lowering is deprecated + # Benchmarkable tests return a Case that only the benchmarkable plugin runs; without + # it pytest discards the Case and the test passes having verified nothing. + error::pytest.PytestReturnNotNoneWarning diff --git a/tests/jax/test_softmax.py b/tests/jax/test_softmax.py index 7af9613538..d61b7b476d 100644 --- a/tests/jax/test_softmax.py +++ b/tests/jax/test_softmax.py @@ -16,6 +16,7 @@ from utils import assert_allclose +from transformer_engine.common.testing import Case, benchmark from transformer_engine.jax.cpp_extensions import is_softmax_kernel_available from transformer_engine.jax.cpp_extensions.attention import AttnSoftmaxType from transformer_engine.jax.softmax import SoftmaxFusionType, softmax @@ -98,15 +99,6 @@ def _setup_inputs(self): case _: raise ValueError(f"Unknown {self.softmax_fusion_type=}") - def test_forward(self): - """ - Test transformer_engine.jax.softmax.softmax fwd rule - """ - self._setup_inputs() - primitive_out = softmax(self.logits, self.mask, self.scale_factor, self.softmax_fusion_type) - reference_out = __class__.reference_softmax(self.logits, self.mask, self.scale_factor) - assert_allclose(primitive_out, reference_out, dtype=self.dtype) - def test_backward(self): """ Test transformer_engine.jax.softmax.softmax bwd rule @@ -148,10 +140,6 @@ class SoftmaxPrimitivesRunner(SoftmaxRunner): Jax Softmax Primitives runner """ - @catch_unsupported - def test_forward(self): - return super().test_forward() - @catch_unsupported def test_backward(self): return super().test_backward() @@ -187,6 +175,8 @@ def test_forward(self): # Run softmax primitives test +# The pinned shape must be one the fused kernel supports, or the benchmark times a raise. +@benchmark("b,s_q,s_kv,h", [(8, 2048, 2048, 16)]) @pytest.mark.parametrize( "b, s_q, s_kv, h", [ @@ -222,9 +212,36 @@ def test_forward(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype): Test forward with parameterized configs """ runner = SoftmaxPrimitivesRunner(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype) - runner.test_forward() + # Resolved here, outside the timed callables; depends only on construction-time fields. + supported = runner._is_support() + + def setup(): + runner._setup_inputs() + return runner + + def evaluate(state): + # Unsupported configs must raise from the primitive rather than compute a result. + if not supported: + with pytest.raises(AssertionError): + softmax(state.logits, state.mask, state.scale_factor, state.softmax_fusion_type) + return None + return softmax(state.logits, state.mask, state.scale_factor, state.softmax_fusion_type) + + def reference(state): + if not supported: + return None + return state.reference_softmax(state.logits, state.mask, state.scale_factor) + + def verify(actual, expected): + # Unsupported configs are checked by the expected raise in evaluate(). + if not supported: + return + assert_allclose(actual, expected, dtype=dtype) + + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify) @staticmethod + @benchmark.skip(reason="returns no Case; keeps the class's axes out of collection") def test_backward(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype): """ Test forward with parameterized configs diff --git a/tests/pytorch/test_fused_rope.py b/tests/pytorch/test_fused_rope.py index 5cc4fc0f8a..951ee8e463 100644 --- a/tests/pytorch/test_fused_rope.py +++ b/tests/pytorch/test_fused_rope.py @@ -5,6 +5,7 @@ import math import torch import pytest +from transformer_engine.common.testing import Case, benchmark from transformer_engine.pytorch.attention.rope import ( RotaryPositionEmbedding, apply_rotary_pos_emb, @@ -29,6 +30,16 @@ def _non_overlapping_grad(output: Union[List[torch.Tensor], torch.Tensor]) -> to return torch.sum(output * t) +# start_positions must stay False here: True would hit the margin==0 skip below and +# drop half of the inherited margin values. +@benchmark("start_positions", [False]) +@benchmark("dtype", [torch.bfloat16]) +@benchmark("seq_length", [8192]) +@benchmark("hidden_size", [128]) +@benchmark("transpose", [None]) +@benchmark("tensor_format", ["sbhd"]) +@benchmark("cp_size", [1, 2]) +@benchmark("interleaved", [True]) @pytest.mark.parametrize("start_positions", [True, False]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) @pytest.mark.parametrize("seq_length", [2048, 4096]) @@ -52,7 +63,7 @@ def test_fused_rope( cp_size: int, interleaved: bool, start_positions: bool, -) -> None: +): if margin == 0 and start_positions == True: # This makes sure that the `start_positions` offsets being applied # are with the maximum length of the rope embeddings. @@ -60,67 +71,97 @@ def test_fused_rope( device = torch.device("cuda:0") batch_size, head_num = 2, 64 - t = torch.rand( - (seq_length - margin, batch_size, head_num, hidden_size), - dtype=dtype, - device=device, - ) - # Get arbitrary offsets to be used with RoPE for all the sequences - start_positions = ( - torch.randint(0, margin, (batch_size,), dtype=torch.int32, device=device) - if start_positions - else None - ) - - if tensor_format == "bshd": - t = t.transpose(0, 1).contiguous() - if transpose: - t = t.transpose(*transpose).contiguous().transpose(*transpose) - t.requires_grad = True + def setup(): + t = torch.rand( + (seq_length - margin, batch_size, head_num, hidden_size), + dtype=dtype, + device=device, + ) - rotary_pos_emb = RotaryPositionEmbedding(hidden_size, rotary_percent, interleaved=interleaved) - emb = rotary_pos_emb(seq_length * cp_size) - assert emb.is_contiguous() + # Get arbitrary offsets to be used with RoPE for all the sequences + sp = ( + torch.randint(0, margin, (batch_size,), dtype=torch.int32, device=device) + if start_positions + else None + ) - for cp_rank in range(cp_size): - # unfused - # The fused kernel computes in float32 internally, so we force the unfused func to use float32 - # for more accurate comparison - output_unfused = apply_rotary_pos_emb( - t.float(), - emb, - tensor_format=tensor_format, - start_positions=start_positions, - interleaved=interleaved, - fused=False, - cp_size=cp_size, - cp_rank=cp_rank, - ).to(dtype) - loss_unfused = loss_func(output_unfused) - loss_unfused.backward() - grad_unfused = t.grad.detach().clone() - t.grad = None + if tensor_format == "bshd": + t = t.transpose(0, 1).contiguous() + if transpose: + t = t.transpose(*transpose).contiguous().transpose(*transpose) + t.requires_grad = True - # fused - output_fused = apply_rotary_pos_emb( - t, - emb, - tensor_format=tensor_format, - start_positions=start_positions, - interleaved=interleaved, - fused=True, - cp_size=cp_size, - cp_rank=cp_rank, + rotary_pos_emb = RotaryPositionEmbedding( + hidden_size, rotary_percent, interleaved=interleaved ) - loss_fused = loss_func(output_fused) - loss_fused.backward() - grad_fused = t.grad.detach().clone() - t.grad = None + emb = rotary_pos_emb(seq_length * cp_size) + assert emb.is_contiguous() + return {"t": t, "start_positions": sp, "emb": emb} + + def reset(state): + # A no-op here (the cp_rank loop already clears `t.grad`); declared to + # exercise the runner's reset contract, whose presence pins this case to + # batchable=False / inner_iterations=1. + state["t"].grad = None + + def evaluate(state): + t, emb, sp = state["t"], state["emb"], state["start_positions"] + outputs = [] + for cp_rank in range(cp_size): + output_fused = apply_rotary_pos_emb( + t, + emb, + tensor_format=tensor_format, + start_positions=sp, + interleaved=interleaved, + fused=True, + cp_size=cp_size, + cp_rank=cp_rank, + ) + loss_fused = loss_func(output_fused) + loss_fused.backward() + grad_fused = t.grad.detach() + # detach() without clone(): the `t.grad = None` below is what keeps the + # capture alias-free -- the next backward() installs a fresh buffer + # instead of accumulating into what `grad_fused` aliases. Removing or + # reordering it silently corrupts every earlier capture. + t.grad = None + outputs.append((output_fused, grad_fused)) + return tuple(outputs) + + def reference(state): + t, emb, sp = state["t"], state["emb"], state["start_positions"] + outputs = [] + for cp_rank in range(cp_size): + # The fused kernel computes in float32 internally, so we force the unfused func to use float32 + # for more accurate comparison + output_unfused = apply_rotary_pos_emb( + t.float(), + emb, + tensor_format=tensor_format, + start_positions=sp, + interleaved=interleaved, + fused=False, + cp_size=cp_size, + cp_rank=cp_rank, + ).to(dtype) + loss_unfused = loss_func(output_unfused) + loss_unfused.backward() + grad_unfused = t.grad.detach() + # detach() without clone(); as in evaluate(), the `t.grad = None` below + # is what keeps the capture alias-free. + t.grad = None + outputs.append((output_unfused, grad_unfused)) + return tuple(outputs) + + def verify(actual, expected): + for (output_fused, grad_fused), (output_unfused, grad_unfused) in zip(actual, expected): + torch.testing.assert_close(output_fused, output_unfused) + torch.testing.assert_close(grad_fused, grad_unfused) + assert output_fused.is_contiguous() - torch.testing.assert_close(output_fused, output_unfused) - torch.testing.assert_close(grad_fused, grad_unfused) - assert output_fused.is_contiguous() + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify, reset=reset) @pytest.mark.parametrize("margin", [10]) diff --git a/transformer_engine/common/testing/__init__.py b/transformer_engine/common/testing/__init__.py new file mode 100644 index 0000000000..985472819e --- /dev/null +++ b/transformer_engine/common/testing/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Unified correctness and benchmarking test utilities.""" + +from .case import Case, CaseSkip + +__all__ = ["Case", "CaseSkip", "benchmark"] # pylint: disable=undefined-all-variable + + +def __getattr__(name): + """Resolve ``benchmark`` lazily so importing this package does not require pytest.""" + if name == "benchmark": + from .decorator import benchmark as _benchmark + + return _benchmark + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/transformer_engine/common/testing/artifacts.py b/transformer_engine/common/testing/artifacts.py new file mode 100644 index 0000000000..5e0ae11048 --- /dev/null +++ b/transformer_engine/common/testing/artifacts.py @@ -0,0 +1,358 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Artifact helpers for benchmarkable runs.""" + +from __future__ import annotations + +import csv +from datetime import datetime, timezone +import json +import os +from pathlib import Path +import platform +import socket +import subprocess +import sys +from typing import Any + +from .case import axis_value +from .device import build_architectures, device_metadata + + +REPORT_SCHEMA_VERSION = "benchmark_report/v1" + +# .../transformer_engine/common/testing/artifacts.py -> testing -> common +# -> transformer_engine -> repo root +REPO_ROOT = Path(__file__).resolve().parents[3] + + +def write_run_artifacts( + output_dir: Path, + records: list[dict[str, Any]], + command: list[str], + selection: dict[str, Any], + sharding: dict[str, Any] | None = None, + report_name: str = "benchmark_report.json", + records_name: str = "benchmark_records.jsonl", + summary_name: str = "benchmark_summary.csv", +) -> dict[str, Path]: + """Write JSON, JSONL and CSV artifacts for one benchmark run.""" + output_dir.mkdir(parents=True, exist_ok=True) + report = build_report(records, command, selection, sharding=sharding) + + report_path = output_dir / report_name + records_path = output_dir / records_name + summary_path = output_dir / summary_name + + _write_json(report_path, report) + _write_jsonl(records_path, records) + _write_summary_csv(summary_path, records) + return {"report": report_path, "records": records_path, "summary": summary_path} + + +def build_report( + records: list[dict[str, Any]], + command: list[str], + selection: dict[str, Any], + sharding: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Build a self-contained machine-readable benchmark report.""" + status_counts: dict[str, int] = {} + for record in records: + status = str(record.get("status", "unknown")) + status_counts[status] = status_counts.get(status, 0) + 1 + + return { + "schema_version": REPORT_SCHEMA_VERSION, + "created_at_utc": datetime.now(timezone.utc).isoformat(), + "command": command, + "selection": selection, + "environment": collect_environment(command), + "sharding": sharding or {"enabled": False}, + "summary": { + "record_count": len(records), + "status_counts": status_counts, + }, + "records": records, + } + + +def collect_environment(command: list[str] | None = None) -> dict[str, Any]: + """Collect stable environment metadata without copying arbitrary environment variables. + + Framework versions are read from already-imported modules; never import one here. + ``build`` holds the architectures embedded in the loaded ``libtransformer_engine.so``. + """ + return { + "python": { + "version": sys.version, + "executable": sys.executable, + }, + "platform": { + "system": platform.system(), + "release": platform.release(), + "machine": platform.machine(), + "processor": platform.processor(), + "python_implementation": platform.python_implementation(), + }, + "host": { + "hostname": socket.gethostname(), + "cpu_count": os.cpu_count(), + }, + "git": _git_metadata(), + "frameworks": _framework_versions(), + "devices": {"cuda": device_metadata()}, + "build": build_architectures(), + "scheduler": _scheduler_metadata(), + "command": command or [], + } + + +def merge_worker_reports( + output_dir: Path, + worker_report_paths: list[Path], + command: list[str], + selection: dict[str, Any], + sharding: dict[str, Any], +) -> dict[str, Path]: + """Merge worker reports into the standard top-level artifact set. + + A worker report that cannot be read is recorded in the output rather than raising. + """ + merged_records: list[dict[str, Any]] = [] + worker_summaries = [] + unreadable_reports = [] + for report_path in sorted(worker_report_paths): + try: + with report_path.open("r", encoding="utf-8") as handle: + report = json.load(handle) + except (OSError, json.JSONDecodeError) as exc: + unreadable_reports.append({"path": str(report_path), "reason": str(exc)}) + continue + worker_summaries.append( + { + "path": str(report_path), + "summary": report.get("summary", {}), + "selection": report.get("selection", {}), + } + ) + merged_records.extend(report.get("records", [])) + + expected_records = _expected_records_from_worker_summaries(worker_summaries) + sharding = dict(sharding) + sharding["worker_reports"] = worker_summaries + sharding["unreadable_worker_reports"] = unreadable_reports + sharding["merge_validation"] = _validate_merged_records( + merged_records, + expected_records=expected_records, + unreadable_reports=unreadable_reports, + ) + return write_run_artifacts( + output_dir, + merged_records, + command, + selection, + sharding=sharding, + ) + + +def record_key(record: dict[str, Any]) -> str: + """Return a deterministic key for comparing benchmark records. + + ``params`` goes through ``axis_value``. Unlike the record writers, this ``json.dumps`` + has no ``default=str``, so a value it did not normalize raises here instead of becoming + a key that differs every run. + """ + key = { + "case_id": record.get("case_id"), + "framework": record.get("framework"), + "operation": record.get("operation"), + "variant": record.get("variant"), + "params": {name: axis_value(value) for name, value in record.get("params", {}).items()}, + } + return json.dumps(key, sort_keys=True, separators=(",", ":")) + + +def _validate_merged_records( + records: list[dict[str, Any]], + expected_records: list[dict[str, Any]] | None = None, + unreadable_reports: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + counts: dict[str, int] = {} + for record in records: + key = record_key(record) + counts[key] = counts.get(key, 0) + 1 + duplicates = sorted(key for key, count in counts.items() if count > 1) + missing = [] + if expected_records is not None: + expected_keys = sorted({record_key(record) for record in expected_records}) + missing = sorted(key for key in expected_keys if key not in counts) + unreadable = unreadable_reports or [] + return { + "record_count": len(records), + "unique_record_count": len(counts), + "duplicate_keys": duplicates, + "missing_keys": missing, + "unreadable_report_count": len(unreadable), + "valid": not duplicates and not missing and not unreadable, + } + + +def _expected_records_from_worker_summaries( + worker_summaries: list[dict[str, Any]], +) -> list[dict[str, Any]]: + expected_records = [] + for worker in worker_summaries: + selection = worker.get("selection", {}) + include_reference = selection.get("include_reference", True) + for case in selection.get("selected_cases", []): + variants = [] + if include_reference and case.get("has_reference"): + variants.append("reference") + variants.append("evaluation") + for variant in variants: + expected_records.append( + { + "case_id": case.get("case_id"), + "framework": case.get("framework"), + "operation": case.get("operation"), + "variant": variant, + "params": case.get("params", {}), + } + ) + return expected_records + + +def _write_json(path: Path, data: dict[str, Any]) -> None: + with path.open("w", encoding="utf-8") as handle: + # default=str: ``params`` hold pytest parametrize values, which are not always + # JSON-native (a JAX dtype is a bare type object, not a serializable instance). + json.dump(data, handle, indent=2, sort_keys=True, default=str) + handle.write("\n") + + +def _write_jsonl(path: Path, records: list[dict[str, Any]]) -> None: + with path.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, sort_keys=True, default=str)) + handle.write("\n") + + +def _write_summary_csv(path: Path, records: list[dict[str, Any]]) -> None: + fields = [ + "status", + "framework", + "case_id", + "variant", + "component", + "operation", + "median_ms", + "mean_ms", + "p95_ms", + "reason", + ] + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + for record in records: + timing = record.get("timing", {}) + writer.writerow( + { + "status": record.get("status"), + "framework": record.get("framework"), + "case_id": record.get("case_id"), + "variant": record.get("variant"), + "component": record.get("component"), + "operation": record.get("operation"), + "median_ms": timing.get("median_ms"), + "mean_ms": timing.get("mean_ms"), + "p95_ms": timing.get("p95_ms"), + "reason": record.get("reason"), + } + ) + + +def _git_metadata() -> dict[str, Any]: + return { + "commit": _run_git(["rev-parse", "HEAD"]), + "branch": _run_git(["rev-parse", "--abbrev-ref", "HEAD"]), + "dirty": bool(_run_git(["status", "--porcelain"])), + } + + +def _run_git(args: list[str]) -> str | None: + try: + result = subprocess.run( + ["git", *args], + check=False, + capture_output=True, + encoding="utf-8", + cwd=REPO_ROOT, + ) + except OSError: + return None + if result.returncode != 0: + return None + return result.stdout.strip() + + +def _framework_versions() -> dict[str, str | None]: + versions: dict[str, str | None] = { + "transformer_engine": _module_version("transformer_engine"), + "torch": _module_version("torch"), + "jax": _module_version("jax"), + "jaxlib": _module_version("jaxlib"), + } + return versions + + +def _module_version(module_name: str) -> str | None: + module = sys.modules.get(module_name) + if module is None: + return None + return getattr(module, "__version__", None) + + +def _scheduler_metadata() -> dict[str, Any]: + names = [ + "CUDA_VISIBLE_DEVICES", + "SLURM_JOB_ID", + "SLURM_JOB_GPUS", + "SLURM_GPUS", + "SLURM_GPUS_ON_NODE", + "SLURM_STEP_GPUS", + ] + metadata = {name: os.environ.get(name) for name in names if os.environ.get(name) is not None} + visible_devices = visible_cuda_devices() + allocated_devices = scheduler_allocated_devices() + metadata.update( + { + "visible_cuda_devices": visible_devices, + "visible_gpu_count": len(visible_devices), + "scheduler_allocated_devices": allocated_devices, + "scheduler_allocated_gpu_count": len(allocated_devices), + } + ) + return metadata + + +def scheduler_allocated_devices() -> list[str]: + """Return the device list the job scheduler allocated to this process.""" + for name in ("SLURM_STEP_GPUS", "SLURM_JOB_GPUS", "CUDA_VISIBLE_DEVICES"): + devices = parse_device_list(os.environ.get(name)) + if devices: + return devices + return [] + + +def visible_cuda_devices() -> list[str]: + """Return the devices named by ``CUDA_VISIBLE_DEVICES``.""" + return parse_device_list(os.environ.get("CUDA_VISIBLE_DEVICES")) + + +def parse_device_list(raw: str | None) -> list[str]: + """Split a comma-separated device list, dropping empty entries.""" + if raw is None: + return [] + return [item.strip() for item in raw.split(",") if item.strip()] diff --git a/transformer_engine/common/testing/case.py b/transformer_engine/common/testing/case.py new file mode 100644 index 0000000000..d38ad9ab58 --- /dev/null +++ b/transformer_engine/common/testing/case.py @@ -0,0 +1,104 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""The Case contract returned by a benchmarkable test, and its axis_value renderer.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +import enum +import os +from typing import Any + +from .declaration import plugin_active + + +def _stable_name(value: Any) -> str | None: + """Return ``value``'s ``__qualname__``, else its ``__name__``, else ``None``.""" + for attribute in ("__qualname__", "__name__"): + name = getattr(value, attribute, None) + if isinstance(name, str): + return name + return None + + +def axis_value(value: Any) -> Any: + """Render one parametrize value canonically for artifact identity.""" + # The enum check must precede the int check: IntEnum is an int. + if isinstance(value, enum.Enum): + return value.name + if isinstance(value, type): + return value.__qualname__ + if value is None or isinstance(value, (bool, int, float, str)): + return value + name = _stable_name(value) + if name is not None: + return name + rendered = str(value) + # An address-bearing rendering is not stable across processes: reject it. + if " at 0x" in rendered: + raise ValueError( + f"Benchmark axis value {rendered!r} has no stable rendering: its repr embeds " + "a memory address, so no two runs would agree on its case_id. Parametrize " + "this axis with an enum or a named function/class instead of a bare object." + ) + return rendered + + +class CaseSkip(Exception): + """Raised by a Case callable when a feature is unavailable; the plugin skips.""" + + +def _require_plugin() -> None: + """Refuse to build a Case inside a pytest session that has no plugin to run it.""" + # pytest sets PYTEST_CURRENT_TEST only while a test is executing, so this fires + # exactly when a Case would be discarded and its test pass having verified nothing. + # It lives here, not in the decorators, because a Case-bearing test need not be + # decorated at all, and a class-level decorator never wraps its methods. + current = os.environ.get("PYTEST_CURRENT_TEST") + if not current or plugin_active(): + return + raise RuntimeError( + f"{current} built a Case, but the nvte-benchmark pytest plugin is not registered," + " so nothing would run it and this test asserts nothing. It autoloads from" + " transformer_engine's entry point, so reinstall transformer_engine (nvte-setup)" + " if pytest is not picking it up, and drop any '-p no:nvte-benchmark'." + ) + + +@dataclass +class Case: + """What a benchmarkable test returns: ``setup`` builds an opaque state object, + ``evaluate`` is the Transformer Engine path, and ``reference`` is the comparison + target in correctness mode and a timed baseline in benchmark mode. + """ + + setup: Callable[[], Any] + evaluate: Callable[[Any], Any] + reference: Callable[[Any], Any] | None = None + verify: Callable[[Any, Any], None] | None = None + reset: Callable[[Any], None] | None = None + batchable: bool = True + time_reference: bool = True + bytes_moved: int | None = None + flops: int | None = None + + # Setting `reset` implies not batchable: ``runner.py`` pins inner_iterations to 1. + + def __post_init__(self) -> None: + """Validate the plugin is present to run this Case, and that ``reference`` and + ``verify`` are set together.""" + _require_plugin() + if self.reference is None and self.verify is not None: + raise ValueError("Case defines verify but no reference to compare against.") + if self.reference is not None and self.verify is None: + raise ValueError( + "Case defines a reference but no verify, and there is no default " + "comparator. Pass a verify built on the framework's tolerance helper " + "(tests/pytorch/utils.py::dtype_tols, tests/jax/utils.py::assert_allclose)." + ) + + def run_verify(self, actual: Any, expected: Any) -> None: + """Compare one evaluate output against one reference output.""" + self.verify(actual, expected) diff --git a/transformer_engine/common/testing/compare.py b/transformer_engine/common/testing/compare.py new file mode 100644 index 0000000000..97b70c067d --- /dev/null +++ b/transformer_engine/common/testing/compare.py @@ -0,0 +1,267 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Compare benchmarkable reports against historical artifacts.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import sys +from typing import Any + +from .artifacts import record_key + + +DEFAULT_RELATIVE_THRESHOLD = 0.05 +DEFAULT_ABSOLUTE_THRESHOLD_MS = 0.01 + + +def compare_reports( + baseline_report: dict[str, Any], + current_report: dict[str, Any], + relative_threshold: float = DEFAULT_RELATIVE_THRESHOLD, + absolute_threshold_ms: float = DEFAULT_ABSOLUTE_THRESHOLD_MS, +) -> dict[str, Any]: + """Compare completed records in two reports.""" + baseline_records = _completed_records_by_key(baseline_report) + current_records = _completed_records_by_key(current_report) + + regressions = [] + improvements = [] + unchanged = [] + incompatible = [] + hardware_compatible, incompatible_reason = _hardware_compatible(baseline_report, current_report) + + for key, current in sorted(current_records.items()): + baseline = baseline_records.get(key) + if baseline is None: + continue + if not hardware_compatible: + incompatible.append({"key": key, "reason": incompatible_reason}) + continue + + baseline_ms = _median_ms(baseline) + current_ms = _median_ms(current) + delta_ms = current_ms - baseline_ms + relative_delta = delta_ms / baseline_ms if baseline_ms > 0 else 0.0 + threshold = _threshold_for_record(current, relative_threshold, absolute_threshold_ms) + entry = { + "key": key, + "case_id": current.get("case_id"), + "variant": current.get("variant"), + "params": current.get("params", {}), + "baseline_median_ms": baseline_ms, + "current_median_ms": current_ms, + "delta_ms": delta_ms, + "relative_delta": relative_delta, + "relative_threshold": threshold["relative"], + "absolute_threshold_ms": threshold["absolute_ms"], + } + if delta_ms > threshold["absolute_ms"] and relative_delta > threshold["relative"]: + regressions.append(entry) + elif delta_ms < -threshold["absolute_ms"] and -relative_delta > threshold["relative"]: + improvements.append(entry) + else: + unchanged.append(entry) + + missing = sorted(key for key in baseline_records if key not in current_records) + new = sorted(key for key in current_records if key not in baseline_records) + return { + "schema_version": "benchmark_comparison/v1", + "summary": { + "baseline_records": len(baseline_records), + "current_records": len(current_records), + "regressions": len(regressions), + "improvements": len(improvements), + "unchanged": len(unchanged), + "missing": len(missing), + "new": len(new), + "incompatible": len(incompatible), + }, + "regressions": regressions, + "improvements": improvements, + "unchanged": unchanged, + "missing": missing, + "new": new, + "incompatible": incompatible, + } + + +def main() -> int: + """Compare two benchmark_report/v1 JSON files and write a comparison report.""" + parser = argparse.ArgumentParser(description="Compare benchmarkable JSON reports.") + parser.add_argument("--baseline", required=True, help="Historical benchmark_report.json") + parser.add_argument("--current", required=True, help="Current benchmark_report.json") + parser.add_argument("--output", required=True, help="Comparison JSON output path") + parser.add_argument( + "--relative-threshold", + type=float, + default=DEFAULT_RELATIVE_THRESHOLD, + help="Default relative regression threshold (default: %(default)s)", + ) + parser.add_argument( + "--absolute-threshold-ms", + type=float, + default=DEFAULT_ABSOLUTE_THRESHOLD_MS, + help="Default absolute regression threshold in ms (default: %(default)s)", + ) + parser.add_argument( + "--fail-on-regression", + action="store_true", + help="Exit nonzero when any regression is detected.", + ) + parser.add_argument( + "--fail-on-missing", + action="store_true", + help=( + "Exit nonzero when a baseline record has no completed counterpart in the current " + "report. A case that starts erroring disappears from the comparison entirely, so " + "without this flag a broken case looks like a clean run." + ), + ) + parser.add_argument( + "--fail-on-incompatible", + action="store_true", + help=( + "Exit nonzero when any record is hardware-incompatible with its baseline " + "counterpart. Every incompatible record short-circuits past the regression test " + "and 'missing' stays zero because the key is still present, so without this flag " + "a baseline from different hardware -- e.g. an H100 baseline against a B200 run -- " + "compares as clean by default." + ), + ) + args = parser.parse_args() + + baseline = _read_json(Path(args.baseline)) + current = _read_json(Path(args.current)) + comparison = compare_reports( + baseline, + current, + relative_threshold=args.relative_threshold, + absolute_threshold_ms=args.absolute_threshold_ms, + ) + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("w", encoding="utf-8") as handle: + json.dump(comparison, handle, indent=2, sort_keys=True) + handle.write("\n") + + if comparison["incompatible"]: + reasons = sorted( + {entry["reason"] for entry in comparison["incompatible"] if entry.get("reason")} + ) + detail = f" ({'; '.join(reasons)})" if reasons else "" + print( + f"Warning: {len(comparison['incompatible'])} record(s) are hardware-incompatible " + f"with their baseline counterpart and were skipped for regression comparison{detail}.", + file=sys.stderr, + ) + + if args.fail_on_regression and comparison["regressions"]: + print( + f"Error: {len(comparison['regressions'])} regression(s) detected.", + file=sys.stderr, + ) + return 2 + if args.fail_on_missing and comparison["missing"]: + print( + f"Error: {len(comparison['missing'])} baseline record(s) missing from the current " + "report.", + file=sys.stderr, + ) + return 3 + if args.fail_on_incompatible and comparison["incompatible"]: + print( + f"Error: {len(comparison['incompatible'])} record(s) are hardware-incompatible " + "with their baseline counterpart.", + file=sys.stderr, + ) + return 4 + return 0 + + +def _read_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def _completed_records_by_key(report: dict[str, Any]) -> dict[str, dict[str, Any]]: + records = {} + for record in report.get("records", []): + if record.get("status") == "completed": + records[record_key(record)] = record + return records + + +def _median_ms(record: dict[str, Any]) -> float: + return float(record.get("timing", {}).get("median_ms", 0.0)) + + +def _threshold_for_record( + record: dict[str, Any], + default_relative: float, + default_absolute_ms: float, +) -> dict[str, float]: + override = record.get("regression_threshold") or {} + return { + "relative": float(override.get("relative", default_relative)), + "absolute_ms": float(override.get("absolute_ms", default_absolute_ms)), + } + + +def _hardware_compatible( + baseline_report: dict[str, Any], + current_report: dict[str, Any], +) -> tuple[bool, str | None]: + """Return whether the two reports were recorded on compatible hardware, plus a reason. + + Missing device metadata on either side means compatibility cannot be established, and is + treated as incompatible rather than as a match. + """ + baseline_device = _primary_device_identity(baseline_report) + current_device = _primary_device_identity(current_report) + if baseline_device is None and current_device is None: + return ( + False, + ( + "neither the baseline nor the current report carries device metadata -- " + "hardware compatibility cannot be established" + ), + ) + if baseline_device is None: + return ( + False, + ( + "the baseline report carries no device metadata -- hardware " + "compatibility cannot be established" + ), + ) + if current_device is None: + return ( + False, + ( + "the current report carries no device metadata -- hardware " + "compatibility cannot be established" + ), + ) + if baseline_device != current_device: + return False, "hardware metadata differs" + return True, None + + +def _primary_device_identity(report: dict[str, Any]) -> tuple[Any, ...] | None: + """Return the primary CUDA device's ``(name, compute_capability)``. + + Returns ``None`` when the report carries no CUDA device metadata. + """ + devices = report.get("environment", {}).get("devices", {}).get("cuda", {}).get("devices", []) + if not devices: + return None + device = devices[0] + return (device.get("name"), tuple(device.get("compute_capability", []))) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transformer_engine/common/testing/declaration.py b/transformer_engine/common/testing/declaration.py new file mode 100644 index 0000000000..eb85006279 --- /dev/null +++ b/transformer_engine/common/testing/declaration.py @@ -0,0 +1,48 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Storage for benchmark axis declarations, kept free of pytest imports.""" + +from __future__ import annotations + +from typing import Any + +DECLARATION_ATTR = "_te_benchmark_axes" + +_PLUGIN_ACTIVE = False +"""Set true by ``plugin.py``'s ``pytest_configure`` once the plugin is registered.""" + + +def set_plugin_active(active: bool) -> None: + """Record whether the benchmarkable pytest plugin is registered in this session.""" + global _PLUGIN_ACTIVE # pylint: disable=global-statement + _PLUGIN_ACTIVE = bool(active) + + +def plugin_active() -> bool: + """Return whether the benchmarkable pytest plugin registered itself this session.""" + return _PLUGIN_ACTIVE + + +def normalize_argnames(argnames: Any) -> str: + """Return a canonical comma-joined key with no spaces, so a declaration and a + ``pytest.mark.parametrize`` mark written with different spacing still match.""" + if isinstance(argnames, str): + names = [name.strip() for name in argnames.split(",") if name.strip()] + else: + names = [str(name).strip() for name in argnames] + return ",".join(names) + + +def record_axis(holder: Any, argnames: Any, values: Any) -> None: + """Attach one benchmark axis declaration to a function, method, or class.""" + declarations = holder.__dict__.get(DECLARATION_ATTR) + if declarations is None: + declarations = {} + setattr(holder, DECLARATION_ATTR, declarations) + declarations[normalize_argnames(argnames)] = list(values) + + +def declared_axes(function: Any) -> dict[str, list]: + """Return the benchmark axis declarations attached to one test function.""" + return dict(getattr(function, DECLARATION_ATTR, None) or {}) diff --git a/transformer_engine/common/testing/decorator.py b/transformer_engine/common/testing/decorator.py new file mode 100644 index 0000000000..1608291b95 --- /dev/null +++ b/transformer_engine/common/testing/decorator.py @@ -0,0 +1,85 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Benchmarkable mode constants, plugin markers, and the ``benchmark`` decorator.""" + +import pytest + +from .declaration import record_axis + +MODE_CORRECTNESS = "correctness" +"""Run the correctness matrix, verified, with nothing timed.""" + +MODE_BENCHMARK = "benchmark" +"""Run only the benchmark matrix, gated by a one-time correctness check per point.""" + +CASE_MARKER = "nvte_case" +"""Case-bearing: the test returns a Case, so the plugin runs it instead of pytest.""" + +BENCHMARK_MARKER = "nvte_benchmark" +"""Benchmark-eligible: the test declared benchmark axes with ``@benchmark(...)``.""" + +SUPPRESS_MARKER = "nvte_no_benchmark" +"""Benchmarking suppressed by ``@benchmark.skip`` or a true ``@benchmark.skipif``.""" + +MARKERS = ( + (CASE_MARKER, "test returns a Case for the nvte-benchmark plugin to run."), + (BENCHMARK_MARKER, "Case-bearing test that declared benchmark axes."), + (SUPPRESS_MARKER, "Case-bearing test that is never a benchmark point."), +) + + +def _apply(holder, *names, **kwargs): + """Apply plugin markers by name to a function, a method, or a class.""" + for name in names: + holder = getattr(pytest.mark, name)(**kwargs)(holder) + return holder + + +class _Benchmark: + """Implements the ``benchmark`` decorator. + + An instance, not a module-level function carrying attributes, so ``skip`` and + ``skipif`` are real bound methods that ``help()``, ``inspect`` and pylint resolve. + """ + + def __call__(self, argnames, values): + """Mark a test Case-bearing and benchmark-eligible, and declare the values one axis + takes in benchmark mode. + + ``argnames`` accepts the same forms as ``pytest.mark.parametrize``. The values are + substituted into an existing ``parametrize`` mark on the test, its class, or its + module; undeclared axes keep their correctness values. Applies to a function, a + method, or a class. + """ + + def decorator(holder): + record_axis(holder, argnames, values) + return _apply(holder, CASE_MARKER, BENCHMARK_MARKER) + + return decorator + + def skip(self, holder=None, *, reason=""): + """Mark a test Case-bearing but never benchmarked. Usable bare or called with a + ``reason``, on a function, a method, or a class.""" + + def decorator(target): + target = _apply(target, CASE_MARKER) + return _apply(target, SUPPRESS_MARKER, reason=reason) + + return decorator if holder is None else decorator(holder) + + def skipif(self, condition, *, reason=""): + """Mark a test Case-bearing, and suppress benchmarking when ``condition`` is true. + + A false condition declares nothing: skip and skipif only ever subtract benchmarking. + """ + + def decorator(target): + target = _apply(target, CASE_MARKER) + return _apply(target, SUPPRESS_MARKER, reason=reason) if condition else target + + return decorator + + +benchmark = _Benchmark() diff --git a/transformer_engine/common/testing/device.py b/transformer_engine/common/testing/device.py new file mode 100644 index 0000000000..edff4fb296 --- /dev/null +++ b/transformer_engine/common/testing/device.py @@ -0,0 +1,190 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""CUDA device access for benchmarkable tests, via cuda-python rather than a framework.""" + +from __future__ import annotations + +import functools +import re +import shutil +import subprocess +from typing import Any + +# Cubin sections from ``cuobjdump --list-elf``, e.g. ``libtransformer_engine.4.sm_103a.cubin``. +# The ``[af]`` group must stay: dropping the Blackwell ``a``/``f`` suffix yields the wrong arch. +_ELF_ARCH = re.compile(r"\.sm_(\d+)([af]?)\.cubin") + + +def _runtime(): + """Return the cuda-python runtime module, or None when it is unavailable.""" + try: + from cuda.bindings import runtime + except ImportError: + try: # cuda-python < 12.8 + from cuda import cudart as runtime + except ImportError: + return None + return runtime + + +def cuda_available() -> bool: + """Return True when at least one CUDA device is visible.""" + runtime = _runtime() + if runtime is None: + return False + error, count = runtime.cudaGetDeviceCount() + return int(error) == 0 and count > 0 + + +def synchronize(output: Any = None) -> None: + """Block until all device work, including ``output``, has completed. + + ``output`` is duck-typed for host-async frameworks; the device synchronization + covers eager ones. Without cuda-python that synchronization is skipped, so timings + taken around this call are not trustworthy. + """ + _block_until_ready(output) + runtime = _runtime() + if runtime is None: + return + (error,) = runtime.cudaDeviceSynchronize() + if int(error) != 0: + # Must raise: a sticky error (illegal access in the kernel under test) + # surfaces here, and swallowing it records a fast, meaningless timing. + raise RuntimeError(f"cudaDeviceSynchronize failed: {error}") + + +def _block_until_ready(value: Any) -> None: + """Block on JAX-style arrays without importing jax.""" + if value is None: + return + if hasattr(value, "block_until_ready"): + value.block_until_ready() + elif isinstance(value, dict): + for item in value.values(): + _block_until_ready(item) + elif isinstance(value, (list, tuple)): + for item in value: + _block_until_ready(item) + + +def profiler_start() -> bool: + """Start CUDA profiler capture, reporting whether it actually started.""" + runtime = _runtime() + if runtime is None: + return False + error = runtime.cudaProfilerStart() + error = error[0] if isinstance(error, tuple) else error + return int(error) == 0 + + +def profiler_stop() -> bool: + """Stop CUDA profiler capture, reporting whether it stopped cleanly.""" + runtime = _runtime() + if runtime is None: + return False + error = runtime.cudaProfilerStop() + error = error[0] if isinstance(error, tuple) else error + return int(error) == 0 + + +def device_metadata() -> dict[str, Any]: + """Describe the visible CUDA devices without consulting any framework.""" + runtime = _runtime() + if runtime is None: + return {"available": False, "reason": "cuda-python is not installed"} + + error, count = runtime.cudaGetDeviceCount() + if int(error) != 0 or count == 0: + return {"available": False, "device_count": 0} + + error, current = runtime.cudaGetDevice() + if int(error) != 0: + return {"available": False, "reason": f"cudaGetDevice failed: {error}"} + + devices = [] + for index in range(count): + error, props = runtime.cudaGetDeviceProperties(index) + if int(error) != 0: + continue + devices.append( + { + "index": index, + "name": bytes(props.name).decode(errors="replace").rstrip("\x00"), + "total_memory": int(props.totalGlobalMem), + "compute_capability": [int(props.major), int(props.minor)], + "multi_processor_count": int(props.multiProcessorCount), + "uuid": bytes(props.uuid.bytes).hex(), + } + ) + return { + "available": True, + "device_count": count, + "current_device": int(current), + "devices": devices, + } + + +def device_architecture() -> str | None: + """Return the current device's arch as ``sm_``. + + Never raises: any failure yields None, which callers must treat as "unknown" rather + than as a proven architecture mismatch. + """ + runtime = _runtime() + if runtime is None: + return None + error, count = runtime.cudaGetDeviceCount() + if int(error) != 0 or count == 0: + return None + error, current = runtime.cudaGetDevice() + if int(error) != 0: + return None + error, props = runtime.cudaGetDeviceProperties(current) + if int(error) != 0: + return None + return f"sm_{int(props.major) * 10 + int(props.minor)}" + + +@functools.lru_cache(maxsize=None) +def build_architectures() -> dict[str, Any]: + """Return the SASS architectures embedded in ``libtransformer_engine.so``. + + Reads the loaded shared object with ``cuobjdump`` rather than trusting + ``NVTE_CUDA_ARCHS``, so it also holds for libraries built elsewhere. Cached because + ``cuobjdump`` takes seconds. Never raises: any failure yields + ``{"available": False, "reason": ...}``, which callers must treat as "unknown" + rather than as a proven architecture mismatch. + """ + from transformer_engine.common import _get_shared_object_file + + try: + so_path = _get_shared_object_file("core") + except (OSError, RuntimeError) as exc: + return {"available": False, "reason": f"could not locate libtransformer_engine.so: {exc}"} + + exe = shutil.which("cuobjdump") + if exe is None: + return {"available": False, "reason": "cuobjdump not found (CUDA toolkit absent)"} + + try: + result = subprocess.run( + [exe, "--list-elf", str(so_path)], + capture_output=True, + encoding="utf-8", + timeout=120, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + return {"available": False, "reason": f"cuobjdump failed: {exc}"} + + if result.returncode != 0: + stderr = (result.stderr or "").strip() + return {"available": False, "reason": f"cuobjdump exited {result.returncode}: {stderr}"} + + archs = sorted( + {f"sm_{m.group(1)}{m.group(2)}" for m in _ELF_ARCH.finditer(result.stdout)}, + key=lambda s: (int(re.sub(r"\D", "", s)), s), + ) + return {"available": True, "library": str(so_path), "cuda_architectures": archs} diff --git a/transformer_engine/common/testing/plugin.py b/transformer_engine/common/testing/plugin.py new file mode 100644 index 0000000000..207a9e8490 --- /dev/null +++ b/transformer_engine/common/testing/plugin.py @@ -0,0 +1,393 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Pytest plugin driving benchmarkable tests.""" + +from __future__ import annotations + +import inspect +import os +from pathlib import Path +import sys +from typing import Any +import warnings + +import pytest + +from .artifacts import write_run_artifacts +from .case import Case, CaseSkip +from .declaration import declared_axes, normalize_argnames, set_plugin_active +from .decorator import ( + BENCHMARK_MARKER, + CASE_MARKER, + MARKERS, + MODE_BENCHMARK, + MODE_CORRECTNESS, + SUPPRESS_MARKER, +) +from .device import build_architectures, cuda_available, device_architecture +from .runner import _run_benchmark_point, _run_correctness +from .timing import TIMING_METHOD + + +def pytest_addoption(parser): + """Register the benchmarkable command-line options, once per parser.""" + # Keyed to the parser, not to a module-level flag: pytest builds a fresh parser per + # pytest.main() call, and a sticky flag would hide --nvte-benchmark from a second run. + if getattr(parser, "_te_benchmarkable_registered", False): + return + parser._te_benchmarkable_registered = True # pylint: disable=protected-access + + group = parser.getgroup("benchmarkable") + group.addoption( + "--nvte-benchmark", + action="store_true", + default=False, + help=( + "Run the benchmark matrix only, gated by a one-time correctness check at " + "each benchmark point." + ), + ) + group.addoption( + "--nvte-benchmark-report-dir", + default=None, + help="Directory for benchmark_report/v1 artifacts.", + ) + group.addoption("--nvte-benchmark-warmup", type=int, default=5) + group.addoption("--nvte-benchmark-iterations", type=int, default=20) + group.addoption("--nvte-benchmark-inner-iterations", type=int, default=1) + group.addoption("--nvte-benchmark-min-run-time", type=float, default=0.0) + group.addoption( + "--nvte-benchmark-no-reference", + action="store_true", + default=False, + help="Skip timing the reference variant.", + ) + + +def _resolve_mode(config) -> str: + """Return ``MODE_BENCHMARK`` if ``--nvte-benchmark`` was passed, else ``MODE_CORRECTNESS``.""" + return MODE_BENCHMARK if config.getoption("--nvte-benchmark") else MODE_CORRECTNESS + + +def pytest_configure(config): + """Resolve the mode before test modules are imported, and register markers.""" + # Ahead of the double-fire guard, so a second registered copy of this hook still + # sets the flag; Case.__post_init__ reads it to detect a plugin-less session. + set_plugin_active(True) + + if getattr(config, "_benchmarkable_configured", False): + return + config._benchmarkable_configured = True # pylint: disable=protected-access + + for name, description in MARKERS: + config.addinivalue_line("markers", f"{name}: {description}") + + if _resolve_mode(config) == MODE_CORRECTNESS: + return + + if config.getoption("--nvte-benchmark-inner-iterations") < 1: + raise pytest.UsageError("--nvte-benchmark-inner-iterations must be at least 1.") + if config.getoption("--nvte-benchmark-iterations") < 1: + raise pytest.UsageError("--nvte-benchmark-iterations must be at least 1.") + + +def _is_case_bearing(node) -> bool: + """Whether the test returns a Case, so this plugin runs it instead of pytest.""" + # get_closest_marker, not node.keywords: keywords also carry every ancestor node's + # bare name, so an unmarked test under a path named nvte_case would match here. + return node.get_closest_marker(CASE_MARKER) is not None + + +def _is_benchmark_eligible(node) -> bool: + """Whether the test declared benchmark axes and nothing suppressed them.""" + # Suppression is a union over the test, its class and its module, and nothing + # re-enables: skip and skipif only ever subtract benchmarking. + if node.get_closest_marker(BENCHMARK_MARKER) is None: + return False + return next(node.iter_markers(name=SUPPRESS_MARKER), None) is None + + +def _substitutions(config) -> dict: + """Per-session record of which test definitions had a benchmark axis substituted.""" + registry = getattr(config, "_benchmarkable_substitutions", None) + if registry is None: + registry = {} + config._benchmarkable_substitutions = registry # pylint: disable=protected-access + return registry + + +def _definition_key(node) -> tuple: + """Key one test definition by the collector and name that ``pytest_generate_tests`` + and every item it generates share.""" + # Not the function object: a method inherited by two classes is one function but two + # definitions, and one class substituting an axis says nothing about the other. + return (node.parent, getattr(node, "originalname", None) or node.name) + + +# trylast is required: pluggy's LIFO ordering would otherwise run this hook before +# pytest's own -k/-m deselection, and the guard must see the final item list. +@pytest.hookimpl(trylast=True) +def pytest_collection_modifyitems(config, items): + """Narrow the session to benchmark points, then run the post-collection guard.""" + if _resolve_mode(config) == MODE_CORRECTNESS: + return + + substituted = _substitutions(config) + kept, dropped, suppressed, unmatched = [], [], 0, 0 + for item in items: + eligible = _is_benchmark_eligible(item) + if eligible and substituted.get(_definition_key(item), False): + kept.append(item) + continue + dropped.append(item) + if eligible: + unmatched += 1 + elif _is_case_bearing(item): + suppressed += 1 + if dropped: + config.hook.pytest_deselected(items=dropped) + items[:] = kept + + _guard_benchmark_collection(config, kept, suppressed, unmatched) + + +def _guard_benchmark_collection(config, kept, suppressed, unmatched) -> None: + """Fail loudly when a requested benchmark mode has nothing valid to time.""" + if not kept: + raise pytest.UsageError(_zero_points_message(suppressed, unmatched)) + + if config.option.collectonly: + return + + # Without cuda-python, device.synchronize() is a silent no-op: every measurement + # would be host launch time only and the report would still claim a device ran. + if not cuda_available(): + raise pytest.UsageError( + "--nvte-benchmark requires a CUDA device reachable through cuda-python; none is " + "available. Install cuda-python (the 'test' extra) or drop --nvte-benchmark." + ) + + _guard_build_architecture() + + +def _zero_points_message(suppressed, unmatched) -> str: + """Explain a zero-point benchmark session without blaming the wrong thing.""" + if suppressed: + detail = ( + f"all {suppressed} Case-bearing test(s) here carry @benchmark.skip or a true" + " @benchmark.skipif." + ) + elif unmatched: + detail = ( + f"{unmatched} benchmark-eligible test(s) here parametrize none of the axes" + " declared on their class." + ) + else: + detail = ( + "check the -k/-m expression, and that a test under the selected paths declares" + " axes with @benchmark(argnames, values)." + ) + return f"--nvte-benchmark collected zero benchmark points: {detail}" + + +def _guard_build_architecture() -> None: + """Fail loudly when the loaded library's embedded SASS cannot run on this device.""" + # Only a proven mismatch may raise: an unavailable build architecture (no + # cuobjdump, or the library was not located) or an unknown device architecture + # leaves the comparison undecidable and must not block the run. + build = build_architectures() + if not build.get("available"): + return + device_arch = device_architecture() + if device_arch is None: + return + + embedded = build["cuda_architectures"] + if device_arch.rstrip("af") in {arch.rstrip("af") for arch in embedded}: + return + + raise pytest.UsageError( + f"{build['library']} was built for " + f"{', '.join(embedded) or '(no architectures embedded)'}, but this device is " + f"{device_arch}, so no kernel can launch. Rebuild with an NVTE_CUDA_ARCHS " + "matching this device (use 100 for sm_100a/sm_103a)." + ) + + +def benchmark_settings(config) -> dict[str, Any]: + """Return the resolved benchmark settings for this session.""" + report_dir = config.getoption("--nvte-benchmark-report-dir") + return { + "mode": _resolve_mode(config), + "warmup": config.getoption("--nvte-benchmark-warmup"), + "iterations": config.getoption("--nvte-benchmark-iterations"), + "inner_iterations": config.getoption("--nvte-benchmark-inner-iterations"), + "min_run_time": config.getoption("--nvte-benchmark-min-run-time"), + "timing_method": TIMING_METHOD, + "report_dir": Path(report_dir) if report_dir else None, + "no_reference": config.getoption("--nvte-benchmark-no-reference"), + } + + +@pytest.hookimpl(hookwrapper=True) +def pytest_generate_tests(metafunc): + """Swap benchmark values into a test's existing parametrize marks. + + Substituting in place preserves the parametrization's positional structure, so ``pytest.param`` + ids and coupled argnames keep working. Class-level declarations are picked up via + ``metafunc.cls``; function-level entries win. Mutates ``own_markers``, not public pytest API + (pytest is pinned to 8.2.1). + """ + saved = [] + # Eligibility is checked before anything else: a suppressed or undeclared sibling + # method must not inherit its class's declarations, nor be failed by the check below. + eligible = _is_benchmark_eligible(metafunc.definition) + try: + if eligible and _resolve_mode(metafunc.config) == MODE_BENCHMARK: + own = declared_axes(metafunc.function) + declarations = dict(own) + if metafunc.cls is not None: + for key, values in declared_axes(metafunc.cls).items(): + declarations.setdefault(key, values) + node = metafunc.definition + while node is not None and declarations: + markers = getattr(node, "own_markers", None) + if markers: + for index, mark in enumerate(markers): + if mark.name != "parametrize": + continue + key = normalize_argnames(mark.args[0]) + if key in declarations: + saved.append((markers, index, mark)) + markers[index] = pytest.mark.parametrize( + mark.args[0], declarations[key] + ).mark + node = getattr(node, "parent", None) + _check_own_declarations_matched(metafunc, own, saved) + _substitutions(metafunc.config)[_definition_key(metafunc.definition)] = bool(saved) + yield + finally: + # Class- and module-level marks are shared, so a substituted mark left in place + # corrupts every sibling test for the rest of the session. The finally must + # therefore span the validation calls and the yield, not just the swap loop. + for markers, index, original in saved: + markers[index] = original + + +def _check_own_declarations_matched(metafunc, own, saved) -> None: + """Fail loudly when an axis declared on this test itself matches no parametrize mark: + a typo would otherwise silently benchmark the correctness values. An axis inherited from + a class is best-effort, since a sibling method need not parametrize it.""" + matched = {normalize_argnames(mark.args[0]) for _, _, mark in saved} + missing = sorted(set(own) - matched) + if not missing: + return + raise pytest.UsageError( + f"{metafunc.definition.nodeid}: @benchmark declared axes {missing}, which match no " + "pytest.mark.parametrize on this test, its class, or its module. Check the " + 'spelling, and declare a coupled group exactly as parametrized ("m,n,k").' + ) + + +@pytest.hookimpl(tryfirst=True) +def pytest_pyfunc_call(pyfuncitem): + """Run a Case-bearing test's Case according to the active mode.""" + if not _is_case_bearing(pyfuncitem): + return None + # An async test cannot return a Case, and leaving it to pytest keeps pytest's own + # warn-and-skip for it byte-identical. + if inspect.iscoroutinefunction(pyfuncitem.obj) or inspect.isasyncgenfunction(pyfuncitem.obj): + return None + + argnames = pyfuncitem._fixtureinfo.argnames # pylint: disable=protected-access + kwargs = {name: pyfuncitem.funcargs[name] for name in argnames} + case = pyfuncitem.obj(**kwargs) + if not isinstance(case, Case): + return _dispose_of_non_case(pyfuncitem, case) + + # CaseSkip from setup() is a coverage skip (unavailable backend or arch), not a + # failure. + if _resolve_mode(pyfuncitem.config) == MODE_CORRECTNESS: + try: + _run_correctness(case) + except CaseSkip as exc: + pytest.skip(str(exc)) + return True + + settings = benchmark_settings(pyfuncitem.config) + try: + records = _run_benchmark_point(case, settings, pyfuncitem) + except CaseSkip as exc: + pytest.skip(str(exc)) + store = getattr(pyfuncitem.config, "_benchmarkable_records", None) + if store is None: + store = [] + pyfuncitem.config._benchmarkable_records = store # pylint: disable=protected-access + store.extend(records) + return True + + +def _dispose_of_non_case(pyfuncitem, result): + """Handle a Case-bearing test that returned no Case, exactly as pytest would.""" + if declared_axes(pyfuncitem.function): + raise TypeError( + f"{pyfuncitem.nodeid}: @benchmark declares axes on this test, so it must return " + f"a Case; got {type(result).__name__}. A test that is Case-bearing but never " + "benchmarked is written @benchmark.skip." + ) + if _resolve_mode(pyfuncitem.config) == MODE_BENCHMARK: + pytest.skip("returns no Case, so it is not a benchmark point.") + # Mirrors _pytest/python.py::pytest_pyfunc_call, which owns this test in every other + # respect: an inherited declaration is a blanket statement, not a per-method claim. + if result is not None: + warnings.warn( + pytest.PytestReturnNotNoneWarning( + f"Expected None, but {pyfuncitem.nodeid} returned {result!r}, which will be " + "an error in a future version of pytest. Did you mean to use `assert` " + "instead of `return`?" + ) + ) + return True + + +def pytest_sessionfinish(session, exitstatus): # pylint: disable=unused-argument + """Write benchmark artifacts once the session completes. Guarded against a second + registered copy of the plugin, since this hook is not ``firstresult``.""" + config = session.config + if getattr(config, "_benchmarkable_report_written", False): + return + config._benchmarkable_report_written = True # pylint: disable=protected-access + records = getattr(config, "_benchmarkable_records", None) + if not records: + return + settings = benchmark_settings(config) + if settings["report_dir"] is None: + print( + f"\nWarning: {len(records)} benchmark record(s) were collected but " + "--nvte-benchmark-report-dir was not set, so they were discarded. Pass " + "--nvte-benchmark-report-dir to write a benchmark_report/v1 artifact.", + file=sys.stderr, + ) + return + + selection = { + "mode": settings["mode"], + "warmup": settings["warmup"], + "iterations": settings["iterations"], + "inner_iterations": settings["inner_iterations"], + "min_run_time": settings["min_run_time"], + "timing_method": settings["timing_method"], + "include_reference": not settings["no_reference"], + "args": list(config.invocation_params.args), + } + # Only argv[0]'s basename is kept: the full launcher path leaks the invoking user's + # home directory into every persisted report, and selection["args"] has the rest. + command = [os.path.basename(sys.argv[0])] + list(sys.argv[1:]) + paths = write_run_artifacts( + settings["report_dir"], + records, + command, + selection, + ) + print(f"\nWrote benchmark report: {paths['report']}") diff --git a/transformer_engine/common/testing/runner.py b/transformer_engine/common/testing/runner.py new file mode 100644 index 0000000000..7b96c1248b --- /dev/null +++ b/transformer_engine/common/testing/runner.py @@ -0,0 +1,166 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Case-execution engine: drives a Case's setup/evaluate/reference/verify cycle.""" + +from __future__ import annotations + +import os +import time + +from .case import Case, axis_value +from .device import synchronize +from .timing import TIMING_METHOD, WallClockSampler, timing_stats + + +def _framework_for(pyfuncitem) -> str: + """Infer the framework label for the report from the test's location under tests/. + + This is a reporting label only; nothing in the execution path branches on it. + """ + path = str(pyfuncitem.path) + if f"{os.sep}jax{os.sep}" in path or path.endswith(f"{os.sep}jax"): + return "jax" + return "pytorch" + + +def _run_correctness(case: Case) -> None: + """Run one setup/evaluate/reference/verify cycle with no timing.""" + state = case.setup() + actual = case.evaluate(state) + if case.reference is None: + return + # ``synchronize(actual)`` must precede ``reset(state)``: ``actual`` may alias + # ``state``, so an eager ``reset`` would race with in-flight ``evaluate`` work. + # ``reset`` must precede ``reference`` so it sees unmutated state. + synchronize(actual) + if case.reset is not None: + case.reset(state) + expected = case.reference(state) + synchronize(expected) + case.run_verify(actual, expected) + + +def _run_benchmark_point(case, settings, pyfuncitem): + """Gate once on correctness, then time evaluate and optionally reference.""" + state = case.setup() + + precondition_verified = False + if case.reference is not None: + actual = case.evaluate(state) + # Same ordering rule as ``_run_correctness``: synchronize before reset. + synchronize(actual) + if case.reset is not None: + case.reset(state) + expected = case.reference(state) + synchronize(expected) + case.run_verify(actual, expected) + precondition_verified = True + + variants = [("evaluation", case.evaluate)] + if case.reference is not None and case.time_reference and not settings["no_reference"]: + variants.insert(0, ("reference", case.reference)) + + # A case needing reset between calls cannot be batched: inner iterations are + # submitted back to back with no chance to reset, so timings would average over + # drifting state. + batchable = case.batchable and case.reset is None + inner = settings["inner_iterations"] if batchable else 1 + records = [] + for name, function in variants: + records.append( + _time_variant( + case, settings, pyfuncitem, name, function, inner, batchable, precondition_verified + ) + ) + return records + + +def _time_variant( + case, settings, pyfuncitem, variant, function, inner, batchable, precondition_verified +): + """Warm up, then time ``function``, with no verification inside the timed loop. + + A ``CaseSkip`` from ``setup()`` here means ``setup()`` is non-deterministic, which + the ``Case`` contract forbids, so it is left to propagate rather than caught. + """ + record = _base_record(pyfuncitem, variant, precondition_verified) + state = case.setup() + for _ in range(settings["warmup"]): + output = function(state) + synchronize(output) + if case.reset is not None: + case.reset(state) + synchronize() + + sampler = WallClockSampler(inner) + samples_ms = [] + start = time.perf_counter() + while ( + len(samples_ms) < settings["iterations"] + or time.perf_counter() - start < settings["min_run_time"] + ): + samples_ms.append(sampler(function, state)) + if case.reset is not None: + case.reset(state) + synchronize() + + stats = timing_stats(samples_ms) + record.update( + { + "status": "completed", + "warmup_iterations": settings["warmup"], + "iterations": len(samples_ms), + "inner_iterations": inner, + "batchable": batchable, + "timing_method": TIMING_METHOD, + "samples_ms": samples_ms, + "timing": stats, + "metrics": _metrics(case, stats["median_ms"]), + } + ) + return record + + +def _metrics(case, median_ms): + """Derive bandwidth and FLOPs metrics from the median timing, when available.""" + metrics = {} + if median_ms > 0: + if case.bytes_moved is not None: + metrics["bandwidth_GBps"] = case.bytes_moved / (median_ms / 1.0e3) / 1.0e9 + if case.flops is not None: + metrics["tflops"] = case.flops / (median_ms / 1.0e3) / 1.0e12 + return metrics + + +def _base_record(pyfuncitem, variant, precondition_verified): + """Build the record fields that are known before timing runs.""" + params = {name: axis_value(v) for name, v in pyfuncitem.callspec.params.items()} + return { + "schema_version": "benchmark_record/v1", + "status": "pending", + "variant": variant, + "framework": _framework_for(pyfuncitem), + "case_id": case_id_for(pyfuncitem, params), + "component": pyfuncitem.module.__name__.rsplit(".", maxsplit=1)[-1], + "operation": pyfuncitem.originalname, + "params": params, + "node_id": pyfuncitem.nodeid, + "precondition_verified": precondition_verified, + "tags": [], + "unit_test": pyfuncitem.nodeid, + "source": pyfuncitem.module.__name__, + "regression_threshold": None, + } + + +def case_id_for(pyfuncitem, params) -> str: + """Build the stable identity for a benchmark record. + + Keyed on module, test function and sorted named axis values rather than the + positional pytest node ID, which would rename cases when an axis is reordered. + ``params`` is already rendered by the caller; re-applying ``axis_value`` here is safe + only because it is idempotent on the primitives that rendering produces. + """ + axes = ".".join(f"{name}{axis_value(params[name])}" for name in sorted(params)) + return f"{pyfuncitem.module.__name__}.{pyfuncitem.originalname}.{axes}" diff --git a/transformer_engine/common/testing/timing.py b/transformer_engine/common/testing/timing.py new file mode 100644 index 0000000000..3fe8141e5c --- /dev/null +++ b/transformer_engine/common/testing/timing.py @@ -0,0 +1,56 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Timing backends for benchmarkable cases.""" + +from __future__ import annotations + +import math +import statistics +import time + +from .device import synchronize + + +TIMING_METHOD = "wall-clock-with-device-sync" + + +class WallClockSampler: + """Time ``inner_iterations`` calls with the host clock and one device sync.""" + + def __init__(self, inner_iterations: int) -> None: + self.inner_iterations = inner_iterations + + def __call__(self, function, state) -> float: + """Return the mean milliseconds per call across the inner iterations.""" + output = None + start = time.perf_counter() + for _ in range(self.inner_iterations): + output = function(state) + synchronize(output) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + return elapsed_ms / self.inner_iterations + + +def timing_stats(samples_ms: list[float]) -> dict[str, float]: + """Compute timing statistics in milliseconds.""" + if not samples_ms: + return { + "median_ms": 0.0, + "mean_ms": 0.0, + "min_ms": 0.0, + "max_ms": 0.0, + "stddev_ms": 0.0, + "p95_ms": 0.0, + } + + ordered = sorted(samples_ms) + p95_index = max(0, math.ceil(0.95 * len(ordered)) - 1) + return { + "median_ms": statistics.median(ordered), + "mean_ms": statistics.fmean(ordered), + "min_ms": ordered[0], + "max_ms": ordered[-1], + "stddev_ms": statistics.pstdev(ordered) if len(ordered) > 1 else 0.0, + "p95_ms": ordered[p95_index], + }