feat: Add Ascend RDMA support

[janssen-v]

Summary

Add Ascend RDMA support and update ACLSHMEM to v1.5.0.rc1.

This PR bumps the ACLSHMEM submodule, updates Ascend tests/tutorials for the corrected aclshmem_finalize() API, restores the shared libshmem_device dispatch proxy surface, and adds Ascend device wrappers for ACLSHMEM RMA, signaling, team queries, barriers, quiet, and fence.

It also updates ModuleProxy.dispatch so proxied Triton builtins advertise _semantic in their generated wrapper signatures. This is required for the Ascend path because triton-ascend follows Triton 3.5-style codegen, which only injects _semantic when it is present in the callable signature. The proxy forwards _semantic only to backend methods that accept it, preserving CUDA, HIP, and MACA compatibility.

Changes

• Bump 3rdparty/shmem to ACLSHMEM v1.5.0.rc1.
• Replace deprecated aclshmem_finialize() usages with aclshmem_finalize().
• Restore the shared libshmem_device dispatch proxy API by:
  • removing explicit _semantic parameters from dispatch stubs
  • returning to unconditional backend module imports
• Update ModuleProxy.dispatch to:
  • advertise _semantic=None in generated wrapper signatures
  • forward _semantic only when the selected backend method declares it
  • preserve the original dispatch stub signature
• Add Ascend ACLSHMEM wrappers for:
  • team queries: team_my_pe, team_n_pes, team_translate_pe
  • barriers: barrier, barrier_vec
  • RMA: getmem, putmem, getmem_nbi, putmem_nbi
  • signaling: putmem_signal, putmem_signal_nbi, signal_op, signal_wait_until
  • ordering: quiet, fence
• Expose ACLSHMEM team, compare, and signal constants through the shared SHMEM proxy API.
• Add Ascend dtype-to-kernel suffix mapping for RMA extern calls.
• Keep Triton pointer types hashable for Ascend extern signature dictionaries.
• Wire Ascend SHMEM extern handling.
• Preserve distributed-to-HIVM side effects using MLIR memory-effect information instead of a hardcoded symbol list.
• Improve Ascend distributed test setup with required environment variables and dist pytest marker registration.

Testing

Added Ascend distributed tests for:

• barriers, vector barriers, and barrier_all
• world-team and sub-team barrier behavior
• putmem / putmem_nbi
• getmem / getmem_nbi
• signal_op
• signal_wait_until

New test files:

• python/triton_dist/test/ascend/test_barrier_ops.py
• python/triton_dist/test/ascend/test_put_get_mem.py
• python/triton_dist/test/ascend/test_signal_op.py

These cover ACLSHMEM team/world barriers, vector barriers, sub-team participation, blocking and non-blocking RMA with quiet, signal set/wait behavior, and rank-to-rank ping-pong signaling.

Notes

• ModuleProxy.dispatch carries the _semantic compatibility fix centrally, avoiding _semantic parameters in every shared libshmem_device dispatch stub.
• ACLSHMEM RMA currently supports the dtype subset listed in RMA_DTYPES, with suffixes mapped through DTYPE_TO_KERNEL_SUFFIX.
• Ascend SHMEM extern handling uses an empty libshmem path because ACLSHMEM device symbols are handled by AscendNPU-IR rather than a separate bitcode library.

[CLAassistant]

All committers have signed the CLA.

[peterjl]

Built this on a single-node Ascend 910 (16 chips, no RDMA NIC) and ran the Ascend tests + tutorials. The 3 tutorials and test_my_pe/num_ranks/symm_at/wait_notify pass. A few things I ran into, all in Ascend-only files:

1. signal_wait_until takes _builder instead of _semantic (libaclshmem_device.py). Every other extern here uses _semantic, and the dispatcher only forwards _semantic, so this one throws at trace time (_semantic argument must be provided...). Just needs _builder → _semantic in the signature, the two tl.cast(...) calls, and the final extern_call.

2. "char" in RMA_DTYPES isn't a valid Triton dtype (libaclshmem_device.py). core.dtype("char") asserts on 3.5.0, which breaks the signature dicts for all RMA ops (putmem/getmem/signal_wait_until). It's also redundant with int8. Dropping the "char" entry from RMA_DTYPES + DTYPE_TO_KERNEL_SUFFIX fixes it.

3. test_barrier_ops.py deadlocks. All four kernels call the barrier inside if sub_vec_id() == 0:. The barrier is a multi-vector-core collective, so with only sub-block 0 entering it hangs on-device. The tutorials call barrier_all() outside that guard and work fine — moving it out of the guard makes the test pass:

val = tl.load(counter_ptr)
libshmem_device.barrier_all()          # all vector cores
if sub_vec_id() == 0:
    remote = libshmem_device.remote_ptr(counter_ptr, (my_rank + 1) % world_size)
    tl.store(remote, val + 1)
libshmem_device.barrier_all()          # all vector cores

Minor:

• jit.py returns {"libshmem": str("")} — str("") is redundant, and {} would match the CUDA branch.
• getmem_nbi/putmem_nbi take qp_id but never forward it.
• test_put_get_mem/test_signal_op hardcode OpEngineType.ROCE, so they fail init on a single node without RDMA (the others use MTE). Might be worth defaulting to MTE / skipping when there's no RDMA link.

Can send a PR with the fixes if that helps.

[peterjl]

Heads up: we'll be merging #178 (internal sync) fairly soon — it touches a fair bit of the tree, so you'll likely need to rebase this branch afterwards. Will give you a ping once it lands.

[janssen-v]

Heads up: we'll be merging #178 (internal sync) fairly soon — it touches a fair bit of the tree, so you'll likely need to rebase this branch afterwards. Will give you a ping once it lands.

Got it, no worries.

[janssen-v]

Built this on a single-node Ascend 910 (16 chips, no RDMA NIC) and ran the Ascend tests + tutorials. The 3 tutorials and test_my_pe/num_ranks/symm_at/wait_notify pass. A few things I ran into, all in Ascend-only files:

1. signal_wait_until takes _builder instead of _semantic (libaclshmem_device.py). Every other extern here uses _semantic, and the dispatcher only forwards _semantic, so this one throws at trace time (_semantic argument must be provided...). Just needs _builder → _semantic in the signature, the two tl.cast(...) calls, and the final extern_call.

2. "char" in RMA_DTYPES isn't a valid Triton dtype (libaclshmem_device.py). core.dtype("char") asserts on 3.5.0, which breaks the signature dicts for all RMA ops (putmem/getmem/signal_wait_until). It's also redundant with int8. Dropping the "char" entry from RMA_DTYPES + DTYPE_TO_KERNEL_SUFFIX fixes it.

3. test_barrier_ops.py deadlocks. All four kernels call the barrier inside if sub_vec_id() == 0:. The barrier is a multi-vector-core collective, so with only sub-block 0 entering it hangs on-device. The tutorials call barrier_all() outside that guard and work fine — moving it out of the guard makes the test pass:

val = tl.load(counter_ptr)
libshmem_device.barrier_all()          # all vector cores
if sub_vec_id() == 0:
    remote = libshmem_device.remote_ptr(counter_ptr, (my_rank + 1) % world_size)
    tl.store(remote, val + 1)
libshmem_device.barrier_all()          # all vector cores

Minor:

• jit.py returns {"libshmem": str("")} — str("") is redundant, and {} would match the CUDA branch.
• getmem_nbi/putmem_nbi take qp_id but never forward it.
• test_put_get_mem/test_signal_op hardcode OpEngineType.ROCE, so they fail init on a single node without RDMA (the others use MTE). Might be worth defaulting to MTE / skipping when there's no RDMA link.

Can send a PR with the fixes if that helps.

Hi Peter, thanks for reviewing our PR.

1. Sorry, I missed this in a rebase. We have fixed this issue in our internal repository. We will be moving our dev repository to track upstream more closely once the triton-ascend dependency is decoupled so that there will be less chances of this happening in the future.
2. Got it.
3. Hmm, I didn't experience this error on both my 910 & 950 test servers. This should pass since regardless of where barrier_all is called from, it will execute a barrier on all ranks. I will check your version later.

Minor:

• Ok
• Currently ACLSHMEM does not support queue pairs, so we don't pass it to the backend. I suppose removing it from the ACLSHMEM dispatch would be fine.
• I'll set the tests to be MTE by default and have ROCE separate.

[janssen-v]

3. test_barrier_ops.py deadlocks. All four kernels call the barrier inside if sub_vec_id() == 0:. The barrier is a multi-vector-core collective, so with only sub-block 0 entering it hangs on-device. The tutorials call barrier_all() outside that guard and work fine — moving it out of the guard makes the test pass:

Upon further review, there seems to be a problem with running the barrier test when Triton-distributed is paired with a recent version of NPU-IR. Prior versions do not have this issue. We will investigate this further.

Also, while both the original and your modified versions of the barrier test pass on my system, moving the barrier outside the guard does make the test run faster end to end so I'll add the change. Thanks!