Multiple updates and refactorings (#280)

csrc/jit/device_runtime.hpp

@@ -17,19 +17,7 @@ class DeviceRuntime {
     static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024;
 
 public:
-#if TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 3)
-    // For PyTorch 2.3+, share the PyTorch cuBLASLt handle
-    DeviceRuntime() = default;
-
-    static cublasLtHandle_t get_cublaslt_handle() {
-        return at::cuda::getCurrentCUDABlasLtHandle();
-    }
-
-    static torch::Tensor get_cublaslt_workspace() {
-        return torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA));
-    }
-#else
-    // Otherwise, create the cuBLASLt handle ourselves
+    // Create the cuBLASLt handle ourselves
     cublasLtHandle_t cublaslt_handle{};
     std::shared_ptr<torch::Tensor> cublaslt_workspace;
 
@@ -49,7 +37,6 @@ public:
     torch::Tensor get_cublaslt_workspace() const {
         return *cublaslt_workspace;
     }
-#endif
 
     std::shared_ptr<cudaDeviceProp> get_prop() {
         if (cached_prop == nullptr) {