添加一个新算子¶

写一个新算子需要在以下六个位置添加实现，表格的顺序也是推荐的动手顺序。

其中 spec 要第一个写：后面五个文件的内容都由它决定，最后也都由它校验。spec 是这条流程的输入，其余五处都是照它写出来的。

#	文件	spec 里由谁指名	内容
1	`src/tileops/manifest/<family>.yaml`	顶层的名字就是算子类名	spec 本身
2	`src/tileops/ops/<family>/…`	`source.op`	算子类，继承 `Op`
2	`src/tileops/ops/__init__.py`	——	导出算子名
3	`src/tileops/kernels/<family>/…`	`source.kernel`	kernel 类，继承 `Kernel`
4	`tests/ops/test_<名字>.py`	`source.test`	与 `ref_api` 的数值比对
5	`benchmarks/ops/bench_<名字>.py`	`source.bench`	benchmark

下文以最简单的矩阵乘 GemmFwdOp 为例走一遍这六处。

动手之前：Ascend 的编译环境¶

🚨 pyproject.toml 声明的 tilelang 里没有 Ascend 后端。 本仓声明的依赖是 tilelang>=0.1.9,<0.2.0，它解析到 tile-ai/tilelang —— 那里面没有任何 Ascend 代码生成。 照声明依赖装完再写 kernel，它在 NPU 上编不出来。

Ascend 支持在 tile-ai/tilelang-ascend。该仓没有 main 分支，按后端路线选：ascendc_pto（AscendC / PTO 代码生成）或 npuir（MLIR 路线）。⚠️ 两条路线要用两个独立环境 —— 一个 Python 进程里 import tilelang 只会解析到其中一个。

⚠️ PYTHONPATH 要追加，不要覆盖：

export PYTHONPATH=/path/to/tilelang-ascend:$PYTHONPATH   # 对
export PYTHONPATH=/path/to/tilelang-ascend               # 错 —— 顶掉 CANN 自己的 python 路径

覆盖之后 CANN 的 $ASCEND_HOME_PATH/python/site-packages 与 .../opp/built-in/op_impl/ai_core/tbe 就没了，编译阶段抛 ModuleNotFoundError: No module named 'tbe' —— 这条报错读起来像「这个算子编不动」，其实是环境。

动手前先跑一遍自检，它能挡掉上面绝大多数问题：

python -c "
import torch, torch_npu, os, tilelang, pathlib
print('torch_npu', torch_npu.__version__, '| CANN', os.environ.get('ASCEND_HOME_PATH'))
print('NPU 可见', torch.npu.device_count())
root = pathlib.Path(tilelang.__file__).parent
n = len(list(root.rglob('*ascend*')))
print('tilelang ascend 相关文件', n, '->', 'OK' if n else '🚨 装的 tilelang 没有 Ascend 后端')"

最后一行为 0 就先别往下写。

第一步：写 spec¶

spec 各字段的含义与写法见读写 manifest，这里只说新算子特有的两件事。

第一件是状态。新算子先写成 status: spec-only，表示接口已经定下来、实现还没有，这时校验只跑 L0（结构检查），不会因为找不到实现而报错。

第二件是 source.kernel_map，这份 spec 里唯一推导不出来的字段。

一个算子可能对应多个 kernel：GEMM 在一般形状下用矩阵乘的 kernel，M 为 1 时退化成矩阵向量乘，换另一个 kernel 更快。kernel_map 就是这几个 kernel 的名单，每个 kernel 起一个名字，对应它的 Kernel 类：

GemmFwdOp:
  ref_api: torch.matmul
  family: gemm
  status: spec-only
  signature:
    inputs:
      a: {dtype: "float16 | bfloat16"}
      b: {dtype: "same_as(a)"}
    outputs:
      d: {dtype: "same_as(a)"}
    params:
      trans_a: {type: bool, default: false}
      trans_b: {type: bool, default: true}
    shape_rules:
      - "d.shape == ((a.shape[1] if trans_a else a.shape[0]), (b.shape[0] if trans_b else b.shape[1]))"
  source:
    kernel: tileops/kernels/gemm/dense.py
    kernel_map:
      gemm_kernel: GemmKernel
      gemv_kernel: GemvKernel
    op: tileops/ops/gemm/gemm.py
    test: tests/ops/test_gemm.py
    bench: benchmarks/ops/bench_gemm.py

算子在运行时按这些名字取 kernel：_eager_forward 里挑出用哪一个，把名字传给 get_or_build_kernel，算子层再从 kernel_map 找到对应的类去构造（见第二步）。外部后端也是照这份名单注册的 —— 它为哪个名字注册 build_kernel，就接管了算子的哪一个 kernel。

名字自己起，但要和 kernel 的用途对得上，而且写进算子代码之后就不该再改：它同时是 spec、算子实现与外部后端三方约定的那个词。这也是它推导不出来的原因 —— 只有写 kernel 的人知道这个算子要分几种情形。

第二步：写算子类¶

算子类继承 Op，是 spec 与 kernel 之间的一层：它按 spec 校验入参、推导输出形状，再取到 kernel 并 launch。先写它，是因为它的内容全部由 spec 决定，而它调用 kernel 的那一行同时定下了 kernel 的构造签名。

类的骨架与四个成员¶

GemmFwdOp 的骨架，方法体略去与本页无关的部分：

class GemmFwdOp(Op):
    def __init__(self, trans_a=False, trans_b=True, kernel_map=None, tune=False):
        self.trans_a, self.trans_b, self.tune = trans_a, trans_b, tune
        self.dispatch_kernel(kernel_map)             # 建立 kernel_map，不可省略

    @property
    def default_kernel_map(self):                    # spec 的 source.kernel_map
        return {"gemm_kernel": GemmKernel, "gemv_kernel": GemvKernel}

    def _infer_output_shapes(self, a_shape, b_shape):
        m = a_shape[1] if self.trans_a else a_shape[0]
        n = b_shape[0] if self.trans_b else b_shape[1]
        return {"d": (m, n)}                         # spec 的 shape_rules

    def forward(self, a, b):
        self._validate_dtypes(a, b)                  # 基类按 spec 生成，直接调用
        m, n, k = self._infer_mnk(a, b)
        a, b = a.contiguous(), b.contiguous()        # 按 spec 声明的形状交给 kernel
        slot = "gemv_kernel" if m == 1 else "gemm_kernel"
        kernel = self.get_or_build_kernel(
            slot,                                    # kernel_map 里的名字
            (a, b),                                  # 外部路径按这些张量查表，后端也收到它们
            key=(m, n, k, a.dtype),                  # 自带 kernel 按什么查表
            build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune),
        )
        return kernel(a, b)

要自己写的是这四个成员，内容都从 spec 来：

#	成员	照 spec 的哪一部分写
1	`__init__`	`signature.params` 的名字与默认值，另加 `kernel_map` 与 `tune`；结尾调用 `self.dispatch_kernel(kernel_map)` 建立本实例的 kernel_map
2	`default_kernel_map`	`source.kernel_map`：名字照抄，取值换成 Kernel 类本身
3	`_infer_output_shapes`	`signature.shape_rules` 里推导输出形状那几条
4	`forward`	`signature.inputs` 的顺序与默认值（可选输入排在必填之后），加上校验、连续化、取 kernel、launch kernel

另有两个成员不用写：_validate_dtypes 与 eval_roofline 由基类在子类定义时照 spec 的 dtype 声明与 roofline 生成并装上，直接调用即可，只有需要特殊行为时才自己覆写。

`get_or_build_kernel`¶

kernel 是编译产物，构造一次要几百毫秒到几秒，而一个算子实例会被反复调用，形状与 dtype 各不相同。算子层因此维护一张记忆表：本次调用要的 kernel 已经构造过就取回来，没有才构造并存进去。get_or_build_kernel 是这张表唯一的入口，也是自带实现与外部后端的分岔点（后端协议里的第二层选择）。

四个参数：

name —— 本次要哪一个 kernel，取值是 kernel_map 里的名字。

slot = "gemv_kernel" if m == 1 else "gemm_kernel"

自带实现按这个名字找到 Kernel 类，外部后端按它找到注册在同名下的 build_kernel。算子分几种情形，kernel_map 就有几个名字。

inputs —— 即将传给 kernel 的那些张量，顺序照 signature.inputs，一个输入占一个位置。

self.get_or_build_kernel(slot, (a, b), ...)                # GEMM：两个必填输入
self.get_or_build_kernel("group_norm", (x, weight, bias), ...)  # 没传的可选输入位置上是 None

外部路径按它查表：设备加上每个位置的 (dtype, shape)。设备也算在内，因为为一块卡编译的产物可能持有那块卡上的资源。后端的 build_kernel 收到的也是它，每个张量转成只有 device、dtype、shape 的 TensorSpec，不含数据。

没传的可选输入要留下位置、值为 None：后端由这个值判断输入传没传，而不是数位置个数 —— 挤掉空位，只给下界的 clamp 与只给上界的就成了同一个描述。

inputs 漏掉当场不报错，装上后端才抛 OpNotAvailableError —— 这个算子于是只能用自带 kernel，外部 target 接管不了（见安装之后：两种状态）。

key —— 自带 kernel 特化在什么上，只有自带这条路用（换成外部后端服务这个算子时这两个参数怎么走，见算子层这一侧的调用）。

key=(m, n, k, a.dtype)                             # GEMM：三个维度加 dtype
key=(self._cache_key(*input_shapes), x.dtype)      # 通用写法

_cache_key 的默认实现取所有输入形状中非静态轴的尺寸，正确但可能过细 —— 一个形状编译一次。kernel 实际只依赖其中几个量时覆写它，把形状投影过去，例如 kernel 把输入当二维处理，就把前面几维乘成一个数。

build —— 怎么构造这个自带 kernel，同样只有自带这条路用。

build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune)

每个 key 只调用一次，所以编译放在这里是安全的。算子完全没有自带实现、只指望外部后端服务时，build 可以不传 —— 那样在没有 target 认领设备时，调用会抛 OpNotAvailableError。返回值可以是一个 Kernel、一组一起构造出来的 Kernel，或一个带着它们的 dataclass —— 后两种适合一次调用要 launch 多个 kernel 的算子。

收尾：编译边界与注册¶

两件事收尾，都是几行的事：

要支持 torch.compile，得多声明一条编译边界：forward 只调用那个不透明算子，校验、取 kernel、launch kernel 挪进 _eager_forward。上面这个算子没有声明，所以 forward 里就是全部工作。做法见接入 torch.compile。
把算子名加进 src/tileops/ops/__init__.py 的导入与 __all__，from tileops.ops import ... 才拿得到它。

第三步：写 kernel¶

kernel 放在 src/tileops/kernels/ 下，用 TileLang 写。本仓的 kernel 是一个用 @register 登记的 build 函数，不是继承基类的 kernel 类 —— 它接过本次调用的张量，返回一个可调用对象：

from .._registry import register

@register("GemmFwdOp")                       # 名字就是算子类名
def build_gemm(a, b, *, trans_a=False, trans_b=True):
    ...                                      # 校验
    return build_gemm_kernel(tuple(a.shape), tuple(b.shape), a.dtype, trans_a, trans_b)

算子层通过 get_or_build_kernel(name, inputs) 要到它，算子层自己从不构造 kernel。

这个 build 函数的签名由 spec 定死：照 signature.inputs 的声明顺序，一个位置参数对应一个输入（没传的可选输入是 None），然后 signature.params 按关键字传。 kernel 的内容不受约束 —— kernel 不读 spec，spec 校验器也不检查它算什么，只登记它的路径。

kernel 本体用 T.Kernel(..., is_npu=True) 起，body 按执行单元分成两种 scope —— T.Scope("C") 是 Cube、T.Scope("V") 是 Vector：

with T.Kernel(launch_blocks, is_npu=True) as (cid, vid):
    with T.Scope("C"):                       # Cube：矩阵乘
        ...
    with T.Scope("V"):                       # Vector：逐元素、归约
        ...

划分参数有一条硬性要求：只有会被编译进生成代码的值才进那个被缓存的 builder。 形状、dtype、布局标志属于这一类，因为生成的代码里它们是常量：循环边界、tile 尺寸、搬运指令的形状都按它们展开。张量本身不进 builder —— 它们由 @register 那层接住, 每次调用只换指针。

分错的代价是重新编译。decode 一步一步往前走，seq_len 每步 +1，batch 随 running set 变化：

# 错：seq_len 进了被缓存的 builder —— 每一步都是一个新 kernel
kernel = build_attn_kernel(batch, seq_len, num_heads, dtype)

# 对：只有编译期常量进 builder，变化的量从张量形状读
kernel = build_attn_kernel(num_heads, head_dim, dtype)
out = kernel(q, k, v)                       # seq_len 从张量形状里读

上一种写法下，get_or_build_kernel 的 key 里带着 seq_len，每步都未命中、每步都编译一次，decode 直接跑不动。

第四步：写测试¶

测试放在 tests/ops/，比对对象就是 spec 的 ref_api，逐点比。形状与 dtype 取 spec 声明的范围，小形状标 smoke 进 PR 检查，大形状标 full 留给 nightly。

骨架用 tests/test_base.py 里的 TestBase 与 FixtureBase，用例写在 PARAMS 里。

如果这个算子有可选输入，传与不传各至少要有一条用例 —— 两侧走的往往是不同的 kernel。

🚨 在 Ascend 上，构造算子时必须显式传 target="ascend"：

op = MyFwdOp(...)
op.target = "ascend"          # ← 这种写法对所有算子都有效

⚠️ 不要用 MyFwdOp(..., target="ascend") —— target 是 Op 的类属性, 只有一部分算子在自己的 __init__ 里也接了这个关键字参数（实测 157 个算子里 81 个接、76 个不接, AbsFwdOp / BmmFwdOp / CosFwdOp 都在不接那一半）。 赋属性的写法处处可用，传构造参数的写法会在近一半算子上抛 TypeError。

tileops.backend.dispatch.detect_target() 对 npu 设备返回 None（None 的语义是「这台硬件没有装外部后端」），所以不指定 target 的算子在第一次 forward 时找不到 kernel。 症状是整个测试文件全红，看起来像 kernel 完全不工作 —— 测试里每一处构造算子都要带上它。

第五步：写 benchmark¶

benchmark 放在 benchmarks/ops/，继承 ManifestBenchmark。形状不自己写，而是经 load_workloads(<算子名>) 从 spec 的 workloads 取 —— 手写形状过不了 L4 校验：

from benchmarks.benchmark_base import ManifestBenchmark, workload_params
from tileops.manifest import load_workloads

_OP_NAME = "GemmFwdOp"

另外至少要记一个非 TileOPs 的基线，否则这一行没有比较对象。基线若需要转换输入，转换的代码留在它自己的计时区间内，不要挪出去。报出来的数字各是什么意思，见benchmark 怎么计时。

第六步：反转实现状态，让算子进入 CI 校验¶

上面五处都写完之后，先跑下面三条命令自查一遍：

python scripts/validate_manifest.py --check-op GemmFwdOp   # spec 与实现一致，五级全跑
python -m pytest tests/ops/test_gemm.py -v                # 数值与 ref_api 一致
python -m pytest benchmarks/ops/bench_gemm.py             # benchmark 能出数

三样都过，再把 spec 的 status 从 spec-only 反转成 implemented。这一改动的效果是让校验从 L0 扩到五级全跑，算子由此进入 CI 的保护范围：往后每次改动，spec 校验器、测试与 nightly benchmark 都会对照 spec 检查一遍。

接下来¶

算子跑起来之后，还有两件可选的事：

让算子能进使用者的编译图 —— 接入 torch.compile。
让它在别的硬件上由别人的 kernel 服务 —— 接入新硬件后端。