Pytorch matmul vs cudaTensorCoreGemm on Jetson Orin NX

Accelerated Computing CUDA CUDA Programming and Performance
,
1

Hello,

I have a Orin NX 16GB development Kit, that has CUDA 12.6.68. I have pyTorch 2.8.0 and numpy 1.21.5 on this system. I was exploring the cudaTensorCoreGemm.cu from the cuda-samples and wanted to compare the result of cudaTensorCoreGemm.cu’s compute_gemm kernel and a pytorch cuBLAS matmul. M,N,K for each is 16K.

I have modified cudaTensorCoreGemm.cu’s M_TILES, N_TILES and K_TILES defines such that each of them is 1024. With cudaTensorCoreGemm.cu and compute_gemm the result is:

Time: 149.821793 ms
TFLOPS: 58.71

I have attached the small pytorch program that I have written to perform a matmul for the same matrix sizes. Matrix A and B are FP16, while matrix C is FP32. However the result observed for this case is:

Time: 1002.24 msec
TFLOPS 8.78

A few things I notice with JTOP running is that:

  1. When using cudaTensorCoreGemm.cu and compute_gemm kernel one of the CPU cores is at 100% usage, no GPU activity is captured by JTOP.

  2. When using mm_cublas, the GPU is utilized as shown in JTOP.

  3. Forcing the cudaTensorCoreGemm.cu to not use compute_gemm but use simple_wmma_gemm, shows GPU activity.

It is not clear to me why there would be such a large discrepancy in results. I am hoping one the experts here could please point out what could be the reason for such behavior.

Any advice/suggestions is greatly appreciated.

Thank you,

vangogh

mm_cublas.txt (1.8 KB)

2

Questions about NVIDIA’s embedded system platforms are generally best asked in the sub-forums dedicated to them:

3

Okay. Thank you.