From bfb4d240a37b1b2a0a7197b5fc722c229c8e6f67 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 4 Aug 2026 17:17:21 +0000 Subject: [PATCH 1/6] Add bulk asynchronous copy primitives Expose the sm_90+ bulk copy engine (cp.async.bulk / cp.reduce.async.bulk) as a pointer-based device API in a new header, mscclpp/bulk_device.hpp. The primitives are channel-free by design. Bulk copies move bytes between global and shared memory and do not care whether the global side is local or peer mapped; channels answer where a peer's memory is and how to synchronize with it. Those concerns compose through a plain pointer. The expert-parallel kernels that motivate this feature gather from arrays of raw peer pointers and use channels only for signal/wait, so binding the primitives to a channel would not serve them. This also matches how SwitchChannel exposes multimem: a public, pointer-based primitive is the foundation. Surface: BulkBarrier load completion, with caller-held phase parity bulkLoad global -> shared, tracked by a barrier bulkStore shared -> global bulkReduceStore shared -> global, accumulating at the destination bulkStoreCommit close the current bulk group bulkStoreWait stores have landed bulkStoreWaitSource source tiles are reusable, stores may be in flight bulkFence order generic shared accesses against the async proxy isBulkSupported() host capability query, alongside isNvlsSupported() Notes on specific choices: - MSCCLPP_BULK_AVAILABLE gates the declarations, so unguarded use on an unsupported target is a compile error rather than a silent no-op. This follows the existing __CUDA_ARCH__ >= 900 call-site guards used for NVLS. - BulkBarrier storage is declared on every target, including host compilation, because host code must size the dynamic shared memory that holds barriers. Only the operations are gated. - expect, arrive and arriveAndExpect are separate so a multi-source gather can accumulate N loads against one barrier and wait once. Every expect for a batch must precede the arrival that completes the arrival count; arriveAndExpect carrying the batch total is documented as the recommended form. - init() includes the proxy fence that publishes the barrier to the async proxy. relaxedInit() omits it so an array of barriers can be set up under one fence, mirroring signal() and relaxedSignal(). - wait() advances the phase, so a barrier is initialized once and reused rather than reinitialized per batch. - bulkStoreWait and bulkStoreWaitSource are distinct because a double-buffered store pipeline needs to refill a tile without draining the store. Tests. test/unit/bulk_tests.cu covers the primitives on a single GPU. The reduction tests seed the destination so they distinguish accumulate from overwrite. test/mp_unit/bulk_pattern_tests.cu adds BulkPatternTest, three multi-rank kernels shaped after expert-parallel dispatch and combine: a staged push, a multi-source pull and reduce that is double buffered across chunks, and the same reduction expressed as a push using bulkReduceStore. Verified on H200 (sm_90): unit_tests 40/40, mp_unit_tests 60/60 at 2 ranks, BulkPatternTest 3/3 at 6 ranks. Guard behavior confirmed: guarded code builds at sm_80 and for multi-arch sm_80+sm_90, unguarded code fails to compile at sm_80, and an unsupported reduction type fails its static_assert. --- docs/cpp_api.rst | 21 ++ include/mscclpp/bulk_device.hpp | 264 ++++++++++++++++++ include/mscclpp/gpu_utils.hpp | 5 + src/core/gpu_utils.cc | 17 ++ test/mp_unit/CMakeLists.txt | 1 + test/mp_unit/bulk_pattern_tests.cu | 432 +++++++++++++++++++++++++++++ test/mp_unit/mp_unit_tests.hpp | 24 ++ test/unit/CMakeLists.txt | 1 + test/unit/bulk_tests.cu | 254 +++++++++++++++++ 9 files changed, 1019 insertions(+) create mode 100644 include/mscclpp/bulk_device.hpp create mode 100644 test/mp_unit/bulk_pattern_tests.cu create mode 100644 test/unit/bulk_tests.cu diff --git a/docs/cpp_api.rst b/docs/cpp_api.rst index a7ebaaf90..f9bb91c52 100644 --- a/docs/cpp_api.rst +++ b/docs/cpp_api.rst @@ -292,6 +292,27 @@ FIFO Device Interfaces Device Utilities ~~~~~~~~~~~~~~~~ +.. doxygendefine:: MSCCLPP_BULK_AVAILABLE + +.. doxygenstruct:: mscclpp::BulkBarrier + :members: + +.. doxygenenum:: mscclpp::BulkRedOp + +.. doxygenfunction:: mscclpp::bulkLoad + +.. doxygenfunction:: mscclpp::bulkStore + +.. doxygenfunction:: mscclpp::bulkReduceStore + +.. doxygenfunction:: mscclpp::bulkStoreCommit + +.. doxygenfunction:: mscclpp::bulkStoreWait + +.. doxygenfunction:: mscclpp::bulkStoreWaitSource + +.. doxygenfunction:: mscclpp::bulkFence + .. doxygenstruct:: mscclpp::DeviceSemaphore :members: diff --git a/include/mscclpp/bulk_device.hpp b/include/mscclpp/bulk_device.hpp new file mode 100644 index 000000000..6f15bba9d --- /dev/null +++ b/include/mscclpp/bulk_device.hpp @@ -0,0 +1,264 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +#ifndef MSCCLPP_BULK_DEVICE_HPP_ +#define MSCCLPP_BULK_DEVICE_HPP_ + +#include +#include + +#include "assert_device.hpp" +#include "device.hpp" +#include "poll_device.hpp" + +/// 1 if bulk asynchronous copy is available on the current device compilation target, 0 otherwise. +/// The declarations in this header exist only where this is 1, so call sites must be guarded by it. +#if defined(MSCCLPP_DEVICE_CUDA) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +#define MSCCLPP_BULK_AVAILABLE 1 +#else +#define MSCCLPP_BULK_AVAILABLE 0 +#endif + +#if MSCCLPP_BULK_AVAILABLE +#include +#endif // MSCCLPP_BULK_AVAILABLE + +namespace mscclpp { + +#if MSCCLPP_BULK_AVAILABLE + +namespace detail { // NOLINT + +template +constexpr bool bulkDependentFalse = false; + +/// Bulk copies require 16-byte aligned addresses and a size that is a multiple of 16 bytes. +MSCCLPP_DEVICE_INLINE bool bulkAligned(const void* ptr) { return (reinterpret_cast(ptr) & 15) == 0; } + +} // namespace detail + +struct BulkBarrier; + +MSCCLPP_DEVICE_INLINE void bulkFence(); +MSCCLPP_DEVICE_INLINE void bulkLoad(void* dstShared, const void* srcGlobal, uint32_t bytes, BulkBarrier& barrier); + +#endif // MSCCLPP_BULK_AVAILABLE + +/// Completion barrier for bulk loads. +/// +/// A bulk load completes asynchronously; its completion is tracked by counting transferred bytes +/// against this barrier. The barrier must live in shared memory and alternates between two phases, +/// so one barrier serves an unbounded number of batches without re-initialization. +/// +/// Per batch the caller declares the bytes to wait for with expect(), signals participation with +/// arrive(), and blocks with wait(); arriveAndExpect() fuses the common single-issuer case. The +/// phase starts at 0 and wait() flips it, so the caller keeps the phase in a register and passes it +/// back on the next batch. +/// +/// @warning Every expect() for a batch must be issued before the arrival that completes the arrival +/// count. Arriving while the outstanding byte count is momentarily zero completes the phase +/// immediately, and wait() then returns over tiles that have not been filled. Issuing a single +/// arriveAndExpect() carrying the batch total avoids this entirely and is the recommended form. +/// +/// The storage is declared on every build target, including host compilation, so that host code can +/// size the shared memory that holds barriers. The operations exist only where MSCCLPP_BULK_AVAILABLE +/// is 1. +struct BulkBarrier { +#if MSCCLPP_BULK_AVAILABLE + /// Initialize the barrier and publish it to the asynchronous proxy. Called by a single thread + /// before any other method, and not while a bulk load tracked by this barrier is in flight. + /// @param arriveCount Number of arrivals that complete a phase. + MSCCLPP_DEVICE_INLINE void init(uint32_t arriveCount = 1) { + relaxedInit(arriveCount); + bulkFence(); + } + + /// Initialize the barrier without publishing it to the asynchronous proxy. + /// + /// The initialization is not visible to bulk copies until a bulkFence() runs on the initializing + /// thread. Use this to initialize several barriers under a single fence; otherwise use init(). + /// @param arriveCount Number of arrivals that complete a phase. + MSCCLPP_DEVICE_INLINE void relaxedInit(uint32_t arriveCount = 1) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(addr()), "r"(arriveCount)); + } + + /// Invalidate the barrier, releasing the underlying hardware state. Called by a single thread + /// before the shared memory holding this barrier is reused for another purpose. + MSCCLPP_DEVICE_INLINE void invalidate() { asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::"r"(addr())); } + + /// Add to the bytes the current phase waits for, without arriving. + /// @param bytes Bytes expected to arrive. + MSCCLPP_DEVICE_INLINE void expect(uint32_t bytes) { + asm volatile("mbarrier.expect_tx.shared::cta.b64 [%0], %1;" ::"r"(addr()), "r"(bytes)); + } + + /// Signal one arrival on the current phase. + MSCCLPP_DEVICE_INLINE void arrive() { asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(addr())); } + + /// Signal one arrival and add to the bytes the current phase waits for. Equivalent to expect() + /// followed by arrive(), issued as a single instruction. + /// @param bytes Bytes expected to arrive. + MSCCLPP_DEVICE_INLINE void arriveAndExpect(uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(addr()), "r"(bytes)); + } + + /// Check whether the given phase has completed, without blocking or advancing the phase. + /// @param phase The phase to check. + /// @return True if the phase has completed. + MSCCLPP_DEVICE_INLINE bool poll(uint32_t phase) { + uint32_t done; + asm volatile("{.reg .pred p; mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2; selp.u32 %0, 1, 0, p;}" + : "=r"(done) + : "r"(addr()), "r"(phase & 1u)); + return done != 0; + } + + /// Wait for the given phase to complete and advance @p phase to the next one. + /// + /// Does not order the delivered data against generic shared memory reads; run bulkFence() on each + /// consuming thread, or a thread barrier that orders it against one that did, before reading the + /// loaded tiles. + /// + /// @param phase The phase to wait for, flipped on return. Starts at 0 and is kept by the caller. + /// @param maxSpinCount The maximum number of spins before asserting. Never assert if negative. + MSCCLPP_DEVICE_INLINE void wait(uint32_t& phase, int64_t maxSpinCount = 10000000) { + POLL_MAYBE_JAILBREAK(!poll(phase), maxSpinCount); + phase ^= 1u; + } + + private: + friend MSCCLPP_DEVICE_INLINE void bulkLoad(void* dstShared, const void* srcGlobal, uint32_t bytes, + BulkBarrier& barrier); + + MSCCLPP_DEVICE_INLINE uint32_t addr() const { return static_cast(__cvta_generic_to_shared(&mbar_)); } +#endif // MSCCLPP_BULK_AVAILABLE + + private: + alignas(8) uint64_t mbar_; +}; + +#if MSCCLPP_BULK_AVAILABLE + +/// Issue an asynchronous bulk load from global memory to shared memory. +/// +/// Returns immediately; completion is tracked by @p barrier, which must expect at least @p bytes for +/// the phase being waited on. Issued by a single thread. Loaded data is visible to generic shared +/// memory reads only after the barrier's phase completes and bulkFence() has run. +/// +/// @param dstShared Destination in shared memory. 16-byte aligned; 128 bytes is faster. +/// @param srcGlobal Source in global memory, local or peer-mapped. 16-byte aligned. +/// @param bytes Bytes to load. A multiple of 16. +/// @param barrier Barrier tracking completion of this load. +MSCCLPP_DEVICE_INLINE void bulkLoad(void* dstShared, const void* srcGlobal, uint32_t bytes, BulkBarrier& barrier) { + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(dstShared), "bulkLoad destination is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(srcGlobal), "bulkLoad source is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE((bytes & 15) == 0, "bulkLoad size is not a multiple of 16 bytes"); + asm volatile("cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + static_cast(__cvta_generic_to_shared(dstShared))), + "l"(srcGlobal), "r"(bytes), "r"(barrier.addr()) + : "memory"); +} + +/// Issue an asynchronous bulk store from shared memory to global memory. +/// +/// Returns immediately; the store joins the calling thread's open bulk group, which is closed by +/// bulkStoreCommit() and drained by bulkStoreWait() or bulkStoreWaitSource(). Issued by a single +/// thread. Data written to @p srcShared by generic shared memory writes is visible to the store only +/// after bulkFence(). +/// +/// @param dstGlobal Destination in global memory, local or peer-mapped. 16-byte aligned. +/// @param srcShared Source in shared memory. 16-byte aligned. +/// @param bytes Bytes to store. A multiple of 16. +MSCCLPP_DEVICE_INLINE void bulkStore(void* dstGlobal, const void* srcShared, uint32_t bytes) { + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(dstGlobal), "bulkStore destination is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(srcShared), "bulkStore source is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE((bytes & 15) == 0, "bulkStore size is not a multiple of 16 bytes"); + asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dstGlobal), + "r"(static_cast(__cvta_generic_to_shared(srcShared))), "r"(bytes) + : "memory"); +} + +/// Reduction operation applied by bulkReduceStore(). +enum class BulkRedOp { Add }; + +/// Issue an asynchronous bulk reduction from shared memory into global memory. +/// +/// Accumulates @p srcShared into @p dstGlobal elementwise, in the copy engine rather than on the +/// SM, so the destination is never read back across the interconnect. Works on peer-mapped +/// destinations, which makes it a remote accumulate. Otherwise behaves exactly like bulkStore(): +/// returns immediately, joins the calling thread's open bulk group, and requires a preceding +/// bulkFence() to make generic writes to @p srcShared visible. +/// +/// Measured on H200 at roughly 90% of the bulkStore() rate for the same payload, from a single +/// issuing thread, so the accumulate is close to free relative to the transfer. +/// +/// @tparam T Element type. Currently `float`, `__nv_bfloat16`, and `uint32_t`. Other types are +/// rejected at compile time rather than silently mapped, because the underlying instruction accepts +/// only certain operation and type combinations. +/// @tparam Op Reduction operation. +/// @param dstGlobal Destination in global memory, local or peer-mapped. 16-byte aligned. +/// @param srcShared Source in shared memory. 16-byte aligned. +/// @param bytes Bytes to reduce. A multiple of 16. +template +MSCCLPP_DEVICE_INLINE void bulkReduceStore(void* dstGlobal, const void* srcShared, uint32_t bytes) { + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(dstGlobal), "bulkReduceStore destination is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE(detail::bulkAligned(srcShared), "bulkReduceStore source is not 16-byte aligned"); + MSCCLPP_ASSERT_DEVICE((bytes & 15) == 0, "bulkReduceStore size is not a multiple of 16 bytes"); + const uint32_t src = static_cast(__cvta_generic_to_shared(srcShared)); + if constexpr (Op == BulkRedOp::Add && std::is_same_v) { + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [%0], [%1], %2;" ::"l"(dstGlobal), + "r"(src), "r"(bytes) + : "memory"); + } else if constexpr (Op == BulkRedOp::Add && std::is_same_v) { + // The instruction requires an explicit .noftz for bf16 addition. + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.bf16 [%0], [%1], %2;" ::"l"(dstGlobal), + "r"(src), "r"(bytes) + : "memory"); + } else if constexpr (Op == BulkRedOp::Add && std::is_same_v) { + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.u32 [%0], [%1], %2;" ::"l"(dstGlobal), + "r"(src), "r"(bytes) + : "memory"); + } else { + static_assert(detail::bulkDependentFalse, "Unsupported bulk reduction type or operation"); + } +} + +/// Close the calling thread's open bulk group, committing every bulkStore() and bulkReduceStore() +/// issued since the previous bulkStoreCommit(). +MSCCLPP_DEVICE_INLINE void bulkStoreCommit() { asm volatile("cp.async.bulk.commit_group;" ::: "memory"); } + +/// Wait until at most @p PendingGroups committed bulk groups have yet to complete. +/// +/// Completion means the data has reached the destination. Wait for this before signaling a peer that +/// the data is available. To reuse the source shared memory only, bulkStoreWaitSource() is cheaper. +/// @tparam PendingGroups Number of most recently committed groups allowed to remain outstanding. +template +MSCCLPP_DEVICE_INLINE void bulkStoreWait() { + asm volatile("cp.async.bulk.wait_group %0;" ::"n"(PendingGroups) : "memory"); +} + +/// Wait until at most @p PendingGroups committed bulk groups may still read their source. +/// +/// Weaker than bulkStoreWait(): it guarantees only that the source shared memory can be overwritten, +/// not that the data has reached the destination. This is what a double-buffered store pipeline +/// needs between refills, and it lets stores stay in flight while the next tile is being staged. +/// @tparam PendingGroups Number of most recently committed groups allowed to still read their source. +template +MSCCLPP_DEVICE_INLINE void bulkStoreWaitSource() { + asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(PendingGroups) : "memory"); +} + +/// Order the calling thread's generic shared memory accesses against bulk copies. +/// +/// Bulk copies reach shared memory through a separate proxy, so they are not ordered against generic +/// accesses by default. Run this after a bulk load completes and before reading its destination, +/// after writing a source and before issuing a bulk store from it, and after initializing a barrier +/// with relaxedInit(). Every thread that touches the shared memory must run it, unless a thread +/// barrier such as __syncthreads() already orders that thread against one that did. +MSCCLPP_DEVICE_INLINE void bulkFence() { asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); } + +#endif // MSCCLPP_BULK_AVAILABLE + +} // namespace mscclpp + +#endif // MSCCLPP_BULK_DEVICE_HPP_ diff --git a/include/mscclpp/gpu_utils.hpp b/include/mscclpp/gpu_utils.hpp index 82fa3ec0a..b705f8830 100644 --- a/include/mscclpp/gpu_utils.hpp +++ b/include/mscclpp/gpu_utils.hpp @@ -312,6 +312,11 @@ inline void gpuMemset(void* ptr, int value, size_t bytes) { detail::gpuMemset(pt /// @return True if NVLink SHARP (NVLS) is supported, false otherwise. bool isNvlsSupported(); +/// Check if bulk asynchronous copy is supported by the current device. +/// +/// @return True if bulk asynchronous copy is supported, false otherwise. +bool isBulkSupported(); + /// Check if ptr is allocaed by cuMemMap. /// @param ptr The pointer to check. /// @return True if the pointer is allocated by cuMemMap, false otherwise. diff --git a/src/core/gpu_utils.cc b/src/core/gpu_utils.cc index 1ce61322c..b6fbaddfa 100644 --- a/src/core/gpu_utils.cc +++ b/src/core/gpu_utils.cc @@ -299,6 +299,23 @@ bool isNvlsSupported() { return false; } +bool isBulkSupported() { + [[maybe_unused]] static bool result = false; + [[maybe_unused]] static bool isChecked = false; +#if !defined(MSCCLPP_USE_ROCM) + if (!isChecked) { + int deviceId; + int major; + MSCCLPP_CUDATHROW(cudaGetDevice(&deviceId)); + MSCCLPP_CUDATHROW(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, deviceId)); + result = (major >= 9); + isChecked = true; + } + return result; +#endif + return false; +} + bool isCuMemMapAllocated([[maybe_unused]] void* ptr) { #if defined(MSCCLPP_USE_ROCM) return false; diff --git a/test/mp_unit/CMakeLists.txt b/test/mp_unit/CMakeLists.txt index d4004e8e6..2f01c3870 100644 --- a/test/mp_unit/CMakeLists.txt +++ b/test/mp_unit/CMakeLists.txt @@ -8,6 +8,7 @@ target_sources(mp_unit_tests PRIVATE communicator_tests.cu port_channel_tests.cu memory_channel_tests.cu + bulk_pattern_tests.cu semaphore_perf_tests.cu switch_channel_tests.cu executor_tests.cc diff --git a/test/mp_unit/bulk_pattern_tests.cu b/test/mp_unit/bulk_pattern_tests.cu new file mode 100644 index 000000000..de5295f25 --- /dev/null +++ b/test/mp_unit/bulk_pattern_tests.cu @@ -0,0 +1,432 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// Bulk-copy usage patterns, shaped after the expert-parallel dispatch and combine kernels. +// +// These are deliberately simplified: fixed token counts, one contribution per rank pair, and plain +// arithmetic instead of quantization or top-k routing. What they preserve is the communication +// shape, which is what the API has to support: +// +// Dispatch local buffer -> shared (transform) -> peer buffer. Push, staged through shared +// memory, pipelined so a tile is refilled while its store is still in flight. +// Combine every peer's buffer -> shared -> reduce -> local output. Pull, many sources into +// one barrier, double-buffered across chunks. +// CombineReduce the same reduction expressed as a push, with the copy engine accumulating +// directly into the peer's memory, so nothing is ever read back. +// +// Note where the channels are and are not used. Bulk data moves through raw peer pointers, exactly +// as the expert-parallel kernels do; the channels carry only signal/wait. That split is the reason +// the bulk primitives are pointer-based rather than channel methods. + +#include +#include +#include +#include +#include + +#include "mp_unit_tests.hpp" + +void BulkPatternTest::SetUp() { + if (gEnv->nRanksPerNode < 2) { + SKIP_TEST(); + } + setNumRanksToUse(gEnv->nRanksPerNode); + CommunicatorTestBase::SetUp(); + rank = gEnv->rank; + worldSize = numRanksToUse; +} + +void BulkPatternTest::TearDown() { + syncHandles.clear(); + syncChannels.clear(); + peerPools.clear(); + remotePoolMemories.clear(); + CommunicatorTestBase::TearDown(); +} + +void BulkPatternTest::setupPeerPools(void* pool, size_t poolBytes) { + const mscclpp::TransportFlags transport = mscclpp::Transport::CudaIpc; + mscclpp::RegisteredMemory localMem = communicator->registerMemory(pool, poolBytes, transport); + + std::vector> connFutures(worldSize); + std::vector> memFutures(worldSize); + for (int r = 0; r < worldSize; ++r) { + if (r == rank) continue; + connFutures[r] = communicator->connect(mscclpp::Transport::CudaIpc, r); + communicator->sendMemory(localMem, r); + memFutures[r] = communicator->recvMemory(r); + } + + peerPools.assign(worldSize, nullptr); + peerPools[rank] = pool; + for (int r = 0; r < worldSize; ++r) { + if (r == rank) continue; + mscclpp::RegisteredMemory remote = memFutures[r].get(); + peerPools[r] = remote.data(); + remotePoolMemories.push_back(remote); + syncChannels.emplace_back(communicator->buildSemaphore(connFutures[r].get(), r).get()); + } + for (const auto& chan : syncChannels) syncHandles.push_back(mscclpp::deviceHandle(chan)); + registeredMemories.push_back(localMem); +} + +#if defined(MSCCLPP_DEVICE_CUDA) + +namespace { + +constexpr int kNumTokens = 16; +constexpr int kHidden = 1024; // floats per token +constexpr uint32_t kTokenBytes = kHidden * sizeof(float); // 4 KB +constexpr uint32_t kChunkBytes = 1024; // staging granularity +constexpr int kChunksPerToken = kTokenBytes / kChunkBytes; +constexpr int kChunkFloats = kChunkBytes / sizeof(float); +constexpr int kStages = 2; + +// Distinct, exactly representable value for (source rank, destination rank, token, element). +MSCCLPP_HOST_DEVICE_INLINE float payload(int src, int dst, int token, int elem) { + return (float)(src * 1000 + dst * 100 + token * 10 + (elem % 10)); +} + +// Offset in floats of the [src][token] row within a receive pool. +MSCCLPP_HOST_DEVICE_INLINE int64_t rowOffset(int src, int token) { + return ((int64_t)src * kNumTokens + token) * kHidden; +} + +} // namespace + +// Signal every peer, then wait on every peer. +// +// Launched as its own kernel rather than folded into the pattern kernels: the ranks have to agree on +// a point that every block has passed, and kernel completion already provides exactly that. Doing it +// inside a kernel would need a cooperative grid launch. +__global__ void kernelPeerBarrier(mscclpp::BaseMemoryChannelDeviceHandle* chans, int nPeers) { + if ((int)threadIdx.x < nPeers) { + chans[threadIdx.x].signal(); + chans[threadIdx.x].wait(); + } +} + +namespace { + +// Peer pool pointers and synchronization handles, staged in device memory for the kernels. +struct DeviceState { + std::shared_ptr> chans; + std::shared_ptr pools; + int nPeers; +}; + +DeviceState uploadPeerState(const std::vector>& handles, + const std::vector& pools) { + DeviceState state; + state.nPeers = (int)handles.size(); + state.chans = mscclpp::detail::gpuCallocShared>(handles.size()); + MSCCLPP_CUDATHROW(cudaMemcpy(state.chans.get(), handles.data(), + handles.size() * sizeof(DeviceHandle), + cudaMemcpyHostToDevice)); + state.pools = mscclpp::detail::gpuCallocShared(pools.size()); + MSCCLPP_CUDATHROW(cudaMemcpy(state.pools.get(), pools.data(), pools.size() * sizeof(void*), cudaMemcpyHostToDevice)); + return state; +} + +} // namespace + +// ------------------------------------------------------------------------------------------------ +// Dispatch: push each local token to every peer, staged and transformed in shared memory. +// +// Per chunk the leader loads into a tile, the block transforms it in place, and the leader pushes it +// out. bulkStoreWaitSource() lets the next load start while the store is still travelling, which is +// the reason that entry point exists separately from bulkStoreWait(). +// ------------------------------------------------------------------------------------------------ +__global__ void kernelDispatch(const float* localTokens, void** peerPools, int rank, int worldSize) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) uint8_t tile[kChunkBytes]; + __shared__ mscclpp::BulkBarrier barrier; + + if (threadIdx.x == 0) barrier.init(); + __syncthreads(); + + uint32_t phase = 0; // initialized once; wait() advances it on every chunk + bool storePending = false; + + // One block per (destination, token) pair, strided over the grid. + const int totalRows = worldSize * kNumTokens; + for (int row = blockIdx.x; row < totalRows; row += gridDim.x) { + const int dst = row / kNumTokens; + const int token = row % kNumTokens; + const float* srcRow = localTokens + (int64_t)token * kHidden; + auto* dstRow = reinterpret_cast(peerPools[dst]) + rowOffset(rank, token) * sizeof(float); + + for (int c = 0; c < kChunksPerToken; ++c) { + const uint32_t off = c * kChunkBytes; + if (threadIdx.x == 0) { + // The previous chunk's store has released the tile by now; make sure before refilling. + if (storePending) mscclpp::bulkStoreWaitSource<0>(); + barrier.arriveAndExpect(kChunkBytes); + mscclpp::bulkLoad(tile, reinterpret_cast(srcRow) + off, kChunkBytes, barrier); + barrier.wait(phase); + } + __syncthreads(); + mscclpp::bulkFence(); // loaded data -> generic reads + + // Stand-in for the quantize/scale step a real dispatch performs while the token is staged. + float* staged = reinterpret_cast(tile); + for (int i = threadIdx.x; i < kChunkFloats; i += blockDim.x) staged[i] += (float)(dst * 100); + + __syncthreads(); + mscclpp::bulkFence(); // generic writes -> the store's read of the tile + if (threadIdx.x == 0) { + mscclpp::bulkStore(dstRow + off, tile, kChunkBytes); + mscclpp::bulkStoreCommit(); + storePending = true; + } + __syncthreads(); + } + } + + // Every push must have landed before any peer is told the data is ready. + if (threadIdx.x == 0) { + mscclpp::bulkStoreWait<0>(); + barrier.invalidate(); + } +#else + (void)localTokens; + (void)peerPools; + (void)rank; + (void)worldSize; +#endif +} + +TEST(BulkPatternTest, Dispatch) { + if (gEnv->rank >= numRanksToUse) return; + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + + const size_t poolFloats = (size_t)worldSize * kNumTokens * kHidden; + std::shared_ptr pool = mscclpp::GpuBuffer(poolFloats).memory(); + std::shared_ptr tokens = mscclpp::GpuBuffer((size_t)kNumTokens * kHidden).memory(); + MSCCLPP_CUDATHROW(cudaMemset(pool.get(), 0, poolFloats * sizeof(float))); + + // Local tokens carry payload(rank, 0, token, elem); the kernel adds dst * 100 while staged, so the + // value that lands on rank dst is payload(rank, dst, token, elem). + std::vector host((size_t)kNumTokens * kHidden); + for (int t = 0; t < kNumTokens; ++t) + for (int h = 0; h < kHidden; ++h) host[(size_t)t * kHidden + h] = payload(gEnv->rank, 0, t, h); + MSCCLPP_CUDATHROW(cudaMemcpy(tokens.get(), host.data(), host.size() * sizeof(float), cudaMemcpyHostToDevice)); + + setupPeerPools(pool.get(), poolFloats * sizeof(float)); + + DeviceState dev = uploadPeerState(syncHandles, peerPools); + + kernelDispatch<<>>(tokens.get(), dev.pools.get(), rank, worldSize); + // Kernel completion is the grid-wide point after which every push has landed; only then is it safe + // to tell the peers, and to read what they pushed here. + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector out(poolFloats); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), pool.get(), poolFloats * sizeof(float), cudaMemcpyDeviceToHost)); + for (int src = 0; src < worldSize; ++src) { + for (int t = 0; t < kNumTokens; ++t) { + for (int h = 0; h < kHidden; h += 97) { // sparse check keeps the assertion count sane + EXPECT_EQ(out[rowOffset(src, t) + h], payload(src, rank, t, h)); + } + } + } +} + +// ------------------------------------------------------------------------------------------------ +// Combine: pull this rank's row from every peer into shared memory and reduce it. +// +// This is the pattern the barrier design exists for. All contributors of a chunk are issued against +// one barrier and waited on once, and two stages let the next chunk's loads be issued before the +// current one is reduced. The two barriers are set up with relaxedInit() under a single bulkFence() +// rather than paying a fence each. +// ------------------------------------------------------------------------------------------------ +__global__ void kernelCombine(float* output, void** peerPools, int rank, int worldSize) { +#if MSCCLPP_BULK_AVAILABLE + extern __shared__ __align__(128) uint8_t shared[]; + // Layout: [stage][contributor][chunk] tiles, then the per-stage barriers. + uint8_t* tiles = shared; + auto* barriers = reinterpret_cast(shared + (size_t)kStages * worldSize * kChunkBytes); + auto tile = [&](int stage, int contributor) { + return tiles + ((size_t)stage * worldSize + contributor) * kChunkBytes; + }; + + if (threadIdx.x == 0) { + for (int s = 0; s < kStages; ++s) barriers[s].relaxedInit(); + mscclpp::bulkFence(); // one fence publishes both barriers to the async proxy + } + __syncthreads(); + + uint32_t phase[kStages] = {0, 0}; + + auto issue = [&](int stage, int token, int chunk) { + barriers[stage].arriveAndExpect(kChunkBytes * worldSize); // whole batch, before any arrival lands + for (int src = 0; src < worldSize; ++src) { + const auto* row = reinterpret_cast(peerPools[src]) + rowOffset(rank, token) * sizeof(float); + mscclpp::bulkLoad(tile(stage, src), row + (size_t)chunk * kChunkBytes, kChunkBytes, barriers[stage]); + } + }; + + for (int token = blockIdx.x; token < kNumTokens; token += gridDim.x) { + if (threadIdx.x == 0) issue(0, token, 0); + for (int c = 0; c < kChunksPerToken; ++c) { + const int stage = c % kStages; + if (threadIdx.x == 0) { + if (c + 1 < kChunksPerToken) issue((c + 1) % kStages, token, c + 1); // prefetch next chunk + barriers[stage].wait(phase[stage]); + } + __syncthreads(); + mscclpp::bulkFence(); + + float* out = output + (int64_t)token * kHidden + (int64_t)c * kChunkFloats; + for (int i = threadIdx.x; i < kChunkFloats; i += blockDim.x) { + float acc = 0.0f; + for (int src = 0; src < worldSize; ++src) acc += reinterpret_cast(tile(stage, src))[i]; + out[i] = acc; + } + __syncthreads(); // this stage's tiles are consumed before they are reissued + } + } + + if (threadIdx.x == 0) { + for (int s = 0; s < kStages; ++s) barriers[s].invalidate(); + } +#else + (void)output; + (void)peerPools; + (void)rank; + (void)worldSize; +#endif +} + +TEST(BulkPatternTest, Combine) { + if (gEnv->rank >= numRanksToUse) return; + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + + const size_t poolFloats = (size_t)worldSize * kNumTokens * kHidden; + std::shared_ptr pool = mscclpp::GpuBuffer(poolFloats).memory(); + std::shared_ptr output = mscclpp::GpuBuffer((size_t)kNumTokens * kHidden).memory(); + + // Each rank publishes, for every destination, the contribution that destination will pull. + std::vector host(poolFloats); + for (int dst = 0; dst < worldSize; ++dst) + for (int t = 0; t < kNumTokens; ++t) + for (int h = 0; h < kHidden; ++h) host[rowOffset(dst, t) + h] = payload(gEnv->rank, dst, t, h); + MSCCLPP_CUDATHROW(cudaMemcpy(pool.get(), host.data(), host.size() * sizeof(float), cudaMemcpyHostToDevice)); + MSCCLPP_CUDATHROW(cudaMemset(output.get(), 0, (size_t)kNumTokens * kHidden * sizeof(float))); + + setupPeerPools(pool.get(), poolFloats * sizeof(float)); + + DeviceState dev = uploadPeerState(syncHandles, peerPools); + + const size_t sharedBytes = (size_t)kStages * worldSize * kChunkBytes + kStages * sizeof(mscclpp::BulkBarrier); + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); // every pool is filled before any pull + kernelCombine<<>>(output.get(), dev.pools.get(), rank, worldSize); + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); // no rank tears down while peers read + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector out((size_t)kNumTokens * kHidden); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), output.get(), out.size() * sizeof(float), cudaMemcpyDeviceToHost)); + for (int t = 0; t < kNumTokens; ++t) { + for (int h = 0; h < kHidden; h += 97) { + float expected = 0.0f; + for (int src = 0; src < worldSize; ++src) expected += payload(src, rank, t, h); + EXPECT_EQ(out[(size_t)t * kHidden + h], expected); + } + } +} + +// ------------------------------------------------------------------------------------------------ +// CombineReduce: the same reduction, expressed as a push. +// +// Instead of every rank pulling and summing, each rank accumulates its own contribution straight +// into the destination's accumulator with bulkReduceStore(). The copy engine performs the add at the +// destination, so the accumulator is never read back across the link and no rank stages anyone +// else's data. The reverse link direction stays idle. +// ------------------------------------------------------------------------------------------------ +__global__ void kernelCombineReduce(const float* contribution, void** peerPools, int rank, int worldSize) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) uint8_t tile[kChunkBytes]; + __shared__ mscclpp::BulkBarrier barrier; + + if (threadIdx.x == 0) barrier.init(); + __syncthreads(); + uint32_t phase = 0; + + const int totalRows = worldSize * kNumTokens; + for (int row = blockIdx.x; row < totalRows; row += gridDim.x) { + const int dst = row / kNumTokens; + const int token = row % kNumTokens; + const auto* srcRow = reinterpret_cast(contribution + rowOffset(dst, token)); + // Every rank accumulates into the same [rank-independent] row of the destination. + auto* dstRow = reinterpret_cast(peerPools[dst]) + rowOffset(0, token) * sizeof(float); + + for (int c = 0; c < kChunksPerToken; ++c) { + const uint32_t off = c * kChunkBytes; + if (threadIdx.x == 0) { + barrier.arriveAndExpect(kChunkBytes); + mscclpp::bulkLoad(tile, srcRow + off, kChunkBytes, barrier); + barrier.wait(phase); + mscclpp::bulkFence(); // staged data -> the reduction's read of the tile + mscclpp::bulkReduceStore(dstRow + off, tile, kChunkBytes); + mscclpp::bulkStoreCommit(); + mscclpp::bulkStoreWait<0>(); // the accumulate must land before the tile is refilled + } + __syncthreads(); + } + } + + if (threadIdx.x == 0) barrier.invalidate(); +#else + (void)contribution; + (void)peerPools; + (void)rank; + (void)worldSize; +#endif +} + +TEST(BulkPatternTest, CombineReduce) { + if (gEnv->rank >= numRanksToUse) return; + if (!mscclpp::isBulkSupported()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + + const size_t poolFloats = (size_t)worldSize * kNumTokens * kHidden; + std::shared_ptr pool = mscclpp::GpuBuffer(poolFloats).memory(); + std::shared_ptr contribution = mscclpp::GpuBuffer(poolFloats).memory(); + MSCCLPP_CUDATHROW(cudaMemset(pool.get(), 0, poolFloats * sizeof(float))); // accumulator starts at zero + + std::vector host(poolFloats); + for (int dst = 0; dst < worldSize; ++dst) + for (int t = 0; t < kNumTokens; ++t) + for (int h = 0; h < kHidden; ++h) host[rowOffset(dst, t) + h] = payload(gEnv->rank, dst, t, h); + MSCCLPP_CUDATHROW(cudaMemcpy(contribution.get(), host.data(), host.size() * sizeof(float), cudaMemcpyHostToDevice)); + + setupPeerPools(pool.get(), poolFloats * sizeof(float)); + + DeviceState dev = uploadPeerState(syncHandles, peerPools); + + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); // every accumulator is zeroed first + kernelCombineReduce<<>>(contribution.get(), dev.pools.get(), rank, worldSize); + kernelPeerBarrier<<<1, 32>>>(dev.chans.get(), dev.nPeers); // every contribution has landed + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector out(poolFloats); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), pool.get(), poolFloats * sizeof(float), cudaMemcpyDeviceToHost)); + for (int t = 0; t < kNumTokens; ++t) { + for (int h = 0; h < kHidden; h += 97) { + float expected = 0.0f; + for (int src = 0; src < worldSize; ++src) expected += payload(src, rank, t, h); + EXPECT_EQ(out[rowOffset(0, t) + h], expected); + } + } +} + +#endif // defined(MSCCLPP_DEVICE_CUDA) diff --git a/test/mp_unit/mp_unit_tests.hpp b/test/mp_unit/mp_unit_tests.hpp index d01857f3f..a75eecac2 100644 --- a/test/mp_unit/mp_unit_tests.hpp +++ b/test/mp_unit/mp_unit_tests.hpp @@ -187,6 +187,30 @@ class MemoryChannelOneToOneTest : public CommunicatorTestBase { std::unordered_map> memorySemaphores; }; +/// Fixture for the bulk-copy usage patterns, shaped after the expert-parallel dispatch and combine +/// kernels: a full mesh where every rank holds a receive pool addressable by every peer, channels +/// used only for synchronization, and the peer pool pointers handed to the kernel as a plain array. +class BulkPatternTest : public CommunicatorTestBase { + protected: + void SetUp() override; + void TearDown() override; + + /// Register @p pool on every rank, exchange it around a full mesh, and build one synchronization + /// channel per peer. Fills peerPools and syncHandles. + void setupPeerPools(void* pool, size_t poolBytes); + + int worldSize = 0; + int rank = 0; + /// Receive pool base address of each rank, indexed by rank. The local entry is the local pool. + /// This mirrors the expert-parallel kernels, which address peers by raw pointer and never route + /// bulk data through a channel. + std::vector peerPools; + /// One channel per peer, used only for signal/wait. + std::vector syncChannels; + std::vector> syncHandles; + std::vector remotePoolMemories; +}; + class SemaphorePerfTest : public CommunicatorTestBase { protected: void SetUp() override; diff --git a/test/unit/CMakeLists.txt b/test/unit/CMakeLists.txt index 2547278b4..e6d09d2d7 100644 --- a/test/unit/CMakeLists.txt +++ b/test/unit/CMakeLists.txt @@ -17,4 +17,5 @@ target_sources(unit_tests PRIVATE gpu_data_types_tests.cu local_channel_tests.cu gpu_ipc_mem_tests.cc + bulk_tests.cu ) diff --git a/test/unit/bulk_tests.cu b/test/unit/bulk_tests.cu new file mode 100644 index 000000000..e328ccec2 --- /dev/null +++ b/test/unit/bulk_tests.cu @@ -0,0 +1,254 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#include +#include +#include +#include + +#include "../framework.hpp" + +#if defined(MSCCLPP_DEVICE_CUDA) + +constexpr uint32_t kTile = 4096; +constexpr uint32_t kElems = kTile / sizeof(int); + +// Load one tile into shared memory, then copy it out. +__global__ void kernelBulkLoad(const int* src, int* dst) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) int tile[kElems]; + __shared__ mscclpp::BulkBarrier barrier; + + uint32_t phase = 0; + if (threadIdx.x == 0) { + barrier.init(); + barrier.arriveAndExpect(kTile); + mscclpp::bulkLoad(tile, src, kTile, barrier); + barrier.wait(phase); + } + __syncthreads(); + mscclpp::bulkFence(); + + for (uint32_t i = threadIdx.x; i < kElems; i += blockDim.x) dst[i] = tile[i]; + + __syncthreads(); + if (threadIdx.x == 0) barrier.invalidate(); +#else + (void)src; + (void)dst; +#endif +} + +// Gather NumSrc tiles into one barrier, then reduce them. +template +__global__ void kernelBulkGather(const int* src, int* dst) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) int tiles[NumSrc][kElems]; + __shared__ mscclpp::BulkBarrier barrier; + + uint32_t phase = 0; + if (threadIdx.x == 0) { + barrier.init(); + barrier.arriveAndExpect(kTile * NumSrc); // whole batch total, before any arrival completes + for (int s = 0; s < NumSrc; ++s) mscclpp::bulkLoad(tiles[s], src + s * kElems, kTile, barrier); + barrier.wait(phase); + } + __syncthreads(); + mscclpp::bulkFence(); + + for (uint32_t i = threadIdx.x; i < kElems; i += blockDim.x) { + int acc = 0; + for (int s = 0; s < NumSrc; ++s) acc += tiles[s][i]; + dst[i] = acc; + } + + __syncthreads(); + if (threadIdx.x == 0) barrier.invalidate(); +#else + (void)src; + (void)dst; +#endif +} + +// Reuse one barrier across NumChunks phases, staging each chunk in and storing it back out. +template +__global__ void kernelBulkRoundTrip(const int* src, int* dst) { +#if MSCCLPP_BULK_AVAILABLE + __shared__ alignas(128) int tile[kElems]; + __shared__ mscclpp::BulkBarrier barrier; + + if (threadIdx.x == 0) barrier.init(); + __syncthreads(); + + uint32_t phase = 0; // one barrier, no re-initialization; wait() advances the phase + for (int c = 0; c < NumChunks; ++c) { + if (threadIdx.x == 0) { + barrier.arriveAndExpect(kTile); + mscclpp::bulkLoad(tile, src + c * kElems, kTile, barrier); + barrier.wait(phase); + } + __syncthreads(); + mscclpp::bulkFence(); + + for (uint32_t i = threadIdx.x; i < kElems; i += blockDim.x) tile[i] += 1; + + __syncthreads(); + mscclpp::bulkFence(); + if (threadIdx.x == 0) { + mscclpp::bulkStore(dst + c * kElems, tile, kTile); + mscclpp::bulkStoreCommit(); + mscclpp::bulkStoreWaitSource(); // tile refillable; the store may still be in flight + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + mscclpp::bulkStoreWait(); // every store has landed + barrier.invalidate(); + } +#else + (void)src; + (void)dst; +#endif +} + +// Accumulate a staged tile into a seeded destination with the copy engine. +template +__global__ void kernelBulkReduce(T* dst, T addend, uint32_t count) { +#if MSCCLPP_BULK_AVAILABLE + extern __shared__ __align__(128) uint8_t raw[]; + T* tile = reinterpret_cast(raw); + for (uint32_t i = threadIdx.x; i < count; i += blockDim.x) tile[i] = addend; + __syncthreads(); + mscclpp::bulkFence(); + if (threadIdx.x == 0) { + mscclpp::bulkReduceStore(dst, tile, count * sizeof(T)); + mscclpp::bulkStoreCommit(); + mscclpp::bulkStoreWait(); + } +#else + (void)dst; + (void)addend; + (void)count; +#endif +} + +static bool bulkAvailable() { + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) return false; + int major = 0; + if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) return false; + return major >= 9; +} + +class BulkTestData { + public: + BulkTestData(int numElems) + : src_(mscclpp::GpuBuffer(numElems).memory()), + dst_(mscclpp::GpuBuffer(numElems).memory()), + numElems_(numElems) { + std::vector host(numElems); + for (int i = 0; i < numElems; ++i) host[i] = i + 1; + MSCCLPP_CUDATHROW(cudaMemcpy(src_.get(), host.data(), numElems * sizeof(int), cudaMemcpyHostToDevice)); + MSCCLPP_CUDATHROW(cudaMemset(dst_.get(), 0, numElems * sizeof(int))); + } + + std::vector result() { + std::vector out(numElems_); + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), dst_.get(), out.size() * sizeof(int), cudaMemcpyDeviceToHost)); + return out; + } + + int* src() { return src_.get(); } + int* dst() { return dst_.get(); } + + private: + std::shared_ptr src_; + std::shared_ptr dst_; + int numElems_; +}; + +TEST(BulkTest, Load) { + if (!bulkAvailable()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + BulkTestData f(kElems); + kernelBulkLoad<<<1, 256>>>(f.src(), f.dst()); + std::vector out = f.result(); + for (uint32_t i = 0; i < kElems; ++i) EXPECT_EQ(out[i], (int)i + 1); +} + +TEST(BulkTest, Gather) { + if (!bulkAvailable()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + constexpr int kNumSrc = 4; + BulkTestData f(kElems * kNumSrc); + kernelBulkGather<<<1, 256>>>(f.src(), f.dst()); + std::vector out = f.result(); + for (uint32_t i = 0; i < kElems; ++i) { + int expected = 0; + for (int s = 0; s < kNumSrc; ++s) expected += s * kElems + i + 1; + EXPECT_EQ(out[i], expected); + } +} + +TEST(BulkTest, RoundTrip) { + if (!bulkAvailable()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + constexpr int kNumChunks = 8; + BulkTestData f(kElems * kNumChunks); + kernelBulkRoundTrip<<<1, 256>>>(f.src(), f.dst()); + std::vector out = f.result(); + for (uint32_t i = 0; i < kElems * kNumChunks; ++i) EXPECT_EQ(out[i], (int)i + 2); +} + +// Each reduction type accumulates into a destination seeded with a known base, so the test +// distinguishes an accumulate from an overwrite. +template +static void reduceStoreTest(float base, float addend) { + constexpr uint32_t kCount = 1024; + std::shared_ptr dst = mscclpp::GpuBuffer(kCount).memory(); + std::vector host(kCount, static_cast(base)); + MSCCLPP_CUDATHROW(cudaMemcpy(dst.get(), host.data(), kCount * sizeof(T), cudaMemcpyHostToDevice)); + + kernelBulkReduce<<<1, 256, kCount * sizeof(T)>>>(dst.get(), static_cast(addend), kCount); + MSCCLPP_CUDATHROW(cudaDeviceSynchronize()); + + std::vector out(kCount); + MSCCLPP_CUDATHROW(cudaMemcpy(out.data(), dst.get(), kCount * sizeof(T), cudaMemcpyDeviceToHost)); + for (uint32_t i = 0; i < kCount; ++i) { + EXPECT_EQ(static_cast(out[i]), base + addend); + } +} + +TEST(BulkTest, ReduceStoreFloat) { + if (!bulkAvailable()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + reduceStoreTest(10.0f, 1.5f); +} + +TEST(BulkTest, ReduceStoreBf16) { + if (!bulkAvailable()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + reduceStoreTest<__nv_bfloat16>(10.0f, 2.0f); +} + +TEST(BulkTest, ReduceStoreUint32) { + if (!bulkAvailable()) { + SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; + return; + } + reduceStoreTest(10.0f, 3.0f); +} + +#endif // defined(MSCCLPP_DEVICE_CUDA) From 2785141dc85f88d058da15fe5ab52824a4158a63 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Tue, 4 Aug 2026 17:26:22 +0000 Subject: [PATCH 2/6] Make bulk reduction element types nameable from host code was included only where MSCCLPP_BULK_AVAILABLE is 1, so __nv_bfloat16 could not be named in a host translation unit and bulkReduceStore<__nv_bfloat16> could not be instantiated from one. The unit test only compiled because gpu_utils.hpp happened to pull the type in. Gate the include on MSCCLPP_DEVICE_CUDA instead, matching switch_channel_device.hpp. Plain host builds without the CUDA toolkit are unaffected. Found while validating on GB200 (sm_100), where the header is used without that incidental include. Verified on GB200 (sm_100, CUDA 13.0, aarch64): 11/11 standalone checks including peer-memory load, store and reduce, and 4-GPU concurrent accumulate into one buffer over 50 iterations. Guard behavior holds under CUDA 13: guarded code builds at sm_80 and multi-arch sm_80+sm_100a, unguarded code fails to compile at sm_80, plain g++ reports sizeof(BulkBarrier)=8. Re-verified on H200 (sm_90): unit_tests 40/40, mp_unit_tests 60/60 at 2 ranks. --- include/mscclpp/bulk_device.hpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/include/mscclpp/bulk_device.hpp b/include/mscclpp/bulk_device.hpp index 6f15bba9d..aeea5664e 100644 --- a/include/mscclpp/bulk_device.hpp +++ b/include/mscclpp/bulk_device.hpp @@ -19,9 +19,11 @@ #define MSCCLPP_BULK_AVAILABLE 0 #endif -#if MSCCLPP_BULK_AVAILABLE +// Included whenever the CUDA toolkit is present, not only where bulk copy is available, so that the +// element types accepted by bulkReduceStore() can be named from host code too. +#if defined(MSCCLPP_DEVICE_CUDA) #include -#endif // MSCCLPP_BULK_AVAILABLE +#endif // defined(MSCCLPP_DEVICE_CUDA) namespace mscclpp { From 52f3ff88aadd170b7b9a7b3ba3b8eea69a9e8618 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Fri, 7 Aug 2026 06:03:39 +0000 Subject: [PATCH 3/6] Fix docs generation, add Python capability query, drop duplicate check Three cleanups found by running things that had not been run. Doxygen strips everything guarded by MSCCLPP_BULK_AVAILABLE, because the macro derives from __CUDA_ARCH__ and doxygen does not define it. Every directive added to cpp_api.rst therefore failed: WARNING: doxygenfunction: Cannot find function "mscclpp::bulkLoad" in doxygen xml output for project "mscclpp" from directory: ./doxygen/xml Add MSCCLPP_BULK_AVAILABLE=1 to PREDEFINED in the Doxyfile, next to the existing MSCCLPP_DEVICE_COMPILE and MSCCLPP_DEVICE_CUDA entries that exist for the same reason. Sphinx now emits no bulk warnings and every symbol renders. Bind isBulkSupported() as is_bulk_supported and export it, matching is_nvls_supported. Replace the unit test's homegrown compute-capability check with isBulkSupported(). The mp_unit tests already used it; having two ways to ask the same question is what the host query exists to avoid. Also document that the cross-device atomicity of bulkReduceStore() is established by measurement rather than by the PTX documentation, so callers do not take it as guaranteed. Verified on H200: docs build clean of bulk warnings and all ten symbols present in the generated HTML; unit_tests 40/40; mp_unit_tests 60/60 at 2 ranks; Python bindings build and mscclpp.is_bulk_supported() returns True. --- docs/Doxyfile | 1 + include/mscclpp/bulk_device.hpp | 6 ++++++ python/csrc/gpu_utils_py.cpp | 1 + python/mscclpp/__init__.py | 2 ++ test/unit/bulk_tests.cu | 20 ++++++-------------- 5 files changed, 16 insertions(+), 14 deletions(-) diff --git a/docs/Doxyfile b/docs/Doxyfile index d17aa6eb1..1709aa5c6 100644 --- a/docs/Doxyfile +++ b/docs/Doxyfile @@ -2188,6 +2188,7 @@ INCLUDE_FILE_PATTERNS = # This tag requires that the tag ENABLE_PREPROCESSING is set to YES. PREDEFINED = __CUDACC__ \ + MSCCLPP_BULK_AVAILABLE=1 \ MSCCLPP_DEVICE_COMPILE \ MSCCLPP_DEVICE_CUDA \ MSCCLPP_DEVICE_HIP \ diff --git a/include/mscclpp/bulk_device.hpp b/include/mscclpp/bulk_device.hpp index aeea5664e..07129bedc 100644 --- a/include/mscclpp/bulk_device.hpp +++ b/include/mscclpp/bulk_device.hpp @@ -194,6 +194,12 @@ enum class BulkRedOp { Add }; /// Measured on H200 at roughly 90% of the bulkStore() rate for the same payload, from a single /// issuing thread, so the accumulate is close to free relative to the transfer. /// +/// @warning When several devices accumulate into the same destination address concurrently, the +/// result depends on the per-element reduction being atomic across peers. That holds in every +/// measurement taken here (exact results on H200 across 8 six-rank runs and on GB200 across 50 +/// four-GPU runs), but it is an empirical result, not a guarantee found in the PTX documentation. +/// Confirm it on the target hardware before relying on it for a collective. +/// /// @tparam T Element type. Currently `float`, `__nv_bfloat16`, and `uint32_t`. Other types are /// rejected at compile time rather than silently mapped, because the underlying instruction accepts /// only certain operation and type combinations. diff --git a/python/csrc/gpu_utils_py.cpp b/python/csrc/gpu_utils_py.cpp index d65275022..c05e5cc33 100644 --- a/python/csrc/gpu_utils_py.cpp +++ b/python/csrc/gpu_utils_py.cpp @@ -113,6 +113,7 @@ static nb::capsule toDlpack(GpuBuffer buffer, std::string dataType, std::v void register_gpu_utils(nb::module_& m) { m.def("is_nvls_supported", &isNvlsSupported); + m.def("is_bulk_supported", &isBulkSupported); nb::enum_(m, "CppGpuBufferGranularity") .value("MultiCastMinimum", GpuBufferGranularity::MultiCastMinimum) diff --git a/python/mscclpp/__init__.py b/python/mscclpp/__init__.py index 094081712..55004eebc 100644 --- a/python/mscclpp/__init__.py +++ b/python/mscclpp/__init__.py @@ -52,6 +52,7 @@ CppRawGpuBuffer as RawGpuBuffer, CppReduceOp as ReduceOp, env, + is_bulk_supported, is_nvls_supported, cpp_npkit as npkit, ) @@ -87,6 +88,7 @@ "ReduceOp", "env", "version", + "is_bulk_supported", "is_nvls_supported", "alloc_shared_physical_cuda", "npkit", diff --git a/test/unit/bulk_tests.cu b/test/unit/bulk_tests.cu index e328ccec2..6731a12b3 100644 --- a/test/unit/bulk_tests.cu +++ b/test/unit/bulk_tests.cu @@ -133,14 +133,6 @@ __global__ void kernelBulkReduce(T* dst, T addend, uint32_t count) { #endif } -static bool bulkAvailable() { - int device = 0; - if (cudaGetDevice(&device) != cudaSuccess) return false; - int major = 0; - if (cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device) != cudaSuccess) return false; - return major >= 9; -} - class BulkTestData { public: BulkTestData(int numElems) @@ -170,7 +162,7 @@ class BulkTestData { }; TEST(BulkTest, Load) { - if (!bulkAvailable()) { + if (!mscclpp::isBulkSupported()) { SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; return; } @@ -181,7 +173,7 @@ TEST(BulkTest, Load) { } TEST(BulkTest, Gather) { - if (!bulkAvailable()) { + if (!mscclpp::isBulkSupported()) { SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; return; } @@ -197,7 +189,7 @@ TEST(BulkTest, Gather) { } TEST(BulkTest, RoundTrip) { - if (!bulkAvailable()) { + if (!mscclpp::isBulkSupported()) { SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; return; } @@ -228,7 +220,7 @@ static void reduceStoreTest(float base, float addend) { } TEST(BulkTest, ReduceStoreFloat) { - if (!bulkAvailable()) { + if (!mscclpp::isBulkSupported()) { SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; return; } @@ -236,7 +228,7 @@ TEST(BulkTest, ReduceStoreFloat) { } TEST(BulkTest, ReduceStoreBf16) { - if (!bulkAvailable()) { + if (!mscclpp::isBulkSupported()) { SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; return; } @@ -244,7 +236,7 @@ TEST(BulkTest, ReduceStoreBf16) { } TEST(BulkTest, ReduceStoreUint32) { - if (!bulkAvailable()) { + if (!mscclpp::isBulkSupported()) { SKIP_TEST() << "Bulk asynchronous copy requires compute capability 9.0 or higher."; return; } From 9b9996cb6d711e5edbb193bb7b004da9fb27b391 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Fri, 7 Aug 2026 06:42:19 +0000 Subject: [PATCH 4/6] Silence unused-private-field warning on ROCm Where MSCCLPP_BULK_AVAILABLE is 0, BulkBarrier has no operations, only the storage declared so host code can size shared memory holding barriers. Clang then warns on every ROCm translation unit that includes the header: include/mscclpp/bulk_device.hpp:139:23: warning: private field 'mbar_' is not used [-Wunused-private-field] Mark the member maybe_unused. No effect where the operations exist. Verified on MI300X (ROCm 7.2, gfx942): warning count for this field 0, and both suites match the merge-base exactly -- unit_tests 32 ran / 23 passed / 9 skipped, mp_unit_tests --filter=-Ib 29 ran / 24 passed / 3 failed. The three failures are CommunicatorTest.BasicWrite, .WriteWithDeviceSemaphores and .WriteWithHostSemaphores, which fail identically on the merge-base; this node has no IB device. Re-verified on H200: unit_tests 40/40, mp_unit_tests 60/60. --- include/mscclpp/bulk_device.hpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/include/mscclpp/bulk_device.hpp b/include/mscclpp/bulk_device.hpp index 07129bedc..996ac7e60 100644 --- a/include/mscclpp/bulk_device.hpp +++ b/include/mscclpp/bulk_device.hpp @@ -136,7 +136,9 @@ struct BulkBarrier { #endif // MSCCLPP_BULK_AVAILABLE private: - alignas(8) uint64_t mbar_; + // Marked maybe_unused because the operations that read it exist only where MSCCLPP_BULK_AVAILABLE + // is 1, while the storage is declared everywhere so host code can size shared memory. + [[maybe_unused]] alignas(8) uint64_t mbar_; }; #if MSCCLPP_BULK_AVAILABLE From 89305a7a71156f1da4c5acfd09752f21f8efc1c3 Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Fri, 7 Aug 2026 16:35:01 +0000 Subject: [PATCH 5/6] Use maybe_unused for architecture-gated kernel parameters Mark kernel parameters [[maybe_unused]] at their declarations instead of adding #else blocks with (void)param casts when bulk copy is unavailable. This keeps the unsupported-target path declarative and removes 25 lines of warning-only code. Verified with nvcc at sm_80, the H200 bulk tests (6/6 unit, 3/3 multi-rank), and the MI300X ROCm build (no unused-parameter warnings; unit tests 23 passed / 9 skipped). --- test/mp_unit/bulk_pattern_tests.cu | 24 ++++++------------------ test/unit/bulk_tests.cu | 21 ++++----------------- 2 files changed, 10 insertions(+), 35 deletions(-) diff --git a/test/mp_unit/bulk_pattern_tests.cu b/test/mp_unit/bulk_pattern_tests.cu index de5295f25..9a0f6b624 100644 --- a/test/mp_unit/bulk_pattern_tests.cu +++ b/test/mp_unit/bulk_pattern_tests.cu @@ -137,7 +137,8 @@ DeviceState uploadPeerState(const std::vector(); barrier.invalidate(); } -#else - (void)localTokens; - (void)peerPools; - (void)rank; - (void)worldSize; #endif } @@ -244,7 +240,8 @@ TEST(BulkPatternTest, Dispatch) { // current one is reduced. The two barriers are set up with relaxedInit() under a single bulkFence() // rather than paying a fence each. // ------------------------------------------------------------------------------------------------ -__global__ void kernelCombine(float* output, void** peerPools, int rank, int worldSize) { +__global__ void kernelCombine([[maybe_unused]] float* output, [[maybe_unused]] void** peerPools, + [[maybe_unused]] int rank, [[maybe_unused]] int worldSize) { #if MSCCLPP_BULK_AVAILABLE extern __shared__ __align__(128) uint8_t shared[]; // Layout: [stage][contributor][chunk] tiles, then the per-stage barriers. @@ -294,11 +291,6 @@ __global__ void kernelCombine(float* output, void** peerPools, int rank, int wor if (threadIdx.x == 0) { for (int s = 0; s < kStages; ++s) barriers[s].invalidate(); } -#else - (void)output; - (void)peerPools; - (void)rank; - (void)worldSize; #endif } @@ -350,7 +342,8 @@ TEST(BulkPatternTest, Combine) { // destination, so the accumulator is never read back across the link and no rank stages anyone // else's data. The reverse link direction stays idle. // ------------------------------------------------------------------------------------------------ -__global__ void kernelCombineReduce(const float* contribution, void** peerPools, int rank, int worldSize) { +__global__ void kernelCombineReduce([[maybe_unused]] const float* contribution, [[maybe_unused]] void** peerPools, + [[maybe_unused]] int rank, [[maybe_unused]] int worldSize) { #if MSCCLPP_BULK_AVAILABLE __shared__ alignas(128) uint8_t tile[kChunkBytes]; __shared__ mscclpp::BulkBarrier barrier; @@ -383,11 +376,6 @@ __global__ void kernelCombineReduce(const float* contribution, void** peerPools, } if (threadIdx.x == 0) barrier.invalidate(); -#else - (void)contribution; - (void)peerPools; - (void)rank; - (void)worldSize; #endif } diff --git a/test/unit/bulk_tests.cu b/test/unit/bulk_tests.cu index 6731a12b3..73c92a823 100644 --- a/test/unit/bulk_tests.cu +++ b/test/unit/bulk_tests.cu @@ -14,7 +14,7 @@ constexpr uint32_t kTile = 4096; constexpr uint32_t kElems = kTile / sizeof(int); // Load one tile into shared memory, then copy it out. -__global__ void kernelBulkLoad(const int* src, int* dst) { +__global__ void kernelBulkLoad([[maybe_unused]] const int* src, [[maybe_unused]] int* dst) { #if MSCCLPP_BULK_AVAILABLE __shared__ alignas(128) int tile[kElems]; __shared__ mscclpp::BulkBarrier barrier; @@ -33,15 +33,12 @@ __global__ void kernelBulkLoad(const int* src, int* dst) { __syncthreads(); if (threadIdx.x == 0) barrier.invalidate(); -#else - (void)src; - (void)dst; #endif } // Gather NumSrc tiles into one barrier, then reduce them. template -__global__ void kernelBulkGather(const int* src, int* dst) { +__global__ void kernelBulkGather([[maybe_unused]] const int* src, [[maybe_unused]] int* dst) { #if MSCCLPP_BULK_AVAILABLE __shared__ alignas(128) int tiles[NumSrc][kElems]; __shared__ mscclpp::BulkBarrier barrier; @@ -64,15 +61,12 @@ __global__ void kernelBulkGather(const int* src, int* dst) { __syncthreads(); if (threadIdx.x == 0) barrier.invalidate(); -#else - (void)src; - (void)dst; #endif } // Reuse one barrier across NumChunks phases, staging each chunk in and storing it back out. template -__global__ void kernelBulkRoundTrip(const int* src, int* dst) { +__global__ void kernelBulkRoundTrip([[maybe_unused]] const int* src, [[maybe_unused]] int* dst) { #if MSCCLPP_BULK_AVAILABLE __shared__ alignas(128) int tile[kElems]; __shared__ mscclpp::BulkBarrier barrier; @@ -106,15 +100,12 @@ __global__ void kernelBulkRoundTrip(const int* src, int* dst) { mscclpp::bulkStoreWait(); // every store has landed barrier.invalidate(); } -#else - (void)src; - (void)dst; #endif } // Accumulate a staged tile into a seeded destination with the copy engine. template -__global__ void kernelBulkReduce(T* dst, T addend, uint32_t count) { +__global__ void kernelBulkReduce([[maybe_unused]] T* dst, [[maybe_unused]] T addend, [[maybe_unused]] uint32_t count) { #if MSCCLPP_BULK_AVAILABLE extern __shared__ __align__(128) uint8_t raw[]; T* tile = reinterpret_cast(raw); @@ -126,10 +117,6 @@ __global__ void kernelBulkReduce(T* dst, T addend, uint32_t count) { mscclpp::bulkStoreCommit(); mscclpp::bulkStoreWait(); } -#else - (void)dst; - (void)addend; - (void)count; #endif } From bf7c736ae53aafa227a36d75e537215d39aac3aa Mon Sep 17 00:00:00 2001 From: Changho Hwang Date: Thu, 13 Aug 2026 05:35:16 +0000 Subject: [PATCH 6/6] Address bulk reduction review feedback --- include/mscclpp/bulk_device.hpp | 20 ++++++++++++-------- test/mp_unit/bulk_pattern_tests.cu | 11 ++++++++--- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/include/mscclpp/bulk_device.hpp b/include/mscclpp/bulk_device.hpp index 996ac7e60..1377db38a 100644 --- a/include/mscclpp/bulk_device.hpp +++ b/include/mscclpp/bulk_device.hpp @@ -193,16 +193,20 @@ enum class BulkRedOp { Add }; /// returns immediately, joins the calling thread's open bulk group, and requires a preceding /// bulkFence() to make generic writes to @p srcShared visible. /// -/// Measured on H200 at roughly 90% of the bulkStore() rate for the same payload, from a single -/// issuing thread, so the accumulate is close to free relative to the transfer. +/// Measured on H200 at roughly 90% of the bulkStore() rate for the same payload from a single +/// issuing thread, a roughly 10% throughput cost in that microbenchmark. The cost visible to a +/// workload depends on its transfer overlap and bottlenecks. /// -/// @warning When several devices accumulate into the same destination address concurrently, the -/// result depends on the per-element reduction being atomic across peers. That holds in every -/// measurement taken here (exact results on H200 across 8 six-rank runs and on GB200 across 50 -/// four-GPU runs), but it is an empirical result, not a guarantee found in the PTX documentation. -/// Confirm it on the target hardware before relying on it for a collective. +/// PTX ISA 9.3 specifies each element-wise reduction as atomic and defaults an omitted scope to +/// `.relaxed.sys`. Earlier PTX documentation, including that shipped with CUDA 12.9 and 13.0, +/// specifies `.relaxed.gpu` semantics and does not guarantee atomicity across devices. Because this +/// API supports those toolkits, concurrent reductions from several devices into the same peer address +/// must be treated as empirical rather than portable; exact results were observed on H200 and GB200, +/// but callers must confirm the behavior on their target toolchain and hardware. /// -/// @tparam T Element type. Currently `float`, `__nv_bfloat16`, and `uint32_t`. Other types are +/// @tparam T Source and destination element type. Currently `float`, `__nv_bfloat16`, and `uint32_t`. +/// The instruction does not perform mixed-precision conversion: `__nv_bfloat16` accumulates in bf16; +/// for fp32 accumulation, convert the source tile and use a float destination. Other types are /// rejected at compile time rather than silently mapped, because the underlying instruction accepts /// only certain operation and type combinations. /// @tparam Op Reduction operation. diff --git a/test/mp_unit/bulk_pattern_tests.cu b/test/mp_unit/bulk_pattern_tests.cu index 9a0f6b624..2fc276e45 100644 --- a/test/mp_unit/bulk_pattern_tests.cu +++ b/test/mp_unit/bulk_pattern_tests.cu @@ -340,7 +340,9 @@ TEST(BulkPatternTest, Combine) { // Instead of every rank pulling and summing, each rank accumulates its own contribution straight // into the destination's accumulator with bulkReduceStore(). The copy engine performs the add at the // destination, so the accumulator is never read back across the link and no rank stages anyone -// else's data. The reverse link direction stays idle. +// else's data. Cross-device atomicity is guaranteed only by PTX ISA 9.3's system-scope semantics; +// with earlier supported toolkits this pattern exercises empirically observed behavior rather than a +// portable guarantee. The reverse link direction stays idle. // ------------------------------------------------------------------------------------------------ __global__ void kernelCombineReduce([[maybe_unused]] const float* contribution, [[maybe_unused]] void** peerPools, [[maybe_unused]] int rank, [[maybe_unused]] int worldSize) { @@ -369,13 +371,16 @@ __global__ void kernelCombineReduce([[maybe_unused]] const float* contribution, mscclpp::bulkFence(); // staged data -> the reduction's read of the tile mscclpp::bulkReduceStore(dstRow + off, tile, kChunkBytes); mscclpp::bulkStoreCommit(); - mscclpp::bulkStoreWait<0>(); // the accumulate must land before the tile is refilled + mscclpp::bulkStoreWaitSource<0>(); // the tile can be refilled once its source read completes } __syncthreads(); } } - if (threadIdx.x == 0) barrier.invalidate(); + if (threadIdx.x == 0) { + mscclpp::bulkStoreWait<0>(); // every accumulate must land before the kernel completes + barrier.invalidate(); + } #endif }