FP6 Speed on A100 80g

[shihaobai]

ENV:
cuda: 12.1
torch: 2.5.0+cu121
python benchmark_fp6.py

Hello, have you tested the performance of the FP6 kernel on the A100? I found that the speed is much slower compared to FP16."

[gau-nernst]

Looks like related to #1092 (the speedup numbers are similar). What is your torchao version? Can you try update torchao or install nightly / from source? Should be fixed in 0.6.1

[shihaobai]

Thanks for your help! My torchao version is torchao-0.7.0.dev20241028+cu121. I tried the 0.6.1 and got the correct performance.

[gau-nernst]

torchao-0.7.0.dev20241028+cu121 should have the correct fix I think. Can you double check that torchao-0.7.0.dev20241028+cu121 is also working correctly?

If everything works as expected, let me know so I can close the issue.

cc @tobiasvanderwerff

[shihaobai]

Thanks for your help. I tried the latest torchao==0.7.0+gitcbd90e38 and it worked correctly. But when i installed the torchao-0.7.0.dev20241028+cu121 again, I encounterd the bug:

ENV:
cuda: 12.1
torch: 2.5.0+cu121
pip install torchao==0.7.0.dev20241028 --index-url https://download.pytorch.org/whl/nightly/cu121

[tobiasvanderwerff]

@shihaobai Have you tried recompiling the C++/CUDA code by running pip install . in the base directory of ao? This might help resolve the error.

[shihaobai]

@tobiasvanderwerff I tried based on the latest commit and it worked correctly.