update

Author: lmxyy

setup.py

@@ -47,7 +47,9 @@ def get_sm_targets() -> list[str]:
         for i in range(torch.cuda.device_count()):
             capability = torch.cuda.get_device_capability(i)
             sm = f"{capability[0]}{capability[1]}"
-            if sm in ["120", "121"] and support_sm120:
+            if sm in ["120"] and support_sm120:
+                sm = sm + "a"
+            if sm in ["121"] and support_sm121:
                 sm = sm + "a"
             assert sm in ["75", "80", "86", "89", "120a", "121a"], f"Unsupported SM {sm}"
             if sm not in ret: