Single Operator Development

This chapter explains how to use regular Triton operators within distributed kernels—how standard operators like tl.load / tl.store / tl.dot combine with dl.* distributed primitives, and what considerations exist for such combinations.

---

1. Core Principle: Remote Pointers Are Just Regular Pointers

Distributed kernels and regular Triton kernels are separated by only one layer: Remote pointers returned by dl.symm_at() are regular Triton pointers in both syntax and semantics.

remote_ptr = dl.symm_at(peer_mem_ptr, target_rank)   # Get remote address
remote_ptrs = remote_ptr + offs_m[:, None] * stride_m + offs_n[None, :] * stride_n
tl.store(remote_ptrs, data, mask=mask)               # Exactly like writing local memory

All regular Triton operators—pointer arithmetic, tl.load, tl.store, tl.atomic_add, tl.dot—use remote pointers exactly the same as local pointers. No special “remote load” operator is needed, nor explicit data movement calls.

This is the foundation of the entire programming model: Cross-card access is expressed as address space extension, not as a new set of communication APIs.

1.1 Three Usage Rules

Rule 1: symm_at can only operate on scalar base addresses

# Correct: first symm_at to get base address, then add offset vector
remote_ptr = dl.symm_at(peer_mem_ptr, target_rank)
remote_ptrs = remote_ptr + offs[:, None] * stride

# Wrong: symm_at doesn't accept block pointer
remote_ptrs = dl.symm_at(peer_mem_ptr + offs[:, None] * stride, target_rank)

The underlying distributed.symm_at operation asserts at IR layer that input must be scalar pointer.

Rule 2: Two equivalent ways to read local symmetric memory

local_ptr = dl.symm_at(peer_mem_ptr, rank)   # Explicit way
local_ptr = peer_mem_ptr                     # Equivalent, recommended

The computation result of symm_at(p, my_rank) is p itself (local base address minus local base address equals 0 offset). Using bare pointer is more concise and one less instruction.

Rule 3: Must go through dl.consume_token when preceded by dl.wait

token = dl.wait(sig_ptr, 1, "npu", "acquire")
data_ptr = dl.consume_token(data_ptr, token)     # Cannot omit
data = tl.load(data_ptr + offs, mask=mask)

See Section 3 for details.

---

2. Usage of Regular Operators in Distributed Kernels

2.1 tl.dot

Completely standard usage, no distributed-related modifications:

accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for block_id_k in range(0, num_k_blocks):
    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)
    accumulator += tl.dot(a, b)
c = accumulator.to(dtype)

The only distributed difference is what the base address of a_ptrs is:

Fusion Mode	a_ptrs Base Address	Description
AllGather + GEMM	peer_mem_ptr	Data already pushed by all peers to local symmetric memory
GEMM + ReduceScatter	a_ptr	Just the local input tensor

Accumulator uniformly uses fp32, finally .to(dtype) to convert back. This is the same convention as single-card GEMM—Cube Engine’s MMA accumulation type is FP32.

Important impact: Whether the kernel contains tl.dot, libshmem_device.barrier_all will implicitly decide launch mode (cv0:1 or cv1:2).

2.2 tl.atomic_add

Core operator for reduction-type communication. Used to accumulate partial results pulled from remote to local output:

c_mask = (offs_cm[:, None] < m_per_rank) & (offs_cn[None, :] < N)
c_temp = tl.load(remote_ptrs)                          # Remote read
tl.atomic_add(c_ptr + c_offs, c_temp, mask=c_mask)     # Local atomic accumulation

atomic_add target should be local memory, source is remote. The reverse (every card atomically adds to the same remote address) causes severe write-side contention.

2.3 tl.load / tl.store

Reading and writing remote pointers is exactly the same as local:

Mask for writing symmetric memory and mask for reading source data are usually different.

Because source data is in input tensor’s global coordinate system, target is in symmetric memory buffer’s coordinate system, and the two coordinate systems have different boundary conditions:

# Read source data: use A's global coordinate system to judge boundary
a = tl.load(a_ptrs, mask=(comm_offs_k[None, :] < K) & comm_msk_m, other=0.0)

# Write symmetric memory: use buffer coordinate system to judge boundary
tl.store(remote_ptrs, a, mask=(comm_offs_k[None, :] < K) & peermem_comm_msk_m)

When mask can be omitted: When data is written to buffer aligned by BLOCK, and buffer itself is also allocated by BLOCK, out-of-bounds cannot occur:

tl.store(peer_mem_ptrs, c)     # Write GEMM result to buffer, no mask needed
...
c_temp = tl.load(remote_ptrs)  # Read from buffer, no mask needed

At this point, the real output boundary is guaranteed by c_mask when finally writing C.

---

3. Connection Between dl.wait and Data Access

3.1 Why consume_token Is Needed

dl.wait returns a token, which must be “injected” into subsequently accessed data pointer through dl.consume_token(ptr, token):

token = dl.wait(sig_ptr, 1, "npu", "acquire")
data_ptr = dl.consume_token(data_ptr, token)
data = tl.load(data_ptr + offs, mask=mask)

The consequence of missing consume_token is not memory reordering, but the entire dl.wait being deleted.

The reason is at IR level: distributed.wait’s memory effect is only Read, if its return value is unused by anyone, this operation is dead code and will be directly eliminated by optimization pass—synchronization completely disappears.

consume_token does nothing at runtime (implementation just returns input as-is, zero overhead), its sole purpose is to establish dependency edge between wait and data access on IR.

---

4. Ascend-Specific Language Extensions

Distributed kernels will use several Ascend platform language extensions, which come from upstream Triton’s Ascend language extension module:

from triton.language.extra.cann.extension import sub_vec_id, sub_vec_num

API	Returns	Description
sub_vec_id()	Vector Core index within this AI Core	Runtime value
sub_vec_num()	Number of Vector Cores per AI Core	constexpr, equals get_aivector_core_num() // get_aicore_num()

Host-side core count query:

from triton.backends.ascend.driver import NPUUtils

aicore_num = NPUUtils().get_aicore_num()
aivec_num  = NPUUtils().get_aivector_core_num()

Don’t hardcode core counts—different product models and variants have different core counts, always query at runtime.

4.1 Typical Usage: Gated Communication Section

subblock_idx = sub_vec_id()
if subblock_idx == 0:
    # Communication section: initiate movement on only one Vector Core
    remote_ptr = dl.symm_at(peer_mem_ptr, target_rank)
    tl.store(remote_ptr + offs, data, mask=mask)

Why only use half the Vector Cores for communication: One AI Core has two Vector Cores, half participating in communication is enough to saturate communication bandwidth, leaving the other for computation is more valuable. This is not wasteful.

---

5. Task Division Conventions

5.1 Grid-Stride Loop

All tile loops in distributed kernels uniformly adopt grid-stride form:

ncore = tl.num_programs(axis=0)
pid   = tl.program_id(axis=0)

for task_idx in range(pid, total_tasks, ncore):
    ...

The benefit is that task count decouples from core count, tail blocks are automatically handled, no extra boundary judgment needed.

5.2 Linearization of Multi-Dimensional Task Numbers

When task space is multi-dimensional (like rank × head × d-block), first linearize then unpack:

for task_idx in range(pid, total_prod_tasks, ncore):
    tmp = task_idx
    rank_loop_id = tmp % rank_size
    tmp //= rank_size
    h_id = tmp % H
    block_id_d = tmp // H
    target_rank = (rank + rank_loop_id) % rank_size

Note the last line: target_rank = (rank + rank_loop_id) % rank_size rather than directly using rank_loop_id. If all cards access targets in order 0, 1, 2, ..., all cards hit the same peer at any moment, causing incast hotspot. After adding own rank for rotation, each card hits different peer at same moment, maximizing link concurrency.

5.3 Task Space Must Be Flattened First

A rule summarized from practice: Don’t distribute cores along only one dimension.

# Wrong: only distribute cores along rank_size dimension: when world_size=2, only 2 cores working
for rank_loop_id in range(logical_core_id, rank_size, n_role_cores):
    ...

# Correct: distribute cores after flattening: all cores have work
total_work = num_blocks_s * total_prod_tasks
for wid in range(logical_core_id, total_work, n_role_cores):
    global_id_s = wid // total_prod_tasks
    task_idx    = wid % total_prod_tasks

---

6. Boundary Handling: Use Mask Not Padding

Distributed kernels uniformly use dynamic mask when handling non-aligned shapes, no padding:

Two-dimensional boundary mask (most common)

offs_s = global_id_s * COMM_BLOCK_S + tl.arange(0, COMM_BLOCK_S)
offs_d = block_id_d * COMM_BLOCK_D + tl.arange(0, COMM_BLOCK_D)
mask = (offs_s < S)[:, None] & (offs_d < D)[None, :]

Row mask broadcast (when column dimension always equals BLOCK, omit column judgment)

offs_s = block_id * BLOCK_S + tl.arange(0, BLOCK_S)
row_mask = offs_s < sequence_length
io_mask = row_mask[:, None]