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 example

From the repository root:

torchrun --nproc-per-node=2 tutorials/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 example 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 main execution flow of this example is as follows:

  1. Call shmem.aclshmem_init with an InitAttr carrying rank, world size, symmetric heap size, and a rendezvous address to set up the symmetric heap.

  2. shmem.aclshmem_create_tensor creates peer-visible staging buffers. Their size depends on the block sizes, world size, and 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.

Note

The example uses tcp://127.0.0.1:8666 as its shmem rendezvous address by default. Two concurrent runs on the same host will collide on that port. Run one instance at a time, or pass a different ASH_MASTER_PORT to each instance.

Run the test suite

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

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. This example 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