Testing Python Operators#

FlagGems uses pytest for operator accuracy testing and performance benchmarking. It leverages Triton's triton.testing.do_bench for kernel-level performance evaluation.

1. Accuracy tests for operators#

To run unit tests on a specific backend like CUDA:

pytest tests/test_${name}.py

The following command runs the tests on CPU:

pytest tests/test_foo.py --ref cpu

2. Accuracy in the context of models#

pytest examples/${name}_test.py

3. Test operator performance#

To test CUDA performance

pytest benchmark/test_foo.py -s

To benchmark the end-to-end performance for operators:

pytest benchmark/test_foo.py -s --ref cpu

4. Command-line options supported by the test suite#

Besides the built-in pytest options (such as -k for keyword filtering, -m for marker-based filtering, and -x for stopping at the first failure), tests/conftest.py registers a few FlagGems-specific options that give you finer control over how accuracy tests run.

OptionValuesDescription
--ref<device> (default), cpuDevice used to compute the reference result the FlagGems output is compared against.
--quickflagShrink the shape/dtype combinations used by tests/accuracy_utils.py, for faster smoke runs.
--recordnone (default), log, jsonRecord test execution details to a log file or a JSON report.
--outputpathDestination file when --record json is used. Defaults to accuracy_result.json.
--collect-markspathInstead of running the tests, dump every collected test's markers to the given YAML file.

For example, to quickly sanity-check an operator on CPU and keep a JSON report:

pytest tests/test_softmax.py --ref cpu --quick --record json --output softmax_result.json

Marker-based filtering#

Every accuracy test is tagged with a marker named after the operator it covers, for example @pytest.mark.fliplr in tests/test_fliplr.py. This lets you run tests for one or more operators regardless of which file they live in:

pytest tests/ -m "fliplr or rms_norm"

Reporting and coverage integration#

FlagGems CI uses the pytest-md-report plugin to turn a test run into a Markdown summary, and wraps the whole invocation with coverage to collect unit-test coverage:

coverage run -m pytest -s tests/test_foo.py \
    --md-report --md-report-verbose=1 --md-report-output=summary.md
coverage html

See Unit-test coverage for more details on the coverage report.