Use driver API

Author: LyricZhao

csrc/jit/handle.hpp

@@ -8,7 +8,7 @@
 
 namespace deep_gemm {
 
-#if CUDART_VERSION >= 12080 and not defined(DG_JIT_USE_DRIVER_API)
+#if CUDART_VERSION >= 12080 and defined(DG_JIT_USE_RUNTIME_API)
 
 // Use CUDA runtime API
 using LibraryHandle = cudaLibrary_t;

setup.py

@@ -28,9 +28,9 @@
     'third-party/cutlass/include/cutlass',
 ]
 
-# Use driver API for older CUDA compatibility
-if int(os.environ.get('DG_JIT_USE_DRIVER_API', '0')):
-    cxx_flags.append('-DDG_JIT_USE_DRIVER_API')
+# Use runtime API
+if int(os.environ.get('DG_JIT_USE_RUNTIME_API', '0')):
+    cxx_flags.append('-DDG_JIT_USE_RUNTIME_API')
 
 
 class CustomBuildPy(build_py):