diff --git a/docs/api/moe.md b/docs/api/moe.md index 7085bd57..9bde06a0 100644 --- a/docs/api/moe.md +++ b/docs/api/moe.md @@ -11,7 +11,8 @@ each token's experts, the ops that move tokens into an expert-contiguous layout back, and the expert GEMMs that run on it. `MoeGroupedGemmFwdOp` is one grouped GEMM; `MoeExpertMLPFwdOp` is the pair of them with the gated activation fused into the first; `FusedMoEExpertsFwdOp` is that MLP with the permutes around it, on the -tight (no-pad) layout. The routing has to produce the layout the GEMM expects. +tight (no-pad) layout, and `IndexedExpertMLPFwdOp` is the backend it picks instead when +the routes are few enough to read the weights once per route rather than once per expert. The routing has to produce the layout the GEMM expects. ## Fused forward @@ -66,3 +67,9 @@ tight (no-pad) layout. The routing has to produce the layout the GEMM expects. show_root_heading: true heading_level: 3 members: ["__init__", "forward"] + +::: tileops.moe.IndexedExpertMLPFwdOp + options: + show_root_heading: true + heading_level: 3 + members: ["__init__", "forward"] diff --git a/docs/backends.md b/docs/backends.md index bc8a256c..a8450f4e 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -182,11 +182,10 @@ d = op(a, b) # every input on one device: a.device == b.devi # two or more True → AmbiguousTargetError, asking for an explicit target= # ── op layer: the one place GemmFwdOp.forward fetches a kernel ─────── -kernel = self.get_or_build_kernel( +kernel = self.kernel_for( "gemm_kernel", # a name from kernel_map (a, b), # the tensors the kernel is about to get, in signature.inputs order - key=(m, n, k, a.dtype), # in-tree only; not used on this call - build=lambda: GemmKernel(m, n, k, a.dtype), # in-tree only; not used on this call + (m, n, k, a.dtype), # what this call is; entry_for reads it, in-tree only ) # ── op layer: look up the external memo table — device, then input signature ── @@ -220,14 +219,14 @@ register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) The op layer calls `build_gemm`; the backend never calls it itself. Importing the backend module only records it in the registry, and the call comes when an op call reaches -`get_or_build_kernel` and misses the external memo table — once per device and input +`kernel_for` and misses the external memo table — once per device and input signature. Whatever it returns, the op layer stores and launches. Four things follow from that: -- **`key` and `build` are the op author's, not a backend's.** They serve the in-tree path - only: `key` decides what the in-tree kernel is looked up on, `build` how it is built. - Neither is used once a target serves the call. +- **`entry_for` is the op author's, not a backend's.** It serves the in-tree path only, + answering with what the in-tree kernel is looked up on and how it is built. Neither + answer is asked for once a target serves the call. - **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. @@ -236,8 +235,8 @@ 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 omit `build`, and then a call - with no target claiming the device raises `OpNotAvailableError`. + `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`. ## Writing a backend that runs {#runnable} diff --git a/docs/backends.zh.md b/docs/backends.zh.md index 8e255cab..05424029 100644 --- a/docs/backends.zh.md +++ b/docs/backends.zh.md @@ -148,11 +148,10 @@ d = op(a, b) # 所有输入必须在同一设备上:a.devi # 两个以上返回 True → 抛 AmbiguousTargetError,要求显式写 target= # ── 算子层:GemmFwdOp.forward 里唯一取 kernel 的那一处 ─────────────── -kernel = self.get_or_build_kernel( +kernel = self.kernel_for( "gemm_kernel", # kernel_map 里的名字 (a, b), # 即将传给 kernel 的张量,顺序照 signature.inputs - key=(m, n, k, a.dtype), # 自带实现用,这次不走 - build=lambda: GemmKernel(m, n, k, a.dtype), # 自带实现用,这次不走 + (m, n, k, a.dtype), # 本次调用是什么;由 entry_for 读,自带实现用,这次不走 ) # ── 算子层:按设备与输入签名查外部记忆表 ───────────────────────────── @@ -184,14 +183,14 @@ def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) ``` -`build_gemm` 由算子层调用,后端自己从不调它:import 后端模块时只是把它登记进注册表,真正被调是在一次调用走到 `get_or_build_kernel`、且外部记忆表未命中的时候,每个「设备 + 输入签名」一次。它返回的可调用对象随后由算子层 launch,也由算子层存进记忆表。 +`build_gemm` 由算子层调用,后端自己从不调它:import 后端模块时只是把它登记进注册表,真正被调是在一次调用走到 `kernel_for`、且外部记忆表未命中的时候,每个「设备 + 输入签名」一次。它返回的可调用对象随后由算子层 launch,也由算子层存进记忆表。 四点对应关系值得记住: -- **`key` 与 `build` 由算子作者写,与后端无关。** 它们只服务自带实现:`key` 决定自带 kernel 按什么查表,`build` 决定它怎么构造。target 选中后端时这两个参数整条不走。 +- **`entry_for` 由算子作者写,与后端无关。** 它只服务自带实现:给出自带 kernel 按什么查表、又怎么构造。target 选中后端时这两个答案都不会被问。 - **张量按位置传,参数按名字传。** `build_kernel(*inputs, **params)`:位置实参是 `TensorSpec`(没传的可选输入是 `None`),关键字实参是 manifest 里 `params` 的名字与本次调用的确定值。 - **一个 `(算子, target)` 只注册一个 builder。** 算子内部分几种情形(GEMM 的 `gemm_kernel` 与 `gemv_kernel`)不会传进来,`build_kernel` 从 `TensorSpec` 自行判断该返回哪个 kernel。 -- **不必自己做记忆。** 同一个设备与输入签名,算子层不会再调第二次;要更细的区分或更少的重建,在 `build_kernel` 内部另加一层缓存。算子完全没有自带实现时 `build` 可以不传,那时没有 target 认领设备,调用直接抛 `OpNotAvailableError`。 +- **不必自己做记忆。** 同一个设备与输入签名,算子层不会再调第二次;要更细的区分或更少的重建,在 `build_kernel` 内部另加一层缓存。算子完全没有自带实现时 `entry_for` 可以不写,那时没有 target 认领设备,调用直接抛 `OpNotAvailableError`。 ## 实现一个可运行的后端 {#runnable} diff --git a/docs/new-op.md b/docs/new-op.md index d0b0d70a..526c3809 100644 --- a/docs/new-op.md +++ b/docs/new-op.md @@ -61,7 +61,7 @@ GemmFwdOp: ``` Those names are how a kernel is asked for at runtime: `_eager_forward` picks one, passes -the name to `get_or_build_kernel`, and the op layer looks the class up in `kernel_map` and +the name to `kernel_for`, and the op layer looks the class up in `kernel_map` and builds it (see [step 2](#op-class)). An external backend registers against the same roster — whichever name it registers a `build_kernel` for is the kernel of the op it takes over. @@ -100,17 +100,20 @@ class GemmFwdOp(Op): self._validate_dtypes(a, b) # generated by the base class m, n, k = self._infer_mnk(a, b) a, b = a.contiguous(), b.contiguous() # handed over as the spec declares it - slot = "gemv_kernel" if m == 1 else "gemm_kernel" - kernel = self.get_or_build_kernel( - slot, # a name from kernel_map - (a, b), # the memo key's tensors, and what a backend receives - key=(m, n, k, a.dtype), # the cache key on the in-tree side - build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune), + role = "gemv_kernel" if m == 1 else "gemm_kernel" + kernel = self.kernel_for( + role, # a name from kernel_map + (a, b), # what a backend is described with + (m, n, k, a.dtype), # what this call is ) return kernel(a, b) + + def entry_for(self, role, call): # the in-tree recipe + m, n, k, dtype = call + return call, lambda: self.kernel_map[role](m, n, k, dtype, tune=self.tune) ``` -Four members to write, each of them from the spec: +Five members to write, the first four of them from the spec: | # | Member | Written from | | --- | --- | --- | @@ -118,38 +121,40 @@ Four members to write, each of them from the spec: | 2 | `default_kernel_map` | `source.kernel_map`: the same names, against the Kernel classes themselves | | 3 | `_infer_output_shapes` | the rules in `signature.shape_rules` that derive an output's shape | | 4 | `forward` | `signature.inputs` — its order and defaults, optional inputs last — plus the validation, the contiguity, fetching the kernel and launching it | +| 5 | `entry_for` | what two calls must share to reuse one kernel, and how that kernel is built | Two more members arrive on their own. When the subclass is defined, the base class synthesises `_validate_dtypes` and `eval_roofline` from the spec's dtype declarations and its `roofline`, so they are there to call — and worth overriding only where the op needs something the spec cannot say. -### `get_or_build_kernel` +### `kernel_for` and `entry_for` 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 keeps a memo table: a kernel this call needs and has built before comes straight back, -and only otherwise is one built and stored. `get_or_build_kernel` is that table's only -entrance, and the point where the in-tree implementation and an external backend part -ways — the second layer of selection in [the backend protocol](backends.md). +and only otherwise is one built and stored. `kernel_for` is that table's only entrance, +and the point where the in-tree implementation and an external backend part ways — the +second layer of selection in [the backend protocol](backends.md). -Its four arguments: +Its three arguments: -**`name`** — which kernel this call wants, as a name from `kernel_map`. +**`role`** — which of this op's kernels the call wants, as a name from `kernel_map`. ```python -slot = "gemv_kernel" if m == 1 else "gemm_kernel" +role = "gemv_kernel" if m == 1 else "gemm_kernel" ``` The in-tree side looks up the Kernel class under that name; a backend looks up the -`build_kernel` it registered under it. An op has as many names as it has cases. +`build_kernel` it registered under it. A role is a memoization bucket, one per kernel the +op runs — never the name of the implementation selection picked for this call. **`inputs`** — the tensors the kernel is about to be handed, in `signature.inputs` order, one slot per input. ```python -self.get_or_build_kernel(slot, (a, b), ...) # GEMM: two required inputs -self.get_or_build_kernel("group_norm", (x, weight, bias), ...) # an absent optional input is None +self.kernel_for(role, (a, b), ...) # GEMM: two required inputs +self.kernel_for("group_norm", (x, weight, bias), ...) # an absent optional input is None ``` The external path keys on it — the device, plus each slot's `(dtype, shape)`. The device @@ -161,36 +166,34 @@ An optional input that was not passed keeps its slot, as `None`; that is what a reads presence off, rather than counting slots. Squeeze the empty slots out, and a clamp with only a lower bound looks exactly like one with only an upper bound. -Omitting `inputs` raises nothing until a backend is installed, and then -`OpNotAvailableError`: the op stays in-tree only, out of reach of any target (see [after -install: two states](backends.md#three-states)). +**`call`** — what this call is, in whatever form the op's own `entry_for` reads. An op +that selects among several implementations passes the record its family defines; an op +with one implementation passes the few values its kernel is built from. -**`key`** — what the in-tree kernel specializes on; the in-tree path only. What becomes of -these last two once a backend serves the op is in [how one call reaches +`entry_for(role, call)` answers with the pair the memo table needs. What becomes of it +once a backend serves the op is in [how one call reaches `build_kernel`](backends.md#from-op-layer). ```python -key=(m, n, k, a.dtype) # GEMM: three dimensions and the dtype -key=(self._cache_key(*input_shapes), x.dtype) # the general form +def entry_for(self, role, call): + m, n, k, dtype = call + return call, lambda: self.kernel_map[role](m, n, k, dtype, tune=self.tune) ``` -The default `_cache_key` takes the sizes of every non-static axis across the inputs — -always correct, but it can over-fragment: one compile per distinct shape. Where the kernel -depends on fewer quantities, override it to project the shape onto those, flattening the -leading dims to one product when the kernel treats its input as 2-D. - -**`build`** — how that in-tree kernel is constructed; the in-tree path only. - -```python -build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune) -``` +The **identity** it returns first is what two calls must share to be one entry: the +construction arguments, plus the device wherever the constructor could produce a different +object on another card. Carry too little and a second dtype reuses the first dtype's +kernel; carry the whole shape where the kernel depends on fewer quantities and you compile +once per distinct shape. -Called once per `key`, which is why compiling belongs here. It may return one Kernel, a -sequence of Kernels built together, or a dataclass carrying them — the last two suit an op -that launches several kernels per call. +The **builder** it returns second runs once per identity, which is why compiling belongs +inside it. It may return one Kernel, a sequence of Kernels built together, or a dataclass +carrying them — the last two suit an op that launches several kernels per call. -An op with no in-tree implementation at all, one written to depend on a backend, may leave -`build` out; a call on a device no target claims then raises `OpNotAvailableError`. +An op that selects among candidate Kernel classes writes no `entry_for` at all: the +default asks the class selection chose, which states its own identity and builder. An op +with no in-tree implementation, written to depend on a backend, leaves both out; a call on +a device no target claims then raises `OpNotAvailableError`. ### Finishing: the compile boundary, and registering @@ -198,8 +201,9 @@ Two things to finish, a few lines each: - **To support `torch.compile`**, declare a compile boundary as well: `forward` only calls the opaque operator, and the validation, the kernel lookup and the launch move into - `_eager_forward`. The op above declares none, so its `forward` holds all the work. How to - declare it is in [bringing an op into torch.compile](torch-compile.md). + `_eager_forward`. The operator itself is generated from the spec, so declaring the + boundary is one class attribute. The op above declares none, so its `forward` holds all + the work. How to declare it is in [bringing an op into torch.compile](torch-compile.md). - **Add the op's name** to the imports and `__all__` in two places: its family's [`src/tileops/ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops), where the class is implemented, and [`src/tileops/.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops), the public @@ -247,8 +251,8 @@ kernel = AttnKernel(num_heads, head_dim, dtype) out = kernel(q, k, v) # seq_len is read off the tensor shapes ``` -With the first form, `seq_len` ends up in `get_or_build_kernel`'s `key`, every step misses, -every step compiles, and decode goes nowhere. +With the first form, `seq_len` ends up in the identity `entry_for` returns, every step +misses, every step compiles, and decode goes nowhere. ## Step 4: write the test diff --git a/docs/new-op.zh.md b/docs/new-op.zh.md index 0df5850c..8f5dc7e6 100644 --- a/docs/new-op.zh.md +++ b/docs/new-op.zh.md @@ -51,7 +51,7 @@ GemmFwdOp: bench: benchmarks/ops/bench_gemm.py ``` -算子在运行时按这些名字取 kernel:`_eager_forward` 里挑出用哪一个,把名字传给 `get_or_build_kernel`,算子层再从 `kernel_map` 找到对应的类去构造(见[第二步](#op-class))。外部后端也是照这份名单注册的 —— 它为哪个名字注册 `build_kernel`,就接管了算子的哪一个 kernel。 +算子在运行时按这些名字取 kernel:`_eager_forward` 里挑出用哪一个,把名字传给 `kernel_for`,算子层再从 `kernel_map` 找到对应的类去构造(见[第二步](#op-class))。外部后端也是照这份名单注册的 —— 它为哪个名字注册 `build_kernel`,就接管了算子的哪一个 kernel。 名字自己起,但要和 kernel 的用途对得上,而且写进算子代码之后就不该再改:它同时是 spec、算子实现与外部后端三方约定的那个词。这也是它推导不出来的原因 —— 只有写 kernel 的人知道这个算子要分几种情形。 @@ -82,17 +82,20 @@ class GemmFwdOp(Op): 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), + role = "gemv_kernel" if m == 1 else "gemm_kernel" + kernel = self.kernel_for( + role, # kernel_map 里的名字 + (a, b), # 后端按这些张量被描述 + (m, n, k, a.dtype), # 本次调用是什么 ) return kernel(a, b) + + def entry_for(self, role, call): # 自带实现的构造方法 + m, n, k, dtype = call + return call, lambda: self.kernel_map[role](m, n, k, dtype, tune=self.tune) ``` -要自己写的是这四个成员,内容都从 spec 来: +要自己写的是这五个成员,前四个内容都从 spec 来: | # | 成员 | 照 spec 的哪一部分写 | | --- | --- | --- | @@ -100,28 +103,29 @@ class GemmFwdOp(Op): | 2 | `default_kernel_map` | `source.kernel_map`:名字照抄,取值换成 Kernel 类本身 | | 3 | `_infer_output_shapes` | `signature.shape_rules` 里推导输出形状那几条 | | 4 | `forward` | `signature.inputs` 的顺序与默认值(可选输入排在必填之后),加上校验、连续化、取 kernel、launch kernel | +| 5 | `entry_for` | 两次调用要共享哪些值才算同一个 kernel,以及这个 kernel 怎么构造 | 另有两个成员不用写:`_validate_dtypes` 与 `eval_roofline` 由基类在子类定义时照 spec 的 dtype 声明与 `roofline` 生成并装上,直接调用即可,只有需要特殊行为时才自己覆写。 -### `get_or_build_kernel` +### `kernel_for` 与 `entry_for` -kernel 是编译产物,构造一次要几百毫秒到几秒,而一个算子实例会被反复调用,形状与 dtype 各不相同。算子层因此维护一张记忆表:本次调用要的 kernel 已经构造过就取回来,没有才构造并存进去。`get_or_build_kernel` 是这张表唯一的入口,也是自带实现与外部后端的分岔点([后端协议](backends.md)里的第二层选择)。 +kernel 是编译产物,构造一次要几百毫秒到几秒,而一个算子实例会被反复调用,形状与 dtype 各不相同。算子层因此维护一张记忆表:本次调用要的 kernel 已经构造过就取回来,没有才构造并存进去。`kernel_for` 是这张表唯一的入口,也是自带实现与外部后端的分岔点([后端协议](backends.md)里的第二层选择)。 -四个参数: +三个参数: -**`name`** —— 本次要哪一个 kernel,取值是 `kernel_map` 里的名字。 +**`role`** —— 本次要这个算子的哪一个 kernel,取值是 `kernel_map` 里的名字。 ```python -slot = "gemv_kernel" if m == 1 else "gemm_kernel" +role = "gemv_kernel" if m == 1 else "gemm_kernel" ``` -自带实现按这个名字找到 Kernel 类,外部后端按它找到注册在同名下的 `build_kernel`。算子分几种情形,`kernel_map` 就有几个名字。 +自带实现按这个名字找到 Kernel 类,外部后端按它找到注册在同名下的 `build_kernel`。role 是记忆表的一个桶,算子跑几个 kernel 就有几个 —— 它不是选择挑中的那个实现的名字。 **`inputs`** —— 即将传给 kernel 的那些张量,顺序照 `signature.inputs`,一个输入占一个位置。 ```python -self.get_or_build_kernel(slot, (a, b), ...) # GEMM:两个必填输入 -self.get_or_build_kernel("group_norm", (x, weight, bias), ...) # 没传的可选输入位置上是 None +self.kernel_for(role, (a, b), ...) # GEMM:两个必填输入 +self.kernel_for("group_norm", (x, weight, bias), ...) # 没传的可选输入位置上是 None ``` 外部路径按它查表:设备加上每个位置的 `(dtype, shape)`。设备也算在内,因为为一块卡编译的产物可能持有那块卡上的资源。后端的 `build_kernel` 收到的也是它,每个张量转成只有 device、dtype、shape 的 `TensorSpec`,不含数据。 @@ -130,22 +134,21 @@ self.get_or_build_kernel("group_norm", (x, weight, bias), ...) # 没传的可 `inputs` 漏掉当场不报错,装上后端才抛 `OpNotAvailableError` —— 这个算子于是只能用自带 kernel,外部 target 接管不了(见[安装之后:两种状态](backends.md#three-states))。 -**`key`** —— 自带 kernel 特化在什么上,只有自带这条路用(换成外部后端服务这个算子时这两个参数怎么走,见[算子层这一侧的调用](backends.md#from-op-layer))。 +**`call`** —— 本次调用是什么,形式由这个算子自己的 `entry_for` 决定。有多个实现可选的算子传家族定义的那条记录;只有一个实现的算子传它构造 kernel 用到的那几个值。 + +`entry_for(role, call)` 返回记忆表要的那一对。换成外部后端服务这个算子时这两者怎么走,见[算子层这一侧的调用](backends.md#from-op-layer)。 ```python -key=(m, n, k, a.dtype) # GEMM:三个维度加 dtype -key=(self._cache_key(*input_shapes), x.dtype) # 通用写法 +def entry_for(self, role, call): + m, n, k, dtype = call + return call, lambda: self.kernel_map[role](m, n, k, dtype, tune=self.tune) ``` -`_cache_key` 的默认实现取所有输入形状中非静态轴的尺寸,正确但可能过细 —— 一个形状编译一次。kernel 实际只依赖其中几个量时覆写它,把形状投影过去,例如 kernel 把输入当二维处理,就把前面几维乘成一个数。 - -**`build`** —— 怎么构造这个自带 kernel,同样只有自带这条路用。 +先返回的**身份**是两次调用要共享什么才算同一条记录:构造参数,再加上设备 —— 只要换一块卡构造出来的对象可能不同就要带上。带少了,第二种 dtype 会复用第一种 dtype 的 kernel;kernel 其实只依赖其中几个量却把整个形状带上,就变成一个形状编译一次。 -```python -build=lambda: self.kernel_map[slot](m, n, k, a.dtype, tune=self.tune) -``` +后返回的**构造方法**每个身份只跑一次,所以编译放在里面是安全的。返回值可以是一个 Kernel、一组一起构造出来的 Kernel,或一个带着它们的 dataclass —— 后两种适合一次调用要 launch 多个 kernel 的算子。 -每个 `key` 只调用一次,所以编译放在这里是安全的。算子完全没有自带实现、只指望外部后端服务时,`build` 可以不传 —— 那样在没有 target 认领设备时,调用会抛 `OpNotAvailableError`。返回值可以是一个 Kernel、一组一起构造出来的 Kernel,或一个带着它们的 dataclass —— 后两种适合一次调用要 launch 多个 kernel 的算子。 +有多个候选 Kernel 类可选的算子根本不写 `entry_for`:默认实现会去问选择挑中的那个类,由它给出自己的身份与构造方法。完全没有自带实现、只指望外部后端的算子两者都不写 —— 那样在没有 target 认领设备时,调用会抛 `OpNotAvailableError`。 ### 收尾:编译边界与注册 @@ -185,7 +188,7 @@ kernel = AttnKernel(num_heads, head_dim, dtype) out = kernel(q, k, v) # seq_len 从张量形状里读 ``` -上一种写法下,`get_or_build_kernel` 的 `key` 里带着 `seq_len`,每步都未命中、每步都编译一次,decode 直接跑不动。 +上一种写法下,`entry_for` 返回的身份里带着 `seq_len`,每步都未命中、每步都编译一次,decode 直接跑不动。 ## 第四步:写测试 diff --git a/docs/torch-compile.md b/docs/torch-compile.md index 8c046e1e..a66ccbe7 100644 --- a/docs/torch-compile.md +++ b/docs/torch-compile.md @@ -21,7 +21,7 @@ layer and `fullgraph=True` works, an empty tuple means the op has not migrated y ```python >>> from tileops.norm import RMSNormFwdOp >>> RMSNormFwdOp.compile_op_names -('tileops::norm_rms_norm_fwd',) +('tileops::normalization_rms_norm_fwd',) ``` An op that has not migrated raises under `fullgraph=True`, and breaks the graph under @@ -48,7 +48,7 @@ block(x, w) ``` `TORCH_LOGS=graph_code` prints the captured graph: one node, -`tileops::norm_rms_norm_fwd`, not the calls inside the kernel. +`tileops::normalization_rms_norm_fwd`, not the calls inside the kernel. ### The five calling conventions @@ -96,43 +96,39 @@ file is ```python class RMSNormFwdOp(Op): - # the operators in the graph that belong to this op - compile_op_names = ("tileops::norm_rms_norm_fwd",) + # one OperatorSpec per operator this op registers; the registration and + # compile_op_names are generated from the manifest entry + compile_boundary: ClassVar[tuple[OperatorSpec, ...]] = (OperatorSpec(),) def _infer_output_shapes(self, x_shape, weight_shape): return {"output": tuple(x_shape)} # the manifest's shape_rules def forward(self, x, weight): - # the only line: call the opaque operator - return _rms_norm_fwd(x, weight, self._instance_key) + # the only line: call the generated operator + return self._wrapped(x, weight, self._instance_key) def _eager_forward(self, x, weight): ... # validate, make contiguous - kernel = self.get_or_build_kernel( - "rms_norm", (x, weight), key=x.dtype, build=..., - ) + kernel = self.kernel_for("rms_norm", (x, weight), x.dtype) return kernel(x, weight) - - -@torch.library.custom_op("tileops::norm_rms_norm_fwd", mutates_args=()) -def _rms_norm_fwd(x, weight, instance_key: str) -> torch.Tensor: - return get_instance(instance_key)._eager_forward(x, weight) - - -@_rms_norm_fwd.register_fake -def _rms_norm_fwd_fake(x, weight, instance_key): - op = get_instance(instance_key) - shapes = op._infer_output_shapes(tuple(x.shape), tuple(weight.shape)) - return x.new_empty(shapes["output"]) ``` +That is the whole declaration. The operator and its fake are generated from the entry: +its tensor arguments are `signature.inputs` in order, what it returns is +`signature.outputs`, which arguments it writes is the inputs marked `mutated: true`, and +each output's dtype is the entry's, or the caller's where the entry marks the output +`caller_stated`. Its name is `tileops::_` — here +`tileops::normalization_rms_norm_fwd` — so no op chooses its own, and `compile_op_names` +cannot disagree with what was registered. An op with a second operator, an in-place or an +`out=` form, adds a second spec saying which argument that one writes. + The layers one call passes through, and where the boundary falls:
Op.__call__resolve the target, unsettle on failure
forwardone line, calls the opaque operator
compile boundary
-
_rms_norm_fwdthe operator body, recovers the instance
+
the generated operatorrecovers the instance
_eager_forwardvalidate, contiguous, kernel, launch
The two violet layers are inside dynamo's trace, and that one line of forward is the last thing it reaches; below the boundary the opaque operator takes over, invisible to the compiler.
@@ -160,7 +156,7 @@ contiguous, while `empty_like` copies the input's strides: a non-contiguous inpu have the fake declare a layout real execution never produces. **Third, the target is resolved twice — once in `Op.__call__`, once in -`get_or_build_kernel`.** When traced code runs `self.x = ...`, dynamo records a pending +`kernel_for`.** When traced code runs `self.x = ...`, dynamo records a pending side effect and applies it only after the whole graph has run, while the opaque node runs before that: a resolution written just outside the node is unreadable inside it. Two things follow, both inside the node: diff --git a/docs/torch-compile.zh.md b/docs/torch-compile.zh.md index d71ad317..3b8a45d6 100644 --- a/docs/torch-compile.zh.md +++ b/docs/torch-compile.zh.md @@ -17,7 +17,7 @@ ```python >>> from tileops.norm import RMSNormFwdOp >>> RMSNormFwdOp.compile_op_names -('tileops::norm_rms_norm_fwd',) +('tileops::normalization_rms_norm_fwd',) ``` 尚未迁移的算子在 `fullgraph=True` 下报错,默认设置下切图。 @@ -41,7 +41,7 @@ w = torch.randn(4096, device="cuda", dtype=torch.float16) block(x, w) ``` -用 `TORCH_LOGS=graph_code` 运行会打印捕获到的图:里面是 `tileops::norm_rms_norm_fwd` 一个节点,不是 kernel 内部的多次调用。 +用 `TORCH_LOGS=graph_code` 运行会打印捕获到的图:里面是 `tileops::normalization_rms_norm_fwd` 一个节点,不是 kernel 内部的多次调用。 ### 调用时要遵守的五条约定 @@ -69,43 +69,32 @@ block(x, w) ```python class RMSNormFwdOp(Op): - # 图中属于这个算子的算子名 - compile_op_names = ("tileops::norm_rms_norm_fwd",) + # 这个算子注册几个 operator 就写几个 OperatorSpec;注册本身与 + # compile_op_names 都由 manifest 条目生成 + compile_boundary: ClassVar[tuple[OperatorSpec, ...]] = (OperatorSpec(),) def _infer_output_shapes(self, x_shape, weight_shape): return {"output": tuple(x_shape)} # manifest 的 shape_rules def forward(self, x, weight): - # 唯一一行:调用那个不透明算子 - return _rms_norm_fwd(x, weight, self._instance_key) + # 唯一一行:调用生成出来的那个算子 + return self._wrapped(x, weight, self._instance_key) def _eager_forward(self, x, weight): ... # 校验、连续化 - kernel = self.get_or_build_kernel( - "rms_norm", (x, weight), key=x.dtype, build=..., - ) + kernel = self.kernel_for("rms_norm", (x, weight), x.dtype) return kernel(x, weight) - - -@torch.library.custom_op("tileops::norm_rms_norm_fwd", mutates_args=()) -def _rms_norm_fwd(x, weight, instance_key: str) -> torch.Tensor: - return get_instance(instance_key)._eager_forward(x, weight) - - -@_rms_norm_fwd.register_fake -def _rms_norm_fwd_fake(x, weight, instance_key): - op = get_instance(instance_key) - shapes = op._infer_output_shapes(tuple(x.shape), tuple(weight.shape)) - return x.new_empty(shapes["output"]) ``` +声明就这么多。operator 与它的 fake 都从条目生成:张量参数是 `signature.inputs` 的顺序,返回什么看 `signature.outputs`,写哪些参数看标了 `mutated: true` 的输入,每个输出的 dtype 取自条目,或者在条目标了 `caller_stated` 时取自调用方。名字是 `tileops::_`,这里就是 `tileops::normalization_rms_norm_fwd` —— 没有算子自己起名字,`compile_op_names` 也就不可能和注册的名字对不上。有第二个 operator(in-place 或 `out=` 形态)的算子再加一个 spec,说明那一个写哪个参数。 + 一次调用经过的各层,以及边界落在哪里:
Op.__call__判定 target,失败则撤销
forward一行,调用不透明算子
编译边界
-
_rms_norm_fwd算子体,取回算子实例
+
生成出来的算子算子体,取回算子实例
_eager_forward校验、连续化、取 kernel、launch kernel
紫色两层在 dynamo 的追踪范围内,forward 那一行是它追到的最后一处;界下由不透明算子接手,编译器看不见。
@@ -119,7 +108,7 @@ def _rms_norm_fwd_fake(x, weight, instance_key): **第二处,fake 用 `x.new_empty(shape)` 构造,而不是 `torch.empty_like(x)`。** fake 返回的张量,形状、dtype 与 stride 三项都必须与真实执行返回的一致;不一致或在追踪期报错,或在运行期按错误布局访问而静默出错。算子体先连续化再写入新分配的输出,真实输出恒为连续,而 `empty_like` 会把入参的 stride 一起复制 —— 非连续输入就让 fake 宣称了一种真实执行不会产出的布局。 -**第三处,target 判定在 `Op.__call__` 与 `get_or_build_kernel` 中各做一次。** 追踪期执行 `self.x = ...`,dynamo 把这次写入记成待办的副作用,等整张图跑完才补上;而不透明节点的执行早于补写,所以节点之外刚写下的判定结果,节点之内读不到。两件事因此都落在节点内部: +**第三处,target 判定在 `Op.__call__` 与 `kernel_for` 中各做一次。** 追踪期执行 `self.x = ...`,dynamo 把这次写入记成待办的副作用,等整张图跑完才补上;而不透明节点的执行早于补写,所以节点之外刚写下的判定结果,节点之内读不到。两件事因此都落在节点内部: - 少了节点内部这一次判定,第一次编译调用会静默用错实现。 - 判定失败时的撤销由做出判定的那一处负责,因为编译产物不保留调用点的 `try/except`。