Fix 20x slowdown of FP6 kernel due to device properties query (#1128)

Fix 20x slowdown of FP6 kernel due to device properties query (#1092) Replace `cudaGetDeviceProperties` with `cudaDeviceGetAttribute` Co-authored-by: Tobias van der Werff <[email protected]>
pytorch · Oct 21, 2024 · 99c8d52 · 99c8d52
1 parent 01cbc09
commit 99c8d52
Showing 1 changed file with 8 additions and 9 deletions.
diff --git a/torchao/csrc/cuda/fp6_llm/fp6_linear.cu b/torchao/csrc/cuda/fp6_llm/fp6_linear.cu
@@ -29,17 +29,16 @@
 inline bool isSM75GPU() {
     int device;
     cudaError_t err = cudaGetDevice(&device);
-    if (err != cudaSuccess) {
-        return false;
-    }
+    if (err != cudaSuccess) return false;
 
-    cudaDeviceProp props;
-    err = cudaGetDeviceProperties(&props, device);
-    if (err != cudaSuccess) {
-        return false;
-    }
+    int major, minor;
+    err = cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device);
+    if (err != cudaSuccess) return false;
+
+    err = cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device);
+    if (err != cudaSuccess) return false;
 
-    return (props.major == 7) && (props.minor == 5);
+    return (major == 7) && (minor == 5);
 }
 
 template<typename TilingConfig, typename OutputDataType, int EXPONENT, int MANTISSA>