fix: 多卡剩余显存计算

ztxz16 · Jul 24, 2024 · 7490802 · 7490802
1 parent b3b3503
commit 7490802
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/src/devices/cuda/fastllm-cuda.cu b/src/devices/cuda/fastllm-cuda.cu
@@ -120,6 +120,9 @@ std::vector <long long> FastllmCudaGetFreeSizes() {
     std::vector <long long> ret;
 
     // 遍历所有设备  
+    int id = -1;
+    cudaGetDevice(&id);
+
     for (int i = 0; i < deviceCount; ++i) {
         cudaDeviceProp prop;
         error = cudaGetDeviceProperties(&prop, i);
@@ -129,6 +132,7 @@ std::vector <long long> FastllmCudaGetFreeSizes() {
             // printf("  Total global memory: %zu bytes\n", prop.totalGlobalMem);
 
             // 获取当前设备的显存使用情况  
+            cudaSetDevice(i);
             size_t free = 0, total = 0;
             cudaMemGetInfo(&free, &total);
             ret.push_back(free);
@@ -138,6 +142,7 @@ std::vector <long long> FastllmCudaGetFreeSizes() {
             printf("cudaGetDeviceProperties returned %d\n-> %s\n", (int)error, cudaGetErrorString(error));
         }
     }
+    cudaSetDevice(id);
     return ret;
 }
 

diff --git a/src/models/basellm.cpp b/src/models/basellm.cpp
@@ -489,9 +489,21 @@ namespace fastllm {
                     long long kvCacheLimit = 16LL << 30;
 #ifdef USE_CUDA
                     auto freeSizes = FastllmCudaGetFreeSizes();
+                    auto dmap = GetDeviceMap();
+                    std::set <int> deviceIds;
+                    for (auto &it : dmap) {
+                        if (StartWith(it.first, "cuda")) {
+                            for (int id : ParseDeviceIds(it.first, "cuda")) {
+                                deviceIds.insert(id);
+                            }
+                        }
+                    }
+                    if (deviceIds.size() == 0) {
+                        deviceIds.insert(0);
+                    }
                     kvCacheLimit = 0;
-                    for (long long i : freeSizes) {
-                        kvCacheLimit += std::max(0LL, i - (2LL << 30));
+                    for (int id : deviceIds) {
+                        kvCacheLimit += std::max(0LL, freeSizes[id] - (2LL << 30));
                     }
 #endif
                     if (model->kvCacheLimit > 0) {
@@ -803,10 +815,14 @@ printf("len = %d, spend = %f s. tokens / s = %f\n", (int)total, spend, (float)to
                 } else {
                     if (context->isEnding) {
                         responseContextDict.RemoveHandle(handleId);
+                        dictLocker.unlock();
+                        dictCV.notify_one();
                         if (context->error == ResponseContextErrorNone) {
                             return -1;
                         } else if (context->error == ResponseContextErrorPromptTooLong) {
                             return -2;
+                        } else {
+                            return -1;
                         }
                     }
                 }
@@ -836,6 +852,8 @@ printf("len = %d, spend = %f s. tokens / s = %f\n", (int)total, spend, (float)to
                 } else {
                     if (context->isEnding) {
                         responseContextDict.RemoveHandle(handleId);
+                        dictLocker.unlock();
+                        dictCV.notify_one();
                         return -1;
                     }
                 }