From b9e897c2c76111eb6e512f588683c5c9572e53bd Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Thu, 1 Oct 2026 12:32:08 +0800 Subject: [PATCH 1/5] [Docs] Describe op->kernel dispatch as kernel interfaces and implementations The dispatch migration replaced the op-level hooks the manual still taught. An op declares `interfaces`, `kernel_for(interface, call)` takes two arguments, and the implementations of an interface state their own availability, applicability, precedence and `entry_for`. The role concept, the op-level `entry_for` and the `kernel_for` path for an op with no interfaces are gone. new-op rewrites the kernel-selection section against that model, with the real `GemmFwdOp`, `GemmFwdInterface` and `RMSNormKernel`. backends gains the three extension granularities, `kernel_map=` and `register_implementation` beside a target, and drops the claim that an op may leave `entry_for` out. torch-compile gives `RMSNormFwdOp` its `interfaces` and its real `_eager_forward` body. manifest renames a composite's kernel roles to kernel keys. The backends pages also stated that no target is used unless a backend claims the device, which the same page contradicts forty lines earlier by documenting `set_default_target`. --- docs/backends.md | 116 ++++++++++++++++++++++++---- docs/backends.zh.md | 68 ++++++++++++++-- docs/manifest.md | 2 +- docs/manifest.zh.md | 2 +- docs/new-op.md | 163 ++++++++++++++++++++++++++++----------- docs/new-op.zh.md | 118 +++++++++++++++++++++------- docs/torch-compile.md | 19 +++-- docs/torch-compile.zh.md | 17 ++-- 8 files changed, 397 insertions(+), 108 deletions(-) diff --git a/docs/backends.md b/docs/backends.md index 95358271..b972a577 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -11,16 +11,95 @@ devices run your kernels. **A backend supplies one thing: something callable that computes this call.** Everything else is the op layer's. -The first half is the work, in the order it is done: the four things to write, the -protocol's four functions, how one call reaches them, a backend that installs and runs as -it stands, how to turn the template into a backend for real hardware, the four rules for -writing a kernel, what each phase may do, and — after install — which state each op is in -and what each error means. +This page is about a target, the largest of the three ways in. The first half is the +work, in the order it is done: the four things to write, the protocol's four functions, +how one call reaches them, a backend that installs and runs as it stands, how to turn the +template into a backend for real hardware, the four rules for writing a kernel, what each +phase may do, and — after install — which state each op is in and what each error means. The second half is why the protocol looks like this: the two layers of selection, the op layer's contract, when a kernel is rebuilt, what a caller can reach for, and what the protocol deliberately leaves out. +## Three ways in {#three-ways} + +How much of an op a package outside TileOPs takes over decides which of three it uses. +The two smaller ones write a kernel class against a [kernel +interface](new-op.md#kernel-selection), the same contract the in-tree +implementations are written against; a target writes a `build_kernel` against the op's +manifest signature instead. + +| # | | `kernel_map=` | `register_implementation` | target | +| --- | --- | --- | --- | --- | +| 1 | Changes | the class running behind one key; which calls that key serves does not change | adds a key, with its own applicability and precedence | every call of the op | +| 2 | Applies to | the one op instance the caller constructed it on | every instance of that op constructed afterwards | every instance that settles on the target | +| 3 | Written against | the kernel interface | the kernel interface | the op's manifest signature | +| 4 | Calls the new class does not serve | an error when that key is selected | still served by the in-tree implementations | none: a target serves them all | + +**`kernel_map=`** is a constructor argument of every op, a mapping from key to class. It +swaps the class behind that key in this instance; the key keeps the registered +implementation's `applies`, `general` and `preferred_over`, and the replacement, like +every implementation, inherits the key's interface and is built through its own +`entry_for`. A selected key whose replacement cannot serve the call is an error, never a +fall back to what it replaced. From +[`tests/test_kernel_dispatch.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_kernel_dispatch.py): + +```python +class _TorchLayerNorm(Kernel, LayerNormFwdInterface): + """A replacement written against ``LayerNormFwdInterface`` alone.""" + + devices = frozenset({torch.device(run_device()).type}) + + def __init__(self, n: int, eps: float) -> None: + super().__init__() + self.n, self.eps = n, eps + + @classmethod + def entry_for(cls, call: LayerNormCall): + return (call.n, call.eps), lambda: cls(call.n, call.eps) + + def forward(self, x, weight, bias): + return F.layer_norm(x.float(), (self.n,), weight.float(), bias.float(), self.eps).to( + x.dtype + ) + + +op = LayerNormFwdOp((32,), kernel_map={"layer_norm": _TorchLayerNorm}, target=BUILTIN) +``` + +A key the op does not have, but another op does, is ignored, which is how a composite op +passes one mapping down to its sub-ops; a key no op has raises +`was given kernel_map keys no op has` at construction. + +**`register_implementation(op, key, implementation)`** adds an implementation instead of +replacing one. `op` is the op's manifest key, `key` the new implementation's dispatch key, +and the interface it joins is the one the class inherits. It declares its own region, so +it needs `preferred_over` where it overlaps an in-tree implementation that is not +`general`: + +```python +class _NarrowTorchLayerNorm(_TorchLayerNorm): + """An added implementation for short rows, which wins over the in-tree one there.""" + + preferred_over = frozenset({"layer_norm"}) + + @classmethod + def applies(cls, call: LayerNormCall) -> bool: + return call.n <= 64 + + +register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) +``` + +Here `n <= 64` goes to `_NarrowTorchLayerNorm` and `n = 1024` stays with the in-tree +`LayerNormKernel`. The addition reaches only op instances constructed after the call. +Registering the same key twice under one op raises `BackendError`; a key an in-tree +implementation already uses raises `reuse keys it has` when an instance is constructed. + +`register_implementation` runs when the backend module is imported, through the same +entry point a target uses; `kernel_map=` registers nothing and is passed by the caller at +construction. What follows is the target. + ## Four things to write | # | What to do | @@ -222,9 +301,9 @@ signature. Whatever it returns, the op layer stores and launches. Four things follow from that: -- **`kernel_for` and `entry_for` are the op author's, not a backend's.** They serve the - in-tree path only: which in-tree kernel to fetch, what it is looked up on and how it is - built. Once a target serves the op, it serves the whole op, and none of them runs. +- **`kernel_for` and the implementations' `entry_for` serve the in-tree path only.** They + decide which in-tree kernel is fetched, what it is looked up on and how it is built. + Once a target serves the op, it serves the whole op, and none of them runs. - **Tensors arrive positionally, params by name.** `build_kernel(*inputs, **params)`: the positional arguments are `TensorSpec`s (`None` for an optional input the call omitted), the keywords the manifest's `params` names with the values this call settled on. @@ -233,8 +312,9 @@ Four things follow from that: `TensorSpec`s which kernel to return. - **No memoisation of its own is needed.** For the same device and input signature the op layer does not call again; for a finer split, or fewer rebuilds, add a cache inside - `build_kernel`. An op with no in-tree implementation may leave `entry_for` out, and then - a call with no target claiming the device raises `OpNotAvailableError`. + `build_kernel`. An op written to depend on a backend declares neither `kernel_types` nor + `interfaces`, and a call on it with no target claiming the device raises + `OpNotAvailableError`. ## Writing a backend that runs {#runnable} @@ -553,9 +633,11 @@ print(load_failures()) | **`build_kernel`** | A function a backend writes per op, one per `(op, target)` | The second layer: it receives a description of this call — each input's device, dtype and shape, plus the op's parameters — and picks, builds and returns a kernel from its own set | **Selection has two layers: TileOPs picks the target, the target picks the -kernel.** The second happens inside `build_kernel`, with no protocol involvement; -there is no kernel-level concept, no capability negotiation and no candidate -filtering. +kernel.** The second happens inside `build_kernel`, with no protocol involvement: on this +path there is no kernel-level concept, no capability negotiation and no candidate +filtering. The candidate filtering TileOPs does run — availability, applicability, +precedence — belongs to the in-tree path and to the two smaller ways in +([three ways in](#three-ways)), which a target bypasses. `detect` answers only which devices belong to the backend, and that is as fine as it gets. **Whether this call is supported — dtype, shape, parameter combination — @@ -592,10 +674,10 @@ third-party backend neither bypasses one nor substitutes its own. The kernels TileOPs ships ([`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)) are the **default implementation**: they have no target name and are not in the registry. -**The default state is no substitution.** With no backend claiming a device, calls on -it run the shipped implementation; only once a backend is installed and has claimed -that device does the op's kernel become the backend's. The protocol has no notion of a -"default target". +**The default state is no substitution.** With no backend installed, no `target=` named +and no process default set, calls run the shipped implementation. Nothing is preconfigured +to substitute: a backend serves an op only once it is installed and either claims the +device through its `detect`, or is named by `target=` or `set_default_target`. ## When a kernel is rebuilt {#memo} diff --git a/docs/backends.zh.md b/docs/backends.zh.md index c9a8f832..d1204e91 100644 --- a/docs/backends.zh.md +++ b/docs/backends.zh.md @@ -6,10 +6,68 @@ TileLang 是多后端 DSL,每种硬件各有一套独立的 kernel,由各自 **后端只提供一件事:一个能算这次调用的可调用对象。** 其余都由算子层负责。 -前半按动手顺序列出要做的事:要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么,以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。 +本页讲的是 target,三种接入方式中范围最大的一种。前半按动手顺序列出要做的事:要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么,以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。 后半说明协议何以如此设计:两层选择、算子层的契约、kernel 的重建条件、调用方可用的接口,以及刻意不支持的情形。 +## 三种接入方式 {#three-ways} + +仓外的包要接管一个算子的多少调用,决定用三种方式中的哪一种。范围较小的两种,写出的 kernel 类依据的是 [kernel 接口](new-op.md#kernel-selection),与自带实现依据的是同一份契约;target 依据的则是算子在 manifest 中的签名,写一个 `build_kernel`。 + +| # | | `kernel_map=` | `register_implementation` | target | +| --- | --- | --- | --- | --- | +| 1 | 改变什么 | 一个 key 背后运行的类;这个 key 服务哪些调用不变 | 新增一个 key,带有它自己的适用范围与优先关系 | 算子的全部调用 | +| 2 | 作用于 | 调用方构造的那一个算子实例 | 注册之后构造的所有该算子实例 | 选中该 target 的算子实例 | +| 3 | 依据的契约 | kernel 接口 | kernel 接口 | 算子在 manifest 中的签名 | +| 4 | 新类不服务的调用 | 该 key 被选中时报错 | 仍由自带实现服务 | 不存在,target 服务全部调用 | + +**`kernel_map=`** 是每个算子构造函数的参数,值是从 key 到类的映射。它只替换这个实例中该 key 背后运行的类:这个 key 仍按原来登记的实现声明的 `applies`、`general` 与 `preferred_over` 参与选择,替换者与其他实现一样继承这个 key 所属的接口,由它自己的 `entry_for` 构造。该 key 被选中而替换者不服务这次调用时报错,不退回被替换的实现。下例取自 [`tests/test_kernel_dispatch.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_kernel_dispatch.py): + +```python +class _TorchLayerNorm(Kernel, LayerNormFwdInterface): + """A replacement written against ``LayerNormFwdInterface`` alone.""" + + devices = frozenset({torch.device(run_device()).type}) + + def __init__(self, n: int, eps: float) -> None: + super().__init__() + self.n, self.eps = n, eps + + @classmethod + def entry_for(cls, call: LayerNormCall): + return (call.n, call.eps), lambda: cls(call.n, call.eps) + + def forward(self, x, weight, bias): + return F.layer_norm(x.float(), (self.n,), weight.float(), bias.float(), self.eps).to( + x.dtype + ) + + +op = LayerNormFwdOp((32,), kernel_map={"layer_norm": _TorchLayerNorm}, target=BUILTIN) +``` + +`kernel_map=` 中本算子没有、而其他算子有的 key 被忽略,复合算子由此把一份映射传给各个子算子;任何算子都没有的 key 在构造时报 `was given kernel_map keys no op has`。 + +**`register_implementation(op, key, implementation)`** 不是替换,而是新增一个实现。`op` 是算子在 manifest 中的 key,`key` 是新实现的名字,它加入哪个接口由这个类继承了哪个接口决定。新实现自己声明适用范围,与某个不是 `general` 的自带实现重叠时要声明 `preferred_over`: + +```python +class _NarrowTorchLayerNorm(_TorchLayerNorm): + """An added implementation for short rows, which wins over the in-tree one there.""" + + preferred_over = frozenset({"layer_norm"}) + + @classmethod + def applies(cls, call: LayerNormCall) -> bool: + return call.n <= 64 + + +register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) +``` + +这里 `n <= 64` 由 `_NarrowTorchLayerNorm` 服务,`n = 1024` 仍由自带的 `LayerNormKernel` 服务。新增只进入注册之后构造的算子实例。同一个算子下重复注册同一个 key 报 `BackendError`;key 与自带实现的 key 相同时,构造实例时报 `reuse keys it has`。 + +`register_implementation` 在后端模块被导入时执行,入口与 target 用的是同一条 entry point;`kernel_map=` 不注册任何东西,由调用方在构造算子时传入。以下各节讲 target。 + ## 写一个后端要做的四件事 | # | 做什么 | @@ -184,10 +242,10 @@ register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) 四点对应关系值得记住: -- **`kernel_for` 与 `entry_for` 由算子作者写,与后端无关。** 它们只服务自带实现:取哪个自带 kernel、按什么查表、又怎么构造。target 选中后端时整个算子由 target 服务,这几处都不会执行。 +- **`kernel_for` 与各实现的 `entry_for` 只服务自带实现。** 它们决定取哪个自带 kernel、按什么查表、又怎么构造。target 选中后端时整个算子由 target 服务,这几处都不会执行。 - **张量按位置传,参数按名字传。** `build_kernel(*inputs, **params)`:位置实参是 `TensorSpec`(没传的可选输入是 `None`),关键字实参是 manifest 里 `params` 的名字与本次调用的确定值。 - **一个 `(算子, target)` 只注册一个 builder。** 自带实现内部分几种 kernel(GEMM 的 `kernel_types` 里有三个)不会传进来,`build_kernel` 从 `TensorSpec` 自行判断该返回哪个 kernel。 -- **不必自己做记忆。** 同一个设备与输入签名,算子层不会再调第二次;要更细的区分或更少的重建,在 `build_kernel` 内部另加一层缓存。算子完全没有自带实现时 `entry_for` 可以不写,那时没有 target 认领设备,调用直接抛 `OpNotAvailableError`。 +- **不必自己做记忆。** 同一个设备与输入签名,算子层不会再调第二次;要更细的区分或更少的重建,在 `build_kernel` 内部另加一层缓存。写来依赖外部后端的算子,`kernel_types` 与 `interfaces` 都不声明,那时没有 target 认领设备,调用直接抛 `OpNotAvailableError`。 ## 实现一个可运行的后端 {#runnable} @@ -443,7 +501,7 @@ print(load_failures()) | **`detect`** | 后端写的一个函数,一个 target 一个 | 第一层怎么选:接收一块 `torch.device`,回答这类设备是不是自己这套 kernel 的目标设备;不是则返回 `False` | | **`build_kernel`** | 后端为某个算子写的一个函数,一组 `(算子, target)` 一个 | 第二层:接收本次调用的描述,即各输入张量的 device、dtype、shape 与算子参数,在自己这套 kernel 中选定一个、构造好并返回 | -**选择分两层:TileOPs 选 target,target 在自己那套 kernel 里选一个。** 第二层发生在 `build_kernel` 内部,协议不参与,也不存在 kernel 一级的概念、能力协商与候选筛选。 +**选择分两层:TileOPs 选 target,target 在自己那套 kernel 里选一个。** 第二层发生在 `build_kernel` 内部,协议不参与:这条路径上不存在 kernel 一级的概念、能力协商与候选筛选。TileOPs 确实执行的候选筛选 —— 可用性、适用范围、优先关系 —— 属于自带实现与范围较小的两种接入方式([三种接入方式](#three-ways)),target 绕过它们。 `detect` 只回答设备的归属,粒度到此为止。**本次调用是否受支持 —— 涉及 dtype、形状与参数组合 —— 由 `build_kernel` 回答**,因为只有它看得到完整的输入描述与参数;不支持时在那里报错。这些判断交给 `detect` 是做不到的,它只拿到一块 `torch.device`。 @@ -471,7 +529,7 @@ TileOPs 不解析 `torch.device`,而是把它原样传给 `detect`。这样做 TileOPs 自带的 kernel([`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels))是**默认实现**:它没有 target 名,也不进注册表。 -**默认状态是不替换。** 没有后端认领某块设备时,这台设备上的调用走自带实现;装上一个后端、并且它认领了这块设备,该算子的 kernel 才换成后端的那一套。协议里因此不存在「默认 target」这个概念。 +**默认状态是不替换。** 没有装后端、没有指名 `target=`、也没有设置进程默认值时,调用走自带实现。协议不预置任何替换:装上一个后端之后,它的 `detect` 认领了这块设备,或者被 `target=` 与 `set_default_target` 指名,该算子的 kernel 才换成后端的那一套。 ## kernel 的重建条件 {#memo} diff --git a/docs/manifest.md b/docs/manifest.md index 230ec6ce..f5c0b120 100644 --- a/docs/manifest.md +++ b/docs/manifest.md @@ -48,7 +48,7 @@ required once the other direction also has an entry — and the validator requir | `signature` | yes | the op's type, below | | `workloads` | yes | the calls tests and benchmarks run | | `roofline` | yes | the cost of one call, specified in [Roofline](design/roofline.md) | -| `composition` | no | for a composite op, its stages in order: the sub-op classes it may hold and its own kernel roles | +| `composition` | no | for a composite op, its stages in order: the sub-op classes it may hold and its own kernel keys | The signature is a function type over named type indices: diff --git a/docs/manifest.zh.md b/docs/manifest.zh.md index cf4cbb4b..ace0ff24 100644 --- a/docs/manifest.zh.md +++ b/docs/manifest.zh.md @@ -33,7 +33,7 @@ TileOPs 的组织方式相反:算子的规格先声明,实现由规格推导 | `signature` | 是 | 算子的类型,见下表 | | `workloads` | 是 | 测试与 benchmark 执行的调用 | | `roofline` | 是 | 一次调用的开销,规范见 [Roofline](design/roofline.md) | -| `composition` | 否 | 复合算子按顺序列出的各阶段:可能持有的子算子类,以及它自己的 kernel 角色 | +| `composition` | 否 | 复合算子按顺序列出的各阶段:可能持有的子算子类,以及它自己的 kernel key | 签名是一个以具名类型变量为参数的函数类型: diff --git a/docs/new-op.md b/docs/new-op.md index 7ffc39e4..6fab1b92 100644 --- a/docs/new-op.md +++ b/docs/new-op.md @@ -73,14 +73,17 @@ What it writes is how a call reaches a kernel. class GemmFwdOp(Op): compile_boundary: ClassVar[bool] = True # optional: claims fullgraph=True kernel_types: ClassVar[Mapping[str, type[Kernel]]] = { - "gemm_tma_kernel": GemmTmaKernel, - "gemm_cp_async_kernel": GemmCpAsyncKernel, - "gemv_kernel": GemvKernel, + "gemm_tma": GemmTmaKernel, + "gemm_cp_async": GemmCpAsyncKernel, + "gemv": GemvKernel, } + interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = {"gemm": GemmFwdInterface} def __init__(self, trans_a=False, trans_b=True, *, target=None, kernel_map=None, tune=False): - self.trans_a, self.trans_b = trans_a, trans_b - self.target, self.tune = target, tune + self.trans_a = trans_a + self.trans_b = trans_b + self.target = target + self.tune = tune self.dispatch_kernel(kernel_map) # installs this instance's kernel map def forward(self, a, b): @@ -89,22 +92,26 @@ class GemmFwdOp(Op): def _eager_forward(self, a, b): # the generated checks have run a, b = a.contiguous(), b.contiguous() # handed over as the spec declares it m, k = (a.shape[1], a.shape[0]) if self.trans_a else a.shape - n = b.shape[0] if self.trans_b else b.shape[1] - kernel = self.kernel_for( - "gemm", # the memoization bucket - (a, b), # the tensors the kernel gets - self._call_spec(m, n, k, a.dtype, a.device), # what this call is + call = GemmCall( # what this call is + m=m, + n=b.shape[0] if self.trans_b else b.shape[1], + k=k, + dtype=a.dtype, + trans_a=self.trans_a, + trans_b=self.trans_b, + device=a.device, ) - return kernel(a, b) + return self.kernel_for("gemm", call)(a, b) ``` | # | Member | Written from | | --- | --- | --- | | 1 | `__init__` | the names, order and defaults in `signature.params`, then `target`, `kernel_map` and `tune`, closing with `self.dispatch_kernel(kernel_map)` | -| 2 | `kernel_types` | the Kernel classes that can serve the op, each under a name; a `kernel_map=` override replaces one by that name | -| 3 | `forward` | `signature.inputs` — its order, optional inputs last with default `None` | -| 4 | `_eager_forward` | contiguity, the call record, fetching the kernel and launching it | -| 5 | `compute_roof` | optional: the GPU-profile unit that prices the op's FLOPs, where it is not CUDA-core fp32 | +| 2 | `kernel_types` | the Kernel classes that can serve the op, each under a key; a `kernel_map=` override replaces one by that key | +| 3 | `interfaces` | one entry per place the op calls a kernel: the name `kernel_for` uses → the `KernelInterface` class the implementations of that place inherit | +| 4 | `forward` | `signature.inputs` — its order, optional inputs last with default `None` | +| 5 | `_eager_forward` | contiguity, the call spec, fetching the kernel and launching it | +| 6 | `compute_roof` | optional: the GPU-profile unit that prices the op's FLOPs, where it is not CUDA-core fp32 | `_infer_output_shapes`, `_validate_dtypes` and `eval_roofline` are generated from the spec and are not written. @@ -113,7 +120,7 @@ An op without a compile boundary writes the body of `_eager_forward` in `forward declaring the boundary moves it behind the generated operator. How that works is in [bringing an op into torch.compile](torch-compile.md). -### `kernel_for`, and choosing among kernels +### `kernel_for`, and choosing among kernels {#kernel-selection} A kernel is a compiled artefact, hundreds of milliseconds to seconds to build, while an op instance is called over and over at different shapes and dtypes. The op layer therefore @@ -121,38 +128,108 @@ keeps a memo table: a kernel this call needs and has built before comes straight and only otherwise is one built and stored. `kernel_for` is that table's only entrance on the in-tree path; a [target](backends.md) serves the whole op instead and never reaches it. -Its three arguments: +Its two arguments: + +- **`interface`** — a key of `interfaces`, naming one place in the op where a kernel is + called. `GemmFwdOp` calls a kernel in one place, so it declares one, `"gemm"`. A second + interface is opened only where the semantics or the call contract changes: + `BatchNormFwdOp` has `batch_norm_fwd_train` and `batch_norm_fwd_infer`, which return + different things. A faster kernel for some shape range or some architecture is not a + new interface — it is another implementation of the one that is there. +- **`call`** — a frozen `CallSpec` subclass carrying the facts of this call that selection + and building read: shapes, the dtype, the op's semantic parameters, and the device. It + has to be the interface's `request` type. Device facts (`arch`, `sm_count`, + `calibration`, `smem_budget`) are not passed in; the dispatcher derives them from + `call.device` on a miss. + +The kernel that comes back is called with the parameters of the interface's abstract +`forward`, in that order. + +An interface is a class in +[`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels), +beside the call spec it names in `request`. Its abstract `forward` is the whole contract +an implementation — in-tree or from a backend — is written against, so its docstring +states each tensor's shape, dtype, layout, device and whether it is written in place: -- **`role`** — the memoization bucket, one per kernel the op runs per call. `GemmFwdOp` - runs one, so it has one role, whichever of its three classes serves the call. -- **`inputs`** — the tensors the kernel is about to be handed, in `signature.inputs` - order, one slot per input. An optional input that was not passed keeps its slot as - `None`. -- **`call`** — what this call is. `GemmCall` carries every fact the GEMM kernels read: - `m`, `n`, `k`, the dtype, the layout, the device. +```python +class GemmFwdInterface(KernelInterface): + """Dense matmul under the ``(trans_a, trans_b)`` layout the call states.""" + + request = GemmCall + + @abstractmethod + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Multiply the two matrices; nothing is written in place. + + Both operands are contiguous on ``call.device`` in ``call.dtype``, which is + ``float16`` or ``bfloat16``; the contraction accumulates in ``float32``. + + Args: + a: ``(call.m, call.k)``, or ``(call.k, call.m)`` when ``call.trans_a``. + b: ``(call.n, call.k)`` when ``call.trans_b``, else ``(call.k, call.n)``. + + Returns: + A new ``(call.m, call.n)`` tensor in ``call.dtype``. + """ +``` -Which class serves a call is decided by the classes, not the op. Each states the region -it serves (`applies`, `refusal`), one is marked `general` for everything else, and two -specialised classes claiming one call is an error, never a silent preference. The chosen -class's `entry_for(call)` returns the **identity** two calls must share to reuse one -kernel, and the **builder** that runs once per identity. Carry too little in the identity -and a second dtype reuses the first dtype's kernel; carry the whole shape where the kernel -depends on fewer quantities and it compiles once per distinct shape. +An implementation is a class inheriting both `Kernel` and one interface, listed in +`kernel_types` under a key. Which implementation serves a call is decided by the +implementations, not by the op, from four declarations: -An op with a single kernel and no call record writes `entry_for(role, call)` on the op -itself and states the identity and builder there, as `RMSNormFwdOp` does: +| # | Declaration | States | Left undeclared | +| --- | --- | --- | --- | +| 1 | `devices`, `supported_archs` | where the implementation runs | CUDA devices, every architecture | +| 2 | `applies(call)`, `refusal(call)` | which calls it serves, stated positively | every call | +| 3 | `general`, `preferred_over` | which implementation wins where two of them serve one call | wins over none | +| 4 | `entry_for(call)` | the build identity, and the builder that runs once per identity | the whole call spec, built by `cls(call)` | + +Availability filters first. Among the implementations that are left and that apply, +`general` loses to every other one and `preferred_over` names the keys its class wins +over. Nothing left raises `no implementation serves this call`, or `OpNotAvailableError` where +no key runs on the call's device type at all; two with no relation between them raise +`dispatch is ambiguous`; declaration order decides nothing. An +implementation never excludes a sibling in its own `applies` — the one that should win +declares `preferred_over`. + +`GemmFwdOp`'s three implementations divide the `"gemm"` interface like this: + +| # | Key | Serves | Declares | +| --- | --- | --- | --- | +| 1 | `gemm_tma` | SM90 shapes whose operands TMA can address | `supported_archs = [90]`, and a `refusal` naming the misalignment | +| 2 | `gemv` | at most two rows contracted over K, where reducing on CUDA cores wins | `supported_archs = [90]`, `applies` through `band_for`, `preferred_over = frozenset({"gemm_tma"})` | +| 3 | `gemm_cp_async` | every shape the other two do not claim, down to a K row spanning one four-byte load | `supported_archs = [80, 86, 89, 90]`, `general = True`, and a `refusal` for a narrower K row | + +`entry_for(call)` returns the **identity** two calls must share to reuse one kernel, and +the **builder** that runs once per identity. Carry too little in the identity and a second +dtype reuses the first dtype's kernel; carry the whole shape where the kernel depends on +fewer quantities and it compiles once per distinct shape. + +An interface with one implementation needs nothing beyond inheriting it. `RMSNormKernel` +is the whole of `RMSNormFwdOp`'s dispatch +([`src/tileops/kernels/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/norm/rms_norm.py)): ```python -def entry_for(self, role, call): # call is the input dtype - n = math.prod(self.normalized_shape) - eps = torch.finfo(torch.float32).eps if self.eps is None else float(self.eps) - return call, lambda: self.kernel_map["rms_norm"](n, eps, call, tune=self.tune) +class RMSNormKernel(Kernel, RMSNormFwdInterface): + supported_archs: list[int] = [80, 86, 89, 90] + + @classmethod + def entry_for(cls, call: LayerNormCall) -> Entry: + identity = (call.n, call.eps, call.dtype) + return identity, lambda: cls(*identity) ``` -An op with no in-tree implementation, written to depend on a backend, leaves out both -`kernel_types` and `entry_for`; a call on a device no target claims then raises +An op defines no `entry_for` of its own and keeps no kernel cache of its own — no dict, no +build guarded on an attribute being unset. Holding what `kernel_for` returned in +`self.kernel` is not one. + +An op with no in-tree implementation, written to depend on a backend, declares neither +`kernel_types` nor `interfaces`; a call on a device no target claims then raises `OpNotAvailableError`. +A backend adds an implementation to an interface, or replaces the class behind one key, +without changing TileOPs; both are in [adding a hardware backend](backends.md). + ### Registering Add the op's name to the imports and `__all__` in two places: its family's @@ -164,10 +241,10 @@ reference cannot collect it. ## Step 3: write the kernel -A kernel class subclasses [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py), lives under [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels), is written in +A kernel class subclasses [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) and the interface it implements, lives under [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels), is written in TileLang, compiles at construction and implements `forward`, which the base class's -`__call__` runs. Its constructor is what its -`entry_for` builder calls, and its call signature is the `kernel(a, b)` of step 2. +`__call__` runs. Its constructor is what its own `entry_for` builder calls, and its +`forward` takes the interface's parameters, the `kernel(a, b)` of step 2. This is the one place of the six the spec does not constrain: a kernel neither reads the spec nor is checked against it. @@ -177,7 +254,7 @@ values compiled into the generated code go in the constructor.** `GemmTmaKernel` them like this: ```python -class GemmTmaKernel(Kernel): +class GemmTmaKernel(Kernel, GemmFwdInterface): def __init__(self, m, n, k, dtype, config=None, tune=False, trans_a=False, trans_b=False, ...): self.kernel = _gemm_kernel(m, n, k, trans_a, trans_b, self.dtype_str, ...) # compiles self.init_config(config, tune) # tile sizes and pipeline depth diff --git a/docs/new-op.zh.md b/docs/new-op.zh.md index 9ff07bd0..48882d53 100644 --- a/docs/new-op.zh.md +++ b/docs/new-op.zh.md @@ -64,14 +64,17 @@ spec 不写文件路径,也不写 kernel:由哪些 kernel 服务这个算子 class GemmFwdOp(Op): compile_boundary: ClassVar[bool] = True # optional: claims fullgraph=True kernel_types: ClassVar[Mapping[str, type[Kernel]]] = { - "gemm_tma_kernel": GemmTmaKernel, - "gemm_cp_async_kernel": GemmCpAsyncKernel, - "gemv_kernel": GemvKernel, + "gemm_tma": GemmTmaKernel, + "gemm_cp_async": GemmCpAsyncKernel, + "gemv": GemvKernel, } + interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = {"gemm": GemmFwdInterface} def __init__(self, trans_a=False, trans_b=True, *, target=None, kernel_map=None, tune=False): - self.trans_a, self.trans_b = trans_a, trans_b - self.target, self.tune = target, tune + self.trans_a = trans_a + self.trans_b = trans_b + self.target = target + self.tune = tune self.dispatch_kernel(kernel_map) # installs this instance's kernel map def forward(self, a, b): @@ -80,49 +83,104 @@ class GemmFwdOp(Op): def _eager_forward(self, a, b): # the generated checks have run a, b = a.contiguous(), b.contiguous() # handed over as the spec declares it m, k = (a.shape[1], a.shape[0]) if self.trans_a else a.shape - n = b.shape[0] if self.trans_b else b.shape[1] - kernel = self.kernel_for( - "gemm", # the memoization bucket - (a, b), # the tensors the kernel gets - self._call_spec(m, n, k, a.dtype, a.device), # what this call is + call = GemmCall( # what this call is + m=m, + n=b.shape[0] if self.trans_b else b.shape[1], + k=k, + dtype=a.dtype, + trans_a=self.trans_a, + trans_b=self.trans_b, + device=a.device, ) - return kernel(a, b) + return self.kernel_for("gemm", call)(a, b) ``` | # | 成员 | 照什么写 | | --- | --- | --- | | 1 | `__init__` | `signature.params` 的名字、顺序与默认值,再加 `target`、`kernel_map`、`tune`;结尾调用 `self.dispatch_kernel(kernel_map)` | -| 2 | `kernel_types` | 能服务这个算子的 Kernel 类,各起一个名字;`kernel_map=` 按这个名字替换其中一个 | -| 3 | `forward` | `signature.inputs` 的顺序,可选输入排在最后、默认 `None` | -| 4 | `_eager_forward` | 连续化、调用记录、取 kernel、launch kernel | -| 5 | `compute_roof` | 可选:给算子 FLOPs 定价的 GPU profile 单元,不是 CUDA core fp32 时才写 | +| 2 | `kernel_types` | 能服务这个算子的 Kernel 类,各起一个 key;`kernel_map=` 按这个 key 替换其中一个 | +| 3 | `interfaces` | 算子调用 kernel 的每一处各占一条:`kernel_for` 用的名字 → 这一处的各实现所继承的 `KernelInterface` 类 | +| 4 | `forward` | `signature.inputs` 的顺序,可选输入排在最后、默认 `None` | +| 5 | `_eager_forward` | 连续化、构造 call spec、取 kernel、launch kernel | +| 6 | `compute_roof` | 可选:给算子 FLOPs 定价的 GPU profile 单元,不是 CUDA core fp32 时才写 | `_infer_output_shapes`、`_validate_dtypes` 与 `eval_roofline` 都照 spec 生成,不用写。 不声明编译边界的算子,把 `_eager_forward` 的内容直接写在 `forward` 里;声明了边界,这些内容挪到生成的 operator 后面。做法见[接入 torch.compile](torch-compile.md)。 -### `kernel_for` 与 kernel 的选择 +### `kernel_for` 与 kernel 的选择 {#kernel-selection} kernel 是编译产物,构造一次要几百毫秒到几秒,而一个算子实例会被反复调用,形状与 dtype 各不相同。算子层因此维护一张记忆表:本次调用要的 kernel 已经构造过就取回来,没有才构造并存进去。`kernel_for` 是自带实现走到这张表的唯一入口;[target](backends.md) 服务的是整个算子,不经过它。 -三个参数: +两个参数: -- **`role`**:记忆表的桶,算子一次调用跑几个 kernel 就有几个。`GemmFwdOp` 只跑一个,所以只有一个 role,不论三个类中哪一个服务这次调用。 -- **`inputs`**:即将传给 kernel 的张量,顺序照 `signature.inputs`,一个输入占一个位置。没传的可选输入留下位置,值为 `None`。 -- **`call`**:本次调用是什么。`GemmCall` 带着 GEMM 各 kernel 要读的全部事实:`m`、`n`、`k`、dtype、布局、设备。 +- **`interface`**:`interfaces` 的一个 key,指算子调用 kernel 的一处。`GemmFwdOp` 只在一处调用 kernel,因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新开一个接口:`BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的东西不同。某个形状范围或某个架构上更快的 kernel 不是新接口,而是已有接口的另一个实现。 +- **`call`**:一个冻结的 `CallSpec` 子类,带着这次调用中选择与构建会读取的事实:形状、dtype、算子的语义参数与设备。它必须是这个接口 `request` 指定的类型。设备事实(`arch`、`sm_count`、`calibration`、`smem_budget`)不由调用方传入,未命中时派发机制从 `call.device` 推出。 -由哪个类服务一次调用,由这些类自己决定,不由算子决定。每个类声明自己服务的范围(`applies`、`refusal`),其中一个标为 `general`,负责其余情形;两个专用类同时认领一次调用会报错,不会静默挑一个。选中的类用 `entry_for(call)` 返回两样东西:两次调用要共享什么才算同一个 kernel 的**身份**,以及每个身份只跑一次的**构造方法**。身份带少了,第二种 dtype 会复用第一种 dtype 的 kernel;kernel 只依赖其中几个量却把整个形状带上,就变成一个形状编译一次。 +取回的 kernel 按接口抽象 `forward` 的参数表、按同样的顺序调用。 -只有一个 kernel、也没有调用记录的算子,在算子类上自己写 `entry_for(role, call)`,在那里给出身份与构造方法,`RMSNormFwdOp` 就是这样: +接口是写在 [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 里的一个类,与它在 `request` 中指名的 call spec 放在一起。它的抽象 `forward` 是各实现(自带的与后端提供的)唯一依据的契约,docstring 因此写明每个张量的形状、dtype、内存布局、设备,以及是否被原地写入: ```python -def entry_for(self, role, call): # call is the input dtype - n = math.prod(self.normalized_shape) - eps = torch.finfo(torch.float32).eps if self.eps is None else float(self.eps) - return call, lambda: self.kernel_map["rms_norm"](n, eps, call, tune=self.tune) +class GemmFwdInterface(KernelInterface): + """Dense matmul under the ``(trans_a, trans_b)`` layout the call states.""" + + request = GemmCall + + @abstractmethod + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Multiply the two matrices; nothing is written in place. + + Both operands are contiguous on ``call.device`` in ``call.dtype``, which is + ``float16`` or ``bfloat16``; the contraction accumulates in ``float32``. + + Args: + a: ``(call.m, call.k)``, or ``(call.k, call.m)`` when ``call.trans_a``. + b: ``(call.n, call.k)`` when ``call.trans_b``, else ``(call.k, call.n)``. + + Returns: + A new ``(call.m, call.n)`` tensor in ``call.dtype``. + """ +``` + +实现是同时继承 `Kernel` 与某一个接口的类,以一个 key 列在 `kernel_types` 中。由哪个实现服务一次调用,由各实现自己决定,不由算子决定,依据以下四项声明: + +| # | 声明 | 说明什么 | 不声明时 | +| --- | --- | --- | --- | +| 1 | `devices`、`supported_archs` | 实现能在哪些设备上运行 | CUDA 设备,全部架构 | +| 2 | `applies(call)`、`refusal(call)` | 实现服务哪些调用,正面写出 | 服务全部调用 | +| 3 | `general`、`preferred_over` | 两个实现都服务同一次调用时谁胜出 | 不胜过任何实现 | +| 4 | `entry_for(call)` | build identity,以及每个 identity 只跑一次的构造方法 | 以整个 call spec 为 identity,用 `cls(call)` 构造 | + +可用性先过滤。在剩下的、且适用的实现中,`general` 低于其他所有实现,`preferred_over` 列出本类胜过的 key。一个不剩时报 `no implementation serves this call`,没有任何 key 能在这次调用的设备类型上运行时报 `OpNotAvailableError`;剩下两个互相没有优先关系时报 `dispatch is ambiguous`。声明顺序不起作用。一个实现不在自己的 `applies` 里排除另一个实现的范围,应当胜出的一方声明 `preferred_over`。 + +`GemmFwdOp` 的三个实现这样分割 `"gemm"` 接口: + +| # | key | 服务 | 声明 | +| --- | --- | --- | --- | +| 1 | `gemm_tma` | 操作数能被 TMA 寻址的 SM90 形状 | `supported_archs = [90]`,以及给出未对齐原因的 `refusal` | +| 2 | `gemv` | 沿 K 规约、最多两行,在 CUDA core 上规约更快的形状 | `supported_archs = [90]`、经 `band_for` 实现的 `applies`、`preferred_over = frozenset({"gemm_tma"})` | +| 3 | `gemm_cp_async` | 其余两个都不认领的全部形状,下限是一行 K 至少占满一次 4 字节读取 | `supported_archs = [80, 86, 89, 90]`、`general = True`,以及拒绝更窄 K 行的 `refusal` | + +`entry_for(call)` 返回两样东西:两次调用要共享什么才算同一个 kernel 的 **build identity**,以及每个 identity 只跑一次的**构造方法**。identity 带少了,第二种 dtype 会复用第一种 dtype 的 kernel;kernel 只依赖其中几个量却把整个形状带上,就变成一个形状编译一次。 + +只有一个实现的接口,除继承接口外不需要别的声明。`RMSNormKernel` 就是 `RMSNormFwdOp` 派发的全部内容([`src/tileops/kernels/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/norm/rms_norm.py)): + +```python +class RMSNormKernel(Kernel, RMSNormFwdInterface): + supported_archs: list[int] = [80, 86, 89, 90] + + @classmethod + def entry_for(cls, call: LayerNormCall) -> Entry: + identity = (call.n, call.eps, call.dtype) + return identity, lambda: cls(*identity) ``` -完全没有自带实现、只依赖外部后端的算子,`kernel_types` 与 `entry_for` 都不写;在没有 target 认领设备时,调用会抛 `OpNotAvailableError`。 +算子自己不写 `entry_for`,也不自建 kernel 缓存:没有缓存字典,也不以某个属性是否已赋值来决定要不要构造。把 `kernel_for` 的返回值存进 `self.kernel` 不算自建缓存。 + +完全没有自带实现、只依赖外部后端的算子,`kernel_types` 与 `interfaces` 都不写;在没有 target 认领设备时,调用会抛 `OpNotAvailableError`。 + +后端可以为一个接口新增实现,也可以替换一个 key 背后的类,这两件事都不改动 TileOPs,见[接入一类新硬件](backends.md)。 ### 注册 @@ -130,14 +188,14 @@ def entry_for(self, role, call): # call is the input dtype ## 第三步:写 kernel -kernel 类继承 [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py),放在 [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 下,用 TileLang 编写,在构造时编译。它实现 `forward`,基类的 `__call__` 会调用它。构造函数由它的 `entry_for` 构造方法调用,调用签名就是第二步里的 `kernel(a, b)`。 +kernel 类继承 [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) 与它实现的那个接口,放在 [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 下,用 TileLang 编写,在构造时编译。它实现 `forward`,基类的 `__call__` 会调用它。构造函数由本类 `entry_for` 给出的构造方法调用,`forward` 接受接口规定的参数,就是第二步里的 `kernel(a, b)`。 它是这六处里唯一不受 spec 约束的一处:kernel 不读 spec,也不对照 spec 检查。 构造参数与调用参数的划分有一条硬性要求:**只有会被编译进生成代码的值才进构造函数。** `GemmTmaKernel` 是这样分的: ```python -class GemmTmaKernel(Kernel): +class GemmTmaKernel(Kernel, GemmFwdInterface): def __init__(self, m, n, k, dtype, config=None, tune=False, trans_a=False, trans_b=False, ...): self.kernel = _gemm_kernel(m, n, k, trans_a, trans_b, self.dtype_str, ...) # compiles self.init_config(config, tune) # tile sizes and pipeline depth @@ -159,7 +217,7 @@ kernel = AttnKernel(num_heads, head_dim, dtype) out = kernel(q, k, v) # seq_len 从张量形状里读 ``` -上一种写法下,`entry_for` 返回的身份里带着 `seq_len`,每步都未命中、每步都编译一次,decode 直接跑不动。 +上一种写法下,`entry_for` 返回的 build identity 里带着 `seq_len`,每步都未命中、每步都编译一次,decode 直接跑不动。 ## 第四步:写测试 diff --git a/docs/torch-compile.md b/docs/torch-compile.md index 2ba0f197..61dbd626 100644 --- a/docs/torch-compile.md +++ b/docs/torch-compile.md @@ -70,8 +70,8 @@ compiled path behave differently from the eager one. that follows it. See [what each phase may do](backends.md#phase-limits). - **A second card may need its own build.** For a call a target serves, the device is part of the kernel's memo key, so the same instance builds again on a second card. An - in-tree kernel is keyed by what its `entry_for` returns, which includes the device only - when the build depends on it. A `target=` named in the constructor + in-tree kernel is keyed by the build identity the selected implementation's `entry_for` + returns, which includes the device only when the build depends on it. A `target=` named in the constructor is honoured on the first compiled call too, and a failed build pins the op to no target. @@ -93,7 +93,7 @@ The code bringing one op in takes: how the boundary is declared, how the fake is written, and why the target is resolved again inside the node. For the tracing, graph breaks and guards below, see [how dynamo works](#how-dynamo-works). -`RMSNormFwdOp` was the first op brought in. Its skeleton, method bodies elided — the full +`RMSNormFwdOp` was the first op brought in. Its skeleton, docstrings elided — the full file is [`src/tileops/ops/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/norm/rms_norm.py): @@ -102,15 +102,22 @@ class RMSNormFwdOp(Op): # the operators, their fakes and compile_op_names are generated from the manifest entry compile_boundary: ClassVar[bool] = True kernel_types: ClassVar[Mapping[str, type[Kernel]]] = {"rms_norm": RMSNormKernel} + interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = {"rms_norm": RMSNormFwdInterface} def forward(self, x, weight=None): # the only line: call the generated operator return self._call_boundary(x, weight) def _eager_forward(self, x, weight=None): - ... # the generated checks have run; make contiguous - kernel = self.kernel_for("rms_norm", (x, weight), x.dtype) - return kernel(x, weight) + weight = None if weight is None else weight.contiguous() + x = x.contiguous() # the generated checks have run + call = LayerNormCall( + device=x.device, + n=math.prod(self.normalized_shape), + eps=torch.finfo(torch.float32).eps if self.eps is None else float(self.eps), + dtype=x.dtype, + ) + return self.kernel_for("rms_norm", call)(x, weight) ``` That is the whole declaration. The operators and their fakes are generated from the diff --git a/docs/torch-compile.zh.md b/docs/torch-compile.zh.md index a35941a9..11dbd8c0 100644 --- a/docs/torch-compile.zh.md +++ b/docs/torch-compile.zh.md @@ -51,7 +51,7 @@ block(x, w) - **不要依赖 stride 原样传递。** 算子不写入的非连续输入在节点内部连续化,算子自己分配的输出恒为连续张量;后续计算需要别的布局,在算子之外自行转换。输出若就是被写入的输入(`alias`)或调用方给的 `out`,沿用那个张量的存储。 - **不能用 meta 张量预热。** 有了边界,传入 meta 或 fake 张量的调用就在 fake 处返回,走不到构造 kernel 那一步。 - **CUDA graph 捕获之前先行预热。** 用真实张量、相同形状至少调用一次:构造 kernel 允许编译,捕获期间只允许查表命中后直接调用。各阶段分别允许执行哪些操作,见[各阶段允许做什么](backends.md#phase-limits)。 -- **换一块卡可能要重新构造。** 由 target 服务的调用,设备是 kernel 记忆键的一部分,同一个实例换到另一块卡上会重新构造一次。自带 kernel 的记忆键是它的 `entry_for` 的返回值,只有构造结果与设备有关时才包含设备。构造函数里指名的 `target=` 在首次编译调用中同样生效;构造失败不会把算子固定到任何 target。 +- **换一块卡可能要重新构造。** 由 target 服务的调用,设备是 kernel 记忆键的一部分,同一个实例换到另一块卡上会重新构造一次。自带 kernel 的记忆键是选中实现的 `entry_for` 返回的 build identity,只有构造结果与设备有关时才包含设备。构造函数里指名的 `target=` 在首次编译调用中同样生效;构造失败不会把算子固定到任何 target。 ### 接入之后成立的三项保证 @@ -65,22 +65,29 @@ block(x, w) 接入一个算子要写的代码:边界怎么声明、fake 怎么写、target 判定为什么要在节点内部重做一次。其中的追踪、切图、guard 见[dynamo 是怎么工作的](#dynamo)。 -`RMSNormFwdOp` 是仓内第一个接入的算子。下面是它的骨架,方法体一律省略,完整代码见 [`src/tileops/ops/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/norm/rms_norm.py): +`RMSNormFwdOp` 是仓内第一个接入的算子。下面是它的骨架,略去 docstring,完整代码见 [`src/tileops/ops/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/norm/rms_norm.py): ```python class RMSNormFwdOp(Op): # the operators, their fakes and compile_op_names are generated from the manifest entry compile_boundary: ClassVar[bool] = True kernel_types: ClassVar[Mapping[str, type[Kernel]]] = {"rms_norm": RMSNormKernel} + interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = {"rms_norm": RMSNormFwdInterface} def forward(self, x, weight=None): # the only line: call the generated operator return self._call_boundary(x, weight) def _eager_forward(self, x, weight=None): - ... # the generated checks have run; make contiguous - kernel = self.kernel_for("rms_norm", (x, weight), x.dtype) - return kernel(x, weight) + weight = None if weight is None else weight.contiguous() + x = x.contiguous() # the generated checks have run + call = LayerNormCall( + device=x.device, + n=math.prod(self.normalized_shape), + eps=torch.finfo(torch.float32).eps if self.eps is None else float(self.eps), + dtype=x.dtype, + ) + return self.kernel_for("rms_norm", call)(x, weight) ``` 声明就这么多。operator 与它的 fake 都从条目生成,每个副作用分支一个 operator:张量参数是 `signature.inputs` 的顺序,返回什么看 `signature.outputs`,写哪些参数恰好是标了 `mutated` 的输入,每个输出的形状与 dtype 取自签名。名字是 `tileops::_`(类名本身以 family 开头时只写一次),这里就是 `tileops::norm_rms_norm_fwd`;一个分支的 operator 若写入某个输入、填写 `buffer: out` 或不产出某个输出,名字后面就按这个顺序分别加上 `_writes_<输入>`、`_out`、`_without_<输出>`。没有算子自己起名字,`compile_op_names` 也就不可能和注册的名字对不上。 From 5ad3ec7268e80f87cb06b985b75357caec235a37 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Thu, 1 Oct 2026 13:09:59 +0800 Subject: [PATCH 2/5] [Docs] State where a one-file family keeps its interface, and the name lint TileOPs #2352 fixed the interface naming rule and added interface-names-lint after this branch was written. --- docs/new-op.md | 4 +++- docs/new-op.zh.md | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/new-op.md b/docs/new-op.md index 6fab1b92..a8c6e867 100644 --- a/docs/new-op.md +++ b/docs/new-op.md @@ -147,7 +147,9 @@ The kernel that comes back is called with the parameters of the interface's abst An interface is a class in [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels), -beside the call spec it names in `request`. Its abstract `forward` is the whole contract +beside the call spec it names in `request`; a family with one kernel file keeps both in +that file. Its name is `{Name}{Fwd|Bwd}Interface`, with variant words before the +direction, which `interface-names-lint` checks. Its abstract `forward` is the whole contract an implementation — in-tree or from a backend — is written against, so its docstring states each tensor's shape, dtype, layout, device and whether it is written in place: diff --git a/docs/new-op.zh.md b/docs/new-op.zh.md index 48882d53..2d92b6ad 100644 --- a/docs/new-op.zh.md +++ b/docs/new-op.zh.md @@ -119,7 +119,7 @@ kernel 是编译产物,构造一次要几百毫秒到几秒,而一个算子 取回的 kernel 按接口抽象 `forward` 的参数表、按同样的顺序调用。 -接口是写在 [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 里的一个类,与它在 `request` 中指名的 call spec 放在一起。它的抽象 `forward` 是各实现(自带的与后端提供的)唯一依据的契约,docstring 因此写明每个张量的形状、dtype、内存布局、设备,以及是否被原地写入: +接口是写在 [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 里的一个类,与它在 `request` 中指名的 call spec 放在一起;只有一个 kernel 文件的 family 把两者都写在那个文件里。名字是 `{Name}{Fwd|Bwd}Interface`,变体词写在方向之前,由 `interface-names-lint` 检查。它的抽象 `forward` 是各实现(自带的与后端提供的)唯一依据的契约,docstring 因此写明每个张量的形状、dtype、内存布局、设备,以及是否被原地写入: ```python class GemmFwdInterface(KernelInterface): From bd89ba9689837ff0a34f9517a8a9f1bc8540a836 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Thu, 1 Oct 2026 13:21:45 +0800 Subject: [PATCH 3/5] [Docs] Give the dispatch prose its subjects back The Chinese dropped subjects and compressed clauses to the point of reading as notes: the dispatcher, the op and the implementation are now named where they act. Three places defined a thing by what it is not and now say what it is. Replaced the invented word for making a tensor contiguous, and the one for the builder, with what they do. --- docs/backends.zh.md | 14 +++++++------- docs/new-op.md | 31 ++++++++++++++++--------------- docs/new-op.zh.md | 24 ++++++++++++------------ docs/torch-compile.zh.md | 2 +- 4 files changed, 36 insertions(+), 35 deletions(-) diff --git a/docs/backends.zh.md b/docs/backends.zh.md index d1204e91..1fd592ac 100644 --- a/docs/backends.zh.md +++ b/docs/backends.zh.md @@ -6,13 +6,13 @@ TileLang 是多后端 DSL,每种硬件各有一套独立的 kernel,由各自 **后端只提供一件事:一个能算这次调用的可调用对象。** 其余都由算子层负责。 -本页讲的是 target,三种接入方式中范围最大的一种。前半按动手顺序列出要做的事:要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么,以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。 +本页讲的是 target,三种接入方式中范围最大的一种。前半部分按实现顺序列出要做的事:要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么,以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。 后半说明协议何以如此设计:两层选择、算子层的契约、kernel 的重建条件、调用方可用的接口,以及刻意不支持的情形。 ## 三种接入方式 {#three-ways} -仓外的包要接管一个算子的多少调用,决定用三种方式中的哪一种。范围较小的两种,写出的 kernel 类依据的是 [kernel 接口](new-op.md#kernel-selection),与自带实现依据的是同一份契约;target 依据的则是算子在 manifest 中的签名,写一个 `build_kernel`。 +一个仓库之外的包要接管某个算子的多少调用,决定了它用三种方式中的哪一种。范围较小的两种,写出的 kernel 类遵守 [kernel 接口](new-op.md#kernel-selection),与自带实现遵守同一份契约;target 遵守的则是算子在 manifest 中的签名,写一个 `build_kernel`。 | # | | `kernel_map=` | `register_implementation` | target | | --- | --- | --- | --- | --- | @@ -48,7 +48,7 @@ op = LayerNormFwdOp((32,), kernel_map={"layer_norm": _TorchLayerNorm}, target=BU `kernel_map=` 中本算子没有、而其他算子有的 key 被忽略,复合算子由此把一份映射传给各个子算子;任何算子都没有的 key 在构造时报 `was given kernel_map keys no op has`。 -**`register_implementation(op, key, implementation)`** 不是替换,而是新增一个实现。`op` 是算子在 manifest 中的 key,`key` 是新实现的名字,它加入哪个接口由这个类继承了哪个接口决定。新实现自己声明适用范围,与某个不是 `general` 的自带实现重叠时要声明 `preferred_over`: +**`register_implementation(op, key, implementation)`** 为一个接口新增一个实现。`op` 是算子在 manifest 中的 key,`key` 是新实现的名字,它加入哪个接口由这个类继承了哪个接口决定。新实现自己声明适用范围,与某个不是 `general` 的自带实现重叠时要声明 `preferred_over`: ```python class _NarrowTorchLayerNorm(_TorchLayerNorm): @@ -64,9 +64,9 @@ class _NarrowTorchLayerNorm(_TorchLayerNorm): register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) ``` -这里 `n <= 64` 由 `_NarrowTorchLayerNorm` 服务,`n = 1024` 仍由自带的 `LayerNormKernel` 服务。新增只进入注册之后构造的算子实例。同一个算子下重复注册同一个 key 报 `BackendError`;key 与自带实现的 key 相同时,构造实例时报 `reuse keys it has`。 +这样 `n <= 64` 的调用由 `_NarrowTorchLayerNorm` 服务,`n = 1024` 仍由自带的 `LayerNormKernel` 服务。新增的实现只进入注册之后构造的算子实例。同一个算子下重复注册同一个 key 报 `BackendError`;key 与自带实现的 key 相同时,构造实例时报 `reuse keys it has`。 -`register_implementation` 在后端模块被导入时执行,入口与 target 用的是同一条 entry point;`kernel_map=` 不注册任何东西,由调用方在构造算子时传入。以下各节讲 target。 +`register_implementation` 在后端模块被导入时执行,入口与 target 用的是同一条 entry point;`kernel_map=` 则由调用方在构造算子时传入,不经过注册。以下各节讲 target。 ## 写一个后端要做的四件事 @@ -245,7 +245,7 @@ register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) - **`kernel_for` 与各实现的 `entry_for` 只服务自带实现。** 它们决定取哪个自带 kernel、按什么查表、又怎么构造。target 选中后端时整个算子由 target 服务,这几处都不会执行。 - **张量按位置传,参数按名字传。** `build_kernel(*inputs, **params)`:位置实参是 `TensorSpec`(没传的可选输入是 `None`),关键字实参是 manifest 里 `params` 的名字与本次调用的确定值。 - **一个 `(算子, target)` 只注册一个 builder。** 自带实现内部分几种 kernel(GEMM 的 `kernel_types` 里有三个)不会传进来,`build_kernel` 从 `TensorSpec` 自行判断该返回哪个 kernel。 -- **不必自己做记忆。** 同一个设备与输入签名,算子层不会再调第二次;要更细的区分或更少的重建,在 `build_kernel` 内部另加一层缓存。写来依赖外部后端的算子,`kernel_types` 与 `interfaces` 都不声明,那时没有 target 认领设备,调用直接抛 `OpNotAvailableError`。 +- **不必自己缓存构建结果。** 设备与输入签名相同时,算子层不会再调用一次 `build_kernel`;需要更细的区分或更少的重建,在 `build_kernel` 内部另加一层缓存。专为外部后端而写的算子,`kernel_types` 与 `interfaces` 都不声明,没有 target 认领设备时调用直接抛 `OpNotAvailableError`。 ## 实现一个可运行的后端 {#runnable} @@ -529,7 +529,7 @@ TileOPs 不解析 `torch.device`,而是把它原样传给 `detect`。这样做 TileOPs 自带的 kernel([`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels))是**默认实现**:它没有 target 名,也不进注册表。 -**默认状态是不替换。** 没有装后端、没有指名 `target=`、也没有设置进程默认值时,调用走自带实现。协议不预置任何替换:装上一个后端之后,它的 `detect` 认领了这块设备,或者被 `target=` 与 `set_default_target` 指名,该算子的 kernel 才换成后端的那一套。 +**默认走自带实现。** 没有装后端、没有指名 `target=`、也没有设置进程默认值时,调用由自带实现服务。装上一个后端之后,要它的 `detect` 认领了这块设备,或者它被 `target=` 与 `set_default_target` 指名,该算子的 kernel 才换成这个后端的那一套。 ## kernel 的重建条件 {#memo} diff --git a/docs/new-op.md b/docs/new-op.md index a8c6e867..2c6076ca 100644 --- a/docs/new-op.md +++ b/docs/new-op.md @@ -134,12 +134,12 @@ Its two arguments: called. `GemmFwdOp` calls a kernel in one place, so it declares one, `"gemm"`. A second interface is opened only where the semantics or the call contract changes: `BatchNormFwdOp` has `batch_norm_fwd_train` and `batch_norm_fwd_infer`, which return - different things. A faster kernel for some shape range or some architecture is not a - new interface — it is another implementation of the one that is there. + different things. A faster kernel for some shape range or some architecture is another + implementation of the interface already there. - **`call`** — a frozen `CallSpec` subclass carrying the facts of this call that selection and building read: shapes, the dtype, the op's semantic parameters, and the device. It - has to be the interface's `request` type. Device facts (`arch`, `sm_count`, - `calibration`, `smem_budget`) are not passed in; the dispatcher derives them from + has to be the interface's `request` type. The dispatcher derives the device facts (`arch`, + `sm_count`, `calibration`, `smem_budget`) itself, from `call.device` on a miss. The kernel that comes back is called with the parameters of the interface's abstract @@ -176,8 +176,8 @@ class GemmFwdInterface(KernelInterface): ``` An implementation is a class inheriting both `Kernel` and one interface, listed in -`kernel_types` under a key. Which implementation serves a call is decided by the -implementations, not by the op, from four declarations: +`kernel_types` under a key. Which implementation serves a call follows from four +declarations the implementations make about themselves; the op makes none of them: | # | Declaration | States | Left undeclared | | --- | --- | --- | --- | @@ -186,21 +186,22 @@ implementations, not by the op, from four declarations: | 3 | `general`, `preferred_over` | which implementation wins where two of them serve one call | wins over none | | 4 | `entry_for(call)` | the build identity, and the builder that runs once per identity | the whole call spec, built by `cls(call)` | -Availability filters first. Among the implementations that are left and that apply, -`general` loses to every other one and `preferred_over` names the keys its class wins -over. Nothing left raises `no implementation serves this call`, or `OpNotAvailableError` where -no key runs on the call's device type at all; two with no relation between them raise -`dispatch is ambiguous`; declaration order decides nothing. An -implementation never excludes a sibling in its own `applies` — the one that should win -declares `preferred_over`. +The dispatcher filters by availability first, then picks the single winner among the +implementations that are left and that apply: `general` loses to every other one, and the +rest are compared by the keys each names in `preferred_over`. Nothing left raises +`no implementation serves this call`, or `OpNotAvailableError` where no key runs on the +call's device type at all; two with no relation between them raise `dispatch is +ambiguous`. Declaration order decides nothing. Where one implementation should give up a +range to another, the one that should win declares `preferred_over`, rather than the +other one excluding that range in its own `applies`. `GemmFwdOp`'s three implementations divide the `"gemm"` interface like this: | # | Key | Serves | Declares | | --- | --- | --- | --- | | 1 | `gemm_tma` | SM90 shapes whose operands TMA can address | `supported_archs = [90]`, and a `refusal` naming the misalignment | -| 2 | `gemv` | at most two rows contracted over K, where reducing on CUDA cores wins | `supported_archs = [90]`, `applies` through `band_for`, `preferred_over = frozenset({"gemm_tma"})` | -| 3 | `gemm_cp_async` | every shape the other two do not claim, down to a K row spanning one four-byte load | `supported_archs = [80, 86, 89, 90]`, `general = True`, and a `refusal` for a narrower K row | +| 2 | `gemv` | shapes of at most two rows contracted over K, where reducing on CUDA cores wins | `supported_archs = [90]`, `applies` through `band_for`, `preferred_over = frozenset({"gemm_tma"})` | +| 3 | `gemm_cp_async` | every shape the other two do not claim, provided a K row spans at least one four-byte load | `supported_archs = [80, 86, 89, 90]`, `general = True`, and a `refusal` for a narrower K row | `entry_for(call)` returns the **identity** two calls must share to reuse one kernel, and the **builder** that runs once per identity. Carry too little in the identity and a second diff --git a/docs/new-op.zh.md b/docs/new-op.zh.md index 2d92b6ad..7a86999e 100644 --- a/docs/new-op.zh.md +++ b/docs/new-op.zh.md @@ -101,8 +101,8 @@ class GemmFwdOp(Op): | 2 | `kernel_types` | 能服务这个算子的 Kernel 类,各起一个 key;`kernel_map=` 按这个 key 替换其中一个 | | 3 | `interfaces` | 算子调用 kernel 的每一处各占一条:`kernel_for` 用的名字 → 这一处的各实现所继承的 `KernelInterface` 类 | | 4 | `forward` | `signature.inputs` 的顺序,可选输入排在最后、默认 `None` | -| 5 | `_eager_forward` | 连续化、构造 call spec、取 kernel、launch kernel | -| 6 | `compute_roof` | 可选:给算子 FLOPs 定价的 GPU profile 单元,不是 CUDA core fp32 时才写 | +| 5 | `_eager_forward` | 把输入变成连续的,构造 call spec,取出 kernel,再调用它 | +| 6 | `compute_roof` | 可选。算子的 FLOPs 按哪个硬件单元的峰值定价,默认是 CUDA core 上的 fp32,用别的单元时才写 | `_infer_output_shapes`、`_validate_dtypes` 与 `eval_roofline` 都照 spec 生成,不用写。 @@ -112,12 +112,12 @@ class GemmFwdOp(Op): kernel 是编译产物,构造一次要几百毫秒到几秒,而一个算子实例会被反复调用,形状与 dtype 各不相同。算子层因此维护一张记忆表:本次调用要的 kernel 已经构造过就取回来,没有才构造并存进去。`kernel_for` 是自带实现走到这张表的唯一入口;[target](backends.md) 服务的是整个算子,不经过它。 -两个参数: +`kernel_for` 接受两个参数: -- **`interface`**:`interfaces` 的一个 key,指算子调用 kernel 的一处。`GemmFwdOp` 只在一处调用 kernel,因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新开一个接口:`BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的东西不同。某个形状范围或某个架构上更快的 kernel 不是新接口,而是已有接口的另一个实现。 -- **`call`**:一个冻结的 `CallSpec` 子类,带着这次调用中选择与构建会读取的事实:形状、dtype、算子的语义参数与设备。它必须是这个接口 `request` 指定的类型。设备事实(`arch`、`sm_count`、`calibration`、`smem_budget`)不由调用方传入,未命中时派发机制从 `call.device` 推出。 +- **`interface`**:`interfaces` 的一个 key,指算子调用 kernel 的一处。`GemmFwdOp` 只在一处调用 kernel,因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新开一个接口:`BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的东西不同。某个形状范围或某个架构上更快的 kernel 是已有接口的另一个实现。 +- **`call`**:一个冻结的 `CallSpec` 子类,带着这次调用中选择与构建会读取的事实:形状、dtype、算子的语义参数与设备。它必须是这个接口 `request` 指定的类型。设备事实(`arch`、`sm_count`、`calibration`、`smem_budget`)由派发机制在未命中时从 `call.device` 推出,调用方不填。 -取回的 kernel 按接口抽象 `forward` 的参数表、按同样的顺序调用。 +算子用接口抽象 `forward` 声明的参数、按同样的顺序调用取回的 kernel。 接口是写在 [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 里的一个类,与它在 `request` 中指名的 call spec 放在一起;只有一个 kernel 文件的 family 把两者都写在那个文件里。名字是 `{Name}{Fwd|Bwd}Interface`,变体词写在方向之前,由 `interface-names-lint` 检查。它的抽象 `forward` 是各实现(自带的与后端提供的)唯一依据的契约,docstring 因此写明每个张量的形状、dtype、内存布局、设备,以及是否被原地写入: @@ -143,26 +143,26 @@ class GemmFwdInterface(KernelInterface): """ ``` -实现是同时继承 `Kernel` 与某一个接口的类,以一个 key 列在 `kernel_types` 中。由哪个实现服务一次调用,由各实现自己决定,不由算子决定,依据以下四项声明: +实现是同时继承 `Kernel` 与某一个接口的类,以一个 key 列在 `kernel_types` 中。一次调用由哪个实现服务,取决于各实现自己的四项声明,算子不参与: -| # | 声明 | 说明什么 | 不声明时 | +| # | 声明 | 说明什么 | 未声明时的默认 | | --- | --- | --- | --- | | 1 | `devices`、`supported_archs` | 实现能在哪些设备上运行 | CUDA 设备,全部架构 | | 2 | `applies(call)`、`refusal(call)` | 实现服务哪些调用,正面写出 | 服务全部调用 | | 3 | `general`、`preferred_over` | 两个实现都服务同一次调用时谁胜出 | 不胜过任何实现 | | 4 | `entry_for(call)` | build identity,以及每个 identity 只跑一次的构造方法 | 以整个 call spec 为 identity,用 `cls(call)` 构造 | -可用性先过滤。在剩下的、且适用的实现中,`general` 低于其他所有实现,`preferred_over` 列出本类胜过的 key。一个不剩时报 `no implementation serves this call`,没有任何 key 能在这次调用的设备类型上运行时报 `OpNotAvailableError`;剩下两个互相没有优先关系时报 `dispatch is ambiguous`。声明顺序不起作用。一个实现不在自己的 `applies` 里排除另一个实现的范围,应当胜出的一方声明 `preferred_over`。 +派发机制先按可用性过滤,再在剩下的、且适用的实现中选出唯一的胜者:`general` 的实现低于其他所有实现,其余按各自 `preferred_over` 列出的 key 比较。一个实现都不剩时报 `no implementation serves this call`;没有任何 key 能在这次调用的设备类型上运行时报 `OpNotAvailableError`;剩下的实现之间互相没有优先关系时报 `dispatch is ambiguous`。实现的声明顺序不影响选择结果。需要让出一段范围时,由应当胜出的一方声明 `preferred_over`,而不是让另一方在自己的 `applies` 里把这段范围排除掉。 `GemmFwdOp` 的三个实现这样分割 `"gemm"` 接口: | # | key | 服务 | 声明 | | --- | --- | --- | --- | | 1 | `gemm_tma` | 操作数能被 TMA 寻址的 SM90 形状 | `supported_archs = [90]`,以及给出未对齐原因的 `refusal` | -| 2 | `gemv` | 沿 K 规约、最多两行,在 CUDA core 上规约更快的形状 | `supported_archs = [90]`、经 `band_for` 实现的 `applies`、`preferred_over = frozenset({"gemm_tma"})` | -| 3 | `gemm_cp_async` | 其余两个都不认领的全部形状,下限是一行 K 至少占满一次 4 字节读取 | `supported_archs = [80, 86, 89, 90]`、`general = True`,以及拒绝更窄 K 行的 `refusal` | +| 2 | `gemv` | 沿 K 规约且最多两行的形状,这种形状在 CUDA core 上规约更快 | `supported_archs = [90]`、经 `band_for` 实现的 `applies`、`preferred_over = frozenset({"gemm_tma"})` | +| 3 | `gemm_cp_async` | 其余两个都不认领的全部形状,前提是一行 K 至少占满一次 4 字节读取 | `supported_archs = [80, 86, 89, 90]`、`general = True`,以及拒绝更窄 K 行的 `refusal` | -`entry_for(call)` 返回两样东西:两次调用要共享什么才算同一个 kernel 的 **build identity**,以及每个 identity 只跑一次的**构造方法**。identity 带少了,第二种 dtype 会复用第一种 dtype 的 kernel;kernel 只依赖其中几个量却把整个形状带上,就变成一个形状编译一次。 +`entry_for(call)` 返回两样东西:两次调用共享它才算同一个 kernel 的 **build identity**,以及每个 identity 只运行一次的**构建函数**。identity 少带一个量,第二种 dtype 就会复用第一种 dtype 的 kernel;kernel 只依赖其中几个量却把整个形状都带上,就变成每个形状各编译一次。 只有一个实现的接口,除继承接口外不需要别的声明。`RMSNormKernel` 就是 `RMSNormFwdOp` 派发的全部内容([`src/tileops/kernels/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/norm/rms_norm.py)): diff --git a/docs/torch-compile.zh.md b/docs/torch-compile.zh.md index 11dbd8c0..b7f338c7 100644 --- a/docs/torch-compile.zh.md +++ b/docs/torch-compile.zh.md @@ -51,7 +51,7 @@ block(x, w) - **不要依赖 stride 原样传递。** 算子不写入的非连续输入在节点内部连续化,算子自己分配的输出恒为连续张量;后续计算需要别的布局,在算子之外自行转换。输出若就是被写入的输入(`alias`)或调用方给的 `out`,沿用那个张量的存储。 - **不能用 meta 张量预热。** 有了边界,传入 meta 或 fake 张量的调用就在 fake 处返回,走不到构造 kernel 那一步。 - **CUDA graph 捕获之前先行预热。** 用真实张量、相同形状至少调用一次:构造 kernel 允许编译,捕获期间只允许查表命中后直接调用。各阶段分别允许执行哪些操作,见[各阶段允许做什么](backends.md#phase-limits)。 -- **换一块卡可能要重新构造。** 由 target 服务的调用,设备是 kernel 记忆键的一部分,同一个实例换到另一块卡上会重新构造一次。自带 kernel 的记忆键是选中实现的 `entry_for` 返回的 build identity,只有构造结果与设备有关时才包含设备。构造函数里指名的 `target=` 在首次编译调用中同样生效;构造失败不会把算子固定到任何 target。 +- **换一块卡可能要重新构造。** 由 target 服务的调用,设备是 kernel 记忆键的一部分,同一个实例换到另一块卡上会重新构造一次。自带 kernel 的记忆键是选中实现的 `entry_for` 返回的 build identity,只有构建结果与设备有关时才带上设备。构造函数里指名的 `target=` 在首次编译调用中同样生效;构造失败不会把算子固定到任何 target。 ### 接入之后成立的三项保证 From eb8660193adb8d37602ab8a1fab741219ecb5100 Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Thu, 1 Oct 2026 13:26:59 +0800 Subject: [PATCH 4/5] [Docs] Answer the prose review Twelve places where the two languages had drifted into each other: a spatial metaphor for a mapping, an implementation that 'reaches' an instance, a coined 'place' for a kernel call, builder translated two different ways and once as constructor, and two defaults defined by what they are not. --- docs/backends.md | 17 ++++++++++------- docs/backends.zh.md | 12 ++++++------ docs/new-op.md | 16 ++++++++-------- docs/new-op.zh.md | 20 ++++++++++---------- 4 files changed, 34 insertions(+), 31 deletions(-) diff --git a/docs/backends.md b/docs/backends.md index b972a577..c6763b1b 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -11,7 +11,8 @@ devices run your kernels. **A backend supplies one thing: something callable that computes this call.** Everything else is the op layer's. -This page is about a target, the largest of the three ways in. The first half is the +This page is about a target, the one of the three extension mechanisms that covers the +whole op. The first half is the work, in the order it is done: the four things to write, the protocol's four functions, how one call reaches them, a backend that installs and runs as it stands, how to turn the template into a backend for real hardware, the four rules for writing a kernel, what each @@ -21,9 +22,10 @@ The second half is why the protocol looks like this: the two layers of selection layer's contract, when a kernel is rebuilt, what a caller can reach for, and what the protocol deliberately leaves out. -## Three ways in {#three-ways} +## Three ways to extend dispatch {#three-ways} -How much of an op a package outside TileOPs takes over decides which of three it uses. +A package outside TileOPs picks one of three mechanisms by how much of an op it takes +over. The two smaller ones write a kernel class against a [kernel interface](new-op.md#kernel-selection), the same contract the in-tree implementations are written against; a target writes a `build_kernel` against the op's @@ -31,13 +33,13 @@ manifest signature instead. | # | | `kernel_map=` | `register_implementation` | target | | --- | --- | --- | --- | --- | -| 1 | Changes | the class running behind one key; which calls that key serves does not change | adds a key, with its own applicability and precedence | every call of the op | +| 1 | Changes | the class registered under one key; which calls that key serves is unchanged | adds a key, with its own applicability and precedence | every call of the op | | 2 | Applies to | the one op instance the caller constructed it on | every instance of that op constructed afterwards | every instance that settles on the target | | 3 | Written against | the kernel interface | the kernel interface | the op's manifest signature | | 4 | Calls the new class does not serve | an error when that key is selected | still served by the in-tree implementations | none: a target serves them all | **`kernel_map=`** is a constructor argument of every op, a mapping from key to class. It -swaps the class behind that key in this instance; the key keeps the registered +replaces the class registered under that key in this instance; the key keeps the registered implementation's `applies`, `general` and `preferred_over`, and the replacement, like every implementation, inherits the key's interface and is built through its own `entry_for`. A selected key whose replacement cannot serve the call is an error, never a @@ -92,7 +94,8 @@ register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerN ``` Here `n <= 64` goes to `_NarrowTorchLayerNorm` and `n = 1024` stays with the in-tree -`LayerNormKernel`. The addition reaches only op instances constructed after the call. +`LayerNormKernel`. The added implementation applies only to op instances constructed +after the call. Registering the same key twice under one op raises `BackendError`; a key an in-tree implementation already uses raises `reuse keys it has` when an instance is constructed. @@ -674,7 +677,7 @@ third-party backend neither bypasses one nor substitutes its own. The kernels TileOPs ships ([`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)) are the **default implementation**: they have no target name and are not in the registry. -**The default state is no substitution.** With no backend installed, no `target=` named +**The in-tree implementations run by default.** With no backend installed, no `target=` named and no process default set, calls run the shipped implementation. Nothing is preconfigured to substitute: a backend serves an op only once it is installed and either claims the device through its `detect`, or is named by `target=` or `set_default_target`. diff --git a/docs/backends.zh.md b/docs/backends.zh.md index 1fd592ac..4d48f8a7 100644 --- a/docs/backends.zh.md +++ b/docs/backends.zh.md @@ -12,16 +12,16 @@ TileLang 是多后端 DSL,每种硬件各有一套独立的 kernel,由各自 ## 三种接入方式 {#three-ways} -一个仓库之外的包要接管某个算子的多少调用,决定了它用三种方式中的哪一种。范围较小的两种,写出的 kernel 类遵守 [kernel 接口](new-op.md#kernel-selection),与自带实现遵守同一份契约;target 遵守的则是算子在 manifest 中的签名,写一个 `build_kernel`。 +一个仓库之外的包按自己的接管范围,在三种方式中选一种。范围较小的两种,写出的 kernel 类遵守 [kernel 接口](new-op.md#kernel-selection),与自带实现遵守同一份契约;target 遵守的则是算子在 manifest 中的签名,写一个 `build_kernel`。 | # | | `kernel_map=` | `register_implementation` | target | | --- | --- | --- | --- | --- | -| 1 | 改变什么 | 一个 key 背后运行的类;这个 key 服务哪些调用不变 | 新增一个 key,带有它自己的适用范围与优先关系 | 算子的全部调用 | +| 1 | 改变什么 | 一个 key 登记的类;这个 key 服务哪些调用不变 | 新增一个 key,带有它自己的适用范围与优先关系 | 算子的全部调用 | | 2 | 作用于 | 调用方构造的那一个算子实例 | 注册之后构造的所有该算子实例 | 选中该 target 的算子实例 | | 3 | 依据的契约 | kernel 接口 | kernel 接口 | 算子在 manifest 中的签名 | | 4 | 新类不服务的调用 | 该 key 被选中时报错 | 仍由自带实现服务 | 不存在,target 服务全部调用 | -**`kernel_map=`** 是每个算子构造函数的参数,值是从 key 到类的映射。它只替换这个实例中该 key 背后运行的类:这个 key 仍按原来登记的实现声明的 `applies`、`general` 与 `preferred_over` 参与选择,替换者与其他实现一样继承这个 key 所属的接口,由它自己的 `entry_for` 构造。该 key 被选中而替换者不服务这次调用时报错,不退回被替换的实现。下例取自 [`tests/test_kernel_dispatch.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_kernel_dispatch.py): +**`kernel_map=`** 是每个算子构造函数的参数,值是从 key 到类的映射。它只把这个实例中该 key 登记的类换掉:这个 key 仍按原来登记的实现声明的 `applies`、`general` 与 `preferred_over` 参与选择,替换者与其他实现一样继承这个 key 所属的接口,由它自己的 `entry_for` 构造。该 key 被选中而替换者不服务这次调用时报错,不退回被替换的实现。下例取自 [`tests/test_kernel_dispatch.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_kernel_dispatch.py): ```python class _TorchLayerNorm(Kernel, LayerNormFwdInterface): @@ -64,9 +64,9 @@ class _NarrowTorchLayerNorm(_TorchLayerNorm): register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) ``` -这样 `n <= 64` 的调用由 `_NarrowTorchLayerNorm` 服务,`n = 1024` 仍由自带的 `LayerNormKernel` 服务。新增的实现只进入注册之后构造的算子实例。同一个算子下重复注册同一个 key 报 `BackendError`;key 与自带实现的 key 相同时,构造实例时报 `reuse keys it has`。 +这样 `n <= 64` 的调用由 `_NarrowTorchLayerNorm` 服务,`n = 1024` 仍由自带的 `LayerNormKernel` 服务。新增的实现只对注册之后构造的算子实例生效。同一个算子下重复注册同一个 key 报 `BackendError`;key 与自带实现的 key 相同时,构造实例时报 `reuse keys it has`。 -`register_implementation` 在后端模块被导入时执行,入口与 target 用的是同一条 entry point;`kernel_map=` 则由调用方在构造算子时传入,不经过注册。以下各节讲 target。 +`register_implementation` 在后端模块被导入时执行,通过 target 所用的同一个 entry point 触发;`kernel_map=` 则由调用方在构造算子时传入,不经过注册。以下各节讲 target。 ## 写一个后端要做的四件事 @@ -242,7 +242,7 @@ register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) 四点对应关系值得记住: -- **`kernel_for` 与各实现的 `entry_for` 只服务自带实现。** 它们决定取哪个自带 kernel、按什么查表、又怎么构造。target 选中后端时整个算子由 target 服务,这几处都不会执行。 +- **`kernel_for` 与各实现的 `entry_for` 只服务自带实现。** 它们决定取哪个自带 kernel、按什么查表、又怎么构造。算子选中某个 target 之后,整个算子由这个 target 服务,这几处都不会执行。 - **张量按位置传,参数按名字传。** `build_kernel(*inputs, **params)`:位置实参是 `TensorSpec`(没传的可选输入是 `None`),关键字实参是 manifest 里 `params` 的名字与本次调用的确定值。 - **一个 `(算子, target)` 只注册一个 builder。** 自带实现内部分几种 kernel(GEMM 的 `kernel_types` 里有三个)不会传进来,`build_kernel` 从 `TensorSpec` 自行判断该返回哪个 kernel。 - **不必自己缓存构建结果。** 设备与输入签名相同时,算子层不会再调用一次 `build_kernel`;需要更细的区分或更少的重建,在 `build_kernel` 内部另加一层缓存。专为外部后端而写的算子,`kernel_types` 与 `interfaces` 都不声明,没有 target 认领设备时调用直接抛 `OpNotAvailableError`。 diff --git a/docs/new-op.md b/docs/new-op.md index 2c6076ca..0c8cf351 100644 --- a/docs/new-op.md +++ b/docs/new-op.md @@ -108,7 +108,7 @@ class GemmFwdOp(Op): | --- | --- | --- | | 1 | `__init__` | the names, order and defaults in `signature.params`, then `target`, `kernel_map` and `tune`, closing with `self.dispatch_kernel(kernel_map)` | | 2 | `kernel_types` | the Kernel classes that can serve the op, each under a key; a `kernel_map=` override replaces one by that key | -| 3 | `interfaces` | one entry per place the op calls a kernel: the name `kernel_for` uses → the `KernelInterface` class the implementations of that place inherit | +| 3 | `interfaces` | one entry per kernel call the op makes: the name `kernel_for` uses → the `KernelInterface` class every implementation of that call inherits | | 4 | `forward` | `signature.inputs` — its order, optional inputs last with default `None` | | 5 | `_eager_forward` | contiguity, the call spec, fetching the kernel and launching it | | 6 | `compute_roof` | optional: the GPU-profile unit that prices the op's FLOPs, where it is not CUDA-core fp32 | @@ -130,14 +130,13 @@ the in-tree path; a [target](backends.md) serves the whole op instead and never Its two arguments: -- **`interface`** — a key of `interfaces`, naming one place in the op where a kernel is - called. `GemmFwdOp` calls a kernel in one place, so it declares one, `"gemm"`. A second - interface is opened only where the semantics or the call contract changes: +- **`interface`** — a key of `interfaces`, naming one kernel call the op makes. + `GemmFwdOp` makes one, so it declares one, `"gemm"`. A second interface is opened only where the semantics or the call contract changes: `BatchNormFwdOp` has `batch_norm_fwd_train` and `batch_norm_fwd_infer`, which return different things. A faster kernel for some shape range or some architecture is another implementation of the interface already there. -- **`call`** — a frozen `CallSpec` subclass carrying the facts of this call that selection - and building read: shapes, the dtype, the op's semantic parameters, and the device. It +- **`call`** — a frozen `CallSpec` subclass carrying the facts needed to select and build + the kernel: shapes, the dtype, the op's semantic parameters, and the device. It has to be the interface's `request` type. The dispatcher derives the device facts (`arch`, `sm_count`, `calibration`, `smem_budget`) itself, from `call.device` on a miss. @@ -195,7 +194,7 @@ ambiguous`. Declaration order decides nothing. Where one implementation should g range to another, the one that should win declares `preferred_over`, rather than the other one excluding that range in its own `applies`. -`GemmFwdOp`'s three implementations divide the `"gemm"` interface like this: +`GemmFwdOp`'s three implementations cover the `"gemm"` interface's calls like this: | # | Key | Serves | Declares | | --- | --- | --- | --- | @@ -230,7 +229,8 @@ An op with no in-tree implementation, written to depend on a backend, declares n `kernel_types` nor `interfaces`; a call on a device no target claims then raises `OpNotAvailableError`. -A backend adds an implementation to an interface, or replaces the class behind one key, +A backend adds an implementation to an interface, or replaces the class registered under +one key, without changing TileOPs; both are in [adding a hardware backend](backends.md). ### Registering diff --git a/docs/new-op.zh.md b/docs/new-op.zh.md index 7a86999e..0cc94e3b 100644 --- a/docs/new-op.zh.md +++ b/docs/new-op.zh.md @@ -99,7 +99,7 @@ class GemmFwdOp(Op): | --- | --- | --- | | 1 | `__init__` | `signature.params` 的名字、顺序与默认值,再加 `target`、`kernel_map`、`tune`;结尾调用 `self.dispatch_kernel(kernel_map)` | | 2 | `kernel_types` | 能服务这个算子的 Kernel 类,各起一个 key;`kernel_map=` 按这个 key 替换其中一个 | -| 3 | `interfaces` | 算子调用 kernel 的每一处各占一条:`kernel_for` 用的名字 → 这一处的各实现所继承的 `KernelInterface` 类 | +| 3 | `interfaces` | 算子发出的每一个 kernel 调用各占一条:`kernel_for` 用的名字 → 这个调用的各实现所继承的 `KernelInterface` 类 | | 4 | `forward` | `signature.inputs` 的顺序,可选输入排在最后、默认 `None` | | 5 | `_eager_forward` | 把输入变成连续的,构造 call spec,取出 kernel,再调用它 | | 6 | `compute_roof` | 可选。算子的 FLOPs 按哪个硬件单元的峰值定价,默认是 CUDA core 上的 fp32,用别的单元时才写 | @@ -114,10 +114,10 @@ kernel 是编译产物,构造一次要几百毫秒到几秒,而一个算子 `kernel_for` 接受两个参数: -- **`interface`**:`interfaces` 的一个 key,指算子调用 kernel 的一处。`GemmFwdOp` 只在一处调用 kernel,因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新开一个接口:`BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的东西不同。某个形状范围或某个架构上更快的 kernel 是已有接口的另一个实现。 -- **`call`**:一个冻结的 `CallSpec` 子类,带着这次调用中选择与构建会读取的事实:形状、dtype、算子的语义参数与设备。它必须是这个接口 `request` 指定的类型。设备事实(`arch`、`sm_count`、`calibration`、`smem_budget`)由派发机制在未命中时从 `call.device` 推出,调用方不填。 +- **`interface`**:`interfaces` 的一个 key,指算子发出的某一个 kernel 调用。`GemmFwdOp` 只发出一个 kernel 调用,因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新开一个接口:`BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的东西不同。某个形状范围或某个架构上更快的 kernel 是已有接口的另一个实现。 +- **`call`**:一个冻结的 `CallSpec` 子类,带着选择实现和构造 kernel 所需的信息:形状、dtype、算子的语义参数与设备。它必须是这个接口 `request` 指定的类型。设备事实(`arch`、`sm_count`、`calibration`、`smem_budget`)由派发机制在未命中时从 `call.device` 推出,调用方不填。 -算子用接口抽象 `forward` 声明的参数、按同样的顺序调用取回的 kernel。 +算子按接口的抽象 `forward` 方法所声明的参数、以同样的顺序调用取回的 kernel。 接口是写在 [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 里的一个类,与它在 `request` 中指名的 call spec 放在一起;只有一个 kernel 文件的 family 把两者都写在那个文件里。名字是 `{Name}{Fwd|Bwd}Interface`,变体词写在方向之前,由 `interface-names-lint` 检查。它的抽象 `forward` 是各实现(自带的与后端提供的)唯一依据的契约,docstring 因此写明每个张量的形状、dtype、内存布局、设备,以及是否被原地写入: @@ -150,11 +150,11 @@ class GemmFwdInterface(KernelInterface): | 1 | `devices`、`supported_archs` | 实现能在哪些设备上运行 | CUDA 设备,全部架构 | | 2 | `applies(call)`、`refusal(call)` | 实现服务哪些调用,正面写出 | 服务全部调用 | | 3 | `general`、`preferred_over` | 两个实现都服务同一次调用时谁胜出 | 不胜过任何实现 | -| 4 | `entry_for(call)` | build identity,以及每个 identity 只跑一次的构造方法 | 以整个 call spec 为 identity,用 `cls(call)` 构造 | +| 4 | `entry_for(call)` | build identity,以及每个 identity 只运行一次的 builder | 以整个 call spec 为 identity,用 `cls(call)` 构造 | 派发机制先按可用性过滤,再在剩下的、且适用的实现中选出唯一的胜者:`general` 的实现低于其他所有实现,其余按各自 `preferred_over` 列出的 key 比较。一个实现都不剩时报 `no implementation serves this call`;没有任何 key 能在这次调用的设备类型上运行时报 `OpNotAvailableError`;剩下的实现之间互相没有优先关系时报 `dispatch is ambiguous`。实现的声明顺序不影响选择结果。需要让出一段范围时,由应当胜出的一方声明 `preferred_over`,而不是让另一方在自己的 `applies` 里把这段范围排除掉。 -`GemmFwdOp` 的三个实现这样分割 `"gemm"` 接口: +`GemmFwdOp` 的三个实现这样分担 `"gemm"` 接口的调用: | # | key | 服务 | 声明 | | --- | --- | --- | --- | @@ -162,9 +162,9 @@ class GemmFwdInterface(KernelInterface): | 2 | `gemv` | 沿 K 规约且最多两行的形状,这种形状在 CUDA core 上规约更快 | `supported_archs = [90]`、经 `band_for` 实现的 `applies`、`preferred_over = frozenset({"gemm_tma"})` | | 3 | `gemm_cp_async` | 其余两个都不认领的全部形状,前提是一行 K 至少占满一次 4 字节读取 | `supported_archs = [80, 86, 89, 90]`、`general = True`,以及拒绝更窄 K 行的 `refusal` | -`entry_for(call)` 返回两样东西:两次调用共享它才算同一个 kernel 的 **build identity**,以及每个 identity 只运行一次的**构建函数**。identity 少带一个量,第二种 dtype 就会复用第一种 dtype 的 kernel;kernel 只依赖其中几个量却把整个形状都带上,就变成每个形状各编译一次。 +`entry_for(call)` 返回两样东西:两次调用共享它才算同一个 kernel 的 **build identity**,以及每个 identity 只运行一次的 **builder**。identity 少带一个量,第二种 dtype 就会复用第一种 dtype 的 kernel;kernel 只依赖其中几个量却把整个形状都带上,就变成每个形状各编译一次。 -只有一个实现的接口,除继承接口外不需要别的声明。`RMSNormKernel` 就是 `RMSNormFwdOp` 派发的全部内容([`src/tileops/kernels/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/norm/rms_norm.py)): +只有一个实现的接口,除继承接口外不需要别的声明。`RMSNormKernel` 是 `RMSNormFwdOp` 唯一的实现([`src/tileops/kernels/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/norm/rms_norm.py)): ```python class RMSNormKernel(Kernel, RMSNormFwdInterface): @@ -180,7 +180,7 @@ class RMSNormKernel(Kernel, RMSNormFwdInterface): 完全没有自带实现、只依赖外部后端的算子,`kernel_types` 与 `interfaces` 都不写;在没有 target 认领设备时,调用会抛 `OpNotAvailableError`。 -后端可以为一个接口新增实现,也可以替换一个 key 背后的类,这两件事都不改动 TileOPs,见[接入一类新硬件](backends.md)。 +后端可以为一个接口新增实现,也可以替换某个 key 登记的类,这两件事都不改动 TileOPs,见[接入一类新硬件](backends.md)。 ### 注册 @@ -188,7 +188,7 @@ class RMSNormKernel(Kernel, RMSNormFwdInterface): ## 第三步:写 kernel -kernel 类继承 [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) 与它实现的那个接口,放在 [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 下,用 TileLang 编写,在构造时编译。它实现 `forward`,基类的 `__call__` 会调用它。构造函数由本类 `entry_for` 给出的构造方法调用,`forward` 接受接口规定的参数,就是第二步里的 `kernel(a, b)`。 +kernel 类继承 [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) 与它实现的那个接口,放在 [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 下,用 TileLang 编写,在构造时编译。它实现 `forward`,基类的 `__call__` 会调用它。构造函数由本类 `entry_for` 返回的 builder 调用,`forward` 接受接口规定的参数,就是第二步里的 `kernel(a, b)`。 它是这六处里唯一不受 spec 约束的一处:kernel 不读 spec,也不对照 spec 检查。 From 66b9c35234c47743333846e652435885e06edcde Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Thu, 1 Oct 2026 13:30:04 +0800 Subject: [PATCH 5/5] [Docs] Answer the second prose round An ungrammatical English clause left by the previous fix, a Chinese sentence that still ranked the three mechanisms by size, implementations attached to a call rather than to the interface they implement, and one sentence corrected in Chinese but not in English. --- docs/backends.md | 4 ++-- docs/backends.zh.md | 2 +- docs/new-op.md | 4 ++-- docs/new-op.zh.md | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/backends.md b/docs/backends.md index c6763b1b..499a10b1 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -11,8 +11,8 @@ devices run your kernels. **A backend supplies one thing: something callable that computes this call.** Everything else is the op layer's. -This page is about a target, the one of the three extension mechanisms that covers the -whole op. The first half is the +This page is about a target, the extension mechanism that covers a whole op. The first +half is the work, in the order it is done: the four things to write, the protocol's four functions, how one call reaches them, a backend that installs and runs as it stands, how to turn the template into a backend for real hardware, the four rules for writing a kernel, what each diff --git a/docs/backends.zh.md b/docs/backends.zh.md index 4d48f8a7..127bb29e 100644 --- a/docs/backends.zh.md +++ b/docs/backends.zh.md @@ -6,7 +6,7 @@ TileLang 是多后端 DSL,每种硬件各有一套独立的 kernel,由各自 **后端只提供一件事:一个能算这次调用的可调用对象。** 其余都由算子层负责。 -本页讲的是 target,三种接入方式中范围最大的一种。前半部分按实现顺序列出要做的事:要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么,以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。 +本页讲的是 target,接管整个算子的那一种接入方式。前半部分按实现顺序列出要做的事:要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么,以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。 后半说明协议何以如此设计:两层选择、算子层的契约、kernel 的重建条件、调用方可用的接口,以及刻意不支持的情形。 diff --git a/docs/new-op.md b/docs/new-op.md index 0c8cf351..dc24f762 100644 --- a/docs/new-op.md +++ b/docs/new-op.md @@ -108,7 +108,7 @@ class GemmFwdOp(Op): | --- | --- | --- | | 1 | `__init__` | the names, order and defaults in `signature.params`, then `target`, `kernel_map` and `tune`, closing with `self.dispatch_kernel(kernel_map)` | | 2 | `kernel_types` | the Kernel classes that can serve the op, each under a key; a `kernel_map=` override replaces one by that key | -| 3 | `interfaces` | one entry per kernel call the op makes: the name `kernel_for` uses → the `KernelInterface` class every implementation of that call inherits | +| 3 | `interfaces` | one entry per kernel call the op makes: the name `kernel_for` uses → the `KernelInterface` class whose implementations serve that call | | 4 | `forward` | `signature.inputs` — its order, optional inputs last with default `None` | | 5 | `_eager_forward` | contiguity, the call spec, fetching the kernel and launching it | | 6 | `compute_roof` | optional: the GPU-profile unit that prices the op's FLOPs, where it is not CUDA-core fp32 | @@ -208,7 +208,7 @@ dtype reuses the first dtype's kernel; carry the whole shape where the kernel de fewer quantities and it compiles once per distinct shape. An interface with one implementation needs nothing beyond inheriting it. `RMSNormKernel` -is the whole of `RMSNormFwdOp`'s dispatch +is the only implementation `RMSNormFwdOp` has ([`src/tileops/kernels/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/norm/rms_norm.py)): ```python diff --git a/docs/new-op.zh.md b/docs/new-op.zh.md index 0cc94e3b..f838054f 100644 --- a/docs/new-op.zh.md +++ b/docs/new-op.zh.md @@ -99,7 +99,7 @@ class GemmFwdOp(Op): | --- | --- | --- | | 1 | `__init__` | `signature.params` 的名字、顺序与默认值,再加 `target`、`kernel_map`、`tune`;结尾调用 `self.dispatch_kernel(kernel_map)` | | 2 | `kernel_types` | 能服务这个算子的 Kernel 类,各起一个 key;`kernel_map=` 按这个 key 替换其中一个 | -| 3 | `interfaces` | 算子发出的每一个 kernel 调用各占一条:`kernel_for` 用的名字 → 这个调用的各实现所继承的 `KernelInterface` 类 | +| 3 | `interfaces` | 算子发出的每一个 kernel 调用各占一条:`kernel_for` 用的名字 → 服务这个调用的各实现所继承的 `KernelInterface` 类 | | 4 | `forward` | `signature.inputs` 的顺序,可选输入排在最后、默认 `None` | | 5 | `_eager_forward` | 把输入变成连续的,构造 call spec,取出 kernel,再调用它 | | 6 | `compute_roof` | 可选。算子的 FLOPs 按哪个硬件单元的峰值定价,默认是 CUDA core 上的 fp32,用别的单元时才写 |