# 全聚合-通用矩阵乘法（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的输出结果一致。