Skip to content

Add switch-channel NVLS cross-rank barrier - #843

Open
RJ Souza (Empyreus) wants to merge 22 commits into
mainfrom
rjsouza/nvls-barrier
Open

Add switch-channel NVLS cross-rank barrier#843
RJ Souza (Empyreus) wants to merge 22 commits into
mainfrom
rjsouza/nvls-barrier

Conversation

@Empyreus

Copy link
Copy Markdown
Contributor

Adds a device-side cross-rank barrier built on NVLS (multimem) atomics

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

(cherry picked from commit 51a8235)
(cherry picked from commit 38494de)
(cherry picked from commit 5f93881)
(cherry picked from commit 921b280)
(cherry picked from commit 1d55a28)

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds a device-side, cross-rank barrier for NVLS (multimem) SwitchChannels and wires it through the executor and Python DSL so kernels can synchronize all ranks in an NVLS group without an O(n²) semaphore mesh.

Changes:

  • Add NVLS-backed barrier state to NvlsConnection/SwitchChannel and expose it via SwitchChannelDeviceHandle::barrier().
  • Introduce a new executor op (MULTI_BARRIER / "gbarrier") with device-side implementation that grid-syncs then issues a single multimem arrival per rank.
  • Add Python DSL support (SwitchChannel.barrier(...)) plus a single-node allgather example that uses the new barrier.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
src/core/switch_channel.cc Creates an auxiliary NVLS multicast for barrier state and stamps barrier pointers onto bound SwitchChannels.
src/core/include/execution_kernel.hpp Implements MULTI_BARRIER device handler (grid-converge → one multimem barrier → release).
src/core/include/execution_common.hpp Adds OperationType::MULTI_BARRIER.
src/core/executor/execution_plan.cc Parses "gbarrier" and maps JSON switch_channel_id into Operation.nvlsInputIndex.
include/mscclpp/switch_channel.hpp Stores barrier resources on NvlsConnection and propagates them into SwitchChannel handles.
include/mscclpp/switch_channel_device.hpp Adds SwitchChannelDeviceHandle::barrier() API and barrier metadata fields.
python/mscclpp/language/internal/types.py Adds DSL instruction enum value for "gbarrier".
python/mscclpp/language/internal/operations.py Adds GroupBarrier operation with instancing semantics collapsed onto one per-rank syncer.
python/mscclpp/language/internal/buffer_access.py Treats group_barrier as an access-sequence breaker like other barriers/nops.
python/mscclpp/language/channel.py Adds SwitchChannel.barrier() plus rank-view convenience wrapper and docstrings.
python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_barrier.py Example program exercising the new switch-native barrier.

Comment thread include/mscclpp/switch_channel_device.hpp Outdated
Comment thread src/core/switch_channel.cc Outdated
Comment thread include/mscclpp/switch_channel.hpp Outdated
Comment thread include/mscclpp/switch_channel_device.hpp Outdated
Comment thread python/mscclpp/language/internal/buffer_access.py Outdated
Comment thread include/mscclpp/switch_channel_device.hpp
Comment thread include/mscclpp/switch_channel.hpp Outdated
Comment thread include/mscclpp/switch_channel.hpp Outdated
Comment thread src/core/switch_channel.cc Outdated
Comment thread src/core/switch_channel.cc Outdated
Comment thread src/core/switch_channel.cc Outdated
Comment thread include/mscclpp/switch_channel.hpp Outdated
Comment thread src/core/include/execution_kernel.hpp Outdated
Comment thread src/core/include/execution_kernel.hpp Outdated
Comment thread src/core/switch_channel.cc Outdated
Comment thread src/core/switch_channel.cc Outdated

@chhwang Changho Hwang (chhwang) left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants