Closed anupambhatnagar closed 3 months ago
Definitely curious what’s going on here. This is less than half of peak bandwidth; by contrast on A100 we get essentially 100% of peak bw
@bertmaher Not sure whether you already figured this out. If you change the kernel signature to:
@triton.jit
def _rms_norm_kernel(
x_ptr,
h1_ptr,
w_ptr,
eps,
stride,
N_COLS: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
INCLUDE_WEIGHT: tl.constexpr,
):
i.e., changing N_COLS
to const, there is pretty big performance boost to:
22 12288.0 2247.299783
23 12800.0 2223.278530
24 13312.0 2266.673025
25 13824.0 2274.271469
26 14336.0 2324.921587
27 14848.0 2327.270964
28 15360.0 2374.492701
29 15872.0 2360.028753
though we still have room to improve.
@scxiao I made the suggested change you mentioned along with some others and now bandwidth is slightly over to 3TB/s. Analyzing the LLVM IR I found that the loads and stores are not being vectorized at the LLVM IR layer. We are working on a fix that we hope to push upstream. https://github.com/ROCmSoftwarePlatform/triton/pull/445
Closing due to inactivity
@anupambhatnagar I know this is a closed issue, the Babelstream results that you pasted above, is it from this repo? https://github.com/UoB-HPC/BabelStream I don't see the Read, Write kernels that you have listed in the results, is this some modified version of Babelstream that you are running
The RMS norm implementation below achieves well below the peak possible memory bandwidth on MI300. The results can be reproduced using the code below.
When benchmarking BabelStream on the same device I achieved much better performance. Any insights on how to achieve better perf on RMS norm will be appreciated. Thank you!
RMS norm implementation
RMS Norm memory bandwidth achieved in GB/sec.
BabelStream Performance