AllGather-GEMM (Allgather Gemm)
In this section, we will use Triton distribute to write a program that fuses the AllGather and general matrix multiplication operators.
AllGather + GEMM (Fused Distributed Kernel)
The AllGather+GEMM fused kernel merges the AllGather communication primitive with GEMM computation into a single kernel launch on Ascend NPU, eliminating the intermediate DDR write between the communication phase and computation phase. In the standard decomposed approach, AllGather first collects all ranks’ sharded input matrices into DDR to form a complete matrix, which is then read from DDR by a separate GEMM kernel. The fused approach keeps the gathered data in on-chip symmetric memory and feeds it directly into the Cube Engine pipeline, achieving pipeline overlap between allgather and cube computation.
Two-Phase Pipeline
The kernel follows a two-phase pipeline within each tile iteration, using sub_vec_id() to assign communication and computation tasks to different sub-blocks on the same AICore:
URMA has two constraints: (1) single core single pe. (2) full block priority. Phase 1: Communication (subblock_idx == 1 and pid < world_size)
Write local A sub-block to remote symmetric memory via URMA
Each rank writes its A shard to all other ranks’ symmetric memory
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's symmetric memory, slot "from rank"
a_ptr + global_id_m * BLOCK_SIZE_M * pvalue * K, # src: entire data segment for target_rank in local A
actual_block_size_m * K * dtype.primitive_bitwidth // 8,
target_rank, # pe: target rank for putmem
)
Phase 2: Communication (subblock_idx == 0 or (subblock_idx == 1 and pid >= world_size)
Write local A sub-block to local symmetric memory
tl.store(peer_ptr + remote_offset, local_a_data, mask=mask)
Phase 3: Computation (after barrier_all, all cube core)
Read gathered A data from local symmetric memory
Execute GEMM: C = A_gathered @ B
Execute tl.dot(a_block, b_block) on Cube Engine The key point is: Phase 1 uses UDMA to gain higher cross-card bandwidth. Phase 2 and the Consumer stage still use MTE because they require fine-grained tile-level parallelism and double-buffered pipeline overlap, which is exactly what UDMA’s limitations don’t allow but MTE naturally supports. After barrier_all(), all cube cores participate in GEMM computation, reading from their respective local symmetric memory buffers the data that has now been gathered. Meanwhile, vector cores can continue handling communication for the next tile.
Double Buffering
The pipeline uses buffer_num (typically 2) to overlap communication and computation across iterations:
Iteration i uses buffer_id = global_id % buffer_num When vector cores write iteration i+1’s data to buffer slot (i+1) % buffer_num, all cube cores use buffer slot i % buffer_num’s data for GEMM computation The barrier_all() between two phases ensures the buffer is fully written before being read for computation This double buffering naturally overlaps communication latency with Cube Engine computation, achieving near full utilization for large K dimensions.
The kernel uses swizzle algorithm to optimize memory access patterns:
from triton_dist.language.extra.ascend.algorithm import dist_swizzle2d_Nz, gemm_swizzle2d_Nz
# Computation: determine the GEMM tile coordinates for this iteration
data_row_idx, data_col_idx = \
gemm_swizzle2d_Nz(iter_id, data_rows, data_cols,
tile_rows, tile_cols, swizzle_offset)
pvalue Parameter
The pvalue parameter controls how many tile rows of A are gathered per communication iteration. Larger pvalue can amortize the barrier_all() overhead per iteration but increases symmetric memory buffer requirements:
Symmetric memory layout: [buffer_num, rank_size * BLOCK_M * pvalue, K] The rank_size factor appears in the M dimension because AllGather collects data from all ranks. Each rank contributes BLOCK_M * pvalue rows, totaling rank_size * BLOCK_M * pvalue rows per buffer slot. Total rows in symmetric memory per buffer slot: rank_size * BLOCK_M * pvalue
Performance Considerations
Barrier overhead: Each iteration incurs one barrier_all() call. Using larger pvalue can reduce the number of iterations and barrier calls, but increases memory requirements.
CV seperation (Cube/Vector Separation)
Since the triton kernel is a mix kernel, it will be split into cube func and vector func in AscendNPU-IR. Specifically, libshmem_device.barrier_all() is a mix core type API, so after splitting the mix kernel it will exist in both vector func and cube func, as shown in the code below.
import torch
import torch_npu
import triton
import triton.language as tl
@triton.jit
def kernel_allgather_gemm(
# Pointers to matrices
a_ptr, # local M*K matrix
b_ptr, # local N*K matrix
c_ptr, # output matrix
peer_mem_ptr, # shared memory pointer
# 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, # amount of M dimension per program
BLOCK_SIZE_N: tl.constexpr, # amount of N dimension per program
BLOCK_SIZE_K: tl.constexpr, # amount of K dimension per program
COMM_BLOCK_SIZE_M: tl.constexpr, # amount of M dimension for communication per program
COMM_BLOCK_SIZE_K: tl.constexpr, # amount of K dimension for communication per program
)
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):
# Using double buffering, reading buffer_id=0 overlaps with writing buffer_id=1
buffer_id = global_id_m % buffer_num
# Amount of M dimension communicated per iteration, pvalue used for subsequent gemm computation's m dimension
actual_block_size_m = BLOCK_SIZE_M * pvalue
# Handle tail block
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 communication performs better than MTE, but URMA cannot operate on local UB, so use part of vector1 cores to transfer local A matrix to peer_mem of other ranks
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's advantage is transferring large data amounts, so no need to split blocks
target_rank, # target_rank: the target rank of putmem.
)
if subblock_idx == 0 or (subblock_idx == 1 and pid >= rank_size):
# URMA communication performs better than MTE, but URMA cannot operate on local UB, so use vector0 and remaining cores of vector1 to write local A matrix to current rank's peer_mem
total_core = ncore + max(ncore - rank_size, 0)
local_pid = pid
if subblock_idx == 1:
local_pid = ncore + pid - rank_size
# Split A matrix into COMM_BLOCK_SIZE_M*COMM_BLOCK_SIZE_K blocks
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
)
# Read data from local A matrix
a = tl.load(
a_ptrs, mask=(comm_offs_k[None, :] < K) & comm_msk_m, other=0.0
)
# Store the read data to local symmetric memory
tl.store(
remote_ptrs, a, mask=(comm_offs_k[None, :] < K) & peermem_comm_msk_m
)
# Wait for all communication to complete, local peer_mem has received A matrix arrays from all ranks
libshmem_device.barrier_all()
num_tiles_m = tl.cdiv(actual_block_size_m, BLOCK_SIZE_M)
# Each iteration computes [BLOCK_SIZE_M,BLOCK_SIZE_N] matrix result and stores to C matrix
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] equals A[BLOCK_SIZE_M, BLOCK_SIZE_K]*B[BLOCK_SIZE_K, BLOCK_SIZE_N] accumulated on K dimension
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)
# Accumulate on K dimension
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
)
# Store C[BLOCK_SIZE_M,BLOCK_SIZE_N] result to C
tl.store(c_ptrs, c, mask=c_mask)
Therefore barrier_all ensures that communication for iter i is complete, after which cube cores begin computation for iter i, while vector cores continue to iter i+1, achieving overlap between communication and computation.
Create a helper function to:
Generate M * K matrix A and K * N matrix B;
Create symmetric memory
Enqueue the above kernel with appropriate grid/block sizes.
def run_test_distributed():
M = 4096
K = 4096
N = 4096
dtype = torch.float16
# Create A matrix (A*B)
A_local = torch.randn([M, K], dtype=dtype).npu()
# Create B matrix (A*B)
B = torch.randn([K, N], dtype=dtype).npu()
# Compute reference value using torch.distributed.all_gather
C_golden = torch_allgather_gemm(A_local, B, world_size)
# Need to pre-allocate output.
C = torch.zeros([M * world_size, N], dtype=dtype).npu()
# Launch grid represents the number of kernel instances running in parallel. In this case, all cube core count
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
# Symmetric memory layout peer_mem shape: [buffer_num, rank_size * BLOCK_M * pvalue, K], the rank_size factor appears in M dimension because AllGather collects data from all ranks. Each rank contributes BLOCK_M * pvalue rows, totaling rank_size * BLOCK_M * pvalue rows per buffer slot
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,)
# Return handle to z.
return C
Use the torch_allgather_gemm function to compute the multiplication of two matrices A and B, and compare with the fused algorithm to test its correctness:
def torch_allgather_gemm(A_local, B, world_size):
"""
- A_local: local A matrix
- B: local matrix
- world_size: total number of ranks
Returns: C_golden computed from all gathered A matrices
"""
# Create a list to hold all gathered A matrices
A_list = [torch.empty_like(A_local) for _ in range(world_size)]
# Gather all A matrices
dist.all_gather(A_list, A_local)
# Concatenate all A matrices along the first dimension
A_golden = torch.cat(A_list, dim=0)
# Compute reference value: C = A_golden * B
C_golden = torch.matmul(A_golden, B)
return C_golden
Compare whether C and C_golden are the same
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)” indicates that the fused algorithm’s output matches PyTorch’s result.