@@ -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)"
@@ -0,0 +1,48 @@
+# 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"
+# The runner's GPU has no memory pools, which is what makes this job cover the
+# allocator's fallback; fail rather than skip those tests if it ever has them.
+$env:EXECUTORCH_CUDA_TEST_REQUIRE_NO_MEMORY_POOLS = "1"
+
+$tests = @(
+    "test_cuda_allocator",
+    "test_cuda_mutable_state",
+    "test_cuda_weight_cache",
+    "test_cuda_sort_rand",
+    "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 }
+}
@@ -17,6 +17,11 @@ on:
       - backends/cuda/**
       - backends/aoti/**
       - extension/cuda/**
+      # Read only by unittest-cuda-runtime-windows; the other jobs filter on
+      # their own paths and stay skipped.
+      - .ci/scripts/test_cuda_runtime_windows.ps1
+      - .ci/scripts/install_nvidia_driver_windows.ps1
+      - tools/cmake/Test.cmake
   workflow_dispatch:
 
 concurrency:
@@ -223,6 +228,41 @@ 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') ||
+      contains(needs.changed-files.outputs.changed-files, 'tools/cmake/Test.cmake') ||
+      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
 
@@ -577,10 +577,11 @@ jobs:
           -v -o "addopts="
 
         cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON
-        cmake --build cmake-out --target test_cuda_allocator test_cuda_allocator_exit test_cuda_mutable_state test_cuda_weight_cache test_cuda_guard test_cuda_stream_guard -j$(nproc)
+        cmake --build cmake-out --target test_cuda_allocator test_cuda_allocator_exit test_cuda_mutable_state test_cuda_weight_cache test_cuda_sort_rand test_cuda_guard test_cuda_stream_guard -j$(nproc)
         ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_mutable_state --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_weight_cache --output-on-failure -V
+        ctest --test-dir cmake-out -R test_cuda_sort_rand --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_guard --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_stream_guard --output-on-failure -V
 
 
@@ -475,4 +475,13 @@ if(BUILD_TESTING)
     EXTRA_LIBS aoti_cuda_backend
   )
   target_compile_definitions(test_cuda_weight_cache PRIVATE CUDA_AVAILABLE=1)
+  # sort.cu and rand.cu are compiled only under this condition (see
+  # _aoti_cuda_shim_sources), so the test that calls them is too.
+  if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER)
+    et_cxx_test(
+      test_cuda_sort_rand SOURCES runtime/test/test_cuda_sort_rand.cpp
+      EXTRA_LIBS aoti_cuda_backend
+    )
+    target_compile_definitions(test_cuda_sort_rand PRIVATE CUDA_AVAILABLE=1)
+  endif()
 endif()
@@ -89,6 +89,15 @@ constexpr int kCudaGraphWarmupSteps = 3;
 constexpr char kWeightSharingAcrossMethods[] = "weight_sharing_across_methods";
 } // anonymous namespace
 
+bool CudaGraphState::start_warmup() {
+  if (!CudaAllocator::memory_pools_supported(-1)) {
+    return false;
+  }
+  phase = CudaGraphPhase::Warmup;
+  warmup_remaining = kCudaGraphWarmupSteps;
+  return true;
+}
+
 // Advances the off-graph KV cache past a step's tokens. Every successful exit
 // of execute() ends here, and only those: a run that failed never wrote its
 // tokens, and committing them would desynchronise the cache's length from the
@@ -539,15 +548,22 @@ class ET_EXPERIMENTAL CudaBackend final
         handle->get_cuda_stream(),
         method_name.c_str());
 
-    // Initialize CUDA graph state if enabled for this method.
+    // Initialize CUDA graph state if enabled for this method. Not on a device
+    // without memory pools, where the method runs without a graph.
     if (should_use_cuda_graph_for_method(method_name)) {
-      handle->cuda_graph_state.phase = CudaGraphPhase::Warmup;
-      handle->cuda_graph_state.warmup_remaining = kCudaGraphWarmupSteps;
-      ET_LOG(
-          Info,
-          "CUDA graph enabled for method '%s' (warmup=%d)",
-          method_name.c_str(),
-          kCudaGraphWarmupSteps);
+      if (handle->cuda_graph_state.start_warmup()) {
+        ET_LOG(
+            Info,
+            "CUDA graph enabled for method '%s' (warmup=%d)",
+            method_name.c_str(),
+            kCudaGraphWarmupSteps);
+      } else {
+        ET_LOG(
+            Info,
+            "CUDA graph requested for method '%s' but this device has no "
+            "memory pools; running without a CUDA graph",
+            method_name.c_str());
+      }
     }
 
     mutable_state_note_handle(handle);
 
@@ -135,6 +135,13 @@ struct CudaGraphState {
 
   CudaGraphState() = default;
 
+  // Starts the warmup that leads to capture, when the current device can
+  // capture: without memory pools the allocator falls back to cudaMalloc,
+  // which a captured graph cannot own, so the phase stays Disabled there.
+  // Returns whether it started. Defined with the backend, which owns the
+  // warmup length.
+  bool start_warmup();
+
   ~CudaGraphState() {
     release();
   }
 
@@ -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/extension/cuda/cuda_allocator.h>
 #include <executorch/runtime/platform/assert.h>
 #include <executorch/runtime/platform/log.h>
 
@@ -20,6 +21,7 @@
 #include <cstdint>
 #include <ctime>
 #include <mutex>
+#include <new>
 #include <vector>
 
 namespace executorch::backends::cuda {
@@ -57,20 +59,37 @@ static std::once_flag g_rng_init_flag;
 
 // Initialize RNG state on the given stream.
 // Must be called during warmup (before graph capture). Subsequent calls
-// 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);
-    RngState h;
-    h.seed = static_cast<unsigned long long>(time(nullptr));
-    h.counter = 0;
-    h.base_scratch = 0;
-    cudaMemcpyAsync(
-        d_rng, &h, sizeof(RngState), cudaMemcpyHostToDevice, stream);
-    // Synchronize to ensure the copy completes before we return
-    // (the host-side RngState `h` is on the stack).
-    cudaStreamSynchronize(stream);
-  });
+// from any thread are no-ops thanks to std::call_once. A failed allocation
+// throws out of the call_once, which leaves it not done, so the next call
+// tries again rather than every later launch using a null state.
+bool ensure_rng_init(cudaStream_t stream) {
+  try {
+    std::call_once(g_rng_init_flag, [&]() {
+      // The current device: the stream callers pass from getCurrentCUDAStream
+      // may have been created on it whatever index they looked up, and the
+      // cudaMalloc fallback on a device without memory pools follows this
+      // index rather than the stream.
+      auto rng =
+          CudaAllocator::allocate_stream_ordered(sizeof(RngState), -1, stream);
+      if (!rng.ok()) {
+        throw std::bad_alloc();
+      }
+      d_rng = static_cast<RngState*>(rng.get());
+      RngState h;
+      h.seed = static_cast<unsigned long long>(time(nullptr));
+      h.counter = 0;
+      h.base_scratch = 0;
+      cudaMemcpyAsync(
+          d_rng, &h, sizeof(RngState), cudaMemcpyHostToDevice, stream);
+      // Synchronize to ensure the copy completes before we return
+      // (the host-side RngState `h` is on the stack).
+      cudaStreamSynchronize(stream);
+    });
+  } catch (const std::bad_alloc&) {
+    ET_LOG(Error, "rand: allocating the RNG state failed");
+    return false;
+  }
+  return true;
 }
 
 // Philox-based randint kernel. Reads its base offset from `rng->base_scratch`
@@ -211,7 +230,10 @@ AOTITorchError aoti_torch_cuda_rand(
       "aoti_torch_cuda_rand: failed to get CUDA stream");
   cudaStream_t stream = stream_result.get();
 
-  ensure_rng_init(stream);
+  ET_CHECK_OR_RETURN_ERROR(
+      ensure_rng_init(stream),
+      MemoryAllocationFailed,
+      "aoti_torch_cuda_rand: failed to allocate the RNG state");
 
   constexpr int kThreads = 256;
   int blocks = static_cast<int>((numel + kThreads - 1) / kThreads);
@@ -271,7 +293,10 @@ AOTITorchError aoti_torch_cuda_randint_low_out(
       "aoti_torch_cuda_randint_low_out: failed to get CUDA stream");
   cudaStream_t stream = stream_result.get();
 
-  ensure_rng_init(stream);
+  ET_CHECK_OR_RETURN_ERROR(
+      ensure_rng_init(stream),
+      MemoryAllocationFailed,
+      "aoti_torch_cuda_randint_low_out: failed to allocate the RNG state");
 
   int64_t range = high - low;
   int64_t* out_data = static_cast<int64_t*>(out->data_ptr());