feat: support libtorch

2025-10-21 09:55:33 +00:00
parent 239112cb4c
commit 301cbc1d75
22 changed files with 652 additions and 245 deletions
@@ -17,6 +17,11 @@ class DeviceRuntime {
    cublasLtHandle_t cublaslt_handle{};
    std::shared_ptr<torch::Tensor> cublaslt_workspace;

+    // cuBLASLt utils
+    static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024;
+    cublasLtHandle_t cublaslt_handle{};
+    std::shared_ptr<torch::Tensor> cublaslt_workspace;
+
 public:
    explicit DeviceRuntime() {
        cublaslt_workspace = std::make_shared<torch::Tensor>(torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA)));