[Bugfix] Release cached cuBLAS workspaces after daemon warmup

docs/features/preload.md

@@ -70,6 +70,11 @@ loading and before binding its socket, then writes the cache:
 vllm preload --model /path/to/model --tensor-parallel-size 4
 ```
 
+After warmup, the daemon releases its communication resources and, on CUDA,
+cached cuBLAS workspaces before accepting engine connections. The model weights
+remain resident for CUDA IPC sharing. This reduces the daemon's GPU memory
+overhead; the amount reclaimed depends on the model and allocator state.
+
 ## How it works
 
 1. `vllm preload` spawns one daemon process per GPU. Each daemon loads its

tests/model_executor/model_loader/test_weight_cache.py

@@ -686,6 +686,42 @@ def test_ipc_loader_copy_mode_reports_no_external_weight_memory():
     assert loader.get_external_weight_memory(None) == 0
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+@create_new_process_for_each_test("spawn")
+def test_daemon_releases_cublas_workspaces():
+    """Startup cleanup releases cached GEMM scratch without freeing live tensors."""
+    import os
+
+    import torch
+
+    from vllm.config import ModelConfig, VllmConfig
+    from vllm.distributed import init_distributed_environment
+    from vllm.model_executor.model_loader.weight_cache.daemon import WeightCacheDaemon
+
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:2"
+    os.environ["TORCH_CUBLAS_WORKSPACE_CACHE"] = "1"
+    config = VllmConfig(model_config=ModelConfig(model=QWEN_CASE.model))
+    address = f"tcp://127.0.0.1:{get_open_port()}"
+    server = WeightCacheDaemon(config, 0, 0, address)
+    torch.accelerator.set_device_index(0)
+    init_distributed_environment(1, 0, address, 0, backend="nccl")
+    torch._C._cuda_clearCublasWorkspaces()
+    matrix = torch.ones(128, 128, device="cuda")
+    result = torch.empty_like(matrix)
+    before = torch.accelerator.memory_allocated()
+    torch.mm(matrix, matrix, out=result)
+    torch.accelerator.synchronize()
+    assert torch.accelerator.memory_allocated() > before
+
+    server.release_distributed()
+    after = torch.accelerator.memory_allocated()
+    torch._C._cuda_clearCublasWorkspaces()
+    assert torch.accelerator.memory_allocated() == after, (
+        "daemon retained BLAS workspaces"
+    )
+    torch.testing.assert_close(result, torch.full_like(result, 128), rtol=0, atol=0)
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
 def test_daemon_reports_exported_weight_bytes():
     """The engine's budget excludes exactly the storages it maps from the

vllm/model_executor/model_loader/weight_cache/daemon.py

@@ -342,11 +342,11 @@ def load_model(self) -> None:
         )
 
     def release_distributed(self) -> None:
-        """Tear down the daemon's process groups after loading and warmup.
+        """Release process groups and BLAS workspaces after loading and warmup.
 
         Serving only exports CUDA IPC handles, so the communicators (NCCL,
-        custom all-reduce and FlashInfer workspaces) would otherwise pin GPU
-        memory outside the engine's budget for the daemon's lifetime.
+        custom all-reduce and FlashInfer workspaces) and cached BLAS scratch
+        would otherwise pin GPU memory for the daemon's lifetime.
         """
         free_before, _ = torch.accelerator.get_memory_info()
         # No rank may tear down a communicator while a peer is still in a
@@ -357,10 +357,17 @@ def release_distributed(self) -> None:
         destroy_model_parallel()
         destroy_distributed_environment()
         gc.collect()
+        if current_platform.is_cuda():
+            # ATen caches these separately from the model's exported storages.
+            clear_workspaces = getattr(torch._C, "_cuda_clearCublasWorkspaces", None)
+            if clear_workspaces is not None:
+                clear_workspaces()
+            else:
+                logger.debug("PyTorch does not expose BLAS workspace cleanup")
         torch.accelerator.empty_cache()
         free_after, _ = torch.accelerator.get_memory_info()
         logger.info(
-            "Weight cache %s daemon rank %d released its process groups (%s GiB freed)",
+            "Weight cache %s daemon rank %d released startup resources (%s GiB freed)",
             self.role,
             self.global_rank,
             format_gib(free_after - free_before),