Low GEMM Performance on Hopper GPU with Small M Shapes

[happierpig]

Hi,

Thank you for the great library! I’m observing some unexpected performance with GEMM on Hopper GPUs when using small M dimensions. I followed the example in example14_autotune.py.

Compared to the PyTorch implementation, the performance is significantly lower — around 30% of the expected TFLOPS and memory bandwidth utilization.

Not sure I am correctly using the API — I would greatly appreciate any suggestions.

Environment:
• GPU: H200
• CUDA: 12.8
• PyTorch: 2.6.0
• nvmath-python: 0.3.0

Benchmark code:

import argparse

import torch
import nvmath
from triton.testing import do_bench


def profile(m, n, k, dtype):
    device = torch.device("cuda")
    assert isinstance(device, torch.device)

    X = torch.randn(m, k, device=device, dtype=dtype)
    Y = torch.randn(n, k, device=device, dtype=dtype)
    
    _torch_gemm = lambda: torch.matmul(X, Y)
    
    mm = nvmath.linalg.advanced.Matmul(X, Y)
    
    mm.plan(preferences={"limit":1000})
    mm.autotune(iterations=1000)
    
    # print(mm.algorithms[0].capabilities)
    _nvmath_gemm = lambda: mm.execute()
    
    t_torch = do_bench(_torch_gemm)
    t_nvmath = do_bench(_nvmath_gemm)
    
    return t_torch, t_nvmath


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="GEMM profile")
    parser.add_argument("--m", type=int, default=4096)
    parser.add_argument("--n", type=int, default=4096)
    parser.add_argument("--k", type=int, default=4096)
    args = parser.parse_args()

    print("Provider,Operation,dtype,m,n,k,Runtime,GB/s,GFLOPs")
    
    for dtype in [torch.float16, torch.bfloat16]:
        t_torch, t_nvmath = profile(args.m, args.n, args.k, dtype)

        m = args.m
        n = args.n
        k = args.k

        torch_mem_bd = 2 * (m * n + n * k + m * k) * 1e3 / t_torch / 1e9
        torch_gflops = 2 * m * n * k * 1e3 / t_torch / 1e9
        nv_mem_bd = 2 * (m * n + n * k + m * k) * 1e3 / t_nvmath / 1e9
        nv_gflops = 2 * m * n * k * 1e3 / t_nvmath / 1e9
        
        print(f"TORCH,0,{dtype},{args.m},{args.n},{args.k},{t_torch},{torch_mem_bd},{torch_gflops}")
        print(f"NVMATH,0,{dtype},{args.m},{args.n},{args.k},{t_nvmath},{nv_mem_bd},{nv_gflops}")

[samaid]

Thank you for bringing this to us, @happierpig. We are on it. On a surface you are doing everything right. We are trying to replicate your results internally. Will come back to you

[szkarpinski]

Hello, @happierpig, and thank you for reporting this issue. I ran your script on H200, but I'm unable to reproduce your issue.

In my experiments I ran your script in a loop for values of M between 100 and 4000 using the following bash line:

for m in {100..4100..100}; do python3 bench.py --m $m; done

and there was no significant difference between nvmath-python and PyTorch. You can try to recreate my experiment in a PyTorch container:

docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it --rm  nvcr.io/nvidia/pytorch:25.01-py3

# Then, inside the container
pip3 install nvmath-python[cu12]
cat > bench.py  # Paste your script and press Ctrl-D
for m in {100..4100..100}; do python3 bench.py --m $m | grep bfloat16; done

To investigate this further, I'd like to ask you for a few more details.

Can you please share the output of nvidia-smi command, and the information about your CPU?

Can you verify that both Torch and nvmath-python actually use CUDA 12.8? If you have some previous CUDA versions installed, non-obvious versions might be linked. Please run
python3 -c "import torch; print(torch.version.cuda)"
and
python3 -c "import nvmath; print(nvmath.bindings.cublasLt.get_version())"

How did you run the script for different M's to generate the plot? Did you do it manually or with a script? Sometimes the performance measurements might be affected by thermal throttling, as the temperature of the GPU can change over time. For example, if you first benchmarked torch for all M's and then nvmath for all M's, it might be the case that the clocks were lower for nvmath-python. However, if you ran them interleaved (like I did with my loop), it's unlikely for thermal effects to affect only nvmath-python, and not PyTorch.

Did you try running your experiment without autotuning? For now, the implementation autotuning is quite simple, and the results can be affected by some random factors (such as changing clocks or temperature). Improving it is on our roadmap. Autotuning is optional, as even without autotuning the cuBLAS heuristics should do a good job in choosing an appropriate algorithm. Maybe disabling the autotuning will help with the performance - and if not, we'll eliminate one factor from our investigation.

I hope that with more details we'll be able to understand the reason for the lower performance. Thank you!