fix cuda>=12.8 build

Author: yanjun.qiu

3rdparty/cutlass

@@ -17,13 +17,15 @@
 
 #pragma once
 
+
 #include "cute/util/debug.hpp"
 #include "cute/util/print.hpp"
 #include "cutlass/detail/helper_macros.hpp"
 #include <cstddef>
 #include <cstring>
 
-#if defined(CUDA_VERSION) && CUDA_VERSION > 12080
+#include <cuda.h>
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 12080
 #include <cuda/atomic>
 #else
 #include <cuda/std/atomic>

include/flux/cuda/cuda_common_device.hpp

@@ -24,7 +24,6 @@
 #include <ostream>
 #include <type_traits>
 #include <cuda_runtime_api.h>
-#include <cuda/std/atomic>
 #include <utility>
 #include "cutlass/detail/helper_macros.hpp"
 #ifdef FLUX_SHM_USE_NVSHMEM
@@ -52,6 +51,13 @@
 #include <nccl.h>
 #endif
 
+#include <cuda.h>
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 12080
+#include <cuda/atomic>
+#else
+#include <cuda/std/atomic>
+#endif
+
 #define NextRank(rank_) (((rank_) + 1) % kLocalWorldSize)
 #define NextLocalRank(rank_, node_) ((((rank_) + 1) % kNumaWorldSize + (node_) * kNumaWorldSize))
 #define PrevLocalRank(rank_, node_) \

src/gemm_rs/reduce_scatter_kernel.hpp

@@ -37,15 +37,20 @@
 
 #pragma once
 
-
 #include "cutlass/cutlass.h"
 #include "cutlass/float8.h"
 #include "cutlass/matrix_coord.h"
 #include "cutlass/complex.h"
 #include "cutlass/gemm/kernel/gemm_transpose_operands.h"
 #include <ctime>
 #include <type_traits>
+
+#include <cuda.h>
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 12080
+#include <cuda/atomic>
+#else
 #include <cuda/std/atomic>
+#endif
 #include "ag_scatter_grouped_problem_visitor.hpp"
 
 /////////////////////////////////////////////////////////////////////////////////////////////////

src/moe_ag_scatter/cutlass_impls/ag_scatter_gemm_grouped_with_absmax.h

@@ -52,8 +52,13 @@
 #include <type_traits>
 
 #include "cutlass/barrier.h"
-#include <cuda/std/atomic>
 
+#include <cuda.h>
+#if defined(CUDA_VERSION) && CUDA_VERSION >= 12080
+#include <cuda/atomic>
+#else
+#include <cuda/std/atomic>
+#endif
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 namespace cutlass {