From 48e02a888149057da192e80657ecacc2135f4dc3 Mon Sep 17 00:00:00 2001 From: PengYoun9 Date: Thu, 17 Sep 2026 10:38:00 +0800 Subject: [PATCH] [Fix][cuDNN] Avoid thread-local workspace pool for ConvEntry (#20361) Use AllocDataSpace and FreeDataSpace instead of AllocWorkspace and FreeWorkspace for cuDNN convolution workspace allocation. This avoids the lifetime dependency on the thread-local workspace pool during process teardown and fixes the shutdown crash. --- src/runtime/extra/contrib/cudnn/cudnn_utils.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/runtime/extra/contrib/cudnn/cudnn_utils.cc b/src/runtime/extra/contrib/cudnn/cudnn_utils.cc index 3edb20dbacbc..f7f02b2f01bb 100644 --- a/src/runtime/extra/contrib/cudnn/cudnn_utils.cc +++ b/src/runtime/extra/contrib/cudnn/cudnn_utils.cc @@ -160,12 +160,14 @@ void ConvEntry::UpdateWorkspace(const size_t wsize) { CleanWorkspace(); } workspace_size = wsize; - workspace = cuda_api->AllocWorkspace(device, workspace_size); + + DLDataType type_hint{kDLUInt, 8, 1}; + workspace = cuda_api->AllocDataSpace(device, wsize, runtime::kTempAllocaAlignment, type_hint); } } void ConvEntry::CleanWorkspace() { - if (workspace) cuda_api->FreeWorkspace(device, workspace); + if (workspace) cuda_api->FreeDataSpace(device, workspace); workspace_size = 0; }