[Common] Benchmarkable Unit Tests

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.

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
 

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"]

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",

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

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