3.3 -> 3.4 TritonBench H100 perf regressions (& improvements)

[davidberard98]

Perf regressions are listed below. The data is collected from https://github.com/pytorch-labs/tritonbench/actions/runs/16035767306 (and for Meta employees, aggregated using https://fburl.com/scuba/pytorch_user_benchmarks/2mp2d6a9)

Note: just look at the "speedup difference ((new-old)/old)" column; apologies for the weird column names / numbers, it's an artifact of the data visualization / comparison tool used to generate the csv. For example, the worst real regression is tritonbench_layer_norm_bwd[x_6656-liger_layer_norm]_speedup, and the value of 4.34 means that the speedup is regressed 434% compared to 3.3, or assuming the torch baseline is stable, the latency has regressed 434%. For the examples below, I've double checked to verify that there are actual regressions in the measured kernels (not just improvements in the baseline).

These are notes about anything w/ > 10% regression

Biggest regressions (note: ignore the first 13 results with infinite regression, they are failing/not running on both versions): https://gist.github.com/davidberard98/af67f9d7be019c353bb4821a07f28bfa

• liger layernorm bwd: compared to torch, the liger kernels appear to do well on small x, and then fall off a cliff. On Triton 3.3, the cliff occurs between 8192 and 8704; on Triton 3.4, it occurs between 4096 and 4608. Therefore, all of the shapes 4608 to 8192 have significant regressions (like 400%+) compared to Triton 3.3. Repro: python run.py --op layer_norm --baseline torch_layer_norm --metrics latency --only liger_layer_norm,torch_layer_norm --bwd --input-id 7 --num-inputs 1. Bisect points to [Backend] Bump to llvm/llvm-project@8957e64a20fc triton-lang/triton#7138, an LLVM pin update (so this is a regression in LLVM). Bisection on llvm points to [NVPTX] Remove Float register classes llvm/llvm-project#140487 as the cause of the regression
• liger embedding bwd: 20-53% regression across a range of shapes. Repro: python run.py --op embedding --baseline torch_embedding --metrics latency --bwd --only liger_embedding,torch_embedding,torch_embedding. Bisects to Update ptxas to version 12.8.93 triton-lang/triton#6149

Can't repro:

• rope bwd: This appears in the dataset but looks fine when I run locally. Possible that this is because the original data was collected on an old version of Liger. Repro: python run.py --op rope --baseline apply_rotary_pos_emb --metrics latency --only liger_rotary_pos_emb,apply_rotary_pos_emb
• low_mem_dropout: these appear to have high variability, reported regression is likely just noise
• flex_attention_fwd: 19.9% regression on x_ (8, 16, 512, 16, 512, 128) - likely noise (baseline appears noisy in some runs) Repro python run.py --op flex_attention --baseline eager --metrics latency,speedup --only compiled,eager

Biggest improvements
https://gist.github.com/davidberard98/ec81ec7a5c035db225053074e7e280d6

• flex attention!
• fp8_gemm_blockwise
• int4_gemm

[xuzhao9]

Curious how did you do the bisect? I wish we could integrate it into the CI for auto-bisection.

[davidberard98]

@xuzhao9 agreed, bisect feature would be amazing!

I used this script:

function finish() {
  pushd ../triton
  git checkout -- scripts/build-llvm-project.sh
}
git apply asdf.diff
trap finish EXIT
if [ -f setup.py ]; then
  pushd .
  make dev-install-llvm
else
  python -m pip install -e python
fi
code=$?
if [ $code -ne 0 ]
then
  exit 125
fi
pushd ../tritonbench
# python run.py --op layer_norm --baseline torch_layer_norm --metrics speedup --only liger_layer_norm,torch_layer_norm --bwd --input-id 7 --num-inputs 1 | python proc.py
python run.py --op embedding --baseline torch_embedding --metrics speedup --bwd --only liger_embedding,torch_embedding,torch_embedding --input-id 0 --num-inputs 1 | python proc.py
code=$?
exit $code

where

• proc.py just parses the torchbench output
• asdf.diff is just some build patches to enable ccache

[Artem-B]

Can you extract PTX before/after the change? It would help figuring out what exactly is causing the trouble.

[davidberard98]

Before (good): https://gist.github.com/davidberard98/9a4b820d39b7d2f89479d59def8f10c4
After (bad): https://gist.github.com/davidberard98/c316464b8c193cf45b8c779a443bafef

Diff: https://www.diffchecker.com/6JSwi95W/

notably:

	.reg .b32 	%r<1184>;
	.reg .b32 	%f<1317>;

changes to

	.reg .b32 	%r<2918>;

where 2918 > (1184 + 1317 = 2501)

and looking at the NCU trace, I think the new kernel is register spilling

[Artem-B]

AFAICT, there are no functional differences in the generated PTX. As the patch intended, we've eliminated a lot of redundant register moves, but the overall structure of the code is literally identical, modulo data moves between integer and FP registers.

I suspect that using integer registers everywhere may have triggered a bug or misoptimization in ptxas.

Next step is to examine the cubin produced from those PTX files and see if something odd pops up. Based on increased number of loads/stores, I suspect that ptxas may have gotten a bit over enthusiastic and started spilling registers into memory.

I think I've seen one case like that triggered by exactly this patch. It helped to curb ptxas' enthusiasm with --register-usage-level=4 which is one step below the default value of 5.

[Artem-B]

.reg .b32 %r<2918>;

This is not a very useful metric. Due to the way LLVM back-end works, that number is effectively the upper bar on the virtual registers used before we even get to generate PTX. A lot of them never materialize, but we can't change the numbering, so the upper bar remains.

If you count the actual unique registers used in the code, there way be way less of them that the number nominally allocated in both cases.
If anything, post-patch the code uses less registers because it no longer has to duplicate data across integer and FP registers.

The culprit here is almost certainly ptxas.

[davidberard98]

@Artem-B thanks for the investigation! I'll plan to file a bug with NVIDIA, unless you or @AlexMaclean have already?

[AlexMaclean]

@Artem-B thanks for the investigation! I'll plan to file a bug with NVIDIA, unless you or @AlexMaclean have already?

If you're able to file a bug that would be a big help! I think you'd probably be able to include more details then me and it may get addressed more rapidly. If you attach the bug number here I'll try to forward it on to the relevant people. Of course, if you prefer I can file one myself.

[Artem-B]

Just in case, which CUDA version do you currently use? Can you check if it's already fixed in the latest 12.9.1. I doubt it, but who knows. We may be dealing with a known (to NVIDIA) and possibly already fixed issue here.