# 快速开始

本页在昇腾 NPU 上运行第一个分布式 kernel。假设你已经完成
[安装](installation.md)和[从源码构建](build.md)。

## 配置运行环境

每个运行 Triton-distributed-ascend 程序的 shell 需要两项环境配置：CANN 变量以及
AscendNPU-IR 二进制文件。

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

将 `{PATH_TO}` 替换为你克隆 AscendNPU-IR 的路径。如果 `PATH` 中缺少这些二进制文件，
kernel 会在编译时而非导入时失败。

## 运行 AllGather-GEMM 教程

在仓库根目录执行：

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

`--nproc-per-node` 设置 rank 数量，每个 NPU 对应一个进程。从 2 开始；这是能够验证通信的
最小配置。不要将其设置为高于 `npu-smi info` 报告的 NPU 数量。

两 rank 运行的预期输出：

```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!
```

每个 rank 将自己的输出与参考结果比对（参考结果来自 `torch.distributed.all_gather` 接
`torch.matmul`），然后各 rank 交换 pass/fail 状态，任何一个 rank 失败都会导致整体失败。

## 示例做了什么

该教程将 AllGather 与 GEMM 融合。每个 rank 持有形状为 `[M, K]` 的分片 `A_local` 和
形状为 `[K, N]` 的完整 `B`，并生成形状为 `[M * world_size, N]` 的 `C`。默认问题规模为
fp16，`M = N = K = 4096`。

该 kernel 没有先聚集所有分片再相乘，而是将两者掩盖：在从 peer 拉取一个 `A` 块时，前一个
块已经在进行乘法。掩盖是本项目的核心，也是 kernel 需要对称内存缓冲区而非普通设备 tensor
的原因。

以下结构需要认识，因为每个教程都遵循它：

1. `shmem.aclshmem_init`，传入包含 rank、world size 和会合地址的 `InitAttr`。此步骤
   设置对称堆。
2. `shmem.aclshmem_create_tensor` 创建对 peer 可见的暂存缓冲区。其大小取决于块大小、
   world size 以及流水线缓冲区数量。
3. Triton kernel，使用 `triton_dist.language` 原语进行通信。
4. `shmem.aclshmem_free_tensor` 和 `shmem.aclshmem_finalize` 用于清理。

进程组设置使用 `hccl` 后端，rank 分配从环境变量 `LOCAL_RANK` 读取，由 `torchrun` 提供。

```{note}
示例使用 `tcp://127.0.0.1:8666` 作为 shmem 会合地址。该地址硬编码在代码中，因此同一台
主机上的两次并发运行会在该端口冲突。一次只运行一个实例，或编辑教程中的 `G_IP_PORT`。
```

## 其他教程

`tutorials/ascend/` 目录涵盖其余原语与掩盖模式：

| 教程 | 主题 |
| --- | --- |
| `01-ascend-allgather-gemm.py` | AllGather 与 GEMM 融合 |
| `02-ascend-gemm-reduce-scatter.py` | GEMM 与 ReduceScatter 融合 |
| `03-ascend-gemm-allreduce-oneshot.py` | GEMM 与 one-shot AllReduce |
| `04-ascend-reverse-all2all` | Reverse All2All |
| `05-ascend-reverse-all2all-barrier.py` | 带 barrier 同步的 Reverse All2All |
| `06-ascend-gemm-ar-notify.py` | 基于 notify 信号的 GEMM AllReduce |
| `07-qkv-alltoall` | attention 的 QKV All2All |
| `08-ascend-transpose-all2all` | 转置 All2All |
| `09-ascend-dispatch-all2all` | MoE dispatch |
| `10-ascend-combine-all2all` | MoE combine |

以相同方式运行任一教程，根据设备数量调整 `--nproc-per-node`：

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

## 运行测试套件

在教程之外，还可在你的硬件上确认原语正常工作：

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

`dist` 标记选择需要多张 NPU 的测试。这些测试覆盖 barrier、wait/notify、对称内存寻址以及
put/get。

## 故障排查

**Kernel 编译错误。** AscendNPU-IR 二进制文件不在 `PATH` 中。按上面所示重新导出。

**`KeyError: 'LOCAL_RANK'`。** 脚本直接用 `python` 运行。这些教程必须通过 `torchrun`
启动。

**启动时挂起。** 通常是前次运行的残留进程占用了会合端口，或 `--nproc-per-node` 设置得
高于可用 NPU 数量。检查残留进程并确认设备数量。

**`aclshmem_init failed`。** shmem 无法设置对称堆。确认 shmem wheel 已安装且当前 shell
已 source CANN。

## 后续步骤

- 阅读 `tutorials/ascend/` 中的教程，了解各通信原语。
- 如果打算提交变更，请参阅[贡献指南](https://gitcode.com/Ascend/Triton-distributed-ascend/blob/master/CONTRIBUTING.md)。