Seperate headers for GPU data types

include/mscclpp/gpu.hpp

@@ -6,8 +6,6 @@
 
 #if defined(__HIP_PLATFORM_AMD__)
 
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp16.h>
 #include <hip/hip_runtime.h>
 
 using cudaError_t = hipError_t;
@@ -92,14 +90,7 @@ constexpr auto CU_MEM_ACCESS_FLAGS_PROT_READWRITE = hipMemAccessFlagsProtReadWri
 #else
 
 #include <cuda.h>
-#include <cuda_fp16.h>
 #include <cuda_runtime.h>
-#if (CUDART_VERSION >= 11000)
-#include <cuda_bf16.h>
-#endif
-#if (CUDART_VERSION >= 11080)
-#include <cuda_fp8.h>
-#endif
 
 #endif
 

include/mscclpp/gpu_data_types.hpp

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#ifndef MSCCLPP_GPU_DATA_TYPES_HPP_
+#define MSCCLPP_GPU_DATA_TYPES_HPP_
+
+#if defined(__HIP_PLATFORM_AMD__)
+
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp16.h>
+
+#else
+
+#include <cuda_fp16.h>
+#if (CUDART_VERSION >= 11000)
+#include <cuda_bf16.h>
+#endif
+#if (CUDART_VERSION >= 11080)
+#include <cuda_fp8.h>
+#endif
+
+#endif
+
+#endif  // MSCCLPP_GPU_DATA_TYPES_HPP_

include/mscclpp/nvls_device.hpp

@@ -5,6 +5,7 @@
 #define MSCCLPP_NVLS_DEVICE_HPP_
 
 #include <mscclpp/gpu.hpp>
+#include <mscclpp/gpu_data_types.hpp>
 #include <type_traits>
 
 #include "device.hpp"

src/include/execution_kernel.hpp

@@ -5,6 +5,7 @@
 #define MSCCLPP_EXECUTION_KERNEL_HPP_
 
 #include <mscclpp/executor.hpp>
+#include <mscclpp/gpu_data_types.hpp>
 #include <mscclpp/packet_device.hpp>
 #include <mscclpp/proxy_channel.hpp>
 #include <mscclpp/sm_channel.hpp>