全聚合-通用矩阵乘法（Allgather Gemm）
在本节中，我们将使用 Triton distribute编写一个全聚合-通用矩阵乘法算子融合的程序。
AllGather + GEMM (Fused Distributed Kernel)
AllGather+GEMM 融合 kernel 将 AllGather 通信原语与 GEMM 计算合并为 Ascend NPU 上的单次 kernel launch，消除了通信阶段与计算阶段之间的中间 DDR 写。在标准的拆分式（decomposed）做法中，AllGather 先将所有 rank 的分片输入矩阵收集到 DDR 中组成完整矩阵，然后由一个独立的 GEMM kernel 从 DDR 读取完整矩阵。而融合做法把 gather 来的数据保留在片上对称内存（symmetric memory）中，直接送入 Cube Engine pipeline，实现 allgather 与 cube 计算的流水线重叠。
Two-Phase Pipeline（两阶段流水线）
kernel 在每个 tile 迭代内遵循两阶段流水线，使用 sub_vec_id() 将通信与计算任务分配到同一 AICore 上的不同 sub-block：
URMA有两条限制（1）单核单pe。（2）整块优先。 Phase 1: Communication (subblock_idx == 1 and pid < world_size)
通过URMA将本地 A sub-block 写入远端对称内存
每个 rank 将自己的 A 分片写入所有其他 rank 的对称内存
for target_rank in range(pid, rank_size, ncore):
if target_rank != rank:
libshmem_device.putmem(
peer_mem_ptr + (buffer_id * buffer_row_size
+ rank * BLOCK_SIZE_M * pvalue) * K, # dst: target_rank 的对称内存, "来自 rank" 的 slot
a_ptr + global_id_m * BLOCK_SIZE_M * pvalue * K, # src: 本地 A 中属于 target_rank 的整段数据
actual_block_size_m * K * dtype.primitive_bitwidth // 8,
target_rank, # pe: putmem 的目标 rank
)
Phase 2: Communication (subblock_idx == 0 or (subblock_idx == 1 and pid >= world_size)
将本地 A sub-block 写入本端对称内存
tl.store(peer_ptr + remote_offset, local_a_data, mask=mask)
Phase 3: Computation (after barrier_all, all cube core)
从本地对称内存读取已 gather 的 A 数据
执行 GEMM: C = A_gathered @ B
在 Cube Engine 上执行 tl.dot(a_block, b_block) 关键点在于： Phase 1 用 UDMA 换取更高的跨卡带宽。Phase 2 和 Consumer 阶段仍用 MTE，因为它们需要细粒度 tile 级并行与双缓冲流水线重叠，这正是 UDMA 的限制所不允许、而 MTE 天然支持的场景。barrier_all() 之后，所有 cube core 参与 GEMM 计算，从各自的本地对称内存 buffer 读取此时已 gather 到位的数据。与此同时，vector 可继续为下一个 tile 处理通信。
Double Buffering（双缓冲）
流水线使用 buffer_num（通常为 2）来跨迭代重叠通信与计算：
迭代 i 使用 buffer_id = global_id % buffer_num 当 vectore core 将迭代 i+1 的数据写入 buffer slot (i+1) % buffer_num 时，所有 cube core 使用 buffer slot i % buffer_num 的数据进行 GEMM 计算 两阶段之间的 barrier_all() 确保 buffer 在被计算读取前已完全写好 这种双缓冲自然地将通信延迟与 Cube Engine 计算重叠，对大 K 维实现近满载利用率。
kernel 使用swizzle 算法来优化内存访问模式：
from triton_dist.language.extra.ascend.algorithm import dist_swizzle2d_Nz, gemm_swizzle2d_Nz
# 计算：确定本次迭代的 GEMM tile 坐标
data_row_idx, data_col_idx = \
gemm_swizzle2d_Nz(iter_id, data_rows, data_cols,
tile_rows, tile_cols, swizzle_offset)
pvalue Parameter（pvalue 参数）
pvalue 参数控制每次通信迭代 gather 多少个 A 的 tile 行。更大的 pvalue 可摊薄每次迭代的 barrier_all() 开销，但会增大对称内存 buffer 的需求：
对称内存布局：[buffer_num, rank_size * BLOCK_M * pvalue, K] M 维出现 rank_size 因子，是因为 AllGather 要从所有 rank 收集数据。每个 rank 贡献 BLOCK_M * pvalue 行，每个 buffer slot 共 rank_size * BLOCK_M * pvalue 行。 每个 buffer slot 中对称内存的总行数：rank_size * BLOCK_M * pvalue
Performance Considerations（性能考量）
Barrier 开销：每次迭代产生一次 barrier_all() 调用。使用更大的 pvalue 可减少迭代数与 barrier 次数，但会增加内存需求。
CV seperation（Cube/Vector 分离）
由于 triton kernel 是 mix kernel（混合 kernel），在 AscendNPU-IR 中会被拆分为 cube func 与 vector func。具体而言，libshmem_device.barrier_all() 是一个 mix core type API，因此拆分 mix kernel 后它会同时存在于 vector func 与 cube func 中，代码如下。
import torch
import torch_npu
import triton
import triton.language as tl
@triton.jit
def kernel_allgather_gemm(
# Pointers to matrices
a_ptr, #本地M*K矩阵
b_ptr, #本地N*K矩阵
c_ptr, #输出矩阵
peer_mem_ptr, #共享内存指针
# Distributed parameters
rank,
rank_size,
buffer_num,
# Matrix dimensions
M,
N,
K,
stride_am,
stride_ak,
stride_bk,
stride_bn,
stride_cm,
stride_cn,
# Meta-parameters
pvalue: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr, #每个程序处理M维度数量
BLOCK_SIZE_N: tl.constexpr, #每个程序处理N维度数量
BLOCK_SIZE_K: tl.constexpr, #每个程序处理K维度数量
COMM_BLOCK_SIZE_M: tl.constexpr, #每个程序通信M维度数量
COMM_BLOCK_SIZE_K: tl.constexpr, #每个程序通信M维度数量
)
dtype = tl.float16
subblock_idx = sub_vec_id()
ncore = tl.num_programs(axis=0)
pid = tl.program_id(axis=0)
num_loops_m = tl.cdiv(M, BLOCK_SIZE_M * pvalue)
num_loops_n = tl.cdiv(N, BLOCK_SIZE_N)
buffer_row_size = BLOCK_SIZE_M * pvalue * rank_size
for global_id_m in range(0, num_loops_m):
#采用双缓冲，buffer_id=0的读取与buffer_id=1的写入overlap
buffer_id = global_id_m % buffer_num
#每个迭代通信M维的数量，pvalue用于后续gemm计算的m维的数量
actual_block_size_m = BLOCK_SIZE_M * pvalue
# 处理尾块
if global_id_m == num_loops_m - 1:
actual_block_size_m = M - global_id_m * BLOCK_SIZE_M * pvalue
num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K)
comm_num_m_blocks = tl.cdiv(actual_block_size_m, COMM_BLOCK_SIZE_M)
comm_num_k_blocks = tl.cdiv(K, COMM_BLOCK_SIZE_K)
if subblock_idx == 1:
#URMA通信性能优于MTE，但URMA不能操作本地UB，因此用vector1的部分核将本地A矩阵传送到其余rank的peer_mem
if pid < rank_size:
for target_rank in range(pid, rank_size, ncore):
if target_rank != rank:
libshmem_device.putmem(
peer_mem_ptr
+ (
buffer_id * buffer_row_size
+ rank * BLOCK_SIZE_M * pvalue
)
* K, # dst: symmetric memory of target_rank, slot from rank.
a_ptr
+ global_id_m
* BLOCK_SIZE_M
* pvalue
* K, # src: Entire data segment of target_rank in local A.
actual_block_size_m * K * dtype.primitive_bitwidth // 8, #urma的优势在于传输大数据量，故不需要分块
target_rank, # target_rank: the target rank of putmem.
)
if subblock_idx == 0 or (subblock_idx == 1 and pid >= rank_size):
#URMA通信性能优于MTE，但URMA不能操作本地UB，因此用vector0与vectore1剩余的核将本地A矩阵写入到当前rank的peer_mem
total_core = ncore + max(ncore - rank_size, 0)
local_pid = pid
if subblock_idx == 1:
local_pid = ncore + pid - rank_size
#将A矩阵切分成COMM_BLOCK_SIZE_M*COMM_BLOCK_SIZE_K
for k in range(
local_pid, comm_num_m_blocks * comm_num_k_blocks, total_core
):
block_id_m = k // comm_num_k_blocks
block_id_k = k % comm_num_k_blocks
comm_row_shape = tl.minimum(
actual_block_size_m - block_id_m * COMM_BLOCK_SIZE_M,
COMM_BLOCK_SIZE_M,
)
remote_ptr = dl.symm_at(peer_mem_ptr, rank)
comm_offs_m = (
tl.arange(0, COMM_BLOCK_SIZE_M)
+ block_id_m * COMM_BLOCK_SIZE_M
+ global_id_m * BLOCK_SIZE_M * pvalue
)
comm_offs_k = (
tl.arange(0, COMM_BLOCK_SIZE_K) + block_id_k * COMM_BLOCK_SIZE_K
)
a_ptrs = a_ptr + (
comm_offs_m[:, None] * stride_am + comm_offs_k[None, :] * stride_ak
)
peermem_comm_offs_m = (
buffer_id * buffer_row_size
+ rank * BLOCK_SIZE_M * pvalue
+ block_id_m * COMM_BLOCK_SIZE_M
+ tl.arange(0, COMM_BLOCK_SIZE_M)
)
remote_ptrs = remote_ptr + (
peermem_comm_offs_m[:, None] * stride_am
+ comm_offs_k[None, :] * stride_ak
)
comm_msk_m = comm_offs_m[:, None] < M
peermem_comm_msk_m = (
peermem_comm_offs_m[:, None]
< buffer_id * buffer_row_size
+ BLOCK_SIZE_M * rank * pvalue
+ block_id_m * COMM_BLOCK_SIZE_M
+ comm_row_shape
)
#从本地A矩阵读出数据
a = tl.load(
a_ptrs, mask=(comm_offs_k[None, :] < K) & comm_msk_m, other=0.0
)
#将读出的数据存入本地对称内存
tl.store(
remote_ptrs, a, mask=(comm_offs_k[None, :] < K) & peermem_comm_msk_m
)
#等待所有通信完成，本地的peer_mem写入了所有rank的A矩阵数组
libshmem_device.barrier_all()
num_tiles_m = tl.cdiv(actual_block_size_m, BLOCK_SIZE_M)
#每次迭代计算[BLOCK_SIZE_M,BLOCK_SIZE_N]矩阵的结果存入C矩阵
for block_id in range(pid, num_tiles_m * num_loops_n * rank_size, ncore):
block_id_m, block_id_n = gemm_swizzle2d_Nz(
block_id,
rank_size * BLOCK_SIZE_M * num_tiles_m,
N,
BLOCK_SIZE_M,
BLOCK_SIZE_N,
)
rank_idx = block_id_m // num_tiles_m
block_id_m = rank_idx * pvalue + (block_id_m % num_tiles_m)
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
matmul_offs_am = (
buffer_id * buffer_row_size
+ block_id_m * BLOCK_SIZE_M
+ tl.arange(0, BLOCK_SIZE_M)
)
matmul_msk_am = matmul_offs_am[:, None] < (
buffer_id * buffer_row_size
+ BLOCK_SIZE_M * rank_idx * pvalue
+ actual_block_size_m
)
offs_bn = block_id_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
msk_n = offs_bn[None, :] < N
#C[BLOCK_SIZE_M,BLOCK_SIZE_N]等于A[BLOCK_SIZE_M, BLOCK_SIZE_K]*B[BLOCK_SIZE_K, BLOCK_SIZE_N]并在K维上累加
for block_id_k in range(0, num_k_blocks):
offs_k = tl.arange(0, BLOCK_SIZE_K) + block_id_k * BLOCK_SIZE_K
a_ptrs = peer_mem_ptr + (
matmul_offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
)
b_ptrs = b_ptr + (
offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
)
a = tl.load(
a_ptrs, mask=(offs_k[None, :] < K) & matmul_msk_am, other=0.0
)
b = tl.load(b_ptrs, mask=(offs_k[:, None] < K) & msk_n, other=0.0)
# 在K维度上累加
accumulator += tl.dot(a, b)
tl.extra.cann.extension.compile_hint(accumulator, "matmul_at_least_once")
c = accumulator.to(dtype)
# -----------------------------------------------------------
# Write back the block of the output matrix C with masks.
offs_cm = (
block_id_m // pvalue * M
+ global_id_m * BLOCK_SIZE_M * pvalue
+ (block_id_m % pvalue) * BLOCK_SIZE_M
+ tl.arange(0, BLOCK_SIZE_M)
)
offs_cn = block_id_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
c_mask = (offs_cm[:, None] < M * (block_id_m // pvalue + 1)) & (
offs_cn[None, :] < N
)
#C[BLOCK_SIZE_M,BLOCK_SIZE_N]结果存到C中
tl.store(c_ptrs, c, mask=c_mask)
因此 barrier_all 会确保 iter i 的通信已完成，随后 cube core 开始 iter i 的计算，与此同时 vector core 继续进入 iter i+1，实现通信与计算的重叠。
创建一个辅助函数用于：
生成M * K的A矩阵与K * N的B矩阵；
创建对称内存
用适当的 grid/block sizes 将上述内核加入队列。
def run_test_distributed():
M = 4096
K = 4096
N = 4096
dtype = torch.float16
#创建A矩阵（A*B)
A_local = torch.randn([M, K], dtype=dtype).npu()
#创建B矩阵（A*B)
B = torch.randn([K, N], dtype=dtype).npu()
# 计算参考值使用torch.distributed.all_gather
C_golden = torch_allgather_gemm(A_local, B, world_size)
# 需要预分配输出。
C = torch.zeros([M * world_size, N], dtype=dtype).npu()
# 启动网格表示并行运行的内核实例的数量。本例是所有cube核数量
ncore = NPUUtils().get_aicore_num()
BLOCK_SIZE_M = 128
BLOCK_SIZE_N = 256
BLOCK_SIZE_K = 256
COMM_BLOCK_SIZE_M = 20
COMM_BLOCK_SIZE_K = 256
buffer_num = 2
pvalue = 4
# 对称内存布局 peer_mem shape: [buffer_num, rank_size * BLOCK_M * pvalue, K]，M 维出现 rank_size 因子，是因为 AllGather 要从所有 rank 收集数据。每个 rank 贡献 BLOCK_M * pvalue 行，每个 buffer slot 共 rank_size * BLOCK_M * pvalue 行
peer_mem_size = (
BLOCK_SIZE_M * pvalue * world_size * buffer_num * max(K, BLOCK_SIZE_K)
)
peer_mem = ash.aclshmem_create_tensor(
[peer_mem_size],
dtype=dtype,
device_id=pe,
)
kernel_allgather_gemm[ncore]( A,
B,
C,
peer_mem,
rank,
rank_size,
buffer_num,
M,
N,
K,
A.stride(0),
A.stride(1),
B.stride(0),
B.stride(1),
C.stride(0),
C.stride(1),
pvalue,
BLOCK_SIZE_M,
BLOCK_SIZE_N,
BLOCK_SIZE_K,
COMM_BLOCK_SIZE_M,
COMM_BLOCK_SIZE_K,)
# 返回 z 的句柄。
return C
使用torch_allgather_gemm函数计算两个A，B矩阵乘，并对比融合算法测试其正确性：
def torch_allgather_gemm(A_local, B, world_size):
"""
- A_local: 本地A矩阵
- B: 本地矩阵
- world_size: rank的总数
Returns: C_golden 是从所有收集到的A矩阵计算得出的
"""
# 创建一个列表来保存所有收集到的A矩阵
A_list = [torch.empty_like(A_local) for _ in range(world_size)]
# 收集所有A矩阵
dist.all_gather(A_list, A_local)
# 将所有A矩阵沿第一维度进行拼接
A_golden = torch.cat(A_list, dim=0)
# 计算参考值: C = A_golden * B
C_golden = torch.matmul(A_golden, B)
return C_golden
比较C与C_golden是否相同
try:
torch.testing.assert_close(C_golden, C, rtol=1e-3, atol=1e-3)
except AssertionError as e:
passed[0] = 0
error_msg = str(e)
raise
# Gather all ranks' pass/fail status
all_passed = [torch.zeros(1, dtype=torch.int32).npu() for _ in range(world_size)]
dist.all_gather(all_passed, passed)
# Print sequentially, one rank at a time
dist.barrier()
for rank_id in range(world_size):
if pe == rank_id:
if all_passed[rank_id].item() == 1:
print(
f"{GREEN}[PASS]{RESET} Rank {pe}: C_golden and C match within tolerances (rtol=1e-3, atol=1e-3).",
flush=True,
)
else:
print(
f"{RED}[FAIL]{RESET} Rank {pe}: C_golden and C do NOT match. Details:\n{error_msg}",
flush=True,
)
dist.barrier()
Out:
[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).
[PASS] Rank 2: C_golden and C match within tolerances (rtol=1e-3, atol=1e-3).
[PASS] Rank 3: C_golden and C match within tolerances (rtol=1e-3, atol=1e-3)
"C_golden and C match within tolerances (rtol=1e-3, atol=1e-3)" 表示融合算法与PyTorch的输出结果一致。