Compile FBGEMM on H100 (#2298)

Summary:
H100 requires a special `9.0a` arch list to gain full kernel compilation.

Pull Request resolved: #2298

Reviewed By: q10

Differential Revision: D58484593

Pulled By: xuzhao9

fbshipit-source-id: d58def67bcba27f0e58c79fac6d7329b267f3de2

userbenchmark/triton/install.py

@@ -10,8 +10,8 @@
 def install_fbgemm():
  cmd = ["pip", "install", "-r", "requirements.txt"]
  subprocess.check_call(cmd, cwd=str(FBGEMM_PATH.resolve()))
- # Build target A100(8.0) or H100(9.0)
- cmd = [sys.executable, "setup.py", "bdist_wheel", "--package_variant=genai", "-DTORCH_CUDA_ARCH_LIST=8.0;9.0"]
+ # Build target A100(8.0) or H100(9.0, 9.0a)
+ cmd = [sys.executable, "setup.py", "bdist_wheel", "--package_variant=genai", "-DTORCH_CUDA_ARCH_LIST=8.0;9.0;9.0a"]
  subprocess.check_call(cmd, cwd=str(FBGEMM_PATH.resolve()))
 
 def test_fbgemm():