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 GEMM profiling guide to Transformer Engine documentation along with a companion benchmark tool (benchmarks/gemm/benchmark_gemm.py). It covers shape derivation for all 12 per-layer GEMMs, benchmarking across BF16/FP8/MXFP8/NVFP4 precisions, and interpretation of speedup results.

• benchmark_gemm.py: 1883-line standalone tool supporting model-config and manual-shape modes, two timing backends (CUDA events and torch.profiler), autocast and pre-quantized measurement modes, Nsight Systems integration, and bar-chart output.
• gemm_profiling.rst: 589-line tutorial with worked examples on B300 and H200; contains a documentation inaccuracy where the pre-quantized examples show an FP8Delayed ms column that the code does not produce.
• speedups.rst + toctree updates: New summary page added to the low-precision-training section; docs/index.rst correctly adds the tutorial to the toctree.

Confidence Score: 4/5

Safe to merge after addressing the FP8Delayed pre-quantized documentation mismatch; the benchmark tool itself works correctly.

The benchmark code correctly excludes FP8Delayed when --pre-quantize is passed, but the tutorial note and both pre-quantized example outputs (B300 and H200) show a FP8Delayed ms column that users will not see when they run the documented commands. This directly misleads anyone following the H200 pre-quantized walkthrough. The rest of the benchmark logic, toctree wiring, and plot generation look correct.

docs/examples/gemm_profiling/gemm_profiling.rst — the pre-quantized sections (note at line 120, B300 output at line 143, H200 output at line 291) need to be reconciled with the actual code behaviour that skips FP8Delayed entirely in pre-quantized mode.

Important Files Changed

Filename	Overview
benchmarks/gemm/benchmark_gemm.py	New 1883-line GEMM benchmark tool supporting BF16, FP8 (Current/Delayed/Block), MXFP8, NVFP4; FP8Delayed is correctly excluded from pre-quantized runs but the docs misrepresent this.
docs/examples/gemm_profiling/gemm_profiling.rst	New 589-line tutorial RST. Pre-quantized B300/H200 examples claim FP8Delayed is included, contradicting the code. H200 speedup ordering also mismatches code iteration order.
docs/features/low_precision_training/speedups.rst	New summary page covering autocast vs pre-quantized results for B300 and H200 with benchmark plots; no logic issues.
docs/features/low_precision_training/index.rst	Adds speedups.rst to the low_precision_training toctree; straightforward change.
docs/index.rst	Adds examples/gemm_profiling/gemm_profiling.rst to the top-level toctree; resolves the prior missing-toctree issue.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[CLI args] --> B{Mode?}
    B -->|model config args| C[run_model_config_benchmarks]
    B -->|--shapes or default| D[run_benchmarks]
    C --> F[compute_gemm_shapes]
    F --> G[_benchmark_single_shape per shape]
    G --> H{pre_quantize?}
    H -->|False| I[autocast: FP8Current/Delayed/Block/MXFP8/FP4]
    H -->|True| J[prequantized: FP8Current/Block/MXFP8/FP4 — FP8Delayed SKIPPED]
    I --> K[print summary + speedup table]
    J --> K
    K --> L[create_model_config_plot]
    D --> M{pre_quantize?}
    M -->|False| N[autocast benchmarks]
    M -->|True| O[prequantized — FP8Delayed SKIPPED]
    N --> P[create_plot]
    O --> P

Loading

_{Reviews (10): Last reviewed commit: "Regenerate GEMM speedup figures with Del..." | 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

[greptile-apps]

--verify-dgrad plot silently uses approximation instead of measured values

When --verify-dgrad is passed, run_model_config_benchmarks benchmarks and records actual Dgrad timings into dgrad_results, and the printed table correctly shows those measured values. However, create_model_config_plot is never given dgrad_results — the call site only passes fprop_results and wgrad_results. Inside the plot function, Fprop+Dgrad bar height is always computed as fp.avg_time_ms * 2 (the approximation), so the chart silently contradicts the table when --verify-dgrad is used.

Fix: add dgrad_results and verify_dgrad parameters to create_model_config_plot, and when verify_dgrad=True, use fprop_ms[j] + dgrad_ms[j] instead of fprop_ms[j] * 2 for each op bar.

[pggPL]

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

I added some comments.

[pggPL]

In the Features section we want to have only text with very short code snippets - one should be able to read it as academic-like handbook and we want it to be concise. The Tutorials and examples is better place for code user should run and more elaborate text.

I see that describing what is real impact of the precisions is super impactful and user reading the docs now is not able to estimate it. So what i propose is:

• add short speedups.rst here section that would have some short introduction + 1-2 graphs/tables and link to gemm_profiling.rst + short description what user can find there,
• add gemm_profiling.rst to examples which would basically be this whole tutorial.

I will elaborate about this in separate comment, but I think we would need 2 pictures/tables in speedups.rst:

• one with Hopper recipes: bf16, fp8 tensorwise, fp8 blockwise,
• one with Blackwell recipes: bf16, fp8 tensorwise, mxfp8, nvfp4
  for some reasonable shape sizes.

And then link to gemm profiling example.

[pggPL]

FP8 block scaling on Blackwell is emulated via MXFP8 tensor cores and we support it mostly for backward compatibility. You can see it in the fp8 blockwise docs page. This recipe is aimed for Hopper.

Also, I see that FP8 tensorwise is omitted, is there a specific reason for that? I think it is still widely used (i think also on Blackwell, but i am not sure).

So I think we should run 2 experiments - 1 for hopper, 1 for blackwell with appropriate recipes.

[jomitchellnv]

Yes -- will do and update accordingly thanks.

[pggPL]

i think this comment should be removed, since we do not aim to run both of them on the same device - one is for hopper, on is for blackwell

[jomitchellnv]

done -- its gone. We replaced it with a device-specific note explaining that FP8 Block targets Hopper and MXFP8/NVFP4 target Blackwell

[pggPL]

This is weird, I think we should run dgrad - even if the number of FLOPs is the same, some tensors are transposed and it may impact the result.

[jomitchellnv]

Yea ok I fixed it --> done

[pggPL]

This is really important section and should not be in appendix imo

[jomitchellnv]

ok for sure

[pggPL]

as commented later, I think we should always benchmark it separately

[jomitchellnv]

done

[pggPL]

is it auto generated by the script?

[jomitchellnv]

yea -- its auto generated

[pggPL]

This should be link, maybe even to github

[pggPL]

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

[greptile-apps]

The "Estimated GEMM Speedups" block only prints results for MXFP8 and NVFP4. When the tool is run with --no-fp8 --no-fp4 (the documented H200 invocation), run_fp8 and run_fp4 are both False, so all three if guards evaluate to False and the section emits nothing — directly contradicting the tutorial output in docs/examples/gemm_profiling/gemm_profiling.rst which shows FP8Delayed vs BF16: 1.69x, FP8Current vs BF16: 1.58x, and FP8Block vs BF16: 1.40x for that exact invocation.

Suggested change

	print("\nEstimated GEMM Speedups:")
	bf16_total = full_model.get("BF16", 0)
	if run_fp8 and bf16_total > 0:
	fp8_total = full_model.get("MXFP8", 0)
	if fp8_total > 0:
	print(f" MXFP8 vs BF16: {bf16_total / fp8_total:.2f}x")
	if run_fp4 and run_fp8:
	fp8_total = full_model.get("MXFP8", 0)
	fp4_total = full_model.get("NVFP4", 0)
	if fp8_total > 0 and fp4_total > 0:
	print(f" NVFP4 vs MXFP8: {fp8_total / fp4_total:.2f}x")
	if run_fp4 and bf16_total > 0:
	fp4_total = full_model.get("NVFP4", 0)
	if fp4_total > 0:
	print(f" NVFP4 vs BF16: {bf16_total / fp4_total:.2f}x")
	print(sep)
	print("\nEstimated GEMM Speedups:")
	bf16_total = full_model.get("BF16", 0)
	if bf16_total > 0:
	for p in precisions[1:]: # all enabled precisions except BF16
	p_total = full_model.get(p, 0)
	if p_total > 0:
	print(f" {p} vs BF16: {bf16_total / p_total:.2f}x")
	if run_fp4 and run_fp8:
	fp8_total = full_model.get("MXFP8", 0)
	fp4_total = full_model.get("NVFP4", 0)
	if fp8_total > 0 and fp4_total > 0:
	print(f" NVFP4 vs MXFP8: {fp8_total / fp4_total:.2f}x")
	print(sep)

[pggPL]

I left some comments

[pggPL]

This can be removed

[jomitchellnv]

done

[pggPL]

I would like to have pictures first because they are the most important

[jomitchellnv]

Done. Moved both example sections (B300 and H200) with their figures above the "Speedup Is Shape-Dependent" text section, so readers see the benchmark plots first.

[pggPL]

I think this is similar to Models with large hidden dimensions and intermediate sizes

[jomitchellnv]

Done. Merged the separate "Batch size and sequence length also matter" bullet into the first bullet — it now
reads "Models with large hidden dimensions, intermediate sizes, and token counts (micro_batch_size *
sequence_length)..." — eliminating the redundancy.

[pggPL]

It's worth mentioning that speedup can be faster and is not extremely fast due to qunatization, but pre-qunatized mode is not introduced by this time.

[jomitchellnv]

Moved the pre-quantized reference out of the body text and into a note callout that explains the concept
(raw kernel throughput, --pre-quantize flag). The Sphinx tabs above also show the pre-quantized graph, so the reader sees the concept visually before the text discusses it.

[pggPL]

maybe insert also the graph with pre-qunaitzation and expose them via sphinx tab

[pggPL]

i think this is important point for people to understand, speedup of nvfp4 is big, but qunatization can diminish it

[jomitchellnv]

Done — both B300 and H200 sections now use .. tabs:: with Autocast and Pre-quantized tabs showing the
corresponding graphs.

Added a .. note:: callout under the B300 example that highlights this explicitly — shows the 1.98x autocast
vs 3.48x kernel-only gap and explains what causes it.

[pggPL]

i think we can remove this seciton

[jomitchellnv]

removed

[pggPL]

this also can be removed

[jomitchellnv]

removed

[greptile-apps]

"Estimated GEMM Speedups" never prints FP8Current/FP8Delayed/FP8Block entries

The three if guards only check run_fp8 (MXFP8) and run_fp4 (NVFP4). There are no corresponding branches for run_fp8_current, run_fp8_delayed, or run_fp8_block. On H200 — the documented use case that passes --no-fp8 --no-fp4 — both flags are False, so all three guards fail and the section emits nothing but its header. This directly contradicts the tutorial in docs/examples/gemm_profiling/gemm_profiling.rst (lines 259–263), which shows FP8Delayed vs BF16: 1.69x, FP8Current vs BF16: 1.58x, and FP8Block vs BF16: 1.40x for exactly that invocation.

[jomitchellnv]

Fixed — the hardcoded MXFP8/NVFP4 speedup branches are replaced with a loop over all active precisions, so it now prints {prec} vs BF16 for every enabled precision. Works correctly on Hopper with --no-fp8 --no-fp4

[greptile-apps]

B300 example output is stale and contradicts the current code

The documented B300 Estimated GEMM Speedups block (lines 97–100) shows only MXFP8 vs BF16, plus cross-precision lines like NVFP4 vs MXFP8: 1.39x that the current code never produces. The actual code at line 1483–1489 of benchmark_gemm.py iterates precisions[1:] and emits each precision vs BF16 only — no cross-precision pairs. Running the command as documented (no --no-fp8-current or --no-fp8-delayed flags) on a B300 would also print FP8Current vs BF16 and FP8Delayed vs BF16, which are absent from the docs. The Fprop table further omits the FP8Block column even though the shown command does not pass --no-fp8-block. This means users following the tutorial will see materially different output than documented.

[pggPL]

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

[pggPL]

This autocast/pre-quantized is not defined and user does not know what that means.

[jomitchellnv]

Good catch. Added a definition block at the top of the page (before any figures) defining both: Autocast = the end-to-end speedup seen in training, including per-step quantization; Pre-quantized = raw GEMM kernel throughput with inputs already in the target format (the hardware upper bound).

[pggPL]

no delayed scaling in H200 pre quantized

[jomitchellnv]

You're right, and this was actually a bug. In pre-quantized mode FP8Delayed had no pre-quantized variant and silently fell back to the autocast te.Linear path, so it shouldn't have appeared at all. Fixed benchmark_gemm.py to skip DelayedScaling under --pre-quantize (it differs from CurrentScaling only in how the scale is computed each step, which pre-quantized mode skips). Regenerating the pre-quantized plots rn so the FP8Delayed bar is gone.

[pggPL]

mentioning --pre-quantize flag, when we do not define script does not make sense. We should explain difference of pre qunatized vs autocast above and remove it.

[jomitchellnv]

Done. Removed the flag reference and the standalone note, and moved the autocast vs pre-quantized explanation to the top of the page as you suggested

[pggPL]

"(format conversion, block scaling, Hadamard transforms)" what format conversion and block scaling mean in this context?

it seems that claude generated it without understanding

[jomitchellnv]

Removed the unclear jargon. Replaced it with plain language: the per-step quantization cost is "converting the input tensors to the low-precision format and computing their scaling factors."

[pggPL]

similar - what is "including Fprop vs Dgrad comparisons"?

[jomitchellnv]

I reworded it but it was a comparison that I added. I now have per operation fprop, dgrad, and wgrad breakdowns

[pggPL]

this sections seems to be super verbose - it says basically "the bigger shape, the bigger speedup" in 15 sentences

[jomitchellnv]

i trimmed it down alot

[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