Adds GEMM Profiling Guide to TE

[jomitchellnv]

Description

Adds a GEMM profiling guide to the Transformer Engine documentation and a companion benchmark tool. The guide
explains how to derive all 12 per-layer GEMM shapes (Fprop, Dgrad, Wgrad) from transformer model
hyperparameters, benchmark them across precisions (BF16, FP8 Block, MXFP8, NVFP4), and interpret the resulting
speedup estimates.

The benchmark tool supports two modes: model config mode (derives shapes automatically from hidden_size,
intermediate_size, etc.) and manual shape mode (explicit MxKxN triplets). It measures both autocast performance
(realistic end-to-end with quantization overhead) and pre-quantized kernel-only throughput, using CUDA events
or torch.profiler timing backends.

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

• Add benchmarks/gemm/benchmark_gemm.py — standalone GEMM benchmark tool supporting BF16, FP8 Block, MXFP8, and
  NVFP4 precisions with autocast and pre-quantized modes, CUDA event and torch.profiler timing, Nsight Systems
  integration, and bar-chart output

• Add docs/features/low_precision_training/gemm_profiling/gemm_profiling.rst — documentation covering GEMM
  shape derivation from model configs, forward/backward pass shape conventions, precision mapping per GEMM pass,
  speedup calculation methodology, and a worked example on B300

• Add benchmark result plots (img/model_config_speedup.png, img/model_config_speedup_prequant.png)

• Update docs/features/low_precision_training/index.rst toctree to include the new guide
  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

[greptile-apps]

Greptile Summary

This PR adds a standalone GEMM benchmarking tool (benchmarks/gemm/benchmark_gemm.py) and accompanying documentation — a tutorial RST, a speedups overview page, and B300/H200 result plots — to help users measure and interpret per-layer GEMM performance across BF16, FP8, MXFP8, and NVFP4 precisions.

• benchmark_gemm.py: 1,883-line benchmark supporting model-config mode (derives all 12 GEMM shapes from hyperparameters) and manual shape mode, with CUDA-events and torch.profiler timing backends, Nsight Systems integration, and bar-chart output.
• docs/examples/gemm_profiling/gemm_profiling.rst: End-to-end tutorial covering shape derivation, forward/backward conventions, worked B300 and H200 examples, and guidance on interpreting speedups.
• docs/features/low_precision_training/speedups.rst + index update: Adds a high-level speedups overview page to the low-precision training section, cross-linked to the full tutorial.

Confidence Score: 4/5

Safe to merge with one fix: the MXFP8 autocast benchmark will crash on Hopper if the user omits --no-fp8.

The MXFP8 autocast function (benchmark_fp8) has no hardware guard and no exception handler, unlike every other non-BF16 benchmark in the file. On Hopper (SM90), invoking the tool without --no-fp8 will attempt MXFP8BlockScaling and, if TE rejects it, crash without any diagnostic message. The pre-quantized MXFP8 path and the NVFP4 paths are both protected; this one gap is the only blocking issue in an otherwise clean addition.

benchmarks/gemm/benchmark_gemm.py — specifically the benchmark_fp8 function and the run_fp8 flag computation in both orchestrators.

Important Files Changed

Filename	Overview
benchmarks/gemm/benchmark_gemm.py	New 1,883-line GEMM benchmark supporting BF16/FP8/MXFP8/NVFP4; FP8Block shape-mode fix confirmed, speedup loop fixed, but benchmark_fp8 (MXFP8 autocast) lacks Blackwell guard and try/except, crashing on Hopper without --no-fp8.
docs/examples/gemm_profiling/gemm_profiling.rst	New tutorial RST (589 lines); H200 speedup block lists precisions in FP8Delayed-first order while the code emits FP8Current first; B300 example output still carries stale speedup lines (previously flagged).
docs/features/low_precision_training/speedups.rst	New 114-line overview page with B300/H200 benchmark tabs; correctly linked in the low_precision_training toctree and uses existing sphinx-tabs extension.
docs/features/low_precision_training/index.rst	Adds speedups.rst to the toctree and fixes a missing trailing newline; straightforward change.
docs/index.rst	Adds examples/gemm_profiling/gemm_profiling.rst to the root toctree, making the tutorial reachable from the Sphinx navigation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[CLI args] --> B{Model config args?}
    B -- Yes --> C[run_model_config_benchmarks]
    B -- No --> D[run_benchmarks\nshape mode]
    C --> E[compute_gemm_shapes\nFprop / Dgrad / Wgrad]
    E --> F[_benchmark_single_shape\nfor each shape]
    D --> F2[shape loop\nfor each MxKxN]
    F --> G{Precision enabled?}
    F2 --> G
    G --> H[benchmark_bf16]
    G --> I[benchmark_fp8_current]
    G --> J[benchmark_fp8_delayed]
    G --> K[benchmark_fp8_block]
    G --> L[benchmark_fp8 ⚠️\nno hardware guard]
    G --> M{is_blackwell_available?}
    M -- Yes --> N[benchmark_fp4]
    M -- No --> O[skip NVFP4]
    L -->|No guard on non-Blackwell| P[⚠️ May crash on Hopper]
    H & I & J & K & N --> Q[GEMMResult]
    Q --> R[Print table + speedup summary]
    R --> S[create_model_config_plot or create_plot]

Loading

_{Reviews (12): Last reviewed commit: "Apply suggestion from @pggPL" | Re-trigger Greptile}

[pggPL]

Hi @jomitchellnv, I see that this PR is open, but "Documentation" job is failing. If you fix it, please ping me and I'll review it.

[jomitchellnv]

@pggPL they should be fixed now I hope

[jomitchellnv]

/te-ci L1 pytorch

[pggPL]

I'm super happy about this change, I think we really need that.

I added some comments.

[pggPL]

One more thought - we should consider adding MoE - grouped gemm.

[pggPL]

I left some comments

[pggPL]

I'm ok with the tutorial, but the speedups part needs polishing

[pggPL]

I've added some comments

[jomitchellnv]

i need to grab an H200 and B300 node to regenerate the plots then it should be ok

[pggPL]

LGTM

[greptile-apps]

MXFP8 autocast path has no hardware guard or exception handler

benchmark_fp8 (MXFP8 autocast) contains no is_blackwell_available() early-return and no try/except block. In the default mode (no --pre-quantize), both run_benchmarks and run_model_config_benchmarks select this function via fp8_fn = ... if pre_quantize else benchmark_fp8. If a Hopper user forgets --no-fp8, the call to te.autocast(enabled=True, recipe=MXFP8BlockScaling(...)) will likely raise, propagating uncaught through the loop and crashing the entire tool.

The documentation explicitly states MXFP8 requires Blackwell, and only NVFP4 has a corresponding guard — benchmark_fp4 returns None early when not is_blackwell_available(), and the orchestrators print a skip notice. The pre-quantized MXFP8 path (benchmark_fp8_prequantized) also has a try/except, so only the autocast variant is unprotected. Adding either an is_blackwell_available() guard at the top of this function or a try/except with a return None fallback would make it consistent with the rest of the codebase.