Add bulk asynchronous copy primitives - #853
Open
Changho Hwang (chhwang) wants to merge 5 commits into
Open
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
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<T,Op> shared -> global, accumulating at the destination bulkStoreCommit close the current bulk group bulkStoreWait<N> stores have landed bulkStoreWaitSource<N> 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.
Changho Hwang (chhwang)
force-pushed
the
feature/tma-bulk-memory-channel
branch
from
August 4, 2026 17:17
c407c26 to
bfb4d24
Compare
<cuda_bf16.h> 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.
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.
Changho Hwang (chhwang)
changed the title
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.
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).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds pointer-based 1-D bulk-copy primitives for NVIDIA sm_90+ in
mscclpp/bulk_device.hpp:BulkBarrierfor reusable async-load completionbulkLoadandbulkStorefor global ↔ shared copiesbulkReduceStore<T>for copy-engine reduction into global or peer memoryisBulkSupported()/is_bulk_supportedcapability queriesThe API operates on global pointers, including peer-mapped pointers. It does not add
MemoryChannelmethods: EP kernels use raw peer-pointer arrays for data movement and channels only for synchronization. The API covers all TMA instructions used by #852 andqinghuazhou/unified_ep_bench_ht_python.Unsupported device targets fail at compile time instead of silently doing nothing.
BulkBarrierstorage remains host-visible so launch code can size dynamic shared memory. The store API separates destination completion from source-tile reuse for pipelined kernels.bulkReduceStorecurrently supportsAddforfloat,__nv_bfloat16, anduint32_t; unsupported combinations fail astatic_assert.Tests
unit_tests40/40;mp_unit_tests60/60; patterns 3/3 at 6 ranksunit_tests40/40; non-IBmp_unit_tests32/32; patterns 3/3 at 4 ranksLimitation
Concurrent reduction from several GPUs into the same peer address produced exact results across repeated H200 and GB200 runs, but cross-peer atomicity was not found in the PTX documentation. The header documents this as empirical rather than guaranteed.