feat: Add support for compute capability 12.1 (Grace Blackwell) (#785)

Author: multippt

nunchaku/utils.py

@@ -219,7 +219,7 @@ def get_precision(
             device = torch.device(device)
         capability = torch.cuda.get_device_capability(0 if device.index is None else device.index)
         sm = f"{capability[0]}{capability[1]}"
-        precision = "fp4" if sm == "120" else "int4"
+        precision = "fp4" if sm in ["120", "121"] else "int4"
     if pretrained_model_name_or_path is not None:
         if precision == "int4":
             if "fp4" in str(pretrained_model_name_or_path):
@@ -305,7 +305,7 @@ def check_hardware_compatibility(quantization_config: dict, device: str | torch.
         device = torch.device(device)
     capability = torch.cuda.get_device_capability(0 if device.index is None else device.index)
     sm = f"{capability[0]}{capability[1]}"
-    if sm == "120":  # you can only use the fp4 models
+    if sm in ["120", "121"]:  # you can only use the fp4 models
         if quantization_config["weight"]["dtype"] != "fp4_e2m1_all":
             raise ValueError('Please use "fp4" quantization for Blackwell GPUs. ')
     elif sm in ["75", "80", "86", "89"]:

setup.py

@@ -46,16 +46,16 @@ 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 == "120" and support_sm120:
-                sm = "120a"
-            assert sm in ["75", "80", "86", "89", "120a"], f"Unsupported SM {sm}"
+            if sm in ["120", "121"] and support_sm120:
+                sm = sm + "a"
+            assert sm in ["75", "80", "86", "89", "120a", "121a"], f"Unsupported SM {sm}"
             if sm not in ret:
                 ret.append(sm)
     else:
         assert install_mode == "ALL"
         ret = ["75", "80", "86", "89"]
         if support_sm120:
-            ret.append("120a")
+            ret.extend(["120a", "121a"])
     return ret