diff --git a/include/mscclpp/switch_channel_device.hpp b/include/mscclpp/switch_channel_device.hpp index b52b65723..7b749f7a9 100644 --- a/include/mscclpp/switch_channel_device.hpp +++ b/include/mscclpp/switch_channel_device.hpp @@ -37,7 +37,11 @@ struct SwitchChannelDeviceHandle { SwitchChannelDeviceHandle::multimemStore(val, reinterpret_cast(mcPtr) + index); } - template + /// Vectorized multimem load+reduce. The optional `AccumT` template parameter selects the + /// accumulator: when `AccumT == __half` and `VectorType` is an FP8 vector type, the + /// `.acc::f16` variant of the instruction is used (faster but lower precision than the + /// default FP32 accumulator). For all other types `AccumT` is ignored. + template MSCCLPP_DEVICE_INLINE static VectorType multimemLoadReduce(VectorType* ptr) { VectorType val; if constexpr (std::is_same_v) { @@ -81,29 +85,71 @@ struct SwitchChannelDeviceHandle { : "l"(ptr) : "memory"); } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.e4m3x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.e4m3x4 %0, [%1];" + : "=r"(val.words[0]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.e4m3x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e4m3x4 {%0,%1}, [%2];" - : "=r"(val.words[0]), "=r"(val.words[1]) - : "l"(ptr) - : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v2.e4m3x4 {%0,%1}, [%2];" + : "=r"(val.words[0]), "=r"(val.words[1]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e4m3x4 {%0,%1}, [%2];" + : "=r"(val.words[0]), "=r"(val.words[1]) + : "l"(ptr) + : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e4m3x4 {%0,%1,%2,%3}, [%4];" - : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) - : "l"(ptr) - : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v4.e4m3x4 {%0,%1,%2,%3}, [%4];" + : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e4m3x4 {%0,%1,%2,%3}, [%4];" + : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) + : "l"(ptr) + : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.e5m2x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.e5m2x4 %0, [%1];" + : "=r"(val.words[0]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.e5m2x4 %0, [%1];" : "=r"(val.words[0]) : "l"(ptr) : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e5m2x4 {%0,%1}, [%2];" - : "=r"(val.words[0]), "=r"(val.words[1]) - : "l"(ptr) - : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v2.e5m2x4 {%0,%1}, [%2];" + : "=r"(val.words[0]), "=r"(val.words[1]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.v2.e5m2x4 {%0,%1}, [%2];" + : "=r"(val.words[0]), "=r"(val.words[1]) + : "l"(ptr) + : "memory"); + } } else if constexpr (std::is_same_v) { - asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e5m2x4 {%0,%1,%2,%3}, [%4];" - : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) - : "l"(ptr) - : "memory"); + if constexpr (std::is_same_v) { + asm("multimem.ld_reduce.relaxed.sys.global.add.acc::f16.v4.e5m2x4 {%0,%1,%2,%3}, [%4];" + : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) + : "l"(ptr) + : "memory"); + } else { + asm("multimem.ld_reduce.relaxed.sys.global.add.v4.e5m2x4 {%0,%1,%2,%3}, [%4];" + : "=r"(val.words[0]), "=r"(val.words[1]), "=r"(val.words[2]), "=r"(val.words[3]) + : "l"(ptr) + : "memory"); + } } else { static_assert(dependentFalse, "Not supported type"); } diff --git a/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu b/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu index 2d71cd638..95ec384bf 100644 --- a/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu +++ b/src/ext/collectives/allreduce/allreduce_nvls_block_pipeline.cu @@ -13,7 +13,7 @@ namespace collective { __device__ DeviceSemaphore deviceSemaphore[NUM_SEMAPHORES]; -template +template __global__ void __launch_bounds__(1024, 1) allreduceNvlsBlockPipeline([[maybe_unused]] const void* src, [[maybe_unused]] void* scratch, [[maybe_unused]] void* dst, @@ -105,7 +105,7 @@ __global__ void __launch_bounds__(1024, 1) deviceSemaphore[oriBid].acquire(); } __syncthreads(); - handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, reduceIterSize, tid, blockDim.x); + handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, reduceIterSize, tid, blockDim.x); __syncthreads(); if (tid == 0) { deviceSemaphore[nBlocksForCopy + bidForReduce * copyReduceRatio + i].release(); @@ -158,24 +158,19 @@ struct NvlsBlockPipelineAdapter { } else if constexpr (std::is_same_v) { // fp8_e4m3b15 is a software-only type with no hardware NVLS support. return cudaErrorNotSupported; - } else -#if defined(__CUDA_ARCH__) // Skip the __CUDA_ARCH__ < 1000 since FP8 has not been supported for NVLS - if constexpr (std::is_same_v || std::is_same_v) { - return cudaErrorNotSupported; - } else -#endif - { - using ChannelType = DeviceHandle; - allreduceNvlsBlockPipeline - <<>>(input, scratch, output, (ChannelType*)memoryChannels, - nvlsChannels, inputSize, scratchBufferSize, rank, nRanksPerNode); - return cudaGetLastError(); - } + } else { + using ChannelType = DeviceHandle; + allreduceNvlsBlockPipeline + <<>>(input, scratch, output, (ChannelType*)memoryChannels, nvlsChannels, + inputSize, scratchBufferSize, rank, nRanksPerNode); + return cudaGetLastError(); + } } }; void AllreduceNvlsBlockPipeline::initialize(std::shared_ptr comm) { nSwitchChannels_ = 8; + fp8NvlsSupported_ = isFp8NvlsSupported(); int nBaseChannels = 64; this->conns_ = setupConnections(comm); // setup semaphores @@ -187,12 +182,15 @@ void AllreduceNvlsBlockPipeline::initialize(std::shared_ptr comm) this->nvlsConnections_ = setupNvlsConnections(comm, nvlsBufferSize_, nSwitchChannels_); } -CommResult AllreduceNvlsBlockPipeline::allreduceKernelFunc(const std::shared_ptr ctx_void, const void* input, - void* output, size_t inputSize, DataType dtype, ReduceOp op, - cudaStream_t stream, int nBlocks, int nThreadsPerBlock, - const std::unordered_map& extras, - DataType accumDtype) { +CommResult AllreduceNvlsBlockPipeline::allreduceKernelFunc( + const std::shared_ptr ctx_void, const void* input, void* output, size_t inputSize, DataType dtype, + ReduceOp op, cudaStream_t stream, int nBlocks, int nThreadsPerBlock, + [[maybe_unused]] const std::unordered_map& extras, DataType accumDtype) { auto ctx = std::static_pointer_cast(ctx_void); + if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) { + WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction."); + return CommResult::CommInvalidArgument; + } AllreduceFunc allreduce = dispatch(op, dtype, accumDtype); if (!allreduce) { WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast(dtype)); diff --git a/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu b/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu index 3bb054dae..385aebd52 100644 --- a/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu +++ b/src/ext/collectives/allreduce/allreduce_nvls_warp_pipeline.cu @@ -11,7 +11,7 @@ namespace mscclpp { namespace collective { -template +template __global__ void __launch_bounds__(1024, 1) allreduceNvlsWarpPipeline([[maybe_unused]] const void* src, [[maybe_unused]] void* scratch, [[maybe_unused]] void* dst, @@ -85,7 +85,8 @@ __global__ void __launch_bounds__(1024, 1) asm volatile("bar.sync %0, %1;" ::"r"(1), "r"((NCOPY_WARPS + NREDUCE_WARPS) * WARP_SIZE) : "memory"); T* mcBuff = (T*)multicastPtr->mcPtr; size_t offset = blockScratchOffset + (it * copyPerIter) % scratchSizePerBlock; - handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, iterSize, tidInReduce, NREDUCE_WARPS * WARP_SIZE); + handleMultiLoadReduceStore(mcBuff, mcBuff, offset, offset, iterSize, tidInReduce, + NREDUCE_WARPS * WARP_SIZE); asm volatile("bar.sync %0, %1;" ::"r"(2), "r"((NRECV_COPY_WARPS + NREDUCE_WARPS) * WARP_SIZE) : "memory"); } if (warpId >= startRecvCopyWid && warpId < endRecvCopyWid) { @@ -121,24 +122,19 @@ struct NvlsWarpPipelineAdapter { } else if constexpr (std::is_same_v) { // fp8_e4m3b15 is a software-only type with no hardware NVLS support. return cudaErrorNotSupported; - } else -#if defined(__CUDA_ARCH__) // Skip the __CUDA_ARCH__ < 1000 since FP8 has not been supported for NVLS - if constexpr (std::is_same_v || std::is_same_v) { - return cudaErrorNotSupported; - } else -#endif - { - using ChannelType = DeviceHandle; - allreduceNvlsWarpPipeline - <<>>(input, scratch, output, (ChannelType*)memoryChannels, - nvlsChannels, inputSize, scratchBufferSize, rank, nRanksPerNode); - return cudaGetLastError(); - } + } else { + using ChannelType = DeviceHandle; + allreduceNvlsWarpPipeline + <<>>(input, scratch, output, (ChannelType*)memoryChannels, nvlsChannels, + inputSize, scratchBufferSize, rank, nRanksPerNode); + return cudaGetLastError(); + } } }; void AllreduceNvlsWarpPipeline::initialize(std::shared_ptr comm) { nSwitchChannels_ = 8; + fp8NvlsSupported_ = isFp8NvlsSupported(); int nBaseChannels = 64; this->conns_ = setupConnections(comm); // setup semaphores @@ -155,6 +151,10 @@ CommResult AllreduceNvlsWarpPipeline::allreduceKernelFunc( ReduceOp op, cudaStream_t stream, int nBlocks, int nThreadsPerBlock, [[maybe_unused]] const std::unordered_map& extras, DataType accumDtype) { auto ctx = std::static_pointer_cast(ctx_void); + if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) { + WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction."); + return CommResult::CommInvalidArgument; + } AllreduceFunc allreduce = dispatch(op, dtype, accumDtype); if (!allreduce) { WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast(dtype)); diff --git a/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu b/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu index e7f2028fa..48ac67ffc 100644 --- a/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu +++ b/src/ext/collectives/allreduce/allreduce_nvls_zero_copy.cu @@ -13,7 +13,7 @@ namespace collective { constexpr int MAX_NBLOCKS = 32; -template +template __global__ void __launch_bounds__(1024, 1) allreduceNvls([[maybe_unused]] mscclpp::DeviceHandle* memoryChannels, [[maybe_unused]] mscclpp::DeviceHandle* multicast, @@ -56,8 +56,8 @@ __global__ void __launch_bounds__(1024, 1) T* src = (T*)multicastPtr->mcPtr; T* dst = (T*)multicastOutPtr->mcPtr; if (curBlockSize > 0) { - handleMultiLoadReduceStore(src, dst, blockOffset + channelInOffset, blockOffset + channelOutOffset, curBlockSize, - threadIdx.x, blockDim.x); + handleMultiLoadReduceStore(src, dst, blockOffset + channelInOffset, blockOffset + channelOutOffset, + curBlockSize, threadIdx.x, blockDim.x); } __syncthreads(); if (threadIdx.x < nPeers) { @@ -80,17 +80,11 @@ struct NvlsAdapter { } else if constexpr (std::is_same_v) { // fp8_e4m3b15 is a software-only type with no hardware NVLS support. return cudaErrorNotSupported; - } else -#if (!defined(__CUDA_ARCH_SPECIFIC__) && !defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) || (__CUDA_ARCH__ < 1000) - if constexpr (std::is_same_v || std::is_same_v) { - return cudaErrorNotSupported; - } else -#endif - { + } else { using ChannelType = DeviceHandle; - allreduceNvls<<>>((ChannelType*)memoryChannels, nvlsChannels, - nvlsOutChannels, channelInOffset, channelOutOffset, - inputSize, rank, nRanksPerNode); + allreduceNvls + <<>>((ChannelType*)memoryChannels, nvlsChannels, nvlsOutChannels, + channelInOffset, channelOutOffset, inputSize, rank, nRanksPerNode); return cudaGetLastError(); } } @@ -102,6 +96,7 @@ void AllreduceNvls::initialize(std::shared_ptr comm) { cudaDeviceProp deviceProp; MSCCLPP_CUDATHROW(cudaGetDeviceProperties(&deviceProp, device)); computeCapabilityMajor_ = deviceProp.major; + fp8NvlsSupported_ = isFp8NvlsSupported(); nSwitchChannels_ = 32; this->conns_ = setupConnections(comm); // setup semaphores @@ -124,6 +119,10 @@ CommResult AllreduceNvls::allreduceKernelFunc(const std::shared_ptr ctx_vo return CommResult::CommInvalidArgument; } auto ctx = std::static_pointer_cast(ctx_void); + if (isNativeFp8DataType(dtype) && !fp8NvlsSupported_) { + WARN("FP8 NVLS allreduce requires device support for FP8 multimem reduction."); + return CommResult::CommInvalidArgument; + } AllreduceFunc allreduce = dispatch(op, dtype, accumDtype); if (!allreduce) { WARN("Unsupported operation or data type for allreduce, dtype=%d", static_cast(dtype)); diff --git a/src/ext/collectives/collective_utils.cc b/src/ext/collectives/collective_utils.cu similarity index 80% rename from src/ext/collectives/collective_utils.cc rename to src/ext/collectives/collective_utils.cu index 016c4a5cc..2868c9795 100644 --- a/src/ext/collectives/collective_utils.cc +++ b/src/ext/collectives/collective_utils.cu @@ -1,16 +1,93 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -#include "collective_utils.hpp" - #include #include #include +#include #include #include +#include "collective_utils.hpp" + namespace mscclpp { namespace collective { + +namespace { + +#if !defined(MSCCLPP_DEVICE_HIP) +__global__ void fp8NvlsSupportProbeKernel(int* supported) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \ + (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) + *supported = 1; +#else + *supported = 0; +#endif +} + +bool detectFp8NvlsSupport() { + AvoidCudaGraphCaptureGuard cgcGuard; + auto supportedDevice = mscclpp::detail::gpuCallocUnique(); + int supportedHost = 0; + auto stream = gpuStreamPool()->getStream(); + + fp8NvlsSupportProbeKernel<<<1, 1, 0, stream>>>(supportedDevice.get()); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + return false; + } + + MSCCLPP_CUDATHROW( + cudaMemcpyAsync(&supportedHost, supportedDevice.get(), sizeof(supportedHost), cudaMemcpyDeviceToHost, stream)); + err = cudaStreamSynchronize(stream); + if (err != cudaSuccess) { + (void)cudaGetLastError(); + return false; + } + return supportedHost != 0; +} +#endif + +} // namespace + +bool isFp8DataType(DataType dtype) { + return dtype == DataType::FLOAT8_E4M3FN || dtype == DataType::FLOAT8_E4M3FNUZ || dtype == DataType::FLOAT8_E5M2 || + dtype == DataType::FLOAT8_E5M2FNUZ || dtype == DataType::FLOAT8_E4M3B15; +} + +bool isNativeFp8DataType(DataType dtype) { +#if defined(__FP8_TYPES_EXIST__) +#if defined(__FP8_E4M3_IS_FNUZ__) + if (dtype == DataType::FLOAT8_E4M3FNUZ) { + return true; + } +#else + if (dtype == DataType::FLOAT8_E4M3FN) { + return true; + } +#endif +#if defined(__FP8_E5M2_IS_FNUZ__) + if (dtype == DataType::FLOAT8_E5M2FNUZ) { + return true; + } +#else + if (dtype == DataType::FLOAT8_E5M2) { + return true; + } +#endif +#endif + return false; +} + +bool isFp8NvlsSupported() { +#if defined(MSCCLPP_DEVICE_HIP) + return false; +#else + static const bool supported = detectFp8NvlsSupport(); + return supported; +#endif +} + std::vector setupRemoteMemories(std::shared_ptr comm, int rank, mscclpp::RegisteredMemory localMemory) { std::vector remoteMemories; diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp index 81b74add4..b408c64cc 100644 --- a/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp +++ b/src/ext/collectives/include/allreduce/allreduce_nvls_block_pipeline.hpp @@ -33,6 +33,7 @@ class AllreduceNvlsBlockPipeline : public AlgorithmBuilder { std::vector baseChannels_; std::vector conns_; std::vector> nvlsConnections_; + bool fp8NvlsSupported_{false}; }; } // namespace collective } // namespace mscclpp diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp index 8f02a8738..2ce3a4fbb 100644 --- a/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp +++ b/src/ext/collectives/include/allreduce/allreduce_nvls_warp_pipeline.hpp @@ -33,6 +33,7 @@ class AllreduceNvlsWarpPipeline : public AlgorithmBuilder { std::vector baseChannels_; std::vector conns_; std::vector> nvlsConnections_; + bool fp8NvlsSupported_{false}; }; } // namespace collective } // namespace mscclpp diff --git a/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp b/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp index d53ea180b..edc6bcebd 100644 --- a/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp +++ b/src/ext/collectives/include/allreduce/allreduce_nvls_zero_copy.hpp @@ -36,6 +36,7 @@ class AllreduceNvls : public AlgorithmBuilder { std::vector> nvlsConnections_; std::vector> nvlsOutConnections_; int computeCapabilityMajor_{0}; + bool fp8NvlsSupported_{false}; }; } // namespace collective diff --git a/src/ext/collectives/include/allreduce/common.hpp b/src/ext/collectives/include/allreduce/common.hpp index 93b18e262..6f9a3d4cc 100644 --- a/src/ext/collectives/include/allreduce/common.hpp +++ b/src/ext/collectives/include/allreduce/common.hpp @@ -36,9 +36,21 @@ MSCCLPP_DEVICE_INLINE constexpr std::size_t calcVectorSize() { } } -template +template MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t srcOffset, size_t dstOffset, size_t size, int tid, int nThreads) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \ + (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) + constexpr bool fp8NvlsArchSupported = true; +#else + constexpr bool fp8NvlsArchSupported = false; +#endif + constexpr bool nativeFp8Type = std::is_same_v || std::is_same_v; + if constexpr (nativeFp8Type && !fp8NvlsArchSupported) { + MSCCLPP_ASSERT_DEVICE(false, "FP8 NVLS multimem reduction is not supported on this architecture"); + return; + } + // nvls can only handle 4 bytes alignment MSCCLPP_ASSERT_DEVICE(size % 4 == 0, "size must be 4 bytes aligned"); constexpr size_t nElem = calcVectorSize(); @@ -54,7 +66,7 @@ MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t src vectorType* src4 = (vectorType*)src; vectorType* dst4 = (vectorType*)dst; for (size_t idx = tid; idx < nVec; idx += nThreads) { - auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce(src4 + srcOffset4 + idx); + auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce(src4 + srcOffset4 + idx); mscclpp::SwitchChannelDeviceHandle::multimemStore(val, dst4 + dstOffset4 + idx); } // handle rest of data @@ -64,7 +76,8 @@ MSCCLPP_DEVICE_INLINE void handleMultiLoadReduceStore(T* src, T* dst, size_t src const size_t startIdx = (srcOffset + processed) / sizeof(restVectorType); const size_t endIdx = (srcOffset + size) / sizeof(restVectorType); for (size_t idx = tid + startIdx; idx < endIdx; idx += nThreads) { - auto val = mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce((restVectorType*)src + idx); + auto val = + mscclpp::SwitchChannelDeviceHandle::multimemLoadReduce((restVectorType*)src + idx); mscclpp::SwitchChannelDeviceHandle::multimemStore(val, (restVectorType*)dst + idx); } } diff --git a/src/ext/collectives/include/collective_utils.hpp b/src/ext/collectives/include/collective_utils.hpp index f705a9d1d..b7f86473d 100644 --- a/src/ext/collectives/include/collective_utils.hpp +++ b/src/ext/collectives/include/collective_utils.hpp @@ -32,6 +32,10 @@ constexpr int MAX_NRANKS_PER_NODE = 8; constexpr int SCRATCH_SIZE = 2 * 1024 * 1024 * 70; // double buffer * 35 thread-blocks * 8 ranks * 256KB = 70MB +bool isFp8DataType(DataType dtype); +bool isNativeFp8DataType(DataType dtype); +bool isFp8NvlsSupported(); + std::vector setupRemoteMemories(std::shared_ptr comm, int rank, RegisteredMemory localMemory); diff --git a/src/ext/nccl/CMakeLists.txt b/src/ext/nccl/CMakeLists.txt index 9767e66f8..463b75506 100644 --- a/src/ext/nccl/CMakeLists.txt +++ b/src/ext/nccl/CMakeLists.txt @@ -13,6 +13,7 @@ target_include_directories(mscclpp_nccl PRIVATE include ${PROJECT_SOURCE_DIR}/include ${PROJECT_SOURCE_DIR}/src/core/include + ${PROJECT_SOURCE_DIR}/src/ext/collectives/include ${GPU_INCLUDE_DIRS} ) target_link_libraries(mscclpp_nccl PUBLIC mscclpp mscclpp_collectives) diff --git a/src/ext/nccl/algorithm_selector.cc b/src/ext/nccl/algorithm_selector.cc index c94aab34a..1ccac65d4 100644 --- a/src/ext/nccl/algorithm_selector.cc +++ b/src/ext/nccl/algorithm_selector.cc @@ -6,6 +6,7 @@ #include #include +#include "collective_utils.hpp" #include "debug.h" namespace mscclpp { @@ -20,24 +21,15 @@ static bool isNvlsSupportedForDataType(const AlgorithmSelectorConfig& config, Da return false; } - const bool isFp8 = dtype == DataType::FLOAT8_E4M3FN || dtype == DataType::FLOAT8_E4M3FNUZ || - dtype == DataType::FLOAT8_E5M2 || dtype == DataType::FLOAT8_E5M2FNUZ; - - if (!isFp8) { + if (!collective::isFp8DataType(dtype)) { return nvlsSupported; } - // FP8 handling #if !defined(__HIP_PLATFORM_AMD__) - // NVLS does not support FP8 on devices with compute capability < 10 - if (config.computeCapability.first < 10) { + if (!collective::isNativeFp8DataType(dtype)) { return false; } -#if (defined(__CUDA_ARCH_SPECIFIC__) || defined(__CUDA_ARCH_FAMILY_SPECIFIC__)) - return true; -#else - return false; -#endif + return nvlsSupported && config.fp8NvlsSupported; #else return nvlsSupported; #endif diff --git a/src/ext/nccl/algorithm_selector.hpp b/src/ext/nccl/algorithm_selector.hpp index c8705f8bf..2048ea05c 100644 --- a/src/ext/nccl/algorithm_selector.hpp +++ b/src/ext/nccl/algorithm_selector.hpp @@ -16,6 +16,7 @@ namespace nccl { struct AlgorithmSelectorConfig { bool symmetricMemory; bool nvlsSupported; + bool fp8NvlsSupported; bool isCuMemMapAllocated; bool inCaptureMode; std::pair computeCapability; diff --git a/src/ext/nccl/nccl.cc b/src/ext/nccl/nccl.cc index 2d6c5f9d8..8fcc1bb15 100644 --- a/src/ext/nccl/nccl.cc +++ b/src/ext/nccl/nccl.cc @@ -20,6 +20,7 @@ #include #include "algorithm_selector.hpp" +#include "collective_utils.hpp" #include "datatype_conversion.hpp" static constexpr auto MSCCLPP_NCCL = mscclpp::LogSubsys::NCCL; @@ -239,6 +240,8 @@ static std::shared_ptr algoSelector( static const bool isNvlsSupported = mscclpp::isNvlsSupported(); static const std::pair deviceComputeCapability = getDeviceComputeCapability(); static const bool ncclSymmetricMemory = mscclpp::env()->ncclSymmetricMemory; + const bool fp8NvlsSupported = + mscclpp::collective::isNativeFp8DataType(request.dtype) ? mscclpp::collective::isFp8NvlsSupported() : false; const bool isCuMemMapAllocated = mscclpp::isCuMemMapAllocated(const_cast(request.inputBuffer)) && mscclpp::isCuMemMapAllocated(request.outputBuffer); @@ -249,6 +252,7 @@ static std::shared_ptr algoSelector( mscclpp::nccl::AlgorithmSelectorConfig config{.symmetricMemory = ncclSymmetricMemory, .nvlsSupported = isNvlsSupported, + .fp8NvlsSupported = fp8NvlsSupported, .isCuMemMapAllocated = isCuMemMapAllocated, .inCaptureMode = inCaptureMode, .computeCapability = deviceComputeCapability,