性能优化指南

一、优化技术

1.1 Double Buffer 流水

使用场景

所有具备"通信阶段 + 计算阶段"两段结构的融合算子。buffer_num 典型取 2。

优化原理

用 buffer_id = global_id % buffer_num 在多个物理 buffer 间轮转：第 i 次迭代的计算读 buffer i % 2 的同时，通信侧正在写 buffer (i+1) % 2，两个阶段在时间上重叠。没有双缓冲时，每个阶段都必须串行等待前一阶段完全结束。 double buffer示意

实现要点

AllGather + GEMM 算子：barrier_all() 放在两阶段之间，确保 buffer 在被计算读取前已写完。
与细粒度信号协同：为每个 buffer 分配独立的信号槽组。相比 barrier_all_vec 要求"buffer 0 的全部 Producer 完成后，buffer 0 的 Consumer 才能开始"，细粒度信号允许 Consumer 在第一个 Producer 信号到达时就开始读，流水更紧密。
环形复用的判定：当 buffer_num < num_blocks_s 时 buffer 会被循环复用，此时必须实现 credit 回传（Consumer 读完后把信号复位），否则第二轮复用时 Producer 会永久阻塞。这是最常见的死锁原因之一，详见 2.1。

1.2 全核利用（Full-Core Utilization）

典型场景：

已按 [vec_num, 1, 1] 启动了全部 Vector Core，但小规模场景性能差；
profiling 的 Block Num 等于 vec_num，但实际在干活的核数少于 vec_num；
for task_idx in range(logical_core_id, total_prod_tasks, n_role_cores) 中的 total_prod_tasks 偏小，导致大量 logical_core_id 拿到空 range 直接退出；
total_prod_tasks 越小，性能越差。

优化原理

一维展平任务分配下，实际参与工作的核数是：

active_cores = min(total_prod_tasks, n_role_cores)
utilization  = active_cores / n_role_cores
满核条件      = total_prod_tasks >= n_role_cores

当任务总数少于角色核数时，多出来的核全程空转。问题不在于 launch 了多少核，而在于可分配的任务粒度不够细。

实现要点：S维展平 针对shape [S, H, D], 错误做法是在每个 S-block 内部各自一维展平——此时每轮的任务数只有 num_blocks_d * H * rank_size，小 H 场景下远小于核数。

正确做法是把 (global_id_s, task_idx) 展平成一个一维工作 ID，让 S 维也参与并行分配：

total_prod_tasks = num_blocks_d * H * rank_size
total_work       = num_blocks_s * total_prod_tasks   # 含 S 维，远大于 n_role_cores

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

"并入哪一维"的决策：

`total_work` 偏小的原因	应并入的维度	展平后的 `total_work`
单 block 内任务少（小 H）	`global_id_s`	`num_blocks_s × total_prod_tasks`
仍然不够（小 S 且小 H）	再并入 `h_id`	`num_blocks_s × rank_size × H × …`
存在多个 D tile	并入 `block_id_d`	`… × num_blocks_d`

核占用率量化示例

以 num_blocks_d = 1、rank_size = 2（即 total_prod_tasks = 2H）、n_role_cores = 36 为例：

H	`total_prod_tasks` = 2H	干活核数（共 36）	空转核数	是否满核
8	16	16	20	否
12	24	24	12	否
24	48	36	0	是
28	56	36	0	是
32	64	36	0	是
40	80	36	0	是
48	96	36	0	是
56	112	36	0	是

上表是按公式推算的核占用率，不是实测耗时。

1.3 细粒度信号同步替代全局 barrier

这是纯通信类算子收益最明显的一项优化，但并非所有算子都适用。

适用性判据

场景	是否适用	理由
纯通信 Producer-Consumer 流水（Reverse All2All）	强烈推荐	无 Cube 计算，细粒度信号避免最慢的 Producer 阻塞所有 Consumer
单 rank kernel	不需要	无跨卡通信
多轮循环流水线	推荐	双缓冲场景下收益更大

判断标准：如果 Consumer 只需等待特定的 Producer task 就能独立工作（而不是必须等全部 Producer 完成），则适用细粒度同步。

优化原理

同步模式	Consumer 等待时间	总流水延迟
`barrier_all_vec`	`max(t_P0, …, t_P{M-1})`	`num_blocks_s × max(t_P)`
`dl.wait` / `dl.notify`	`t_Pi`（只等自己依赖的那一个）	`num_blocks_s × avg(t_P + t_C)`

因此：Producer 延迟分布越不均匀（例如不同 rank 之间跨卡带宽存在差异），收益越大；延迟完全均匀时两者接近。

实现要点：Producer 5 步 / Consumer 6 步

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)

信号值语义约定：

0 = Consumer 已读完，Producer 可以写入
非 0 = Producer 已写完，Consumer 可以读取

Consumer 读完后必须 notify(rank, 0) 把槽位复位，否则下一轮 Producer 的 wait(waitValue=0) 会永久阻塞。

由于NPU架构不保证多核并发写入同一cacheline（DataCache）（64B）的访存一致性，信号内存布局（flat int64 数组，每槽 64 字节）：

signal_mem_ptr
  + buffer_id   * (num_blocks_d * H * rank_size) * 8    # 双缓冲分组
  + rank_factor * (num_blocks_d * H) * 8                # rank 分组
  + task_idx    * 8                                     # task 槽位

平台既定语义（这些是硬事实，不要重新推导）：

事实	工程含义
`libshmem_device.putmem` 是同步的	调用返回即数据已到达远端。UDMA 路径下 notify 前不需要 fence
`libshmem_device.fence()` 保证前序 load/store 完成	MTE 路径必须严格遵循 `store → fence() → notify`
`dl.wait` 是非消耗语义	只等待、不复位。多个 waiter 等同一个槽是安全的

waiter / 复位者数量约束

由于 dl.wait 不消耗信号，约束落在复位侧：

槽位类型	复位者数量	waiter 数量
需要 credit 回传（buffer 环形复用）	恰好 1 个	工程上要求也恰好 1 个
不复位（UDMA 跨卡、无 buffer 复用）	无	任意多，无约束

1.4 Swizzle 循环变换

使用场景

分布式 GEMM 的 L2 命中率偏低；
顺序 rank 访问导致跨卡带宽未被充分利用；
所有核同时访问同一 L2 cache line 产生 contention；
朴素 row-major tile 调度导致 K 维复用差。

优化原理

Swizzle 同时解决两个正交的问题：

L2 Cache Contention——row-major 顺序下，相邻 tile 访问 B 矩阵的不相交行，L2 对每个新行都要 evict 并重新加载。
跨 Rank 带宽不均衡——朴素调度下所有核先访问 rank 0、再访问 rank 1……任意时刻只有一条卡间链路是活跃的。

Nz 模式把 N（列）维与 Z（rank）维交织，同时保持 M 维的局部性；并且在奇数 N group 上反转 M 方向，构成 zigzag（boustrophedon）遍历，使时间局部性窗口翻倍。

两个核心函数

# GEMM 阶段：优化 tile 访问顺序
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
)

# 通信阶段：决定 tile 发往哪个 rank
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 内部做了两级 rank 置换：

Stride 置换：rank_idx = (rank_idx * rank_stride) % rank_stride + (rank_idx * rank_stride) // rank_size
Data-shift 置换：rank_idx = (rank_idx + data_tile_idx) % rank_size

效果是：rank_size = 8、comm_npu_split = 1 时，内层 8 次迭代同时指向 8 个不同的 rank，8 条卡间链路同时饱和，而不是逐条串行使用。

1.5 UDMA 编程

适用范围：目标硬件为 Ascend 950，且 aclshmem 以 data_op_engine_type = ash.OpEngineType.UDMA 初始化。

使用场景

存在跨卡的整块数据搬运，且搬运范围在 kernel 启动前已知，可以用单次 putmem / getmem 一次性表达。

非法使用场景

本地 shmem ↔ 本地 GM 之间的搬运（必须走 MTE 路径）；

三条硬约束

约束 1 — 地址类型不对称

接口	dst 要求	src 要求
`putmem(dst, src, bytes, pe)`	必须是远端 shmem 地址	本地地址即可（shmem 或普通 GM）
`getmem(dst, src, bytes, pe)`	必须是本地地址（shmem 或 GM）	必须是远端 shmem 地址

注意 putmem 的 dst 传的是远端 rank 上与本地指针数值相同的对称偏移地址，不是 dl.symm_at() 返回的映射指针。这与 MTE 链路的写法根本不同，是最容易搞错的一点。

约束 2 — 同一个 pe 必须由单核发起

根本原因：quiet() 接口对每个 pe 只保留了一个 signal 槽位，无法区分多个并发请求。

约束 3 — 整块传输，避免逐 tile 调用

每次 UDMA 调用都有固有的发起与轮询开销。把一次大块传输拆成多次小块调用，会带来多次调用的额外开销。

二、反模式与性能问题

2.1 从 barrier_all 迁移到 wait/notify 后死锁或信号越界

细粒度同步收益可观，但改造过程往往存在问题。排查时先看"失效时机"，它比失效现象更能定位根因：

失效时机	优先怀疑的根因
小 shape 全部通过，大 shape 死等	credit 回传缺失（只有 `num_blocks_s > buffer_num` 时 buffer 才会复用）
第 2 轮 buffer 复用起开始死等	credit 被多个复位者擦掉（lost signal）
首次 launch 正常，多次 launch 后出错	host 侧漏了 `signal_mem.fill_(0)`，或清零操作未用 `dist.barrier()` 对齐
grid 调大后反而死锁	导致aicore分批调度，Consumer 未驻留，Producer 等不到 credit
踩内存	槽地址公式与 host 侧分配大小漂移

定位手段

死等场景下没有任何输出，需要主动埋点。Host 侧看门狗超时后 dump 信号量的非零分布：

信号槽停留值	含义
停在 1	Consumer 没有 wait 到——槽地址算错，或该槽根本没有 waiter
停在 0	Producer 没有 notify——notify 循环与 store 循环的范围不一致
复用槽第 2 轮后才挂	credit 回传缺失，或被多个复位者擦掉

另外加一行 total_vec 断言：

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

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

其他高发问题

漏掉 dl.consume_token——编译器会认定 dl.wait 没有后续依赖，从而把整条 dl.wait 当作死代码消除掉。这不是"乱序执行"，而是同步彻底消失。

# 正确：先全部 store，fence 之后再全部 notify
for task_idx in range(start, total, step):
    tl.store(...)
libshmem_device.fence()
for task_idx in range(start, total, step):   # 三个参数与上面完全一致
    dl.notify(...)

三、各融合算子调优案例

性能优化指南

一、优化技术

1.1 Double Buffer 流水

1.2 全核利用（Full-Core Utilization）

1.3 细粒度信号同步替代全局 barrier

1.4 Swizzle 循环变换

1.5 UDMA 编程

二、反模式与性能问题

2.1 从 barrier_all 迁移到 wait/notify 后死锁或信号越界

三、各融合算子调优案例

3.1 AllGather + GEMM

3.2 GEMM + ReduceScatter

3.3 GEMM + AllReduce（One-Shot）

3.4 Reverse All2All