Important: This only speed up inference for now!!!
If you’ve trained or run inference on an AMD GPU with ROCm, you’ve probably run into this: there isn’t just one way to run a matmul, a conv3d, or a group norm. There are several — stock ROCm (rocBLAS, hipBLASLt, MIOpen), [aiter], [Composable Kernel], the ROCm fork of [TransformerEngine], and sometimes hand-rolled HIP kernels. Each one is faster than the others for some shape, dtype, and op combination — and slower for others.
Simple approach:
The core trick is embarrassingly simple, which is why it works:
1. amd_tuned_torch monkeypatches** a handful of torch / torch.nn.functional ops.
2. The first time it sees a call with a given shape and dtype, it benchmarks every available kernel candidate for that op — including stock PyTorch itself.
3. It caches the winner, keyed by shape.
4. Every subsequent call with that shape dispatches straight to the cached winner — zero benchmarking overhead after the first hit.
This will allow to use newer versions of kernel to be used with older version of pytorch and depending of your used pytorch version increase your performance .
If nothing beats stock for a given shape, it falls back to stock. So the worst case is “same as vanilla PyTorch plus one extra benchmark call the first time a shape appears.” There’s no scenario where this makes your model slower in steady state.
To use it just import it after torch .
import amd_tuned_torch# patches torch / torch.nn.functional on import
amd_tuned_torch.disable() # opt out at runtime
amd_tuned_torch.enable() # opt back in
Set AMD_TUNED_TORCH_AUTOPATCH=0 if you want to import the package without it touching anything.
** monkeypatches means overwriting the default used function in pytorch so every call affected.
What’s currently patched
- F.linear, torch.matmul, torch.bmm
- F.conv2d, F.conv3d
- F.group_norm
- F.scaled_dot_product_attention, F.rms_norm, F.gelu, F.silu (optional, TransformerEngine-backed)
TransformerEngine does bring currently no benefits and is by default disabled.
