microbenchmarks: use pytest as execution backend

[matthiasdiener]

Description

• Migrates the microbenchmarks from standalone argparse scripts to a pytest runner; each benchmark becomes a (parametrized) pytest, replacing the run_benchmarks() driver
• Shape/precision/backend/direction sweeps are parametrize axes, filterable via -k/-m (e.g. -k "mxfp8 and hipkittens", -k "bwd and QKV")
• Inspired by pytest-benchmark and similar functionality in usv (but does not add any extra dependencies)
• Reuses the torch.utils.benchmark timing, rotating input buffers, and CSV schema (with added columns)
• Advantages:
  • Benchmark selection/filtering (e.g., "run all mxfp8 GEMM benchmarks") without code changes
  • Direct integration of non-default backends (triton, ck_tile, hipkittens)
  • More consistent console output
  • Potential future benefit: pass/fail regression gating
• Example usage:

  # pytest: MXFP8 GEMM, HipKittens backend, backward only, save CSV + per-sample data
  python -m pytest benchmark_gemm.py -k "mxfp8 and hipkittens and bwd" --csv --csv-samples
  
  # standalone (still runs through pytest): Llama3 shapes, Triton backend, no input rotation
  python benchmark_gemm.py -k "Llama3 and triton" --no-rotating --csv

  $ python benchmark_casting.py -k "Qwen and 72B and not dequantize and E4M3"  -v
  ========================================================================= test session starts =========================================================================
  platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0 -- /opt/venv/bin/python
  cachedir: .pytest_cache
  rootdir: /dockerx/TransformerEngine
  configfile: pyproject.toml
  collected 200 items / 192 deselected / 8 selected
  
  benchmark_casting.py::test_cast[Qwen2.5-72B-FP8-E4M3-quantize-default-M1024] PASSED                                                                             [ 12%]
  benchmark_casting.py::test_cast[Qwen2.5-72B-FP8-E4M3-quantize-default-M2048] PASSED                                                                             [ 25%]
  benchmark_casting.py::test_cast[Qwen2.5-72B-FP8-E4M3-quantize-default-M4096] PASSED                                                                             [ 37%]
  benchmark_casting.py::test_cast[Qwen2.5-72B-FP8-E4M3-quantize-default-M8192] PASSED                                                                             [ 50%]
  benchmark_casting.py::test_cast[Qwen2.5-72B-FP8-E4M3-quantize-triton-M1024] PASSED                                                                              [ 62%]
  benchmark_casting.py::test_cast[Qwen2.5-72B-FP8-E4M3-quantize-triton-M2048] PASSED                                                                              [ 75%]
  benchmark_casting.py::test_cast[Qwen2.5-72B-FP8-E4M3-quantize-triton-M4096] PASSED                                                                              [ 87%]
  benchmark_casting.py::test_cast[Qwen2.5-72B-FP8-E4M3-quantize-triton-M8192] PASSED                                                                              [100%]
  
  
  benchmark: 8 tests
  
  | Name                                        | Min (ms) | Median (ms) | Max (ms) | StdDev (ms) |   Throughput |
  | :------------------------------------------ | -------: | ----------: | -------: | ----------: | -----------: |
  | Qwen2.5-72B-FP8-E4M3-quantize-default-M1024 |   0.0103 |      0.0103 |   0.0103 |      0.0000 | 2448.02 GB/s |
  | Qwen2.5-72B-FP8-E4M3-quantize-default-M2048 |   0.0196 |      0.0196 |   0.0196 |      0.0000 | 2571.83 GB/s |
  | Qwen2.5-72B-FP8-E4M3-quantize-default-M4096 |   0.0384 |      0.0385 |   0.0385 |      0.0000 | 2617.13 GB/s |
  | Qwen2.5-72B-FP8-E4M3-quantize-default-M8192 |   0.0725 |      0.0726 |   0.0726 |      0.0000 | 2775.04 GB/s |
  | Qwen2.5-72B-FP8-E4M3-quantize-triton-M1024  |   0.0271 |      0.0273 |   0.0289 |      0.0005 |  914.73 GB/s |
  | Qwen2.5-72B-FP8-E4M3-quantize-triton-M2048  |   0.0271 |      0.0272 |   0.0274 |      0.0001 | 1851.49 GB/s |
  | Qwen2.5-72B-FP8-E4M3-quantize-triton-M4096  |   0.0347 |      0.0347 |   0.0347 |      0.0000 | 2899.96 GB/s |
  | Qwen2.5-72B-FP8-E4M3-quantize-triton-M8192  |   0.0615 |      0.0616 |   0.0616 |      0.0000 | 3269.54 GB/s |

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[Copilot]

Pull request overview

This PR migrates the microbenchmark suite under benchmarks/microbenchmarks/ from standalone argparse runners to a pytest-driven execution model, enabling selection/filtering via pytest (-k, markers) while still producing CSV/samples/kernel-profile outputs for downstream benchmarking workflows.

Changes:

• Add pytest integration (conftest.py) with a microbench fixture and session-level result collection/printing/output writing.
• Introduce shared pytest-oriented helpers in utils.py (case recording, CSV/samples writers, formatted summary table).
• Convert several benchmark modules (GEMM, grouped GEMM, casting, normalization) to parametrized pytest tests with backend/direction sweep axes.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
benchmarks/microbenchmarks/utils.py	Adds pytest execution support helpers (result store, output writers, summary formatter) and refactors rotating-buffer configuration.
benchmarks/microbenchmarks/conftest.py	Implements pytest hooks/options and the microbench fixture to run/record benchmarks and emit outputs.
benchmarks/microbenchmarks/benchmark_gemm.py	Converts dense GEMM benchmark to parametrized pytest tests with backend + direction axes.
benchmarks/microbenchmarks/benchmark_grouped_gemm.py	Converts grouped GEMM benchmark to parametrized pytest tests with backend + direction axes and skip rules.
benchmarks/microbenchmarks/benchmark_casting.py	Converts casting benchmark to parametrized pytest tests with backend axis and case generation refactor.
benchmarks/microbenchmarks/benchmark_normalization.py	Converts normalization benchmark to parametrized pytest tests with backend axis and case generation refactor.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.