# 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.**

```python
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**

```python
# 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**

```python
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`**

```python
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:

```python
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:

```python
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:

```python
# 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:

```python
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)`:

```python
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:

```python
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:

```python
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

```python
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:

```python
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:

```python
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.**

```python
# 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)**

```python
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)

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