[SymmMem] Remove redundant dist.barrier in Triton NVSHMEM tests & add device‐side signal_op support #156684
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/156684
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ No FailuresAs of commit a3c5209 with merge base 455dfd2 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
oncall: distributed
codingwithsurya
changed the title
codingwithsurya
added
the
release notes: distributed (c10d)
codingwithsurya
force-pushed
the
gh/codingwithsurya/6/head
branch
2 times, most recently
from
373dc9d to
def902f
Compare
kwen2501
added
the
release notes: distributed (symm_mem)
| signal: tl.constexpr, | ||
| sig_op: tl.constexpr, | ||
| peer: tl.constexpr, |
There was a problem hiding this comment.
Should these three operands be constexpr? Can they be regular variables? Just to avoid re-compilation.
There was a problem hiding this comment.
talked offline; these make sense to be regular variables
There was a problem hiding this comment.
i'm going to update this on top PR of this stack.
|
|
||
| peer = 1 - rank | ||
| NVSHMEM_CMP_EQ = 0 # from nvshmem.h | ||
| NVSHMEM_SIGNAL_SET = 0 # atomic set operation |
There was a problem hiding this comment.
I see we can also pass in NVSHMEM_SIGNAL_ADD here? (https://docs.nvidia.com/nvshmem/api/gen/api/signal.html#available-signal-operators )
Do we see any scenario where it could be useful? @kwen2501
There was a problem hiding this comment.
yes, when the sender wants to monotonically update a counter as versioning info to the receiver.
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 4 mandatory check(s) failed. The first few are:
Dig deeper by viewing the failures on hud |
|
some checks need to be rerun. pull-rebasing. |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
…d file (#156685) ## Summary Moved the Triton-specific NVSHMEM tests in `test_nvshmem.py` into a dedicated `test_nvshmem_triton.py` file. Also put the shared Triton JIT kernels at the top-level of new file for reusability. ## Testing ```bash TORCH_SYMMMEM=NVSHMEM python test/distributed/test_nvshmem.py TORCH_SYMMMEM=NVSHMEM python test/distributed/test_nvshmem_triton.py ``` All 16 original tests pass with no functionality changes. Pull Request resolved: #156685 Approved by: https://github.com/mandroid6, https://github.com/kwen2501 ghstack dependencies: #156684
4 participants
Summary
This PR removes unnecessary
dist.barriercalls up in our Triton NVSHMEM test suite and adds signal_op support, which is a lightweight device-side signaling mechanism. Added test for this in ourwait_untilkernel and correspondingcore.externwrapper.Why did we drop the
dist.barrier()calls?We dropped the host‐side dist.barrier() in all Triton NVSHMEM tests (except the raw put/get cases) because every other test already uses NVSHMEM collectives or device‐side sync primitives (fence/quiet/signal/wait), making the extra barrier redundant. This keeps synchronization entirely on the GPU and leverages NVSHMEM’s native ordering guarantees for clearer, more efficient tests.
test_triton_wait_untilupdateput_kernelwrites the data, launchessignal_op_kernelto atomically set Rank 0's flag vianvshmemx_signal_opdist.barrier()and simply callswait_until_kernelto spin-wait on the device flag, then asserts data correctnessTesting
Stack from ghstack (oldest at bottom):
cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k