Updated the handling of CUDAGeneratorImpl.h to new path

apex/contrib/csrc/fmha/src/fmha.h

@@ -30,7 +30,7 @@
 #include <cuda.h>
 #include <vector>
 
-#ifdef OLD_GENERATOR_PATH
+#if !defined(NEW_GENERATOR_PATH)
 #include <ATen/CUDAGeneratorImpl.h>
 #else
 #include <ATen/cuda/CUDAGeneratorImpl.h>

apex/contrib/csrc/multihead_attn/dropout.h

 #include <ATen/ATen.h>
 
-#ifdef OLD_GENERATOR_PATH
+#if !defined(NEW_GENERATOR_PATH)
 #include <ATen/CUDAGeneratorImpl.h>
 #else
 #include <ATen/cuda/CUDAGeneratorImpl.h>

apex/contrib/csrc/multihead_attn/softmax.h

@@ -3,7 +3,7 @@
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
 #include <curand_kernel.h>
 
-#ifdef OLD_GENERATOR_PATH
+#if !defined(NEW_GENERATOR_PATH)
 #include <ATen/CUDAGeneratorImpl.h>
 #else
 #include <ATen/cuda/CUDAGeneratorImpl.h>

apex/contrib/csrc/transducer/transducer_joint_kernel.cu

@@ -5,7 +5,7 @@
 #include <torch/extension.h>
 #include <ATen/AccumulateType.h>
 
-#ifdef OLD_GENERATOR_PATH
+#if !defined(NEW_GENERATOR_PATH)
 #include <ATen/CUDAGeneratorImpl.h>
 #else
 #include <ATen/cuda/CUDAGeneratorImpl.h>

setup.py

@@ -366,8 +366,8 @@ if "--deprecated_fused_lamb" in sys.argv or "--cuda_ext" in sys.argv:
 # See https://github.com/pytorch/pytorch/pull/70650
 generator_flag = []
 torch_dir = torch.__path__[0]
-if os.path.exists(os.path.join(torch_dir, "include", "ATen", "CUDAGeneratorImpl.h")):
-    generator_flag = ["-DOLD_GENERATOR_PATH"]
+if os.path.exists(os.path.join(torch_dir, "include", "ATen", "cuda", "CUDAGeneratorImpl.h")):
+    generator_flag = ["-DNEW_GENERATOR_PATH"]
 
 if "--fast_layer_norm" in sys.argv:
     sys.argv.remove("--fast_layer_norm")