Quick Start

This page runs your first distributed kernel on Ascend NPUs. It assumes you have completed Installation and Build from Source.

Set up the runtime environment

Every shell that runs a Triton-distributed-ascend program needs two things on the environment: the CANN variables and the AscendNPU-IR binaries.

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export PATH={PATH_TO}/AscendNPU-IR/build/bin:$PATH

Replace {PATH_TO} with wherever you cloned AscendNPU-IR. Without this on PATH, kernels fail at compile time rather than at import.

Run the AllGather-GEMM tutorial

From the repository root:

torchrun --nproc-per-node=2 tutorials/ascend/01-ascend-allgather-gemm.py

--nproc-per-node sets the number of ranks, one process per NPU. Use 2 to start; it is the smallest configuration that exercises communication. Do not set it higher than the number of NPUs reported by npu-smi info.

Expected output for a two-rank run:

[INFO] Rank 0 of 2 initialized
[INFO] Rank 1 of 2 initialized
[PASS] Rank 0: C_golden and C match within tolerances (rtol=1e-3, atol=1e-3).
[PASS] Rank 1: C_golden and C match within tolerances (rtol=1e-3, atol=1e-3).
[INFO] Test passed successfully for 2 ranks!

Each rank checks its own output against a reference built from torch.distributed.all_gather followed by a local torch.matmul, then the ranks exchange pass/fail status so a failure on any rank fails the whole run.

What the example does

The tutorial fuses an AllGather with a GEMM. Every rank holds a shard A_local of shape [M, K] and a full copy of B of shape [K, N], and produces C of shape [M * world_size, N]. The default problem size is fp16 with M = N = K = 4096.

Rather than gathering all shards and then multiplying, the kernel overlaps the two: while one block of A is being fetched from a peer, the previous block is already being multiplied. That overlap is the point of the project, and it is why the kernel needs a symmetric memory buffer instead of ordinary device tensors.

The structure to recognize, since every tutorial follows it:

1. shmem.aclshmem_init with an InitAttr carrying rank, world size, and a rendezvous address. This sets up the symmetric heap.
2. shmem.aclshmem_create_tensor for the peer-visible staging buffer. Its size depends on the block sizes, the world size, and the number of pipeline buffers.
3. The Triton kernel, using triton_dist.language primitives for the communication.
4. shmem.aclshmem_free_tensor and shmem.aclshmem_finalize for teardown.

Process group setup uses the hccl backend, and rank assignment reads LOCAL_RANK from the environment, which torchrun provides.

The example uses `tcp://127.0.0.1:8666` as its shmem rendezvous address. It is hardcoded, so
two concurrent runs on the same host will collide on that port. Run one at a time, or edit
`G_IP_PORT` in the tutorial.

Other tutorials

The tutorials/ascend/ directory covers the rest of the primitives and overlap patterns:

Tutorial	Topic
01-ascend-allgather-gemm.py	AllGather fused with GEMM
02-ascend-gemm-reduce-scatter.py	GEMM fused with ReduceScatter
03-ascend-gemm-allreduce-oneshot.py	GEMM with one-shot AllReduce
04-ascend-reverse-all2all	Reverse All2All
05-ascend-reverse-all2all-barrier.py	Reverse All2All with barrier synchronization
06-ascend-gemm-ar-notify.py	GEMM AllReduce with notify-based signaling
07-qkv-alltoall	QKV All2All for attention
08-ascend-transpose-all2all	Transposing All2All
09-ascend-dispatch-all2all	MoE dispatch
10-ascend-combine-all2all	MoE combine

Run any of them the same way, adjusting --nproc-per-node to your device count:

torchrun --nproc-per-node=2 tutorials/ascend/02-ascend-gemm-reduce-scatter.py

Run the test suite

To confirm the primitives work on your hardware beyond the tutorials:

pytest python/triton_dist/test/ascend/ -m dist

The dist marker selects the tests that need multiple NPUs. These cover barriers, wait/notify, symmetric memory addressing, and put/get.

Troubleshooting

Kernel compilation errors. The AscendNPU-IR binaries are not on PATH. Re-export it as shown above.

KeyError: 'LOCAL_RANK'. The script was run directly with python. These tutorials must be launched through torchrun.

Hangs at startup. Usually a stale process holding the rendezvous port from a previous run, or --nproc-per-node set higher than the number of available NPUs. Check for leftover processes and confirm your device count.

aclshmem_init failed. shmem could not set up the symmetric heap. Confirm the shmem wheel is installed and CANN is sourced in the current shell.

Next steps

• Read through the tutorials in tutorials/ascend/ for the communication primitives.
• See the Contributing guide if you plan to submit changes.