[Common] Benchmarkable Unit Tests - #3408
Draft
denera wants to merge 1 commit into
Draft
Conversation
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 <adener@nvidia.com>
denera
force-pushed
the
common/benchmarkable-unit-test-infrastructure
branch
from
August 19, 2026 23:42
3df6140 to
1803f57
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR introduces a new
transformer_engine.common.testingsubmodule that defines a new unified standard for implementing Transformer Engine unit tests that are also benchmarkable to avoid duplicate implementations between separate test and benchmark implementations.tests/jax/test_softmax.pyandtests/pytorch/test_fused_rose.pyhave been ported as demonstrative examples, but a complete porting of all existing unit tests as well as a re-write of existing benchmarking scripts are deferred to follow-up PRs.Type of change
Checklist: