Cherry-pick b2fdf9c4 from upstream Apex and resolve conflicts (#68)

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

@@ -30,7 +30,12 @@
 #include <cuda.h>
 #include <vector>
 
+#ifdef OLD_GENERATOR_PATH
+#include <ATen/CUDAGeneratorImpl.h>
+#else
 #include <ATen/cuda/CUDAGeneratorImpl.h>
+#endif
+
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
 
 #include <fmha_utils.h>

apex/contrib/csrc/multihead_attn/dropout.h

 #include <ATen/ATen.h>
 
-#ifdef OLD_GENERATOR
-#include <ATen/CUDAGenerator.h>
+#ifdef OLD_GENERATOR_PATH
+#include <ATen/CUDAGeneratorImpl.h>
 #else
 #include <ATen/cuda/CUDAGeneratorImpl.h>
 #endif
@@ -178,15 +178,10 @@ void apex_fused_dropout_cuda(scalar_t const *inputs, scalar_t *outputs,
   std::pair<uint64_t, uint64_t> rng_engine_inputs;
   {
     // See Note [Acquire lock when using random generators]
-#ifdef OLD_GENERATOR
-    std::lock_guard<std::mutex> lock(gen->mutex_);
-    rng_engine_inputs = gen->philox_engine_inputs(counter_offset);
-#else
     std::lock_guard<std::mutex> lock(gen.mutex());
     rng_engine_inputs =
         at::check_generator<at::CUDAGeneratorImpl>(gen)->philox_engine_inputs(
             counter_offset);
-#endif
   }
 
   apex_fused_dropout_kernel<scalar_t, accscalar_t, IndexType>
@@ -219,15 +214,10 @@ void apex_dropout_add_cuda(scalar_t const *inputs, scalar_t const *add_inputs,
   std::pair<uint64_t, uint64_t> rng_engine_inputs;
   {
     // See Note [Acquire lock when using random generators]
-#ifdef OLD_GENERATOR
-    std::lock_guard<std::mutex> lock(gen->mutex_);
-    rng_engine_inputs = gen->philox_engine_inputs(counter_offset);
-#else
     std::lock_guard<std::mutex> lock(gen.mutex());
     rng_engine_inputs =
         at::check_generator<at::CUDAGeneratorImpl>(gen)->philox_engine_inputs(
             counter_offset);
-#endif
   }
 
   apex_dropout_add_kernel<scalar_t, accscalar_t, IndexType>

apex/contrib/csrc/multihead_attn/softmax.h

 #pragma once
 #include "philox.h"
-#include <ATen/cuda/CUDAGeneratorImpl.h>
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
 #include <curand_kernel.h>
 
+#ifdef OLD_GENERATOR_PATH
+#include <ATen/CUDAGeneratorImpl.h>
+#else
+#include <ATen/cuda/CUDAGeneratorImpl.h>
+#endif
+
 #include <assert.h>
 #include <cfloat>
 #include <cmath>

apex/contrib/csrc/transducer/transducer_joint_kernel.cu

@@ -4,7 +4,13 @@
 
 #include <torch/extension.h>
 #include <ATen/AccumulateType.h>
+
+#ifdef OLD_GENERATOR_PATH
+#include <ATen/CUDAGeneratorImpl.h>
+#else
 #include <ATen/cuda/CUDAGeneratorImpl.h>
+#endif
+
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
 #include <c10/macros/Macros.h>

setup.py

@@ -362,11 +362,12 @@ if "--deprecated_fused_lamb" in sys.argv or "--cuda_ext" in sys.argv:
                           include_dirs=[os.path.join(this_dir, 'csrc')],
                           extra_compile_args = nvcc_args_fused_lamb if not IS_ROCM_PYTORCH else hipcc_args_fused_lamb))
 
-# Check, if ATen/CUDAGenerator.h is found, otherwise use the new ATen/CUDAGeneratorImpl.h, due to breaking change in https://github.com/pytorch/pytorch/pull/36026
+# Check, if ATen/CUDAGeneratorImpl.h is found, otherwise use ATen/cuda/CUDAGeneratorImpl.h
+# 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', 'CUDAGenerator.h')):
-    generator_flag = ['-DOLD_GENERATOR']
+if os.path.exists(os.path.join(torch_dir, "include", "ATen", "CUDAGeneratorImpl.h")):
+    generator_flag = ["-DOLD_GENERATOR_PATH"]
 
 if "--fast_layer_norm" in sys.argv:
     sys.argv.remove("--fast_layer_norm")