add global clear_workspaces

Author: lvhan028

example/test_linear.py

@@ -116,3 +116,5 @@ def dequantize(qweight, qzeros, scales, group_size: int = 128):
 print(f'abs_diff {abs_diff:4f}, '
       f'rel_diff {rel_diff:4f}, '
       f'outliers {outliers:4f}')
+
+tm.Linear.clear_workspaces()

src/turbomind/api/python/bind.cpp

@@ -342,5 +342,6 @@ PYBIND11_MODULE(_turbomind_ext, m)
             auto _out   = TorchTensorToTurbomindTensor(out);
             auto stream = reinterpret_cast<cudaStream_t>(stream_id);
             return self->forward(*_in, *_out, stream);
-        });
+        })
+        .def_static("clear_workspaces", &turbomind::Linear::clearWorkspaces);
 }

src/turbomind/api/python/linear.cc

@@ -283,6 +283,10 @@ struct Linear::Impl {
         }
     }
 
+    static void clearWorkspaces() {
+        workspace_cache_.clear();
+    }
+
 private:
     static gemm::Workspace& getWorkspace(int device_id, cudaStream_t stream)
     {
@@ -295,9 +299,13 @@ struct Linear::Impl {
         }
 
         // create a new workspace if cache missed
-        auto workspace = std::shared_ptr<gemm::Workspace>(new gemm::Workspace, [](gemm::Workspace* p) {
-            cudaFreeAsync(p->barriers, 0);
-            cudaFreeAsync(p->partials, 0);
+        auto workspace = std::shared_ptr<gemm::Workspace>(new gemm::Workspace, [device_id](gemm::Workspace* p) {
+            int old{};
+            check_cuda_error(cudaGetDevice(&old));
+            check_cuda_error(cudaSetDevice(device_id));
+            check_cuda_error(cudaFree(p->barriers));
+            check_cuda_error(cudaFree(p->partials));
+            check_cuda_error(cudaSetDevice(old));
         });
 
         workspace->barriers_size = gemm::Gemm::kBarriersSize;
@@ -349,4 +357,9 @@ void Linear::forward(const Tensor& in, Tensor& out, cudaStream_t stream)
 {
     impl_->forward(in, out, stream);
 }
+
+void Linear::clearWorkspaces() {
+    Linear::Impl::clearWorkspaces();
+}
+
 }  // namespace turbomind

src/turbomind/api/python/linear.h

@@ -28,6 +28,8 @@ class Linear {
     void forward(const Tensor& in, Tensor& out, cudaStream_t stream = nullptr);
     ~Linear() {}
 
+    static void clearWorkspaces();
+
 private:
     struct Impl;
     std::shared_ptr<Impl> impl_;

turbomind/linear.py

@@ -188,3 +188,7 @@ def __call__(self, x: torch.Tensor):
 
     def to_half(x: torch.Tensor):
         return x.to(torch.half)
+
+    @classmethod
+    def clear_workspaces(cls):
+        return _turbomind_ext.Linear.clear_workspaces()