Fall back to cudaMalloc on CUDA devices without memory pools

.ci/scripts/install_nvidia_driver_windows.ps1

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Installs the NVIDIA display driver on a Windows GPU runner. The runner images carry the
+# GPU but no driver, so without this nvcuda.dll is missing and every CUDA call fails with
+# cudaErrorInsufficientDriver. The same driver, from the same place, that PyTorch's Windows
+# CUDA smoke tests install first (.ci/pytorch/windows/internal/driver_update.bat in
+# pytorch/pytorch), so both projects test against one driver.
+$ErrorActionPreference = "Stop"
+$version = "580.88"
+$installer = Join-Path $env:RUNNER_TEMP "$version-data-center-tesla-desktop-win10-win11-64bit-dch-international.exe"
+$url = "https://ossci-windows.s3.amazonaws.com/$(Split-Path $installer -Leaf)"
+
+$nvcuda = Join-Path $env:SystemRoot "System32\nvcuda.dll"
+if (Test-Path $nvcuda) {
+    Write-Host "NVIDIA driver already present: $((Get-Item $nvcuda).VersionInfo.FileVersion)"
+    exit 0
+}
+curl.exe --retry 3 -fsSL $url --output $installer
+if ($LASTEXITCODE -ne 0) { throw "downloading the NVIDIA driver from $url failed" }
+$process = Start-Process -FilePath $installer -ArgumentList "-s", "-noreboot" -Wait -PassThru
+Remove-Item $installer -ErrorAction SilentlyContinue
+if ($process.ExitCode -ne 0) { throw "the NVIDIA driver installer exited with $($process.ExitCode)" }
+if (-not (Test-Path $nvcuda)) { throw "the NVIDIA driver installed but $nvcuda is missing" }
+Write-Host "NVIDIA driver $version installed: $((Get-Item $nvcuda).VersionInfo.FileVersion)"

.ci/scripts/test_cuda_runtime_windows.ps1

@@ -0,0 +1,44 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Builds and runs the CUDA backend's runtime C++ tests on a Windows GPU runner, the
+# counterpart of the unittest-cuda-runtime job in cuda.yml. The Windows runners carry
+# data-center GPUs in TCC mode, which have no memory pools, so this is also what covers the
+# allocator's synchronous fallback; the Linux job covers the pool path.
+$ErrorActionPreference = "Stop"
+
+# Windows PowerShell 5.1 does not stop on a failing native command.
+function Invoke-Native {
+    param([Parameter(Mandatory = $true)][scriptblock]$Command)
+    & $Command
+    if ($LASTEXITCODE -ne 0) {
+        throw "exit code ${LASTEXITCODE}: $Command"
+    }
+}
+
+& "$PSScriptRoot\install_nvidia_driver_windows.ps1"
+
+$cudaNvcc = Join-Path $env:CUDA_HOME "bin\nvcc.exe"
+if (-not (Test-Path $cudaNvcc)) { throw "CUDA compiler not found at '$cudaNvcc'" }
+$env:CUDACXX = $cudaNvcc
+$env:PATH = "$env:CUDA_HOME\bin\x64;$env:CUDA_HOME\bin;$env:PATH"
+
+$tests = @(
+    "test_cuda_allocator",
+    "test_cuda_mutable_state",
+    "test_cuda_weight_cache",
+    "test_cuda_guard",
+    "test_cuda_stream_guard"
+)
+$numCores = [Math]::Max([Environment]::ProcessorCount - 1, 1)
+Invoke-Native {
+    cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON -DCMAKE_CXX_STANDARD=20 `
+        -T "cuda=$env:CUDA_HOME" "-DCMAKE_CUDA_COMPILER=$cudaNvcc" "-DCUDAToolkit_ROOT=$env:CUDA_HOME"
+}
+Invoke-Native { cmake --build cmake-out --config Release -j $numCores --target $tests }
+foreach ($test in $tests) {
+    Invoke-Native { ctest --test-dir cmake-out -C Release -R "^$test`$" --output-on-failure -V }
+}

.github/workflows/cuda-windows.yml

@@ -223,6 +223,40 @@ jobs:
           .ci/scripts/test_model_e2e_windows.ps1 -Device cuda-windows -HfModel '${{ matrix.model_repo }}/${{ matrix.model_name }}' -QuantName '${{ matrix.quant }}' -ModelDir \$artifactDir -ExpectedCudaVersion '13.0'
         }"
 
+  # The counterpart of cuda.yml's unittest-cuda-runtime on Windows. The Windows GPU runners
+  # carry data-center GPUs in TCC mode, which have no memory pools, so this also covers the
+  # allocator's synchronous fallback that the Linux job cannot reach.
+  unittest-cuda-runtime-windows:
+    name: unittest-cuda-runtime-windows
+    needs: [changed-files, run-decision]
+    if: |
+      contains(needs.changed-files.outputs.changed-files, 'backends/cuda') ||
+      contains(needs.changed-files.outputs.changed-files, 'backends/aoti') ||
+      contains(needs.changed-files.outputs.changed-files, 'extension/cuda') ||
+      contains(needs.changed-files.outputs.changed-files, '.github/workflows/cuda-windows.yml') ||
+      contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_cuda_runtime_windows.ps1') ||
+      contains(needs.changed-files.outputs.changed-files, '.ci/scripts/install_nvidia_driver_windows.ps1') ||
+      needs.run-decision.outputs.is-full-run == 'true'
+    uses: pytorch/test-infra/.github/workflows/windows_job.yml@main
+    with:
+      timeout: 120
+      runner: windows.g5.4xlarge.nvidia.gpu
+      gpu-arch-type: cuda
+      gpu-arch-version: "13.0"
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      script: |
+        git config --global http.sslBackend openssl
+        git submodule update --init --recursive
+        conda init powershell
+        powershell -Command "& {
+          Set-PSDebug -Trace 1
+          \$ErrorActionPreference = 'Stop'
+          \$env:CUDA_HOME = 'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0'
+          \$env:CUDA_PATH = \$env:CUDA_HOME
+          .ci/scripts/setup-windows.ps1
+          .ci/scripts/test_cuda_runtime_windows.ps1
+        }"
+
   delete-model-cuda-windows-artifacts:
     name: delete-model-cuda-windows-artifacts
     # The exports exist only to hand the models to the Windows job above, about

backends/cuda/runtime/cuda_allocator.cpp

@@ -121,6 +121,38 @@ cudaMemPool_t mem_pool_for(int device) {
   return pool;
 }
 
+// Whether `device` supports the stream-ordered allocator. Devices without
+// memory pools, such as data-center GPUs in TCC mode on Windows, reject
+// cudaMallocAsync and cudaMemPoolCreate with cudaErrorNotSupported, so their
+// allocations take the synchronous path instead. Cached because it is asked on
+// every allocation. An attribute query that fails is treated as supported,
+// which keeps the stream-ordered path and its own error reporting.
+bool memory_pools_supported(int device) {
+  static std::mutex mutex;
+  static std::unordered_map<int, bool> cache;
+  const std::lock_guard<std::mutex> lock(mutex);
+  const auto it = cache.find(device);
+  if (it != cache.end()) {
+    return it->second;
+  }
+  int value = 0;
+  const cudaError_t err =
+      cudaDeviceGetAttribute(&value, cudaDevAttrMemoryPoolsSupported, device);
+  if (err != cudaSuccess) {
+    (void)cudaGetLastError();
+  }
+  const bool supported = err != cudaSuccess || value != 0;
+  if (!supported) {
+    ET_LOG(
+        Info,
+        "CUDA device %d does not support memory pools; stream-ordered "
+        "allocations use cudaMalloc and cudaFree instead.",
+        device);
+  }
+  cache.emplace(device, supported);
+  return supported;
+}
+
 #endif // !EXECUTORCH_USE_HIP
 
 Error copy_impl(
@@ -426,6 +458,13 @@ Result<void*> CudaAllocator::allocate_async(
     (void)cudaGetLastError();
     stream_device = -1;
   }
+  const int target = device >= 0 ? device : stream_device;
+  if (target >= 0 && !memory_pools_supported(target)) {
+    // Without memory pools there is no stream-ordered allocation to make. A
+    // synchronous allocation is usable from any stream at once, so the caller's
+    // ordering still holds; deallocate_async frees it to match.
+    return CudaAllocator::instance().allocate(nbytes, index);
+  }
   cudaMemPool_t pool =
       (device >= 0 && device == stream_device) ? mem_pool_for(device) : nullptr;
   if (device >= 0) {
@@ -460,6 +499,23 @@ void CudaAllocator::deallocate_async(
     return;
   }
 
+#if !defined(EXECUTORCH_USE_HIP)
+  // Decided by the device the memory lives on, the same answer allocate_async
+  // reached for it. cudaFree waits for all work on the device first, so work
+  // still using the block on `stream` finishes before it is released.
+  cudaPointerAttributes attributes{};
+  if (cudaPointerGetAttributes(&attributes, ptr) == cudaSuccess) {
+    if (attributes.type == cudaMemoryTypeDevice && attributes.device >= 0 &&
+        !memory_pools_supported(attributes.device)) {
+      CudaAllocator::instance().deallocate(
+          ptr, static_cast<DeviceIndex>(attributes.device));
+      return;
+    }
+  } else {
+    (void)cudaGetLastError();
+  }
+#endif
+
   cudaError_t err = cudaFreeAsync(ptr, stream);
   if (err != cudaSuccess) {
     ET_LOG(
@@ -535,7 +591,11 @@ void CudaAllocator::release_cached_memory(DeviceIndex index) {
     // Device scoped, unlike everything else in this function: it releases
     // unused graph memory cached by every user of the device, so another
     // library in this process pays to map its own graph allocations again.
-    // Nothing breaks, since only unused blocks go.
+    // Nothing breaks, since only unused blocks go. Graph memory is served by
+    // memory pools, so a device without them has none to release.
+    if (!memory_pools_supported(device)) {
+      continue;
+    }
     const cudaError_t graph_err = cudaDeviceGraphMemTrim(device);
     if (graph_err != cudaSuccess) {
       ET_LOG(

backends/cuda/runtime/shims/rand.cu

@@ -11,6 +11,7 @@
 #include <executorch/backends/aoti/slim/cuda/guard.h>
 #include <executorch/backends/aoti/slim/factory/empty.h>
 #include <executorch/backends/aoti/slim/util/size_util.h>
+#include <executorch/backends/cuda/runtime/cuda_allocator.h>
 #include <executorch/runtime/platform/assert.h>
 #include <executorch/runtime/platform/log.h>
 
@@ -60,7 +61,14 @@ static std::once_flag g_rng_init_flag;
 // from any thread are no-ops thanks to std::call_once.
 void ensure_rng_init(cudaStream_t stream) {
   std::call_once(g_rng_init_flag, [&]() {
-    cudaMallocAsync(&d_rng, sizeof(RngState), stream);
+    // Through the backend allocator, which falls back to cudaMalloc on a device
+    // without memory pools, where cudaMallocAsync is not supported.
+    auto rng = CudaAllocator::allocate_async(sizeof(RngState), -1, stream);
+    if (!rng.ok()) {
+      ET_LOG(Error, "rand: allocating the RNG state failed");
+      return;
+    }
+    d_rng = static_cast<RngState*>(rng.get());
     RngState h;
     h.seed = static_cast<unsigned long long>(time(nullptr));
     h.counter = 0;

backends/cuda/runtime/shims/sort.cu

@@ -17,6 +17,7 @@
 #include <new>
 
 #include <executorch/backends/aoti/utils.h>
+#include <executorch/backends/cuda/runtime/cuda_allocator.h>
 #include <executorch/backends/cuda/runtime/shims/memory.h>
 #include <executorch/backends/cuda/runtime/shims/sort.h>
 #include <executorch/backends/aoti/slim/cuda/guard.h>
@@ -111,21 +112,22 @@ void launch_permute(
 }
 
 // Stream-ordered scratch for thrust. With par_nosync this keeps each slice sort
-// from blocking on cudaMalloc/cudaFree and synchronizing the stream.
+// from blocking on cudaMalloc/cudaFree and synchronizing the stream. Through the
+// backend allocator, which falls back to those on a device without memory pools.
 struct StreamOrderedAllocator {
   using value_type = char;
 
   char* allocate(std::ptrdiff_t bytes) {
-    void* ptr = nullptr;
-    if (cudaMallocAsync(&ptr, static_cast<size_t>(bytes), stream) !=
-        cudaSuccess) {
+    auto ptr =
+        CudaAllocator::allocate_async(static_cast<size_t>(bytes), -1, stream);
+    if (!ptr.ok()) {
       throw std::bad_alloc();
     }
-    return static_cast<char*>(ptr);
+    return static_cast<char*>(ptr.get());
   }
 
   void deallocate(char* ptr, size_t) {
-    (void)cudaFreeAsync(ptr, stream);
+    CudaAllocator::deallocate_async(ptr, -1, stream);
   }
 
   cudaStream_t stream;
@@ -318,14 +320,21 @@ AOTITorchError aoti_torch_cuda_sort_stable(
       inner_size *= input_sizes[d];
     }
 
-    ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMallocAsync(
-        &temp_values_buf,
-        static_cast<size_t>(total_elements * elem_size),
-        stream));
-    ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMallocAsync(
-        &temp_indices_buf,
-        static_cast<size_t>(total_elements) * sizeof(int64_t),
-        stream));
+    auto values = CudaAllocator::allocate_async(
+        static_cast<size_t>(total_elements * elem_size), -1, stream);
+    ET_CHECK_OR_RETURN_ERROR(
+        values.ok(),
+        MemoryAllocationFailed,
+        "sort: allocating the transposed values failed");
+    temp_values_buf = values.get();
+    auto indices = CudaAllocator::allocate_async(
+        static_cast<size_t>(total_elements) * sizeof(int64_t), -1, stream);
+    if (!indices.ok()) {
+      CudaAllocator::deallocate_async(temp_values_buf, -1, stream);
+      ET_LOG(Error, "sort: allocating the transposed indices failed");
+      return Error::MemoryAllocationFailed;
+    }
+    temp_indices_buf = indices.get();
 
     // Gather: [outer, sort, inner] → [outer, inner, sort]
     launch_permute(
@@ -450,8 +459,8 @@ AOTITorchError aoti_torch_cuda_sort_stable(
         stream);
     ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR();
 
-    ET_CUDA_CHECK_OR_RETURN_ERROR(cudaFreeAsync(temp_values_buf, stream));
-    ET_CUDA_CHECK_OR_RETURN_ERROR(cudaFreeAsync(temp_indices_buf, stream));
+    CudaAllocator::deallocate_async(temp_values_buf, -1, stream);
+    CudaAllocator::deallocate_async(temp_indices_buf, -1, stream);
   }
 
   return Error::Ok;