快速开始

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

配置运行环境

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

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 教程

在仓库根目录执行：

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

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

两 rank 运行的预期输出：

[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 提供。

备注

示例使用 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：

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

运行测试套件

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

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/ 中的教程，了解各通信原语。

  • 如果打算提交变更，请参阅贡献指南。