分布式架构设计

术语缩写表

编译与 IR

缩写	英文全称	中文说明
IR	Intermediate Representation	中间表示
AST	Abstract Syntax Tree	抽象语法树
MLIR	Multi-Level Intermediate Representation	多级中间表示，LLVM 生态下可扩展的编译器基础设施，Distributed Dialect / HIVM Dialect 均构建于其上
TTIR	Triton IR	Triton 方言层面的中间表示，也是编译流程中一个 stage 的名字
LLVM	（项目专名，不再对应缩写展开）	GPU 路径的下沉目标 IR / 后端基础设施
HIVM	Ascend 后端核心 IR 方言名	`AscendNPU-IR`/`bishengir` 代码中作为专有名词使用，暂未见官方缩写全称展开
DPS	Destination-Passing Style	目标传递风格：通过预先分配的 init/输出值来传递计算结果的 IR 设计模式
CSE	Common Subexpression Elimination	公共子表达式消除
LICM	Loop-Invariant Code Motion	循环不变代码外提
`TT_*`（如 `TT_Type`/`TT_Ptr`/`TT_IntLike`）	Triton Type	TableGen 中 Triton 方言类型约束的前缀

硬件与执行单元

缩写	英文全称	中文说明
NPU	Neural Processing Unit	神经网络处理器，本文指华为昇腾 AI 处理器
GPU	Graphics Processing Unit	图形处理器
GM	Global Memory	全局内存，Device 侧全局可寻址内存
UB	Unified Buffer	统一缓冲区，AICore 片上临时缓冲区
AIC	AI Cube Core	矩阵单元，负责矩阵运算
AIV	AI Vector Core	向量单元，负责向量计算与低时延通信
PE	Processing Element	处理单元，对应一个 rank / 进程 / 设备参与方

通信引擎与协议

缩写	英文全称	中文说明
SHMEM	Symmetric Hierarchical Memory / Shared Memory	对称内存通信库，本文指 `shmem`/`aclshmem` 对称堆通信库
ACLSHMEM	ACL（AscendCL）+ SHMEM	本文对称堆通信接口/符号的前缀命名
MTE	Memory Transfer Engine	内存传输引擎，昇腾 AICore 侧数据搬运引擎
SDMA	System Direct Memory Access	系统直接内存访问
UDMA	UB Direct Memory Access	统一直接内存访问
RDMA	Remote Direct Memory Access	远程直接内存访问
RoCE / ROCE	RDMA over Converged Ethernet	基于融合以太网的 RDMA
TLS	Transport Layer Security	传输层安全协议
QP	Queue Pair	队列对
SQ	Send Queue	发送队列
SQE	Send Queue Element	发送队列元素
WQE	Work Queue Element	工作队列元素
SGE	Scatter/Gather Element	分散/聚合元素

其他

缩写	英文全称	中文说明
SIMD	Single Instruction Multiple Data	单指令多数据
SDK	Software Development Kit	软件开发工具包

本章节介绍 Triton-distributed-ascend 在 Ascend 平台上的整体架构：从 Python 分布式原语到 NPU 硬件执行，中间经历了哪些 IR 层次、由哪些组件负责、以及对称内存运行时如何支撑跨卡通信。

一、整体架构

1.1 分层设计

整体架构图

核心设计理念：Distributed Dialect 是平台无关的中间表示。各平台通过不同的 Conversion Pass 将其映射到各自的通信库，业务侧的 dl.* 代码不需要因平台而改写。

Ascend 平台的落点是 hivm.custom 操作——Distributed Dialect 的每个操作被转换为一条带符号名的 custom op，符号在设备端模板库中静态存在，由后端在链接期解析。编译期结束后不需要任何运行时的 module patch，这与 A/B 路径有本质区别（见 6.3）。

1.2 与 GPU 路径的架构差异

	GPU	Ascend
通信库符号如何绑定	编译出二进制后，host 侧把 SHMEM 上下文指针 patch 进 device module	编译期生成 `hivm.custom` 符号，链接期静态解析
device 端如何找到通信状态	通过被 patch 进去的全局变量	位于架构固定的 GM 地址，无需传参
运行时初始化开销	每次加载 module 都要 patch	无

这个差异决定了 Ascend 上 dl.rank() 这类操作可以做成无额外参数的纯函数——设备状态的地址是编译期已知的常量。

二、Distributed Dialect

2.1 Dialect 定义

Distributed Dialect 的 dependentDialects 被刻意置空，不依赖任何后端方言。这是它能同时下沉到 LLVM（GPU 路径）和 HIVM（Ascend 路径）的前提。

Dialect 名：distributed
C++ namespace：::mlir::triton::distributed

2.2 七个分布式op

Op	operands	results	traits
`distributed.wait`	`barrierPtr`(TT_PtrLike), `numBarriers`(TT_IntLike), `waitValue`(TT_Int), `scope`, `semantic`	`token`(TT_IntLike)	`MemoryEffectsOpInterface`, `TypesMatchWith`
`distributed.consume_token`	`input`(TT_Type 或 TT_TensorDescType), `token`(TT_IntLike)	同 `input` 类型	`Elementwise`, `MemoryEffects`, `InferTypeOpInterface`
`distributed.get_rank`	`axis`(I32)	`result`(I32)	`Pure`
`distributed.get_num_ranks`	`axis`(I32)	`result`(I32)	`Pure`
`distributed.symm_at`	`symmAddr`(TT_Ptr), `rank`(I32)	`remoteAddr`(TT_Ptr，同类型)	`MemoryEffects`, `TypesMatchWith`
`distributed.notify`	`sigAddr`(TT_Ptr), `signalVal`(I64), `rank`(I32), `sigOp`, `commScope`	无	`MemoryEffects`
`distributed.extern_call`	`srcs`(Variadic), `libname`/`libpath`/`symbol`(StrAttr), `pure`(BoolAttr)	`result`(Variadic)	`MemoryEffects`, `ConditionallySpeculatable`

2.3 属性枚举

SignalOp：

值	名称	语义
1	SET	将信号值设置为指定值
2	ADD	将指定值加到当前信号值

2.4 内存副作用模型

副作用不在 TableGen 中声明，而是在 C++ 侧实现 MemoryEffectsOpInterface：

Op	副作用
`WaitOp`	只有 `Read`
`ConsumeTokenOp`	空——纯粹的 IR 依赖锚点
`SymmAtOp`	只有 `Read`
`NotifyOp`	`Read` + `Write`
`ExternCallOp`	`pure=true` 时无副作用，否则 `Write` + `Read`

三、ConvertTritonDistributedToHIVM Pass

3.1 Pass 概览

Pass 名：convert-triton-distributed-to-hivm
作用域：ModuleOp
依赖方言：hivm::HIVMDialect、triton::distributed::DistributedDialect

3.2 驱动逻辑

Pass 的 runOnOperation() 做四件事：

扫描是否存在 triton::DotOp / DotScaledOp，得到 existDot 标志。没有 tl.dot 就说明这是纯 AIV kernel，后续 core type 推导要用到这个信息。
注册七个distributed op的重写模式。
调用 applyPatternsAndFoldGreedily ，其余 Triton IR 要原样保留给下游的单卡编译阶段。

3.3 aclshmem接口OP生成

转换模板类为每个操作生成一个符号名：

源 op	对应aclshmem接口名
`SymmAtOp`	`aclshmem_ptr_<typename>`，如 `aclshmem_ptr_float`
`GetRankOp`	`aclshmem_my_pe`
`GetNumRanksOp`	`aclshmem_n_pes`
`NotifyOp`（i32 信号）	`aclshmemx_signal_op`
`NotifyOp`（ui64 / i64 信号）	`aclshmem_uint64_p` / `aclshmem_int64_p`
`ConsumeTokenOp`	`aclshmem_consume_token_<typename>`，如 `aclshmem_consume_token_float_ptr_1d`
`WaitOp`	`aclshmem_wait_<typename>`，如 `aclshmem_wait_int32`
`ExternCallOp`	沿用 op 自带的 `symbol` 属性

3.4 属性设置

TCoreType（决定这条操作在哪类核上执行）：

推导顺序是：先查前缀映射表，未命中则按 existDot 决定。

符号前缀	Core Type
`aclshmem_barrier_all`	CUBE_AND_VECTOR
`aclshmemx_barrier_all_vec`	VECTOR
`aclshmem_putmem` / `getmem` / `putmem_nbi` / `getmem_nbi` / `putmem_signal` / `putmem_signal_nbi`	VECTOR
其他（未命中前缀表）	`existDot ? CUBE_AND_VECTOR : VECTOR`

aclshmem_barrier_all 被标为 CUBE_AND_VECTOR 这一点很重要：它意味着 barrier 在 kernel 被拆分为 cube func 与 vector func 之后同时存在于两个 func 中，这正是计算-通信重叠得以成立的机制。

其余属性：

属性	取值	说明
`PIPE`	`PIPE_S`	见下方说明
`VFMode`	`SIMD`	无条件设置
`hivm.is_distributed`	UnitAttr	下游标记，供内存作用域推导、块指针分析、mix kernel 拆分等多个后续阶段识别分布式操作
`symbol`	StrAttr	裸符号名
`no_side_effect`	UnitAttr（条件）	仅当源 op 的内存副作用为空时设置
`gm_addr_args_indices`	DenseI32Array	收集「既是指针类型、又是 `tt.func` 入口块参数」的操作数下标。张量指针与 `tt.addptr` 的结果都不计入

关于 PIPE 的现状说明：符号到流水线的映射表在代码中已声明但当前为空，因此所有 distributed 操作实际都落到 PIPE_S（标量流水）。这是一个预留的扩展点，而非"按操作类型分配流水线"的既成设计。阅读代码或做性能分析时请以此为准。

3.5 结果构造与替换

对每个 RankedTensorType 结果，先创建 tensor::EmptyOp 作为 DPS（destination-passing style）的 init 值，再构造 hivm::CustomOp。最后：无结果则 eraseOp，有结果则 replaceOp。

四、编译链路

4.1 完整链路与各阶段入口

  Python AST → Triton IR（用 DistributedOpBuilder 生成 distributed.* 操作）
  │     builder 绑定：create_distributed_wait / consume_token / get_rank /
  │                   get_num_ranks / symm_at / notify / extern_call
  ↓
【stage "ttir"】
        标准 Triton 优化：inliner / combine / canonicalize / cse / licm / loop_unroll
  ↓
【stage "ttadapter"】     ← ★ 本项目的唯一插入点 ★
        add_convert_triton_distributed_to_hivm(pm)
        —— 在单卡编译的所有后续 pass 之前执行，
           distributed 操作在进入 TTIR→Linalg 转换前就已变成 hivm.custom
  ↓
【stage "npubin"】
        单卡 NPU 编译与二进制生成
        （TTIR → Linalg IR → AscendNPU IR → 机器码，
          由三方依赖 triton-ascend 完成，本文不展开其内部实现）
  ↓
  Ascend NPU 可执行 kernel

关键点：本项目对编译流程的介入只有一处——在 ttadapter 阶段的最开始插入 ConvertTritonDistributedToHIVM Pass。此后所有 distributed 操作都已变成 hivm.custom，下游的单卡编译能力由三方依赖 triton-ascend 提供。

distributed 模块在编译器中是可选依赖——通过 try / except ImportError 加载，未构建分布式支持时自动跳过该 Pass。

4.2 符号如何变成真实代码（AscendNPU-IR）

这是理解整个架构的最后一环：Pass 只生成了一个字符串符号名，它是怎么变成可执行代码的？

第一步：符号名 mangle

std::string prefix = concreteOp.getSymbol();
if (!hasMemrefInArgOrRet()) { prefix = "_mlir_ciface_" + prefix; }
return prefix + callNameMangleSuffix(op);

即 symbol = "aclshmem_my_pe" → 链接名 _mlir_ciface_aclshmem_my_pe。

第二步：模板库中的静态实现

Pass 生成的每一个符号，在设备端模板库中都有对应的 _mlir_ciface_ 实现。几个代表性的例子：

_mlir_ciface_aclshmem_my_pe / _mlir_ciface_aclshmem_n_pes —— 直接读设备状态
ACLSHMEM_PTR_WRAPPER 宏 —— 展开出 12 种类型的 aclshmem_ptr_<T>，把 aclshmem_ptr() 的结果包回 memref 结构
ACLSHMEM_P_WRAPPER 宏 —— 展开出 aclshmem_int64_p / uint64_p 等
ACLSHMEM_WAIT_WRAPPER 宏 —— 展开出 aclshmem_wait_int32/64 等

dl.wait 由于NPU架构不保证多核并发写入同一个cacheline（DataCache）的访存一致性（64B），遍历多个 barrier 时按 64 字节硬编码跨步。这就是应用层 SIGNAL_SLOT_STRIDE = 64 / sizeof(dtype) 的由来——int32 信号是 16 个元素，int64 信号是 8 个元素。

CONSUME_TOKEN_*_WRAPPER 的实现体则就是把输入原样返回，是给编译器制造数据依赖的屏障，无任何运行时开销。

第三步：编译为 bitcode 并链接

模板库源码会被编译为四个变体（aic / aiv / mix_aiv / mix_aic）的 bitcode，在最终生成 kernel 二进制时链入。

这就是"为什么 Pass 只需要生成一个字符串符号名"的完整答案：符号在模板库里静态存在，由后端在链接期解析。

五、对称堆运行时

对称堆：所有 rank 上大小、相对布局一致的一段设备内存，堆内偏移在各 rank 间镜像对齐。对称张量：从对称堆分配出的普通 torch 张量，可直接参与运算和传参。二者的价值在于——任意 rank 都能用"本地偏移 + 目标 PE 堆基址"算出该张量在其他 rank 上的地址（见 5.3 symm_at），无需运行时地址交换即可跨卡读写。

5.1 Host 端初始化

Ascend 的对称堆初始化由用户代码直接调用 shmem Python 包完成（惯例 import shmem as ash），分五步：

ash.set_conf_store_tls(False, "")                                  # 1. TLS 配置
attributes = ash.InitAttr()                                        # 2. 构造初始化属性
attributes.my_rank = rank
attributes.n_ranks = world_size
attributes.local_mem_size = 1024 * 1024 * 1024                     # 3. 对称堆大小
attributes.ip_port = "tcp://127.0.0.1:8666"                        #    bootstrap 地址
attributes.option_attr.data_op_engine_type = ash.OpEngineType.MTE  # 4. 传输引擎
ash.aclshmem_init(attributes)                                      # 5. 初始化

传输引擎枚举：MTE / SDMA / ROCE / UDMA，对应四条不同的数据通路。

5.2 对称张量分配

peer_mem = ash.aclshmem_create_tensor(shape, dtype=torch.float16, device_id=rank)
...
ash.aclshmem_free_tensor(peer_mem)

内部流程是 calc_nbytes → aclshmem_malloc → construct_tensor_from_ptr，返回一个正常的 torch 张量，可以直接参与 PyTorch 运算与 kernel 传参。

约束：aclshmem_malloc / aclshmem_free 必须在所有进程同步调用，且分配或释放相同大小的内存。

5.3 `symm_at` 的远端地址解析（核心机制）

答案：查表 + 偏移算术的混合方案，两者缺一不可。

ACLSHMEM_DEVICE __gm__ void *aclshmem_ptr(__gm__ void *ptr, int pe)
{
    __gm__ aclshmem_device_host_state_t *device_state = aclshmemi_get_state();
    ptrdiff_t offset = (uintptr_t)ptr - (uintptr_t)device_state->heap_base;
    uintptr_t remote_ptr = (uintptr_t)device_state->p2p_device_heap_base[pe] + offset;
    return (__gm__ void *)remote_ptr;
}

三个步骤：

取设备状态；
用本地指针减去本地堆基址，回退到"对称偏移"；
加上目标 PE 的堆基址，得到远端地址。

设备状态结构中有三套并行的基址表——p2p_device_heap_base、rdma_device_heap_base、sdma_device_heap_base，运行时按目标 PE 的传输方式位掩码选择使用哪一套。

六、通信原语的运行时实现

6.1 机制总表

操作	底层机制
signal SET（P2P / MTE）	远端 GM 普通标量写 + cacheline 回写指令，非硬件原子、非 doorbell
signal ADD（P2P）	硬件原子加 + UB→GM DataCopy，在写端做 read-modify-write
signal（跨节点 ROCE）	RDMA write + quiet，或 RDMA atomic
wait / wait_until	本地自旋忙等，每轮先失效 cacheline 再读
barrier	软件算法（集中式 pull），由 signal_set + spin 组合而成
UDMA	真 doorbell：填 WQE/SGE 到 SQ ring → 写doorbell寄存器

6.2 signal / notify

SET 路径的实现是两步：

普通标量存储 *addr = val；
dcci_cacheline(addr) —— 把数据缓存刷回 GM，注释直言其目的是 "flush data cache to GM after signal to ensure it is visible to other ranks"。

ADD 路径则要走一遍 UB：置值到 UB → 设置/等待 MTE3 标志 → 开启硬件原子加模式 → copy_ub2gm → 关闭原子模式。

aclshmem_<T>_p（点写）的机制与 signal SET 同构：aclshmem_ptr 求地址 + 标量存储 + cacheline 回写。

6.3 wait 的轮询机制

六个比较谓词（_eq / _ne / _gt / _ge / _lt / _le）形态完全一致：

do { dcci_cacheline((__gm__ uint8_t *)sig_addr); } while (*sig_addr != cmp_val);

没有硬件阻塞或中断机制，正确性完全依赖每轮循环打掉 L2 cacheline。dcci_cacheline 内部用 DataCacheCleanAndInvalid，前后各夹一条空的 __asm__ __volatile__("") 防止被编译器优化掉。

这也解释了细粒度同步的一个特性：dl.wait 是非消耗语义——它只轮询、不修改信号值，因此多个 waiter 等同一个槽是安全的。约束落在复位侧（谁负责把信号写回 0）。

6.4 barrier_all 的实现

aclshmem_barrier_all() 最终落到一个集中式 pull barrier：

每个 PE 只写自己的 flag（signal_set）；
然后通过 aclshmem_ptr(sync_pool, remote_pe) 直接自旋读对端的 flag——不需要对方主动 push；
工作按可用 Vector Core 数分摊，复杂度 O(N/K)。

一个防死锁的关键细节：barrier 专用的等待谓词除了接受 == cmp_val，还接受 == cmp_val + 1，注释说明是 "in case when peer pe enters next barrier"——防止快的 PE 已经进入下一轮 barrier 导致慢的 PE 永远等不到。

6.5 UDMA

UDMA 不是一组新的 dl.* 操作，而是 host 初始化时选定的传输引擎。 同一份 libshmem_device.putmem() 调用，在 MTE 与 UDMA 两种配置下走的是完全不同的底层路径。

UDMA 的设备端发送流程是真正的 doorbell 机制：

QP fetch → slot 计算 → 填 SQE → 填 SGE → dcci_cachelines → 敲doorbell寄存器

七、跨平台分发机制

7.1 分发结构

dl.wait / dl.notify / dl.symm_at / dl.rank / dl.num_ranks / dl.consume_token
   → 静态，无 proxy → _builder.create_* → distributed 方言 → (Ascend) HIVM Pass
   ※ 唯一的后端条件分支：dl.notify 的 dtype 校验

libshmem_device.<f>  → ModuleProxy 动态分发
   → (Ascend) libaclshmem_device.<f> → extern_call → ExternCallOp → aclshmem_* 符号

核心的 dl.* 原语走静态路径——它们直接创建 MLIR 操作，平台差异完全由 Conversion Pass 承担，Python 层没有分支。这是跨平台可移植性的基础。

libshmem_device.* 走动态路径——因为不同平台的 SHMEM 库函数集不完全相同，需要运行时按平台选择模块。

7.2 ModuleProxy

分发的唯一原语是 ModuleProxy，逻辑很简单：

构造时对 [(predicate, module), ...] 逐一求值，断言恰好有一个 predicate 为真；
__getattr__ 把属性访问透传到选中的模块；
dispatch 装饰器让被装饰函数在调用时才按 func.__name__ 去 active module 查找同名函数。

7.3 后端信息获取

后端判定全部基于 shutil.which(<厂商 smi 工具>) ：

函数	判据
`is_ascend()`	`shutil.which("npu-smi")`
`is_cuda()`	`shutil.which("nvidia-smi")`
`is_hip()`	`shutil.which("rocm-smi")`
`is_maca()`	`shutil.which("mx-smi")`

7.4 Ascend 分支清单

位置	Ascend 分支做了什么
`libshmem_device.py`	把 `ascend.libaclshmem_device` 加入 proxy 列表
`language_extra.py`	把 `ascend.language_extra` 加入 proxy 列表（当前为空模块）
`distributed_ops.py`	`dl.notify` dtype 校验：Ascend 允许 int32（set + add），int64/uint64 仅 set；其他平台只允许 int64/uint64
`utils.py`	导入期 Ascend 分支为空，不导入任何厂商 SDK
编译器	在 `ttadapter` 阶段插入 HIVM Pass
pybind 层	构建 Ascend 版本时挂载 `ascend_passes` 子模块