从PyTorch到AscendC:借助KernelSwift拿下昇腾算子实操工坊2w+现金激励
作者:ZMZ,PZL from DeepLink Group @Shanghai AI Lab
从一段 PyTorch 算子语义描述,到昇腾硬件上的高性能实现,距离有多远?
过去,这条链路要靠手工走完:拆解语义、手写 Kernel、设计 dispatch、对照 Profile 反复调优,往往要以周为单位计算成本。
本文记录一次团队代表参与昇腾社区算子实操工坊的真实提交经历:以 CHEMM(单精度复数 Hermitian 矩阵乘)算子开发为例,使用 KernelSwift 从一段 PyTorch 语义描述出发,自动生成可评测的 AscendC 实现,再按九类场景做 dispatch 调优,1000 项精度、200 项性能全部通过,并经过昇腾社区算子实操工坊任务验收,最终斩获2w+任务激励。

KernelSwift 是上海AI实验室 DeepLink 团队推出的一款智能算子迁移系统,用 AI 自动化能力解决多元芯片的算子适配难题,让高性能算子开发从 “手工编码” 走向 “一键生成”。开发者可以先用清晰的参考实现表达计算语义,再让平台围绕目标硬件迭代候选实现,这能够大幅缩短从需求到可评测代码的开发周期。
文章将分享参赛者从算子语义到高性能实现的完整链路,并说明在过程中如何验证生成代码的正确性和性能,最后总结调优经验。
CHEMM 算子功能及实现难点
先看这次的提交对象。CHEMM 是 BLAS Level 3 中的单精度复数 Hermitian 矩阵乘法算子,执行:

其中, A A A 满足共轭对称关系 A = A H A=A^H A=AH,需读取指定的上三角或下三角,对角线虚部按零处理; B 、 C B、C B、C 为一般复数矩阵, α 、 β \alpha、\beta α、β 为复数标量。接口定义
它适用于 Hermitian 矩阵作用于一组复数向量的计算,例如量子计算中的哈密顿矩阵与多状态相乘、信号处理中的复数协方差矩阵与滤波权重相乘,以及科学计算中的分块迭代求解。CHEMM 是这些计算中的基础运算,并不单独完成特征值求解或矩阵分解。
高性能实现主要有四类难点:
-
结构化访存。 一个三角直接读取,另一个三角需要转置索引和共轭处理;还要兼容左右乘、列主序、非紧凑 leading dimension 和尾块。
-
复数计算映射。 需要将复数乘法转化为硬件擅长的实数矩阵运算,同时控制数据重排和中间结果的开销。
-
性能与精度平衡。 要借助 3M 方法以三个实数矩阵乘积重建复数结果,减少乘积数量,但加减重建可能放大数值抵消误差,因此还需要风险检测及补偿重算。
-
不同尺寸的调度。 大矩阵关注计算流水和数据复用,小矩阵、窄矩阵则可能受 kernel 启动、pack、combine 和同步开销主导,需要按形状选择策略。
参赛者的实现过程分为三个阶段,先根据任务要求对待实现算子的不同场景进行分类并构造特定的例子,然后将对应的 pytorch 实现提交到 KernelSwift 获取到优化后的 kernel 实现,接下来对实现的各个kernel拼接起来,提交到昇腾社区等待测试与验收。
实现的最终算子效果如下:
| 项目 | 结果 |
|---|---|
| 精度与接口用例 | 1000/1000 通过 |
| 性能用例 | 200/200 通过 |
| GPU_us / NPU_us 最小值 | 0.407995 |
它是怎么一步步跑出来的?下文将详细介绍平台使用流程及调优经验。
KernelSwift 操作流程
- 进入平台 https://deeplink.org.cn/kernelswift ,完成登录,点击左侧“新建任务”。

- 在 model.py 编辑器中提供 class Model、get_init_inputs() 和 get_inputs()。Model.forward 定义参考计算;get_inputs 给出可复现的测试张量;get_init_inputs 返回构造参数。下列代码与本次提交的计算语义一致,可直接作为 model.py:
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super().__init__()
def forward(self, a, b, c, alpha, beta):
lower = torch.tril(a, diagonal=-1)
h = lower + lower.conj().transpose(-2, -1)
h = h + torch.diag_embed(a.diagonal().real).to(a.dtype)
return alpha * torch.matmul(h, b) + beta * c
def get_inputs():
torch.manual_seed(42)
m, n = 128, 64
a = torch.randn(m, m, dtype=torch.complex64)
b = torch.randn(m, n, dtype=torch.complex64)
c = torch.randn(m, n, dtype=torch.complex64)
alpha = torch.tensor(0.75 + 0.25j, dtype=torch.complex64)
beta = torch.tensor(-0.5 + 0.125j, dtype=torch.complex64)
return [a, b, c, alpha, beta]
def get_init_inputs():
return []
- 填写任务名,这里使用的例子名称是 chemm;生成模式选择“生成 Clike 算子”;适配硬件选择“华为昇腾”。必须核对模式,避免误用默认的 Triton 选项。本次 Clike 表单未展示可编辑的最大轮次项,提交后的任务页显示上限为 200 轮;其他模式的轮次设置不能直接套用。

-
点击“优化”,确认进入任务详情页,并保存任务链接及输入代码。提交成功仅表示任务已创建,不等于精度或性能已通过。
-
在任务详情跟踪环境初始化、算子预编译、精度验证与性能调优,必要时使用刷新按钮。完成后检查生成源码是否确为 AscendC,并记录硬件型号、软件栈、编译选项、正确性报告与性能报告。
-
导出最终代码,在独立环境复测,补充多尺寸和边界用例,再将源码、运行方法、测试报告与性能对比。提交的参考代码优先表达语义,并结合目标的场景构造对应的输入输出数据形状,实际获得算子加速效果也需要由实际生成和编译和测试结果确认。提交的任务如下图。

KernelSwift 的优势
工具能力最终要落到算子优化的实际流程中才能体现价值。手写第一版内核时,通常需要同时把握算子语义与硬件特性;此后每轮修改都伴随重新编译、精度验证与性能复测,反馈链路一长,调优便容易退化为反复试错;更换硬件或后端时,既有工作往往需要重做;在特定算子与 shape 上积累的调优经验,也难以直接迁移至其他场景。
KernelSwift 将这些环节纳入统一的任务流程,其核心能力与实际作用如下:
算子自动优化,降低起步成本。 输入标准 PyTorch 算子代码,即可生成适配目标硬件的高性能算子。算子语义与测试输入继续以 PyTorch 表达,减少手工搭建首版内核的工作量。
任务全流程管理,形成反馈闭环。 代码生成、编译、精度验证与性能调优被串联在同一流程内,候选实现能够及时获得执行反馈。需说明的是,进度百分比仅表示流程状态,不代表最终正确率或加速收益。
生成模式与硬件选择,支撑后端适配。 同一份高层计算描述可作为不同后端的适配起点;各后端的具体支持范围仍以实际任务结果为准。
算子广场与经验总结,支持复用与协作。 算子广场可复用已优化的高加速比算子,并支持查看、对比与二次优化;优化经验自动总结使后续任务参考既有结果,逐步提升算子性能基线。操作手册另介绍了任务查看与代码对比能力,便于保存实验记录、开展二次优化;本例实际使用的是新建任务与任务详情功能。

实践过程
KernelSwift可以生成高性能的算子,但是对于当前要实现的算子CHEMM有一些特定的要求,给定的测试集中包含了200个不同类别的测试用例,参赛者将测试用例分成了下面九类,并为每一个类别实现一个高性能的Kernel,最终通过dispatch的方式,组合成一个完整的算子。
| 场景 | 尺寸 M×N | 用途 |
|---|---|---|
| left/right × lower/upper regular | 128×64 | 四种数学语义;四者不能互相替代 |
| left_lower_small | 32×32 | 小矩阵启动成本与直接计算候选 |
| right_upper_narrow | 120×16 | 右乘窄输出、小 A 复用 |
| left_upper_tail_padded | 129×65 | 尾块与 lda/ldb/ldc 非紧凑 |
| right_lower_tail_padded | 65×129 | 右乘尾块与非紧凑布局 |
| left_lower_deep_k | 1024×32 | 左乘长归约、split-K 候选 |
| right_upper_deep_k | 32×1024 | 右乘长归约、split-K 候选 |
| left_lower_beta_zero | 128×64 | β=0 测试数据 |
| right_upper_alpha_zero | 128×64 | α=0 测试数据 |
另外,对于部分测试数据存在精度问题,参赛者实现的功能中还加入了检查浮点位精度的模式,因此dispatch分两层:先选择Cube或ordered/SIMT路径,再在Cube内选择分块、split-K、pack 和 combine 策略。实现的host测dispatch基本逻辑如下:
def chemm_dispatch(side, uplo, M, N, alpha, A, lda, B, ldb, beta, C, ldc):
validate_arguments(...)
if M == 0 or N == 0:
return SUCCESS
K = M if side == LEFT else N
workload = classify_workload(side, uplo, M, N, K, lda, ldb, ldc)
# 当前实现对这类标量采用保守计算路径
ordered_scalar = (
imag(alpha) != 0
or abs(real(alpha)) > 1
or abs(real(beta)) > 1
or abs(imag(beta)) > 1
)
# 实际源码对虚部/零值还检查了浮点位模式
narrow_cube_exception = (
workload.family == RIGHT_VERY_NARROW
and alpha == 1
and beta == 0
)
use_cube = (
not ordered_scalar
and K <= 4096
and 64 <= M <= 16384
and (N >= 64 or narrow_cube_exception)
and N <= 16384
and alpha != 0
and (side == LEFT or imag(alpha) == 0)
)
if use_cube:
return run_cube_pipeline(workload, ...)
return run_ordered_or_simt(workload, ordered_scalar, ...)
七条调优经验总结
KernelSwift具有自学习能力,会对优化后的算子进行经验总结为后续优化提供更好的知识储备。参赛者从本次使用KernelSwift实现并优化CHEMM过程中提炼了部分优化经验,可总结为以下七条:
-
融合 Hermitian 取数。按输出块计算时,根据索引选择 A[i,k] 或 conj(A[k,i]),对角线虚部归零,探索避免完整 H 的全局内存物化。
-
分解复数矩阵乘。探索由实数矩阵乘及向量运算组合实现。是否使用 Cube、采用何种精度及布局,要结合目标芯片能力、误差要求和搬运成本实测决定。
-
融合 α/β 与输出写回。在乘积累加后完成复数缩放和 βC 加法,减少中间结果与额外启动;α=0、β=0 等特例可探索专门路径,但必须保留正确的输入读取语义。
-
调整分块和多核分工。围绕 M/N/K 分块、片上存储容量、数据复用与尾块处理寻找合适配置;对小矩阵关注启动和调度成本,对大矩阵关注计算吞吐及带宽。
-
探索搬运与计算重叠。针对实际 AscendC 实现检查缓冲复用、双缓冲及同步开销,在满足资源约束和依赖关系的前提下减少等待。
-
按形状选择策略。当前 128×64 只是起点,应补充窄矩阵、非对齐尺寸和更大规模。
-
算子在不同shape下瓶颈点不相同,因此参赛者还是加入了一些手动调优的动作,比如:使用dispatch方式,通过多kernel的形式获得全shape下的高性能算子。
如果你喜欢我们的内容,欢迎点赞、收藏、关注我们
也欢迎在评论区交流你的实现与性能数据。
参考资料
- KernelSwift 平台:https://deeplink.org.cn/kernelswift
- KernelSwift智能算子迁移系统操作手册:https://aicarrier.feishu.cn/wiki/QtDbw0ADiinu9zk0sHbc81aQnGg
更多推荐



所有评论(0)