Performance Optimization Guide

1. Optimization Techniques

1.1 Double Buffer Pipeline

Use Case

All fusion operators with “communication phase + computation phase” two-stage structure. buffer_num typically takes 2.

Optimization Principle

Use buffer_id = global_id % buffer_num to rotate among multiple physical buffers: computation of iteration i reads buffer i % 2 while communication side is writing buffer (i+1) % 2, two phases overlap in time. Without double buffering, each phase must serially wait for previous phase to completely finish. double buffer diagram

Implementation Key Points

  • AllGather + GEMM operator: barrier_all() placed between two phases, ensures buffer is written before being read by computation.

  • Coordinating with fine-grained signals: Allocate independent signal slot groups for each buffer. Compared to barrier_all_vec requiring “all Producers of buffer 0 complete before Consumer of buffer 0 can start”, fine-grained signals allow Consumer to start reading when first Producer signal arrives, tighter pipeline.

  • Ring reuse determination: When buffer_num < num_blocks_s, buffers are cyclically reused, at this time must implement credit return (Consumer resets signal after reading), otherwise second round reuse causes Producer to block permanently. This is one of the most common deadlock causes, see Section 2.1.

1.2 Full-Core Utilization

Typical scenario:

  • Already launched all Vector Cores with [vec_num, 1, 1], but small-scale scenario performance poor;

  • Profiling’s Block Num equals vec_num, but actual working cores fewer than vec_num;

  • for task_idx in range(logical_core_id, total_prod_tasks, n_role_cores) where total_prod_tasks is small, causing many logical_core_id to get empty range and exit directly;

  • The smaller total_prod_tasks, the worse performance.

Optimization Principle

Under one-dimensional flattened task distribution, actual participating cores are:

active_cores = min(total_prod_tasks, n_role_cores)
utilization  = active_cores / n_role_cores
Full core condition = total_prod_tasks >= n_role_cores

When total task count is less than role core count, excess cores idle throughout. The problem is not how many cores launched, but assignable task granularity is not fine enough.

Implementation Key Points: S-dimension Flattening For shape [S, H, D], Wrong approach is to flatten one-dimensionally within each S-block separately—at this point each round’s task count is only num_blocks_d * H * rank_size, far less than core count in small H scenarios.

Correct approach is to flatten (global_id_s, task_idx) into a one-dimensional work ID, letting S dimension also participate in parallel distribution:

total_prod_tasks = num_blocks_d * H * rank_size
total_work       = num_blocks_s * total_prod_tasks   # Includes S dimension, far greater than n_role_cores

for wid in range(logical_core_id, total_work, n_role_cores):
    global_id_s = wid // total_prod_tasks            # Unpack S-block
    task_idx    = wid % total_prod_tasks             # Unpack (rank, h, d)
    buffer_id   = global_id_s % buffer_num

“Which dimension to merge” decision:

Reason total_work is small

Dimension to merge

Flattened total_work

Few tasks per single block (small H)

global_id_s

num_blocks_s × total_prod_tasks

Still insufficient (small S and small H)

Further merge h_id

num_blocks_s × rank_size × H × …

Multiple D tiles exist

Merge block_id_d

… × num_blocks_d

Core Occupancy Quantification Example

Taking num_blocks_d = 1, rank_size = 2 (i.e., total_prod_tasks = 2H), n_role_cores = 36 as example:

H

total_prod_tasks = 2H

Working cores (out of 36)

Idle cores

Full core?

8

16

16

20

No

12

24

24

12

No

24

48

36

0

Yes

28

56

36

0

Yes

32

64

36

0

Yes

40

80

36

0

Yes

48

96

36

0

Yes

56

112

36

0

Yes

Above table is calculated core occupancy by formula, not measured time.

1.3 Fine-Grained Signal Synchronization Replacing Global Barrier

This is the optimization with most obvious benefit for pure communication operators, but not applicable to all operators.

Applicability Criterion

Scenario

Applicable?

Reason

Pure communication Producer-Consumer pipeline (Reverse All2All)

Strongly recommended

No Cube computation, fine-grained signals avoid slowest Producer blocking all Consumers

Single rank kernel

Not needed

No cross-card communication

Multi-round loop pipeline

Recommended

Greater benefit in double-buffer scenario

Judgment criterion: If Consumer only needs to wait for specific Producer task to work independently (rather than must wait for all Producers to complete), then applicable fine-grained synchronization.

Optimization Principle

Synchronization Mode

Consumer Wait Time

Total Pipeline Latency

barrier_all_vec

max(t_P0, …, t_P{M-1})

num_blocks_s × max(t_P)

dl.wait / dl.notify

t_Pi (only wait for dependency)

num_blocks_s × avg(t_P + t_C)

Therefore: The more uneven Producer latency distribution (e.g., cross-card bandwidth differences between different ranks), the greater benefit; when latency is completely uniform, the two are close.

Implementation Key Points: Producer 5 Steps / Consumer 6 Steps

Producer:  wait(waitValue=0) → consume_token → tl.store → fence() → notify(target_rank)
Consumer:  wait(acquire)     → consume_token → tl.load  → tl.store → fence() → notify(rank, 0)

Signal value semantic convention:

  • 0 = Consumer has finished reading, Producer can write

  • Non-0 = Producer has finished writing, Consumer can read

Consumer must notify(rank, 0) to reset slot after reading, otherwise next round Producer’s wait(waitValue=0) will block permanently.

Since NPU architecture does not guarantee memory consistency for multi-core concurrent writes to the same cacheline (DataCache) (64B), signal memory layout (flat int64 array, each slot 64 bytes):

signal_mem_ptr
  + buffer_id   * (num_blocks_d * H * rank_size) * 8    # Double-buffer grouping
  + rank_factor * (num_blocks_d * H) * 8                # rank grouping
  + task_idx    * 8                                     # task slot

Platform-Defined Semantics (these are hard facts, don’t re-derive):

Fact

Engineering Implication

libshmem_device.putmem is synchronous

Returned when data has arrived at remote. Under UDMA path, no need for fence before notify

libshmem_device.fence() guarantees completion of preceding load/store

MTE path must strictly follow store → fence() → notify

dl.wait is non-consuming semantics

Only waits, doesn’t reset. Multiple waiters on same slot is safe

Waiter / Resetter Count Constraints

Since dl.wait doesn’t consume signal, constraint falls on reset side:

Slot Type

Resetter Count

Waiter Count

Needs credit return (buffer ring reuse)

Exactly 1

Engineeringly required also exactly 1

No reset (UDMA cross-card, no buffer reuse)

None

Arbitrary many, no constraint

1.4 Swizzle Loop Transformation

Use Case

  • L2 hit rate of distributed GEMM is low;

  • Sequential rank access causes cross-card bandwidth underutilization;

  • All cores simultaneously accessing same L2 cache line produces contention;

  • Naive row-major tile scheduling causes poor K dimension reuse.

Optimization Principle

Swizzle simultaneously solves two orthogonal problems:

  1. L2 Cache Contention—under row-major order, adjacent tiles access non-overlapping rows of B matrix, L2 must evict and reload for each new row.

  2. Cross-Rank Bandwidth Imbalance—under naive scheduling, all cores first access rank 0, then rank 1… At any moment only one inter-card link is active.

Nz pattern interleaves N (column) dimension with Z (rank) dimension while maintaining M dimension locality; and reverses M direction on odd N groups, constructing zigzag (boustrophedon) traversal, doubling temporal locality window.

Two Core Functions

# GEMM phase: optimize tile access order
data_row_idx, data_col_idx = gemm_swizzle2d_Nz(
    iter_id, data_row_shape, data_col_shape,
    tile_row_shape, tile_col_shape, swizzle_offset=7
)

# Communication phase: decide which rank tile goes to
row, col, rank_idx, comm_row_size, comm_col_size = dist_swizzle2d_Nz(
    iter_id, rank_size, data_row_shape, data_col_shape,
    tile_row_shape, tile_col_shape, comm_npu_split=1
)

dist_swizzle2d_Nz internally performs two-level rank permutation:

  • Stride permutation: rank_idx = (rank_idx * rank_stride) % rank_stride + (rank_idx * rank_stride) // rank_size

  • Data-shift permutation: rank_idx = (rank_idx + data_tile_idx) % rank_size

Effect: When rank_size = 8, comm_npu_split = 1, inner 8 iterations simultaneously point to 8 different ranks, 8 inter-card links simultaneously saturated, rather than serially using one by one.

1.5 UDMA Programming

Applicable scope: Target hardware is Ascend 950, and aclshmem initialized with data_op_engine_type = ash.OpEngineType.UDMA.

Use Case

  • Cross-card whole block data movement exists, and movement range is known before kernel launch, can be expressed with single putmem / getmem at once.

Illegal Use Case

  • Movement between local shmem ↔ local GM (must go through MTE path);

Three Hard Constraints

Constraint 1 — Address Type Asymmetry

Interface

dst Requirement

src Requirement

putmem(dst, src, bytes, pe)

Must be remote shmem address

Local address is fine (shmem or regular GM)

getmem(dst, src, bytes, pe)

Must be local address (shmem or GM)

Must be remote shmem address

Note putmem’s dst receives symmetric offset address on remote rank with same pointer value as local, not the mapped pointer returned by dl.symm_at(). This is fundamentally different from MTE link usage, the most error-prone point.

Constraint 2 — Same pe Must Be Initiated by Single Core

Root cause: quiet() interface only retains one signal slot per pe, cannot distinguish multiple concurrent requests.

Constraint 3 — Whole Block Transfer, Avoid Per-Tile Calls

Each UDMA call has inherent initiation and polling overhead. Splitting one large block transfer into multiple small block calls brings additional overhead of multiple calls.

2. Anti-Patterns and Performance Issues

2.1 Deadlock or Signal Out-of-Bounds After Migration from barrier_all to wait/notify

Fine-grained synchronization benefit is considerable, but transformation process often has issues. When troubleshooting, first look at “failure timing”, it locates root cause better than failure phenomenon:

Failure Timing

Suspected Root Cause

All small shapes pass, large shapes hang

Credit return missing (only when num_blocks_s > buffer_num do buffers reuse)

Starts hanging from 2nd round buffer reuse

Credit erased by multiple resetters (lost signal)

First launch normal, errors after multiple launches

Host side missed signal_mem.fill_(0), or zero operation not aligned with dist.barrier()

Deadlock after increasing grid

Causes aicore batch scheduling, Consumer not resident, Producer can’t get credit

Memory corruption

Slot address formula drifts from host-side allocation size

Locating Means

No output in hanging scenarios, need to actively plant debugging points. Host-side watchdog timeout then dump signal’s non-zero distribution:

Signal Slot Stuck Value

Meaning

Stuck at 1

Consumer didn’t wait—slot address calculated wrong, or slot has no waiter at all

Stuck at 0

Producer didn’t notify—notify loop and store loop range mismatch

Hangs after 2nd round reuse slot

Credit return missing, or erased by multiple resetters

Additionally add one line total_vec assertion:

# In kernel
if global_vec_id == 0:
    tl.store(dbg_ptr, total_vec)

# Host side
assert dbg.item() == NPUUtils().get_aivector_core_num(), \
    f"launch mode mismatch: total_vec={dbg.item()}"

Other High-Frequency Issues

  • Missing dl.consume_token—compiler will consider dl.wait has no subsequent dependency, thus eliminate entire dl.wait as dead code. This is not “out-of-order execution”, but synchronization completely disappears.

    # Correct: first all store, after fence then all notify
for task_idx in range(start, total, step):
    tl.store(...)
libshmem_device.fence()
for task_idx in range(start, total, step):   # Three parameters completely same as above
    dl.notify(...)

3. Tuning Case Studies for Various Fusion Operators

3.1 AllGather + GEMM

AllGather + GEMM Operator Practice

3.2 GEMM + ReduceScatter

GEMM + ReduceScatter Operator Practice

3.3 GEMM + AllReduce (One-Shot)

GEMM + AllReduce one-shot Operator Practice

3.4 Reverse All2All

Reverse All2All Operator Practice