# Quick Start

This page runs your first distributed kernel on Ascend NPUs. It assumes you have completed
[Installation](installation.md) and [Build from Source](build.md).

## 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.

```bash
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:

```bash
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:

```text
[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.

```{note}
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:

```bash
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:

```bash
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](https://gitcode.com/Ascend/Triton-distributed-ascend/blob/master/CONTRIBUTING.md)
  if you plan to submit changes.