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.
• Update triton-ascend submodule repository link in .gitmodules
• 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.
• Add gm_addr_args_indices attribute to HIVM custom operators to track function parameter pointers.

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
• putmem_signal / putmem_signal_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
• python/triton_dist/test/ascend/test_gm_addr_args_indices.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!

[peterjl]

All tests have been passed with the following configuration:

• Triton-distributed：d36c35d2ae89
• triton-ascend：bfd8f55d124f
• SHMEM：81c95bad8b94
• AscendNPU-IR：05c7506f5741
• llvm-project：487bf69401b9

[peterjl]

Also, I found three issues:

1. [Blocking] Keep update = none for triton-ascend

Removing it makes NVIDIA/AMD's recursive submodule initialization enter triton-ascend and its GitCode-hosted AscendNPU-IR dependency. Our GitHub CI cannot access GitCode reliably, so unrelated jobs may fail or stall.

Please keep update = none. Ascend builds can opt in using the existing main-branch command:

git submodule update --init --checkout --depth=1 \
    3rdparty/triton-ascend 3rdparty/shmem

There is currently no Ascend-specific build script under scripts/; this command in docs/build.md is the existing entry point.

2. [Medium] signal_wait_until advertises invalid non-int32 overloads

The tested int32 path works. However, the wrapper registers every RMA_DTYPES dtype while always casting the comparison value to int32; NPU-IR expects (pointer<T>, int32, T) -> void. This makes float and 64-bit overloads incorrectly typed, and fp16/bf16 have no generated NPU-IR wait wrapper.

Please either expose only the supported int32 form, or add matching per-dtype signatures and tests.

3. [Minor] Remove the Chinese comment

test_gm_addr_args_indices.py:34 contains # triton + 内置 dialect. Please change it to English.

[janssen-v]

Also, I found three issues:

1. [Blocking] Keep update = none for triton-ascend

Removing it makes NVIDIA/AMD's recursive submodule initialization enter triton-ascend and its GitCode-hosted AscendNPU-IR dependency. Our GitHub CI cannot access GitCode reliably, so unrelated jobs may fail or stall.

Ah, I didn't notice that NV/AMD builds initialized submodules recursively and thought updating triton-ascend to point to the in-tree GitHub repo would have been enough. update = none has been reinstated in this PR; will also follow-up with the triton-ascend guys to see if they would be amenable to switching their NPU-IR submodule to point to the GitHub mirror instead.

2. [Medium] signal_wait_until advertises invalid non-int32 overloads

The tested int32 path works. However, the wrapper registers every RMA_DTYPES dtype while always casting the comparison value to int32; NPU-IR expects (pointer<T>, int32, T) -> void. This makes float and 64-bit overloads incorrectly typed, and fp16/bf16 have no generated NPU-IR wait wrapper.

Please either expose only the supported int32 form, or add matching per-dtype signatures and tests.

Good catch, it seems that wait_until falls under P2P sync instead of RMA in Ascend SHMEM and has a different set of supported types. Upon closer review, I also found that our tests for signal_wait_until didn't exercise the return value for this op, which isn't currently supported in the variable typed wait_until.

I've modified the Ascend implementation to dispatch signal_wait_until to the fixed int32 aclshmem_signal_wait_until API for now while we find a good way to support return values for the variable typed wait_until. The signal_op test has also been modified to exercise the return value.

3. [Minor] Remove the Chinese comment

test_gm_addr_args_indices.py:34 contains # triton + 内置 dialect. Please change it to English.

Updated

[peterjl]

Thanks, this matches my expectation. Restricting signal_wait_until to the fixed int32 API and adding coverage for its return value addresses my concern. The update = none and comment fixes also look good.

I will rerun the Ascend test suite against the latest head and report the result.

[janssen-v]

Thanks, this matches my expectation. Restricting signal_wait_until to the fixed int32 API and adding coverage for its return value addresses my concern. The update = none and comment fixes also look good.

I will rerun the Ascend test suite against the latest head and report the result.

Please run the suite against NPU-IR at 971145a0 on my fork since the changes to the template are still waiting to be merged.

[peterjl]

LGTM